Fix forwarding rpc status and improve api consistency

Author: kralicky

clientconn.go

@@ -50,12 +50,20 @@ func (cc *ClientConn) invoke(
 	callOpts ...grpc.CallOption,
 ) error {
 	var serviceName, methodName string
-	if method != Forward {
+	forwarding := method == Forward
+	if !forwarding {
 		var err error
 		serviceName, methodName, err = parseQualifiedMethod(method)
 		if err != nil {
 			return err
 		}
+	} else {
+		if _, ok := req.(*RPC); !ok {
+			return fmt.Errorf("[totem] invalid forwarding request type: %T (expected *totem.RPC)", req)
+		}
+		if _, ok := reply.(*RPC); !ok {
+			return fmt.Errorf("[totem] invalid forwarding response type: %T (expected *totem.RPC)", reply)
+		}
 	}
 
 	md, ok := metadata.FromOutgoingContext(ctx)
@@ -127,18 +135,20 @@ func (cc *ClientConn) invoke(
 	select {
 	case rpc := <-future:
 		resp := rpc.GetResponse()
-		stat := resp.GetStatus()
-		if err := stat.Err(); err != nil {
-			cc.logger.Debug("received reply with error", "tag", rpc.Tag,
-				"method", method,
-				"error", err)
+		if !forwarding {
+			// don't unpack response status from replies to forwarded messages
+			stat := resp.GetStatus()
+			if err := stat.Err(); err != nil {
+				cc.logger.Debug("received reply with error", "tag", rpc.Tag,
+					"method", method,
+					"error", err)
 
-			recordErrorStatus(span, stat)
-			return err
+				recordErrorStatus(span, stat)
+				return err
+			}
 		}
 
-		cc.logger.Debug("received reply", "tag", rpc.Tag,
-			"method", method)
+		cc.logger.Debug("received reply", "tag", rpc.Tag, "method", method)
 
 		recordSuccess(span)
 		cc.metrics.TrackSvcTxLatency(serviceName, methodName, time.Since(startTime))
@@ -158,6 +168,7 @@ func (cc *ClientConn) invoke(
 			reply.Content = &RPC_Response{
 				Response: resp,
 			}
+			reply.Metadata = FromMD(rpc.Metadata.ToMD())
 		case protoadapt.MessageV2:
 			if err := proto.Unmarshal(resp.GetResponse(), reply); err != nil {
 				cc.logger.Error("received malformed response message", "tag", rpc.Tag,

stream.go

@@ -591,10 +591,6 @@ func (sh *StreamController) handleRequest(ctx context.Context, msg *RPC, md meta
 				info := first.MethodInfo[method]
 				switch {
 				case !info.IsServerStream && !info.IsClientStream:
-					// very important to copy the message here, otherwise the tag
-					// will be overwritten, and we need to preserve it to reply to
-					// the original request
-					//  todo: does the above still apply?
 					response, err := invoker.Invoke(ctx, msg)
 					if err != nil {
 						recordError(span, err)

totem_suite_test.go

@@ -12,18 +12,24 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	"github.com/onsi/gomega/format"
 	"go.opentelemetry.io/contrib/propagators/autoprop"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/exporters/jaeger"
 	"go.opentelemetry.io/otel/sdk/resource"
 	tracesdk "go.opentelemetry.io/otel/sdk/trace"
 	semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
 
+	"google.golang.org/genproto/googleapis/rpc/errdetails"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/encoding/prototext"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/testing/protocmp"
 	"google.golang.org/protobuf/types/known/emptypb"
 
 	. "github.com/kralicky/totem/test"
@@ -86,6 +92,7 @@ type testServer struct {
 }
 
 func (ts *testServer) TestStream(stream Test_TestStreamServer) error {
+	defer GinkgoRecover()
 	defer ts.wg.Done()
 	return ts.testCase.ServerHandler(stream)
 }
@@ -237,10 +244,16 @@ type errorServer struct {
 	requestLimiter
 }
 
-func (s *errorServer) Error(ctx context.Context, err *ErrorRequest) (*emptypb.Empty, error) {
+func (s *errorServer) Error(ctx context.Context, req *ErrorRequest) (*emptypb.Empty, error) {
 	defer s.requestLimiter.Tick()
-	if err.ReturnError {
-		return nil, status.Error(codes.Aborted, "error")
+	err, _ := status.New(codes.Aborted, "error").WithDetails(&errdetails.ErrorInfo{Reason: "reason", Domain: "domain", Metadata: map[string]string{"key": "value"}})
+
+	// if err := grpc.SetHeader(ctx, metadata.Pairs("errorKey", "errorValue")); err != nil {
+	// 	panic(err)
+	// }
+
+	if req.ReturnError {
+		return nil, err.Err()
 	} else {
 		return &emptypb.Empty{}, nil
 	}
@@ -308,3 +321,43 @@ func (s *countServer) Count(in *Number, stream Count_CountServer) error {
 	}
 	return nil
 }
+
+func ProtoEqual(expected proto.Message) *ProtoMatcher {
+	return &ProtoMatcher{
+		Expected: expected,
+	}
+}
+
+type ProtoMatcher struct {
+	Expected proto.Message
+}
+
+func (matcher *ProtoMatcher) Match(actual any) (success bool, err error) {
+	if actual == nil && matcher.Expected == nil {
+		return false, fmt.Errorf("Refusing to compare <nil> to <nil>.\nBe explicit and use BeNil() instead.  This is to avoid mistakes where both sides of an assertion are erroneously uninitialized")
+	}
+	if _, ok := actual.(proto.Message); !ok {
+		return false, fmt.Errorf("ProtoMatcher expects a proto.Message. Got:\n%s", format.Object(actual, 1))
+	}
+	return proto.Equal(actual.(proto.Message), matcher.Expected), nil
+}
+
+func (matcher *ProtoMatcher) FailureMessage(actual any) (message string) {
+	diff := cmp.Diff(actual.(proto.Message), matcher.Expected, protocmp.Transform())
+	return fmt.Sprintf("Expected\n%s\n%s\n%s\ndiff:\n%s",
+		format.IndentString(prototext.Format(actual.(proto.Message)), 1),
+		"to equal",
+		format.IndentString(prototext.Format(matcher.Expected), 1),
+		diff,
+	)
+}
+
+func (matcher *ProtoMatcher) NegatedFailureMessage(actual any) (message string) {
+	diff := cmp.Diff(actual.(proto.Message), matcher.Expected, protocmp.Transform())
+	return fmt.Sprintf("Expected\n%s\n%s\n%s\ndiff:\n%s",
+		format.IndentString(prototext.Format(actual.(proto.Message)), 1),
+		"not to equal",
+		format.IndentString(prototext.Format(matcher.Expected), 1),
+		diff,
+	)
+}

totem_test.go

@@ -9,27 +9,35 @@ import (
 	sync "sync"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	"github.com/kralicky/totem"
 	"github.com/kralicky/totem/test"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	"github.com/onsi/gomega/format"
+	"google.golang.org/genproto/googleapis/rpc/errdetails"
+	statuspb "google.golang.org/genproto/googleapis/rpc/status"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/encoding/prototext"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/testing/protocmp"
+	"google.golang.org/protobuf/types/known/anypb"
 	"google.golang.org/protobuf/types/known/durationpb"
 )
 
 var (
 	timeout = time.Second * 6
 )
 
-var _ = Describe("Test", func() {
+var _ = FDescribe("Test", func() {
 	It("should work with two different servers", func() {
 		a, b := make(chan struct{}), make(chan struct{})
 		tc := testCase{
 			ServerHandler: func(stream test.Test_TestStreamServer) error {
+				GinkgoHelper()
 				ts, err := totem.NewServer(stream)
 				if err != nil {
 					return err
@@ -184,7 +192,9 @@ var _ = Describe("Test", func() {
 					return err
 				}
 				incSrv := incrementServer{}
+				errSrv := errorServer{}
 				test.RegisterIncrementServer(ts, &incSrv)
+				test.RegisterErrorServer(ts, &errSrv)
 				_, errC := ts.Serve()
 
 				return <-errC
@@ -212,72 +222,57 @@ var _ = Describe("Test", func() {
 				ctx, ca := context.WithTimeout(context.Background(), timeout)
 				defer ca()
 				err = cc.Invoke(ctx, totem.Forward, req, reply)
-				if err != nil {
-					return err
-				}
+				Expect(err).NotTo(HaveOccurred())
 
 				respValue := &test.Number{}
 				err = proto.Unmarshal(reply.GetResponse().GetResponse(), respValue)
 
 				if err != nil {
 					return err
 				}
-				if respValue.GetValue() != 1235 {
-					return fmt.Errorf("expected 1235, got %d", respValue.GetValue())
-				}
+				Expect(respValue.GetValue()).To(Equal(int64(1235)))
 
-				close(done)
-				return <-errC
-			},
-		}
-		tc.Run(done)
-	})
-
-	It("should forward raw RPCs and receive regular proto messages", func() {
-		done := make(chan struct{})
-		tc := testCase{
-			ServerHandler: func(stream test.Test_TestStreamServer) error {
-				ts, err := totem.NewServer(stream)
-				if err != nil {
-					return err
-				}
-				incSrv := incrementServer{}
-				test.RegisterIncrementServer(ts, &incSrv)
-				_, errC := ts.Serve()
-
-				return <-errC
-			},
-			ClientHandler: func(stream test.Test_TestStreamClient) error {
-				ts, err := totem.NewServer(stream)
-				if err != nil {
-					return err
-				}
-
-				cc, errC := ts.Serve()
-
-				reqBytes, _ := proto.Marshal(&test.Number{
-					Value: 1234,
-				})
-				req := &totem.RPC{
-					ServiceName: "test.Increment",
-					MethodName:  "Inc",
+				errReq := &test.ErrorRequest{ReturnError: true}
+				errReqBytes, _ := proto.Marshal(errReq)
+				req = &totem.RPC{
+					ServiceName: "test.Error",
+					MethodName:  "Error",
 					Content: &totem.RPC_Request{
-						Request: reqBytes,
+						Request: errReqBytes,
 					},
 				}
-				reply := &test.Number{}
-
-				ctx, ca := context.WithTimeout(context.Background(), timeout)
-				defer ca()
+				reply = &totem.RPC{}
 				err = cc.Invoke(ctx, totem.Forward, req, reply)
-				if err != nil {
-					return err
+				Expect(err).To(BeNil())
+				errinfo, _ := anypb.New(&errdetails.ErrorInfo{
+					Reason:   "reason",
+					Domain:   "domain",
+					Metadata: map[string]string{"key": "value"},
+				})
+				expected := &totem.RPC{
+					Content: &totem.RPC_Response{
+						Response: &totem.Response{
+							Response: nil,
+							StatusProto: &statuspb.Status{
+								Code:    int32(codes.Aborted),
+								Message: "error",
+								Details: []*anypb.Any{
+									errinfo,
+								},
+							},
+						},
+					},
+					// Metadata: totem.FromMD(metadata.Pairs("errorKey", "errorValue")),
 				}
-
-				if reply.GetValue() != 1235 {
-					return fmt.Errorf("expected 1235, got %d", reply.GetValue())
+				if !proto.Equal(reply, expected) {
+					diff := cmp.Diff(reply, expected, protocmp.Transform())
+					return fmt.Errorf("Expected\n%s\n%s\n%s\ndiff:\n%s",
+						format.IndentString(prototext.Format(reply), 1),
+						"to equal",
+						format.IndentString(prototext.Format(expected), 1),
+						diff,
+					)
 				}
-
 				close(done)
 				return <-errC
 			},
@@ -1093,6 +1088,7 @@ var _ = Describe("Test", func() {
 })
 
 func checkIncrement(cc grpc.ClientConnInterface) {
+	GinkgoHelper()
 	incClient := test.NewIncrementClient(cc)
 	ctx := metadata.AppendToOutgoingContext(context.Background(), "test", "increment")
 	result, err := incClient.Inc(ctx, &test.Number{
@@ -1103,6 +1099,7 @@ func checkIncrement(cc grpc.ClientConnInterface) {
 }
 
 func checkDecrement(cc grpc.ClientConnInterface) {
+	GinkgoHelper()
 	decClient := test.NewDecrementClient(cc)
 	ctx := metadata.AppendToOutgoingContext(context.Background(), "test", "decrement")
 	result, err := decClient.Dec(ctx, &test.Number{
@@ -1113,6 +1110,7 @@ func checkDecrement(cc grpc.ClientConnInterface) {
 }
 
 func checkMultiply(cc grpc.ClientConnInterface) {
+	GinkgoHelper()
 	mulClient := test.NewMultiplyClient(cc)
 	ctx := metadata.AppendToOutgoingContext(context.Background(), "test", "multiply")
 	result, err := mulClient.Mul(ctx, &test.Operands{
@@ -1124,6 +1122,7 @@ func checkMultiply(cc grpc.ClientConnInterface) {
 }
 
 func checkHash(cc grpc.ClientConnInterface) {
+	GinkgoHelper()
 	hashClient := test.NewHashClient(cc)
 	ctx := metadata.AppendToOutgoingContext(context.Background(), "test", "hash")
 	result, err := hashClient.Hash(ctx, &test.String{