Fix phantom WLM rejections and stats for coordinator and shard search tasks

[LilyCaroline17]

Description

Coordinator-level WLM search tasks were never actually rejected after resource limits were reached, and every attempted rejection produced a phantom entry in _wlm/stats where a task would be considered rejected when it actually ran to completion. This is because the admission check ran in WorkloadGroupRequestOperationListener.onRequestStart, inside a CompositeListener that swallows exceptions, so the OpenSearchRejectedExecutionException never reached the client, while rejectIfNeeded had already bumped the counters.

This change moves the rejectIfNeeded call from WorkloadGroupRequestOperationListener.onRequestStart to TransportSearchAction.executeRequest where the rejection is returned to the client with an onFailure call. Counters now are incremented only for real rejected requests.

Another issue with phantom entries in _wlm/stats was discovered where rejected tasks for shard search tasks were also counted as completions as setWorkloadGroupId, which tags the task via isWorkloadGroupSet, ran before rejectIfNeeded in WorkloadManagementTransportInterceptor.messageReceived. A rejected task was therefore tagged, so WorkloadGroupService.onTaskCompleted counted it in total_completions, meaning that a rejected task was counted as a completion.

The change for this moves the rejectIfNeeded call before setWorkloadGroupId, so a rejected task is never tagged and never counted as a completion.

Related Issues

