Enforce cluster-level WLM request throttling (shared_limit)

server/src/main/java/org/opensearch/action/search/TransportSearchAction.java

@@ -39,6 +39,7 @@
 import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsRequest;
 import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsResponse;
 import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.ContextPreservingActionListener;
 import org.opensearch.action.support.HandledTransportAction;
 import org.opensearch.action.support.IndicesOptions;
 import org.opensearch.action.support.TimeoutTaskCancellationUtility;
@@ -68,7 +69,6 @@
 import org.opensearch.core.common.breaker.CircuitBreaker;
 import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.core.common.io.stream.Writeable;
-import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
 import org.opensearch.core.index.Index;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.core.indices.breaker.CircuitBreakerService;
@@ -485,25 +485,71 @@ void executeRequest(
             // 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). Principal header
-                // is null unless the security plugin's extractor set it.
-                try {
-                    String principal = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER);
-                    Releasable throttlePermit = workloadGroupService.acquireThrottleOrReject(
-                        ((WorkloadGroupTask) task).getWorkloadGroupId(),
-                        principal
-                    );
-                    if (throttlePermit != null) {
-                        updatedListener = ActionListener.runBefore(updatedListener, throttlePermit::close);
-                    }
-                } catch (OpenSearchRejectedExecutionException e) {
-                    updatedListener.onFailure(e);
-                    return;
-                }
+                // Two-tier throttle admission (node-local + cluster-level shared). Runs before onRequestStart so a
+                // rejection doesn't leak the request gauges (decremented only on request end/failure, which the early
+                // return skips). The node-local tier is granted synchronously here (zero added latency); only an
+                // overflow to the shared tier does an asynchronous owner round-trip, so this call may complete inline
+                // or on a transport thread. Principal header is null unless the security plugin's extractor set it.
+                final String principal = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER);
+                final ActionListener<SearchResponse> outerListener = updatedListener;
+                // Preserve the request's thread context (workload group + principal headers, etc.) so that when the
+                // shared-tier admission callback runs on a transport thread, downstream search setup sees them.
+                final ActionListener<Releasable> admissionListener = ContextPreservingActionListener.wrapPreservingContext(
+                    ActionListener.wrap(throttlePermit -> {
+                        ActionListener<SearchResponse> proceedListener = throttlePermit == null
+                            ? outerListener
+                            : ActionListener.runBefore(outerListener, throttlePermit::close);
+                        proceedWithSearch(
+                            task,
+                            originalSearchRequest,
+                            searchAsyncActionProvider,
+                            proceedListener,
+                            searchRequestContext,
+                            timeProvider,
+                            requestSpan
+                        );
+                    }, outerListener::onFailure),
+                    threadPool.getThreadContext()
+                );
+                workloadGroupService.acquireThrottlePermit(((WorkloadGroupTask) task).getWorkloadGroupId(), principal, admissionListener);
+            } else {
+                proceedWithSearch(
+                    task,
+                    originalSearchRequest,
+                    searchAsyncActionProvider,
+                    updatedListener,
+                    searchRequestContext,
+                    timeProvider,
+                    requestSpan
+                );
             }
+        }
+    }
 
