diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/SleuthReactorProperties.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/SleuthReactorProperties.java index 881f0530f..0c22d1eaf 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/SleuthReactorProperties.java +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/SleuthReactorProperties.java @@ -89,9 +89,7 @@ public class SleuthReactorProperties { /** * Decorates on each operator, will be less performing, but logging will always * contain the tracing entries in each operator. - * @deprecated to be removed in Sleuth 4.0.0 */ - @Deprecated DECORATE_ON_EACH, /** diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfiguration.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfiguration.java index e54d09928..e81b12d70 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfiguration.java +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfiguration.java @@ -48,7 +48,9 @@ import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.core.env.ConfigurableEnvironment; +import static org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.SLEUTH_REACTOR_EXECUTOR_SERVICE_KEY; import static org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration.SLEUTH_TRACE_REACTOR_KEY; +import static org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth.onLastOperatorForOnEachInstrumentation; /** * {@link org.springframework.boot.autoconfigure.EnableAutoConfiguration @@ -125,12 +127,18 @@ class HooksRefresher implements ApplicationListener } Hooks.resetOnEachOperator(SLEUTH_TRACE_REACTOR_KEY); Hooks.resetOnLastOperator(SLEUTH_TRACE_REACTOR_KEY); + Hooks.resetOnLastOperator(SLEUTH_REACTOR_EXECUTOR_SERVICE_KEY); switch (this.reactorProperties.getInstrumentationType()) { case DECORATE_ON_EACH: if (log.isTraceEnabled()) { log.trace("Decorating onEach operator instrumentation"); } - Hooks.onEachOperator(SLEUTH_TRACE_REACTOR_KEY, ReactorSleuth.scopePassingSpanOperator(this.context)); + Hooks.onEachOperator(SLEUTH_TRACE_REACTOR_KEY, + ReactorSleuth.onEachOperatorForOnEachInstrumentation(this.context)); + Hooks.onLastOperator(SLEUTH_TRACE_REACTOR_KEY, + ReactorSleuth.onLastOperatorForOnEachInstrumentation(this.context)); + Schedulers.onScheduleHook(TraceReactorAutoConfiguration.SLEUTH_REACTOR_EXECUTOR_SERVICE_KEY, + ReactorSleuth.scopePassingOnScheduleHook(this.context)); break; case DECORATE_ON_LAST: if (log.isTraceEnabled()) { @@ -178,15 +186,23 @@ class HookRegisteringBeanDefinitionRegistryPostProcessor implements BeanDefiniti } else if (property == SleuthReactorProperties.InstrumentationType.DECORATE_ON_EACH) { decorateOnEach(springContext); + decorateOnLast(onLastOperatorForOnEachInstrumentation(springContext)); + decorateScheduler(springContext); } else if (property == SleuthReactorProperties.InstrumentationType.DECORATE_ON_LAST) { decorateOnLast(ReactorSleuth.scopePassingSpanOperator(springContext)); + decorateScheduler(springContext); } else if (property == SleuthReactorProperties.InstrumentationType.MANUAL) { decorateOnLast(ReactorSleuth.springContextSpanOperator(springContext)); } } + private static void decorateScheduler(ConfigurableApplicationContext springContext) { + Schedulers.onScheduleHook(TraceReactorAutoConfiguration.SLEUTH_REACTOR_EXECUTOR_SERVICE_KEY, + ReactorSleuth.scopePassingOnScheduleHook(springContext)); + } + private static void decorateOnLast(Function, ? extends Publisher> function) { if (log.isTraceEnabled()) { log.trace("Decorating onLast operator instrumentation"); @@ -198,7 +214,8 @@ class HookRegisteringBeanDefinitionRegistryPostProcessor implements BeanDefiniti if (log.isTraceEnabled()) { log.trace("Decorating onEach operator instrumentation"); } - Hooks.onEachOperator(SLEUTH_TRACE_REACTOR_KEY, ReactorSleuth.scopePassingSpanOperator(springContext)); + Hooks.onEachOperator(SLEUTH_TRACE_REACTOR_KEY, + ReactorSleuth.onEachOperatorForOnEachInstrumentation(springContext)); } @Override @@ -208,7 +225,7 @@ class HookRegisteringBeanDefinitionRegistryPostProcessor implements BeanDefiniti } Hooks.resetOnEachOperator(SLEUTH_TRACE_REACTOR_KEY); Hooks.resetOnLastOperator(SLEUTH_TRACE_REACTOR_KEY); - Schedulers.removeExecutorServiceDecorator(TraceReactorAutoConfiguration.SLEUTH_REACTOR_EXECUTOR_SERVICE_KEY); + Schedulers.resetOnScheduleHook(TraceReactorAutoConfiguration.SLEUTH_REACTOR_EXECUTOR_SERVICE_KEY); } } diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfigurationAccessorConfiguration.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfigurationAccessorConfiguration.java index 64faaea9d..eea405a1a 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfigurationAccessorConfiguration.java +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfigurationAccessorConfiguration.java @@ -43,7 +43,7 @@ public final class TraceReactorAutoConfigurationAccessorConfiguration { } Hooks.resetOnEachOperator(SLEUTH_TRACE_REACTOR_KEY); Hooks.resetOnLastOperator(SLEUTH_TRACE_REACTOR_KEY); - Schedulers.removeExecutorServiceDecorator(SLEUTH_REACTOR_EXECUTOR_SERVICE_KEY); + Schedulers.resetOnScheduleHook(SLEUTH_REACTOR_EXECUTOR_SERVICE_KEY); } public static void setup(ConfigurableApplicationContext context) { diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/web/TraceWebFluxConfiguration.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/web/TraceWebFluxConfiguration.java index 0b64a2435..65a2fd9ec 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/web/TraceWebFluxConfiguration.java +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/web/TraceWebFluxConfiguration.java @@ -18,6 +18,7 @@ package org.springframework.cloud.sleuth.autoconfig.instrument.web; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication; +import org.springframework.cloud.sleuth.CurrentTraceContext; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.http.HttpServerHandler; import org.springframework.cloud.sleuth.instrument.web.TraceWebFilter; @@ -38,8 +39,8 @@ class TraceWebFluxConfiguration { @Bean TraceWebFilter traceFilter(Tracer tracer, HttpServerHandler httpServerHandler, - SleuthWebProperties sleuthWebProperties) { - TraceWebFilter traceWebFilter = new TraceWebFilter(tracer, httpServerHandler); + CurrentTraceContext currentTraceContext, SleuthWebProperties sleuthWebProperties) { + TraceWebFilter traceWebFilter = new TraceWebFilter(tracer, httpServerHandler, currentTraceContext); traceWebFilter.setOrder(sleuthWebProperties.getFilterOrder()); return traceWebFilter; } diff --git a/spring-cloud-sleuth-instrumentation/pom.xml b/spring-cloud-sleuth-instrumentation/pom.xml index 6e6b84cf4..91118b215 100644 --- a/spring-cloud-sleuth-instrumentation/pom.xml +++ b/spring-cloud-sleuth-instrumentation/pom.xml @@ -166,6 +166,11 @@ archunit-junit5 test + + io.projectreactor + reactor-test + test + diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/annotation/ReactorSleuthMethodInvocationProcessor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/annotation/ReactorSleuthMethodInvocationProcessor.java index fcaeee33a..76c0a34d5 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/annotation/ReactorSleuthMethodInvocationProcessor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/annotation/ReactorSleuthMethodInvocationProcessor.java @@ -36,6 +36,7 @@ import org.springframework.cloud.sleuth.TraceContext; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.annotation.ContinueSpan; import org.springframework.cloud.sleuth.annotation.NewSpan; +import org.springframework.cloud.sleuth.instrument.reactor.TraceContextPropagator; import org.springframework.util.StringUtils; /** @@ -146,7 +147,7 @@ public class ReactorSleuthMethodInvocationProcessor extends AbstractSleuthMethod } - private static final class MonoSpan extends MonoOperator { + private static final class MonoSpan extends MonoOperator implements TraceContextPropagator { final Span span; @@ -189,6 +190,14 @@ public class ReactorSleuthMethodInvocationProcessor extends AbstractSleuthMethod } } + @Override + public Object scanUnsafe(Attr key) { + if (key == Attr.RUN_STYLE) { + return Attr.RunStyle.SYNC; + } + return super.scanUnsafe(key); + } + } private static final class SpanSubscriber implements CoreSubscriber, Subscription, Scannable { 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 new file mode 100644 index 000000000..b50dc152b --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelper.java @@ -0,0 +1,476 @@ +/* + * Copyright 2013-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +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; + +/** + * Helper class for reactor ON_EACH instrumentation. Reduces number of sleuth operators + * time in the reactive chain during Assembly by using {@code Scannable.Attr.RUN_STYLE}. + * Below is the example of what it does
{@code
+ *	Mono.just(0) // (-)
+ * 	.map(it -> 1) // (-)
+ * 	.flatMap(it ->
+ * 		Mono.just(it) // (-)
+ * 	) // (-)
+ * 	.flatMapMany(it ->
+ * 		asyncPublisher(it) // (+)(*)
+ * 	) // (-)
+ * 	.filter(it -> true) // (-)
+ * 	.publishOn(Schedulers.single()) //(+)
+ * 	.map(it -> it) // (-)
+ * 	.map(it -> it * 10) // (-)
+ * 	.filter(it -> true) // (-)
+ * 	.scan((l, r) -> l + r) // (-)
+ * 	.doOnNext(it -> { // (-)
+ * 		//log
+ * 	})
+ * 	.doFirst(() -> { // (-)
+ * 		//log
+ * 	})
+ * 	.doFinally(signalType -> { // (-)
+ * 		//log
+ * 	})
+ * 	.subscribeOn(Schedulers.parallel()) //(+)
+ * 	.subscribe();//(*)
+ * 	(*) - captures tracing context if it differs from what was captured before at subscription and propagates it.
+ *  As far as check is performed on Subscriber Context it allocates Publisher to access it.
+ * 	(+) - is ASYNC need to decorate Publisher
+ * 	(-) - is SYNC no need to decorate Publisher
+ *}
So it creates 13 (out of 16) less Publisher (+ Subscriber and all their stuff) + * on assembly time comparing to the original logic (decorate every Publisher and checking + * only at subscription time). + *

+ * It also handles the case when onEachHook was not applied for some operator in the chain + * (It could be possible if custom operators or Processor are used) by checkin source + * source Publisher.

{@code
+ * 	processor/customOperator // (?) is async and hook is not applied
+ * 	.map(it -> ...) // (+) is SYNC but should add hook as previous Processor/operator does not use hooks
+ * 	.doOnNext(it -> { // (-) is SYNC no need to wrap
+ * 		//log
+ * 	})
+ * 	.subscribe();
+ *}
+ */ +final class ReactorHooksHelper { + + // 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; + static { + Class c; + try { + c = Class.forName("reactor.core.publisher.SourceProducer"); + } + catch (ClassNotFoundException e) { + c = Void.class; + } + sourceProducerClass = c; + } + + private ReactorHooksHelper() { + } + + /** + * Determines whether to decorate input publisher with + * {@code ScopePassingSpanOperator}. + * @param p publisher to check + * @return returns true if input publisher or one of its source publishers is not + * {@code RunStyle.SYNC} and there is no {@link TraceContextPropagator} between input + * publisher and not SYNC publisher. + */ + public static boolean shouldDecorate(Publisher p) { + Assert.notNull(p, "source Publisher is null"); + Publisher current = p; + while (true) { + if (current == null) { + // is start of the chain, Publisher without source or foreign Publisher + return true; + } + if (current instanceof Fuseable.ScalarCallable) { + return false; + } + if (isTraceContextPropagator(current)) { + return false; + } + if (!isLifter(current)) { + if (!isSync(current)) { + return true; + } + + if (isSourceProducer(current)) { + return false; + } + } + + current = getParent(current); + } + } + + private static boolean isTraceContextPropagator(Publisher current) { + return current instanceof TraceContextPropagator; + } + + private static boolean isSourceProducer(Publisher p) { + return sourceProducerClass.isInstance(p); + } + + private static boolean isSync(Publisher p) { + return !(p instanceof Processor) + && Scannable.Attr.RunStyle.SYNC == Scannable.from(p).scan(Scannable.Attr.RUN_STYLE); + } + + @Nullable + private static Publisher getParent(Publisher publisher) { + Object parent = Scannable.from(publisher).scanUnsafe(Scannable.Attr.PARENT); + if (parent instanceof Publisher) { + return (Publisher) parent; + } + return null; + } + + private static boolean isReactorCorePublisher(String pubClassName) { + return pubClassName != null && pubClassName.startsWith("reactor.core.publisher"); + } + + private static boolean isLifter(Publisher current) { + if (current == null) { + return false; + } + String name = current.getClass().getName(); + return isReactorCorePublisher(name) && ((name.endsWith("Lift") || name.endsWith("LiftFuseable"))); + } + + /** + * 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}. + */ + 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; + } + + 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 a3b0987a1..d9b72fbab 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 @@ -16,7 +16,9 @@ package org.springframework.cloud.sleuth.instrument.reactor; +import java.util.function.BiFunction; import java.util.function.Function; +import java.util.function.Predicate; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -26,7 +28,9 @@ import org.reactivestreams.Subscription; import reactor.core.CoreSubscriber; import reactor.core.Fuseable; import reactor.core.Scannable; +import reactor.core.publisher.Hooks; import reactor.core.publisher.Operators; +import reactor.util.annotation.Nullable; import reactor.util.context.Context; import org.springframework.cloud.sleuth.CurrentTraceContext; @@ -78,13 +82,38 @@ public abstract class ReactorSleuth { LazyBean lazyCurrentTraceContext = LazyBean.create(springContext, CurrentTraceContext.class); - return Operators.liftPublisher((p, sub) -> { - // We don't scope scalar results as they happen in an instant. This prevents - // excessive overhead when using Flux/Mono #just, #empty, #error, etc. - if (p instanceof Fuseable.ScalarCallable) { - return sub; - } + return Operators.liftPublisher(p -> !(p instanceof Fuseable.ScalarCallable), + (BiFunction) liftFunction(springContext, lazyCurrentTraceContext)); + } + /** + * Creates scope passing span operator which applies only to not + * {@code Scannable.Attr.RunStyle.SYNC} {@code Publisher}s. Used by + * {@code InstrumentationType#DECORATE_ON_EACH} + * @param springContext the Spring context. + * @param an arbitrary type that is left unchanged by the span operator. + * @return operator to apply to {@link Hooks#onEachOperator(Function)}. + */ + public static Function, ? extends Publisher> onEachOperatorForOnEachInstrumentation( + ConfigurableApplicationContext springContext) { + if (log.isTraceEnabled()) { + log.trace("Scope passing operator [" + springContext + "]"); + } + + // keep a reference outside the lambda so that any caching will be visible to + // all publishers + LazyBean lazyCurrentTraceContext = LazyBean.create(springContext, + CurrentTraceContext.class); + + Predicate publisherPredicate = ReactorHooksHelper::shouldDecorate; + BiFunction lifter = liftFunction(springContext, lazyCurrentTraceContext); + + return ReactorHooksHelper.liftPublisher(publisherPredicate, lifter); + } + + static BiFunction, ? extends CoreSubscriber> liftFunction( + ConfigurableApplicationContext springContext, LazyBean lazyCurrentTraceContext) { + return (p, sub) -> { if (!springContext.isActive()) { if (log.isTraceEnabled()) { String message = "Spring Context [" + springContext @@ -114,7 +143,7 @@ public abstract class ReactorSleuth { return sub; } - context = contextWithBeans(context, springContext, sub); + context = contextWithBeans(context, springContext); if (log.isTraceEnabled()) { log.trace("Spring context [" + springContext + "], Reactor context [" + context + "], name [" + name(sub) + "]"); @@ -129,15 +158,12 @@ public abstract class ReactorSleuth { log.trace("Creating a scope passing span subscriber with Reactor Context " + "[" + context + "] and name [" + name(sub) + "]"); } - // if (runStyle == Scannable.Attr.RunStyle.SYNC) { - // return sub; - // } + return new ScopePassingSpanSubscriber<>(sub, context, currentTraceContext, parent); - }); + }; } - private static Context contextWithBeans(Context context, ConfigurableApplicationContext springContext, - CoreSubscriber sub) { + private static Context contextWithBeans(Context context, ConfigurableApplicationContext springContext) { if (!context.hasKey(Tracer.class)) { context = context.put(Tracer.class, springContext.getBean(Tracer.class)); } @@ -158,20 +184,57 @@ public abstract class ReactorSleuth { if (log.isTraceEnabled()) { log.trace("Spring Context passing operator [" + springContext + "]"); } - return Operators.liftPublisher((p, sub) -> { + return Operators.liftPublisher(p -> { // We don't scope scalar results as they happen in an instant. This prevents // excessive overhead when using Flux/Mono #just, #empty, #error, etc. - if (p instanceof Fuseable.ScalarCallable) { + return !(p instanceof Fuseable.ScalarCallable) && springContext.isActive(); + }, (p, sub) -> { + Context ctxBefore = context(sub); + Context context = contextWithBeans(ctxBefore, springContext); + if (context == ctxBefore) { return sub; } - if (!springContext.isActive()) { - return sub; - } - final Context context = contextWithBeans(context(sub), springContext, sub); return new SleuthContextOperator<>(context, sub); }); } + /** + * Creates tracing context capturing reactor operator. Used by + * {@code InstrumentationType#DECORATE_ON_EACH}. + * @param springContext the Spring context. + * @param an arbitrary type that is left unchanged by the span operator. + * @return operator to apply to {@link Hooks#onLastOperator(Function)} for + * {@code InstrumentationType#DECORATE_ON_EACH} + */ + public static Function, ? extends Publisher> onLastOperatorForOnEachInstrumentation( + ConfigurableApplicationContext springContext) { + LazyBean lazyCurrentTraceContext = LazyBean.create(springContext, + CurrentTraceContext.class); + + BiFunction, ? extends CoreSubscriber> scopePassingSpanSubscriber = liftFunction( + springContext, lazyCurrentTraceContext); + BiFunction, ? extends CoreSubscriber> skipIfNoTraceCtx = ( + pub, sub) -> { + // lazyCurrentTraceContext.get() is not null here. see predicate bellow + TraceContext traceContext = lazyCurrentTraceContext.get().context(); + if (context(sub).getOrDefault(TraceContext.class, null) == traceContext) { + return sub; + } + return scopePassingSpanSubscriber.apply(pub, sub); + }; + + return ReactorHooksHelper.liftPublisher(p -> { + boolean addContext = !(p instanceof Fuseable.ScalarCallable) && springContext.isActive(); + if (addContext) { + CurrentTraceContext currentTraceContext = lazyCurrentTraceContext.get(); + if (currentTraceContext != null) { + addContext = currentTraceContext.context() != null; + } + } + return addContext; + }, skipIfNoTraceCtx); + } + private static Context context(CoreSubscriber sub) { try { return sub.currentContext(); @@ -199,9 +262,30 @@ public abstract class ReactorSleuth { return fallback.context(); } + public static Function scopePassingOnScheduleHook( + ConfigurableApplicationContext springContext) { + LazyBean lazyCurrentTraceContext = LazyBean.create(springContext, + CurrentTraceContext.class); + return delegate -> { + if (springContext.isActive()) { + final CurrentTraceContext currentTraceContext = lazyCurrentTraceContext.get(); + if (currentTraceContext == null) { + return delegate; + } + final TraceContext traceContext = currentTraceContext.context(); + return () -> { + try (CurrentTraceContext.Scope scope = currentTraceContext.maybeScope(traceContext)) { + delegate.run(); + } + }; + } + return delegate; + }; + } + } -class SleuthContextOperator implements Subscription, CoreSubscriber { +class SleuthContextOperator implements Subscription, CoreSubscriber, Scannable { private final Context context; @@ -250,4 +334,13 @@ class SleuthContextOperator implements Subscription, CoreSubscriber { return this.context; } + @Nullable + @Override + public Object scanUnsafe(Attr key) { + if (key == Attr.RUN_STYLE) { + return Attr.RunStyle.SYNC; + } + return null; + } + } diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java index a0d8bcf61..64a6e19f5 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java @@ -44,7 +44,7 @@ final class ScopePassingSpanSubscriber implements SpanSubscription, Scanna private final CurrentTraceContext currentTraceContext; - private final TraceContext parent; + final TraceContext parent; private Subscription s; @@ -108,11 +108,15 @@ final class ScopePassingSpanSubscriber implements SpanSubscription, Scanna return this.context; } + @reactor.util.annotation.Nullable @Override public Object scanUnsafe(Attr key) { if (key == Attr.PARENT) { return this.s; } + else if (key == Attr.RUN_STYLE) { + return Attr.RunStyle.SYNC; + } else { return key == Attr.ACTUAL ? this.subscriber : null; } diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceContextPropagator.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceContextPropagator.java new file mode 100644 index 000000000..8afbd8b9e --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceContextPropagator.java @@ -0,0 +1,25 @@ +/* + * Copyright 2013-2020 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.reactor; + +/** + * Only for internal use. Marker interface for Publisher and Subscriber that propagate + * Tracing context into execution Thread. + */ +public interface TraceContextPropagator { + +} diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java index 434f4aa54..4bcf69080 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java @@ -28,12 +28,14 @@ import reactor.core.publisher.MonoOperator; import reactor.util.annotation.Nullable; import reactor.util.context.Context; +import org.springframework.cloud.sleuth.CurrentTraceContext; import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.TraceContext; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.http.HttpServerHandler; import org.springframework.cloud.sleuth.http.HttpServerRequest; import org.springframework.cloud.sleuth.http.HttpServerResponse; +import org.springframework.cloud.sleuth.instrument.reactor.TraceContextPropagator; import org.springframework.core.Ordered; import org.springframework.http.HttpStatus; import org.springframework.http.server.reactive.ServerHttpRequest; @@ -69,11 +71,14 @@ public class TraceWebFilter implements WebFilter, Ordered { private final HttpServerHandler handler; + private final CurrentTraceContext currentTraceContext; + private int order; - public TraceWebFilter(Tracer tracer, HttpServerHandler handler) { + public TraceWebFilter(Tracer tracer, HttpServerHandler handler, CurrentTraceContext currentTraceContext) { this.tracer = tracer; this.handler = handler; + this.currentTraceContext = currentTraceContext; } @Override @@ -105,7 +110,7 @@ public class TraceWebFilter implements WebFilter, Ordered { this.order = order; } - private static class MonoWebFilterTrace extends MonoOperator { + private static class MonoWebFilterTrace extends MonoOperator implements TraceContextPropagator { final ServerWebExchange exchange; @@ -119,11 +124,14 @@ public class TraceWebFilter implements WebFilter, Ordered { final boolean initialTracePresent; + final CurrentTraceContext currentTraceContext; + MonoWebFilterTrace(Mono source, ServerWebExchange exchange, boolean initialTracePresent, TraceWebFilter parent) { super(source); this.tracer = parent.tracer; this.handler = parent.handler; + this.currentTraceContext = parent.currentTraceContext; this.exchange = exchange; this.span = exchange.getAttribute(TRACE_REQUEST_ATTR); this.initialTracePresent = initialTracePresent; @@ -132,7 +140,18 @@ public class TraceWebFilter implements WebFilter, Ordered { @Override public void subscribe(CoreSubscriber subscriber) { Context context = contextWithoutInitialSpan(subscriber.currentContext()); - this.source.subscribe(new WebFilterTraceSubscriber(subscriber, context, findOrCreateSpan(context), this)); + Span span = findOrCreateSpan(context); + try (CurrentTraceContext.Scope scope = this.currentTraceContext.maybeScope(span.context())) { + this.source.subscribe(new WebFilterTraceSubscriber(subscriber, context, span, this)); + } + } + + @Override + public Object scanUnsafe(Attr key) { + if (key == Attr.RUN_STYLE) { + return Attr.RunStyle.SYNC; + } + return super.scanUnsafe(key); } private Context contextWithoutInitialSpan(Context context) { diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java index f12eafc2a..8a858586f 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java @@ -26,6 +26,7 @@ import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.reactivestreams.Subscription; import reactor.core.CoreSubscriber; +import reactor.core.Scannable; import reactor.core.publisher.Mono; import reactor.util.annotation.Nullable; import reactor.util.context.Context; @@ -37,6 +38,7 @@ import org.springframework.cloud.sleuth.TraceContext; import org.springframework.cloud.sleuth.http.HttpClientHandler; import org.springframework.cloud.sleuth.http.HttpClientRequest; import org.springframework.cloud.sleuth.http.HttpClientResponse; +import org.springframework.cloud.sleuth.instrument.reactor.TraceContextPropagator; import org.springframework.context.ConfigurableApplicationContext; import org.springframework.web.reactive.function.client.ClientRequest; import org.springframework.web.reactive.function.client.ClientResponse; @@ -139,7 +141,8 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { return this.handler; } - private static final class MonoWebClientTrace extends Mono { + private static final class MonoWebClientTrace extends Mono + implements Scannable, TraceContextPropagator { final ExchangeFunction next; @@ -175,13 +178,22 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { .subscribe(new TraceWebClientSubscriber(subscriber, context, span, parent, this)); } + @Nullable + @Override + public Object scanUnsafe(Scannable.Attr key) { + if (key == Scannable.Attr.RUN_STYLE) { + return Scannable.Attr.RunStyle.SYNC; + } + return null; + } + } /** * Subscriber for WebClient. */ static final class TraceWebClientSubscriber extends AtomicReference - implements CoreSubscriber { + implements CoreSubscriber, Scannable { final CoreSubscriber actual; @@ -277,6 +289,14 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { return this.context; } + @Override + public Object scanUnsafe(Scannable.Attr key) { + if (key == Scannable.Attr.RUN_STYLE) { + return Scannable.Attr.RunStyle.SYNC; + } + return null; + } + } static class TraceWebClientSubscription implements Subscription { 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 new file mode 100644 index 000000000..82d7bd25f --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelperTests.java @@ -0,0 +1,367 @@ +/* + * Copyright 2013-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +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; + +class ReactorHooksHelperTests { + + @BeforeEach + public void setUp() { + Hooks.resetOnEachOperator(); + Hooks.resetOnLastOperator(); + Schedulers.resetOnScheduleHooks(); + } + + @Test + public void shouldDecorateWhenSyncSourceThanShouldNotDecorate() { + boolean actual = ReactorHooksHelper.shouldDecorate(Mono.just(1).map(Function.identity())); + assertThat(actual).isFalse(); + } + + @Test + public void shouldDecorateWhenSyncSourceProducerThanShouldNotDecorate() { + // Mono.fromSupplier() is instance of SourceProduced and is not + // Fuseable.ScalarCallable like Mono.just() + Mono syncSource = Mono.fromSupplier(() -> 1); + boolean actual = ReactorHooksHelper.shouldDecorate(syncSource.map(Function.identity())); + assertThat(actual).isFalse(); + } + + @Test + public void shouldDecorateWhenNotSyncSourceProducerThanShouldDecorate() { + Mono asyncSource = Mono.delay(Duration.ofMillis(10)); + boolean actual = ReactorHooksHelper.shouldDecorate(asyncSource.map(Function.identity())); + assertThat(actual).isTrue(); + } + + @Test + public void shouldDecorateWhenProcessorSourceThanShouldNotDecorate() { + Mono asyncSource = Sinks.one().asMono(); + boolean actual = ReactorHooksHelper.shouldDecorate(asyncSource.map(Function.identity())); + assertThat(actual).isTrue(); + } + + @Test + public void shouldDecorateWhenAsyncSourceAndLifterThanShouldDecorate() { + Mono asyncSourceWithLifter = Mono.delay(Duration.ofMillis(10)).transform(Operators.liftPublisher( + new BiFunction, CoreSubscriber>() { + @Override + public CoreSubscriber apply(Publisher pub, CoreSubscriber sub) { + return sub; + } + })); + boolean actual = ReactorHooksHelper.shouldDecorate(asyncSourceWithLifter.map(Function.identity())); + assertThat(actual).isTrue(); + } + + @Test + public void shouldDecorateWhenSyncSourceAndLifterThanShouldNotDecorate() { + Mono syncSourceWithLifter = Mono.fromSupplier(() -> 1L).transform(Operators.liftPublisher( + new BiFunction, CoreSubscriber>() { + @Override + public CoreSubscriber apply(Publisher pub, CoreSubscriber sub) { + return sub; + } + })); + boolean actual = ReactorHooksHelper.shouldDecorate(syncSourceWithLifter.map(Function.identity())); + assertThat(actual).isFalse(); + } + + @Test + public void shouldDecorateWhenAsyncSourceAndTraceContextPropagatorThanShouldNotDecorate() { + Function, ? extends Publisher> traceContextPropagator = ReactorHooksHelper + .liftPublisher(publisher -> true, (publisher, coreSubscriber) -> coreSubscriber); + Mono syncSourceWithLifter = Mono.delay(Duration.ofMillis(10)).transform(traceContextPropagator); + boolean actual = ReactorHooksHelper.shouldDecorate(syncSourceWithLifter.map(Function.identity())); + assertThat(actual).isFalse(); + } + + @Test + public void shouldDecorateWhenNullSourceProducerThanError() { + assertThatCode(() -> ReactorHooksHelper.shouldDecorate(null)).isInstanceOf(IllegalArgumentException.class) + .hasMessage("source Publisher is null"); + } + + @Test + public void shouldDecorateWhenOperatorWithoutRunStyleThanShouldDecorate() { + Mono source = Mono.just(1).as(CustomMonoWithoutRunStyleOperator::new); + boolean actual = ReactorHooksHelper.shouldDecorate(source.map(Function.identity())); + 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) { + super(source); + } + + @Override + public void subscribe(CoreSubscriber actual) { + source.subscribe(actual); + } + + @Nullable + @Override + public Object scanUnsafe(Attr key) { + if (Attr.RUN_STYLE == key) { + return null; + } + return super.scanUnsafe(key); + } + + } + + static class PlusOneSubscriber extends BaseSubscriber { + + CoreSubscriber actual; + + PlusOneSubscriber(CoreSubscriber actual) { + this.actual = actual; + } + + @Override + public Context currentContext() { + return actual.currentContext(); + } + + @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); + } + + } + +} 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 eb42b4786..23d17a1b5 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 @@ -38,7 +38,7 @@ 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 @@ -78,7 +78,7 @@ public abstract class FlowsScopePassingSpanSubscriberTests { springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext); springContext.refresh(); - Function, ? extends Publisher> transformer = scopePassingSpanOperator( + Function, ? extends Publisher> transformer = onEachOperatorForOnEachInstrumentation( this.springContext); try (CurrentTraceContext.Scope ws = currentTraceContext().newScope(context())) { diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberSpringBootTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberSpringBootTests.java index 0c67be78d..cf2d14c22 100644 --- a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberSpringBootTests.java +++ b/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 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 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 spanInOperation = new AtomicReference<>(); + + Sinks.One 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 fromMono = Mono.fromCallable(this.currentTraceContext::context); diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java index 75a29a486..7df4424ec 100644 --- a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java +++ b/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