unify tracing and logging

grpcutil/middlewares/tracing/tracing_middleware.go

@@ -12,26 +12,44 @@ import (
 )
 
 type TracingMiddleware struct {
+	logger *zerolog.Logger
 }
 
-func NewTracingMiddleware() *TracingMiddleware {
-	return &TracingMiddleware{}
+type TracingHook struct{}
+
+type TracingFields struct{}
+
+func (h TracingHook) Run(e *zerolog.Event, level zerolog.Level, msg string) {
+	ctx := e.GetCtx()
+
+	e.Fields(header.KnownContextValueStrings(ctx))
+
+	serviceMethod, ok := grpc.Method(ctx)
+	if ok {
+		e.Str("service", serviceMethod)
+	}
+
+	tracingFields := ctx.Value(TracingFields{})
+	e.Fields(tracingFields)
+}
+
+func NewTracingMiddleware(logger *zerolog.Logger) *TracingMiddleware {
+	return &TracingMiddleware{
+		logger: logger,
+	}
 }
 
 var _ grpcutil.Middleware = &TracingMiddleware{}
 
 func (m *TracingMiddleware) Unary() grpc.UnaryServerInterceptor {
 	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-		apiLogger := zerolog.Ctx(ctx).With().
-			Fields(header.KnownContextValueStrings(ctx)).
-			Logger()
+		apiLogger := m.logger.Hook(TracingHook{})
+		apiLoggerCtx := apiLogger.WithContext(ctx)
 
 		apiLogger.Trace().Interface("request", req).Msg("grpc call start")
 
-		newCtx := apiLogger.WithContext(ctx)
-
 		start := time.Now()
-		result, err := handler(newCtx, req)
+		result, err := handler(apiLoggerCtx, req)
 		apiLogger.Trace().Dur("duration-ms", time.Since(start)).Msg("grpc call end")
 
 		return result, err
@@ -41,17 +59,13 @@ func (m *TracingMiddleware) Unary() grpc.UnaryServerInterceptor {
 func (m *TracingMiddleware) Stream() grpc.StreamServerInterceptor {
 	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
 		ctx := stream.Context()
+		apiLogger := m.logger.Hook(TracingHook{})
+		apiLoggerCtx := apiLogger.WithContext(ctx)
 
-		apiLogger := zerolog.Ctx(ctx)
-
-		apiLogger.Trace().
-			Fields(header.KnownContextValueStrings(ctx)).
-			Msg("grpc stream call")
-
-		newCtx := apiLogger.WithContext(ctx)
+		apiLogger.Trace().Msg("grpc stream call")
 
 		wrapped := grpcmiddleware.WrapServerStream(stream)
-		wrapped.WrappedContext = newCtx
+		wrapped.WrappedContext = apiLoggerCtx
 
 		return handler(srv, wrapped)
 	}