Fix Mono failure tests

open-telemetry · Aug 8, 2021 · e56b099 · e56b099
1 parent bc9f077
commit e56b099
Show file tree

Hide file tree

Showing 4 changed files with 78 additions and 11 deletions.
diff --git a/...in/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/server/AdviceUtils.java b/...in/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/server/AdviceUtils.java
@@ -9,10 +9,14 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.tracer.ClassNames;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.springframework.web.server.ServerWebExchange;
 import reactor.core.publisher.Mono;
 
 public class AdviceUtils {
 
+  public static final String CONTEXT_ATTRIBUTE = AdviceUtils.class.getName() + ".Context";
+
   public static String spanNameForHandler(Object handler) {
     String className = ClassNames.simpleName(handler.getClass());
     int lambdaIdx = className.indexOf("$$Lambda$");
@@ -23,9 +27,16 @@ public static String spanNameForHandler(Object handler) {
     return className + ".handle";
   }
 
-  public static <T> Mono<T> setPublisherSpan(Mono<T> mono, Context context) {
-    return mono.doOnError(throwable -> instrumenter().end(context, null, null, throwable))
-        .doOnSuccess(t -> instrumenter().end(context, null, null, null))
-        .doOnCancel(() -> instrumenter().end(context, null, null, null));
+  public static <T> Mono<T> end(Mono<T> mono, ServerWebExchange exchange) {
+    return mono.doOnError(throwable -> end(exchange, throwable))
+        .doOnSuccess(t -> end(exchange, null))
+        .doOnCancel(() -> end(exchange, null));
+  }
+
+  private static void end(ServerWebExchange exchange, @Nullable Throwable throwable) {
+    Context context = (Context) exchange.getAttributes().get(AdviceUtils.CONTEXT_ATTRIBUTE);
+    if (context != null) {
+      instrumenter().end(context, null, null, throwable);
+    }
   }
 }
diff --git a/...try/javaagent/instrumentation/spring/webflux/server/DispatcherHandlerInstrumentation.java b/...try/javaagent/instrumentation/spring/webflux/server/DispatcherHandlerInstrumentation.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.webflux.server;
+
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.springframework.web.server.ServerWebExchange;
+import reactor.core.publisher.Mono;
+
+public class DispatcherHandlerInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("org.springframework.web.reactive.DispatcherHandler");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(isPublic())
+            .and(named("handle"))
+            .and(takesArgument(0, named("org.springframework.web.server.ServerWebExchange")))
+            .and(takesArguments(1)),
+        this.getClass().getName() + "$HandleAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class HandleAdvice {
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void methodExit(
+        @Advice.Thrown Throwable throwable,
+        @Advice.Argument(0) ServerWebExchange exchange,
+        @Advice.Return(readOnly = false) Mono<Void> mono) {
+      if (mono != null) {
+        mono = AdviceUtils.end(mono, exchange);
+      }
+    }
+  }
+}
diff --git a/...emetry/javaagent/instrumentation/spring/webflux/server/HandlerAdapterInstrumentation.java b/...emetry/javaagent/instrumentation/spring/webflux/server/HandlerAdapterInstrumentation.java
@@ -32,7 +32,6 @@
 import org.springframework.web.reactive.HandlerMapping;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.util.pattern.PathPattern;
-import reactor.core.publisher.Mono;
 
 public class HandlerAdapterInstrumentation implements TypeInstrumentation {
 
@@ -115,7 +114,6 @@ public static void methodEnter(
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
         @Advice.Argument(0) ServerWebExchange exchange,
-        @Advice.Return(readOnly = false) Mono<Void> mono,
         @Advice.Thrown Throwable throwable,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
@@ -124,11 +122,13 @@ public static void methodExit(
       }
       scope.close();
 
-      if (mono != null) {
-        mono = AdviceUtils.setPublisherSpan(mono, context);
-        // span finished in SpanFinishingSubscriber
-      } else {
+      if (throwable != null) {
         instrumenter().end(context, null, null, throwable);
+      } else {
+        exchange.getAttributes().put(AdviceUtils.CONTEXT_ATTRIBUTE, context);
+        // span finished by wrapped Mono in DispatcherHandlerInstrumentation
+        // the Mono is already wrapped at this point, but doesn't read the CONTEXT_ATTRIBUTE until
+        // the Mono is resolved, which is after this point
       }
     }
   }

diff --git a/...y/javaagent/instrumentation/spring/webflux/server/WebfluxServerInstrumentationModule.java b/...y/javaagent/instrumentation/spring/webflux/server/WebfluxServerInstrumentationModule.java
@@ -21,6 +21,9 @@ public WebfluxServerInstrumentationModule() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return asList(new HandlerAdapterInstrumentation(), new RouterFunctionInstrumentation());
+    return asList(
+        new DispatcherHandlerInstrumentation(),
+        new HandlerAdapterInstrumentation(),
+        new RouterFunctionInstrumentation());
   }
 }