From f3d083f8d3028012c5988f3b43333036b744a832 Mon Sep 17 00:00:00 2001
From: Ankur Shrivastava <ankur@ankurs.com>
Date: Sat, 4 Apr 2026 00:49:14 +0800
Subject: [PATCH 1/2] perf: add SetTraceIdWithValue, cache OTEL span lookup
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add SetTraceIdWithValue that returns (context.Context, string), avoiding
a separate GetTraceId call after SetTraceId. SetTraceId delegates to it
with //go:inline for zero overhead on existing callers.

Cache oteltrace.SpanFromContext result in a local variable — called once
instead of up to 3 times per invocation.
---
 notifier/notifier.go      | 34 ++++++++++++++++----------
 notifier/notifier_test.go | 50 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 12 deletions(-)

diff --git a/notifier/notifier.go b/notifier/notifier.go
index 3ccf16a..1bdd4ce 100644
--- a/notifier/notifier.go
+++ b/notifier/notifier.go
@@ -530,16 +530,18 @@ 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.
+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 +553,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 +566,21 @@ 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
+//
+//go:inline
+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")
+	}
+}

From b91a78db46c107b0fa1c22e056e78cf7f0844dab Mon Sep 17 00:00:00 2001
From: Ankur Shrivastava <ankur@ankurs.com>
Date: Sat, 4 Apr 2026 00:57:49 +0800
Subject: [PATCH 2/2] fix: remove invalid //go:inline directive, improve
 SetTraceIdWithValue doc

Address review feedback:
- Remove //go:inline (not a valid Go compiler directive)
- Add doc note that callers must use the returned context
---
 notifier/notifier.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/notifier/notifier.go b/notifier/notifier.go
index 1bdd4ce..8325679 100644
--- a/notifier/notifier.go
+++ b/notifier/notifier.go
@@ -532,6 +532,8 @@ func SetRelease(rel string) {
 
 // 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()
@@ -576,8 +578,6 @@ func SetTraceIdWithValue(ctx context.Context) (context.Context, string) {
 // 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
-//
-//go:inline
 func SetTraceId(ctx context.Context) context.Context {
 	ctx, _ = SetTraceIdWithValue(ctx)
 	return ctx