diff --git a/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java index 466ff9cdf7..d0a4bd1423 100644 --- a/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java @@ -28,7 +28,11 @@ public class PlayWSClientInstrumentation extends BasePlayWSClientInstrumentation DECORATE.onRequest(span, request); propagate().inject(span, request, SETTER); - asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); + if (asyncHandler instanceof StreamedAsyncHandler) { + asyncHandler = new StreamedAsyncHandlerWrapper((StreamedAsyncHandler) asyncHandler, span); + } else { + asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); + } return span; } diff --git a/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/StreamedAsyncHandlerWrapper.java b/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/StreamedAsyncHandlerWrapper.java new file mode 100644 index 0000000000..da10db2a42 --- /dev/null +++ b/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/StreamedAsyncHandlerWrapper.java @@ -0,0 +1,20 @@ +package datadog.trace.instrumentation.playws1; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import org.reactivestreams.Publisher; +import play.shaded.ahc.org.asynchttpclient.handler.StreamedAsyncHandler; + +public class StreamedAsyncHandlerWrapper extends AsyncHandlerWrapper + implements StreamedAsyncHandler { + private final StreamedAsyncHandler streamedDelegate; + + public StreamedAsyncHandlerWrapper(final StreamedAsyncHandler delegate, final AgentSpan span) { + super(delegate, span); + streamedDelegate = delegate; + } + + @Override + public State onStream(final Publisher publisher) { + return streamedDelegate.onStream(publisher); + } +} diff --git a/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java index d78f7f4595..68e1cdf4b4 100644 --- a/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java @@ -28,7 +28,11 @@ public class PlayWSClientInstrumentation extends BasePlayWSClientInstrumentation DECORATE.onRequest(span, request); propagate().inject(span, request, SETTER); - asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); + if (asyncHandler instanceof StreamedAsyncHandler) { + asyncHandler = new StreamedAsyncHandlerWrapper((StreamedAsyncHandler) asyncHandler, span); + } else { + asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); + } return span; } diff --git a/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/StreamedAsyncHandlerWrapper.java b/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/StreamedAsyncHandlerWrapper.java new file mode 100644 index 0000000000..6e0c7dc518 --- /dev/null +++ b/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/StreamedAsyncHandlerWrapper.java @@ -0,0 +1,20 @@ +package datadog.trace.instrumentation.playws21; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import org.reactivestreams.Publisher; +import play.shaded.ahc.org.asynchttpclient.handler.StreamedAsyncHandler; + +public class StreamedAsyncHandlerWrapper extends AsyncHandlerWrapper + implements StreamedAsyncHandler { + private final StreamedAsyncHandler streamedDelegate; + + public StreamedAsyncHandlerWrapper(final StreamedAsyncHandler delegate, final AgentSpan span) { + super(delegate, span); + streamedDelegate = delegate; + } + + @Override + public State onStream(final Publisher publisher) { + return streamedDelegate.onStream(publisher); + } +} diff --git a/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java index 2bcbb6ac3d..25bed2f257 100644 --- a/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java @@ -28,7 +28,11 @@ public class PlayWSClientInstrumentation extends BasePlayWSClientInstrumentation DECORATE.onRequest(span, request); propagate().inject(span, request, SETTER); - asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); + if (asyncHandler instanceof StreamedAsyncHandler) { + asyncHandler = new StreamedAsyncHandlerWrapper((StreamedAsyncHandler) asyncHandler, span); + } else { + asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); + } return span; } diff --git a/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/StreamedAsyncHandlerWrapper.java b/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/StreamedAsyncHandlerWrapper.java new file mode 100644 index 0000000000..d4861367f1 --- /dev/null +++ b/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/StreamedAsyncHandlerWrapper.java @@ -0,0 +1,20 @@ +package datadog.trace.instrumentation.playws2; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import org.reactivestreams.Publisher; +import play.shaded.ahc.org.asynchttpclient.handler.StreamedAsyncHandler; + +public class StreamedAsyncHandlerWrapper extends AsyncHandlerWrapper + implements StreamedAsyncHandler { + private final StreamedAsyncHandler streamedDelegate; + + public StreamedAsyncHandlerWrapper(final StreamedAsyncHandler delegate, final AgentSpan span) { + super(delegate, span); + streamedDelegate = delegate; + } + + @Override + public State onStream(final Publisher publisher) { + return streamedDelegate.onStream(publisher); + } +} diff --git a/dd-java-agent/instrumentation/play-ws/src/main/java/datadog/trace/instrumentation/playws/BasePlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws/src/main/java/datadog/trace/instrumentation/playws/BasePlayWSClientInstrumentation.java index 9e0e3267a0..9ec2142581 100644 --- a/dd-java-agent/instrumentation/play-ws/src/main/java/datadog/trace/instrumentation/playws/BasePlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws/src/main/java/datadog/trace/instrumentation/playws/BasePlayWSClientInstrumentation.java @@ -50,6 +50,7 @@ public abstract class BasePlayWSClientInstrumentation extends Instrumenter.Defau "datadog.trace.instrumentation.playws.PlayWSClientDecorator", "datadog.trace.instrumentation.playws.HeadersInjectAdapter", packageName + ".AsyncHandlerWrapper", + packageName + ".StreamedAsyncHandlerWrapper" }; } }