Skip to content

Websocket: ensure that trace context is propagated from the handshake #8619

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 31, 2025
Merged

Websocket: ensure that trace context is propagated from the handshake #8619

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f67f0a8
Websocket: ensure that trace context is propagated from the handshake
amarziali Mar 26, 2025
ef2a423
Merge branch 'master' into andrea.marziali/websocket-propagate
amarziali Mar 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions ...p/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
88 changes: 73 additions & 15 deletions ...-agent/instrumentation/websocket/javax-websocket-1.0/src/test/groovy/WebsocketTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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()
}
}
}
17 changes: 17 additions & 0 deletions dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -846,4 +847,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);
}
}
}
2 changes: 2 additions & 0 deletions dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,4 +133,6 @@ public HashMap<String, String> createTagMap() {
fillTagMap(result);
return result;
}

public abstract void updateAndLockDecisionMaker(PropagationTags source);
}
18 changes: 17 additions & 1 deletion dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ static class PTags extends PropagationTags {
// tags that don't require any modifications and propagated as-is
private final List<TagElement> tagPairs;

private final boolean canChangeDecisionMaker;
private boolean canChangeDecisionMaker;

// extracted decision maker tag for easier updates
private volatile TagValue decisionMakerTagValue;
Expand All @@ -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}, <code>null</code> if not set.
*/
private volatile TagValue traceIdHighOrderBitsHexTagValue;

/**
* The original <a href="https://www.w3.org/TR/trace-context/#tracestate-header">W3C tracestate
* header</a> value.
*/
protected volatile String tracestate;

/**
* The {@link PTagsFactory#PROPAGATION_ERROR_TAG_KEY propagation tag error} value, {@code null
* if no error while parsing header}.
Expand Down Expand Up @@ -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);
}
}
}
}
}
4 changes: 4 additions & 0 deletions internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,10 @@ default AgentSpan asAgentSpan() {
return this;
}

default void copyPropagationAndBaggage(final AgentSpan source) {
Copy link
Contributor

@datadog-datadog-prod-us1 datadog-datadog-prod-us1 bot Mar 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚪ Code Quality Violation

Suggested change
default void copyPropagationAndBaggage(final AgentSpan source) {
default void copyPropagationAndBaggage(AgentSpan source) {
Avoid declaring as final as it may not be used (...read more)

The rule "Avoid useless final type in interface method" advises against the unnecessary use of the final keyword in the method parameters of an interface. In Java, the final keyword is used to denote that a variable cannot be changed once assigned. However, in the context of an interface method, this is redundant as the value of the parameter cannot be changed within the method anyway.

The importance of this rule lies in the clarity and simplicity of code. Unnecessary use of final in this context can lead to confusion for those reading the code, as it suggests that there may be a specific reason for its use when there is not. It can also clutter the code, making it less readable.

Good coding practices to avoid this rule violation include simply not using the final keyword in the method parameters of an interface. This does not affect the functionality of the code, but it makes it cleaner and easier to understand. For example, instead of writing void process(final Object arg);, you can write void process(Object arg);. This maintains the same functionality but improves the readability of the code.

View in Datadog  Leave us feedback  Documentation

// no op default
}

@Override
default Context storeInto(Context context) {
return context.with(SPAN_KEY, this);
Expand Down