Simplify webflux mono handling
MODIFY
instrumentation/spring/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/server/AdviceUtils.java
instrumentation/spring/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/server/AdviceUtils.java
@@ -7,19 +7,12 @@ package io.opentelemetry.javaagent.instrumentation.spring.webflux.server;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.StatusCode;
-import io.opentelemetry.context.Scope;
+import io.opentelemetry.context.Context;
import io.opentelemetry.instrumentation.api.tracer.ClassNames;
import java.util.Map;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Function;
-import org.reactivestreams.Publisher;
-import org.reactivestreams.Subscription;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
-import reactor.core.CoreSubscriber;
import reactor.core.publisher.Mono;
-import reactor.core.publisher.Operators;
-import reactor.util.context.Context;
public class AdviceUtils {
@@ -35,20 +28,10 @@ public class AdviceUtils {
return className + ".handle";
}
- public static <T> Mono<T> setPublisherSpan(
- Mono<T> mono, io.opentelemetry.context.Context context) {
- return mono.<T>transform(finishSpanNextOrError(context));
- }
-
- /**
- * Idea for this has been lifted from https://github.com/reactor/reactor-core/issues/947. Newer
- * versions of reactor-core have easier way to access context but we want to support older
- * versions.
- */
- public static <T> Function<? super Publisher<T>, ? extends Publisher<T>> finishSpanNextOrError(
- io.opentelemetry.context.Context context) {
- return Operators.lift(
- (scannable, subscriber) -> new SpanFinishingSubscriber<>(subscriber, context));
+ public static <T> Mono<T> setPublisherSpan(Mono<T> mono, Context context) {
+ return mono.doOnError(t -> finishSpanIfPresent(context, t))
+ .doOnSuccess(x -> finishSpanIfPresent(context, null))
+ .doOnCancel(() -> finishSpanIfPresent(context, null));
}
public static void finishSpanIfPresent(ServerWebExchange exchange, Throwable throwable) {
@@ -63,7 +46,7 @@ public class AdviceUtils {
}
}
- static void finishSpanIfPresent(io.opentelemetry.context.Context context, Throwable throwable) {
+ static void finishSpanIfPresent(Context context, Throwable throwable) {
if (context != null) {
Span span = Span.fromContext(context);
if (throwable != null) {
@@ -76,73 +59,7 @@ public class AdviceUtils {
private static void finishSpanIfPresentInAttributes(
Map<String, Object> attributes, Throwable throwable) {
- io.opentelemetry.context.Context context =
- (io.opentelemetry.context.Context) attributes.remove(CONTEXT_ATTRIBUTE);
+ Context context = (Context) attributes.remove(CONTEXT_ATTRIBUTE);
finishSpanIfPresent(context, throwable);
}
-
- public static class SpanFinishingSubscriber<T> implements CoreSubscriber<T>, Subscription {
-
- private final CoreSubscriber<? super T> subscriber;
- private final io.opentelemetry.context.Context otelContext;
- private final Context context;
- private final AtomicBoolean completed = new AtomicBoolean();
- private Subscription subscription;
-
- public SpanFinishingSubscriber(
- CoreSubscriber<? super T> subscriber, io.opentelemetry.context.Context otelContext) {
- this.subscriber = subscriber;
- this.otelContext = otelContext;
- context = subscriber.currentContext().put(Span.class, otelContext);
- }
-
- @Override
- public void onSubscribe(Subscription subscription) {
- this.subscription = subscription;
- try (Scope ignored = otelContext.makeCurrent()) {
- subscriber.onSubscribe(this);
- }
- }
-
- @Override
- public void onNext(T t) {
- try (Scope ignored = otelContext.makeCurrent()) {
- subscriber.onNext(t);
- }
- }
-
- @Override
- public void onError(Throwable t) {
- if (completed.compareAndSet(false, true)) {
- finishSpanIfPresent(otelContext, t);
- }
- subscriber.onError(t);
- }
-
- @Override
- public void onComplete() {
- if (completed.compareAndSet(false, true)) {
- finishSpanIfPresent(otelContext, null);
- }
- subscriber.onComplete();
- }
-
- @Override
- public Context currentContext() {
- return context;
- }
-
- @Override
- public void request(long n) {
- subscription.request(n);
- }
-
- @Override
- public void cancel() {
- if (completed.compareAndSet(false, true)) {
- finishSpanIfPresent(otelContext, null);
- }
- subscription.cancel();
- }
- }
}

Simplify webflux mono handling (#3816)