Add tracing support for StreamingRestChannel

[Hailong-am]

Description

Add tracing support for StreamingRestChannel, streaming operation start with prepareResponse and end with sendChunk(lastChunk), the combination will be in same span

Related Issues

opensearch-project/OpenSearch-Dashboards#11090

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.

Summary by CodeRabbit

Bug Fixes

• Fixed thread context restoration for streaming HTTP requests in Reactor Netty 4 transport, ensuring proper context handling and cleanup during streaming operations.

Tests

• Added integration tests to validate thread context restoration behavior for streaming responses.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

📝 Walkthrough

Walkthrough

This PR fixes thread context restoration for streaming HTTP responses in Reactor Netty 4. It introduces thread context storage during streaming lifecycle, ensures proper cleanup after streaming completes, and adds integration testing with mock streaming handlers to validate the fix.

Changes

Cohort / File(s)	Summary
Documentation CHANGELOG.md	Added changelog entry documenting the thread context restoration fix for streaming responses.
Test Implementation plugins/transport-reactor-netty4/src/internalClusterTest/java/.../ReactorNetty4ThreadContextRestorationIT.java	Added integration test class with MockStreamingPlugin (registers mock REST streaming endpoint) and MockStreamingRestHandler (simulates ML streaming behavior with SSE-like chunks and thread context restoration). Test validates that thread context is properly restored during streaming and no duplicate opaque ID errors occur.
HTTP Transport plugins/transport-reactor-netty4/src/main/java/.../ReactorNetty4HttpServerTransport.java	Modified streaming response handling to store thread context at start and ensure cleanup after streaming completes by chaining finalization logic.
REST Channel Handling server/src/main/java/org/opensearch/rest/RestController.java, server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableRestChannel.java	Added TraceableRestChannel unwrapping in RestController's streaming dispatch path to correctly access underlying StreamingRestChannel; added public unwrap() accessor method to TraceableRestChannel to expose the delegate channel.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• Fix X-Opaque-Id header propagation (along with other response headers) for streaming Reactor Netty 4 transport #20371 — Addresses streaming HTTP response header handling (X-Opaque-Id) and thread-context behavior in Reactor Netty 4 streaming.
• Keep track and release Reactor Netty 4 Transport accepted Http Channels during the Node shutdown #20106 — Modifies ReactorNetty4HttpServerTransport and streaming request handling in the reactor-netty4 transport layer.

Suggested labels

bug

Suggested reviewers

• andrross
• cwperks
• reta

Poem

🐰 A thread's context was lost in the stream,
But now with our fix, it's restored to gleam!
Through Netty we store, then carefully free,
The context flows clean—no duplicates we see! ✨

🚥 Pre-merge checks | ❌ 3

❌ Failed checks (2 warnings, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 15.79% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	⚠️ Warning	The PR title 'Add tracing support for StreamingRestChannel' does not accurately reflect the main objectives, which focus on fixing thread context restoration for streaming, not adding tracing support.	Update the title to clearly describe the primary change: 'Fix thread context restoration for streaming' or similar, to match the actual PR objectives and implementation.
Description check	❓ Inconclusive	The PR description provides minimal detail about what change achieves and lacks specificity about implementation approach.	Expand the description with details on: (1) how thread context restoration works, (2) which components were modified and why, (3) the specific problem this solves with streaming operations, and (4) any potential side effects or considerations.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

❌ Gradle check result for ebff832: 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]

✅ Gradle check result for c35d432: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.35%. Comparing base (967c809) to head (a25529f).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #20361      +/-   ##
============================================
+ Coverage     73.23%   73.35%   +0.12%     
- Complexity    71953    72008      +55     
============================================
  Files          5795     5796       +1     
  Lines        329248   329264      +16     
  Branches      47410    47412       +2     
============================================
+ Hits         241122   241539     +417     
+ Misses        68841    68359     -482     
- Partials      19285    19366      +81     

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

[coderabbitai]

Actionable comments posted: 2

🤖 Fix all issues with AI agents

In
@plugins/transport-reactor-netty4/src/internalClusterTest/java/org/opensearch/http/reactor/netty4/ReactorNetty4ThreadContextRestorationIT.java:
- Around line 73-81: The method featureFlagSettings() in
ReactorNetty4ThreadContextRestorationIT is intended to override the base
implementation but lacks the @Override annotation; add the @Override annotation
above the featureFlagSettings() method declaration to match other overrides and
enable compile-time signature checks.
- Around line 147-175: The code currently double-closes storedContext by using
try-with-resources (try (ThreadContext.StoredContext context = storedContext) {
... }) which auto-closes, and also calling storedContext.close() in doFinally;
remove the try-with-resources construct: replace the try(...) block with a plain
block that calls storedContext.restore() (use storedContext.restore() instead of
context.restore()) and remove the corresponding automatic close, leaving the
doFinally(signalType -> storedContext.close()) as the sole cleanup.

