Skip to content

Commit b2cc87f

Browse files
retareta
authored
Fix the issue with DefaultSpanScope restoring wrong span in the TracerContextStorage upon detach (opensearch-project#11316) (opensearch-project#11567)
Signed-off-by: Andriy Redko <andriy.redko@aiven.io> (cherry picked from commit 00dd577)
1 parent dc88929 commit b2cc87f

File tree

3 files changed

+94
-4
lines changed

3 files changed

+94
-4
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
9494
- Fix SuggestSearch.testSkipDuplicates by forcing refresh when indexing its test documents ([#11068](https://github.com/opensearch-project/OpenSearch/pull/11068))
9595
- [BUG] Fix the thread context that is not properly cleared and messes up the traces ([#10873](https://github.com/opensearch-project/OpenSearch/pull/10873))
9696
- Handle canMatchSearchAfter for frozen context scenario ([#11249](https://github.com/opensearch-project/OpenSearch/pull/11249))
97+
- Fix the issue with DefaultSpanScope restoring wrong span in the TracerContextStorage upon detach ([#11316](https://github.com/opensearch-project/OpenSearch/issues/11316))
9798
- Remove shadowJar from `lang-painless` module publication ([#11369](https://github.com/opensearch-project/OpenSearch/issues/11369))
9899
- Fix remote shards balancer and remove unused variables ([#11167](https://github.com/opensearch-project/OpenSearch/pull/11167))
99100
- Fix bug where replication lag grows post primary relocation ([#11238](https://github.com/opensearch-project/OpenSearch/pull/11238))

‎libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
class DefaultSpanScope implements SpanScope {
2222
private final Span span;
2323
private final SpanScope previousSpanScope;
24+
private final Span beforeSpan;
2425
private static final ThreadLocal<SpanScope> spanScopeThreadLocal = new ThreadLocal<>();
2526
private final TracerContextStorage<String, Span> tracerContextStorage;
2627

@@ -29,8 +30,14 @@ class DefaultSpanScope implements SpanScope {
2930
* @param span span
3031
* @param previousSpanScope before attached span scope.
3132
*/
32-
private DefaultSpanScope(Span span, SpanScope previousSpanScope, TracerContextStorage<String, Span> tracerContextStorage) {
33+
private DefaultSpanScope(
34+
Span span,
35+
final Span beforeSpan,
36+
SpanScope previousSpanScope,
37+
TracerContextStorage<String, Span> tracerContextStorage
38+
) {
3339
this.span = Objects.requireNonNull(span);
40+
this.beforeSpan = beforeSpan;
3441
this.previousSpanScope = previousSpanScope;
3542
this.tracerContextStorage = tracerContextStorage;
3643
}
@@ -43,7 +50,8 @@ private DefaultSpanScope(Span span, SpanScope previousSpanScope, TracerContextSt
4350
*/
4451
public static SpanScope create(Span span, TracerContextStorage<String, Span> tracerContextStorage) {
4552
final SpanScope beforeSpanScope = spanScopeThreadLocal.get();
46-
SpanScope newSpanScope = new DefaultSpanScope(span, beforeSpanScope, tracerContextStorage);
53+
final Span beforeSpan = tracerContextStorage.get(TracerContextStorage.CURRENT_SPAN);
54+
SpanScope newSpanScope = new DefaultSpanScope(span, beforeSpan, beforeSpanScope, tracerContextStorage);
4755
return newSpanScope;
4856
}
4957

@@ -61,8 +69,8 @@ public SpanScope attach() {
6169

6270
private void detach() {
6371
spanScopeThreadLocal.set(previousSpanScope);
64-
if (previousSpanScope != null) {
65-
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, previousSpanScope.getSpan());
72+
if (beforeSpan != null) {
73+
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, beforeSpan);
6674
} else {
6775
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, null);
6876
}

‎server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,87 @@ public void run() {
145145
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));
146146
}
147147

148+
public void testNoThreadContextToPreserve() throws InterruptedException, ExecutionException, TimeoutException {
149+
final Runnable r = new Runnable() {
150+
@Override
151+
public void run() {
152+
assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));
153+
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));
154+
155+
final Span local1 = tracer.startSpan(SpanCreationContext.internal().name("test-local-1"));
156+
try (SpanScope localScope = tracer.withSpanInScope(local1)) {
157+
try (StoredContext ignored = threadContext.stashContext()) {
158+
assertThat(local1.getParentSpan(), is(nullValue()));
159+
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local1));
160+
}
161+
}
162+
163+
final Span local2 = tracer.startSpan(SpanCreationContext.internal().name("test-local-2"));
164+
try (SpanScope localScope = tracer.withSpanInScope(local2)) {
165+
try (StoredContext ignored = threadContext.stashContext()) {
166+
assertThat(local2.getParentSpan(), is(nullValue()));
167+
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local2));
168+
}
169+
}
170+
171+
final Span local3 = tracer.startSpan(SpanCreationContext.internal().name("test-local-3"));
172+
try (SpanScope localScope = tracer.withSpanInScope(local3)) {
173+
try (StoredContext ignored = threadContext.stashContext()) {
174+
assertThat(local3.getParentSpan(), is(nullValue()));
175+
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local3));
176+
}
177+
}
178+
}
179+
};
180+
181+
executorService.submit(threadContext.preserveContext(r)).get(1, TimeUnit.SECONDS);
182+
183+
assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));
184+
}
185+
186+
public void testPreservingContextThreadContextMultipleSpans() throws InterruptedException, ExecutionException, TimeoutException {
187+
final Span span = tracer.startSpan(SpanCreationContext.internal().name("test"));
188+
189+
try (SpanScope scope = tracer.withSpanInScope(span)) {
190+
final Runnable r = new Runnable() {
191+
@Override
192+
public void run() {
193+
assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(not(nullValue())));
194+
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(span));
195+
196+
final Span local1 = tracer.startSpan(SpanCreationContext.internal().name("test-local-1"));
197+
try (SpanScope localScope = tracer.withSpanInScope(local1)) {
198+
try (StoredContext ignored = threadContext.stashContext()) {
199+
assertThat(local1.getParentSpan(), is(span));
200+
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local1));
201+
}
202+
}
203+
204+
final Span local2 = tracer.startSpan(SpanCreationContext.internal().name("test-local-2"));
205+
try (SpanScope localScope = tracer.withSpanInScope(local2)) {
206+
try (StoredContext ignored = threadContext.stashContext()) {
207+
assertThat(local2.getParentSpan(), is(span));
208+
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local2));
209+
}
210+
}
211+
212+
final Span local3 = tracer.startSpan(SpanCreationContext.internal().name("test-local-3"));
213+
try (SpanScope localScope = tracer.withSpanInScope(local3)) {
214+
try (StoredContext ignored = threadContext.stashContext()) {
215+
assertThat(local3.getParentSpan(), is(span));
216+
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local3));
217+
}
218+
}
219+
}
220+
};
221+
222+
executorService.submit(threadContext.preserveContext(r)).get(1, TimeUnit.SECONDS);
223+
}
224+
225+
assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(not(nullValue())));
226+
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));
227+
}
228+
148229
public void testPreservingContextAndStashingThreadContext() throws InterruptedException, ExecutionException, TimeoutException {
149230
final Span span = tracer.startSpan(SpanCreationContext.internal().name("test"));
150231

0 commit comments

Comments
 (0)