fix: Server return with Attachment (apache#2641)

Signed-off-by: YarBor <[email protected]>
YarBor · Mar 31, 2024 · 5ff9818 · 5ff9818
1 parent 55a0639
commit 5ff9818
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/protocol/triple/client.go b/protocol/triple/client.go
@@ -69,6 +69,11 @@ func (cm *clientManager) callUnary(ctx context.Context, method string, req, resp
 	triResp := tri.NewResponse(resp)
 	if err := triClient.CallUnary(ctx, triReq, triResp); err != nil {
 		return err
+	} else {
+		i := ctx.Value(constant.AttachmentKey)
+		for k, vs := range triResp.Header() {
+			i.(map[string]interface{})[k] = vs
+		}
 	}
 	return nil
 }

diff --git a/protocol/triple/server.go b/protocol/triple/server.go
@@ -246,11 +246,12 @@ func (s *Server) handleServiceWithInfo(interfaceName string, invoker protocol.In
 					res := invoker.Invoke(ctx, invo)
 					// todo(DMwangnima): modify InfoInvoker to get a unified processing logic
 					// please refer to server/InfoInvoker.Invoke()
-					if triResp, ok := res.Result().(*tri.Response); ok {
-						return triResp, res.Error()
+					triResp, ok := res.Result().(*tri.Response)
+					if !ok {
+						// please refer to proxy/proxy_factory/ProxyInvoker.Invoke
+						triResp = tri.NewResponse([]interface{}{res.Result()})
 					}
-					// please refer to proxy/proxy_factory/ProxyInvoker.Invoke
-					triResp := tri.NewResponse([]interface{}{res.Result()})
+					addReturnAttachments(triResp, res)
 					return triResp, res.Error()
 				},
 				opts...,
@@ -305,6 +306,14 @@ func (s *Server) handleServiceWithInfo(interfaceName string, invoker protocol.In
 	}
 }
 
+func addReturnAttachments(resp *tri.Response, res protocol.Result) {
+	for k, vs := range res.Attachments() {
+		for _, v := range vs.([]string) {
+			resp.Header().Add(k, v)
+		}
+	}
+}
+
 func (s *Server) saveServiceInfo(interfaceName string, info *server.ServiceInfo) {
 	ret := grpc.ServiceInfo{}
 	ret.Methods = make([]grpc.MethodInfo, 0, len(info.Methods))

diff --git a/server/server.go b/server/server.go
@@ -33,6 +33,7 @@ import (
 
 import (
 	"dubbo.apache.org/dubbo-go/v3/common"
+	"dubbo.apache.org/dubbo-go/v3/common/constant"
 	"dubbo.apache.org/dubbo-go/v3/metadata"
 	"dubbo.apache.org/dubbo-go/v3/protocol"
 	registry_exposed "dubbo.apache.org/dubbo-go/v3/registry/exposed_tmp"
@@ -106,6 +107,9 @@ func (ii *infoInvoker) Invoke(ctx context.Context, invocation protocol.Invocatio
 		res, err := method.MethodFunc(ctx, args, ii.svc)
 		result.SetResult(res)
 		result.SetError(err)
+		if attach, ok := ctx.Value(constant.AttachmentKey).(map[string]interface{}); ok {
+			result.SetAttachments(attach)
+		}
 		return result
 	}
 	result.SetError(fmt.Errorf("no match method for %s", name))