Skip to content

Commit

Permalink
[sdk-tracing] Allow samplers to set TraceState for propagation-only s…
Browse files Browse the repository at this point in the history 
…pans (#6058)
  • Loading branch information
@CodeBlanch
CodeBlanch authored Jan 15, 2025
1 parent 17bbd83 commit 3d1f74d
Show file tree
Hide file tree
Showing 3 changed files with 131 additions and 84 deletions.
4 changes: 4 additions & 0 deletions src/OpenTelemetry/CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
35 changes: 17 additions & 18 deletions src/OpenTelemetry/Trace/TracerProviderSdk.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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;
}
}
}
176 changes: 110 additions & 66 deletions test/OpenTelemetry.Tests/Trace/SamplersTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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]
Expand All @@ -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);
},
};
Expand All @@ -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);
}
}

Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 3d1f74d

Please sign in to comment.