Enhance AsyncHandlerData to include Instrumenter

Author: steverao

instrumentation/async-http-client/async-http-client-1-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/asynchttpclient/common/AsyncHandlerData.java

@@ -7,18 +7,26 @@
 
 import com.google.auto.value.AutoValue;
 import com.ning.http.client.Request;
+import com.ning.http.client.Response;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 
 @AutoValue
 public abstract class AsyncHandlerData {
 
-  public static AsyncHandlerData create(Context parentContext, Context context, Request request) {
-    return new AutoValue_AsyncHandlerData(parentContext, context, request);
+  public static AsyncHandlerData create(
+      Context parentContext, 
+      Context context, 
+      Request request, 
+      Instrumenter<Request, Response> instrumenter) {
+    return new AutoValue_AsyncHandlerData(parentContext, context, request, instrumenter);
   }
 
   public abstract Context getParentContext();
 
   public abstract Context getContext();
 
   public abstract Request getRequest();
+
+  public abstract Instrumenter<Request, Response> getInstrumenter();
 }

instrumentation/async-http-client/async-http-client-1-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/asynchttpclient/common/InstrumenterContextKey.java

@@ -63,10 +63,7 @@ public static Scope onEnter(
       }
       asyncHandlerDataField.set(handler, null);
 
-      // Get instrumenter from the data context - we'll store it there
-      @SuppressWarnings("unchecked")
-      Instrumenter<Request, Response> instrumenter = 
-          (Instrumenter<Request, Response>) data.getContext().get(InstrumenterContextKey.KEY);
+      Instrumenter<Request, Response> instrumenter = data.getInstrumenter();
       if (instrumenter != null) {
         instrumenter.end(data.getContext(), data.getRequest(), response, null);
       }
@@ -96,10 +93,8 @@ public static Scope onEnter(
       }
       asyncHandlerDataField.set(handler, null);
 
-      // Get instrumenter from the data context - we'll store it there
-      @SuppressWarnings("unchecked")
-      Instrumenter<Request, Response> instrumenter = 
-          (Instrumenter<Request, Response>) data.getContext().get(InstrumenterContextKey.KEY);
+      // Get instrumenter directly from data - much cleaner!
+      Instrumenter<Request, Response> instrumenter = data.getInstrumenter();
       if (instrumenter != null) {
         instrumenter.end(data.getContext(), data.getRequest(), null, throwable);
       }

instrumentation/async-http-client/async-http-client-1-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/asynchttpclient/common/ResponseInstrumentation.java

@@ -21,7 +21,6 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import javax.annotation.Nullable;
 import io.opentelemetry.javaagent.instrumentation.asynchttpclient.common.AsyncHandlerData;
-import io.opentelemetry.javaagent.instrumentation.asynchttpclient.common.InstrumenterContextKey;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -57,9 +56,7 @@ public static Scope onEnter(
       }
 
       Context context = instrumenter().start(parentContext, request);
-      // Store instrumenter in context so ResponseInstrumentation can access it
-      Context contextWithInstrumenter = context.with(InstrumenterContextKey.KEY, instrumenter());
-      ASYNC_HANDLER_DATA.set(handler, AsyncHandlerData.create(parentContext, contextWithInstrumenter, request));
+      ASYNC_HANDLER_DATA.set(handler, AsyncHandlerData.create(parentContext, context, request, instrumenter()));
       return context.makeCurrent();
     }
 

instrumentation/async-http-client/async-http-client-1.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/asynchttpclient/v1_8/AsyncHttpProviderInstrumentation.java

@@ -19,7 +19,6 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import io.opentelemetry.javaagent.instrumentation.asynchttpclient.common.AsyncHandlerData;
-import io.opentelemetry.javaagent.instrumentation.asynchttpclient.common.InstrumenterContextKey;
 import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -55,9 +54,7 @@ public static Scope onEnter(
       }
 
       Context context = instrumenter().start(parentContext, request);
-      // Store instrumenter in context so ResponseInstrumentation can access it
-      Context contextWithInstrumenter = context.with(InstrumenterContextKey.KEY, instrumenter());
-      ASYNC_HANDLER_DATA.set(handler, AsyncHandlerData.create(parentContext, contextWithInstrumenter, request));
+      ASYNC_HANDLER_DATA.set(handler, AsyncHandlerData.create(parentContext, context, request, instrumenter()));
       return context.makeCurrent();
     }