java-http-client

Author: SylvainJuge

instrumentation/java-http-client/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/javahttpclient/HttpClientInstrumentation.java

@@ -28,7 +28,9 @@
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 import java.util.concurrent.CompletableFuture;
+import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -65,87 +67,114 @@ public void transform(TypeTransformer transformer) {
         HttpClientInstrumentation.class.getName() + "$SendAsyncAdvice");
   }
 
-  @SuppressWarnings("unused")
-  public static class SendAdvice {
+  public static class AdviceLocals {
+    private final Context context;
+    private final Context parentContext;
+    private final Scope scope;
+    private final CallDepth callDepth;
+
+    private AdviceLocals(Context parentContext, Context context, Scope scope, CallDepth callDepth) {
+      this.parentContext = parentContext;
+      this.context = context;
+      this.scope = scope;
+      this.callDepth = callDepth;
+    }
 
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void methodEnter(
-        @Advice.Argument(value = 0) HttpRequest httpRequest,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
+    @Nullable
+    public static AdviceLocals start(HttpRequest request) {
+      return start(request, null);
+    }
+
+    private static AdviceLocals start(HttpRequest request, CallDepth callDepth) {
       Context parentContext = currentContext();
-      if (!instrumenter().shouldStart(parentContext, httpRequest)) {
-        return;
+      if (!instrumenter().shouldStart(parentContext, request)) {
+        return null;
+      }
+
+      Context context = instrumenter().start(parentContext, request);
+      return new AdviceLocals(parentContext, context, context.makeCurrent(), callDepth);
+    }
+
+    public void end(HttpRequest request, HttpResponse<?> response, Throwable throwable) {
+      scope.close();
+      instrumenter().end(context, request, response, throwable);
+    }
+
+    public static AdviceLocals startWithCallDepth(HttpRequest httpRequest) {
+      CallDepth callDepth = CallDepth.forClass(HttpClient.class);
+      if (callDepth.getAndIncrement() > 0) {
+        return new AdviceLocals(null, null, null, callDepth);
+      }
+      return start(httpRequest, callDepth);
+    }
+
+    public boolean endWithCallDepth(HttpRequest httpRequest, Throwable throwable) {
+
+      if (callDepth.decrementAndGet() > 0) {
+        return false;
+      }
+
+      scope.close();
+      if (throwable != null) {
+        instrumenter().end(context, httpRequest, null, throwable);
+        return false;
       }
+      return true;
+    }
 
-      context = instrumenter().start(parentContext, httpRequest);
-      scope = context.makeCurrent();
+    public CompletableFuture<HttpResponse<?>> wrapFuture(
+        CompletableFuture<HttpResponse<?>> future, HttpRequest httpRequest) {
+      future = future.whenComplete(new ResponseConsumer(instrumenter(), context, httpRequest));
+      return CompletableFutureWrapper.wrap(future, parentContext);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class SendAdvice {
+
+    @Nullable
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AdviceLocals methodEnter(@Advice.Argument(value = 0) HttpRequest httpRequest) {
+      return AdviceLocals.start(httpRequest);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
         @Advice.Argument(0) HttpRequest httpRequest,
         @Advice.Return HttpResponse<?> httpResponse,
         @Advice.Thrown Throwable throwable,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      if (scope == null) {
-        return;
-      }
+        @Advice.Enter @Nullable AdviceLocals locals) {
 
-      scope.close();
-      instrumenter().end(context, httpRequest, httpResponse, throwable);
+      if (locals != null) {
+        locals.end(httpRequest, httpResponse, throwable);
+      }
     }
   }
 
   @SuppressWarnings("unused")
   public static class SendAsyncAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void methodEnter(
-        @Advice.Argument(value = 0) HttpRequest httpRequest,
-        @Advice.Local("otelCallDepth") CallDepth callDepth,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelParentContext") Context parentContext,
-        @Advice.Local("otelScope") Scope scope) {
-      callDepth = CallDepth.forClass(HttpClient.class);
-      if (callDepth.getAndIncrement() > 0) {
-        return;
-      }
-
-      parentContext = currentContext();
-      if (!instrumenter().shouldStart(parentContext, httpRequest)) {
-        return;
-      }
-
-      context = instrumenter().start(parentContext, httpRequest);
-      scope = context.makeCurrent();
+    public static AdviceLocals methodEnter(@Advice.Argument(value = 0) HttpRequest httpRequest) {
+      return AdviceLocals.startWithCallDepth(httpRequest);
     }
 
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void methodExit(
+    public static CompletableFuture<HttpResponse<?>> methodExit(
         @Advice.Argument(value = 0) HttpRequest httpRequest,
-        @Advice.Return(readOnly = false) CompletableFuture<HttpResponse<?>> future,
+        @Advice.Return CompletableFuture<HttpResponse<?>> originalFuture,
         @Advice.Thrown Throwable throwable,
-        @Advice.Local("otelCallDepth") CallDepth callDepth,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelParentContext") Context parentContext,
-        @Advice.Local("otelScope") Scope scope) {
-      if (callDepth.decrementAndGet() > 0) {
-        return;
-      }
+        @Advice.Enter AdviceLocals locals) {
 
-      if (scope == null) {
-        return;
-      }
+      @SuppressWarnings("UnnecessaryLocalVariable")
+      CompletableFuture<HttpResponse<?>> future = originalFuture;
 
-      scope.close();
-      if (throwable != null) {
-        instrumenter().end(context, httpRequest, null, throwable);
-      } else {
-        future = future.whenComplete(new ResponseConsumer(instrumenter(), context, httpRequest));
-        future = CompletableFutureWrapper.wrap(future, parentContext);
+      if (!locals.endWithCallDepth(httpRequest, throwable)) {
+        return future;
       }
+
+      return locals.wrapFuture(future, httpRequest);
     }
   }
 }

instrumentation/java-http-client/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/javahttpclient/HttpClientInstrumentationModule.java

@@ -10,10 +10,12 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import java.util.List;
 
 @AutoService(InstrumentationModule.class)
-public class HttpClientInstrumentationModule extends InstrumentationModule {
+public class HttpClientInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public HttpClientInstrumentationModule() {
     super("java-http-client");
   }
@@ -22,4 +24,9 @@ public HttpClientInstrumentationModule() {
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new HttpClientInstrumentation(), new HttpHeadersInstrumentation());
   }
+
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
 }

instrumentation/java-http-client/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/javahttpclient/HttpHeadersInstrumentation.java

@@ -16,6 +16,7 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import java.net.http.HttpHeaders;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -38,9 +39,10 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class HeadersAdvice {
 
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void methodExit(@Advice.Return(readOnly = false) HttpHeaders headers) {
-      headers = setter().inject(headers, Context.current());
+    public static HttpHeaders methodExit(@Advice.Return HttpHeaders headers) {
+      return setter().inject(headers, Context.current());
     }
   }
 }