fix(sampling): Force manual sampling on locked traces (#4084)

Co-authored-by: benjamin.debernardi <[email protected]>

Author: genesor

ddtrace/tracer/span.go

@@ -408,6 +408,21 @@ func (s *Span) setSamplingPriorityLocked(priority int, sampler samplernames.Samp
 	s.context.setSamplingPriority(priority, sampler)
 }
 
+// forceSetSamplingPriorityLocked updates the sampling priority.
+// If the trace is locked, the sampling priority is forced to the given value.
+//
+// This function is should only be used when applying a manual keep or drop decision.
+func (s *Span) forceSetSamplingPriorityLocked(priority int, sampler samplernames.SamplerName) {
+	// We don't lock spans when flushing, so we could have a data race when
+	// modifying a span as it's being flushed. This protects us against that
+	// race, since spans are marked `finished` before we flush them.
+	if s.finished {
+		return
+	}
+	s.setMetric(keySamplingPriority, float64(priority))
+	s.context.forceSetSamplingPriority(priority, sampler)
+}
+
 // setTagError sets the error tag. It accounts for various valid scenarios.
 // This method is not safe for concurrent use.
 func (s *Span) setTagError(value interface{}, cfg errorConfig) {
@@ -549,11 +564,11 @@ func (s *Span) setTagBool(key string, v bool) {
 		}
 	case ext.ManualDrop:
 		if v {
-			s.setSamplingPriorityLocked(ext.PriorityUserReject, samplernames.Manual)
+			s.forceSetSamplingPriorityLocked(ext.PriorityUserReject, samplernames.Manual)
 		}
 	case ext.ManualKeep:
 		if v {
-			s.setSamplingPriorityLocked(ext.PriorityUserKeep, samplernames.Manual)
+			s.forceSetSamplingPriorityLocked(ext.PriorityUserKeep, samplernames.Manual)
 		}
 	default:
 		if v {

ddtrace/tracer/span_test.go

@@ -571,6 +571,52 @@ func TestTraceManualKeepAndManualDrop(t *testing.T) {
 			span.SetTag(scenario.tag, true)
 			assert.Equal(t, scenario.keep, shouldKeep(span))
 		})
+		t.Run(fmt.Sprintf("%s/upstream-drop-locked", scenario.tag), func(t *testing.T) {
+			tracer, err := newTracer()
+			defer tracer.Stop()
+			assert.NoError(t, err)
+
+			spanCtx := &SpanContext{
+				traceID: traceIDFrom64Bits(42),
+				spanID:  42,
+				trace:   newTrace(),
+			}
+
+			// Set sampling priority (0 = drop decision from upstream) & lock the trace
+			// mimicking inheriting a trace from an upstream service with a drop decision.
+			spanCtx.setSamplingPriority(ext.PriorityAutoReject, samplernames.Unknown)
+			spanCtx.trace.setLocked(true)
+
+			span := tracer.StartSpan("child span with sampling decision", ChildOf(spanCtx))
+			span.SetTag(scenario.tag, true)
+
+			// The sampling decision should be applied as manual sampling takes
+			// precedence over propagated decision
+			assert.Equal(t, scenario.keep, shouldKeep(span))
+		})
+		t.Run(fmt.Sprintf("%s/upstream-keep-locked", scenario.tag), func(t *testing.T) {
+			tracer, err := newTracer()
+			defer tracer.Stop()
+			assert.NoError(t, err)
+
+			spanCtx := &SpanContext{
+				traceID: traceIDFrom64Bits(42),
+				spanID:  42,
+				trace:   newTrace(),
+			}
+
+			// Set sampling priority (1 = keep decision from upstream) & lock the trace
+			// mimicking inheriting a trace from an upstream service with a keep decision.
+			spanCtx.setSamplingPriority(ext.PriorityAutoKeep, samplernames.Unknown)
+			spanCtx.trace.setLocked(true)
+
+			span := tracer.StartSpan("child span with sampling decision", ChildOf(spanCtx))
+			span.SetTag(scenario.tag, true)
+
+			// The sampling decision should be applied as manual sampling takes
+			// precedence over propagated decision
+			assert.Equal(t, scenario.keep, shouldKeep(span))
+		})
 	}
 }
 

ddtrace/tracer/spancontext.go

@@ -274,6 +274,17 @@ func (c *SpanContext) setSamplingPriority(p int, sampler samplernames.SamplerNam
 	}
 }
 
+// forceSetSamplingPriority sets (and forces if the trace is locked) the sampling priority and decision maker (based on `sampler`).
+func (c *SpanContext) forceSetSamplingPriority(p int, sampler samplernames.SamplerName) {
+	if c.trace == nil {
+		c.trace = newTrace()
+	}
+	if c.trace.forceSetSamplingPriority(p, sampler) {
+		// the trace's sampling priority or sampler was updated: mark this as updated
+		c.updated = true
+	}
+}
+
 func (c *SpanContext) SamplingPriority() (p int, ok bool) {
 	if c == nil || c.trace == nil {
 		return 0, false
@@ -396,6 +407,14 @@ func (t *trace) setSamplingPriority(p int, sampler samplernames.SamplerName) boo
 	return t.setSamplingPriorityLocked(p, sampler)
 }
 
+// forceSetSamplingPriority forces the sampling priority and the decision maker
+// and returns true if it was modified.
+func (t *trace) forceSetSamplingPriority(p int, sampler samplernames.SamplerName) bool {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.setSamplingPriorityLockedWithForce(p, sampler, true)
+}
+
 func (t *trace) keep() {
 	atomic.CompareAndSwapUint32((*uint32)(&t.samplingDecision), uint32(decisionNone), uint32(decisionKeep))
 }
@@ -421,8 +440,13 @@ func samplerToDM(sampler samplernames.SamplerName) string {
 	return "-" + strconv.Itoa(int(sampler))
 }
 
-func (t *trace) setSamplingPriorityLocked(p int, sampler samplernames.SamplerName) bool {
-	if t.locked {
+// setSamplingPriority sets the sampling priority and the decision maker
+// and returns true if it was modified.
+//
+// The force parameter is used to bypass the locked sampling decision check
+// when setting the sampling priority. This is used to apply a manual keep or drop decision.
+func (t *trace) setSamplingPriorityLockedWithForce(p int, sampler samplernames.SamplerName, force bool) bool {
+	if t.locked && !force {
 		return false
 	}
 
@@ -455,6 +479,10 @@ func (t *trace) setSamplingPriorityLocked(p int, sampler samplernames.SamplerNam
 	return updatedPriority
 }
 
+func (t *trace) setSamplingPriorityLocked(p int, sampler samplernames.SamplerName) bool {
+	return t.setSamplingPriorityLockedWithForce(p, sampler, false)
+}
+
 func (t *trace) isLocked() bool {
 	t.mu.RLock()
 	defer t.mu.RUnlock()