Add per-node request throttling to workload management

[dzane17]

Description

Adds per-node request throttling to workload management: a workload group can cap how many search requests one node admits for it concurrently, rejecting overflow with a 429.

Config is a new nested throttling object — a Settings bag like the existing settings field, so a per-key null clears one field and an absent key keeps its current value:

PUT _wlm/workload_group
{
  "name": "analytics",
  "resiliency_mode": "enforced",
  "resource_limits": { "cpu": 0.5 },
  "throttling": { "attribute": "group", "node_limit": 10 }
}

• attribute — the dimension the limit is keyed by: group (one bucket for the group), or username / role .
• node_limit — per-node concurrent in-flight allowance per bucket.

_search, _count, _msearch sub-searches and _search/scroll continuations all draw on the same budget. A nested coordinator search issued by the rewrite phase (a terms lookup with a subquery) is not charged a second permit for a bucket its own request already holds. Rejections surface as total_throttled on _wlm/stats and a TOTAL_THROTTLED column on _list/wlm_stats.

How the limit is enforced. Each node throttles in isolation. There is no cross-node coordination, no shared counter, and no extra network hop on the request path — a node decides admission from its own in-memory state. Each node keeps a ConcurrentHashMap of bucket key to in-flight count, incremented on admission and decremented when the request finishes; a bucket's entry is dropped once it drains to zero, so idle buckets cost nothing and cardinality follows live traffic rather than configuration.

node_limit is therefore a per-node ceiling, not a cluster-wide one: a 3-node cluster with node_limit: 10 admits up to 10 concurrent requests for the group per node, so up to 30 in aggregate. Sizing it means reasoning about one node's capacity. (A cluster-wide pool is the follow-up described at the end.)

The bucket key is <workload_group_id>:<attribute>:<attribute_value>, where the value depends on the attribute:

• group — the literal group, so every request tagged to the workload group shares one bucket per node.
• username / role — the caller's value for that subfield, read from the principal the security plugin's extractor supplied, so each principal gets its own bucket per node. When a caller has several values for the subfield (a user in many roles), the request is charged to the lexicographically smallest one, so the same caller always lands in the same bucket rather than drawing a fresh allowance per request.

Two behaviors to know before sizing a limit:

• Each _msearch sub-search takes its own permit, so a multi-search with more sub-requests than node_limit throttles itself — a 429 on the individual item inside an enclosing 200.
• An empty "throttling": {} disables throttling, exactly as "settings": {} clears settings. Use "throttling": null to disable explicitly; omit the field to leave it untouched.

First of two parts: the cluster-wide shared tier is excluded, and shared_limit is not an accepted key, so it cannot be set with no effect.

Related Issues

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 🔍

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: Introduce throttle tracker and settings validation Relevant files: server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleSettings.java server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java Sub-PR theme: Add total_throttled to stats and REST output Relevant files: server/src/main/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsAction.java server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupState.java server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupStats.java server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStatsTests.java Sub-PR theme: Wire throttle admission into search transport actions Relevant files: server/src/main/java/org/opensearch/action/search/TransportSearchAction.java server/src/main/java/org/opensearch/action/search/TransportSearchScrollAction.java server/src/main/java/org/opensearch/action/search/StreamTransportSearchAction.java server/src/main/java/org/opensearch/indices/IndicesService.java server/src/main/java/org/opensearch/search/SearchService.java

