diff --git a/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmClusterThrottlingIT.java b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmClusterThrottlingIT.java
new file mode 100644
index 0000000000000..84da62a996194
--- /dev/null
+++ b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmClusterThrottlingIT.java
@@ -0,0 +1,386 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.wlm;
+
+import org.apache.logging.log4j.LogManager;
+import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest;
+import org.opensearch.action.admin.cluster.wlm.WlmStatsAction;
+import org.opensearch.action.admin.cluster.wlm.WlmStatsRequest;
+import org.opensearch.action.admin.cluster.wlm.WlmStatsResponse;
+import org.opensearch.action.index.IndexResponse;
+import org.opensearch.action.search.SearchRequestBuilder;
+import org.opensearch.action.search.SearchResponse;
+import org.opensearch.action.support.WriteRequest;
+import org.opensearch.cluster.metadata.WorkloadGroup;
+import org.opensearch.common.action.ActionFuture;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
+import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType;
+import org.opensearch.plugins.Plugin;
+import org.opensearch.plugins.PluginsService;
+import org.opensearch.rule.RuleAttribute;
+import org.opensearch.rule.RuleFrameworkPlugin;
+import org.opensearch.rule.RulePersistenceServiceRegistry;
+import org.opensearch.rule.RuleRoutingServiceRegistry;
+import org.opensearch.rule.action.CreateRuleAction;
+import org.opensearch.rule.action.CreateRuleRequest;
+import org.opensearch.rule.autotagging.AutoTaggingRegistry;
+import org.opensearch.rule.autotagging.FeatureType;
+import org.opensearch.rule.autotagging.Rule;
+import org.opensearch.script.MockScriptPlugin;
+import org.opensearch.script.Script;
+import org.opensearch.script.ScriptType;
+import org.opensearch.search.lookup.LeafFieldsLookup;
+import org.opensearch.test.OpenSearchIntegTestCase;
+import org.opensearch.wlm.MutableWorkloadGroupFragment;
+import org.opensearch.wlm.ResourceType;
+import org.opensearch.wlm.WorkloadGroupThrottleSettings;
+import org.opensearch.wlm.WorkloadManagementSettings;
+import org.opensearch.wlm.stats.WlmStats;
+import org.opensearch.wlm.stats.WorkloadGroupStats.WorkloadGroupStatsHolder;
+import org.joda.time.Instant;
+import org.junit.After;
+import org.junit.Before;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+import java.util.function.ToLongFunction;
+
+import static org.opensearch.index.query.QueryBuilders.scriptQuery;
+import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
+
+/**
+ * End-to-end integration test for cluster-level ({@code shared_limit}) WLM request throttling across a multi-node
+ * cluster. Whereas {@link WlmNodeThrottlingIT} proves the per-node tier, this verifies the distributed shared tier:
+ * a bucket's cluster-wide in-flight count is enforced by a single owner node regardless of which coordinator a
+ * request lands on, so concurrent searches spread across coordinators still cap at {@code shared_limit}.
+ *
+ * The group sets only {@code shared_limit} (no {@code node_limit}), so every admitted request must consult the
+ * bucket owner — exercising both the local-owner short-circuit and the cross-node acquire RPC depending on where
+ * the coordinator sits relative to the owner.
+ */
+@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 3, numClientNodes = 0, supportsDedicatedMasters = false)
+public class WlmClusterThrottlingIT extends OpenSearchIntegTestCase {
+
+ private static final TimeValue TIMEOUT = new TimeValue(30, TimeUnit.SECONDS);
+ private static final String PUT = "PUT";
+
+ @Override
+ protected Collection> nodePlugins() {
+ List> plugins = new ArrayList<>(super.nodePlugins());
+ plugins.add(WlmAutoTaggingIT.TestWorkloadManagementPlugin.class);
+ plugins.add(RuleFrameworkPlugin.class);
+ plugins.add(ClusterScriptedBlockPlugin.class);
+ return plugins;
+ }
+
+ @Before
+ public void registerFeatureTypeIfMissingOnAllNodes() {
+ AutoTaggingRegistry.featureTypesRegistryMap.remove(WorkloadGroupFeatureType.NAME);
+ FeatureType featureType = WlmAutoTaggingIT.TestWorkloadManagementPlugin.featureType;
+ AutoTaggingRegistry.registerFeatureType(featureType);
+
+ for (String node : internalCluster().getNodeNames()) {
+ RulePersistenceServiceRegistry persistenceRegistry = internalCluster().getInstance(RulePersistenceServiceRegistry.class, node);
+ RuleRoutingServiceRegistry routingRegistry = internalCluster().getInstance(RuleRoutingServiceRegistry.class, node);
+ try {
+ routingRegistry.getRuleRoutingService(featureType);
+ } catch (IllegalArgumentException ex) {
+ persistenceRegistry.register(featureType, WlmAutoTaggingIT.TestWorkloadManagementPlugin.rulePersistenceService);
+ routingRegistry.register(featureType, WlmAutoTaggingIT.TestWorkloadManagementPlugin.ruleRoutingService);
+ }
+ }
+ }
+
+ @After
+ public void clearWlmModeSetting() throws Exception {
+ Settings.Builder builder = Settings.builder().putNull(WorkloadManagementSettings.WLM_MODE_SETTING.getKey());
+ assertAcked(client().admin().cluster().prepareUpdateSettings().setPersistentSettings(builder).get());
+ }
+
+ public void testClusterWideCeilingHoldsAcrossCoordinators() throws Exception {
+ String workloadGroupId = "wlm_shared_throttle_group";
+ String ruleId = "wlm_shared_throttle_rule";
+ String indexName = "shared_throttle_index";
+
+ setWlmMode("enabled");
+
+ // shared_limit = 2, node_limit unset -> a cluster-wide ceiling of 2 in-flight for the whole group,
+ // enforced by the bucket owner no matter which coordinator receives the request.
+ WorkloadGroup workloadGroup = createSharedThrottledGroup("shared_test_group", workloadGroupId, 2);
+ updateWorkloadGroupInClusterState(PUT, workloadGroup);
+
+ assertBusy(() -> {
+ boolean present = client().admin()
+ .cluster()
+ .prepareState()
+ .get()
+ .getState()
+ .metadata()
+ .workloadGroups()
+ .containsKey(workloadGroupId);
+ assertTrue("workload group not yet applied in cluster state", present);
+ }, 30, TimeUnit.SECONDS);
+
+ FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME);
+ createRule(ruleId, "shared throttle rule", indexName, featureType, workloadGroupId);
+ indexDocument(indexName);
+
+ // Wait for rule propagation on EVERY coordinator this test drives, not just a random one. The rule refresh is
+ // asynchronous and can reach nodes at different times; a search on an un-refreshed coordinator stays untagged
+ // and bypasses admit(), which would break the ceiling assertions below. Poll each named node's own client
+ // until a search through it is tagged to the group (its completions advance).
+ for (String node : internalCluster().getNodeNames()) {
+ assertBusy(() -> {
+ long before = getCompletions(workloadGroupId);
+ try {
+ client(node).prepareSearch(indexName).setQuery(QueryBuilders.matchAllQuery()).get();
+ } catch (Exception e) {
+ // Each probe consults the shared tier (shared_limit=2, node_limit unset) and briefly holds a permit;
+ // remote releases are fire-and-forget, so a burst of probes can transiently exhaust the budget and
+ // throttle this one. That is not a propagation failure — turn the 429 into an AssertionError so
+ // assertBusy retries it (assertBusy only re-runs on AssertionError; a raw 429 would escape). Any
+ // other exception is a real failure and propagates.
+ assertFalse("transient throttle during propagation probe — retry: " + e, hasRejectedExecutionCause(e));
+ throw e;
+ }
+ long after = getCompletions(workloadGroupId);
+ assertTrue("search via [" + node + "] not yet tagged to the throttled workload group", after > before);
+ }, 30, TimeUnit.SECONDS);
+ }
+
+ List plugins = initBlockFactory();
+
+ List coordinators = new ArrayList<>(List.of(internalCluster().getNodeNames()));
+ // Fill the cluster-wide budget of 2 using two DIFFERENT coordinators.
+ ActionFuture first = blockingSearchVia(coordinators.get(0), indexName).execute();
+ ActionFuture second = blockingSearchVia(coordinators.get(1 % coordinators.size()), indexName).execute();
+ awaitBlockedCount(plugins, 2);
+
+ long throttledBefore = getThrottled(workloadGroupId);
+
+ // A third concurrent search on yet another coordinator must be rejected: the cluster-wide ceiling of 2 is
+ // reached, and the owner (a single node) sees the count regardless of coordinator.
+ String thirdCoordinator = coordinators.get(2 % coordinators.size());
+ Throwable rejection = expectThrows(Throwable.class, () -> blockingSearchVia(thirdCoordinator, indexName).get());
+ assertTrue(
+ "Expected an OpenSearchRejectedExecutionException in the cause chain but was: " + rejection,
+ hasRejectedExecutionCause(rejection)
+ );
+ assertEquals("total_throttled should increment by exactly one", throttledBefore + 1, getThrottled(workloadGroupId));
+
+ // Release the blocks; the two admitted searches complete successfully.
+ disableBlocks(plugins);
+ assertNotNull(first.actionGet(TIMEOUT));
+ assertNotNull(second.actionGet(TIMEOUT));
+
+ // Once the budget drains, a new search is admitted again. The two fills may have released REMOTELY (when the
+ // bucket owner is a third node, release is a fire-and-forget RPC with no happens-before to this re-admission),
+ // so the owner's in-flight count may not have dropped the instant actionGet() above returned. Poll: retry a
+ // non-blocking search until it is admitted (tolerating a transient 429 during the drain window), rather than
+ // assuming the drain is immediately visible.
+ assertBusy(() -> {
+ try {
+ client(thirdCoordinator).prepareSearch(indexName).setQuery(QueryBuilders.matchAllQuery()).get();
+ } catch (Exception e) {
+ assertFalse("still draining (transient 429) — retry", hasRejectedExecutionCause(e));
+ throw e;
+ }
+ }, 30, TimeUnit.SECONDS);
+ }
+
+ // Helpers
+
+ private static boolean hasRejectedExecutionCause(Throwable t) {
+ for (Throwable cur = t; cur != null; cur = cur.getCause()) {
+ if (cur instanceof OpenSearchRejectedExecutionException) {
+ return true;
+ }
+ if (cur.getCause() == cur) {
+ break;
+ }
+ }
+ return false;
+ }
+
+ private long getCompletions(String groupId) throws Exception {
+ return sumAcrossNodes(groupId, WorkloadGroupStatsHolder::getCompletions);
+ }
+
+ private long getThrottled(String groupId) throws Exception {
+ return sumAcrossNodes(groupId, WorkloadGroupStatsHolder::getThrottled);
+ }
+
+ // Sums a per-group stat across all nodes using the typed WlmStats response (no brittle string parsing).
+ private long sumAcrossNodes(String groupId, ToLongFunction extractor) throws Exception {
+ WlmStatsRequest request = new WlmStatsRequest(null, new HashSet<>(Collections.singletonList(groupId)), null);
+ WlmStatsResponse response = client().execute(WlmStatsAction.INSTANCE, request).get();
+ long total = 0;
+ for (WlmStats nodeStats : response.getNodes()) {
+ WorkloadGroupStatsHolder holder = nodeStats.getWorkloadGroupStats().getStats().get(groupId);
+ if (holder != null) {
+ total += extractor.applyAsLong(holder);
+ }
+ }
+ return total;
+ }
+
+ private SearchRequestBuilder blockingSearchVia(String nodeName, String indexName) {
+ return client(nodeName).prepareSearch(indexName)
+ .setQuery(
+ scriptQuery(new Script(ScriptType.INLINE, "mockscript", ClusterScriptedBlockPlugin.SCRIPT_NAME, Collections.emptyMap()))
+ );
+ }
+
+ private List initBlockFactory() {
+ List plugins = new ArrayList<>();
+ for (PluginsService pluginsService : internalCluster().getDataNodeInstances(PluginsService.class)) {
+ plugins.addAll(pluginsService.filterPlugins(ClusterScriptedBlockPlugin.class));
+ }
+ for (ClusterScriptedBlockPlugin plugin : plugins) {
+ plugin.reset();
+ plugin.enableBlock();
+ }
+ return plugins;
+ }
+
+ private void awaitBlockedCount(List plugins, int expected) throws Exception {
+ assertBusy(() -> {
+ int blocked = 0;
+ for (ClusterScriptedBlockPlugin plugin : plugins) {
+ blocked += plugin.hits.get();
+ }
+ assertEquals("expected exactly " + expected + " searches blocked in-flight", expected, blocked);
+ }, 30, TimeUnit.SECONDS);
+ }
+
+ private void disableBlocks(List plugins) {
+ for (ClusterScriptedBlockPlugin plugin : plugins) {
+ plugin.disableBlock();
+ }
+ }
+
+ private void createRule(String ruleId, String ruleName, String indexPattern, FeatureType featureType, String workloadGroupId)
+ throws Exception {
+ Rule rule = new Rule(
+ ruleId,
+ ruleName,
+ Map.of(RuleAttribute.INDEX_PATTERN, Set.of(indexPattern)),
+ featureType,
+ workloadGroupId,
+ Instant.now().toString()
+ );
+ client().execute(CreateRuleAction.INSTANCE, new CreateRuleRequest(rule)).get();
+ }
+
+ private void setWlmMode(String mode) throws Exception {
+ Settings.Builder settings = Settings.builder().put("wlm.workload_group.mode", mode);
+ ClusterUpdateSettingsRequest request = new ClusterUpdateSettingsRequest().persistentSettings(settings);
+ client().admin().cluster().updateSettings(request).get();
+ }
+
+ private WorkloadGroup createSharedThrottledGroup(String name, String id, int sharedLimit) {
+ Settings throttling = Settings.builder()
+ .put(WorkloadGroupThrottleSettings.ATTRIBUTE.getKey(), "group")
+ .put(WorkloadGroupThrottleSettings.SHARED_LIMIT.getKey(), sharedLimit)
+ .build();
+ return new WorkloadGroup(
+ name,
+ id,
+ new MutableWorkloadGroupFragment(
+ MutableWorkloadGroupFragment.ResiliencyMode.SOFT,
+ Map.of(ResourceType.CPU, 0.9, ResourceType.MEMORY, 0.9),
+ Settings.EMPTY,
+ throttling
+ ),
+ Instant.now().getMillis()
+ );
+ }
+
+ private void indexDocument(String indexName) {
+ assertAcked(
+ client().admin()
+ .indices()
+ .prepareCreate(indexName)
+ // One shard so each search produces exactly one blocking script hit, making the in-flight count exact.
+ // Throttle admission happens on the coordinator that receives the request, independent of where the
+ // shard lives, so a single shard is sufficient to exercise the cross-coordinator ceiling.
+ .setSettings(Settings.builder().put("index.number_of_shards", 1).put("index.number_of_replicas", 0))
+ );
+ IndexResponse response = client().prepareIndex(indexName)
+ .setId("1")
+ .setSource(Map.of("field", "value"))
+ .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
+ .get();
+ assertEquals(org.opensearch.action.DocWriteResponse.Result.CREATED, response.getResult());
+ }
+
+ private void updateWorkloadGroupInClusterState(String method, WorkloadGroup workloadGroup) throws InterruptedException {
+ WlmAutoTaggingIT.ExceptionCatchingListener listener = new WlmAutoTaggingIT.ExceptionCatchingListener();
+ client().execute(
+ WlmAutoTaggingIT.TestClusterUpdateTransportAction.ACTION,
+ new WlmAutoTaggingIT.TestClusterUpdateRequest(workloadGroup, method),
+ listener
+ );
+ boolean completed = listener.getLatch().await(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+ assertTrue("cluster-state update did not complete in time", completed);
+ if (listener.getException() != null) {
+ throw new AssertionError("cluster-state update failed", listener.getException());
+ }
+ }
+
+ /**
+ * Test script plugin that blocks during the query phase until released, keeping a search in-flight on whichever
+ * coordinator dispatched it.
+ */
+ public static class ClusterScriptedBlockPlugin extends MockScriptPlugin {
+ static final String SCRIPT_NAME = "cluster_search_block";
+
+ private final AtomicInteger hits = new AtomicInteger();
+ private final AtomicBoolean shouldBlock = new AtomicBoolean(true);
+
+ public void reset() {
+ hits.set(0);
+ }
+
+ public void disableBlock() {
+ shouldBlock.set(false);
+ }
+
+ public void enableBlock() {
+ shouldBlock.set(true);
+ }
+
+ @Override
+ public Map, Object>> pluginScripts() {
+ return Collections.singletonMap(SCRIPT_NAME, params -> {
+ LeafFieldsLookup fieldsLookup = (LeafFieldsLookup) params.get("_fields");
+ LogManager.getLogger(WlmClusterThrottlingIT.class).info("Blocking on the document {}", fieldsLookup.get("_id"));
+ hits.incrementAndGet();
+ try {
+ assertBusy(() -> assertFalse(shouldBlock.get()));
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ return true;
+ });
+ }
+ }
+}
diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
index 5785688fa244c..0d781c597224e 100644
--- a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
+++ b/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 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 admissionListener = ContextPreservingActionListener.wrapPreservingContext(
+ ActionListener.wrap(throttlePermit -> {
+ ActionListener 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 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 listener;
diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java
index 832a9721ce7e5..a7104e5a45862 100644
--- a/server/src/main/java/org/opensearch/node/Node.java
+++ b/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 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 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();
diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java b/server/src/main/java/org/opensearch/wlm/NodeThrottleTracker.java
similarity index 61%
rename from server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java
rename to server/src/main/java/org/opensearch/wlm/NodeThrottleTracker.java
index 1c8ededc7f5dd..428098dd30e74 100644
--- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java
+++ b/server/src/main/java/org/opensearch/wlm/NodeThrottleTracker.java
@@ -10,7 +10,6 @@
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;
@@ -18,7 +17,9 @@
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).
*
* 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
@@ -26,36 +27,43 @@
* number of concurrently active buckets rather than the total population of users/roles.
*
* 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 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);
}
diff --git a/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java b/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java
new file mode 100644
index 0000000000000..8a005c1cfddf8
--- /dev/null
+++ b/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java
@@ -0,0 +1,139 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.opensearch.common.annotation.ExperimentalApi;
+
+import java.util.Iterator;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.LongSupplier;
+
+/**
+ * Cluster-level tier of workload-group request throttling: the authoritative, cluster-wide in-flight counter for a
+ * throttle bucket, enforcing the bucket's {@code shared_limit}. Its node-local sibling is {@link NodeThrottleTracker},
+ * and {@code WorkloadGroupService.acquireThrottlePermit} composes the two (local first, then shared on overflow). On the node that
+ * owns a bucket (chosen by {@link ThrottleOwnerSelector}), this holds one live-lease set per active bucket.
+ *
+ * Unlike the node tier, acquire and release happen on different nodes (a coordinator acquires from the owner, then
+ * releases via a separate RPC), so each permit is a lease with an id and a TTL rather than a self-releasing
+ * {@link org.opensearch.common.lease.Releasable}.
+ *
+ * Correctness under concurrency and node churn:
+ *
+ * - Atomic admission. The check-count-then-add step runs inside {@link ConcurrentHashMap#compute},
+ * which holds the per-key bin lock, so many coordinators racing on the same bucket can never push the live
+ * count over the limit.
+ * - Count = size of the live-lease map, never a bare integer. This makes {@link #release} and
+ * {@link #sweepExpired} idempotent: removing by lease id is a no-op if the lease is already gone (already
+ * swept, or the bucket was remapped to a different owner), so a sweep racing a late release cannot
+ * double-decrement, and a stray release after ring remap cannot corrupt a live count.
+ * - TTL leases. Every lease carries an expiry. A coordinator that crashes (or whose release RPC is
+ * lost) leaves a lease that {@link #sweepExpired} reclaims after the TTL, so a bucket cannot get wedged at
+ * its limit forever. A rejoining owner starts empty — no phantom leases.
+ *
+ * A bucket entry exists only while it has at least one live lease, so memory scales with concurrently-active
+ * buckets, not the total user/role population.
+ */
+@ExperimentalApi
+public class SharedThrottleTracker {
+
+ private final Map> leasesByBucket = new ConcurrentHashMap<>();
+ private final LongSupplier nanoTimeSupplier;
+
+ public SharedThrottleTracker() {
+ this(System::nanoTime);
+ }
+
+ public SharedThrottleTracker(LongSupplier nanoTimeSupplier) {
+ this.nanoTimeSupplier = nanoTimeSupplier;
+ }
+
+ /**
+ * Attempts to admit one request into a bucket under its cluster-wide shared limit.
+ *
+ * @param bucketKey the throttle bucket identifier
+ * @param sharedLimit the maximum concurrent in-flight requests allowed across the cluster for this bucket
+ * @param leaseId a globally-unique lease id minted by the requesting coordinator
+ * @param ttlNanos how long the lease may live before {@link #sweepExpired} may reclaim it
+ * @return {@code true} if the lease was granted (live count was below the limit), {@code false} if the bucket
+ * is already at the limit
+ */
+ public boolean tryAcquire(String bucketKey, int sharedLimit, String leaseId, long ttlNanos) {
+ final long expiresAt = nanoTimeSupplier.getAsLong() + ttlNanos;
+ final boolean[] granted = new boolean[1];
+ leasesByBucket.compute(bucketKey, (k, leases) -> {
+ if (leases == null) {
+ leases = new ConcurrentHashMap<>();
+ }
+ // Only prune when the bucket looks full. An expired lease can change the outcome only when we would
+ // otherwise reject; below the limit there is a free slot regardless, so we skip the O(size) scan and keep
+ // the common under-limit acquire O(1). A saturated bucket still self-heals its stuck leases right here, at
+ // the one moment it matters. (Leases in an idle bucket are reclaimed by the periodic sweep instead.)
+ if (leases.size() >= sharedLimit) {
+ pruneExpired(leases);
+ }
+ if (leases.size() < sharedLimit) {
+ leases.put(leaseId, expiresAt);
+ granted[0] = true;
+ }
+ // Never leave an empty map behind (keeps the live set == active buckets).
+ return leases.isEmpty() ? null : leases;
+ });
+ return granted[0];
+ }
+
+ /**
+ * Releases a previously-granted lease. Removing by lease id makes this idempotent and safe: an unknown id
+ * (already swept, double release, or a lease that belonged to a previous owner of a since-remapped bucket) is
+ * a no-op and can never decrement a live lease that belongs to some other request.
+ *
+ * @param bucketKey the throttle bucket identifier
+ * @param leaseId the lease id returned to the coordinator at acquire time
+ */
+ public void release(String bucketKey, String leaseId) {
+ leasesByBucket.computeIfPresent(bucketKey, (k, leases) -> {
+ leases.remove(leaseId);
+ return leases.isEmpty() ? null : leases;
+ });
+ }
+
+ /**
+ * Reclaims all expired leases across every bucket. Intended to be called periodically by the owning service.
+ * Safe to run concurrently with {@link #tryAcquire}/{@link #release} thanks to per-key {@code compute}.
+ */
+ public void sweepExpired() {
+ for (String bucketKey : leasesByBucket.keySet()) {
+ leasesByBucket.computeIfPresent(bucketKey, (k, leases) -> {
+ pruneExpired(leases);
+ return leases.isEmpty() ? null : leases;
+ });
+ }
+ }
+
+ // Current live (non-expired-at-read-time) count for a bucket. Package-private for tests.
+ int inFlight(String bucketKey) {
+ Map leases = leasesByBucket.get(bucketKey);
+ return leases == null ? 0 : leases.size();
+ }
+
+ // Number of buckets currently holding at least one lease. Package-private for tests.
+ int activeBuckets() {
+ return leasesByBucket.size();
+ }
+
+ private void pruneExpired(Map leases) {
+ final long now = nanoTimeSupplier.getAsLong();
+ for (Iterator> it = leases.entrySet().iterator(); it.hasNext();) {
+ if (it.next().getValue() <= now) {
+ it.remove();
+ }
+ }
+ }
+}
diff --git a/server/src/main/java/org/opensearch/wlm/ThrottleOwnerSelector.java b/server/src/main/java/org/opensearch/wlm/ThrottleOwnerSelector.java
new file mode 100644
index 0000000000000..c311e6ec34bce
--- /dev/null
+++ b/server/src/main/java/org/opensearch/wlm/ThrottleOwnerSelector.java
@@ -0,0 +1,141 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.opensearch.Version;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.node.DiscoveryNodes;
+import org.opensearch.cluster.routing.Murmur3HashFunction;
+import org.opensearch.common.annotation.ExperimentalApi;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.SortedMap;
+import java.util.TreeMap;
+
+/**
+ * Maps a throttle bucket to the single node that authoritatively owns its cluster-level counter, using a
+ * consistent-hash ring built from cluster state. Every node computes the same ring from the same
+ * {@link DiscoveryNodes}, so any coordinator can locate a bucket's owner with no election or lookup, and node
+ * join/leave only remaps roughly {@code 1/N} of buckets.
+ *
+ * Design choices that matter:
+ *
+ * - Data nodes only. Cluster-manager nodes are deliberately excluded — a domain may run zero or one
+ * dedicated cluster-manager, so they are not a dependable ownership pool.
+ * - Version-filtered. During a rolling upgrade, nodes older than {@link #MIN_OWNER_VERSION} (which do
+ * not run the shared-throttle handler) are excluded from the ring, so every bucket lands on a node that can
+ * actually enforce it rather than silently going unthrottled. This trades extra remapping during the
+ * upgrade window for real enforcement.
+ * - Virtual nodes. Each physical node is placed at {@link #VIRTUAL_NODES} points on the ring so buckets
+ * spread evenly and a departing node's share redistributes across the survivors rather than piling onto one
+ * neighbour.
+ * - Empty ring ⇒ no owner. When no eligible node exists (bootstrap, coordinating-only topologies, or a
+ * fleet with no upgraded data nodes yet), {@link #ownerFor} returns empty and the caller fails open.
+ *
+ * Instances are immutable snapshots; {@link WorkloadGroupSharedThrottleService} rebuilds one whenever the
+ * discovery-node set changes.
+ *
+ * Transient ceiling breach on rebalance. When a bucket remaps to a new owner (node join/leave), the previous
+ * owner still holds leases for that bucket's in-flight requests while the new owner starts counting from zero, so the
+ * bucket's cluster-wide in-flight count can briefly exceed {@code shared_limit} (up to ~2x) until the old leases drain
+ * or TTL out. This is inherent to stateless consistent hashing and is accepted for this experimental feature.
+ */
+@ExperimentalApi
+public class ThrottleOwnerSelector {
+
+ /**
+ * First version whose nodes run the shared-throttle transport handler and may therefore own buckets. Nodes below
+ * this are excluded from the ring so a bucket never maps to an owner that would answer the acquire RPC with an
+ * unknown-action failure (which fails open, silently disabling {@code shared_limit} for that bucket).
+ *
+ * TODO(merge): set this to the actual release that first ships these handlers before upstreaming. On this feature
+ * branch it is a placeholder; if the handlers land in a build after {@code V_3_7_0}, this must be bumped to that
+ * version, otherwise a rolling upgrade with older-but-same-major data nodes would map buckets to nodes lacking the
+ * handler. (Tracked alongside the throttle-settings version-gate TODO.)
+ */
+ public static final Version MIN_OWNER_VERSION = Version.V_3_7_0;
+
+ private static final int VIRTUAL_NODES = 128;
+
+ // Ring position -> node id. TreeMap gives the ceiling/first lookup that consistent hashing needs.
+ private final SortedMap ring;
+ private final Map nodesById;
+
+ private ThrottleOwnerSelector(SortedMap ring, Map nodesById) {
+ this.ring = ring;
+ this.nodesById = nodesById;
+ }
+
+ /**
+ * Builds a ring from the current discovery nodes, including only data nodes at or above
+ * {@link #MIN_OWNER_VERSION}.
+ */
+ public static ThrottleOwnerSelector fromDiscoveryNodes(DiscoveryNodes discoveryNodes) {
+ final SortedMap ring = new TreeMap<>();
+ final Map nodesById = new HashMap<>();
+ for (DiscoveryNode node : discoveryNodes.getDataNodes().values()) {
+ if (node.getVersion().onOrAfter(MIN_OWNER_VERSION) == false) {
+ continue;
+ }
+ nodesById.put(node.getId(), node);
+ for (int i = 0; i < VIRTUAL_NODES; i++) {
+ // Hash the node id with a replica suffix to scatter its virtual points around the ring.
+ ring.put(Murmur3HashFunction.hash(node.getId() + "#" + i), node.getId());
+ }
+ }
+ return new ThrottleOwnerSelector(ring, nodesById);
+ }
+
+ /**
+ * @return the node that owns the given bucket, or empty if the ring has no eligible node (fail open).
+ */
+ public Optional ownerFor(String bucketKey) {
+ if (ring.isEmpty()) {
+ return Optional.empty();
+ }
+ final int hash = Murmur3HashFunction.hash(bucketKey);
+ // First ring point at or after the bucket's hash, wrapping around to the first point.
+ SortedMap tail = ring.tailMap(hash);
+ Integer point = tail.isEmpty() ? ring.firstKey() : tail.firstKey();
+ return Optional.ofNullable(nodesById.get(ring.get(point)));
+ }
+
+ /**
+ * The eligible owner nodes in this snapshot as a set of full {@link DiscoveryNode}s. Used to decide whether a
+ * rebuild is needed: comparing whole nodes (not just persistent ids) detects a same-id restart, where the node
+ * keeps its persistent id but gets a new ephemeral id/address — otherwise the ring would keep the stale node and
+ * {@code nodeConnected} would fail open for its buckets indefinitely.
+ */
+ Set eligibleNodeSet() {
+ return new HashSet<>(nodesById.values());
+ }
+
+ /**
+ * Computes the eligible owner set for a candidate {@link DiscoveryNodes} without building the full ring
+ * (no virtual-node hashing). Lets a caller cheaply check whether membership actually changed before paying to
+ * rebuild the ring on every cluster-state update.
+ */
+ static Set eligibleNodeSet(DiscoveryNodes discoveryNodes) {
+ Set eligible = new HashSet<>();
+ for (DiscoveryNode node : discoveryNodes.getDataNodes().values()) {
+ if (node.getVersion().onOrAfter(MIN_OWNER_VERSION)) {
+ eligible.add(node);
+ }
+ }
+ return eligible;
+ }
+
+ boolean isEmpty() {
+ return ring.isEmpty();
+ }
+}
diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java
index b45f3fb81bddc..a7e385f74ae00 100644
--- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java
+++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java
@@ -19,6 +19,7 @@
import org.opensearch.common.lease.Releasable;
import org.opensearch.common.lifecycle.AbstractLifecycleComponent;
import org.opensearch.common.settings.Settings;
+import org.opensearch.core.action.ActionListener;
import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
import org.opensearch.monitor.jvm.JvmStats;
import org.opensearch.monitor.process.ProcessProbe;
@@ -62,7 +63,9 @@ public class WorkloadGroupService extends AbstractLifecycleComponent
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.
- private final WorkloadGroupThrottleTracker throttleTracker = new WorkloadGroupThrottleTracker();
+ private final NodeThrottleTracker throttleTracker = new NodeThrottleTracker();
+ // Cluster-level (shared_limit) tier; late-bound after the transport service exists. Null => only the local tier.
+ private volatile WorkloadGroupSharedThrottleService sharedThrottleService;
public WorkloadGroupService(
WorkloadGroupTaskCancellationService taskCancellationService,
@@ -317,73 +320,184 @@ public void rejectIfNeeded(String workloadGroupId) {
}
/**
- * Acquires one node-level throttle permit for the request, or returns {@code null} (nothing to release) when the
- * request is not throttled: WLM disabled, default/unknown group, no {@code node_limit}, or no resolvable bucket
- * (see {@link #resolveThrottleBucketKey}). The bucket depends on the group's throttle {@code attribute}.
+ * Late-binds the cluster-level ({@code shared_limit}) throttle tier. It is constructed after the transport
+ * service (which it needs for the owner round-trip), so it cannot be a constructor argument. When unset, only
+ * the node-local tier applies and a shared-only config fails open.
+ */
+ public void setSharedThrottleService(WorkloadGroupSharedThrottleService sharedThrottleService) {
+ this.sharedThrottleService = sharedThrottleService;
+ }
+
+ /**
+ * Two-tier throttle admission for a search request. Notifies {@code listener} with:
+ *
+ * - a non-null {@link Releasable} — admitted; close it exactly once on request completion to release the
+ * permit (works for both a node-local and a cluster-level shared permit);
+ * - {@code null} — admitted but not throttle-tracked (throttling disabled/not configured for this request,
+ * request not attributable to a bucket, or a fail-open path); nothing to release;
+ * - {@link ActionListener#onFailure} with an {@link OpenSearchRejectedExecutionException} (HTTP 429) — the
+ * bucket is at its limit.
+ *
+ * The node-local tier is checked synchronously (zero added latency on the common path); only an overflow to the
+ * shared tier does the asynchronous owner round-trip, so the calling thread is never blocked. The listener may
+ * therefore be invoked inline or, for the shared-tier overflow, on a transport thread.
*
* @param workloadGroupId the workload group the request is assigned to
- * @param principal the raw {@code WORKLOAD_GROUP_PRINCIPAL_HEADER} value, or {@code null} (see resolver)
- * @return a permit to close on request completion, or {@code null} if not throttled
- * @throws OpenSearchRejectedExecutionException if the bucket is already at its node limit
+ * @param principal the raw {@code WORKLOAD_GROUP_PRINCIPAL_HEADER} value, or {@code null}
+ * @param listener receives the permit / null / 429
*/
- public Releasable acquireThrottleOrReject(String workloadGroupId, String principal) {
+ public void acquireThrottlePermit(String workloadGroupId, String principal, ActionListener listener) {
+ final ThrottlePlan plan;
+ try {
+ plan = resolveThrottlePlan(workloadGroupId, principal);
+ } catch (Exception e) {
+ // A bug in the throttle-resolution path must never fail an otherwise-valid search, so fail open.
+ logger.warn("Skipping throttle for workload group [" + workloadGroupId + "] due to an error", e);
+ listener.onResponse(null);
+ return;
+ }
+ if (plan == null) {
+ listener.onResponse(null); // not throttled / not attributable -> fail open
+ return;
+ }
+
+ // Node-local tier: synchronous, no cross-node coordination. Granting here is the zero-latency common path.
+ if (plan.nodeLimit != WorkloadGroupThrottleSettings.UNSET_LIMIT) {
+ Releasable localPermit = throttleTracker.tryAcquire(plan.bucketKey, plan.nodeLimit);
+ if (localPermit != null) {
+ listener.onResponse(localPermit);
+ return;
+ }
+ // Local allowance exhausted -> fall through to the shared pool if one is configured.
+ }
+
+ // Cluster-level shared tier (asynchronous owner round-trip).
+ if (plan.sharedLimit != WorkloadGroupThrottleSettings.UNSET_LIMIT && sharedThrottleService != null) {
+ sharedThrottleService.acquireAsync(plan.bucketKey, plan.sharedLimit, ActionListener.wrap(listener::onResponse, e -> {
+ if (e instanceof OpenSearchRejectedExecutionException) {
+ // Recompose the denial with the human-readable group name/attribute (the shared tier only has the
+ // opaque bucket key), and count it.
+ incrementThrottled(workloadGroupId);
+ listener.onFailure(sharedTierRejection(plan));
+ } else {
+ listener.onFailure(e);
+ }
+ }));
+ return;
+ }
+
+ if (plan.nodeLimit != WorkloadGroupThrottleSettings.UNSET_LIMIT && plan.sharedLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) {
+ // Node-only config with the local allowance exhausted and no shared tier to overflow to -> reject.
+ incrementThrottled(workloadGroupId);
+ listener.onFailure(nodeTierRejection(plan));
+ } else {
+ // Either a shared tier was configured but is unavailable (not yet wired), or a shared-only config with no
+ // wired tier. Fail open rather than reject, consistent with every other shared-tier-unavailable path.
+ listener.onResponse(null);
+ }
+ }
+
+ // User-facing 429 for the node-local tier. Names the group (and attribute, if keyed) and the per-node limit knob.
+ private static OpenSearchRejectedExecutionException nodeTierRejection(ThrottlePlan plan) {
+ return new OpenSearchRejectedExecutionException(
+ "Request throttled: " + plan.describeTarget() + " reached its per-node limit of " + plan.nodeLimit + " concurrent requests."
+ );
+ }
+
+ // User-facing 429 for the cluster-level shared tier. Same shape as the node message, naming the cluster-wide limit.
+ private static OpenSearchRejectedExecutionException sharedTierRejection(ThrottlePlan plan) {
+ return new OpenSearchRejectedExecutionException(
+ "Request throttled: "
+ + plan.describeTarget()
+ + " reached its cluster-wide limit of "
+ + plan.sharedLimit
+ + " concurrent requests."
+ );
+ }
+
+ // Records a throttle rejection. Uses the raw state map, not the DEFAULT-fallback accessor, so a not-yet-registered
+ // group isn't misattributed to DEFAULT, and never lets a stats failure swallow the 429.
+ private void incrementThrottled(String workloadGroupId) {
+ try {
+ WorkloadGroupState workloadGroupState = workloadGroupsStateAccessor.getWorkloadGroupStateMap().get(workloadGroupId);
+ if (workloadGroupState != null) {
+ workloadGroupState.totalThrottled.inc();
+ }
+ } catch (Exception statsException) {
+ logger.warn("Failed to record throttle stat for workload group [" + workloadGroupId + "]", statsException);
+ }
+ }
+
+ /**
+ * Resolves the throttle configuration for a request into a bucket key plus the node and shared limits, or
+ * {@code null} when the request is not throttled: WLM disabled, default/unknown group, throttling not configured
+ * (both limits unset), or the request can't be attributed to a bucket (e.g. username/role with no principal).
+ */
+ private ThrottlePlan resolveThrottlePlan(String workloadGroupId, String principal) {
if (workloadManagementSettings.getWlmMode() != WlmMode.ENABLED) {
return null;
}
if (workloadGroupId == null || workloadGroupId.equals(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get())) {
return null;
}
- try {
- WorkloadGroup workloadGroup = getWorkloadGroupById(workloadGroupId);
- if (workloadGroup == null) {
- return null;
- }
- Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling();
- int nodeLimit = WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling);
- if (nodeLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) {
- return null;
- }
- // A null bucket means the request can't be attributed (e.g. username/role with no principal) -> fail open.
- String bucketKey = resolveThrottleBucketKey(
- workloadGroupId,
- WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling),
- principal
- );
- if (bucketKey == null) {
- return null;
- }
- try {
- return throttleTracker.acquire(bucketKey, nodeLimit);
- } catch (OpenSearchRejectedExecutionException e) {
- // Record the rejection without ever letting a stats failure swallow the 429. Use the raw state map,
- // not the DEFAULT-fallback accessor, so a not-yet-registered group isn't misattributed to DEFAULT.
- try {
- WorkloadGroupState workloadGroupState = workloadGroupsStateAccessor.getWorkloadGroupStateMap().get(workloadGroupId);
- if (workloadGroupState != null) {
- workloadGroupState.totalThrottled.inc();
- }
- } catch (Exception statsException) {
- logger.warn("Failed to record throttle stat for workload group [" + workloadGroupId + "]", statsException);
- }
- throw e;
- }
- } catch (OpenSearchRejectedExecutionException e) {
- throw e; // the intended 429
- } catch (Exception e) {
- // A bug in the throttle path must never fail an otherwise-valid search, so fail open.
- logger.warn("Skipping node-level throttle for workload group [" + workloadGroupId + "] due to an error", e);
+ WorkloadGroup workloadGroup = getWorkloadGroupById(workloadGroupId);
+ if (workloadGroup == null) {
return null;
}
+ Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling();
+ int nodeLimit = WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling);
+ int sharedLimit = WorkloadGroupThrottleSettings.SHARED_LIMIT.get(throttling);
+ if (nodeLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT && sharedLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) {
+ return null; // throttling not configured
+ }
+ String attribute = WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling);
+ String value = resolveAttributeValue(attribute, principal); // null for whole-group; the principal value otherwise
+ String bucketKey = buildBucketKey(workloadGroupId, attribute, value);
+ if (bucketKey == null) {
+ return null; // can't attribute the request to a bucket -> fail open
+ }
+ return new ThrottlePlan(bucketKey, nodeLimit, sharedLimit, workloadGroup.getName(), attribute, value);
+ }
+
+ // The resolved throttle configuration for a single request. Carries the human-readable group name and the throttle
+ // dimension (attribute + resolved value) purely so a rejection message can name them; the opaque bucketKey remains
+ // the identity used by both tracker tiers.
+ private static class ThrottlePlan {
+ final String bucketKey;
+ final int nodeLimit;
+ final int sharedLimit;
+ final String groupName;
+ final String attribute; // "group" | "username" | "role"
+ final String value; // the resolved principal value for username/role; null for whole-group throttling
+
+ ThrottlePlan(String bucketKey, int nodeLimit, int sharedLimit, String groupName, String attribute, String value) {
+ this.bucketKey = bucketKey;
+ this.nodeLimit = nodeLimit;
+ this.sharedLimit = sharedLimit;
+ this.groupName = groupName;
+ this.attribute = attribute;
+ this.value = value;
+ }
+
+ // "workload group [analytics]" or "workload group [analytics] for username [alice]".
+ String describeTarget() {
+ String base = "workload group [" + groupName + "]";
+ if ("group".equals(attribute) || value == null) {
+ return base;
+ }
+ return base + " for " + attribute + " [" + value + "]";
+ }
}
/**
- * Builds the throttle bucket key: {@code :group} for whole-group throttling, or
- * {@code ::} keyed by the matching principal subfield for {@code username}/{@code role}.
- * Returns {@code null} (fail open, not throttled) when the principal is absent or lacks a token for the subfield.
+ * Resolves the throttle attribute's value for a request: {@code null} for whole-group throttling
+ * ({@code attribute == "group"}), otherwise the value of the matching {@code username}/{@code role} subfield
+ * parsed out of the principal header. Returns {@code null} (fail open, not throttled) when a keyed attribute has
+ * no usable principal token.
*/
- private String resolveThrottleBucketKey(String workloadGroupId, String attribute, String principal) {
+ private String resolveAttributeValue(String attribute, String principal) {
if ("group".equals(attribute)) {
- return workloadGroupId + ":group";
+ return null;
}
if (principal == null || principal.isEmpty()) {
return null;
@@ -394,13 +508,28 @@ private String resolveThrottleBucketKey(String workloadGroupId, String attribute
if (trimmed.startsWith(subfieldPrefix)) {
String value = trimmed.substring(subfieldPrefix.length());
if (value.isEmpty() == false) {
- return workloadGroupId + ":" + attribute + ":" + value;
+ return value;
}
}
}
return null;
}
+ /**
+ * Assembles the opaque bucket key used by both tracker tiers: {@code :group} for whole-group throttling,
+ * or {@code ::} for a {@code username}/{@code role} bucket. Returns {@code null} when a
+ * keyed attribute has no resolved value (fail open).
+ */
+ private String buildBucketKey(String workloadGroupId, String attribute, String value) {
+ if ("group".equals(attribute)) {
+ return workloadGroupId + ":group";
+ }
+ if (value == null) {
+ return null;
+ }
+ return workloadGroupId + ":" + attribute + ":" + value;
+ }
+
private double getNormalisedRejectionThreshold(double limit, ResourceType resourceType) {
if (resourceType == ResourceType.CPU) {
return limit * workloadManagementSettings.getNodeLevelCpuRejectionThreshold();
diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java
new file mode 100644
index 0000000000000..7ee5d10a5f782
--- /dev/null
+++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java
@@ -0,0 +1,417 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.cluster.ClusterChangedEvent;
+import org.opensearch.cluster.ClusterStateListener;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.node.DiscoveryNodes;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.UUIDs;
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.common.lease.Releasable;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
+import org.opensearch.core.transport.TransportResponse;
+import org.opensearch.threadpool.Scheduler;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportException;
+import org.opensearch.transport.TransportRequest;
+import org.opensearch.transport.TransportRequestOptions;
+import org.opensearch.transport.TransportResponseHandler;
+import org.opensearch.transport.TransportService;
+
+import java.io.IOException;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * Cluster-level ({@code shared_limit}) throttle tier. Each throttle bucket has one authoritative in-flight counter
+ * living on the node that owns it per a consistent-hash ring ({@link ThrottleOwnerSelector}); this service is both
+ * the coordinator-side client that acquires/releases a shared permit for an incoming request and the
+ * owner-side host of the {@link SharedThrottleTracker} that answers those requests for the buckets it owns.
+ *
+ * Latency: the acquire is asynchronous so the calling (transport) thread never blocks on the round-trip — critical
+ * under the very bursts throttling defends against. When the coordinator itself owns the bucket, the acquire
+ * short-circuits to a direct in-memory tracker call with no network hop at all.
+ *
+ * Availability: any failure to reach the owner (timeout, disconnect, missing handler on an old node, or an empty
+ * ring) fails open — the request is admitted with no shared permit. This is a deliberate, un-toggleable
+ * choice: a single unreachable owner must not turn a network blip into a cluster-wide rejection storm for its
+ * share of buckets. A stuck lease is instead reclaimed by the owner's TTL sweep.
+ */
+@ExperimentalApi
+public class WorkloadGroupSharedThrottleService implements ClusterStateListener {
+
+ /** Internal transport action names (not client-facing REST). */
+ public static final String ACQUIRE_ACTION_NAME = "internal:wlm/throttle/shared/acquire";
+ public static final String RELEASE_ACTION_NAME = "internal:wlm/throttle/shared/release";
+
+ // Fixed operational constants. Deliberately not cluster settings: they are internal-coordination knobs an operator
+ // would never need to tune, and fail-open + a generous TTL make them safe as constants. Promote to a setting later
+ // (non-breaking) only if a real deployment need emerges.
+ //
+ // Timeout for the acquire round-trip to a bucket's owner. Small so a slow/unreachable owner fails open quickly
+ // rather than adding latency to the request path.
+ static final TimeValue ACQUIRE_TIMEOUT = TimeValue.timeValueMillis(200);
+ // Time-to-live for a lease on the owner. Reclaims a lease left behind by a crashed coordinator or a lost release
+ // RPC. Set generously above the typical request duration so a still-running request's lease is not reclaimed early.
+ // Accepted tradeoff: leases are NOT renewed, so a search that runs longer than this TTL has its lease reclaimed
+ // while still executing, which can transiently admit one-or-more requests beyond shared_limit for that bucket
+ // (self-correcting, fail-open direction — same flavor as the ring-rebalance transient breach). Long-running search
+ // is uncommon (no default search timeout, but heavy aggs/scripts can exceed it); if it becomes a problem the fix is
+ // lease renewal or deriving the TTL from a request deadline, not a larger constant.
+ static final long LEASE_TTL_NANOS = TimeValue.timeValueMinutes(5).nanos();
+ // How often the owner sweeps expired leases. This is a pure memory-hygiene backstop, NOT a correctness mechanism:
+ // tryAcquire() already prunes a bucket's expired leases on every acquire, so an active bucket self-heals and can
+ // never over-reject due to a stuck lease. The sweep only reclaims the map entry for a bucket that abandoned a lease
+ // and then went completely idle (rejects nothing). An entry can't be reclaimed before its TTL elapses anyway, so
+ // sweeping faster than the TTL is pointless; run it infrequently.
+ static final TimeValue SWEEP_INTERVAL = TimeValue.timeValueMinutes(5);
+
+ private static final Logger logger = LogManager.getLogger(WorkloadGroupSharedThrottleService.class);
+
+ private final ClusterService clusterService;
+ private final ThreadPool threadPool;
+ private final TransportService transportService;
+ private final SharedThrottleTracker tracker;
+
+ // Immutable ring snapshot, swapped wholesale on discovery-node changes so readers never see a half-built ring.
+ private final AtomicReference ring = new AtomicReference<>();
+ private volatile Scheduler.Cancellable sweepTask;
+
+ public WorkloadGroupSharedThrottleService(ClusterService clusterService, ThreadPool threadPool, TransportService transportService) {
+ this(clusterService, threadPool, transportService, new SharedThrottleTracker());
+ }
+
+ public WorkloadGroupSharedThrottleService(
+ ClusterService clusterService,
+ ThreadPool threadPool,
+ TransportService transportService,
+ SharedThrottleTracker tracker
+ ) {
+ this.clusterService = clusterService;
+ this.threadPool = threadPool;
+ this.transportService = transportService;
+ this.tracker = tracker;
+ // Start with an empty ring (fail-open) rather than reading cluster state here: during node construction the
+ // ClusterApplierService has no state yet. The ring is populated by the first clusterChanged whose eligible
+ // owner set differs from this empty ring — which includes a single-node cluster's very first applied state
+ // (its node set does not "change" relative to the coordinator-seeded initial state, so we must NOT gate on
+ // nodesChanged()).
+ this.ring.set(ThrottleOwnerSelector.fromDiscoveryNodes(DiscoveryNodes.EMPTY_NODES));
+ clusterService.addListener(this);
+ transportService.registerRequestHandler(
+ ACQUIRE_ACTION_NAME,
+ ThreadPool.Names.SAME,
+ AcquireRequest::new,
+ (request, channel, task) -> channel.sendResponse(handleAcquire(request))
+ );
+ transportService.registerRequestHandler(
+ RELEASE_ACTION_NAME,
+ ThreadPool.Names.SAME,
+ ReleaseRequest::new,
+ (request, channel, task) -> {
+ tracker.release(request.bucketKey, request.leaseId);
+ channel.sendResponse(TransportResponse.Empty.INSTANCE);
+ }
+ );
+ }
+
+ /** Starts the periodic TTL sweep. Idempotent-safe to call once at node start. */
+ public void start() {
+ // The ring is populated by clusterChanged (see the constructor), not here: at node start the initial cluster
+ // state may not be applied yet, so reading clusterService.state() here can throw "initial cluster state not
+ // set yet". This method only schedules the memory-hygiene sweep.
+ sweepTask = threadPool.scheduleWithFixedDelay(() -> {
+ try {
+ tracker.sweepExpired();
+ } catch (Exception e) {
+ logger.warn("Shared throttle TTL sweep failed", e);
+ }
+ }, SWEEP_INTERVAL, ThreadPool.Names.GENERIC);
+ }
+
+ public void stop() {
+ if (sweepTask != null) {
+ sweepTask.cancel();
+ }
+ }
+
+ @Override
+ public void clusterChanged(ClusterChangedEvent event) {
+ // Rebuild whenever the eligible-owner set differs from the current ring, rather than gating on
+ // event.nodesChanged(). nodesChanged() is a delta vs the previous applied state, and the coordinator seeds the
+ // initial applied state already containing the local node — so on a single-node cluster (or any cluster whose
+ // membership is stable after this service starts) nodesChanged() is never true and the ring would stay empty,
+ // silently disabling shared throttling.
+ //
+ // Compare the eligible-node set FIRST, without building the ring (cheap: a version-filtered data-node scan, no
+ // virtual-node hashing), so an unchanged membership — the common case, including every cluster-state update
+ // while WLM is disabled — pays nothing. Compare whole DiscoveryNodes, not just persistent ids, so a same-id
+ // restart (new ephemeral id/address) is treated as a change; otherwise the ring would keep the stale node and
+ // its buckets would fail open forever.
+ Set candidate = ThrottleOwnerSelector.eligibleNodeSet(event.state().nodes());
+ if (ring.get().eligibleNodeSet().equals(candidate) == false) {
+ ring.set(ThrottleOwnerSelector.fromDiscoveryNodes(event.state().nodes()));
+ }
+ }
+
+ /**
+ * Asynchronously acquires one cluster-level shared permit for {@code bucketKey}, notifying {@code listener} with:
+ *
+ * - a non-null {@link Releasable} — granted; close it on request completion to release the shared permit;
+ * - {@code null} — fail open: admit the request without a shared permit (owner unreachable, empty
+ * ring, or any error). Nothing to release.
+ *
+ * A denial (bucket at its shared limit) is signalled via {@link ActionListener#onFailure} with a
+ * message-less {@link OpenSearchRejectedExecutionException} — a pure "denied" marker. This service holds
+ * only the opaque bucket key, not the human-readable group/attribute, so {@link WorkloadGroupService} recomposes
+ * the user-facing 429 message; the exception's type (not its text) is what distinguishes a denial from a
+ * transport error (which also arrives via {@code onFailure} but must pass through as fail-open, not a 429).
+ * The listener may be invoked inline (local-owner short-circuit) or on a transport thread.
+ */
+ public void acquireAsync(String bucketKey, int sharedLimit, ActionListener listener) {
+ final ThrottleOwnerSelector currentRing = ring.get();
+ final DiscoveryNode owner = currentRing.ownerFor(bucketKey).orElse(null);
+ if (owner == null) {
+ listener.onResponse(null); // empty ring -> fail open
+ return;
+ }
+
+ final String leaseId = leaseId();
+ final long ttlNanos = LEASE_TTL_NANOS;
+
+ // Local-owner short-circuit: this coordinator owns the bucket, so hit the tracker directly with no network hop.
+ if (owner.getId().equals(clusterService.localNode().getId())) {
+ if (tracker.tryAcquire(bucketKey, sharedLimit, leaseId, ttlNanos)) {
+ listener.onResponse(releaseLocal(bucketKey, leaseId));
+ } else {
+ listener.onFailure(deniedMarker());
+ }
+ return;
+ }
+
+ // Owner known-disconnected: fail open immediately rather than sending into a dead socket and waiting out
+ // ACQUIRE_TIMEOUT. This is an in-memory connection-map lookup (ConcurrentHashMap.containsKey), not a network
+ // call, so it costs nothing. It collapses the bulk of the post-node-drop detection window; the tiny remaining
+ // sub-window (socket dead but not yet flagged) still hits the timeout and is handled by handleException below.
+ if (transportService.nodeConnected(owner) == false) {
+ listener.onResponse(null);
+ return;
+ }
+
+ final TransportRequestOptions options = TransportRequestOptions.builder().withTimeout(ACQUIRE_TIMEOUT).build();
+ transportService.sendRequest(
+ owner,
+ ACQUIRE_ACTION_NAME,
+ new AcquireRequest(bucketKey, sharedLimit, leaseId, ttlNanos),
+ options,
+ new TransportResponseHandler() {
+ @Override
+ public AcquireResponse read(StreamInput in) throws IOException {
+ return new AcquireResponse(in);
+ }
+
+ @Override
+ public void handleResponse(AcquireResponse response) {
+ if (response.granted) {
+ listener.onResponse(releaseRemote(owner, bucketKey, leaseId));
+ } else {
+ listener.onFailure(deniedMarker());
+ }
+ }
+
+ @Override
+ public void handleException(TransportException exp) {
+ // Owner unreachable / timed out -> fail open. The owner may have GRANTED this lease before the
+ // response was lost (e.g. the acquire arrived but the reply timed out), which would otherwise
+ // occupy a shared slot until its TTL and cause false 429s once connectivity recovers. Send a
+ // best-effort release for this leaseId to reclaim it immediately; release-by-id is idempotent, so
+ // if no lease was created it is a harmless no-op. Then admit (fail open).
+ logger.debug("Shared throttle acquire to owner [{}] for bucket [{}] failed; failing open", owner.getId(), bucketKey);
+ sendRelease(owner, bucketKey, leaseId);
+ listener.onResponse(null);
+ }
+
+ @Override
+ public String executor() {
+ return ThreadPool.Names.SAME;
+ }
+ }
+ );
+ }
+
+ // Owner-side admission. Package-private for tests.
+ AcquireResponse handleAcquire(AcquireRequest request) {
+ boolean granted = tracker.tryAcquire(request.bucketKey, request.sharedLimit, request.leaseId, request.ttlNanos);
+ return new AcquireResponse(granted);
+ }
+
+ private Releasable releaseLocal(String bucketKey, String leaseId) {
+ return releaseOnce(() -> tracker.release(bucketKey, leaseId));
+ }
+
+ private Releasable releaseRemote(DiscoveryNode owner, String bucketKey, String leaseId) {
+ return releaseOnce(() -> sendRelease(owner, bucketKey, leaseId));
+ }
+
+ // Fire-and-forget RELEASE RPC to the bucket owner. Bounded by the same timeout as acquire so a half-open
+ // connection can't leave the response handler pending until the connection is torn down. A lost release is not
+ // fatal — the owner's TTL sweep reclaims the lease — so failures are logged at debug only.
+ private void sendRelease(DiscoveryNode owner, String bucketKey, String leaseId) {
+ final TransportRequestOptions options = TransportRequestOptions.builder().withTimeout(ACQUIRE_TIMEOUT).build();
+ transportService.sendRequest(
+ owner,
+ RELEASE_ACTION_NAME,
+ new ReleaseRequest(bucketKey, leaseId),
+ options,
+ new TransportResponseHandler() {
+ @Override
+ public TransportResponse.Empty read(StreamInput in) {
+ return TransportResponse.Empty.INSTANCE;
+ }
+
+ @Override
+ public void handleResponse(TransportResponse.Empty response) {}
+
+ @Override
+ public void handleException(TransportException exp) {
+ logger.debug(
+ "Shared throttle release to owner [{}] for bucket [{}] failed; TTL will reclaim",
+ owner.getId(),
+ bucketKey
+ );
+ }
+
+ @Override
+ public String executor() {
+ return ThreadPool.Names.SAME;
+ }
+ }
+ );
+ }
+
+ // Guards against a double close (e.g. onRequestEnd and onRequestFailure) releasing twice.
+ private static Releasable releaseOnce(Runnable release) {
+ final AtomicBoolean released = new AtomicBoolean(false);
+ return () -> {
+ if (released.compareAndSet(false, true)) {
+ release.run();
+ }
+ };
+ }
+
+ private static String leaseId() {
+ return UUIDs.base64UUID();
+ }
+
+ // Message-less "denied" marker: the bucket is at its shared limit. Carries no user text because this service has
+ // only the opaque bucket key; WorkloadGroupService recomposes the user-facing 429 with the group name/attribute.
+ // The exception TYPE (OpenSearchRejectedExecutionException) is the signal — the orchestrator uses it to tell a
+ // denial (recompose as a 429) apart from a transport error (pass through as fail-open).
+ private static OpenSearchRejectedExecutionException deniedMarker() {
+ return new OpenSearchRejectedExecutionException();
+ }
+
+ // Package-private accessor for tests.
+ SharedThrottleTracker tracker() {
+ return tracker;
+ }
+
+ ThrottleOwnerSelector ring() {
+ return ring.get();
+ }
+
+ /**
+ * Acquire RPC: a coordinator asks the bucket's owner to admit one request under {@code sharedLimit}.
+ */
+ public static class AcquireRequest extends TransportRequest {
+ final String bucketKey;
+ final int sharedLimit;
+ final String leaseId;
+ final long ttlNanos;
+
+ AcquireRequest(String bucketKey, int sharedLimit, String leaseId, long ttlNanos) {
+ this.bucketKey = bucketKey;
+ this.sharedLimit = sharedLimit;
+ this.leaseId = leaseId;
+ this.ttlNanos = ttlNanos;
+ }
+
+ AcquireRequest(StreamInput in) throws IOException {
+ super(in);
+ this.bucketKey = in.readString();
+ this.sharedLimit = in.readVInt();
+ this.leaseId = in.readString();
+ this.ttlNanos = in.readVLong();
+ }
+
+ @Override
+ public void writeTo(StreamOutput out) throws IOException {
+ super.writeTo(out);
+ out.writeString(bucketKey);
+ out.writeVInt(sharedLimit);
+ out.writeString(leaseId);
+ out.writeVLong(ttlNanos);
+ }
+ }
+
+ /**
+ * Acquire RPC response: whether the owner granted a shared permit.
+ */
+ public static class AcquireResponse extends TransportResponse {
+ final boolean granted;
+
+ AcquireResponse(boolean granted) {
+ this.granted = granted;
+ }
+
+ AcquireResponse(StreamInput in) throws IOException {
+ this.granted = in.readBoolean();
+ }
+
+ @Override
+ public void writeTo(StreamOutput out) throws IOException {
+ out.writeBoolean(granted);
+ }
+ }
+
+ /**
+ * Release RPC: fire-and-forget, tells the owner to drop a previously-granted lease.
+ */
+ public static class ReleaseRequest extends TransportRequest {
+ final String bucketKey;
+ final String leaseId;
+
+ ReleaseRequest(String bucketKey, String leaseId) {
+ this.bucketKey = bucketKey;
+ this.leaseId = leaseId;
+ }
+
+ ReleaseRequest(StreamInput in) throws IOException {
+ super(in);
+ this.bucketKey = in.readString();
+ this.leaseId = in.readString();
+ }
+
+ @Override
+ public void writeTo(StreamOutput out) throws IOException {
+ super.writeTo(out);
+ out.writeString(bucketKey);
+ out.writeString(leaseId);
+ }
+ }
+}
diff --git a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java
index a7535984ceca0..66b5ea3d7bf82 100644
--- a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java
+++ b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java
@@ -164,6 +164,7 @@
import org.opensearch.common.Nullable;
import org.opensearch.common.SetOnce;
import org.opensearch.common.cache.module.CacheModule;
+import org.opensearch.common.lease.Releasable;
import org.opensearch.common.network.NetworkModule;
import org.opensearch.common.settings.ClusterSettings;
import org.opensearch.common.settings.IndexScopedSettings;
@@ -297,6 +298,8 @@
import static org.hamcrest.Matchers.iterableWithSize;
import static org.hamcrest.Matchers.lessThanOrEqualTo;
import static org.hamcrest.Matchers.notNullValue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@@ -2381,6 +2384,13 @@ public void onFailure(final Exception e) {
);
SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory =
new SearchRequestOperationsCompositeListenerFactory();
+ // admit() is void+async: a default mock would never invoke the listener, leaving every search hung.
+ // Stub it to admit immediately with no throttle permit (the not-throttled path), matching production.
+ WorkloadGroupService workloadGroupService = mock(WorkloadGroupService.class);
+ doAnswer(invocation -> {
+ invocation.>getArgument(2).onResponse(null);
+ return null;
+ }).when(workloadGroupService).acquireThrottlePermit(any(), any(), any());
actions.put(
SearchAction.INSTANCE,
new TransportSearchAction(
@@ -2411,7 +2421,7 @@ public void onFailure(final Exception e) {
NoopTracer.INSTANCE,
new TaskResourceTrackingService(settings, clusterSettings, threadPool),
mockIndicesService,
- mock(WorkloadGroupService.class)
+ workloadGroupService
)
);
actions.put(
diff --git a/server/src/test/java/org/opensearch/wlm/NodeThrottleTrackerTests.java b/server/src/test/java/org/opensearch/wlm/NodeThrottleTrackerTests.java
new file mode 100644
index 0000000000000..80d1bbaa90a9a
--- /dev/null
+++ b/server/src/test/java/org/opensearch/wlm/NodeThrottleTrackerTests.java
@@ -0,0 +1,120 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.opensearch.common.lease.Releasable;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class NodeThrottleTrackerTests extends OpenSearchTestCase {
+
+ public void testAcquireUnderLimitSucceeds() {
+ NodeThrottleTracker tracker = new NodeThrottleTracker();
+ Releasable p1 = tracker.tryAcquire("bucket", 2);
+ Releasable p2 = tracker.tryAcquire("bucket", 2);
+ assertNotNull(p1);
+ assertNotNull(p2);
+ assertEquals(2, tracker.inFlight("bucket"));
+ p1.close();
+ p2.close();
+ }
+
+ public void testAcquireAtLimitReturnsNull() {
+ NodeThrottleTracker tracker = new NodeThrottleTracker();
+ assertNotNull(tracker.tryAcquire("bucket", 1));
+ // at the limit -> tryAcquire returns null (production then falls through to the shared tier or rejects)
+ assertNull(tracker.tryAcquire("bucket", 1));
+ // declined acquire must not leave the count inflated
+ assertEquals(1, tracker.inFlight("bucket"));
+ }
+
+ public void testReleaseFreesAPermit() {
+ NodeThrottleTracker tracker = new NodeThrottleTracker();
+ Releasable p = tracker.tryAcquire("bucket", 1);
+ assertNotNull(p);
+ // at the limit
+ assertNull(tracker.tryAcquire("bucket", 1));
+ p.close();
+ // permit freed -> a fresh acquire now succeeds
+ Releasable p2 = tracker.tryAcquire("bucket", 1);
+ assertNotNull(p2);
+ assertEquals(1, tracker.inFlight("bucket"));
+ p2.close();
+ }
+
+ public void testDrainToZeroRemovesBucketThenReacquire() {
+ NodeThrottleTracker tracker = new NodeThrottleTracker();
+ Releasable p = tracker.tryAcquire("bucket", 5);
+ assertNotNull(p);
+ assertEquals(1, tracker.inFlight("bucket"));
+ p.close();
+ assertEquals(0, tracker.inFlight("bucket"));
+ // re-acquiring after the bucket drained (and was removed) works and starts from 1
+ Releasable p2 = tracker.tryAcquire("bucket", 5);
+ assertNotNull(p2);
+ assertEquals(1, tracker.inFlight("bucket"));
+ p2.close();
+ }
+
+ public void testReleaseIsIdempotent() {
+ NodeThrottleTracker tracker = new NodeThrottleTracker();
+ Releasable p = tracker.tryAcquire("bucket", 5);
+ assertNotNull(p);
+ p.close();
+ p.close(); // double close must not decrement twice
+ assertEquals(0, tracker.inFlight("bucket"));
+ }
+
+ public void testBucketsAreIndependent() {
+ NodeThrottleTracker tracker = new NodeThrottleTracker();
+ assertNotNull(tracker.tryAcquire("a", 1));
+ // "a" is full but "b" is a separate bucket
+ assertNull(tracker.tryAcquire("a", 1));
+ Releasable pb = tracker.tryAcquire("b", 1);
+ assertNotNull(pb);
+ assertEquals(1, tracker.inFlight("a"));
+ assertEquals(1, tracker.inFlight("b"));
+ pb.close();
+ }
+
+ public void testConcurrentAcquireAdmitsExactlyLimit() throws Exception {
+ // Many threads race to acquire the same bucket at once. tryAcquire must admit exactly nodeLimit of them —
+ // never fewer (the over-decline race where all racers observe an inflated shared count and roll back) and
+ // never more (over-admit). Each caller decides on its own captured post-increment value.
+ final NodeThrottleTracker tracker = new NodeThrottleTracker();
+ final int limit = 10;
+ final int threads = 64;
+ final CountDownLatch start = new CountDownLatch(1);
+ final AtomicInteger granted = new AtomicInteger(0);
+ Thread[] workers = new Thread[threads];
+ for (int i = 0; i < threads; i++) {
+ workers[i] = new Thread(() -> {
+ try {
+ start.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return;
+ }
+ if (tracker.tryAcquire("hot", limit) != null) {
+ granted.incrementAndGet();
+ }
+ });
+ workers[i].start();
+ }
+ start.countDown();
+ for (Thread w : workers) {
+ w.join(TimeUnit.SECONDS.toMillis(30));
+ }
+ assertEquals("must admit exactly the limit under contention", limit, granted.get());
+ assertEquals(limit, tracker.inFlight("hot"));
+ }
+}
diff --git a/server/src/test/java/org/opensearch/wlm/SharedThrottleTrackerTests.java b/server/src/test/java/org/opensearch/wlm/SharedThrottleTrackerTests.java
new file mode 100644
index 0000000000000..bcc12b663e61a
--- /dev/null
+++ b/server/src/test/java/org/opensearch/wlm/SharedThrottleTrackerTests.java
@@ -0,0 +1,152 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+
+public class SharedThrottleTrackerTests extends OpenSearchTestCase {
+
+ private static final long TTL = TimeUnit.MINUTES.toNanos(5);
+
+ public void testGrantsUpToLimitThenDenies() {
+ SharedThrottleTracker tracker = new SharedThrottleTracker();
+ assertTrue(tracker.tryAcquire("b", 2, "l1", TTL));
+ assertTrue(tracker.tryAcquire("b", 2, "l2", TTL));
+ assertFalse("third acquire must be denied at limit 2", tracker.tryAcquire("b", 2, "l3", TTL));
+ assertEquals(2, tracker.inFlight("b"));
+ }
+
+ public void testReleaseFreesASlot() {
+ SharedThrottleTracker tracker = new SharedThrottleTracker();
+ assertTrue(tracker.tryAcquire("b", 1, "l1", TTL));
+ assertFalse(tracker.tryAcquire("b", 1, "l2", TTL));
+ tracker.release("b", "l1");
+ assertEquals(0, tracker.inFlight("b"));
+ assertTrue("slot freed after release", tracker.tryAcquire("b", 1, "l3", TTL));
+ }
+
+ public void testReleaseByUnknownLeaseIsNoOp() {
+ SharedThrottleTracker tracker = new SharedThrottleTracker();
+ assertTrue(tracker.tryAcquire("b", 2, "l1", TTL));
+ // Unknown lease (already swept, double release, or belonged to a previous owner of a remapped bucket).
+ tracker.release("b", "does-not-exist");
+ tracker.release("unknown-bucket", "l1");
+ assertEquals("a stray release must not decrement a live lease", 1, tracker.inFlight("b"));
+ }
+
+ public void testDoubleReleaseDecrementsOnce() {
+ SharedThrottleTracker tracker = new SharedThrottleTracker();
+ assertTrue(tracker.tryAcquire("b", 2, "l1", TTL));
+ assertTrue(tracker.tryAcquire("b", 2, "l2", TTL));
+ tracker.release("b", "l1");
+ tracker.release("b", "l1"); // second release of the same lease id is a no-op
+ assertEquals(1, tracker.inFlight("b"));
+ }
+
+ public void testDrainedBucketIsRemoved() {
+ SharedThrottleTracker tracker = new SharedThrottleTracker();
+ assertTrue(tracker.tryAcquire("b", 5, "l1", TTL));
+ assertEquals(1, tracker.activeBuckets());
+ tracker.release("b", "l1");
+ assertEquals("bucket entry removed when it drains to zero", 0, tracker.activeBuckets());
+ }
+
+ public void testExpiredLeaseIsReclaimedBySweep() {
+ AtomicLong clock = new AtomicLong(0);
+ SharedThrottleTracker tracker = new SharedThrottleTracker(clock::get);
+ assertTrue(tracker.tryAcquire("b", 1, "l1", 100));
+ assertFalse(tracker.tryAcquire("b", 1, "l2", 100)); // at limit
+ clock.set(101); // lease l1 has now expired
+ tracker.sweepExpired();
+ assertEquals(0, tracker.inFlight("b"));
+ assertTrue("slot reclaimed after TTL sweep", tracker.tryAcquire("b", 1, "l3", 100));
+ }
+
+ public void testExpiredLeaseReclaimedLazilyOnAcquire() {
+ AtomicLong clock = new AtomicLong(0);
+ SharedThrottleTracker tracker = new SharedThrottleTracker(clock::get);
+ assertTrue(tracker.tryAcquire("b", 1, "l1", 100));
+ clock.set(101); // l1 expired; a fresh acquire should prune it and be granted without an explicit sweep
+ assertTrue(tracker.tryAcquire("b", 1, "l2", 100));
+ assertEquals(1, tracker.inFlight("b"));
+ }
+
+ public void testSaturatedBucketWithExpiredLeaseIsReclaimedOnAcquire() {
+ // Prune now runs only when the bucket is at/over its limit; verify a fully-saturated bucket with one expired
+ // lease still admits the next acquire (the case the "prune only when full" optimization must not break).
+ AtomicLong clock = new AtomicLong(0);
+ SharedThrottleTracker tracker = new SharedThrottleTracker(clock::get);
+ assertTrue(tracker.tryAcquire("b", 2, "l1", 100)); // expires at 100
+ assertTrue(tracker.tryAcquire("b", 2, "l2", 500)); // expires at 500
+ assertFalse("bucket is full", tracker.tryAcquire("b", 2, "l3", 100));
+ clock.set(200); // l1 expired, l2 still live
+ // Bucket is at the limit -> acquire prunes l1 and grants; l2 remains.
+ assertTrue(tracker.tryAcquire("b", 2, "l4", 500));
+ assertEquals(2, tracker.inFlight("b"));
+ }
+
+ public void testSweepRacingLateReleaseDoesNotUnderCount() {
+ AtomicLong clock = new AtomicLong(0);
+ SharedThrottleTracker tracker = new SharedThrottleTracker(clock::get);
+ assertTrue(tracker.tryAcquire("b", 5, "l1", 100));
+ assertTrue(tracker.tryAcquire("b", 5, "l2", 100));
+ clock.set(101); // both expired
+ tracker.sweepExpired(); // removes l1 and l2
+ // A late release for an already-swept lease must be a no-op, not push the count negative.
+ tracker.release("b", "l1");
+ assertEquals(0, tracker.inFlight("b"));
+ // fresh lease unaffected
+ clock.set(150);
+ assertTrue(tracker.tryAcquire("b", 5, "l3", 100));
+ assertEquals(1, tracker.inFlight("b"));
+ }
+
+ public void testConcurrentAcquireNeverExceedsLimit() throws Exception {
+ final SharedThrottleTracker tracker = new SharedThrottleTracker();
+ final int limit = 10;
+ final int threads = 64;
+ final CountDownLatch start = new CountDownLatch(1);
+ final AtomicInteger granted = new AtomicInteger(0);
+ Thread[] workers = new Thread[threads];
+ for (int i = 0; i < threads; i++) {
+ final int id = i;
+ workers[i] = new Thread(() -> {
+ try {
+ start.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return;
+ }
+ if (tracker.tryAcquire("hot", limit, "lease-" + id, TTL)) {
+ granted.incrementAndGet();
+ }
+ });
+ workers[i].start();
+ }
+ start.countDown();
+ for (Thread w : workers) {
+ w.join();
+ }
+ assertEquals("exactly limit grants under contention", limit, granted.get());
+ assertEquals(limit, tracker.inFlight("hot"));
+ }
+
+ public void testBucketsAreIndependent() {
+ SharedThrottleTracker tracker = new SharedThrottleTracker();
+ assertTrue(tracker.tryAcquire("a", 1, "la", TTL));
+ assertTrue("different bucket has its own budget", tracker.tryAcquire("b", 1, "lb", TTL));
+ assertFalse(tracker.tryAcquire("a", 1, "la2", TTL));
+ assertEquals(2, tracker.activeBuckets());
+ }
+}
diff --git a/server/src/test/java/org/opensearch/wlm/ThrottleOwnerSelectorTests.java b/server/src/test/java/org/opensearch/wlm/ThrottleOwnerSelectorTests.java
new file mode 100644
index 0000000000000..815f577f35cdd
--- /dev/null
+++ b/server/src/test/java/org/opensearch/wlm/ThrottleOwnerSelectorTests.java
@@ -0,0 +1,141 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.opensearch.Version;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.node.DiscoveryNodeRole;
+import org.opensearch.cluster.node.DiscoveryNodes;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+public class ThrottleOwnerSelectorTests extends OpenSearchTestCase {
+
+ private DiscoveryNode dataNode(String id, Version version) {
+ return new DiscoveryNode(
+ "name_" + id,
+ id,
+ buildNewFakeTransportAddress(),
+ Collections.emptyMap(),
+ Set.of(DiscoveryNodeRole.DATA_ROLE),
+ version
+ );
+ }
+
+ private DiscoveryNode managerOnlyNode(String id) {
+ return new DiscoveryNode(
+ "name_" + id,
+ id,
+ buildNewFakeTransportAddress(),
+ Collections.emptyMap(),
+ Set.of(DiscoveryNodeRole.CLUSTER_MANAGER_ROLE),
+ Version.CURRENT
+ );
+ }
+
+ private DiscoveryNodes nodesOf(DiscoveryNode... nodes) {
+ DiscoveryNodes.Builder builder = DiscoveryNodes.builder();
+ for (DiscoveryNode node : nodes) {
+ builder.add(node);
+ }
+ builder.localNodeId(nodes[0].getId());
+ return builder.build();
+ }
+
+ public void testEmptyRingWhenNoDataNodes() {
+ ThrottleOwnerSelector selector = ThrottleOwnerSelector.fromDiscoveryNodes(nodesOf(managerOnlyNode("m1")));
+ assertTrue(selector.isEmpty());
+ assertFalse("empty ring must yield no owner (fail open)", selector.ownerFor("group-a:group").isPresent());
+ }
+
+ public void testClusterManagerOnlyNodesExcluded() {
+ ThrottleOwnerSelector selector = ThrottleOwnerSelector.fromDiscoveryNodes(
+ nodesOf(dataNode("d1", Version.CURRENT), managerOnlyNode("m1"))
+ );
+ // Only the data node is eligible; the manager-only node must never be an owner.
+ for (int i = 0; i < 200; i++) {
+ assertEquals("d1", selector.ownerFor("bucket-" + i).orElseThrow().getId());
+ }
+ }
+
+ public void testOldVersionNodesExcludedFromRing() {
+ DiscoveryNode newNode = dataNode("new", Version.CURRENT);
+ DiscoveryNode oldNode = dataNode("old", Version.V_3_6_0); // below MIN_OWNER_VERSION
+ ThrottleOwnerSelector selector = ThrottleOwnerSelector.fromDiscoveryNodes(nodesOf(newNode, oldNode));
+ assertEquals(1, selector.eligibleNodeSet().size());
+ for (int i = 0; i < 200; i++) {
+ assertEquals("only the upgraded node may own buckets", "new", selector.ownerFor("bucket-" + i).orElseThrow().getId());
+ }
+ }
+
+ public void testDeterministicForSameNodeSet() {
+ DiscoveryNodes nodes = nodesOf(dataNode("a", Version.CURRENT), dataNode("b", Version.CURRENT), dataNode("c", Version.CURRENT));
+ ThrottleOwnerSelector s1 = ThrottleOwnerSelector.fromDiscoveryNodes(nodes);
+ ThrottleOwnerSelector s2 = ThrottleOwnerSelector.fromDiscoveryNodes(nodes);
+ for (int i = 0; i < 500; i++) {
+ String key = "bucket-" + i;
+ assertEquals(s1.ownerFor(key).orElseThrow().getId(), s2.ownerFor(key).orElseThrow().getId());
+ }
+ }
+
+ public void testDistributionAcrossNodesIsReasonablyBalanced() {
+ ThrottleOwnerSelector selector = ThrottleOwnerSelector.fromDiscoveryNodes(
+ nodesOf(dataNode("a", Version.CURRENT), dataNode("b", Version.CURRENT), dataNode("c", Version.CURRENT))
+ );
+ Map counts = new HashMap<>();
+ int n = 3000;
+ for (int i = 0; i < n; i++) {
+ String owner = selector.ownerFor("bucket-" + i).orElseThrow().getId();
+ counts.merge(owner, 1, Integer::sum);
+ }
+ // With 128 virtual nodes each, every physical node should get a non-trivial share (well above 10%).
+ assertEquals(3, counts.size());
+ for (int c : counts.values()) {
+ assertTrue("each node should own a meaningful share, saw " + counts, c > n / 10);
+ }
+ }
+
+ public void testRebalanceOnlyRemapsSmallFraction() {
+ DiscoveryNode a = dataNode("a", Version.CURRENT);
+ DiscoveryNode b = dataNode("b", Version.CURRENT);
+ DiscoveryNode c = dataNode("c", Version.CURRENT);
+ DiscoveryNode d = dataNode("d", Version.CURRENT);
+ ThrottleOwnerSelector before = ThrottleOwnerSelector.fromDiscoveryNodes(nodesOf(a, b, c, d));
+ ThrottleOwnerSelector after = ThrottleOwnerSelector.fromDiscoveryNodes(nodesOf(a, b, c)); // d leaves
+
+ int n = 5000;
+ int remapped = 0;
+ for (int i = 0; i < n; i++) {
+ String key = "bucket-" + i;
+ String ownerBefore = before.ownerFor(key).orElseThrow().getId();
+ String ownerAfter = after.ownerFor(key).orElseThrow().getId();
+ if (ownerBefore.equals(ownerAfter) == false) {
+ remapped++;
+ }
+ }
+ // Removing 1 of 4 nodes should remap roughly 1/4 of keys. Allow generous slack, but it must be far from "all".
+ double fraction = (double) remapped / n;
+ assertTrue("remap fraction should be well under half, was " + fraction, fraction < 0.45);
+ // Keys not owned by the departed node 'd' should mostly keep their owner.
+ for (int i = 0; i < n; i++) {
+ String key = "bucket-" + i;
+ if (before.ownerFor(key).orElseThrow().getId().equals("d") == false) {
+ assertEquals(
+ "survivor-owned keys must not move",
+ before.ownerFor(key).orElseThrow().getId(),
+ after.ownerFor(key).orElseThrow().getId()
+ );
+ }
+ }
+ }
+}
diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java
index 18ea170d86e95..f8fabb8ff0b85 100644
--- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java
+++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java
@@ -13,11 +13,15 @@
import org.opensearch.cluster.ClusterState;
import org.opensearch.cluster.metadata.Metadata;
import org.opensearch.cluster.metadata.WorkloadGroup;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.node.DiscoveryNodeRole;
+import org.opensearch.cluster.node.DiscoveryNodes;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.lease.Releasable;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.unit.TimeValue;
import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.core.action.ActionListener;
import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
import org.opensearch.search.backpressure.trackers.NodeDuressTrackers;
import org.opensearch.tasks.Task;
@@ -36,6 +40,7 @@
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BooleanSupplier;
import org.mockito.ArgumentCaptor;
@@ -488,6 +493,25 @@ private void stubClusterStateWithGroup(WorkloadGroup wg) {
when(metadata.workloadGroups()).thenReturn(Map.of(wg.get_id(), wg));
}
+ /**
+ * Drives the async {@link WorkloadGroupService#acquireThrottlePermit} back into the synchronous "permit or 429"
+ * contract these node-tier tests assert on. With no shared throttle service wired, it completes inline: responds
+ * with a permit (or null when not throttled) or fails with the 429. Rethrows the failure so {@code expectThrows}
+ * works.
+ */
+ private Releasable acquireThrottlePermitSync(WorkloadGroupService service, String workloadGroupId, String principal) {
+ AtomicReference permit = new AtomicReference<>();
+ AtomicReference failure = new AtomicReference<>();
+ service.acquireThrottlePermit(workloadGroupId, principal, ActionListener.wrap(permit::set, failure::set));
+ if (failure.get() != null) {
+ if (failure.get() instanceof RuntimeException) {
+ throw (RuntimeException) failure.get();
+ }
+ throw new RuntimeException(failure.get());
+ }
+ return permit.get();
+ }
+
private WorkloadGroup throttledGroup(String id, Settings throttling) {
return new WorkloadGroup(
id + "-name",
@@ -506,12 +530,12 @@ public void testAcquireThrottleReturnsNullWhenNodeLimitUnset() {
when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED);
mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1");
stubClusterStateWithGroup(throttledGroup("wg-1", Settings.EMPTY)); // throttling not configured
- assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", null));
+ assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null));
}
public void testAcquireThrottleReturnsNullWhenWlmDisabled() {
when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED);
- assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", null));
+ assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null));
}
public void testAcquireThrottleRejectsAtLimitAndIncrementsStat() {
@@ -520,15 +544,19 @@ public void testAcquireThrottleRejectsAtLimitAndIncrementsStat() {
Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build();
stubClusterStateWithGroup(throttledGroup("wg-1", throttling));
- Releasable permit = workloadGroupService.acquireThrottleOrReject("wg-1", null); // first admit succeeds
+ Releasable permit = acquireThrottlePermitSync(workloadGroupService, "wg-1", null); // first admit succeeds
assertNotNull(permit);
// second admit hits node_limit of 1 -> 429 + total_throttled incremented
- expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject("wg-1", null));
+ OpenSearchRejectedExecutionException e = expectThrows(
+ OpenSearchRejectedExecutionException.class,
+ () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", null)
+ );
+ assertEquals("Request throttled: workload group [wg-1-name] reached its per-node limit of 1 concurrent requests.", e.getMessage());
assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled());
// releasing the first permit frees the slot so a subsequent acquire succeeds
permit.close();
- assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", null));
+ assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null));
}
public void testAcquireThrottleUsernameKeepsPerUserBuckets() {
@@ -537,22 +565,26 @@ public void testAcquireThrottleUsernameKeepsPerUserBuckets() {
Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 1).build();
stubClusterStateWithGroup(throttledGroup("wg-1", throttling));
- // alice takes her single slot; a second alice request is rejected.
- Releasable alice = workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice");
+ // alice takes her single slot; a second alice request is rejected — and the message names the username.
+ Releasable alice = acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|alice");
assertNotNull(alice);
- expectThrows(
+ OpenSearchRejectedExecutionException e = expectThrows(
OpenSearchRejectedExecutionException.class,
- () -> workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice")
+ () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|alice")
+ );
+ assertEquals(
+ "Request throttled: workload group [wg-1-name] for username [alice] reached its per-node limit of 1 concurrent requests.",
+ e.getMessage()
);
assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled());
// bob is a different bucket, so he is admitted even while alice is at her limit.
- Releasable bob = workloadGroupService.acquireThrottleOrReject("wg-1", "username|bob");
+ Releasable bob = acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|bob");
assertNotNull(bob);
// releasing alice frees her bucket
alice.close();
- assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice"));
+ assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|alice"));
}
public void testAcquireThrottleUsernameWithCommaDoesNotCollide() {
@@ -567,12 +599,12 @@ public void testAcquireThrottleUsernameWithCommaDoesNotCollide() {
// user "a" is a genuinely different principal
String userA = "username|a";
- Releasable ab = workloadGroupService.acquireThrottleOrReject("wg-1", userAB); // fills "a,b" bucket
+ Releasable ab = acquireThrottlePermitSync(workloadGroupService, "wg-1", userAB); // fills "a,b" bucket
assertNotNull(ab);
// user "a" must NOT be treated as the same bucket as "a,b" -> still admitted
- assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", userA));
+ assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", userA));
// a second "a,b" request hits the "a,b" bucket limit -> rejected
- expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject("wg-1", userAB));
+ expectThrows(OpenSearchRejectedExecutionException.class, () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", userAB));
}
public void testAcquireThrottleRolePicksMatchingSubfieldFromMultiTokenPrincipal() {
@@ -583,10 +615,10 @@ public void testAcquireThrottleRolePicksMatchingSubfieldFromMultiTokenPrincipal(
// A principal header may carry both subfields; the role bucket must key off the role token only.
String delim = WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER;
- assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice" + delim + "role|admin"));
+ assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|alice" + delim + "role|admin"));
expectThrows(
OpenSearchRejectedExecutionException.class,
- () -> workloadGroupService.acquireThrottleOrReject("wg-1", "username|bob" + delim + "role|admin")
+ () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|bob" + delim + "role|admin")
);
assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled());
}
@@ -598,9 +630,9 @@ public void testAcquireThrottleFailsOpenWhenPrincipalMissingForUsername() {
stubClusterStateWithGroup(throttledGroup("wg-1", throttling));
// No principal (e.g. security plugin not installed) or no matching subfield -> not throttled (fail open).
- assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", null));
- assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", ""));
- assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", "role|admin")); // no username token
+ assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null));
+ assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", ""));
+ assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", "role|admin")); // no username token
assertEquals(0, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled());
}
@@ -628,9 +660,9 @@ public void testAcquireThrottleStillRejectsWhenStatUpdateFails() {
new HashSet<>()
);
- assertNotNull(serviceWithNullState.acquireThrottleOrReject("wg-1", null)); // first admit fills the single slot
+ assertNotNull(acquireThrottlePermitSync(serviceWithNullState, "wg-1", null)); // first admit fills the single slot
// second acquire is over the limit; a null state must not let the stat update swallow the 429
- expectThrows(OpenSearchRejectedExecutionException.class, () -> serviceWithNullState.acquireThrottleOrReject("wg-1", null));
+ expectThrows(OpenSearchRejectedExecutionException.class, () -> acquireThrottlePermitSync(serviceWithNullState, "wg-1", null));
// accessor whose state-map lookup throws must also still propagate the 429
WorkloadGroupsStateAccessor throwingStateAccessor = Mockito.mock(WorkloadGroupsStateAccessor.class);
@@ -646,8 +678,8 @@ public void testAcquireThrottleStillRejectsWhenStatUpdateFails() {
new HashSet<>()
);
- assertNotNull(serviceWithThrowingState.acquireThrottleOrReject("wg-1", null)); // fills the single slot
- expectThrows(OpenSearchRejectedExecutionException.class, () -> serviceWithThrowingState.acquireThrottleOrReject("wg-1", null));
+ assertNotNull(acquireThrottlePermitSync(serviceWithThrowingState, "wg-1", null)); // fills the single slot
+ expectThrows(OpenSearchRejectedExecutionException.class, () -> acquireThrottlePermitSync(serviceWithThrowingState, "wg-1", null));
}
/**
@@ -662,8 +694,8 @@ public void testAcquireThrottleDoesNotMisattributeToDefaultDuringRegistrationLag
// DEFAULT group state exists, but wg-1 is NOT yet registered (registration lag).
mockWorkloadGroupsStateAccessor.addNewWorkloadGroup(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get());
- assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); // fills the single slot
- expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject("wg-1", null));
+ assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); // fills the single slot
+ expectThrows(OpenSearchRejectedExecutionException.class, () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", null));
// the rejection must NOT have landed on the DEFAULT group
assertEquals(
@@ -673,6 +705,70 @@ public void testAcquireThrottleDoesNotMisattributeToDefaultDuringRegistrationLag
);
}
+ /**
+ * The headline two-tier path: node_limit=1 + shared_limit=1 with a shared throttle service wired in. The first
+ * request takes the local slot, the second overflows and takes the shared slot, and the third is rejected (both
+ * tiers full) with total_throttled incremented. Uses a single-data-node ring so the shared acquire resolves to
+ * this node and runs via the local-owner short-circuit (no real network).
+ */
+ public void testAdmitFallsThroughNodeTierToSharedTier() {
+ when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED);
+ mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1");
+
+ Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).put("shared_limit", 1).build();
+ WorkloadGroup group = throttledGroup("wg-1", throttling);
+
+ // Cluster state must expose both the group metadata and a single data node (so this node owns every bucket).
+ DiscoveryNode localNode = new DiscoveryNode(
+ "local",
+ "local",
+ buildNewFakeTransportAddress(),
+ Collections.emptyMap(),
+ Set.of(DiscoveryNodeRole.DATA_ROLE),
+ org.opensearch.Version.CURRENT
+ );
+ DiscoveryNodes nodes = DiscoveryNodes.builder().add(localNode).localNodeId("local").build();
+ ClusterState clusterState = Mockito.mock(ClusterState.class);
+ Metadata metadata = Mockito.mock(Metadata.class);
+ when(mockClusterService.state()).thenReturn(clusterState);
+ when(mockClusterService.localNode()).thenReturn(localNode);
+ when(clusterState.metadata()).thenReturn(metadata);
+ when(clusterState.nodes()).thenReturn(nodes);
+ when(metadata.workloadGroups()).thenReturn(Map.of(group.get_id(), group));
+
+ WorkloadGroupSharedThrottleService sharedService = new WorkloadGroupSharedThrottleService(
+ mockClusterService,
+ mockThreadPool,
+ Mockito.mock(org.opensearch.transport.TransportService.class)
+ );
+ // Populate the ring (the constructor starts empty; a nodes-changed event builds it — matches production).
+ ClusterState previous = Mockito.mock(ClusterState.class);
+ when(previous.nodes()).thenReturn(DiscoveryNodes.EMPTY_NODES);
+ sharedService.clusterChanged(new ClusterChangedEvent("test", clusterState, previous));
+ workloadGroupService.setSharedThrottleService(sharedService);
+
+ // 1st: local node_limit slot.
+ Releasable local = acquireThrottlePermitSync(workloadGroupService, "wg-1", null);
+ assertNotNull(local);
+ // 2nd: local exhausted -> falls through to shared_limit slot.
+ Releasable shared = acquireThrottlePermitSync(workloadGroupService, "wg-1", null);
+ assertNotNull(shared);
+ // 3rd: both tiers full -> 429 from the shared tier (with the cluster-wide message), and total_throttled bumped.
+ OpenSearchRejectedExecutionException e = expectThrows(
+ OpenSearchRejectedExecutionException.class,
+ () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", null)
+ );
+ assertEquals(
+ "Request throttled: workload group [wg-1-name] reached its cluster-wide limit of 1 concurrent requests.",
+ e.getMessage()
+ );
+ assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled());
+
+ // Releasing the shared permit frees the shared slot so a subsequent request is admitted again.
+ shared.close();
+ assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null));
+ }
+
public void testShouldSBPHandle() {
SearchTask task = createMockTaskWithResourceStats(SearchTask.class, 100, 200, 0, 12);
WorkloadGroupState workloadGroupState = new WorkloadGroupState();
diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java
new file mode 100644
index 0000000000000..792b94a58f750
--- /dev/null
+++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java
@@ -0,0 +1,282 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.opensearch.cluster.ClusterChangedEvent;
+import org.opensearch.cluster.ClusterState;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.node.DiscoveryNodeRole;
+import org.opensearch.cluster.node.DiscoveryNodes;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.lease.Releasable;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportService;
+
+import java.util.Collections;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.mockito.Mockito;
+
+import static org.mockito.Mockito.when;
+
+public class WorkloadGroupSharedThrottleServiceTests extends OpenSearchTestCase {
+
+ private ClusterService clusterService;
+ private ThreadPool threadPool;
+ private TransportService transportService;
+ private DiscoveryNode localNode;
+
+ @Override
+ public void setUp() throws Exception {
+ super.setUp();
+ localNode = new DiscoveryNode(
+ "local",
+ "local",
+ buildNewFakeTransportAddress(),
+ Collections.emptyMap(),
+ Set.of(DiscoveryNodeRole.DATA_ROLE),
+ org.opensearch.Version.CURRENT
+ );
+ clusterService = Mockito.mock(ClusterService.class);
+ threadPool = Mockito.mock(ThreadPool.class);
+ transportService = Mockito.mock(TransportService.class);
+
+ ClusterState state = Mockito.mock(ClusterState.class);
+ DiscoveryNodes singleDataNode = DiscoveryNodes.builder().add(localNode).localNodeId("local").build();
+ when(state.nodes()).thenReturn(singleDataNode);
+ when(clusterService.state()).thenReturn(state);
+ when(clusterService.localNode()).thenReturn(localNode);
+ }
+
+ private WorkloadGroupSharedThrottleService newService() {
+ // Single data node => this node owns every bucket => acquire uses the local short-circuit (no real network).
+ WorkloadGroupSharedThrottleService service = new WorkloadGroupSharedThrottleService(clusterService, threadPool, transportService);
+ // The ring is empty until a cluster-state change populates it (matches production: the constructor does not
+ // read cluster state). Deliver a nodesChanged event with the current nodes.
+ deliverNodesChanged(service, clusterService.state().nodes());
+ return service;
+ }
+
+ private void deliverNodesChanged(WorkloadGroupSharedThrottleService service, DiscoveryNodes nodes) {
+ ClusterState previous = Mockito.mock(ClusterState.class);
+ when(previous.nodes()).thenReturn(DiscoveryNodes.EMPTY_NODES);
+ ClusterState current = Mockito.mock(ClusterState.class);
+ when(current.nodes()).thenReturn(nodes);
+ ClusterChangedEvent event = new ClusterChangedEvent("test", current, previous);
+ service.clusterChanged(event);
+ }
+
+ private static Releasable awaitGrant(WorkloadGroupSharedThrottleService service, String bucket, int limit) {
+ AtomicReference permit = new AtomicReference<>();
+ AtomicReference failure = new AtomicReference<>();
+ service.acquireAsync(bucket, limit, ActionListener.wrap(permit::set, failure::set));
+ if (failure.get() != null) {
+ throw failure.get() instanceof RuntimeException re ? re : new RuntimeException(failure.get());
+ }
+ return permit.get();
+ }
+
+ public void testRingPopulatesWhenNodeSetUnchangedVsPreviousState() {
+ // Regression for the single-node no-op: the coordinator seeds the initial applied state already containing the
+ // local node, so the first real clusterChanged has previous.nodes() == current.nodes() and nodesChanged() is
+ // false. The ring must still populate (we compare eligible owner sets, not the delta), otherwise shared
+ // throttling silently does nothing on a single-node cluster.
+ WorkloadGroupSharedThrottleService service = new WorkloadGroupSharedThrottleService(clusterService, threadPool, transportService);
+ DiscoveryNodes nodes = DiscoveryNodes.builder().add(localNode).localNodeId("local").build();
+ ClusterState previous = Mockito.mock(ClusterState.class);
+ when(previous.nodes()).thenReturn(nodes); // same node set as current -> nodesChanged() == false
+ ClusterState current = Mockito.mock(ClusterState.class);
+ when(current.nodes()).thenReturn(nodes);
+ ClusterChangedEvent event = new ClusterChangedEvent("test", current, previous);
+ assertFalse("precondition: this event reports no node change", event.nodesChanged());
+
+ service.clusterChanged(event);
+
+ // shared_limit=1 must now actually enforce (grant then reject), not fail open.
+ assertNotNull(awaitGrant(service, "b", 1));
+ expectThrows(OpenSearchRejectedExecutionException.class, () -> awaitGrant(service, "b", 1));
+ }
+
+ public void testSameIdRestartRebuildsRingWithFreshNode() {
+ // A node can restart keeping its persistent id but with a new ephemeral id/address. Comparing only persistent
+ // ids would treat this as "no change" and keep the stale DiscoveryNode in the ring, so nodeConnected(stale)
+ // would fail open forever. The ring must rebuild and hold the FRESH node object.
+ WorkloadGroupSharedThrottleService service = new WorkloadGroupSharedThrottleService(clusterService, threadPool, transportService);
+
+ DiscoveryNode first = new DiscoveryNode(
+ "n1",
+ "n1",
+ buildNewFakeTransportAddress(),
+ Collections.emptyMap(),
+ Set.of(DiscoveryNodeRole.DATA_ROLE),
+ org.opensearch.Version.CURRENT
+ );
+ deliverNodesChanged(service, DiscoveryNodes.builder().add(first).localNodeId("n1").build());
+ assertSame("ring should hold the first node instance", first, service.ring().ownerFor("b").orElseThrow());
+
+ // Same persistent id "n1", but a brand-new DiscoveryNode instance (fresh ephemeral id + transport address).
+ DiscoveryNode restarted = new DiscoveryNode(
+ "n1",
+ "n1",
+ buildNewFakeTransportAddress(),
+ Collections.emptyMap(),
+ Set.of(DiscoveryNodeRole.DATA_ROLE),
+ org.opensearch.Version.CURRENT
+ );
+ assertNotEquals("restarted node must not equal the old instance", first, restarted);
+ deliverNodesChanged(service, DiscoveryNodes.builder().add(restarted).localNodeId("n1").build());
+
+ assertSame("ring must have rebuilt with the restarted node instance", restarted, service.ring().ownerFor("b").orElseThrow());
+ }
+
+ public void testLocalOwnerGrantsThenDeniesAtLimit() {
+ WorkloadGroupSharedThrottleService service = newService();
+ Releasable p1 = awaitGrant(service, "b", 1);
+ assertNotNull(p1);
+ expectThrows(OpenSearchRejectedExecutionException.class, () -> awaitGrant(service, "b", 1));
+ // release frees the shared slot
+ p1.close();
+ assertNotNull(awaitGrant(service, "b", 1));
+ }
+
+ public void testDoubleCloseReleasesOnce() {
+ WorkloadGroupSharedThrottleService service = newService();
+ Releasable p = awaitGrant(service, "b", 2);
+ assertNotNull(awaitGrant(service, "b", 2)); // second slot
+ p.close();
+ p.close(); // must not double-release
+ // one slot still held, so only one more grant is available
+ assertNotNull(awaitGrant(service, "b", 2));
+ expectThrows(OpenSearchRejectedExecutionException.class, () -> awaitGrant(service, "b", 2));
+ }
+
+ public void testEmptyRingFailsOpen() {
+ // Cluster with no eligible data node (coordinating/manager-only) => empty ring => fail open (null permit).
+ DiscoveryNode managerOnly = new DiscoveryNode(
+ "m",
+ "m",
+ buildNewFakeTransportAddress(),
+ Collections.emptyMap(),
+ Set.of(DiscoveryNodeRole.CLUSTER_MANAGER_ROLE),
+ org.opensearch.Version.CURRENT
+ );
+ ClusterState state = Mockito.mock(ClusterState.class);
+ when(state.nodes()).thenReturn(DiscoveryNodes.builder().add(managerOnly).localNodeId("m").build());
+ when(clusterService.state()).thenReturn(state);
+ when(clusterService.localNode()).thenReturn(managerOnly);
+
+ WorkloadGroupSharedThrottleService service = newService();
+ AtomicReference permit = new AtomicReference<>();
+ AtomicReference called = new AtomicReference<>(false);
+ service.acquireAsync("b", 1, ActionListener.wrap(p -> {
+ called.set(true);
+ permit.set(p);
+ }, e -> fail("must not fail: " + e)));
+ assertTrue("listener must be invoked inline", called.get());
+ assertNull("fail-open yields a null permit (admit, untracked)", permit.get());
+ }
+
+ public void testDisconnectedOwnerFailsOpenWithoutSendingRequest() {
+ // Ring owner is a remote data node (not the local node), so this is NOT the local short-circuit path.
+ DiscoveryNode remote = new DiscoveryNode(
+ "remote",
+ "remote",
+ buildNewFakeTransportAddress(),
+ Collections.emptyMap(),
+ Set.of(DiscoveryNodeRole.DATA_ROLE),
+ org.opensearch.Version.CURRENT
+ );
+ ClusterState state = Mockito.mock(ClusterState.class);
+ // Only the remote node is a data node, so every bucket hashes to it; local is a coordinating-only node.
+ when(state.nodes()).thenReturn(DiscoveryNodes.builder().add(remote).localNodeId("local").build());
+ when(clusterService.state()).thenReturn(state);
+ when(clusterService.localNode()).thenReturn(localNode);
+ // The owner is known-disconnected.
+ when(transportService.nodeConnected(remote)).thenReturn(false);
+
+ WorkloadGroupSharedThrottleService service = newService();
+ AtomicReference permit = new AtomicReference<>();
+ AtomicReference called = new AtomicReference<>(false);
+ service.acquireAsync("b", 1, ActionListener.wrap(p -> {
+ called.set(true);
+ permit.set(p);
+ }, e -> fail("must not fail: " + e)));
+
+ // The listener completed inline with a null permit. This proves the disconnected pre-check fired: on the
+ // transport path the listener is only invoked from a sendRequest callback, but transportService is a plain
+ // mock whose sendRequest is a no-op that never calls back — so an inline null response can only come from the
+ // nodeConnected()==false fast-fail branch, not from a request sent into the dead socket.
+ assertTrue("listener must be invoked inline (no RTT to a dead node)", called.get());
+ assertNull("disconnected owner -> fail open with a null permit", permit.get());
+ }
+
+ public void testHandleAcquireIsOwnerSideAdmission() {
+ WorkloadGroupSharedThrottleService service = newService();
+ // Directly exercise the owner-side handler (what a remote coordinator's RPC would hit).
+ assertTrue(
+ service.handleAcquire(
+ new WorkloadGroupSharedThrottleService.AcquireRequest("b", 1, "lease-1", WorkloadGroupSharedThrottleService.LEASE_TTL_NANOS)
+ ).granted
+ );
+ assertFalse(
+ service.handleAcquire(
+ new WorkloadGroupSharedThrottleService.AcquireRequest("b", 1, "lease-2", WorkloadGroupSharedThrottleService.LEASE_TTL_NANOS)
+ ).granted
+ );
+ assertEquals(1, service.tracker().inFlight("b"));
+ }
+
+ public void testAcquireRequestSerializationRoundTrip() throws Exception {
+ WorkloadGroupSharedThrottleService.AcquireRequest original = new WorkloadGroupSharedThrottleService.AcquireRequest(
+ "grp1:username:alice",
+ 42,
+ "lease-xyz",
+ 123_456_789L
+ );
+ WorkloadGroupSharedThrottleService.AcquireRequest copy = copyWriteable(
+ original,
+ writableRegistry(),
+ WorkloadGroupSharedThrottleService.AcquireRequest::new
+ );
+ assertEquals(original.bucketKey, copy.bucketKey);
+ assertEquals(original.sharedLimit, copy.sharedLimit);
+ assertEquals(original.leaseId, copy.leaseId);
+ assertEquals(original.ttlNanos, copy.ttlNanos);
+ }
+
+ public void testAcquireResponseSerializationRoundTrip() throws Exception {
+ for (boolean granted : new boolean[] { true, false }) {
+ WorkloadGroupSharedThrottleService.AcquireResponse original = new WorkloadGroupSharedThrottleService.AcquireResponse(granted);
+ WorkloadGroupSharedThrottleService.AcquireResponse copy = copyWriteable(
+ original,
+ writableRegistry(),
+ WorkloadGroupSharedThrottleService.AcquireResponse::new
+ );
+ assertEquals(granted, copy.granted);
+ }
+ }
+
+ public void testReleaseRequestSerializationRoundTrip() throws Exception {
+ WorkloadGroupSharedThrottleService.ReleaseRequest original = new WorkloadGroupSharedThrottleService.ReleaseRequest(
+ "grp1:group",
+ "lease-abc"
+ );
+ WorkloadGroupSharedThrottleService.ReleaseRequest copy = copyWriteable(
+ original,
+ writableRegistry(),
+ WorkloadGroupSharedThrottleService.ReleaseRequest::new
+ );
+ assertEquals(original.bucketKey, copy.bucketKey);
+ assertEquals(original.leaseId, copy.leaseId);
+ }
+}
diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTransportTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTransportTests.java
new file mode 100644
index 0000000000000..0fec6160e6b65
--- /dev/null
+++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTransportTests.java
@@ -0,0 +1,227 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.opensearch.Version;
+import org.opensearch.cluster.ClusterChangedEvent;
+import org.opensearch.cluster.ClusterState;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.node.DiscoveryNodes;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.lease.Releasable;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
+import org.opensearch.telemetry.tracing.noop.NoopTracer;
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.test.transport.MockTransportService;
+import org.opensearch.threadpool.TestThreadPool;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.ConnectTransportException;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * Exercises the cross-node (remote owner) acquire/release path of {@link WorkloadGroupSharedThrottleService}
+ * over a real two-node {@link MockTransportService} harness. Mockito cannot be used to stub the acquire round-trip
+ * because {@code TransportService#sendRequest(node, action, request, options, handler)} is {@code final}; instead we
+ * stand up two genuine transports so the {@code internal:wlm/throttle/shared/acquire} and {@code .../release}
+ * RPCs travel over the wire and hit the owner-side handlers for real.
+ */
+public class WorkloadGroupSharedThrottleServiceTransportTests extends OpenSearchTestCase {
+
+ private ThreadPool threadPool;
+ private MockTransportService coordinatorTransport;
+ private MockTransportService ownerTransport;
+
+ private ClusterService coordinatorClusterService;
+ private ClusterService ownerClusterService;
+
+ private WorkloadGroupSharedThrottleService coordinatorService;
+ private WorkloadGroupSharedThrottleService ownerService;
+
+ private DiscoveryNode coordinatorNode;
+ private DiscoveryNode ownerNode;
+
+ // A bucket key whose ring owner is the remote (owner) node, forcing the coordinator down the transport path.
+ private String remoteKey;
+
+ @Override
+ public void setUp() throws Exception {
+ super.setUp();
+ threadPool = new TestThreadPool(getClass().getName());
+
+ // Two real transports. Default node roles include DATA at Version.CURRENT (>= MIN_OWNER_VERSION), so both
+ // nodes are eligible ring owners.
+ coordinatorTransport = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool, NoopTracer.INSTANCE);
+ ownerTransport = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool, NoopTracer.INSTANCE);
+ coordinatorTransport.start();
+ coordinatorTransport.acceptIncomingRequests();
+ ownerTransport.start();
+ ownerTransport.acceptIncomingRequests();
+ // The coordinator must be able to reach the owner for the acquire/release RPCs.
+ coordinatorTransport.connectToNode(ownerTransport.getLocalDiscoNode());
+
+ coordinatorNode = coordinatorTransport.getLocalDiscoNode();
+ ownerNode = ownerTransport.getLocalDiscoNode();
+
+ // Each service sees itself as the local node.
+ coordinatorClusterService = mock(ClusterService.class);
+ when(coordinatorClusterService.localNode()).thenReturn(coordinatorNode);
+ ownerClusterService = mock(ClusterService.class);
+ when(ownerClusterService.localNode()).thenReturn(ownerNode);
+
+ // Build one service per transport so BOTH register their acquire/release handlers on their own transport.
+ coordinatorService = new WorkloadGroupSharedThrottleService(coordinatorClusterService, threadPool, coordinatorTransport);
+ ownerService = new WorkloadGroupSharedThrottleService(ownerClusterService, threadPool, ownerTransport);
+
+ // Drive an identical membership (both nodes as data nodes) into both services so they build the same ring and
+ // agree on a single deterministic owner per bucket.
+ DiscoveryNodes bothNodes = DiscoveryNodes.builder()
+ .add(coordinatorNode)
+ .add(ownerNode)
+ .localNodeId(coordinatorNode.getId())
+ .build();
+ deliverNodes(coordinatorService, bothNodes);
+ deliverNodes(ownerService, bothNodes);
+
+ // Find a bucket key that the coordinator's ring maps to the REMOTE owner node, so acquireAsync takes the
+ // transport path rather than the local short-circuit. Guard the whole harness against a mis-built ring.
+ for (int i = 0; i < 10000 && remoteKey == null; i++) {
+ String candidate = "bucket-" + i;
+ DiscoveryNode owner = coordinatorService.ring().ownerFor(candidate).orElse(null);
+ if (owner != null && owner.getId().equals(ownerNode.getId())) {
+ remoteKey = candidate;
+ }
+ }
+ assertNotNull("could not find a bucket owned by the remote owner node", remoteKey);
+ assertEquals(
+ "chosen bucket must be owned by the remote node (transport path), not the coordinator",
+ ownerNode.getId(),
+ coordinatorService.ring().ownerFor(remoteKey).orElseThrow().getId()
+ );
+ }
+
+ @Override
+ public void tearDown() throws Exception {
+ coordinatorTransport.close();
+ ownerTransport.close();
+ ThreadPool.terminate(threadPool, 10, TimeUnit.SECONDS);
+ super.tearDown();
+ }
+
+ private void deliverNodes(WorkloadGroupSharedThrottleService service, DiscoveryNodes nodes) {
+ ClusterState previous = mock(ClusterState.class);
+ when(previous.nodes()).thenReturn(DiscoveryNodes.EMPTY_NODES);
+ ClusterState current = mock(ClusterState.class);
+ when(current.nodes()).thenReturn(nodes);
+ service.clusterChanged(new ClusterChangedEvent("test", current, previous));
+ }
+
+ /**
+ * Happy path across nodes: the coordinator asks the remote owner for a permit, the owner grants it over the wire,
+ * and closing the returned {@link Releasable} sends the fire-and-forget RELEASE RPC that drains the owner's tracker.
+ */
+ public void testRemoteAcquireGrantedThenReleaseRoundTrips() throws Exception {
+ CapturingListener listener = new CapturingListener();
+ coordinatorService.acquireAsync(remoteKey, 1, listener);
+ listener.await();
+
+ assertNull("granted acquire must not fail", listener.failure.get());
+ assertTrue("listener must have been notified via onResponse", listener.responded.get());
+ Releasable permit = listener.response.get();
+ assertNotNull("remote owner granted a permit, so a non-null Releasable is expected", permit);
+ // The grant was recorded on the OWNER node's tracker (the acquire handler ran there before responding).
+ assertEquals("owner tracker must hold exactly one in-flight lease", 1, ownerService.tracker().inFlight(remoteKey));
+
+ // Closing the permit fires the RELEASE RPC (fire-and-forget), which the owner applies asynchronously.
+ permit.close();
+ assertBusy(() -> assertEquals("release RPC must drain the owner's in-flight count", 0, ownerService.tracker().inFlight(remoteKey)));
+ }
+
+ /**
+ * When the remote owner is already at its shared limit for the bucket, the owner denies the acquire over the wire
+ * and the coordinator surfaces an {@link OpenSearchRejectedExecutionException} via {@code onFailure} (a 429).
+ */
+ public void testRemoteAcquireDeniedReturns429() throws Exception {
+ // Pre-fill the owner's tracker to the limit directly so the next acquire is denied at the source.
+ assertTrue(ownerService.tracker().tryAcquire(remoteKey, 1, "pre", WorkloadGroupSharedThrottleService.LEASE_TTL_NANOS));
+ assertEquals(1, ownerService.tracker().inFlight(remoteKey));
+
+ CapturingListener listener = new CapturingListener();
+ coordinatorService.acquireAsync(remoteKey, 1, listener);
+ listener.await();
+
+ assertFalse("denied acquire must not invoke onResponse", listener.responded.get());
+ assertNotNull("denied acquire must invoke onFailure", listener.failure.get());
+ assertTrue(
+ "denial must be an OpenSearchRejectedExecutionException (429), was: " + listener.failure.get(),
+ listener.failure.get() instanceof OpenSearchRejectedExecutionException
+ );
+ }
+
+ /**
+ * If the acquire RPC to the owner fails at send time, the coordinator must FAIL OPEN — admit the request with a
+ * {@code null} permit via {@code onResponse}, never {@code onFailure}. We keep the owner connected (so the
+ * in-memory {@code nodeConnected} pre-check passes) and inject a send behavior that throws, forcing the failure
+ * through {@code handleException} rather than the disconnected pre-check.
+ */
+ public void testFailOpenWhenOwnerTransportUnavailable() throws Exception {
+ assertTrue(
+ "precondition: owner must be connected so we exercise the send path, not the pre-check",
+ coordinatorTransport.nodeConnected(ownerNode)
+ );
+ coordinatorTransport.addSendBehavior(ownerTransport, (connection, requestId, action, request, options) -> {
+ if (WorkloadGroupSharedThrottleService.ACQUIRE_ACTION_NAME.equals(action)) {
+ // Simulate the send blowing up; TransportService routes this to the handler's handleException.
+ throw new ConnectTransportException(connection.getNode(), "simulated acquire send failure");
+ }
+ connection.sendRequest(requestId, action, request, options);
+ });
+
+ CapturingListener listener = new CapturingListener();
+ coordinatorService.acquireAsync(remoteKey, 1, listener);
+ listener.await();
+
+ assertTrue("fail-open must invoke onResponse, not onFailure", listener.responded.get());
+ assertNull("fail-open on transport error must yield a null permit (admit, untracked)", listener.response.get());
+ assertNull("fail-open must never propagate a failure to the listener", listener.failure.get());
+ }
+
+ /** Captures the outcome of an async acquire, distinguishing onResponse(null) from an unfired listener via a latch. */
+ private static final class CapturingListener implements ActionListener {
+ final CountDownLatch latch = new CountDownLatch(1);
+ final AtomicReference response = new AtomicReference<>();
+ final AtomicReference failure = new AtomicReference<>();
+ final AtomicBoolean responded = new AtomicBoolean(false);
+
+ @Override
+ public void onResponse(Releasable releasable) {
+ response.set(releasable);
+ responded.set(true);
+ latch.countDown();
+ }
+
+ @Override
+ public void onFailure(Exception e) {
+ failure.set(e);
+ latch.countDown();
+ }
+
+ void await() throws InterruptedException {
+ assertTrue("listener was not invoked within the timeout", latch.await(10, TimeUnit.SECONDS));
+ }
+ }
+}
diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java
deleted file mode 100644
index 475760834af5a..0000000000000
--- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java
+++ /dev/null
@@ -1,81 +0,0 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- *
- * The OpenSearch Contributors require contributions made to
- * this file be licensed under the Apache-2.0 license or a
- * compatible open source license.
- */
-
-package org.opensearch.wlm;
-
-import org.opensearch.common.lease.Releasable;
-import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
-import org.opensearch.test.OpenSearchTestCase;
-
-public class WorkloadGroupThrottleTrackerTests extends OpenSearchTestCase {
-
- public void testAcquireUnderLimitSucceeds() {
- WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker();
- Releasable p1 = tracker.acquire("bucket", 2);
- Releasable p2 = tracker.acquire("bucket", 2);
- assertEquals(2, tracker.inFlight("bucket"));
- p1.close();
- p2.close();
- }
-
- public void testAcquireAtLimitRejects() {
- WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker();
- tracker.acquire("bucket", 1);
- OpenSearchRejectedExecutionException e = expectThrows(
- OpenSearchRejectedExecutionException.class,
- () -> tracker.acquire("bucket", 1)
- );
- assertTrue(e.getMessage().contains("throttle limit reached"));
- assertFalse(e.getMessage().contains("bucket"));
- // rejected acquire must not leave the count inflated
- assertEquals(1, tracker.inFlight("bucket"));
- }
-
- public void testReleaseFreesAPermit() {
- WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker();
- Releasable p = tracker.acquire("bucket", 1);
- // at the limit
- expectThrows(OpenSearchRejectedExecutionException.class, () -> tracker.acquire("bucket", 1));
- p.close();
- // permit freed -> a fresh acquire now succeeds
- Releasable p2 = tracker.acquire("bucket", 1);
- assertEquals(1, tracker.inFlight("bucket"));
- p2.close();
- }
-
- public void testDrainToZeroRemovesBucketThenReacquire() {
- WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker();
- Releasable p = tracker.acquire("bucket", 5);
- assertEquals(1, tracker.inFlight("bucket"));
- p.close();
- assertEquals(0, tracker.inFlight("bucket"));
- // re-acquiring after the bucket drained (and was removed) works and starts from 1
- Releasable p2 = tracker.acquire("bucket", 5);
- assertEquals(1, tracker.inFlight("bucket"));
- p2.close();
- }
-
- public void testReleaseIsIdempotent() {
- WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker();
- Releasable p = tracker.acquire("bucket", 5);
- p.close();
- p.close(); // double close must not decrement twice
- assertEquals(0, tracker.inFlight("bucket"));
- }
-
- public void testBucketsAreIndependent() {
- WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker();
- tracker.acquire("a", 1);
- // "a" is full but "b" is a separate bucket
- expectThrows(OpenSearchRejectedExecutionException.class, () -> tracker.acquire("a", 1));
- Releasable pb = tracker.acquire("b", 1);
- assertEquals(1, tracker.inFlight("a"));
- assertEquals(1, tracker.inFlight("b"));
- pb.close();
- }
-}