Colisweb · cverdier · Jan 28, 2021 · Jan 28, 2021
diff --git a/dd-java-agent/instrumentation/http4s-blaze-0.21/http4s-blaze-0.21.gradle b/dd-java-agent/instrumentation/http4s-blaze-0.21/http4s-blaze-0.21.gradle
@@ -0,0 +1,39 @@
+// Set properties before any plugins get loaded
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_1_8
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+apply from: "$rootDir/gradle/test-with-scala.gradle"
+
+muzzle {
+  pass {
+    group = 'org.http4s'
+    module = 'http4s-blaze-server_2.12'
+    versions = "[0.21.0,)"
+  }
+  pass {
+    group = 'org.http4s'
+    module = 'http4s-blaze-server_2.13'
+    versions = "[0.21.0,)"
+  }
+}
+
+configurations {
+  testArtifacts
+}
+
+dependencies {
+  main_java8CompileOnly group: 'org.http4s', name: 'http4s-blaze-server_2.12', version: '0.21.0'
+  main_java8CompileOnly group: 'com.google.auto.service', name: 'auto-service-annotations', version: '1.0-rc7'
+
+  testCompile project(':dd-java-agent:instrumentation:trace-annotation')
+  testCompile project(':dd-java-agent:instrumentation:scala-concurrent')
+
+//  testCompile group: 'com.typesafe.akka', name: 'akka-http_2.11', version: '10.0.0'
+//  testCompile group: 'com.squareup.okhttp3', name: 'okhttp', version: '3.6.0'
+//  testCompile project(':dd-java-agent:instrumentation:trace-annotation')
+//  testCompile project(':dd-java-agent:instrumentation:akka-concurrent')
+//  testCompile project(':dd-java-agent:instrumentation:scala-promise:scala-promise-2.10')
+//  testCompile project(':dd-java-agent:instrumentation:scala-promise:scala-promise-2.13')
+}
diff --git a/...ain/java/datadog/trace/instrumentation/http4sblaze/BlazeServerBuilderInstrumentation.java b/...ain/java/datadog/trace/instrumentation/http4sblaze/BlazeServerBuilderInstrumentation.java
@@ -0,0 +1,146 @@
+package datadog.trace.instrumentation.http4sblaze;
+
+import cats.data.Kleisli;
+import cats.data.Kleisli$;
+import cats.effect.ConcurrentEffect;
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import lombok.extern.slf4j.Slf4j;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.http4s.Request;
+import org.http4s.Response;
+import scala.Function1;
+import scala.util.Either;
+
+import java.util.Map;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.*;
+import static datadog.trace.instrumentation.http4sblaze.Http4sServerDecorator.DECORATOR;
+import static datadog.trace.instrumentation.http4sblaze.Http4sServerDecorator.HTTP4S_BLAZE_REQUEST;
+import static datadog.trace.instrumentation.http4sblaze.Http4sServerHeaders.HEADERS;
+import static java.util.Collections.singletonMap;
+
+@Slf4j
+@AutoService(Instrumenter.class)
+public class BlazeServerBuilderInstrumentation extends Instrumenter.Tracing {
+
+  public BlazeServerBuilderInstrumentation() {
+    super("http4s-blaze", "http4s-blaze-server");
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return named("org.http4s.server.blaze.BlazeServerBuilder");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[]{
+      getClass().getName() + "$DatadogWrapperHelper",
+      getClass().getName() + "$WithHttpAppAdvice",
+      getClass().getName() + "$WithHttpAppAdvice$RequestHandler",
+      getClass().getName() + "$WithHttpAppAdvice$ResponseHandler",
+      packageName + ".Http4sServerDecorator",
+      packageName + ".Http4sServerHeaders",
+      packageName + ".Http4sUriAdapter",
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+      named("withHttpApp"),
+      getClass().getName() + "$WithHttpAppAdvice");
+  }
+
+  public static class WithHttpAppAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static <F> void enter(
+      @Advice.Argument(value = 0, readOnly = false)
+        Kleisli<F, Request<F>, Response<F>> httpApp,
+      @Advice.FieldValue(value = "F", readOnly = true) final ConcurrentEffect<F> effect
+    ) {
+      final RequestHandler<F> requestHandler = new RequestHandler<F>(effect, httpApp);
+      httpApp = Kleisli$.MODULE$.apply(requestHandler);
+    }
+
+    public static class RequestHandler<F> implements Function1<Request<F>, F> {
+      private final ConcurrentEffect<F> effect;
+      private final Kleisli<F, Request<F>, Response<F>> baseApp;
+
+      public RequestHandler(ConcurrentEffect<F> effect, Kleisli<F, Request<F>, Response<F>> baseApp) {
+        this.effect = effect;
+        this.baseApp = baseApp;
+      }
+
+      public F apply(Request<F> request) {
+        final AgentScope scope = DatadogWrapperHelper.createSpan(request);
+        final ResponseHandler<F> responseHandler = new ResponseHandler<F>(effect, scope);
+        return effect.flatMap(
+          effect.attempt(baseApp.apply(request)),
+          responseHandler
+        );
+      }
+    }
+
+    public static class ResponseHandler<F> implements Function1<Either<Throwable, Response<F>>, F> {
+      private final ConcurrentEffect<F> effect;
+      private final AgentScope scope;
+
+      public ResponseHandler(ConcurrentEffect<F> effect, AgentScope scope) {
+        this.effect = effect;
+        this.scope = scope;
+      }
+
+      public F apply(Either<Throwable, Response<F>> result) {
+        if (result.isRight()) {
+          final Response<F> response = result.right().get();
+          DatadogWrapperHelper.finishSpan(scope.span(), response);
+          return effect.pure(response);
+        } else {
+          final Throwable error = result.left().get();
+          DatadogWrapperHelper.finishSpan(scope.span(), error);
+          return effect.raiseError(error);
+        }
+      }
+    }
+  }
+
+  public static class DatadogWrapperHelper {
+    public static AgentScope createSpan(final Request<?> request) {
+      final AgentSpan.Context extractedContext = propagate().extract(request, HEADERS);
+      final AgentSpan span = startSpan(HTTP4S_BLAZE_REQUEST, extractedContext);
+      span.setMeasured(true);
+
+      DECORATOR.afterStart(span);
+      DECORATOR.onConnection(span, request);
+      DECORATOR.onRequest(span, request);
+
+      final AgentScope scope = activateSpan(span);
+      scope.setAsyncPropagation(true);
+      return scope;
+    }
+
+    public static void finishSpan(final AgentSpan span, final Response<?> response) {
+      DECORATOR.onResponse(span, response);
+      DECORATOR.beforeFinish(span);
+
+      span.finish();
+    }
+
+    public static void finishSpan(final AgentSpan span, final Throwable t) {
+      DECORATOR.onError(span, t);
+      span.setTag(Tags.HTTP_STATUS, 500);
+      DECORATOR.beforeFinish(span);
+
+      span.finish();
+    }
+  }
+}
diff --git a/...e-0.21/src/main/java/datadog/trace/instrumentation/http4sblaze/Http4sServerDecorator.java b/...e-0.21/src/main/java/datadog/trace/instrumentation/http4sblaze/Http4sServerDecorator.java
@@ -0,0 +1,35 @@
+package datadog.trace.instrumentation.http4sblaze;
+
+import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
+import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator;
+import org.http4s.Request;
+import org.http4s.Response;
+
+public class Http4sServerDecorator extends HttpServerDecorator<Request<?>, Request<?>, Response<?>> {
+
+  public static final CharSequence HTTP4S_BLAZE_REQUEST = UTF8BytesString.createConstant("http4s-blaze.request");
+  public static final CharSequence HTTP4S_BLAZE_SERVER = UTF8BytesString.createConstant("http4s-blaze-server");
+  public static final Http4sServerDecorator DECORATOR = new Http4sServerDecorator();
+
+  @Override
+  protected String[] instrumentationNames() { return new String[] {"http4s-blaze", "http4s-blaze-server"}; }
+
+  @Override
+  protected CharSequence component() { return HTTP4S_BLAZE_SERVER; }
+
+  @Override
+  protected String method(Request<?> request) { return request.method().name(); }
+
+  @Override
+  protected URIDataAdapter url(Request<?> request) { return new Http4sUriAdapter(request.uri()); }
+
+  @Override
+  protected String peerHostIP(Request<?> request) { return null; }
+
+  @Override
+  protected int peerPort(Request<?> request) { return 0; }
+
+  @Override
+  protected int status(Response<?> response) { return response.status().code(); }
+}
diff --git a/...aze-0.21/src/main/java/datadog/trace/instrumentation/http4sblaze/Http4sServerHeaders.java b/...aze-0.21/src/main/java/datadog/trace/instrumentation/http4sblaze/Http4sServerHeaders.java
@@ -0,0 +1,21 @@
+package datadog.trace.instrumentation.http4sblaze;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+import org.http4s.Header;
+import org.http4s.Request;
+import scala.collection.JavaConverters;
+
+public class Http4sServerHeaders implements AgentPropagation.ContextVisitor<Request<?>> {
+
+  public static final Http4sServerHeaders HEADERS = new Http4sServerHeaders();
+
+  @Override
+  public void forEachKey(Request<?> carrier, AgentPropagation.KeyClassifier classifier) {
+    final scala.collection.immutable.List<Header> headers = carrier.headers();
+    for (final Header header : JavaConverters.asJavaCollection(headers)) {
+      if (!classifier.accept(header.name().value().toLowerCase(), header.value())) {
+        return;
+      }
+    }
+  }
+}
diff --git a/...-blaze-0.21/src/main/java/datadog/trace/instrumentation/http4sblaze/Http4sUriAdapter.java b/...-blaze-0.21/src/main/java/datadog/trace/instrumentation/http4sblaze/Http4sUriAdapter.java
@@ -0,0 +1,52 @@
+package datadog.trace.instrumentation.http4sblaze;
+
+import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
+import org.http4s.Uri;
+import scala.Option;
+
+public class Http4sUriAdapter implements URIDataAdapter {
+
+  private final Uri uri;
+
+  public Http4sUriAdapter(Uri uri) {
+    this.uri = uri;
+  }
+
+  @Override
+  public String scheme() {
+    Option<Uri.Scheme> scheme = uri.scheme();
+    if (scheme.isEmpty()) return "";
+    return scheme.get().value();
+  }
+
+  @Override
+  public String host() {
+    Option<Uri.Host> host = uri.host();
+    if (host.isEmpty()) return "";
+    return host.get().value();
+  }
+
+  @Override
+  public int port() {
+    Option<Object> port = uri.port();
+    if (port.isEmpty()) return 0;
+    return (int) port.get();
+  }
+
+  @Override
+  public String path() {
+    return uri.path();
+  }
+
+  @Override
+  public String fragment() {
+    Option<String> fragment = uri.fragment();
+    if (fragment.isEmpty()) return "";
+    return fragment.get();
+  }
+
+  @Override
+  public String query() {
+    return uri.query().renderString();
+  }
+}