Repository: elastic/apm-agent-go
Pull request #1441 (2023-05-17 10:04:57)

Handle context flags when creating remote transactions and spans
We need to properly handle traceflags in the context when creating transactions and spans that come from distributed services, or everything is considered as unsampled (and we don't get any data if there is a remote context).


Patch:
diff --git a/module/apmotel/tracer.go b/module/apmotel/tracer.go
index 3f638b41..e3ec0672 100644
--- a/module/apmotel/tracer.go
+++ b/module/apmotel/tracer.go
@@ -88,8 +88,9 @@ func (t *tracer) Start(ctx context.Context, spanName string, opts ...trace.SpanS
 			})
 			s.tx = parent.tx
 			s.spanContext = trace.NewSpanContext(trace.SpanContextConfig{
-				TraceID: trace.TraceID(s.span.TraceContext().Trace),
-				SpanID:  trace.SpanID(s.span.TraceContext().Span),
+				TraceID:    trace.TraceID(s.span.TraceContext().Trace),
+				SpanID:     trace.SpanID(s.span.TraceContext().Span),
+				TraceFlags: trace.TraceFlags(0).WithSampled(s.tx.Sampled()),
 			})
 			return trace.ContextWithSpan(ctx, s), s
 		}
@@ -98,14 +99,16 @@ func (t *tracer) Start(ctx context.Context, spanName string, opts ...trace.SpanS
 	var tranOpts apm.TransactionOptions
 	if psc.HasTraceID() && psc.HasSpanID() {
 		tranOpts.TraceContext = apm.TraceContext{
-			Trace: [16]byte(psc.TraceID()),
-			Span:  [8]byte(psc.SpanID()),
+			Trace:   [16]byte(psc.TraceID()),
+			Span:    [8]byte(psc.SpanID()),
+			Options: apm.TraceOptions(0).WithRecorded(psc.IsSampled()),
 		}
 	}
 	s.tx = t.provider.apmTracer.StartTransactionOptions(spanName, "", tranOpts)
 	s.spanContext = trace.NewSpanContext(trace.SpanContextConfig{
-		TraceID: trace.TraceID(s.tx.TraceContext().Trace),
-		SpanID:  trace.SpanID(s.tx.TraceContext().Span),
+		TraceID:    trace.TraceID(s.tx.TraceContext().Trace),
+		SpanID:     trace.SpanID(s.tx.TraceContext().Span),
+		TraceFlags: trace.TraceFlags(0).WithSampled(s.tx.Sampled()),
 	})
 
 	return trace.ContextWithSpan(ctx, s), s


Test patch:
diff --git a/module/apmotel/tracer_test.go b/module/apmotel/tracer_test.go
index b2d16df4..811fccbb 100644
--- a/module/apmotel/tracer_test.go
+++ b/module/apmotel/tracer_test.go
@@ -55,26 +55,60 @@ func TestTracerStartTransaction(t *testing.T) {
 	assert.Nil(t, s.(*span).span)
 
 	assert.True(t, trace.SpanContextFromContext(ctx).IsValid())
+	assert.True(t, trace.SpanContextFromContext(ctx).IsSampled())
 }
 
 func TestTracerStartTransactionWithParentContext(t *testing.T) {
-	tp, err := NewTracerProvider()
-	assert.NoError(t, err)
-	tracer := newTracer(tp.(*tracerProvider))
-
-	ctx := context.Background()
-	psc := trace.NewSpanContext(trace.SpanContextConfig{
-		TraceID: [16]byte{1},
-		SpanID:  [8]byte{42},
-	})
-	ctx = trace.ContextWithSpanContext(context.Background(), psc)
-
-	ctx, s := tracer.Start(ctx, "name")
-
-	assert.NotNil(t, s.(*span).tx)
-	assert.Nil(t, s.(*span).span)
-
-	assert.True(t, trace.SpanContextFromContext(ctx).IsValid())
+	for _, tt := range []struct {
+		name string
+
+		spanContext     trace.SpanContext
+		expectedSampled bool
+	}{
+		{
+			name: "with an empty span context",
+
+			expectedSampled: true,
+		},
+		{
+			name: "with a sampled span context",
+			spanContext: trace.NewSpanContext(trace.SpanContextConfig{
+				TraceID:    [16]byte{1},
+				SpanID:     [8]byte{42},
+				TraceFlags: trace.TraceFlags(0).WithSampled(true),
+			}),
+
+			expectedSampled: true,
+		},
+		{
+			name: "with a non-sampled span context",
+			spanContext: trace.NewSpanContext(trace.SpanContextConfig{
+				TraceID:    [16]byte{1},
+				SpanID:     [8]byte{42},
+				TraceFlags: trace.TraceFlags(0),
+			}),
+
+			expectedSampled: false,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+
+			tp, err := NewTracerProvider()
+			assert.NoError(t, err)
+			tracer := newTracer(tp.(*tracerProvider))
+
+			ctx := context.Background()
+			ctx = trace.ContextWithSpanContext(context.Background(), tt.spanContext)
+
+			ctx, s := tracer.Start(ctx, "name")
+
+			assert.NotNil(t, s.(*span).tx)
+			assert.Nil(t, s.(*span).span)
+
+			assert.True(t, trace.SpanContextFromContext(ctx).IsValid())
+			assert.Equal(t, trace.SpanContextFromContext(ctx).IsSampled(), tt.expectedSampled)
+		})
+	}
 }
 
 func TestTracerStartChildSpan(t *testing.T) {
@@ -88,21 +122,64 @@ func TestTracerStartChildSpan(t *testing.T) {
 
 	assert.Equal(t, ps.(*span).tx, cs.(*span).tx)
 	assert.NotNil(t, cs.(*span).span)
+
+	assert.True(t, trace.SpanContextFromContext(ctx).IsValid())
 }
 
 func TestTracerStartChildSpanFromTransactionInContext(t *testing.T) {
-	apmTracer, _ := transporttest.NewRecorderTracer()
-	tp, err := NewTracerProvider(WithAPMTracer(apmTracer))
-	assert.NoError(t, err)
-	tracer := newTracer(tp.(*tracerProvider))
-
-	ctx := context.Background()
-	tx := apmTracer.StartTransaction("parent", "")
-	ctx = apm.ContextWithTransaction(context.Background(), tx)
-	ctx, cs := tracer.Start(ctx, "childSpan")
-
-	assert.Equal(t, tx, cs.(*span).tx)
-	assert.NotNil(t, cs.(*span).span)
+	for _, tt := range []struct {
+		name  string
+		txOpt apm.TransactionOptions
+
+		expectedSampled bool
+	}{
+		{
+			name: "with an empty trace context",
+
+			expectedSampled: true,
+		},
+		{
+			name: "with a sampled transaction",
+			txOpt: apm.TransactionOptions{
+				TraceContext: apm.TraceContext{
+					Trace:   apm.TraceID{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
+					Options: apm.TraceOptions(0).WithRecorded(true),
+				},
+			},
+
+			expectedSampled: true,
+		},
+		{
+			name: "with a non-sampled transaction",
+			txOpt: apm.TransactionOptions{
+				TraceContext: apm.TraceContext{
+					Trace:   apm.TraceID{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
+					Options: apm.TraceOptions(0),
+				},
+			},
+
+			expectedSampled: false,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+
+			apmTracer, _ := transporttest.NewRecorderTracer()
+			tp, err := NewTracerProvider(WithAPMTracer(apmTracer))
+			assert.NoError(t, err)
+			tracer := newTracer(tp.(*tracerProvider))
+
+			ctx := context.Background()
+			tx := apmTracer.StartTransactionOptions("parent", "", tt.txOpt)
+			ctx = apm.ContextWithTransaction(context.Background(), tx)
+			ctx, cs := tracer.Start(ctx, "childSpan")
+
+			assert.Equal(t, tx, cs.(*span).tx)
+			assert.NotNil(t, cs.(*span).span)
+
+			assert.True(t, trace.SpanContextFromContext(ctx).IsValid())
+			assert.Equal(t, tt.expectedSampled, trace.SpanContextFromContext(ctx).IsSampled())
+		})
+	}
 }
 
 func TestTracerStartChildSpanWithNewRoot(t *testing.T) {