Propagate context to self-observability measurements in sdk/trace (#7209)

Ensures metric functionality that integrates with trace context (e.g.
exemplars) correctly receive the trace context and anything else the
user has passed.

Author: MrAlias

CHANGELOG.md

@@ -50,7 +50,7 @@ The next release will require at least [Go 1.24].
   The package contains semantic conventions from the `v1.36.0` version of the OpenTelemetry Semantic Conventions.
   See the [migration documentation](./semconv/v1.36.0/MIGRATION.md) for information on how to upgrade from `go.opentelemetry.io/otel/semconv/v1.34.0.`(#7032)
 - Add experimental self-observability span and batch span processor metrics in `go.opentelemetry.io/otel/sdk/trace`.
-  Check the `go.opentelemetry.io/otel/sdk/trace/internal/x` package documentation for more information. (#7027, #6393)
+  Check the `go.opentelemetry.io/otel/sdk/trace/internal/x` package documentation for more information. (#7027, #6393, #7209)
 - Add native histogram exemplar support in `go.opentelemetry.io/otel/exporters/prometheus`. (#6772)
 - Add experimental self-observability log metrics in `go.opentelemetry.io/otel/sdk/log`.
   Check the `go.opentelemetry.io/otel/sdk/log/internal/x` package documentation for more information. (#7121)

sdk/trace/span.go

@@ -509,7 +509,10 @@ func (s *recordingSpan) End(options ...trace.SpanEndOption) {
 				attrSamplingResult = s.tracer.spanLiveMetric.AttrSpanSamplingResult(otelconv.SpanSamplingResultRecordOnly)
 			}
 
-			s.tracer.spanLiveMetric.Add(context.Background(), -1, attrSamplingResult)
+			// Add the span to the context to ensure the metric is recorded
+			// with the correct span context.
+			ctx := trace.ContextWithSpan(context.Background(), s)
+			s.tracer.spanLiveMetric.Add(ctx, -1, attrSamplingResult)
 		}()
 	}
 

sdk/trace/trace_test.go

@@ -2306,8 +2306,15 @@ func TestSelfObservability(t *testing.T) {
 						},
 					},
 				}
+
 				got := scopeMetrics()
-				metricdatatest.AssertEqual(t, want, got, metricdatatest.IgnoreTimestamp())
+				metricdatatest.AssertEqual(
+					t,
+					want,
+					got,
+					metricdatatest.IgnoreTimestamp(),
+					metricdatatest.IgnoreExemplars(),
+				)
 
 				span.End()
 
@@ -2361,7 +2368,13 @@ func TestSelfObservability(t *testing.T) {
 					},
 				}
 				got = scopeMetrics()
-				metricdatatest.AssertEqual(t, want, got, metricdatatest.IgnoreTimestamp())
+				metricdatatest.AssertEqual(
+					t,
+					want,
+					got,
+					metricdatatest.IgnoreTimestamp(),
+					metricdatatest.IgnoreExemplars(),
+				)
 			},
 		},
 		{
@@ -2404,7 +2417,13 @@ func TestSelfObservability(t *testing.T) {
 				}
 
 				got := scopeMetrics()
-				metricdatatest.AssertEqual(t, want, got, metricdatatest.IgnoreTimestamp())
+				metricdatatest.AssertEqual(
+					t,
+					want,
+					got,
+					metricdatatest.IgnoreTimestamp(),
+					metricdatatest.IgnoreExemplars(),
+				)
 			},
 		},
 		{
@@ -2465,7 +2484,13 @@ func TestSelfObservability(t *testing.T) {
 				}
 
 				got := scopeMetrics()
-				metricdatatest.AssertEqual(t, want, got, metricdatatest.IgnoreTimestamp())
+				metricdatatest.AssertEqual(
+					t,
+					want,
+					got,
+					metricdatatest.IgnoreTimestamp(),
+					metricdatatest.IgnoreExemplars(),
+				)
 			},
 		},
 		{
@@ -2535,7 +2560,13 @@ func TestSelfObservability(t *testing.T) {
 					},
 				}
 				got := scopeMetrics()
-				metricdatatest.AssertEqual(t, want, got, metricdatatest.IgnoreTimestamp())
+				metricdatatest.AssertEqual(
+					t,
+					want,
+					got,
+					metricdatatest.IgnoreTimestamp(),
+					metricdatatest.IgnoreExemplars(),
+				)
 			},
 		},
 		{
@@ -2607,7 +2638,13 @@ func TestSelfObservability(t *testing.T) {
 				}
 
 				got := scopeMetrics()
-				metricdatatest.AssertEqual(t, want, got, metricdatatest.IgnoreTimestamp())
+				metricdatatest.AssertEqual(
+					t,
+					want,
+					got,
+					metricdatatest.IgnoreTimestamp(),
+					metricdatatest.IgnoreExemplars(),
+				)
 
 				childSpan.End()
 				parentSpan.End()
@@ -2674,7 +2711,13 @@ func TestSelfObservability(t *testing.T) {
 				}
 
 				got = scopeMetrics()
-				metricdatatest.AssertEqual(t, want, got, metricdatatest.IgnoreTimestamp())
+				metricdatatest.AssertEqual(
+					t,
+					want,
+					got,
+					metricdatatest.IgnoreTimestamp(),
+					metricdatatest.IgnoreExemplars(),
+				)
 			},
 		},
 	}
@@ -2700,6 +2743,98 @@ func TestSelfObservability(t *testing.T) {
 	}
 }
 
