From d359901f17242e5a7898feedefe7973284a43221 Mon Sep 17 00:00:00 2001 From: Eric Wennerberg Date: Thu, 31 Oct 2024 14:40:25 +0100 Subject: [PATCH] fix(cloudtrace): re-add trace logger fields for zap Commit 0aeada65cb3721447da341d679bce38d2dd3d139 accidentally removed trace field on log statements written with the request logger. This commit reverts parts of that commit, such that zap request logger still has the trace field populated. --- cloudtrace/middleware.go | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/cloudtrace/middleware.go b/cloudtrace/middleware.go index ce12b381..70c0f4b9 100644 --- a/cloudtrace/middleware.go +++ b/cloudtrace/middleware.go @@ -5,6 +5,8 @@ import ( "net/http" "go.einride.tech/cloudrunner/cloudstream" + "go.einride.tech/cloudrunner/cloudzap" + "go.uber.org/zap" "google.golang.org/grpc" "google.golang.org/grpc/metadata" ) @@ -34,6 +36,7 @@ func (i *Middleware) GRPCServerUnaryInterceptor( } ctx = i.withOutgoingRequestTracing(ctx, values[0]) ctx = i.withInternalContext(ctx, values[0]) + ctx = i.withLogTracing(ctx, values[0]) return handler(ctx, req) } @@ -55,6 +58,7 @@ func (i *Middleware) GRPCStreamServerInterceptor( ctx := ss.Context() ctx = i.withOutgoingRequestTracing(ctx, values[0]) ctx = i.withInternalContext(ctx, values[0]) + ctx = i.withLogTracing(ctx, values[0]) return handler(srv, cloudstream.NewContextualServerStream(ctx, ss)) } @@ -69,6 +73,7 @@ func (i *Middleware) HTTPServer(next http.Handler) http.Handler { w.Header().Set(ContextHeader, header) ctx := i.withOutgoingRequestTracing(r.Context(), header) ctx = i.withInternalContext(ctx, header) + ctx = i.withLogTracing(ctx, header) next.ServeHTTP(w, r.WithContext(ctx)) }) } @@ -84,3 +89,22 @@ func (i *Middleware) withInternalContext(ctx context.Context, header string) con } return SetContext(ctx, result) } + +func (i *Middleware) withLogTracing(ctx context.Context, header string) context.Context { + var traceContext Context + if err := traceContext.UnmarshalString(header); err != nil { + return ctx + } + if i.TraceHook != nil { + ctx = i.TraceHook(ctx, traceContext) + } + fields := make([]zap.Field, 0, 3) + fields = append(fields, cloudzap.Trace(i.ProjectID, traceContext.TraceID)) + if traceContext.SpanID != "" { + fields = append(fields, cloudzap.SpanID(traceContext.SpanID)) + } + if traceContext.Sampled { + fields = append(fields, cloudzap.TraceSampled(traceContext.Sampled)) + } + return cloudzap.WithLoggerFields(ctx, fields...) +}