Ensure that transient ThreadContext headers with propagators survive restore

server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java

@@ -293,10 +293,25 @@ public StoredContext newStoredContext(boolean preserveResponseHeaders, Collectio
         final ThreadContextStruct newContext = threadLocal.get();
 
         return () -> {
+            ThreadContextStruct current = threadLocal.get();
+            ThreadContextStruct enrichedOriginalContext = originalContext;
+            final Map<String, Object> propagated = propagateTransients(current.transientHeaders, current.isSystemContext);
+            if (!propagated.isEmpty()) {
+                Map<String, Object> merged = new HashMap<>(originalContext.transientHeaders);
+                propagated.forEach(merged::putIfAbsent);
+                enrichedOriginalContext = new ThreadContextStruct(
+                    enrichedOriginalContext.requestHeaders,
+                    enrichedOriginalContext.responseHeaders,
+                                            merged,
+                    enrichedOriginalContext.persistentHeaders,
+                    enrichedOriginalContext.isSystemContext,
+                    enrichedOriginalContext.warningHeadersSize
+                                        );
+            }
             if (preserveResponseHeaders && threadLocal.get() != newContext) {
-                threadLocal.set(originalContext.putResponseHeaders(threadLocal.get().responseHeaders));
+                threadLocal.set(enrichedOriginalContext.putResponseHeaders(threadLocal.get().responseHeaders));
             } else {
-                threadLocal.set(originalContext);
+                threadLocal.set(enrichedOriginalContext);
             }
         };
     }

server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java

@@ -56,7 +56,7 @@ public Map<String, Object> transients(Map<String, Object> source) {
         final Map<String, Object> transients = new HashMap<>();
         if (source.containsKey(CURRENT_SPAN)) {
             final SpanReference current = (SpanReference) source.get(CURRENT_SPAN);
-            if (current != null) {
+            if (current != null && current.getSpan() != null) {
                 transients.put(CURRENT_SPAN, new SpanReference(current.getSpan()));
             }
         }

server/src/test/java/org/opensearch/common/util/concurrent/ThreadContextTests.java

@@ -851,4 +851,47 @@ protected void doRun() throws Exception {
             }
         };
     }
+
+    //We are simulating behavior that happens in Netty4HttpRequestHeaderVerifier
+    //It take a snapshot of state and stores in CONTEXT_TO_RESTORE and
+    // later tries to restore the same in SecurityFilter. Any transients added in between are lost
+    public void testPropagatedTransientsAreRestored() {
+        ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
+        final String PROPAGATED_KEY = "test_propagated_transient";
+        final Object PROPAGATED_VALUE = new Object();
+
+        // Register a propagator that declares PROPAGATED_KEY as a transient to carry across stashes.
+        threadContext.registerThreadContextStatePropagator(new ThreadContextStatePropagator() {
+            @Override
+            @SuppressWarnings("removal")
+            public Map<String, Object> transients(Map<String, Object> source) {
+                if (source.containsKey(PROPAGATED_KEY)) {
+                    return Collections.singletonMap(PROPAGATED_KEY, source.get(PROPAGATED_KEY));
+                }
+                return Collections.emptyMap();
+            }
+
+            @Override
+            @SuppressWarnings("removal")
+            public Map<String, String> headers(Map<String, Object> source) {
+                return Collections.emptyMap();
+            }
+        });
+
+        ThreadContext.StoredContext storedContext = null;
+        try(ThreadContext.StoredContext sc = threadContext.newStoredContext(false, Collections.emptyList())) {
+            //now we add something to original thread
+            // Simulate the tracing infrastructure writing CURRENT_SPAN into the stashed context.
+            storedContext = sc;
+            threadContext.putTransient(PROPAGATED_KEY, PROPAGATED_VALUE);
+        }
+        catch(Exception e) {
+            //unlikey to get exception, if we got one, test should fail
+            throw e;
+        }
+        //storedContext would have closed. Now we restore and after that, our original thread should have it
+        storedContext.restore();
+        //we should be able to find the key now
+        assertEquals(threadContext.getTransient(PROPAGATED_KEY), PROPAGATED_VALUE);
+    }
 }

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

@@ -22,7 +22,9 @@
 import org.junit.After;
 import org.junit.Before;
 
+import java.util.HashMap;
 import java.util.Optional;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -269,4 +271,12 @@ public void testSpanNotPropagatedToChildSystemThreadContext() {
         assertThat(threadContext.getTransient(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(not(nullValue())));
         assertThat(threadContextStorage.get(ThreadContextBasedTracerContextStorage.CURRENT_SPAN), is(nullValue()));
     }
+
+    public void testNullSpanWithinSpanReference() {
+        SpanReference spanReference = new SpanReference(null);
+        Map<String, Object> source = new HashMap<>();
+        source.put(ThreadContextBasedTracerContextStorage.CURRENT_SPAN, spanReference);
+        ThreadContextBasedTracerContextStorage context = (ThreadContextBasedTracerContextStorage) threadContextStorage;
+        assertTrue(context.transients(source).isEmpty());
+    }
 }