From fd0e10f0cb071bb5c5e4ce49e9313a0a35e385ca Mon Sep 17 00:00:00 2001 From: Laplie Anderson Date: Tue, 3 Mar 2020 17:45:21 -0500 Subject: [PATCH] Returned a StreamedAsyncHandler instance where appropriate --- .../playws1/PlayWSClientInstrumentation.java | 6 +++++- .../playws1/StreamedAsyncHandlerWrapper.java | 20 +++++++++++++++++++ .../playws21/PlayWSClientInstrumentation.java | 6 +++++- .../playws21/StreamedAsyncHandlerWrapper.java | 20 +++++++++++++++++++ .../playws2/PlayWSClientInstrumentation.java | 6 +++++- .../playws2/StreamedAsyncHandlerWrapper.java | 20 +++++++++++++++++++ .../BasePlayWSClientInstrumentation.java | 1 + 7 files changed, 76 insertions(+), 3 deletions(-) create mode 100644 dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/StreamedAsyncHandlerWrapper.java create mode 100644 dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/StreamedAsyncHandlerWrapper.java create mode 100644 dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/StreamedAsyncHandlerWrapper.java 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" }; } }