handle double wrapping

This commit is contained in:
Roman Matiushchenko
2021-01-24 20:50:14 +02:00
parent e5329a9bb2
commit 615f70c03a
5 changed files with 153 additions and 26 deletions

View File

@@ -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 {

View File

@@ -143,7 +143,7 @@ final class ReactorHooksHelper {
}
}
private static boolean isTraceContextPropagator(Publisher<?> current) {
static boolean isTraceContextPropagator(Publisher<?> current) {
return current instanceof TraceContextPropagator;
}

View File

@@ -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();

View File

@@ -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;
}
}

View File

@@ -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())));
}
}