Extracts LazyBean and backports reactor code to use it (#1542)

There are numerous places where code defensively guards access to
BeanFactory methods. This centralizes the code, first using in reactor
as that's the more performance sensitive.

This also weaves in feedback from #1541 given by @simonbasle and @robotmrv

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

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2013-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.sleuth.instrument.reactor;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.lang.Nullable;
+
+/**
+ * Avoids calling the expensive {@link ConfigurableApplicationContext#getBean(Class)} many
+ * times or throwing an exception.
+ */
+final class LazyBean<T> {
+
+	// spring-jcl uses commons-logging, so do we.
+	private static final Log log = LogFactory.getLog(LazyBean.class);
+
+	final ConfigurableApplicationContext springContext;
+
+	final Class<T> requiredType;
+
+	T value;
+
+	LazyBean(ConfigurableApplicationContext springContext, Class<T> requiredType) {
+		this.springContext = springContext;
+		this.requiredType = requiredType;
+	}
+
+	/**
+	 * Attempts to provision from the underlying bean factory, if not already provisioned.
+	 * @return the bean value or null if there was an exception getting it.
+	 */
+	@Nullable
+	T get() {
+		if (this.value != null) {
+			return this.value;
+		}
+
+		try {
+			this.value = springContext.getBean(requiredType);
+		}
+		catch (Exception ex) {
+			if (log.isDebugEnabled()) {
+				log.debug("Spring context [" + springContext + "] error getting ["
+						+ requiredType + "].", ex);
+			}
+		}
+		return this.value;
+	}
+
+}

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

@@ -16,8 +16,6 @@
 
 package org.springframework.cloud.sleuth.instrument.reactor;
 
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 
 import brave.Tracing;
@@ -70,10 +68,14 @@ public abstract class ReactorSleuth {
 			log.trace("Scope passing operator [" + springContext + "]");
 		}
 
+		// keep a reference outside the lambda so that any caching will be visible to
+		// all publishers
+		LazyBean<CurrentTraceContext> lazyCurrentTraceContext = new LazyBean<>(
+				springContext, CurrentTraceContext.class);
+
 		return Operators.liftPublisher((p, sub) -> {
-			// While supply of scalar types may be deferred, we don't currently scope
-			// production of values in a trace context. This prevents excessive overhead
-			// when using constant results such as Flux/Mono #just, #empty, #error
+			// We don't scope scalar results as they happen in an instant. This prevents
+			// excessive overhead when using Flux/Mono #just, #empty, #error, etc.
 			if (p instanceof Fuseable.ScalarCallable) {
 				return sub;
 			}
@@ -88,13 +90,36 @@ public abstract class ReactorSleuth {
 				return sub;
 			}
 
+			Context context = sub.currentContext();
+
 			if (log.isTraceEnabled()) {
-				log.trace("Spring Context [" + springContext
-						+ "] Creating a scope passing span subscriber with Reactor Context "
-						+ "[" + sub.currentContext() + "] and name [" + name(sub) + "]");
+				log.trace("Spring context [" + springContext + "], Reactor context ["
+						+ context + "], name [" + name(sub) + "]");
 			}
 
-			return scopePassingSpanSubscription(springContext, sub);
+			// Try to get the current trace context bean, lenient when there are problems
+			CurrentTraceContext currentTraceContext = lazyCurrentTraceContext.get();
+			if (currentTraceContext == null) {
+				if (log.isTraceEnabled()) {
+					log.trace("Spring Context [" + springContext
+							+ "] did not return a CurrentTraceContext. Reactor Context is ["
+							+ sub.currentContext() + "] and name is [" + name(sub) + "]");
+				}
+				assert false; // should never happen, but don't break.
+				return sub;
+			}
+
+			TraceContext parent = traceContext(context, currentTraceContext);
+			if (parent == null) {
+				return sub; // no need to scope a null parent
+			}
+
+			if (log.isTraceEnabled()) {
+				log.trace("Creating a scope passing span subscriber with Reactor Context "
+						+ "[" + context + "] and name [" + name(sub) + "]");
+			}
+			return new ScopePassingSpanSubscriber<>(sub, context, currentTraceContext,
+					parent);
 		});
 	}
 
@@ -102,25 +127,14 @@ public abstract class ReactorSleuth {
 		return Scannable.from(sub).name();
 	}
 
