Fix FlightOutboundHandler clearing caller's ThreadContext

[Hailong-am]

Description

FlightOutboundHandler uses stashContext() to capture the thread context before dispatching work to a plain ExecutorService. However, stashContext() clears the calling thread's context as a side effect. Since the context is only restored on the executor thread, the caller loses its ThreadContext (security headers, request metadata, etc.) after calling sendResponseBatch(), completeStream(), or sendErrorResponse().

This replaces the manual stashContext() / restore() pattern with ThreadContext.preserveContext(), which is the idiomatic OpenSearch mechanism for cross-thread context propagation. It:

• Captures context at wrap time via newStoredContext(false) without clearing the caller
• At run time, stashes the executor thread's context, restores the captured context, runs the task, then cleans up

Changes

FlightOutboundHandler.java:

• Replace stashContext() with preserveContext() wrapping the executor lambdas in sendResponseBatch(), completeStream(), and sendErrorResponse()
• Remove StoredContext field from BatchTask record
• Remove manual storedContext.restore() calls from processBatchTask(), processCompleteTask(), processErrorTask()

FlightOutboundHandlerTests.java (new):

• Unit tests verifying caller's ThreadContext is preserved after each method call
• Tests context propagation to executor thread via preserveContext()
• Tests context preservation across multiple batch sends

FlightOutboundHandlerContextPropagationTests.java (new):

• Integration tests using real FlightTransport and StreamTransportService
• End-to-end verification of context propagation through stream response batches
• Tests context preservation through error response paths

FlightTransportChannelTests.java:

• Updated BatchTask constructor call to match new signature (removed storedContext parameter)

Related Issues

Resolves #21166

#19403

Check List

• New functionality includes testing
• New functionality has been documented
• API changes companion pull request - N/A
• Commits are signed per the DCO using --signoff

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certification of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit b0dcda6)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 Multiple PR themes Sub-PR theme: Replace stashContext with preserveContext in FlightOutboundHandler Relevant files: plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportChannelTests.java Sub-PR theme: Add unit and integration tests for ThreadContext preservation Relevant files: plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerContextPropagationTests.java

