Do not apply adaptive sampler to child spans (#457)

* Do not apply adaptive sampler to child spans Signed-off-by: Yuri Shkuro <[email protected]> * Cleanup Signed-off-by: Yuri Shkuro <[email protected]> * DRY Signed-off-by: Yuri Shkuro <[email protected]> * Avoid named return values Signed-off-by: Yuri Shkuro <[email protected]>
jaegertracing · Oct 28, 2019 · 8fad375 · 8fad375
1 parent 790ac7d
commit 8fad375
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 13 deletions.
diff --git a/propagation_test.go b/propagation_test.go
@@ -101,10 +101,11 @@ func TestSpanPropagator(t *testing.T) {
 			t.Fatalf("%d: ParentID %d does not match expectation %d", i, a, e)
 		} else {
 			// Prepare for comparison.
-			sp.context.spanID, sp.context.parentID = exp.context.SpanID(), 0
 			sp.duration, sp.startTime = exp.duration, exp.startTime
 		}
-		assert.Equal(t, exp.context, sp.context, formatName)
+		assert.Equal(t, exp.context.traceID, sp.context.traceID, formatName)
+		assert.Equal(t, exp.context.Flags(), sp.context.Flags(), formatName)
+		assert.Equal(t, exp.context.baggage, sp.context.baggage, formatName)
 		assert.Equal(t, "span.kind", sp.tags[0].key)
 		assert.Equal(t, expTags, sp.tags[1:] /*skip span.kind tag*/, formatName)
 		assert.Empty(t, sp.logs, formatName)

diff --git a/sampler.go b/sampler.go
@@ -361,18 +361,25 @@ func (s *AdaptiveSampler) IsSampled(id TraceID, operation string) (bool, []Tag)
 	return false, nil
 }
 
+func (s *AdaptiveSampler) trySampling(span *Span, operationName string) (bool, []Tag) {
+	samplerV1 := s.getSamplerForOperation(operationName)
+	var sampled bool
+	var tags []Tag
+	if span.context.samplingState.isLocalRootSpan(span.context.spanID) {
+		sampled, tags = samplerV1.IsSampled(span.context.TraceID(), operationName)
+	}
+	return sampled, tags
+}
+
 // OnCreateSpan implements OnCreateSpan of SamplerV2.
 func (s *AdaptiveSampler) OnCreateSpan(span *Span) SamplingDecision {
-	operationName := span.OperationName()
-	samplerV1 := s.getSamplerForOperation(operationName)
-	sampled, tags := samplerV1.IsSampled(span.context.TraceID(), operationName)
+	sampled, tags := s.trySampling(span, span.OperationName())
 	return SamplingDecision{Sample: sampled, Retryable: true, Tags: tags}
 }
 
 // OnSetOperationName implements OnSetOperationName of SamplerV2.
 func (s *AdaptiveSampler) OnSetOperationName(span *Span, operationName string) SamplingDecision {
-	samplerV1 := s.getSamplerForOperation(operationName)
-	sampled, tags := samplerV1.IsSampled(span.context.TraceID(), operationName)
+	sampled, tags := s.trySampling(span, operationName)
 	return SamplingDecision{Sample: sampled, Retryable: false, Tags: tags}
 }
 

