From 3d198306ea107900579bbfaa77f4e0b91ddb37c7 Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Fri, 14 Apr 2023 13:35:35 +0200 Subject: [PATCH] Properly clears the scopes (#2280) without this fix when calling `maybeScope(null)` or `withSpan(null)` or `newScope(null)` we ended up with creating a null scope that would allocate additional resources and not clear things. with this fix we're clearing all the scopes and removing thread locals --- .../cloud/sleuth/CurrentTraceContext.java | 7 + .../springframework/cloud/sleuth/Tracer.java | 7 + .../bridge/BraveCurrentTraceContext.java | 65 ++++++-- .../sleuth/brave/bridge/BraveTracer.java | 21 ++- .../bridge/BraveCurrentTraceContextTests.java | 152 ++++++++++++++++++ .../sleuth/brave/bridge/BraveTracerTests.java | 112 +++++++++++++ 6 files changed, 344 insertions(+), 20 deletions(-) create mode 100644 spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveCurrentTraceContextTests.java create mode 100644 spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveTracerTests.java diff --git a/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/CurrentTraceContext.java b/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/CurrentTraceContext.java index aaece67cb..e7ba6fea0 100644 --- a/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/CurrentTraceContext.java +++ b/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(); diff --git a/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/Tracer.java b/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/Tracer.java index 6e2ee26a9..7b2374ebb 100644 --- a/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/Tracer.java +++ b/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/Tracer.java @@ -170,6 +170,13 @@ public interface Tracer extends BaggageManager { */ interface SpanInScope extends Closeable { + /** + * Noop instance. + */ + SpanInScope NOOP = () -> { + + }; + @Override void close(); diff --git a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveCurrentTraceContext.java b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveCurrentTraceContext.java index 2bc97d580..f7990e396 100644 --- a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveCurrentTraceContext.java +++ b/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.Callable; 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 scopes = new ThreadLocal<>(); + final brave.propagation.CurrentTraceContext delegate; public BraveCurrentTraceContext(brave.propagation.CurrentTraceContext delegate) { @@ -40,20 +38,36 @@ public class BraveCurrentTraceContext implements CurrentTraceContext { @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 @@ public class BraveCurrentTraceContext implements CurrentTraceContext { } +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; diff --git a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveTracer.java b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveTracer.java index d0de994f8..6aa04962b 100644 --- a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveTracer.java +++ b/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.Span; 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 class BraveTracer implements Tracer { @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 class BraveTracer implements Tracer { 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); + } } } diff --git a/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveCurrentTraceContextTests.java b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveCurrentTraceContextTests.java new file mode 100644 index 000000000..9f70be6cb --- /dev/null +++ b/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(); + } + +} diff --git a/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveTracerTests.java b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveTracerTests.java new file mode 100644 index 000000000..36a73d345 --- /dev/null +++ b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveTracerTests.java @@ -0,0 +1,112 @@ +/* + * 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.Tracing; +import brave.context.slf4j.MDCScopeDecorator; +import brave.propagation.ThreadLocalCurrentTraceContext; +import org.junit.jupiter.api.Test; +import org.slf4j.MDC; + +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.Tracer; + +import static org.assertj.core.api.BDDAssertions.then; + +class BraveTracerTests { + + ThreadLocalCurrentTraceContext currentTraceContext = ThreadLocalCurrentTraceContext.newBuilder() + .addScopeDecorator(MDCScopeDecorator.newBuilder().build()).build(); + + Tracing tracing = Tracing.newBuilder().currentTraceContext(currentTraceContext).build(); + + BraveCurrentTraceContext braveCurrentTraceContext = new BraveCurrentTraceContext(currentTraceContext); + + BraveTracer braveTracer = new BraveTracer(tracing.tracer(), braveCurrentTraceContext, new BraveBaggageManager()); + + @Test + void should_clear_any_thread_locals_and_scopes_when_null_context_passed_to_with_span() { + Span span = braveTracer.nextSpan(); + Tracer.SpanInScope newScope = braveTracer.withSpan(span.start()); + then(braveTracer.currentSpan()).isEqualTo(span); + + Tracer.SpanInScope noopScope = braveTracer.withSpan(null); + + thenThreadLocalsGotCleared(braveCurrentTraceContext, noopScope); + newScope.close(); + thenThreadLocalsGotCleared(braveCurrentTraceContext, noopScope); + } + + @Test + void should_clear_any_thread_locals_and_scopes_when_null_context_passed_to_with_span_with_nested_scopes() { + Span nextSpan1 = braveTracer.nextSpan(); + try (Tracer.SpanInScope scope1 = braveTracer.withSpan(nextSpan1.start())) { + then(braveTracer.currentSpan()).isEqualTo(nextSpan1); + thenMdcEntriesArePresent(nextSpan1.context()); + Span nextSpan2 = braveTracer.nextSpan(); + try (Tracer.SpanInScope scope2 = braveTracer.withSpan(nextSpan2.start())) { + then(braveTracer.currentSpan()).isEqualTo(nextSpan2); + thenMdcEntriesArePresent(nextSpan2.context()); + Span nextSpan3 = braveTracer.nextSpan(); + try (Tracer.SpanInScope scope3 = braveTracer.withSpan(nextSpan3.start())) { + then(braveTracer.currentSpan()).isEqualTo(nextSpan3); + thenMdcEntriesArePresent(nextSpan3.context()); + try (Tracer.SpanInScope nullScope = braveTracer.withSpan(null)) { + // This closes all scopes and MDC entries + then(braveTracer.currentSpan()).isNull(); + then(currentTraceContext.get()).isNull(); + } + // We have nothing to revert to since the nullScope is ignoring + // everything there was before + then(braveTracer.currentSpan()).isNull(); + then(currentTraceContext.get()).isNull(); + thenMdcEntriesAreMissing(); + nextSpan3.end(); + } + then(braveTracer.currentSpan()).isNull(); + then(currentTraceContext.get()).isNull(); + thenMdcEntriesAreMissing(); + nextSpan2.end(); + } + then(braveTracer.currentSpan()).isNull(); + then(currentTraceContext.get()).isNull(); + thenMdcEntriesAreMissing(); + nextSpan1.end(); + } + + then(currentTraceContext.get()).isNull(); + then(braveTracer.currentSpan()).isNull(); + then(MDC.getCopyOfContextMap()).isEmpty(); + } + + private static void thenMdcEntriesArePresent(org.springframework.cloud.sleuth.TraceContext traceContext) { + then(MDC.get("traceId")).isEqualTo(traceContext.traceId()); + then(MDC.get("spanId")).isEqualTo(traceContext.spanId()); + } + + private static void thenMdcEntriesAreMissing() { + then(MDC.getCopyOfContextMap()).isEmpty(); + } + + private void thenThreadLocalsGotCleared(BraveCurrentTraceContext braveCurrentTraceContext, + Tracer.SpanInScope scope) { + then(scope).isSameAs(Tracer.SpanInScope.NOOP); + then(braveTracer.currentSpan()).isNull(); + then(braveCurrentTraceContext.context()).isNull(); + } + +}