🧹 Nitpick comments (2)

plugins/transport-reactor-netty4/src/internalClusterTest/java/org/opensearch/http/reactor/netty4/ReactorNetty4ThreadContextRestorationIT.java (2)

109-109: Unused logger instance.

The logger field is declared but never used in the class.

♻️ Remove unused logger or use it for error logging

Either remove the unused logger:

-        private final Logger logger = LogManager.getLogger(MockStreamingRestHandler.class);

Or use it in the empty catch block (see next comment).

---

180-182: Empty catch block silently swallows errors.

The catch block has only a comment. Use the declared logger to log the error for debugging.

♻️ Log the error

                             } catch (Exception e) {
-                                // Log error
+                                logger.warn("Failed to send error chunk", e);
                             }

📜 Review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 733931b and 328583f.

📒 Files selected for processing (5)

• CHANGELOG.md
• plugins/transport-reactor-netty4/src/internalClusterTest/java/org/opensearch/http/reactor/netty4/ReactorNetty4ThreadContextRestorationIT.java
• plugins/transport-reactor-netty4/src/main/java/org/opensearch/http/reactor/netty4/ReactorNetty4HttpServerTransport.java
• server/src/main/java/org/opensearch/rest/RestController.java
• server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableRestChannel.java

🧰 Additional context used

🧠 Learnings (4)

📓 Common learnings

Learnt from: reta
Repo: opensearch-project/OpenSearch PR: 20017
File: plugins/transport-reactor-netty4/src/test/java/org/opensearch/http/reactor/netty4/ssl/SecureReactorNetty4HttpServerTransportTests.java:256-256
Timestamp: 2025-12-12T18:40:08.452Z
Learning: In the OpenSearch ReactorNetty4 secure HTTP transport tests (plugins/transport-reactor-netty4/src/test/java/org/opensearch/http/reactor/netty4/ssl/SecureReactorNetty4HttpServerTransportTests.java), URI limit validation has been moved from the protocol layer to the transport layer, making it protocol-agnostic. The random protocol selection in ReactorHttpClient.https(settings) is intentional to ensure all tests validate correct behavior across HTTP/1.1, HTTP/2, and HTTP/3.

📚 Learning: 2025-12-12T18:40:08.452Z

Learnt from: reta
Repo: opensearch-project/OpenSearch PR: 20017
File: plugins/transport-reactor-netty4/src/test/java/org/opensearch/http/reactor/netty4/ssl/SecureReactorNetty4HttpServerTransportTests.java:256-256
Timestamp: 2025-12-12T18:40:08.452Z
Learning: In the OpenSearch ReactorNetty4 secure HTTP transport tests (plugins/transport-reactor-netty4/src/test/java/org/opensearch/http/reactor/netty4/ssl/SecureReactorNetty4HttpServerTransportTests.java), URI limit validation has been moved from the protocol layer to the transport layer, making it protocol-agnostic. The random protocol selection in ReactorHttpClient.https(settings) is intentional to ensure all tests validate correct behavior across HTTP/1.1, HTTP/2, and HTTP/3.

Applied to files:

• plugins/transport-reactor-netty4/src/internalClusterTest/java/org/opensearch/http/reactor/netty4/ReactorNetty4ThreadContextRestorationIT.java
• plugins/transport-reactor-netty4/src/main/java/org/opensearch/http/reactor/netty4/ReactorNetty4HttpServerTransport.java

📚 Learning: 2025-12-12T13:31:51.234Z

Learnt from: andriyredko
Repo: opensearch-project/OpenSearch PR: 20017
File: plugins/transport-reactor-netty4/src/test/java/org/opensearch/http/reactor/netty4/ssl/SecureReactorNetty4HttpServerTransportTests.java:642-646
Timestamp: 2025-12-12T13:31:51.234Z
Learning: In the OpenSearch ReactorNetty4 secure HTTP transport tests (plugins/transport-reactor-netty4/src/test/java/org/opensearch/http/reactor/netty4/ssl/SecureReactorNetty4HttpServerTransportTests.java), the createBuilderWithPort() helper intentionally uses randomBoolean() for the HTTP/3 enabled setting to ensure all tests validate correct behavior with both HTTP/3 enabled and disabled configurations.

Applied to files:

• plugins/transport-reactor-netty4/src/internalClusterTest/java/org/opensearch/http/reactor/netty4/ReactorNetty4ThreadContextRestorationIT.java

📚 Learning: 2025-12-13T20:16:15.318Z

