Skip to content

Commit 3d19830

Browse files
marcingrzejszczakmarcingrzejszczak
authored
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
1 parent d05059a commit 3d19830

File tree

6 files changed

+344
-20
lines changed

6 files changed

+344
-20
lines changed

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

+7
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,13 @@ public interface CurrentTraceContext {
9494
*/
9595
interface Scope extends Closeable {
9696

97+
/**
98+
* Noop instance.
99+
*/
100+
Scope NOOP = () -> {
101+
102+
};
103+
97104
@Override
98105
void close();
99106

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

+7
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,13 @@ default CurrentTraceContext currentTraceContext() {
170170
*/
171171
interface SpanInScope extends Closeable {
172172

173+
/**
174+
* Noop instance.
175+
*/
176+
SpanInScope NOOP = () -> {
177+
178+
};
179+
173180
@Override
174181
void close();
175182

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

+51-14
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2013-2021 the original author or authors.
2+
* Copyright 2013-2023 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -20,17 +20,15 @@
2020
import java.util.concurrent.Executor;
2121
import java.util.concurrent.ExecutorService;
2222

23+
import brave.propagation.ThreadLocalCurrentTraceContext;
24+
2325
import org.springframework.cloud.sleuth.CurrentTraceContext;
2426
import org.springframework.cloud.sleuth.TraceContext;
2527

26-
/**
27-
* Brave implementation of a {@link CurrentTraceContext}.
28-
*
29-
* @author Marcin Grzejszczak
30-
* @since 3.0.0
31-
*/
3228
public class BraveCurrentTraceContext implements CurrentTraceContext {
3329

30+
final ThreadLocal<Scope> scopes = new ThreadLocal<>();
31+
3432
final brave.propagation.CurrentTraceContext delegate;
3533

3634
public BraveCurrentTraceContext(brave.propagation.CurrentTraceContext delegate) {
@@ -40,20 +38,36 @@ public BraveCurrentTraceContext(brave.propagation.CurrentTraceContext delegate)
4038
@Override
4139
public TraceContext context() {
4240
brave.propagation.TraceContext context = this.delegate.get();
41+
return context == null ? null : new BraveTraceContext(context);
42+
}
43+
44+
@Override
45+
public CurrentTraceContext.Scope newScope(TraceContext context) {
4346
if (context == null) {
44-
return null;
47+
clearScopes();
48+
return Scope.NOOP;
4549
}
46-
return new BraveTraceContext(context);
50+
return new RevertingScope(this, new BraveScope(this.delegate.newScope(BraveTraceContext.toBrave(context))));
4751
}
4852

4953
@Override
50-
public Scope newScope(TraceContext context) {
51-
return new BraveScope(this.delegate.newScope(BraveTraceContext.toBrave(context)));
54+
public CurrentTraceContext.Scope maybeScope(TraceContext context) {
55+
if (context == null) {
56+
clearScopes();
57+
return Scope.NOOP;
58+
}
59+
return new RevertingScope(this, new BraveScope(this.delegate.maybeScope(BraveTraceContext.toBrave(context))));
5260
}
5361

54-
@Override
55-
public Scope maybeScope(TraceContext context) {
56-
return new BraveScope(this.delegate.maybeScope(BraveTraceContext.toBrave(context)));
62+
private void clearScopes() {
63+
Scope current = this.scopes.get();
64+
while (current != null) {
65+
current.close();
66+
current = this.scopes.get();
67+
}
68+
if (this.delegate instanceof ThreadLocalCurrentTraceContext) {
69+
((ThreadLocalCurrentTraceContext) this.delegate).clear();
70+
}
5771
}
5872

5973
@Override
@@ -86,6 +100,29 @@ static CurrentTraceContext fromBrave(brave.propagation.CurrentTraceContext conte
86100

87101
}
88102

103+
class RevertingScope implements CurrentTraceContext.Scope {
104+
105+
private final BraveCurrentTraceContext currentTraceContext;
106+
107+
private final CurrentTraceContext.Scope previous;
108+
109+
private final CurrentTraceContext.Scope current;
110+
111+
RevertingScope(BraveCurrentTraceContext currentTraceContext, CurrentTraceContext.Scope current) {
112+
this.currentTraceContext = currentTraceContext;
113+
this.previous = this.currentTraceContext.scopes.get();
114+
this.current = current;
115+
this.currentTraceContext.scopes.set(this);
116+
}
117+
118+
@Override
119+
public void close() {
120+
this.current.close();
121+
this.currentTraceContext.scopes.set(this.previous);
122+
}
123+
124+
}
125+
89126
class BraveScope implements CurrentTraceContext.Scope {
90127

91128
private final brave.propagation.CurrentTraceContext.Scope delegate;

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

+15-6
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
package org.springframework.cloud.sleuth.brave.bridge;
1818

19+
import java.io.Closeable;
20+
import java.io.IOException;
1921
import java.util.Map;
2022

2123
import brave.propagation.TraceContextOrSamplingFlags;
@@ -27,7 +29,6 @@
2729
import org.springframework.cloud.sleuth.SpanCustomizer;
2830
import org.springframework.cloud.sleuth.TraceContext;
2931
import org.springframework.cloud.sleuth.Tracer;
30-
import org.springframework.cloud.sleuth.docs.AssertingSpan;
3132

3233
/**
3334
* Brave implementation of a {@link Tracer}.
@@ -70,8 +71,11 @@ public Span nextSpan(Span parent) {
7071

7172
@Override
7273
public SpanInScope withSpan(Span span) {
73-
return new BraveSpanInScope(
74-
tracer.withSpanInScope(span == null ? null : ((BraveSpan) AssertingSpan.unwrap(span)).delegate));
74+
if (span == null) {
75+
currentTraceContext.maybeScope(null);
76+
return SpanInScope.NOOP;
77+
}
78+
return new BraveSpanInScope(currentTraceContext.maybeScope(span.context()));
7579
}
7680

7781
@Override
@@ -142,15 +146,20 @@ public CurrentTraceContext currentTraceContext() {
142146

143147
class BraveSpanInScope implements Tracer.SpanInScope {
144148

145-
final brave.Tracer.SpanInScope delegate;
149+
final Closeable delegate;
146150

147-
BraveSpanInScope(brave.Tracer.SpanInScope delegate) {
151+
BraveSpanInScope(Closeable delegate) {
148152
this.delegate = delegate;
149153
}
150154

151155
@Override
152156
public void close() {
153-
this.delegate.close();
157+
try {
158+
this.delegate.close();
159+
}
160+
catch (IOException e) {
161+
throw new RuntimeException(e);
162+
}
154163
}
155164

156165
}

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

+152
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
/*
2+
* Copyright 2013-2023 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.cloud.sleuth.brave.bridge;
18+
19+
import brave.context.slf4j.MDCScopeDecorator;
20+
import brave.propagation.ThreadLocalCurrentTraceContext;
21+
import brave.propagation.TraceContext;
22+
import org.junit.jupiter.api.Test;
23+
import org.slf4j.MDC;
24+
25+
import org.springframework.cloud.sleuth.CurrentTraceContext;
26+
27+
import static org.assertj.core.api.BDDAssertions.then;
28+
29+
class BraveCurrentTraceContextTests {
30+
31+
ThreadLocalCurrentTraceContext currentTraceContext = ThreadLocalCurrentTraceContext.newBuilder()
32+
.addScopeDecorator(MDCScopeDecorator.newBuilder().build()).build();
33+
34+
@Test
35+
void should_clear_any_thread_locals_and_scopes_when_null_context_passed_to_new_scope() {
36+
BraveCurrentTraceContext braveCurrentTraceContext = new BraveCurrentTraceContext(currentTraceContext);
37+
brave.propagation.CurrentTraceContext.Scope newScope = currentTraceContext
38+
.newScope(TraceContext.newBuilder().traceId(12345678).spanId(12345678).build());
39+
then(currentTraceContext.get()).isNotNull();
40+
41+
CurrentTraceContext.Scope scope = braveCurrentTraceContext.newScope(null);
42+
43+
thenThreadLocalsGotCleared(braveCurrentTraceContext, scope);
44+
newScope.close();
45+
thenThreadLocalsGotCleared(braveCurrentTraceContext, scope);
46+
}
47+
48+
@Test
49+
void should_clear_any_thread_locals_and_scopes_when_null_context_passed_to_new_scope_with_nested_scopes() {
50+
BraveCurrentTraceContext braveCurrentTraceContext = new BraveCurrentTraceContext(currentTraceContext);
51+
52+
try (CurrentTraceContext.Scope scope1 = braveCurrentTraceContext.newScope(
53+
BraveTraceContext.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345670).build()))) {
54+
then(currentTraceContext.get()).isNotNull();
55+
thenMdcEntriesArePresent();
56+
try (CurrentTraceContext.Scope scope2 = braveCurrentTraceContext.newScope(BraveTraceContext
57+
.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345671).build()))) {
58+
then(currentTraceContext.get()).isNotNull();
59+
thenMdcEntriesArePresent();
60+
try (CurrentTraceContext.Scope scope3 = braveCurrentTraceContext.newScope(BraveTraceContext
61+
.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345672).build()))) {
62+
then(currentTraceContext.get()).isNotNull();
63+
thenMdcEntriesArePresent();
64+
try (CurrentTraceContext.Scope nullScope = braveCurrentTraceContext.newScope(null)) {
65+
// This closes all scopes and MDC entries
66+
then(currentTraceContext.get()).isNull();
67+
}
68+
// We have nothing to revert to since the nullScope is ignoring
69+
// everything there was before
70+
then(currentTraceContext.get()).isNull();
71+
thenMdcEntriesAreMissing();
72+
}
73+
then(currentTraceContext.get()).isNull();
74+
thenMdcEntriesAreMissing();
75+
}
76+
then(currentTraceContext.get()).isNull();
77+
thenMdcEntriesAreMissing();
78+
}
79+
80+
then(currentTraceContext.get()).isNull();
81+
then(braveCurrentTraceContext.scopes.get()).isNull();
82+
then(MDC.getCopyOfContextMap()).isEmpty();
83+
}
84+
85+
@Test
86+
void should_clear_any_thread_locals_and_scopes_when_null_context_passed_to_maybe_scope_with_nested_scopes() {
87+
BraveCurrentTraceContext braveCurrentTraceContext = new BraveCurrentTraceContext(currentTraceContext);
88+
89+
try (CurrentTraceContext.Scope scope1 = braveCurrentTraceContext.maybeScope(
90+
BraveTraceContext.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345670).build()))) {
91+
then(currentTraceContext.get()).isNotNull();
92+
thenMdcEntriesArePresent();
93+
try (CurrentTraceContext.Scope scope2 = braveCurrentTraceContext.maybeScope(BraveTraceContext
94+
.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345671).build()))) {
95+
then(currentTraceContext.get()).isNotNull();
96+
thenMdcEntriesArePresent();
97+
try (CurrentTraceContext.Scope scope3 = braveCurrentTraceContext.maybeScope(BraveTraceContext
98+
.fromBrave(TraceContext.newBuilder().traceId(12345678).spanId(12345672).build()))) {
99+
then(currentTraceContext.get()).isNotNull();
100+
thenMdcEntriesArePresent();
101+
try (CurrentTraceContext.Scope nullScope = braveCurrentTraceContext.maybeScope(null)) {
102+
// This closes all scopes and MDC entries
103+
then(currentTraceContext.get()).isNull();
104+
}
105+
// We have nothing to revert to since the nullScope is ignoring
106+
// everything there was before
107+
then(currentTraceContext.get()).isNull();
108+
thenMdcEntriesAreMissing();
109+
}
110+
then(currentTraceContext.get()).isNull();
111+
thenMdcEntriesAreMissing();
112+
}
113+
then(currentTraceContext.get()).isNull();
114+
thenMdcEntriesAreMissing();
115+
}
116+
117+
then(currentTraceContext.get()).isNull();
118+
then(braveCurrentTraceContext.scopes.get()).isNull();
119+
then(MDC.getCopyOfContextMap()).isEmpty();
120+
}
121+
122+
private static void thenMdcEntriesArePresent() {
123+
then(MDC.get("traceId")).isEqualTo("0000000000bc614e");
124+
then(MDC.get("spanId")).isNotEmpty();
125+
}
126+
127+
private static void thenMdcEntriesAreMissing() {
128+
then(MDC.getCopyOfContextMap()).isEmpty();
129+
}
130+
131+
@Test
132+
void should_clear_any_thread_locals_and_scopes_when_null_context_passed_to_maybe_scope() {
133+
BraveCurrentTraceContext braveCurrentTraceContext = new BraveCurrentTraceContext(currentTraceContext);
134+
brave.propagation.CurrentTraceContext.Scope maybeScope = currentTraceContext
135+
.newScope(TraceContext.newBuilder().traceId(12345678).spanId(12345678).build());
136+
then(currentTraceContext.get()).isNotNull();
137+
138+
CurrentTraceContext.Scope scope = braveCurrentTraceContext.maybeScope(null);
139+
140+
thenThreadLocalsGotCleared(braveCurrentTraceContext, scope);
141+
maybeScope.close();
142+
thenThreadLocalsGotCleared(braveCurrentTraceContext, scope);
143+
}
144+
145+
private void thenThreadLocalsGotCleared(BraveCurrentTraceContext braveCurrentTraceContext,
146+
CurrentTraceContext.Scope scope) {
147+
then(scope).isSameAs(CurrentTraceContext.Scope.NOOP);
148+
then(currentTraceContext.get()).isNull();
149+
then(braveCurrentTraceContext.scopes.get()).isNull();
150+
}
151+
152+
}

0 commit comments

Comments
 (0)