Using lazy span subscriber only when context wasn't refreshed

fixes gh-949 gh-813

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java

@@ -20,13 +20,17 @@ import java.util.function.Function;
 import java.util.function.Predicate;
 
 import brave.Tracing;
-import org.reactivestreams.Publisher;
-import org.springframework.beans.factory.BeanFactory;
+import reactor.core.CoreSubscriber;
 import reactor.core.Fuseable;
 import reactor.core.Scannable;
 import reactor.core.publisher.ConnectableFlux;
 import reactor.core.publisher.Operators;
 import reactor.util.context.Context;
+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;
 
 /**
  * Reactive Span pointcuts factories
@@ -36,6 +40,8 @@ import reactor.util.context.Context;
  */
 public abstract class ReactorSleuth {
 
+	private static final Log log = LogFactory.getLog(ReactorSleuth.class);
+
 	/**
 	 * Return a span operator pointcut given a {@link BeanFactory}. This can be used in reactor
 	 * via {@link reactor.core.publisher.Flux#transform(Function)}, {@link
@@ -55,6 +61,10 @@ public abstract class ReactorSleuth {
 			if(scannable instanceof Fuseable && sub instanceof Fuseable.QueueSubscription){
 				return sub;
 			}
+			if (log.isTraceEnabled()) {
+				log.trace("Creating a lazy span subscriber with context "
+						+ "[" + sub.currentContext() + "] and name [" + scannable.name() + "]");
+			}
 			return new LazySpanSubscriber<T>(
 					new SpanSubscriptionProvider(
 						beanFactory,
@@ -78,6 +88,7 @@ public abstract class ReactorSleuth {
 	 *
 	 * @return a new lazy span operator pointcut
 	 */
+	@SuppressWarnings("unchecked")
 	public static <T> Function<? super Publisher<T>, ? extends Publisher<T>> scopePassingSpanOperator(
 			BeanFactory beanFactory) {
 		return Operators.lift(POINTCUT_FILTER, ((scannable, sub) -> {
@@ -85,26 +96,51 @@ public abstract class ReactorSleuth {
 			if(scannable instanceof Fuseable && sub instanceof Fuseable.QueueSubscription){
 				return sub;
 			}
+			if (contextRefreshed(beanFactory)) {
+				if (log.isTraceEnabled()) {
+					log.trace("Spring Context 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 is not yet refreshed, falling back to lazy span subscriber. "
+						+ "Reactor Context is [" + sub.currentContext() + "] and name is [" + scannable.name() + "]");
+			}
 			return new LazySpanSubscriber<T>(
-					new SpanSubscriptionProvider(
-							beanFactory,
-							sub,
-							sub.currentContext(),
-							scannable.name()) {
-						@Override SpanSubscription newCoreSubscriber(Tracing tracing) {
-							return new ScopePassingSpanSubscriber<T>(
-									sub,
-									sub != null ? sub.currentContext() : Context.empty(),
-									tracing);
-						}
-					}
+					scopePassingSpanSubscription(beanFactory, scannable, sub)
 			);
 		}));
 	}
 
+	private static boolean contextRefreshed(BeanFactory beanFactory) {
+		try {
+			return beanFactory.getBean(ApplicationContextRefreshedListener.class).isRefreshed();
+		} catch (NoSuchBeanDefinitionException e) {
+			return false;
+		}
+	}
+
+	private static <T> SpanSubscriptionProvider scopePassingSpanSubscription(
+			BeanFactory beanFactory, Scannable scannable, CoreSubscriber<? super T> sub) {
+		return new SpanSubscriptionProvider(
+				beanFactory,
+				sub,
+				sub.currentContext(),
+				scannable.name()) {
+			@Override SpanSubscription newCoreSubscriber(Tracing tracing) {
+				return new ScopePassingSpanSubscriber<T>(
+						sub,
+						sub != null ? sub.currentContext() : Context.empty(),
+						tracing);
+			}
+		};
+	}
+
 	private static final Predicate<Scannable> POINTCUT_FILTER =
 			s ->  !(s instanceof ConnectableFlux) && !(s instanceof Fuseable.ScalarCallable) && s.isScanAvailable();
 
 	private ReactorSleuth() {
 	}
-}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java

@@ -21,9 +21,11 @@ import java.util.concurrent.atomic.AtomicBoolean;
 import brave.Span;
 import brave.Tracer;
 import brave.Tracing;
+import reactor.util.context.Context;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
-import reactor.util.context.Context;
 
 /**
  * A trace representation of the {@link Subscriber} that always
@@ -34,6 +36,8 @@ import reactor.util.context.Context;
  */
 final class ScopePassingSpanSubscriber<T> extends AtomicBoolean implements SpanSubscription<T> {
 
+	private static final Log log = LogFactory.getLog(ScopePassingSpanSubscriber.class);
+
 	private final Span span;
 	private final Subscriber<? super T> subscriber;
 	private final Context context;
@@ -48,6 +52,9 @@ final class ScopePassingSpanSubscriber<T> extends AtomicBoolean implements SpanS
 		this.span = root;
 		this.context = ctx != null && root != null ? ctx.put(Span.class, root) :
 				ctx != null ? ctx : Context.empty();
+		if (log.isTraceEnabled()) {
+			log.trace("Root span [" + root + "], context [" + this.context + "]");
+		}
 	}
 
 	@Override public void onSubscribe(Subscription subscription) {

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriptionProvider.java

@@ -19,9 +19,11 @@ package org.springframework.cloud.sleuth.instrument.reactor;
 import java.util.function.Supplier;
 
 import brave.Tracing;
+import reactor.util.context.Context;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.reactivestreams.Subscriber;
 import org.springframework.beans.factory.BeanFactory;
-import reactor.util.context.Context;
 
 /**
  * Supplier to lazily start a {@link SpanSubscription}
@@ -30,6 +32,8 @@ import reactor.util.context.Context;
  */
 class SpanSubscriptionProvider<T> implements Supplier<SpanSubscription<T>> {
 
+	private static final Log log = LogFactory.getLog(SpanSubscriptionProvider.class);
+
 	final BeanFactory beanFactory;
 	final Subscriber<? super T> subscriber;
 	final Context context;
@@ -43,6 +47,9 @@ class SpanSubscriptionProvider<T> implements Supplier<SpanSubscription<T>> {
 		this.subscriber = subscriber;
 		this.context = context;
 		this.name = name;
+		if (log.isTraceEnabled()) {
+			log.trace("Context [" + context + "], name [" + name + "]");
+		}
 	}
 
 	@Override public SpanSubscription<T> get() {

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java

@@ -16,11 +16,15 @@
 
 package org.springframework.cloud.sleuth.instrument.reactor;
 
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.function.Supplier;
 import javax.annotation.PreDestroy;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
 
 import brave.Tracing;
+import reactor.core.publisher.Hooks;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
@@ -35,11 +39,10 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
 import org.springframework.cloud.sleuth.instrument.async.TraceableScheduledExecutorService;
 import org.springframework.cloud.sleuth.instrument.web.TraceWebFluxAutoConfiguration;
+import org.springframework.context.ApplicationListener;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import reactor.core.publisher.Hooks;
-import reactor.core.publisher.Mono;
-import reactor.core.scheduler.Schedulers;
+import org.springframework.context.event.ContextRefreshedEvent;
 
 /**
  * {@link org.springframework.boot.autoconfigure.EnableAutoConfiguration Auto-configuration}
@@ -84,6 +87,10 @@ public class TraceReactorAutoConfiguration {
 		static HookRegisteringBeanDefinitionRegistryPostProcessor traceHookRegisteringBeanDefinitionRegistryPostProcessor() {
 			return new HookRegisteringBeanDefinitionRegistryPostProcessor();
 		}
+
+		@Bean ApplicationContextRefreshedListener traceApplicationContextRefreshedListener() {
+			return new ApplicationContextRefreshedListener();
+		}
 	}
 }
 
@@ -121,4 +128,19 @@ class HookRegisteringBeanDefinitionRegistryPostProcessor implements
 			}
 		};
 	}
-}
+}
+
+class ApplicationContextRefreshedListener implements
+		ApplicationListener<ContextRefreshedEvent> {
+
+	AtomicBoolean refreshed = new AtomicBoolean();
+
+	@Override
+	public void onApplicationEvent(ContextRefreshedEvent contextRefreshedEvent) {
+		this.refreshed.set(true);
+	}
+
+	boolean isRefreshed() {
+		return this.refreshed.get();
+	}
+}