Support for gRPC Server Response Headers and Trailers, adding ClientAfter handler

examples/addsvc/cmd/addsvc/main.go

@@ -222,7 +222,7 @@ func main() {
 			return
 		}
 
-		srv := addsvc.MakeGRPCServer(ctx, endpoints, tracer, logger)
+		srv := addsvc.MakeGRPCServer(endpoints, tracer, logger)
 		s := grpc.NewServer()
 		pb.RegisterAddServer(s, srv)
 

examples/addsvc/transport_grpc.go

@@ -16,20 +16,18 @@ import (
 )
 
 // MakeGRPCServer makes a set of endpoints available as a gRPC AddServer.
-func MakeGRPCServer(ctx context.Context, endpoints Endpoints, tracer stdopentracing.Tracer, logger log.Logger) pb.AddServer {
+func MakeGRPCServer(endpoints Endpoints, tracer stdopentracing.Tracer, logger log.Logger) pb.AddServer {
 	options := []grpctransport.ServerOption{
 		grpctransport.ServerErrorLogger(logger),
 	}
 	return &grpcServer{
 		sum: grpctransport.NewServer(
-			ctx,
 			endpoints.SumEndpoint,
 			DecodeGRPCSumRequest,
 			EncodeGRPCSumResponse,
 			append(options, grpctransport.ServerBefore(opentracing.FromGRPCRequest(tracer, "Sum", logger)))...,
 		),
 		concat: grpctransport.NewServer(
-			ctx,
 			endpoints.ConcatEndpoint,
 			DecodeGRPCConcatRequest,
 			EncodeGRPCConcatResponse,

transport/grpc/_grpc_test/client.go

@@ -0,0 +1,50 @@
+package test
+
+import (
+	"context"
+
+	"google.golang.org/grpc"
+
+	"github.com/go-kit/kit/endpoint"
+	grpctransport "github.com/go-kit/kit/transport/grpc"
+	"github.com/go-kit/kit/transport/grpc/_grpc_test/pb"
+)
+
+type clientBinding struct {
+	test endpoint.Endpoint
+}
+
+func (c *clientBinding) Test(ctx context.Context, a string, b int64) (context.Context, string, error) {
+	response, err := c.test(ctx, TestRequest{A: a, B: b})
+	if err != nil {
+		return nil, "", err
+	}
+	r := response.(*TestResponse)
+	return r.Ctx, r.V, nil
+}
+
+func NewClient(cc *grpc.ClientConn) Service {
+	return &clientBinding{
+		test: grpctransport.NewClient(
+			cc,
+			"pb.Test",
+			"Test",
+			encodeRequest,
+			decodeResponse,
+			&pb.TestResponse{},
+			grpctransport.ClientBefore(
+				injectCorrelationID,
+			),
+			grpctransport.ClientBefore(
+				displayClientRequestHeaders,
+			),
+			grpctransport.ClientAfter(
+				displayClientResponseHeaders,
+				displayClientResponseTrailers,
+			),
+			grpctransport.ClientAfter(
+				extractConsumedCorrelationID,
+			),
+		).Endpoint(),
+	}
+}

transport/grpc/_grpc_test/context_metadata.go

@@ -0,0 +1,136 @@
+package test
+
+import (
+	"context"
+	"fmt"
+
+	"google.golang.org/grpc/metadata"
+)
+
+type metaContext string
+
+const (
+	correlationID     metaContext = "correlation-id"
+	responseHDR       metaContext = "my-response-header"
+	responseTRLR      metaContext = "my-response-trailer"
+	correlationIDTRLR metaContext = "correlation-id-consumed"
+)
+
+/* client before functions */
+
+func injectCorrelationID(ctx context.Context, md *metadata.MD) context.Context {
+	if hdr, ok := ctx.Value(correlationID).(string); ok {
+		fmt.Printf("\tClient found correlationID %q in context, set metadata header\n", hdr)
+		(*md)[string(correlationID)] = append((*md)[string(correlationID)], hdr)
+	}
+	return ctx
+}
+
+func displayClientRequestHeaders(ctx context.Context, md *metadata.MD) context.Context {
+	if len(*md) > 0 {
+		fmt.Println("\tClient >> Request Headers:")
+		for key, val := range *md {
+			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
+		}
+	}
+	return ctx
+}
+
+/* server before functions */
+
+func extractCorrelationID(ctx context.Context, md *metadata.MD) context.Context {
+	if hdr, ok := (*md)[string(correlationID)]; ok {
+		cID := hdr[len(hdr)-1]
+		ctx = context.WithValue(ctx, correlationID, cID)
+		fmt.Printf("\tServer received correlationID %q in metadata header, set context\n", cID)
+	}
+	return ctx
+}
+
+func displayServerRequestHeaders(ctx context.Context, md *metadata.MD) context.Context {
+	if len(*md) > 0 {
+		fmt.Println("\tServer << Request Headers:")
+		for key, val := range *md {
+			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
+		}
+	}
+	return ctx
+}
+
+/* server after functions */
+
+func injectResponseHeader(ctx context.Context, md *metadata.MD, _ *metadata.MD) {
+	*md = metadata.Join(*md, metadata.Pairs(string(responseHDR), "has-a-value"))
+}
+
+func displayServerResponseHeaders(ctx context.Context, md *metadata.MD, _ *metadata.MD) {
+	if len(*md) > 0 {
+		fmt.Println("\tServer >> Response Headers:")
+		for key, val := range *md {
+			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
+		}
+	}
+}
+
+func injectResponseTrailer(ctx context.Context, _ *metadata.MD, md *metadata.MD) {
+	*md = metadata.Join(*md, metadata.Pairs(string(responseTRLR), "has-a-value-too"))
+}
+
+func injectConsumedCorrelationID(ctx context.Context, _ *metadata.MD, md *metadata.MD) {
+	if hdr, ok := ctx.Value(correlationID).(string); ok {
+		fmt.Printf("\tServer found correlationID %q in context, set consumed trailer\n", hdr)
+		*md = metadata.Join(*md, metadata.Pairs(string(correlationIDTRLR), hdr))
+	}
+}
+
+func displayServerResponseTrailers(ctx context.Context, _ *metadata.MD, md *metadata.MD) {
+	if len(*md) > 0 {
+		fmt.Println("\tServer >> Response Trailers:")
+		for key, val := range *md {
+			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
+		}
+	}
+}
+
+/* client after functions */
+
+func displayClientResponseHeaders(ctx context.Context, md metadata.MD, _ metadata.MD) context.Context {
+	if len(md) > 0 {
+		fmt.Println("\tClient << Response Headers:")
+		for key, val := range md {
+			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
+		}
+	}
+	return ctx
+}
+
+func displayClientResponseTrailers(ctx context.Context, _ metadata.MD, md metadata.MD) context.Context {
+	if len(md) > 0 {
+		fmt.Println("\tClient << Response Trailers:")
+		for key, val := range md {
+			fmt.Printf("\t\t%s: %s\n", key, val[len(val)-1])
+		}
+	}
+	return ctx
+}
+
+func extractConsumedCorrelationID(ctx context.Context, _ metadata.MD, md metadata.MD) context.Context {
+	if hdr, ok := md[string(correlationIDTRLR)]; ok {
+		fmt.Printf("\tClient received consumed correlationID %q in metadata trailer, set context\n", hdr[len(hdr)-1])
+		ctx = context.WithValue(ctx, correlationIDTRLR, hdr[len(hdr)-1])
+	}
+	return ctx
+}
+
+/* CorrelationID context handlers */
+
+func SetCorrelationID(ctx context.Context, v string) context.Context {
+	return context.WithValue(ctx, correlationID, v)
+}
+
+func GetConsumedCorrelationID(ctx context.Context) string {
+	if trlr, ok := ctx.Value(correlationIDTRLR).(string); ok {
+		return trlr
+	}
+	return ""
+}

transport/grpc/_grpc_test/pb/generate.go

@@ -0,0 +1,3 @@
+package pb
+
+//go:generate protoc test.proto --go_out=plugins=grpc:.