Optimize Reactor Instrumentation (#1122)

- Use onLastOperator instead of onEachOperator
- Do not instrument scalar publishers
- Revisit ReactorSleuthMethodInvocationProcessor to reduce ops overhead
- Use only one operator for the MethodInvocationProcessor
- Warning Behavior Change: @NewSpan will defer Span creation/context

fixes https://github.com/spring-cloud/spring-cloud-sleuth/issues/1098

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/annotation/ReactorSleuthMethodInvocationProcessor.java

@@ -17,17 +17,22 @@
 package org.springframework.cloud.sleuth.annotation;
 
 import java.lang.reflect.Method;
-import java.util.function.Consumer;
-import java.util.function.Supplier;
 
 import brave.Span;
 import brave.Tracer;
+import brave.Tracing;
+import brave.propagation.CurrentTraceContext;
 import org.aopalliance.intercept.MethodInvocation;
 import org.reactivestreams.Publisher;
-import org.springframework.util.StringUtils;
+import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
+import reactor.core.Scannable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
-import reactor.core.publisher.SignalType;
+import reactor.util.annotation.Nullable;
+import reactor.util.context.Context;
+
+import org.springframework.util.StringUtils;
 
 /**
  * Method Invocation Processor for Reactor.
@@ -40,6 +45,15 @@ class ReactorSleuthMethodInvocationProcessor
 
 	private NonReactorSleuthMethodInvocationProcessor nonReactorSleuthMethodInvocationProcessor;
 
+	Tracing tracing;
+
+	Tracing tracing() {
+		if (this.tracing == null) {
+			this.tracing = this.beanFactory.getBean(Tracing.class);
+		}
+		return this.tracing;
+	}
+
 	@Override
 	public Object process(MethodInvocation invocation, NewSpan newSpan,
 			ContinueSpan continueSpan) throws Throwable {
@@ -53,79 +67,271 @@ class ReactorSleuthMethodInvocationProcessor
 		}
 	}
 
+	@SuppressWarnings("unchecked")
 	private Object proceedUnderReactorSpan(MethodInvocation invocation, NewSpan newSpan,
 			ContinueSpan continueSpan) throws Throwable {
 		Span spanPrevious = tracer().currentSpan();
 		// in case of @ContinueSpan and no span in tracer we start new span and should
 		// close it on completion
-		boolean startNewSpan = newSpan != null || spanPrevious == null;
 		Span span;
-		if (startNewSpan) {
-			span = tracer().nextSpan();
-			newSpanParser().parse(invocation, newSpan, span);
+		if (newSpan != null || spanPrevious == null) {
+			span = null;
 		}
 		else {
 			span = spanPrevious;
 		}
+
 		String log = log(continueSpan);
-		boolean hasLog = StringUtils.hasText(log);
-		try (Tracer.SpanInScope ws = tracer().withSpanInScope(span)) {
-			Publisher<?> publisher = (Publisher) invocation.proceed();
-			Mono<Span> startSpan = Mono.defer(() -> withSpanInScope(span, () -> {
-				if (startNewSpan) {
-					span.start();
-				}
-				before(invocation, span, log, hasLog);
-				return Mono.just(span);
-			}));
-			if (publisher instanceof Mono) {
-				return startSpan
-						.flatMap(spanStarted -> ((Mono<?>) publisher)
-								.doOnError(onFailureReactor(log, hasLog, spanStarted))
-								.doFinally(afterReactor(startNewSpan, log, hasLog,
-										spanStarted)))
-						// put span in context so it can be used by
-						// ScopePassingSpanSubscriber
-						.subscriberContext(context -> context.put(Span.class, span));
-			}
-			else if (publisher instanceof Flux) {
-				return startSpan
-						.flatMapMany(spanStarted -> ((Flux<?>) publisher)
-								.doOnError(onFailureReactor(log, hasLog, spanStarted))
-								.doFinally(afterReactor(startNewSpan, log, hasLog,
-										spanStarted)))
-						// put span in context so it can be used by
-						// ScopePassingSpanSubscriber
-						.subscriberContext(context -> context.put(Span.class, span));
+		Publisher<?> publisher = (Publisher) invocation.proceed();
+
+		if (publisher instanceof Mono) {
+			return new MonoSpan((Mono<Object>) publisher,
+					this,
+					newSpan,
+					span,
+					invocation,
+					log);
+		}
+		else if (publisher instanceof Flux) {
+			return new FluxSpan((Flux<Object>) publisher,
+					this,
+					newSpan,
+					span,
+					invocation,
+					log);
+		}
+		else {
+			throw new IllegalArgumentException("Unexpected type of publisher: " + publisher.getClass());
+		}
+	}
+
+	private static final class FluxSpan extends Flux<Object> implements Scannable {
+
+		final Flux<Object>                           source;
+		final Span                                   span;
+		final MethodInvocation                       invocation;
+		final String                                 log;
+		final boolean                                hasLog;
+		final ReactorSleuthMethodInvocationProcessor processor;
+		final NewSpan                                newSpan;
+
+		FluxSpan(Flux<Object> source,
+				ReactorSleuthMethodInvocationProcessor processor,
+				NewSpan newSpan,
+				@Nullable Span span,
+				MethodInvocation invocation,
+				String log) {
+			this.source = source;
+			this.span = span;
+			this.newSpan = newSpan;
+			this.invocation = invocation;
+			this.log = log;
+			this.hasLog = StringUtils.hasText(log);
+			this.processor = processor;
+		}
+
+		@Override
+		public void subscribe(CoreSubscriber<? super Object> actual) {
+			Span span;
+			Tracer tracer = this.processor.tracer();
+			if (this.span == null) {
+				span = tracer.nextSpan();
+				this.processor.newSpanParser().parse(invocation, newSpan, span);
+				span.start();
 			}
 			else {
-				throw new IllegalArgumentException(
-						"Unexpected type of publisher: " + publisher.getClass());
+				span = this.span;
+			}
+			try (Tracer.SpanInScope ws = tracer.withSpanInScope(span)) {
+				this.source.subscribe(new SpanSubscriber(actual,
+						this.processor,
+						this.invocation,
+						this.span == null,
+						span,
+						this.log,
+						this.hasLog));
+			}
+		}
+
+		@Nullable
+		public Object scanUnsafe(Attr key) {
+			if (key == Attr.PARENT) {
+				return this.source;
+			}
+			else {
+				return null;
 			}
 		}
 	}
 
-	private <T> T withSpanInScope(Span span, Supplier<T> supplier) {
-		try (Tracer.SpanInScope ws1 = tracer().withSpanInScope(span)) {
-			return supplier.get();
+	private static final class MonoSpan extends Mono<Object> implements Scannable {
+
+		final Mono<Object>                           source;
+		final Span                                   span;
+		final MethodInvocation                       invocation;
+		final String                                 log;
+		final boolean                                hasLog;
+		final ReactorSleuthMethodInvocationProcessor processor;
+		final NewSpan                                newSpan;
+
+		MonoSpan(Mono<Object> source,
+				ReactorSleuthMethodInvocationProcessor processor,
+				NewSpan newSpan,
+				@Nullable Span span,
+				MethodInvocation invocation,
+				String log) {
+			this.source = source;
+			this.processor = processor;
+			this.newSpan = newSpan;
+			this.span = span;
+			this.invocation = invocation;
+			this.log = log;
+			this.hasLog = StringUtils.hasText(log);
+		}
+
+		@Override
+		public void subscribe(CoreSubscriber<? super Object> actual) {
+			Span span;
+			Tracer tracer = this.processor.tracer();
+			if (this.span == null) {
+				span = tracer.nextSpan();
+				this.processor.newSpanParser().parse(invocation, newSpan, span);
+				span.start();
+			}
+			else {
+				span = this.span;
+			}
+			try (Tracer.SpanInScope ws = tracer.withSpanInScope(span)) {
+				this.source.subscribe(new SpanSubscriber(actual,
+						this.processor,
+						this.invocation,
+						this.span == null,
+						span,
+						this.log,
+						this.hasLog));
+			}
+		}
+
+		@Nullable
+		public Object scanUnsafe(Attr key) {
+			if (key == Attr.PARENT) {
+				return this.source;
+			}
+			else {
+				return null;
+			}
 		}
 	}
 
-	private Consumer<SignalType> afterReactor(boolean isNewSpan, String log,
-			boolean hasLog, Span span) {
-		return signalType -> {
-			try (Tracer.SpanInScope ws = tracer().withSpanInScope(span)) {
-				after(span, isNewSpan, log, hasLog);
-			}
-		};
-	}
+	private static final class SpanSubscriber implements CoreSubscriber<Object>,
+	                                                     Subscription,
+	                                                     Scannable {
 
-	private Consumer<Throwable> onFailureReactor(String log, boolean hasLog, Span span) {
-		return throwable -> {
-			try (Tracer.SpanInScope ws = tracer().withSpanInScope(span)) {
-				onFailure(span, log, hasLog, throwable);
+		final CoreSubscriber<? super Object>         actual;
+		final boolean                                isNewSpan;
+		final Span                                   span;
+		final String                                 log;
+		final boolean                                hasLog;
+		final CurrentTraceContext                    currentTraceContext;
+		final ReactorSleuthMethodInvocationProcessor processor;
+		final Context                                context;
+
+		Subscription parent;
+
+		SpanSubscriber(CoreSubscriber<? super Object> actual,
+				ReactorSleuthMethodInvocationProcessor processor,
+				MethodInvocation invocation,
+				boolean isNewSpan,
+				Span span,
+				String log,
+				boolean hasLog) {
+			this.actual = actual;
+			this.isNewSpan = isNewSpan;
+			this.span = span;
+			this.log = log;
+			this.hasLog = hasLog;
+			this.processor = processor;
+
+			this.currentTraceContext = processor.tracing().currentTraceContext();
+			this.context = actual.currentContext().put(Span.class, span);
+
+			processor.before(invocation, this.span, this.log, this.hasLog);
+		}
+
+		@Override
+		public void request(long n) {
+			try (CurrentTraceContext.Scope scope = this.currentTraceContext
+					.maybeScope(this.span.context())) {
+				this.parent.request(n);
 			}
-		};
+		}
+
+		@Override
+		public void cancel() {
+			try (CurrentTraceContext.Scope scope = this.currentTraceContext
+					.maybeScope(this.span.context())) {
+				this.parent.cancel();
+			}
+			finally {
+				this.processor.after(this.span, this.isNewSpan, this.log, this.hasLog);
+			}
+		}
+
+		@Override
+		public Context currentContext() {
+			return context;
+		}
+
+		@Override
+		public void onSubscribe(Subscription subscription) {
+			this.parent = subscription;
+			try (CurrentTraceContext.Scope scope = this.currentTraceContext
+					.maybeScope(this.span.context())) {
+				this.actual.onSubscribe(this);
+			}
+		}
+
+		@Override
+		public void onNext(Object o) {
+			try (CurrentTraceContext.Scope scope = this.currentTraceContext
+					.maybeScope(this.span.context())) {
+				this.actual.onNext(o);
+			}
+		}
+
+		@Override
+		public void onError(Throwable error) {
+			try (CurrentTraceContext.Scope scope = this.currentTraceContext
+					.maybeScope(this.span.context())) {
+				this.processor.onFailure(this.span, this.log, this.hasLog, error);
+				this.actual.onError(error);
+			}
+			finally {
+				this.processor.after(this.span, this.isNewSpan, this.log, this.hasLog);
+			}
+		}
+
+		@Override
+		public void onComplete() {
+			try (CurrentTraceContext.Scope scope = this.currentTraceContext
+					.maybeScope(this.span.context())) {
+				this.actual.onComplete();
+			}
+			finally {
+				this.processor.after(this.span, this.isNewSpan, this.log, this.hasLog);
+			}
+		}
+
+		@Override
+		public Object scanUnsafe(Attr key) {
+			if (key == Attr.ACTUAL) {
+				return this.actual;
+			}
+			if (key == Attr.PARENT) {
+				return this.parent;
+			}
+			return null;
+		}
 	}
 
 	private boolean isReactorReturnType(Class<?> returnType) {

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java

@@ -16,21 +16,20 @@
 
 package org.springframework.cloud.sleuth.instrument.reactor;
 
+import java.util.function.BooleanSupplier;
 import java.util.function.Function;
 
+import brave.Span;
 import brave.Tracing;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.reactivestreams.Publisher;
 import org.springframework.beans.factory.BeanFactory;
-import org.springframework.beans.factory.NoSuchBeanDefinitionException;
-import org.springframework.context.ApplicationContext;
 import org.springframework.context.ConfigurableApplicationContext;
 
 import reactor.core.CoreSubscriber;
+import reactor.core.Fuseable;
 import reactor.core.Scannable;
-import reactor.core.publisher.ConnectableFlux;
-import reactor.core.publisher.GroupedFlux;
 import reactor.core.publisher.Operators;
 import reactor.util.context.Context;
 
@@ -51,7 +50,7 @@ public abstract class ReactorSleuth {
 	 * Return a span operator pointcut given a {@link Tracing}. This can be used in
 	 * reactor via {@link reactor.core.publisher.Flux#transform(Function)},
 	 * {@link reactor.core.publisher.Mono#transform(Function)},
-	 * {@link reactor.core.publisher.Hooks#onEachOperator(Function)} or
+	 * {@link reactor.core.publisher.Hooks#onLastOperator(Function)} or
 	 * {@link reactor.core.publisher.Hooks#onLastOperator(Function)}. The Span operator
 	 * pointcut will pass the Scope of the Span without ever creating any new spans.
 	 * @param beanFactory - {@link BeanFactory}
@@ -60,58 +59,63 @@ public abstract class ReactorSleuth {
 	 */
 	@SuppressWarnings("unchecked")
 	public static <T> Function<? super Publisher<T>, ? extends Publisher<T>> scopePassingSpanOperator(
-			ConfigurableApplicationContext beanFactory) {
+			BeanFactory beanFactory) {
 		if (log.isTraceEnabled()) {
 			log.trace("Scope passing operator [" + beanFactory + "]");
 		}
-		return (sourcePub -> {
-			// TODO: Remove this once Reactor 3.1.8 is released
-			// do the checks directly on actual original Publisher
-			if (sourcePub instanceof ConnectableFlux // Operators.lift can't handle that
-					|| sourcePub instanceof GroupedFlux // Operators.lift can't handle
-														// that
-			) {
-				return sourcePub;
-			}
-			// no more POINTCUT_FILTER since mechanism is broken
-			Function<? super Publisher<T>, ? extends Publisher<T>> lift = Operators
-					.lift((scannable, sub) -> {
-						// rest of the logic unchanged...
-						if (beanFactory.isActive()) {
-							if (log.isTraceEnabled()) {
-								log.trace("Spring Context [" + beanFactory
-										+ "] already refreshed. Creating a scope "
-										+ "passing span subscriber with Reactor Context "
-										+ "[" + sub.currentContext() + "] and name ["
-										+ scannable.name() + "]");
-							}
-							return scopePassingSpanSubscription(beanFactory, scannable,
-									sub).get();
-						}
-						if (log.isTraceEnabled()) {
-							log.trace("Spring Context [" + beanFactory
-									+ "] is not yet refreshed, falling back to lazy span subscriber. Reactor Context is ["
-									+ sub.currentContext() + "] and name is ["
-									+ scannable.name() + "]");
-						}
-						return new LazySpanSubscriber<T>(scopePassingSpanSubscription(
-								beanFactory, scannable, sub));
-					});
 
-			return lift.apply(sourcePub);
+		//Adapt if lazy bean factory
+		BooleanSupplier isActive =
+				beanFactory instanceof ConfigurableApplicationContext ?
+						((ConfigurableApplicationContext) beanFactory)::isActive :
+						() -> true;
+
+		return Operators.liftPublisher((p, sub) -> {
+			//if Flux/Mono #just, #empty, #error
+			if (p instanceof Fuseable.ScalarCallable) {
+				return sub;
+			}
+			Scannable scannable = Scannable.from(p);
+			// rest of the logic unchanged...
+			if (isActive.getAsBoolean()) {
+				if (log.isTraceEnabled()) {
+					log.trace("Spring Context [" + beanFactory
+							+ "] already refreshed. Creating a scope "
+							+ "passing span subscriber with Reactor Context "
+							+ "[" + sub.currentContext() + "] and name ["
+							+ scannable.name() + "]");
+				}
+
+				return scopePassingSpanSubscription(beanFactory.getBean(Tracing.class), sub);
+			}
+			if (log.isTraceEnabled()) {
+				log.trace("Spring Context [" + beanFactory
+						+ "] is not yet refreshed, falling back to lazy span subscriber. Reactor Context is ["
+						+ sub.currentContext() + "] and name is ["
+						+ scannable.name() + "]");
+			}
+			return new LazySpanSubscriber<>(lazyScopePassingSpanSubscription(beanFactory, scannable, sub));
 		});
 	}
 
-	private static <T> SpanSubscriptionProvider<T> scopePassingSpanSubscription(
+	static <T> SpanSubscriptionProvider<T> lazyScopePassingSpanSubscription(
 			BeanFactory beanFactory, Scannable scannable, CoreSubscriber<? super T> sub) {
-		return new SpanSubscriptionProvider<T>(beanFactory, sub, sub.currentContext(),
-				scannable.name()) {
-			@Override
-			SpanSubscription newCoreSubscriber(Tracing tracing) {
-				return new ScopePassingSpanSubscriber<T>(sub,
-						sub != null ? sub.currentContext() : Context.empty(), tracing);
-			}
-		};
+		return new SpanSubscriptionProvider<>(beanFactory, sub, sub.currentContext(), scannable.name());
+	}
+
+
+	static <T> CoreSubscriber<? super T> scopePassingSpanSubscription(
+			Tracing tracing, CoreSubscriber<? super T> sub) {
+
+		Context context = sub.currentContext();
+
+		Span root = context.hasKey(Span.class) ? context.get(Span.class) : tracing.tracer().currentSpan();
+		if (root != null) {
+			return new ScopePassingSpanSubscriber<>(sub, context, tracing, root);
+		}
+		else {
+			return sub; //no need to trace
+		}
 	}
 
 }

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java

@@ -16,6 +16,8 @@
 
 package org.springframework.cloud.sleuth.instrument.reactor;
 
+import javax.annotation.Nullable;
+
 import brave.Span;
 import brave.Tracer;
 import brave.Tracing;
@@ -25,6 +27,7 @@ import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
+import reactor.core.Scannable;
 import reactor.util.context.Context;
 
 /**
@@ -34,7 +37,7 @@ import reactor.util.context.Context;
  * @author Marcin Grzejszczak
  * @since 2.0.0
  */
-final class ScopePassingSpanSubscriber<T> implements SpanSubscription<T> {
+final class ScopePassingSpanSubscriber<T> implements SpanSubscription<T>, Scannable {
 
 	private static final Log log = LogFactory.getLog(ScopePassingSpanSubscriber.class);
 
@@ -46,18 +49,13 @@ final class ScopePassingSpanSubscriber<T> implements SpanSubscription<T> {
 
 	private final TraceContext traceContext;
 
-	private final Tracer tracer;
-
 	private Subscription s;
 
 	ScopePassingSpanSubscriber(Subscriber<? super T> subscriber, Context ctx,
-			Tracing tracing) {
+			Tracing tracing, @Nullable Span root) {
 		this.subscriber = subscriber;
-		this.tracer = tracing.tracer();
 		this.currentTraceContext = tracing.currentTraceContext();
-		Span root = ctx != null
-				? ctx.hasKey(Span.class) ? ctx.get(Span.class) : this.tracer.currentSpan()
-				: null;
+
 		this.traceContext = root == null ? null : root.context();
 		this.context = ctx != null && root != null ? ctx.put(Span.class, root)
 				: ctx != null ? ctx : Context.empty();
@@ -121,4 +119,12 @@ final class ScopePassingSpanSubscriber<T> implements SpanSubscription<T> {
 		return this.context;
 	}
 
+	@Override
+	public Object scanUnsafe(Attr key) {
+		if (key == Attr.PARENT) {
+			return this.s;
+		} else {
+			return key == Attr.ACTUAL ? this.subscriber : null;
+		}
+	}
 }

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriptionProvider.java

@@ -18,6 +18,7 @@ package org.springframework.cloud.sleuth.instrument.reactor;
 
 import java.util.function.Supplier;
 
+import brave.Span;
 import brave.Tracing;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -31,7 +32,7 @@ import reactor.util.context.Context;
  * @param <T> type of returned subscription
  * @author Marcin Grzejszczak
  */
-class SpanSubscriptionProvider<T> implements Supplier<SpanSubscription<T>> {
+final class SpanSubscriptionProvider<T> implements Supplier<SpanSubscription<T>> {
 
 	private static final Log log = LogFactory.getLog(SpanSubscriptionProvider.class);
 
@@ -63,7 +64,8 @@ class SpanSubscriptionProvider<T> implements Supplier<SpanSubscription<T>> {
 	}
 
 	SpanSubscription<T> newCoreSubscriber(Tracing tracing) {
-		return new ScopePassingSpanSubscriber<>(this.subscriber, this.context, tracing);
+		Span root = context.hasKey(Span.class) ? context.get(Span.class) : tracing.tracer().currentSpan();
+		return new ScopePassingSpanSubscriber<>(this.subscriber, this.context, tracing, root);
 	}
 
 	private Tracing tracing() {

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java

@@ -85,7 +85,7 @@ public class TraceReactorAutoConfiguration {
 			if (log.isTraceEnabled()) {
 				log.trace("Cleaning up hooks");
 			}
-			Hooks.resetOnEachOperator(SLEUTH_TRACE_REACTOR_KEY);
+			Hooks.resetOnLastOperator(SLEUTH_TRACE_REACTOR_KEY);
 			Schedulers.resetFactory();
 		}
 
@@ -115,7 +115,7 @@ class HookRegisteringBeanDefinitionRegistryPostProcessor
 	}
 
 	void setupHooks(BeanFactory beanFactory) {
-		Hooks.onEachOperator(
+		Hooks.onLastOperator(
 				TraceReactorAutoConfiguration.TraceReactorConfiguration.SLEUTH_TRACE_REACTOR_KEY,
 				ReactorSleuth.scopePassingSpanOperator(this.context));
 		Schedulers.setFactory(factoryInstance(beanFactory));

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectFluxTests.java

@@ -590,8 +590,7 @@ public class SleuthSpanCreatorAspectFluxTests {
 
 		@Override
 		public Flux<Long> newSpanInTraceContext() {
-			Long id = id(tracer);
-			return Flux.defer(() -> Flux.just(id));
+			return Flux.defer(() -> Flux.just(id(tracer)));
 		}
 
 		@Override

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java

@@ -22,6 +22,8 @@ import brave.Tracing;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.junit.MockitoJUnitRunner;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.BaseSubscriber;
 import reactor.util.context.Context;
 
 import static org.assertj.core.api.BDDAssertions.then;
@@ -36,16 +38,17 @@ public class ScopePassingSpanSubscriberTests {
 
 	@Test
 	public void should_propagate_current_context() {
-		ScopePassingSpanSubscriber subscriber = new ScopePassingSpanSubscriber(null,
-				Context.of("foo", "bar"), this.tracing);
+		ScopePassingSpanSubscriber<?> subscriber = new ScopePassingSpanSubscriber<>(null,
+				Context.of("foo", "bar"), this.tracing, null);
 
 		then((String) subscriber.currentContext().get("foo")).isEqualTo("bar");
 	}
 
 	@Test
 	public void should_set_empty_context_when_context_is_null() {
-		ScopePassingSpanSubscriber subscriber = new ScopePassingSpanSubscriber(null, null,
-				this.tracing);
+		ScopePassingSpanSubscriber<?> subscriber = new ScopePassingSpanSubscriber<>(null
+				, null,
+				this.tracing, null);
 
 		then(subscriber.currentContext().isEmpty()).isTrue();
 	}
@@ -55,8 +58,9 @@ public class ScopePassingSpanSubscriberTests {
 		Span span = this.tracing.tracer().nextSpan();
 		try (Tracer.SpanInScope ws = this.tracing.tracer()
 				.withSpanInScope(span.start())) {
-			ScopePassingSpanSubscriber subscriber = new ScopePassingSpanSubscriber(null,
-					Context.empty(), this.tracing);
+			CoreSubscriber<?> subscriber =
+					ReactorSleuth.scopePassingSpanSubscription(tracing, new BaseSubscriber<Object>() {
+					});
 
 			then(subscriber.currentContext().get(Span.class)).isEqualTo(span);
 		}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriberTests.java

@@ -17,31 +17,34 @@
 package org.springframework.cloud.sleuth.instrument.reactor;
 
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
 
 import brave.Span;
 import brave.Tracer;
 import brave.sampler.Sampler;
-import reactor.core.publisher.BaseSubscriber;
-import reactor.core.publisher.Flux;
-import reactor.core.publisher.Hooks;
-import reactor.core.publisher.Mono;
-import reactor.core.publisher.Operators;
-import reactor.core.scheduler.Schedulers;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.awaitility.Awaitility;
-import org.junit.AfterClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Hooks;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
+
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.test.context.junit4.SpringRunner;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.BDDAssertions.then;
 
 @RunWith(SpringRunner.class)
@@ -53,6 +56,9 @@ public class SpanSubscriberTests {
 	@Autowired
 	Tracer tracer;
 
+	@Autowired
+	ConfigurableApplicationContext factory;
+
 	@Test
 	public void should_pass_tracing_info_when_using_reactor() {
 		Span span = this.tracer.nextSpan().name("foo").start();
@@ -82,10 +88,10 @@ public class SpanSubscriberTests {
 
 		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) {
 			Mono.just(1).flatMap(d -> Flux.just(d + 1).collectList().map(p -> p.get(0)))
-					.map(d -> d + 1).map((d) -> {
-						spanInOperation.set(this.tracer.currentSpan());
-						return d + 1;
-					}).map(d -> d + 1).subscribe(System.out::println);
+			    .map(d -> d + 1).map((d) -> {
+				spanInOperation.set(this.tracer.currentSpan());
+				return d + 1;
+			}).map(d -> d + 1).subscribe(System.out::println);
 		}
 		finally {
 			span.finish();
@@ -98,47 +104,90 @@ public class SpanSubscriberTests {
 	@Test
 	public void should_not_trace_scalar_flows() {
 		Span span = this.tracer.nextSpan().name("foo").start();
-		final AtomicReference<Subscription> spanInOperation = new AtomicReference<>();
 		log.info("Hello");
 
+		//Disable global hooks for local hook testing
+		Hooks.resetOnLastOperator();
+
 		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) {
-			Mono.just(1).subscribe(new BaseSubscriber<Integer>() {
-				@Override
-				protected void hookOnSubscribe(Subscription subscription) {
-					spanInOperation.set(subscription);
-				}
-			});
 
-			then(this.tracer.currentSpan()).isNotNull();
+			Function<? super Publisher<Integer>, ? extends Publisher<Integer>> transformer =
+					ReactorSleuth.scopePassingSpanOperator(factory);
 
-			Mono.<Integer>error(new Exception()).subscribe(new BaseSubscriber<Integer>() {
+			Subscriber<Object> assertNoSpanSubscriber = new CoreSubscriber<Object>() {
 				@Override
-				protected void hookOnSubscribe(Subscription subscription) {
-					spanInOperation.set(subscription);
+				public void onSubscribe(Subscription s) {
+					s.request(Long.MAX_VALUE);
+					assertThat(s).isNotInstanceOf(ScopePassingSpanSubscriber.class);
 				}
 
 				@Override
-				protected void hookOnError(Throwable throwable) {
+				public void onNext(Object o) {
+
 				}
-			});
 
-			then(this.tracer.currentSpan()).isNotNull();
-
-			Mono.<Integer>empty().subscribe(new BaseSubscriber<Integer>() {
 				@Override
-				protected void hookOnSubscribe(Subscription subscription) {
-					spanInOperation.set(subscription);
-				}
-			});
+				public void onError(Throwable t) {
+
+				}
+
+				@Override
+				public void onComplete() {
+
+				}
+			};
+
+			Subscriber<Object> assertSpanSubscriber = new CoreSubscriber<Object>() {
+				@Override
+				public void onSubscribe(Subscription s) {
+					s.request(Long.MAX_VALUE);
+					assertThat(s).isInstanceOf(ScopePassingSpanSubscriber.class);
+				}
+
+				@Override
+				public void onNext(Object o) {
+
+				}
+
+				@Override
+				public void onError(Throwable t) {
+
+				}
+
+				@Override
+				public void onComplete() {
+
+				}
+			};
+			transformer.apply(Mono.just(1).hide())
+			           .subscribe(assertSpanSubscriber);
+
+			transformer.apply(Mono.just(1))
+			           .subscribe(assertNoSpanSubscriber);
+
+			transformer.apply(Mono.<Integer>error(new Exception()).hide())
+			           .subscribe(assertSpanSubscriber);
+
+
+			transformer.apply(Mono.error(new Exception()))
+			           .subscribe(assertNoSpanSubscriber);
+
+			transformer.apply(Mono.<Integer>empty().hide())
+			           .subscribe(assertSpanSubscriber);
+
+			transformer.apply(Mono.empty())
+			           .subscribe(assertNoSpanSubscriber);
+
+
 
-			then(this.tracer.currentSpan()).isNotNull();
-			then(spanInOperation.get()).isEqualTo(Operators.emptySubscription());
 		}
 		finally {
 			span.finish();
 		}
 
-		then(this.tracer.currentSpan()).isNull();
+		Awaitility.await().untilAsserted(() -> {
+			then(this.tracer.currentSpan()).isNull();
+		});
 	}
 
 	@Test
@@ -149,16 +198,16 @@ public class SpanSubscriberTests {
 
 		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) {
 			Flux.just(1, 2, 3).publishOn(Schedulers.single()).log("reactor.1")
-					.map(d -> d + 1).map(d -> d + 1)
-					.publishOn(Schedulers.newSingle("secondThread")).log("reactor.2")
-					.map((d) -> {
-						spanInOperation.set(this.tracer.currentSpan());
-						return d + 1;
-					}).map(d -> d + 1).blockLast();
+			    .map(d -> d + 1).map(d -> d + 1)
+			    .publishOn(Schedulers.newSingle("secondThread")).log("reactor.2")
+			    .map((d) -> {
+				    spanInOperation.set(this.tracer.currentSpan());
+				    return d + 1;
+			    }).map(d -> d + 1).blockLast();
 
 			Awaitility.await().untilAsserted(() -> {
-				then(spanInOperation.get().context().spanId())
-						.isEqualTo(span.context().spanId());
+				then(spanInOperation.get().context().traceId())
+						.isEqualTo(span.context().traceId());
 			});
 			then(this.tracer.currentSpan()).isEqualTo(span);
 		}
@@ -171,15 +220,15 @@ public class SpanSubscriberTests {
 
 		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(foo2)) {
 			Flux.just(1, 2, 3).publishOn(Schedulers.single()).log("reactor.")
-					.map(d -> d + 1).map(d -> d + 1).map((d) -> {
-						spanInOperation.set(this.tracer.currentSpan());
-						return d + 1;
-					}).map(d -> d + 1).blockLast();
+			    .map(d -> d + 1).map(d -> d + 1).map((d) -> {
+				spanInOperation.set(this.tracer.currentSpan());
+				return d + 1;
+			}).map(d -> d + 1).blockLast();
 
 			then(this.tracer.currentSpan()).isEqualTo(foo2);
 			// parent cause there's an async span in the meantime
-			then(spanInOperation.get().context().spanId())
-					.isEqualTo(foo2.context().spanId());
+			then(spanInOperation.get().context().traceId())
+					.isEqualTo(foo2.context().traceId());
 		}
 		finally {
 			foo2.finish();
@@ -194,11 +243,11 @@ public class SpanSubscriberTests {
 		log.info("Hello");
 		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(parentSpan)) {
 			final Long spanId = Mono.fromCallable(tracer::currentSpan)
-					.map(span -> span.context().spanId()).block();
+			                        .map(span -> span.context().spanId()).block();
 			then(spanId).isNotNull();
 
 			final Long secondSpanId = Mono.fromCallable(tracer::currentSpan)
-					.map(span -> span.context().spanId()).block();
+			                              .map(span -> span.context().spanId()).block();
 			then(secondSpanId).isEqualTo(spanId); // different trace ids here
 		}
 	}
@@ -211,19 +260,19 @@ public class SpanSubscriberTests {
 
 		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(initSpan)) {
 			Mono.fromCallable(tracer::currentSpan).map(span -> span.context().spanId())
-					.doOnNext(spanInOperation::set)
-					.zipWith(Mono.fromCallable(tracer::currentSpan)
-							.map(span -> span.context().spanId())
-							.doOnNext(spanInZipOperation::set))
-					.block();
+			    .doOnNext(spanInOperation::set)
+			    .zipWith(Mono.fromCallable(tracer::currentSpan)
+			                 .map(span -> span.context().spanId())
+			                 .doOnNext(spanInZipOperation::set))
+			    .block();
 		}
 
 		then(spanInZipOperation).hasValue(initSpan.context().spanId()); // ok here
 		then(spanInOperation).hasValue(initSpan.context().spanId()); // Expecting
-																		// <AtomicReference[null]>
-																		// to have value:
-																		// <1L> but did
-																		// not.
+		// <AtomicReference[null]>
+		// to have value:
+		// <1L> but did
+		// not.
 	}
 
 	// #646
@@ -234,8 +283,8 @@ public class SpanSubscriberTests {
 
 		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(initSpan)) {
 			Mono.just("value1")
-					.flatMap(request -> Mono.just("value2").then(Mono.just("foo")))
-					.map(a -> "qwe").block();
+			    .flatMap(request -> Mono.just("value2").then(Mono.just("foo")))
+			    .map(a -> "qwe").block();
 		}
 	}
 
@@ -247,19 +296,13 @@ public class SpanSubscriberTests {
 
 		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(initSpan)) {
 			Mono.subscriberContext()
-					.map(context -> tracer.currentSpan().context().spanId())
-					.doOnNext(spanInSubscriberContext::set).block();
+			    .map(context -> tracer.currentSpan().context().spanId())
+			    .doOnNext(spanInSubscriberContext::set).block();
 		}
 
 		then(spanInSubscriberContext).hasValue(initSpan.context().spanId()); // ok here
 	}
 
-	@AfterClass
-	public static void cleanup() {
-		Hooks.resetOnEachOperator();
-		Schedulers.resetFactory();
-	}
-
 	@EnableAutoConfiguration
 	@Configuration
 	static class Config {

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfigurationAccessorConfiguration.java

@@ -18,7 +18,7 @@ public class TraceReactorAutoConfigurationAccessorConfiguration {
 			log.trace("Cleaning up hooks");
 		}
 		new TraceReactorAutoConfiguration.TraceReactorConfiguration().cleanupHooks();
-		Hooks.resetOnEachOperator();
+		Hooks.resetOnLastOperator();
 		Hooks.resetOnLastOperator();
 		Schedulers.resetFactory();
 	}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java

@@ -544,17 +544,15 @@ class TestBean {
 	@NewSpan(name = "newSpanInTraceContext")
 	public Mono<Long> newSpanInTraceContext() {
 		log.info("New Span in Trace Context");
-		Long span = tracer.currentSpan().context().spanId();
-		return Mono.defer(() -> Mono.just(span));
+		return Mono.defer(() -> Mono.just(tracer.currentSpan().context().spanId()));
 	}
 
 	@NewSpan(name = "newSpanInSubscriberContext")
 	public Mono<Long> newSpanInSubscriberContext() {
 		log.info("New Span in Subscriber Context");
-		Long span = tracer.currentSpan().context().spanId();
 		return Mono.subscriberContext()
 				.doOnSuccess(context -> log.info("New Span in deferred Trace Context"))
-				.flatMap(context -> Mono.defer(() -> Mono.just(span)));
+				.flatMap(context -> Mono.defer(() -> Mono.just(tracer.currentSpan().context().spanId())));
 	}
 
 }