diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/annotation/ReactorSleuthMethodInvocationProcessor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/annotation/ReactorSleuthMethodInvocationProcessor.java index 1c2fa7c9f..db52db558 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/annotation/ReactorSleuthMethodInvocationProcessor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/annotation/ReactorSleuthMethodInvocationProcessor.java @@ -17,17 +17,22 @@ package org.springframework.cloud.sleuth.annotation; import java.lang.reflect.Method; -import java.util.function.Consumer; -import java.util.function.Supplier; import brave.Span; import brave.Tracer; +import brave.Tracing; +import brave.propagation.CurrentTraceContext; import org.aopalliance.intercept.MethodInvocation; import org.reactivestreams.Publisher; -import org.springframework.util.StringUtils; +import org.reactivestreams.Subscription; +import reactor.core.CoreSubscriber; +import reactor.core.Scannable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import reactor.core.publisher.SignalType; +import reactor.util.annotation.Nullable; +import reactor.util.context.Context; + +import org.springframework.util.StringUtils; /** * Method Invocation Processor for Reactor. @@ -40,6 +45,15 @@ class ReactorSleuthMethodInvocationProcessor private NonReactorSleuthMethodInvocationProcessor nonReactorSleuthMethodInvocationProcessor; + Tracing tracing; + + Tracing tracing() { + if (this.tracing == null) { + this.tracing = this.beanFactory.getBean(Tracing.class); + } + return this.tracing; + } + @Override public Object process(MethodInvocation invocation, NewSpan newSpan, ContinueSpan continueSpan) throws Throwable { @@ -53,79 +67,271 @@ class ReactorSleuthMethodInvocationProcessor } } + @SuppressWarnings("unchecked") private Object proceedUnderReactorSpan(MethodInvocation invocation, NewSpan newSpan, ContinueSpan continueSpan) throws Throwable { Span spanPrevious = tracer().currentSpan(); // in case of @ContinueSpan and no span in tracer we start new span and should // close it on completion - boolean startNewSpan = newSpan != null || spanPrevious == null; Span span; - if (startNewSpan) { - span = tracer().nextSpan(); - newSpanParser().parse(invocation, newSpan, span); + if (newSpan != null || spanPrevious == null) { + span = null; } else { span = spanPrevious; } + String log = log(continueSpan); - boolean hasLog = StringUtils.hasText(log); - try (Tracer.SpanInScope ws = tracer().withSpanInScope(span)) { - Publisher publisher = (Publisher) invocation.proceed(); - Mono startSpan = Mono.defer(() -> withSpanInScope(span, () -> { - if (startNewSpan) { - span.start(); - } - before(invocation, span, log, hasLog); - return Mono.just(span); - })); - if (publisher instanceof Mono) { - return startSpan - .flatMap(spanStarted -> ((Mono) publisher) - .doOnError(onFailureReactor(log, hasLog, spanStarted)) - .doFinally(afterReactor(startNewSpan, log, hasLog, - spanStarted))) - // put span in context so it can be used by - // ScopePassingSpanSubscriber - .subscriberContext(context -> context.put(Span.class, span)); - } - else if (publisher instanceof Flux) { - return startSpan - .flatMapMany(spanStarted -> ((Flux) publisher) - .doOnError(onFailureReactor(log, hasLog, spanStarted)) - .doFinally(afterReactor(startNewSpan, log, hasLog, - spanStarted))) - // put span in context so it can be used by - // ScopePassingSpanSubscriber - .subscriberContext(context -> context.put(Span.class, span)); + Publisher publisher = (Publisher) invocation.proceed(); + + if (publisher instanceof Mono) { + return new MonoSpan((Mono) publisher, + this, + newSpan, + span, + invocation, + log); + } + else if (publisher instanceof Flux) { + return new FluxSpan((Flux) publisher, + this, + newSpan, + span, + invocation, + log); + } + else { + throw new IllegalArgumentException("Unexpected type of publisher: " + publisher.getClass()); + } + } + + private static final class FluxSpan extends Flux implements Scannable { + + final Flux source; + final Span span; + final MethodInvocation invocation; + final String log; + final boolean hasLog; + final ReactorSleuthMethodInvocationProcessor processor; + final NewSpan newSpan; + + FluxSpan(Flux source, + ReactorSleuthMethodInvocationProcessor processor, + NewSpan newSpan, + @Nullable Span span, + MethodInvocation invocation, + String log) { + this.source = source; + this.span = span; + this.newSpan = newSpan; + this.invocation = invocation; + this.log = log; + this.hasLog = StringUtils.hasText(log); + this.processor = processor; + } + + @Override + public void subscribe(CoreSubscriber actual) { + Span span; + Tracer tracer = this.processor.tracer(); + if (this.span == null) { + span = tracer.nextSpan(); + this.processor.newSpanParser().parse(invocation, newSpan, span); + span.start(); } else { - throw new IllegalArgumentException( - "Unexpected type of publisher: " + publisher.getClass()); + span = this.span; + } + try (Tracer.SpanInScope ws = tracer.withSpanInScope(span)) { + this.source.subscribe(new SpanSubscriber(actual, + this.processor, + this.invocation, + this.span == null, + span, + this.log, + this.hasLog)); + } + } + + @Nullable + public Object scanUnsafe(Attr key) { + if (key == Attr.PARENT) { + return this.source; + } + else { + return null; } } } - private T withSpanInScope(Span span, Supplier supplier) { - try (Tracer.SpanInScope ws1 = tracer().withSpanInScope(span)) { - return supplier.get(); + private static final class MonoSpan extends Mono implements Scannable { + + final Mono source; + final Span span; + final MethodInvocation invocation; + final String log; + final boolean hasLog; + final ReactorSleuthMethodInvocationProcessor processor; + final NewSpan newSpan; + + MonoSpan(Mono source, + ReactorSleuthMethodInvocationProcessor processor, + NewSpan newSpan, + @Nullable Span span, + MethodInvocation invocation, + String log) { + this.source = source; + this.processor = processor; + this.newSpan = newSpan; + this.span = span; + this.invocation = invocation; + this.log = log; + this.hasLog = StringUtils.hasText(log); + } + + @Override + public void subscribe(CoreSubscriber actual) { + Span span; + Tracer tracer = this.processor.tracer(); + if (this.span == null) { + span = tracer.nextSpan(); + this.processor.newSpanParser().parse(invocation, newSpan, span); + span.start(); + } + else { + span = this.span; + } + try (Tracer.SpanInScope ws = tracer.withSpanInScope(span)) { + this.source.subscribe(new SpanSubscriber(actual, + this.processor, + this.invocation, + this.span == null, + span, + this.log, + this.hasLog)); + } + } + + @Nullable + public Object scanUnsafe(Attr key) { + if (key == Attr.PARENT) { + return this.source; + } + else { + return null; + } } } - private Consumer afterReactor(boolean isNewSpan, String log, - boolean hasLog, Span span) { - return signalType -> { - try (Tracer.SpanInScope ws = tracer().withSpanInScope(span)) { - after(span, isNewSpan, log, hasLog); - } - }; - } + private static final class SpanSubscriber implements CoreSubscriber, + Subscription, + Scannable { - private Consumer onFailureReactor(String log, boolean hasLog, Span span) { - return throwable -> { - try (Tracer.SpanInScope ws = tracer().withSpanInScope(span)) { - onFailure(span, log, hasLog, throwable); + final CoreSubscriber actual; + final boolean isNewSpan; + final Span span; + final String log; + final boolean hasLog; + final CurrentTraceContext currentTraceContext; + final ReactorSleuthMethodInvocationProcessor processor; + final Context context; + + Subscription parent; + + SpanSubscriber(CoreSubscriber actual, + ReactorSleuthMethodInvocationProcessor processor, + MethodInvocation invocation, + boolean isNewSpan, + Span span, + String log, + boolean hasLog) { + this.actual = actual; + this.isNewSpan = isNewSpan; + this.span = span; + this.log = log; + this.hasLog = hasLog; + this.processor = processor; + + this.currentTraceContext = processor.tracing().currentTraceContext(); + this.context = actual.currentContext().put(Span.class, span); + + processor.before(invocation, this.span, this.log, this.hasLog); + } + + @Override + public void request(long n) { + try (CurrentTraceContext.Scope scope = this.currentTraceContext + .maybeScope(this.span.context())) { + this.parent.request(n); } - }; + } + + @Override + public void cancel() { + try (CurrentTraceContext.Scope scope = this.currentTraceContext + .maybeScope(this.span.context())) { + this.parent.cancel(); + } + finally { + this.processor.after(this.span, this.isNewSpan, this.log, this.hasLog); + } + } + + @Override + public Context currentContext() { + return context; + } + + @Override + public void onSubscribe(Subscription subscription) { + this.parent = subscription; + try (CurrentTraceContext.Scope scope = this.currentTraceContext + .maybeScope(this.span.context())) { + this.actual.onSubscribe(this); + } + } + + @Override + public void onNext(Object o) { + try (CurrentTraceContext.Scope scope = this.currentTraceContext + .maybeScope(this.span.context())) { + this.actual.onNext(o); + } + } + + @Override + public void onError(Throwable error) { + try (CurrentTraceContext.Scope scope = this.currentTraceContext + .maybeScope(this.span.context())) { + this.processor.onFailure(this.span, this.log, this.hasLog, error); + this.actual.onError(error); + } + finally { + this.processor.after(this.span, this.isNewSpan, this.log, this.hasLog); + } + } + + @Override + public void onComplete() { + try (CurrentTraceContext.Scope scope = this.currentTraceContext + .maybeScope(this.span.context())) { + this.actual.onComplete(); + } + finally { + this.processor.after(this.span, this.isNewSpan, this.log, this.hasLog); + } + } + + @Override + public Object scanUnsafe(Attr key) { + if (key == Attr.ACTUAL) { + return this.actual; + } + if (key == Attr.PARENT) { + return this.parent; + } + return null; + } } private boolean isReactorReturnType(Class returnType) { diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java index bde24fc1c..1be32c46e 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java @@ -16,21 +16,20 @@ package org.springframework.cloud.sleuth.instrument.reactor; +import java.util.function.BooleanSupplier; import java.util.function.Function; +import brave.Span; import brave.Tracing; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.reactivestreams.Publisher; import org.springframework.beans.factory.BeanFactory; -import org.springframework.beans.factory.NoSuchBeanDefinitionException; -import org.springframework.context.ApplicationContext; import org.springframework.context.ConfigurableApplicationContext; import reactor.core.CoreSubscriber; +import reactor.core.Fuseable; import reactor.core.Scannable; -import reactor.core.publisher.ConnectableFlux; -import reactor.core.publisher.GroupedFlux; import reactor.core.publisher.Operators; import reactor.util.context.Context; @@ -51,7 +50,7 @@ public abstract class ReactorSleuth { * Return a span operator pointcut given a {@link Tracing}. This can be used in * reactor via {@link reactor.core.publisher.Flux#transform(Function)}, * {@link reactor.core.publisher.Mono#transform(Function)}, - * {@link reactor.core.publisher.Hooks#onEachOperator(Function)} or + * {@link reactor.core.publisher.Hooks#onLastOperator(Function)} or * {@link reactor.core.publisher.Hooks#onLastOperator(Function)}. The Span operator * pointcut will pass the Scope of the Span without ever creating any new spans. * @param beanFactory - {@link BeanFactory} @@ -60,58 +59,63 @@ public abstract class ReactorSleuth { */ @SuppressWarnings("unchecked") public static Function, ? extends Publisher> scopePassingSpanOperator( - ConfigurableApplicationContext beanFactory) { + BeanFactory beanFactory) { if (log.isTraceEnabled()) { log.trace("Scope passing operator [" + beanFactory + "]"); } - return (sourcePub -> { - // TODO: Remove this once Reactor 3.1.8 is released - // do the checks directly on actual original Publisher - if (sourcePub instanceof ConnectableFlux // Operators.lift can't handle that - || sourcePub instanceof GroupedFlux // Operators.lift can't handle - // that - ) { - return sourcePub; - } - // no more POINTCUT_FILTER since mechanism is broken - Function, ? extends Publisher> lift = Operators - .lift((scannable, sub) -> { - // rest of the logic unchanged... - if (beanFactory.isActive()) { - if (log.isTraceEnabled()) { - log.trace("Spring Context [" + beanFactory - + "] already refreshed. Creating a scope " - + "passing span subscriber with Reactor Context " - + "[" + sub.currentContext() + "] and name [" - + scannable.name() + "]"); - } - return scopePassingSpanSubscription(beanFactory, scannable, - sub).get(); - } - if (log.isTraceEnabled()) { - log.trace("Spring Context [" + beanFactory - + "] is not yet refreshed, falling back to lazy span subscriber. Reactor Context is [" - + sub.currentContext() + "] and name is [" - + scannable.name() + "]"); - } - return new LazySpanSubscriber(scopePassingSpanSubscription( - beanFactory, scannable, sub)); - }); - return lift.apply(sourcePub); + //Adapt if lazy bean factory + BooleanSupplier isActive = + beanFactory instanceof ConfigurableApplicationContext ? + ((ConfigurableApplicationContext) beanFactory)::isActive : + () -> true; + + return Operators.liftPublisher((p, sub) -> { + //if Flux/Mono #just, #empty, #error + if (p instanceof Fuseable.ScalarCallable) { + return sub; + } + Scannable scannable = Scannable.from(p); + // rest of the logic unchanged... + if (isActive.getAsBoolean()) { + if (log.isTraceEnabled()) { + log.trace("Spring Context [" + beanFactory + + "] already refreshed. Creating a scope " + + "passing span subscriber with Reactor Context " + + "[" + sub.currentContext() + "] and name [" + + scannable.name() + "]"); + } + + return scopePassingSpanSubscription(beanFactory.getBean(Tracing.class), sub); + } + if (log.isTraceEnabled()) { + log.trace("Spring Context [" + beanFactory + + "] is not yet refreshed, falling back to lazy span subscriber. Reactor Context is [" + + sub.currentContext() + "] and name is [" + + scannable.name() + "]"); + } + return new LazySpanSubscriber<>(lazyScopePassingSpanSubscription(beanFactory, scannable, sub)); }); } - private static SpanSubscriptionProvider scopePassingSpanSubscription( + static SpanSubscriptionProvider lazyScopePassingSpanSubscription( BeanFactory beanFactory, Scannable scannable, CoreSubscriber sub) { - return new SpanSubscriptionProvider(beanFactory, sub, sub.currentContext(), - scannable.name()) { - @Override - SpanSubscription newCoreSubscriber(Tracing tracing) { - return new ScopePassingSpanSubscriber(sub, - sub != null ? sub.currentContext() : Context.empty(), tracing); - } - }; + return new SpanSubscriptionProvider<>(beanFactory, sub, sub.currentContext(), scannable.name()); + } + + + static CoreSubscriber scopePassingSpanSubscription( + Tracing tracing, CoreSubscriber sub) { + + Context context = sub.currentContext(); + + Span root = context.hasKey(Span.class) ? context.get(Span.class) : tracing.tracer().currentSpan(); + if (root != null) { + return new ScopePassingSpanSubscriber<>(sub, context, tracing, root); + } + else { + return sub; //no need to trace + } } } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java index 96cc954c0..8925cecf1 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java @@ -16,6 +16,8 @@ package org.springframework.cloud.sleuth.instrument.reactor; +import javax.annotation.Nullable; + import brave.Span; import brave.Tracer; import brave.Tracing; @@ -25,6 +27,7 @@ import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.reactivestreams.Subscriber; import org.reactivestreams.Subscription; +import reactor.core.Scannable; import reactor.util.context.Context; /** @@ -34,7 +37,7 @@ import reactor.util.context.Context; * @author Marcin Grzejszczak * @since 2.0.0 */ -final class ScopePassingSpanSubscriber implements SpanSubscription { +final class ScopePassingSpanSubscriber implements SpanSubscription, Scannable { private static final Log log = LogFactory.getLog(ScopePassingSpanSubscriber.class); @@ -46,18 +49,13 @@ final class ScopePassingSpanSubscriber implements SpanSubscription { private final TraceContext traceContext; - private final Tracer tracer; - private Subscription s; ScopePassingSpanSubscriber(Subscriber subscriber, Context ctx, - Tracing tracing) { + Tracing tracing, @Nullable Span root) { this.subscriber = subscriber; - this.tracer = tracing.tracer(); this.currentTraceContext = tracing.currentTraceContext(); - Span root = ctx != null - ? ctx.hasKey(Span.class) ? ctx.get(Span.class) : this.tracer.currentSpan() - : null; + this.traceContext = root == null ? null : root.context(); this.context = ctx != null && root != null ? ctx.put(Span.class, root) : ctx != null ? ctx : Context.empty(); @@ -121,4 +119,12 @@ final class ScopePassingSpanSubscriber implements SpanSubscription { return this.context; } + @Override + public Object scanUnsafe(Attr key) { + if (key == Attr.PARENT) { + return this.s; + } else { + return key == Attr.ACTUAL ? this.subscriber : null; + } + } } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriptionProvider.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriptionProvider.java index 917717109..be85c8b59 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriptionProvider.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriptionProvider.java @@ -18,6 +18,7 @@ package org.springframework.cloud.sleuth.instrument.reactor; import java.util.function.Supplier; +import brave.Span; import brave.Tracing; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -31,7 +32,7 @@ import reactor.util.context.Context; * @param type of returned subscription * @author Marcin Grzejszczak */ -class SpanSubscriptionProvider implements Supplier> { +final class SpanSubscriptionProvider implements Supplier> { private static final Log log = LogFactory.getLog(SpanSubscriptionProvider.class); @@ -63,7 +64,8 @@ class SpanSubscriptionProvider implements Supplier> { } SpanSubscription newCoreSubscriber(Tracing tracing) { - return new ScopePassingSpanSubscriber<>(this.subscriber, this.context, tracing); + Span root = context.hasKey(Span.class) ? context.get(Span.class) : tracing.tracer().currentSpan(); + return new ScopePassingSpanSubscriber<>(this.subscriber, this.context, tracing, root); } private Tracing tracing() { diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java index 24c3cf540..7b0a1b23a 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java @@ -85,7 +85,7 @@ public class TraceReactorAutoConfiguration { if (log.isTraceEnabled()) { log.trace("Cleaning up hooks"); } - Hooks.resetOnEachOperator(SLEUTH_TRACE_REACTOR_KEY); + Hooks.resetOnLastOperator(SLEUTH_TRACE_REACTOR_KEY); Schedulers.resetFactory(); } @@ -115,7 +115,7 @@ class HookRegisteringBeanDefinitionRegistryPostProcessor } void setupHooks(BeanFactory beanFactory) { - Hooks.onEachOperator( + Hooks.onLastOperator( TraceReactorAutoConfiguration.TraceReactorConfiguration.SLEUTH_TRACE_REACTOR_KEY, ReactorSleuth.scopePassingSpanOperator(this.context)); Schedulers.setFactory(factoryInstance(beanFactory)); diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectFluxTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectFluxTests.java index 0a14055bf..0c5801dfb 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectFluxTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectFluxTests.java @@ -590,8 +590,7 @@ public class SleuthSpanCreatorAspectFluxTests { @Override public Flux newSpanInTraceContext() { - Long id = id(tracer); - return Flux.defer(() -> Flux.just(id)); + return Flux.defer(() -> Flux.just(id(tracer))); } @Override diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java index 0e6016733..03c867e79 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java @@ -22,6 +22,8 @@ import brave.Tracing; import org.junit.Test; import org.junit.runner.RunWith; import org.mockito.junit.MockitoJUnitRunner; +import reactor.core.CoreSubscriber; +import reactor.core.publisher.BaseSubscriber; import reactor.util.context.Context; import static org.assertj.core.api.BDDAssertions.then; @@ -36,16 +38,17 @@ public class ScopePassingSpanSubscriberTests { @Test public void should_propagate_current_context() { - ScopePassingSpanSubscriber subscriber = new ScopePassingSpanSubscriber(null, - Context.of("foo", "bar"), this.tracing); + ScopePassingSpanSubscriber subscriber = new ScopePassingSpanSubscriber<>(null, + Context.of("foo", "bar"), this.tracing, null); then((String) subscriber.currentContext().get("foo")).isEqualTo("bar"); } @Test public void should_set_empty_context_when_context_is_null() { - ScopePassingSpanSubscriber subscriber = new ScopePassingSpanSubscriber(null, null, - this.tracing); + ScopePassingSpanSubscriber subscriber = new ScopePassingSpanSubscriber<>(null + , null, + this.tracing, null); then(subscriber.currentContext().isEmpty()).isTrue(); } @@ -55,8 +58,9 @@ public class ScopePassingSpanSubscriberTests { Span span = this.tracing.tracer().nextSpan(); try (Tracer.SpanInScope ws = this.tracing.tracer() .withSpanInScope(span.start())) { - ScopePassingSpanSubscriber subscriber = new ScopePassingSpanSubscriber(null, - Context.empty(), this.tracing); + CoreSubscriber subscriber = + ReactorSleuth.scopePassingSpanSubscription(tracing, new BaseSubscriber() { + }); then(subscriber.currentContext().get(Span.class)).isEqualTo(span); } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriberTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriberTests.java index 25eb98ac4..205884669 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriberTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriberTests.java @@ -17,31 +17,34 @@ package org.springframework.cloud.sleuth.instrument.reactor; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; import brave.Span; import brave.Tracer; import brave.sampler.Sampler; -import reactor.core.publisher.BaseSubscriber; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Hooks; -import reactor.core.publisher.Mono; -import reactor.core.publisher.Operators; -import reactor.core.scheduler.Schedulers; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.awaitility.Awaitility; -import org.junit.AfterClass; import org.junit.Test; import org.junit.runner.RunWith; import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; import org.reactivestreams.Subscription; +import reactor.core.CoreSubscriber; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Hooks; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.context.ConfigurableApplicationContext; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.test.context.junit4.SpringRunner; +import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.BDDAssertions.then; @RunWith(SpringRunner.class) @@ -53,6 +56,9 @@ public class SpanSubscriberTests { @Autowired Tracer tracer; + @Autowired + ConfigurableApplicationContext factory; + @Test public void should_pass_tracing_info_when_using_reactor() { Span span = this.tracer.nextSpan().name("foo").start(); @@ -82,10 +88,10 @@ public class SpanSubscriberTests { try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) { Mono.just(1).flatMap(d -> Flux.just(d + 1).collectList().map(p -> p.get(0))) - .map(d -> d + 1).map((d) -> { - spanInOperation.set(this.tracer.currentSpan()); - return d + 1; - }).map(d -> d + 1).subscribe(System.out::println); + .map(d -> d + 1).map((d) -> { + spanInOperation.set(this.tracer.currentSpan()); + return d + 1; + }).map(d -> d + 1).subscribe(System.out::println); } finally { span.finish(); @@ -98,47 +104,90 @@ public class SpanSubscriberTests { @Test public void should_not_trace_scalar_flows() { Span span = this.tracer.nextSpan().name("foo").start(); - final AtomicReference spanInOperation = new AtomicReference<>(); log.info("Hello"); + //Disable global hooks for local hook testing + Hooks.resetOnLastOperator(); + try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) { - Mono.just(1).subscribe(new BaseSubscriber() { - @Override - protected void hookOnSubscribe(Subscription subscription) { - spanInOperation.set(subscription); - } - }); - then(this.tracer.currentSpan()).isNotNull(); + Function, ? extends Publisher> transformer = + ReactorSleuth.scopePassingSpanOperator(factory); - Mono.error(new Exception()).subscribe(new BaseSubscriber() { + Subscriber assertNoSpanSubscriber = new CoreSubscriber() { @Override - protected void hookOnSubscribe(Subscription subscription) { - spanInOperation.set(subscription); + public void onSubscribe(Subscription s) { + s.request(Long.MAX_VALUE); + assertThat(s).isNotInstanceOf(ScopePassingSpanSubscriber.class); } @Override - protected void hookOnError(Throwable throwable) { + public void onNext(Object o) { + } - }); - then(this.tracer.currentSpan()).isNotNull(); - - Mono.empty().subscribe(new BaseSubscriber() { @Override - protected void hookOnSubscribe(Subscription subscription) { - spanInOperation.set(subscription); - } - }); + public void onError(Throwable t) { + + } + + @Override + public void onComplete() { + + } + }; + + Subscriber assertSpanSubscriber = new CoreSubscriber() { + @Override + public void onSubscribe(Subscription s) { + s.request(Long.MAX_VALUE); + assertThat(s).isInstanceOf(ScopePassingSpanSubscriber.class); + } + + @Override + public void onNext(Object o) { + + } + + @Override + public void onError(Throwable t) { + + } + + @Override + public void onComplete() { + + } + }; + transformer.apply(Mono.just(1).hide()) + .subscribe(assertSpanSubscriber); + + transformer.apply(Mono.just(1)) + .subscribe(assertNoSpanSubscriber); + + transformer.apply(Mono.error(new Exception()).hide()) + .subscribe(assertSpanSubscriber); + + + transformer.apply(Mono.error(new Exception())) + .subscribe(assertNoSpanSubscriber); + + transformer.apply(Mono.empty().hide()) + .subscribe(assertSpanSubscriber); + + transformer.apply(Mono.empty()) + .subscribe(assertNoSpanSubscriber); + + - then(this.tracer.currentSpan()).isNotNull(); - then(spanInOperation.get()).isEqualTo(Operators.emptySubscription()); } finally { span.finish(); } - then(this.tracer.currentSpan()).isNull(); + Awaitility.await().untilAsserted(() -> { + then(this.tracer.currentSpan()).isNull(); + }); } @Test @@ -149,16 +198,16 @@ public class SpanSubscriberTests { try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) { Flux.just(1, 2, 3).publishOn(Schedulers.single()).log("reactor.1") - .map(d -> d + 1).map(d -> d + 1) - .publishOn(Schedulers.newSingle("secondThread")).log("reactor.2") - .map((d) -> { - spanInOperation.set(this.tracer.currentSpan()); - return d + 1; - }).map(d -> d + 1).blockLast(); + .map(d -> d + 1).map(d -> d + 1) + .publishOn(Schedulers.newSingle("secondThread")).log("reactor.2") + .map((d) -> { + spanInOperation.set(this.tracer.currentSpan()); + return d + 1; + }).map(d -> d + 1).blockLast(); Awaitility.await().untilAsserted(() -> { - then(spanInOperation.get().context().spanId()) - .isEqualTo(span.context().spanId()); + then(spanInOperation.get().context().traceId()) + .isEqualTo(span.context().traceId()); }); then(this.tracer.currentSpan()).isEqualTo(span); } @@ -171,15 +220,15 @@ public class SpanSubscriberTests { try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(foo2)) { Flux.just(1, 2, 3).publishOn(Schedulers.single()).log("reactor.") - .map(d -> d + 1).map(d -> d + 1).map((d) -> { - spanInOperation.set(this.tracer.currentSpan()); - return d + 1; - }).map(d -> d + 1).blockLast(); + .map(d -> d + 1).map(d -> d + 1).map((d) -> { + spanInOperation.set(this.tracer.currentSpan()); + return d + 1; + }).map(d -> d + 1).blockLast(); then(this.tracer.currentSpan()).isEqualTo(foo2); // parent cause there's an async span in the meantime - then(spanInOperation.get().context().spanId()) - .isEqualTo(foo2.context().spanId()); + then(spanInOperation.get().context().traceId()) + .isEqualTo(foo2.context().traceId()); } finally { foo2.finish(); @@ -194,11 +243,11 @@ public class SpanSubscriberTests { log.info("Hello"); try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(parentSpan)) { final Long spanId = Mono.fromCallable(tracer::currentSpan) - .map(span -> span.context().spanId()).block(); + .map(span -> span.context().spanId()).block(); then(spanId).isNotNull(); final Long secondSpanId = Mono.fromCallable(tracer::currentSpan) - .map(span -> span.context().spanId()).block(); + .map(span -> span.context().spanId()).block(); then(secondSpanId).isEqualTo(spanId); // different trace ids here } } @@ -211,19 +260,19 @@ public class SpanSubscriberTests { try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(initSpan)) { Mono.fromCallable(tracer::currentSpan).map(span -> span.context().spanId()) - .doOnNext(spanInOperation::set) - .zipWith(Mono.fromCallable(tracer::currentSpan) - .map(span -> span.context().spanId()) - .doOnNext(spanInZipOperation::set)) - .block(); + .doOnNext(spanInOperation::set) + .zipWith(Mono.fromCallable(tracer::currentSpan) + .map(span -> span.context().spanId()) + .doOnNext(spanInZipOperation::set)) + .block(); } then(spanInZipOperation).hasValue(initSpan.context().spanId()); // ok here then(spanInOperation).hasValue(initSpan.context().spanId()); // Expecting - // - // to have value: - // <1L> but did - // not. + // + // to have value: + // <1L> but did + // not. } // #646 @@ -234,8 +283,8 @@ public class SpanSubscriberTests { try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(initSpan)) { Mono.just("value1") - .flatMap(request -> Mono.just("value2").then(Mono.just("foo"))) - .map(a -> "qwe").block(); + .flatMap(request -> Mono.just("value2").then(Mono.just("foo"))) + .map(a -> "qwe").block(); } } @@ -247,19 +296,13 @@ public class SpanSubscriberTests { try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(initSpan)) { Mono.subscriberContext() - .map(context -> tracer.currentSpan().context().spanId()) - .doOnNext(spanInSubscriberContext::set).block(); + .map(context -> tracer.currentSpan().context().spanId()) + .doOnNext(spanInSubscriberContext::set).block(); } then(spanInSubscriberContext).hasValue(initSpan.context().spanId()); // ok here } - @AfterClass - public static void cleanup() { - Hooks.resetOnEachOperator(); - Schedulers.resetFactory(); - } - @EnableAutoConfiguration @Configuration static class Config { diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfigurationAccessorConfiguration.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfigurationAccessorConfiguration.java index 9e43884d4..fa66082e6 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfigurationAccessorConfiguration.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfigurationAccessorConfiguration.java @@ -18,7 +18,7 @@ public class TraceReactorAutoConfigurationAccessorConfiguration { log.trace("Cleaning up hooks"); } new TraceReactorAutoConfiguration.TraceReactorConfiguration().cleanupHooks(); - Hooks.resetOnEachOperator(); + Hooks.resetOnLastOperator(); Hooks.resetOnLastOperator(); Schedulers.resetFactory(); } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java index 6641be206..e1d556f59 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java @@ -544,17 +544,15 @@ class TestBean { @NewSpan(name = "newSpanInTraceContext") public Mono newSpanInTraceContext() { log.info("New Span in Trace Context"); - Long span = tracer.currentSpan().context().spanId(); - return Mono.defer(() -> Mono.just(span)); + return Mono.defer(() -> Mono.just(tracer.currentSpan().context().spanId())); } @NewSpan(name = "newSpanInSubscriberContext") public Mono newSpanInSubscriberContext() { log.info("New Span in Subscriber Context"); - Long span = tracer.currentSpan().context().spanId(); return Mono.subscriberContext() .doOnSuccess(context -> log.info("New Span in deferred Trace Context")) - .flatMap(context -> Mono.defer(() -> Mono.just(span))); + .flatMap(context -> Mono.defer(() -> Mono.just(tracer.currentSpan().context().spanId()))); } } \ No newline at end of file