DRAFT: APM AI TOOLKIT New Integration: Feign Client Instrumentation

dd-java-agent/instrumentation/feign-core/build.gradle

@@ -0,0 +1,21 @@
+muzzle {
+  pass {
+    group = "io.github.openfeign"
+    module = "feign-core"
+    versions = "[10.0.0,)"
+    assertInverse = true
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+addTestSuiteForDir('latestDepTest', 'test')
+
+dependencies {
+  compileOnly group: 'io.github.openfeign', name: 'feign-core', version: '10.0.0'
+
+  testImplementation group: 'io.github.openfeign', name: 'feign-core', version: '10.0.0'
+  testImplementation project(':dd-java-agent:testing')
+
+  latestDepTestImplementation group: 'io.github.openfeign', name: 'feign-core', version: '+'
+}

dd-java-agent/instrumentation/feign-core/src/main/java/datadog/trace/instrumentation/feign/FeignAsyncClientAdvice.java

@@ -0,0 +1,85 @@
+package datadog.trace.instrumentation.feign;
+
+import static datadog.context.Context.current;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.feign.FeignDecorator.DECORATE;
+import static datadog.trace.instrumentation.feign.FeignDecorator.FEIGN_REQUEST;
+import static datadog.trace.instrumentation.feign.FeignHeadersInjectAdapter.SETTER;
+
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import feign.Client;
+import feign.Request;
+import feign.Response;
+import java.util.concurrent.CompletableFuture;
+import net.bytebuddy.asm.Advice;
+
+public class FeignAsyncClientAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope onEnter(@Advice.Argument(0) final Request request) {
+    final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Client.class);
+    if (callDepth > 0) {
+      return null;
+    }
+
+    final AgentSpan span = startSpan(FEIGN_REQUEST);
+    DECORATE.afterStart(span);
+    DECORATE.onRequest(span, request);
+    DECORATE.injectContext(current().with(span), request, SETTER);
+
+    return activateSpan(span);
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void onExit(
+      @Advice.Enter final AgentScope scope,
+      @Advice.Return final CompletableFuture<Response> future,
+      @Advice.Thrown final Throwable throwable) {
+    if (scope == null) {
+      return;
+    }
+    final AgentSpan span = scope.span();
+
+    if (throwable != null) {
+      try {
+        DECORATE.onError(span, throwable);
+        DECORATE.beforeFinish(span);
+      } finally {
+        scope.close();
+        span.finish();
+        CallDepthThreadLocalMap.reset(Client.class);
+      }
+      return;
+    }
+
+    scope.close();
+
+    if (future != null) {
+      future.whenComplete(
+          (response, error) -> {
+            try {
+              if (response != null) {
+                DECORATE.onResponse(span, response);
+              }
+              if (error != null) {
+                DECORATE.onError(span, error);
+              }
+              DECORATE.beforeFinish(span);
+            } finally {
+              span.finish();
+              CallDepthThreadLocalMap.reset(Client.class);
+            }
+          });
+    } else {
+      try {
+        DECORATE.beforeFinish(span);
+      } finally {
+        span.finish();
+        CallDepthThreadLocalMap.reset(Client.class);
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/feign-core/src/main/java/datadog/trace/instrumentation/feign/FeignAsyncClientInstrumentation.java

@@ -0,0 +1,50 @@
+package datadog.trace.instrumentation.feign;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public class FeignAsyncClientInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+
+  public FeignAsyncClientInstrumentation() {
+    super("feign");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "feign.AsyncClient";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named("feign.AsyncClient"));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".FeignDecorator",
+      packageName + ".FeignHeadersInjectAdapter",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("execute"))
+            .and(takesArgument(0, named("feign.Request"))),
+        packageName + ".FeignAsyncClientAdvice");
+  }
+}

dd-java-agent/instrumentation/feign-core/src/main/java/datadog/trace/instrumentation/feign/FeignClientAdvice.java

@@ -0,0 +1,56 @@
+package datadog.trace.instrumentation.feign;
+
+import static datadog.context.Context.current;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.feign.FeignDecorator.DECORATE;
+import static datadog.trace.instrumentation.feign.FeignDecorator.FEIGN_REQUEST;
+import static datadog.trace.instrumentation.feign.FeignHeadersInjectAdapter.SETTER;
+
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import feign.Client;
+import feign.Request;
+import feign.Response;
+import net.bytebuddy.asm.Advice;
+
+public class FeignClientAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope onEnter(@Advice.Argument(0) final Request request) {
+    final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Client.class);
+    if (callDepth > 0) {
+      return null;
+    }
+
+    final AgentSpan span = startSpan(FEIGN_REQUEST);
+    DECORATE.afterStart(span);
+    DECORATE.onRequest(span, request);
+    DECORATE.injectContext(current().with(span), request, SETTER);
+
+    return activateSpan(span);
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void onExit(
+      @Advice.Enter final AgentScope scope,
+      @Advice.Return final Response response,
+      @Advice.Thrown final Throwable throwable) {
+    if (scope == null) {
+      return;
+    }
+    final AgentSpan span = scope.span();
+    try {
+      if (response != null) {
+        DECORATE.onResponse(span, response);
+      }
+      DECORATE.onError(span, throwable);
+      DECORATE.beforeFinish(span);
+    } finally {
+      scope.close();
+      span.finish();
+      CallDepthThreadLocalMap.reset(Client.class);
+    }
+  }
+}

