handle double wrapping
This commit is contained in:
@@ -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<CurrentTraceContext> provider = LazyBean.create(context, CurrentTraceContext.class);
|
||||
|
||||
assertThatCode(() -> provider.getOrError()).isInstanceOf(NoSuchBeanDefinitionException.class);
|
||||
}
|
||||
|
||||
@Configuration(proxyBeanMethods = false)
|
||||
static class BasicConfig {
|
||||
|
||||
|
||||
@@ -143,7 +143,7 @@ final class ReactorHooksHelper {
|
||||
}
|
||||
}
|
||||
|
||||
private static boolean isTraceContextPropagator(Publisher<?> current) {
|
||||
static boolean isTraceContextPropagator(Publisher<?> current) {
|
||||
return current instanceof TraceContextPropagator;
|
||||
}
|
||||
|
||||
|
||||
@@ -82,8 +82,10 @@ public abstract class ReactorSleuth {
|
||||
LazyBean<CurrentTraceContext> lazyCurrentTraceContext = LazyBean.create(springContext,
|
||||
CurrentTraceContext.class);
|
||||
|
||||
LazyBean<Tracer> 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<CurrentTraceContext> lazyCurrentTraceContext = LazyBean.create(springContext,
|
||||
CurrentTraceContext.class);
|
||||
|
||||
Predicate<Publisher> publisherPredicate = ReactorHooksHelper::shouldDecorate;
|
||||
BiFunction lifter = liftFunction(springContext, lazyCurrentTraceContext);
|
||||
LazyBean<Tracer> lazyTracer = LazyBean.create(springContext, Tracer.class);
|
||||
|
||||
return ReactorHooksHelper.liftPublisher(publisherPredicate, lifter);
|
||||
@SuppressWarnings("rawtypes")
|
||||
Predicate<Publisher> shouldDecorate = ReactorHooksHelper::shouldDecorate;
|
||||
@SuppressWarnings("rawtypes")
|
||||
BiFunction<Publisher, ? super CoreSubscriber<? super T>, ? extends CoreSubscriber<? super T>> lifter = liftFunction(
|
||||
springContext, lazyCurrentTraceContext, lazyTracer);
|
||||
|
||||
return ReactorHooksHelper.liftPublisher(shouldDecorate, lifter);
|
||||
}
|
||||
|
||||
static <O> BiFunction<Publisher, ? super CoreSubscriber<? super O>, ? extends CoreSubscriber<? super O>> liftFunction(
|
||||
ConfigurableApplicationContext springContext, LazyBean<CurrentTraceContext> lazyCurrentTraceContext) {
|
||||
ConfigurableApplicationContext springContext, LazyBean<CurrentTraceContext> lazyCurrentTraceContext,
|
||||
LazyBean<Tracer> 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 <T> Context contextWithBeans(Context context, ConfigurableApplicationContext springContext) {
|
||||
private static <T> Context contextWithBeans(Context context, LazyBean<Tracer> tracer,
|
||||
LazyBean<CurrentTraceContext> 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<Tracer> lazyTracer = LazyBean.create(springContext, Tracer.class);
|
||||
LazyBean<CurrentTraceContext> 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<CurrentTraceContext> lazyCurrentTraceContext = LazyBean.create(springContext,
|
||||
CurrentTraceContext.class);
|
||||
LazyBean<Tracer> lazyTracer = LazyBean.create(springContext, Tracer.class);
|
||||
|
||||
BiFunction<Publisher, ? super CoreSubscriber<? super T>, ? extends CoreSubscriber<? super T>> scopePassingSpanSubscriber = liftFunction(
|
||||
springContext, lazyCurrentTraceContext);
|
||||
springContext, lazyCurrentTraceContext, lazyTracer);
|
||||
|
||||
BiFunction<Publisher, ? super CoreSubscriber<? super T>, ? extends CoreSubscriber<? super T>> 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();
|
||||
|
||||
@@ -55,12 +55,8 @@ public final class LazyBean<T> {
|
||||
*/
|
||||
@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<T> {
|
||||
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;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -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<Mono<Integer>> sourceSupplier) {
|
||||
springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext);
|
||||
springContext.registerBean(Tracer.class, () -> Mockito.mock(Tracer.class));
|
||||
springContext.refresh();
|
||||
|
||||
Hooks.onEachOperator(HOOK_KEY, ReactorSleuth.onEachOperatorForOnEachInstrumentation(springContext));
|
||||
Hooks.onLastOperator(HOOK_KEY, ReactorSleuth.onLastOperatorForOnEachInstrumentation(springContext));
|
||||
|
||||
AtomicBoolean once = new AtomicBoolean();
|
||||
Hooks.onLastOperator("test", p -> {
|
||||
// check only first onLast Hook
|
||||
if (once.compareAndSet(false, true)) {
|
||||
assertThat(p).isInstanceOf(TraceContextPropagator.class);
|
||||
Object parent = Scannable.from(p).scanUnsafe(Scannable.Attr.PARENT);
|
||||
assertThat(parent).isNotInstanceOf(TraceContextPropagator.class);
|
||||
}
|
||||
return p;
|
||||
});
|
||||
try (CurrentTraceContext.Scope ws = currentTraceContext().newScope(context())) {
|
||||
Mono<Integer> source = sourceSupplier.get();
|
||||
|
||||
source.subscribe((Subscriber<? super Integer>) new CoreSubscriber<Integer>() {
|
||||
@Override
|
||||
public void onSubscribe(Subscription subscription) {
|
||||
assertThat(subscription).isInstanceOf(ScopePassingSpanSubscriber.class);
|
||||
ScopePassingSpanSubscriber<Integer> spanSubscriber = (ScopePassingSpanSubscriber) subscription;
|
||||
Object parent = spanSubscriber.scanUnsafe(Scannable.Attr.PARENT);
|
||||
|
||||
assertThat(parent).isInstanceOf(Subscriber.class).isNotInstanceOf(ScopePassingSpanSubscriber.class);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onNext(Integer integer) {
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onError(Throwable throwable) {
|
||||
throw Exceptions.propagate(throwable);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onComplete() {
|
||||
}
|
||||
});
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
private static Stream<Arguments> should_not_double_wrap_async_publisher_Args() {
|
||||
return Stream.of(
|
||||
// async source is hidden by Mono.defer()
|
||||
Arguments.of("hidden by defer",
|
||||
(Supplier) () -> Mono
|
||||
.defer(() -> Mono.fromSupplier(() -> 1).hide().subscribeOn(Schedulers.parallel()))),
|
||||
// async source is directly accessible during subscription
|
||||
Arguments.of("directly accessible",
|
||||
(Supplier) () -> Mono.fromSupplier(() -> 1).hide().subscribeOn(Schedulers.parallel())));
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user