Spring webflux: add user spans as children of the controller span

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

@@ -8,13 +8,16 @@
 import static io.opentelemetry.javaagent.instrumentation.spring.webflux.v5_0.server.WebfluxSingletons.instrumenter;
 
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
 import javax.annotation.Nullable;
 import org.springframework.web.server.ServerWebExchange;
+import reactor.core.CoreSubscriber;
 import reactor.core.publisher.Mono;
 
 public final class AdviceUtils {
 
-  public static final String ON_SPAN_END = AdviceUtils.class.getName() + ".Context";
+  public static final String ON_SPAN_END = AdviceUtils.class.getName() + ".OnSpanEnd";
+  public static final String CONTEXT = AdviceUtils.class.getName() + ".Context";
 
   public static void registerOnSpanEnd(
       ServerWebExchange exchange, Context context, Object handler) {
@@ -38,10 +41,35 @@ private static void end(ServerWebExchange exchange, @Nullable Throwable throwabl
     }
   }
 
+  public static <T> Mono<T> wrapMono(Mono<T> mono, Context context) {
+    if (context == null) {
+      return mono;
+    }
+    return new ContextMono<>(mono, context);
+  }
+
   @FunctionalInterface
   interface OnSpanEnd {
     void end(Throwable throwable);
   }
 
+  private static class ContextMono<T> extends Mono<T> {
+
+    private final Mono<T> delegate;
+    private final Context parentContext;
+
+    ContextMono(Mono<T> delegate, Context parentContext) {
+      this.delegate = delegate;
+      this.parentContext = parentContext;
+    }
+
+    @Override
+    public void subscribe(CoreSubscriber<? super T> actual) {
+      try (Scope ignored = parentContext.makeCurrent()) {
+        delegate.subscribe(actual);
+      }
+    }
+  }
+
   private AdviceUtils() {}
 }

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

@@ -35,6 +35,11 @@ public void transform(TypeTransformer transformer) {
             .and(takesArgument(0, named("org.springframework.web.server.ServerWebExchange")))
             .and(takesArguments(1)),
         this.getClass().getName() + "$HandleAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(named("handleResult"))
+            .and(takesArgument(0, named("org.springframework.web.server.ServerWebExchange"))),
+        this.getClass().getName() + "$HandleResultAdvice");
   }
 
   @SuppressWarnings("unused")
@@ -53,4 +58,15 @@ public static void methodExit(
       }
     }
   }
+
+  @SuppressWarnings("unused")
+  public static class HandleResultAdvice {
+
+    @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));
+    }
+  }
 }

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

@@ -26,7 +26,9 @@
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.springframework.web.reactive.HandlerResult;
 import org.springframework.web.server.ServerWebExchange;
+import reactor.core.publisher.Mono;
 
 public class HandlerAdapterInstrumentation implements TypeInstrumentation {
 
@@ -86,6 +88,7 @@ public static void methodEnter(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
+        @Advice.Return(readOnly = false) Mono<HandlerResult> mono,
         @Advice.Argument(0) ServerWebExchange exchange,
         @Advice.Argument(1) Object handler,
         @Advice.Thrown Throwable throwable,
@@ -99,6 +102,8 @@ public static void methodExit(
       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

instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/SpringWebfluxTest.java

@@ -290,7 +290,7 @@ void getAsyncResponseTest(Parameter parameter) {
                 },
                 span ->
                     span.hasName("tracedMethod")
-                        .hasParent(trace.getSpan(0))
+                        .hasParent(trace.getSpan(1))
                         .hasTotalAttributeCount(0)));
   }
 
@@ -410,7 +410,7 @@ void createSpanDuringHandlerFunctionTest(Parameter parameter) {
                 },
                 span ->
                     span.hasName("tracedMethod")
-                        .hasParent(trace.getSpan(parameter.annotatedMethod != null ? 0 : 1))
+                        .hasParent(trace.getSpan(1))
                         .hasTotalAttributeCount(0)));
   }
 

instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/server/base/ControllerSpringWebFluxServerTest.java

@@ -74,7 +74,6 @@ protected SpanDataAssert assertHandlerSpan(
   @Override
   protected void configure(HttpServerTestOptions options) {
     super.configure(options);
-    options.setHasHandlerAsControllerParentSpan(unused -> false);
     // TODO (trask) it seems like in this case ideally the controller span (which ends when the
     // Mono that the controller returns completes) should end before the server span (which needs
     // the result of the Mono)

instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/server/base/DelayedHandlerSpringWebFluxServerTest.java

@@ -5,7 +5,6 @@
 
 package server.base;
 
-import io.opentelemetry.instrumentation.testing.junit.http.HttpServerTestOptions;
 import io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint;
 import java.time.Duration;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -58,10 +57,4 @@ protected Mono<ServerResponse> wrapResponse(
                       }));
     }
   }
-
-  @Override
-  protected void configure(HttpServerTestOptions options) {
-    super.configure(options);
-    options.setHasHandlerAsControllerParentSpan(unused -> false);
-  }
 }

testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/AbstractHttpServerTest.java

@@ -571,15 +571,11 @@ protected void assertHighConcurrency(int count) {
                   span -> assertHandlerSpan(span, "GET", endpoint).hasParent(trace.getSpan(1)));
             }
 
-            int parentIndex = spanAssertions.size() - 2;
-            if (options.hasHandlerAsControllerParentSpan.test(endpoint)) {
-              parentIndex = parentIndex + 1;
-            }
-            int finalParentIndex = parentIndex;
+            int parentIndex = spanAssertions.size() - 1;
             spanAssertions.add(
                 span ->
                     assertIndexedControllerSpan(span, requestId)
-                        .hasParent(trace.getSpan(finalParentIndex)));
+                        .hasParent(trace.getSpan(parentIndex)));
 
             trace.hasSpansSatisfyingExactly(spanAssertions);
           });
@@ -646,8 +642,7 @@ protected void assertTheTraces(
 
             if (endpoint != NOT_FOUND) {
               int parentIndex = 0;
-              if (options.hasHandlerSpan.test(endpoint)
-                  && options.hasHandlerAsControllerParentSpan.test(endpoint)) {
+              if (options.hasHandlerSpan.test(endpoint)) {
                 parentIndex = spanAssertions.size() - 1;
               }
               int finalParentIndex = parentIndex;

testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/HttpServerTestOptions.java

@@ -52,7 +52,6 @@ public final class HttpServerTestOptions {
   Predicate<ServerEndpoint> hasHandlerSpan = unused -> false;
   Predicate<ServerEndpoint> hasResponseSpan = unused -> false;
   Predicate<ServerEndpoint> hasErrorPageSpans = unused -> false;
-  Predicate<ServerEndpoint> hasHandlerAsControllerParentSpan = unused -> true;
   Predicate<ServerEndpoint> hasResponseCustomizer = unused -> false;
 
   Predicate<ServerEndpoint> hasExceptionOnServerSpan = endpoint -> !hasHandlerSpan.test(endpoint);
@@ -142,13 +141,6 @@ public HttpServerTestOptions setHasErrorPageSpans(Predicate<ServerEndpoint> hasE
     return this;
   }
 
-  @CanIgnoreReturnValue
-  public HttpServerTestOptions setHasHandlerAsControllerParentSpan(
-      Predicate<ServerEndpoint> hasHandlerAsControllerParentSpan) {
-    this.hasHandlerAsControllerParentSpan = hasHandlerAsControllerParentSpan;
-    return this;
-  }
-
   @CanIgnoreReturnValue
   public HttpServerTestOptions setHasExceptionOnServerSpan(
       Predicate<ServerEndpoint> hasExceptionOnServerSpan) {