From 2635b84b7ebec5a548fb6e92454b378038394740 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Mon, 28 Sep 2026 16:32:39 -0700 Subject: [PATCH 01/23] Add Jetty 12 native WebSocket tracing and fix upgrade span completion --- .../JettyWebSocketUpgradeInstrumentation.java | 45 ++ .../WebSocketUpgradeAdvice.java | 23 + .../jetty-websocket-12.0/build.gradle | 21 + .../jetty-websocket-12.0/gradle.lockfile | 38 +- .../jetty12/Jetty12NativeWebSocketModule.java | 67 +++ .../jetty12/NativeMethodHandleWrappers.java | 306 ++++++++++ .../jetty12/NativeWebSocketAdvices.java | 59 ++ .../src/test/groovy/JettyEndpoints.groovy | 156 ++++++ .../groovy/JettyWebSocketUpgradeTest.groovy | 60 ++ .../src/test/groovy/JettyWebsocketTest.groovy | 525 ++++++++++++++++++ 10 files changed, 1285 insertions(+), 15 deletions(-) create mode 100644 dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java create mode 100644 dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/WebSocketUpgradeAdvice.java create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeWebSocketAdvices.java create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyEndpoints.groovy create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebSocketUpgradeTest.groovy create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebsocketTest.groovy diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java new file mode 100644 index 00000000000..03eb291ed66 --- /dev/null +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java @@ -0,0 +1,45 @@ +package datadog.trace.instrumentation.jetty_client12; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static java.util.Collections.singletonMap; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.util.Map; + +@AutoService(InstrumenterModule.class) +public class JettyWebSocketUpgradeInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + public JettyWebSocketUpgradeInstrumentation() { + super("jetty-client"); + } + + @Override + public String instrumentedType() { + return "org.eclipse.jetty.websocket.core.client.CoreClientUpgradeRequest"; + } + + @Override + public String[] helperClassNames() { + return new String[] {packageName + ".JettyClientDecorator"}; + } + + @Override + public Map contextStore() { + return singletonMap("org.eclipse.jetty.client.Request", AgentSpan.class.getName()); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("upgrade") + .and(takesArguments(2)) + .and(takesArgument(0, named("org.eclipse.jetty.client.Response"))) + .and(takesArgument(1, named("org.eclipse.jetty.io.EndPoint"))), + packageName + ".WebSocketUpgradeAdvice"); + } +} diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/WebSocketUpgradeAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/WebSocketUpgradeAdvice.java new file mode 100644 index 00000000000..d051ae0f708 --- /dev/null +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/WebSocketUpgradeAdvice.java @@ -0,0 +1,23 @@ +package datadog.trace.instrumentation.jetty_client12; + +import static datadog.trace.instrumentation.jetty_client12.JettyClientDecorator.DECORATE; + +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import net.bytebuddy.asm.Advice; +import org.eclipse.jetty.client.Request; +import org.eclipse.jetty.client.Response; + +public class WebSocketUpgradeAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void afterUpgrade(@Advice.Argument(0) Response response) { + AgentSpan span = + InstrumentationContext.get(Request.class, AgentSpan.class).get(response.getRequest()); + if (span != null) { + // Successful upgrades bypass the request's response completion listeners. + DECORATE.onResponse(span, response); + DECORATE.beforeFinish(span); + span.finish(); + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle index 295e18d7d71..d7613f481b2 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle @@ -3,6 +3,13 @@ plugins { } muzzle { + pass { + name = 'jetty-websocket-12-native' + group = 'org.eclipse.jetty.websocket' + module = 'jetty-websocket-jetty-server' + versions = "[12,12.0.17]" + javaVersion = "17" + } pass { name = 'jetty-websocket-12ee8' group = 'org.eclipse.jetty.ee8.websocket' @@ -30,6 +37,14 @@ testJvmConstraints { minJavaVersion = JavaVersion.VERSION_17 } +tracerJava { + addSourceSetFor(JavaVersion.VERSION_17) +} + +tasks.named("compileMain_java17Java", JavaCompile) { + configureCompiler(it, 17, JavaVersion.VERSION_1_8, "Uses Jetty 12 APIs (Java 17+) at Java 8 bytecode") +} + addTestSuiteForDir("latestDepTest", "test") ["compileTestGroovy", "compileLatestDepTestGroovy"].each { name -> @@ -39,16 +54,22 @@ addTestSuiteForDir("latestDepTest", "test") } dependencies { + main_java17CompileOnly 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0' implementation project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-10.0") testImplementation group: 'org.eclipse.jetty.ee8.websocket', name: 'jetty-ee8-websocket-javax-server', version: '12.0.0' testImplementation group: 'org.eclipse.jetty.ee9.websocket', name: 'jetty-ee9-websocket-jakarta-server', version: '12.0.0' testImplementation group: 'org.eclipse.jetty.ee10.websocket', name: 'jetty-ee10-websocket-jakarta-server', version: '12.0.0' + testImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0' + testImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.0' //TODO: jetty-12.1.0 is still alpha but wraps MethodHandle class into a MethodHolder class. // Today that is not stable but we'll need to port those advices to support that once the code base will be a bit more stable latestDepTestImplementation group: 'org.eclipse.jetty.ee8.websocket', name: 'jetty-ee8-websocket-javax-server', version: '12.0.17' latestDepTestImplementation group: 'org.eclipse.jetty.ee9.websocket', name: 'jetty-ee9-websocket-jakarta-server', version: '12.0.17' latestDepTestImplementation group: 'org.eclipse.jetty.ee10.websocket', name: 'jetty-ee10-websocket-jakarta-server', version: '12.0.17' + latestDepTestImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.17' + latestDepTestImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.17' + testRuntimeOnly project(":dd-java-agent:instrumentation:jetty:jetty-client:jetty-client-12.0") testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-10.0") testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-11.0") testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:javax-websocket-1.0") diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile index bf30846eadd..d56f5eeeec0 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile @@ -6,11 +6,11 @@ cafe.cryptography:curve25519-elisabeth:0.1.0=latestDepTestRuntimeClasspath,testR cafe.cryptography:ed25519-elisabeth:0.1.0=latestDepTestRuntimeClasspath,testRuntimeClasspath ch.qos.logback:logback-classic:1.2.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath ch.qos.logback:logback-core:1.2.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.blogspot.mydailyjava:weak-lock-free:0.17=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +com.blogspot.mydailyjava:weak-lock-free:0.17=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath com.datadoghq.okhttp3:okhttp:3.12.15=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath com.datadoghq.okio:okio:1.17.6=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.datadoghq:dd-instrument-java:0.0.5=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath -com.datadoghq:dd-javac-plugin-client:0.2.2=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +com.datadoghq:dd-instrument-java:0.0.5=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +com.datadoghq:dd-javac-plugin-client:0.2.2=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath com.datadoghq:java-dogstatsd-client:4.4.5=latestDepTestRuntimeClasspath,testRuntimeClasspath com.datadoghq:sketches-java:0.8.3=latestDepTestRuntimeClasspath,testRuntimeClasspath com.github.javaparser:javaparser-core:3.25.6=codenarc @@ -68,8 +68,8 @@ jakarta.websocket:jakarta.websocket-client-api:2.1.1=latestDepTestCompileClasspa javax.servlet:javax.servlet-api:3.1.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath jaxen:jaxen:2.0.6=spotbugs junit:junit:4.13.2=latestDepTestRuntimeClasspath,testRuntimeClasspath -net.bytebuddy:byte-buddy-agent:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath -net.bytebuddy:byte-buddy:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +net.bytebuddy:byte-buddy-agent:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +net.bytebuddy:byte-buddy:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath net.java.dev.jna:jna-platform:5.8.0=latestDepTestRuntimeClasspath,testRuntimeClasspath net.java.dev.jna:jna:5.8.0=latestDepTestRuntimeClasspath,testRuntimeClasspath net.sf.saxon:Saxon-HE:12.10=spotbugs @@ -155,29 +155,37 @@ org.eclipse.jetty.toolchain:jetty-javax-websocket-api:1.1.2=latestDepTestCompile org.eclipse.jetty.toolchain:jetty-servlet-api:4.0.6=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-client:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-api:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-api:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-alpn-client:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-alpn-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-client:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-ee:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty:jetty-http:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty:jetty-http:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-http:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty:jetty-io:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty:jetty-io:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-io:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-jndi:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-jndi:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-plus:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-security:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-security:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty:jetty-server:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty:jetty-server:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-session:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-session:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty:jetty-util:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty:jetty-util:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-util:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-xml:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-xml:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath @@ -211,16 +219,16 @@ org.ow2.asm:asm-tree:9.5=testCompileClasspath org.ow2.asm:asm-tree:9.7.1=latestDepTestCompileClasspath org.ow2.asm:asm-util:9.10.1=spotbugs org.ow2.asm:asm-util:9.7.1=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.ow2.asm:asm:9.10.1=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,spotbugs,testCompileClasspath,testRuntimeClasspath +org.ow2.asm:asm:9.10.1=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleTooling,runtimeClasspath,spotbugs,testCompileClasspath,testRuntimeClasspath org.slf4j:jcl-over-slf4j:1.7.30=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.slf4j:jul-to-slf4j:1.7.30=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.slf4j:log4j-over-slf4j:1.7.30=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.slf4j:slf4j-api:1.7.30=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,spotbugs,spotbugsSlf4j,testCompileClasspath,testRuntimeClasspath +org.slf4j:slf4j-api:1.7.30=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,spotbugs,spotbugsSlf4j,testCompileClasspath,testRuntimeClasspath org.slf4j:slf4j-simple:2.0.17=spotbugsSlf4j -org.snakeyaml:snakeyaml-engine:2.9=buildTimeInstrumentationPlugin,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testRuntimeClasspath +org.snakeyaml:snakeyaml-engine:2.9=buildTimeInstrumentationPlugin,latestDepTestRuntimeClasspath,main_java17RuntimeClasspath,muzzleTooling,runtimeClasspath,testRuntimeClasspath org.spockframework:spock-bom:2.4-groovy-3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.spockframework:spock-core:2.4-groovy-3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.tabletest:tabletest-junit:1.2.2=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.tabletest:tabletest-parser:1.2.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.xmlresolver:xmlresolver:5.3.3=spotbugs -empty=spotbugsPlugins +empty=main_java17AnnotationProcessor,spotbugsPlugins diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java new file mode 100644 index 00000000000..0376df38a67 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java @@ -0,0 +1,67 @@ +package datadog.trace.instrumentation.websocket.jetty12; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static java.util.Collections.singletonMap; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; +import java.util.Map; + +@AutoService(InstrumenterModule.class) +public class Jetty12NativeWebSocketModule extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + public Jetty12NativeWebSocketModule() { + super("jetty", "jetty-websocket", "websocket"); + } + + @Override + protected boolean defaultEnabled() { + return InstrumenterConfig.get().isWebsocketTracingEnabled(); + } + + @Override + public String instrumentedType() { + return "org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler"; + } + + @Override + public String muzzleDirective() { + return "jetty-websocket-12-native"; + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".NativeMethodHandleWrappers", + packageName + ".NativeMethodHandleWrappers$ReceiveContexts", + packageName + ".NativeMethodHandleWrappers$BinaryMessage", + packageName + ".NativeMethodHandleWrappers$ReceiveCallback" + }; + } + + @Override + public Map contextStore() { + return singletonMap( + "org.eclipse.jetty.websocket.core.CoreSession", + packageName + ".NativeMethodHandleWrappers$ReceiveContexts"); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("onOpen") + .and(takesArguments(2)) + .and(takesArgument(0, named("org.eclipse.jetty.websocket.core.CoreSession"))), + packageName + ".NativeWebSocketAdvices$OpenAdvice"); + transformer.applyAdvice( + named("createMessageSink").and(takesArguments(4)), + packageName + ".NativeWebSocketAdvices$MessageSinkAdvice"); + transformer.applyAdvice( + named("notifyOnClose").and(takesArguments(2)), + packageName + ".NativeWebSocketAdvices$CloseAdvice"); + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java new file mode 100644 index 00000000000..dbb1c58f118 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java @@ -0,0 +1,306 @@ +package datadog.trace.instrumentation.websocket.jetty12; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; +import static java.lang.invoke.MethodHandles.dropArguments; +import static java.lang.invoke.MethodHandles.insertArguments; + +import datadog.context.ContextScope; +import datadog.trace.api.Config; +import datadog.trace.bootstrap.ContextStore; +import datadog.trace.bootstrap.ExceptionLogger; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import datadog.trace.util.MethodHandles; +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodType; +import java.nio.ByteBuffer; +import java.util.HashSet; +import java.util.Set; +import org.eclipse.jetty.websocket.api.Callback; +import org.eclipse.jetty.websocket.api.Session; +import org.eclipse.jetty.websocket.core.CoreSession; + +/** + * Wraps Jetty's endpoint handles after argument normalization, preserving their invocation types. + */ +public class NativeMethodHandleWrappers { + private static final MethodHandles LOOKUP = + new MethodHandles(NativeMethodHandleWrappers.class.getClassLoader()); + private static final MethodHandle TEXT = + LOOKUP.method( + NativeMethodHandleWrappers.class, + "onText", + MethodHandle.class, + HandlerContext.Receiver.class, + boolean.class, + String.class, + boolean.class); + private static final MethodHandle BINARY = + LOOKUP.method( + NativeMethodHandleWrappers.class, + "onBinary", + MethodHandle.class, + ReceiveContexts.class, + boolean.class, + ByteBuffer.class, + boolean.class, + Callback.class); + private static final MethodHandle CLOSE = + LOOKUP.method( + NativeMethodHandleWrappers.class, + "onClose", + MethodHandle.class, + HandlerContext.Receiver.class, + Session.class, + int.class, + String.class); + + private static HandlerContext.Receiver context(AgentSpan span, CoreSession session) { + if (Config.get().isWebsocketMessagesInheritSampling()) { + span.forceSamplingDecision(); + } + return new HandlerContext.Receiver( + span.getLocalRootSpan(), Integer.toHexString(System.identityHashCode(session))); + } + + public static MethodHandle wrapMessage( + MethodHandle delegate, + AgentSpan span, + CoreSession session, + ContextStore contextStore) { + MethodType type = delegate.type(); + Class payload = type.parameterType(0); + if (payload != String.class && payload != ByteBuffer.class) { + return delegate; + } + boolean partial = type.parameterCount() > 1 && type.parameterType(1) == boolean.class; + MethodHandle normalized = partial ? delegate : dropArguments(delegate, 1, boolean.class); + ReceiveContexts contexts = contextStore.get(session); + if (contexts == null) { + contexts = contextStore.getOrPut(session, new ReceiveContexts()); + } + HandlerContext.Receiver context = context(span, session); + contexts.add(context, payload == String.class); + MethodHandle wrapper = + payload == String.class + ? insertArguments(TEXT, 0, normalized, context, partial) + : insertArguments(BINARY, 0, normalized, contexts, partial); + return partial ? wrapper : insertArguments(wrapper, 1, true); + } + + public static MethodHandle wrapClose(MethodHandle delegate, AgentSpan span, CoreSession session) { + // Annotated endpoints have a Session argument; listeners do not. Jetty binds it during onOpen. + if (delegate.type().parameterType(0) != Session.class) { + delegate = dropArguments(delegate, 0, Session.class); + } + return insertArguments(CLOSE, 0, delegate, context(span, session)); + } + + private static ContextScope startMessage( + HandlerContext.Receiver context, Object data, boolean partial) { + try { + synchronized (context) { + return activateSpan(DECORATE.startInboundFrameSpan(context, data, partial)); + } + } catch (Throwable t) { + ExceptionLogger.LOGGER.debug("Unable to start native Jetty WebSocket span", t); + return null; + } + } + + public static void onText( + MethodHandle delegate, + HandlerContext.Receiver context, + boolean partial, + String payload, + boolean last) + throws Throwable { + boolean finish = last; + try (ContextScope ignored = startMessage(context, payload, partial)) { + try { + delegate.invokeExact(payload, last); + } catch (Throwable t) { + finish = true; + synchronized (context) { + DECORATE.onError(context.getWebsocketSpan(), t); + } + throw t; + } + } finally { + if (finish) { + synchronized (context) { + DECORATE.onFrameEnd(context); + } + } + } + } + + public static void onBinary( + MethodHandle delegate, + ReceiveContexts contexts, + boolean partial, + ByteBuffer payload, + boolean last, + Callback callback) + throws Throwable { + ReceiveCallback wrapped; + ContextScope scope; + synchronized (contexts) { + BinaryMessage message = contexts.startBinaryMessage(last); + scope = startMessage(message, payload, partial); + wrapped = new ReceiveCallback(callback, contexts, message); + } + try (ContextScope ignored = scope) { + try { + delegate.invokeExact(payload, last, (Callback) wrapped); + } catch (Throwable t) { + wrapped.onFailure(t); + throw t; + } + } finally { + wrapped.onHandlerExit(); + } + } + + public static class ReceiveContexts { + private HandlerContext.Receiver text; + private HandlerContext.Receiver binary; + private BinaryMessage currentBinary; + private final Set pendingBinary = new HashSet<>(); + + public void add(HandlerContext.Receiver context, boolean isText) { + if (isText) { + text = context; + } else { + binary = context; + } + } + + public synchronized BinaryMessage startBinaryMessage(boolean last) { + BinaryMessage message = currentBinary; + if (message == null) { + message = new BinaryMessage(binary); + pendingBinary.add(message); + } + message.pendingCallbacks++; + message.complete = last; + // Jetty can deliver the next message before this message's callbacks complete. + currentBinary = last ? null : message; + return message; + } + + public synchronized void finish() { + finish(text); + for (BinaryMessage message : pendingBinary) { + finish(message); + } + pendingBinary.clear(); + currentBinary = null; + } + + private static void finish(HandlerContext.Receiver context) { + if (context != null) { + synchronized (context) { + DECORATE.onFrameEnd(context); + } + } + } + } + + public static class BinaryMessage extends HandlerContext.Receiver { + private int pendingCallbacks; + private boolean complete; + + public BinaryMessage(HandlerContext.Receiver context) { + super(context.getHandshakeSpan(), context.getSessionId()); + } + } + + public static class ReceiveCallback implements Callback { + private final Callback delegate; + private final ReceiveContexts contexts; + private final BinaryMessage message; + private boolean completed; + private boolean handlerExited; + private boolean released; + + public ReceiveCallback(Callback delegate, ReceiveContexts contexts, BinaryMessage message) { + this.delegate = delegate; + this.contexts = contexts; + this.message = message; + } + + @Override + public void succeed() { + synchronized (contexts) { + completed = true; + finish(); + } + delegate.succeed(); + } + + @Override + public void fail(Throwable failure) { + onFailure(failure); + delegate.fail(failure); + } + + public void onFailure(Throwable failure) { + synchronized (contexts) { + completed = true; + message.complete = true; + if (contexts.currentBinary == message) { + contexts.currentBinary = null; + } + if (message.getWebsocketSpan() != null) { + DECORATE.onError(message.getWebsocketSpan(), failure); + } + finish(); + } + } + + public void onHandlerExit() { + synchronized (contexts) { + handlerExited = true; + finish(); + } + } + + private void finish() { + // A synchronous callback may complete before the handler throws or closes its scope. + if (handlerExited && completed && !released) { + released = true; + if (--message.pendingCallbacks == 0 && message.complete) { + DECORATE.onFrameEnd(message); + contexts.pendingBinary.remove(message); + } + } + } + } + + public static void onClose( + MethodHandle delegate, + HandlerContext.Receiver context, + Session session, + int code, + String reason) + throws Throwable { + ContextScope scope = null; + try { + scope = activateSpan(DECORATE.startInboundCloseSpan(context, reason, code)); + } catch (Throwable t) { + ExceptionLogger.LOGGER.debug("Unable to start native Jetty WebSocket close span", t); + } + try (ContextScope ignored = scope) { + try { + delegate.invokeExact(session, code, reason); + } catch (Throwable t) { + DECORATE.onError(context.getWebsocketSpan(), t); + throw t; + } + } finally { + DECORATE.onFrameEnd(context); + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeWebSocketAdvices.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeWebSocketAdvices.java new file mode 100644 index 00000000000..cc921cb1645 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeWebSocketAdvices.java @@ -0,0 +1,59 @@ +package datadog.trace.instrumentation.websocket.jetty12; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; + +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.instrumentation.websocket.jetty12.NativeMethodHandleWrappers.ReceiveContexts; +import java.lang.invoke.MethodHandle; +import net.bytebuddy.asm.Advice; +import org.eclipse.jetty.websocket.common.WebSocketSession; +import org.eclipse.jetty.websocket.core.Behavior; +import org.eclipse.jetty.websocket.core.CoreSession; + +public class NativeWebSocketAdvices { + public static class OpenAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter( + @Advice.Argument(0) CoreSession session, + @Advice.FieldValue(value = "closeHandle", readOnly = false) MethodHandle closeHandle) { + AgentSpan span = activeSpan(); + if (span != null && session.getBehavior() == Behavior.SERVER && closeHandle != null) { + closeHandle = NativeMethodHandleWrappers.wrapClose(closeHandle, span, session); + } + } + } + + public static class MessageSinkAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter( + @Advice.Argument(1) WebSocketSession session, + @Advice.Argument(value = 2, readOnly = false) MethodHandle handle) { + AgentSpan span = activeSpan(); + if (span != null + && session.getCoreSession().getBehavior() == Behavior.SERVER + && handle != null) { + handle = + NativeMethodHandleWrappers.wrapMessage( + handle, + span, + session.getCoreSession(), + InstrumentationContext.get(CoreSession.class, ReceiveContexts.class)); + } + } + } + + public static class CloseAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter(@Advice.FieldValue("session") WebSocketSession session) { + if (session != null) { + ReceiveContexts contexts = + InstrumentationContext.get(CoreSession.class, ReceiveContexts.class) + .remove(session.getCoreSession()); + if (contexts != null) { + contexts.finish(); + } + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyEndpoints.groovy b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyEndpoints.groovy new file mode 100644 index 00000000000..3eb1027e650 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyEndpoints.groovy @@ -0,0 +1,156 @@ +import datadog.trace.bootstrap.instrumentation.api.AgentSpan +import org.eclipse.jetty.websocket.api.Callback +import org.eclipse.jetty.websocket.api.Session +import org.eclipse.jetty.websocket.api.annotations.OnWebSocketClose +import org.eclipse.jetty.websocket.api.annotations.OnWebSocketMessage +import org.eclipse.jetty.websocket.api.annotations.OnWebSocketOpen +import org.eclipse.jetty.websocket.api.annotations.WebSocket + +import java.nio.ByteBuffer + +import static java.nio.charset.StandardCharsets.UTF_8 +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan + +class JettyEndpoints { + static class EndpointEvents { + Session session + final List messages = [] + final List finalFragments = [] + final List messageSpans = [] + boolean failMessages + boolean failClose + boolean deferCallback + boolean failCallback + boolean failAfterCallback + Callback pendingCallback + int closeCode + String closeReason + + void recordMessage(String message) { + messages.add(message) + messageSpans.add(activeSpan()) + if (failMessages) { + throw new IllegalStateException("handler failed") + } + } + + void recordClose(int statusCode, String reason) { + closeCode = statusCode + closeReason = reason + if (failClose) { + throw new IllegalStateException("handler failed") + } + } + + void completeBinary(Callback callback) { + if (deferCallback) { + pendingCallback = callback + } else if (failCallback) { + callback.fail(new IllegalStateException("callback failed")) + } else { + callback.succeed() + } + if (failAfterCallback) { + throw new IllegalStateException("handler failed") + } + } + } + + static class ListenerEndpoint extends EndpointEvents implements Session.Listener.AutoDemanding { + @Override + void onWebSocketOpen(Session session) { + this.session = session + } + + @Override + void onWebSocketClose(int statusCode, String reason) { + recordClose(statusCode, reason) + } + } + + static class FullListener extends ListenerEndpoint { + @Override + void onWebSocketText(String message) { + recordMessage(message) + } + + @Override + void onWebSocketBinary(ByteBuffer payload, Callback callback) { + recordMessage(UTF_8.decode(payload).toString()) + completeBinary(callback) + } + } + + static class PartialListener extends ListenerEndpoint { + @Override + void onWebSocketPartialText(String payload, boolean last) { + finalFragments.add(last) + recordMessage(payload) + } + + @Override + void onWebSocketPartialBinary(ByteBuffer payload, boolean last, Callback callback) { + finalFragments.add(last) + recordMessage(UTF_8.decode(payload).toString()) + completeBinary(callback) + } + } + + static class PojoEndpoint extends EndpointEvents { + @OnWebSocketOpen + void onOpen(Session session) { + this.session = session + } + + @OnWebSocketClose + void onClose(int statusCode, String reason) { + recordClose(statusCode, reason) + } + } + + @WebSocket + static class PojoFullEndpoint extends PojoEndpoint { + @OnWebSocketMessage + void onText(Session session, String payload) { + assert session == this.session + recordMessage(payload) + } + + @OnWebSocketMessage + void onBinary(ByteBuffer payload, Callback callback) { + recordMessage(UTF_8.decode(payload).toString()) + completeBinary(callback) + } + } + + @WebSocket + static class PojoPartialEndpoint extends PojoEndpoint { + @OnWebSocketMessage + void onText(String payload, boolean last) { + finalFragments.add(last) + recordMessage(payload) + } + + @OnWebSocketMessage + void onBinary(Session session, ByteBuffer payload, boolean last, Callback callback) { + assert session == this.session + finalFragments.add(last) + recordMessage(UTF_8.decode(payload).toString()) + completeBinary(callback) + } + } + + @WebSocket + static class NoCloseEndpoint extends EndpointEvents { + @OnWebSocketMessage + void onText(String payload, boolean last) { + recordMessage(payload) + } + + @OnWebSocketMessage + void onBinary(ByteBuffer payload, boolean last, Callback callback) { + recordMessage(UTF_8.decode(payload).toString()) + completeBinary(callback) + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebSocketUpgradeTest.groovy b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebSocketUpgradeTest.groovy new file mode 100644 index 00000000000..4473ab257a5 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebSocketUpgradeTest.groovy @@ -0,0 +1,60 @@ +import datadog.trace.agent.test.InstrumentationSpecification +import datadog.trace.api.DDSpanTypes +import org.eclipse.jetty.server.Server +import org.eclipse.jetty.server.ServerConnector +import org.eclipse.jetty.server.handler.ContextHandler +import org.eclipse.jetty.websocket.api.Session +import org.eclipse.jetty.websocket.client.WebSocketClient +import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler + +import static java.util.concurrent.TimeUnit.SECONDS + +class JettyWebSocketUpgradeTest extends InstrumentationSpecification { + + def "HTTP client span finishes when the WebSocket upgrade succeeds"() { + setup: + def server = new Server(0) + def context = new ContextHandler("/") + server.handler = context + context.handler = WebSocketUpgradeHandler.from(server, context).configure { container -> + container.addMapping("/upgrade", { request, response, callback -> new Endpoint() }) + } + def client = new WebSocketClient() + server.start() + client.start() + def uri = URI.create("ws://localhost:${((ServerConnector) server.connectors[0]).localPort}/upgrade") + + when: + def session = client.connect(new Endpoint(), uri).get(5, SECONDS) + + then: + session.open + // The HTTP handshake must be reported before the WebSocket connection closes. + assertTraces(1) { + trace(1) { + span { + operationName "http.request" + resourceName "GET /upgrade" + spanType DDSpanTypes.HTTP_CLIENT + parent() + errored false + tags(false) { + "component" "jetty-client" + "span.kind" "client" + "http.status_code" 101 + } + } + } + } + + cleanup: + try { + client.stop() + } finally { + server.stop() + } + } + + static class Endpoint implements Session.Listener.AutoDemanding { + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebsocketTest.groovy b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebsocketTest.groovy new file mode 100644 index 00000000000..13134a58b7d --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebsocketTest.groovy @@ -0,0 +1,525 @@ +import datadog.trace.agent.test.InstrumentationSpecification +import datadog.trace.api.DDSpanTypes +import datadog.trace.core.DDSpan +import org.eclipse.jetty.util.Callback +import org.eclipse.jetty.websocket.api.UpgradeRequest +import org.eclipse.jetty.websocket.api.WebSocketContainer +import org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler +import org.eclipse.jetty.websocket.core.CloseStatus +import org.eclipse.jetty.websocket.core.Behavior +import org.eclipse.jetty.websocket.core.CoreSession +import org.eclipse.jetty.websocket.core.Frame +import org.eclipse.jetty.websocket.core.OpCode +import org.eclipse.jetty.websocket.core.WebSocketComponents +import org.eclipse.jetty.websocket.server.internal.ServerFrameHandlerFactory + +import java.util.concurrent.ExecutionException + +import static datadog.trace.agent.test.base.HttpServerTest.websocketCloseSpan +import static datadog.trace.agent.test.base.HttpServerTest.websocketReceiveSpan +import static datadog.trace.agent.test.utils.TraceUtils.basicSpan +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan +import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_URL +import static java.util.concurrent.TimeUnit.SECONDS + +class JettyWebsocketTest extends InstrumentationSpecification { + + def "test native jetty advices with endpoint class #endpoint.class and message type #msgType"() { + setup: + def url = "ws://inmemory/test/param" + def frameHandler = createFrameHandler(endpoint) + def messageCallback = new Callback.Completable() + def closeCallback = new Callback.Completable() + + when: + openFrameHandler(frameHandler) + frameHandler.onFrame(new Frame(msgType == "text" ? OpCode.TEXT : OpCode.BINARY, "hello world"), messageCallback) + messageCallback.get(5, SECONDS) + frameHandler.onFrame(CloseStatus.toFrame(CloseStatus.NORMAL, "bye"), closeCallback) + closeCallback.get(5, SECONDS) + + then: + endpoint.session != null + endpoint.messages == ["hello world"] + endpoint.messageSpans*.operationName*.toString() == ["websocket.receive"] + activeSpan() == null + endpoint.closeCode == 1000 + endpoint.closeReason == "bye" + assertTraces(3) { + DDSpan handshake + trace(1) { + handshake = span(0) + basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): url]) + } + trace(1) { + websocketReceiveSpan(it, handshake, msgType, 11) + } + trace(1) { + websocketCloseSpan(it, handshake, false, 1000, "bye") + } + } + + where: + endpoint | msgType + new JettyEndpoints.FullListener() | "text" + new JettyEndpoints.FullListener() | "binary" + new JettyEndpoints.PartialListener() | "text" + new JettyEndpoints.PartialListener() | "binary" + new JettyEndpoints.PojoFullEndpoint() | "text" + new JettyEndpoints.PojoFullEndpoint() | "binary" + new JettyEndpoints.PojoPartialEndpoint() | "text" + new JettyEndpoints.PojoPartialEndpoint() | "binary" + } + + def "fragmented #msgType messages share a span for #endpoint.class"() { + setup: + def frameHandler = createFrameHandler(endpoint) + openFrameHandler(frameHandler) + def opcode = msgType == "text" ? OpCode.TEXT : OpCode.BINARY + + when: + deliver(frameHandler, new Frame(opcode, "hello ").setFin(false)) + + then: + activeSpan() == null + endpoint.messages == ["hello "] + TEST_WRITER.size() == 1 + + when: + deliver(frameHandler, new Frame(OpCode.CONTINUATION, "world")) + deliver(frameHandler, new Frame(opcode, "again")) + + then: + endpoint.messages == ["hello ", "world", "again"] + endpoint.finalFragments == [false, true, true] + endpoint.messageSpans[0].is(endpoint.messageSpans[1]) + !endpoint.messageSpans[1].is(endpoint.messageSpans[2]) + activeSpan() == null + assertTraces(3) { + DDSpan handshake + trace(1) { + handshake = span(0) + basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): "ws://inmemory/test/param"]) + } + trace(1) { + websocketReceiveSpan(it, handshake, msgType, 11, 2) + } + trace(1) { + websocketReceiveSpan(it, handshake, msgType, 5) + } + } + + where: + endpoint | msgType + new JettyEndpoints.PartialListener() | "text" + new JettyEndpoints.PartialListener() | "binary" + new JettyEndpoints.PojoPartialEndpoint() | "text" + new JettyEndpoints.PojoPartialEndpoint() | "binary" + } + + def "handler failure marks #operation as errored and closes the scope"() { + setup: + def frameHandler = createFrameHandler(endpoint) + openFrameHandler(frameHandler) + + when: + deliver(frameHandler, frame) + + then: + def error = thrown(ExecutionException) + error.cause != null + activeSpan() == null + assertTraces(2) { + trace(1) { + basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): "ws://inmemory/test/param"]) + } + trace(1) { + span { + operationName operation + resourceName "websocket /test/param" + spanType DDSpanTypes.WEBSOCKET + errored true + ignoreSpanLinks() + tags(false) { + errorTags(IllegalStateException, "handler failed") + } + } + } + } + + where: + endpoint | frame | operation + new JettyEndpoints.FullListener(failMessages: true) | new Frame(OpCode.TEXT, "hello") | "websocket.receive" + new JettyEndpoints.FullListener(failMessages: true) | new Frame(OpCode.BINARY, "hello") | "websocket.receive" + new JettyEndpoints.PartialListener(failMessages: true) | new Frame(OpCode.TEXT, "hello").setFin(false) | "websocket.receive" + new JettyEndpoints.PojoFullEndpoint(failClose: true) | CloseStatus.toFrame(CloseStatus.NORMAL, "bye") | "websocket.close" + } + + def "does not trace native messages for #behavior with handshake tracing #traced"() { + setup: + def endpoint = new JettyEndpoints.FullListener() + def frameHandler = createFrameHandler(endpoint) + + when: + openFrameHandler(frameHandler, behavior, traced) + deliver(frameHandler, new Frame(OpCode.TEXT, "hello")) + deliver(frameHandler, CloseStatus.toFrame(CloseStatus.NORMAL, "bye")) + + then: + endpoint.messages == ["hello"] + endpoint.messageSpans == [null] + endpoint.closeCode == 1000 + activeSpan() == null + assertTraces(traced ? 1 : 0) { + if (traced) { + trace(1) { + basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): "ws://inmemory/test/param"]) + } + } + } + + where: + behavior | traced + Behavior.SERVER | false + Behavior.CLIENT | true + } + + def "binary callback failure marks the receive span for #endpoint.class"() { + setup: + endpoint.failCallback = true + def frameHandler = createFrameHandler(endpoint) + openFrameHandler(frameHandler) + + when: + deliver(frameHandler, new Frame(OpCode.BINARY, "hello").setFin(last)) + + then: + def error = thrown(ExecutionException) + error.cause.message == "callback failed" + activeSpan() == null + TEST_WRITER.waitForTraces(2) + endpoint.messageSpans[0].error + endpoint.messageSpans[0].getTag("error.message") == "callback failed" + + where: + endpoint | last + new JettyEndpoints.FullListener() | true + new JettyEndpoints.PartialListener() | true + new JettyEndpoints.PartialListener() | false + new JettyEndpoints.PojoFullEndpoint() | true + new JettyEndpoints.PojoPartialEndpoint() | false + } + + def "deferred binary callback completes the receive span with failure #fail"() { + setup: + def endpoint = new JettyEndpoints.FullListener(deferCallback: true) + def frameHandler = createFrameHandler(endpoint) + def callback = new Callback.Completable() + openFrameHandler(frameHandler) + + when: + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello"), callback) + + then: + endpoint.pendingCallback != null + !callback.done + TEST_WRITER.size() == 1 + activeSpan() == null + + when: + def completion = java.util.concurrent.CompletableFuture.runAsync { + if (fail) { + endpoint.pendingCallback.fail(new IllegalStateException("callback failed")) + } else { + endpoint.pendingCallback.succeed() + } + assert activeSpan() == null + } + completion.get(5, SECONDS) + + then: + callback.done + callback.completedExceptionally == fail + TEST_WRITER.waitForTraces(2) + endpoint.messageSpans[0].isError() == fail + activeSpan() == null + + where: + fail << [false, true] + } + + def "pending binary messages have independent spans for #endpoint.class with reverse completion #reverse"() { + setup: + endpoint.deferCallback = true + def frameHandler = createFrameHandler(endpoint) + openFrameHandler(frameHandler) + def callbacks = [new Callback.Completable(), new Callback.Completable()] + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello"), callbacks[0]) + def first = endpoint.pendingCallback + frameHandler.onFrame(new Frame(OpCode.BINARY, "again"), callbacks[1]) + def second = endpoint.pendingCallback + + expect: + !endpoint.messageSpans[0].is(endpoint.messageSpans[1]) + callbacks.every { !it.done } + TEST_WRITER.size() == 1 + activeSpan() == null + + when: + if (reverse) { + second.fail(new IllegalStateException("second callback failed")) + } else { + first.succeed() + } + + then: + TEST_WRITER.waitForTraces(2) + TEST_WRITER[1][0].is(endpoint.messageSpans[reverse ? 1 : 0]) + !callbacks[reverse ? 0 : 1].done + + when: + if (reverse) { + first.succeed() + } else { + second.fail(new IllegalStateException("second callback failed")) + } + + then: + TEST_WRITER.waitForTraces(3) + TEST_WRITER[2][0].is(endpoint.messageSpans[reverse ? 0 : 1]) + !endpoint.messageSpans[0].isError() + endpoint.messageSpans[1].isError() + endpoint.messageSpans[1].getTag("error.message") == "second callback failed" + callbacks[0].done && !callbacks[0].completedExceptionally + callbacks[1].completedExceptionally + activeSpan() == null + + where: + endpoint | reverse + new JettyEndpoints.FullListener() | false + new JettyEndpoints.FullListener() | true + new JettyEndpoints.PartialListener() | false + new JettyEndpoints.PartialListener() | true + new JettyEndpoints.PojoFullEndpoint() | false + new JettyEndpoints.PojoFullEndpoint() | true + new JettyEndpoints.PojoPartialEndpoint() | false + new JettyEndpoints.PojoPartialEndpoint() | true + } + + def "partial binary message waits for earlier fragment callbacks with failure #fail"() { + setup: + def endpoint = new JettyEndpoints.PartialListener(deferCallback: true) + def frameHandler = createFrameHandler(endpoint) + openFrameHandler(frameHandler) + def first = new Callback.Completable() + def last = new Callback.Completable() + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello ").setFin(false), first) + def firstCallback = endpoint.pendingCallback + frameHandler.onFrame(new Frame(OpCode.CONTINUATION, "world"), last) + + when: + endpoint.pendingCallback.succeed() + last.get(5, SECONDS) + + then: + !first.done + TEST_WRITER.size() == 1 + endpoint.messageSpans[0].is(endpoint.messageSpans[1]) + + when: + if (fail) { + firstCallback.fail(new IllegalStateException("first fragment failed")) + } else { + firstCallback.succeed() + } + + then: + first.done + first.completedExceptionally == fail + TEST_WRITER.waitForTraces(2) + endpoint.messageSpans[0].isError() == fail + if (fail) { + assert endpoint.messageSpans[0].getTag("error.message") == "first fragment failed" + } + endpoint.messageSpans[0].getTag("websocket.message.length") == 11 + endpoint.messageSpans[0].getTag("websocket.message.frames") == 2 + activeSpan() == null + + where: + fail << [false, true] + } + + def "handler failure after callback success still marks the receive span"() { + setup: + def endpoint = new JettyEndpoints.FullListener(failAfterCallback: true) + def frameHandler = createFrameHandler(endpoint) + openFrameHandler(frameHandler) + + when: + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello"), new Callback.Completable()) + + then: + TEST_WRITER.waitForTraces(2) + endpoint.messageSpans[0].error + endpoint.messageSpans[0].getTag("error.message") == "handler failed" + activeSpan() == null + } + + def "termination finishes a fragmented #msgType receive span with close frame #closeFrame and close handler #closeHandler"() { + setup: + def endpoint = closeHandler ? new JettyEndpoints.PartialListener() : new JettyEndpoints.NoCloseEndpoint() + def frameHandler = createFrameHandler(endpoint) + openFrameHandler(frameHandler) + deliver(frameHandler, new Frame(msgType == "text" ? OpCode.TEXT : OpCode.BINARY, "hello").setFin(false)) + assert TEST_WRITER.size() == 1 + + when: + if (closeFrame) { + deliver(frameHandler, CloseStatus.toFrame(CloseStatus.NORMAL, "bye")) + } else { + def callback = new Callback.Completable() + frameHandler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), callback) + callback.get(5, SECONDS) + } + + then: + activeSpan() == null + assertTraces(closeHandler ? 3 : 2) { + DDSpan handshake + trace(1) { + handshake = span(0) + basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): "ws://inmemory/test/param"]) + } + trace(1) { + websocketReceiveSpan(it, handshake, msgType, 5) + } + if (closeHandler) { + trace(1) { + websocketCloseSpan(it, handshake, false, 1000, "bye") + } + } + } + + where: + msgType | closeFrame | closeHandler + "text" | true | true + "binary" | true | true + "text" | false | true + "binary" | false | true + "text" | true | false + "binary" | true | false + "text" | false | false + "binary" | false | false + } + + def "deferred partial binary success keeps the span open until the final callback"() { + setup: + def endpoint = new JettyEndpoints.PartialListener(deferCallback: true) + def frameHandler = createFrameHandler(endpoint) + openFrameHandler(frameHandler) + def first = new Callback.Completable() + def last = new Callback.Completable() + + when: + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello ").setFin(false), first) + endpoint.pendingCallback.succeed() + first.get(5, SECONDS) + frameHandler.onFrame(new Frame(OpCode.CONTINUATION, "world"), last) + + then: + !last.done + TEST_WRITER.size() == 1 + endpoint.messageSpans[0].is(endpoint.messageSpans[1]) + + when: + endpoint.pendingCallback.succeed() + last.get(5, SECONDS) + + then: + assertTraces(2) { + DDSpan handshake + trace(1) { + handshake = span(0) + basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): "ws://inmemory/test/param"]) + } + trace(1) { + websocketReceiveSpan(it, handshake, "binary", 11, 2) + } + } + activeSpan() == null + } + + def "termination finishes all pending binary messages with final fragment #last"() { + setup: + def endpoint = new JettyEndpoints.PartialListener(deferCallback: true) + def frameHandler = createFrameHandler(endpoint) + openFrameHandler(frameHandler) + def previousMessage = new Callback.Completable() + frameHandler.onFrame(new Frame(OpCode.BINARY, "previous"), previousMessage) + def previousCallback = endpoint.pendingCallback + def message = new Callback.Completable() + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello").setFin(last), message) + + expect: + TEST_WRITER.size() == 1 + !endpoint.messageSpans[0].is(endpoint.messageSpans[1]) + + when: + def closed = new Callback.Completable() + frameHandler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), closed) + closed.get(5, SECONDS) + + then: + TEST_WRITER.waitForTraces(4) + TEST_WRITER.flatten().containsAll(endpoint.messageSpans) + + when: + endpoint.pendingCallback.succeed() + message.get(5, SECONDS) + previousCallback.succeed() + previousMessage.get(5, SECONDS) + + then: + TEST_WRITER.size() == 4 + activeSpan() == null + + where: + last << [false, true] + } + + private JettyWebSocketFrameHandler createFrameHandler(Object endpoint) { + def factory = new ServerFrameHandlerFactory(Stub(WebSocketContainer), new WebSocketComponents()) + def frameHandler = factory.newJettyFrameHandler(endpoint) + frameHandler.setUpgradeRequest(Stub(UpgradeRequest) { + getRequestURI() >> URI.create("ws://inmemory/test/param") + }) + return frameHandler + } + + private void openFrameHandler(JettyWebSocketFrameHandler frameHandler, Behavior connectionBehavior = Behavior.SERVER, boolean traced = true) { + def session = new CoreSession.Empty() { + @Override + Behavior getBehavior() { + return connectionBehavior + } + } + def openCallback = new Callback.Completable() + if (traced) { + runUnderTrace("parent") { + activeSpan().setTag(HTTP_URL, "ws://inmemory/test/param") + frameHandler.onOpen(session, openCallback) + openCallback.get(5, SECONDS) + } + } else { + frameHandler.onOpen(session, openCallback) + openCallback.get(5, SECONDS) + } + } + + private static void deliver(JettyWebSocketFrameHandler frameHandler, Frame frame) { + def callback = new Callback.Completable() + frameHandler.onFrame(frame, callback) + callback.get(5, SECONDS) + } +} From f8975b9bd8cd81855f84ac919e1721c4c2eaf39b Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Mon, 28 Sep 2026 17:49:25 -0700 Subject: [PATCH 02/23] Convert Jetty 12 WebSocket tests to Java and JUnit 5 --- .../jetty-websocket-12.0/build.gradle | 7 + .../jetty-websocket-12.0/gradle.lockfile | 3 +- .../src/test/groovy/JettyEndpoints.groovy | 156 ----- .../groovy/JettyWebSocketUpgradeTest.groovy | 60 -- .../src/test/groovy/JettyWebsocketTest.groovy | 525 --------------- .../src/test/java/JettyEndpoints.java | 159 +++++ .../test/java/JettyWebSocketUpgradeTest.java | 69 ++ .../src/test/java/JettyWebsocketTest.java | 607 ++++++++++++++++++ 8 files changed, 844 insertions(+), 742 deletions(-) delete mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyEndpoints.groovy delete mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebSocketUpgradeTest.groovy delete mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebsocketTest.groovy create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle index d7613f481b2..7af4629a730 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle @@ -53,7 +53,14 @@ addTestSuiteForDir("latestDepTest", "test") } } +["compileTestJava", "compileLatestDepTestJava"].each { name -> + tasks.named(name, JavaCompile) { + configureCompiler(it, 17, JavaVersion.VERSION_1_8, "Uses Jetty 12 APIs (Java 17+) at Java 8 bytecode") + } +} + dependencies { + testImplementation libs.bundles.mockito main_java17CompileOnly 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0' implementation project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-10.0") testImplementation group: 'org.eclipse.jetty.ee8.websocket', name: 'jetty-ee8-websocket-javax-server', version: '12.0.0' diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile index d56f5eeeec0..7e31f1cf65c 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile @@ -206,7 +206,8 @@ org.junit.platform:junit-platform-runner:1.14.1=latestDepTestRuntimeClasspath,te org.junit.platform:junit-platform-suite-api:1.14.1=latestDepTestRuntimeClasspath,testRuntimeClasspath org.junit.platform:junit-platform-suite-commons:1.14.1=latestDepTestRuntimeClasspath,testRuntimeClasspath org.junit:junit-bom:5.14.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.mockito:mockito-core:4.4.0=latestDepTestRuntimeClasspath,testRuntimeClasspath +org.mockito:mockito-core:4.4.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.mockito:mockito-junit-jupiter:4.4.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.objenesis:objenesis:3.3=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.opentest4j:opentest4j:1.3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.ow2.asm:asm-analysis:9.10.1=spotbugs diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyEndpoints.groovy b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyEndpoints.groovy deleted file mode 100644 index 3eb1027e650..00000000000 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyEndpoints.groovy +++ /dev/null @@ -1,156 +0,0 @@ -import datadog.trace.bootstrap.instrumentation.api.AgentSpan -import org.eclipse.jetty.websocket.api.Callback -import org.eclipse.jetty.websocket.api.Session -import org.eclipse.jetty.websocket.api.annotations.OnWebSocketClose -import org.eclipse.jetty.websocket.api.annotations.OnWebSocketMessage -import org.eclipse.jetty.websocket.api.annotations.OnWebSocketOpen -import org.eclipse.jetty.websocket.api.annotations.WebSocket - -import java.nio.ByteBuffer - -import static java.nio.charset.StandardCharsets.UTF_8 -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan - -class JettyEndpoints { - static class EndpointEvents { - Session session - final List messages = [] - final List finalFragments = [] - final List messageSpans = [] - boolean failMessages - boolean failClose - boolean deferCallback - boolean failCallback - boolean failAfterCallback - Callback pendingCallback - int closeCode - String closeReason - - void recordMessage(String message) { - messages.add(message) - messageSpans.add(activeSpan()) - if (failMessages) { - throw new IllegalStateException("handler failed") - } - } - - void recordClose(int statusCode, String reason) { - closeCode = statusCode - closeReason = reason - if (failClose) { - throw new IllegalStateException("handler failed") - } - } - - void completeBinary(Callback callback) { - if (deferCallback) { - pendingCallback = callback - } else if (failCallback) { - callback.fail(new IllegalStateException("callback failed")) - } else { - callback.succeed() - } - if (failAfterCallback) { - throw new IllegalStateException("handler failed") - } - } - } - - static class ListenerEndpoint extends EndpointEvents implements Session.Listener.AutoDemanding { - @Override - void onWebSocketOpen(Session session) { - this.session = session - } - - @Override - void onWebSocketClose(int statusCode, String reason) { - recordClose(statusCode, reason) - } - } - - static class FullListener extends ListenerEndpoint { - @Override - void onWebSocketText(String message) { - recordMessage(message) - } - - @Override - void onWebSocketBinary(ByteBuffer payload, Callback callback) { - recordMessage(UTF_8.decode(payload).toString()) - completeBinary(callback) - } - } - - static class PartialListener extends ListenerEndpoint { - @Override - void onWebSocketPartialText(String payload, boolean last) { - finalFragments.add(last) - recordMessage(payload) - } - - @Override - void onWebSocketPartialBinary(ByteBuffer payload, boolean last, Callback callback) { - finalFragments.add(last) - recordMessage(UTF_8.decode(payload).toString()) - completeBinary(callback) - } - } - - static class PojoEndpoint extends EndpointEvents { - @OnWebSocketOpen - void onOpen(Session session) { - this.session = session - } - - @OnWebSocketClose - void onClose(int statusCode, String reason) { - recordClose(statusCode, reason) - } - } - - @WebSocket - static class PojoFullEndpoint extends PojoEndpoint { - @OnWebSocketMessage - void onText(Session session, String payload) { - assert session == this.session - recordMessage(payload) - } - - @OnWebSocketMessage - void onBinary(ByteBuffer payload, Callback callback) { - recordMessage(UTF_8.decode(payload).toString()) - completeBinary(callback) - } - } - - @WebSocket - static class PojoPartialEndpoint extends PojoEndpoint { - @OnWebSocketMessage - void onText(String payload, boolean last) { - finalFragments.add(last) - recordMessage(payload) - } - - @OnWebSocketMessage - void onBinary(Session session, ByteBuffer payload, boolean last, Callback callback) { - assert session == this.session - finalFragments.add(last) - recordMessage(UTF_8.decode(payload).toString()) - completeBinary(callback) - } - } - - @WebSocket - static class NoCloseEndpoint extends EndpointEvents { - @OnWebSocketMessage - void onText(String payload, boolean last) { - recordMessage(payload) - } - - @OnWebSocketMessage - void onBinary(ByteBuffer payload, boolean last, Callback callback) { - recordMessage(UTF_8.decode(payload).toString()) - completeBinary(callback) - } - } -} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebSocketUpgradeTest.groovy b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebSocketUpgradeTest.groovy deleted file mode 100644 index 4473ab257a5..00000000000 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebSocketUpgradeTest.groovy +++ /dev/null @@ -1,60 +0,0 @@ -import datadog.trace.agent.test.InstrumentationSpecification -import datadog.trace.api.DDSpanTypes -import org.eclipse.jetty.server.Server -import org.eclipse.jetty.server.ServerConnector -import org.eclipse.jetty.server.handler.ContextHandler -import org.eclipse.jetty.websocket.api.Session -import org.eclipse.jetty.websocket.client.WebSocketClient -import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler - -import static java.util.concurrent.TimeUnit.SECONDS - -class JettyWebSocketUpgradeTest extends InstrumentationSpecification { - - def "HTTP client span finishes when the WebSocket upgrade succeeds"() { - setup: - def server = new Server(0) - def context = new ContextHandler("/") - server.handler = context - context.handler = WebSocketUpgradeHandler.from(server, context).configure { container -> - container.addMapping("/upgrade", { request, response, callback -> new Endpoint() }) - } - def client = new WebSocketClient() - server.start() - client.start() - def uri = URI.create("ws://localhost:${((ServerConnector) server.connectors[0]).localPort}/upgrade") - - when: - def session = client.connect(new Endpoint(), uri).get(5, SECONDS) - - then: - session.open - // The HTTP handshake must be reported before the WebSocket connection closes. - assertTraces(1) { - trace(1) { - span { - operationName "http.request" - resourceName "GET /upgrade" - spanType DDSpanTypes.HTTP_CLIENT - parent() - errored false - tags(false) { - "component" "jetty-client" - "span.kind" "client" - "http.status_code" 101 - } - } - } - } - - cleanup: - try { - client.stop() - } finally { - server.stop() - } - } - - static class Endpoint implements Session.Listener.AutoDemanding { - } -} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebsocketTest.groovy b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebsocketTest.groovy deleted file mode 100644 index 13134a58b7d..00000000000 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/groovy/JettyWebsocketTest.groovy +++ /dev/null @@ -1,525 +0,0 @@ -import datadog.trace.agent.test.InstrumentationSpecification -import datadog.trace.api.DDSpanTypes -import datadog.trace.core.DDSpan -import org.eclipse.jetty.util.Callback -import org.eclipse.jetty.websocket.api.UpgradeRequest -import org.eclipse.jetty.websocket.api.WebSocketContainer -import org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler -import org.eclipse.jetty.websocket.core.CloseStatus -import org.eclipse.jetty.websocket.core.Behavior -import org.eclipse.jetty.websocket.core.CoreSession -import org.eclipse.jetty.websocket.core.Frame -import org.eclipse.jetty.websocket.core.OpCode -import org.eclipse.jetty.websocket.core.WebSocketComponents -import org.eclipse.jetty.websocket.server.internal.ServerFrameHandlerFactory - -import java.util.concurrent.ExecutionException - -import static datadog.trace.agent.test.base.HttpServerTest.websocketCloseSpan -import static datadog.trace.agent.test.base.HttpServerTest.websocketReceiveSpan -import static datadog.trace.agent.test.utils.TraceUtils.basicSpan -import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan -import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_URL -import static java.util.concurrent.TimeUnit.SECONDS - -class JettyWebsocketTest extends InstrumentationSpecification { - - def "test native jetty advices with endpoint class #endpoint.class and message type #msgType"() { - setup: - def url = "ws://inmemory/test/param" - def frameHandler = createFrameHandler(endpoint) - def messageCallback = new Callback.Completable() - def closeCallback = new Callback.Completable() - - when: - openFrameHandler(frameHandler) - frameHandler.onFrame(new Frame(msgType == "text" ? OpCode.TEXT : OpCode.BINARY, "hello world"), messageCallback) - messageCallback.get(5, SECONDS) - frameHandler.onFrame(CloseStatus.toFrame(CloseStatus.NORMAL, "bye"), closeCallback) - closeCallback.get(5, SECONDS) - - then: - endpoint.session != null - endpoint.messages == ["hello world"] - endpoint.messageSpans*.operationName*.toString() == ["websocket.receive"] - activeSpan() == null - endpoint.closeCode == 1000 - endpoint.closeReason == "bye" - assertTraces(3) { - DDSpan handshake - trace(1) { - handshake = span(0) - basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): url]) - } - trace(1) { - websocketReceiveSpan(it, handshake, msgType, 11) - } - trace(1) { - websocketCloseSpan(it, handshake, false, 1000, "bye") - } - } - - where: - endpoint | msgType - new JettyEndpoints.FullListener() | "text" - new JettyEndpoints.FullListener() | "binary" - new JettyEndpoints.PartialListener() | "text" - new JettyEndpoints.PartialListener() | "binary" - new JettyEndpoints.PojoFullEndpoint() | "text" - new JettyEndpoints.PojoFullEndpoint() | "binary" - new JettyEndpoints.PojoPartialEndpoint() | "text" - new JettyEndpoints.PojoPartialEndpoint() | "binary" - } - - def "fragmented #msgType messages share a span for #endpoint.class"() { - setup: - def frameHandler = createFrameHandler(endpoint) - openFrameHandler(frameHandler) - def opcode = msgType == "text" ? OpCode.TEXT : OpCode.BINARY - - when: - deliver(frameHandler, new Frame(opcode, "hello ").setFin(false)) - - then: - activeSpan() == null - endpoint.messages == ["hello "] - TEST_WRITER.size() == 1 - - when: - deliver(frameHandler, new Frame(OpCode.CONTINUATION, "world")) - deliver(frameHandler, new Frame(opcode, "again")) - - then: - endpoint.messages == ["hello ", "world", "again"] - endpoint.finalFragments == [false, true, true] - endpoint.messageSpans[0].is(endpoint.messageSpans[1]) - !endpoint.messageSpans[1].is(endpoint.messageSpans[2]) - activeSpan() == null - assertTraces(3) { - DDSpan handshake - trace(1) { - handshake = span(0) - basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): "ws://inmemory/test/param"]) - } - trace(1) { - websocketReceiveSpan(it, handshake, msgType, 11, 2) - } - trace(1) { - websocketReceiveSpan(it, handshake, msgType, 5) - } - } - - where: - endpoint | msgType - new JettyEndpoints.PartialListener() | "text" - new JettyEndpoints.PartialListener() | "binary" - new JettyEndpoints.PojoPartialEndpoint() | "text" - new JettyEndpoints.PojoPartialEndpoint() | "binary" - } - - def "handler failure marks #operation as errored and closes the scope"() { - setup: - def frameHandler = createFrameHandler(endpoint) - openFrameHandler(frameHandler) - - when: - deliver(frameHandler, frame) - - then: - def error = thrown(ExecutionException) - error.cause != null - activeSpan() == null - assertTraces(2) { - trace(1) { - basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): "ws://inmemory/test/param"]) - } - trace(1) { - span { - operationName operation - resourceName "websocket /test/param" - spanType DDSpanTypes.WEBSOCKET - errored true - ignoreSpanLinks() - tags(false) { - errorTags(IllegalStateException, "handler failed") - } - } - } - } - - where: - endpoint | frame | operation - new JettyEndpoints.FullListener(failMessages: true) | new Frame(OpCode.TEXT, "hello") | "websocket.receive" - new JettyEndpoints.FullListener(failMessages: true) | new Frame(OpCode.BINARY, "hello") | "websocket.receive" - new JettyEndpoints.PartialListener(failMessages: true) | new Frame(OpCode.TEXT, "hello").setFin(false) | "websocket.receive" - new JettyEndpoints.PojoFullEndpoint(failClose: true) | CloseStatus.toFrame(CloseStatus.NORMAL, "bye") | "websocket.close" - } - - def "does not trace native messages for #behavior with handshake tracing #traced"() { - setup: - def endpoint = new JettyEndpoints.FullListener() - def frameHandler = createFrameHandler(endpoint) - - when: - openFrameHandler(frameHandler, behavior, traced) - deliver(frameHandler, new Frame(OpCode.TEXT, "hello")) - deliver(frameHandler, CloseStatus.toFrame(CloseStatus.NORMAL, "bye")) - - then: - endpoint.messages == ["hello"] - endpoint.messageSpans == [null] - endpoint.closeCode == 1000 - activeSpan() == null - assertTraces(traced ? 1 : 0) { - if (traced) { - trace(1) { - basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): "ws://inmemory/test/param"]) - } - } - } - - where: - behavior | traced - Behavior.SERVER | false - Behavior.CLIENT | true - } - - def "binary callback failure marks the receive span for #endpoint.class"() { - setup: - endpoint.failCallback = true - def frameHandler = createFrameHandler(endpoint) - openFrameHandler(frameHandler) - - when: - deliver(frameHandler, new Frame(OpCode.BINARY, "hello").setFin(last)) - - then: - def error = thrown(ExecutionException) - error.cause.message == "callback failed" - activeSpan() == null - TEST_WRITER.waitForTraces(2) - endpoint.messageSpans[0].error - endpoint.messageSpans[0].getTag("error.message") == "callback failed" - - where: - endpoint | last - new JettyEndpoints.FullListener() | true - new JettyEndpoints.PartialListener() | true - new JettyEndpoints.PartialListener() | false - new JettyEndpoints.PojoFullEndpoint() | true - new JettyEndpoints.PojoPartialEndpoint() | false - } - - def "deferred binary callback completes the receive span with failure #fail"() { - setup: - def endpoint = new JettyEndpoints.FullListener(deferCallback: true) - def frameHandler = createFrameHandler(endpoint) - def callback = new Callback.Completable() - openFrameHandler(frameHandler) - - when: - frameHandler.onFrame(new Frame(OpCode.BINARY, "hello"), callback) - - then: - endpoint.pendingCallback != null - !callback.done - TEST_WRITER.size() == 1 - activeSpan() == null - - when: - def completion = java.util.concurrent.CompletableFuture.runAsync { - if (fail) { - endpoint.pendingCallback.fail(new IllegalStateException("callback failed")) - } else { - endpoint.pendingCallback.succeed() - } - assert activeSpan() == null - } - completion.get(5, SECONDS) - - then: - callback.done - callback.completedExceptionally == fail - TEST_WRITER.waitForTraces(2) - endpoint.messageSpans[0].isError() == fail - activeSpan() == null - - where: - fail << [false, true] - } - - def "pending binary messages have independent spans for #endpoint.class with reverse completion #reverse"() { - setup: - endpoint.deferCallback = true - def frameHandler = createFrameHandler(endpoint) - openFrameHandler(frameHandler) - def callbacks = [new Callback.Completable(), new Callback.Completable()] - frameHandler.onFrame(new Frame(OpCode.BINARY, "hello"), callbacks[0]) - def first = endpoint.pendingCallback - frameHandler.onFrame(new Frame(OpCode.BINARY, "again"), callbacks[1]) - def second = endpoint.pendingCallback - - expect: - !endpoint.messageSpans[0].is(endpoint.messageSpans[1]) - callbacks.every { !it.done } - TEST_WRITER.size() == 1 - activeSpan() == null - - when: - if (reverse) { - second.fail(new IllegalStateException("second callback failed")) - } else { - first.succeed() - } - - then: - TEST_WRITER.waitForTraces(2) - TEST_WRITER[1][0].is(endpoint.messageSpans[reverse ? 1 : 0]) - !callbacks[reverse ? 0 : 1].done - - when: - if (reverse) { - first.succeed() - } else { - second.fail(new IllegalStateException("second callback failed")) - } - - then: - TEST_WRITER.waitForTraces(3) - TEST_WRITER[2][0].is(endpoint.messageSpans[reverse ? 0 : 1]) - !endpoint.messageSpans[0].isError() - endpoint.messageSpans[1].isError() - endpoint.messageSpans[1].getTag("error.message") == "second callback failed" - callbacks[0].done && !callbacks[0].completedExceptionally - callbacks[1].completedExceptionally - activeSpan() == null - - where: - endpoint | reverse - new JettyEndpoints.FullListener() | false - new JettyEndpoints.FullListener() | true - new JettyEndpoints.PartialListener() | false - new JettyEndpoints.PartialListener() | true - new JettyEndpoints.PojoFullEndpoint() | false - new JettyEndpoints.PojoFullEndpoint() | true - new JettyEndpoints.PojoPartialEndpoint() | false - new JettyEndpoints.PojoPartialEndpoint() | true - } - - def "partial binary message waits for earlier fragment callbacks with failure #fail"() { - setup: - def endpoint = new JettyEndpoints.PartialListener(deferCallback: true) - def frameHandler = createFrameHandler(endpoint) - openFrameHandler(frameHandler) - def first = new Callback.Completable() - def last = new Callback.Completable() - frameHandler.onFrame(new Frame(OpCode.BINARY, "hello ").setFin(false), first) - def firstCallback = endpoint.pendingCallback - frameHandler.onFrame(new Frame(OpCode.CONTINUATION, "world"), last) - - when: - endpoint.pendingCallback.succeed() - last.get(5, SECONDS) - - then: - !first.done - TEST_WRITER.size() == 1 - endpoint.messageSpans[0].is(endpoint.messageSpans[1]) - - when: - if (fail) { - firstCallback.fail(new IllegalStateException("first fragment failed")) - } else { - firstCallback.succeed() - } - - then: - first.done - first.completedExceptionally == fail - TEST_WRITER.waitForTraces(2) - endpoint.messageSpans[0].isError() == fail - if (fail) { - assert endpoint.messageSpans[0].getTag("error.message") == "first fragment failed" - } - endpoint.messageSpans[0].getTag("websocket.message.length") == 11 - endpoint.messageSpans[0].getTag("websocket.message.frames") == 2 - activeSpan() == null - - where: - fail << [false, true] - } - - def "handler failure after callback success still marks the receive span"() { - setup: - def endpoint = new JettyEndpoints.FullListener(failAfterCallback: true) - def frameHandler = createFrameHandler(endpoint) - openFrameHandler(frameHandler) - - when: - frameHandler.onFrame(new Frame(OpCode.BINARY, "hello"), new Callback.Completable()) - - then: - TEST_WRITER.waitForTraces(2) - endpoint.messageSpans[0].error - endpoint.messageSpans[0].getTag("error.message") == "handler failed" - activeSpan() == null - } - - def "termination finishes a fragmented #msgType receive span with close frame #closeFrame and close handler #closeHandler"() { - setup: - def endpoint = closeHandler ? new JettyEndpoints.PartialListener() : new JettyEndpoints.NoCloseEndpoint() - def frameHandler = createFrameHandler(endpoint) - openFrameHandler(frameHandler) - deliver(frameHandler, new Frame(msgType == "text" ? OpCode.TEXT : OpCode.BINARY, "hello").setFin(false)) - assert TEST_WRITER.size() == 1 - - when: - if (closeFrame) { - deliver(frameHandler, CloseStatus.toFrame(CloseStatus.NORMAL, "bye")) - } else { - def callback = new Callback.Completable() - frameHandler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), callback) - callback.get(5, SECONDS) - } - - then: - activeSpan() == null - assertTraces(closeHandler ? 3 : 2) { - DDSpan handshake - trace(1) { - handshake = span(0) - basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): "ws://inmemory/test/param"]) - } - trace(1) { - websocketReceiveSpan(it, handshake, msgType, 5) - } - if (closeHandler) { - trace(1) { - websocketCloseSpan(it, handshake, false, 1000, "bye") - } - } - } - - where: - msgType | closeFrame | closeHandler - "text" | true | true - "binary" | true | true - "text" | false | true - "binary" | false | true - "text" | true | false - "binary" | true | false - "text" | false | false - "binary" | false | false - } - - def "deferred partial binary success keeps the span open until the final callback"() { - setup: - def endpoint = new JettyEndpoints.PartialListener(deferCallback: true) - def frameHandler = createFrameHandler(endpoint) - openFrameHandler(frameHandler) - def first = new Callback.Completable() - def last = new Callback.Completable() - - when: - frameHandler.onFrame(new Frame(OpCode.BINARY, "hello ").setFin(false), first) - endpoint.pendingCallback.succeed() - first.get(5, SECONDS) - frameHandler.onFrame(new Frame(OpCode.CONTINUATION, "world"), last) - - then: - !last.done - TEST_WRITER.size() == 1 - endpoint.messageSpans[0].is(endpoint.messageSpans[1]) - - when: - endpoint.pendingCallback.succeed() - last.get(5, SECONDS) - - then: - assertTraces(2) { - DDSpan handshake - trace(1) { - handshake = span(0) - basicSpan(it, "parent", "/test/param", null, null, [(HTTP_URL): "ws://inmemory/test/param"]) - } - trace(1) { - websocketReceiveSpan(it, handshake, "binary", 11, 2) - } - } - activeSpan() == null - } - - def "termination finishes all pending binary messages with final fragment #last"() { - setup: - def endpoint = new JettyEndpoints.PartialListener(deferCallback: true) - def frameHandler = createFrameHandler(endpoint) - openFrameHandler(frameHandler) - def previousMessage = new Callback.Completable() - frameHandler.onFrame(new Frame(OpCode.BINARY, "previous"), previousMessage) - def previousCallback = endpoint.pendingCallback - def message = new Callback.Completable() - frameHandler.onFrame(new Frame(OpCode.BINARY, "hello").setFin(last), message) - - expect: - TEST_WRITER.size() == 1 - !endpoint.messageSpans[0].is(endpoint.messageSpans[1]) - - when: - def closed = new Callback.Completable() - frameHandler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), closed) - closed.get(5, SECONDS) - - then: - TEST_WRITER.waitForTraces(4) - TEST_WRITER.flatten().containsAll(endpoint.messageSpans) - - when: - endpoint.pendingCallback.succeed() - message.get(5, SECONDS) - previousCallback.succeed() - previousMessage.get(5, SECONDS) - - then: - TEST_WRITER.size() == 4 - activeSpan() == null - - where: - last << [false, true] - } - - private JettyWebSocketFrameHandler createFrameHandler(Object endpoint) { - def factory = new ServerFrameHandlerFactory(Stub(WebSocketContainer), new WebSocketComponents()) - def frameHandler = factory.newJettyFrameHandler(endpoint) - frameHandler.setUpgradeRequest(Stub(UpgradeRequest) { - getRequestURI() >> URI.create("ws://inmemory/test/param") - }) - return frameHandler - } - - private void openFrameHandler(JettyWebSocketFrameHandler frameHandler, Behavior connectionBehavior = Behavior.SERVER, boolean traced = true) { - def session = new CoreSession.Empty() { - @Override - Behavior getBehavior() { - return connectionBehavior - } - } - def openCallback = new Callback.Completable() - if (traced) { - runUnderTrace("parent") { - activeSpan().setTag(HTTP_URL, "ws://inmemory/test/param") - frameHandler.onOpen(session, openCallback) - openCallback.get(5, SECONDS) - } - } else { - frameHandler.onOpen(session, openCallback) - openCallback.get(5, SECONDS) - } - } - - private static void deliver(JettyWebSocketFrameHandler frameHandler, Frame frame) { - def callback = new Callback.Completable() - frameHandler.onFrame(frame, callback) - callback.get(5, SECONDS) - } -} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java new file mode 100644 index 00000000000..96c570daa89 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java @@ -0,0 +1,159 @@ +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.jupiter.api.Assertions.assertSame; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import org.eclipse.jetty.websocket.api.Callback; +import org.eclipse.jetty.websocket.api.Session; +import org.eclipse.jetty.websocket.api.annotations.OnWebSocketClose; +import org.eclipse.jetty.websocket.api.annotations.OnWebSocketMessage; +import org.eclipse.jetty.websocket.api.annotations.OnWebSocketOpen; +import org.eclipse.jetty.websocket.api.annotations.WebSocket; + +public class JettyEndpoints { + public static class EndpointEvents { + Session session; + final List messages = new ArrayList<>(); + final List finalFragments = new ArrayList<>(); + final List messageSpans = new ArrayList<>(); + boolean failMessages; + boolean failClose; + boolean deferCallback; + boolean failCallback; + boolean failAfterCallback; + Callback pendingCallback; + int closeCode; + String closeReason; + + void recordMessage(String message) { + messages.add(message); + messageSpans.add(activeSpan()); + if (failMessages) { + throw new IllegalStateException("handler failed"); + } + } + + void recordClose(int statusCode, String reason) { + closeCode = statusCode; + closeReason = reason; + if (failClose) { + throw new IllegalStateException("handler failed"); + } + } + + void completeBinary(Callback callback) { + if (deferCallback) { + pendingCallback = callback; + } else if (failCallback) { + callback.fail(new IllegalStateException("callback failed")); + } else { + callback.succeed(); + } + if (failAfterCallback) { + throw new IllegalStateException("handler failed"); + } + } + } + + public static class ListenerEndpoint extends EndpointEvents + implements Session.Listener.AutoDemanding { + @Override + public void onWebSocketOpen(Session session) { + this.session = session; + } + + @Override + public void onWebSocketClose(int statusCode, String reason) { + recordClose(statusCode, reason); + } + } + + public static class FullListener extends ListenerEndpoint { + @Override + public void onWebSocketText(String message) { + recordMessage(message); + } + + @Override + public void onWebSocketBinary(ByteBuffer payload, Callback callback) { + recordMessage(UTF_8.decode(payload).toString()); + completeBinary(callback); + } + } + + public static class PartialListener extends ListenerEndpoint { + @Override + public void onWebSocketPartialText(String payload, boolean last) { + finalFragments.add(last); + recordMessage(payload); + } + + @Override + public void onWebSocketPartialBinary(ByteBuffer payload, boolean last, Callback callback) { + finalFragments.add(last); + recordMessage(UTF_8.decode(payload).toString()); + completeBinary(callback); + } + } + + public static class PojoEndpoint extends EndpointEvents { + @OnWebSocketOpen + public void onOpen(Session session) { + this.session = session; + } + + @OnWebSocketClose + public void onClose(int statusCode, String reason) { + recordClose(statusCode, reason); + } + } + + @WebSocket + public static class PojoFullEndpoint extends PojoEndpoint { + @OnWebSocketMessage + public void onText(Session session, String payload) { + assertSame(this.session, session); + recordMessage(payload); + } + + @OnWebSocketMessage + public void onBinary(ByteBuffer payload, Callback callback) { + recordMessage(UTF_8.decode(payload).toString()); + completeBinary(callback); + } + } + + @WebSocket + public static class PojoPartialEndpoint extends PojoEndpoint { + @OnWebSocketMessage + public void onText(String payload, boolean last) { + finalFragments.add(last); + recordMessage(payload); + } + + @OnWebSocketMessage + public void onBinary(Session session, ByteBuffer payload, boolean last, Callback callback) { + assertSame(this.session, session); + finalFragments.add(last); + recordMessage(UTF_8.decode(payload).toString()); + completeBinary(callback); + } + } + + @WebSocket + public static class NoCloseEndpoint extends EndpointEvents { + @OnWebSocketMessage + public void onText(String payload, boolean last) { + recordMessage(payload); + } + + @OnWebSocketMessage + public void onBinary(ByteBuffer payload, boolean last, Callback callback) { + recordMessage(UTF_8.decode(payload).toString()); + completeBinary(callback); + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java new file mode 100644 index 00000000000..96025ee2e10 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java @@ -0,0 +1,69 @@ +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static java.util.concurrent.TimeUnit.SECONDS; +import static java.util.regex.Pattern.compile; +import static java.util.regex.Pattern.quote; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.api.DDSpanTypes; +import datadog.trace.core.DDSpan; +import java.net.URI; +import org.eclipse.jetty.server.Server; +import org.eclipse.jetty.server.ServerConnector; +import org.eclipse.jetty.server.handler.ContextHandler; +import org.eclipse.jetty.websocket.api.Session; +import org.eclipse.jetty.websocket.client.WebSocketClient; +import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler; +import org.junit.jupiter.api.Test; + +class JettyWebSocketUpgradeTest extends AbstractInstrumentationTest { + @Test + void httpClientSpanFinishesWhenWebSocketUpgradeSucceeds() throws Exception { + Server server = new Server(0); + WebSocketClient client = new WebSocketClient(); + try { + ContextHandler context = new ContextHandler("/"); + server.setHandler(context); + context.setHandler( + WebSocketUpgradeHandler.from(server, context) + .configure( + container -> + container.addMapping( + "/upgrade", (request, response, callback) -> new Endpoint()))); + server.start(); + client.start(); + URI uri = + URI.create( + "ws://localhost:" + + ((ServerConnector) server.getConnectors()[0]).getLocalPort() + + "/upgrade"); + + Session session = client.connect(new Endpoint(), uri).get(5, SECONDS); + + assertTrue(session.isOpen()); + // The HTTP handshake must be reported before the WebSocket connection closes. + assertTraces( + trace( + span() + .operationName(compile(quote("http.request"))) + .resourceName(compile(quote("GET /upgrade"))) + .type(DDSpanTypes.HTTP_CLIENT) + .root() + .error(false))); + DDSpan handshake = writer.get(0).get(0); + assertEquals("jetty-client", handshake.getTag("component").toString()); + assertEquals("client", handshake.getTag("span.kind")); + assertEquals(101, handshake.getTag("http.status_code")); + } finally { + try { + client.stop(); + } finally { + server.stop(); + } + } + } + + public static class Endpoint implements Session.Listener.AutoDemanding {} +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java new file mode 100644 index 00000000000..8bd563b3464 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java @@ -0,0 +1,607 @@ +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TagsMatcher.defaultTags; +import static datadog.trace.agent.test.assertions.TagsMatcher.includes; +import static datadog.trace.agent.test.assertions.TagsMatcher.tag; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace; +import static datadog.trace.api.DDTags.DECISION_MAKER_INHERITED; +import static datadog.trace.api.DDTags.DECISION_MAKER_RESOURCE; +import static datadog.trace.api.DDTags.DECISION_MAKER_SERVICE; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_CLOSE_CODE; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_CLOSE_REASON; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_FRAMES; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_LENGTH; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_RECEIVE_TIME; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_TYPE; +import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_URL; +import static datadog.trace.test.junit.utils.assertions.Matchers.is; +import static datadog.trace.test.junit.utils.assertions.Matchers.isNull; +import static datadog.trace.test.junit.utils.assertions.Matchers.matches; +import static java.util.Arrays.asList; +import static java.util.Collections.singletonList; +import static java.util.concurrent.TimeUnit.SECONDS; +import static java.util.regex.Pattern.compile; +import static java.util.regex.Pattern.quote; +import static java.util.stream.Collectors.toList; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.agent.test.assertions.SpanLinkMatcher; +import datadog.trace.agent.test.assertions.SpanMatcher; +import datadog.trace.agent.test.assertions.TagsMatcher; +import datadog.trace.api.DDSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink; +import datadog.trace.bootstrap.instrumentation.api.SpanAttributes; +import datadog.trace.core.DDSpan; +import java.net.URI; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import org.eclipse.jetty.util.Callback; +import org.eclipse.jetty.websocket.api.UpgradeRequest; +import org.eclipse.jetty.websocket.api.WebSocketContainer; +import org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler; +import org.eclipse.jetty.websocket.core.Behavior; +import org.eclipse.jetty.websocket.core.CloseStatus; +import org.eclipse.jetty.websocket.core.CoreSession; +import org.eclipse.jetty.websocket.core.Frame; +import org.eclipse.jetty.websocket.core.OpCode; +import org.eclipse.jetty.websocket.core.WebSocketComponents; +import org.eclipse.jetty.websocket.server.internal.ServerFrameHandlerFactory; +import org.junit.jupiter.api.Test; +import org.tabletest.junit.TableTest; +import org.tabletest.junit.TypeConverter; + +public class JettyWebsocketTest extends AbstractInstrumentationTest { + private static final String URL = "ws://inmemory/test/param"; + + @TableTest({ + "scenario | endpoint | msgType", + "full text | full | text ", + "full binary | full | binary ", + "partial text | partial | text ", + "partial binary | partial | binary ", + "POJO full text | pojoFull | text ", + "POJO full binary | pojoFull | binary ", + "POJO partial text | pojoPartial | text ", + "POJO partial binary | pojoPartial | binary " + }) + void nativeJettyAdvices(JettyEndpoints.EndpointEvents endpoint, String msgType) throws Exception { + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + Callback.Completable messageCallback = new Callback.Completable(); + Callback.Completable closeCallback = new Callback.Completable(); + + openFrameHandler(frameHandler); + frameHandler.onFrame(new Frame(opcode(msgType), "hello world"), messageCallback); + messageCallback.get(5, SECONDS); + frameHandler.onFrame(CloseStatus.toFrame(CloseStatus.NORMAL, "bye"), closeCallback); + closeCallback.get(5, SECONDS); + + assertNotNull(endpoint.session); + assertEquals(singletonList("hello world"), endpoint.messages); + assertEquals(1, endpoint.messageSpans.size()); + assertEquals("websocket.receive", endpoint.messageSpans.get(0).getOperationName().toString()); + assertNull(activeSpan()); + assertEquals(1000, endpoint.closeCode); + assertEquals("bye", endpoint.closeReason); + DDSpan handshake = handshake(); + assertTraces( + trace(handshakeSpan()), + trace(receiveSpan(handshake, msgType, 11, 1)), + trace(closeSpan(handshake))); + } + + @TableTest({ + "scenario | endpoint | msgType", + "partial text | partial | text ", + "partial binary | partial | binary ", + "POJO partial text | pojoPartial | text ", + "POJO partial binary | pojoPartial | binary " + }) + void fragmentedMessagesShareSpan(JettyEndpoints.EndpointEvents endpoint, String msgType) + throws Exception { + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + byte opcode = opcode(msgType); + + deliver(frameHandler, new Frame(opcode, "hello ").setFin(false)); + + assertNull(activeSpan()); + assertEquals(singletonList("hello "), endpoint.messages); + assertEquals(1, writer.size()); + + deliver(frameHandler, new Frame(OpCode.CONTINUATION, "world")); + deliver(frameHandler, new Frame(opcode, "again")); + + assertEquals(asList("hello ", "world", "again"), endpoint.messages); + assertEquals(asList(false, true, true), endpoint.finalFragments); + assertSame(endpoint.messageSpans.get(0), endpoint.messageSpans.get(1)); + assertNotSame(endpoint.messageSpans.get(1), endpoint.messageSpans.get(2)); + assertNull(activeSpan()); + DDSpan handshake = handshake(); + assertTraces( + trace(handshakeSpan()), + trace(receiveSpan(handshake, msgType, 11, 2)), + trace(receiveSpan(handshake, msgType, 5, 1))); + } + + @TableTest({ + "scenario | endpoint | msgType | last | operation ", + "full text | full | text | true | websocket.receive", + "full binary | full | binary | true | websocket.receive", + "partial text | partial | text | false | websocket.receive", + "POJO close | pojoFull | | true | websocket.close " + }) + void handlerFailureMarksSpanAndClosesScope( + JettyEndpoints.EndpointEvents endpoint, String msgType, boolean last, String operation) + throws Exception { + endpoint.failClose = msgType == null; + endpoint.failMessages = msgType != null; + Frame frame = + endpoint.failClose + ? CloseStatus.toFrame(CloseStatus.NORMAL, "bye") + : new Frame(opcode(msgType), "hello").setFin(last); + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + + ExecutionException error = + assertThrows(ExecutionException.class, () -> deliver(frameHandler, frame)); + + assertNotNull(error.getCause()); + assertNull(activeSpan()); + assertTraces( + trace(handshakeSpan()), + trace( + span() + .operationName(compile(quote(operation))) + .resourceName(compile(quote("websocket /test/param"))) + .type(DDSpanTypes.WEBSOCKET) + .error())); + DDSpan failedSpan = writer.get(1).get(0); + assertEquals(IllegalStateException.class.getName(), failedSpan.getTag("error.type")); + assertEquals("handler failed", failedSpan.getTag("error.message")); + assertInstanceOf(String.class, failedSpan.getTag("error.stack")); + } + + @TableTest({ + "scenario | behavior | traced", + "untraced server | SERVER | false ", + "traced client | CLIENT | true " + }) + void doesNotTraceNativeMessages(Behavior behavior, boolean traced) throws Exception { + JettyEndpoints.FullListener endpoint = new JettyEndpoints.FullListener(); + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + + openFrameHandler(frameHandler, behavior, traced); + deliver(frameHandler, new Frame(OpCode.TEXT, "hello")); + deliver(frameHandler, CloseStatus.toFrame(CloseStatus.NORMAL, "bye")); + + assertEquals(singletonList("hello"), endpoint.messages); + assertEquals(singletonList(null), endpoint.messageSpans); + assertEquals(1000, endpoint.closeCode); + assertNull(activeSpan()); + if (traced) { + assertTraces(trace(handshakeSpan())); + } else { + assertTraces(); + } + } + + @TableTest({ + "scenario | endpoint | last ", + "full | full | true ", + "partial final | partial | true ", + "partial nonfinal | partial | false", + "POJO full | pojoFull | true ", + "POJO partial nonfinal | pojoPartial | false" + }) + void binaryCallbackFailureMarksReceiveSpan(JettyEndpoints.EndpointEvents endpoint, boolean last) + throws Exception { + endpoint.failCallback = true; + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + + ExecutionException error = + assertThrows( + ExecutionException.class, + () -> deliver(frameHandler, new Frame(OpCode.BINARY, "hello").setFin(last))); + + assertEquals("callback failed", error.getCause().getMessage()); + assertNull(activeSpan()); + writer.waitForTraces(2); + assertTrue(endpoint.messageSpans.get(0).isError()); + assertEquals("callback failed", endpoint.messageSpans.get(0).getTag("error.message")); + } + + @TableTest({ + "scenario | fail ", + "success | false", + "failure | true " + }) + void deferredBinaryCallbackCompletesReceiveSpan(boolean fail) throws Exception { + JettyEndpoints.FullListener endpoint = new JettyEndpoints.FullListener(); + endpoint.deferCallback = true; + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + Callback.Completable callback = new Callback.Completable(); + openFrameHandler(frameHandler); + + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello"), callback); + + assertNotNull(endpoint.pendingCallback); + assertFalse(callback.isDone()); + assertEquals(1, writer.size()); + assertNull(activeSpan()); + + CompletableFuture completion = + CompletableFuture.runAsync( + () -> { + if (fail) { + endpoint.pendingCallback.fail(new IllegalStateException("callback failed")); + } else { + endpoint.pendingCallback.succeed(); + } + assertNull(activeSpan()); + }); + completion.get(5, SECONDS); + + assertTrue(callback.isDone()); + assertEquals(fail, callback.isCompletedExceptionally()); + writer.waitForTraces(2); + assertEquals(fail, endpoint.messageSpans.get(0).isError()); + assertNull(activeSpan()); + } + + @TableTest({ + "scenario | endpoint | reverse", + "full forward | full | false ", + "full reverse | full | true ", + "partial forward | partial | false ", + "partial reverse | partial | true ", + "POJO full forward | pojoFull | false ", + "POJO full reverse | pojoFull | true ", + "POJO partial forward | pojoPartial | false ", + "POJO partial reverse | pojoPartial | true " + }) + void pendingBinaryMessagesHaveIndependentSpans( + JettyEndpoints.EndpointEvents endpoint, boolean reverse) throws Exception { + endpoint.deferCallback = true; + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + Callback.Completable[] callbacks = {new Callback.Completable(), new Callback.Completable()}; + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello"), callbacks[0]); + org.eclipse.jetty.websocket.api.Callback first = endpoint.pendingCallback; + frameHandler.onFrame(new Frame(OpCode.BINARY, "again"), callbacks[1]); + org.eclipse.jetty.websocket.api.Callback second = endpoint.pendingCallback; + + assertNotSame(endpoint.messageSpans.get(0), endpoint.messageSpans.get(1)); + for (Callback.Completable callback : callbacks) { + assertFalse(callback.isDone()); + } + assertEquals(1, writer.size()); + assertNull(activeSpan()); + + if (reverse) { + second.fail(new IllegalStateException("second callback failed")); + } else { + first.succeed(); + } + + writer.waitForTraces(2); + assertSame(endpoint.messageSpans.get(reverse ? 1 : 0), writer.get(1).get(0)); + assertFalse(callbacks[reverse ? 0 : 1].isDone()); + + if (reverse) { + first.succeed(); + } else { + second.fail(new IllegalStateException("second callback failed")); + } + + writer.waitForTraces(3); + assertSame(endpoint.messageSpans.get(reverse ? 0 : 1), writer.get(2).get(0)); + assertFalse(endpoint.messageSpans.get(0).isError()); + assertTrue(endpoint.messageSpans.get(1).isError()); + assertEquals("second callback failed", endpoint.messageSpans.get(1).getTag("error.message")); + assertTrue(callbacks[0].isDone()); + assertFalse(callbacks[0].isCompletedExceptionally()); + assertTrue(callbacks[1].isCompletedExceptionally()); + assertNull(activeSpan()); + } + + @TableTest({ + "scenario | fail ", + "success | false", + "failure | true " + }) + void partialBinaryMessageWaitsForEarlierFragmentCallbacks(boolean fail) throws Exception { + JettyEndpoints.PartialListener endpoint = new JettyEndpoints.PartialListener(); + endpoint.deferCallback = true; + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + Callback.Completable first = new Callback.Completable(); + Callback.Completable last = new Callback.Completable(); + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello ").setFin(false), first); + org.eclipse.jetty.websocket.api.Callback firstCallback = endpoint.pendingCallback; + frameHandler.onFrame(new Frame(OpCode.CONTINUATION, "world"), last); + + endpoint.pendingCallback.succeed(); + last.get(5, SECONDS); + + assertFalse(first.isDone()); + assertEquals(1, writer.size()); + assertSame(endpoint.messageSpans.get(0), endpoint.messageSpans.get(1)); + + if (fail) { + firstCallback.fail(new IllegalStateException("first fragment failed")); + } else { + firstCallback.succeed(); + } + + assertTrue(first.isDone()); + assertEquals(fail, first.isCompletedExceptionally()); + writer.waitForTraces(2); + assertEquals(fail, endpoint.messageSpans.get(0).isError()); + if (fail) { + assertEquals("first fragment failed", endpoint.messageSpans.get(0).getTag("error.message")); + } + assertEquals(11L, endpoint.messageSpans.get(0).getTag(WEBSOCKET_MESSAGE_LENGTH)); + assertEquals(2L, endpoint.messageSpans.get(0).getTag(WEBSOCKET_MESSAGE_FRAMES)); + assertNull(activeSpan()); + } + + @Test + void handlerFailureAfterCallbackSuccessStillMarksReceiveSpan() throws Exception { + JettyEndpoints.FullListener endpoint = new JettyEndpoints.FullListener(); + endpoint.failAfterCallback = true; + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello"), new Callback.Completable()); + + writer.waitForTraces(2); + assertTrue(endpoint.messageSpans.get(0).isError()); + assertEquals("handler failed", endpoint.messageSpans.get(0).getTag("error.message")); + assertNull(activeSpan()); + } + + @TableTest({ + "scenario | msgType | closeFrame | closeHandler", + "text close frame with handler | text | true | true ", + "binary close frame with handler | binary | true | true ", + "text closed with handler | text | false | true ", + "binary closed with handler | binary | false | true ", + "text close frame no handler | text | true | false ", + "binary close frame no handler | binary | true | false ", + "text closed no handler | text | false | false ", + "binary closed no handler | binary | false | false " + }) + void terminationFinishesFragmentedReceiveSpan( + String msgType, boolean closeFrame, boolean closeHandler) throws Exception { + JettyEndpoints.EndpointEvents endpoint = + closeHandler ? new JettyEndpoints.PartialListener() : new JettyEndpoints.NoCloseEndpoint(); + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + deliver(frameHandler, new Frame(opcode(msgType), "hello").setFin(false)); + assertEquals(1, writer.size()); + + if (closeFrame) { + deliver(frameHandler, CloseStatus.toFrame(CloseStatus.NORMAL, "bye")); + } else { + Callback.Completable callback = new Callback.Completable(); + frameHandler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), callback); + callback.get(5, SECONDS); + } + + assertNull(activeSpan()); + DDSpan handshake = handshake(); + if (closeHandler) { + assertTraces( + trace(handshakeSpan()), + trace(receiveSpan(handshake, msgType, 5, 1)), + trace(closeSpan(handshake))); + } else { + assertTraces(trace(handshakeSpan()), trace(receiveSpan(handshake, msgType, 5, 1))); + } + } + + @Test + void deferredPartialBinarySuccessKeepsSpanOpenUntilFinalCallback() throws Exception { + JettyEndpoints.PartialListener endpoint = new JettyEndpoints.PartialListener(); + endpoint.deferCallback = true; + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + Callback.Completable first = new Callback.Completable(); + Callback.Completable last = new Callback.Completable(); + + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello ").setFin(false), first); + endpoint.pendingCallback.succeed(); + first.get(5, SECONDS); + frameHandler.onFrame(new Frame(OpCode.CONTINUATION, "world"), last); + + assertFalse(last.isDone()); + assertEquals(1, writer.size()); + assertSame(endpoint.messageSpans.get(0), endpoint.messageSpans.get(1)); + + endpoint.pendingCallback.succeed(); + last.get(5, SECONDS); + + assertTraces(trace(handshakeSpan()), trace(receiveSpan(handshake(), "binary", 11, 2))); + assertNull(activeSpan()); + } + + @TableTest({ + "scenario | last ", + "nonfinal | false", + "final | true " + }) + void terminationFinishesAllPendingBinaryMessages(boolean last) throws Exception { + JettyEndpoints.PartialListener endpoint = new JettyEndpoints.PartialListener(); + endpoint.deferCallback = true; + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + Callback.Completable previousMessage = new Callback.Completable(); + frameHandler.onFrame(new Frame(OpCode.BINARY, "previous"), previousMessage); + org.eclipse.jetty.websocket.api.Callback previousCallback = endpoint.pendingCallback; + Callback.Completable message = new Callback.Completable(); + frameHandler.onFrame(new Frame(OpCode.BINARY, "hello").setFin(last), message); + + assertEquals(1, writer.size()); + assertNotSame(endpoint.messageSpans.get(0), endpoint.messageSpans.get(1)); + + Callback.Completable closed = new Callback.Completable(); + frameHandler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), closed); + closed.get(5, SECONDS); + + writer.waitForTraces(4); + assertTrue( + writer.stream().flatMap(List::stream).collect(toList()).containsAll(endpoint.messageSpans)); + + endpoint.pendingCallback.succeed(); + message.get(5, SECONDS); + previousCallback.succeed(); + previousMessage.get(5, SECONDS); + + assertEquals(4, writer.size()); + assertNull(activeSpan()); + } + + @TypeConverter + public static JettyEndpoints.EndpointEvents endpoint(String name) { + switch (name) { + case "full": + return new JettyEndpoints.FullListener(); + case "partial": + return new JettyEndpoints.PartialListener(); + case "pojoFull": + return new JettyEndpoints.PojoFullEndpoint(); + case "pojoPartial": + return new JettyEndpoints.PojoPartialEndpoint(); + default: + throw new IllegalArgumentException("Unknown endpoint: " + name); + } + } + + private static byte opcode(String msgType) { + return "text".equals(msgType) ? OpCode.TEXT : OpCode.BINARY; + } + + private static JettyWebSocketFrameHandler createFrameHandler(Object endpoint) { + ServerFrameHandlerFactory factory = + new ServerFrameHandlerFactory(mock(WebSocketContainer.class), new WebSocketComponents()); + JettyWebSocketFrameHandler frameHandler = factory.newJettyFrameHandler(endpoint); + UpgradeRequest request = mock(UpgradeRequest.class); + when(request.getRequestURI()).thenReturn(URI.create(URL)); + frameHandler.setUpgradeRequest(request); + return frameHandler; + } + + private static void openFrameHandler(JettyWebSocketFrameHandler frameHandler) throws Exception { + openFrameHandler(frameHandler, Behavior.SERVER, true); + } + + private static void openFrameHandler( + JettyWebSocketFrameHandler frameHandler, Behavior connectionBehavior, boolean traced) + throws Exception { + CoreSession session = + new CoreSession.Empty() { + @Override + public Behavior getBehavior() { + return connectionBehavior; + } + }; + Callback.Completable openCallback = new Callback.Completable(); + if (traced) { + runUnderTrace( + "parent", + () -> { + activeSpan().setTag(HTTP_URL, URL); + frameHandler.onOpen(session, openCallback); + openCallback.get(5, SECONDS); + return null; + }); + } else { + frameHandler.onOpen(session, openCallback); + openCallback.get(5, SECONDS); + } + } + + private static void deliver(JettyWebSocketFrameHandler frameHandler, Frame frame) + throws Exception { + Callback.Completable callback = new Callback.Completable(); + frameHandler.onFrame(frame, callback); + callback.get(5, SECONDS); + } + + private static DDSpan handshake() throws Exception { + writer.waitForTraces(1); + return writer.get(0).get(0); + } + + private static SpanMatcher handshakeSpan() { + return span() + .root() + .operationName(compile(quote("parent"))) + .resourceName(compile(quote("/test/param"))) + .tags(defaultTags(), tag(HTTP_URL, is(URL))); + } + + private static SpanMatcher receiveSpan( + DDSpan handshake, String msgType, long length, long frames) { + return websocketSpan( + handshake, + "websocket.receive", + tag(WEBSOCKET_MESSAGE_TYPE, matches(quote(msgType))), + tag(WEBSOCKET_MESSAGE_LENGTH, is(length)), + tag(WEBSOCKET_MESSAGE_FRAMES, is(frames)), + // Full-message callbacks do not record a receive-time tag. + includes(WEBSOCKET_MESSAGE_RECEIVE_TIME)); + } + + private static SpanMatcher closeSpan(DDSpan handshake) { + return websocketSpan( + handshake, + "websocket.close", + tag(WEBSOCKET_MESSAGE_TYPE, isNull()), + tag(WEBSOCKET_MESSAGE_LENGTH, isNull()), + tag(WEBSOCKET_MESSAGE_FRAMES, isNull()), + tag(WEBSOCKET_CLOSE_CODE, is(1000)), + tag(WEBSOCKET_CLOSE_REASON, is("bye"))); + } + + private static SpanMatcher websocketSpan( + DDSpan handshake, String operation, TagsMatcher... extraTags) { + List tags = new ArrayList<>(asList(extraTags)); + tags.add(defaultTags()); + tags.add(tag("span.kind", is("consumer"))); + tags.add(tag("component", matches("websocket"))); + tags.add(tag("peer.hostname", isNull())); + tags.add(tag(DECISION_MAKER_INHERITED, is(1))); + tags.add(tag(DECISION_MAKER_SERVICE, is(handshake.getServiceName()))); + tags.add(tag(DECISION_MAKER_RESOURCE, matches(quote(handshake.getResourceName().toString())))); + byte flags = + handshake.getSamplingPriority() > 0 + ? AgentSpanLink.SAMPLED_FLAG + : AgentSpanLink.DEFAULT_FLAGS; + return span() + .root() + .operationName(compile(quote(operation))) + .resourceName(compile(quote("websocket /test/param"))) + .type(DDSpanTypes.WEBSOCKET) + .links( + SpanLinkMatcher.to(handshake) + .traceFlags(flags) + .attributes(SpanAttributes.builder().put("dd.kind", "executed_from").build())) + .tags(tags.toArray(new TagsMatcher[0])); + } +} From c78ea6e674a492fb104233ece670c922b4d1fd2e Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Mon, 28 Sep 2026 20:05:45 -0700 Subject: [PATCH 03/23] Refactor Jetty 12 WebSocket advice into instrumentation classes and simplify the build configuration --- .../jetty-client-12.0/build.gradle | 4 -- .../JettyWebSocketUpgradeInstrumentation.java | 21 ++++++- .../WebSocketUpgradeAdvice.java | 23 ------- .../jetty-websocket-12.0/build.gradle | 10 +-- .../jetty-websocket-12.0/gradle.lockfile | 36 +++++------ .../jetty12/Jetty12NativeWebSocketModule.java | 61 +++++++++++++++++-- .../jetty12/NativeMethodHandleWrappers.java | 0 .../jetty12/NativeWebSocketAdvices.java | 59 ------------------ 8 files changed, 96 insertions(+), 118 deletions(-) delete mode 100644 dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/WebSocketUpgradeAdvice.java rename dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/{java17 => java}/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java (100%) delete mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeWebSocketAdvices.java diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle index 947bcc529b0..6d8131826da 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle @@ -18,10 +18,6 @@ tracerJava { addTestSuiteForDir('latestDepTest', 'test') -tasks.named("compileMain_java17Java", JavaCompile) { - configureCompiler(it, JavaVersion.VERSION_17) -} - configurations.matching { it.name.startsWith('test') || it.name.startsWith('latestDepTest') }.configureEach { it.resolutionStrategy { force group: 'org.slf4j', name: 'slf4j-api', version: libs.versions.slf4j.get() diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java index 03eb291ed66..a9143488457 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.jetty_client12; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.instrumentation.jetty_client12.JettyClientDecorator.DECORATE; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; @@ -8,8 +9,12 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.util.Map; +import net.bytebuddy.asm.Advice; +import org.eclipse.jetty.client.Request; +import org.eclipse.jetty.client.Response; @AutoService(InstrumenterModule.class) public class JettyWebSocketUpgradeInstrumentation extends InstrumenterModule.Tracing @@ -40,6 +45,20 @@ public void methodAdvice(MethodTransformer transformer) { .and(takesArguments(2)) .and(takesArgument(0, named("org.eclipse.jetty.client.Response"))) .and(takesArgument(1, named("org.eclipse.jetty.io.EndPoint"))), - packageName + ".WebSocketUpgradeAdvice"); + getClass().getName() + "$WebSocketUpgradeAdvice"); + } + + public static class WebSocketUpgradeAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void afterUpgrade(@Advice.Argument(0) Response response) { + AgentSpan span = + InstrumentationContext.get(Request.class, AgentSpan.class).get(response.getRequest()); + if (span != null) { + // Successful upgrades bypass the request's response completion listeners. + DECORATE.onResponse(span, response); + DECORATE.beforeFinish(span); + span.finish(); + } + } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/WebSocketUpgradeAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/WebSocketUpgradeAdvice.java deleted file mode 100644 index d051ae0f708..00000000000 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/WebSocketUpgradeAdvice.java +++ /dev/null @@ -1,23 +0,0 @@ -package datadog.trace.instrumentation.jetty_client12; - -import static datadog.trace.instrumentation.jetty_client12.JettyClientDecorator.DECORATE; - -import datadog.trace.bootstrap.InstrumentationContext; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import net.bytebuddy.asm.Advice; -import org.eclipse.jetty.client.Request; -import org.eclipse.jetty.client.Response; - -public class WebSocketUpgradeAdvice { - @Advice.OnMethodExit(suppress = Throwable.class) - public static void afterUpgrade(@Advice.Argument(0) Response response) { - AgentSpan span = - InstrumentationContext.get(Request.class, AgentSpan.class).get(response.getRequest()); - if (span != null) { - // Successful upgrades bypass the request's response completion listeners. - DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); - span.finish(); - } - } -} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle index 7af4629a730..c11ad0c6812 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle @@ -37,14 +37,6 @@ testJvmConstraints { minJavaVersion = JavaVersion.VERSION_17 } -tracerJava { - addSourceSetFor(JavaVersion.VERSION_17) -} - -tasks.named("compileMain_java17Java", JavaCompile) { - configureCompiler(it, 17, JavaVersion.VERSION_1_8, "Uses Jetty 12 APIs (Java 17+) at Java 8 bytecode") -} - addTestSuiteForDir("latestDepTest", "test") ["compileTestGroovy", "compileLatestDepTestGroovy"].each { name -> @@ -61,7 +53,7 @@ addTestSuiteForDir("latestDepTest", "test") dependencies { testImplementation libs.bundles.mockito - main_java17CompileOnly 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0' + compileOnly 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0' implementation project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-10.0") testImplementation group: 'org.eclipse.jetty.ee8.websocket', name: 'jetty-ee8-websocket-javax-server', version: '12.0.0' testImplementation group: 'org.eclipse.jetty.ee9.websocket', name: 'jetty-ee9-websocket-jakarta-server', version: '12.0.0' diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile index 7e31f1cf65c..216f0d87215 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile @@ -6,11 +6,11 @@ cafe.cryptography:curve25519-elisabeth:0.1.0=latestDepTestRuntimeClasspath,testR cafe.cryptography:ed25519-elisabeth:0.1.0=latestDepTestRuntimeClasspath,testRuntimeClasspath ch.qos.logback:logback-classic:1.2.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath ch.qos.logback:logback-core:1.2.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.blogspot.mydailyjava:weak-lock-free:0.17=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +com.blogspot.mydailyjava:weak-lock-free:0.17=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath com.datadoghq.okhttp3:okhttp:3.12.15=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath com.datadoghq.okio:okio:1.17.6=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.datadoghq:dd-instrument-java:0.0.5=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath -com.datadoghq:dd-javac-plugin-client:0.2.2=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +com.datadoghq:dd-instrument-java:0.0.5=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +com.datadoghq:dd-javac-plugin-client:0.2.2=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath com.datadoghq:java-dogstatsd-client:4.4.5=latestDepTestRuntimeClasspath,testRuntimeClasspath com.datadoghq:sketches-java:0.8.3=latestDepTestRuntimeClasspath,testRuntimeClasspath com.github.javaparser:javaparser-core:3.25.6=codenarc @@ -68,8 +68,8 @@ jakarta.websocket:jakarta.websocket-client-api:2.1.1=latestDepTestCompileClasspa javax.servlet:javax.servlet-api:3.1.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath jaxen:jaxen:2.0.6=spotbugs junit:junit:4.13.2=latestDepTestRuntimeClasspath,testRuntimeClasspath -net.bytebuddy:byte-buddy-agent:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath -net.bytebuddy:byte-buddy:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +net.bytebuddy:byte-buddy-agent:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +net.bytebuddy:byte-buddy:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath net.java.dev.jna:jna-platform:5.8.0=latestDepTestRuntimeClasspath,testRuntimeClasspath net.java.dev.jna:jna:5.8.0=latestDepTestRuntimeClasspath,testRuntimeClasspath net.sf.saxon:Saxon-HE:12.10=spotbugs @@ -155,37 +155,37 @@ org.eclipse.jetty.toolchain:jetty-javax-websocket-api:1.1.2=latestDepTestCompile org.eclipse.jetty.toolchain:jetty-servlet-api:4.0.6=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-client:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-jetty-api:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-api:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-jetty-api:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-alpn-client:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-alpn-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-client:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-ee:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty:jetty-http:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty:jetty-http:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-http:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty:jetty-io:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty:jetty-io:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-io:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-jndi:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-jndi:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-plus:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-security:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-security:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty:jetty-server:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty:jetty-server:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-session:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-session:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty:jetty-util:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty:jetty-util:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-util:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-xml:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-xml:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath @@ -220,16 +220,16 @@ org.ow2.asm:asm-tree:9.5=testCompileClasspath org.ow2.asm:asm-tree:9.7.1=latestDepTestCompileClasspath org.ow2.asm:asm-util:9.10.1=spotbugs org.ow2.asm:asm-util:9.7.1=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.ow2.asm:asm:9.10.1=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleTooling,runtimeClasspath,spotbugs,testCompileClasspath,testRuntimeClasspath +org.ow2.asm:asm:9.10.1=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,spotbugs,testCompileClasspath,testRuntimeClasspath org.slf4j:jcl-over-slf4j:1.7.30=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.slf4j:jul-to-slf4j:1.7.30=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.slf4j:log4j-over-slf4j:1.7.30=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.slf4j:slf4j-api:1.7.30=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,main_java17CompileClasspath,main_java17RuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,spotbugs,spotbugsSlf4j,testCompileClasspath,testRuntimeClasspath +org.slf4j:slf4j-api:1.7.30=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,spotbugs,spotbugsSlf4j,testCompileClasspath,testRuntimeClasspath org.slf4j:slf4j-simple:2.0.17=spotbugsSlf4j -org.snakeyaml:snakeyaml-engine:2.9=buildTimeInstrumentationPlugin,latestDepTestRuntimeClasspath,main_java17RuntimeClasspath,muzzleTooling,runtimeClasspath,testRuntimeClasspath +org.snakeyaml:snakeyaml-engine:2.9=buildTimeInstrumentationPlugin,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testRuntimeClasspath org.spockframework:spock-bom:2.4-groovy-3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.spockframework:spock-core:2.4-groovy-3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.tabletest:tabletest-junit:1.2.2=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.tabletest:tabletest-parser:1.2.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.xmlresolver:xmlresolver:5.3.3=spotbugs -empty=main_java17AnnotationProcessor,spotbugsPlugins +empty=spotbugsPlugins diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java index 0376df38a67..0d4ff8dad96 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.websocket.jetty12; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; @@ -9,7 +10,15 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.instrumentation.websocket.jetty12.NativeMethodHandleWrappers.ReceiveContexts; +import java.lang.invoke.MethodHandle; import java.util.Map; +import net.bytebuddy.asm.Advice; +import org.eclipse.jetty.websocket.common.WebSocketSession; +import org.eclipse.jetty.websocket.core.Behavior; +import org.eclipse.jetty.websocket.core.CoreSession; @AutoService(InstrumenterModule.class) public class Jetty12NativeWebSocketModule extends InstrumenterModule.Tracing @@ -56,12 +65,56 @@ public void methodAdvice(MethodTransformer transformer) { named("onOpen") .and(takesArguments(2)) .and(takesArgument(0, named("org.eclipse.jetty.websocket.core.CoreSession"))), - packageName + ".NativeWebSocketAdvices$OpenAdvice"); + getClass().getName() + "$OpenAdvice"); transformer.applyAdvice( named("createMessageSink").and(takesArguments(4)), - packageName + ".NativeWebSocketAdvices$MessageSinkAdvice"); + getClass().getName() + "$MessageSinkAdvice"); transformer.applyAdvice( - named("notifyOnClose").and(takesArguments(2)), - packageName + ".NativeWebSocketAdvices$CloseAdvice"); + named("notifyOnClose").and(takesArguments(2)), getClass().getName() + "$CloseAdvice"); + } + + public static class OpenAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter( + @Advice.Argument(0) CoreSession session, + @Advice.FieldValue(value = "closeHandle", readOnly = false) MethodHandle closeHandle) { + AgentSpan span = activeSpan(); + if (span != null && session.getBehavior() == Behavior.SERVER && closeHandle != null) { + closeHandle = NativeMethodHandleWrappers.wrapClose(closeHandle, span, session); + } + } + } + + public static class MessageSinkAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter( + @Advice.Argument(1) WebSocketSession session, + @Advice.Argument(value = 2, readOnly = false) MethodHandle handle) { + AgentSpan span = activeSpan(); + if (span != null + && session.getCoreSession().getBehavior() == Behavior.SERVER + && handle != null) { + handle = + NativeMethodHandleWrappers.wrapMessage( + handle, + span, + session.getCoreSession(), + InstrumentationContext.get(CoreSession.class, ReceiveContexts.class)); + } + } + } + + public static class CloseAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter(@Advice.FieldValue("session") WebSocketSession session) { + if (session != null) { + ReceiveContexts contexts = + InstrumentationContext.get(CoreSession.class, ReceiveContexts.class) + .remove(session.getCoreSession()); + if (contexts != null) { + contexts.finish(); + } + } + } } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java similarity index 100% rename from dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java rename to dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeWebSocketAdvices.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeWebSocketAdvices.java deleted file mode 100644 index cc921cb1645..00000000000 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java17/datadog/trace/instrumentation/websocket/jetty12/NativeWebSocketAdvices.java +++ /dev/null @@ -1,59 +0,0 @@ -package datadog.trace.instrumentation.websocket.jetty12; - -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; - -import datadog.trace.bootstrap.InstrumentationContext; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import datadog.trace.instrumentation.websocket.jetty12.NativeMethodHandleWrappers.ReceiveContexts; -import java.lang.invoke.MethodHandle; -import net.bytebuddy.asm.Advice; -import org.eclipse.jetty.websocket.common.WebSocketSession; -import org.eclipse.jetty.websocket.core.Behavior; -import org.eclipse.jetty.websocket.core.CoreSession; - -public class NativeWebSocketAdvices { - public static class OpenAdvice { - @Advice.OnMethodEnter(suppress = Throwable.class) - public static void onEnter( - @Advice.Argument(0) CoreSession session, - @Advice.FieldValue(value = "closeHandle", readOnly = false) MethodHandle closeHandle) { - AgentSpan span = activeSpan(); - if (span != null && session.getBehavior() == Behavior.SERVER && closeHandle != null) { - closeHandle = NativeMethodHandleWrappers.wrapClose(closeHandle, span, session); - } - } - } - - public static class MessageSinkAdvice { - @Advice.OnMethodEnter(suppress = Throwable.class) - public static void onEnter( - @Advice.Argument(1) WebSocketSession session, - @Advice.Argument(value = 2, readOnly = false) MethodHandle handle) { - AgentSpan span = activeSpan(); - if (span != null - && session.getCoreSession().getBehavior() == Behavior.SERVER - && handle != null) { - handle = - NativeMethodHandleWrappers.wrapMessage( - handle, - span, - session.getCoreSession(), - InstrumentationContext.get(CoreSession.class, ReceiveContexts.class)); - } - } - } - - public static class CloseAdvice { - @Advice.OnMethodEnter(suppress = Throwable.class) - public static void onEnter(@Advice.FieldValue("session") WebSocketSession session) { - if (session != null) { - ReceiveContexts contexts = - InstrumentationContext.get(CoreSession.class, ReceiveContexts.class) - .remove(session.getCoreSession()); - if (contexts != null) { - contexts.finish(); - } - } - } - } -} From e5375ba6b39308d7e72a9b149f3994c188391297 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 29 Sep 2026 10:12:28 -0700 Subject: [PATCH 04/23] Fix Jetty 12 WebSocket tracing for boxed-Void handlers Normalize text, binary, and close delegate return types to primitive void. Simplify binary state synchronization under ReceiveContexts, preserve separate text receiver locking, and resolve SpotBugs synchronization warnings. Replace the binary template receiver with immutable handshake metadata. Narrow the production dependency to jetty-websocket-jetty-common. Add boxed-Void handler coverage and real-server receive tests verifying links to the HTTP handshake span. --- .../jetty-websocket-12.0/build.gradle | 3 +- .../jetty12/NativeMethodHandleWrappers.java | 75 +++++++------- .../src/test/java/JettyEndpoints.java | 60 ++++++++++++ .../test/java/JettyWebSocketUpgradeTest.java | 97 ++++++++++++++++++- .../src/test/java/JettyWebsocketTest.java | 38 +++++--- 5 files changed, 219 insertions(+), 54 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle index c11ad0c6812..e24844fce66 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle @@ -53,7 +53,7 @@ addTestSuiteForDir("latestDepTest", "test") dependencies { testImplementation libs.bundles.mockito - compileOnly 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0' + compileOnly 'org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.0.0' implementation project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-10.0") testImplementation group: 'org.eclipse.jetty.ee8.websocket', name: 'jetty-ee8-websocket-javax-server', version: '12.0.0' testImplementation group: 'org.eclipse.jetty.ee9.websocket', name: 'jetty-ee9-websocket-jakarta-server', version: '12.0.0' @@ -69,6 +69,7 @@ dependencies { latestDepTestImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.17' testRuntimeOnly project(":dd-java-agent:instrumentation:jetty:jetty-client:jetty-client-12.0") + testRuntimeOnly project(":dd-java-agent:instrumentation:jetty:jetty-server:jetty-server-12.0") testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-10.0") testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:jetty-websocket:jetty-websocket-11.0") testRuntimeOnly project(":dd-java-agent:instrumentation:websocket:javax-websocket-1.0") diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java index dbb1c58f118..5aa2811161c 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java @@ -75,21 +75,21 @@ public static MethodHandle wrapMessage( return delegate; } boolean partial = type.parameterCount() > 1 && type.parameterType(1) == boolean.class; + delegate = delegate.asType(type.changeReturnType(void.class)); MethodHandle normalized = partial ? delegate : dropArguments(delegate, 1, boolean.class); ReceiveContexts contexts = contextStore.get(session); if (contexts == null) { - contexts = contextStore.getOrPut(session, new ReceiveContexts()); + contexts = contextStore.getOrPut(session, new ReceiveContexts(span, session)); } - HandlerContext.Receiver context = context(span, session); - contexts.add(context, payload == String.class); MethodHandle wrapper = payload == String.class - ? insertArguments(TEXT, 0, normalized, context, partial) + ? contexts.wrapText(normalized, partial) : insertArguments(BINARY, 0, normalized, contexts, partial); return partial ? wrapper : insertArguments(wrapper, 1, true); } public static MethodHandle wrapClose(MethodHandle delegate, AgentSpan span, CoreSession session) { + delegate = delegate.asType(delegate.type().changeReturnType(void.class)); // Annotated endpoints have a Session argument; listeners do not. Jetty binds it during onOpen. if (delegate.type().parameterType(0) != Session.class) { delegate = dropArguments(delegate, 0, Session.class); @@ -100,9 +100,7 @@ public static MethodHandle wrapClose(MethodHandle delegate, AgentSpan span, Core private static ContextScope startMessage( HandlerContext.Receiver context, Object data, boolean partial) { try { - synchronized (context) { - return activateSpan(DECORATE.startInboundFrameSpan(context, data, partial)); - } + return activateSpan(DECORATE.startInboundFrameSpan(context, data, partial)); } catch (Throwable t) { ExceptionLogger.LOGGER.debug("Unable to start native Jetty WebSocket span", t); return null; @@ -117,7 +115,11 @@ public static void onText( boolean last) throws Throwable { boolean finish = last; - try (ContextScope ignored = startMessage(context, payload, partial)) { + ContextScope scope; + synchronized (context) { + scope = startMessage(context, payload, partial); + } + try (ContextScope ignored = scope) { try { delegate.invokeExact(payload, last); } catch (Throwable t) { @@ -147,7 +149,15 @@ public static void onBinary( ReceiveCallback wrapped; ContextScope scope; synchronized (contexts) { - BinaryMessage message = contexts.startBinaryMessage(last); + BinaryMessage message = contexts.currentBinary; + if (message == null) { + message = new BinaryMessage(contexts.handshakeSpan, contexts.sessionId); + contexts.pendingBinary.add(message); + } + message.pendingCallbacks++; + message.complete = last; + // Jetty can deliver the next message before this message's callbacks complete. + contexts.currentBinary = last ? null : message; scope = startMessage(message, payload, partial); wrapped = new ReceiveCallback(callback, contexts, message); } @@ -164,56 +174,47 @@ public static void onBinary( } public static class ReceiveContexts { + private final AgentSpan handshakeSpan; + private final String sessionId; private HandlerContext.Receiver text; - private HandlerContext.Receiver binary; private BinaryMessage currentBinary; private final Set pendingBinary = new HashSet<>(); - public void add(HandlerContext.Receiver context, boolean isText) { - if (isText) { - text = context; - } else { - binary = context; + public ReceiveContexts(AgentSpan span, CoreSession session) { + if (Config.get().isWebsocketMessagesInheritSampling()) { + span.forceSamplingDecision(); } + handshakeSpan = span.getLocalRootSpan(); + sessionId = Integer.toHexString(System.identityHashCode(session)); } - public synchronized BinaryMessage startBinaryMessage(boolean last) { - BinaryMessage message = currentBinary; - if (message == null) { - message = new BinaryMessage(binary); - pendingBinary.add(message); + private synchronized MethodHandle wrapText(MethodHandle delegate, boolean partial) { + if (text == null) { + text = new HandlerContext.Receiver(handshakeSpan, sessionId); } - message.pendingCallbacks++; - message.complete = last; - // Jetty can deliver the next message before this message's callbacks complete. - currentBinary = last ? null : message; - return message; + return insertArguments(TEXT, 0, delegate, text, partial); } public synchronized void finish() { - finish(text); + if (text != null) { + synchronized (text) { + DECORATE.onFrameEnd(text); + } + } for (BinaryMessage message : pendingBinary) { - finish(message); + DECORATE.onFrameEnd(message); } pendingBinary.clear(); currentBinary = null; } - - private static void finish(HandlerContext.Receiver context) { - if (context != null) { - synchronized (context) { - DECORATE.onFrameEnd(context); - } - } - } } public static class BinaryMessage extends HandlerContext.Receiver { private int pendingCallbacks; private boolean complete; - public BinaryMessage(HandlerContext.Receiver context) { - super(context.getHandshakeSpan(), context.getSessionId()); + public BinaryMessage(AgentSpan handshakeSpan, String sessionId) { + super(handshakeSpan, sessionId); } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java index 96c570daa89..89296e029c4 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java @@ -143,6 +143,66 @@ public void onBinary(Session session, ByteBuffer payload, boolean last, Callback } } + @WebSocket + public static class BoxedFullEndpoint extends EndpointEvents { + @OnWebSocketOpen + public Void onOpen(Session session) { + this.session = session; + return null; + } + + @OnWebSocketMessage + public Void onText(Session session, String payload) { + assertSame(this.session, session); + recordMessage(payload); + return null; + } + + @OnWebSocketMessage + public Void onBinary(ByteBuffer payload, Callback callback) { + recordMessage(UTF_8.decode(payload).toString()); + completeBinary(callback); + return null; + } + + @OnWebSocketClose + public Void onClose(int statusCode, String reason) { + recordClose(statusCode, reason); + return null; + } + } + + @WebSocket + public static class BoxedPartialEndpoint extends EndpointEvents { + @OnWebSocketOpen + public void onOpen(Session session) { + this.session = session; + } + + @OnWebSocketMessage + public Void onText(String payload, boolean last) { + finalFragments.add(last); + recordMessage(payload); + return null; + } + + @OnWebSocketMessage + public Void onBinary(Session session, ByteBuffer payload, boolean last, Callback callback) { + assertSame(this.session, session); + finalFragments.add(last); + recordMessage(UTF_8.decode(payload).toString()); + completeBinary(callback); + return null; + } + + @OnWebSocketClose + public Void onClose(Session session, int statusCode, String reason) { + assertSame(this.session, session); + recordClose(statusCode, reason); + return null; + } + } + @WebSocket public static class NoCloseEndpoint extends EndpointEvents { @OnWebSocketMessage diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java index 96025ee2e10..4e0cf65b704 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java @@ -1,22 +1,33 @@ +import static datadog.trace.agent.test.assertions.SpanLinkMatcher.to; import static datadog.trace.agent.test.assertions.SpanMatcher.span; import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_LENGTH; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_TYPE; +import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.Collections.singletonList; import static java.util.concurrent.TimeUnit.SECONDS; import static java.util.regex.Pattern.compile; import static java.util.regex.Pattern.quote; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import datadog.trace.agent.test.AbstractInstrumentationTest; import datadog.trace.api.DDSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink; +import datadog.trace.bootstrap.instrumentation.api.SpanAttributes; import datadog.trace.core.DDSpan; import java.net.URI; +import java.util.List; import org.eclipse.jetty.server.Server; import org.eclipse.jetty.server.ServerConnector; import org.eclipse.jetty.server.handler.ContextHandler; +import org.eclipse.jetty.websocket.api.Callback; import org.eclipse.jetty.websocket.api.Session; import org.eclipse.jetty.websocket.client.WebSocketClient; import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler; import org.junit.jupiter.api.Test; +import org.tabletest.junit.TableTest; class JettyWebSocketUpgradeTest extends AbstractInstrumentationTest { @Test @@ -51,8 +62,9 @@ void httpClientSpanFinishesWhenWebSocketUpgradeSucceeds() throws Exception { .resourceName(compile(quote("GET /upgrade"))) .type(DDSpanTypes.HTTP_CLIENT) .root() - .error(false))); - DDSpan handshake = writer.get(0).get(0); + .error(false)), + trace(span().type(DDSpanTypes.HTTP_SERVER).error(false))); + DDSpan handshake = handshake("client"); assertEquals("jetty-client", handshake.getTag("component").toString()); assertEquals("client", handshake.getTag("span.kind")); assertEquals(101, handshake.getTag("http.status_code")); @@ -65,5 +77,86 @@ void httpClientSpanFinishesWhenWebSocketUpgradeSucceeds() throws Exception { } } + @TableTest({ + "scenario | endpointName | messageType", + "listener text | full | text ", + "listener binary | full | binary ", + "boxed text | boxedFull | text ", + "boxed binary | boxedFull | binary " + }) + void receivedMessageLinksToServerHandshake(String endpointName, String messageType) + throws Exception { + JettyEndpoints.EndpointEvents endpoint = JettyWebsocketTest.endpoint(endpointName); + Server server = new Server(0); + WebSocketClient client = new WebSocketClient(); + try { + ContextHandler context = new ContextHandler("/"); + server.setHandler(context); + context.setHandler( + WebSocketUpgradeHandler.from(server, context) + .configure( + container -> + container.addMapping("/receive", (request, response, callback) -> endpoint))); + server.start(); + client.start(); + URI uri = + URI.create( + "ws://localhost:" + + ((ServerConnector) server.getConnectors()[0]).getLocalPort() + + "/receive"); + Session session = client.connect(new Endpoint(), uri).get(5, SECONDS); + writer.waitForTraces(2); + DDSpan handshake = handshake("server"); + assertEquals(101, handshake.getTag("http.status_code")); + + Callback.Completable sent = new Callback.Completable(); + if ("text".equals(messageType)) { + session.sendText("hello", sent); + } else { + session.sendBinary(UTF_8.encode("hello"), sent); + } + sent.get(5, SECONDS); + + assertTraces( + trace(span().type(DDSpanTypes.HTTP_CLIENT).error(false)), + trace(span().type(DDSpanTypes.HTTP_SERVER).error(false)), + trace( + span() + .root() + .operationName(compile(quote("websocket.receive"))) + .resourceName(compile(quote("websocket /receive"))) + .type(DDSpanTypes.WEBSOCKET) + .error(false) + .links( + to(handshake) + .traceFlags( + handshake.getSamplingPriority() > 0 + ? AgentSpanLink.SAMPLED_FLAG + : AgentSpanLink.DEFAULT_FLAGS) + .attributes( + SpanAttributes.builder().put("dd.kind", "executed_from").build())))); + assertEquals(singletonList("hello"), endpoint.messages); + assertEquals(1, endpoint.messageSpans.size()); + assertSame(writer.get(2).get(0), endpoint.messageSpans.get(0)); + assertEquals( + messageType, endpoint.messageSpans.get(0).getTag(WEBSOCKET_MESSAGE_TYPE).toString()); + assertEquals(5L, endpoint.messageSpans.get(0).getTag(WEBSOCKET_MESSAGE_LENGTH)); + } finally { + try { + client.stop(); + } finally { + server.stop(); + } + } + } + + private static DDSpan handshake(String kind) { + return writer.stream() + .flatMap(List::stream) + .filter(s -> kind.equals(s.getTag("span.kind"))) + .findFirst() + .orElseThrow(() -> new AssertionError("Missing " + kind + " handshake span")); + } + public static class Endpoint implements Session.Listener.AutoDemanding {} } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java index 8bd563b3464..f83acdc7ae1 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java @@ -68,15 +68,19 @@ public class JettyWebsocketTest extends AbstractInstrumentationTest { private static final String URL = "ws://inmemory/test/param"; @TableTest({ - "scenario | endpoint | msgType", - "full text | full | text ", - "full binary | full | binary ", - "partial text | partial | text ", - "partial binary | partial | binary ", - "POJO full text | pojoFull | text ", - "POJO full binary | pojoFull | binary ", - "POJO partial text | pojoPartial | text ", - "POJO partial binary | pojoPartial | binary " + "scenario | endpoint | msgType", + "full text | full | text ", + "full binary | full | binary ", + "partial text | partial | text ", + "partial binary | partial | binary ", + "POJO full text | pojoFull | text ", + "POJO full binary | pojoFull | binary ", + "POJO partial text | pojoPartial | text ", + "POJO partial binary | pojoPartial | binary ", + "boxed full text | boxedFull | text ", + "boxed full binary | boxedFull | binary ", + "boxed partial text | boxedPartial | text ", + "boxed partial binary | boxedPartial | binary " }) void nativeJettyAdvices(JettyEndpoints.EndpointEvents endpoint, String msgType) throws Exception { JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); @@ -104,11 +108,13 @@ void nativeJettyAdvices(JettyEndpoints.EndpointEvents endpoint, String msgType) } @TableTest({ - "scenario | endpoint | msgType", - "partial text | partial | text ", - "partial binary | partial | binary ", - "POJO partial text | pojoPartial | text ", - "POJO partial binary | pojoPartial | binary " + "scenario | endpoint | msgType", + "partial text | partial | text ", + "partial binary | partial | binary ", + "POJO partial text | pojoPartial | text ", + "POJO partial binary | pojoPartial | binary ", + "boxed partial text | boxedPartial | text ", + "boxed partial binary | boxedPartial | binary " }) void fragmentedMessagesShareSpan(JettyEndpoints.EndpointEvents endpoint, String msgType) throws Exception { @@ -487,6 +493,10 @@ public static JettyEndpoints.EndpointEvents endpoint(String name) { return new JettyEndpoints.PojoFullEndpoint(); case "pojoPartial": return new JettyEndpoints.PojoPartialEndpoint(); + case "boxedFull": + return new JettyEndpoints.BoxedFullEndpoint(); + case "boxedPartial": + return new JettyEndpoints.BoxedPartialEndpoint(); default: throw new IllegalArgumentException("Unknown endpoint: " + name); } From 7c4584ca6cad7e42e5ff1f4c616aa1a5a9d345ff Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 29 Sep 2026 13:17:35 -0700 Subject: [PATCH 05/23] Move Jetty WebSocket upgrade test into client instrumentation module --- .../jetty-client-12.0/build.gradle | 5 ++ .../jetty-client-12.0/gradle.lockfile | 16 ++++ .../test/java/JettyWebSocketUpgradeTest.java | 76 +++++++++++++++++++ ....java => JettyWebSocketHandshakeTest.java} | 51 +------------ 4 files changed, 98 insertions(+), 50 deletions(-) create mode 100644 dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/test/java/JettyWebSocketUpgradeTest.java rename dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/{JettyWebSocketUpgradeTest.java => JettyWebSocketHandshakeTest.java} (72%) diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle index 6d8131826da..3f7550a3b20 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle @@ -36,5 +36,10 @@ dependencies { } testImplementation project(':dd-java-agent:instrumentation:jetty:jetty-util-9.4.31') testImplementation group: 'org.eclipse.jetty', name: 'jetty-client', version: '12.0.0' + testImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.0' + testImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0' + testRuntimeOnly project(':dd-java-agent:instrumentation:jetty:jetty-server:jetty-server-12.0') latestDepTestImplementation group: 'org.eclipse.jetty', name: 'jetty-client', version: '12.+' + latestDepTestImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.+' + latestDepTestImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.+' } diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/gradle.lockfile b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/gradle.lockfile index 6af252614e2..fc2bf9e655b 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/gradle.lockfile +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/gradle.lockfile @@ -82,6 +82,20 @@ org.codenarc:CodeNarc:3.7.0=codenarc org.dom4j:dom4j:2.2.0=spotbugs org.eclipse.jetty.compression:jetty-compression-common:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.compression:jetty-compression-gzip:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-client:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-client:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-common:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-server:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-api:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-api:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-alpn-client:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-alpn-client:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-client:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath @@ -90,6 +104,8 @@ org.eclipse.jetty:jetty-http:12.0.0=main_java17CompileClasspath,testCompileClass org.eclipse.jetty:jetty-http:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-io:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-io:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-server:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty:jetty-server:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-util:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty:jetty-util:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.gmetrics:GMetrics:2.1.0=codenarc diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/test/java/JettyWebSocketUpgradeTest.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/test/java/JettyWebSocketUpgradeTest.java new file mode 100644 index 00000000000..5f81bc6bf97 --- /dev/null +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/test/java/JettyWebSocketUpgradeTest.java @@ -0,0 +1,76 @@ +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static java.util.concurrent.TimeUnit.SECONDS; +import static java.util.regex.Pattern.compile; +import static java.util.regex.Pattern.quote; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.api.DDSpanTypes; +import datadog.trace.core.DDSpan; +import java.net.URI; +import java.util.List; +import org.eclipse.jetty.server.Server; +import org.eclipse.jetty.server.ServerConnector; +import org.eclipse.jetty.server.handler.ContextHandler; +import org.eclipse.jetty.websocket.api.Session; +import org.eclipse.jetty.websocket.client.WebSocketClient; +import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler; +import org.junit.jupiter.api.Test; + +class JettyWebSocketUpgradeTest extends AbstractInstrumentationTest { + @Test + void httpClientSpanFinishesWhenWebSocketUpgradeSucceeds() throws Exception { + Server server = new Server(0); + WebSocketClient client = new WebSocketClient(); + try { + ContextHandler context = new ContextHandler("/"); + server.setHandler(context); + context.setHandler( + WebSocketUpgradeHandler.from(server, context) + .configure( + container -> + container.addMapping( + "/upgrade", (request, response, callback) -> new Endpoint()))); + server.start(); + client.start(); + URI uri = + URI.create( + "ws://localhost:" + + ((ServerConnector) server.getConnectors()[0]).getLocalPort() + + "/upgrade"); + + Session session = client.connect(new Endpoint(), uri).get(5, SECONDS); + + assertTrue(session.isOpen()); + // The HTTP handshake must be reported before the WebSocket connection closes. + assertTraces( + trace( + span() + .operationName(compile(quote("http.request"))) + .resourceName(compile(quote("GET /upgrade"))) + .type(DDSpanTypes.HTTP_CLIENT) + .root() + .error(false)), + trace(span().type(DDSpanTypes.HTTP_SERVER).error(false))); + DDSpan handshake = + writer.stream() + .flatMap(List::stream) + .filter(s -> "client".equals(s.getTag("span.kind"))) + .findFirst() + .orElseThrow(() -> new AssertionError("Missing client handshake span")); + assertEquals("jetty-client", handshake.getTag("component").toString()); + assertEquals("client", handshake.getTag("span.kind")); + assertEquals(101, handshake.getTag("http.status_code")); + } finally { + try { + client.stop(); + } finally { + server.stop(); + } + } + } + + public static class Endpoint implements Session.Listener.AutoDemanding {} +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java similarity index 72% rename from dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java rename to dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java index 4e0cf65b704..b6e44278ae6 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketUpgradeTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java @@ -10,7 +10,6 @@ import static java.util.regex.Pattern.quote; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; import datadog.trace.agent.test.AbstractInstrumentationTest; import datadog.trace.api.DDSpanTypes; @@ -26,57 +25,9 @@ import org.eclipse.jetty.websocket.api.Session; import org.eclipse.jetty.websocket.client.WebSocketClient; import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler; -import org.junit.jupiter.api.Test; import org.tabletest.junit.TableTest; -class JettyWebSocketUpgradeTest extends AbstractInstrumentationTest { - @Test - void httpClientSpanFinishesWhenWebSocketUpgradeSucceeds() throws Exception { - Server server = new Server(0); - WebSocketClient client = new WebSocketClient(); - try { - ContextHandler context = new ContextHandler("/"); - server.setHandler(context); - context.setHandler( - WebSocketUpgradeHandler.from(server, context) - .configure( - container -> - container.addMapping( - "/upgrade", (request, response, callback) -> new Endpoint()))); - server.start(); - client.start(); - URI uri = - URI.create( - "ws://localhost:" - + ((ServerConnector) server.getConnectors()[0]).getLocalPort() - + "/upgrade"); - - Session session = client.connect(new Endpoint(), uri).get(5, SECONDS); - - assertTrue(session.isOpen()); - // The HTTP handshake must be reported before the WebSocket connection closes. - assertTraces( - trace( - span() - .operationName(compile(quote("http.request"))) - .resourceName(compile(quote("GET /upgrade"))) - .type(DDSpanTypes.HTTP_CLIENT) - .root() - .error(false)), - trace(span().type(DDSpanTypes.HTTP_SERVER).error(false))); - DDSpan handshake = handshake("client"); - assertEquals("jetty-client", handshake.getTag("component").toString()); - assertEquals("client", handshake.getTag("span.kind")); - assertEquals(101, handshake.getTag("http.status_code")); - } finally { - try { - client.stop(); - } finally { - server.stop(); - } - } - } - +class JettyWebSocketHandshakeTest extends AbstractInstrumentationTest { @TableTest({ "scenario | endpointName | messageType", "listener text | full | text ", From 6dacd01e478b487f7169ce73187a1916d1a2d607 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 29 Sep 2026 13:57:18 -0700 Subject: [PATCH 06/23] Extract shared Jetty WebSocket endpoint factory and clarify test names --- .../src/test/java/JettyEndpoints.java | 21 ++++++++++++++++ .../java/JettyWebSocketHandshakeTest.java | 21 ++++++++-------- ...ocketTest.java => JettyWebSocketTest.java} | 25 +++---------------- 3 files changed, 35 insertions(+), 32 deletions(-) rename dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/{JettyWebsocketTest.java => JettyWebSocketTest.java} (97%) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java index 89296e029c4..bb9981664e8 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java @@ -12,8 +12,29 @@ import org.eclipse.jetty.websocket.api.annotations.OnWebSocketMessage; import org.eclipse.jetty.websocket.api.annotations.OnWebSocketOpen; import org.eclipse.jetty.websocket.api.annotations.WebSocket; +import org.tabletest.junit.TypeConverter; public class JettyEndpoints { + @TypeConverter + public static EndpointEvents endpoint(String name) { + switch (name) { + case "full": + return new FullListener(); + case "partial": + return new PartialListener(); + case "pojoFull": + return new PojoFullEndpoint(); + case "pojoPartial": + return new PojoPartialEndpoint(); + case "boxedFull": + return new BoxedFullEndpoint(); + case "boxedPartial": + return new BoxedPartialEndpoint(); + default: + throw new IllegalArgumentException("Unknown endpoint: " + name); + } + } + public static class EndpointEvents { Session session; final List messages = new ArrayList<>(); diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java index b6e44278ae6..59ebe618f59 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java @@ -26,18 +26,19 @@ import org.eclipse.jetty.websocket.client.WebSocketClient; import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler; import org.tabletest.junit.TableTest; +import org.tabletest.junit.TypeConverterSources; +@TypeConverterSources(JettyEndpoints.class) class JettyWebSocketHandshakeTest extends AbstractInstrumentationTest { @TableTest({ - "scenario | endpointName | messageType", - "listener text | full | text ", - "listener binary | full | binary ", - "boxed text | boxedFull | text ", - "boxed binary | boxedFull | binary " + "scenario | endpoint | messageType", + "listener text | full | text ", + "listener binary | full | binary ", + "boxed text | boxedFull | text ", + "boxed binary | boxedFull | binary " }) - void receivedMessageLinksToServerHandshake(String endpointName, String messageType) - throws Exception { - JettyEndpoints.EndpointEvents endpoint = JettyWebsocketTest.endpoint(endpointName); + void receivedMessageLinksToServerHandshake( + JettyEndpoints.EndpointEvents endpoint, String messageType) throws Exception { Server server = new Server(0); WebSocketClient client = new WebSocketClient(); try { @@ -55,7 +56,7 @@ void receivedMessageLinksToServerHandshake(String endpointName, String messageTy "ws://localhost:" + ((ServerConnector) server.getConnectors()[0]).getLocalPort() + "/receive"); - Session session = client.connect(new Endpoint(), uri).get(5, SECONDS); + Session session = client.connect(new ClientEndpoint(), uri).get(5, SECONDS); writer.waitForTraces(2); DDSpan handshake = handshake("server"); assertEquals(101, handshake.getTag("http.status_code")); @@ -109,5 +110,5 @@ private static DDSpan handshake(String kind) { .orElseThrow(() -> new AssertionError("Missing " + kind + " handshake span")); } - public static class Endpoint implements Session.Listener.AutoDemanding {} + public static class ClientEndpoint implements Session.Listener.AutoDemanding {} } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java similarity index 97% rename from dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java rename to dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java index f83acdc7ae1..6023a3de14b 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebsocketTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java @@ -62,9 +62,10 @@ import org.eclipse.jetty.websocket.server.internal.ServerFrameHandlerFactory; import org.junit.jupiter.api.Test; import org.tabletest.junit.TableTest; -import org.tabletest.junit.TypeConverter; +import org.tabletest.junit.TypeConverterSources; -public class JettyWebsocketTest extends AbstractInstrumentationTest { +@TypeConverterSources(JettyEndpoints.class) +public class JettyWebSocketTest extends AbstractInstrumentationTest { private static final String URL = "ws://inmemory/test/param"; @TableTest({ @@ -482,26 +483,6 @@ void terminationFinishesAllPendingBinaryMessages(boolean last) throws Exception assertNull(activeSpan()); } - @TypeConverter - public static JettyEndpoints.EndpointEvents endpoint(String name) { - switch (name) { - case "full": - return new JettyEndpoints.FullListener(); - case "partial": - return new JettyEndpoints.PartialListener(); - case "pojoFull": - return new JettyEndpoints.PojoFullEndpoint(); - case "pojoPartial": - return new JettyEndpoints.PojoPartialEndpoint(); - case "boxedFull": - return new JettyEndpoints.BoxedFullEndpoint(); - case "boxedPartial": - return new JettyEndpoints.BoxedPartialEndpoint(); - default: - throw new IllegalArgumentException("Unknown endpoint: " + name); - } - } - private static byte opcode(String msgType) { return "text".equals(msgType) ? OpCode.TEXT : OpCode.BINARY; } From af473db66f4a4a09c8b53e76fe53a458c2f55496 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 29 Sep 2026 14:10:34 -0700 Subject: [PATCH 07/23] Fix Jetty 12 WebSocket Java tests to target Java 17 --- .../websocket/jetty-websocket/jetty-websocket-12.0/build.gradle | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle index e24844fce66..53055dff588 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle @@ -47,7 +47,7 @@ addTestSuiteForDir("latestDepTest", "test") ["compileTestJava", "compileLatestDepTestJava"].each { name -> tasks.named(name, JavaCompile) { - configureCompiler(it, 17, JavaVersion.VERSION_1_8, "Uses Jetty 12 APIs (Java 17+) at Java 8 bytecode") + configureCompiler(it, JavaVersion.VERSION_17) } } From 63206fddf7d7f81df956263cbb4b3243ad3621fd Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Thu, 1 Oct 2026 21:29:49 -0700 Subject: [PATCH 08/23] Validate Jetty WebSocket upgrade method compatibility with muzzle --- .../jetty-client/jetty-client-12.0/build.gradle | 7 +++++++ .../jetty-client-12.0/gradle.lockfile | 4 ++-- .../JettyHttpClientInstrumentation.java | 5 +++++ .../JettyWebSocketUpgradeInstrumentation.java | 16 ++++++++++++++++ 4 files changed, 30 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle index 3f7550a3b20..67606a3ed9b 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle @@ -10,6 +10,12 @@ muzzle { javaVersion = "17" assertInverse = true } + pass { + group = "org.eclipse.jetty.websocket" + module = "jetty-websocket-core-client" + versions = "[12,)" + javaVersion = "17" + } } tracerJava { @@ -26,6 +32,7 @@ configurations.matching { it.name.startsWith('test') || it.name.startsWith('late dependencies { main_java17CompileOnly group: 'org.eclipse.jetty', name: 'jetty-client', version: '12.0.0' + main_java17CompileOnly group: 'org.eclipse.jetty.websocket', name: 'jetty-websocket-core-client', version: '12.0.0' // to test conflicts testImplementation(project(':dd-java-agent:instrumentation:jetty:jetty-client:jetty-client-9.1')) testImplementation(project(':dd-java-agent:instrumentation:jetty:jetty-client:jetty-client-10.0')) diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/gradle.lockfile b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/gradle.lockfile index fc2bf9e655b..e47bbbc5151 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/gradle.lockfile +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/gradle.lockfile @@ -82,9 +82,9 @@ org.codenarc:CodeNarc:3.7.0=codenarc org.dom4j:dom4j:2.2.0=spotbugs org.eclipse.jetty.compression:jetty-compression-common:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.compression:jetty-compression-gzip:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-core-client:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-client:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-client:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.0=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.0=main_java17CompileClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-common:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.0=testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-server:12.1.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyHttpClientInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyHttpClientInstrumentation.java index 82c1f39967b..1203a4d71cc 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyHttpClientInstrumentation.java @@ -27,6 +27,11 @@ public JettyHttpClientInstrumentation() { super("jetty-client"); } + @Override + public String muzzleDirective() { + return "jetty-client"; + } + @Override public String instrumentedType() { return "org.eclipse.jetty.client.transport.HttpRequest"; diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java index a9143488457..7e796941c34 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java @@ -15,6 +15,8 @@ import net.bytebuddy.asm.Advice; import org.eclipse.jetty.client.Request; import org.eclipse.jetty.client.Response; +import org.eclipse.jetty.io.EndPoint; +import org.eclipse.jetty.websocket.core.client.CoreClientUpgradeRequest; @AutoService(InstrumenterModule.class) public class JettyWebSocketUpgradeInstrumentation extends InstrumenterModule.Tracing @@ -23,6 +25,11 @@ public JettyWebSocketUpgradeInstrumentation() { super("jetty-client"); } + @Override + public String muzzleDirective() { + return "jetty-websocket-core-client"; + } + @Override public String instrumentedType() { return "org.eclipse.jetty.websocket.core.client.CoreClientUpgradeRequest"; @@ -60,5 +67,14 @@ public static void afterUpgrade(@Advice.Argument(0) Response response) { span.finish(); } } + + /** + * Lets Muzzle fail CI if the upgrade method is removed or its signature changes, instead of + * silently skipping instrumentation. + */ + private void muzzleCheck( + CoreClientUpgradeRequest request, Response response, EndPoint endPoint) { + request.upgrade(response, endPoint); + } } } From 46d005b9d6afef58dda1f2f548841752afb19123 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Fri, 2 Oct 2026 14:10:18 -0700 Subject: [PATCH 09/23] =?UTF-8?q?Trace=20messages=20sent=20through=20Jetty?= =?UTF-8?q?=2012=E2=80=99s=20native=20WebSocket=20API?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../JettyWebSocketUpgradeInstrumentation.java | 34 +- .../jetty12/Jetty12NativeWebSocketModule.java | 61 ++- .../websocket/jetty12/NativeSendContext.java | 176 ++++++++ .../jetty12/NativeSessionInstrumentation.java | 122 +++++ .../src/test/java/JettyEndpoints.java | 4 + .../java/JettyWebSocketHandshakeTest.java | 41 +- .../src/test/java/JettyWebSocketSendTest.java | 418 ++++++++++++++++++ 7 files changed, 824 insertions(+), 32 deletions(-) create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSessionInstrumentation.java create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java index 7e796941c34..4a131369502 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java @@ -1,12 +1,14 @@ package datadog.trace.instrumentation.jetty_client12; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.instrumentation.jetty_client12.JettyClientDecorator.DECORATE; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.InstrumentationContext; @@ -56,15 +58,33 @@ public void methodAdvice(MethodTransformer transformer) { } public static class WebSocketUpgradeAdvice { - @Advice.OnMethodExit(suppress = Throwable.class) - public static void afterUpgrade(@Advice.Argument(0) Response response) { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static ContextScope beforeUpgrade(@Advice.Argument(0) Response response) { AgentSpan span = InstrumentationContext.get(Request.class, AgentSpan.class).get(response.getRequest()); - if (span != null) { - // Successful upgrades bypass the request's response completion listeners. - DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); - span.finish(); + return span == null ? null : activateSpan(span); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void afterUpgrade( + @Advice.Argument(0) Response response, + @Advice.Enter ContextScope scope, + @Advice.Thrown Throwable failure) { + AgentSpan span = + InstrumentationContext.get(Request.class, AgentSpan.class).get(response.getRequest()); + try { + if (span != null && failure == null) { + // Successful upgrades bypass the request's response completion listeners. + DECORATE.onResponse(span, response); + DECORATE.beforeFinish(span); + } + } finally { + if (scope != null) { + scope.close(); + } + if (span != null && failure == null) { + span.finish(); + } } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java index 0d4ff8dad96..fa8043d50f9 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java @@ -2,7 +2,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; -import static java.util.Collections.singletonMap; +import static java.util.Arrays.asList; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; @@ -14,6 +14,8 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.instrumentation.websocket.jetty12.NativeMethodHandleWrappers.ReceiveContexts; import java.lang.invoke.MethodHandle; +import java.util.HashMap; +import java.util.List; import java.util.Map; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.websocket.common.WebSocketSession; @@ -21,8 +23,7 @@ import org.eclipse.jetty.websocket.core.CoreSession; @AutoService(InstrumenterModule.class) -public class Jetty12NativeWebSocketModule extends InstrumenterModule.Tracing - implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { +public class Jetty12NativeWebSocketModule extends InstrumenterModule.Tracing { public Jetty12NativeWebSocketModule() { super("jetty", "jetty-websocket", "websocket"); } @@ -33,8 +34,8 @@ protected boolean defaultEnabled() { } @Override - public String instrumentedType() { - return "org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler"; + public List typeInstrumentations() { + return asList(new FrameHandlerInstrumentation(), new NativeSessionInstrumentation()); } @Override @@ -48,29 +49,45 @@ public String[] helperClassNames() { packageName + ".NativeMethodHandleWrappers", packageName + ".NativeMethodHandleWrappers$ReceiveContexts", packageName + ".NativeMethodHandleWrappers$BinaryMessage", - packageName + ".NativeMethodHandleWrappers$ReceiveCallback" + packageName + ".NativeMethodHandleWrappers$ReceiveCallback", + packageName + ".NativeSendContext", + packageName + ".NativeSendContext$Message", + packageName + ".NativeSendContext$SendCallback" }; } @Override public Map contextStore() { - return singletonMap( + Map stores = new HashMap<>(); + stores.put( "org.eclipse.jetty.websocket.core.CoreSession", packageName + ".NativeMethodHandleWrappers$ReceiveContexts"); + stores.put( + "org.eclipse.jetty.websocket.common.WebSocketSession", packageName + ".NativeSendContext"); + return stores; } - @Override - public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( - named("onOpen") - .and(takesArguments(2)) - .and(takesArgument(0, named("org.eclipse.jetty.websocket.core.CoreSession"))), - getClass().getName() + "$OpenAdvice"); - transformer.applyAdvice( - named("createMessageSink").and(takesArguments(4)), - getClass().getName() + "$MessageSinkAdvice"); - transformer.applyAdvice( - named("notifyOnClose").and(takesArguments(2)), getClass().getName() + "$CloseAdvice"); + public static class FrameHandlerInstrumentation + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + @Override + public String instrumentedType() { + return "org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler"; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("onOpen") + .and(takesArguments(2)) + .and(takesArgument(0, named("org.eclipse.jetty.websocket.core.CoreSession"))), + Jetty12NativeWebSocketModule.class.getName() + "$OpenAdvice"); + transformer.applyAdvice( + named("createMessageSink").and(takesArguments(4)), + Jetty12NativeWebSocketModule.class.getName() + "$MessageSinkAdvice"); + transformer.applyAdvice( + named("notifyOnClose").and(takesArguments(2)), + Jetty12NativeWebSocketModule.class.getName() + "$CloseAdvice"); + } } public static class OpenAdvice { @@ -108,6 +125,12 @@ public static class CloseAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void onEnter(@Advice.FieldValue("session") WebSocketSession session) { if (session != null) { + NativeSendContext sends = + InstrumentationContext.get(WebSocketSession.class, NativeSendContext.class) + .remove(session); + if (sends != null) { + sends.finish(); + } ReceiveContexts contexts = InstrumentationContext.get(CoreSession.class, ReceiveContexts.class) .remove(session.getCoreSession()); diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java new file mode 100644 index 00000000000..6cdb690dc86 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java @@ -0,0 +1,176 @@ +package datadog.trace.instrumentation.websocket.jetty12; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; +import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_BINARY; +import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_TEXT; + +import datadog.context.ContextScope; +import datadog.trace.api.Config; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; +import java.nio.ByteBuffer; +import java.util.HashSet; +import java.util.Set; +import org.eclipse.jetty.websocket.api.Callback; +import org.eclipse.jetty.websocket.core.Behavior; +import org.eclipse.jetty.websocket.core.CoreSession; + +/** Tracks message boundaries independently of asynchronous send completion. */ +public class NativeSendContext { + private final AgentSpan handshakeSpan; + private final String sessionId; + private final Set pending = new HashSet<>(); + private Message partialMessage; + private boolean closed; + + public NativeSendContext(AgentSpan span, CoreSession session) { + if (Config.get().isWebsocketMessagesInheritSampling()) { + span.forceSamplingDecision(); + } + handshakeSpan = session.getBehavior() == Behavior.CLIENT ? span : span.getLocalRootSpan(); + sessionId = Integer.toHexString(System.identityHashCode(session)); + } + + public synchronized SendCallback start( + Object payload, boolean binary, boolean partial, boolean last, Callback callback) { + if (closed) { + return null; + } + CharSequence type = binary ? MESSAGE_TYPE_BINARY : MESSAGE_TYPE_TEXT; + int size = + payload == null + ? 0 + : binary ? ((ByteBuffer) payload).remaining() : ((String) payload).length(); + Message message = partial ? partialMessage : null; + if (message == null || !type.equals(message.getMessageType())) { + message = new Message(handshakeSpan, sessionId); + pending.add(message); + } + AgentSpan span = DECORATE.startOutboundFrameSpan(message, type, size); + message.pendingCallbacks++; + message.complete = last; + if (partial && (partialMessage == null || partialMessage == message)) { + partialMessage = last ? null : message; + } + return new SendCallback(this, message, span, callback); + } + + public synchronized void finish() { + closed = true; + for (Message message : pending) { + message.finished = true; + DECORATE.onFrameEnd(message); + } + pending.clear(); + partialMessage = null; + } + + public static class Message extends HandlerContext.Sender { + private int pendingCallbacks; + private boolean complete; + private boolean finished; + + public Message(AgentSpan handshakeSpan, String sessionId) { + super(handshakeSpan, sessionId); + } + } + + public static class SendCallback implements Callback { + private final NativeSendContext context; + private final Message message; + private final AgentSpan span; + private final Callback delegate; + private boolean completed; + private boolean methodExited; + private boolean released; + + public SendCallback( + NativeSendContext context, Message message, AgentSpan span, Callback delegate) { + this.context = context; + this.message = message; + this.span = span; + this.delegate = delegate; + } + + public AgentSpan span() { + return span; + } + + @Override + public void succeed() { + complete(null); + } + + @Override + public void fail(Throwable failure) { + complete(failure); + } + + private void complete(Throwable failure) { + synchronized (context) { + onError(failure); + } + try (ContextScope ignored = activateSpan(span)) { + try { + if (delegate != null) { + if (failure == null) { + delegate.succeed(); + } else { + delegate.fail(failure); + } + } + } catch (Throwable t) { + synchronized (context) { + onError(t); + } + throw t; + } + } finally { + synchronized (context) { + completed = true; + finish(); + } + } + } + + public void onMethodExit(ContextScope scope, Throwable failure) { + synchronized (context) { + try { + onError(failure); + } finally { + if (scope != null) { + scope.close(); + } + if (failure != null) { + completed = true; + } + methodExited = true; + finish(); + } + } + } + + private void onError(Throwable failure) { + if (failure != null && !message.finished) { + DECORATE.onError(span, failure); + message.complete = true; + if (context.partialMessage == message) { + context.partialMessage = null; + } + } + } + + private void finish() { + // Synchronous completion must wait for the send method's scope to close. + if (completed && methodExited && !released) { + released = true; + if (--message.pendingCallbacks == 0 && message.complete && !message.finished) { + message.finished = true; + DECORATE.onFrameEnd(message); + context.pending.remove(message); + } + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSessionInstrumentation.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSessionInstrumentation.java new file mode 100644 index 00000000000..dff67181390 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSessionInstrumentation.java @@ -0,0 +1,122 @@ +package datadog.trace.instrumentation.websocket.jetty12; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static net.bytebuddy.matcher.ElementMatchers.isConstructor; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import datadog.context.ContextScope; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.instrumentation.websocket.jetty12.NativeSendContext.SendCallback; +import net.bytebuddy.asm.Advice; +import org.eclipse.jetty.websocket.api.Callback; +import org.eclipse.jetty.websocket.common.WebSocketSession; + +public class NativeSessionInstrumentation + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + @Override + public String instrumentedType() { + return "org.eclipse.jetty.websocket.common.WebSocketSession"; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice(isConstructor(), getClass().getName() + "$ConstructAdvice"); + transformer.applyAdvice( + namedOneOf("sendText", "sendBinary") + .and(takesArguments(2)) + .and(takesArgument(1, named("org.eclipse.jetty.websocket.api.Callback"))), + getClass().getName() + "$SendAdvice"); + transformer.applyAdvice( + namedOneOf("sendPartialText", "sendPartialBinary") + .and(takesArguments(3)) + .and(takesArgument(1, boolean.class)) + .and(takesArgument(2, named("org.eclipse.jetty.websocket.api.Callback"))), + getClass().getName() + "$PartialSendAdvice"); + } + + public static class ConstructAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void after(@Advice.This WebSocketSession session) { + AgentSpan span = activeSpan(); + if (span != null) { + InstrumentationContext.get(WebSocketSession.class, NativeSendContext.class) + .put(session, new NativeSendContext(span, session.getCoreSession())); + } + } + } + + public static class SendAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static ContextScope before( + @Advice.This WebSocketSession session, + @Advice.Origin("#m") String method, + @Advice.Argument(0) Object payload, + @Advice.Argument(value = 1, readOnly = false) Callback callback, + @Advice.Local("send") SendCallback send) { + NativeSendContext context = + InstrumentationContext.get(WebSocketSession.class, NativeSendContext.class).get(session); + if (context == null) { + return null; + } + send = context.start(payload, "sendBinary".equals(method), false, true, callback); + if (send == null) { + return null; + } + callback = send; + return activateSpan(send.span()); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void after( + @Advice.Enter ContextScope scope, + @Advice.Local("send") SendCallback send, + @Advice.Thrown Throwable failure) { + if (send != null) { + send.onMethodExit(scope, failure); + } + } + } + + public static class PartialSendAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static ContextScope before( + @Advice.This WebSocketSession session, + @Advice.Origin("#m") String method, + @Advice.Argument(0) Object payload, + @Advice.Argument(1) boolean last, + @Advice.Argument(value = 2, readOnly = false) Callback callback, + @Advice.Local("send") SendCallback send) { + NativeSendContext context = + InstrumentationContext.get(WebSocketSession.class, NativeSendContext.class).get(session); + if (context == null) { + return null; + } + boolean binary = "sendPartialBinary".equals(method); + send = context.start(payload, binary, true, last, callback); + if (send == null) { + return null; + } + // Jetty accepts null callbacks for binary sends, but partial text sends throw. + if (callback != null || binary) { + callback = send; + } + return activateSpan(send.span()); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void after( + @Advice.Enter ContextScope scope, + @Advice.Local("send") SendCallback send, + @Advice.Thrown Throwable failure) { + if (send != null) { + send.onMethodExit(scope, failure); + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java index bb9981664e8..7a313f4965f 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java @@ -41,6 +41,7 @@ public static class EndpointEvents { final List finalFragments = new ArrayList<>(); final List messageSpans = new ArrayList<>(); boolean failMessages; + boolean echoMessages; boolean failClose; boolean deferCallback; boolean failCallback; @@ -52,6 +53,9 @@ public static class EndpointEvents { void recordMessage(String message) { messages.add(message); messageSpans.add(activeSpan()); + if (echoMessages) { + session.sendText(message, Callback.NOOP); + } if (failMessages) { throw new IllegalStateException("handler failed"); } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java index 59ebe618f59..75e0ccd4b32 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java @@ -1,5 +1,6 @@ import static datadog.trace.agent.test.assertions.SpanLinkMatcher.to; import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TraceMatcher.SORT_BY_START_TIME; import static datadog.trace.agent.test.assertions.TraceMatcher.trace; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_LENGTH; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_TYPE; @@ -9,15 +10,16 @@ import static java.util.regex.Pattern.compile; import static java.util.regex.Pattern.quote; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertSame; import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.agent.test.assertions.SpanMatcher; import datadog.trace.api.DDSpanTypes; import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink; import datadog.trace.bootstrap.instrumentation.api.SpanAttributes; import datadog.trace.core.DDSpan; import java.net.URI; import java.util.List; +import java.util.concurrent.CompletableFuture; import org.eclipse.jetty.server.Server; import org.eclipse.jetty.server.ServerConnector; import org.eclipse.jetty.server.handler.ContextHandler; @@ -37,8 +39,9 @@ class JettyWebSocketHandshakeTest extends AbstractInstrumentationTest { "boxed text | boxedFull | text ", "boxed binary | boxedFull | binary " }) - void receivedMessageLinksToServerHandshake( + void messagesLinkToHandshakesInBothDirections( JettyEndpoints.EndpointEvents endpoint, String messageType) throws Exception { + endpoint.echoMessages = true; Server server = new Server(0); WebSocketClient client = new WebSocketClient(); try { @@ -56,7 +59,8 @@ void receivedMessageLinksToServerHandshake( "ws://localhost:" + ((ServerConnector) server.getConnectors()[0]).getLocalPort() + "/receive"); - Session session = client.connect(new ClientEndpoint(), uri).get(5, SECONDS); + ClientEndpoint clientEndpoint = new ClientEndpoint(); + Session session = client.connect(clientEndpoint, uri).get(5, SECONDS); writer.waitForTraces(2); DDSpan handshake = handshake("server"); assertEquals(101, handshake.getTag("http.status_code")); @@ -68,11 +72,14 @@ void receivedMessageLinksToServerHandshake( session.sendBinary(UTF_8.encode("hello"), sent); } sent.get(5, SECONDS); + assertEquals("hello", clientEndpoint.reply.get(5, SECONDS)); assertTraces( trace(span().type(DDSpanTypes.HTTP_CLIENT).error(false)), trace(span().type(DDSpanTypes.HTTP_SERVER).error(false)), + trace(sendSpan(handshake("client")).root()), trace( + SORT_BY_START_TIME, span() .root() .operationName(compile(quote("websocket.receive"))) @@ -86,10 +93,10 @@ void receivedMessageLinksToServerHandshake( ? AgentSpanLink.SAMPLED_FLAG : AgentSpanLink.DEFAULT_FLAGS) .attributes( - SpanAttributes.builder().put("dd.kind", "executed_from").build())))); + SpanAttributes.builder().put("dd.kind", "executed_from").build())), + sendSpan(handshake).childOfPrevious())); assertEquals(singletonList("hello"), endpoint.messages); assertEquals(1, endpoint.messageSpans.size()); - assertSame(writer.get(2).get(0), endpoint.messageSpans.get(0)); assertEquals( messageType, endpoint.messageSpans.get(0).getTag(WEBSOCKET_MESSAGE_TYPE).toString()); assertEquals(5L, endpoint.messageSpans.get(0).getTag(WEBSOCKET_MESSAGE_LENGTH)); @@ -110,5 +117,27 @@ private static DDSpan handshake(String kind) { .orElseThrow(() -> new AssertionError("Missing " + kind + " handshake span")); } - public static class ClientEndpoint implements Session.Listener.AutoDemanding {} + private static SpanMatcher sendSpan(DDSpan handshake) { + return span() + .operationName(compile(quote("websocket.send"))) + .resourceName(compile(quote("websocket /receive"))) + .type(DDSpanTypes.WEBSOCKET) + .error(false) + .links( + to(handshake) + .traceFlags( + handshake.getSamplingPriority() > 0 + ? AgentSpanLink.SAMPLED_FLAG + : AgentSpanLink.DEFAULT_FLAGS) + .attributes(SpanAttributes.builder().put("dd.kind", "resuming").build())); + } + + public static class ClientEndpoint implements Session.Listener.AutoDemanding { + final CompletableFuture reply = new CompletableFuture<>(); + + @Override + public void onWebSocketText(String message) { + reply.complete(message); + } + } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java new file mode 100644 index 00000000000..accc19f8c3d --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java @@ -0,0 +1,418 @@ +import static datadog.trace.agent.test.assertions.SpanLinkMatcher.to; +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TagsMatcher.defaultTags; +import static datadog.trace.agent.test.assertions.TagsMatcher.tag; +import static datadog.trace.agent.test.assertions.TraceMatcher.SORT_BY_START_TIME; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_FRAMES; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_LENGTH; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_TYPE; +import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_URL; +import static datadog.trace.test.junit.utils.assertions.Matchers.is; +import static datadog.trace.test.junit.utils.assertions.Matchers.matches; +import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.concurrent.TimeUnit.SECONDS; +import static java.util.regex.Pattern.compile; +import static java.util.regex.Pattern.quote; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.agent.test.assertions.SpanMatcher; +import datadog.trace.api.DDSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink; +import datadog.trace.bootstrap.instrumentation.api.SpanAttributes; +import datadog.trace.core.DDSpan; +import java.net.URI; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import org.eclipse.jetty.util.Callback; +import org.eclipse.jetty.websocket.api.Session; +import org.eclipse.jetty.websocket.api.UpgradeRequest; +import org.eclipse.jetty.websocket.api.WebSocketContainer; +import org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler; +import org.eclipse.jetty.websocket.core.Behavior; +import org.eclipse.jetty.websocket.core.CloseStatus; +import org.eclipse.jetty.websocket.core.CoreSession; +import org.eclipse.jetty.websocket.core.Frame; +import org.eclipse.jetty.websocket.core.WebSocketComponents; +import org.eclipse.jetty.websocket.server.internal.ServerFrameHandlerFactory; +import org.junit.jupiter.api.Test; +import org.tabletest.junit.TableTest; + +class JettyWebSocketSendTest extends AbstractInstrumentationTest { + @TableTest({ + "scenario | binary | partial | synchronous", + "text async | false | false | false ", + "binary async | true | false | false ", + "text sync | false | false | true ", + "binary sync | true | false | true ", + "partial text | false | true | false ", + "partial binary | true | true | false " + }) + void sendsFinishOnCallback(boolean binary, boolean partial, boolean synchronous) + throws Exception { + Connection connection = new Connection(true); + connection.synchronous = synchronous; + RecordingCallback callback = new RecordingCallback(); + connection.send(binary, partial, "hello", true, callback); + assertNull(activeSpan()); + AgentSpan sent = connection.spans.get(0); + assertEquals("websocket.send", sent.getOperationName().toString()); + if (!synchronous) { + assertEquals(1, writer.size()); + Thread completion = new Thread(() -> connection.callbacks.get(0).succeeded()); + completion.start(); + completion.join(); + } + callback.get(5, SECONDS); + assertSame(sent, callback.span); + assertNull(activeSpan()); + assertTraces(trace(handshakeSpan()), trace(sendSpan(connection, binary, 5, 1).root())); + } + + @TableTest({ + "scenario | binary", + "text | false ", + "binary | true " + }) + void fragmentedAndOverlappingMessages(boolean binary) throws Exception { + Connection connection = new Connection(true); + connection.send(binary, true, "hello ", false, null); + connection.send(binary, true, "world", true, null); + connection.send(binary, false, "again", true, null); + assertSame(connection.spans.get(0), connection.spans.get(1)); + assertNotSame(connection.spans.get(1), connection.spans.get(2)); + + connection.callbacks.get(1).succeeded(); + assertEquals(1, writer.size()); + connection.callbacks.get(0).succeeded(); + connection.callbacks.get(2).succeeded(); + assertNull(activeSpan()); + assertTraces( + trace(handshakeSpan()), + trace(sendSpan(connection, binary, 11, 2).root()), + trace(sendSpan(connection, binary, 5, 1).root())); + } + + @Test + void sendUsesApplicationParentAndRestoresCallbackContext() throws Exception { + Connection connection = new Connection(true); + RecordingCallback callback = new RecordingCallback(); + runUnderTrace( + "application", + () -> { + AgentSpan parent = activeSpan(); + connection.send(false, false, "hello", true, callback); + assertSame(parent, activeSpan()); + assertEquals(parent.getSpanId(), ((DDSpan) connection.spans.get(0)).getParentId()); + return null; + }); + runUnderTrace( + "completion", + () -> { + AgentSpan parent = activeSpan(); + connection.callbacks.get(0).succeeded(); + assertSame(parent, activeSpan()); + return null; + }); + assertSame(connection.spans.get(0), callback.span); + assertTraces( + trace(handshakeSpan()), + trace( + SORT_BY_START_TIME, + span().operationName("application").root(), + sendSpan(connection, false, 5, 1).childOfPrevious()), + trace(span().operationName("completion").root())); + } + + @TableTest({ + "scenario | failureMode | partial", + "sync failure | fail | false ", + "throw | throw | false ", + "success then throw | successThenThrow | false ", + "fragment failure | fail | true ", + "fragment throw | throw | true " + }) + void synchronousFailuresFinishSpan(String failureMode, boolean partial) throws Exception { + Connection connection = new Connection(true); + connection.failureMode = failureMode; + RecordingCallback callback = new RecordingCallback(); + if ("fail".equals(failureMode)) { + connection.send(false, partial, "hello", !partial, callback); + assertSame(connection.failure, callback.failure); + } else { + assertSame( + connection.failure, + assertThrows( + IllegalStateException.class, + () -> connection.send(false, partial, "hello", !partial, callback))); + } + assertNull(activeSpan()); + assertTraces( + trace(handshakeSpan()), + trace( + span() + .operationName(compile(quote("websocket.send"))) + .type(DDSpanTypes.WEBSOCKET) + .error())); + assertEquals("send failed", connection.spans.get(0).getTag("error.message")); + // A late callback must not finish the span again. + connection.callbacks.get(0).succeeded(); + assertEquals(2, writer.size()); + } + + @TableTest({ + "scenario | partial", + "full | false ", + "partial | true " + }) + void asynchronousFailureClearsMessage(boolean partial) throws Exception { + Connection connection = new Connection(true); + RecordingCallback callback = new RecordingCallback(); + connection.send(false, partial, "hello", !partial, callback); + connection.callbacks.get(0).failed(connection.failure); + assertSame(connection.failure, callback.failure); + assertSame(connection.spans.get(0), callback.span); + connection.send(false, partial, "again", true, null); + connection.callbacks.get(1).succeeded(); + assertNotSame(connection.spans.get(0), connection.spans.get(1)); + assertNull(activeSpan()); + assertTraces( + trace(handshakeSpan()), + trace( + span() + .operationName(compile(quote("websocket.send"))) + .type(DDSpanTypes.WEBSOCKET) + .error()), + trace(sendSpan(connection, false, 5, 1).root())); + } + + @Test + void callbackFailureClosesScopeAndFinishesSpan() throws Exception { + Connection connection = new Connection(true); + connection.send( + false, + false, + "hello", + true, + new org.eclipse.jetty.websocket.api.Callback() { + @Override + public void succeed() { + throw connection.failure; + } + + @Override + public void fail(Throwable failure) {} + }); + assertSame( + connection.failure, + assertThrows(IllegalStateException.class, () -> connection.callbacks.get(0).succeeded())); + assertNull(activeSpan()); + assertTraces( + trace(handshakeSpan()), + trace( + span() + .operationName(compile(quote("websocket.send"))) + .type(DDSpanTypes.WEBSOCKET) + .error())); + } + + @Test + void closeFinishesIncompleteAndPendingSends() throws Exception { + Connection connection = new Connection(true); + connection.send(false, true, "hello", false, null); + connection.callbacks.get(0).succeeded(); + assertEquals(1, writer.size()); + connection.send(false, false, "again", true, null); + Callback.Completable closed = new Callback.Completable(); + connection.handler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), closed); + closed.get(5, SECONDS); + writer.waitForTraces(3); + assertEquals( + 2, + writer.stream() + .flatMap(List::stream) + .filter(s -> "websocket.send".equals(s.getOperationName().toString())) + .count()); + int traces = writer.size(); + connection.callbacks.forEach(Callback::succeeded); + assertEquals(traces, writer.size()); + assertNull(activeSpan()); + } + + @Test + void partialTextPreservesNullCallbackFailure() throws Exception { + Connection connection = new Connection(true); + assertThrows( + NullPointerException.class, () -> connection.session.sendPartialText("hello", true, null)); + assertNull(activeSpan()); + assertTraces( + trace(handshakeSpan()), + trace( + span() + .operationName(compile(quote("websocket.send"))) + .type(DDSpanTypes.WEBSOCKET) + .error())); + } + + @Test + void emptyBinaryMessageKeepsBinaryType() throws Exception { + Connection connection = new Connection(true); + connection.session.sendBinary(null, null); + connection.callbacks.get(0).succeeded(); + assertTraces(trace(handshakeSpan()), trace(sendSpan(connection, true, 0, 1).root())); + } + + @Test + void untracedSessionAndControlFramesDoNotCreateSendSpans() throws Exception { + Connection untraced = new Connection(false); + untraced.synchronous = true; + untraced.send(false, false, "hello", true, null); + assertNull(untraced.spans.get(0)); + assertEquals(0, writer.size()); + Connection traced = new Connection(true); + traced.synchronous = true; + traced.session.sendPing(ByteBuffer.allocate(0), org.eclipse.jetty.websocket.api.Callback.NOOP); + traced.session.sendPong(ByteBuffer.allocate(0), org.eclipse.jetty.websocket.api.Callback.NOOP); + assertNull(traced.spans.get(0)); + assertNull(traced.spans.get(1)); + assertTraces(trace(handshakeSpan())); + } + + private static SpanMatcher handshakeSpan() { + return span().root().operationName("handshake"); + } + + private static SpanMatcher sendSpan( + Connection connection, boolean binary, long size, long frames) { + return span() + .operationName(compile(quote("websocket.send"))) + .resourceName(compile(quote("websocket /send"))) + .type(DDSpanTypes.WEBSOCKET) + .error(false) + .links( + to(connection.handshake) + .traceFlags( + connection.handshake.getSamplingPriority() > 0 + ? AgentSpanLink.SAMPLED_FLAG + : AgentSpanLink.DEFAULT_FLAGS) + .attributes(SpanAttributes.builder().put("dd.kind", "resuming").build())) + .tags( + defaultTags(), + tag("component", matches("websocket")), + tag("span.kind", is("producer")), + tag(WEBSOCKET_MESSAGE_TYPE, matches(binary ? "binary" : "text")), + tag(WEBSOCKET_MESSAGE_LENGTH, is(size)), + tag(WEBSOCKET_MESSAGE_FRAMES, is(frames))); + } + + private static class RecordingCallback + extends org.eclipse.jetty.websocket.api.Callback.Completable { + AgentSpan span; + Throwable failure; + + @Override + public void succeed() { + span = activeSpan(); + super.succeed(); + } + + @Override + public void fail(Throwable failure) { + span = activeSpan(); + this.failure = failure; + super.fail(failure); + } + } + + private static class Connection extends CoreSession.Empty { + final List callbacks = new ArrayList<>(); + final List spans = new ArrayList<>(); + final IllegalStateException failure = new IllegalStateException("send failed"); + final JettyWebSocketFrameHandler handler; + final Session session; + DDSpan handshake; + boolean synchronous; + String failureMode = ""; + + Connection(boolean traced) throws Exception { + ServerFrameHandlerFactory factory = + new ServerFrameHandlerFactory(mock(WebSocketContainer.class), new WebSocketComponents()); + handler = factory.newJettyFrameHandler(new Session.Listener.AutoDemanding() {}); + UpgradeRequest request = mock(UpgradeRequest.class); + when(request.getRequestURI()).thenReturn(URI.create("ws://inmemory/send")); + handler.setUpgradeRequest(request); + Callback.Completable opened = new Callback.Completable(); + if (traced) { + runUnderTrace( + "handshake", + () -> { + handshake = (DDSpan) activeSpan(); + handshake.setResourceName("/send"); + handshake.setTag(HTTP_URL, "ws://inmemory/send"); + handler.onOpen(this, opened); + return null; + }); + } else { + handler.onOpen(this, opened); + } + opened.get(5, SECONDS); + session = handler.getSession(); + } + + @Override + public Behavior getBehavior() { + return Behavior.SERVER; + } + + @Override + public void sendFrame(Frame frame, Callback callback, boolean batch) { + callbacks.add(callback); + spans.add(activeSpan()); + if ("fail".equals(failureMode)) { + callback.failed(failure); + } else if ("throw".equals(failureMode)) { + throw failure; + } else if ("successThenThrow".equals(failureMode)) { + callback.succeeded(); + throw failure; + } else if (synchronous) { + callback.succeeded(); + } + } + + void send( + boolean binary, + boolean partial, + String text, + boolean last, + org.eclipse.jetty.websocket.api.Callback callback) { + if (binary) { + ByteBuffer payload = UTF_8.encode(text); + if (partial) { + session.sendPartialBinary(payload, last, callback); + } else { + session.sendBinary(payload, callback); + } + assertEquals(text.length(), payload.remaining()); + } else if (partial) { + session.sendPartialText( + text, + last, + callback == null ? org.eclipse.jetty.websocket.api.Callback.NOOP : callback); + } else { + session.sendText(text, callback); + } + } + } +} From de874e56db5eeff0b9afe6d1302f596851f96d5c Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Fri, 2 Oct 2026 16:12:13 -0700 Subject: [PATCH 10/23] Fix lost WebSocket send errors after Jetty connection closure --- .../websocket/jetty12/NativeSendContext.java | 26 ++++++--- .../src/test/java/JettyWebSocketSendTest.java | 56 +++++++++++++++---- 2 files changed, 63 insertions(+), 19 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java index 6cdb690dc86..365ce5619da 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java @@ -11,6 +11,7 @@ import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; import java.nio.ByteBuffer; import java.util.HashSet; +import java.util.Iterator; import java.util.Set; import org.eclipse.jetty.websocket.api.Callback; import org.eclipse.jetty.websocket.core.Behavior; @@ -58,11 +59,16 @@ public synchronized SendCallback start( public synchronized void finish() { closed = true; - for (Message message : pending) { - message.finished = true; - DECORATE.onFrameEnd(message); + for (Iterator iterator = pending.iterator(); iterator.hasNext(); ) { + Message message = iterator.next(); + // Close ends the message, but pending callbacks still determine the send outcome. + message.complete = true; + if (message.pendingCallbacks == 0) { + message.finished = true; + DECORATE.onFrameEnd(message); + iterator.remove(); + } } - pending.clear(); partialMessage = null; } @@ -152,11 +158,13 @@ public void onMethodExit(ContextScope scope, Throwable failure) { } private void onError(Throwable failure) { - if (failure != null && !message.finished) { - DECORATE.onError(span, failure); - message.complete = true; - if (context.partialMessage == message) { - context.partialMessage = null; + synchronized (context) { + if (failure != null && !message.finished) { + DECORATE.onError(span, failure); + message.complete = true; + if (context.partialMessage == message) { + context.partialMessage = null; + } } } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java index accc19f8c3d..fedf737c0e9 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java @@ -228,7 +228,7 @@ public void fail(Throwable failure) {} } @Test - void closeFinishesIncompleteAndPendingSends() throws Exception { + void closeFinishesIncompleteSendsAndWaitsForPendingCallbacks() throws Exception { Connection connection = new Connection(true); connection.send(false, true, "hello", false, null); connection.callbacks.get(0).succeeded(); @@ -237,17 +237,53 @@ void closeFinishesIncompleteAndPendingSends() throws Exception { Callback.Completable closed = new Callback.Completable(); connection.handler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), closed); closed.get(5, SECONDS); - writer.waitForTraces(3); - assertEquals( - 2, - writer.stream() - .flatMap(List::stream) - .filter(s -> "websocket.send".equals(s.getOperationName().toString())) - .count()); - int traces = writer.size(); + assertTraces(trace(handshakeSpan()), trace(sendSpan(connection, false, 5, 1).root())); + connection.callbacks.get(1).succeeded(); + assertTraces( + trace(handshakeSpan()), + trace(sendSpan(connection, false, 5, 1).root()), + trace(sendSpan(connection, false, 5, 1).root())); connection.callbacks.forEach(Callback::succeeded); - assertEquals(traces, writer.size()); + assertEquals(3, writer.size()); + assertNull(activeSpan()); + } + + @TableTest({ + "scenario | partial | failed", + "full success | false | false ", + "full failure | false | true ", + "partial success | true | false ", + "partial failure | true | true " + }) + void peerClosePreservesPendingSendOutcome(boolean partial, boolean failed) throws Exception { + Connection connection = new Connection(true); + RecordingCallback callback = new RecordingCallback(); + connection.send(false, partial, "hello", !partial, callback); + Callback.Completable closed = new Callback.Completable(); + connection.handler.onFrame(new CloseStatus(CloseStatus.NORMAL, "bye").toFrame(), closed); + closed.get(5, SECONDS); + assertEquals(1, writer.size()); + + if (failed) { + connection.callbacks.get(0).failed(connection.failure); + assertSame(connection.failure, callback.failure); + assertEquals("send failed", connection.spans.get(0).getTag("error.message")); + } else { + connection.callbacks.get(0).succeeded(); + callback.get(5, SECONDS); + } + assertSame(connection.spans.get(0), callback.span); assertNull(activeSpan()); + SpanMatcher expected = + failed + ? span() + .operationName(compile(quote("websocket.send"))) + .type(DDSpanTypes.WEBSOCKET) + .error() + : sendSpan(connection, false, 5, 1); + assertTraces(trace(handshakeSpan()), trace(expected.root())); + connection.callbacks.get(0).succeeded(); + assertEquals(2, writer.size()); } @Test From 333e30151765f24e5181aa33f9190df5deafdfc1 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 6 Oct 2026 13:35:56 -0700 Subject: [PATCH 11/23] Track Jetty 12.0 patch releases in latest-dependency tests and Muzzle. --- .../jetty-websocket-12.0/build.gradle | 21 ++-- .../jetty-websocket-12.0/gradle.lockfile | 103 +++++++++--------- 2 files changed, 61 insertions(+), 63 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle index 53055dff588..7a3eb2175a8 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle @@ -7,28 +7,28 @@ muzzle { name = 'jetty-websocket-12-native' group = 'org.eclipse.jetty.websocket' module = 'jetty-websocket-jetty-server' - versions = "[12,12.0.17]" + versions = "[12,12.1)" javaVersion = "17" } pass { name = 'jetty-websocket-12ee8' group = 'org.eclipse.jetty.ee8.websocket' module = 'jetty-ee8-websocket-javax-server' - versions = "[12,12.0.17]" + versions = "[12,12.1)" javaVersion = "17" } pass { name = 'jetty-websocket-12ee9' group = 'org.eclipse.jetty.ee9.websocket' module = 'jetty-ee9-websocket-jakarta-server' - versions = "[12,12.0.17]" + versions = "[12,12.1)" javaVersion = "17" } pass { name = 'jetty-websocket-12ee10' group = 'org.eclipse.jetty.ee10.websocket' module = 'jetty-ee10-websocket-jakarta-server' - versions = "[12,12.0.17]" + versions = "[12,12.1)" javaVersion = "17" } } @@ -60,13 +60,12 @@ dependencies { testImplementation group: 'org.eclipse.jetty.ee10.websocket', name: 'jetty-ee10-websocket-jakarta-server', version: '12.0.0' testImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0' testImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.0' - //TODO: jetty-12.1.0 is still alpha but wraps MethodHandle class into a MethodHolder class. - // Today that is not stable but we'll need to port those advices to support that once the code base will be a bit more stable - latestDepTestImplementation group: 'org.eclipse.jetty.ee8.websocket', name: 'jetty-ee8-websocket-javax-server', version: '12.0.17' - latestDepTestImplementation group: 'org.eclipse.jetty.ee9.websocket', name: 'jetty-ee9-websocket-jakarta-server', version: '12.0.17' - latestDepTestImplementation group: 'org.eclipse.jetty.ee10.websocket', name: 'jetty-ee10-websocket-jakarta-server', version: '12.0.17' - latestDepTestImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.17' - latestDepTestImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.17' + //TODO: jetty-12.1.0 wraps MethodHandle class into a MethodHolder class. + latestDepTestImplementation group: 'org.eclipse.jetty.ee8.websocket', name: 'jetty-ee8-websocket-javax-server', version: '12.0.+' + latestDepTestImplementation group: 'org.eclipse.jetty.ee9.websocket', name: 'jetty-ee9-websocket-jakarta-server', version: '12.0.+' + latestDepTestImplementation group: 'org.eclipse.jetty.ee10.websocket', name: 'jetty-ee10-websocket-jakarta-server', version: '12.0.+' + latestDepTestImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.+' + latestDepTestImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.+' testRuntimeOnly project(":dd-java-agent:instrumentation:jetty:jetty-client:jetty-client-12.0") testRuntimeOnly project(":dd-java-agent:instrumentation:jetty:jetty-server:jetty-server-12.0") diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile index 216f0d87215..93820f505f9 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile @@ -94,101 +94,102 @@ org.codehaus.groovy:groovy:3.0.25=latestDepTestCompileClasspath,latestDepTestRun org.codenarc:CodeNarc:3.7.0=codenarc org.dom4j:dom4j:2.2.0=spotbugs org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-jakarta-client:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-jakarta-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-jakarta-client:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-jakarta-common:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-jakarta-common:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-jakarta-common:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-jakarta-server:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-jakarta-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-jakarta-server:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-servlet:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-servlet:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee10.websocket:jetty-ee10-websocket-servlet:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee10:jetty-ee10-annotations:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee10:jetty-ee10-annotations:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee10:jetty-ee10-annotations:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee10:jetty-ee10-plus:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee10:jetty-ee10-plus:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee10:jetty-ee10-plus:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee10:jetty-ee10-servlet:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee10:jetty-ee10-servlet:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee10:jetty-ee10-servlet:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee10:jetty-ee10-webapp:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee10:jetty-ee10-webapp:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee10:jetty-ee10-webapp:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-javax-client:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-javax-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-javax-client:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-javax-common:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-javax-common:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-javax-common:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-javax-server:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-javax-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-javax-server:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-servlet:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-servlet:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee8.websocket:jetty-ee8-websocket-servlet:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee8:jetty-ee8-annotations:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee8:jetty-ee8-annotations:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee8:jetty-ee8-annotations:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee8:jetty-ee8-nested:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee8:jetty-ee8-nested:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee8:jetty-ee8-nested:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee8:jetty-ee8-plus:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee8:jetty-ee8-plus:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee8:jetty-ee8-plus:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee8:jetty-ee8-security:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee8:jetty-ee8-security:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee8:jetty-ee8-security:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee8:jetty-ee8-servlet:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee8:jetty-ee8-servlet:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee8:jetty-ee8-servlet:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee8:jetty-ee8-webapp:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee8:jetty-ee8-webapp:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee8:jetty-ee8-webapp:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-jakarta-client:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-jakarta-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-jakarta-client:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-jakarta-common:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-jakarta-common:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-jakarta-common:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-jakarta-server:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-jakarta-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-jakarta-server:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-servlet:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-servlet:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee9.websocket:jetty-ee9-websocket-servlet:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee9:jetty-ee9-annotations:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee9:jetty-ee9-annotations:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee9:jetty-ee9-annotations:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee9:jetty-ee9-nested:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee9:jetty-ee9-nested:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee9:jetty-ee9-nested:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee9:jetty-ee9-plus:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee9:jetty-ee9-plus:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee9:jetty-ee9-plus:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee9:jetty-ee9-security:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee9:jetty-ee9-security:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee9:jetty-ee9-security:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee9:jetty-ee9-servlet:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee9:jetty-ee9-servlet:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee9:jetty-ee9-servlet:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.ee9:jetty-ee9-webapp:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.ee9:jetty-ee9-webapp:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.ee9:jetty-ee9-webapp:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.toolchain:jetty-jakarta-servlet-api:5.0.2=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.toolchain:jetty-jakarta-websocket-api:2.0.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.eclipse.jetty.toolchain:jetty-javax-websocket-api:1.1.2=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.toolchain:jetty-servlet-api:4.0.6=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.toolchain:jetty-servlet-api:4.0.6=testCompileClasspath,testRuntimeClasspath +org.eclipse.jetty.toolchain:jetty-servlet-api:4.0.9=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-client:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-core-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-client:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-common:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-core-server:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-jetty-api:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-jetty-api:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-api:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-common:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-alpn-client:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty:jetty-alpn-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-alpn-client:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-client:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty:jetty-client:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty:jetty-ee:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-client:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-ee:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-http:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty:jetty-http:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-http:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-io:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty:jetty-io:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-io:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-jndi:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty:jetty-jndi:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath -org.eclipse.jetty:jetty-plus:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-jndi:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-plus:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-security:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty:jetty-security:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-security:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-server:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty:jetty-server:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-server:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-session:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty:jetty-session:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-session:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-util:12.0.0=compileClasspath,testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty:jetty-util:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-util:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.eclipse.jetty:jetty-xml:12.0.0=testCompileClasspath,testRuntimeClasspath -org.eclipse.jetty:jetty-xml:12.0.17=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath +org.eclipse.jetty:jetty-xml:12.0.40=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath org.gmetrics:GMetrics:2.1.0=codenarc org.hamcrest:hamcrest-core:1.3=latestDepTestRuntimeClasspath,testRuntimeClasspath org.hamcrest:hamcrest:3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath @@ -212,12 +213,10 @@ org.objenesis:objenesis:3.3=latestDepTestCompileClasspath,latestDepTestRuntimeCl org.opentest4j:opentest4j:1.3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.ow2.asm:asm-analysis:9.10.1=spotbugs org.ow2.asm:asm-analysis:9.7.1=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.ow2.asm:asm-commons:9.10.1=latestDepTestRuntimeClasspath,spotbugs,testRuntimeClasspath +org.ow2.asm:asm-commons:9.10.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,spotbugs,testRuntimeClasspath org.ow2.asm:asm-commons:9.5=testCompileClasspath -org.ow2.asm:asm-commons:9.7.1=latestDepTestCompileClasspath -org.ow2.asm:asm-tree:9.10.1=latestDepTestRuntimeClasspath,spotbugs,testRuntimeClasspath +org.ow2.asm:asm-tree:9.10.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,spotbugs,testRuntimeClasspath org.ow2.asm:asm-tree:9.5=testCompileClasspath -org.ow2.asm:asm-tree:9.7.1=latestDepTestCompileClasspath org.ow2.asm:asm-util:9.10.1=spotbugs org.ow2.asm:asm-util:9.7.1=latestDepTestRuntimeClasspath,testRuntimeClasspath org.ow2.asm:asm:9.10.1=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,spotbugs,testCompileClasspath,testRuntimeClasspath From f8e0fdd00ebbd6d6829d0ffb43c837574e80331f Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 6 Oct 2026 13:35:56 -0700 Subject: [PATCH 12/23] Trace Jetty 12 native streaming WebSocket handlers Wrap Reader and InputStream handlers to create receive spans and record handler failures. Cover fragmented messages and scope cleanup in tests. --- .../jetty12/NativeMethodHandleWrappers.java | 38 ++++++++- .../src/test/java/JettyEndpoints.java | 20 +++++ .../src/test/java/JettyWebSocketTest.java | 84 +++++++++++++++++-- 3 files changed, 130 insertions(+), 12 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java index 5aa2811161c..deed6c966f3 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java @@ -12,6 +12,8 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; import datadog.trace.util.MethodHandles; +import java.io.InputStream; +import java.io.Reader; import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodType; import java.nio.ByteBuffer; @@ -55,6 +57,13 @@ public class NativeMethodHandleWrappers { Session.class, int.class, String.class); + private static final MethodHandle STREAM = + LOOKUP.method( + NativeMethodHandleWrappers.class, + "onStream", + MethodHandle.class, + ReceiveContexts.class, + Object.class); private static HandlerContext.Receiver context(AgentSpan span, CoreSession session) { if (Config.get().isWebsocketMessagesInheritSampling()) { @@ -71,16 +80,22 @@ public static MethodHandle wrapMessage( ContextStore contextStore) { MethodType type = delegate.type(); Class payload = type.parameterType(0); - if (payload != String.class && payload != ByteBuffer.class) { + boolean streaming = payload == Reader.class || payload == InputStream.class; + if (payload != String.class && payload != ByteBuffer.class && !streaming) { return delegate; } - boolean partial = type.parameterCount() > 1 && type.parameterType(1) == boolean.class; delegate = delegate.asType(type.changeReturnType(void.class)); - MethodHandle normalized = partial ? delegate : dropArguments(delegate, 1, boolean.class); ReceiveContexts contexts = contextStore.get(session); if (contexts == null) { contexts = contextStore.getOrPut(session, new ReceiveContexts(span, session)); } + if (streaming) { + return insertArguments( + STREAM, 0, delegate.asType(MethodType.methodType(void.class, Object.class)), contexts) + .asType(delegate.type()); + } + boolean partial = type.parameterCount() > 1 && type.parameterType(1) == boolean.class; + MethodHandle normalized = partial ? delegate : dropArguments(delegate, 1, boolean.class); MethodHandle wrapper = payload == String.class ? contexts.wrapText(normalized, partial) @@ -138,6 +153,23 @@ public static void onText( } } + public static void onStream(MethodHandle delegate, ReceiveContexts contexts, Object payload) + throws Throwable { + // Jetty dispatches streams to a worker and requires consumption before the handler returns. + HandlerContext.Receiver context = + new HandlerContext.Receiver(contexts.handshakeSpan, contexts.sessionId); + try (ContextScope ignored = startMessage(context, payload, false)) { + try { + delegate.invokeExact(payload); + } catch (Throwable t) { + DECORATE.onError(context.getWebsocketSpan(), t); + throw t; + } + } finally { + DECORATE.onFrameEnd(context); + } + } + public static void onBinary( MethodHandle delegate, ReceiveContexts contexts, diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java index 7a313f4965f..47ce5d9bd79 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java @@ -3,6 +3,10 @@ import static org.junit.jupiter.api.Assertions.assertSame; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.io.IOException; +import java.io.InputStream; +import java.io.Reader; +import java.io.StringWriter; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.List; @@ -151,6 +155,22 @@ public void onBinary(ByteBuffer payload, Callback callback) { } } + @WebSocket + public static class StreamingEndpoint extends PojoEndpoint { + @OnWebSocketMessage + public void onText(Session session, Reader payload) throws IOException { + assertSame(this.session, session); + StringWriter message = new StringWriter(); + payload.transferTo(message); + recordMessage(message.toString()); + } + + @OnWebSocketMessage + public void onBinary(InputStream payload) throws IOException { + recordMessage(new String(payload.readAllBytes(), UTF_8)); + } + } + @WebSocket public static class PojoPartialEndpoint extends PojoEndpoint { @OnWebSocketMessage diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java index 6023a3de14b..f14e543dc23 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java @@ -1,5 +1,6 @@ import static datadog.trace.agent.test.assertions.SpanMatcher.span; import static datadog.trace.agent.test.assertions.TagsMatcher.defaultTags; +import static datadog.trace.agent.test.assertions.TagsMatcher.error; import static datadog.trace.agent.test.assertions.TagsMatcher.includes; import static datadog.trace.agent.test.assertions.TagsMatcher.tag; import static datadog.trace.agent.test.assertions.TraceMatcher.trace; @@ -20,6 +21,7 @@ import static datadog.trace.test.junit.utils.assertions.Matchers.matches; import static java.util.Arrays.asList; import static java.util.Collections.singletonList; +import static java.util.concurrent.Executors.newSingleThreadExecutor; import static java.util.concurrent.TimeUnit.SECONDS; import static java.util.regex.Pattern.compile; import static java.util.regex.Pattern.quote; @@ -49,6 +51,7 @@ import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; import org.eclipse.jetty.util.Callback; import org.eclipse.jetty.websocket.api.UpgradeRequest; import org.eclipse.jetty.websocket.api.WebSocketContainer; @@ -68,6 +71,56 @@ public class JettyWebSocketTest extends AbstractInstrumentationTest { private static final String URL = "ws://inmemory/test/param"; + @TableTest({ + "scenario | msgType | fragmented | fail ", + "reader | text | false | false", + "input stream | binary | false | false", + "fragmented reader | text | true | false", + "fragmented input stream | binary | true | false", + "reader failure | text | false | true ", + "input stream failure | binary | false | true ", + "fragmented reader failure | text | true | true ", + "fragmented stream failure | binary | true | true " + }) + void streamingMessagesHaveReceiveSpan(String msgType, boolean fragmented, boolean fail) + throws Exception { + ExecutorService executor = newSingleThreadExecutor(); + try { + WebSocketComponents components = + new WebSocketComponents(null, null, null, null, null, executor); + JettyEndpoints.StreamingEndpoint endpoint = new JettyEndpoints.StreamingEndpoint(); + endpoint.failMessages = fail; + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler, Behavior.SERVER, true, components); + + if (fragmented) { + deliver(frameHandler, new Frame(opcode(msgType), "hello ").setFin(false)); + deliver(frameHandler, new Frame(OpCode.CONTINUATION, "world")); + } else { + deliver(frameHandler, new Frame(opcode(msgType), "hello world")); + } + // Queue behind Jetty's handler to check cleanup on the dispatch thread as well. + executor.submit(() -> assertNull(activeSpan())).get(5, SECONDS); + + assertEquals(singletonList("hello world"), endpoint.messages); + assertEquals(1, endpoint.messageSpans.size()); + assertNotNull(endpoint.messageSpans.get(0)); + assertEquals("websocket.receive", endpoint.messageSpans.get(0).getOperationName().toString()); + assertNull(activeSpan()); + TagsMatcher[] errorTags = + fail + ? new TagsMatcher[] {error(IllegalStateException.class, "handler failed")} + : new TagsMatcher[0]; + assertTraces( + trace(handshakeSpan()), + trace(receiveSpan(handshake(), msgType, 0, 1, errorTags).error(fail))); + assertSame(endpoint.messageSpans.get(0), writer.get(1).get(0)); + } finally { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, SECONDS)); + } + } + @TableTest({ "scenario | endpoint | msgType", "full text | full | text ", @@ -504,12 +557,26 @@ private static void openFrameHandler(JettyWebSocketFrameHandler frameHandler) th private static void openFrameHandler( JettyWebSocketFrameHandler frameHandler, Behavior connectionBehavior, boolean traced) throws Exception { + openFrameHandler(frameHandler, connectionBehavior, traced, null); + } + + private static void openFrameHandler( + JettyWebSocketFrameHandler frameHandler, + Behavior connectionBehavior, + boolean traced, + WebSocketComponents components) + throws Exception { CoreSession session = new CoreSession.Empty() { @Override public Behavior getBehavior() { return connectionBehavior; } + + @Override + public WebSocketComponents getWebSocketComponents() { + return components; + } }; Callback.Completable openCallback = new Callback.Completable(); if (traced) { @@ -548,15 +615,14 @@ private static SpanMatcher handshakeSpan() { } private static SpanMatcher receiveSpan( - DDSpan handshake, String msgType, long length, long frames) { - return websocketSpan( - handshake, - "websocket.receive", - tag(WEBSOCKET_MESSAGE_TYPE, matches(quote(msgType))), - tag(WEBSOCKET_MESSAGE_LENGTH, is(length)), - tag(WEBSOCKET_MESSAGE_FRAMES, is(frames)), - // Full-message callbacks do not record a receive-time tag. - includes(WEBSOCKET_MESSAGE_RECEIVE_TIME)); + DDSpan handshake, String msgType, long length, long frames, TagsMatcher... extraTags) { + List tags = new ArrayList<>(asList(extraTags)); + tags.add(tag(WEBSOCKET_MESSAGE_TYPE, matches(quote(msgType)))); + tags.add(tag(WEBSOCKET_MESSAGE_LENGTH, is(length))); + tags.add(tag(WEBSOCKET_MESSAGE_FRAMES, is(frames))); + // Full-message callbacks do not record a receive-time tag. + tags.add(includes(WEBSOCKET_MESSAGE_RECEIVE_TIME)); + return websocketSpan(handshake, "websocket.receive", tags.toArray(new TagsMatcher[0])); } private static SpanMatcher closeSpan(DDSpan handshake) { From 3c32223f0e489a5c0a39e254aea377251da7f6e6 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 6 Oct 2026 16:12:39 -0700 Subject: [PATCH 13/23] Fix Jetty 12 WebSocket handshake spans for swallowed upgrade failures --- .../jetty-client-12.0/build.gradle | 1 + .../JettyWebSocketUpgradeInstrumentation.java | 14 ++- .../test/java/JettyWebSocketUpgradeTest.java | 99 ++++++++++++++++--- 3 files changed, 99 insertions(+), 15 deletions(-) diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle index 67606a3ed9b..30384ca8744 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/build.gradle @@ -45,6 +45,7 @@ dependencies { testImplementation group: 'org.eclipse.jetty', name: 'jetty-client', version: '12.0.0' testImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.0.0' testImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-server:12.0.0' + testImplementation libs.bundles.mockito testRuntimeOnly project(':dd-java-agent:instrumentation:jetty:jetty-server:jetty-server-12.0') latestDepTestImplementation group: 'org.eclipse.jetty', name: 'jetty-client', version: '12.+' latestDepTestImplementation 'org.eclipse.jetty.websocket:jetty-websocket-jetty-client:12.+' diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java index 4a131369502..064b71279c5 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyWebSocketUpgradeInstrumentation.java @@ -14,6 +14,9 @@ import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.util.Map; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.client.Request; import org.eclipse.jetty.client.Response; @@ -68,14 +71,23 @@ public static ContextScope beforeUpgrade(@Advice.Argument(0) Response response) @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void afterUpgrade( @Advice.Argument(0) Response response, + @Advice.FieldValue("futureCoreSession") CompletableFuture futureCoreSession, @Advice.Enter ContextScope scope, @Advice.Thrown Throwable failure) { AgentSpan span = InstrumentationContext.get(Request.class, AgentSpan.class).get(response.getRequest()); try { if (span != null && failure == null) { - // Successful upgrades bypass the request's response completion listeners. + // Upgrades bypass response completion listeners even when endpoint failures are + // swallowed. DECORATE.onResponse(span, response); + try { + futureCoreSession.getNow(null); + } catch (CompletionException e) { + DECORATE.onError(span, e.getCause()); + } catch (CancellationException e) { + DECORATE.onError(span, e); + } DECORATE.beforeFinish(span); } } finally { diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/test/java/JettyWebSocketUpgradeTest.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/test/java/JettyWebSocketUpgradeTest.java index 5f81bc6bf97..05be8732fee 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/test/java/JettyWebSocketUpgradeTest.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/test/java/JettyWebSocketUpgradeTest.java @@ -4,18 +4,31 @@ import static java.util.regex.Pattern.compile; import static java.util.regex.Pattern.quote; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.AdditionalAnswers.delegatesTo; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; import datadog.trace.agent.test.AbstractInstrumentationTest; import datadog.trace.api.DDSpanTypes; import datadog.trace.core.DDSpan; import java.net.URI; import java.util.List; +import java.util.concurrent.ExecutionException; +import org.eclipse.jetty.client.Response; +import org.eclipse.jetty.io.Connection; +import org.eclipse.jetty.io.EndPoint; import org.eclipse.jetty.server.Server; import org.eclipse.jetty.server.ServerConnector; import org.eclipse.jetty.server.handler.ContextHandler; import org.eclipse.jetty.websocket.api.Session; import org.eclipse.jetty.websocket.client.WebSocketClient; +import org.eclipse.jetty.websocket.core.FrameHandler; +import org.eclipse.jetty.websocket.core.client.CoreClientUpgradeRequest; +import org.eclipse.jetty.websocket.core.client.WebSocketCoreClient; import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler; import org.junit.jupiter.api.Test; @@ -25,21 +38,8 @@ void httpClientSpanFinishesWhenWebSocketUpgradeSucceeds() throws Exception { Server server = new Server(0); WebSocketClient client = new WebSocketClient(); try { - ContextHandler context = new ContextHandler("/"); - server.setHandler(context); - context.setHandler( - WebSocketUpgradeHandler.from(server, context) - .configure( - container -> - container.addMapping( - "/upgrade", (request, response, callback) -> new Endpoint()))); - server.start(); + URI uri = startServer(server); client.start(); - URI uri = - URI.create( - "ws://localhost:" - + ((ServerConnector) server.getConnectors()[0]).getLocalPort() - + "/upgrade"); Session session = client.connect(new Endpoint(), uri).get(5, SECONDS); @@ -72,5 +72,76 @@ void httpClientSpanFinishesWhenWebSocketUpgradeSucceeds() throws Exception { } } + @Test + void httpClientSpanRecordsFailureSwallowedByWebSocketUpgrade() throws Exception { + Server server = new Server(0); + WebSocketCoreClient client = new WebSocketCoreClient(); + try { + URI uri = startServer(server); + client.start(); + IllegalStateException upgradeFailure = new IllegalStateException("Endpoint upgrade failed"); + FrameHandler frameHandler = mock(FrameHandler.class); + CoreClientUpgradeRequest request = + new CoreClientUpgradeRequest(client, uri) { + @Override + public FrameHandler getFrameHandler() { + return frameHandler; + } + + @Override + public void upgrade(Response response, EndPoint endPoint) { + EndPoint failingEndPoint = mock(EndPoint.class, delegatesTo(endPoint)); + doThrow(upgradeFailure).when(failingEndPoint).upgrade(any(Connection.class)); + super.upgrade(response, failingEndPoint); + } + }; + + ExecutionException failure = + assertThrows(ExecutionException.class, () -> client.connect(request).get(5, SECONDS)); + assertSame(upgradeFailure, failure.getCause()); + + assertTraces( + trace( + span() + .operationName(compile(quote("http.request"))) + .resourceName(compile(quote("GET /upgrade"))) + .type(DDSpanTypes.HTTP_CLIENT) + .root() + .error(true)), + trace(span().type(DDSpanTypes.HTTP_SERVER).error(false))); + DDSpan handshake = + writer.stream() + .flatMap(List::stream) + .filter(s -> "client".equals(s.getTag("span.kind"))) + .findFirst() + .orElseThrow(() -> new AssertionError("Missing client handshake span")); + assertEquals(101, handshake.getTag("http.status_code")); + assertEquals(IllegalStateException.class.getName(), handshake.getTag("error.type")); + assertEquals(upgradeFailure.getMessage(), handshake.getTag("error.message")); + } finally { + try { + client.stop(); + } finally { + server.stop(); + } + } + } + + private static URI startServer(Server server) throws Exception { + ContextHandler context = new ContextHandler("/"); + server.setHandler(context); + context.setHandler( + WebSocketUpgradeHandler.from(server, context) + .configure( + container -> + container.addMapping( + "/upgrade", (request, response, callback) -> new Endpoint()))); + server.start(); + return URI.create( + "ws://localhost:" + + ((ServerConnector) server.getConnectors()[0]).getLocalPort() + + "/upgrade"); + } + public static class Endpoint implements Session.Listener.AutoDemanding {} } From 32b71409b97051f014027c19225ef820d05da0ac Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 6 Oct 2026 17:44:23 -0700 Subject: [PATCH 14/23] Fix lost receive errors after Jetty WebSocket peer close `ReceiveContexts.finish()` finishes and resets every pending binary span when a close frame arrives. If a deferred receive callback subsequently fails, `onFailure()` finds no span, so failed processing is reported as successful. Keep pending receives open until their callbacks resolve, matching the send-side handling. --- .../jetty12/NativeMethodHandleWrappers.java | 12 +++- .../src/test/java/JettyWebSocketTest.java | 65 ++++++++++++++----- 2 files changed, 59 insertions(+), 18 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java index deed6c966f3..e3ef1cfae98 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java @@ -18,6 +18,7 @@ import java.lang.invoke.MethodType; import java.nio.ByteBuffer; import java.util.HashSet; +import java.util.Iterator; import java.util.Set; import org.eclipse.jetty.websocket.api.Callback; import org.eclipse.jetty.websocket.api.Session; @@ -233,10 +234,15 @@ public synchronized void finish() { DECORATE.onFrameEnd(text); } } - for (BinaryMessage message : pendingBinary) { - DECORATE.onFrameEnd(message); + for (Iterator iterator = pendingBinary.iterator(); iterator.hasNext(); ) { + BinaryMessage message = iterator.next(); + // Close ends the message, but pending callbacks still determine the receive outcome. + message.complete = true; + if (message.pendingCallbacks == 0) { + DECORATE.onFrameEnd(message); + iterator.remove(); + } } - pendingBinary.clear(); currentBinary = null; } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java index f14e543dc23..38a413d43ad 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java @@ -25,7 +25,6 @@ import static java.util.concurrent.TimeUnit.SECONDS; import static java.util.regex.Pattern.compile; import static java.util.regex.Pattern.quote; -import static java.util.stream.Collectors.toList; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -501,12 +500,23 @@ void deferredPartialBinarySuccessKeepsSpanOpenUntilFinalCallback() throws Except } @TableTest({ - "scenario | last ", - "nonfinal | false", - "final | true " + "scenario | endpoint | last | closeFrame | fail ", + "full close frame success | full | true | true | false", + "full close frame failure | full | true | true | true ", + "partial final close frame success | partial | true | true | false", + "partial final close frame failure | partial | true | true | true ", + "partial close frame success | partial | false | true | false", + "partial close frame failure | partial | false | true | true ", + "full closed success | full | true | false | false", + "full closed failure | full | true | false | true ", + "partial final closed success | partial | true | false | false", + "partial final closed failure | partial | true | false | true ", + "partial closed success | partial | false | false | false", + "partial closed failure | partial | false | false | true " }) - void terminationFinishesAllPendingBinaryMessages(boolean last) throws Exception { - JettyEndpoints.PartialListener endpoint = new JettyEndpoints.PartialListener(); + void terminationWaitsForPendingBinaryCallbacks( + JettyEndpoints.EndpointEvents endpoint, boolean last, boolean closeFrame, boolean fail) + throws Exception { endpoint.deferCallback = true; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); openFrameHandler(frameHandler); @@ -519,20 +529,45 @@ void terminationFinishesAllPendingBinaryMessages(boolean last) throws Exception assertEquals(1, writer.size()); assertNotSame(endpoint.messageSpans.get(0), endpoint.messageSpans.get(1)); - Callback.Completable closed = new Callback.Completable(); - frameHandler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), closed); - closed.get(5, SECONDS); + if (closeFrame) { + deliver(frameHandler, CloseStatus.toFrame(CloseStatus.NORMAL, "bye")); + } else { + Callback.Completable closed = new Callback.Completable(); + frameHandler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), closed); + closed.get(5, SECONDS); + } - writer.waitForTraces(4); - assertTrue( - writer.stream().flatMap(List::stream).collect(toList()).containsAll(endpoint.messageSpans)); + DDSpan handshake = handshake(); + assertTraces(trace(handshakeSpan()), trace(closeSpan(handshake))); + assertFalse(message.isDone()); + assertFalse(previousMessage.isDone()); + + if (fail) { + endpoint.pendingCallback.fail(new IllegalStateException("callback failed after close")); + } else { + endpoint.pendingCallback.succeed(); + message.get(5, SECONDS); + } + assertTrue(message.isDone()); + assertEquals(fail, message.isCompletedExceptionally()); + TagsMatcher[] errorTags = + fail + ? new TagsMatcher[] {error(IllegalStateException.class, "callback failed after close")} + : new TagsMatcher[0]; + SpanMatcher expected = receiveSpan(handshake, "binary", 5, 1, errorTags).error(fail); + assertTraces(trace(handshakeSpan()), trace(expected), trace(closeSpan(handshake))); + assertSame(endpoint.messageSpans.get(1), writer.get(2).get(0)); + assertFalse(previousMessage.isDone()); - endpoint.pendingCallback.succeed(); - message.get(5, SECONDS); previousCallback.succeed(); previousMessage.get(5, SECONDS); - assertEquals(4, writer.size()); + assertTraces( + trace(handshakeSpan()), + trace(receiveSpan(handshake, "binary", 8, 1)), + trace(expected), + trace(closeSpan(handshake))); + assertSame(endpoint.messageSpans.get(0), writer.get(3).get(0)); assertNull(activeSpan()); } From 572a50eca9d90fb25ede16d7da88ae4fb917b7e5 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Wed, 7 Oct 2026 09:29:36 -0700 Subject: [PATCH 15/23] Add Muzzle guard and tests to prevent partial instrumentation on Jetty 12.1+ --- .../jetty-websocket-12.0/build.gradle | 1 + ...ty12NativeFrameHandlerInstrumentation.java | 101 ++++++++++++++++++ ... Jetty12NativeSessionInstrumentation.java} | 2 +- .../jetty12/Jetty12NativeWebSocketModule.java | 89 +-------------- ...ttyEndpoints.java => NativeEndpoints.java} | 2 +- ...java => NativeWebSocketHandshakeTest.java} | 6 +- ...Test.java => NativeWebSocketSendTest.java} | 2 +- ...cketTest.java => NativeWebSocketTest.java} | 35 +++--- 8 files changed, 129 insertions(+), 109 deletions(-) create mode 100644 dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeFrameHandlerInstrumentation.java rename dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/{NativeSessionInstrumentation.java => Jetty12NativeSessionInstrumentation.java} (98%) rename dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/{JettyEndpoints.java => NativeEndpoints.java} (99%) rename dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/{JettyWebSocketHandshakeTest.java => NativeWebSocketHandshakeTest.java} (96%) rename dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/{JettyWebSocketSendTest.java => NativeWebSocketSendTest.java} (99%) rename dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/{JettyWebSocketTest.java => NativeWebSocketTest.java} (95%) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle index 7a3eb2175a8..09f90f5d2fb 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle @@ -9,6 +9,7 @@ muzzle { module = 'jetty-websocket-jetty-server' versions = "[12,12.1)" javaVersion = "17" + assertInverse = true } pass { name = 'jetty-websocket-12ee8' diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeFrameHandlerInstrumentation.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeFrameHandlerInstrumentation.java new file mode 100644 index 00000000000..94f87cd04f3 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeFrameHandlerInstrumentation.java @@ -0,0 +1,101 @@ +package datadog.trace.instrumentation.websocket.jetty12; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import static org.eclipse.jetty.websocket.core.util.InvokerUtils.bindTo; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.instrumentation.websocket.jetty12.NativeMethodHandleWrappers.ReceiveContexts; +import java.lang.invoke.MethodHandle; +import net.bytebuddy.asm.Advice; +import org.eclipse.jetty.websocket.common.WebSocketSession; +import org.eclipse.jetty.websocket.core.Behavior; +import org.eclipse.jetty.websocket.core.CoreSession; + +public class Jetty12NativeFrameHandlerInstrumentation + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + @Override + public String instrumentedType() { + return "org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler"; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("onOpen") + .and(takesArguments(2)) + .and(takesArgument(0, named("org.eclipse.jetty.websocket.core.CoreSession"))), + getClass().getName() + "$OpenAdvice"); + transformer.applyAdvice( + named("createMessageSink").and(takesArguments(4)), + getClass().getName() + "$MessageSinkAdvice"); + transformer.applyAdvice( + named("notifyOnClose").and(takesArguments(2)), getClass().getName() + "$CloseAdvice"); + } + + public static class OpenAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter( + @Advice.Argument(0) CoreSession session, + @Advice.FieldValue(value = "closeHandle", readOnly = false) MethodHandle closeHandle) { + AgentSpan span = activeSpan(); + if (span != null && session.getBehavior() == Behavior.SERVER && closeHandle != null) { + closeHandle = NativeMethodHandleWrappers.wrapClose(closeHandle, span, session); + } + } + + /** + * Rejects the native module on Jetty 12.1. Without this check, {@code + * Jetty12NativeSessionInstrumentation} installs, but {@code + * Jetty12NativeFrameHandlerInstrumentation} fails due to incompatible {@code MethodHolder} + * signatures. The missing close cleanup leaves spans for successful partial sends with {@code + * last=false} unfinished when the connection closes. + */ + private static void muzzleCheck(MethodHandle handle, Object receiver) { + bindTo(handle, receiver); + } + } + + public static class MessageSinkAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter( + @Advice.Argument(1) WebSocketSession session, + @Advice.Argument(value = 2, readOnly = false) MethodHandle handle) { + AgentSpan span = activeSpan(); + if (span != null + && session.getCoreSession().getBehavior() == Behavior.SERVER + && handle != null) { + handle = + NativeMethodHandleWrappers.wrapMessage( + handle, + span, + session.getCoreSession(), + InstrumentationContext.get(CoreSession.class, ReceiveContexts.class)); + } + } + } + + public static class CloseAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter(@Advice.FieldValue("session") WebSocketSession session) { + if (session != null) { + NativeSendContext sends = + InstrumentationContext.get(WebSocketSession.class, NativeSendContext.class) + .remove(session); + if (sends != null) { + sends.finish(); + } + ReceiveContexts contexts = + InstrumentationContext.get(CoreSession.class, ReceiveContexts.class) + .remove(session.getCoreSession()); + if (contexts != null) { + contexts.finish(); + } + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSessionInstrumentation.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeSessionInstrumentation.java similarity index 98% rename from dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSessionInstrumentation.java rename to dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeSessionInstrumentation.java index dff67181390..81753360362 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSessionInstrumentation.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeSessionInstrumentation.java @@ -17,7 +17,7 @@ import org.eclipse.jetty.websocket.api.Callback; import org.eclipse.jetty.websocket.common.WebSocketSession; -public class NativeSessionInstrumentation +public class Jetty12NativeSessionInstrumentation implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { @Override public String instrumentedType() { diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java index fa8043d50f9..eeb3e88c5d5 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeWebSocketModule.java @@ -1,26 +1,14 @@ package datadog.trace.instrumentation.websocket.jetty12; -import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static java.util.Arrays.asList; -import static net.bytebuddy.matcher.ElementMatchers.takesArgument; -import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.InstrumenterConfig; -import datadog.trace.bootstrap.InstrumentationContext; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import datadog.trace.instrumentation.websocket.jetty12.NativeMethodHandleWrappers.ReceiveContexts; -import java.lang.invoke.MethodHandle; import java.util.HashMap; import java.util.List; import java.util.Map; -import net.bytebuddy.asm.Advice; -import org.eclipse.jetty.websocket.common.WebSocketSession; -import org.eclipse.jetty.websocket.core.Behavior; -import org.eclipse.jetty.websocket.core.CoreSession; @AutoService(InstrumenterModule.class) public class Jetty12NativeWebSocketModule extends InstrumenterModule.Tracing { @@ -35,7 +23,8 @@ protected boolean defaultEnabled() { @Override public List typeInstrumentations() { - return asList(new FrameHandlerInstrumentation(), new NativeSessionInstrumentation()); + return asList( + new Jetty12NativeFrameHandlerInstrumentation(), new Jetty12NativeSessionInstrumentation()); } @Override @@ -66,78 +55,4 @@ public Map contextStore() { "org.eclipse.jetty.websocket.common.WebSocketSession", packageName + ".NativeSendContext"); return stores; } - - public static class FrameHandlerInstrumentation - implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { - @Override - public String instrumentedType() { - return "org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler"; - } - - @Override - public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( - named("onOpen") - .and(takesArguments(2)) - .and(takesArgument(0, named("org.eclipse.jetty.websocket.core.CoreSession"))), - Jetty12NativeWebSocketModule.class.getName() + "$OpenAdvice"); - transformer.applyAdvice( - named("createMessageSink").and(takesArguments(4)), - Jetty12NativeWebSocketModule.class.getName() + "$MessageSinkAdvice"); - transformer.applyAdvice( - named("notifyOnClose").and(takesArguments(2)), - Jetty12NativeWebSocketModule.class.getName() + "$CloseAdvice"); - } - } - - public static class OpenAdvice { - @Advice.OnMethodEnter(suppress = Throwable.class) - public static void onEnter( - @Advice.Argument(0) CoreSession session, - @Advice.FieldValue(value = "closeHandle", readOnly = false) MethodHandle closeHandle) { - AgentSpan span = activeSpan(); - if (span != null && session.getBehavior() == Behavior.SERVER && closeHandle != null) { - closeHandle = NativeMethodHandleWrappers.wrapClose(closeHandle, span, session); - } - } - } - - public static class MessageSinkAdvice { - @Advice.OnMethodEnter(suppress = Throwable.class) - public static void onEnter( - @Advice.Argument(1) WebSocketSession session, - @Advice.Argument(value = 2, readOnly = false) MethodHandle handle) { - AgentSpan span = activeSpan(); - if (span != null - && session.getCoreSession().getBehavior() == Behavior.SERVER - && handle != null) { - handle = - NativeMethodHandleWrappers.wrapMessage( - handle, - span, - session.getCoreSession(), - InstrumentationContext.get(CoreSession.class, ReceiveContexts.class)); - } - } - } - - public static class CloseAdvice { - @Advice.OnMethodEnter(suppress = Throwable.class) - public static void onEnter(@Advice.FieldValue("session") WebSocketSession session) { - if (session != null) { - NativeSendContext sends = - InstrumentationContext.get(WebSocketSession.class, NativeSendContext.class) - .remove(session); - if (sends != null) { - sends.finish(); - } - ReceiveContexts contexts = - InstrumentationContext.get(CoreSession.class, ReceiveContexts.class) - .remove(session.getCoreSession()); - if (contexts != null) { - contexts.finish(); - } - } - } - } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeEndpoints.java similarity index 99% rename from dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java rename to dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeEndpoints.java index 47ce5d9bd79..3c3084b7f36 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeEndpoints.java @@ -18,7 +18,7 @@ import org.eclipse.jetty.websocket.api.annotations.WebSocket; import org.tabletest.junit.TypeConverter; -public class JettyEndpoints { +public class NativeEndpoints { @TypeConverter public static EndpointEvents endpoint(String name) { switch (name) { diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketHandshakeTest.java similarity index 96% rename from dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java rename to dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketHandshakeTest.java index 75e0ccd4b32..d353a4b085e 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketHandshakeTest.java @@ -30,8 +30,8 @@ import org.tabletest.junit.TableTest; import org.tabletest.junit.TypeConverterSources; -@TypeConverterSources(JettyEndpoints.class) -class JettyWebSocketHandshakeTest extends AbstractInstrumentationTest { +@TypeConverterSources(NativeEndpoints.class) +class NativeWebSocketHandshakeTest extends AbstractInstrumentationTest { @TableTest({ "scenario | endpoint | messageType", "listener text | full | text ", @@ -40,7 +40,7 @@ class JettyWebSocketHandshakeTest extends AbstractInstrumentationTest { "boxed binary | boxedFull | binary " }) void messagesLinkToHandshakesInBothDirections( - JettyEndpoints.EndpointEvents endpoint, String messageType) throws Exception { + NativeEndpoints.EndpointEvents endpoint, String messageType) throws Exception { endpoint.echoMessages = true; Server server = new Server(0); WebSocketClient client = new WebSocketClient(); diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java similarity index 99% rename from dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java rename to dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java index fedf737c0e9..c8909a0d17a 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java @@ -49,7 +49,7 @@ import org.junit.jupiter.api.Test; import org.tabletest.junit.TableTest; -class JettyWebSocketSendTest extends AbstractInstrumentationTest { +class NativeWebSocketSendTest extends AbstractInstrumentationTest { @TableTest({ "scenario | binary | partial | synchronous", "text async | false | false | false ", diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java similarity index 95% rename from dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java rename to dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java index 38a413d43ad..46379902cb6 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java @@ -66,8 +66,8 @@ import org.tabletest.junit.TableTest; import org.tabletest.junit.TypeConverterSources; -@TypeConverterSources(JettyEndpoints.class) -public class JettyWebSocketTest extends AbstractInstrumentationTest { +@TypeConverterSources(NativeEndpoints.class) +public class NativeWebSocketTest extends AbstractInstrumentationTest { private static final String URL = "ws://inmemory/test/param"; @TableTest({ @@ -87,7 +87,7 @@ void streamingMessagesHaveReceiveSpan(String msgType, boolean fragmented, boolea try { WebSocketComponents components = new WebSocketComponents(null, null, null, null, null, executor); - JettyEndpoints.StreamingEndpoint endpoint = new JettyEndpoints.StreamingEndpoint(); + NativeEndpoints.StreamingEndpoint endpoint = new NativeEndpoints.StreamingEndpoint(); endpoint.failMessages = fail; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); openFrameHandler(frameHandler, Behavior.SERVER, true, components); @@ -135,7 +135,8 @@ void streamingMessagesHaveReceiveSpan(String msgType, boolean fragmented, boolea "boxed partial text | boxedPartial | text ", "boxed partial binary | boxedPartial | binary " }) - void nativeJettyAdvices(JettyEndpoints.EndpointEvents endpoint, String msgType) throws Exception { + void nativeJettyAdvices(NativeEndpoints.EndpointEvents endpoint, String msgType) + throws Exception { JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); Callback.Completable messageCallback = new Callback.Completable(); Callback.Completable closeCallback = new Callback.Completable(); @@ -169,7 +170,7 @@ void nativeJettyAdvices(JettyEndpoints.EndpointEvents endpoint, String msgType) "boxed partial text | boxedPartial | text ", "boxed partial binary | boxedPartial | binary " }) - void fragmentedMessagesShareSpan(JettyEndpoints.EndpointEvents endpoint, String msgType) + void fragmentedMessagesShareSpan(NativeEndpoints.EndpointEvents endpoint, String msgType) throws Exception { JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); openFrameHandler(frameHandler); @@ -204,7 +205,7 @@ void fragmentedMessagesShareSpan(JettyEndpoints.EndpointEvents endpoint, String "POJO close | pojoFull | | true | websocket.close " }) void handlerFailureMarksSpanAndClosesScope( - JettyEndpoints.EndpointEvents endpoint, String msgType, boolean last, String operation) + NativeEndpoints.EndpointEvents endpoint, String msgType, boolean last, String operation) throws Exception { endpoint.failClose = msgType == null; endpoint.failMessages = msgType != null; @@ -240,7 +241,7 @@ void handlerFailureMarksSpanAndClosesScope( "traced client | CLIENT | true " }) void doesNotTraceNativeMessages(Behavior behavior, boolean traced) throws Exception { - JettyEndpoints.FullListener endpoint = new JettyEndpoints.FullListener(); + NativeEndpoints.FullListener endpoint = new NativeEndpoints.FullListener(); JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); openFrameHandler(frameHandler, behavior, traced); @@ -266,7 +267,7 @@ void doesNotTraceNativeMessages(Behavior behavior, boolean traced) throws Except "POJO full | pojoFull | true ", "POJO partial nonfinal | pojoPartial | false" }) - void binaryCallbackFailureMarksReceiveSpan(JettyEndpoints.EndpointEvents endpoint, boolean last) + void binaryCallbackFailureMarksReceiveSpan(NativeEndpoints.EndpointEvents endpoint, boolean last) throws Exception { endpoint.failCallback = true; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); @@ -290,7 +291,7 @@ void binaryCallbackFailureMarksReceiveSpan(JettyEndpoints.EndpointEvents endpoin "failure | true " }) void deferredBinaryCallbackCompletesReceiveSpan(boolean fail) throws Exception { - JettyEndpoints.FullListener endpoint = new JettyEndpoints.FullListener(); + NativeEndpoints.FullListener endpoint = new NativeEndpoints.FullListener(); endpoint.deferCallback = true; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); Callback.Completable callback = new Callback.Completable(); @@ -334,7 +335,7 @@ void deferredBinaryCallbackCompletesReceiveSpan(boolean fail) throws Exception { "POJO partial reverse | pojoPartial | true " }) void pendingBinaryMessagesHaveIndependentSpans( - JettyEndpoints.EndpointEvents endpoint, boolean reverse) throws Exception { + NativeEndpoints.EndpointEvents endpoint, boolean reverse) throws Exception { endpoint.deferCallback = true; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); openFrameHandler(frameHandler); @@ -384,7 +385,7 @@ void pendingBinaryMessagesHaveIndependentSpans( "failure | true " }) void partialBinaryMessageWaitsForEarlierFragmentCallbacks(boolean fail) throws Exception { - JettyEndpoints.PartialListener endpoint = new JettyEndpoints.PartialListener(); + NativeEndpoints.PartialListener endpoint = new NativeEndpoints.PartialListener(); endpoint.deferCallback = true; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); openFrameHandler(frameHandler); @@ -421,7 +422,7 @@ void partialBinaryMessageWaitsForEarlierFragmentCallbacks(boolean fail) throws E @Test void handlerFailureAfterCallbackSuccessStillMarksReceiveSpan() throws Exception { - JettyEndpoints.FullListener endpoint = new JettyEndpoints.FullListener(); + NativeEndpoints.FullListener endpoint = new NativeEndpoints.FullListener(); endpoint.failAfterCallback = true; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); openFrameHandler(frameHandler); @@ -447,8 +448,10 @@ void handlerFailureAfterCallbackSuccessStillMarksReceiveSpan() throws Exception }) void terminationFinishesFragmentedReceiveSpan( String msgType, boolean closeFrame, boolean closeHandler) throws Exception { - JettyEndpoints.EndpointEvents endpoint = - closeHandler ? new JettyEndpoints.PartialListener() : new JettyEndpoints.NoCloseEndpoint(); + NativeEndpoints.EndpointEvents endpoint = + closeHandler + ? new NativeEndpoints.PartialListener() + : new NativeEndpoints.NoCloseEndpoint(); JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); openFrameHandler(frameHandler); deliver(frameHandler, new Frame(opcode(msgType), "hello").setFin(false)); @@ -476,7 +479,7 @@ void terminationFinishesFragmentedReceiveSpan( @Test void deferredPartialBinarySuccessKeepsSpanOpenUntilFinalCallback() throws Exception { - JettyEndpoints.PartialListener endpoint = new JettyEndpoints.PartialListener(); + NativeEndpoints.PartialListener endpoint = new NativeEndpoints.PartialListener(); endpoint.deferCallback = true; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); openFrameHandler(frameHandler); @@ -515,7 +518,7 @@ void deferredPartialBinarySuccessKeepsSpanOpenUntilFinalCallback() throws Except "partial closed failure | partial | false | false | true " }) void terminationWaitsForPendingBinaryCallbacks( - JettyEndpoints.EndpointEvents endpoint, boolean last, boolean closeFrame, boolean fail) + NativeEndpoints.EndpointEvents endpoint, boolean last, boolean closeFrame, boolean fail) throws Exception { endpoint.deferCallback = true; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); From 449cb65d65595c6aa8f034e84e4144d9dc1e312f Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Wed, 7 Oct 2026 11:35:09 -0700 Subject: [PATCH 16/23] Fix Jetty EE9 Muzzle selection and verify EE modules reject Jetty 12.1+ Correct the EE9 directive name so Muzzle validates its instrumentation. Add inverse assertions for EE8, EE9, and EE10 to verify rejection outside the supported [12,12.1) range. --- .../jetty-websocket/jetty-websocket-12.0/build.gradle | 3 +++ .../jetty12/Jetty12EE9JakartaPojoWebsocketModule.java | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle index 09f90f5d2fb..338969d26f1 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle @@ -17,6 +17,7 @@ muzzle { module = 'jetty-ee8-websocket-javax-server' versions = "[12,12.1)" javaVersion = "17" + assertInverse = true } pass { name = 'jetty-websocket-12ee9' @@ -24,6 +25,7 @@ muzzle { module = 'jetty-ee9-websocket-jakarta-server' versions = "[12,12.1)" javaVersion = "17" + assertInverse = true } pass { name = 'jetty-websocket-12ee10' @@ -31,6 +33,7 @@ muzzle { module = 'jetty-ee10-websocket-jakarta-server' versions = "[12,12.1)" javaVersion = "17" + assertInverse = true } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12EE9JakartaPojoWebsocketModule.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12EE9JakartaPojoWebsocketModule.java index cf55ba3af4d..06232842e10 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12EE9JakartaPojoWebsocketModule.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12EE9JakartaPojoWebsocketModule.java @@ -15,7 +15,7 @@ public Jetty12EE9JakartaPojoWebsocketModule() { @Override public String muzzleDirective() { - return "jetty-websocket-12e9"; + return "jetty-websocket-12ee9"; } @Override From 0e4b096f22affe5a440e67e1a0b7d393c9b6f284 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Wed, 7 Oct 2026 17:17:16 -0700 Subject: [PATCH 17/23] Add tests for native Jetty WebSocket client tracing Cover client receive spans, handshake links, and callback lifecycles. Require an active parent when starting client send spans. Leave 29 tests failing pending implementation. --- .../java/NativeWebSocketHandshakeTest.java | 81 ++++++---- .../test/java/NativeWebSocketSendTest.java | 91 ++++++++++- .../src/test/java/NativeWebSocketTest.java | 149 ++++++++++++------ 3 files changed, 247 insertions(+), 74 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketHandshakeTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketHandshakeTest.java index d353a4b085e..a76028a2f85 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketHandshakeTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketHandshakeTest.java @@ -2,6 +2,8 @@ import static datadog.trace.agent.test.assertions.SpanMatcher.span; import static datadog.trace.agent.test.assertions.TraceMatcher.SORT_BY_START_TIME; import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_LENGTH; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.WEBSOCKET_MESSAGE_TYPE; import static java.nio.charset.StandardCharsets.UTF_8; @@ -10,10 +12,13 @@ import static java.util.regex.Pattern.compile; import static java.util.regex.Pattern.quote; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import datadog.trace.agent.test.AbstractInstrumentationTest; import datadog.trace.agent.test.assertions.SpanMatcher; import datadog.trace.api.DDSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink; import datadog.trace.bootstrap.instrumentation.api.SpanAttributes; import datadog.trace.core.DDSpan; @@ -60,41 +65,49 @@ void messagesLinkToHandshakesInBothDirections( + ((ServerConnector) server.getConnectors()[0]).getLocalPort() + "/receive"); ClientEndpoint clientEndpoint = new ClientEndpoint(); - Session session = client.connect(clientEndpoint, uri).get(5, SECONDS); + Session session = + runUnderTrace( + "application.connect", () -> client.connect(clientEndpoint, uri).get(5, SECONDS)); writer.waitForTraces(2); DDSpan handshake = handshake("server"); assertEquals(101, handshake.getTag("http.status_code")); - Callback.Completable sent = new Callback.Completable(); - if ("text".equals(messageType)) { - session.sendText("hello", sent); - } else { - session.sendBinary(UTF_8.encode("hello"), sent); - } - sent.get(5, SECONDS); + runUnderTrace( + "application.send", + () -> { + Callback.Completable sent = new Callback.Completable(); + if ("text".equals(messageType)) { + session.sendText("hello", sent); + } else { + session.sendBinary(UTF_8.encode("hello"), sent); + } + sent.get(5, SECONDS); + return null; + }); assertEquals("hello", clientEndpoint.reply.get(5, SECONDS)); + assertNotNull(clientEndpoint.messageSpan, "The client reply handler needs a receive span"); + assertEquals("websocket.receive", clientEndpoint.messageSpan.getOperationName().toString()); + assertEquals("consumer", clientEndpoint.messageSpan.getTag("span.kind")); + assertNull(activeSpan()); assertTraces( - trace(span().type(DDSpanTypes.HTTP_CLIENT).error(false)), + trace( + SORT_BY_START_TIME, + span().operationName("application.connect").root(), + span().type(DDSpanTypes.HTTP_CLIENT).childOfPrevious().error(false)), trace(span().type(DDSpanTypes.HTTP_SERVER).error(false)), - trace(sendSpan(handshake("client")).root()), trace( SORT_BY_START_TIME, - span() - .root() - .operationName(compile(quote("websocket.receive"))) - .resourceName(compile(quote("websocket /receive"))) - .type(DDSpanTypes.WEBSOCKET) - .error(false) - .links( - to(handshake) - .traceFlags( - handshake.getSamplingPriority() > 0 - ? AgentSpanLink.SAMPLED_FLAG - : AgentSpanLink.DEFAULT_FLAGS) - .attributes( - SpanAttributes.builder().put("dd.kind", "executed_from").build())), - sendSpan(handshake).childOfPrevious())); + span().operationName("application.send").root(), + messageSpan(handshake("client"), "send").childOfPrevious()), + trace( + SORT_BY_START_TIME, + messageSpan(handshake, "receive").root(), + messageSpan(handshake, "send").childOfPrevious()), + trace( + SORT_BY_START_TIME, + messageSpan(handshake("client"), "receive").root(), + span().operationName("client.handler").childOfPrevious())); assertEquals(singletonList("hello"), endpoint.messages); assertEquals(1, endpoint.messageSpans.size()); assertEquals( @@ -117,9 +130,9 @@ private static DDSpan handshake(String kind) { .orElseThrow(() -> new AssertionError("Missing " + kind + " handshake span")); } - private static SpanMatcher sendSpan(DDSpan handshake) { + private static SpanMatcher messageSpan(DDSpan handshake, String operation) { return span() - .operationName(compile(quote("websocket.send"))) + .operationName(compile(quote("websocket." + operation))) .resourceName(compile(quote("websocket /receive"))) .type(DDSpanTypes.WEBSOCKET) .error(false) @@ -129,15 +142,25 @@ private static SpanMatcher sendSpan(DDSpan handshake) { handshake.getSamplingPriority() > 0 ? AgentSpanLink.SAMPLED_FLAG : AgentSpanLink.DEFAULT_FLAGS) - .attributes(SpanAttributes.builder().put("dd.kind", "resuming").build())); + .attributes( + SpanAttributes.builder() + .put("dd.kind", "receive".equals(operation) ? "executed_from" : "resuming") + .build())); } public static class ClientEndpoint implements Session.Listener.AutoDemanding { final CompletableFuture reply = new CompletableFuture<>(); + AgentSpan messageSpan; @Override public void onWebSocketText(String message) { - reply.complete(message); + messageSpan = activeSpan(); + runUnderTrace( + "client.handler", + () -> { + reply.complete(message); + return null; + }); } } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java index c8909a0d17a..26bb6a04cb0 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java @@ -17,6 +17,7 @@ import static java.util.regex.Pattern.compile; import static java.util.regex.Pattern.quote; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; @@ -50,6 +51,88 @@ import org.tabletest.junit.TableTest; class NativeWebSocketSendTest extends AbstractInstrumentationTest { + @TableTest({ + "scenario | behavior | binary | partial | parent", + "client text without parent | CLIENT | false | false | false ", + "client binary without parent | CLIENT | true | false | false ", + "client partial text without parent | CLIENT | false | true | false ", + "client partial binary no parent | CLIENT | true | true | false ", + "client text with parent | CLIENT | false | false | true ", + "client binary with parent | CLIENT | true | false | true ", + "client partial text with parent | CLIENT | false | true | true ", + "client partial binary with parent | CLIENT | true | true | true ", + "server text without parent | SERVER | false | false | false ", + "server binary without parent | SERVER | true | false | false ", + "server partial text without parent | SERVER | false | true | false ", + "server partial binary no parent | SERVER | true | true | false " + }) + void clientSendsRequireActiveParent( + Behavior behavior, boolean binary, boolean partial, boolean parent) throws Exception { + Connection connection = new Connection(true, behavior); + RecordingCallback callback = new RecordingCallback(); + if (parent) { + runUnderTrace( + "application", + () -> { + AgentSpan application = activeSpan(); + connection.send(binary, partial, "hello", true, callback); + assertSame(application, activeSpan()); + return null; + }); + } else { + connection.send(binary, partial, "hello", true, callback); + } + connection.callbacks.get(0).succeeded(); + callback.get(5, SECONDS); + assertNull(activeSpan()); + AgentSpan sent = connection.spans.get(0); + assertSame(sent, callback.span); + if (behavior == Behavior.CLIENT && !parent) { + assertNull(sent, "Client sends without an active parent must not start a trace"); + assertTraces(trace(handshakeSpan())); + } else if (parent) { + assertNotNull(sent); + assertTraces( + trace(handshakeSpan()), + trace( + SORT_BY_START_TIME, + span().operationName("application").root(), + sendSpan(connection, binary, 5, 1).childOfPrevious())); + } else { + assertNotNull(sent); + assertTraces(trace(handshakeSpan()), trace(sendSpan(connection, binary, 5, 1).root())); + } + } + + @TableTest({ + "scenario | binary", + "text | false ", + "binary | true " + }) + void clientPartialSendKeepsParentUntilLastFragment(boolean binary) throws Exception { + Connection connection = new Connection(true, Behavior.CLIENT); + runUnderTrace( + "application", + () -> { + connection.send(binary, true, "hello ", false, null); + return null; + }); + assertNull(activeSpan()); + connection.send(binary, true, "world", true, null); + connection.callbacks.get(1).succeeded(); + connection.callbacks.get(0).succeeded(); + + assertNotNull(connection.spans.get(0)); + assertSame(connection.spans.get(0), connection.spans.get(1)); + assertNull(activeSpan()); + assertTraces( + trace(handshakeSpan()), + trace( + SORT_BY_START_TIME, + span().operationName("application").root(), + sendSpan(connection, binary, 11, 2).childOfPrevious())); + } + @TableTest({ "scenario | binary | partial | synchronous", "text async | false | false | false ", @@ -377,11 +460,17 @@ private static class Connection extends CoreSession.Empty { final IllegalStateException failure = new IllegalStateException("send failed"); final JettyWebSocketFrameHandler handler; final Session session; + final Behavior behavior; DDSpan handshake; boolean synchronous; String failureMode = ""; Connection(boolean traced) throws Exception { + this(traced, Behavior.SERVER); + } + + Connection(boolean traced, Behavior behavior) throws Exception { + this.behavior = behavior; ServerFrameHandlerFactory factory = new ServerFrameHandlerFactory(mock(WebSocketContainer.class), new WebSocketComponents()); handler = factory.newJettyFrameHandler(new Session.Listener.AutoDemanding() {}); @@ -408,7 +497,7 @@ private static class Connection extends CoreSession.Empty { @Override public Behavior getBehavior() { - return Behavior.SERVER; + return behavior; } @Override diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java index 46379902cb6..8fd278ee6f9 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java @@ -3,6 +3,7 @@ import static datadog.trace.agent.test.assertions.TagsMatcher.error; import static datadog.trace.agent.test.assertions.TagsMatcher.includes; import static datadog.trace.agent.test.assertions.TagsMatcher.tag; +import static datadog.trace.agent.test.assertions.TraceMatcher.SORT_BY_START_TIME; import static datadog.trace.agent.test.assertions.TraceMatcher.trace; import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace; import static datadog.trace.api.DDTags.DECISION_MAKER_INHERITED; @@ -71,18 +72,26 @@ public class NativeWebSocketTest extends AbstractInstrumentationTest { private static final String URL = "ws://inmemory/test/param"; @TableTest({ - "scenario | msgType | fragmented | fail ", - "reader | text | false | false", - "input stream | binary | false | false", - "fragmented reader | text | true | false", - "fragmented input stream | binary | true | false", - "reader failure | text | false | true ", - "input stream failure | binary | false | true ", - "fragmented reader failure | text | true | true ", - "fragmented stream failure | binary | true | true " + "scenario | behavior | msgType | fragmented | fail ", + "reader | SERVER | text | false | false", + "input stream | SERVER | binary | false | false", + "fragmented reader | SERVER | text | true | false", + "fragmented input stream | SERVER | binary | true | false", + "reader failure | SERVER | text | false | true ", + "input stream failure | SERVER | binary | false | true ", + "fragmented reader failure | SERVER | text | true | true ", + "fragmented stream failure | SERVER | binary | true | true ", + "client reader | CLIENT | text | false | false", + "client input stream | CLIENT | binary | false | false", + "client fragmented reader | CLIENT | text | true | false", + "client fragmented input stream | CLIENT | binary | true | false", + "client reader failure | CLIENT | text | false | true ", + "client input stream failure | CLIENT | binary | false | true ", + "client fragmented reader failure | CLIENT | text | true | true ", + "client fragmented stream failure | CLIENT | binary | true | true " }) - void streamingMessagesHaveReceiveSpan(String msgType, boolean fragmented, boolean fail) - throws Exception { + void streamingMessagesHaveReceiveSpan( + Behavior behavior, String msgType, boolean fragmented, boolean fail) throws Exception { ExecutorService executor = newSingleThreadExecutor(); try { WebSocketComponents components = @@ -90,7 +99,7 @@ void streamingMessagesHaveReceiveSpan(String msgType, boolean fragmented, boolea NativeEndpoints.StreamingEndpoint endpoint = new NativeEndpoints.StreamingEndpoint(); endpoint.failMessages = fail; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); - openFrameHandler(frameHandler, Behavior.SERVER, true, components); + openFrameHandler(frameHandler, behavior, true, components); if (fragmented) { deliver(frameHandler, new Frame(opcode(msgType), "hello ").setFin(false)); @@ -162,24 +171,29 @@ void nativeJettyAdvices(NativeEndpoints.EndpointEvents endpoint, String msgType) } @TableTest({ - "scenario | endpoint | msgType", - "partial text | partial | text ", - "partial binary | partial | binary ", - "POJO partial text | pojoPartial | text ", - "POJO partial binary | pojoPartial | binary ", - "boxed partial text | boxedPartial | text ", - "boxed partial binary | boxedPartial | binary " + "scenario | behavior | endpoint | msgType", + "partial text | SERVER | partial | text ", + "partial binary | SERVER | partial | binary ", + "POJO partial text | SERVER | pojoPartial | text ", + "POJO partial binary | SERVER | pojoPartial | binary ", + "boxed partial text | SERVER | boxedPartial | text ", + "boxed partial binary | SERVER | boxedPartial | binary ", + "client partial text | CLIENT | partial | text ", + "client partial binary | CLIENT | partial | binary ", + "client POJO partial text | CLIENT | pojoPartial | text ", + "client POJO partial binary | CLIENT | pojoPartial | binary " }) - void fragmentedMessagesShareSpan(NativeEndpoints.EndpointEvents endpoint, String msgType) - throws Exception { + void fragmentedMessagesShareSpan( + Behavior behavior, NativeEndpoints.EndpointEvents endpoint, String msgType) throws Exception { JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); - openFrameHandler(frameHandler); + openFrameHandler(frameHandler, behavior, true); byte opcode = opcode(msgType); deliver(frameHandler, new Frame(opcode, "hello ").setFin(false)); assertNull(activeSpan()); assertEquals(singletonList("hello "), endpoint.messages); + assertNotNull(endpoint.messageSpans.get(0)); assertEquals(1, writer.size()); deliver(frameHandler, new Frame(OpCode.CONTINUATION, "world")); @@ -198,14 +212,21 @@ void fragmentedMessagesShareSpan(NativeEndpoints.EndpointEvents endpoint, String } @TableTest({ - "scenario | endpoint | msgType | last | operation ", - "full text | full | text | true | websocket.receive", - "full binary | full | binary | true | websocket.receive", - "partial text | partial | text | false | websocket.receive", - "POJO close | pojoFull | | true | websocket.close " + "scenario | behavior | endpoint | msgType | last | operation ", + "full text | SERVER | full | text | true | websocket.receive", + "full binary | SERVER | full | binary | true | websocket.receive", + "partial text | SERVER | partial | text | false | websocket.receive", + "POJO close | SERVER | pojoFull | | true | websocket.close ", + "client full text | CLIENT | full | text | true | websocket.receive", + "client full binary | CLIENT | full | binary | true | websocket.receive", + "client partial text | CLIENT | partial | text | false | websocket.receive" }) void handlerFailureMarksSpanAndClosesScope( - NativeEndpoints.EndpointEvents endpoint, String msgType, boolean last, String operation) + Behavior behavior, + NativeEndpoints.EndpointEvents endpoint, + String msgType, + boolean last, + String operation) throws Exception { endpoint.failClose = msgType == null; endpoint.failMessages = msgType != null; @@ -214,13 +235,16 @@ void handlerFailureMarksSpanAndClosesScope( ? CloseStatus.toFrame(CloseStatus.NORMAL, "bye") : new Frame(opcode(msgType), "hello").setFin(last); JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); - openFrameHandler(frameHandler); + openFrameHandler(frameHandler, behavior, true); ExecutionException error = assertThrows(ExecutionException.class, () -> deliver(frameHandler, frame)); assertNotNull(error.getCause()); assertNull(activeSpan()); + if (msgType != null) { + assertNotNull(endpoint.messageSpans.get(0)); + } assertTraces( trace(handshakeSpan()), trace( @@ -236,15 +260,15 @@ void handlerFailureMarksSpanAndClosesScope( } @TableTest({ - "scenario | behavior | traced", - "untraced server | SERVER | false ", - "traced client | CLIENT | true " + "scenario | behavior", + "untraced server | SERVER ", + "untraced client | CLIENT " }) - void doesNotTraceNativeMessages(Behavior behavior, boolean traced) throws Exception { + void doesNotTraceNativeMessagesWithoutHandshake(Behavior behavior) throws Exception { NativeEndpoints.FullListener endpoint = new NativeEndpoints.FullListener(); JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); - openFrameHandler(frameHandler, behavior, traced); + openFrameHandler(frameHandler, behavior, false); deliver(frameHandler, new Frame(OpCode.TEXT, "hello")); deliver(frameHandler, CloseStatus.toFrame(CloseStatus.NORMAL, "bye")); @@ -252,11 +276,7 @@ void doesNotTraceNativeMessages(Behavior behavior, boolean traced) throws Except assertEquals(singletonList(null), endpoint.messageSpans); assertEquals(1000, endpoint.closeCode); assertNull(activeSpan()); - if (traced) { - assertTraces(trace(handshakeSpan())); - } else { - assertTraces(); - } + assertTraces(); } @TableTest({ @@ -286,16 +306,19 @@ void binaryCallbackFailureMarksReceiveSpan(NativeEndpoints.EndpointEvents endpoi } @TableTest({ - "scenario | fail ", - "success | false", - "failure | true " + "scenario | behavior | fail ", + "success | SERVER | false", + "failure | SERVER | true ", + "client success | CLIENT | false", + "client failure | CLIENT | true " }) - void deferredBinaryCallbackCompletesReceiveSpan(boolean fail) throws Exception { + void deferredBinaryCallbackCompletesReceiveSpan(Behavior behavior, boolean fail) + throws Exception { NativeEndpoints.FullListener endpoint = new NativeEndpoints.FullListener(); endpoint.deferCallback = true; JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); Callback.Completable callback = new Callback.Completable(); - openFrameHandler(frameHandler); + openFrameHandler(frameHandler, behavior, true); frameHandler.onFrame(new Frame(OpCode.BINARY, "hello"), callback); @@ -318,11 +341,49 @@ void deferredBinaryCallbackCompletesReceiveSpan(boolean fail) throws Exception { assertTrue(callback.isDone()); assertEquals(fail, callback.isCompletedExceptionally()); + assertNotNull(endpoint.messageSpans.get(0)); writer.waitForTraces(2); assertEquals(fail, endpoint.messageSpans.get(0).isError()); assertNull(activeSpan()); } + @TableTest({ + "scenario | endpoint | msgType", + "client text | full | text ", + "client binary | full | binary ", + "client POJO text | pojoFull | text ", + "client POJO binary | pojoFull | binary " + }) + void clientReceivesLinkToHandshakeInsteadOfApplicationRoot( + NativeEndpoints.EndpointEvents endpoint, String msgType) throws Exception { + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + runUnderTrace( + "application", + () -> { + openFrameHandler(frameHandler, Behavior.CLIENT, true); + return null; + }); + assertNull(activeSpan()); + + deliver(frameHandler, new Frame(opcode(msgType), "hello")); + + assertEquals(singletonList("hello"), endpoint.messages); + assertNotNull(endpoint.messageSpans.get(0)); + assertNull(activeSpan()); + writer.waitForTraces(2); + DDSpan handshake = + writer.get(0).stream() + .filter(s -> "parent".contentEquals(s.getOperationName())) + .findFirst() + .orElseThrow(() -> new AssertionError("Missing handshake span")); + assertTraces( + trace( + SORT_BY_START_TIME, + span().operationName("application").root(), + span().operationName("parent").childOfPrevious()), + trace(receiveSpan(handshake, msgType, 5, 1))); + } + @TableTest({ "scenario | endpoint | reverse", "full forward | full | false ", From 5e43379d3047c33b4effe01b51efbd2db698574a Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Wed, 7 Oct 2026 17:23:55 -0700 Subject: [PATCH 18/23] Trace native Jetty WebSocket client receives Link client receive spans to the HTTP upgrade span. Require an active parent when starting client send spans, while preserving the span across subsequent fragments. --- .../Jetty12NativeFrameHandlerInstrumentation.java | 4 +--- .../websocket/jetty12/NativeMethodHandleWrappers.java | 3 ++- .../websocket/jetty12/NativeSendContext.java | 9 ++++++++- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeFrameHandlerInstrumentation.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeFrameHandlerInstrumentation.java index 94f87cd04f3..2dd238d8a59 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeFrameHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/Jetty12NativeFrameHandlerInstrumentation.java @@ -66,9 +66,7 @@ public static void onEnter( @Advice.Argument(1) WebSocketSession session, @Advice.Argument(value = 2, readOnly = false) MethodHandle handle) { AgentSpan span = activeSpan(); - if (span != null - && session.getCoreSession().getBehavior() == Behavior.SERVER - && handle != null) { + if (span != null && handle != null) { handle = NativeMethodHandleWrappers.wrapMessage( handle, diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java index e3ef1cfae98..11a7e4d4865 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java @@ -22,6 +22,7 @@ import java.util.Set; import org.eclipse.jetty.websocket.api.Callback; import org.eclipse.jetty.websocket.api.Session; +import org.eclipse.jetty.websocket.core.Behavior; import org.eclipse.jetty.websocket.core.CoreSession; /** @@ -217,7 +218,7 @@ public ReceiveContexts(AgentSpan span, CoreSession session) { if (Config.get().isWebsocketMessagesInheritSampling()) { span.forceSamplingDecision(); } - handshakeSpan = span.getLocalRootSpan(); + handshakeSpan = session.getBehavior() == Behavior.CLIENT ? span : span.getLocalRootSpan(); sessionId = Integer.toHexString(System.identityHashCode(session)); } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java index 365ce5619da..6d90910970a 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.websocket.jetty12; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_BINARY; import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_TEXT; @@ -21,6 +22,7 @@ public class NativeSendContext { private final AgentSpan handshakeSpan; private final String sessionId; + private final boolean client; private final Set pending = new HashSet<>(); private Message partialMessage; private boolean closed; @@ -29,7 +31,8 @@ public NativeSendContext(AgentSpan span, CoreSession session) { if (Config.get().isWebsocketMessagesInheritSampling()) { span.forceSamplingDecision(); } - handshakeSpan = session.getBehavior() == Behavior.CLIENT ? span : span.getLocalRootSpan(); + client = session.getBehavior() == Behavior.CLIENT; + handshakeSpan = client ? span : span.getLocalRootSpan(); sessionId = Integer.toHexString(System.identityHashCode(session)); } @@ -45,6 +48,10 @@ public synchronized SendCallback start( : binary ? ((ByteBuffer) payload).remaining() : ((String) payload).length(); Message message = partial ? partialMessage : null; if (message == null || !type.equals(message.getMessageType())) { + // Later fragments keep the span created under the first fragment's parent. + if (client && activeSpan() == null) { + return null; + } message = new Message(handshakeSpan, sessionId); pending.add(message); } From 24c6afeb4315c8829ade1df4dcc05fa31978ddf7 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Wed, 7 Oct 2026 17:41:27 -0700 Subject: [PATCH 19/23] Reduce native Jetty WebSocket tracing overhead Let callbacks own pending messages and retain only the unfinished fragment sequence in connection state. Remove redundant error-handling locks. --- .../jetty12/NativeMethodHandleWrappers.java | 15 +++----- .../websocket/jetty12/NativeSendContext.java | 36 ++++++++----------- 2 files changed, 18 insertions(+), 33 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java index 11a7e4d4865..4c4594101dd 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java @@ -17,9 +17,6 @@ import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodType; import java.nio.ByteBuffer; -import java.util.HashSet; -import java.util.Iterator; -import java.util.Set; import org.eclipse.jetty.websocket.api.Callback; import org.eclipse.jetty.websocket.api.Session; import org.eclipse.jetty.websocket.core.Behavior; @@ -186,7 +183,6 @@ public static void onBinary( BinaryMessage message = contexts.currentBinary; if (message == null) { message = new BinaryMessage(contexts.handshakeSpan, contexts.sessionId); - contexts.pendingBinary.add(message); } message.pendingCallbacks++; message.complete = last; @@ -212,7 +208,6 @@ public static class ReceiveContexts { private final String sessionId; private HandlerContext.Receiver text; private BinaryMessage currentBinary; - private final Set pendingBinary = new HashSet<>(); public ReceiveContexts(AgentSpan span, CoreSession session) { if (Config.get().isWebsocketMessagesInheritSampling()) { @@ -235,16 +230,15 @@ public synchronized void finish() { DECORATE.onFrameEnd(text); } } - for (Iterator iterator = pendingBinary.iterator(); iterator.hasNext(); ) { - BinaryMessage message = iterator.next(); - // Close ends the message, but pending callbacks still determine the receive outcome. + // Callbacks own completed messages; only the unfinished fragment sequence needs closing. + BinaryMessage message = currentBinary; + currentBinary = null; + if (message != null) { message.complete = true; if (message.pendingCallbacks == 0) { DECORATE.onFrameEnd(message); - iterator.remove(); } } - currentBinary = null; } } @@ -313,7 +307,6 @@ private void finish() { released = true; if (--message.pendingCallbacks == 0 && message.complete) { DECORATE.onFrameEnd(message); - contexts.pendingBinary.remove(message); } } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java index 6d90910970a..a53c2153951 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java @@ -11,9 +11,6 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; import java.nio.ByteBuffer; -import java.util.HashSet; -import java.util.Iterator; -import java.util.Set; import org.eclipse.jetty.websocket.api.Callback; import org.eclipse.jetty.websocket.core.Behavior; import org.eclipse.jetty.websocket.core.CoreSession; @@ -23,7 +20,6 @@ public class NativeSendContext { private final AgentSpan handshakeSpan; private final String sessionId; private final boolean client; - private final Set pending = new HashSet<>(); private Message partialMessage; private boolean closed; @@ -53,7 +49,6 @@ public synchronized SendCallback start( return null; } message = new Message(handshakeSpan, sessionId); - pending.add(message); } AgentSpan span = DECORATE.startOutboundFrameSpan(message, type, size); message.pendingCallbacks++; @@ -66,17 +61,16 @@ public synchronized SendCallback start( public synchronized void finish() { closed = true; - for (Iterator iterator = pending.iterator(); iterator.hasNext(); ) { - Message message = iterator.next(); - // Close ends the message, but pending callbacks still determine the send outcome. + // Callbacks own completed messages; only the unfinished fragment sequence needs closing. + Message message = partialMessage; + partialMessage = null; + if (message != null) { message.complete = true; if (message.pendingCallbacks == 0) { message.finished = true; DECORATE.onFrameEnd(message); - iterator.remove(); } } - partialMessage = null; } public static class Message extends HandlerContext.Sender { @@ -121,9 +115,7 @@ public void fail(Throwable failure) { } private void complete(Throwable failure) { - synchronized (context) { - onError(failure); - } + onError(failure); try (ContextScope ignored = activateSpan(span)) { try { if (delegate != null) { @@ -134,9 +126,7 @@ private void complete(Throwable failure) { } } } catch (Throwable t) { - synchronized (context) { - onError(t); - } + onError(t); throw t; } } finally { @@ -148,10 +138,10 @@ private void complete(Throwable failure) { } public void onMethodExit(ContextScope scope, Throwable failure) { - synchronized (context) { - try { - onError(failure); - } finally { + try { + onError(failure); + } finally { + synchronized (context) { if (scope != null) { scope.close(); } @@ -165,8 +155,11 @@ public void onMethodExit(ContextScope scope, Throwable failure) { } private void onError(Throwable failure) { + if (failure == null) { + return; + } synchronized (context) { - if (failure != null && !message.finished) { + if (!message.finished) { DECORATE.onError(span, failure); message.complete = true; if (context.partialMessage == message) { @@ -183,7 +176,6 @@ private void finish() { if (--message.pendingCallbacks == 0 && message.complete && !message.finished) { message.finished = true; DECORATE.onFrameEnd(message); - context.pending.remove(message); } } } From a66ef77d353ef3a22f59cf041ed756f4e5173b63 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Wed, 7 Oct 2026 17:41:27 -0700 Subject: [PATCH 20/23] Simplify native Jetty WebSocket tracing tests Split mixed scenarios and remove redundant endpoint cases. Verify that rejected fragments preserve the unfinished message for close cleanup. --- .../test/java/NativeWebSocketSendTest.java | 130 ++++++++++++------ .../src/test/java/NativeWebSocketTest.java | 90 ++++++------ 2 files changed, 136 insertions(+), 84 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java index 26bb6a04cb0..6f8bbaf03e8 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java @@ -17,6 +17,7 @@ import static java.util.regex.Pattern.compile; import static java.util.regex.Pattern.quote; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertNull; @@ -46,62 +47,65 @@ import org.eclipse.jetty.websocket.core.CoreSession; import org.eclipse.jetty.websocket.core.Frame; import org.eclipse.jetty.websocket.core.WebSocketComponents; +import org.eclipse.jetty.websocket.core.exception.ProtocolException; import org.eclipse.jetty.websocket.server.internal.ServerFrameHandlerFactory; import org.junit.jupiter.api.Test; import org.tabletest.junit.TableTest; class NativeWebSocketSendTest extends AbstractInstrumentationTest { @TableTest({ - "scenario | behavior | binary | partial | parent", - "client text without parent | CLIENT | false | false | false ", - "client binary without parent | CLIENT | true | false | false ", - "client partial text without parent | CLIENT | false | true | false ", - "client partial binary no parent | CLIENT | true | true | false ", - "client text with parent | CLIENT | false | false | true ", - "client binary with parent | CLIENT | true | false | true ", - "client partial text with parent | CLIENT | false | true | true ", - "client partial binary with parent | CLIENT | true | true | true ", - "server text without parent | SERVER | false | false | false ", - "server binary without parent | SERVER | true | false | false ", - "server partial text without parent | SERVER | false | true | false ", - "server partial binary no parent | SERVER | true | true | false " + "scenario | binary | partial", + "text | false | false ", + "binary | true | false ", + "partial text | false | true ", + "partial binary | true | true " }) - void clientSendsRequireActiveParent( - Behavior behavior, boolean binary, boolean partial, boolean parent) throws Exception { - Connection connection = new Connection(true, behavior); + void clientSendsWithoutParentAreNotTraced(boolean binary, boolean partial) throws Exception { + Connection connection = new Connection(true, Behavior.CLIENT); RecordingCallback callback = new RecordingCallback(); - if (parent) { - runUnderTrace( - "application", - () -> { - AgentSpan application = activeSpan(); - connection.send(binary, partial, "hello", true, callback); - assertSame(application, activeSpan()); - return null; - }); - } else { - connection.send(binary, partial, "hello", true, callback); - } + + connection.send(binary, partial, "hello", true, callback); + connection.callbacks.get(0).succeeded(); + callback.get(5, SECONDS); + + assertNull(connection.spans.get(0)); + assertNull(callback.span); + assertNull(activeSpan()); + assertTraces(trace(handshakeSpan())); + } + + @TableTest({ + "scenario | binary | partial", + "text | false | false ", + "binary | true | false ", + "partial text | false | true ", + "partial binary | true | true " + }) + void clientSendsUseActiveParent(boolean binary, boolean partial) throws Exception { + Connection connection = new Connection(true, Behavior.CLIENT); + RecordingCallback callback = new RecordingCallback(); + + runUnderTrace( + "application", + () -> { + AgentSpan application = activeSpan(); + connection.send(binary, partial, "hello", true, callback); + assertSame(application, activeSpan()); + return null; + }); connection.callbacks.get(0).succeeded(); callback.get(5, SECONDS); + assertNull(activeSpan()); AgentSpan sent = connection.spans.get(0); + assertNotNull(sent); assertSame(sent, callback.span); - if (behavior == Behavior.CLIENT && !parent) { - assertNull(sent, "Client sends without an active parent must not start a trace"); - assertTraces(trace(handshakeSpan())); - } else if (parent) { - assertNotNull(sent); - assertTraces( - trace(handshakeSpan()), - trace( - SORT_BY_START_TIME, - span().operationName("application").root(), - sendSpan(connection, binary, 5, 1).childOfPrevious())); - } else { - assertNotNull(sent); - assertTraces(trace(handshakeSpan()), trace(sendSpan(connection, binary, 5, 1).root())); - } + assertTraces( + trace(handshakeSpan()), + trace( + SORT_BY_START_TIME, + span().operationName("application").root(), + sendSpan(connection, binary, 5, 1).childOfPrevious())); } @TableTest({ @@ -142,11 +146,12 @@ void clientPartialSendKeepsParentUntilLastFragment(boolean binary) throws Except "partial text | false | true | false ", "partial binary | true | true | false " }) - void sendsFinishOnCallback(boolean binary, boolean partial, boolean synchronous) - throws Exception { + void serverSendsWithoutParentFinishOnCallback( + boolean binary, boolean partial, boolean synchronous) throws Exception { Connection connection = new Connection(true); connection.synchronous = synchronous; RecordingCallback callback = new RecordingCallback(); + assertNull(activeSpan()); connection.send(binary, partial, "hello", true, callback); assertNull(activeSpan()); AgentSpan sent = connection.spans.get(0); @@ -331,6 +336,41 @@ void closeFinishesIncompleteSendsAndWaitsForPendingCallbacks() throws Exception assertNull(activeSpan()); } + @TableTest({ + "scenario | binary", + "text rejects binary | false ", + "binary rejects text | true " + }) + void rejectedPartialSendPreservesUnfinishedMessage(boolean binary) throws Exception { + Connection connection = new Connection(true); + connection.send(binary, true, "hello", false, null); + connection.callbacks.get(0).succeeded(); + AgentSpan unfinished = connection.spans.get(0); + RecordingCallback rejected = new RecordingCallback(); + + connection.send(!binary, true, "rejected", false, rejected); + + assertInstanceOf(ProtocolException.class, rejected.failure); + assertNotSame(unfinished, rejected.span); + SpanMatcher failedSend = + span() + .operationName(compile(quote("websocket.send"))) + .type(DDSpanTypes.WEBSOCKET) + .root() + .error(); + assertTraces(trace(handshakeSpan()), trace(failedSend)); + + Callback.Completable closed = new Callback.Completable(); + connection.handler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), closed); + closed.get(5, SECONDS); + + assertTraces( + trace(handshakeSpan()), + trace(sendSpan(connection, binary, 5, 1).root()), + trace(failedSend)); + assertNull(activeSpan()); + } + @TableTest({ "scenario | partial | failed", "full success | false | false ", diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java index 8fd278ee6f9..7e6dbbddcdf 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java @@ -212,28 +212,19 @@ void fragmentedMessagesShareSpan( } @TableTest({ - "scenario | behavior | endpoint | msgType | last | operation ", - "full text | SERVER | full | text | true | websocket.receive", - "full binary | SERVER | full | binary | true | websocket.receive", - "partial text | SERVER | partial | text | false | websocket.receive", - "POJO close | SERVER | pojoFull | | true | websocket.close ", - "client full text | CLIENT | full | text | true | websocket.receive", - "client full binary | CLIENT | full | binary | true | websocket.receive", - "client partial text | CLIENT | partial | text | false | websocket.receive" + "scenario | behavior | endpoint | msgType | last ", + "full text | SERVER | full | text | true ", + "full binary | SERVER | full | binary | true ", + "partial text | SERVER | partial | text | false", + "client full text | CLIENT | full | text | true ", + "client full binary | CLIENT | full | binary | true ", + "client partial text | CLIENT | partial | text | false" }) - void handlerFailureMarksSpanAndClosesScope( - Behavior behavior, - NativeEndpoints.EndpointEvents endpoint, - String msgType, - boolean last, - String operation) + void messageHandlerFailureMarksSpanAndClosesScope( + Behavior behavior, NativeEndpoints.EndpointEvents endpoint, String msgType, boolean last) throws Exception { - endpoint.failClose = msgType == null; - endpoint.failMessages = msgType != null; - Frame frame = - endpoint.failClose - ? CloseStatus.toFrame(CloseStatus.NORMAL, "bye") - : new Frame(opcode(msgType), "hello").setFin(last); + endpoint.failMessages = true; + Frame frame = new Frame(opcode(msgType), "hello").setFin(last); JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); openFrameHandler(frameHandler, behavior, true); @@ -242,14 +233,41 @@ void handlerFailureMarksSpanAndClosesScope( assertNotNull(error.getCause()); assertNull(activeSpan()); - if (msgType != null) { - assertNotNull(endpoint.messageSpans.get(0)); - } + assertNotNull(endpoint.messageSpans.get(0)); + assertTraces( + trace(handshakeSpan()), + trace( + receiveSpan( + handshake(), + msgType, + 5, + 1, + error(IllegalStateException.class, "handler failed")) + .error())); + assertInstanceOf(String.class, writer.get(1).get(0).getTag("error.stack")); + } + + @Test + void closeHandlerFailureMarksSpanAndClosesScope() throws Exception { + NativeEndpoints.PojoFullEndpoint endpoint = new NativeEndpoints.PojoFullEndpoint(); + endpoint.failClose = true; + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + + ExecutionException error = + assertThrows( + ExecutionException.class, + () -> deliver(frameHandler, CloseStatus.toFrame(CloseStatus.NORMAL, "bye"))); + + assertNotNull(error.getCause()); + assertNull(activeSpan()); + assertEquals(1000, endpoint.closeCode); + assertEquals("bye", endpoint.closeReason); assertTraces( trace(handshakeSpan()), trace( span() - .operationName(compile(quote(operation))) + .operationName(compile(quote("websocket.close"))) .resourceName(compile(quote("websocket /test/param"))) .type(DDSpanTypes.WEBSOCKET) .error())); @@ -280,12 +298,10 @@ void doesNotTraceNativeMessagesWithoutHandshake(Behavior behavior) throws Except } @TableTest({ - "scenario | endpoint | last ", - "full | full | true ", - "partial final | partial | true ", - "partial nonfinal | partial | false", - "POJO full | pojoFull | true ", - "POJO partial nonfinal | pojoPartial | false" + "scenario | endpoint | last ", + "full | full | true ", + "partial final | partial | true ", + "partial nonfinal | partial | false" }) void binaryCallbackFailureMarksReceiveSpan(NativeEndpoints.EndpointEvents endpoint, boolean last) throws Exception { @@ -385,15 +401,11 @@ void clientReceivesLinkToHandshakeInsteadOfApplicationRoot( } @TableTest({ - "scenario | endpoint | reverse", - "full forward | full | false ", - "full reverse | full | true ", - "partial forward | partial | false ", - "partial reverse | partial | true ", - "POJO full forward | pojoFull | false ", - "POJO full reverse | pojoFull | true ", - "POJO partial forward | pojoPartial | false ", - "POJO partial reverse | pojoPartial | true " + "scenario | endpoint | reverse", + "full forward | full | false ", + "full reverse | full | true ", + "partial forward | partial | false ", + "partial reverse | partial | true " }) void pendingBinaryMessagesHaveIndependentSpans( NativeEndpoints.EndpointEvents endpoint, boolean reverse) throws Exception { From a39a803377c6d6bed92db73318abeffa95c69d23 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 13 Oct 2026 10:58:36 -0700 Subject: [PATCH 21/23] Fix Jetty 12 WebSocket text receive spans finishing before handler exit Defer receive span completion until active text handlers exit, preserving handler duration and errors when the connection closes concurrently. --- .../jetty12/NativeMethodHandleWrappers.java | 39 ++++++---- .../src/test/java/NativeEndpoints.java | 4 + .../src/test/java/NativeWebSocketTest.java | 76 +++++++++++++++++++ 3 files changed, 105 insertions(+), 14 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java index 4c4594101dd..3529f86027a 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java @@ -33,7 +33,7 @@ public class NativeMethodHandleWrappers { NativeMethodHandleWrappers.class, "onText", MethodHandle.class, - HandlerContext.Receiver.class, + ReceiveContexts.class, boolean.class, String.class, boolean.class); @@ -123,14 +123,17 @@ private static ContextScope startMessage( public static void onText( MethodHandle delegate, - HandlerContext.Receiver context, + ReceiveContexts contexts, boolean partial, String payload, boolean last) throws Throwable { + HandlerContext.Receiver context; boolean finish = last; ContextScope scope; - synchronized (context) { + synchronized (contexts) { + context = contexts.text; + contexts.activeTextHandlers++; scope = startMessage(context, payload, partial); } try (ContextScope ignored = scope) { @@ -138,16 +141,16 @@ public static void onText( delegate.invokeExact(payload, last); } catch (Throwable t) { finish = true; - synchronized (context) { + synchronized (contexts) { DECORATE.onError(context.getWebsocketSpan(), t); } throw t; } } finally { - if (finish) { - synchronized (context) { - DECORATE.onFrameEnd(context); - } + synchronized (contexts) { + contexts.activeTextHandlers--; + contexts.textFinishRequested |= finish; + contexts.finishTextIfIdle(); } } } @@ -207,6 +210,8 @@ public static class ReceiveContexts { private final AgentSpan handshakeSpan; private final String sessionId; private HandlerContext.Receiver text; + private int activeTextHandlers; + private boolean textFinishRequested; private BinaryMessage currentBinary; public ReceiveContexts(AgentSpan span, CoreSession session) { @@ -221,15 +226,12 @@ private synchronized MethodHandle wrapText(MethodHandle delegate, boolean partia if (text == null) { text = new HandlerContext.Receiver(handshakeSpan, sessionId); } - return insertArguments(TEXT, 0, delegate, text, partial); + return insertArguments(TEXT, 0, delegate, this, partial); } public synchronized void finish() { - if (text != null) { - synchronized (text) { - DECORATE.onFrameEnd(text); - } - } + textFinishRequested = true; + finishTextIfIdle(); // Callbacks own completed messages; only the unfinished fragment sequence needs closing. BinaryMessage message = currentBinary; currentBinary = null; @@ -240,6 +242,15 @@ public synchronized void finish() { } } } + + /** Must be called with this instance's monitor held. */ + private void finishTextIfIdle() { + // Connection close may run while a text handler still owns the receive span. + if (textFinishRequested && activeTextHandlers == 0) { + textFinishRequested = false; + DECORATE.onFrameEnd(text); + } + } } public static class BinaryMessage extends HandlerContext.Receiver { diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeEndpoints.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeEndpoints.java index 3c3084b7f36..b916eea5351 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeEndpoints.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeEndpoints.java @@ -51,12 +51,16 @@ public static class EndpointEvents { boolean failCallback; boolean failAfterCallback; Callback pendingCallback; + Runnable onMessage; int closeCode; String closeReason; void recordMessage(String message) { messages.add(message); messageSpans.add(activeSpan()); + if (onMessage != null) { + onMessage.run(); + } if (echoMessages) { session.sendText(message, Callback.NOOP); } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java index 7e6dbbddcdf..e42b43d28e2 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketTest.java @@ -26,6 +26,8 @@ import static java.util.concurrent.TimeUnit.SECONDS; import static java.util.regex.Pattern.compile; import static java.util.regex.Pattern.quote; +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -508,6 +510,80 @@ void handlerFailureAfterCallbackSuccessStillMarksReceiveSpan() throws Exception assertNull(activeSpan()); } + @TableTest({ + "scenario | endpoint | last | fail ", + "full success | full | true | false", + "full failure | full | true | true ", + "partial final success | partial | true | false", + "partial final failure | partial | true | true ", + "partial non-final success | partial | false | false", + "partial non-final failure | partial | false | true " + }) + void terminationWaitsForActiveTextHandler( + NativeEndpoints.EndpointEvents endpoint, boolean last, boolean fail) throws Exception { + CompletableFuture handlerEntered = new CompletableFuture<>(); + CompletableFuture releaseHandler = new CompletableFuture<>(); + endpoint.onMessage = + () -> { + handlerEntered.complete((DDSpan) activeSpan()); + assertDoesNotThrow(() -> releaseHandler.get(5, SECONDS), "Text handler was not released"); + if (fail) { + throw new IllegalStateException("handler failed after close"); + } + }; + JettyWebSocketFrameHandler frameHandler = createFrameHandler(endpoint); + openFrameHandler(frameHandler); + Callback.Completable message = new Callback.Completable(); + ExecutorService executor = newSingleThreadExecutor(); + try { + CompletableFuture handling = + CompletableFuture.runAsync( + () -> { + frameHandler.onFrame(new Frame(OpCode.TEXT, "hello").setFin(last), message); + assertNull(activeSpan()); + }, + executor); + DDSpan receive = handlerEntered.get(5, SECONDS); + assertNotNull(receive); + + Callback.Completable closed = new Callback.Completable(); + frameHandler.onClosed(new CloseStatus(CloseStatus.NORMAL, "bye"), closed); + closed.get(5, SECONDS); + boolean finishedDuringHandler = receive.isFinished(); + + releaseHandler.complete(null); + handling.get(5, SECONDS); + if (fail) { + ExecutionException failure = + assertThrows(ExecutionException.class, () -> message.get(5, SECONDS)); + assertEquals("handler failed after close", failure.getCause().getMessage()); + } else { + message.get(5, SECONDS); + } + assertAll( + () -> assertFalse(finishedDuringHandler, "Receive span finished before handler exit"), + () -> assertTrue(receive.isFinished(), "Receive span unfinished after handler exit"), + () -> assertEquals(fail, receive.isError()), + () -> + assertEquals( + fail ? "handler failed after close" : null, receive.getTag("error.message"))); + DDSpan handshake = handshake(); + TagsMatcher[] errorTags = + fail + ? new TagsMatcher[] {error(IllegalStateException.class, "handler failed after close")} + : new TagsMatcher[0]; + assertTraces( + trace(handshakeSpan()), + trace(receiveSpan(handshake, "text", 5, 1, errorTags).error(fail)), + trace(closeSpan(handshake))); + assertSame(receive, writer.get(2).get(0)); + } finally { + releaseHandler.complete(null); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, SECONDS)); + } + } + @TableTest({ "scenario | msgType | closeFrame | closeHandler", "text close frame with handler | text | true | true ", From 66e742995956d90d1d15963dc0cdf3c8f0e2deed Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 13 Oct 2026 11:52:08 -0700 Subject: [PATCH 22/23] Align Jetty 12 WebSocket client send tracing Create a root send span when no parent is active and reuse it across fragments, preserving the handshake link and complete message counts. --- .../websocket/jetty12/NativeSendContext.java | 9 +--- .../test/java/NativeWebSocketSendTest.java | 54 +++++++++++++++++-- 2 files changed, 51 insertions(+), 12 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java index a53c2153951..5de015bc0ca 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeSendContext.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.websocket.jetty12; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE; import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_BINARY; import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_TEXT; @@ -19,7 +18,6 @@ public class NativeSendContext { private final AgentSpan handshakeSpan; private final String sessionId; - private final boolean client; private Message partialMessage; private boolean closed; @@ -27,8 +25,7 @@ public NativeSendContext(AgentSpan span, CoreSession session) { if (Config.get().isWebsocketMessagesInheritSampling()) { span.forceSamplingDecision(); } - client = session.getBehavior() == Behavior.CLIENT; - handshakeSpan = client ? span : span.getLocalRootSpan(); + handshakeSpan = session.getBehavior() == Behavior.CLIENT ? span : span.getLocalRootSpan(); sessionId = Integer.toHexString(System.identityHashCode(session)); } @@ -44,10 +41,6 @@ public synchronized SendCallback start( : binary ? ((ByteBuffer) payload).remaining() : ((String) payload).length(); Message message = partial ? partialMessage : null; if (message == null || !type.equals(message.getMessageType())) { - // Later fragments keep the span created under the first fragment's parent. - if (client && activeSpan() == null) { - return null; - } message = new Message(handshakeSpan, sessionId); } AgentSpan span = DECORATE.startOutboundFrameSpan(message, type, size); diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java index 6f8bbaf03e8..b6f9782a0d4 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/NativeWebSocketSendTest.java @@ -60,7 +60,7 @@ class NativeWebSocketSendTest extends AbstractInstrumentationTest { "partial text | false | true ", "partial binary | true | true " }) - void clientSendsWithoutParentAreNotTraced(boolean binary, boolean partial) throws Exception { + void clientSendsWithoutParentCreateRootSpans(boolean binary, boolean partial) throws Exception { Connection connection = new Connection(true, Behavior.CLIENT); RecordingCallback callback = new RecordingCallback(); @@ -68,10 +68,10 @@ void clientSendsWithoutParentAreNotTraced(boolean binary, boolean partial) throw connection.callbacks.get(0).succeeded(); callback.get(5, SECONDS); - assertNull(connection.spans.get(0)); - assertNull(callback.span); + assertNotNull(connection.spans.get(0)); + assertSame(connection.spans.get(0), callback.span); assertNull(activeSpan()); - assertTraces(trace(handshakeSpan())); + assertTraces(trace(handshakeSpan()), trace(sendSpan(connection, binary, 5, 1).root())); } @TableTest({ @@ -137,6 +137,52 @@ void clientPartialSendKeepsParentUntilLastFragment(boolean binary) throws Except sendSpan(connection, binary, 11, 2).childOfPrevious())); } + @TableTest({ + "scenario | binary", + "text | false ", + "binary | true " + }) + void clientPartialSendWithoutInitialParentKeepsRootSpan(boolean binary) throws Exception { + Connection connection = new Connection(true, Behavior.CLIENT); + RecordingCallback callback = new RecordingCallback(); + + assertNull(activeSpan()); + connection.send(binary, true, "hello ", false, null); + connection.callbacks.get(0).succeeded(); + + runUnderTrace( + "application", + () -> { + AgentSpan parent = activeSpan(); + connection.send(binary, true, "world", true, callback); + assertSame(parent, activeSpan()); + return null; + }); + connection.callbacks.get(1).succeeded(); + callback.get(5, SECONDS); + + assertNotNull(connection.spans.get(0)); + assertSame(connection.spans.get(0), connection.spans.get(1)); + assertSame(connection.spans.get(0), callback.span); + assertNull(activeSpan()); + + runUnderTrace( + "next", + () -> { + connection.send(binary, true, "again", true, null); + return null; + }); + connection.callbacks.get(2).succeeded(); + assertTraces( + trace(handshakeSpan()), + trace(sendSpan(connection, binary, 11, 2).root()), + trace(span().operationName("application").root()), + trace( + SORT_BY_START_TIME, + span().operationName("next").root(), + sendSpan(connection, binary, 5, 1).childOfPrevious())); + } + @TableTest({ "scenario | binary | partial | synchronous", "text async | false | false | false ", From 1e670c3cfed4eafd5859afe956bdf65d7fe5ebcd Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 13 Oct 2026 14:21:21 -0700 Subject: [PATCH 23/23] Fix SpotBugs warnings in Jetty 12 WebSocket receive tracing Move text receive state updates into synchronized methods so SpotBugs recognizes the locking, preserving span completion behavior. --- .../jetty12/NativeMethodHandleWrappers.java | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java index 3529f86027a..c37d60215e4 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java @@ -128,30 +128,20 @@ public static void onText( String payload, boolean last) throws Throwable { - HandlerContext.Receiver context; boolean finish = last; - ContextScope scope; - synchronized (contexts) { - context = contexts.text; - contexts.activeTextHandlers++; - scope = startMessage(context, payload, partial); - } + ContextScope scope = contexts.startText(payload, partial); try (ContextScope ignored = scope) { try { delegate.invokeExact(payload, last); } catch (Throwable t) { finish = true; synchronized (contexts) { - DECORATE.onError(context.getWebsocketSpan(), t); + DECORATE.onError(contexts.text.getWebsocketSpan(), t); } throw t; } } finally { - synchronized (contexts) { - contexts.activeTextHandlers--; - contexts.textFinishRequested |= finish; - contexts.finishTextIfIdle(); - } + contexts.endText(finish); } } @@ -229,6 +219,17 @@ private synchronized MethodHandle wrapText(MethodHandle delegate, boolean partia return insertArguments(TEXT, 0, delegate, this, partial); } + private synchronized ContextScope startText(String payload, boolean partial) { + activeTextHandlers++; + return startMessage(text, payload, partial); + } + + private synchronized void endText(boolean finish) { + activeTextHandlers--; + textFinishRequested |= finish; + finishTextIfIdle(); + } + public synchronized void finish() { textFinishRequested = true; finishTextIfIdle(); @@ -243,8 +244,7 @@ public synchronized void finish() { } } - /** Must be called with this instance's monitor held. */ - private void finishTextIfIdle() { + private synchronized void finishTextIfIdle() { // Connection close may run while a text handler still owns the receive span. if (textFinishRequested && activeTextHandlers == 0) { textFinishRequested = false;