From dc8feae4690995bcd44f7f6fdde505a7f6614faa Mon Sep 17 00:00:00 2001 From: liurui Date: Thu, 3 Sep 2026 15:14:40 +0800 Subject: [PATCH] Honor single-trace WebSocket configuration --- .../decorator/WebsocketDecorator.java | 13 ++++++------- .../agent/test/base/HttpServerTest.groovy | 4 +++- .../src/test/groovy/WebsocketTest.groovy | 18 ++++++++++-------- .../src/test/groovy/WebsocketTest.groovy | 18 ++++++++++-------- 4 files changed, 29 insertions(+), 24 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 df2f4b6ddac..35c308d11d8 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 @@ -151,8 +151,8 @@ private AgentSpan onFrameStart( final AgentSpan handshakeSpan = handlerContext.getHandshakeSpan(); boolean inheritSampling = config.isWebsocketMessagesInheritSampling(); boolean useDedicatedTraces = config.isWebsocketMessagesSeparateTraces(); - if (traceStarter) { - if (useDedicatedTraces) { + if (useDedicatedTraces) { + if (traceStarter) { wsSpan = startSpan(WEBSOCKET.toString(), operationName, null); if (inheritSampling) { wsSpan.copyPropagationAndBaggage(handshakeSpan); @@ -161,10 +161,10 @@ private AgentSpan onFrameStart( wsSpan.setTag(DECISION_MAKER_RESOURCE, handshakeSpan.getResourceName()); } } else { - wsSpan = startSpan(WEBSOCKET.toString(), operationName, handshakeSpan.spanContext()); + wsSpan = startSpan(WEBSOCKET.toString(), operationName); } } else { - wsSpan = startSpan(WEBSOCKET.toString(), operationName); + wsSpan = startSpan(WEBSOCKET.toString(), operationName, handshakeSpan.spanContext()); } handlerContext.setWebsocketSpan(wsSpan); afterStart(wsSpan); @@ -178,9 +178,8 @@ private AgentSpan onFrameStart( if (config.isWebsocketTagSessionId()) { wsSpan.setTag(WEBSOCKET_SESSION_ID, handlerContext.getSessionId()); } - if (useDedicatedTraces || !traceStarter) { - // the link is not added if the user wants to have receive frames on the same trace as the - // handshake + if (useDedicatedTraces) { + // The link is not added if the user wants all frames on the same trace as the handshake. wsSpan.addLink( SpanLink.from( inheritSampling diff --git a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy index 46abd76a6da..0a4ca4e292d 100644 --- a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy +++ b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy @@ -2462,6 +2462,8 @@ abstract class HttpServerTest extends WithHttpServer { } if (traceStarter && Config.get().isWebsocketMessagesSeparateTraces()) { parent() + } else if (!Config.get().isWebsocketMessagesSeparateTraces()) { + childOf(handshake) } else { if (parentSpan != null) { childOf(parentSpan) @@ -2470,7 +2472,7 @@ abstract class HttpServerTest extends WithHttpServer { } } spanType(DDSpanTypes.WEBSOCKET) - if (Config.get().isWebsocketMessagesSeparateTraces() || !traceStarter) { + if (Config.get().isWebsocketMessagesSeparateTraces()) { links { link(handshake, linkFlags, linkAttributes) } diff --git a/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/WebsocketTest.groovy b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/WebsocketTest.groovy index 2fc09db5002..e4f9bc11db0 100644 --- a/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/WebsocketTest.groovy +++ b/dd-java-agent/instrumentation/websocket/jakarta-websocket-2.0/src/test/groovy/WebsocketTest.groovy @@ -431,7 +431,7 @@ class WebsocketTest extends InstrumentationSpecification { } } - def "test close and receive on same handshake trace"() { + def "test all messages and close on same handshake trace"() { setup: injectSysConfig(TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES, "false") when: @@ -446,14 +446,19 @@ class WebsocketTest extends InstrumentationSpecification { session.close() } then: - // in reality we have 3 traces but since the handshake finishes soon, the trace structure writer is collecting 5 chunks - assertTraces(5, { + // In reality we have 3 traces, but finished handshake traces are reported in separate chunks. + assertTraces(7, { DDSpan serverHandshake, clientHandshake trace(1) { basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) clientHandshake = span(0) } - + trace(1) { + websocketSendSpan(it, clientHandshake, "text", 5, 1, clientHandshake) + } + trace(1) { + websocketCloseSpan(it, clientHandshake, true, 1000, null, clientHandshake) + } trace(1) { basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) serverHandshake = span(0) @@ -464,11 +469,8 @@ class WebsocketTest extends InstrumentationSpecification { trace(1) { websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }, serverHandshake) } - trace(3) { - sortSpansByStart() + trace(1) { basicSpan(it, "parent") - websocketSendSpan(it, clientHandshake, "text", 5, 1, span(0)) - websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) } }) } 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 f0b1170737e..688652f5b09 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 @@ -434,7 +434,7 @@ class WebsocketTest extends InstrumentationSpecification { } } - def "test close and receive on same handshake trace"() { + def "test all messages and close on same handshake trace"() { setup: injectSysConfig(TRACE_WEBSOCKET_MESSAGES_SEPARATE_TRACES, "false") when: @@ -449,14 +449,19 @@ class WebsocketTest extends InstrumentationSpecification { session.close() } then: - // in reality we have 3 traces but since the handshake finishes soon, the trace structure writer is collecting 5 chunks - assertTraces(5, { + // In reality we have 3 traces, but finished handshake traces are reported in separate chunks. + assertTraces(7, { DDSpan serverHandshake, clientHandshake trace(1) { basicSpan(it, "http.request", "GET /test", null, null, handshakeTags(url)) clientHandshake = span(0) } - + trace(1) { + websocketSendSpan(it, clientHandshake, "text", 5, 1, clientHandshake) + } + trace(1) { + websocketCloseSpan(it, clientHandshake, true, 1000, null, clientHandshake) + } trace(1) { basicSpan(it, "servlet.request", "GET /test", null, null, handshakeTags(url)) serverHandshake = span(0) @@ -467,11 +472,8 @@ class WebsocketTest extends InstrumentationSpecification { trace(1) { websocketCloseSpan(it, serverHandshake, false, 1000, { it == null || it == 'no reason given' }, serverHandshake) } - trace(3) { - sortSpansByStart() + trace(1) { basicSpan(it, "parent") - websocketSendSpan(it, clientHandshake, "text", 5, 1, span(0)) - websocketCloseSpan(it, clientHandshake, true, 1000, null, span(0)) } }) }