Skip to content

Enforce node-level WLM request throttling (node_limit)#9

Open
LilyCaroline17 wants to merge 1 commit into
dzane17:3.7-wlm-throttlingfrom
LilyCaroline17:node-throttling-enforcement
Open

Enforce node-level WLM request throttling (node_limit)#9
LilyCaroline17 wants to merge 1 commit into
dzane17:3.7-wlm-throttlingfrom
LilyCaroline17:node-throttling-enforcement

Conversation

@LilyCaroline17

@LilyCaroline17 LilyCaroline17 commented Jul 22, 2026

Copy link
Copy Markdown

Description

Implements node-level enforcement for WLM Request Throttling (# 21064), building on the throttle settings persisted in #8. This PR admits requests against a per-node in-flight cap and rejects over-limit requests with a 429.

Mechanism

On request admission the coordinator acquires one "permit" from a node-local per-bucket in-flight counter; the "permit" is released when the request completes (success or failure). If a request is over the limit, it triggers a OpenSearchRejectedExecutionException (429) (this mirrors IndexingPressure: atomic increment then rollback acquire + Releasable release with a double-close guard). The counter exists only while a bucket has >= 1 in-flight request (created on first acquire, removed when it drains to 0), so memory scales with active buckets rather than the user/role population. The throttling happens for a request at the entry of TransportSearchAction.

Buckets

The bucket a request counts against depends on the group's throttle attribute:

  • group - the whole group shares one bucket
  • username / role - one bucket per principal value

The principal is supplied by the security plugin's auto-tagging principal extractor. WLM action filter propagates it into the thread context, and enforcement reads it in core. When no principal is available, the request is not throttled rather than rejected.

Example

PUT _wlm/workload_group/
{
  "name": "analytics",
  "resiliency_mode": "enforced",
  "resource_limits": { "cpu": 0.3, "memory": 0.4 },
  "throttling": {
    "attribute": "username",
    "node_limit": 10
  } 
} 

A user firing more than 10 concurrent searches on a node gets one 429 per over-limit request; each distinct username is throttled independently, and total_throttled increments per rejection.

Stats

Adds total_throttled (cumulative rejections per workload group) to _wlm/stats.

Testing

Added unit tests (WorkloadGroupThrottleTrackerTests, WorkloadGroupServiceTests, AutoTaggingActionFilterTests, WorkloadGroupThreadContextStatePropagatorTests) and an end-to-end integration test (WlmNodeThrottlingIT) covering group and username throttling on a live node.

Manually verified end-to-end against a live single-node cluster with the security plugin installed: group / username / role throttling each admit exactly node_limit concurrent requests and reject the rest with 429. mixed-attribute concurrent load enforces every bucket independently; and permits release on request success, failure, and rejection.

Related Issues

Related to # 2106

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.

Signed-off-by: Emily Guo <emilyguo@amazon.com>
@LilyCaroline17
LilyCaroline17 force-pushed the node-throttling-enforcement branch from 212b289 to 657d6b4 Compare July 22, 2026 18:04
try {
return throttleTracker.acquire(bucketKey, nodeLimit);
} catch (OpenSearchRejectedExecutionException e) {
workloadGroupsStateAccessor.getWorkloadGroupState(workloadGroupId).totalThrottled.inc();

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(major) This stat update can silently turn a 429 into an admit.

getWorkloadGroupState(...) reads WorkloadGroupsStateAccessor.workloadGroupStateMap, which is a plain HashMap, from the search-admission thread — concurrently with clusterChanged() mutating it (putIfAbsent/remove) on the cluster-applier thread. An unlucky concurrent read can throw a RuntimeException (or NPE), which is not an OpenSearchRejectedExecutionException, so it skips the throw e on the next line and is caught by the fail-open catch (Exception) below → the method returns null → the request that should have been rejected is admitted. Rare and fail-open, but it silently defeats the limit under concurrency.

Cheapest, most robust fix: make WorkloadGroupsStateAccessor.workloadGroupStateMap a ConcurrentHashMap (kills this whole race class). As defense-in-depth, also wrap just the counter increment so it can never swallow the rejection:

catch (OpenSearchRejectedExecutionException e) {
    try {
        WorkloadGroupState s = workloadGroupsStateAccessor.getWorkloadGroupState(workloadGroupId);
        if (s != null) s.totalThrottled.inc();
    } catch (Exception statsEx) {
        logger.warn("failed to record throttle stat", statsEx);
    }
    throw e;
}

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, good catch!

try {
return throttleTracker.acquire(bucketKey, nodeLimit);
} catch (OpenSearchRejectedExecutionException e) {
workloadGroupsStateAccessor.getWorkloadGroupState(workloadGroupId).totalThrottled.inc();

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(minor) Rejection may be counted against the DEFAULT group during the state-registration lag.

ClusterApplierService publishes state() before invoking listeners, so on this node the coordinator can read a newly-created group's node_limit and reject a request before this node's clusterChanged() has called addNewWorkloadGroup(id). At that instant getWorkloadGroupState(id) falls through its getOrDefault(..., DEFAULT) and increments total_throttled on the DEFAULT group instead of the real one. Stats-only and transient, but it misattributes counts.

Suggest reading via getWorkloadGroupStateMap().get(id) here and skipping the increment when it's null, rather than using the DEFAULT-fallback accessor for stat writes.

public static List<String> PROPAGATED_HEADERS = List.of(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER);
public static List<String> PROPAGATED_HEADERS = List.of(
WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER,
WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(minor) Principal header is broadcast cluster-wide though only the origin coordinator consumes it.

Adding WORKLOAD_GROUP_PRINCIPAL_HEADER to the propagated headers() sends the principal (username/role — potentially PII) to every data node and, for cross-cluster search, to remote clusters. But enforcement only ever reads it on the origin coordinator (in-process, before fan-out). That's unnecessary fan-out of sensitive identity data.

Suggest keeping the principal in transients() only (node-local), or stripping it before dispatch, so it never leaves the coordinator.

return null;
}
String subfieldPrefix = attribute + "|";
for (String token : principal.split(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER)) {

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The principal is joined with , in AutoTaggingActionFilter and split back on , here, with no escaping. If a username can contain a comma, username|a,b splits into username|a and b, so it's bucketed as user a — colliding with a different real user a. Note the security extractor deliberately escapes its own | delimiter (unescapePipe), but there's no equivalent handling for this ,. Worth confirming whether principal values can contain a comma, and escaping it on both sides if so.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants