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);
+	}
+
+}