From 553eaa576c7f22cbd7a9eac7aaa5bfff973f98f9 Mon Sep 17 00:00:00 2001 From: Lauri Tulmin Date: Thu, 28 Sep 2023 21:25:34 +0300 Subject: [PATCH] Spring webflux: add user spans as children of the controller span (#9572) --- .../webflux/v5_0/server/AdviceUtils.java | 30 ++++++++++++++++++- .../DispatcherHandlerInstrumentation.java | 16 ++++++++++ .../server/HandlerAdapterInstrumentation.java | 5 ++++ .../v5_0/server/SpringWebfluxTest.java | 4 +-- .../ControllerSpringWebFluxServerTest.java | 1 - ...DelayedHandlerSpringWebFluxServerTest.java | 7 ----- .../junit/http/AbstractHttpServerTest.java | 11 ++----- .../junit/http/HttpServerTestOptions.java | 8 ----- 8 files changed, 55 insertions(+), 27 deletions(-) diff --git a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/AdviceUtils.java b/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/AdviceUtils.java index bcc31c9856..a7af1a299f 100644 --- a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/AdviceUtils.java +++ b/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 @@ package io.opentelemetry.javaagent.instrumentation.spring.webflux.v5_0.server; 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 @@ public final class AdviceUtils { } } + public static Mono wrapMono(Mono mono, Context context) { + if (context == null) { + return mono; + } + return new ContextMono<>(mono, context); + } + @FunctionalInterface interface OnSpanEnd { void end(Throwable throwable); } + private static class ContextMono extends Mono { + + private final Mono delegate; + private final Context parentContext; + + ContextMono(Mono delegate, Context parentContext) { + this.delegate = delegate; + this.parentContext = parentContext; + } + + @Override + public void subscribe(CoreSubscriber actual) { + try (Scope ignored = parentContext.makeCurrent()) { + delegate.subscribe(actual); + } + } + } + private AdviceUtils() {} } diff --git a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/DispatcherHandlerInstrumentation.java b/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/DispatcherHandlerInstrumentation.java index 03de0b20d4..118b0312e7 100644 --- a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/DispatcherHandlerInstrumentation.java +++ b/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 class DispatcherHandlerInstrumentation implements TypeInstrumentation { .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 class DispatcherHandlerInstrumentation implements TypeInstrumentation { } } } + + @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 mono) { + mono = AdviceUtils.wrapMono(mono, exchange.getAttribute(AdviceUtils.CONTEXT)); + } + } } diff --git a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/HandlerAdapterInstrumentation.java b/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/HandlerAdapterInstrumentation.java index f63b92422c..4c65ab4c64 100644 --- a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/HandlerAdapterInstrumentation.java +++ b/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 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.reactive.HandlerResult; import org.springframework.web.server.ServerWebExchange; +import reactor.core.publisher.Mono; public class HandlerAdapterInstrumentation implements TypeInstrumentation { @@ -86,6 +88,7 @@ public class HandlerAdapterInstrumentation implements TypeInstrumentation { @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( + @Advice.Return(readOnly = false) Mono mono, @Advice.Argument(0) ServerWebExchange exchange, @Advice.Argument(1) Object handler, @Advice.Thrown Throwable throwable, @@ -99,6 +102,8 @@ public class HandlerAdapterInstrumentation implements TypeInstrumentation { 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 diff --git a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/SpringWebfluxTest.java b/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/SpringWebfluxTest.java index f825cbf59f..0b3579bf7c 100644 --- a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/v5_0/server/SpringWebfluxTest.java +++ b/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 @@ public class SpringWebfluxTest { }, span -> span.hasName("tracedMethod") - .hasParent(trace.getSpan(0)) + .hasParent(trace.getSpan(1)) .hasTotalAttributeCount(0))); } @@ -410,7 +410,7 @@ public class SpringWebfluxTest { }, span -> span.hasName("tracedMethod") - .hasParent(trace.getSpan(parameter.annotatedMethod != null ? 0 : 1)) + .hasParent(trace.getSpan(1)) .hasTotalAttributeCount(0))); } diff --git a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/server/base/ControllerSpringWebFluxServerTest.java b/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/server/base/ControllerSpringWebFluxServerTest.java index 0d6cfde5ed..54bbe7360b 100644 --- a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/server/base/ControllerSpringWebFluxServerTest.java +++ b/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/server/base/ControllerSpringWebFluxServerTest.java @@ -74,7 +74,6 @@ public abstract class ControllerSpringWebFluxServerTest extends SpringWebFluxSer @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) diff --git a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/server/base/DelayedHandlerSpringWebFluxServerTest.java b/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/server/base/DelayedHandlerSpringWebFluxServerTest.java index 0fdd674797..b889416e8c 100644 --- a/instrumentation/spring/spring-webflux/spring-webflux-5.0/javaagent/src/test/java/server/base/DelayedHandlerSpringWebFluxServerTest.java +++ b/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 @@ public class DelayedHandlerSpringWebFluxServerTest extends HandlerSpringWebFluxS })); } } - - @Override - protected void configure(HttpServerTestOptions options) { - super.configure(options); - options.setHasHandlerAsControllerParentSpan(unused -> false); - } } diff --git a/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/AbstractHttpServerTest.java b/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/AbstractHttpServerTest.java index ac2544fac4..01ded41360 100644 --- a/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/AbstractHttpServerTest.java +++ b/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/AbstractHttpServerTest.java @@ -571,15 +571,11 @@ public abstract class AbstractHttpServerTest extends AbstractHttpServerU 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 @@ public abstract class AbstractHttpServerTest extends AbstractHttpServerU 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; diff --git a/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/HttpServerTestOptions.java b/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/HttpServerTestOptions.java index 2eca911628..2671232dc4 100644 --- a/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/HttpServerTestOptions.java +++ b/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/HttpServerTestOptions.java @@ -52,7 +52,6 @@ public final class HttpServerTestOptions { Predicate hasHandlerSpan = unused -> false; Predicate hasResponseSpan = unused -> false; Predicate hasErrorPageSpans = unused -> false; - Predicate hasHandlerAsControllerParentSpan = unused -> true; Predicate hasResponseCustomizer = unused -> false; Predicate hasExceptionOnServerSpan = endpoint -> !hasHandlerSpan.test(endpoint); @@ -142,13 +141,6 @@ public final class HttpServerTestOptions { return this; } - @CanIgnoreReturnValue - public HttpServerTestOptions setHasHandlerAsControllerParentSpan( - Predicate hasHandlerAsControllerParentSpan) { - this.hasHandlerAsControllerParentSpan = hasHandlerAsControllerParentSpan; - return this; - } - @CanIgnoreReturnValue public HttpServerTestOptions setHasExceptionOnServerSpan( Predicate hasExceptionOnServerSpan) {