Improve wrapping jetty response listener (#3831)

open-telemetry · Aug 13, 2021 · 392fc5a · 392fc5a
1 parent 539f104
commit 392fc5a
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 19 deletions.
diff --git a/.../io/opentelemetry/instrumentation/jetty/httpclient/v9_2/internal/JettyClientWrapUtil.java b/.../io/opentelemetry/instrumentation/jetty/httpclient/v9_2/internal/JettyClientWrapUtil.java
@@ -9,10 +9,22 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import java.lang.reflect.Proxy;
+import java.util.ArrayList;
 import java.util.List;
 import org.eclipse.jetty.client.api.Response;
 
 public final class JettyClientWrapUtil {
+  private static final Class<?>[] listenerInterfaces = {
+    Response.CompleteListener.class,
+    Response.FailureListener.class,
+    Response.SuccessListener.class,
+    Response.AsyncContentListener.class,
+    Response.ContentListener.class,
+    Response.HeadersListener.class,
+    Response.HeaderListener.class,
+    Response.BeginListener.class
+  };
 
   private JettyClientWrapUtil() {}
 
@@ -33,17 +45,29 @@ public static List<Response.ResponseListener> wrapResponseListeners(
 
   private static Response.ResponseListener wrapTheListener(
       Response.ResponseListener listener, Context context) {
-    Response.ResponseListener wrappedListener = listener;
-    if (listener instanceof Response.CompleteListener
-        && !(listener instanceof JettyHttpClient9TracingInterceptor)) {
-      wrappedListener =
-          (Response.CompleteListener)
-              result -> {
-                try (Scope ignored = context.makeCurrent()) {
-                  ((Response.CompleteListener) listener).onComplete(result);
-                }
-              };
+    if (listener == null || listener instanceof JettyHttpClient9TracingInterceptor) {
+      return listener;
     }
-    return wrappedListener;
+
+    Class<?> listenerClass = listener.getClass();
+    List<Class<?>> interfaces = new ArrayList<>();
+    for (Class<?> type : listenerInterfaces) {
+      if (type.isInstance(listener)) {
+        interfaces.add(type);
+      }
+    }
+    if (interfaces.isEmpty()) {
+      return listener;
+    }
+
+    return (Response.ResponseListener)
+        Proxy.newProxyInstance(
+            listenerClass.getClassLoader(),
+            interfaces.toArray(new Class[0]),
+            (proxy, method, args) -> {
+              try (Scope ignored = context.makeCurrent()) {
+                return method.invoke(listener, args);
+              }
+            });
   }
 }
diff --git a/...vy/io/opentelemetry/instrumentation/jetty/httpclient/v9_2/AbstractJettyClient9Test.groovy b/...vy/io/opentelemetry/instrumentation/jetty/httpclient/v9_2/AbstractJettyClient9Test.groovy
@@ -18,6 +18,7 @@ import org.eclipse.jetty.client.api.Result
 import org.eclipse.jetty.http.HttpMethod
 import org.eclipse.jetty.util.ssl.SslContextFactory
 import spock.lang.Shared
+import spock.lang.Unroll
 
 abstract class AbstractJettyClient9Test extends HttpClientTest<Request> {
 
@@ -34,7 +35,6 @@ abstract class AbstractJettyClient9Test extends HttpClientTest<Request> {
   Request jettyRequest = null
 
   def setupSpec() {
-
     //Start the main Jetty HttpClient and a https client
     client.start()
 
@@ -46,7 +46,6 @@ abstract class AbstractJettyClient9Test extends HttpClientTest<Request> {
 
   @Override
   Request buildRequest(String method, URI uri, Map<String, String> headers) {
-
     HttpClient theClient = uri.scheme == 'https' ? httpsClient : client
 
     Request request = theClient.newRequest(uri)
@@ -78,25 +77,21 @@ abstract class AbstractJettyClient9Test extends HttpClientTest<Request> {
   }
 
   private static class JettyClientListener implements Request.FailureListener, Response.FailureListener {
-
     volatile Throwable failure
 
     @Override
     void onFailure(Request requestF, Throwable failure) {
       this.failure = failure
-
     }
 
     @Override
     void onFailure(Response responseF, Throwable failure) {
       this.failure = failure
     }
-
   }
 
   @Override
   void sendRequestWithCallback(Request request, String method, URI uri, Map<String, String> headers, AbstractHttpClientTest.RequestResult requestResult) {
-
     JettyClientListener jcl = new JettyClientListener()
 
     request.onRequestFailure(jcl)
@@ -108,7 +103,6 @@ abstract class AbstractJettyClient9Test extends HttpClientTest<Request> {
     request.send(new Response.CompleteListener() {
       @Override
       void onComplete(Result result) {
-
         if (jcl.failure != null) {
           requestResult.complete(jcl.failure)
           return
@@ -119,7 +113,6 @@ abstract class AbstractJettyClient9Test extends HttpClientTest<Request> {
     })
   }
 
-
   @Override
   boolean testRedirects() {
     false
@@ -135,4 +128,27 @@ abstract class AbstractJettyClient9Test extends HttpClientTest<Request> {
     super.httpAttributes(uri) + extra
   }
 
+  @Unroll
+  def "test content of #method request #url"() {
+    when:
+    def request = buildRequest(method, url, [:])
+    def response = request.send()
+
+    then:
+    response.status == 200
+    response.getContentAsString() == "Hello."
+
+    assertTraces(1) {
+      trace(0, 2) {
+        clientSpan(it, 0, null, method, url)
+        serverSpan(it, 1, span(0))
+      }
+    }
+
+    where:
+    path << ["/success"]
+
+    method = "GET"
+    url = resolveAddress(path)
+  }
 }