Skip to content

Commit 419da21

Browse files
mccullsmcculls
authored
Always record force keep decisions in any future distributed trace propagation (#10308)
1 parent 9ec266e commit 419da21

File tree

4 files changed

+47
-39
lines changed

4 files changed

+47
-39
lines changed

‎dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -530,10 +530,9 @@ public void forceKeep(byte samplingMechanism) {
530530
private void forceKeepThisSpan(byte samplingMechanism) {
531531
// if the user really wants to keep this trace chunk, we will let them,
532532
// even if the old sampling priority and mechanism have already propagated
533-
if (SAMPLING_PRIORITY_UPDATER.getAndSet(this, PrioritySampling.USER_KEEP)
534-
== PrioritySampling.UNSET) {
535-
propagationTags.updateTraceSamplingPriority(PrioritySampling.USER_KEEP, samplingMechanism);
536-
}
533+
SAMPLING_PRIORITY_UPDATER.set(this, PrioritySampling.USER_KEEP);
534+
// record force keep decision for future distributed trace propagation
535+
propagationTags.forceKeep(samplingMechanism);
537536
}
538537

539538
public void addPropagatedTraceSource(final int value) {

‎dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ public interface Factory {
5757
*/
5858
public abstract void updateTraceSamplingPriority(int samplingPriority, int samplingMechanism);
5959

60+
public abstract void forceKeep(int samplingMechanism);
61+
6062
public abstract int getSamplingPriority();
6163

6264
public abstract void updateTraceOrigin(CharSequence origin);

‎dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java‎

Lines changed: 34 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -185,37 +185,46 @@ static PTags withError(PTagsFactory factory, String error) {
185185
public void updateTraceSamplingPriority(int samplingPriority, int samplingMechanism) {
186186
if (samplingPriority != PrioritySampling.UNSET && canChangeDecisionMaker
187187
|| samplingMechanism == SamplingMechanism.EXTERNAL_OVERRIDE) {
188-
if (this.samplingPriority != samplingPriority) {
189-
// This should invalidate any cached w3c header
190-
clearCachedHeader(W3C);
188+
doUpdateTraceSamplingPriority(samplingPriority, samplingMechanism);
189+
}
190+
}
191+
192+
@Override
193+
public void forceKeep(int samplingMechanism) {
194+
doUpdateTraceSamplingPriority(PrioritySampling.USER_KEEP, samplingMechanism);
195+
}
196+
197+
private void doUpdateTraceSamplingPriority(int samplingPriority, int samplingMechanism) {
198+
if (this.samplingPriority != samplingPriority) {
199+
// This should invalidate any cached w3c header
200+
clearCachedHeader(W3C);
201+
}
202+
this.samplingPriority = samplingPriority;
203+
if (samplingPriority > 0) {
204+
// TODO should try to keep the old sampling mechanism if we override the value?
205+
if (samplingMechanism == SamplingMechanism.EXTERNAL_OVERRIDE) {
206+
// There is no specific value for the EXTERNAL_OVERRIDE, so say that it's the DEFAULT
207+
samplingMechanism = SamplingMechanism.DEFAULT;
191208
}
192-
this.samplingPriority = samplingPriority;
193-
if (samplingPriority > 0) {
194-
// TODO should try to keep the old sampling mechanism if we override the value?
195-
if (samplingMechanism == SamplingMechanism.EXTERNAL_OVERRIDE) {
196-
// There is no specific value for the EXTERNAL_OVERRIDE, so say that it's the DEFAULT
197-
samplingMechanism = SamplingMechanism.DEFAULT;
198-
}
199-
// Protect against possible SamplingMechanism.UNKNOWN (-1) that doesn't comply with the
200-
// format
201-
if (samplingMechanism >= 0) {
202-
TagValue newDM = TagValue.from("-" + samplingMechanism);
203-
if (!newDM.equals(decisionMakerTagValue)) {
204-
// This should invalidate any cached w3c and datadog header
205-
clearCachedHeader(DATADOG);
206-
clearCachedHeader(W3C);
207-
}
208-
decisionMakerTagValue = newDM;
209-
}
210-
} else {
211-
// Drop the decision maker tag
212-
if (decisionMakerTagValue != null) {
209+
// Protect against possible SamplingMechanism.UNKNOWN (-1) that doesn't comply with the
210+
// format
211+
if (samplingMechanism >= 0) {
212+
TagValue newDM = TagValue.from("-" + samplingMechanism);
213+
if (!newDM.equals(decisionMakerTagValue)) {
213214
// This should invalidate any cached w3c and datadog header
214215
clearCachedHeader(DATADOG);
215216
clearCachedHeader(W3C);
216217
}
217-
decisionMakerTagValue = null;
218+
decisionMakerTagValue = newDM;
219+
}
220+
} else {
221+
// Drop the decision maker tag
222+
if (decisionMakerTagValue != null) {
223+
// This should invalidate any cached w3c and datadog header
224+
clearCachedHeader(DATADOG);
225+
clearCachedHeader(W3C);
218226
}
227+
decisionMakerTagValue = null;
219228
}
220229
}
221230

‎dd-trace-core/src/test/groovy/datadog/trace/core/DDSpanContextPropagationTagsTest.groovy‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -99,11 +99,10 @@ class DDSpanContextPropagationTagsTest extends DDCoreSpecification {
9999
dd.createTagMap() == tagMap
100100

101101
where:
102-
priority | header | newHeader | tagMap
103-
UNSET | "_dd.p.usr=123" | "_dd.p.dm=-4,_dd.p.usr=123" | ["_dd.p.dm": "-4", "_dd.p.usr": "123"]
104-
// decision has already been made, propagate as-is
105-
SAMPLER_KEEP | "_dd.p.dm=9bf3439f2f-1,_dd.p.usr=123" | "_dd.p.dm=9bf3439f2f-1,_dd.p.usr=123" | ["_dd.p.dm": "9bf3439f2f-1", "_dd.p.usr": "123"]
106-
SAMPLER_KEEP | "_dd.p.usr=123" | "_dd.p.usr=123" | ["_dd.p.usr": "123"]
102+
priority | header | newHeader | tagMap
103+
UNSET | "_dd.p.usr=123" | "_dd.p.dm=-4,_dd.p.usr=123" | ["_dd.p.dm": "-4", "_dd.p.usr": "123"]
104+
SAMPLER_KEEP | "_dd.p.dm=9bf3439f2f-1,_dd.p.usr=123" | "_dd.p.dm=-4,_dd.p.usr=123" | ["_dd.p.dm": "-4", "_dd.p.usr": "123"]
105+
SAMPLER_KEEP | "_dd.p.usr=123" | "_dd.p.dm=-4,_dd.p.usr=123" | ["_dd.p.dm": "-4", "_dd.p.usr": "123"]
107106
}
108107

109108
def "forceKeep trace PropagationTags #priority #header"() {
@@ -127,10 +126,9 @@ class DDSpanContextPropagationTagsTest extends DDCoreSpecification {
127126
ddRoot.createTagMap() == rootTagMap
128127

129128
where:
130-
priority | header | rootHeader | rootTagMap
131-
UNSET | "_dd.p.usr=123" | "_dd.p.dm=-4,_dd.p.usr=123" | ["_dd.p.dm": "-4", "_dd.p.usr": "123"]
132-
// decision has already been made, propagate as-is
133-
SAMPLER_KEEP | "_dd.p.dm=9bf3439f2f-1,_dd.p.usr=123" | "_dd.p.dm=9bf3439f2f-1,_dd.p.usr=123" | ["_dd.p.dm": "9bf3439f2f-1", "_dd.p.usr": "123"]
134-
SAMPLER_KEEP | "_dd.p.usr=123" | "_dd.p.usr=123" | ["_dd.p.usr": "123"]
129+
priority | header | rootHeader | rootTagMap
130+
UNSET | "_dd.p.usr=123" | "_dd.p.dm=-4,_dd.p.usr=123" | ["_dd.p.dm": "-4", "_dd.p.usr": "123"]
131+
SAMPLER_KEEP | "_dd.p.dm=9bf3439f2f-1,_dd.p.usr=123" | "_dd.p.dm=-4,_dd.p.usr=123" | ["_dd.p.dm": "-4", "_dd.p.usr": "123"]
132+
SAMPLER_KEEP | "_dd.p.usr=123" | "_dd.p.dm=-4,_dd.p.usr=123" | ["_dd.p.dm": "-4", "_dd.p.usr": "123"]
135133
}
136134
}

0 commit comments

Comments
 (0)