Skip to content

Commit

Permalink
Fix the issue with DefaultSpanScope restoring wrong span in the Trace…
Browse files Browse the repository at this point in the history
…rContextStorage upon detach (#11316) (#11567)

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>
(cherry picked from commit 00dd577)
  • Loading branch information
reta committed Dec 11, 2023
1 parent dc88929 commit b2cc87f
Show file tree
Hide file tree
Showing 3 changed files with 94 additions and 4 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
- Fix SuggestSearch.testSkipDuplicates by forcing refresh when indexing its test documents ([#11068](https://github.com/opensearch-project/OpenSearch/pull/11068))
- [BUG] Fix the thread context that is not properly cleared and messes up the traces ([#10873](https://github.com/opensearch-project/OpenSearch/pull/10873))
- Handle canMatchSearchAfter for frozen context scenario ([#11249](https://github.com/opensearch-project/OpenSearch/pull/11249))
- Fix the issue with DefaultSpanScope restoring wrong span in the TracerContextStorage upon detach ([#11316](https://github.com/opensearch-project/OpenSearch/issues/11316))
- Remove shadowJar from `lang-painless` module publication ([#11369](https://github.com/opensearch-project/OpenSearch/issues/11369))
- Fix remote shards balancer and remove unused variables ([#11167](https://github.com/opensearch-project/OpenSearch/pull/11167))
- Fix bug where replication lag grows post primary relocation ([#11238](https://github.com/opensearch-project/OpenSearch/pull/11238))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
class DefaultSpanScope implements SpanScope {
private final Span span;
private final SpanScope previousSpanScope;
private final Span beforeSpan;
private static final ThreadLocal<SpanScope> spanScopeThreadLocal = new ThreadLocal<>();
private final TracerContextStorage<String, Span> tracerContextStorage;

Expand All @@ -29,8 +30,14 @@ class DefaultSpanScope implements SpanScope {
* @param span span
* @param previousSpanScope before attached span scope.
*/
private DefaultSpanScope(Span span, SpanScope previousSpanScope, TracerContextStorage<String, Span> tracerContextStorage) {
private DefaultSpanScope(
Span span,
final Span beforeSpan,
SpanScope previousSpanScope,
TracerContextStorage<String, Span> tracerContextStorage
) {
this.span = Objects.requireNonNull(span);
this.beforeSpan = beforeSpan;
this.previousSpanScope = previousSpanScope;
this.tracerContextStorage = tracerContextStorage;
}
Expand All @@ -43,7 +50,8 @@ private DefaultSpanScope(Span span, SpanScope previousSpanScope, TracerContextSt
*/
public static SpanScope create(Span span, TracerContextStorage<String, Span> tracerContextStorage) {
final SpanScope beforeSpanScope = spanScopeThreadLocal.get();
SpanScope newSpanScope = new DefaultSpanScope(span, beforeSpanScope, tracerContextStorage);
final Span beforeSpan = tracerContextStorage.get(TracerContextStorage.CURRENT_SPAN);
SpanScope newSpanScope = new DefaultSpanScope(span, beforeSpan, beforeSpanScope, tracerContextStorage);
return newSpanScope;
}

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

private void detach() {
spanScopeThreadLocal.set(previousSpanScope);
if (previousSpanScope != null) {
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, previousSpanScope.getSpan());
if (beforeSpan != null) {
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, beforeSpan);
} else {
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, null);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,87 @@ public void run() {
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));
}

public void testNoThreadContextToPreserve() throws InterruptedException, ExecutionException, TimeoutException {
final Runnable r = new Runnable() {
@Override
public void run() {
assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));

final Span local1 = tracer.startSpan(SpanCreationContext.internal().name("test-local-1"));
try (SpanScope localScope = tracer.withSpanInScope(local1)) {
try (StoredContext ignored = threadContext.stashContext()) {
assertThat(local1.getParentSpan(), is(nullValue()));
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local1));
}
}

final Span local2 = tracer.startSpan(SpanCreationContext.internal().name("test-local-2"));
try (SpanScope localScope = tracer.withSpanInScope(local2)) {
try (StoredContext ignored = threadContext.stashContext()) {
assertThat(local2.getParentSpan(), is(nullValue()));
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local2));
}
}

final Span local3 = tracer.startSpan(SpanCreationContext.internal().name("test-local-3"));
try (SpanScope localScope = tracer.withSpanInScope(local3)) {
try (StoredContext ignored = threadContext.stashContext()) {
assertThat(local3.getParentSpan(), is(nullValue()));
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local3));
}
}
}
};

executorService.submit(threadContext.preserveContext(r)).get(1, TimeUnit.SECONDS);

assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));
}

public void testPreservingContextThreadContextMultipleSpans() throws InterruptedException, ExecutionException, TimeoutException {
final Span span = tracer.startSpan(SpanCreationContext.internal().name("test"));

try (SpanScope scope = tracer.withSpanInScope(span)) {
final Runnable r = new Runnable() {
@Override
public void run() {
assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(not(nullValue())));
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(span));

final Span local1 = tracer.startSpan(SpanCreationContext.internal().name("test-local-1"));
try (SpanScope localScope = tracer.withSpanInScope(local1)) {
try (StoredContext ignored = threadContext.stashContext()) {
assertThat(local1.getParentSpan(), is(span));
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local1));
}
}

final Span local2 = tracer.startSpan(SpanCreationContext.internal().name("test-local-2"));
try (SpanScope localScope = tracer.withSpanInScope(local2)) {
try (StoredContext ignored = threadContext.stashContext()) {
assertThat(local2.getParentSpan(), is(span));
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local2));
}
}

final Span local3 = tracer.startSpan(SpanCreationContext.internal().name("test-local-3"));
try (SpanScope localScope = tracer.withSpanInScope(local3)) {
try (StoredContext ignored = threadContext.stashContext()) {
assertThat(local3.getParentSpan(), is(span));
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(local3));
}
}
}
};

executorService.submit(threadContext.preserveContext(r)).get(1, TimeUnit.SECONDS);
}

assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(not(nullValue())));
assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));
}

public void testPreservingContextAndStashingThreadContext() throws InterruptedException, ExecutionException, TimeoutException {
final Span span = tracer.startSpan(SpanCreationContext.internal().name("test"));

Expand Down

0 comments on commit b2cc87f

Please sign in to comment.