dd-java-agent/instrumentation/feign-core/src/main/java/datadog/trace/instrumentation/feign/FeignClientInstrumentation.java

@@ -0,0 +1,50 @@
+package datadog.trace.instrumentation.feign;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public class FeignClientInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+
+  public FeignClientInstrumentation() {
+    super("feign");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "feign.Client";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named("feign.Client"));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".FeignDecorator",
+      packageName + ".FeignHeadersInjectAdapter",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("execute"))
+            .and(takesArgument(0, named("feign.Request"))),
+        packageName + ".FeignClientAdvice");
+  }
+}

dd-java-agent/instrumentation/feign-core/src/main/java/datadog/trace/instrumentation/feign/FeignDecorator.java

@@ -0,0 +1,80 @@
+package datadog.trace.instrumentation.feign;
+
+import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
+import feign.Request;
+import feign.Response;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collection;
+import java.util.Map;
+
+public class FeignDecorator extends HttpClientDecorator<Request, Response> {
+
+  public static final CharSequence FEIGN = UTF8BytesString.create("feign");
+  public static final FeignDecorator DECORATE = new FeignDecorator();
+  public static final CharSequence FEIGN_REQUEST =
+      UTF8BytesString.create(DECORATE.operationName());
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"feign"};
+  }
+
+  @Override
+  protected String service() {
+    return null;
+  }
+
+  @Override
+  protected CharSequence component() {
+    return FEIGN;
+  }
+
+  @Override
+  protected String method(final Request request) {
+    return request.httpMethod().name();
+  }
+
+  @Override
+  protected URI url(final Request request) {
+    try {
+      return new URI(request.url());
+    } catch (URISyntaxException e) {
+      return null;
+    }
+  }
+
+  @Override
+  protected int status(final Response response) {
+    return response.status();
+  }
+
+  @Override
+  protected String getRequestHeader(Request request, String headerName) {
+    Map<String, Collection<String>> headers = request.headers();
+    if (headers != null) {
+      for (Map.Entry<String, Collection<String>> entry : headers.entrySet()) {
+        if (entry.getKey().equalsIgnoreCase(headerName)) {
+          Collection<String> values = entry.getValue();
+          if (values != null && !values.isEmpty()) {
+            return values.iterator().next();
+          }
+        }
+      }
+    }
+    return null;
+  }
+
+  @Override
+  protected String getResponseHeader(Response response, String headerName) {
+    Map<String, Collection<String>> headers = response.headers();
+    if (headers != null) {
+      Collection<String> values = headers.get(headerName);
+      if (values != null && !values.isEmpty()) {
+        return values.iterator().next();
+      }
+    }
+    return null;
+  }
+}

dd-java-agent/instrumentation/feign-core/src/main/java/datadog/trace/instrumentation/feign/FeignHeadersInjectAdapter.java

@@ -0,0 +1,17 @@
+package datadog.trace.instrumentation.feign;
+
+import datadog.context.propagation.CarrierSetter;
+import feign.Request;
+import java.util.Collections;
+import javax.annotation.ParametersAreNonnullByDefault;
+
+@ParametersAreNonnullByDefault
+public class FeignHeadersInjectAdapter implements CarrierSetter<Request> {
+
+  public static final FeignHeadersInjectAdapter SETTER = new FeignHeadersInjectAdapter();
+
+  @Override
+  public void set(final Request carrier, final String key, final String value) {
+    carrier.headers().put(key, Collections.singletonList(value));
+  }
+}