ddtrace/tracer: allow changes of priority even when the root is non-local

ddtrace/tracer/span_test.go

@@ -8,13 +8,15 @@ package tracer
 import (
 	"errors"
 	"fmt"
+	"math"
 	"os"
 	"strings"
 	"sync/atomic"
 	"testing"
 	"time"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/samplernames"
 
 	"github.com/DataDog/datadog-agent/pkg/obfuscate"
 	"github.com/stretchr/testify/assert"
@@ -348,6 +350,33 @@ func TestSpanSetTagError(t *testing.T) {
 	})
 }
 
+func TestTraceManualKeepAndManualDrop(t *testing.T) {
+	for _, scenario := range []struct {
+		tag  string
+		keep bool
+		p    int // priority
+	}{
+		{ext.ManualKeep, true, 0},
+		{ext.ManualDrop, false, 1},
+	} {
+		t.Run(fmt.Sprintf("%s/local", scenario.tag), func(t *testing.T) {
+			tracer := newTracer()
+			span := tracer.newRootSpan("root span", "my service", "my resource")
+			span.SetTag(scenario.tag, true)
+			assert.Equal(t, scenario.keep, shouldKeep(span))
+		})
+
+		t.Run(fmt.Sprintf("%s/non-local", scenario.tag), func(t *testing.T) {
+			tracer := newTracer()
+			spanCtx := &spanContext{traceID: 42, spanID: 42}
+			spanCtx.setSamplingPriority("", scenario.p, samplernames.Upstream, math.NaN())
+			span := tracer.StartSpan("non-local root span", ChildOf(spanCtx)).(*span)
+			span.SetTag(scenario.tag, true)
+			assert.Equal(t, scenario.keep, shouldKeep(span))
+		})
+	}
+}
+
 func TestSpanSetDatadogTags(t *testing.T) {
 	assert := assert.New(t)
 

ddtrace/tracer/spancontext.go

@@ -221,11 +221,6 @@ func (t *trace) setSamplingPriorityLocked(service string, p int, sampler sampler
 	if t.locked {
 		return
 	}
-	if t.root == nil {
-		// this trace is distributed (no local root); modifications
-		// to the sampling priority are not allowed.
-		t.locked = true
-	}
 	if t.priority == nil {
 		t.priority = new(float64)
 	}

ddtrace/tracer/spancontext_test.go

@@ -242,21 +242,6 @@ func TestSpanFinishPriority(t *testing.T) {
 	assert.Fail("span not found")
 }
 
-func TestTracePriorityLocked(t *testing.T) {
-	assert := assert.New(t)
-	ddHeaders := TextMapCarrier(map[string]string{
-		DefaultTraceIDHeader:  "2",
-		DefaultParentIDHeader: "2",
-		DefaultPriorityHeader: "2",
-	})
-
-	ctx, err := NewPropagator(nil).Extract(ddHeaders)
-	assert.Nil(err)
-	sctx, ok := ctx.(*spanContext)
-	assert.True(ok)
-	assert.True(sctx.trace.locked)
-}
-
 func TestNewSpanContext(t *testing.T) {
 	t.Run("basic", func(t *testing.T) {
 		span := &span{