diff --git a/sampler_test.go b/sampler_test.go
@@ -20,6 +20,7 @@ import (
 	"sync"
 	"testing"
 
+	"github.com/opentracing/opentracing-go"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -273,6 +274,48 @@ func TestMaxOperations(t *testing.T) {
 	assert.Equal(t, testProbabilisticExpectedTags, decision.Tags)
 }
 
+func TestAdaptiveSamplerDoesNotApplyToChildrenSpans(t *testing.T) {
+	strategies := &sampling.PerOperationSamplingStrategies{
+		DefaultSamplingProbability:       0,
+		DefaultLowerBoundTracesPerSecond: 0,
+		PerOperationStrategies: []*sampling.OperationSamplingStrategy{
+			{
+				Operation:             "op1",
+				ProbabilisticSampling: &sampling.ProbabilisticSamplingStrategy{SamplingRate: 0.0},
+			},
+			{
+				Operation:             "op2",
+				ProbabilisticSampling: &sampling.ProbabilisticSamplingStrategy{SamplingRate: 1.0},
+			},
+		},
+	}
+
+	sampler, err := NewAdaptiveSampler(strategies, 1)
+	assert.NoError(t, err)
+	tracer, closer := NewTracer("service", sampler, NewNullReporter())
+	defer closer.Close()
+
+	_ = tracer.StartSpan("op1") // exhaust lower bound sampler once
+	span1 := tracer.StartSpan("op1")
+	assert.False(t, span1.Context().(SpanContext).IsSampled(), "op1 should not be sampled on root span")
+	span1.SetOperationName("op2")
+	assert.True(t, span1.Context().(SpanContext).IsSampled(), "op2 should be sampled on root span")
+
+	span2 := tracer.StartSpan("op2")
+	assert.True(t, span2.Context().(SpanContext).IsSampled(), "op2 should be sampled on root span")
+
+	parent := tracer.StartSpan("op1")
+	assert.False(t, parent.Context().(SpanContext).IsSampled(), "parent span should not be sampled")
+	assert.False(t, parent.Context().(SpanContext).IsSamplingFinalized(), "parent span should not be finalized")
+
+	child := tracer.StartSpan("op2", opentracing.ChildOf(parent.Context()))
+	assert.False(t, child.Context().(SpanContext).IsSampled(), "child span should not be sampled even with op2")
+	assert.False(t, child.Context().(SpanContext).IsSamplingFinalized(), "child span should not be finalized")
+	child.SetOperationName("op2")
+	assert.False(t, child.Context().(SpanContext).IsSampled(), "op2 should not be sampled on the child span")
+	assert.True(t, child.Context().(SpanContext).IsSamplingFinalized(), "child span should be finalized after setOperationName")
+}
+
 func TestAdaptiveSampler_lockRaceCondition(t *testing.T) {
 	agent, remoteSampler, _ := initAgent(t)
 	defer agent.Close()
@@ -290,18 +333,18 @@ func TestAdaptiveSampler_lockRaceCondition(t *testing.T) {
 	// Overwrite the sampler with an adaptive sampler
 	remoteSampler.sampler = adaptiveSampler
 
+	tracer, closer := NewTracer("service", remoteSampler, NewNullReporter())
+	defer closer.Close()
+
 	var wg sync.WaitGroup
 	defer wg.Wait()
 	wg.Add(2)
 
 	isSampled := func(t *testing.T, remoteSampler *RemotelyControlledSampler, numOperations int, operationNamePrefix string) {
 		for i := 0; i < numOperations; i++ {
 			runtime.Gosched()
-			span := &Span{
-				operationName: fmt.Sprintf("%s%d", operationNamePrefix, i),
-			}
-			decision := remoteSampler.OnCreateSpan(span)
-			assert.True(t, decision.Sample)
+			span := tracer.StartSpan(fmt.Sprintf("%s%d", operationNamePrefix, i))
+			assert.True(t, span.Context().(SpanContext).IsSampled())
 		}
 	}
 

diff --git a/span_context.go b/span_context.go
@@ -84,11 +84,18 @@ type samplingState struct {
 	// like SetOperationName / SetTag, and the spans will remain writable.
 	final atomic.Bool
 
+	// localRootSpan stores the SpanID of the first span created in this process for a given trace.
+	localRootSpan SpanID
+
 	// extendedState allows samplers to keep intermediate state.
 	// The keys and values in this map are completely opaque: interface{} -> interface{}.
 	extendedState sync.Map
 }
 
+func (s *samplingState) isLocalRootSpan(id SpanID) bool {
+	return id == s.localRootSpan
+}
+
 func (s *samplingState) setFlag(newFlag int32) {
 	swapped := false
 	for !swapped {

diff --git a/tracer.go b/tracer.go
@@ -74,6 +74,7 @@ type Tracer struct {
 // NewTracer creates Tracer implementation that reports tracing to Jaeger.
 // The returned io.Closer can be used in shutdown hooks to ensure that the internal
 // queue of the Reporter is drained and all buffered spans are submitted to collectors.
+// TODO (breaking change) return *Tracer only, without closer.
 func NewTracer(
 	serviceName string,
 	sampler Sampler,
@@ -272,7 +273,9 @@ func (t *Tracer) startSpanWithOptions(
 			}
 			ctx.spanID = SpanID(ctx.traceID.Low)
 			ctx.parentID = 0
-			ctx.samplingState = &samplingState{}
+			ctx.samplingState = &samplingState{
+				localRootSpan: ctx.spanID,
+			}
 			if hasParent && parent.isDebugIDContainerOnly() && t.isDebugAllowed(operationName) {
 				ctx.samplingState.setDebugAndSampled()
 				internalTags = append(internalTags, Tag{key: JaegerDebugHeader, value: parent.debugID})
@@ -290,6 +293,7 @@ func (t *Tracer) startSpanWithOptions(
 			ctx.samplingState = parent.samplingState
 			if parent.remote {
 				ctx.samplingState.setFinal()
+				ctx.samplingState.localRootSpan = ctx.spanID
 			}
 		}
 		if hasParent {