chore: apply the same wrapper fix as in open-telemetry#3831, but this…

… time for request listners
robododge · Aug 13, 2021 · 085cc0d · 085cc0d
1 parent 392fc5a
commit 085cc0d
Showing 1 changed file with 38 additions and 17 deletions.
diff --git a/...ry/instrumentation/jetty/httpclient/v9_2/internal/JettyHttpClient9TracingInterceptor.java b/...ry/instrumentation/jetty/httpclient/v9_2/internal/JettyHttpClient9TracingInterceptor.java
@@ -10,6 +10,8 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
+import java.lang.reflect.Proxy;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.ListIterator;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -37,6 +39,16 @@ public class JettyHttpClient9TracingInterceptor
   private static final Logger logger =
       LoggerFactory.getLogger(JettyHttpClient9TracingInterceptor.class);
 
+  private static final Class<?>[] requestlistenerInterfaces = {
+    Request.BeginListener.class,
+    Request.FailureListener.class,
+    Request.SuccessListener.class,
+    Request.HeadersListener.class,
+    Request.ContentListener.class,
+    Request.CommitListener.class,
+    Request.QueuedListener.class
+  };
+
   @Nullable private Context context;
 
   @Nullable
@@ -79,26 +91,35 @@ public void attachToRequest(Request jettyRequest) {
   private void wrapRequestListeners(List<Request.RequestListener> requestListeners) {
 
     ListIterator<Request.RequestListener> iterator = requestListeners.listIterator();
+
     while (iterator.hasNext()) {
-      Request.RequestListener requestListener = iterator.next();
-      if (requestListener instanceof Request.FailureListener) {
-        iterator.set(
-            (Request.FailureListener)
-                (request, throwable) -> {
-                  try (Scope ignore = context.makeCurrent()) {
-                    ((Request.FailureListener) requestListener).onFailure(request, throwable);
-                  }
-                });
+      List<Class<?>> interfaces = new ArrayList<>();
+      Request.RequestListener listener = iterator.next();
+
+      Class<?> listenerClass = listener.getClass();
+
+      for (Class<?> type : requestlistenerInterfaces) {
+        if (type.isInstance(listener)) {
+          interfaces.add(type);
+        }
       }
-      if (requestListener instanceof Request.BeginListener) {
-        iterator.set(
-            (Request.FailureListener)
-                (request, throwable) -> {
-                  try (Scope ignore = context.makeCurrent()) {
-                    ((Request.BeginListener) requestListener).onBegin(request);
-                  }
-                });
+
+      if (interfaces.isEmpty()) {
+        continue;
       }
+
+      Request.RequestListener proxiedListner =
+          (Request.RequestListener)
+              Proxy.newProxyInstance(
+                  listenerClass.getClassLoader(),
+                  interfaces.toArray(new Class[0]),
+                  (proxy, method, args) -> {
+                    try (Scope ignored = context.makeCurrent()) {
+                      return method.invoke(listener, args);
+                    }
+                  });
+
+      iterator.set(proxiedListner);
     }
   }