gcp: Surface trace propagation functions on context (#40)

markusthoemmes · web-flow · commit 517105f6eff9 · 2025-07-08T08:04:47.000-07:00
Users of clog have to be able to reliably put trace information into the
context to make the logger pick it up correctly. The current
implementation uses a string, which is prone to collisions.

This surfaces context manipulation functions instead that operate on a
private type and allow for easier integration.
diff --git a/gcp/setup.go b/gcp/setup.go
@@ -41,7 +41,7 @@ func (h *Handler) Enabled(ctx context.Context, level slog.Level) bool {
 }
 
 func (h *Handler) Handle(ctx context.Context, rec slog.Record) error {
-	if trace := traceFromContext(ctx); trace != "" {
+	if trace := TraceFromContext(ctx); trace != "" {
 		rec = rec.Clone()
 		// Add trace ID	to the record so it is correlated with the request log
 		// See https://cloud.google.com/trace/docs/trace-log-integration
diff --git a/gcp/trace.go b/gcp/trace.go
@@ -101,14 +101,25 @@ func WithCloudTraceContext(h http.Handler) http.Handler {
 			if traceID != "" {
 				trace = fmt.Sprintf("projects/%s/traces/%s", projectID, traceID)
 			}
-			r = r.WithContext(context.WithValue(r.Context(), "trace", trace))
+			r = r.WithContext(WithTrace(r.Context(), trace))
 		}
 		h.ServeHTTP(w, r)
 	})
 }
 
-func traceFromContext(ctx context.Context) string {
-	trace := ctx.Value("trace")
+type traceKey struct{}
+
+// WithTrace adds a trace information to the context.
+func WithTrace(ctx context.Context, trace string) context.Context {
+	if trace == "" {
+		return ctx
+	}
+	return context.WithValue(ctx, traceKey{}, trace)
+}
+
+// TraceFromContext retrieves the trace information from the context.
+func TraceFromContext(ctx context.Context) string {
+	trace := ctx.Value(traceKey{})
 	if trace == nil {
 		return ""
 	}
diff --git a/gcp/trace_test.go b/gcp/trace_test.go
@@ -42,15 +42,9 @@ func TestTrace(t *testing.T) {
 					t.Error("got empty trace context header, want non-empty")
 				}
 
-				traceCtx := ctx.Value("trace")
-				if traceCtx == nil {
-					if c.wantTrace != "" {
-						t.Fatalf("want %s, not found", c.wantTrace)
-					}
-				} else {
-					if traceCtx != c.wantTrace {
-						t.Fatalf("got %s, want %s", traceCtx, c.wantTrace)
-					}
+				traceCtx := TraceFromContext(ctx)
+				if traceCtx != c.wantTrace {
+					t.Fatalf("got %s, want %s", traceCtx, c.wantTrace)
 				}
 			}))
 			srv := httptest.NewServer(h)

Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ func (h *Handler) Enabled(ctx context.Context, level slog.Level) bool {`
`41`	`41`	`}`
`42`	`42`
`43`	`43`	`func (h *Handler) Handle(ctx context.Context, rec slog.Record) error {`
`44`		`- if trace := traceFromContext(ctx); trace != "" {`
	`44`	`+ if trace := TraceFromContext(ctx); trace != "" {`
`45`	`45`	`rec = rec.Clone()`
`46`	`46`	`// Add trace ID to the record so it is correlated with the request log`
`47`	`47`	`// See https://cloud.google.com/trace/docs/trace-log-integration`