Properly clears the scopes

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

@@ -94,6 +94,13 @@ public interface CurrentTraceContext {
 	 */
 	interface Scope extends Closeable {
 
+		/**
+		 * Noop instance.
+		 */
+		Scope NOOP = () -> {
+
+		};
+
 		@Override
 		void close();
 

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

@@ -170,6 +170,13 @@ default CurrentTraceContext currentTraceContext() {
 	 */
 	interface SpanInScope extends Closeable {
 
+		/**
+		 * Noop instance.
+		 */
+		SpanInScope NOOP = () -> {
+
+		};
+
 		@Override
 		void close();
 

spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveCurrentTraceContext.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 the original author or authors.
+ * Copyright 2013-2023 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.
@@ -20,17 +20,15 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 
+import brave.propagation.ThreadLocalCurrentTraceContext;
+
 import org.springframework.cloud.sleuth.CurrentTraceContext;
 import org.springframework.cloud.sleuth.TraceContext;
 
-/**
- * Brave implementation of a {@link CurrentTraceContext}.
- *
- * @author Marcin Grzejszczak
- * @since 3.0.0
- */
 public class BraveCurrentTraceContext implements CurrentTraceContext {
 
+	final ThreadLocal<Scope> scopes = new ThreadLocal<>();
+
 	final brave.propagation.CurrentTraceContext delegate;
 
 	public BraveCurrentTraceContext(brave.propagation.CurrentTraceContext delegate) {
@@ -40,20 +38,36 @@ public BraveCurrentTraceContext(brave.propagation.CurrentTraceContext delegate)
 	@Override
 	public TraceContext context() {
 		brave.propagation.TraceContext context = this.delegate.get();
+		return context == null ? null : new BraveTraceContext(context);
+	}
+
+	@Override
+	public CurrentTraceContext.Scope newScope(TraceContext context) {
 		if (context == null) {
-			return null;
+			clearScopes();
+			return Scope.NOOP;
 		}
-		return new BraveTraceContext(context);
+		return new RevertingScope(this, new BraveScope(this.delegate.newScope(BraveTraceContext.toBrave(context))));
 	}
 
 	@Override
-	public Scope newScope(TraceContext context) {
-		return new BraveScope(this.delegate.newScope(BraveTraceContext.toBrave(context)));
+	public CurrentTraceContext.Scope maybeScope(TraceContext context) {
+		if (context == null) {
+			clearScopes();
+			return Scope.NOOP;
+		}
+		return new RevertingScope(this, new BraveScope(this.delegate.maybeScope(BraveTraceContext.toBrave(context))));
 	}
 
-	@Override
-	public Scope maybeScope(TraceContext context) {
-		return new BraveScope(this.delegate.maybeScope(BraveTraceContext.toBrave(context)));
+	private void clearScopes() {
+		Scope current = this.scopes.get();
+		while (current != null) {
+			current.close();
+			current = this.scopes.get();
+		}
+		if (this.delegate instanceof ThreadLocalCurrentTraceContext) {
+			((ThreadLocalCurrentTraceContext) this.delegate).clear();
+		}
 	}
 
 	@Override
@@ -86,6 +100,29 @@ static CurrentTraceContext fromBrave(brave.propagation.CurrentTraceContext conte
 
 }
 
+class RevertingScope implements CurrentTraceContext.Scope {
+
+	private final BraveCurrentTraceContext currentTraceContext;
+
+	private final CurrentTraceContext.Scope previous;
+
+	private final CurrentTraceContext.Scope current;
+
+	RevertingScope(BraveCurrentTraceContext currentTraceContext, CurrentTraceContext.Scope current) {
+		this.currentTraceContext = currentTraceContext;
+		this.previous = this.currentTraceContext.scopes.get();
+		this.current = current;
+		this.currentTraceContext.scopes.set(this);
+	}
+
+	@Override
+	public void close() {
+		this.current.close();
+		this.currentTraceContext.scopes.set(this.previous);
+	}
+
+}
+
 class BraveScope implements CurrentTraceContext.Scope {
 
 	private final brave.propagation.CurrentTraceContext.Scope delegate;

spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveTracer.java

@@ -16,6 +16,8 @@
 
 package org.springframework.cloud.sleuth.brave.bridge;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.util.Map;
 
 import brave.propagation.TraceContextOrSamplingFlags;
@@ -27,7 +29,6 @@
 import org.springframework.cloud.sleuth.SpanCustomizer;
 import org.springframework.cloud.sleuth.TraceContext;
 import org.springframework.cloud.sleuth.Tracer;
-import org.springframework.cloud.sleuth.docs.AssertingSpan;
 
 /**
  * Brave implementation of a {@link Tracer}.
@@ -70,8 +71,11 @@ public Span nextSpan(Span parent) {
 
 	@Override
 	public SpanInScope withSpan(Span span) {
-		return new BraveSpanInScope(
-				tracer.withSpanInScope(span == null ? null : ((BraveSpan) AssertingSpan.unwrap(span)).delegate));
+		if (span == null) {
+			currentTraceContext.maybeScope(null);
+			return SpanInScope.NOOP;
+		}
+		return new BraveSpanInScope(currentTraceContext.maybeScope(span.context()));
 	}
 
 	@Override
@@ -142,15 +146,20 @@ public CurrentTraceContext currentTraceContext() {
 
 class BraveSpanInScope implements Tracer.SpanInScope {
 
-	final brave.Tracer.SpanInScope delegate;
+	final Closeable delegate;
 
-	BraveSpanInScope(brave.Tracer.SpanInScope delegate) {
+	BraveSpanInScope(Closeable delegate) {
 		this.delegate = delegate;
 	}
 
 	@Override
 	public void close() {
-		this.delegate.close();
+		try {
+			this.delegate.close();
+		}
+		catch (IOException e) {
+			throw new RuntimeException(e);
+		}
 	}
 
 }

spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveCurrentTraceContextTests.java

@@ -0,0 +1,152 @@
+/*
+ * Copyright 2013-2023 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.brave.bridge;
+
+import brave.context.slf4j.MDCScopeDecorator;
+import brave.propagation.ThreadLocalCurrentTraceContext;
+import brave.propagation.TraceContext;
+import org.junit.jupiter.api.Test;
+import org.slf4j.MDC;
+
+import org.springframework.cloud.sleuth.CurrentTraceContext;
+
+import static org.assertj.core.api.BDDAssertions.then;
+
+class BraveCurrentTraceContextTests {
+
+	ThreadLocalCurrentTraceContext currentTraceContext = ThreadLocalCurrentTraceContext.newBuilder()
+			.addScopeDecorator(MDCScopeDecorator.newBuilder().build()).build();
+
+	@Test
+	void should_clear_any_thread_locals_and_scopes_when_null_context_passed_to_new_scope() {
+		BraveCurrentTraceContext braveCurrentTraceContext = new BraveCurrentTraceContext(currentTraceContext);
+		brave.propagation.CurrentTraceContext.Scope newScope = currentTraceContext
+				.newScope(TraceContext.newBuilder().traceId(12345678).spanId(12345678).build());
+		then(currentTraceContext.get()).isNotNull();
+
+		CurrentTraceContext.Scope scope = braveCurrentTraceContext.newScope(null);
+
+		thenThreadLocalsGotCleared(braveCurrentTraceContext, scope);
+		newScope.close();
+		thenThreadLocalsGotCleared(braveCurrentTraceContext, scope);
+	}
+
+	@Test
+	void should_clear_any_thread_locals_and_scopes_when_null_context_passed_to_new_scope_with_nested_scopes() {
+		BraveCurrentTraceContext braveCurrentTraceContext = new BraveCurrentTraceContext(currentTraceContext);
+
+		try (CurrentTraceContext.Scope scope1 = braveCurrentTraceContext.newScope(
+				BraveTraceContext.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345670).build()))) {
+			then(currentTraceContext.get()).isNotNull();
+			thenMdcEntriesArePresent();
+			try (CurrentTraceContext.Scope scope2 = braveCurrentTraceContext.newScope(BraveTraceContext
+					.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345671).build()))) {
+				then(currentTraceContext.get()).isNotNull();
+				thenMdcEntriesArePresent();
+				try (CurrentTraceContext.Scope scope3 = braveCurrentTraceContext.newScope(BraveTraceContext
+						.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345672).build()))) {
+					then(currentTraceContext.get()).isNotNull();
+					thenMdcEntriesArePresent();
+					try (CurrentTraceContext.Scope nullScope = braveCurrentTraceContext.newScope(null)) {
+						// This closes all scopes and MDC entries
+						then(currentTraceContext.get()).isNull();
+					}
+					// We have nothing to revert to since the nullScope is ignoring
+					// everything there was before
+					then(currentTraceContext.get()).isNull();
+					thenMdcEntriesAreMissing();
+				}
+				then(currentTraceContext.get()).isNull();
+				thenMdcEntriesAreMissing();
+			}
+			then(currentTraceContext.get()).isNull();
+			thenMdcEntriesAreMissing();
+		}
+
+		then(currentTraceContext.get()).isNull();
+		then(braveCurrentTraceContext.scopes.get()).isNull();
+		then(MDC.getCopyOfContextMap()).isEmpty();
+	}
+
+	@Test
+	void should_clear_any_thread_locals_and_scopes_when_null_context_passed_to_maybe_scope_with_nested_scopes() {
+		BraveCurrentTraceContext braveCurrentTraceContext = new BraveCurrentTraceContext(currentTraceContext);
+
+		try (CurrentTraceContext.Scope scope1 = braveCurrentTraceContext.maybeScope(
+				BraveTraceContext.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345670).build()))) {
+			then(currentTraceContext.get()).isNotNull();
+			thenMdcEntriesArePresent();
+			try (CurrentTraceContext.Scope scope2 = braveCurrentTraceContext.maybeScope(BraveTraceContext
+					.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345671).build()))) {
+				then(currentTraceContext.get()).isNotNull();
+				thenMdcEntriesArePresent();
+				try (CurrentTraceContext.Scope scope3 = braveCurrentTraceContext.maybeScope(BraveTraceContext
+						.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345672).build()))) {
+					then(currentTraceContext.get()).isNotNull();
+					thenMdcEntriesArePresent();
+					try (CurrentTraceContext.Scope nullScope = braveCurrentTraceContext.maybeScope(null)) {
+						// This closes all scopes and MDC entries
+						then(currentTraceContext.get()).isNull();
+					}
+					// We have nothing to revert to since the nullScope is ignoring
+					// everything there was before
+					then(currentTraceContext.get()).isNull();
+					thenMdcEntriesAreMissing();
+				}
+				then(currentTraceContext.get()).isNull();
+				thenMdcEntriesAreMissing();
+			}
+			then(currentTraceContext.get()).isNull();
+			thenMdcEntriesAreMissing();
+		}
+
+		then(currentTraceContext.get()).isNull();
+		then(braveCurrentTraceContext.scopes.get()).isNull();
+		then(MDC.getCopyOfContextMap()).isEmpty();
+	}
+
+	private static void thenMdcEntriesArePresent() {
+		then(MDC.get("traceId")).isEqualTo("0000000000bc614e");
+		then(MDC.get("spanId")).isNotEmpty();
+	}
+
+	private static void thenMdcEntriesAreMissing() {
+		then(MDC.getCopyOfContextMap()).isEmpty();
+	}
+
+	@Test
+	void should_clear_any_thread_locals_and_scopes_when_null_context_passed_to_maybe_scope() {
+		BraveCurrentTraceContext braveCurrentTraceContext = new BraveCurrentTraceContext(currentTraceContext);
+		brave.propagation.CurrentTraceContext.Scope maybeScope = currentTraceContext
+				.newScope(TraceContext.newBuilder().traceId(12345678).spanId(12345678).build());
+		then(currentTraceContext.get()).isNotNull();
+
+		CurrentTraceContext.Scope scope = braveCurrentTraceContext.maybeScope(null);
+
+		thenThreadLocalsGotCleared(braveCurrentTraceContext, scope);
+		maybeScope.close();
+		thenThreadLocalsGotCleared(braveCurrentTraceContext, scope);
+	}
+
+	private void thenThreadLocalsGotCleared(BraveCurrentTraceContext braveCurrentTraceContext,
+			CurrentTraceContext.Scope scope) {
+		then(scope).isSameAs(CurrentTraceContext.Scope.NOOP);
+		then(currentTraceContext.get()).isNull();
+		then(braveCurrentTraceContext.scopes.get()).isNull();
+	}
+
+}