优化文件传输

opms_libary/micro_srv/micro_srv.go

@@ -2,8 +2,11 @@ package micro_srv
 
 import (
 	"context"
+	"dashoo.cn/opms_libary/dynamic"
+	"dashoo.cn/opms_libary/multipart"
 	"encoding/json"
 	"errors"
+	"github.com/gogf/gf/util/gconv"
 	"io/ioutil"
 	"net"
 	"os"
@@ -13,9 +16,7 @@ import (
 	"time"
 
 	"dashoo.cn/common_definition/comm_def"
-	"dashoo.cn/opms_libary/dynamic"
 	"dashoo.cn/opms_libary/gtoken"
-	"dashoo.cn/opms_libary/multipart"
 	"dashoo.cn/opms_libary/myerrors"
 	"dashoo.cn/opms_libary/request"
 	"github.com/gogf/gf/encoding/gbase64"
@@ -24,7 +25,6 @@ import (
 	"github.com/gogf/gf/frame/g"
 	"github.com/gogf/gf/net/ghttp"
 	"github.com/gogf/gf/text/gstr"
-	"github.com/gogf/gf/util/gconv"
 	"github.com/rcrowley/go-metrics"
 	consulclient "github.com/rpcxio/rpcx-consul/client"
 	"github.com/rpcxio/rpcx-consul/serverplugin"
@@ -129,55 +129,69 @@ func streamHandler(conn net.Conn, args *share.StreamServiceArgs) {
 	args.Meta["userInfo"] = resp.DataString()
 	ctx = context.WithValue(ctx, share.ReqMetaDataKey, args.Meta)
 
-	fileName := args.Meta["fileName"]
-	//获取文件后缀
-	suffix := path.Ext(fileName)
-	tmpFile, err := ioutil.TempFile(os.TempDir(), "multipart-*"+suffix)
-	if err != nil {
-		g.Log().Error(err)
-		return
-	}
-	defer os.Remove(tmpFile.Name())
-
-	fileSize := args.Meta["fileSize"]
-	size, _ := strconv.Atoi(fileSize)
-
-	total := 0
-	buf := make([]byte, 4096)
+	var array = []*multipart.MultipartFile{}
 	for {
-		n, _ := conn.Read(buf)
-		total += n
-
-		_, err = tmpFile.Write(buf)
+		// 读取文件名长度
+		buf := make([]byte, 3)
+		conn.Read(buf)
+		length, _ := strconv.Atoi(strings.TrimSpace(string(buf)))
+
+		// 读取文件名
+		fileHeader := make([]byte, length)
+		_, err = conn.Read(fileHeader)
+
+		headers := strings.Split(string(fileHeader), " ")
+		paramName := headers[0]
+		fileName := headers[1]
+		fileSize := headers[2]
+
+		//获取文件后缀
+		suffix := path.Ext(fileName)
+		tmpFile, err := ioutil.TempFile(os.TempDir(), "multipart-*"+suffix)
 		if err != nil {
 			g.Log().Error(err)
 			return
 		}
+		defer os.Remove(tmpFile.Name())
 
-		//如果实际总接受字节数与客户端给的要传输字节数相等，说明传输完毕
-		if total == size {
-			result := make([]byte, 0)
-			className, _ := args.Meta["reqService"]
-			methodName, _ := args.Meta["reqMethod"]
-			message := new(dynamic.Message)
-			message.ClassName = className
-			message.MethodName = methodName
-			message.Payload = &multipart.MultipartFile{FileName: fileName, FileSize: gconv.Int64(fileSize), File: tmpFile, Meta: args.Meta}
-			rsp, err := dynamic.Invoker.HandleInvoker(ctx, message)
-			if err != nil {
-				resp := make(map[string]interface{})
-				resp["code"] = 500
-				resp["data"] = err.Error()
-				result, _ = json.Marshal(resp)
-			} else {
-				result, _ = json.Marshal(rsp.Payload)
-			}
+		size, _ := strconv.Atoi(fileSize)
 
-			conn.Write(result)
+		buf = make([]byte, size)
+		conn.Read(buf)
+		tmpFile.Write(buf)
+
+		// 判断是否结束
+		isEnd := make([]byte, 1)
+		_, err = conn.Read(isEnd)
+		if err != nil || isEnd[0] == '1' {
+			file := &multipart.MultipartFile{FileName: paramName, FileSize: gconv.Int64(fileSize), File: tmpFile, Meta: args.Meta}
+			_ = append(array, file)
+			continue
+		}
+		if isEnd[0] == '2' {
 			break
 		}
 	}
 
+	result := make([]byte, 0)
+	className, _ := args.Meta["reqService"]
+	methodName, _ := args.Meta["reqMethod"]
+	message := new(dynamic.Message)
+	message.ClassName = className
+	message.MethodName = methodName
+	message.Payload = array
+	rsp, err := dynamic.Invoker.HandleInvoker(ctx, message)
+	if err != nil {
+		resp := make(map[string]interface{})
+		resp["code"] = 500
+		resp["data"] = err.Error()
+		result, _ = json.Marshal(resp)
+	} else {
+		result, _ = json.Marshal(rsp.Payload)
+	}
+
+	conn.Write(result)
+
 	conn.Close()
 
 }