WIP request ids + trace context

rpc/grpcclientidentity/grpchelper/authlistener_grpc_adaptor_wrapper.go

@@ -4,14 +4,18 @@ package grpchelper
 
 import (
 	"context"
+	"fmt"
+	"strings"
 	"time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/keepalive"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/zrepl/zrepl/logger"
 	"github.com/zrepl/zrepl/rpc/grpcclientidentity"
 	"github.com/zrepl/zrepl/rpc/netadaptor"
+	"github.com/zrepl/zrepl/tracing"
 	"github.com/zrepl/zrepl/transport"
 )
 
@@ -28,15 +32,40 @@ type Logger = logger.Logger
 
 // ClientConn is an easy-to-use wrapper around the Dialer and TransportCredentials interface
 // to produce a grpc.ClientConn
-func ClientConn(cn transport.Connecter, log Logger) *grpc.ClientConn {
+func ClientConn(cn transport.Connecter, log Logger, genReqID func() string) *grpc.ClientConn {
 	ka := grpc.WithKeepaliveParams(keepalive.ClientParameters{
 		Time:                StartKeepalivesAfterInactivityDuration,
 		Timeout:             KeepalivePeerTimeout,
 		PermitWithoutStream: true,
 	})
+	unaryIntcpt := grpc.WithUnaryInterceptor(func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+
+		traceStack := tracing.GetStack(ctx)
+		if len(traceStack) == 0 {
+			panic("implementation error: expecting trace stack")
+		}
+		reqId := genReqID() + " " + strings.Join(traceStack, " <> ")
+
+		l := log
+		l = l.WithField("reqID", reqId)
+		l = l.WithField("reqT", fmt.Sprintf("%T", req))
+		l = l.WithField("repT", fmt.Sprintf("%T", reply))
+		l = l.WithField("ctxT", fmt.Sprintf("%T", ctx))
+		l.Debug("unary intercepted")
+
+		ctx = metadata.AppendToOutgoingContext(ctx, "zrepl-req-id", reqId)
+
+		err := invoker(ctx, method, req, reply, cc, opts...)
+
+		l = log.WithField("repT", fmt.Sprintf("%T", reply))
+		l = log.WithField("errT", fmt.Sprintf("%T", err))
+		l.Debug("reply intercepted")
+
+		return err
+	})
 	dialerOption := grpc.WithDialer(grpcclientidentity.NewDialer(log, cn))
 	cred := grpc.WithTransportCredentials(grpcclientidentity.NewTransportCredentials(log))
-	cc, err := grpc.DialContext(context.Background(), "doesn't matter done by dialer", dialerOption, cred, ka)
+	cc, err := grpc.DialContext(context.Background(), "doesn't matter done by dialer", dialerOption, cred, ka, unaryIntcpt)
 	if err != nil {
 		log.WithError(err).Error("cannot create gRPC client conn (non-blocking)")
 		// It's ok to panic here: the we call grpc.DialContext without the
@@ -50,7 +79,7 @@ func ClientConn(cn transport.Connecter, log Logger) *grpc.ClientConn {
 }
 
 // NewServer is a convenience interface around the TransportCredentials and Interceptors interface.
-func NewServer(authListener transport.AuthenticatedListener, clientIdentityKey interface{}, logger grpcclientidentity.Logger, ctxInterceptor grpcclientidentity.ContextInterceptor) (srv *grpc.Server, serve func() error) {
+func NewServer(authListener transport.AuthenticatedListener, clientIdentityKey interface{}, logger grpcclientidentity.Logger, interceptor grpcclientidentity.ContextInterceptor, pre grpcclientidentity.PreHandlerInspector, post grpcclientidentity.PostHandlerInspector) (srv *grpc.Server, serve func() error) {
 	ka := grpc.KeepaliveParams(keepalive.ServerParameters{
 		Time:    StartKeepalivesAfterInactivityDuration,
 		Timeout: KeepalivePeerTimeout,
@@ -60,8 +89,17 @@ func NewServer(authListener transport.AuthenticatedListener, clientIdentityKey i
 		PermitWithoutStream: true,
 	})
 	tcs := grpcclientidentity.NewTransportCredentials(logger)
-	unary, stream := grpcclientidentity.NewInterceptors(logger, clientIdentityKey, ctxInterceptor)
-	srv = grpc.NewServer(grpc.Creds(tcs), grpc.UnaryInterceptor(unary), grpc.StreamInterceptor(stream), ka, ep)
+	unary, stream := grpcclientidentity.NewInterceptors(logger, clientIdentityKey, interceptor, pre, post)
+	unary2 := func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
+		md, ok := metadata.FromIncomingContext(ctx)
+		if !ok {
+			logger.Error("no request id in incoming request")
+		} else {
+			logger.Info(fmt.Sprintf("req-id = %v", md.Get("zrepl-req-id")))
+		}
+		return unary(ctx, req, info, handler)
+	}
+	srv = grpc.NewServer(grpc.Creds(tcs), grpc.UnaryInterceptor(unary2), grpc.StreamInterceptor(stream), ka, ep)
 
 	serve = func() error {
 		if err := srv.Serve(netadaptor.New(authListener, logger)); err != nil {