Skip to content

Commit bfdfd2b

Browse files
amarzialiamarziali
authored
Websocket: ensure that trace context is propagated from the handshake (#8619)
1 parent b187a03 commit bfdfd2b

File tree

6 files changed

+114
-16
lines changed

6 files changed

+114
-16
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ private AgentSpan onFrameStart(
138138
if (useDedicatedTraces) {
139139
wsSpan = startSpan(WEBSOCKET.toString(), operationName, null);
140140
if (inheritSampling) {
141+
wsSpan.copyPropagationAndBaggage(handshakeSpan);
141142
wsSpan.setTag(DECISION_MAKER_INHERITED, 1);
142143
wsSpan.setTag(DECISION_MAKER_SERVICE, handshakeSpan.getServiceName());
143144
wsSpan.setTag(DECISION_MAKER_RESOURCE, handshakeSpan.getResourceName());

‎dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/WebsocketTest.groovy‎

Lines changed: 73 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,11 @@
1-
import static datadog.trace.agent.test.base.HttpServerTest.someBytes
2-
import static datadog.trace.agent.test.base.HttpServerTest.websocketCloseSpan
3-
import static datadog.trace.agent.test.base.HttpServerTest.websocketReceiveSpan
4-
import static datadog.trace.agent.test.base.HttpServerTest.websocketSendSpan
5-
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
6-
import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
7-
import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CLASSES_EXCLUDE
8-
import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_ENABLED
9-
import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING
10-
import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES
11-
import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_TAG_SESSION_ID
12-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan
13-
141
import datadog.trace.agent.test.AgentTestRunner
152
import datadog.trace.api.DDTags
3+
import datadog.trace.api.sampling.PrioritySampling
164
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
175
import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags
186
import datadog.trace.bootstrap.instrumentation.api.Tags
197
import datadog.trace.core.DDSpan
8+
import datadog.trace.core.propagation.ExtractedContext
209
import net.bytebuddy.utility.RandomString
2110
import org.glassfish.tyrus.container.inmemory.InMemoryClientContainer
2211
import org.glassfish.tyrus.server.TyrusServerConfiguration
@@ -29,6 +18,19 @@ import javax.websocket.server.ServerApplicationConfig
2918
import javax.websocket.server.ServerEndpointConfig
3019
import java.nio.ByteBuffer
3120

21+
import static datadog.trace.agent.test.base.HttpServerTest.someBytes
22+
import static datadog.trace.agent.test.base.HttpServerTest.websocketCloseSpan
23+
import static datadog.trace.agent.test.base.HttpServerTest.websocketReceiveSpan
24+
import static datadog.trace.agent.test.base.HttpServerTest.websocketSendSpan
25+
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
26+
import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
27+
import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CLASSES_EXCLUDE
28+
import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_ENABLED
29+
import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING
30+
import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES
31+
import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_TAG_SESSION_ID
32+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan
33+
3234
class WebsocketTest extends AgentTestRunner {
3335

3436
@Override
@@ -38,8 +40,8 @@ class WebsocketTest extends AgentTestRunner {
3840
injectSysConfig(TRACE_CLASSES_EXCLUDE, "EndpointWrapper")
3941
}
4042

41-
def createHandshakeSpan(String spanName, String url) {
42-
def span = TEST_TRACER.startSpan("test", spanName, null)
43+
def createHandshakeSpan(String spanName, String url, Object parentContext = null) {
44+
def span = TEST_TRACER.startSpan("test", spanName, parentContext)
4345
handshakeTags(url).each { span.setTag(it.key, it.value) }
4446
span.finish()
4547
span
@@ -574,4 +576,60 @@ class WebsocketTest extends AgentTestRunner {
574576
}
575577
})
576578
}
579+
580+
def "test trace state is inherited"() {
581+
when:
582+
String url = "ws://inmemory/test"
583+
def clientHandshake = createHandshakeSpan("http.request", url) //simulate client span
584+
clientHandshake.setSamplingPriority(PrioritySampling.SAMPLER_DROP) // simulate sampler drop
585+
def serverHandshake = createHandshakeSpan("servlet.request", url,
586+
new ExtractedContext(clientHandshake.context().getTraceId(), clientHandshake.context().getSpanId(), clientHandshake.context().getSamplingPriority(),
587+
"test", 0, ["example_baggage": "test"], null, null, null, null, null)) // simulate server span
588+
def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()),
589+
clientHandshake, serverHandshake, url)
590+
591+
runUnderTrace("parent") {
592+
session.getBasicRemote().sendText("Hello")
593+
session.close()
594+
}
595+
then:
596+
def ht = handshakeTags(url)
597+
assertTraces(5, {
598+
trace(1) {
599+
basicSpan(it, "http.request", "GET /test", null, null, ht)
600+
}
601+
trace(1) {
602+
span {
603+
operationName "servlet.request"
604+
resourceName "GET /test"
605+
childOf(clientHandshake as DDSpan)
606+
tags {
607+
for (def entry : ht) {
608+
tag(entry.key, entry.value)
609+
}
610+
defaultTags(true)
611+
}
612+
}
613+
}
614+
trace(3) {
615+
sortSpansByStart()
616+
basicSpan(it, "parent")
617+
websocketSendSpan(it, clientHandshake as DDSpan, "text", 5, 1, span(0))
618+
websocketCloseSpan(it, clientHandshake as DDSpan, true, 1000, null, span(0))
619+
}
620+
trace(1) {
621+
websocketReceiveSpan(it, serverHandshake as DDSpan, "text", 5, 1)
622+
}
623+
trace(1) {
624+
websocketCloseSpan(it, serverHandshake as DDSpan, false, 1000, { it == null || it == 'no reason given' })
625+
}
626+
})
627+
// check that the handshake trace state is inherited
628+
TEST_WRITER.flatten().findAll { span -> (span as DDSpan).getSpanType() == "websocket" && (span as DDSpan).getParentId() == 0}.each {
629+
assert (it as DDSpan).getSamplingPriority() == serverHandshake.getSamplingPriority()
630+
assert (it as DDSpan).getOrigin() == serverHandshake.context().getOrigin()
631+
assert (it as DDSpan).getBaggage() == serverHandshake.context().getBaggageItems()
632+
assert !(it as DDSpan).getBaggage().isEmpty()
633+
}
634+
}
577635
}

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.core;
22

33
import static datadog.trace.api.DDTags.TRACE_START_TIME;
4+
import static datadog.trace.api.sampling.SamplingMechanism.DEFAULT;
45
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.RECORD_END_TO_END_DURATION_MS;
56
import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_STATUS;
67
import static java.util.concurrent.TimeUnit.MICROSECONDS;
@@ -846,4 +847,20 @@ public boolean isOutbound() {
846847
Object spanKind = context.getTag(Tags.SPAN_KIND);
847848
return Tags.SPAN_KIND_CLIENT.equals(spanKind) || Tags.SPAN_KIND_PRODUCER.equals(spanKind);
848849
}
850+
851+
@Override
852+
public void copyPropagationAndBaggage(final AgentSpan source) {
853+
if (source instanceof DDSpan) {
854+
final DDSpanContext sourceSpanContext = ((DDSpan) source).context();
855+
// align the sampling priority for this span context
856+
setSamplingPriority(sourceSpanContext.getSamplingPriority(), DEFAULT);
857+
// the sampling mechanism determine the dm tag hence we need to override and lock the current
858+
// ptags
859+
context
860+
.getPropagationTags()
861+
.updateAndLockDecisionMaker(sourceSpanContext.getPropagationTags());
862+
context.setOrigin(sourceSpanContext.getOrigin());
863+
sourceSpanContext.getBaggageItems().forEach(context::setBaggageItem);
864+
}
865+
}
849866
}

‎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
@@ -133,4 +133,6 @@ public HashMap<String, String> createTagMap() {
133133
fillTagMap(result);
134134
return result;
135135
}
136+
137+
public abstract void updateAndLockDecisionMaker(PropagationTags source);
136138
}

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ static class PTags extends PropagationTags {
7676
// tags that don't require any modifications and propagated as-is
7777
private final List<TagElement> tagPairs;
7878

79-
private final boolean canChangeDecisionMaker;
79+
private boolean canChangeDecisionMaker;
8080

8181
// extracted decision maker tag for easier updates
8282
private volatile TagValue decisionMakerTagValue;
@@ -90,18 +90,22 @@ static class PTags extends PropagationTags {
9090
private volatile int samplingPriority;
9191
private volatile CharSequence origin;
9292
private volatile String[] headerCache = null;
93+
9394
/** The high-order 64 bits of the trace id. */
9495
private volatile long traceIdHighOrderBits;
96+
9597
/**
9698
* The zero-padded lower-case 16 character hexadecimal representation of the high-order 64 bits
9799
* of the trace id, wrapped into a {@link TagValue}, <code>null</code> if not set.
98100
*/
99101
private volatile TagValue traceIdHighOrderBitsHexTagValue;
102+
100103
/**
101104
* The original <a href="https://www.w3.org/TR/trace-context/#tracestate-header">W3C tracestate
102105
* header</a> value.
103106
*/
104107
protected volatile String tracestate;
108+
105109
/**
106110
* The {@link PTagsFactory#PROPAGATION_ERROR_TAG_KEY propagation tag error} value, {@code null
107111
* if no error while parsing header}.
@@ -404,5 +408,17 @@ public void updateW3CTracestate(String tracestate) {
404408
String getError() {
405409
return this.error;
406410
}
411+
412+
@Override
413+
public void updateAndLockDecisionMaker(PropagationTags source) {
414+
if (source instanceof PTags) {
415+
canChangeDecisionMaker = false;
416+
decisionMakerTagValue = ((PTags) source).getDecisionMakerTagValue();
417+
if (decisionMakerTagValue != null) {
418+
clearCachedHeader(DATADOG);
419+
clearCachedHeader(W3C);
420+
}
421+
}
422+
}
407423
}
408424
}

‎internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,10 @@ default AgentSpan asAgentSpan() {
192192
return this;
193193
}
194194

195+
default void copyPropagationAndBaggage(final AgentSpan source) {
196+
// no op default
197+
}
198+
195199
@Override
196200
default Context storeInto(Context context) {
197201
return context.with(SPAN_KEY, this);

0 commit comments

Comments
 (0)