tests(grpc): adds tests for inPayload parsing.

middleware/grpc/server.go

@@ -48,9 +48,9 @@ func WithServerInPayloadParser(parser func(*stats.InPayload, zipkin.Span)) Serve
 	}
 }
 
-// WithserverInTrailerParser adds a parser for the stats.InTrailer to be able to access
+// WithServerInTrailerParser adds a parser for the stats.InTrailer to be able to access
 // the request trailer
-func WithserverInTrailerParser(parser func(*stats.InTrailer, zipkin.Span)) ServerOption {
+func WithServerInTrailerParser(parser func(*stats.InTrailer, zipkin.Span)) ServerOption {
 	return func(h *serverHandler) {
 		h.handleRPCParser.inTrailer = parser
 	}
@@ -64,6 +64,30 @@ func WithServerInHeaderParser(parser func(*stats.InHeader, zipkin.Span)) ServerO
 	}
 }
 
+// WithServerOutPayloadParser adds a parser for the stats.OutPayload to be able to access
+// the response payload
+func WithServerOutPayloadParser(parser func(*stats.OutPayload, zipkin.Span)) ServerOption {
+	return func(h *serverHandler) {
+		h.handleRPCParser.outPayload = parser
+	}
+}
+
+// WithServerOutTrailerParser adds a parser for the stats.OutTrailer to be able to access
+// the response trailer
+func WithServerOutTrailerParser(parser func(*stats.OutTrailer, zipkin.Span)) ServerOption {
+	return func(h *serverHandler) {
+		h.handleRPCParser.outTrailer = parser
+	}
+}
+
+// WithServerOutHeaderParser adds a parser for the stats.OutHeader to be able to access
+// the response payload
+func WithServerOutHeaderParser(parser func(*stats.OutHeader, zipkin.Span)) ServerOption {
+	return func(h *serverHandler) {
+		h.handleRPCParser.outHeader = parser
+	}
+}
+
 // NewServerHandler returns a stats.Handler which can be used with grpc.WithStatsHandler to add
 // tracing to a gRPC server. The gRPC method name is used as the span name and by default the only
 // tags are the gRPC status code if the call fails. Use ServerTags to add additional tags that

middleware/grpc/server_parser_test.go

@@ -19,13 +19,14 @@ import (
 	"testing"
 
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/stats"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/openzipkin/zipkin-go"
 	zipkingrpc "github.com/openzipkin/zipkin-go/middleware/grpc"
 	"github.com/openzipkin/zipkin-go/model"
 	service "github.com/openzipkin/zipkin-go/proto/testing"
+	"google.golang.org/protobuf/encoding/protojson"
 )
 
 func TestGRPCServerCreatesASpanAndContext(t *testing.T) {
@@ -77,22 +78,26 @@ func TestGRPCServerCreatesASpanAndContext(t *testing.T) {
 	}
 }
 
-func TestGRPCServerCanAccessToHeaders(t *testing.T) {
+func TestGRPCServerCanAccessToPayload(t *testing.T) {
 	tracer, flusher := createTracer(false)
 
 	s := grpc.NewServer(
 		grpc.StatsHandler(
 			zipkingrpc.NewServerHandler(
 				tracer,
 				zipkingrpc.ServerTags(map[string]string{"default": "tag"}),
-				zipkingrpc.WithServerInHeaderParser(func(inHeader *stats.InHeader, span zipkin.Span) {
-					if want, have := "test_value", inHeader.Header.Get("test_key")[0]; want != have {
-						t.Errorf("unexpected metadata value in header, want: %q, have %q", want, have)
+				zipkingrpc.WithServerInPayloadParser(func(inPayload *stats.InPayload, span zipkin.Span) {
+					m, ok := inPayload.Payload.(proto.Message)
+					if !ok {
+						t.Fatal("failed to cast the payload as a proto.Message")
 					}
-				}),
-				zipkingrpc.WithServerInTrailerParser(func(inTrailer *stats.InTrailer, span zipkin.Span) {
-					if want, have := "test_value", inTrailer.Trailer.Get("test_key")[0]; want != have {
-						t.Errorf("unexpected metadata value in header, want: %q, have %q", want, have)
+					mb, err := protojson.MarshalOptions{}.Marshal(m)
+					if err != nil {
+						t.Fatalf("unexpected error: %v", err)
+					}
+
+					if want, have := `{"payload":"Hello"}`, string(mb); want != have {
+						t.Errorf("incorrect payload: want %q, have %q", want, have)
 					}
 				}),
 			),
@@ -118,7 +123,6 @@ func TestGRPCServerCanAccessToHeaders(t *testing.T) {
 
 	client := service.NewHelloServiceClient(conn)
 
-	ctx = metadata.AppendToOutgoingContext(ctx, "test_key", "test_value")
 	_, err = client.Hello(ctx, &service.HelloRequest{
 		Payload: "Hello",
 	})
@@ -130,9 +134,4 @@ func TestGRPCServerCanAccessToHeaders(t *testing.T) {
 	if want, have := 1, len(spans); want != have {
 		t.Errorf("unexpected number of spans, want %d, have %d", want, have)
 	}
-
-	span := spans[0]
-	if want, have := model.Server, span.Kind; want != have {
-		t.Errorf("unexpected kind, want %q, have %q", want, have)
-	}
 }

middleware/grpc/shared.go

@@ -28,9 +28,12 @@ import (
 )
 
 type handleRPCParser struct {
-	inPayload func(*stats.InPayload, zipkin.Span)
-	inTrailer func(*stats.InTrailer, zipkin.Span)
-	inHeader  func(*stats.InHeader, zipkin.Span)
+	inPayload  func(*stats.InPayload, zipkin.Span)
+	inTrailer  func(*stats.InTrailer, zipkin.Span)
+	inHeader   func(*stats.InHeader, zipkin.Span)
+	outPayload func(*stats.OutPayload, zipkin.Span)
+	outTrailer func(*stats.OutTrailer, zipkin.Span)
+	outHeader  func(*stats.OutHeader, zipkin.Span)
 }
 
 // A RPCHandler can be registered using WithClientRPCHandler or WithServerRPCHandler to intercept calls to HandleRPC of
@@ -59,6 +62,18 @@ func handleRPC(ctx context.Context, rs stats.RPCStats, h handleRPCParser) {
 		if h.inTrailer != nil {
 			h.inTrailer(rs, span)
 		}
+	case *stats.OutPayload:
+		if h.outPayload != nil {
+			h.outPayload(rs, span)
+		}
+	case *stats.OutHeader:
+		if h.outHeader != nil {
+			h.outHeader(rs, span)
+		}
+	case *stats.OutTrailer:
+		if h.outTrailer != nil {
+			h.outTrailer(rs, span)
+		}
 	case *stats.End:
 		s, ok := status.FromError(rs.Error)
 		// rs.Error should always be convertable to a status, this is just a defensive check.