Enforce node-level WLM request throttling (node_limit)#9
Conversation
Signed-off-by: Emily Guo <emilyguo@amazon.com>
212b289 to
657d6b4
Compare
| try { | ||
| return throttleTracker.acquire(bucketKey, nodeLimit); | ||
| } catch (OpenSearchRejectedExecutionException e) { | ||
| workloadGroupsStateAccessor.getWorkloadGroupState(workloadGroupId).totalThrottled.inc(); |
There was a problem hiding this comment.
(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;
}| try { | ||
| return throttleTracker.acquire(bucketKey, nodeLimit); | ||
| } catch (OpenSearchRejectedExecutionException e) { | ||
| workloadGroupsStateAccessor.getWorkloadGroupState(workloadGroupId).totalThrottled.inc(); |
There was a problem hiding this comment.
(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 |
There was a problem hiding this comment.
(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)) { |
There was a problem hiding this comment.
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.
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 mirrorsIndexingPressure: atomic increment then rollback acquire +Releasablerelease 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 bucketusername/role- one bucket per principal valueThe 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
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_throttledincrements 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_limitconcurrent 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.