From e25af914836276d3185c42188a420e8231767f90 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Mon, 28 Sep 2026 16:32:39 -0700 Subject: [PATCH 01/10] 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 087453fea5d6c3292ad6338d4811cd27c47069a6 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Mon, 28 Sep 2026 17:49:25 -0700 Subject: [PATCH 02/10] 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 6e5b2c453e368e25bc4930ec2a891a38db6fda7c Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Mon, 28 Sep 2026 20:05:45 -0700 Subject: [PATCH 03/10] 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 2bc4277d3ec5bf0927adfab3aea6ae30e8c8a5db Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 29 Sep 2026 10:12:28 -0700 Subject: [PATCH 04/10] 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 5be991e226a9feda5574147ded1e091885e78d35 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 29 Sep 2026 13:17:35 -0700 Subject: [PATCH 05/10] 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 b99a2334e4317c61eb4c2ebe9c76c1bc8493c2cd Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 29 Sep 2026 13:57:18 -0700 Subject: [PATCH 06/10] 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 1f751f6278444436af17d0bf702925795dfeff45 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Tue, 29 Sep 2026 14:10:34 -0700 Subject: [PATCH 07/10] 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 77b4362c602665190e3b3a3fd1bd5ed1317da7a6 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Thu, 1 Oct 2026 21:29:49 -0700 Subject: [PATCH 08/10] 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 38b819be5c3a33adb92068e55c73354addece411 Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Fri, 2 Oct 2026 14:10:18 -0700 Subject: [PATCH 09/10] =?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 5052a0c55041919d76ad54d60073f5b6dd90e8aa Mon Sep 17 00:00:00 2001 From: Yury Gribkov Date: Fri, 2 Oct 2026 16:12:13 -0700 Subject: [PATCH 10/10] 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