+// ctxKeyT is a custom context value type used for testing context propagation.
+type ctxKeyT string
+
+// ctxKey is a context key used to store and retrieve values in the context.
+var ctxKey = ctxKeyT("testKey")
+
+func TestSelfObservabilityContextPropagation(t *testing.T) {
+	t.Setenv("OTEL_GO_X_SELF_OBSERVABILITY", "True")
+	prev := otel.GetMeterProvider()
+	t.Cleanup(func() { otel.SetMeterProvider(prev) })
+
+	// Approximate number of expected measuresments. This is not a strict
+	// requirement, but it should be enough to ensure no backpressure.
+	const count = 3 * 2 // 3 measurements per span, 2 spans (parent and child).
+	ctxCh, fltr := filterFn(count)
+
+	const want = "testValue"
+	n := make(chan int)
+	go func() {
+		// Validate the span context is propagated to all measurements by
+		// testing the context passed to the registered exemplar filter. This
+		// filter receives the measurement context in the standard metric SDK
+		// that we have registered.
+
+		// Count of how many contexts were received.
+		var count int
+
+		for ctx := range ctxCh {
+			count++
+
+			s := trace.SpanFromContext(ctx)
+
+			// All spans should have a valid span context. This should be
+			// passed to the measurements in all cases.
+			isValid := s.SpanContext().IsValid()
+			assert.True(t, isValid, "Context should have a valid span")
+
+			if s.IsRecording() {
+				// Check if the context value is propagated correctly for Span
+				// starts. The Span end operation does not receive any user
+				// context so do not check this if the span is not recording
+				// (i.e. end operation).
+
+				got := ctx.Value(ctxKey)
+				assert.Equal(t, want, got, "Context value not propagated")
+			}
+		}
+		n <- count
+	}()
+
+	// At least one reader is required to not get a no-op MeterProvider and
+	// short-circuit any instrumentation measurements.
+	r := metric.NewManualReader()
+	mp := metric.NewMeterProvider(
+		metric.WithExemplarFilter(fltr),
+		metric.WithReader(r),
+	)
+	otel.SetMeterProvider(mp)
+
+	tp := NewTracerProvider()
+
+	wrap := func(parentCtx context.Context, name string, fn func(context.Context)) {
+		const tracer = "TestSelfObservabilityContextPropagation"
+		ctx, s := tp.Tracer(tracer).Start(parentCtx, name)
+		defer s.End()
+		fn(ctx)
+	}
+
+	ctx := context.WithValue(context.Background(), ctxKey, want)
+	wrap(ctx, "parent", func(ctx context.Context) {
+		wrap(ctx, "child", func(context.Context) {})
+	})
+
+	require.NoError(t, tp.Shutdown(context.Background()))
+
+	// The TracerProvider shutdown returned, no more measurements will be sent
+	// to the exemplar filter.
+	close(ctxCh)
+
+	assert.Positive(t, <-n, "Expected at least 1 context propagations")
+}
+
+// filterFn returns a channel that receives contexts passed to the returned
+// exemplar filter function.
+func filterFn(n int) (chan context.Context, func(ctx context.Context) bool) {
+	out := make(chan context.Context, n)
+	return out, func(ctx context.Context) bool {
+		out <- ctx
+		return true
+	}
+}
+
 // RecordingOnly creates a Sampler that samples no traces, but enables recording.
 // The created sampler maintains any tracestate from the parent span context.
 func RecordingOnly() Sampler {

sdk/trace/tracer.go

@@ -52,6 +52,7 @@ func (tr *tracer) Start(
 	}
 
 	s := tr.newSpan(ctx, name, &config)
+	newCtx := trace.ContextWithSpan(ctx, s)
 	if tr.selfObservabilityEnabled {
 		// Check if the span has a parent span and set the origin attribute accordingly.
 		var attrParentOrigin attribute.KeyValue
@@ -78,20 +79,21 @@ func (tr *tracer) Start(
 			attrSamplingResult = tr.spanStartedMetric.AttrSpanSamplingResult(otelconv.SpanSamplingResultDrop)
 		}
 
-		tr.spanStartedMetric.Add(context.Background(), 1, attrParentOrigin, attrSamplingResult)
+		tr.spanStartedMetric.Add(newCtx, 1, attrParentOrigin, attrSamplingResult)
 	}
 
 	if rw, ok := s.(ReadWriteSpan); ok && s.IsRecording() {
 		sps := tr.provider.getSpanProcessors()
 		for _, sp := range sps {
+			// Use original context.
 			sp.sp.OnStart(ctx, rw)
 		}
 	}
 	if rtt, ok := s.(runtimeTracer); ok {
-		ctx = rtt.runtimeTrace(ctx)
+		newCtx = rtt.runtimeTrace(newCtx)
 	}
 
-	return trace.ContextWithSpan(ctx, s), s
+	return newCtx, s
 }
 
 type runtimeTracer interface {
@@ -147,11 +149,12 @@ func (tr *tracer) newSpan(ctx context.Context, name string, config *trace.SpanCo
 	if !isRecording(samplingResult) {
 		return tr.newNonRecordingSpan(sc)
 	}
-	return tr.newRecordingSpan(psc, sc, name, samplingResult, config)
+	return tr.newRecordingSpan(ctx, psc, sc, name, samplingResult, config)
 }
 
 // newRecordingSpan returns a new configured recordingSpan.
 func (tr *tracer) newRecordingSpan(
+	ctx context.Context,
 	psc, sc trace.SpanContext,
 	name string,
 	sr SamplingResult,
@@ -199,7 +202,10 @@ func (tr *tracer) newRecordingSpan(
 			attrSamplingResult = tr.spanLiveMetric.AttrSpanSamplingResult(otelconv.SpanSamplingResultRecordOnly)
 		}
 
-		tr.spanLiveMetric.Add(context.Background(), 1, attrSamplingResult)
+		// Propagate any existing values from the context with the new span to
+		// the measurement context.
+		ctx = trace.ContextWithSpan(ctx, s)
+		tr.spanLiveMetric.Add(ctx, 1, attrSamplingResult)
 	}
 
 	return s