webflux

Author: SylvainJuge

instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/DispatcherHandlerInstrumentation.java

@@ -14,6 +14,7 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.springframework.web.server.ServerWebExchange;
@@ -45,28 +46,31 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class HandleAdvice {
 
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void methodExit(
+    public static Mono<Void> methodExit(
         @Advice.Thrown Throwable throwable,
         @Advice.Argument(0) ServerWebExchange exchange,
-        @Advice.Return(readOnly = false) Mono<Void> mono) {
+        @Advice.Return Mono<Void> originalMono) {
+      Mono<Void> mono = originalMono;
       if (mono != null) {
         // note: it seems like this code should go in @OnMethodExit of
         // HandlerAdapterInstrumentation.HandleAdvice instead, but for some reason "GET to bad
         // endpoint annotation API fail Mono" test fails with that placement
         mono = AdviceUtils.end(mono, exchange);
       }
+      return mono;
     }
   }
 
   @SuppressWarnings("unused")
   public static class HandleResultAdvice {
 
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void methodExit(
-        @Advice.Argument(0) ServerWebExchange exchange,
-        @Advice.Return(readOnly = false) Mono<Void> mono) {
-      mono = AdviceUtils.wrapMono(mono, exchange.getAttribute(AdviceUtils.CONTEXT));
+    public static Mono<Void> methodExit(
+        @Advice.Argument(0) ServerWebExchange exchange, @Advice.Return Mono<Void> mono) {
+      return AdviceUtils.wrapMono(mono, exchange.getAttribute(AdviceUtils.CONTEXT));
     }
   }
 }

instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/HandlerAdapterInstrumentation.java

@@ -23,7 +23,9 @@
 import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRouteSource;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+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;
 import org.springframework.web.reactive.HandlerResult;
@@ -58,12 +60,40 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class HandleAdvice {
 
+    public static class AdviceScope {
+      private final Context context;
+      private final Scope scope;
+
+      public AdviceScope(Context context, Scope scope) {
+        this.context = context;
+        this.scope = scope;
+      }
+
+      public Mono<HandlerResult> exit(
+          Throwable throwable,
+          ServerWebExchange exchange,
+          Object handler,
+          Mono<HandlerResult> mono) {
+        scope.close();
+
+        if (throwable != null) {
+          instrumenter().end(context, handler, null, throwable);
+        } else {
+          mono = AdviceUtils.wrapMono(mono, context);
+          exchange.getAttributes().put(AdviceUtils.CONTEXT, context);
+          AdviceUtils.registerOnSpanEnd(exchange, context, handler);
+          // span finished by wrapped Mono in DispatcherHandlerInstrumentation
+          // the Mono is already wrapped at this point, but doesn't read the ON_SPAN_END until
+          // the Mono is resolved, which is after this point
+        }
+        return mono;
+      }
+    }
+
+    @Nullable
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void methodEnter(
-        @Advice.Argument(0) ServerWebExchange exchange,
-        @Advice.Argument(1) Object handler,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
+    public static AdviceScope methodEnter(
+        @Advice.Argument(0) ServerWebExchange exchange, @Advice.Argument(1) Object handler) {
 
       Context parentContext = Context.current();
 
@@ -75,40 +105,30 @@ public static void methodEnter(
           parentContext, HttpServerRouteSource.NESTED_CONTROLLER, httpRouteGetter(), exchange);
 
       if (handler == null) {
-        return;
+        return null;
       }
 
       if (!instrumenter().shouldStart(parentContext, handler)) {
-        return;
+        return null;
       }
-
-      context = instrumenter().start(parentContext, handler);
-      scope = context.makeCurrent();
+      Context context = instrumenter().start(parentContext, handler);
+      return new AdviceScope(context, context.makeCurrent());
     }
 
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void methodExit(
-        @Advice.Return(readOnly = false) Mono<HandlerResult> mono,
+    public static Mono<HandlerResult> methodExit(
+        @Advice.Return Mono<HandlerResult> mono,
         @Advice.Argument(0) ServerWebExchange exchange,
         @Advice.Argument(1) Object handler,
         @Advice.Thrown Throwable throwable,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      if (scope == null) {
-        return;
-      }
-      scope.close();
-
-      if (throwable != null) {
-        instrumenter().end(context, handler, null, throwable);
-      } else {
-        mono = AdviceUtils.wrapMono(mono, context);
-        exchange.getAttributes().put(AdviceUtils.CONTEXT, context);
-        AdviceUtils.registerOnSpanEnd(exchange, context, handler);
-        // span finished by wrapped Mono in DispatcherHandlerInstrumentation
-        // the Mono is already wrapped at this point, but doesn't read the ON_SPAN_END until
-        // the Mono is resolved, which is after this point
+        @Advice.Enter @Nullable AdviceScope adviceScope) {
+
+      if (adviceScope == null) {
+        return mono;
       }
+
+      return adviceScope.exit(throwable, exchange, handler, mono);
     }
   }
 }

instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/RouterFunctionInstrumentation.java

@@ -18,6 +18,7 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.springframework.web.reactive.function.server.HandlerFunction;
@@ -61,14 +62,17 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class RouteAdvice {
 
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void methodExit(
+    public static Mono<HandlerFunction<?>> methodExit(
         @Advice.This RouterFunction<?> thiz,
-        @Advice.Return(readOnly = false) Mono<HandlerFunction<?>> result,
+        @Advice.Return Mono<HandlerFunction<?>> originalResult,
         @Advice.Thrown Throwable throwable) {
+      Mono<HandlerFunction<?>> result = originalResult;
       if (throwable == null) {
         result = result.doOnNext(new RouteOnSuccess(thiz));
       }
+      return result;
     }
   }
 }

instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/WebfluxServerInstrumentationModule.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 WebfluxServerInstrumentationModule extends InstrumentationModule {
+public class WebfluxServerInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
 
   public WebfluxServerInstrumentationModule() {
     super("spring-webflux", "spring-webflux-5.0", "spring-webflux-server");
@@ -26,4 +28,9 @@ public List<TypeInstrumentation> typeInstrumentations() {
         new HandlerAdapterInstrumentation(),
         new RouterFunctionInstrumentation());
   }
+
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
 }

instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/reactornetty/ContextHandlerInstrumentation.java

@@ -14,6 +14,7 @@
 import io.opentelemetry.instrumentation.netty.v4_1.internal.ServerContexts;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -35,6 +36,7 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class CreateOperationsAdvice {
 
+    @Nullable
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static Scope onEnter(@Advice.Argument(0) Channel channel) {
       // set context to the first unprocessed request
@@ -46,7 +48,7 @@ public static Scope onEnter(@Advice.Argument(0) Channel channel) {
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void onExit(@Advice.Enter Scope scope) {
+    public static void onExit(@Advice.Enter @Nullable Scope scope) {
       if (scope != null) {
         scope.close();
       }

instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/reactornetty/HttpTrafficHandlerInstrumentation.java

@@ -14,6 +14,7 @@
 import io.opentelemetry.instrumentation.netty.v4_1.internal.ServerContexts;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -35,6 +36,7 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class RunAdvice {
 
+    @Nullable
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static Scope onEnter(
         @Advice.FieldValue("ctx") ChannelHandlerContext channelHandlerContext) {
@@ -47,7 +49,7 @@ public static Scope onEnter(
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void onExit(@Advice.Enter Scope scope) {
+    public static void onExit(@Advice.Enter @Nullable Scope scope) {
       if (scope != null) {
         scope.close();
       }