From 615f70c03a9baf1eed3c102eb303fb0263fd9c14 Mon Sep 17 00:00:00 2001 From: Roman Matiushchenko Date: Sun, 24 Jan 2021 20:50:14 +0200 Subject: [PATCH] handle double wrapping --- .../cloud/sleuth/internal/LazyBeanTests.java | 11 +++ .../reactor/ReactorHooksHelper.java | 2 +- .../instrument/reactor/ReactorSleuth.java | 63 ++++++++++---- .../cloud/sleuth/internal/LazyBean.java | 21 +++-- .../FlowsScopePassingSpanSubscriberTests.java | 82 ++++++++++++++++++- 5 files changed, 153 insertions(+), 26 deletions(-) diff --git a/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/internal/LazyBeanTests.java b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/internal/LazyBeanTests.java index c39352cf2..ec32d8298 100644 --- a/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/internal/LazyBeanTests.java +++ b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/internal/LazyBeanTests.java @@ -20,10 +20,12 @@ import brave.propagation.CurrentTraceContext; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; +import org.springframework.beans.factory.NoSuchBeanDefinitionException; import org.springframework.context.annotation.AnnotationConfigApplicationContext; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; +import static org.assertj.core.api.Assertions.assertThatCode; import static org.assertj.core.api.BDDAssertions.then; public class LazyBeanTests { @@ -54,6 +56,15 @@ public class LazyBeanTests { then(provider.get()).isNull(); } + @Test + public void should_throw_error_when_no_basic_type() { + context.refresh(); + + LazyBean provider = LazyBean.create(context, CurrentTraceContext.class); + + assertThatCode(() -> provider.getOrError()).isInstanceOf(NoSuchBeanDefinitionException.class); + } + @Configuration(proxyBeanMethods = false) static class BasicConfig { diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelper.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelper.java index b50dc152b..5f7731a22 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelper.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorHooksHelper.java @@ -143,7 +143,7 @@ final class ReactorHooksHelper { } } - private static boolean isTraceContextPropagator(Publisher current) { + static boolean isTraceContextPropagator(Publisher current) { return current instanceof TraceContextPropagator; } 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 d9b72fbab..cf96f09a7 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 @@ -82,8 +82,10 @@ public abstract class ReactorSleuth { LazyBean lazyCurrentTraceContext = LazyBean.create(springContext, CurrentTraceContext.class); + LazyBean lazyTracer = LazyBean.create(springContext, Tracer.class); + return Operators.liftPublisher(p -> !(p instanceof Fuseable.ScalarCallable), - (BiFunction) liftFunction(springContext, lazyCurrentTraceContext)); + (BiFunction) liftFunction(springContext, lazyCurrentTraceContext, lazyTracer)); } /** @@ -105,14 +107,20 @@ public abstract class ReactorSleuth { LazyBean lazyCurrentTraceContext = LazyBean.create(springContext, CurrentTraceContext.class); - Predicate publisherPredicate = ReactorHooksHelper::shouldDecorate; - BiFunction lifter = liftFunction(springContext, lazyCurrentTraceContext); + LazyBean lazyTracer = LazyBean.create(springContext, Tracer.class); - return ReactorHooksHelper.liftPublisher(publisherPredicate, lifter); + @SuppressWarnings("rawtypes") + Predicate shouldDecorate = ReactorHooksHelper::shouldDecorate; + @SuppressWarnings("rawtypes") + BiFunction, ? extends CoreSubscriber> lifter = liftFunction( + springContext, lazyCurrentTraceContext, lazyTracer); + + return ReactorHooksHelper.liftPublisher(shouldDecorate, lifter); } static BiFunction, ? extends CoreSubscriber> liftFunction( - ConfigurableApplicationContext springContext, LazyBean lazyCurrentTraceContext) { + ConfigurableApplicationContext springContext, LazyBean lazyCurrentTraceContext, + LazyBean lazyTracer) { return (p, sub) -> { if (!springContext.isActive()) { if (log.isTraceEnabled()) { @@ -143,17 +151,25 @@ public abstract class ReactorSleuth { return sub; } - context = contextWithBeans(context, springContext); - if (log.isTraceEnabled()) { - log.trace("Spring context [" + springContext + "], Reactor context [" + context + "], name [" - + name(sub) + "]"); - } - TraceContext parent = traceContext(context, currentTraceContext); if (parent == null) { return sub; // no need to scope a null parent } + // Handle scenarios such as Mono.defer + if (sub instanceof ScopePassingSpanSubscriber) { + ScopePassingSpanSubscriber scopePassing = (ScopePassingSpanSubscriber) sub; + if (scopePassing.parent.equals(parent)) { + return sub; // don't double-wrap + } + } + + context = contextWithBeans(context, lazyTracer, lazyCurrentTraceContext); + if (log.isTraceEnabled()) { + log.trace("Spring context [" + springContext + "], Reactor context [" + context + "], name [" + + name(sub) + "]"); + } + if (log.isTraceEnabled()) { log.trace("Creating a scope passing span subscriber with Reactor Context " + "[" + context + "] and name [" + name(sub) + "]"); @@ -163,12 +179,13 @@ public abstract class ReactorSleuth { }; } - private static Context contextWithBeans(Context context, ConfigurableApplicationContext springContext) { + private static Context contextWithBeans(Context context, LazyBean tracer, + LazyBean currentTraceContext) { if (!context.hasKey(Tracer.class)) { - context = context.put(Tracer.class, springContext.getBean(Tracer.class)); + context = context.put(Tracer.class, tracer.getOrError()); } if (!context.hasKey(CurrentTraceContext.class)) { - context = context.put(CurrentTraceContext.class, springContext.getBean(CurrentTraceContext.class)); + context = context.put(CurrentTraceContext.class, currentTraceContext.getOrError()); } return context; } @@ -184,13 +201,18 @@ public abstract class ReactorSleuth { if (log.isTraceEnabled()) { log.trace("Spring Context passing operator [" + springContext + "]"); } + + LazyBean lazyTracer = LazyBean.create(springContext, Tracer.class); + LazyBean lazyCurrentTraceContext = LazyBean.create(springContext, + CurrentTraceContext.class); + 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. return !(p instanceof Fuseable.ScalarCallable) && springContext.isActive(); }, (p, sub) -> { Context ctxBefore = context(sub); - Context context = contextWithBeans(ctxBefore, springContext); + Context context = contextWithBeans(ctxBefore, lazyTracer, lazyCurrentTraceContext); if (context == ctxBefore) { return sub; } @@ -210,9 +232,11 @@ public abstract class ReactorSleuth { ConfigurableApplicationContext springContext) { LazyBean lazyCurrentTraceContext = LazyBean.create(springContext, CurrentTraceContext.class); + LazyBean lazyTracer = LazyBean.create(springContext, Tracer.class); BiFunction, ? extends CoreSubscriber> scopePassingSpanSubscriber = liftFunction( - springContext, lazyCurrentTraceContext); + springContext, lazyCurrentTraceContext, lazyTracer); + BiFunction, ? extends CoreSubscriber> skipIfNoTraceCtx = ( pub, sub) -> { // lazyCurrentTraceContext.get() is not null here. see predicate bellow @@ -224,6 +248,13 @@ public abstract class ReactorSleuth { }; return ReactorHooksHelper.liftPublisher(p -> { + /* + * this prevent double decoration when last operator in the chain is not SYNC + * like {@code Mono.fromSuppler(() -> ...).subscribeOn(Schedulers.parallel())} + */ + if (ReactorHooksHelper.isTraceContextPropagator(p)) { + return false; + } boolean addContext = !(p instanceof Fuseable.ScalarCallable) && springContext.isActive(); if (addContext) { CurrentTraceContext currentTraceContext = lazyCurrentTraceContext.get(); diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/internal/LazyBean.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/internal/LazyBean.java index 845f5a4f5..328ef61ed 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/internal/LazyBean.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/internal/LazyBean.java @@ -55,12 +55,8 @@ public final class LazyBean { */ @Nullable public T get() { - if (this.value != null) { - return this.value; - } - try { - this.value = springContext.getBean(requiredType); + return getOrError(); } catch (Exception ex) { if (log.isDebugEnabled()) { @@ -70,4 +66,19 @@ public final class LazyBean { return this.value; } + /** + * Attempts to provision from the underlying bean factory, if not already provisioned. + * @return the bean value. This variant does not catch exception. + */ + public T getOrError() { + T bean = this.value; + if (bean != null) { + return bean; + } + + bean = springContext.getBean(requiredType); + this.value = bean; + return bean; + } + } 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 23d17a1b5..ecca4b380 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 @@ -17,23 +17,33 @@ package org.springframework.cloud.sleuth.instrument.reactor; import java.util.Objects; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; +import java.util.function.Supplier; +import java.util.stream.Stream; import org.assertj.core.presentation.StandardRepresentation; import org.awaitility.Awaitility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.mockito.Mockito; import org.reactivestreams.Publisher; import org.reactivestreams.Subscriber; import org.reactivestreams.Subscription; import reactor.core.CoreSubscriber; +import reactor.core.Exceptions; +import reactor.core.Scannable; import reactor.core.publisher.Hooks; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; import org.springframework.cloud.sleuth.CurrentTraceContext; import org.springframework.cloud.sleuth.TraceContext; +import org.springframework.cloud.sleuth.Tracer; import org.springframework.context.annotation.AnnotationConfigApplicationContext; import static org.assertj.core.api.Assertions.assertThat; @@ -45,6 +55,8 @@ import static org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth. */ public abstract class FlowsScopePassingSpanSubscriberTests { + static final String HOOK_KEY = "org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration"; + static { // AssertJ will recognise QueueSubscription implements queue and try to invoke // iterator. That's not allowed, and will cause an exception @@ -61,16 +73,17 @@ public abstract class FlowsScopePassingSpanSubscriberTests { @BeforeEach public void setup() { - Hooks.resetOnEachOperator( - "org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration"); - Hooks.resetOnLastOperator( - "org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration"); + Hooks.resetOnEachOperator(); + Hooks.resetOnLastOperator(); Schedulers.resetOnScheduleHooks(); } @AfterEach public void close() { springContext.close(); + Hooks.resetOnEachOperator(); + Hooks.resetOnLastOperator(); + Schedulers.resetOnScheduleHooks(); } @Test @@ -144,4 +157,65 @@ public abstract class FlowsScopePassingSpanSubscriberTests { Awaitility.await().untilAsserted(() -> then(currentTraceContext().context()).isNull()); } + @ParameterizedTest + @MethodSource("should_not_double_wrap_async_publisher_Args") + public void should_not_double_wrap_async_publisher(String name, Supplier> sourceSupplier) { + springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext); + springContext.registerBean(Tracer.class, () -> Mockito.mock(Tracer.class)); + springContext.refresh(); + + Hooks.onEachOperator(HOOK_KEY, ReactorSleuth.onEachOperatorForOnEachInstrumentation(springContext)); + Hooks.onLastOperator(HOOK_KEY, ReactorSleuth.onLastOperatorForOnEachInstrumentation(springContext)); + + AtomicBoolean once = new AtomicBoolean(); + Hooks.onLastOperator("test", p -> { + // check only first onLast Hook + if (once.compareAndSet(false, true)) { + assertThat(p).isInstanceOf(TraceContextPropagator.class); + Object parent = Scannable.from(p).scanUnsafe(Scannable.Attr.PARENT); + assertThat(parent).isNotInstanceOf(TraceContextPropagator.class); + } + return p; + }); + try (CurrentTraceContext.Scope ws = currentTraceContext().newScope(context())) { + Mono source = sourceSupplier.get(); + + source.subscribe((Subscriber) new CoreSubscriber() { + @Override + public void onSubscribe(Subscription subscription) { + assertThat(subscription).isInstanceOf(ScopePassingSpanSubscriber.class); + ScopePassingSpanSubscriber spanSubscriber = (ScopePassingSpanSubscriber) subscription; + Object parent = spanSubscriber.scanUnsafe(Scannable.Attr.PARENT); + + assertThat(parent).isInstanceOf(Subscriber.class).isNotInstanceOf(ScopePassingSpanSubscriber.class); + } + + @Override + public void onNext(Integer integer) { + } + + @Override + public void onError(Throwable throwable) { + throw Exceptions.propagate(throwable); + } + + @Override + public void onComplete() { + } + }); + + } + } + + private static Stream should_not_double_wrap_async_publisher_Args() { + return Stream.of( + // async source is hidden by Mono.defer() + Arguments.of("hidden by defer", + (Supplier) () -> Mono + .defer(() -> Mono.fromSupplier(() -> 1).hide().subscribeOn(Schedulers.parallel()))), + // async source is directly accessible during subscription + Arguments.of("directly accessible", + (Supplier) () -> Mono.fromSupplier(() -> 1).hide().subscribeOn(Schedulers.parallel()))); + } + }