Fix the issue with DefaultSpanScope restoring wrong span in the Trace…

CHANGELOG.md

@@ -158,6 +158,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Delegating CachingWeightWrapper#count to internal weight object ([#10543](https://github.com/opensearch-project/OpenSearch/pull/10543))
 - Fix per request latency last phase not tracked ([#10934](https://github.com/opensearch-project/OpenSearch/pull/10934))
 - Fix for stuck update action in a bulk with `retry_on_conflict` property ([#11152](https://github.com/opensearch-project/OpenSearch/issues/11152))
+- Fix the issue with DefaultSpanScope restoring wrong span in the TracerContextStorage upon detach ([#11316](https://github.com/opensearch-project/OpenSearch/issues/11316))
 
 ### Security
 

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