Merge branch 'optimize-reactor-hooks' of https://github.com/robotmrv/spring-cloud-sleuth into robotmrv-optimize-reactor-hooks

tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/FlowsScopePassingSpanSubscriberTests.java

@@ -17,34 +17,46 @@
 package org.springframework.cloud.sleuth.instrument.reactor;
 
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.stream.Stream;
 
 import org.assertj.core.presentation.StandardRepresentation;
 import org.awaitility.Awaitility;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.mockito.Mockito;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import reactor.core.CoreSubscriber;
+import reactor.core.Exceptions;
+import reactor.core.Scannable;
 import reactor.core.publisher.Hooks;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
 
 import org.springframework.cloud.sleuth.CurrentTraceContext;
 import org.springframework.cloud.sleuth.TraceContext;
+import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.BDDAssertions.then;
-import static org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth.scopePassingSpanOperator;
+import static org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth.onEachOperatorForOnEachInstrumentation;
 
 /**
  * @author Marcin Grzejszczak
  */
 public abstract class FlowsScopePassingSpanSubscriberTests {
 
+	static final String HOOK_KEY = "org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration";
+
 	static {
 		// AssertJ will recognise QueueSubscription implements queue and try to invoke
 		// iterator. That's not allowed, and will cause an exception
@@ -61,16 +73,17 @@ public abstract class FlowsScopePassingSpanSubscriberTests {
 
 	@BeforeEach
 	public void setup() {
-		Hooks.resetOnEachOperator(
-				"org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration");
-		Hooks.resetOnLastOperator(
-				"org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration");
+		Hooks.resetOnEachOperator();
+		Hooks.resetOnLastOperator();
 		Schedulers.resetOnScheduleHooks();
 	}
 
 	@AfterEach
 	public void close() {
 		springContext.close();
+		Hooks.resetOnEachOperator();
+		Hooks.resetOnLastOperator();
+		Schedulers.resetOnScheduleHooks();
 	}
 
 	@Test
@@ -78,7 +91,7 @@ public abstract class FlowsScopePassingSpanSubscriberTests {
 		springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext);
 		springContext.refresh();
 
-		Function<? super Publisher<Integer>, ? extends Publisher<Integer>> transformer = scopePassingSpanOperator(
+		Function<? super Publisher<Integer>, ? extends Publisher<Integer>> transformer = onEachOperatorForOnEachInstrumentation(
 				this.springContext);
 
 		try (CurrentTraceContext.Scope ws = currentTraceContext().newScope(context())) {
@@ -144,4 +157,65 @@ public abstract class FlowsScopePassingSpanSubscriberTests {
 		Awaitility.await().untilAsserted(() -> then(currentTraceContext().context()).isNull());
 	}
 
+	@ParameterizedTest
+	@MethodSource("should_not_double_wrap_async_publisher_Args")
+	public void should_not_double_wrap_async_publisher(String name, Supplier<Mono<Integer>> sourceSupplier) {
+		springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext);
+		springContext.registerBean(Tracer.class, () -> Mockito.mock(Tracer.class));
+		springContext.refresh();
+
+		Hooks.onEachOperator(HOOK_KEY, ReactorSleuth.onEachOperatorForOnEachInstrumentation(springContext));
+		Hooks.onLastOperator(HOOK_KEY, ReactorSleuth.onLastOperatorForOnEachInstrumentation(springContext));
+
+		AtomicBoolean once = new AtomicBoolean();
+		Hooks.onLastOperator("test", p -> {
+			// check only first onLast Hook
+			if (once.compareAndSet(false, true)) {
+				assertThat(p).isInstanceOf(TraceContextPropagator.class);
+				Object parent = Scannable.from(p).scanUnsafe(Scannable.Attr.PARENT);
+				assertThat(parent).isNotInstanceOf(TraceContextPropagator.class);
+			}
+			return p;
+		});
+		try (CurrentTraceContext.Scope ws = currentTraceContext().newScope(context())) {
+			Mono<Integer> source = sourceSupplier.get();
+
+			source.subscribe((Subscriber<? super Integer>) new CoreSubscriber<Integer>() {
+				@Override
+				public void onSubscribe(Subscription subscription) {
+					assertThat(subscription).isInstanceOf(ScopePassingSpanSubscriber.class);
+					ScopePassingSpanSubscriber<Integer> spanSubscriber = (ScopePassingSpanSubscriber) subscription;
+					Object parent = spanSubscriber.scanUnsafe(Scannable.Attr.PARENT);
+
+					assertThat(parent).isInstanceOf(Subscriber.class).isNotInstanceOf(ScopePassingSpanSubscriber.class);
+				}
+
+				@Override
+				public void onNext(Integer integer) {
+				}
+
+				@Override
+				public void onError(Throwable throwable) {
+					throw Exceptions.propagate(throwable);
+				}
+
+				@Override
+				public void onComplete() {
+				}
+			});
+
+		}
+	}
+
+	private static Stream<Arguments> should_not_double_wrap_async_publisher_Args() {
+		return Stream.of(
+				// async source is hidden by Mono.defer()
+				Arguments.of("hidden by defer",
+						(Supplier) () -> Mono
+								.defer(() -> Mono.fromSupplier(() -> 1).hide().subscribeOn(Schedulers.parallel()))),
+				// async source is directly accessible during subscription
+				Arguments.of("directly accessible",
+						(Supplier) () -> Mono.fromSupplier(() -> 1).hide().subscribeOn(Schedulers.parallel())));
+	}
+
 }

tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberSpringBootTests.java

@@ -20,10 +20,14 @@ import java.time.Duration;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.awaitility.Awaitility;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
+import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
 
 import org.springframework.beans.factory.annotation.Autowired;
@@ -45,10 +49,26 @@ public abstract class ScopePassingSpanSubscriberSpringBootTests {
 	@Autowired
 	CurrentTraceContext currentTraceContext;
 
+	private Scheduler subscribeScheduler;
+
+	private Scheduler secondScheduler;
+
 	protected abstract TraceContext context();
 
 	protected abstract TraceContext context2();
 
+	@BeforeEach
+	void setUp() {
+		subscribeScheduler = Schedulers.newSingle("subscribeThread2");
+		secondScheduler = Schedulers.newSingle("secondThread");
+	}
+
+	@AfterEach
+	void tearDown() {
+		subscribeScheduler.dispose();
+		secondScheduler.dispose();
+	}
+
 	@Test
 	public void should_pass_tracing_info_when_using_reactor() {
 		final AtomicReference<TraceContext> spanInOperation = new AtomicReference<>();
@@ -88,7 +108,7 @@ public abstract class ScopePassingSpanSubscriberSpringBootTests {
 
 		try (CurrentTraceContext.Scope ws = this.currentTraceContext.newScope(context())) {
 			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) -> {
+					.publishOn(secondScheduler).log("reactor.2").map((d) -> {
 						spanInOperation.set(this.currentTraceContext.context());
 						return d + 1;
 					}).map(d -> d + 1).blockLast();
@@ -113,6 +133,71 @@ public abstract class ScopePassingSpanSubscriberSpringBootTests {
 		then(this.currentTraceContext.context()).isNull();
 	}
 
+	@Test
+	public void should_pass_tracing_info_into_sources_when_using_reactor_async() {
+		final AtomicReference<TraceContext> spanInOperation = new AtomicReference<>();
+
+		try (CurrentTraceContext.Scope ws = this.currentTraceContext.newScope(context())) {
+			Mono.fromSupplier(() -> {
+				spanInOperation.set(this.currentTraceContext.context());
+				return 1;
+			}).publishOn(Schedulers.single()).log("reactor.1").map(d -> d + 1).map(d -> d + 1)
+					.publishOn(secondScheduler).log("reactor.2").map((d) -> d + 1).map(d -> d + 1)
+					.subscribeOn(subscribeScheduler).block();
+
+			Awaitility.await().untilAsserted(() -> then(spanInOperation.get()).isEqualTo(context()));
+			then(this.currentTraceContext.context()).isEqualTo(context());
+		}
+
+		then(this.currentTraceContext.context()).isNull();
+
+		try (CurrentTraceContext.Scope ws = this.currentTraceContext.newScope(context2())) {
+			Mono.fromCallable(() -> {
+				spanInOperation.set(this.currentTraceContext.context());
+				return 1;
+			}).log("reactor.").map(d -> d + 1).map(d -> d + 1).map(d -> d + 1).subscribeOn(subscribeScheduler).block();
+
+			then(this.currentTraceContext.context()).isEqualTo(context2());
+			then(spanInOperation.get()).isEqualTo(context2());
+		}
+
+		then(this.currentTraceContext.context()).isNull();
+	}
+
+	@Test
+	public void should_pass_tracing_info_when_using_reactor_async_processor() {
+		final AtomicReference<TraceContext> spanInOperation = new AtomicReference<>();
+
+		Sinks.One<Integer> one = Sinks.one();
+		try (CurrentTraceContext.Scope ws = this.currentTraceContext.newScope(context())) {
+			one.asMono()
+					// hook is not applied for Processors so first operator after it will
+					// not be decorated with brave context
+					.map(d -> d + 1).map((d) -> {
+						spanInOperation.set(this.currentTraceContext.context());
+						return d + 1;
+					}).map(d -> d + 1).doOnSubscribe(subscription -> {
+						final Thread thread = new Thread(() -> {
+							try {
+								Thread.sleep(300);
+							}
+							catch (InterruptedException e) {
+								e.printStackTrace();
+							}
+							one.tryEmitValue(0);
+						});
+						thread.setName("async_processor_source");
+						thread.setDaemon(true);
+						thread.start();
+					}).subscribeOn(Schedulers.boundedElastic()).block();
+
+			Awaitility.await().untilAsserted(() -> then(spanInOperation.get()).isEqualTo(context()));
+			then(this.currentTraceContext.context()).isEqualTo(context());
+		}
+
+		then(this.currentTraceContext.context()).isNull();
+	}
+
 	@Test
 	public void onlyConsidersContextDuringSubscribe() {
 		Mono<TraceContext> fromMono = Mono.fromCallable(this.currentTraceContext::context);

tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java

@@ -143,7 +143,7 @@ public abstract class ScopePassingSpanSubscriberTests {
 				"org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration");
 		Hooks.resetOnLastOperator(
 				"org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration");
-		Schedulers.removeExecutorServiceDecorator("s;euth");
+		Schedulers.resetOnScheduleHook("sleuth");
 	}
 
 	@AfterEach