⚡ Recommended focus areas for review Null StoredContext The BatchTask record now always receives null as the error field in sendResponseBatch and completeStream, and null as the last parameter in all three methods. The null passed as the last constructor argument (previously storedContext, now error in sendResponseBatch/completeStream) should be verified to not cause NullPointerExceptions in processErrorTask or processBatchTask when task.error() is accessed. null Missing Context Assertion In testSendResponseBatchPropagatesContextToExecutorThread, the test verifies the latch completes but never actually asserts that the captured header value on the executor thread equals HEADER_VALUE. The test name implies context propagation is verified, but no assertion on the propagated header value is made inside the executor lambda. public void testSendResponseBatchPropagatesContextToExecutorThread() throws Exception { ThreadContext threadContext = threadPool.getThreadContext(); threadContext.putHeader(HEADER_KEY, HEADER_VALUE); CountDownLatch latch = new CountDownLatch(1); // Use a mock executor that runs the preserveContext-wrapped runnable ExecutorService mockExecutor = mock(ExecutorService.class); doAnswer(invocation -> { Runnable command = invocation.getArgument(0); executor.execute(() -> { command.run(); // After the preserveContext wrapper runs, capture the header // The wrapper stashes the executor thread context, restores caller's, runs, then restores executor's latch.countDown(); }); return null; }).when(mockExecutor).execute(any(Runnable.class)); when(mockFlightChannel.getExecutor()).thenReturn(mockExecutor); handler.sendResponseBatch( Version.CURRENT, Collections.emptySet(), mockFlightChannel, mock(FlightTransportChannel.class), 1L, "test-action", mock(TransportResponse.class), false, false ); assertTrue("Executor task should complete", latch.await(5, TimeUnit.SECONDS)); } Race Condition In testThreadContextPropagatedThroughStreamResponseBatch, the assertions on capturedHeaderOnServer and responseCount are made after handlerLatch.await(), but the request handler runs on ThreadPool.Names.SAME. If the handler completes before the latch is counted down, there could be ordering issues. Additionally, the capturedHeaderOnServer is set after sendResponseBatch but the latch is only released after handleStreamResponse completes — it's worth verifying there's no race between the server handler thread and the client response handler thread. public void testThreadContextPropagatedThroughStreamResponseBatch() throws InterruptedException { String action = "internal:test/context-propagation"; CountDownLatch handlerLatch = new CountDownLatch(1); AtomicInteger responseCount = new AtomicInteger(0); AtomicReference<Exception> handlerException = new AtomicReference<>(); AtomicReference<String> capturedHeaderOnServer = new AtomicReference<>(); streamTransportService.registerRequestHandler(action, ThreadPool.Names.SAME, TestRequest::new, (request, channel, task) -> { try { // Set a header in the request handler's thread context threadPool.getThreadContext().putHeader(CONTEXT_HEADER, CONTEXT_VALUE); // Verify context is set before sending batch assertEquals(CONTEXT_VALUE, threadPool.getThreadContext().getHeader(CONTEXT_HEADER)); channel.sendResponseBatch(new TestResponse("Response 1")); // Verify the caller's context is preserved after sendResponseBatch capturedHeaderOnServer.set(threadPool.getThreadContext().getHeader(CONTEXT_HEADER)); channel.sendResponseBatch(new TestResponse("Response 2")); // Verify context is still preserved after second batch assertEquals(CONTEXT_VALUE, threadPool.getThreadContext().getHeader(CONTEXT_HEADER)); channel.completeStream(); // Verify context is still preserved after completeStream assertEquals(CONTEXT_VALUE, threadPool.getThreadContext().getHeader(CONTEXT_HEADER)); } catch (Exception e) { try { channel.sendResponse(e); } catch (IOException ignored) {} } }); TestRequest testRequest = new TestRequest(); TransportRequestOptions options = TransportRequestOptions.builder().withType(TransportRequestOptions.Type.STREAM).build(); StreamTransportResponseHandler<TestResponse> responseHandler = new StreamTransportResponseHandler<TestResponse>() { @Override public void handleStreamResponse(StreamTransportResponse<TestResponse> streamResponse) { try (streamResponse) { try { while (streamResponse.nextResponse() != null) { responseCount.incrementAndGet(); } } catch (Exception e) { handlerException.set(e); } } catch (Exception ignored) {} finally { handlerLatch.countDown(); } } @Override public void handleException(TransportException exp) { handlerException.set(exp); handlerLatch.countDown(); } @Override public String executor() { return ThreadPool.Names.SAME; } @Override public TestResponse read(StreamInput in) throws IOException { return new TestResponse(in); } }; streamTransportService.sendRequest(remoteNode, action, testRequest, options, responseHandler); assertTrue(handlerLatch.await(TIMEOUT_SEC, TimeUnit.SECONDS)); assertEquals(2, responseCount.get()); assertNull("No exception expected but got: " + handlerException.get(), handlerException.get()); assertEquals( "Thread context header should be preserved on the server handler thread after sendResponseBatch", CONTEXT_VALUE, capturedHeaderOnServer.get() ); }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to b0dcda6
Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Await async task before asserting thread context The test sets up a latch that counts down when onResponseSent is called, but never waits on the latch before asserting the thread context. Since the executor runs asynchronously, the assertion may execute before the executor task completes, making the test unreliable. Add latch.await(5, TimeUnit.SECONDS) before the assertion to ensure the async task has finished. plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java [71-99] doAnswer(invocation -> { latch.countDown(); return null; }).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); handler.sendResponseBatch( - ... + Version.CURRENT, + Collections.emptySet(), + mockFlightChannel, + mock(FlightTransportChannel.class), + 1L, + "test-action", + mock(TransportResponse.class), + false, + false ); + +assertTrue("Executor task should complete", latch.await(5, TimeUnit.SECONDS)); // Verify the caller's thread context is NOT cleared assertEquals( "Caller's thread context should be preserved after sendResponseBatch", HEADER_VALUE, threadContext.getHeader(HEADER_KEY) ); Suggestion importance[1-10]: 7 __ Why: The test creates a latch and sets up mockListener.onResponseSent to count it down, but never calls latch.await() before asserting the thread context. Since sendResponseBatch submits work asynchronously to the executor, the assertion may run before the executor task completes, making the test unreliable. Adding latch.await(5, TimeUnit.SECONDS) before the assertion would make the test deterministic.	Medium
	Synchronize before asserting async context preservation The completeStream method submits work to an async executor, so the assertion immediately after the call may race with the executor task. Without synchronization (e.g., a latch triggered by mockListener.onResponseSent), the test does not reliably verify behavior after the async task completes. Add a latch and wait for the executor task to finish before asserting. plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java [101-115] public void testCompleteStreamPreservesCallerThreadContext() throws Exception { ThreadContext threadContext = threadPool.getThreadContext(); threadContext.putHeader(HEADER_KEY, HEADER_VALUE); + + CountDownLatch latch = new CountDownLatch(1); + doAnswer(invocation -> { + latch.countDown(); + return null; + }).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); handler.completeStream( Version.CURRENT, Collections.emptySet(), mockFlightChannel, mock(FlightTransportChannel.class), 1L, "test-action" ); + assertTrue("Executor task should complete", latch.await(5, TimeUnit.SECONDS)); assertEquals("Caller's thread context should be preserved after completeStream", HEADER_VALUE, threadContext.getHeader(HEADER_KEY)); } Suggestion importance[1-10]: 6 __ Why: The completeStream method submits work asynchronously to the executor, so the assertEquals immediately after the call may race with the executor task. Adding a latch triggered by mockListener.onResponseSent and awaiting it before the assertion would make the test reliable and deterministic.	Low