-	private static Map<ConfigurableApplicationContext, CurrentTraceContext> CACHE = new ConcurrentHashMap<>();
-
-	static <T> CoreSubscriber<? super T> scopePassingSpanSubscription(
-			ConfigurableApplicationContext springContext, CoreSubscriber<? super T> sub) {
-		CurrentTraceContext currentTraceContext = CACHE.computeIfAbsent(springContext,
-				springContext1 -> springContext1.getBean(CurrentTraceContext.class));
-		Context context = sub.currentContext();
-
-		TraceContext parent = context.getOrDefault(TraceContext.class, null);
-		if (parent == null) {
-			parent = currentTraceContext.get();
-		}
-		if (parent != null) {
-			return new ScopePassingSpanSubscriber<>(sub, context, currentTraceContext,
-					parent);
-		}
-		else {
-			return sub; // no need to trace
+	/**
+	 * Like {@link CurrentTraceContext#get()}, except it first checks the reactor context.
+	 */
+	static TraceContext traceContext(Context context, CurrentTraceContext fallback) {
+		if (context.hasKey(TraceContext.class)) {
+			return context.get(TraceContext.class);
 		}
+		return fallback.get();
 	}
 
 }

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

@@ -54,8 +54,7 @@ final class ScopePassingSpanSubscriber<T> implements SpanSubscription<T>, Scanna
 		this.subscriber = subscriber;
 		this.currentTraceContext = currentTraceContext;
 		this.parent = parent;
-		this.context = ctx != null && parent != null ? ctx.put(TraceContext.class, parent)
-				: ctx != null ? ctx : Context.empty();
+		this.context = parent != null ? ctx.put(TraceContext.class, parent) : ctx;
 		if (log.isTraceEnabled()) {
 			log.trace("Parent span [" + parent + "], context [" + this.context + "]");
 		}
@@ -81,7 +80,6 @@ final class ScopePassingSpanSubscriber<T> implements SpanSubscription<T>, Scanna
 		try (Scope scope = this.currentTraceContext.maybeScope(this.parent)) {
 			this.s.cancel();
 		}
-
 	}
 
 	@Override

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

@@ -109,7 +109,7 @@ final class SpanSubscriber<T> extends AtomicBoolean implements SpanSubscription<
 			// no additional cleaning is required cause we operate on scopes
 			if (log.isTraceEnabled()) {
 				log.trace("Request after cleaning. Current span [{}]",
-						this.currentTraceContext.get());
+						this.span.context());
 			}
 		}
 	}

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

@@ -45,6 +45,7 @@ import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.env.ConfigurableEnvironment;
 
+import static org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth.scopePassingSpanOperator;
 import static org.springframework.cloud.sleuth.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration.SLEUTH_TRACE_REACTOR_KEY;
 
 /**
@@ -154,14 +155,14 @@ class HooksRefresher implements ApplicationListener<RefreshScopeRefreshedEvent>
 				log.trace("Decorating onEach operator instrumentation");
 			}
 			Hooks.onEachOperator(SLEUTH_TRACE_REACTOR_KEY,
-					ReactorSleuth.scopePassingSpanOperator(this.context));
+					scopePassingSpanOperator(this.context));
 		}
 		else {
 			if (log.isTraceEnabled()) {
 				log.trace("Decorating onLast operator instrumentation");
 			}
 			Hooks.onLastOperator(SLEUTH_TRACE_REACTOR_KEY,
-					ReactorSleuth.scopePassingSpanOperator(this.context));
+					scopePassingSpanOperator(this.context));
 		}
 	}
 
@@ -198,14 +199,14 @@ class HookRegisteringBeanDefinitionRegistryPostProcessor
 				log.trace("Decorating onEach operator instrumentation");
 			}
 			Hooks.onEachOperator(SLEUTH_TRACE_REACTOR_KEY,
-					ReactorSleuth.scopePassingSpanOperator(springContext));
+					scopePassingSpanOperator(this.springContext));
 		}
 		else {
 			if (log.isTraceEnabled()) {
 				log.trace("Decorating onLast operator instrumentation");
 			}
 			Hooks.onLastOperator(SLEUTH_TRACE_REACTOR_KEY,
-					ReactorSleuth.scopePassingSpanOperator(springContext));
+					scopePassingSpanOperator(this.springContext));
 		}
 		Schedulers.setExecutorServiceDecorator(
 				TraceReactorAutoConfiguration.SLEUTH_REACTOR_EXECUTOR_SERVICE_KEY,

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/LazyBeanTests.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2013-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.sleuth.instrument.reactor;
+
+import brave.propagation.CurrentTraceContext;
+import org.junit.Test;
+
+import org.springframework.context.ConfigurableApplicationContext;
+
+import static org.assertj.core.api.BDDAssertions.then;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class LazyBeanTests {
+
+	@Test
+	public void should_return_null_when_exception_thrown_upon_bean_retrieval() {
+		ConfigurableApplicationContext springContext = mock(
+				ConfigurableApplicationContext.class);
+
+		when(springContext.getBean(CurrentTraceContext.class))
+				.thenThrow(new IllegalStateException());
+
+		LazyBean<CurrentTraceContext> provider = new LazyBean<>(springContext,
+				CurrentTraceContext.class);
+
+		then(provider.get()).isNull();
+	}
+
+}