improve unit test make sure thread context propagate to executor thread

[Hailong-am]

Description

Follow up for #21167 (comment) to add assertion to make sure thread context propagate to executor thread

Related Issues

Resolves #[Issue number to be closed when this PR is merged]

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

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 Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 018f585)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Incomplete Assertion In testSendResponseBatchPreservesCallerThreadContext, the CountDownLatch was removed and replaced with a simple doAnswer that returns null. There is no longer any synchronization mechanism to wait for the executor thread to complete before the test ends. This means the test may pass trivially without actually verifying that the response was sent or that the context was preserved, since the executor thread may not have run yet when assertions are checked. doAnswer(invocation -> 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 ); // Verify the caller's thread context is NOT cleared assertEquals( Missing Latch in First Test The first test testSendResponseBatchPreservesCallerThreadContext no longer waits for the async executor task to complete (the CountDownLatch was removed). Without a latch or other synchronization, the test cannot reliably verify that the context was preserved on the executor thread, and may pass even if the executor task never runs. public void testSendResponseBatchPreservesCallerThreadContext() throws Exception { ThreadContext threadContext = threadPool.getThreadContext(); threadContext.putHeader(HEADER_KEY, HEADER_VALUE); doAnswer(invocation -> 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 ); // Verify the caller's thread context is NOT cleared assertEquals(

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 018f585

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Add synchronization for async executor task completion In testSendResponseBatchPreservesCallerThreadContext, the CountDownLatch was removed, so there is no synchronization mechanism to wait for the async executor task to complete before the test ends. Without a latch or similar barrier, the test may pass vacuously if assertions run before the executor thread finishes. Consider adding a latch or other synchronization to ensure the async work completes before assertions. plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java [80-84] -doAnswer(invocation -> null).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); +CountDownLatch latch = new CountDownLatch(1); +doAnswer(invocation -> { + latch.countDown(); + return null; +}).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); handler.sendResponseBatch( Version.CURRENT, Collections.emptySet(), Suggestion importance[1-10]: 6 __ Why: This is a valid concern — removing the CountDownLatch from testSendResponseBatchPreservesCallerThreadContext could lead to a race condition where the test completes before the async executor task finishes. Adding synchronization ensures the test reliably waits for the async work to complete.	Low
	Capture header on executor thread, not listener callback The onResponseSent callback may not be invoked on the executor thread — it could be called on the calling thread after the executor task completes. If the goal is to verify that the thread context is propagated to the executor thread specifically, the header should be captured inside a task submitted to the executor, not inside the listener callback. Otherwise, the test may pass even if context propagation to the executor thread is broken. plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java [148-152] doAnswer(invocation -> { - capturedHeader.set(threadPool.getThreadContext().getHeader(HEADER_KEY)); - latch.countDown(); + // Verify the header is captured on the executor thread + executor.submit(() -> { + capturedHeader.set(threadPool.getThreadContext().getHeader(HEADER_KEY)); + latch.countDown(); + }); return null; }).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); Suggestion importance[1-10]: 5 __ Why: The suggestion raises a valid concern about whether onResponseSent is called on the executor thread. However, the comment in the PR code explicitly states the callback "runs within the preserveContext wrapper on the executor thread," suggesting the author intentionally captures the header there. The improved code introduces additional async complexity that may not be appropriate.	Low

---

Previous suggestions

Suggestions up to commit a3dea65

Category	Suggestion	Impact
Possible issue	Restore synchronization for async test completion The testSendResponseBatchPreservesCallerThreadContext test removed the CountDownLatch that was used to wait for the async executor task to complete before the test method returns. Without it, the test may pass trivially or produce flaky results because the assertions may run before the async work finishes. A latch or other synchronization mechanism should be retained to ensure the test waits for the async operation to complete. plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java [80-84] -doAnswer(invocation -> null).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); +CountDownLatch latch = new CountDownLatch(1); +doAnswer(invocation -> { + latch.countDown(); + return null; +}).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); handler.sendResponseBatch( Version.CURRENT, Collections.emptySet(), Suggestion importance[1-10]: 7 __ Why: This is a valid concern — removing the CountDownLatch from testSendResponseBatchPreservesCallerThreadContext could make the test non-deterministic since there's no synchronization to wait for the async executor task to complete before the test method returns. The improved code correctly restores the latch-based synchronization.	Medium
	Verify context on actual executor thread The onResponseSent callback may not necessarily run on the executor thread — it could be invoked on the calling thread after the executor task completes. To reliably verify that the thread context is propagated to the executor thread, the header should be captured inside the actual executor-submitted task, not in the listener callback. Consider capturing the header in a task submitted to the executor or wrapping the assertion within the actual execution path on the executor thread. plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java [148-152] doAnswer(invocation -> { capturedHeader.set(threadPool.getThreadContext().getHeader(HEADER_KEY)); latch.countDown(); return null; }).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); +// Note: To truly verify executor thread context propagation, consider +// intercepting the Runnable submitted to the executor and capturing +// the header inside that runnable's execution. + Suggestion importance[1-10]: 4 __ Why: The suggestion raises a valid concern about whether onResponseSent runs on the executor thread, but the 'improved_code' only adds a comment without actually fixing the issue. The suggestion doesn't provide a concrete fix, making it more of an observation than an actionable improvement.	Low

Suggestions up to commit 931ed4b

Category	Suggestion	Impact
Possible issue	Fix incorrect matcher type in mock setup The mock is set up with any(Exception.class) as the third argument matcher, but onResponseSent is likely called with a TransportResponse object (not an Exception). This means the mock will never match and the latch.countDown() will never be called, causing the test to time out. It should use any(TransportResponse.class) to match the actual invocation, consistent with the other test method. plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java [148] -}).when(mockListener).onResponseSent(anyLong(), anyString(), any(Exception.class)); +}).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); Suggestion importance[1-10]: 9 __ Why: The mock uses any(Exception.class) but onResponseSent is called with a TransportResponse object, so the mock will never match and latch.countDown() will never be called, causing the test to time out. This is a critical bug that would cause the test to fail, and the fix is consistent with the other test method that correctly uses any(TransportResponse.class).	High