---

Previous suggestions

Suggestions up to commit cb38601

Category	Suggestion	Impact
Possible issue	Await async task completion before asserting context The latch is created and the mockListener is set up to count it down, but latch.await() is never called before the assertion. This means the test does not wait for the async executor task to complete before asserting, which could lead to a race condition where the executor task (which uses preserveContext) hasn't run yet and potentially interfered with the context. Add latch.await(5, TimeUnit.SECONDS) before the assertion to ensure the async task has completed. plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java [70-98] -public void testSendResponseBatchPreservesCallerThreadContext() throws Exception { - ThreadContext threadContext = threadPool.getThreadContext(); - threadContext.putHeader(HEADER_KEY, HEADER_VALUE); +handler.sendResponseBatch( + Version.CURRENT, + Collections.emptySet(), + mockFlightChannel, + mock(FlightTransportChannel.class), + 1L, + "test-action", + mock(TransportResponse.class), + false, + false +); - CountDownLatch latch = new CountDownLatch(1); - doAnswer(invocation -> { - latch.countDown(); - return null; - }).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); +assertTrue("Executor task should complete", latch.await(5, TimeUnit.SECONDS)); - handler.sendResponseBatch( - ... - ); +// Verify the caller's thread context is NOT cleared +assertEquals( + "Caller's thread context should be preserved after sendResponseBatch", + HEADER_VALUE, + threadContext.getHeader(HEADER_KEY) +); - // Verify the caller's thread context is NOT cleared - assertEquals( - "Caller's thread context should be preserved after sendResponseBatch", - HEADER_VALUE, - threadContext.getHeader(HEADER_KEY) - ); -} - Suggestion importance[1-10]: 7 __ Why: The latch is set up but latch.await() is never called before the assertion, creating a potential race condition where the async executor task could interfere with the thread context before the assertion runs. This is a real correctness issue in the test.	Medium
General	Add assertion for context propagation in executor thread The test testSendResponseBatchPropagatesContextToExecutorThread does not actually assert that the context header is visible inside the executor thread. The latch.countDown() is called after command.run() but outside the runnable, so there is no capture or assertion of the propagated header value inside the executor thread. Add an AtomicReference to capture the header value inside command.run() and assert it equals HEADER_VALUE after the latch awaits. plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java [141-174] public void testSendResponseBatchPropagatesContextToExecutorThread() throws Exception { ThreadContext threadContext = threadPool.getThreadContext(); threadContext.putHeader(HEADER_KEY, HEADER_VALUE); CountDownLatch latch = new CountDownLatch(1); + AtomicReference<String> capturedHeader = new AtomicReference<>(); - // Use a mock executor that runs the preserveContext-wrapped runnable ExecutorService mockExecutor = mock(ExecutorService.class); doAnswer(invocation -> { Runnable command = invocation.getArgument(0); executor.execute(() -> { command.run(); - // After the preserveContext wrapper runs, capture the header - // The wrapper stashes the executor thread context, restores caller's, runs, then restores executor's + capturedHeader.set(threadPool.getThreadContext().getHeader(HEADER_KEY)); latch.countDown(); }); return null; }).when(mockExecutor).execute(any(Runnable.class)); when(mockFlightChannel.getExecutor()).thenReturn(mockExecutor); - ... + + handler.sendResponseBatch( + Version.CURRENT, + Collections.emptySet(), + mockFlightChannel, + mock(FlightTransportChannel.class), + 1L, + "test-action", + mock(TransportResponse.class), + false, + false + ); + assertTrue("Executor task should complete", latch.await(5, TimeUnit.SECONDS)); + assertEquals("Context header should be propagated to executor thread", HEADER_VALUE, capturedHeader.get()); } Suggestion importance[1-10]: 6 __ Why: The test testSendResponseBatchPropagatesContextToExecutorThread only verifies the task completes but never asserts that the context header was actually propagated to the executor thread, making it an incomplete test for the stated purpose. Adding a capturedHeader assertion would make the test meaningful.	Low

[Hailong-am]

@rishabhmaurya would you mind to help to review this fix?

[github-actions]

❌ Gradle check result for cb38601: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit b0dcda6

[github-actions]

✅ Gradle check result for b0dcda6: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.23%. Comparing base (9bfcc1d) to head (b0dcda6).
⚠️ Report is 19 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21167      +/-   ##
============================================
+ Coverage     73.10%   73.23%   +0.12%     
- Complexity    73213    73255      +42     
============================================
  Files          5968     5968              
  Lines        334539   334539              
  Branches      48174    48170       -4     
============================================
+ Hits         244572   244997     +425     
+ Misses        70421    69943     -478     
- Partials      19546    19599      +53     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Hailong-am]

