handle context flags when creating remote transactions and spans (#1441)

elastic · May 17, 2023 · 9a7527b · 9a7527b
1 parent 9d8dd37
commit 9a7527b
Show file tree

Hide file tree

Showing 2 changed files with 115 additions and 35 deletions.
diff --git 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

diff --git 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) {