Suggestions up to commit 4d860a6

Category	Suggestion	Impact
Possible issue	Fix incorrect argument matcher type in mock setup The mock is set up with any(Exception.class) as the third argument matcher, but onResponseSent is called with a TransportResponse object (not an Exception). This means the mock will never match and capturedHeader will remain null, causing the latch to never count down and the test to time out. The matcher should be any(TransportResponse.class) to be consistent with the actual call and the other test method. plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightOutboundHandlerTests.java [148] -}).when(mockListener).onResponseSent(anyLong(), anyString(), any(Exception.class)); +}).when(mockListener).onResponseSent(anyLong(), anyString(), any(TransportResponse.class)); Suggestion importance[1-10]: 9 __ Why: The mock uses any(Exception.class) but onResponseSent is called with a TransportResponse object, so the mock will never match. This would cause capturedHeader to remain null and the latch to never count down, resulting in a test timeout failure.	High

[github-actions]

Persistent review updated to latest commit 931ed4b

[github-actions]

Persistent review updated to latest commit a3dea65

[github-actions]

❌ Gradle check result for a3dea65: null

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]

❌ Gradle check result for a3dea65: 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?

[Hailong-am]

test failure not related to this change. https://build.ci.opensearch.org/job/gradle-check/74403/console

Tests with failures:
 - org.opensearch.remotestore.RemoteIndexRecoveryIT.testOngoingRecoveryAndClusterManagerFailOver

5779 tests completed, 1 failed, 72 skipped

[github-actions]

Persistent review updated to latest commit 018f585

[github-actions]

✅ Gradle check result for 018f585: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.16%. Comparing base (1bb2c64) to head (018f585).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21200      +/-   ##
============================================
- Coverage     73.18%   73.16%   -0.03%     
+ Complexity    72939    72883      -56     
============================================
  Files          5888     5888              
  Lines        333169   333169              
  Branches      48058    48058              
============================================
- Hits         243820   243751      -69     
- Misses        69855    69922      +67     
- Partials      19494    19496       +2     

☔ 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.