Skip to content

Commit d5a6721

Browse files
MilanGarnierMilanGarnier
committed
fix(otlp): resolve tracestate during span export
1 parent 927edb4 commit d5a6721

9 files changed

Lines changed: 180 additions & 15 deletions

File tree

‎dd-trace-core/src/main/java/datadog/trace/core/otlp/trace/OtlpTraceJson.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,9 @@ public static void writeSpan(
5656
writer.name("traceId").value(hexTraceId(span.getTraceId()));
5757
writer.name("spanId").value(hexSpanId(span.getSpanId()));
5858

59-
String tracestate = propagationTags.getW3CTracestate();
59+
int samplingPriority = span.samplingPriority();
60+
// TODO Cache the effective tracestate once per trace.
61+
String tracestate = propagationTags.getW3CTracestate(samplingPriority);
6062
if (tracestate != null) {
6163
writer.name("traceState").value(tracestate);
6264
}
@@ -66,7 +68,7 @@ public static void writeSpan(
6668
}
6769

6870
int traceFlags = NO_TRACE_FLAGS;
69-
if (span.samplingPriority() > 0) {
71+
if (samplingPriority > 0) {
7072
traceFlags |= SAMPLED_TRACE_FLAG;
7173
}
7274
if (span.spanContext().isRemote()) {

‎dd-trace-core/src/main/java/datadog/trace/core/otlp/trace/OtlpTraceProto.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,9 @@ public static int recordSpanMessage(
9090
writeTag(buf, 2, LEN_WIRE_TYPE);
9191
writeSpanId(buf, span.getSpanId());
9292

93-
String tracestate = propagationTags.getW3CTracestate();
93+
int samplingPriority = span.samplingPriority();
94+
// TODO Cache the effective tracestate once per trace.
95+
String tracestate = propagationTags.getW3CTracestate(samplingPriority);
9496
if (tracestate != null) {
9597
writeTag(buf, 3, LEN_WIRE_TYPE);
9698
writeString(buf, tracestate);
@@ -102,7 +104,7 @@ public static int recordSpanMessage(
102104
}
103105

104106
int traceFlags = NO_TRACE_FLAGS;
105-
if (span.samplingPriority() > 0) {
107+
if (samplingPriority > 0) {
106108
traceFlags |= SAMPLED_TRACE_FLAG;
107109
}
108110
if (span.spanContext().isRemote()) {

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,11 @@ public abstract void updateOtelTraceState(
9090
*/
9191
public abstract String getW3CTracestate();
9292

93+
/** Gets the W3C tracestate with OTel sampling state resolved for the supplied priority. */
94+
public String getW3CTracestate(int samplingPriority) {
95+
return getW3CTracestate();
96+
}
97+
9398
/**
9499
* Stores the original <a href="https://www.w3.org/TR/trace-context/#tracestate-header">W3C
95100
* tracestate header</a> value.

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -234,8 +234,12 @@ public void forceKeep(int samplingMechanism) {
234234
}
235235

236236
private void doUpdateTraceSamplingPriority(int samplingPriority, int samplingMechanism) {
237+
boolean initialUnknownPriority =
238+
this.samplingPriority == PrioritySampling.UNSET
239+
&& samplingMechanism == SamplingMechanism.UNKNOWN;
237240
boolean removeOtelProbability =
238-
samplingMechanism != SamplingMechanism.EXTERNAL_OVERRIDE
241+
!initialUnknownPriority
242+
&& samplingMechanism != SamplingMechanism.EXTERNAL_OVERRIDE
239243
&& !isProbabilitySamplingMechanism(samplingMechanism);
240244
if (this.samplingPriority != samplingPriority) {
241245
// This should invalidate any cached w3c header
@@ -580,6 +584,16 @@ public String getW3CTracestate() {
580584
return this.tracestate;
581585
}
582586

587+
@Override
588+
public String getW3CTracestate(int samplingPriority) {
589+
OtelTraceState current = otelTraceState;
590+
OtelTraceState resolved = resolveOtelTraceState(W3C, samplingPriority);
591+
if (current == resolved) {
592+
return tracestate;
593+
}
594+
return W3CPTagsCodec.updateOtelTraceState(this, resolved);
595+
}
596+
583597
@Override
584598
public void updateW3CTracestate(String tracestate) {
585599
setW3CTracestate(tracestate, W3CPTagsCodec.extractOtelTraceState(tracestate));

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

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -356,7 +356,7 @@ protected int appendSuffix(
356356
size = 0;
357357
}
358358
// Append the managed OTel member and all other non-Datadog list-members
359-
if (appendOtelAndVendorMembers(sb, ptags, otelTraceState, size != 0)) {
359+
if (appendOtelAndVendorMembers(sb, ptags, otelTraceState, size != 0, false)) {
360360
// We don't care about the total size in bytes here, but only the fact that we added something
361361
// that should be returned
362362
size = Math.max(size, EMPTY_SIZE + 1);
@@ -763,7 +763,11 @@ private static int cleanUpAndAppendUnknown(StringBuilder sb, W3CPTags w3CPTags,
763763
}
764764

765765
private static boolean appendOtelAndVendorMembers(
766-
StringBuilder sb, PTags ptags, OtelTraceState otelTraceState, boolean hasDatadogMember) {
766+
StringBuilder sb,
767+
PTags ptags,
768+
OtelTraceState otelTraceState,
769+
boolean hasDatadogMember,
770+
boolean preserveDatadogMember) {
767771
String original = ptags.tracestate;
768772
int remainingMembers = MAX_MEMBER_COUNT - (hasDatadogMember ? 1 : 0);
769773
int otherMemberPosition = 0;
@@ -777,10 +781,14 @@ private static boolean appendOtelAndVendorMembers(
777781
if (memberEnd < 0) {
778782
memberEnd = len;
779783
}
780-
boolean managedMember =
781-
original.startsWith(DATADOG_MEMBER_KEY, memberStart)
782-
|| original.startsWith(OTEL_MEMBER_KEY, memberStart);
783-
if (!managedMember) {
784+
boolean datadogMember = original.startsWith(DATADOG_MEMBER_KEY, memberStart);
785+
boolean otelMember = original.startsWith(OTEL_MEMBER_KEY, memberStart);
786+
if (datadogMember && preserveDatadogMember) {
787+
int end = stripTrailingOWC(original, memberStart, memberEnd);
788+
appendMember(sb, original, memberStart, end);
789+
remainingMembers--;
790+
memberAppended = true;
791+
} else if (!datadogMember && !otelMember) {
784792
if (otelTraceState != null
785793
&& !otelTraceStateAppended
786794
&& otelTraceState.getInheritedPosition() == otherMemberPosition) {
@@ -810,6 +818,17 @@ private static boolean appendOtelAndVendorMembers(
810818
return memberAppended;
811819
}
812820

821+
static String updateOtelTraceState(PTags ptags, OtelTraceState otelTraceState) {
822+
String original = ptags.tracestate;
823+
int capacity = original == null ? 0 : original.length();
824+
if (otelTraceState != null) {
825+
capacity += OTEL_MEMBER_KEY.length() + otelTraceState.length() + 1;
826+
}
827+
StringBuilder updated = new StringBuilder(capacity);
828+
appendOtelAndVendorMembers(updated, ptags, otelTraceState, false, true);
829+
return updated.length() == 0 ? null : updated.toString();
830+
}
831+
813832
private static void appendMember(StringBuilder sb, String member, int start, int end) {
814833
if (sb.length() != 0) {
815834
sb.append(MEMBER_SEPARATOR);

‎dd-trace-core/src/test/java/datadog/trace/core/otlp/trace/OtlpTraceJsonCollectorTest.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND;
44
import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_SERVER;
5+
import static datadog.trace.common.sampling.RuleBasedTraceSampler.SAMPLING_RULE_RATE;
56
import static datadog.trace.core.DDSpanContext.SPAN_SAMPLING_MECHANISM_TAG;
67
import static datadog.trace.core.otlp.common.OtlpCommonJson.hexSpanId;
78
import static datadog.trace.core.otlp.common.OtlpCommonJson.hexTraceId;
@@ -153,6 +154,30 @@ void spanTraceStateIncludedWhenPropagated() throws IOException {
153154
assertEquals("vendor=state", parsedSpan.get("traceState"));
154155
}
155156

157+
@Test
158+
void spanTraceStateIncludesLocalOtelSamplingState() throws IOException {
159+
AgentSpan agentSpan = TRACER.startSpan("test", "op.otel-tracestate");
160+
agentSpan.setResourceName("op.otel-tracestate");
161+
((DDSpan) agentSpan)
162+
.setSamplingPriority(
163+
PrioritySampling.USER_KEEP,
164+
SAMPLING_RULE_RATE,
165+
0.5,
166+
SamplingMechanism.LOCAL_USER_RULE,
167+
true);
168+
agentSpan.finish();
169+
170+
OtlpTraceJsonCollector collector = new OtlpTraceJsonCollector();
171+
collector.addTrace(asList((CoreSpan<?>) agentSpan));
172+
Map<String, Object> parsedSpan = onlySpan(collector.collectTraces());
173+
174+
String traceState = (String) parsedSpan.get("traceState");
175+
assertTrue(traceState.startsWith("ot=rv:"));
176+
assertTrue(traceState.contains(";th:8"));
177+
assertFalse(traceState.contains("dd="));
178+
assertEquals(SAMPLED_TRACE_FLAG, ((Number) parsedSpan.get("flags")).intValue());
179+
}
180+
156181
@Test
157182
void spanFlagsOmittedWhenNotSampled() throws IOException {
158183
AgentSpan agentSpan = TRACER.startSpan("test", "op.noflags");

‎dd-trace-core/src/test/java/datadog/trace/core/otlp/trace/OtlpTraceProtoTest.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_INTERNAL;
77
import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_PRODUCER;
88
import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_SERVER;
9+
import static datadog.trace.common.sampling.RuleBasedTraceSampler.SAMPLING_RULE_RATE;
910
import static datadog.trace.core.otlp.common.OtlpTraceFlags.SAMPLED_TRACE_FLAG;
1011
import static java.util.Arrays.asList;
1112
import static java.util.Arrays.copyOfRange;
@@ -676,6 +677,59 @@ void testSpanOrderInTracePreserved() throws IOException {
676677
"spans must appear in trace order in the payload");
677678
}
678679

680+
@Test
681+
void spanTraceStateIncludesLocalOtelSamplingState() throws IOException {
682+
AgentSpan agentSpan = TRACER.startSpan("test", "op.otel-tracestate");
683+
agentSpan.setResourceName("op.otel-tracestate");
684+
((DDSpan) agentSpan)
685+
.setSamplingPriority(
686+
PrioritySampling.USER_KEEP,
687+
SAMPLING_RULE_RATE,
688+
0.5,
689+
SamplingMechanism.LOCAL_USER_RULE,
690+
true);
691+
agentSpan.finish();
692+
693+
OtlpTraceProtoCollector collector = new OtlpTraceProtoCollector();
694+
collector.addTrace(asList((DDSpan) agentSpan));
695+
696+
String traceState = parseFirstSpanTraceState(collector.collectTraces());
697+
assertTrue(traceState.startsWith("ot=rv:"));
698+
assertTrue(traceState.contains(";th:8"));
699+
assertFalse(traceState.contains("dd="));
700+
}
701+
702+
private static String parseFirstSpanTraceState(OtlpPayload payload) throws IOException {
703+
CodedInputStream tracesData = CodedInputStream.newInstance(payload.getContent());
704+
tracesData.readTag();
705+
CodedInputStream resourceSpans = tracesData.readBytes().newCodedInput();
706+
707+
while (!resourceSpans.isAtEnd()) {
708+
int tag = resourceSpans.readTag();
709+
if (WireFormat.getTagFieldNumber(tag) != 2) {
710+
resourceSpans.skipField(tag);
711+
continue;
712+
}
713+
CodedInputStream scopeSpans = resourceSpans.readBytes().newCodedInput();
714+
while (!scopeSpans.isAtEnd()) {
715+
int scopeTag = scopeSpans.readTag();
716+
if (WireFormat.getTagFieldNumber(scopeTag) != 2) {
717+
scopeSpans.skipField(scopeTag);
718+
continue;
719+
}
720+
CodedInputStream span = scopeSpans.readBytes().newCodedInput();
721+
while (!span.isAtEnd()) {
722+
int spanTag = span.readTag();
723+
if (WireFormat.getTagFieldNumber(spanTag) == 3) {
724+
return span.readString();
725+
}
726+
span.skipField(spanTag);
727+
}
728+
}
729+
}
730+
return null;
731+
}
732+
679733
private static List<String> parseSpanNamesFromPayload(OtlpPayload payload) throws IOException {
680734
CodedInputStream tracesData = CodedInputStream.newInstance(payload.getContent());
681735
tracesData.readTag(); // field 1: TracesData.resource_spans

‎dd-trace-core/src/test/java/datadog/trace/core/propagation/OtelTraceStatePropagationTest.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import static datadog.trace.api.sampling.PrioritySampling.USER_KEEP;
55
import static datadog.trace.api.sampling.SamplingMechanism.LOCAL_USER_RULE;
66
import static datadog.trace.api.sampling.SamplingMechanism.MANUAL;
7+
import static datadog.trace.api.sampling.SamplingMechanism.UNKNOWN;
78
import static datadog.trace.core.propagation.PropagationTags.HeaderType.W3C;
89
import static org.junit.jupiter.api.Assertions.assertEquals;
910
import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -99,11 +100,27 @@ void manualKeepRetainsInheritedRandomnessAndUnknownFields() {
99100

100101
propagationTags.forceKeep(MANUAL);
101102

103+
assertEquals(
104+
"foo=bar,ot=rv:" + RV + ";th:" + TH + ";future:value,other=state",
105+
propagationTags.getW3CTracestate());
102106
assertEquals(
103107
"dd=s:2;t.dm:-4,ot=rv:" + RV + ";future:value,foo=bar,other=state",
104108
propagationTags.headerValue(W3C));
105109
}
106110

111+
@Test
112+
void initialUnknownPriorityDoesNotOverrideInheritedProbabilityDecision() {
113+
PropagationTags propagationTags =
114+
PropagationTags.factory()
115+
.fromHeaderValue(W3C, "foo=bar,ot=rv:" + RV + ";th:" + TH + ";future:value");
116+
117+
propagationTags.updateTraceSamplingPriority(USER_KEEP, UNKNOWN);
118+
119+
assertEquals(
120+
"dd=s:2,foo=bar,ot=rv:" + RV + ";th:" + TH + ";future:value",
121+
propagationTags.headerValue(W3C));
122+
}
123+
107124
@Test
108125
void manualKeepRemovesLocallyGeneratedRandomness() {
109126
PropagationTags propagationTags = PropagationTags.factory().empty();
@@ -131,6 +148,28 @@ void serializesSamplingStateFromOnePrioritySnapshot() {
131148
assertEquals("dd=s:0,ot=rv:" + RV, propagationTags.headerValue(W3C, null, SAMPLER_DROP));
132149
}
133150

151+
@Test
152+
void effectiveTracestateOnlyUpdatesOtelMember() {
153+
PropagationTags propagationTags =
154+
PropagationTags.factory()
155+
.fromHeaderValue(W3C, "dd=s:1;t.dm:-3,vendor=state,ot=rv:" + RV + ";th:" + TH);
156+
157+
assertEquals(
158+
"dd=s:1;t.dm:-3,ot=rv:" + RV + ",vendor=state",
159+
propagationTags.getW3CTracestate(SAMPLER_DROP));
160+
}
161+
162+
@Test
163+
void effectiveTracestateAddsOnlyLocallyGeneratedOtelMember() {
164+
PropagationTags propagationTags = PropagationTags.factory().empty();
165+
propagationTags.updateTraceSamplingPriority(USER_KEEP, LOCAL_USER_RULE);
166+
propagationTags.updateOtelTraceState(TRACE_ID, SAMPLE_RATE_0_5, true);
167+
168+
assertEquals(
169+
"ot=rv:" + GENERATED_RV + ";th:" + THRESHOLD_0_5,
170+
propagationTags.getW3CTracestate(USER_KEEP));
171+
}
172+
134173
@Test
135174
void updatingW3cTraceStateCapturesOtelStateForMixedHeaderExtraction() {
136175
PropagationTags propagationTags =

‎dd-trace-core/src/test/java/datadog/trace/core/propagation/W3COtelTraceStateContinuationTest.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
import static java.util.Arrays.asList;
1111
import static java.util.Collections.emptyMap;
1212
import static org.junit.jupiter.api.Assertions.assertEquals;
13-
import static org.junit.jupiter.api.Assertions.assertFalse;
1413
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
1514
import static org.junit.jupiter.api.Assertions.assertTrue;
1615
import static org.mockito.Mockito.mock;
@@ -65,12 +64,18 @@ void removesMalformedOtelRandomValueButKeepsThreshold() {
6564
}
6665

6766
@ParameterizedTest
68-
@CsvSource({"0, 01, ef284ace7a91e1, 00", "2, 00, 00000000000000, 01"})
67+
@CsvSource({
68+
"0, 01, ef284ace7a91e1, 00, false",
69+
"2, 00, 00000000000000, 01, false",
70+
"0, 00, 00000000000000, 00, true",
71+
"2, 01, ef284ace7a91e1, 01, true"
72+
})
6973
void alignsOtelStateWithSampledFlagAfterCompoundExtraction(
7074
int datadogSamplingPriority,
7175
String inboundTraceFlags,
7276
String otelRandomValue,
73-
String outboundTraceFlags) {
77+
String outboundTraceFlags,
78+
boolean thresholdExpected) {
7479
String inboundTracestate = DD_MEMBER + ",ot=rv:" + otelRandomValue + ";th:" + OTEL_THRESHOLD;
7580
String traceParent = TRACE_PARENT.substring(0, TRACE_PARENT.length() - 2) + inboundTraceFlags;
7681
Map<String, String> inboundHeaders =
@@ -108,7 +113,7 @@ void alignsOtelStateWithSampledFlagAfterCompoundExtraction(
108113
String outboundTracestate = outboundHeaders.get(TRACE_STATE_KEY);
109114
assertTrue(outboundTracestate.startsWith("dd=s:" + datadogSamplingPriority));
110115
assertTrue(outboundTracestate.contains("ot=rv:" + otelRandomValue));
111-
assertFalse(outboundTracestate.contains("th:" + OTEL_THRESHOLD));
116+
assertEquals(thresholdExpected, outboundTracestate.contains("th:" + OTEL_THRESHOLD));
112117
} finally {
113118
tracer.close();
114119
}

0 commit comments

Comments
 (0)