From 6c8d2e1fc3f984e2e6f976486bc9c85fc47d0eaa Mon Sep 17 00:00:00 2001
From: Steve Dodge <steve@robododge.com>
Date: Fri, 13 Aug 2021 18:45:18 -0500
Subject: [PATCH] chore: apply the same wrapper fix as in open-telemetry#3831,
 but this time for request listners (#3833)

---
 .../JettyHttpClient9TracingInterceptor.java   | 55 +++++++++++++------
 1 file changed, 38 insertions(+), 17 deletions(-)

diff --git a/instrumentation/jetty-httpclient/jetty-httpclient-9.2/library/src/main/java/io/opentelemetry/instrumentation/jetty/httpclient/v9_2/internal/JettyHttpClient9TracingInterceptor.java b/instrumentation/jetty-httpclient/jetty-httpclient-9.2/library/src/main/java/io/opentelemetry/instrumentation/jetty/httpclient/v9_2/internal/JettyHttpClient9TracingInterceptor.java
index ce3ff3c5e758..dfa8b8c85aec 100644
--- a/instrumentation/jetty-httpclient/jetty-httpclient-9.2/library/src/main/java/io/opentelemetry/instrumentation/jetty/httpclient/v9_2/internal/JettyHttpClient9TracingInterceptor.java
+++ b/instrumentation/jetty-httpclient/jetty-httpclient-9.2/library/src/main/java/io/opentelemetry/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);
     }
   }