perf: add SetTraceIdWithValue, cache OTEL span lookup

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

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")
+	}
+}