[sdk-tracing] Allow samplers to set TraceState for propagation-only spans (open-telemetry#6058)

Author: CodeBlanch

src/OpenTelemetry/CHANGELOG.md

@@ -12,6 +12,10 @@ Notes](../../RELEASENOTES.md).
   now lead to unique metrics.
   ([#5982](https://github.com/open-telemetry/opentelemetry-dotnet/pull/5982))
 
+* Fixed a bug in tracing where `TraceState` set by a custom `Sampler` is not
+  applied when creating propagation-only spans.
+  ([#6058](https://github.com/open-telemetry/opentelemetry-dotnet/pull/6058))
+
 ## 1.11.0-rc.1
 
 Released 2024-Dec-11

src/OpenTelemetry/Trace/TracerProviderSdk.cs

@@ -237,7 +237,7 @@ internal TracerProviderSdk(
         else if (this.sampler is AlwaysOffSampler)
         {
             activityListener.Sample = (ref ActivityCreationOptions<ActivityContext> options) =>
-                !Sdk.SuppressInstrumentation ? PropagateOrIgnoreData(options.Parent) : ActivitySamplingResult.None;
+                !Sdk.SuppressInstrumentation ? PropagateOrIgnoreData(ref options) : ActivitySamplingResult.None;
             this.getRequestedDataAction = this.RunGetRequestedDataAlwaysOffSampler;
         }
         else
@@ -493,47 +493,46 @@ private static ActivitySamplingResult ComputeActivitySamplingResult(
         {
             SamplingDecision.RecordAndSample => ActivitySamplingResult.AllDataAndRecorded,
             SamplingDecision.RecordOnly => ActivitySamplingResult.AllData,
-            _ => ActivitySamplingResult.PropagationData,
+            _ => PropagateOrIgnoreData(ref options),
         };
 
-        if (activitySamplingResult != ActivitySamplingResult.PropagationData)
+        if (activitySamplingResult > ActivitySamplingResult.PropagationData)
         {
             foreach (var att in samplingResult.Attributes)
             {
                 options.SamplingTags.Add(att.Key, att.Value);
             }
+        }
 
+        if (activitySamplingResult != ActivitySamplingResult.None
+            && samplingResult.TraceStateString != null)
+        {
             // https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/sdk.md#sampler
             // Spec requires clearing Tracestate if empty Tracestate is returned.
             // Since .NET did not have this capability, it'll break
             // existing samplers if we did that. So the following is
             // adopted to remain spec-compliant and backward compat.
             // The behavior is:
-            // if sampler returns null, its treated as if it has no intend
+            // if sampler returns null, its treated as if it has not intended
             // to change Tracestate. Existing SamplingResult ctors will put null as default TraceStateString,
             // so all existing samplers will get this behavior.
             // if sampler returns non-null, then it'll be used as the
             // new value for Tracestate
             // A sampler can return string.Empty if it intends to clear the state.
-            if (samplingResult.TraceStateString != null)
-            {
-                options = options with { TraceState = samplingResult.TraceStateString };
-            }
-
-            return activitySamplingResult;
+            options = options with { TraceState = samplingResult.TraceStateString };
         }
 
-        return PropagateOrIgnoreData(options.Parent);
+        return activitySamplingResult;
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private static ActivitySamplingResult PropagateOrIgnoreData(in ActivityContext parentContext)
+    private static ActivitySamplingResult PropagateOrIgnoreData(ref ActivityCreationOptions<ActivityContext> options)
     {
-        var isRootSpan = parentContext.TraceId == default;
+        var isRootSpan = options.Parent.TraceId == default;
 
         // If it is the root span or the parent is remote select PropagationData so the trace ID is preserved
         // even if no activity of the trace is recorded (sampled per OpenTelemetry parlance).
-        return (isRootSpan || parentContext.IsRemote)
+        return (isRootSpan || options.Parent.IsRemote)
             ? ActivitySamplingResult.PropagationData
             : ActivitySamplingResult.None;
     }
@@ -606,11 +605,11 @@ private void RunGetRequestedDataOtherSampler(Activity activity)
             {
                 activity.SetTag(att.Key, att.Value);
             }
+        }
 
-            if (samplingResult.TraceStateString != null)
-            {
-                activity.TraceStateString = samplingResult.TraceStateString;
-            }
+        if (samplingResult.TraceStateString != null)
+        {
+            activity.TraceStateString = samplingResult.TraceStateString;
         }
     }
 }

test/OpenTelemetry.Tests/Trace/SamplersTest.cs

@@ -100,37 +100,7 @@ public void TracerProviderSdkSamplerAttributesAreAppliedToLegacyActivity(Samplin
     [InlineData(SamplingDecision.RecordAndSample)]
     public void SamplersCanModifyTraceStateOnLegacyActivity(SamplingDecision samplingDecision)
     {
-        var existingTraceState = "a=1,b=2";
-        var newTraceState = "a=1,b=2,c=3,d=4";
-        var testSampler = new TestSampler
-        {
-            SamplingAction = (samplingParams) =>
-            {
-                Assert.Equal(existingTraceState, samplingParams.ParentContext.TraceState);
-                return new SamplingResult(samplingDecision, newTraceState);
-            },
-        };
-
-        var operationNameForLegacyActivity = Utils.GetCurrentMethodName();
-        using var tracerProvider = Sdk.CreateTracerProviderBuilder()
-                    .SetSampler(testSampler)
-                    .AddLegacySource(operationNameForLegacyActivity)
-                    .Build();
-
-        using var parentActivity = new Activity("Foo");
-        parentActivity.TraceStateString = existingTraceState;
-        parentActivity.Start();
-
-        using var activity = new Activity(operationNameForLegacyActivity);
-        activity.Start();
-        Assert.NotNull(activity);
-        if (samplingDecision != SamplingDecision.Drop)
-        {
-            Assert.Equal(newTraceState, activity.TraceStateString);
-        }
-
-        activity.Stop();
-        parentActivity.Stop();
+        RunLegacyActivitySamplerTest(samplingDecision, samplerTraceState: "a=1,b=2,c=3,d=4");
     }
 
     [Theory]
@@ -139,36 +109,7 @@ public void SamplersCanModifyTraceStateOnLegacyActivity(SamplingDecision samplin
     [InlineData(SamplingDecision.RecordAndSample)]
     public void SamplersDoesNotImpactTraceStateWhenUsingNullLegacyActivity(SamplingDecision samplingDecision)
     {
-        var existingTraceState = "a=1,b=2";
-        var testSampler = new TestSampler
-        {
-            SamplingAction = (samplingParams) =>
-            {
-                Assert.Equal(existingTraceState, samplingParams.ParentContext.TraceState);
-                return new SamplingResult(samplingDecision);
-            },
-        };
-
-        var operationNameForLegacyActivity = Utils.GetCurrentMethodName();
-        using var tracerProvider = Sdk.CreateTracerProviderBuilder()
-                    .SetSampler(testSampler)
-                    .AddLegacySource(operationNameForLegacyActivity)
-                    .Build();
-
-        using var parentActivity = new Activity("Foo");
-        parentActivity.TraceStateString = existingTraceState;
-        parentActivity.Start();
-
-        using var activity = new Activity(operationNameForLegacyActivity);
-        activity.Start();
-        Assert.NotNull(activity);
-        if (samplingDecision != SamplingDecision.Drop)
-        {
-            Assert.Equal(existingTraceState, activity.TraceStateString);
-        }
-
-        activity.Stop();
-        parentActivity.Stop();
+        RunLegacyActivitySamplerTest(samplingDecision, samplerTraceState: null);
     }
 
     [Theory]
@@ -183,7 +124,15 @@ public void SamplersCanModifyTraceState(SamplingDecision sampling)
         {
             SamplingAction = (samplingParams) =>
             {
-                Assert.Equal(parentTraceState, samplingParams.ParentContext.TraceState);
+                if (samplingParams.Name == "root")
+                {
+                    Assert.Equal(parentTraceState, samplingParams.ParentContext.TraceState);
+                }
+                else
+                {
+                    Assert.Equal(newTraceState, samplingParams.ParentContext.TraceState);
+                }
+
                 return new SamplingResult(sampling, newTraceState);
             },
         };
@@ -195,13 +144,54 @@ public void SamplersCanModifyTraceState(SamplingDecision sampling)
             .SetSampler(testSampler)
             .Build();
 
-        var parentContext = new ActivityContext(ActivityTraceId.CreateRandom(), ActivitySpanId.CreateRandom(), ActivityTraceFlags.Recorded, parentTraceState, true);
+        // Note: Remote parent is set as recorded
+        var parentContext = new ActivityContext(ActivityTraceId.CreateRandom(), ActivitySpanId.CreateRandom(), ActivityTraceFlags.Recorded, parentTraceState, isRemote: true);
+
+        using var root = activitySource.StartActivity("root", ActivityKind.Server, parentContext);
+
+        // Note: We always create a root even for Drop. When dropping the
+        // created root is for propagation only
+        Assert.NotNull(root);
+        Assert.Equal(newTraceState, root.TraceStateString);
+
+        if (sampling == SamplingDecision.RecordAndSample)
+        {
+            Assert.True(root.Recorded);
+            Assert.True(root.IsAllDataRequested);
+        }
+        else if (sampling == SamplingDecision.RecordOnly)
+        {
+            // TODO: Update this when repo consumes DS v10.
+            // Note: Seems to be a bug in DiagnosticSource. Root in this case
+            // inherits context from the remote parent and Recorded doesn't get
+            // cleared. This should be fixed in .NET 10:
+            // https://github.com/dotnet/runtime/pull/111289
+            // Assert.False(root.Recorded);
+
+            Assert.True(root.IsAllDataRequested);
+        }
+        else
+        {
+            // TODO: Update this when repo consumes DS v10.
+            // Note: Seems to be a bug in DiagnosticSource. Root in this case
+            // inherits context from the remote parent and Recorded doesn't get
+            // cleared. This should be fixed in .NET 10:
+            // https://github.com/dotnet/runtime/pull/111289
+            // Assert.False(root.Recorded);
+
+            Assert.False(root.IsAllDataRequested);
+        }
+
+        using var child = activitySource.StartActivity("child", ActivityKind.Server);
 
-        using var activity = activitySource.StartActivity("root", ActivityKind.Server, parentContext);
         if (sampling != SamplingDecision.Drop)
         {
-            Assert.NotNull(activity);
-            Assert.Equal(newTraceState, activity.TraceStateString);
+            Assert.NotNull(child);
+            Assert.Equal(newTraceState, child.TraceStateString);
+        }
+        else
+        {
+            Assert.Null(child);
         }
     }
 
@@ -260,6 +250,60 @@ public void SamplerExceptionBubblesUpTest()
         Assert.Throws<InvalidOperationException>(() => activitySource.StartActivity("ThrowingSampler"));
     }
 
+    private static void RunLegacyActivitySamplerTest(SamplingDecision samplingDecision, string? samplerTraceState)
+    {
+        var existingTraceState = "a=1,b=2";
+
+        var operationNameForLegacyActivity = Utils.GetCurrentMethodName();
+
+        var testSampler = new TestSampler
+        {
+            SamplingAction = (samplingParams) =>
+            {
+                Assert.Equal(samplingParams.Name, operationNameForLegacyActivity);
+                Assert.Equal(existingTraceState, samplingParams.ParentContext.TraceState);
+                return new SamplingResult(samplingDecision, samplerTraceState);
+            },
+        };
+
+        using var tracerProvider = Sdk.CreateTracerProviderBuilder()
+            .SetSampler(testSampler)
+            .AddLegacySource(operationNameForLegacyActivity)
+            .Build();
+
+        using var parentActivity = new Activity("Foo");
+        parentActivity.TraceStateString = existingTraceState;
+        parentActivity.Start();
+
+        using var childActivity = new Activity(operationNameForLegacyActivity);
+        childActivity.Start();
+
+        if (samplerTraceState != null)
+        {
+            Assert.Equal(samplerTraceState, childActivity.TraceStateString);
+        }
+        else
+        {
+            Assert.Equal(existingTraceState, childActivity.TraceStateString);
+        }
+
+        if (samplingDecision == SamplingDecision.RecordAndSample)
+        {
+            Assert.True(childActivity.Recorded);
+            Assert.True(childActivity.IsAllDataRequested);
+        }
+        else if (samplingDecision == SamplingDecision.RecordOnly)
+        {
+            Assert.False(childActivity.Recorded);
+            Assert.True(childActivity.IsAllDataRequested);
+        }
+        else
+        {
+            Assert.False(childActivity.Recorded);
+            Assert.False(childActivity.IsAllDataRequested);
+        }
+    }
+
     private sealed class ThrowingSampler : Sampler
     {
         public override SamplingResult ShouldSample(in SamplingParameters samplingParameters)