⚡ Recommended focus areas for review Silent Exception Swallow In the catch block at the end of doExecute, the failure is delivered to throttledListener.onFailure(e), but if throttledListener == listener (i.e., WLM path did not wrap it, e.g., the task was not a WorkloadGroupTask), this is fine. However, when the throttle throws OpenSearchRejectedExecutionException from acquireThrottleOrReject, the throw happens inside the try block, so it is caught here and delivered via throttledListener.onFailure. Since the throttle acquire failed, throttledListener was never wrapped with runAfter, so listener.onFailure is called correctly. This is OK, but note that if acquireThrottleOrReject throws AFTER partially setting throttledListener (it doesn't in current code), releasing wouldn't happen. Minor: verify the ordering — permit is assigned only if non-null before wrapping, so a throw after acquire but before assignment is not possible. Acceptable as-is, but the control flow is subtle and should be commented. protected void doExecute(Task task, SearchScrollRequest request, ActionListener<SearchResponse> listener) { // Holds the throttle permit release once one is acquired, so every exit below (including the catch) frees it. ActionListener<SearchResponse> throttledListener = listener; try { if (task instanceof WorkloadGroupTask) { ((WorkloadGroupTask) task).setWorkloadGroupId(threadPool.getThreadContext()); // A scroll continuation occupies the node like any other search, so it draws on the same node-level // budget. Exempting it would make node_limit evadable by appending ?scroll= to a query. // A scroll continuation issues no nested coordinator search of its own, so there is no ancestor bucket // to inherit; see TransportSearchAction#bucketsHeldByAncestors. Releasable throttlePermit = workloadGroupService.acquireThrottleOrReject((WorkloadGroupTask) task, Set.of()); if (throttlePermit != null) { throttledListener = ActionListener.runAfter(throttledListener, throttlePermit::close); } } ParsedScrollId scrollId = request.parseScrollId(); Runnable action; switch (scrollId.getType()) { case ParsedScrollId.QUERY_THEN_FETCH_TYPE: action = new SearchScrollQueryThenFetchAsyncAction( logger, clusterService, searchTransportService, searchPhaseController, request, (SearchTask) task, scrollId, throttledListener ); break; case ParsedScrollId.QUERY_AND_FETCH_TYPE: // TODO can we get rid of this? action = new SearchScrollQueryAndFetchAsyncAction( logger, clusterService, searchTransportService, searchPhaseController, request, (SearchTask) task, scrollId, throttledListener ); break; default: throw new IllegalArgumentException("Scroll id type [" + scrollId.getType() + "] unrecognized"); } action.run(); } catch (Exception e) { throttledListener.onFailure(e); } Wire Format Compatibility The throttling field is version-gated with Version.V_3_9_0 on both the read (StreamInput ctor) and write (writeTo) sides, which is symmetric. However, WorkloadGroupStats.WorkloadGroupStatsHolder also gates throttled (a VLong) on V_3_9_0 on both sides — good. Verify that Version.V_3_9_0 is the first unreleased version at the time of merge; if 3.9 has already been released and this change lands post-release, the gate must be bumped to the next unreleased version, otherwise mixed-version clusters where a peer at 3.9.0 predates this change will see stream desync on stats and workload group cluster state. // throttling is newer than settings, so it needs its own gate: a 3.7/3.8 peer writes only settings, and // reading a throttling bag that was never written would desync the stream for every field after it. // Decode "not on the wire" as null, not Settings.EMPTY: this class doubles as the partial update fragment, where // an empty bag is the explicit "clear all throttling" gesture, so EMPTY here would make any update routed // through a pre-3.9 node silently wipe the group's throttling config. Null means "field absent, keep existing"; // WorkloadGroup's constructor normalizes it to EMPTY for a full object. if (in.getVersion().onOrAfter(Version.V_3_9_0)) { throttling = Settings.readOptionalSettingsFromStream(in); } else { throttling = null; } } Listener Ordering Risk updatedListener is wrapped with ActionListener.runAfter(updatedListener, throttlePermit::close) before onRequestStart is called and before listener is later derived. If any code between the wrap and the eventual completion throws synchronously and is NOT routed through listener (which wraps updatedListener), the permit could leak. In particular, searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext) runs after the wrap; if it throws, control exits executeRequest without invoking either updatedListener.onFailure or listener.onFailure, and the throttle permit is never released. Consider wrapping the onRequestStart call in try/catch that routes failures through updatedListener.onFailure, or acquire the permit after onRequestStart. // or HTTP header (HTTP header will be deprecated once ActionFilter is implemented) if (task instanceof WorkloadGroupTask) { ((WorkloadGroupTask) task).setWorkloadGroupId(threadPool.getThreadContext()); // Node-level throttle admission. Runs before onRequestStart so a rejection doesn't leak the request // gauges (decremented only on request end/failure, which the early return skips). The principal is null // unless the WLM auto-tagging filter set it from the security plugin's extractor. try { Releasable throttlePermit = workloadGroupService.acquireThrottleOrReject( (WorkloadGroupTask) task, bucketsHeldByAncestors(task) ); if (throttlePermit != null) { // runAfter, not runBefore: it releases in a finally, so a failure to release cannot turn a // successful search into a client-visible error, and the slot is held until the response has // actually been handed downstream rather than freed just before it. updatedListener = ActionListener.runAfter(updatedListener, throttlePermit::close); } } catch (OpenSearchRejectedExecutionException e) { updatedListener.onFailure(e); return; } } searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext); TOCTOU on Version Check validateThrottlingIsEnforceable reads clusterState.nodes().getMinNodeVersion() at the transport action entry, but the actual cluster-state update is applied later on the cluster-manager thread. Between validation and application, an older node could join, at which point the throttling config would still be silently dropped when serialized to that node. This is a narrow race but real; the invariant the docstring promises ("would otherwise return a 200 for a config that silently never takes effect") can still be violated. A follow-up guard inside the state-update task (rejecting or logging on apply if minNodeVersion regressed) would close it. public static void validateThrottlingIsEnforceable(Settings throttling, ClusterState clusterState) { if (throttling == null || throttling.isEmpty()) { return; } Version minNodeVersion = clusterState.nodes().getMinNodeVersion(); if (minNodeVersion.before(Version.V_3_9_0)) { throw new IllegalArgumentException( "workload group throttling requires every node to be on " + Version.V_3_9_0 + " or later, but the oldest node in the cluster is on " + minNodeVersion + ". The throttling config would be silently dropped; complete the upgrade first." ); } // ATTRIBUTE.get returns "" (its default), not null, when the key is absent -- which is the normal shape of a // partial update that only changes the limit. Only an explicitly principal-keyed attribute is checked here; the // merged config is validated separately. String attribute = WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling); if (attribute == null || attribute.isEmpty() || "group".equals(attribute)) { return; } try { FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME); if (featureType.getAllowedAttributesRegistry().containsKey(WorkloadManagementPlugin.PRINCIPAL_ATTRIBUTE_NAME) == false) { throw new IllegalArgumentException( "throttling attribute [" + attribute + "] needs a principal attribute provider (the security plugin) to be installed, otherwise the " + "limit can never be enforced. Use attribute [group] instead." ); } } catch (ResourceNotFoundException e) { // Feature type not registered on this node yet. Skip rather than reject a config that is probably fine -- // the throttle path fails open anyway, so a false rejection here is worse than a missed warning. logger.debug("WLM feature type not registered; skipping principal-attribute check for throttling config", e); } } Concurrent Bucket Race In WorkloadGroupThrottleTracker.tryAcquire, when the increment observed exceeds nodeLimit, release(bucketKey, counter) is called which decrements and then in a separate compute may evict the bucket if <= 0. Because the decrement is outside compute, another thread that just entered the initial compute and incremented the same counter to 1 could have its counter evicted by this rollback path if it briefly sees <= 0 — but the eviction check re-reads existing.get() inside compute, so a concurrent acquire that already incremented would keep it. This appears safe, but the invariant depends on the decrement being visible before the compute() eviction check reads it; on weak memory models, the ordering is guaranteed only because AtomicInteger.decrementAndGet is a full barrier. Worth a comment stating this explicitly; the current comment mentions the invariant but not the memory-ordering dependency. * As of now this is a stub and main implementation PR will be raised soon.Coming PR will collate these changes with core WorkloadGroupService changes * @opensearch.experimental */ public class WorkloadGroupService extends AbstractLifecycleComponent implements ClusterStateListener, TaskResourceTrackingService.TaskCompletionListener { private static final Logger logger = LogManager.getLogger(WorkloadGroupService.class); private final WorkloadGroupTaskCancellationService taskCancellationService; private volatile Scheduler.Cancellable scheduledFuture; private final ThreadPool threadPool; private final ClusterService clusterService; private final WorkloadManagementSettings workloadManagementSettings; private Set<WorkloadGroup> activeWorkloadGroups; private final Set<WorkloadGroup> deletedWorkloadGroups; private final NodeDuressTrackers nodeDuressTrackers; private final WorkloadGroupsStateAccessor workloadGroupsStateAccessor; // Node-local in-flight throttle counters, keyed by throttle bucket. No cross-node coordination in this tier.

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Validate merged throttling on updates On updates, getThrottling() returns only the incoming partial fragment, so validateThrottlingIsEnforceable sees keys like just node_limit (attribute absent) and skips the principal-attribute check even when the merged config will end up keyed by username/role. Validate the merged result (existing group's throttling overlaid with the incoming fragment) so an update cannot bypass the check that create enforces. plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java [74-82] try { - WorkloadGroupPersistenceService.validateThrottlingIsEnforceable( - request.getmMutableWorkloadGroupFragment().getThrottling(), - clusterState - ); + WorkloadGroup existing = clusterState.metadata().workloadGroups().get(request.getName()); + Settings mergedThrottling = existing != null + ? WorkloadGroup.updateExistingWorkloadGroup(existing, request.getmMutableWorkloadGroupFragment()) + .getMutableWorkloadGroupFragment() + .getThrottling() + : request.getmMutableWorkloadGroupFragment().getThrottling(); + WorkloadGroupPersistenceService.validateThrottlingIsEnforceable(mergedThrottling, clusterState); } catch (Exception e) { listener.onFailure(e); return; } Suggestion importance[1-10]: 7 __ Why: Valid concern: on partial updates, only the incoming fragment is validated, which could allow a merged config keyed by username/role to bypass the principal-attribute check that create enforces. The suggested fix is reasonable though the exact API call may need adjustment.	Medium
	Prevent throttle permit leak on sync throw If PipelinedRequest searchRequest assignment or extractParentTask throws before this try/catch, listener is not yet initialized and the synchronous throw will propagate uncaught while the throttle permit is still held (it's only wired to updatedListener via runAfter, but never invoked because no failure/response reaches it). Wrap the broader region (from where the permit is acquired) in a try/catch that routes any unexpected throw through updatedListener.onFailure so the permit is always released. server/src/main/java/org/opensearch/action/search/TransportSearchAction.java [562-568] try { searchRequest.transformRequest(requestTransformListener); } catch (Exception e) { - // Same listener the asynchronous failure path uses above, so a synchronous throw and an async failure - // are reported identically; it wraps updatedListener, so the throttle permit is still released. - listener.onFailure(e); + // Route through updatedListener so the throttle permit release (wired via runAfter) fires even on a + // synchronous throw before `listener` was constructed. + updatedListener.onFailure(e); } Suggestion importance[1-10]: 6 __ Why: The suggestion identifies a potential permit leak if code between permit acquisition and the try/catch throws synchronously. However, the analysis is partially incorrect since listener is initialized in the try block above; still, using updatedListener for consistency and safety is a valid improvement.	Low
General	Guard against empty throttle attribute ATTRIBUTE.get(throttling) returns "" (the setting's default) when the attribute key is absent, and resolveThrottleAttributeValue("", principal) will fall through and iterate the principal looking for "|"-prefixed tokens, potentially matching unintended values (or returning null). Since validateMergedConfig requires attribute when node_limit is set, treat an empty/missing attribute as "misconfigured -> fail open" explicitly rather than relying on downstream behavior. server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java [383-385] String attribute = WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling); +if (attribute == null || attribute.isEmpty()) { + return null; +} // A null value means the request can't be attributed (e.g. username/role with no principal) -> fail open. String attributeValue = resolveThrottleAttributeValue(attribute, principal); Suggestion importance[1-10]: 5 __ Why: A defensive check for empty attribute is reasonable, though validateMergedConfig should have prevented this state. The resolveThrottleAttributeValue method would fall through with "|" prefix search but likely return null, so the practical impact is minimal.	Low

[github-actions]

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