Vertx: wrap internal routes to let the context propagate for blocking…

… handlers (#7563) * Vertx: wrap internal routes to let the context propagate for blocking handlers * Extend the wrapping advices
DataDog · Sep 4, 2024 · 235737f · 235737f
1 parent 86d01e6
commit 235737f
Show file tree

Hide file tree

Showing 6 changed files with 53 additions and 50 deletions.
diff --git a/...main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteHandlerInstrumentation.java b/...main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteHandlerInstrumentation.java
@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.vertx_3_4.server;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -35,7 +36,7 @@ public String instrumentedType() {
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         isMethod()
-            .and(named("handler"))
+            .and(namedOneOf("handler", "blockingHandler"))
             .and(isPublic())
             .and(takesArgument(0, named("io.vertx.core.Handler"))),
         packageName + ".RouteHandlerWrapperAdvice");

diff --git a/...3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteHandlerWrapper.java b/...3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteHandlerWrapper.java
@@ -9,43 +9,54 @@
 import datadog.trace.api.gateway.Flow;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import io.vertx.core.Handler;
 import io.vertx.ext.web.RoutingContext;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
+import io.vertx.ext.web.impl.RouteImpl;
+import io.vertx.ext.web.impl.RouterImpl;
 
 public class RouteHandlerWrapper implements Handler<RoutingContext> {
-  private static final Logger log = LoggerFactory.getLogger(RouteHandlerWrapper.class);
   static final String PARENT_SPAN_CONTEXT_KEY = AgentSpan.class.getName() + ".parent";
   static final String HANDLER_SPAN_CONTEXT_KEY = AgentSpan.class.getName() + ".handler";
   static final String ROUTE_CONTEXT_KEY = "dd." + Tags.HTTP_ROUTE;
 
   private final Handler<RoutingContext> actual;
+  private final boolean spanStarter;
 
   public RouteHandlerWrapper(final Handler<RoutingContext> handler) {
     actual = handler;
+    // When mounting a sub router, the handler is a lambda in either RouterImpl or RouteImpl, so
+    // this skips that. This prevents routers from creating a span during handling. In the event
+    // a route is not found, without this code, a span would be created for the router when it
+    // shouldn't
+    String name = handler.getClass().getName();
+    spanStarter =
+        !(name.startsWith(RouterImpl.class.getName())
+            || name.startsWith(RouteImpl.class.getName()));
   }
 
   @Override
   public void handle(final RoutingContext routingContext) {
     AgentSpan span = routingContext.get(HANDLER_SPAN_CONTEXT_KEY);
     Flow.Action.RequestBlockingAction rba = null;
-    if (span == null) {
-      AgentSpan parentSpan = activeSpan();
-      routingContext.put(PARENT_SPAN_CONTEXT_KEY, parentSpan);
+    if (spanStarter) {
+      if (span == null) {
+        AgentSpan parentSpan = activeSpan();
+        routingContext.put(PARENT_SPAN_CONTEXT_KEY, parentSpan);
 
-      span = startSpan(INSTRUMENTATION_NAME);
-      routingContext.put(HANDLER_SPAN_CONTEXT_KEY, span);
+        span = startSpan(INSTRUMENTATION_NAME);
+        routingContext.put(HANDLER_SPAN_CONTEXT_KEY, span);
 
-      routingContext.response().endHandler(new EndHandlerWrapper(routingContext));
-      DECORATE.afterStart(span);
-      span.setResourceName(DECORATE.className(actual.getClass()));
-    }
-
-    updateRoutingContextWithRoute(routingContext);
+        routingContext.response().endHandler(new EndHandlerWrapper(routingContext));
+        DECORATE.afterStart(span);
+        span.setResourceName(DECORATE.className(actual.getClass()));
+      }
 
-    try (final AgentScope scope = activateSpan(span)) {
+      updateRoutingContextWithRoute(routingContext);
+    }
+    try (final AgentScope scope =
+        span != null ? activateSpan(span) : AgentTracer.NoopAgentScope.INSTANCE) {
       scope.setAsyncPropagation(true);
       try {
         actual.handle(routingContext);

diff --git a/...c/main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteHandlerWrapperAdvice.java b/...c/main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteHandlerWrapperAdvice.java
@@ -2,22 +2,12 @@
 
 import io.vertx.core.Handler;
 import io.vertx.ext.web.RoutingContext;
-import io.vertx.ext.web.impl.RouteImpl;
-import io.vertx.ext.web.impl.RouterImpl;
 import net.bytebuddy.asm.Advice;
 
 public class RouteHandlerWrapperAdvice {
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static void wrapHandler(
       @Advice.Argument(value = 0, readOnly = false) Handler<RoutingContext> handler) {
-    // When mounting a sub router, the handler is a lambda in either RouterImpl or RouteImpl, so
-    // this skips that. This prevents routers from creating a span during handling. In the event
-    // a route is not found, without this code, a span would be created for the router when it
-    // shouldn't
-    String name = handler.getClass().getName();
-    if (!(name.startsWith(RouterImpl.class.getName())
-        || name.startsWith(RouteImpl.class.getName()))) {
-      handler = new RouteHandlerWrapper(handler);
-    }
+    handler = new RouteHandlerWrapper(handler);
   }
 }
diff --git a/...main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerInstrumentation.java b/...main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerInstrumentation.java
@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.vertx_4_0.server;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -41,7 +42,7 @@ public String instrumentedType() {
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         isMethod()
-            .and(named("handler"))
+            .and(namedOneOf("handler", "blockingHandler"))
             .and(isPublic())
             .and(takesArgument(0, named("io.vertx.core.Handler"))),
         packageName + ".RouteHandlerWrapperAdvice");

diff --git a/...4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerWrapper.java b/...4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerWrapper.java
@@ -8,42 +8,49 @@
 
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import io.vertx.core.Handler;
 import io.vertx.ext.web.RoutingContext;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
+import io.vertx.ext.web.impl.RouteImpl;
 
 public class RouteHandlerWrapper implements Handler<RoutingContext> {
-  private static final Logger log = LoggerFactory.getLogger(RouteHandlerWrapper.class);
   static final String PARENT_SPAN_CONTEXT_KEY = AgentSpan.class.getName() + ".parent";
   static final String HANDLER_SPAN_CONTEXT_KEY = AgentSpan.class.getName() + ".handler";
   static final String ROUTE_CONTEXT_KEY = "dd." + Tags.HTTP_ROUTE;
 
   private final Handler<RoutingContext> actual;
+  private final boolean spanStarter;
 
   public RouteHandlerWrapper(final Handler<RoutingContext> handler) {
     actual = handler;
+    // When mounting a sub router, the handler is a method reference to the routers handleContext
+    // method this skips that. This prevents routers from creating a span during handling. In the
+    // event a route is not found, without this code, a span would be created for the router when
+    // it shouldn't
+    spanStarter = !handler.getClass().getName().startsWith(RouteImpl.class.getName());
   }
 
   @Override
   public void handle(final RoutingContext routingContext) {
     AgentSpan span = routingContext.get(HANDLER_SPAN_CONTEXT_KEY);
-    if (span == null) {
-      AgentSpan parentSpan = activeSpan();
-      routingContext.put(PARENT_SPAN_CONTEXT_KEY, parentSpan);
+    if (spanStarter) {
+      if (span == null) {
+        AgentSpan parentSpan = activeSpan();
+        routingContext.put(PARENT_SPAN_CONTEXT_KEY, parentSpan);
 
-      span = startSpan(INSTRUMENTATION_NAME);
-      routingContext.put(HANDLER_SPAN_CONTEXT_KEY, span);
+        span = startSpan(INSTRUMENTATION_NAME);
+        routingContext.put(HANDLER_SPAN_CONTEXT_KEY, span);
 
-      routingContext.response().endHandler(new EndHandlerWrapper(routingContext));
-      DECORATE.afterStart(span);
-      span.setResourceName(DECORATE.className(actual.getClass()));
+        routingContext.response().endHandler(new EndHandlerWrapper(routingContext));
+        DECORATE.afterStart(span);
+        span.setResourceName(DECORATE.className(actual.getClass()));
+      }
+      updateRoutingContextWithRoute(routingContext);
     }
 
-    updateRoutingContextWithRoute(routingContext);
-
-    try (final AgentScope scope = activateSpan(span)) {
+    try (final AgentScope scope =
+        span != null ? activateSpan(span) : AgentTracer.NoopAgentScope.INSTANCE) {
       scope.setAsyncPropagation(true);
       try {
         actual.handle(routingContext);

diff --git a/...c/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerWrapperAdvice.java b/...c/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerWrapperAdvice.java
@@ -2,19 +2,12 @@
 
 import io.vertx.core.Handler;
 import io.vertx.ext.web.RoutingContext;
-import io.vertx.ext.web.impl.RouteImpl;
 import net.bytebuddy.asm.Advice;
 
 public class RouteHandlerWrapperAdvice {
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static void wrapHandler(
       @Advice.Argument(value = 0, readOnly = false) Handler<RoutingContext> handler) {
-    // When mounting a sub router, the handler is a method reference to the routers handleContext
-    // method this skips that. This prevents routers from creating a span during handling. In the
-    // event a route is not found, without this code, a span would be created for the router when
-    // it shouldn't
-    if (!handler.getClass().getName().startsWith(RouteImpl.class.getName())) {
-      handler = new RouteHandlerWrapper(handler);
-    }
+    handler = new RouteHandlerWrapper(handler);
   }
 }