From 81322756a45a64269a952437a1d79ad9d2d2e38d Mon Sep 17 00:00:00 2001 From: robotmrv Date: Tue, 8 Jun 2021 14:23:59 +0300 Subject: [PATCH] simplify ON_EACH reactor instrumentation (#1969) * simplify ON_EACH reactor instrumentation remove custom Publisher decorators that were mostly copy of reactor-core Lift Publishers and use new LIFTER Scannable Attribute instead * change constant scope * add note about test constant origin --- .../reactor/ReactorHooksHelper.java | 380 ++++-------------- .../instrument/reactor/ReactorSleuth.java | 9 +- .../reactor/ReactorHooksHelperTests.java | 237 ++--------- .../FlowsScopePassingSpanSubscriberTests.java | 9 +- 4 files changed, 115 insertions(+), 520 deletions(-) diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelper.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelper.java index dd759fbbd..65821a78b 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelper.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelper.java @@ -18,24 +18,11 @@ package org.springframework.cloud.sleuth.instrument.reactor; import java.util.Objects; import java.util.function.BiFunction; -import java.util.function.Consumer; -import java.util.function.Function; -import java.util.function.Predicate; import org.reactivestreams.Processor; import org.reactivestreams.Publisher; -import reactor.core.CoreSubscriber; -import reactor.core.Disposable; import reactor.core.Fuseable; import reactor.core.Scannable; -import reactor.core.publisher.ConnectableFlux; -import reactor.core.publisher.Flux; -import reactor.core.publisher.FluxOperator; -import reactor.core.publisher.GroupedFlux; -import reactor.core.publisher.Mono; -import reactor.core.publisher.MonoOperator; -import reactor.core.publisher.Operators; -import reactor.core.publisher.ParallelFlux; import reactor.util.annotation.Nullable; import org.springframework.util.Assert; @@ -87,9 +74,13 @@ import org.springframework.util.Assert; * }) * .subscribe(); *} + * + * @author Roman Matiushchenko */ final class ReactorHooksHelper { + static final String LIFTER_NAME = "org.springframework.cloud.sleuth.instrument.reactor.ReactorHooksHelper.ScopePassingLifter"; + // need a way to determine SYNC sources to not add redundant scope passing decorator // most of reactor-core SYNC sources are marked with SourceProducer interface static final Class sourceProducerClass; @@ -131,6 +122,10 @@ final class ReactorHooksHelper { } if (!isSync(current)) { + boolean isLifter = getLifterName(current) != null; + if (isLifter) { + return shouldDecorateLifter(current); + } return true; } @@ -142,8 +137,41 @@ final class ReactorHooksHelper { } } - static boolean isTraceContextPropagator(Publisher current) { - return current instanceof TraceContextPropagator; + /** + * xxxLift Publishers get their RunStyle from source Publisher. So need to check + * whether current chain was decorated with scope passing operator. + * @param p first not sync lifter Publisher in the chain. + * @return {@code true} if Publisher chain does not contain lifter with + * {@link #LIFTER_NAME} name. + */ + private static boolean shouldDecorateLifter(Publisher p) { + Publisher current = getParent(p); + while (true) { + if (current == null) { + // is start of the chain, Publisher without source or foreign Publisher + return true; + } + String lifterName = getLifterName(current); + if (isScopePassingLifter(lifterName)) { + return false; + } + if (lifterName == null) { + return true; + } + current = getParent(current); + } + } + + private static boolean isScopePassingLifter(String lifterName) { + return lifterName == LIFTER_NAME; + } + + private static String getLifterName(Publisher current) { + return Scannable.from(current).scan(Scannable.Attr.LIFTER); + } + + public static boolean isTraceContextPropagator(Publisher current) { + return current instanceof TraceContextPropagator || isScopePassingLifter(getLifterName(current)); } private static boolean isSourceProducer(Publisher p) { @@ -165,299 +193,39 @@ final class ReactorHooksHelper { } /** - * Decorates {@link Publisher} with {@link TraceContextPropagator} {@code Publisher}. - * Mostly it is a copy of reactor-core logic from - * {@code reactor.core.publisher.Operators#liftPublisher()}. - * @param filter the Predicate that the raw Publisher must pass for the transformation - * to occur - * @param lifter the {@link BiFunction} taking the raw {@link Publisher} and - * {@link CoreSubscriber}. The function must return a receiving {@link CoreSubscriber} - * that will immediately subscribe to the {@link Publisher}. - * @param the input and output type. - * @return a new {@link Function}. + * @param name function name. + * @param delegate delegate function. + * @param the type of the first argument to the function + * @param the type of the second argument to the function + * @param the type of the result of the function + * @return function that {@link Object#toString()} returns provided name which is used + * as a value of {@link Scannable.Attr#LIFTER} attribute. */ - public static Function, ? extends Publisher> liftPublisher(Predicate filter, - BiFunction, ? extends CoreSubscriber> lifter) { - Assert.notNull(lifter, "lifter is null"); - return publisher -> { - if (filter != null && !filter.test(publisher)) { - return (Publisher) publisher; - } + static BiFunction named(String name, BiFunction delegate) { + return new NamedLifter<>(name, delegate); + } + + static class NamedLifter implements BiFunction { + + private final BiFunction delegate; + + private final String name; + + NamedLifter(String name, BiFunction delegate) { + this.name = Objects.requireNonNull(name, "name"); + this.delegate = Objects.requireNonNull(delegate, "delegate"); + } + + @Override + public R apply(T t, U u) { + return delegate.apply(t, u); + } + + @Override + public String toString() { + return name; + } - if (publisher instanceof Mono) { - return new SleuthMonoLift<>(publisher, lifter); - } - if (publisher instanceof ParallelFlux) { - return new SleuthParallelLift<>((ParallelFlux) publisher, lifter); - } - if (publisher instanceof ConnectableFlux) { - return new SleuthConnectableLift<>((ConnectableFlux) publisher, lifter); - } - if (publisher instanceof GroupedFlux) { - return new SleuthGroupedLift<>((GroupedFlux) publisher, lifter); - } - return new SleuthFluxLift<>(publisher, lifter); - }; - } - -} - -/** - * Copy of {@code reactor.core.publisher.MonoLift} which implements - * {@link TraceContextPropagator}. - */ -final class SleuthMonoLift extends MonoOperator implements TraceContextPropagator { - - final BiFunction, ? extends CoreSubscriber> lifter; - - SleuthMonoLift(Publisher p, - BiFunction, ? extends CoreSubscriber> lifter) { - super(Mono.from(p)); - this.lifter = lifter; - } - - @Override - public void subscribe(CoreSubscriber actual) { - CoreSubscriber input = actual; - try { - input = Objects.requireNonNull(lifter.apply(source, actual), "Lifted subscriber MUST NOT be null"); - - source.subscribe(input); - } - catch (Throwable e) { - Operators.reportThrowInSubscribe(input, e); - return; - } - } - - @Override - public Object scanUnsafe(Attr key) { - if (key == Attr.RUN_STYLE) { - return Attr.RunStyle.SYNC; - } - return super.scanUnsafe(key); - } - -} - -/** - * Copy of {@code reactor.core.publisher.FluxLift} which implements - * {@link TraceContextPropagator}. - */ -final class SleuthFluxLift extends FluxOperator implements TraceContextPropagator { - - final BiFunction, ? extends CoreSubscriber> lifter; - - SleuthFluxLift(Publisher p, - BiFunction, ? extends CoreSubscriber> lifter) { - super(Flux.from(p)); - this.lifter = lifter; - } - - @Override - public void subscribe(CoreSubscriber actual) { - CoreSubscriber input = actual; - try { - input = Objects.requireNonNull(lifter.apply(source, actual), "Lifted subscriber MUST NOT be null"); - - source.subscribe(input); - } - catch (Throwable e) { - Operators.reportThrowInSubscribe(input, e); - return; - } - } - - @Override - public Object scanUnsafe(Attr key) { - if (key == Attr.RUN_STYLE) { - return Attr.RunStyle.SYNC; - } - return super.scanUnsafe(key); - } - -} - -/** - * Copy of {@code reactor.core.publisher.ConnectableLift} which implements - * {@link TraceContextPropagator}. - */ -final class SleuthConnectableLift extends ConnectableFlux implements Scannable, TraceContextPropagator { - - final ConnectableFlux source; - - final BiFunction, ? extends CoreSubscriber> lifter; - - SleuthConnectableLift(ConnectableFlux p, - BiFunction, ? extends CoreSubscriber> lifter) { - this.source = Objects.requireNonNull(p, "source"); - this.lifter = lifter; - } - - @Override - public int getPrefetch() { - return source.getPrefetch(); - } - - @Override - public void connect(Consumer cancelSupport) { - this.source.connect(cancelSupport); - } - - @Override - @Nullable - public Object scanUnsafe(Attr key) { - if (key == Attr.PREFETCH) { - return source.getPrefetch(); - } - if (key == Attr.PARENT) { - return source; - } - if (key == Attr.RUN_STYLE) { - return Attr.RunStyle.SYNC; - } - return null; - } - - @Override - public void subscribe(CoreSubscriber actual) { - CoreSubscriber input = actual; - try { - input = Objects.requireNonNull(lifter.apply(source, actual), "Lifted subscriber MUST NOT be null"); - - source.subscribe(input); - } - catch (Throwable e) { - Operators.reportThrowInSubscribe(input, e); - return; - } - } - -} - -/** - * Copy of {@code reactor.core.publisher.GroupedLift} which implements - * {@link TraceContextPropagator}. - */ -final class SleuthGroupedLift extends GroupedFlux implements Scannable, TraceContextPropagator { - - final BiFunction, ? extends CoreSubscriber> lifter; - - final GroupedFlux source; - - SleuthGroupedLift(GroupedFlux p, - BiFunction, ? extends CoreSubscriber> lifter) { - this.source = Objects.requireNonNull(p, "source"); - this.lifter = lifter; - } - - @Override - public int getPrefetch() { - return source.getPrefetch(); - } - - @Override - public K key() { - return source.key(); - } - - @Override - @Nullable - public Object scanUnsafe(Attr key) { - if (key == Attr.PARENT) { - return source; - } - if (key == Attr.PREFETCH) { - return getPrefetch(); - } - if (key == Attr.RUN_STYLE) { - return Attr.RunStyle.SYNC; - } - - return null; - } - - @Override - public String stepName() { - if (source instanceof Scannable) { - return Scannable.from(source).stepName(); - } - return Scannable.super.stepName(); - } - - @Override - public void subscribe(CoreSubscriber actual) { - CoreSubscriber input = lifter.apply(source, actual); - - Objects.requireNonNull(input, "Lifted subscriber MUST NOT be null"); - - source.subscribe(input); - } - -} - -/** - * Copy of {@code reactor.core.publisher.ParallelLift} which implements - * {@link TraceContextPropagator}. - */ -final class SleuthParallelLift extends ParallelFlux implements Scannable, TraceContextPropagator { - - final BiFunction, ? extends CoreSubscriber> lifter; - - final ParallelFlux source; - - SleuthParallelLift(ParallelFlux p, - BiFunction, ? extends CoreSubscriber> lifter) { - this.source = Objects.requireNonNull(p, "source"); - this.lifter = lifter; - } - - @Override - public int getPrefetch() { - return source.getPrefetch(); - } - - @Override - public int parallelism() { - return source.parallelism(); - } - - @Override - @Nullable - public Object scanUnsafe(Attr key) { - if (key == Attr.PARENT) { - return source; - } - if (key == Attr.PREFETCH) { - return getPrefetch(); - } - if (key == Attr.RUN_STYLE) { - return Attr.RunStyle.SYNC; - } - - return null; - } - - @Override - public String stepName() { - if (source instanceof Scannable) { - return Scannable.from(source).stepName(); - } - return Scannable.super.stepName(); - } - - @Override - public void subscribe(CoreSubscriber[] s) { - @SuppressWarnings("unchecked") - CoreSubscriber[] subscribers = new CoreSubscriber[parallelism()]; - - int i = 0; - while (i < subscribers.length) { - subscribers[i] = Objects.requireNonNull(lifter.apply(source, s[i]), "Lifted subscriber MUST NOT be null"); - i++; - } - - source.subscribe(subscribers); } } diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java index 82dd63265..db6a784a0 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java @@ -39,10 +39,13 @@ import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.internal.LazyBean; import org.springframework.context.ConfigurableApplicationContext; +import static org.springframework.cloud.sleuth.instrument.reactor.ReactorHooksHelper.named; + /** * Reactive Span pointcuts factories. * * @author Stephane Maldini + * @author Roman Matiushchenko * @since 2.0.0 */ // TODO: this is public as it is used out of package, but unlikely intended to be @@ -115,7 +118,7 @@ public abstract class ReactorSleuth { BiFunction, ? extends CoreSubscriber> lifter = liftFunction( springContext, lazyCurrentTraceContext, lazyTracer); - return ReactorHooksHelper.liftPublisher(shouldDecorate, lifter); + return Operators.liftPublisher(shouldDecorate, named(ReactorHooksHelper.LIFTER_NAME, lifter)); } static BiFunction, ? extends CoreSubscriber> liftFunction( @@ -247,7 +250,7 @@ public abstract class ReactorSleuth { return scopePassingSpanSubscriber.apply(pub, sub); }; - return ReactorHooksHelper.liftPublisher(p -> { + return Operators.liftPublisher(p -> { /* * this prevent double decoration when last operator in the chain is not SYNC * like {@code Mono.fromSuppler(() -> ...).subscribeOn(Schedulers.parallel())} @@ -263,7 +266,7 @@ public abstract class ReactorSleuth { } } return addContext; - }, skipIfNoTraceCtx); + }, named(ReactorHooksHelper.LIFTER_NAME, skipIfNoTraceCtx)); } private static Context context(CoreSubscriber sub) { diff --git a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelperTests.java b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelperTests.java index 82d7bd25f..f180076c7 100644 --- a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelperTests.java +++ b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelperTests.java @@ -17,36 +17,24 @@ package org.springframework.cloud.sleuth.instrument.reactor; import java.time.Duration; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import java.util.function.Function; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -import org.mockito.Mockito; import org.reactivestreams.Publisher; -import org.reactivestreams.Subscription; import reactor.core.CoreSubscriber; -import reactor.core.Disposable; -import reactor.core.Scannable; -import reactor.core.publisher.BaseSubscriber; -import reactor.core.publisher.ConnectableFlux; -import reactor.core.publisher.Flux; -import reactor.core.publisher.GroupedFlux; import reactor.core.publisher.Hooks; import reactor.core.publisher.Mono; import reactor.core.publisher.MonoOperator; import reactor.core.publisher.Operators; -import reactor.core.publisher.ParallelFlux; import reactor.core.publisher.Sinks; import reactor.core.scheduler.Schedulers; -import reactor.test.StepVerifier; import reactor.util.annotation.Nullable; -import reactor.util.context.Context; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatCode; +import static org.springframework.cloud.sleuth.instrument.reactor.ReactorHooksHelper.named; class ReactorHooksHelperTests { @@ -114,8 +102,16 @@ class ReactorHooksHelperTests { @Test public void shouldDecorateWhenAsyncSourceAndTraceContextPropagatorThanShouldNotDecorate() { - Function, ? extends Publisher> traceContextPropagator = ReactorHooksHelper - .liftPublisher(publisher -> true, (publisher, coreSubscriber) -> coreSubscriber); + Mono syncSourceWithLifter = Mono.delay(Duration.ofMillis(10)).as(TraceContextPropagatorOperator::new); + boolean actual = ReactorHooksHelper.shouldDecorate(syncSourceWithLifter.map(Function.identity())); + assertThat(actual).isFalse(); + } + + @Test + public void shouldDecorateWhenAsyncSourceAndScopePassingLifterThanShouldNotDecorate() { + Function, ? extends Publisher> traceContextPropagator = Operators.liftPublisher( + publisher -> true, + named(ReactorHooksHelper.LIFTER_NAME, (publisher, coreSubscriber) -> coreSubscriber)); Mono syncSourceWithLifter = Mono.delay(Duration.ofMillis(10)).transform(traceContextPropagator); boolean actual = ReactorHooksHelper.shouldDecorate(syncSourceWithLifter.map(Function.identity())); assertThat(actual).isFalse(); @@ -134,182 +130,13 @@ class ReactorHooksHelperTests { assertThat(actual).isTrue(); } - @Test - public void liftPublisherWhenFilterDiscardsThenReturnSource() { - Mono source = Mono.just(1); - final Function, ? extends Publisher> transformer = ReactorHooksHelper - .liftPublisher(p -> false, (p, s) -> s); - Mono transformed = source.transform(transformer); - assertThat(source).isSameAs(transformed); - } - - @Test - public void liftPublisherWhenFilterIsNullThenDecorate() { - Mono source = Mono.just(1); - final Function, ? extends Publisher> transformer = ReactorHooksHelper - .liftPublisher(null, (p, s) -> s); - Mono transformed = source.transform(transformer); - assertThat(source).isNotSameAs(transformed); - assertThat(transformed).isInstanceOf(SleuthMonoLift.class); - } - - @Test - public void liftPublisherWhenSourceMonoThenDecorateWithMono() { - Mono source = Mono.just(1); - final Function, ? extends Publisher> transformer = ReactorHooksHelper - .liftPublisher(p -> true, (p, s) -> new PlusOneSubscriber(s)); - Mono transformed = source.transform(transformer); - assertThat(source).isNotSameAs(transformed); - assertThat(transformed).isInstanceOf(SleuthMonoLift.class).isInstanceOf(TraceContextPropagator.class); - - StepVerifier.create(transformed).expectSubscription().expectNext(2).expectComplete() - .verify(Duration.ofSeconds(5)); - } - - @Test - public void liftPublisherWhenSourceFluxThenDecorateWithFlux() { - Flux source = Flux.just(1, 2, 3); - final Function, ? extends Publisher> transformer = ReactorHooksHelper - .liftPublisher(p -> true, (p, s) -> new PlusOneSubscriber(s)); - Flux transformed = source.transform(transformer); - assertThat(source).isNotSameAs(transformed); - assertThat(transformed).isInstanceOf(SleuthFluxLift.class).isInstanceOf(TraceContextPropagator.class); - - StepVerifier.create(transformed).expectSubscription().expectNext(2).expectNext(3).expectNext(4).expectComplete() - .verify(Duration.ofSeconds(5)); - } - - @SuppressWarnings({ "rawtypes", "unchecked" }) - @Test - public void liftPublisherWhenSourceParallelFluxThenDecorateWithParallelFlux() { - ParallelFlux source = Flux.just(1, 2, 3).parallel(); - - Function function = ReactorHooksHelper.liftPublisher(p -> true, - (p, s) -> (CoreSubscriber) new PlusOneSubscriber(s)); - - ParallelFlux transformed = source.transform(function); - assertThat(source).isNotSameAs(transformed); - assertThat(transformed).isInstanceOf(SleuthParallelLift.class).isInstanceOf(TraceContextPropagator.class); - Flux sorted = transformed.map(it -> it).sequential().sort(); - StepVerifier.create(sorted).expectSubscription().expectNext(2).expectNext(3).expectNext(4).expectComplete() - .verify(Duration.ofSeconds(5)); - } - - @SuppressWarnings({ "rawtypes", "unchecked" }) - @Test - public void liftPublisherWhenSourceConnectableFluxThenDecorateWithConnectableFlux() { - final AtomicBoolean sourceCanceled = new AtomicBoolean(); - ConnectableFlux source = Flux.just(1, 2, 3).doOnCancel(() -> sourceCanceled.set(true)).replay(); - - Function function = ReactorHooksHelper.liftPublisher(p -> true, - (p, s) -> (CoreSubscriber) new PlusOneSubscriber(s)); - - Flux transformed = source.transform(function); - assertThat(source).isNotSameAs(transformed); - assertThat(transformed).isInstanceOf(SleuthConnectableLift.class).isInstanceOf(TraceContextPropagator.class); - final AtomicReference cancelSource = new AtomicReference<>(); - Flux autoConnect = ((ConnectableFlux) transformed).autoConnect(1, cancelSource::set); - - StepVerifier.create(autoConnect).expectSubscription().expectNext(2).expectNext(3).expectNext(4).thenCancel() - .verify(Duration.ofSeconds(5)); - - assertThat(cancelSource.get()).isNotNull(); - cancelSource.get().dispose(); - assertThat(sourceCanceled).isTrue(); - } - - @SuppressWarnings({ "rawtypes", "unchecked" }) - @Test - public void liftPublisherWhenSourceGroupedFluxThenDecorateWithGroupedFlux() { - Flux> source = Flux.just(1, 2, 3).groupBy(i -> i % 2); - - Function function = ReactorHooksHelper.liftPublisher(p -> true, - (p, s) -> (CoreSubscriber) new PlusOneSubscriber(s)); - - GroupedFlux grouped = source.filter(it -> it.key() == 1).blockFirst(); - Flux transformed = grouped.transform(function); - assertThat(grouped).isNotSameAs(transformed); - assertThat(transformed).isInstanceOf(SleuthGroupedLift.class).isInstanceOf(TraceContextPropagator.class); - - assertThat(((GroupedFlux) transformed).key()).isEqualTo(1); - - StepVerifier.create(transformed).expectSubscription().expectNext(2).expectNext(4).expectComplete() - .verify(Duration.ofSeconds(5)); - } - - @SuppressWarnings({ "rawtypes", "unchecked" }) - @Test - public void sleuthMonoLiftWhenScanRunStyleThenSync() { - final Mono source = Mockito.mock(Mono.class); - - SleuthMonoLift lifter = new SleuthMonoLift(source, (p, s) -> s); - - assertThat(lifter.scanUnsafe(Scannable.Attr.RUN_STYLE)).isEqualTo(Scannable.Attr.RunStyle.SYNC); - assertThat(lifter.scanUnsafe(Scannable.Attr.PARENT)).isEqualTo(source); - assertThat(lifter.scanUnsafe(Scannable.Attr.PREFETCH)).isEqualTo(Integer.MAX_VALUE); - } - - @SuppressWarnings({ "rawtypes", "unchecked" }) - @Test - public void testSleuthFluxLiftScanUnsafe() { - final Flux source = Mockito.mock(Flux.class); - - SleuthFluxLift lifter = new SleuthFluxLift(source, (p, s) -> s); - - assertThat(lifter.scanUnsafe(Scannable.Attr.RUN_STYLE)).isEqualTo(Scannable.Attr.RunStyle.SYNC); - assertThat(lifter.scanUnsafe(Scannable.Attr.PARENT)).isEqualTo(source); - assertThat(lifter.scanUnsafe(Scannable.Attr.PREFETCH)).isEqualTo(lifter.getPrefetch()); - } - - @SuppressWarnings({ "rawtypes", "unchecked" }) - @Test - public void testSleuthConnectableLiftScanUnsafe() { - final ConnectableFlux source = Mockito.mock(ConnectableFlux.class); - int sourcePrefetch = 1; - Mockito.when(source.getPrefetch()).thenReturn(sourcePrefetch); - - SleuthConnectableLift lifter = new SleuthConnectableLift(source, (p, s) -> s); - - assertThat(lifter.scanUnsafe(Scannable.Attr.RUN_STYLE)).isEqualTo(Scannable.Attr.RunStyle.SYNC); - assertThat(lifter.scanUnsafe(Scannable.Attr.PARENT)).isEqualTo(source); - assertThat(lifter.scanUnsafe(Scannable.Attr.PREFETCH)).isEqualTo(sourcePrefetch); - } - - @SuppressWarnings({ "rawtypes", "unchecked" }) - @Test - public void testSleuthGroupedLiftFluxScanUnsafe() { - final GroupedFlux source = Mockito.mock(GroupedFlux.class); - int sourcePrefetch = 1; - Mockito.when(source.getPrefetch()).thenReturn(sourcePrefetch); - - SleuthGroupedLift lifter = new SleuthGroupedLift(source, (p, s) -> s); - - assertThat(lifter.scanUnsafe(Scannable.Attr.RUN_STYLE)).isEqualTo(Scannable.Attr.RunStyle.SYNC); - assertThat(lifter.scanUnsafe(Scannable.Attr.PARENT)).isEqualTo(source); - assertThat(lifter.scanUnsafe(Scannable.Attr.PREFETCH)).isEqualTo(sourcePrefetch); - } - - @SuppressWarnings({ "rawtypes", "unchecked" }) - @Test - public void testSleuthParallelLiftScanUnsafe() { - final ParallelFlux source = Mockito.mock(ParallelFlux.class); - int sourcePrefetch = 1; - Mockito.when(source.getPrefetch()).thenReturn(sourcePrefetch); - - SleuthParallelLift lifter = new SleuthParallelLift(source, (p, s) -> s); - - assertThat(lifter.scanUnsafe(Scannable.Attr.RUN_STYLE)).isEqualTo(Scannable.Attr.RunStyle.SYNC); - assertThat(lifter.scanUnsafe(Scannable.Attr.PARENT)).isEqualTo(source); - assertThat(lifter.scanUnsafe(Scannable.Attr.PREFETCH)).isEqualTo(sourcePrefetch); - } - static class CustomMonoWithoutRunStyleOperator extends MonoOperator { /** * Build a {@link MonoOperator} wrapper around the passed parent {@link Publisher} * @param source the {@link Publisher} to decorate */ - protected CustomMonoWithoutRunStyleOperator(Mono source) { + CustomMonoWithoutRunStyleOperator(Mono source) { super(source); } @@ -329,37 +156,29 @@ class ReactorHooksHelperTests { } - static class PlusOneSubscriber extends BaseSubscriber { + static class TraceContextPropagatorOperator extends MonoOperator implements TraceContextPropagator { - CoreSubscriber actual; - - PlusOneSubscriber(CoreSubscriber actual) { - this.actual = actual; + /** + * Build a {@link MonoOperator} wrapper around the passed parent {@link Publisher} + * @param source the {@link Publisher} to decorate + */ + TraceContextPropagatorOperator(Mono source) { + super(source); } @Override - public Context currentContext() { - return actual.currentContext(); + public void subscribe(CoreSubscriber actual) { + source.subscribe(actual); } + @Nullable @Override - protected void hookOnSubscribe(Subscription subscription) { - actual.onSubscribe(subscription); - } - - @Override - protected void hookOnNext(Integer value) { - actual.onNext(value + 1); - } - - @Override - protected void hookOnComplete() { - actual.onComplete(); - } - - @Override - protected void hookOnError(Throwable throwable) { - actual.onError(throwable); + public Object scanUnsafe(Attr key) { + // just to check that it pass by instanceof + if (Attr.RUN_STYLE == key) { + return null; + } + return super.scanUnsafe(key); } } diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/FlowsScopePassingSpanSubscriberTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/FlowsScopePassingSpanSubscriberTests.java index 57d2b2cbe..4553041a5 100644 --- a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/FlowsScopePassingSpanSubscriberTests.java +++ b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/FlowsScopePassingSpanSubscriberTests.java @@ -57,6 +57,11 @@ public abstract class FlowsScopePassingSpanSubscriberTests { static final String HOOK_KEY = "org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration"; + // Cannot access + // org.springframework.cloud.sleuth.instrument.reactor.ReactorHooksHelper.LIFTER_NAME + // directly from this test so copy its value + static final String LIFTER_NAME = "org.springframework.cloud.sleuth.instrument.reactor.ReactorHooksHelper.ScopePassingLifter"; + static { // AssertJ will recognise QueueSubscription implements queue and try to invoke // iterator. That's not allowed, and will cause an exception @@ -171,9 +176,9 @@ public abstract class FlowsScopePassingSpanSubscriberTests { Hooks.onLastOperator("test", p -> { // check only first onLast Hook if (once.compareAndSet(false, true)) { - assertThat(p).isInstanceOf(TraceContextPropagator.class); + assertThat(Scannable.from(p).scanUnsafe(Scannable.Attr.LIFTER)).isEqualTo(LIFTER_NAME); Object parent = Scannable.from(p).scanUnsafe(Scannable.Attr.PARENT); - assertThat(parent).isNotInstanceOf(TraceContextPropagator.class); + assertThat(Scannable.from(parent).scanUnsafe(Scannable.Attr.LIFTER)).isNotEqualTo(LIFTER_NAME); } return p; });