From 7b055c0cbaf110df5de61a8007dbd5d71b42e6b1 Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Mon, 11 Dec 2023 09:52:19 -0500 Subject: [PATCH] Fix the issue with DefaultSpanScope restoring wrong span in the TracerContextStorage upon detach (#11316) Signed-off-by: Andriy Redko (cherry picked from commit 00dd577a0211cb7f1bcb75291ed1424cda215b64) --- CHANGELOG.md | 1 + .../telemetry/tracing/DefaultSpanScope.java | 16 +++- ...ContextBasedTracerContextStorageTests.java | 81 +++++++++++++++++++ 3 files changed, 94 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a5b811d62abc..fd67301a052ed 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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)) diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java index decbf49f795c4..93600da510977 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java +++ b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java @@ -21,6 +21,7 @@ class DefaultSpanScope implements SpanScope { private final Span span; private final SpanScope previousSpanScope; + private final Span beforeSpan; private static final ThreadLocal spanScopeThreadLocal = new ThreadLocal<>(); private final TracerContextStorage tracerContextStorage; @@ -29,8 +30,14 @@ class DefaultSpanScope implements SpanScope { * @param span span * @param previousSpanScope before attached span scope. */ - private DefaultSpanScope(Span span, SpanScope previousSpanScope, TracerContextStorage tracerContextStorage) { + private DefaultSpanScope( + Span span, + final Span beforeSpan, + SpanScope previousSpanScope, + TracerContextStorage tracerContextStorage + ) { this.span = Objects.requireNonNull(span); + this.beforeSpan = beforeSpan; this.previousSpanScope = previousSpanScope; this.tracerContextStorage = tracerContextStorage; } @@ -43,7 +50,8 @@ private DefaultSpanScope(Span span, SpanScope previousSpanScope, TracerContextSt */ public static SpanScope create(Span span, TracerContextStorage 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; } @@ -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); } diff --git a/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java b/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java index 3a98a67b53920..ee816aa5f596d 100644 --- a/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java +++ b/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java @@ -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"));