Resolves #22541

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 d82976c)

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 Skipped onRequestStart may leak listener state On coordinator rejection, the code calls updatedListener.onFailure(e) and returns before invoking searchRequestContext.getSearchRequestOperationsListener().onRequestStart(...). Any listener whose onRequestEnd/onRequestFailure counterpart is invoked via the ActionListener chain (or elsewhere) will now observe an end without a matching start. Verify that no operation listener registered via SearchRequestOperationsCompositeListenerFactory (metrics, tracing, slow-log, resource tracking) relies on onRequestStart having been called - otherwise counters/spans could underflow or leak. Consider whether the admission check should happen before SearchRequestContext is even constructed, or whether skipping only WLM's listener (not all) is more appropriate. if (task instanceof WorkloadGroupTask workloadGroupTask) { // Coordinator-task admission point. Runs before onRequestStart (keeps the in-flight gauge balanced) and // before setWorkloadGroupId (so a rejected task is not counted in total_completions on task completion). try { workloadGroupService.rejectIfNeeded( threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER) ); } catch (OpenSearchRejectedExecutionException e) { updatedListener.onFailure(e); return; } workloadGroupTask.setWorkloadGroupId(threadPool.getThreadContext()); } searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext); Header lookup vs task-stored id The previous code rejected using ((WorkloadGroupTask) task).getWorkloadGroupId() after setWorkloadGroupId populated it from the thread context (with fallback to the default group id when the header is absent). The new code reads WORKLOAD_GROUP_ID_HEADER directly from the thread context, which returns null when the header is missing. If rejectIfNeeded behaves differently for a null id versus the default workload group id resolved by setWorkloadGroupId, rejection semantics for shard-level requests without the header will silently change. if (isSearchWorkloadRequest(task)) { // Reject before setWorkloadGroupId so a rejected task is not tagged and thus not counted as a phantom completion. workloadGroupService.rejectIfNeeded(threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER)); ((WorkloadGroupTask) task).setWorkloadGroupId(threadPool.getThreadContext()); }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to d82976c

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Prevent telemetry leak on rejection path The rejection path calls updatedListener.onFailure(e) and returns, but the enclosing try/finally around this block (which likely closes requestSpan / releases resources created just above) may be bypassed depending on structure. More importantly, the SearchRequestContext was constructed but onRequestStart was skipped — verify that no listener state (e.g., request counters/spans) created during context construction leaks on the rejection path, and consider invoking an equivalent onRequestFailure/cleanup so telemetry stays balanced. server/src/main/java/org/opensearch/action/search/TransportSearchAction.java [494-506] if (task instanceof WorkloadGroupTask workloadGroupTask) { - // Coordinator-task admission point. Runs before onRequestStart (keeps the in-flight gauge balanced) and - // before setWorkloadGroupId (so a rejected task is not counted in total_completions on task completion). try { workloadGroupService.rejectIfNeeded( threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER) ); } catch (OpenSearchRejectedExecutionException e) { + // Ensure no partially-initialized telemetry/listener state is leaked on rejection. updatedListener.onFailure(e); return; } workloadGroupTask.setWorkloadGroupId(threadPool.getThreadContext()); } Suggestion importance[1-10]: 5 __ Why: The suggestion raises a valid concern about potential telemetry/listener state leaks on the rejection path since SearchRequestContext is constructed but onRequestStart is skipped. However, it primarily asks the author to verify rather than providing a concrete fix, and the improved_code only adds a comment without functional changes.	Low

---

Previous suggestions

Suggestions up to commit 09b367a

Category	Suggestion	Impact
General	Close tracing span on rejection path The rejection path calls updatedListener.onFailure(e) but the surrounding code likely has already registered request-start hooks or opened a tracing span (requestSpan) that will now be leaked, since onRequestStart was moved after this block. Ensure the tracing span and any other resources opened prior to admission are properly closed/failed on rejection to avoid resource leaks. server/src/main/java/org/opensearch/action/search/TransportSearchAction.java [494-506] if (task instanceof WorkloadGroupTask workloadGroupTask) { - // Coordinator-task admission point. Runs before onRequestStart (keeps the in-flight gauge balanced) and - // before setWorkloadGroupId (so a rejected task is not counted in total_completions on task completion). try { workloadGroupService.rejectIfNeeded( threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER) ); } catch (OpenSearchRejectedExecutionException e) { - updatedListener.onFailure(e); + try { + requestSpan.close(); + } finally { + updatedListener.onFailure(e); + } return; } workloadGroupTask.setWorkloadGroupId(threadPool.getThreadContext()); } Suggestion importance[1-10]: 5 __ Why: The suggestion raises a valid concern about potential resource leaks (tracing span) on the rejection path, but it's speculative without confirmation that requestSpan requires explicit closing here, and TraceableActionListener may already handle span closure via updatedListener.onFailure.	Low

Suggestions up to commit f66d002

Category	Suggestion	Impact
Possible issue	Guard against null workload group service The admission check is only invoked when workloadGroupService is non-null; since it's now @Inject-ed as a required dependency, ensure it is non-null, otherwise a NullPointerException will be thrown here for any coordinator search task. Add a null-guard to preserve backward compatibility when WLM is disabled/absent. server/src/main/java/org/opensearch/action/search/TransportSearchAction.java [494-503] if (task instanceof WorkloadGroupTask) { ((WorkloadGroupTask) task).setWorkloadGroupId(threadPool.getThreadContext()); // Coordinator-task admission point; must run before onRequestStart to keep in-flight gauges balanced. - try { - workloadGroupService.rejectIfNeeded(((WorkloadGroupTask) task).getWorkloadGroupId()); - } catch (OpenSearchRejectedExecutionException e) { - updatedListener.onFailure(e); - return; + if (workloadGroupService != null) { + try { + workloadGroupService.rejectIfNeeded(((WorkloadGroupTask) task).getWorkloadGroupId()); + } catch (OpenSearchRejectedExecutionException e) { + updatedListener.onFailure(e); + return; + } } } Suggestion importance[1-10]: 3 __ Why: Since workloadGroupService is injected as a required dependency via @Inject and is a final field, it should not be null in normal operation. The null-guard is defensive but not clearly necessary given the injection contract.	Low

[github-actions]

✅ Gradle check result for f66d002: SUCCESS

[codecov]

Codecov Report

❌ Patch coverage is 90.90909% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 71.61%. Comparing base (baa324b) to head (d82976c).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
...pensearch/action/search/TransportSearchAction.java	90.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff            @@
##               main   #22839   +/-   ##
=========================================
  Coverage     71.60%   71.61%           
+ Complexity    77358    77351    -7     
=========================================
  Files          6170     6170           
  Lines        359710   359714    +4     
  Branches      52460    52460           
=========================================
+ Hits         257583   257605   +22     
+ Misses        81693    81656   -37     
- Partials      20434    20453   +19     

☔ View full report in Codecov by Harness.
📢 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.

[dzane17]

Right now a rejected request will count toward both total_rejections and total_completions.

You might be able to skip total_completions increment by callling setWorkloadGroupId() after the rejection check.

if (task instanceof WorkloadGroupTask workloadGroupTask) {
    // Admission must precede onRequestStart (in-flight gauge) and setWorkloadGroupId
    // (so a rejected task is not counted in total_completions).
    try {
        workloadGroupService.rejectIfNeeded(threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER));
    } catch (OpenSearchRejectedExecutionException e) {
        updatedListener.onFailure(e);
        return;
    }
    workloadGroupTask.setWorkloadGroupId(threadPool.getThreadContext());
}

[LilyCaroline17]

Oh, I see what you mean. onTaskCompleted in WorkloadGroupService would include it in the total_completions counter since setWorkloadGroupId sets isWorkloadGroupSet to true. Good catch, thanks! Updating now.

[github-actions]

Persistent review updated to latest commit 09b367a

[github-actions]

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

Persistent review updated to latest commit d82976c

[github-actions]

✅ Gradle check result for d82976c: SUCCESS