Learnt from: reta
Repo: opensearch-project/OpenSearch PR: 20017
File: modules/transport-netty4/src/main/java/org/opensearch/http/netty4/Netty4Http3ServerTransport.java:101-123
Timestamp: 2025-12-13T20:16:15.318Z
Learning: In OpenSearch, only one HTTP transport implementation can be active and loaded at a time, so duplicate setting definitions (such as h3.max_stream_local_length, h3.max_stream_remote_length, and h3.max_streams) across different transport implementations like Netty4Http3ServerTransport and ReactorNetty4HttpServerTransport will not cause setting registration conflicts.

Applied to files:

• plugins/transport-reactor-netty4/src/internalClusterTest/java/org/opensearch/http/reactor/netty4/ReactorNetty4ThreadContextRestorationIT.java
• plugins/transport-reactor-netty4/src/main/java/org/opensearch/http/reactor/netty4/ReactorNetty4HttpServerTransport.java

🧬 Code graph analysis (2)

plugins/transport-reactor-netty4/src/internalClusterTest/java/org/opensearch/http/reactor/netty4/ReactorNetty4ThreadContextRestorationIT.java (2)

server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java (1)

• FeatureFlagSettings (21-43)

server/src/main/java/org/opensearch/common/util/FeatureFlags.java (1)

• FeatureFlags (29-341)

server/src/main/java/org/opensearch/rest/RestController.java (1)

server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableRestChannel.java (1)

• TraceableRestChannel (27-118)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (21)

• GitHub Check: gradle-check
• GitHub Check: precommit (25, windows-latest)
• GitHub Check: precommit (21, windows-2025, true)
• GitHub Check: precommit (21, ubuntu-latest)
• GitHub Check: precommit (21, ubuntu-24.04-arm)
• GitHub Check: precommit (25, macos-15-intel)
• GitHub Check: precommit (25, ubuntu-latest)
• GitHub Check: precommit (21, windows-latest)
• GitHub Check: precommit (25, ubuntu-24.04-arm)
• GitHub Check: precommit (25, macos-15)
• GitHub Check: precommit (21, macos-15-intel)
• GitHub Check: precommit (21, macos-15)
• GitHub Check: assemble (25, ubuntu-24.04-arm)
• GitHub Check: assemble (25, ubuntu-latest)
• GitHub Check: Analyze (java)
• GitHub Check: assemble (21, ubuntu-latest)
• GitHub Check: assemble (25, windows-latest)
• GitHub Check: assemble (21, ubuntu-24.04-arm)
• GitHub Check: assemble (21, windows-latest)
• GitHub Check: detect-breaking-change
• GitHub Check: Mend Security Check

🔇 Additional comments (5)

CHANGELOG.md (1)

33-33: LGTM!

The changelog entry follows the established format and is correctly placed under the Fixed section.

plugins/transport-reactor-netty4/src/main/java/org/opensearch/http/reactor/netty4/ReactorNetty4HttpServerTransport.java (1)

385-400: Thread context restoration approach looks correct.

The StoredContext is properly created before streaming begins and closed after streaming completes via then(Mono.fromRunnable(storedContext::close)). The then() operator executes after both successful completion and error signals, ensuring cleanup happens in all cases.

One consideration: newStoredContext(false) means response headers won't be preserved. Verify this aligns with the intended behavior for streaming responses.

server/src/main/java/org/opensearch/rest/RestController.java (1)

349-352: LGTM!

The unwrapping logic correctly handles the case where TraceableRestChannel wraps a StreamingRestChannel. This ensures the streaming detection works properly when telemetry tracing is enabled.

server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableRestChannel.java (1)

62-68: LGTM!

The unwrap() method provides a clean way to access the underlying delegate channel, enabling proper streaming channel detection in RestController.

plugins/transport-reactor-netty4/src/internalClusterTest/java/org/opensearch/http/reactor/netty4/ReactorNetty4ThreadContextRestorationIT.java (1)

213-240: Test logic looks sound for verifying thread context restoration.

The test sends multiple streaming requests and verifies that no duplicate X-Opaque-Id errors occur, which validates that thread context is properly restored between requests.

[github-actions]

❌ Gradle check result for 328583f: 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]

❌ Gradle check result for b08c361: 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]

❌ Gradle check result for 800b406: 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]

❌ Gradle check result for ca9c8f2: 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]

❌ Gradle check result for d66c812: 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]

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

[reta]

@Hailong-am sorry, conflicts popped up, could you please resolve them? thank you

[github-actions]

❌ Gradle check result for 717b0bc: 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]

❌ Gradle check result for a25529f: 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]

✅ Gradle check result for a25529f: SUCCESS