[Backport] [2.x] Fix the issue with DefaultSpanScope restoring wrong span in the TracerContextStorage upon detach (#11316)

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))

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<SpanScope> spanScopeThreadLocal = new ThreadLocal<>();
     private final TracerContextStorage<String, Span> 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<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;
     }
@@ -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;
     }
 
@@ -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);
         }

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"));