diff --git a/notifier/notifier.go b/notifier/notifier.go index 3ccf16a..8325679 100644 --- a/notifier/notifier.go +++ b/notifier/notifier.go @@ -530,16 +530,20 @@ func SetRelease(rel string) { sentryRelease = rel } -// SetTraceId updates the traceID based on context values -// if no trace id is found then it will create one and update the context -// You should use the context returned by this function instead of the one passed -func SetTraceId(ctx context.Context) context.Context { +// SetTraceIdWithValue is like SetTraceId but also returns the resolved trace ID, +// avoiding a separate GetTraceId call. +// Callers must use the returned context, not the original ctx, so the stored +// trace ID is preserved in options and log context. +func SetTraceIdWithValue(ctx context.Context) (context.Context, string) { + span := oteltrace.SpanFromContext(ctx) + hasSpan := span.SpanContext().IsValid() + if traceID := GetTraceId(ctx); traceID != "" { // Trace ID already set — ensure it's linked to the OTEL span. - if span := oteltrace.SpanFromContext(ctx); span.SpanContext().IsValid() { + if hasSpan { span.SetAttributes(otelattr.String("coldbrew.trace_id", traceID)) } - return ctx + return ctx, traceID } var traceID string // Check gRPC metadata first — client-supplied trace ID takes priority. @@ -551,10 +555,8 @@ func SetTraceId(ctx context.Context) context.Context { } } // Fall back to OTEL span trace ID. - if strings.TrimSpace(traceID) == "" { - if span := oteltrace.SpanFromContext(ctx); span.SpanContext().IsValid() { - traceID = span.SpanContext().TraceID().String() - } + if strings.TrimSpace(traceID) == "" && hasSpan { + traceID = span.SpanContext().TraceID().String() } // Last resort: generate UUID. if strings.TrimSpace(traceID) == "" { @@ -566,11 +568,19 @@ func SetTraceId(ctx context.Context) context.Context { } // Link the resolved trace ID to the OTEL span as an attribute // so ColdBrew correlation ID and distributed trace are connected. - if span := oteltrace.SpanFromContext(ctx); span.SpanContext().IsValid() { + if hasSpan { span.SetAttributes(otelattr.String("coldbrew.trace_id", traceID)) } ctx = loggers.AddToLogContext(ctx, "trace", traceID) - return options.AddToOptions(ctx, tracerID, traceID) + return options.AddToOptions(ctx, tracerID, traceID), traceID +} + +// SetTraceId updates the traceID based on context values +// if no trace id is found then it will create one and update the context +// You should use the context returned by this function instead of the one passed +func SetTraceId(ctx context.Context) context.Context { + ctx, _ = SetTraceIdWithValue(ctx) + return ctx } // GetTraceId fetches traceID from context diff --git a/notifier/notifier_test.go b/notifier/notifier_test.go index 469b2e0..7232a07 100644 --- a/notifier/notifier_test.go +++ b/notifier/notifier_test.go @@ -181,3 +181,53 @@ func TestSetTraceId_NoSpan_NoPanic(t *testing.T) { t.Error("expected a generated trace ID") } } + +func TestSetTraceIdWithValue_ReturnsTraceID(t *testing.T) { + ctx, traceID := SetTraceIdWithValue(context.Background()) + if traceID == "" { + t.Fatal("expected a non-empty trace ID") + } + if got := GetTraceId(ctx); got != traceID { + t.Errorf("GetTraceId = %q, want %q", got, traceID) + } +} + +func TestSetTraceIdWithValue_ExistingTraceID(t *testing.T) { + ctx := options.AddToOptions(context.Background(), tracerID, "existing-id") + ctx, traceID := SetTraceIdWithValue(ctx) + if traceID != "existing-id" { + t.Errorf("expected existing-id, got %q", traceID) + } + if got := GetTraceId(ctx); got != "existing-id" { + t.Errorf("GetTraceId = %q, want existing-id", got) + } +} + +func TestSetTraceIdWithValue_SetsOTELAttribute(t *testing.T) { + exporter := setupTestTracer(t) + ctx, span := otel.Tracer("test").Start(context.Background(), "test-span") + + ctx, traceID := SetTraceIdWithValue(ctx) + span.End() + + if traceID == "" { + t.Fatal("expected a non-empty trace ID") + } + if got := GetTraceId(ctx); got != traceID { + t.Errorf("GetTraceId = %q, want %q", got, traceID) + } + + spans := exporter.GetSpans() + if len(spans) == 0 { + t.Fatal("expected at least one span") + } + found := false + for _, attr := range spans[0].Attributes { + if string(attr.Key) == "coldbrew.trace_id" && attr.Value.AsString() == traceID { + found = true + } + } + if !found { + t.Error("coldbrew.trace_id attribute not found on span") + } +}