Fixes invalid ThreadLocalSpan stacking and tracing context leaks; fixes gh-2064; fixes gh-2108

spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/SpanAndScope.java

@@ -56,8 +56,8 @@ public class SpanAndScope implements Closeable {
 
 	@Override
 	public void close() {
-		if (log.isTraceEnabled()) {
-			log.trace("Closing span [" + this.span + "]");
+		if (log.isDebugEnabled()) {
+			log.debug("Closing span [" + this.span + "], scope is not null [" + (this.scope != null) + "]");
 		}
 		if (this.scope != null) {
 			this.scope.close();

spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/ThreadLocalSpan.java

@@ -16,12 +16,7 @@
 
 package org.springframework.cloud.sleuth;
 
-import java.util.Deque;
-import java.util.NoSuchElementException;
-import java.util.concurrent.LinkedBlockingDeque;
-
-import org.apache.commons.logging.Log;
-import org.apache.commons.logging.LogFactory;
+import java.util.ArrayDeque;
 
 /**
  * Represents a {@link Span} stored in thread local.
@@ -31,11 +26,7 @@ import org.apache.commons.logging.LogFactory;
  */
 public class ThreadLocalSpan {
 
-	private static final Log log = LogFactory.getLog(ThreadLocalSpan.class);
-
-	private final ThreadLocal<SpanAndScope> threadLocalSpan = new ThreadLocal<>();
-
-	private final Deque<SpanAndScope> spans = new LinkedBlockingDeque<>();
+	private final ThreadLocal<ArrayDeque<SpanAndScope>> currentSpanInScopeStack = new ThreadLocal<>();
 
 	private final Tracer tracer;
 
@@ -50,18 +41,14 @@ public class ThreadLocalSpan {
 	public void set(Span span) {
 		Tracer.SpanInScope spanInScope = this.tracer.withSpan(span);
 		SpanAndScope newSpanAndScope = new SpanAndScope(span, spanInScope);
-		SpanAndScope scope = this.threadLocalSpan.get();
-		if (scope != null) {
-			this.spans.addFirst(scope);
-		}
-		this.threadLocalSpan.set(newSpanAndScope);
+		getCurrentSpanInScopeStack().addFirst(newSpanAndScope);
 	}
 
 	/**
 	 * @return currently stored span and scope
 	 */
 	public SpanAndScope get() {
-		return this.threadLocalSpan.get();
+		return getCurrentSpanInScopeStack().peekFirst();
 	}
 
 	/**
@@ -69,22 +56,22 @@ public class ThreadLocalSpan {
 	 * current thread local.
 	 */
 	public void remove() {
-		this.threadLocalSpan.remove();
-		if (this.spans.isEmpty()) {
+		SpanAndScope spanAndScope = getCurrentSpanInScopeStack().pollFirst();
+		if (spanAndScope == null) {
 			return;
 		}
-		try {
-			SpanAndScope span = this.spans.removeFirst();
-			if (log.isDebugEnabled()) {
-				log.debug("Took span [" + span + "] from thread local");
-			}
-			this.threadLocalSpan.set(span);
-		}
-		catch (NoSuchElementException ex) {
-			if (log.isTraceEnabled()) {
-				log.trace("Failed to remove a span from the queue", ex);
-			}
+		if (spanAndScope.getScope() != null) {
+			spanAndScope.getScope().close();
 		}
 	}
 
+	private ArrayDeque<SpanAndScope> getCurrentSpanInScopeStack() {
+		ArrayDeque<SpanAndScope> stack = this.currentSpanInScopeStack.get();
+		if (stack == null) {
+			stack = new ArrayDeque<>();
+			this.currentSpanInScopeStack.set(stack);
+		}
+		return stack;
+	}
+
 }

spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/docs/AssertingSpan.java

@@ -153,6 +153,20 @@ public interface AssertingSpan extends Span {
 		return new ImmutableAssertingSpan(documentedSpan, span);
 	}
 
+	/**
+	 * @param documentedSpan span configuration
+	 * @param span span to wrap in assertions
+	 * @return asserting span
+	 */
+	static AssertingSpan continueSpan(DocumentedSpan documentedSpan, Span span) {
+		AssertingSpan assertingSpan = of(documentedSpan, span);
+		if (assertingSpan == null) {
+			return null;
+		}
+		((ImmutableAssertingSpan) assertingSpan).isStarted = true;
+		return assertingSpan;
+	}
+
 	/**
 	 * Returns the underlying delegate. Used when casting is necessary.
 	 * @param span span to check for wrapping

spring-cloud-sleuth-api/src/test/java/org/springframework/cloud/sleuth/ThreadLocalSpanTests.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2013-2021 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;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.BDDMockito;
+
+import static org.assertj.core.api.BDDAssertions.then;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.BDDMockito.given;
+
+class ThreadLocalSpanTests {
+
+	Tracer tracer = BDDMockito.mock(Tracer.class);
+
+	@BeforeEach
+	void setup() {
+		given(this.tracer.withSpan(any())).willReturn(() -> {
+
+		});
+	}
+
+	@Test
+	void should_properly_stack_spans() {
+		// given
+		ThreadLocalSpan threadLocalSpan = new ThreadLocalSpan(tracer);
+		then(threadLocalSpan.get()).isNull();
+
+		// when - Span 1
+		Span span = span();
+		threadLocalSpan.set(span);
+		// then - Span 1
+		then(threadLocalSpan.get().getSpan()).isSameAs(span);
+
+		// when - Span 2
+		Span secondSpan = span();
+		threadLocalSpan.set(secondSpan);
+		// then - Span 2
+		then(threadLocalSpan.get().getSpan()).isSameAs(secondSpan);
+
+		// expect - Span 1
+		threadLocalSpan.remove();
+		then(threadLocalSpan.get().getSpan()).isSameAs(span);
+
+		// expect - null
+		threadLocalSpan.remove();
+		then(threadLocalSpan.get()).isNull();
+	}
+
+	private Span span() {
+		return BDDMockito.mock(Span.class);
+	}
+
+}

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptor.java

@@ -16,14 +16,14 @@
 
 package org.springframework.cloud.sleuth.instrument.messaging;
 
-import java.util.NoSuchElementException;
 import java.util.Set;
-import java.util.concurrent.LinkedBlockingDeque;
 import java.util.function.Function;
 
 import org.springframework.aop.support.AopUtils;
 import org.springframework.beans.BeansException;
 import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.SpanAndScope;
+import org.springframework.cloud.sleuth.ThreadLocalSpan;
 import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.propagation.Propagator;
 import org.springframework.context.ApplicationContext;
@@ -90,7 +90,7 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept
 	private static final Class<?> directWithAttributesChannelClass = ClassUtils.isPresent(STREAM_DIRECT_CHANNEL, null)
 			? ClassUtils.resolveClassName(STREAM_DIRECT_CHANNEL, null) : null;
 
-	private final ThreadLocalSpan threadLocalSpan = new ThreadLocalSpan();
+	private final ThreadLocalSpan threadLocalSpan;
 
 	private final Tracer tracer;
 
@@ -115,6 +115,7 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept
 		this.extractor = getter;
 		this.remoteServiceNameMapper = remoteServiceNameMapper;
 		this.messageSpanCustomizer = messageSpanCustomizer;
+		this.threadLocalSpan = new ThreadLocalSpan(tracer);
 	}
 
 	@Override
@@ -148,8 +149,7 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept
 	}
 
 	private void setSpanInScope(Span span) {
-		Tracer.SpanInScope spanInScope = this.tracer.withSpan(span);
-		this.threadLocalSpan.set(new SpanAndScope(span, spanInScope));
+		this.threadLocalSpan.set(span);
 		log.debug(() -> "Put span in scope " + span);
 	}
 
@@ -308,8 +308,8 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept
 		if (spanAndScope == null) {
 			return;
 		}
-		Span span = spanAndScope.span;
-		Tracer.SpanInScope scope = spanAndScope.scope;
+		Span span = spanAndScope.getSpan();
+		Tracer.SpanInScope scope = spanAndScope.getScope();
 		if (span.isNoop()) {
 			log.debug(() -> "Span " + span + " is noop - will stop the scope");
 			scope.close();
@@ -354,57 +354,4 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept
 		return message;
 	}
 
-	private static class SpanAndScope {
-
-		final Span span;
-
-		final Tracer.SpanInScope scope;
-
-		SpanAndScope(Span span, Tracer.SpanInScope scope) {
-			this.span = span;
-			this.scope = scope;
-		}
-
-	}
-
-	private static class ThreadLocalSpan {
-
-		private static final LogAccessor log = new LogAccessor(ThreadLocalSpan.class);
-
-		private final ThreadLocal<SpanAndScope> threadLocalSpan = new ThreadLocal<>();
-
-		private final LinkedBlockingDeque<SpanAndScope> spans = new LinkedBlockingDeque<>();
-
-		ThreadLocalSpan() {
-		}
-
-		void set(SpanAndScope spanAndScope) {
-			SpanAndScope scope = this.threadLocalSpan.get();
-			if (scope != null) {
-				this.spans.addFirst(scope);
-			}
-			this.threadLocalSpan.set(spanAndScope);
-		}
-
-		SpanAndScope get() {
-			return this.threadLocalSpan.get();
-		}
-
-		void remove() {
-			this.threadLocalSpan.remove();
-			if (this.spans.isEmpty()) {
-				return;
-			}
-			try {
-				SpanAndScope span = this.spans.removeFirst();
-				log.debug(() -> "Took span [" + span + "] from thread local");
-				this.threadLocalSpan.set(span);
-			}
-			catch (NoSuchElementException ex) {
-				log.trace(ex, () -> "Failed to remove a span from the queue");
-			}
-		}
-
-	}
-
 }

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/tx/TraceReactiveTransactionManager.java

@@ -26,6 +26,7 @@ import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.SpanAndScope;
 import org.springframework.cloud.sleuth.TraceContext;
 import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.cloud.sleuth.docs.AssertingSpan;
 import org.springframework.transaction.ReactiveTransaction;
 import org.springframework.transaction.ReactiveTransactionManager;
 import org.springframework.transaction.TransactionDefinition;
@@ -72,10 +73,10 @@ public class TraceReactiveTransactionManager implements ReactiveTransactionManag
 	public Mono<ReactiveTransaction> getReactiveTransaction(TransactionDefinition definition)
 			throws TransactionException {
 		return Mono.deferContextual(contextView -> this.delegate.getReactiveTransaction(definition).map(tx -> {
-			Span span = SleuthTxSpan.TX_SPAN.wrap(span(contextView));
+			Span span = AssertingSpan.continueSpan(SleuthTxSpan.TX_SPAN, SleuthTxSpan.TX_SPAN.wrap(span(contextView)));
 			if (tx.isNewTransaction() || span == null) {
 				if (log.isDebugEnabled()) {
-					log.debug("New transaction is required");
+					log.debug("New transaction is required, span in context [" + span + "]");
 				}
 				if (span == null) {
 					span = SleuthTxSpan.TX_SPAN.wrap(tracer().nextSpan()).name(SleuthTxSpan.TX_SPAN.getName()).start();
@@ -132,11 +133,12 @@ public class TraceReactiveTransactionManager implements ReactiveTransactionManag
 		SpanAndScope spanAndScope = reactiveTransaction.spanAndScope;
 		Span span = spanAndScope.getSpan();
 		if (log.isDebugEnabled()) {
-			log.debug("Commiting the transaction for span [" + spanAndScope + "]");
+			log.debug("Committing the transaction for span [" + spanAndScope + "]");
 		}
+		spanAndScope.getScope().close(); // Otherwise we have a leak
 		return this.delegate.commit(reactiveTransaction.delegate)
 				// TODO: Fix me when this is resolved in Reactor
-				// .doOnSubscribe(__ -> scope.close())
+				// .doOnSubscribe(__ -> spanAndScope.getScope().close())
 				.doOnError(span::error).doOnSuccess(signalType -> spanAndScope.close());
 	}
 
@@ -151,6 +153,7 @@ public class TraceReactiveTransactionManager implements ReactiveTransactionManag
 		if (log.isDebugEnabled()) {
 			log.debug("Rolling back the transaction for span [" + spanAndScope + "]");
 		}
+		spanAndScope.getScope().close(); // Otherwise we have a leak
 		return this.delegate.rollback(reactiveTransaction.delegate)
 				// TODO: Fix me when this is resolved in Reactor
 				// .doOnSubscribe(__ -> scope.close())

tests/brave/spring-cloud-sleuth-instrumentation-r2dbc-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/r2dbc/R2dbcIntegrationTests.java

@@ -28,7 +28,7 @@ import org.springframework.test.context.TestPropertySource;
 
 @SpringBootTest
 @ContextConfiguration(classes = R2dbcIntegrationTests.Config.class)
-@TestPropertySource(properties = "logging.level.org.springframework.cloud=TRACE")
+@TestPropertySource(properties = "logging.level.org.springframework.cloud=DEBUG")
 public class R2dbcIntegrationTests extends org.springframework.cloud.sleuth.instrument.r2dbc.R2dbcIntegrationTests {
 
 	@Configuration(proxyBeanMethods = false)

tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/r2dbc/R2dbcIntegrationTests.java

@@ -41,7 +41,8 @@ import org.springframework.test.context.TestPropertySource;
 import static org.assertj.core.api.BDDAssertions.then;
 
 @ContextConfiguration(classes = R2dbcIntegrationTests.TestConfig.class)
-@TestPropertySource(properties = "spring.application.name=MyApplication")
+@TestPropertySource(properties = { "spring.application.name=MyApplication",
+		"spring.sleuth.reactor.instrumentation-type=decorate_queues" })
 public abstract class R2dbcIntegrationTests {
 
 	@Autowired
@@ -66,8 +67,7 @@ public abstract class R2dbcIntegrationTests {
 				.collect(Collectors.toList());
 		then(spanNames.stream().filter("tx"::equalsIgnoreCase).collect(Collectors.toList())).hasSize(2);
 		then(remoteServiceNames.stream().filter("h2"::equalsIgnoreCase).collect(Collectors.toList())).hasSize(9);
-		// TODO: First fix the SpanAndScope stacking
-		// then(tracer.currentSpan()).isNull();
+		then(tracer.currentSpan()).isNull();
 	}
 
 	@Configuration(proxyBeanMethods = false)