Breaking change brought in by PR https://github.com/opensearch-project/OpenSearch/pull/21039/changes#diff-e13f61d0a17c4541e43c1ade6f9b452696340e76d728970322fa3b85dd9b5aaa, Not related to my change.

[rishabhmaurya]

logically it makes sense. I will check the code in details later in the day.
I'm curious what is your use case?

[Hailong-am]

logically it makes sense. I will check the code in details later in the day. I'm curious what is your use case?

@rishabhmaurya thanks. My case is we are building streaming experience for chat and sending back tool call result to user in the middle of whole chat flow, so the flow is tool A called and send back toolA's result and continue reasoning and call other tools until we got a final answer. To achieve sending tool result back we leverage FlightOutBoundHandler and we still want to keep the thread context as it contains information which need to tool call.

[rishabhmaurya]

@Hailong-am thanks for the explanation, makes a lot of sense. Usually in search case, the thread executes an action and rest of the logic is handled by action listener once response is received. I'm taking a look now

[rishabhmaurya]

@Hailong-am are you restoring the preserved context once response is received so that response handler has the context? From the initial looks of it, it doesn't seems like that's the case.

[Hailong-am]

@rishabhmaurya Yes, the context is restored on the executor thread before the response is processed. Here's the flow:

preserveContext() wraps the runnable with a ContextPreservingRunnable which, at execution time:

1. Stashes the executor thread's existing context (clean slate)
2. Restores the caller's captured context onto the executor thread
3. Runs the task — so processBatchTask / processCompleteTask / processErrorTask all see the correct ThreadContext
4. Restores the executor thread's original context when done

This matters because getHeaderBuffer() reads from threadPool.getThreadContext() to serialize headers into the Flight response (line 288). With preserveContext, that call sees the caller's context, so the response carries the correct headers.

The reason we need this explicitly is that FlightServerChannel uses a plain JDK ExecutorService, not OpenSearch's ThreadPool.executor() which wraps runnables with context propagation automatically.

---

Old Flow (stashContext) — broken

Caller Thread                          Executor Thread
─────────────                          ───────────────
ThreadContext: {key: "value"}
        │
        ▼
stashContext()
  ├─ captures snapshot ──────────────▶ StoredContext
  └─ CLEARS caller's context ✗
        │
ThreadContext: {} ← EMPTY!             
        │                              
        ▼                              
submit task to executor ──────────────▶ executor.execute(() -> {
        │                                  │
        ▼                                  ▼
caller continues with                  storedContext.restore()
EMPTY context ✗                        ThreadContext: {key: "value"} ✓
                                           │
                                           ▼
                                       processBatchTask()
                                       getHeaderBuffer() ← sees correct context
                                           │
                                           ▼
                                       task.close()
                                       storedContext.close()
                                       ThreadContext: ??? ← stale, no cleanup

Problems:

1. Caller loses context immediately after stashContext()
2. Executor thread's pre-existing context not cleaned up before restore()
3. Executor thread context not restored after task completes

---

New Flow (preserveContext) — fixed

Caller Thread                          Executor Thread
─────────────                          ───────────────
ThreadContext: {key: "value"}
        │
        ▼
preserveContext(runnable)
  └─ newStoredContext(false)
     captures snapshot ──────────────▶ ContextPreservingRunnable
     does NOT clear caller ✓               │
        │                                  │
ThreadContext: {key: "value"} ✓            │
        │                                  │
        ▼                                  │
submit wrapped task to executor ──────────▶ executor.execute(wrapped -> {
        │                                  │
        ▼                                  ▼
caller continues with                  1. stashContext()
INTACT context ✓                          └─ saves executor's own context
                                          └─ clears executor thread (clean slate)
                                           │
                                           ▼
                                       2. restore() caller's captured context
                                          ThreadContext: {key: "value"} ✓
                                           │
                                           ▼
                                       3. processBatchTask()
                                          getHeaderBuffer() ← sees correct context
                                           │
                                           ▼
                                       4. try-with-resources closes stashed context
                                          └─ executor thread's original context restored
                                          ThreadContext: {original executor ctx} ✓

Fixed:

1. Caller's context preserved throughout
2. Executor thread starts clean before restoring caller's context
3. Executor thread's own context restored after task completes

[rishabhmaurya]

@Hailong-am LGTM, thanks for fixing it. One minor comment on the test.