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..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 { @@ -18,10 +24,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() @@ -30,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')) @@ -40,5 +43,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..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,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=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=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 +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/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 new file mode 100644 index 00000000000..4a131369502 --- /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,100 @@ +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; +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; +import org.eclipse.jetty.io.EndPoint; +import org.eclipse.jetty.websocket.core.client.CoreClientUpgradeRequest; + +@AutoService(InstrumenterModule.class) +public class JettyWebSocketUpgradeInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + 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"; + } + + @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"))), + getClass().getName() + "$WebSocketUpgradeAdvice"); + } + + public static class WebSocketUpgradeAdvice { + @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()); + 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(); + } + } + } + + /** + * 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); + } + } +} 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/build.gradle b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/build.gradle index 295e18d7d71..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 @@ -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' @@ -38,17 +45,31 @@ addTestSuiteForDir("latestDepTest", "test") } } +["compileTestJava", "compileLatestDepTestJava"].each { name -> + tasks.named(name, JavaCompile) { + configureCompiler(it, JavaVersion.VERSION_17) + } +} + dependencies { + testImplementation libs.bundles.mockito + 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' 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: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/gradle.lockfile b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/gradle.lockfile index bf30846eadd..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 @@ -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=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=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=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=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=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=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=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=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=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 @@ -198,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/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..fa8043d50f9 --- /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,143 @@ +package datadog.trace.instrumentation.websocket.jetty12; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static java.util.Arrays.asList; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.instrumentation.websocket.jetty12.NativeMethodHandleWrappers.ReceiveContexts; +import java.lang.invoke.MethodHandle; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import net.bytebuddy.asm.Advice; +import org.eclipse.jetty.websocket.common.WebSocketSession; +import org.eclipse.jetty.websocket.core.Behavior; +import org.eclipse.jetty.websocket.core.CoreSession; + +@AutoService(InstrumenterModule.class) +public class Jetty12NativeWebSocketModule extends InstrumenterModule.Tracing { + public Jetty12NativeWebSocketModule() { + super("jetty", "jetty-websocket", "websocket"); + } + + @Override + protected boolean defaultEnabled() { + return InstrumenterConfig.get().isWebsocketTracingEnabled(); + } + + @Override + public List typeInstrumentations() { + return asList(new FrameHandlerInstrumentation(), new NativeSessionInstrumentation()); + } + + @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", + packageName + ".NativeSendContext", + packageName + ".NativeSendContext$Message", + packageName + ".NativeSendContext$SendCallback" + }; + } + + @Override + public Map contextStore() { + 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; + } + + public static class FrameHandlerInstrumentation + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + @Override + public String instrumentedType() { + return "org.eclipse.jetty.websocket.common.JettyWebSocketFrameHandler"; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("onOpen") + .and(takesArguments(2)) + .and(takesArgument(0, named("org.eclipse.jetty.websocket.core.CoreSession"))), + Jetty12NativeWebSocketModule.class.getName() + "$OpenAdvice"); + transformer.applyAdvice( + named("createMessageSink").and(takesArguments(4)), + Jetty12NativeWebSocketModule.class.getName() + "$MessageSinkAdvice"); + transformer.applyAdvice( + named("notifyOnClose").and(takesArguments(2)), + Jetty12NativeWebSocketModule.class.getName() + "$CloseAdvice"); + } + } + + public static class OpenAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter( + @Advice.Argument(0) CoreSession session, + @Advice.FieldValue(value = "closeHandle", readOnly = false) MethodHandle closeHandle) { + AgentSpan span = activeSpan(); + if (span != null && session.getBehavior() == Behavior.SERVER && closeHandle != null) { + closeHandle = NativeMethodHandleWrappers.wrapClose(closeHandle, span, session); + } + } + } + + public static class MessageSinkAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter( + @Advice.Argument(1) WebSocketSession session, + @Advice.Argument(value = 2, readOnly = false) MethodHandle handle) { + AgentSpan span = activeSpan(); + if (span != null + && session.getCoreSession().getBehavior() == Behavior.SERVER + && handle != null) { + handle = + NativeMethodHandleWrappers.wrapMessage( + handle, + span, + session.getCoreSession(), + InstrumentationContext.get(CoreSession.class, ReceiveContexts.class)); + } + } + } + + public static class CloseAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter(@Advice.FieldValue("session") WebSocketSession session) { + if (session != null) { + NativeSendContext sends = + InstrumentationContext.get(WebSocketSession.class, NativeSendContext.class) + .remove(session); + if (sends != null) { + sends.finish(); + } + ReceiveContexts contexts = + InstrumentationContext.get(CoreSession.class, ReceiveContexts.class) + .remove(session.getCoreSession()); + if (contexts != null) { + contexts.finish(); + } + } + } + } +} diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/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 new file mode 100644 index 00000000000..5aa2811161c --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/main/java/datadog/trace/instrumentation/websocket/jetty12/NativeMethodHandleWrappers.java @@ -0,0 +1,307 @@ +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; + delegate = delegate.asType(type.changeReturnType(void.class)); + MethodHandle normalized = partial ? delegate : dropArguments(delegate, 1, boolean.class); + ReceiveContexts contexts = contextStore.get(session); + if (contexts == null) { + contexts = contextStore.getOrPut(session, new ReceiveContexts(span, session)); + } + MethodHandle wrapper = + payload == String.class + ? 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); + } + return insertArguments(CLOSE, 0, delegate, context(span, session)); + } + + private static ContextScope startMessage( + HandlerContext.Receiver context, Object data, boolean partial) { + try { + 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; + ContextScope scope; + synchronized (context) { + scope = startMessage(context, payload, partial); + } + try (ContextScope ignored = scope) { + 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.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); + } + 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 final AgentSpan handshakeSpan; + private final String sessionId; + private HandlerContext.Receiver text; + private BinaryMessage currentBinary; + private final Set pendingBinary = new HashSet<>(); + + public ReceiveContexts(AgentSpan span, CoreSession session) { + if (Config.get().isWebsocketMessagesInheritSampling()) { + span.forceSamplingDecision(); + } + handshakeSpan = span.getLocalRootSpan(); + sessionId = Integer.toHexString(System.identityHashCode(session)); + } + + private synchronized MethodHandle wrapText(MethodHandle delegate, boolean partial) { + if (text == null) { + text = new HandlerContext.Receiver(handshakeSpan, sessionId); + } + return insertArguments(TEXT, 0, delegate, text, partial); + } + + public synchronized void finish() { + if (text != null) { + synchronized (text) { + DECORATE.onFrameEnd(text); + } + } + for (BinaryMessage message : pendingBinary) { + DECORATE.onFrameEnd(message); + } + pendingBinary.clear(); + currentBinary = null; + } + } + + public static class BinaryMessage extends HandlerContext.Receiver { + private int pendingCallbacks; + private boolean complete; + + public BinaryMessage(AgentSpan handshakeSpan, String sessionId) { + super(handshakeSpan, sessionId); + } + } + + 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/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..365ce5619da --- /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,184 @@ +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.Iterator; +import java.util.Set; +import org.eclipse.jetty.websocket.api.Callback; +import org.eclipse.jetty.websocket.core.Behavior; +import org.eclipse.jetty.websocket.core.CoreSession; + +/** 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 (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(); + } + } + 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) { + synchronized (context) { + 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 new file mode 100644 index 00000000000..7a313f4965f --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyEndpoints.java @@ -0,0 +1,244 @@ +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; +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<>(); + final List finalFragments = new ArrayList<>(); + final List messageSpans = new ArrayList<>(); + boolean failMessages; + boolean echoMessages; + 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 (echoMessages) { + session.sendText(message, Callback.NOOP); + } + 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 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 + 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/JettyWebSocketHandshakeTest.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java new file mode 100644 index 00000000000..75e0ccd4b32 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketHandshakeTest.java @@ -0,0 +1,143 @@ +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; +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 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; +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.tabletest.junit.TableTest; +import org.tabletest.junit.TypeConverterSources; + +@TypeConverterSources(JettyEndpoints.class) +class JettyWebSocketHandshakeTest extends AbstractInstrumentationTest { + @TableTest({ + "scenario | endpoint | messageType", + "listener text | full | text ", + "listener binary | full | binary ", + "boxed text | boxedFull | text ", + "boxed binary | boxedFull | binary " + }) + void messagesLinkToHandshakesInBothDirections( + JettyEndpoints.EndpointEvents endpoint, String messageType) throws Exception { + endpoint.echoMessages = true; + 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"); + 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")); + + 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); + 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"))) + .resourceName(compile(quote("websocket /receive"))) + .type(DDSpanTypes.WEBSOCKET) + .error(false) + .links( + to(handshake) + .traceFlags( + handshake.getSamplingPriority() > 0 + ? AgentSpanLink.SAMPLED_FLAG + : AgentSpanLink.DEFAULT_FLAGS) + .attributes( + SpanAttributes.builder().put("dd.kind", "executed_from").build())), + sendSpan(handshake).childOfPrevious())); + assertEquals(singletonList("hello"), endpoint.messages); + assertEquals(1, endpoint.messageSpans.size()); + 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")); + } + + 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..fedf737c0e9 --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketSendTest.java @@ -0,0 +1,454 @@ +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 closeFinishesIncompleteSendsAndWaitsForPendingCallbacks() 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); + 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(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 + 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); + } + } + } +} 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..6023a3de14b --- /dev/null +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-12.0/src/test/java/JettyWebSocketTest.java @@ -0,0 +1,598 @@ +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.TypeConverterSources; + +@TypeConverterSources(JettyEndpoints.class) +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 ", + "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); + 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 ", + "boxed partial text | boxedPartial | text ", + "boxed partial binary | boxedPartial | 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()); + } + + 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])); + } +}