Spring webflux: add user spans as children of the controller span (#9572)

This commit is contained in:
Lauri Tulmin 2023-09-28 21:25:34 +03:00 committed by GitHub
parent c3ef68ff88
commit 553eaa576c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 55 additions and 27 deletions

View File

@ -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 <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() {}
}

View File

@ -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<Void> mono) {
mono = AdviceUtils.wrapMono(mono, exchange.getAttribute(AdviceUtils.CONTEXT));
}
}
}

View File

@ -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<HandlerResult> 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

View File

@ -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)));
}

View File

@ -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)

View File

@ -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);
}
}

View File

@ -571,15 +571,11 @@ public abstract class AbstractHttpServerTest<SERVER> 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<SERVER> 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;

View File

@ -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 final class HttpServerTestOptions {
return this;
}
@CanIgnoreReturnValue
public HttpServerTestOptions setHasHandlerAsControllerParentSpan(
Predicate<ServerEndpoint> hasHandlerAsControllerParentSpan) {
this.hasHandlerAsControllerParentSpan = hasHandlerAsControllerParentSpan;
return this;
}
@CanIgnoreReturnValue
public HttpServerTestOptions setHasExceptionOnServerSpan(
Predicate<ServerEndpoint> hasExceptionOnServerSpan) {