-            searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);
+    /**
+     * Runs the search once throttle admission has granted (or been skipped). Extracted from {@link #executeRequest}
+     * so it can be invoked either inline (node-local grant / not throttled) or asynchronously from the shared-tier
+     * acquire callback. Re-enters the request span scope because the async callback may run on a transport thread
+     * where the caller's span scope is no longer active.
+     */
+    private void proceedWithSearch(
+        Task task,
+        SearchRequest originalSearchRequest,
+        SearchAsyncActionProvider searchAsyncActionProvider,
+        ActionListener<SearchResponse> updatedListener,
+        SearchRequestContext searchRequestContext,
+        SearchTimeProvider timeProvider,
+        Span requestSpan
+    ) {
+        try (final SpanScope ignore = tracer.withSpanInScope(requestSpan)) {
+            try {
+                searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);
+            } catch (Exception e) {
+                // Route through updatedListener (which may hold a throttle permit via runBefore) so a failure here
+                // releases the permit rather than leaking it. A node-local permit has no TTL to reclaim it.
+                updatedListener.onFailure(e);
+                return;
+            }
 
             PipelinedRequest searchRequest;
             ActionListener<SearchResponse> listener;

server/src/main/java/org/opensearch/node/Node.java

@@ -309,6 +309,7 @@
 import org.opensearch.usage.UsageService;
 import org.opensearch.watcher.ResourceWatcherService;
 import org.opensearch.wlm.WorkloadGroupService;
+import org.opensearch.wlm.WorkloadGroupSharedThrottleService;
 import org.opensearch.wlm.WorkloadGroupsStateAccessor;
 import org.opensearch.wlm.WorkloadManagementSettings;
 import org.opensearch.wlm.WorkloadManagementTransportInterceptor;
@@ -1423,6 +1424,16 @@ protected Node(final Environment initialEnvironment, Collection<PluginInfo> clas
                 )
                 : Optional.empty();
 
+            // Cluster-level (shared_limit) WLM throttle tier. Built here because it needs the transport service for
+            // the owner round-trip; late-bound into the WorkloadGroupService (constructed earlier) so the two-tier
+            // admit() path can reach it.
+            final WorkloadGroupSharedThrottleService workloadGroupSharedThrottleService = new WorkloadGroupSharedThrottleService(
+                clusterService,
+                threadPool,
+                transportService
+            );
+            workloadGroupService.setSharedThrottleService(workloadGroupSharedThrottleService);
+
             TopNSearchTasksLogger taskConsumer = new TopNSearchTasksLogger(settings, settingsModule.getClusterSettings());
             transportService.getTaskManager().registerTaskResourceConsumer(taskConsumer);
             streamTransportService.ifPresent(service -> service.getTaskManager().registerTaskResourceConsumer(taskConsumer));
@@ -1709,6 +1720,7 @@ protected Node(final Environment initialEnvironment, Collection<PluginInfo> clas
                 b.bind(TaskResourceTrackingService.class).toInstance(taskResourceTrackingService);
                 b.bind(SearchBackpressureService.class).toInstance(searchBackpressureService);
                 b.bind(WorkloadGroupService.class).toInstance(workloadGroupService);
+                b.bind(WorkloadGroupSharedThrottleService.class).toInstance(workloadGroupSharedThrottleService);
                 b.bind(AdmissionControlService.class).toInstance(admissionControlService);
                 b.bind(UsageService.class).toInstance(usageService);
                 b.bind(AggregationUsageService.class).toInstance(searchModule.getValuesSourceRegistry().getUsageService());
@@ -1963,6 +1975,7 @@ public Node start() throws NodeValidationException {
         nodeService.getSearchBackpressureService().start();
         nodeService.getTaskCancellationMonitoringService().start();
         injector.getInstance(WorkloadGroupService.class).start();
+        injector.getInstance(WorkloadGroupSharedThrottleService.class).start();
 
         final ClusterService clusterService = injector.getInstance(ClusterService.class);
 
@@ -2153,6 +2166,7 @@ private Node stop() {
         injector.getInstance(NodeResourceUsageTracker.class).stop();
         injector.getInstance(ResourceUsageCollectorService.class).stop();
         injector.getInstance(WorkloadGroupService.class).stop();
+        injector.getInstance(WorkloadGroupSharedThrottleService.class).stop();
         nodeService.getMonitorService().stop();
         nodeService.getSearchBackpressureService().stop();
         injector.getInstance(GatewayService.class).stop();

server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java → server/src/main/java/org/opensearch/wlm/NodeThrottleTracker.java

@@ -10,52 +10,60 @@
 
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.common.lease.Releasable;
-import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
 
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
- * Tracks the number of in-flight requests per throttle bucket on a single node and enforces a per-node cap.
+ * Node-local tier of workload-group request throttling: tracks in-flight requests per throttle bucket on a single
+ * node and enforces the bucket's {@code node_limit}. Its cluster-wide sibling is {@link SharedThrottleTracker}, and
+ * {@code WorkloadGroupService.acquireThrottlePermit} composes the two (local first, then shared on overflow).
  * <p>
  * A bucket is identified by an opaque key (see {@code WorkloadGroupService} for how the key is built from a
  * workload group and its throttle attribute). A counter exists only while a bucket has at least one in-flight
  * request: it is created on first acquire and removed when it drains back to zero, so memory scales with the
  * number of concurrently active buckets rather than the total population of users/roles.
  * <p>
  * This tier is fully local — no cross-node coordination — mirroring the acquire/rollback + {@link Releasable}
- * release shape of {@link org.opensearch.index.IndexingPressure}.
+ * release shape of {@link org.opensearch.index.IndexingPressure}. Because acquire and release happen in the same
+ * process, a permit is a self-releasing {@link Releasable} and needs no lease id or TTL (contrast
+ * {@link SharedThrottleTracker}, whose distributed acquire/release requires both). When a two-tier throttle is
+ * configured, an exhausted local tier falls through to the cluster-level {@link SharedThrottleTracker} rather than
+ * rejecting outright; see {@link #tryAcquire}, which returns {@code null} on exhaustion instead of throwing.
  */
 @ExperimentalApi
-public class WorkloadGroupThrottleTracker {
+public class NodeThrottleTracker {
 
     private final Map<String, AtomicInteger> inFlightByBucket = new ConcurrentHashMap<>();
 
     /**
-     * Attempts to admit one request into the given bucket under the per-node limit.
+     * Attempts to admit one request into the given bucket under the per-node limit, returning {@code null} instead
+     * of throwing when the bucket is already at the limit. This is the fall-through-friendly form: a caller that
+     * also has a cluster-level shared pool can try the shared pool on a {@code null} result rather than rejecting.
      *
      * @param bucketKey the throttle bucket identifier
      * @param nodeLimit the maximum concurrent in-flight requests this node may admit for the bucket
-     * @return a {@link Releasable} that decrements the bucket's in-flight count exactly once when closed
-     * @throws OpenSearchRejectedExecutionException (HTTP 429) if the bucket is already at the limit
+     * @return a {@link Releasable} that decrements the bucket's in-flight count exactly once when closed, or
+     *         {@code null} if the per-node limit is already reached (nothing was acquired, nothing to release)
      */
-    public Releasable acquire(String bucketKey, int nodeLimit) {
+    public Releasable tryAcquire(String bucketKey, int nodeLimit) {
         // Create-and-increment atomically with respect to the decrement-and-remove in release(), so a concurrent
-        // release draining a bucket to 0 can never orphan the counter this acquire is about to use.
+        // release draining a bucket to 0 can never orphan the counter this acquire is about to use. Capture THIS
+        // caller's post-increment value from inside compute() and decide on it — reading counter.get() afterwards
+        // would observe other concurrent acquirers' increments too, causing spurious over-declines (all racers seeing
+        // an inflated count and rolling back, admitting zero when a slot was free).
+        final int[] countAfterIncrement = new int[1];
         AtomicInteger counter = inFlightByBucket.compute(bucketKey, (k, existing) -> {
             AtomicInteger c = existing != null ? existing : new AtomicInteger(0);
-            c.incrementAndGet();
+            countAfterIncrement[0] = c.incrementAndGet();
             return c;
         });
-        if (counter.get() > nodeLimit) {
-            // Over the cap: roll back this increment and reject. (Reading get() after compute may over-reject under
-            // concurrent acquires on the same bucket, but never admits over the limit — the safe direction.)
+        if (countAfterIncrement[0] > nodeLimit) {
+            // Over the cap: roll back this increment and decline. Never admits over the limit.
             release(bucketKey, counter);
-            throw new OpenSearchRejectedExecutionException(
-                "Node-level workload group throttle limit reached: " + nodeLimit + " concurrent in-flight requests"
-            );
+            return null;
         }
         return releaseOnce(bucketKey, counter);
     }