From f67f0a8f661d63b241111da712999958d5372b1c Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 26 Mar 2025 15:58:56 +0100 Subject: [PATCH] Websocket: ensure that trace context is propagated from the handshake --- .../decorator/WebsocketDecorator.java | 1 + .../src/test/groovy/WebsocketTest.groovy | 88 +++++++++++++++---- .../main/java/datadog/trace/core/DDSpan.java | 17 ++++ .../core/propagation/PropagationTags.java | 2 + .../core/propagation/ptags/PTagsFactory.java | 18 +++- .../instrumentation/api/AgentSpan.java | 4 + 6 files changed, 114 insertions(+), 16 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java index cdc758829d9..2123caa395d 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java @@ -138,6 +138,7 @@ private AgentSpan onFrameStart( if (useDedicatedTraces) { wsSpan = startSpan(WEBSOCKET.toString(), operationName, null); if (inheritSampling) { + wsSpan.copyPropagationAndBaggage(handshakeSpan); wsSpan.setTag(DECISION_MAKER_INHERITED, 1); wsSpan.setTag(DECISION_MAKER_SERVICE, handshakeSpan.getServiceName()); wsSpan.setTag(DECISION_MAKER_RESOURCE, handshakeSpan.getResourceName()); diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/WebsocketTest.groovy b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/WebsocketTest.groovy index 948ba636ce2..eff2cb53483 100644 --- a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/WebsocketTest.groovy +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/WebsocketTest.groovy @@ -1,22 +1,11 @@ -import static datadog.trace.agent.test.base.HttpServerTest.someBytes -import static datadog.trace.agent.test.base.HttpServerTest.websocketCloseSpan -import static datadog.trace.agent.test.base.HttpServerTest.websocketReceiveSpan -import static datadog.trace.agent.test.base.HttpServerTest.websocketSendSpan -import static datadog.trace.agent.test.utils.TraceUtils.basicSpan -import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace -import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CLASSES_EXCLUDE -import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_ENABLED -import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING -import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES -import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_TAG_SESSION_ID -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan - import datadog.trace.agent.test.AgentTestRunner import datadog.trace.api.DDTags +import datadog.trace.api.sampling.PrioritySampling import datadog.trace.bootstrap.instrumentation.api.AgentSpan import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags import datadog.trace.bootstrap.instrumentation.api.Tags import datadog.trace.core.DDSpan +import datadog.trace.core.propagation.ExtractedContext import net.bytebuddy.utility.RandomString import org.glassfish.tyrus.container.inmemory.InMemoryClientContainer import org.glassfish.tyrus.server.TyrusServerConfiguration @@ -29,6 +18,19 @@ import javax.websocket.server.ServerApplicationConfig import javax.websocket.server.ServerEndpointConfig import java.nio.ByteBuffer +import static datadog.trace.agent.test.base.HttpServerTest.someBytes +import static datadog.trace.agent.test.base.HttpServerTest.websocketCloseSpan +import static datadog.trace.agent.test.base.HttpServerTest.websocketReceiveSpan +import static datadog.trace.agent.test.base.HttpServerTest.websocketSendSpan +import static datadog.trace.agent.test.utils.TraceUtils.basicSpan +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CLASSES_EXCLUDE +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_ENABLED +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_INHERIT_SAMPLING +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_TAG_SESSION_ID +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan + class WebsocketTest extends AgentTestRunner { @Override @@ -38,8 +40,8 @@ class WebsocketTest extends AgentTestRunner { injectSysConfig(TRACE_CLASSES_EXCLUDE, "EndpointWrapper") } - def createHandshakeSpan(String spanName, String url) { - def span = TEST_TRACER.startSpan("test", spanName, null) + def createHandshakeSpan(String spanName, String url, Object parentContext = null) { + def span = TEST_TRACER.startSpan("test", spanName, parentContext) handshakeTags(url).each { span.setTag(it.key, it.value) } span.finish() span @@ -574,4 +576,60 @@ class WebsocketTest extends AgentTestRunner { } }) } + + def "test trace state is inherited"() { + when: + String url = "ws://inmemory/test" + def clientHandshake = createHandshakeSpan("http.request", url) //simulate client span + clientHandshake.setSamplingPriority(PrioritySampling.SAMPLER_DROP) // simulate sampler drop + def serverHandshake = createHandshakeSpan("servlet.request", url, + new ExtractedContext(clientHandshake.context().getTraceId(), clientHandshake.context().getSpanId(), clientHandshake.context().getSamplingPriority(), + "test", 0, ["example_baggage": "test"], null, null, null, null, null)) // simulate server span + def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()), + clientHandshake, serverHandshake, url) + + runUnderTrace("parent") { + session.getBasicRemote().sendText("Hello") + session.close() + } + then: + def ht = handshakeTags(url) + assertTraces(5, { + trace(1) { + basicSpan(it, "http.request", "GET /test", null, null, ht) + } + trace(1) { + span { + operationName "servlet.request" + resourceName "GET /test" + childOf(clientHandshake as DDSpan) + tags { + for (def entry : ht) { + tag(entry.key, entry.value) + } + defaultTags(true) + } + } + } + trace(3) { + sortSpansByStart() + basicSpan(it, "parent") + websocketSendSpan(it, clientHandshake as DDSpan, "text", 5, 1, span(0)) + websocketCloseSpan(it, clientHandshake as DDSpan, true, 1000, null, span(0)) + } + trace(1) { + websocketReceiveSpan(it, serverHandshake as DDSpan, "text", 5, 1) + } + trace(1) { + websocketCloseSpan(it, serverHandshake as DDSpan, false, 1000, { it == null || it == 'no reason given' }) + } + }) + // check that the handshake trace state is inherited + TEST_WRITER.flatten().findAll { span -> (span as DDSpan).getSpanType() == "websocket" && (span as DDSpan).getParentId() == 0}.each { + assert (it as DDSpan).getSamplingPriority() == serverHandshake.getSamplingPriority() + assert (it as DDSpan).getOrigin() == serverHandshake.context().getOrigin() + assert (it as DDSpan).getBaggage() == serverHandshake.context().getBaggageItems() + assert !(it as DDSpan).getBaggage().isEmpty() + } + } } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java index 771d22e3758..725af1a4750 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java @@ -1,6 +1,7 @@ package datadog.trace.core; import static datadog.trace.api.DDTags.TRACE_START_TIME; +import static datadog.trace.api.sampling.SamplingMechanism.DEFAULT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.RECORD_END_TO_END_DURATION_MS; import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_STATUS; import static java.util.concurrent.TimeUnit.MICROSECONDS; @@ -856,4 +857,20 @@ public boolean isOutbound() { Object spanKind = context.getTag(Tags.SPAN_KIND); return Tags.SPAN_KIND_CLIENT.equals(spanKind) || Tags.SPAN_KIND_PRODUCER.equals(spanKind); } + + @Override + public void copyPropagationAndBaggage(final AgentSpan source) { + if (source instanceof DDSpan) { + final DDSpanContext sourceSpanContext = ((DDSpan) source).context(); + // align the sampling priority for this span context + setSamplingPriority(sourceSpanContext.getSamplingPriority(), DEFAULT); + // the sampling mechanism determine the dm tag hence we need to override and lock the current + // ptags + context + .getPropagationTags() + .updateAndLockDecisionMaker(sourceSpanContext.getPropagationTags()); + context.setOrigin(sourceSpanContext.getOrigin()); + sourceSpanContext.getBaggageItems().forEach(context::setBaggageItem); + } + } } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java index 7d8ad25a44d..a510ff54dfd 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java @@ -133,4 +133,6 @@ public HashMap createTagMap() { fillTagMap(result); return result; } + + public abstract void updateAndLockDecisionMaker(PropagationTags source); } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java index 942fa5b2821..1cfb52fbf1b 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java @@ -76,7 +76,7 @@ static class PTags extends PropagationTags { // tags that don't require any modifications and propagated as-is private final List tagPairs; - private final boolean canChangeDecisionMaker; + private boolean canChangeDecisionMaker; // extracted decision maker tag for easier updates private volatile TagValue decisionMakerTagValue; @@ -90,18 +90,22 @@ static class PTags extends PropagationTags { private volatile int samplingPriority; private volatile CharSequence origin; private volatile String[] headerCache = null; + /** The high-order 64 bits of the trace id. */ private volatile long traceIdHighOrderBits; + /** * The zero-padded lower-case 16 character hexadecimal representation of the high-order 64 bits * of the trace id, wrapped into a {@link TagValue}, null if not set. */ private volatile TagValue traceIdHighOrderBitsHexTagValue; + /** * The original W3C tracestate * header value. */ protected volatile String tracestate; + /** * The {@link PTagsFactory#PROPAGATION_ERROR_TAG_KEY propagation tag error} value, {@code null * if no error while parsing header}. @@ -404,5 +408,17 @@ public void updateW3CTracestate(String tracestate) { String getError() { return this.error; } + + @Override + public void updateAndLockDecisionMaker(PropagationTags source) { + if (source instanceof PTags) { + canChangeDecisionMaker = false; + decisionMakerTagValue = ((PTags) source).getDecisionMakerTagValue(); + if (decisionMakerTagValue != null) { + clearCachedHeader(DATADOG); + clearCachedHeader(W3C); + } + } + } } } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java index 412ac7c85cc..e9b21f797b9 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java @@ -192,6 +192,10 @@ default AgentSpan asAgentSpan() { return this; } + default void copyPropagationAndBaggage(final AgentSpan source) { + // no op default + } + @Override default Context storeInto(Context context) { return context.with(SPAN_KEY, this);