From 41c972cc9933aa8d6af570adf952b4fe15546ee0 Mon Sep 17 00:00:00 2001 From: David Zane Date: Fri, 28 Aug 2026 16:27:02 -0700 Subject: [PATCH 1/2] Queue WLM requests denied by a throttle limit Signed-off-by: David Zane --- .../plugin/wlm/WlmClusterThrottlingIT.java | 115 ++++ .../opensearch/plugin/wlm/WlmQueueingIT.java | 413 ++++++++++++ .../action/search/TransportSearchAction.java | 2 +- .../cluster/metadata/WorkloadGroup.java | 20 +- .../main/java/org/opensearch/node/Node.java | 11 + .../wlm/MutableWorkloadGroupFragment.java | 70 +- .../opensearch/wlm/SharedThrottleTracker.java | 18 +- .../opensearch/wlm/WorkloadGroupQueue.java | 248 +++++++ .../wlm/WorkloadGroupQueueService.java | 396 +++++++++++ .../wlm/WorkloadGroupQueueSettings.java | 131 ++++ .../opensearch/wlm/WorkloadGroupService.java | 264 +++++++- .../WorkloadGroupSharedThrottleService.java | 619 +++++++++++++++++- .../org/opensearch/wlm/WorkloadGroupTask.java | 60 ++ .../wlm/stats/WorkloadGroupState.java | 83 +++ .../wlm/stats/WorkloadGroupStats.java | 133 +++- .../cluster/wlm/WlmStatsResponseTests.java | 7 + .../cluster/metadata/WorkloadGroupTests.java | 103 ++- .../wlm/SharedThrottleTrackerTests.java | 24 + .../wlm/WorkloadGroupQueueServiceTests.java | 136 ++++ .../wlm/WorkloadGroupQueueSettingsTests.java | 81 +++ .../wlm/WorkloadGroupQueueTests.java | 228 +++++++ .../wlm/WorkloadGroupServiceTests.java | 187 +++++- ...rkloadGroupSharedThrottleServiceTests.java | 128 ++++ .../wlm/WorkloadGroupTaskTests.java | 29 + .../opensearch/wlm/stats/WlmStatsTests.java | 2 +- .../wlm/stats/WorkloadGroupStateTests.java | 23 + .../wlm/stats/WorkloadGroupStatsTests.java | 102 ++- 27 files changed, 3550 insertions(+), 83 deletions(-) create mode 100644 plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmQueueingIT.java create mode 100644 server/src/main/java/org/opensearch/wlm/WorkloadGroupQueue.java create mode 100644 server/src/main/java/org/opensearch/wlm/WorkloadGroupQueueService.java create mode 100644 server/src/main/java/org/opensearch/wlm/WorkloadGroupQueueSettings.java create mode 100644 server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueServiceTests.java create mode 100644 server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueSettingsTests.java create mode 100644 server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueTests.java 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 index 84da62a996194..8be040074a1e6 100644 --- 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 @@ -42,6 +42,7 @@ import org.opensearch.test.OpenSearchIntegTestCase; import org.opensearch.wlm.MutableWorkloadGroupFragment; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadGroupQueueSettings; import org.opensearch.wlm.WorkloadGroupThrottleSettings; import org.opensearch.wlm.WorkloadManagementSettings; import org.opensearch.wlm.stats.WlmStats; @@ -206,6 +207,91 @@ public void testClusterWideCeilingHoldsAcrossCoordinators() throws Exception { }, 30, TimeUnit.SECONDS); } + /** + * Shared-tier QUEUEING via owner-push, on a shared-ONLY group (no {@code node_limit}). This is the exact scenario + * that the enqueue-first fix protects: a request denied at the shared limit is PARKED on its coordinator, and the + * bucket owner pushes it a grant when a slot frees — with no node-tier drain to fall back on. Because the group is + * shared-only, the ONLY way the parked request can ever complete is the cross-node owner-push path. + *

+ * The enqueue-first ordering matters here: the request is placed in the coordinator's queue BEFORE the shared + * acquire is sent, so a grant (or owner-push) can never arrive to find an empty queue and deregister the waiter, + * which would otherwise strand the request (there is no queue timeout to rescue it). + */ + public void testSharedTierQueuesAndDrainsViaOwnerPush() throws Exception { + String workloadGroupId = "wlm_shared_queue_group"; + String ruleId = "wlm_shared_queue_rule"; + String indexName = "shared_queue_index"; + + setWlmMode("enabled"); + + // shared_limit=1, node_limit unset, queue.size_per_bucket=5: one request runs cluster-wide; the rest PARK and drain via + // owner-push. No node tier exists, so owner-push is the sole drain path. + WorkloadGroup workloadGroup = createSharedThrottledQueueingGroup("shared_queue_test_group", workloadGroupId, 1, 5); + 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 queue rule", indexName, featureType, workloadGroupId); + indexDocument(indexName); + + // Wait for rule propagation on every coordinator this test drives (same rationale as the throttling test). + for (String node : internalCluster().getNodeNames()) { + assertBusy(() -> { + long before = getCompletions(workloadGroupId); + try { + client(node).prepareSearch(indexName).setQuery(QueryBuilders.matchAllQuery()).get(); + } catch (Exception e) { + 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 workload group", after > before); + }, 30, TimeUnit.SECONDS); + } + + List plugins = initBlockFactory(); + List coordinators = new ArrayList<>(List.of(internalCluster().getNodeNames())); + + // First search fills the single cluster-wide shared slot and blocks in-flight. + ActionFuture first = blockingSearchVia(coordinators.get(0), indexName).execute(); + awaitBlockedCount(plugins, 1); + + long throttledBefore = getThrottled(workloadGroupId); + long totalQueuedBefore = getTotalQueued(workloadGroupId); + + // Second search on a DIFFERENT coordinator: the shared limit is reached, so instead of a 429 it must be PARKED + // (enqueue-first) and registered for owner-push. It is NOT throttled. + ActionFuture second = blockingSearchVia(coordinators.get(1 % coordinators.size()), indexName).execute(); + assertBusy( + () -> assertEquals("second search should be parked in the queue", 1, getQueuedCurrent(workloadGroupId)), + 30, + TimeUnit.SECONDS + ); + assertEquals("a parked request must not be counted as throttled", throttledBefore, getThrottled(workloadGroupId)); + assertEquals("the parked request should be counted as queued", totalQueuedBefore + 1, getTotalQueued(workloadGroupId)); + + // Release the blocks. The first completes and frees the single shared slot; the owner pushes a grant to the + // coordinator holding the parked second search, which then drains and completes. With no node tier, this can + // ONLY happen via cross-node owner-push — the path the enqueue-first fix keeps race-free. + disableBlocks(plugins); + assertNotNull("the first (blocking) search must complete", first.actionGet(TIMEOUT)); + assertNotNull("the parked second search must be admitted via owner-push and complete", second.actionGet(TIMEOUT)); + + // The queue drains back to empty. + assertBusy(() -> assertEquals("queue must drain to empty", 0, getQueuedCurrent(workloadGroupId)), 30, TimeUnit.SECONDS); + } + // Helpers private static boolean hasRejectedExecutionCause(Throwable t) { @@ -228,6 +314,14 @@ private long getThrottled(String groupId) throws Exception { return sumAcrossNodes(groupId, WorkloadGroupStatsHolder::getThrottled); } + private long getQueuedCurrent(String groupId) throws Exception { + return sumAcrossNodes(groupId, WorkloadGroupStatsHolder::getQueuedCurrent); + } + + private long getTotalQueued(String groupId) throws Exception { + return sumAcrossNodes(groupId, WorkloadGroupStatsHolder::getQueued); + } + // 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); @@ -314,6 +408,27 @@ private WorkloadGroup createSharedThrottledGroup(String name, String id, int sha ); } + // Shared-only throttling (no node_limit) WITH queueing enabled, so a denied request parks and drains via owner-push. + private WorkloadGroup createSharedThrottledQueueingGroup(String name, String id, int sharedLimit, int queueSizePerBucket) { + Settings throttling = Settings.builder() + .put(WorkloadGroupThrottleSettings.ATTRIBUTE.getKey(), "group") + .put(WorkloadGroupThrottleSettings.SHARED_LIMIT.getKey(), sharedLimit) + .build(); + Settings queue = Settings.builder().put(WorkloadGroupQueueSettings.SIZE_PER_BUCKET.getKey(), queueSizePerBucket).build(); + return new WorkloadGroup( + name, + id, + new MutableWorkloadGroupFragment( + MutableWorkloadGroupFragment.ResiliencyMode.SOFT, + Map.of(ResourceType.CPU, 0.9, ResourceType.MEMORY, 0.9), + Settings.EMPTY, + throttling, + queue + ), + Instant.now().getMillis() + ); + } + private void indexDocument(String indexName) { assertAcked( client().admin() diff --git a/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmQueueingIT.java b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmQueueingIT.java new file mode 100644 index 0000000000000..5f9c4a01a0d42 --- /dev/null +++ b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmQueueingIT.java @@ -0,0 +1,413 @@ +/* + * 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.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.WorkloadGroupQueueSettings; +import org.opensearch.wlm.WorkloadGroupThrottleSettings; +import org.opensearch.wlm.WorkloadManagementSettings; +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 static org.opensearch.index.query.QueryBuilders.scriptQuery; +import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.greaterThanOrEqualTo; + +/** + * End-to-end integration test for WLM request QUEUEING on top of node-level throttling. With {@code node_limit=1} and + * {@code queue.size_per_bucket>0}, a second concurrent search that would be rejected by the throttle is instead PARKED in the + * queue (holding no thread) and admitted once the first search completes and frees the permit — exercising the real + * coordinator admission + node-completion drain path, not a mocked service. + */ +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 1, numClientNodes = 0, supportsDedicatedMasters = false) +public class WlmQueueingIT 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(ScriptedBlockPlugin.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 testSecondConcurrentSearchQueuedThenAdmitted() throws Exception { + String workloadGroupId = "wlm_queue_group"; + String ruleId = "wlm_queue_rule"; + String indexName = "queue_index"; + + setWlmMode("enabled"); + + // node_limit=1 with queueing enabled (size 5): the 2nd concurrent search parks instead of 429ing. + WorkloadGroup workloadGroup = createQueueingWorkloadGroup("queue_test_group", workloadGroupId, 1, 5); + updateWorkloadGroupInClusterState(PUT, workloadGroup); + + FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME); + createRule(ruleId, "queue rule", indexName, featureType, workloadGroupId); + + indexDocument(indexName); + + // Wait for rule propagation: a (non-blocking) search must be tagged to the group before the scenario. + assertBusy(() -> { + int before = getStat(workloadGroupId, "total_completions"); + client().prepareSearch(indexName).setQuery(org.opensearch.index.query.QueryBuilders.matchAllQuery()).get(); + int after = getStat(workloadGroupId, "total_completions"); + assertTrue("Expected search to be tagged to the workload group", after > before); + }, 30, TimeUnit.SECONDS); + + List plugins = initBlockFactory(); + + // First search blocks in the query phase, holding the single permit. + ActionFuture firstBlocked = blockingSearch(indexName).execute(); + awaitForBlock(plugins); + + int queuedBefore = getStat(workloadGroupId, "total_queued"); + int throttledBefore = getStat(workloadGroupId, "total_throttled"); + + // Second search while the first is in-flight: it must be QUEUED (parked), not rejected. Run it async and assert + // it gets parked (total_queued increments) and does NOT fail with a 429. + ActionFuture secondQueued = blockingSearch(indexName).execute(); + assertBusy(() -> { + assertEquals("second search should be parked in the queue", queuedBefore + 1, getStat(workloadGroupId, "total_queued")); + }, 30, TimeUnit.SECONDS); + // It was queued, not throttle-rejected. + assertEquals("a queued request must not be counted as throttled", throttledBefore, getStat(workloadGroupId, "total_throttled")); + // queued_current reflects the one parked request. + assertEquals("one request should currently be queued", 1, getStat(workloadGroupId, "queued_current")); + + // Release the block. The first search completes, frees the permit, and the node-completion drain admits the + // parked second search — which then also runs (and completes, since blocks are now disabled). + disableBlocks(plugins); + assertNotNull(firstBlocked.actionGet(TIMEOUT)); + assertNotNull("the queued search must be admitted and complete once a permit frees", secondQueued.actionGet(TIMEOUT)); + + // The queue drains back to empty. + assertBusy(() -> assertEquals("queue must drain to empty", 0, getStat(workloadGroupId, "queued_current")), 30, TimeUnit.SECONDS); + + // The admitted-from-queue request contributed to the queue-wait aggregate: exactly one recorded wait. The + // magnitude is not asserted (the test releases the block as soon as the park is observed, so the actual wait + // can be sub-millisecond; the arithmetic of sum/mean/max is covered deterministically in unit tests). Total and + // max must be internally consistent: max <= total when there is a single sample, and both non-negative. + int waitCount = getStat(workloadGroupId, "queue_wait_count"); + int totalWait = getStat(workloadGroupId, "total_queue_wait_millis"); + int maxWait = getStat(workloadGroupId, "max_queue_wait_millis"); + assertEquals("one admitted request should have a recorded queue wait", 1, waitCount); + assertThat("total queue wait must be non-negative", totalWait, greaterThanOrEqualTo(0)); + assertThat("max queue wait must be non-negative", maxWait, greaterThanOrEqualTo(0)); + assertEquals("with a single sample, max equals total", totalWait, maxWait); + } + + public void testQueueFullRejectsWith429() throws Exception { + String workloadGroupId = "wlm_queuefull_group"; + String ruleId = "wlm_queuefull_rule"; + String indexName = "queuefull_index"; + + setWlmMode("enabled"); + + // node_limit=1, queue.size_per_bucket=1: 1 running + 1 queued is the ceiling; a 3rd concurrent search is rejected. + WorkloadGroup workloadGroup = createQueueingWorkloadGroup("queuefull_test_group", workloadGroupId, 1, 1); + updateWorkloadGroupInClusterState(PUT, workloadGroup); + + FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME); + createRule(ruleId, "queuefull rule", indexName, featureType, workloadGroupId); + + indexDocument(indexName); + + assertBusy(() -> { + int before = getStat(workloadGroupId, "total_completions"); + client().prepareSearch(indexName).setQuery(org.opensearch.index.query.QueryBuilders.matchAllQuery()).get(); + int after = getStat(workloadGroupId, "total_completions"); + assertTrue("Expected search to be tagged to the workload group", after > before); + }, 30, TimeUnit.SECONDS); + + List plugins = initBlockFactory(); + + // First search blocks, holding the permit. + ActionFuture firstBlocked = blockingSearch(indexName).execute(); + awaitForBlock(plugins); + + // Second search parks (fills the single queue slot). + ActionFuture secondQueued = blockingSearch(indexName).execute(); + assertBusy( + () -> assertEquals("second search should be parked", 1, getStat(workloadGroupId, "queued_current")), + 30, + TimeUnit.SECONDS + ); + + int queueRejectionsBefore = getStat(workloadGroupId, "total_queue_rejections"); + + // Third concurrent search: permit taken, queue full -> 429. + Throwable rejection = expectThrows(Throwable.class, () -> blockingSearch(indexName).get()); + assertTrue( + "Expected an OpenSearchRejectedExecutionException in the cause chain but was: " + rejection, + hasRejectedExecutionCause(rejection) + ); + assertEquals( + "queue-full rejection should increment total_queue_rejections", + queueRejectionsBefore + 1, + getStat(workloadGroupId, "total_queue_rejections") + ); + + // Release; the first and the queued second both complete. + disableBlocks(plugins); + assertNotNull(firstBlocked.actionGet(TIMEOUT)); + assertNotNull(secondQueued.actionGet(TIMEOUT)); + } + + // 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 int getStat(String groupId, String fieldName) throws Exception { + WlmStatsRequest request = new WlmStatsRequest(null, new HashSet<>(Collections.singletonList(groupId)), null); + WlmStatsResponse response = client().execute(WlmStatsAction.INSTANCE, request).get(); + return extractStatField(response.toString(), groupId, fieldName); + } + + private int extractStatField(String responseBody, String workloadGroupId, String fieldName) { + int total = 0; + String groupKey = "\"" + workloadGroupId + "\""; + String field = "\"" + fieldName + "\""; + int index = 0; + while ((index = responseBody.indexOf(groupKey, index)) != -1) { + int groupStart = responseBody.indexOf("{", index); + int fieldIndex = responseBody.indexOf(field, groupStart); + if (fieldIndex == -1) break; + int colonIndex = responseBody.indexOf(":", fieldIndex); + int commaIndex = responseBody.indexOf(",", colonIndex); + int braceIndex = responseBody.indexOf("}", colonIndex); + int end = (commaIndex == -1 || (braceIndex != -1 && braceIndex < commaIndex)) ? braceIndex : commaIndex; + String numberStr = responseBody.substring(colonIndex + 1, end).trim(); + total += Integer.parseInt(numberStr); + index = end; + } + return total; + } + + private SearchRequestBuilder blockingSearch(String indexName) { + return client().prepareSearch(indexName) + .setQuery(scriptQuery(new Script(ScriptType.INLINE, "mockscript", ScriptedBlockPlugin.SCRIPT_NAME, Collections.emptyMap()))); + } + + private List initBlockFactory() { + List plugins = new ArrayList<>(); + for (PluginsService pluginsService : internalCluster().getDataNodeInstances(PluginsService.class)) { + plugins.addAll(pluginsService.filterPlugins(ScriptedBlockPlugin.class)); + } + for (ScriptedBlockPlugin plugin : plugins) { + plugin.reset(); + plugin.enableBlock(); + } + return plugins; + } + + private void awaitForBlock(List plugins) throws Exception { + assertBusy(() -> { + int blocked = 0; + for (ScriptedBlockPlugin plugin : plugins) { + blocked += plugin.hits.get(); + } + assertThat(blocked, greaterThan(0)); + }); + } + + private void disableBlocks(List plugins) { + for (ScriptedBlockPlugin 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 createQueueingWorkloadGroup(String name, String id, int nodeLimit, int queueSizePerBucket) { + Settings throttling = Settings.builder() + .put(WorkloadGroupThrottleSettings.ATTRIBUTE.getKey(), "group") + .put(WorkloadGroupThrottleSettings.NODE_LIMIT.getKey(), nodeLimit) + .build(); + Settings queue = Settings.builder().put(WorkloadGroupQueueSettings.SIZE_PER_BUCKET.getKey(), queueSizePerBucket).build(); + return new WorkloadGroup( + name, + id, + new MutableWorkloadGroupFragment( + MutableWorkloadGroupFragment.ResiliencyMode.SOFT, + Map.of(ResourceType.CPU, 0.9, ResourceType.MEMORY, 0.9), + Settings.EMPTY, + throttling, + queue + ), + Instant.now().getMillis() + ); + } + + private void indexDocument(String indexName) { + assertAcked( + client().admin() + .indices() + .prepareCreate(indexName) + .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. + */ + public static class ScriptedBlockPlugin extends MockScriptPlugin { + static final String SCRIPT_NAME = "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(WlmQueueingIT.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 0d781c597224e..a72265992b860 100644 --- a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java +++ b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java @@ -511,7 +511,7 @@ void executeRequest( }, outerListener::onFailure), threadPool.getThreadContext() ); - workloadGroupService.acquireThrottlePermit(((WorkloadGroupTask) task).getWorkloadGroupId(), principal, admissionListener); + workloadGroupService.acquireThrottlePermit((WorkloadGroupTask) task, principal, admissionListener); } else { proceedWithSearch( task, diff --git a/server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java index 92ff011422f8b..8635d465d6114 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java @@ -78,18 +78,25 @@ public WorkloadGroup(String name, String _id, MutableWorkloadGroupFragment mutab // Drop null-valued "clear" keys before storage (meaningful only during an update merge, not on create). Settings normalizedSettings = stripClearMarkers(mutableWorkloadGroupFragment.getSettings()); Settings normalizedThrottling = stripClearMarkers(mutableWorkloadGroupFragment.getThrottling()); + Settings normalizedQueue = stripClearMarkers(mutableWorkloadGroupFragment.getQueue()); if (normalizedSettings.equals(mutableWorkloadGroupFragment.getSettings()) == false - || normalizedThrottling.equals(mutableWorkloadGroupFragment.getThrottling()) == false) { + || normalizedThrottling.equals(mutableWorkloadGroupFragment.getThrottling()) == false + || normalizedQueue.equals(mutableWorkloadGroupFragment.getQueue()) == false) { mutableWorkloadGroupFragment = new MutableWorkloadGroupFragment( mutableWorkloadGroupFragment.getResiliencyMode(), mutableWorkloadGroupFragment.getResourceLimits(), normalizedSettings, - normalizedThrottling + normalizedThrottling, + normalizedQueue ); } // Cross-field checks on the merged throttling config (attribute required with a limit; ceiling must be >= 1). WorkloadGroupThrottleSettings.validateMergedConfig(mutableWorkloadGroupFragment.getThrottling()); + // A queue with no throttle limit has nothing to queue: queueing engages only on a throttle denial. + if (mutableWorkloadGroupFragment.getQueue().isEmpty() == false && mutableWorkloadGroupFragment.getThrottling().isEmpty()) { + throw new IllegalArgumentException("queue requires a throttle limit; set throttling.node_limit or throttling.shared_limit"); + } this.name = name; this._id = _id; @@ -127,10 +134,14 @@ public static WorkloadGroup updateExistingWorkloadGroup( existingGroup.getMutableWorkloadGroupFragment().getThrottling(), mutableWorkloadGroupFragment.getThrottling() ); + final Settings updatedQueue = mergeSettings( + existingGroup.getMutableWorkloadGroupFragment().getQueue(), + mutableWorkloadGroupFragment.getQueue() + ); return new WorkloadGroup( existingGroup.getName(), existingGroup.get_id(), - new MutableWorkloadGroupFragment(mode, updatedResourceLimits, updatedSettings, updatedThrottling), + new MutableWorkloadGroupFragment(mode, updatedResourceLimits, updatedSettings, updatedThrottling, updatedQueue), Instant.now().getMillis() ); } @@ -336,7 +347,8 @@ public static Builder fromXContent(XContentParser parser) throws IOException { mutableWorkloadGroupFragment1.parseField(parser, fieldName); } else if (token == XContentParser.Token.VALUE_NULL) { if (fieldName.equals(MutableWorkloadGroupFragment.SETTINGS_STRING) - || fieldName.equals(MutableWorkloadGroupFragment.THROTTLING_STRING)) { + || fieldName.equals(MutableWorkloadGroupFragment.THROTTLING_STRING) + || fieldName.equals(MutableWorkloadGroupFragment.QUEUE_STRING)) { mutableWorkloadGroupFragment1.parseField(parser, fieldName); } } diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index a7104e5a45862..d433884526f57 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -308,6 +308,7 @@ import org.opensearch.transport.client.node.NodeClient; import org.opensearch.usage.UsageService; import org.opensearch.watcher.ResourceWatcherService; +import org.opensearch.wlm.WorkloadGroupQueueService; import org.opensearch.wlm.WorkloadGroupService; import org.opensearch.wlm.WorkloadGroupSharedThrottleService; import org.opensearch.wlm.WorkloadGroupsStateAccessor; @@ -1434,6 +1435,16 @@ protected Node(final Environment initialEnvironment, Collection clas ); workloadGroupService.setSharedThrottleService(workloadGroupSharedThrottleService); + // Coordinator-local request queues: park a throttle-denied search instead of rejecting it, and admit it + // when a permit frees (node-completion drain, owner-push grant, or the WorkloadGroupService backstop sweep). + final WorkloadGroupQueueService workloadGroupQueueService = new WorkloadGroupQueueService( + threadPool, + workloadGroupsStateAccessor + ); + workloadGroupService.setQueueService(workloadGroupQueueService); + // Owner-push grant -> admit one queued request against the reserved shared permit. + workloadGroupSharedThrottleService.setGrantConsumer(workloadGroupQueueService::admitWithPermit); + TopNSearchTasksLogger taskConsumer = new TopNSearchTasksLogger(settings, settingsModule.getClusterSettings()); transportService.getTaskManager().registerTaskResourceConsumer(taskConsumer); streamTransportService.ifPresent(service -> service.getTaskManager().registerTaskResourceConsumer(taskConsumer)); diff --git a/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java b/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java index 1d6b6b332eac1..31fcbb36554b8 100644 --- a/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java +++ b/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java @@ -37,16 +37,19 @@ public class MutableWorkloadGroupFragment extends AbstractDiffable resourceLimits; private Settings settings; private Settings throttling; + private Settings queue; public static final List acceptedFieldNames = List.of( RESILIENCY_MODE_STRING, RESOURCE_LIMITS_STRING, SETTINGS_STRING, - THROTTLING_STRING + THROTTLING_STRING, + QUEUE_STRING ); public MutableWorkloadGroupFragment() {} @@ -67,14 +70,26 @@ public MutableWorkloadGroupFragment( Map resourceLimits, Settings settings, Settings throttling + ) { + this(resiliencyMode, resourceLimits, settings, throttling, Settings.EMPTY); + } + + public MutableWorkloadGroupFragment( + ResiliencyMode resiliencyMode, + Map resourceLimits, + Settings settings, + Settings throttling, + Settings queue ) { validateResourceLimits(resourceLimits); WorkloadGroupSearchSettings.validate(settings); WorkloadGroupThrottleSettings.validate(throttling); + WorkloadGroupQueueSettings.validate(queue); this.resiliencyMode = resiliencyMode; this.resourceLimits = resourceLimits; this.settings = settings != null ? settings : Settings.EMPTY; this.throttling = throttling != null ? throttling : Settings.EMPTY; + this.queue = queue != null ? queue : Settings.EMPTY; } public MutableWorkloadGroupFragment(StreamInput in) throws IOException { @@ -85,9 +100,13 @@ public MutableWorkloadGroupFragment(StreamInput in) throws IOException { } String updatedResiliencyMode = in.readOptionalString(); resiliencyMode = updatedResiliencyMode == null ? null : ResiliencyMode.fromName(updatedResiliencyMode); + // TODO(merge): this block gained a third read (queue) on the queueing branch while staying gated at V_3_7_0, so + // it is shape-compatible only while queueing ships in the SAME release as throttling. If they split, give queue + // its own version gate — see the matching note on the write path below. if (in.getVersion().onOrAfter(Version.V_3_7_0)) { settings = Settings.readOptionalSettingsFromStream(in); throttling = Settings.readOptionalSettingsFromStream(in); + queue = Settings.readOptionalSettingsFromStream(in); } else if (in.getVersion().onOrAfter(Version.V_3_6_0)) { // Legacy 3.6 format: read and discard (experimental API, no backward compat guarantee) boolean isNull = in.readBoolean(); @@ -96,9 +115,11 @@ public MutableWorkloadGroupFragment(StreamInput in) throws IOException { } settings = Settings.EMPTY; throttling = Settings.EMPTY; + queue = Settings.EMPTY; } else { settings = Settings.EMPTY; throttling = Settings.EMPTY; + queue = Settings.EMPTY; } } @@ -152,6 +173,18 @@ public Settings parseField(XContentParser parser) throws IOException { } } + static class QueueParser implements FieldParser { + public Settings parseField(XContentParser parser) throws IOException { + // "queue": null means clear all queue config (disable queueing) + if (parser.currentToken() == XContentParser.Token.VALUE_NULL) { + return Settings.EMPTY; + } + Settings queue = Settings.fromXContent(parser); + WorkloadGroupQueueSettings.validate(queue); + return queue; + } + } + static class FieldParserFactory { static Optional> fieldParserFor(String fieldName) { return switch (fieldName) { @@ -159,6 +192,7 @@ static Optional> fieldParserFor(String fieldName) { case RESOURCE_LIMITS_STRING -> Optional.of(new ResourceLimitsParser()); case SETTINGS_STRING -> Optional.of(new SearchSettingsParser()); case THROTTLING_STRING -> Optional.of(new ThrottlingParser()); + case QUEUE_STRING -> Optional.of(new QueueParser()); default -> Optional.empty(); }; } @@ -214,6 +248,21 @@ static Optional> fieldParserFor(String fieldName) { } catch (IOException e) { throw new IllegalStateException("writing error encountered for the field " + THROTTLING_STRING); } + }, QUEUE_STRING, (builder) -> { + try { + // Like throttling, queue config is omitted entirely when unset. + Settings q = queue != null ? queue : Settings.EMPTY; + if (q.isEmpty() == false) { + builder.startObject(QUEUE_STRING); + if (q.hasValue(WorkloadGroupQueueSettings.SIZE_PER_BUCKET.getKey())) { + builder.field(WorkloadGroupQueueSettings.SIZE_PER_BUCKET.getKey(), WorkloadGroupQueueSettings.SIZE_PER_BUCKET.get(q)); + } + builder.endObject(); + } + return null; + } catch (IOException e) { + throw new IllegalStateException("writing error encountered for the field " + QUEUE_STRING); + } }); private static void writeSettingsFields(XContentBuilder builder, Settings s) throws IOException { @@ -241,6 +290,7 @@ public void parseField(XContentParser parser, String field) { case RESOURCE_LIMITS_STRING -> setResourceLimits((Map) value); case SETTINGS_STRING -> setSettings((Settings) value); case THROTTLING_STRING -> setThrottling((Settings) value); + case QUEUE_STRING -> setQueue((Settings) value); } } catch (IllegalArgumentException e) { throw e; @@ -265,8 +315,10 @@ public void writeTo(StreamOutput out) throws IOException { out.writeOptionalString(resiliencyMode == null ? null : resiliencyMode.getName()); if (out.getVersion().onOrAfter(Version.V_3_7_0)) { Settings.writeOptionalSettingsToStream(settings, out); - // TODO: when upstreamed, gate throttling behind its own release version to support mixed-build clusters. + // TODO(merge): when upstreamed, gate throttling/queue behind their own release version to support mixed-build + // clusters. Today both ride the V_3_7_0 gate, which only works because they ship in the same release. Settings.writeOptionalSettingsToStream(throttling, out); + Settings.writeOptionalSettingsToStream(queue, out); } else if (out.getVersion().onOrAfter(Version.V_3_6_0)) { // Legacy 3.6 format: write empty map (experimental API, settings not preserved across versions) out.writeBoolean(false); @@ -297,12 +349,13 @@ public boolean equals(Object o) { return Objects.equals(resiliencyMode, that.resiliencyMode) && Objects.equals(resourceLimits, that.resourceLimits) && Objects.equals(settings, that.settings) - && Objects.equals(throttling, that.throttling); + && Objects.equals(throttling, that.throttling) + && Objects.equals(queue, that.queue); } @Override public int hashCode() { - return Objects.hash(resiliencyMode, resourceLimits, settings, throttling); + return Objects.hash(resiliencyMode, resourceLimits, settings, throttling, queue); } public ResiliencyMode getResiliencyMode() { @@ -321,6 +374,10 @@ public Settings getThrottling() { return throttling; } + public Settings getQueue() { + return queue; + } + /** * This enum models the different WorkloadGroup resiliency modes * SOFT - means that this workload group can consume more than workload group resource limits if node is not in duress @@ -371,4 +428,9 @@ void setThrottling(Settings throttling) { this.throttling = throttling != null ? throttling : Settings.EMPTY; } + void setQueue(Settings queue) { + WorkloadGroupQueueSettings.validate(queue); + this.queue = queue != null ? queue : Settings.EMPTY; + } + } diff --git a/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java b/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java index 090d34e981bfc..054bd93acf07f 100644 --- a/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java +++ b/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java @@ -10,7 +10,9 @@ import org.opensearch.common.annotation.ExperimentalApi; +import java.util.ArrayList; import java.util.Iterator; +import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.LongAdder; @@ -138,15 +140,29 @@ public void release(String bucketKey, String permitId) { /** * Reclaims all expired permits 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}. + * + * @return the bucket keys for which at least one permit was reclaimed, i.e. those that just gained free capacity. + * The caller must drive owner-push for these: an expiring permit is the only free-slot signal available when + * the holder crashed or its release RPC was lost, so ignoring it strands a waiting coordinator's parked + * request (there is no queue timeout to rescue it). Empty when nothing expired, which is the common case. */ - public void sweepExpired() { + public List sweepExpired() { final long now = nanoTimeSupplier.getAsLong(); + final List freedBuckets = new ArrayList<>(); for (String bucketKey : permitsByBucket.keySet()) { permitsByBucket.computeIfPresent(bucketKey, (k, bucket) -> { + final int before = bucket.permits.size(); pruneExpired(bucket, now); + if (bucket.permits.size() < before) { + // At least one slot just freed for this bucket. Reclaiming a permit here is the ONLY signal for a + // holder that crashed or whose release RPC was lost — there is no release RPC to drive owner-push — + // so the caller must be told, or a coordinator's parked request can strand with free capacity. + freedBuckets.add(bucketKey); + } return bucket.permits.isEmpty() ? null : bucket; }); } + return freedBuckets; } // Current live (non-expired-at-read-time) count for a bucket. Package-private for tests. diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupQueue.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupQueue.java new file mode 100644 index 0000000000000..9156fa2fe7d42 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupQueue.java @@ -0,0 +1,248 @@ +/* + * 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 org.opensearch.common.lease.Releasable; +import org.opensearch.core.action.ActionListener; + +import java.util.LinkedHashSet; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +/** + * A bounded, per-coordinator request queue for a single workload group. When a search is denied by a throttle limit, + * instead of an immediate 429 the coordinator parks it here and admits it once a permit frees (node-tier completion or + * a cluster-tier owner grant). + *

+ * The queue is partitioned into a per-bucket FIFO ({@code byBucket}) so a permit freed for one bucket wakes a request + * waiting on that same bucket, and a heavily-queued bucket cannot head-of-line-block another. Capacity is bounded on + * two axes (see {@link #offer}): a per-bucket cap ({@code queue.size_per_bucket}, for cross-principal fairness) and a + * fixed per-group total ({@link WorkloadGroupQueueSettings#MAX_GROUP_QUEUE_DEPTH}, a footprint backstop). A parked + * request holds no thread — only its {@link ActionListener} and open client connection — so the queue's footprint + * (heap + sockets) is what those two limits bound; {@code depth} tracks the group total for the fixed ceiling. + *

+ * Thread-safety: {@code depth} is a shared {@link AtomicInteger}; each bucket's {@link LinkedHashSet} is guarded by the + * caller holding this group's per-bucket lock (a {@code KeyedLock} in {@link WorkloadGroupQueueService}). Callers must + * never invoke a parked listener while holding that lock. + *

+ * Each bucket's container is a {@link LinkedHashSet}: insertion order gives the per-bucket FIFO (head = oldest), while + * membership-based removal is O(1) average. This matters because a request can be removed from the middle on + * cancellation ({@link #remove}); an {@code ArrayDeque} would make that O(n), so a mass-cancellation storm on one + * bucket (each cancel firing an independent remove) would be O(n^2). {@link QueuedRequest} has no {@code equals}/ + * {@code hashCode} override, so identity semantics hold and distinct requests never collide. + */ +@ExperimentalApi +public class WorkloadGroupQueue { + + /** + * A parked request: the held listener (already context-preserving, wrapped upstream), its bucket key, the owning + * task, and the handle that deregisters its task-cancellation callback once it is admitted or evicted. + *

+ * A parked request has no wall-clock deadline. There is deliberately no queue timeout: legitimate queue wait is + * unbounded (it grows with backlog depth over throughput), so any fixed cap would eventually cancel healthy, + * still-connected requests under a large slow burst. A client bounds its own wait with + * {@code cancel_after_time_interval} (or a disconnect); either cancels the task, which evicts the entry promptly. + */ + @ExperimentalApi + public static class QueuedRequest { + final ActionListener listener; + final String bucketKey; + final WorkloadGroupTask task; + // Relative-clock instant (nanos) the request was parked, for observability (queue-wait metric). Stable across + // the request's queue lifetime — the request is never re-parked, so this is the true total-wait basis. + final long enqueueNanos; + // Set after construction (the callback needs the enqueued reference); deregisters the cancellation callback. + volatile Releasable cancellationHandle; + + public QueuedRequest(ActionListener listener, String bucketKey, WorkloadGroupTask task, long enqueueNanos) { + this.listener = listener; + this.bucketKey = bucketKey; + this.task = task; + this.enqueueNanos = enqueueNanos; + } + + public ActionListener listener() { + return listener; + } + + public String bucketKey() { + return bucketKey; + } + + public WorkloadGroupTask task() { + return task; + } + + /** Time parked so far, in nanos, as of {@code nowNanos} (relative clock). Never negative. */ + long waitNanos(long nowNanos) { + long w = nowNanos - enqueueNanos; + return w < 0 ? 0 : w; + } + + void releaseCancellationHandle() { + Releasable handle = cancellationHandle; + if (handle != null) { + handle.close(); + } + } + } + + private final Map> byBucket = new ConcurrentHashMap<>(); + private final AtomicInteger depth = new AtomicInteger(0); + private final AtomicLong peak = new AtomicLong(0); + + public WorkloadGroupQueue() {} + + /** + * Attempts to park a request. Must be called while holding the per-bucket lock for {@code req.bucketKey}. Enforces + * TWO limits and returns {@code false} (the caller then rejects with a 429) if either is hit: + *

    + *
  1. per-bucket — the request's own bucket already holds {@code sizePerBucket} parked requests + * ({@code sizePerBucket <= 0} means queueing is disabled). This is the user-facing {@code queue.size_per_bucket} + * knob, giving cross-principal fairness while the group total is below the ceiling: one bucket cannot + * consume another's per-bucket capacity. Note the fairness is bounded, not absolute — once the group total + * reaches the ceiling below, admission is first-come-first-served across buckets, so a high-cardinality flood + * can still crowd out a well-behaved bucket that has not yet filled its own allowance.
    2. + *
  2. per-group total — the group already holds {@link WorkloadGroupQueueSettings#MAX_GROUP_QUEUE_DEPTH} + * parked requests across all buckets on this coordinator. A fixed, non-configurable footprint backstop against + * attacker-controlled bucket cardinality (username/role buckets are derived from the request principal).
    3. + *
+ * {@code sizePerBucket} is passed per call (not stored) so a live {@code queue.size_per_bucket} change takes effect + * immediately: a decrease stops admitting to a bucket once it is at/above the new cap (already-parked requests are + * not evicted); an increase widens capacity at once. + *

+ * The per-bucket depth is read (not created) before reserving the shared group counter, so a group-ceiling rejection + * never leaves an empty bucket set behind — preserving the invariant that a present bucket key has a live waiter. + * + * @param req the request to park + * @param sizePerBucket the group's current {@code queue.size_per_bucket} + * @return {@code true} if the request was enqueued, {@code false} if rejected (disabled / bucket full / group full) + */ + boolean offer(QueuedRequest req, int sizePerBucket) { + if (sizePerBucket <= 0) { + return false; // queueing disabled + } + // Per-bucket cap. Read under this bucket's lock (held by the caller), so the bucket set is stable for this key. + LinkedHashSet existing = byBucket.get(req.bucketKey); + if (existing != null && existing.size() >= sizePerBucket) { + return false; // this bucket's queue is full + } + // Fixed per-group backstop on the TOTAL across all buckets. Reserve first; only touch the bucket set once the + // slot is secured, so the set and the depth counter never disagree (and no empty set is left on rejection). + int updated = depth.incrementAndGet(); + if (updated > WorkloadGroupQueueSettings.MAX_GROUP_QUEUE_DEPTH) { + depth.decrementAndGet(); + return false; // group-wide ceiling hit + } + peak.accumulateAndGet(updated, Math::max); + byBucket.computeIfAbsent(req.bucketKey, k -> new LinkedHashSet<>()).add(req); + return true; + } + + /** + * Returns the oldest parked request for {@code bucketKey} without removing it, or {@code null} if none. Must be + * called while holding the per-bucket lock. + */ + QueuedRequest peekOldest(String bucketKey) { + LinkedHashSet bucket = byBucket.get(bucketKey); + if (bucket == null) { + return null; + } + java.util.Iterator it = bucket.iterator(); + return it.hasNext() ? it.next() : null; // head = oldest (insertion order) + } + + /** + * Removes and returns the oldest parked request for {@code bucketKey}, or {@code null} if none. Must be called + * while holding the per-bucket lock. Decrements the shared depth and prunes the bucket entry when it empties. + */ + QueuedRequest pollOldest(String bucketKey) { + LinkedHashSet bucket = byBucket.get(bucketKey); + if (bucket == null) { + return null; + } + java.util.Iterator it = bucket.iterator(); + QueuedRequest req = null; + if (it.hasNext()) { + req = it.next(); // head = oldest (insertion order) + it.remove(); + depth.decrementAndGet(); + } + if (bucket.isEmpty()) { + byBucket.remove(bucketKey); + } + return req; + } + + /** + * Removes and returns every parked request in {@code bucketKey} whose task is cancelled, leaving all survivors in + * place. Must be called while holding the per-bucket lock. Depth is decremented for each removed request. + *

+ * This is the backstop sweep's cleanup: a defense-in-depth complement to the per-request cancellation callback + * (which normally evicts a cancelled entry immediately), catching any cancelled task the callback missed. There is + * no time-based eviction — a still-live parked request is never removed here regardless of how long it has waited; + * a corrupt/dead but uncancelled entry self-heals when it drains to the head and fails on execution. + */ + java.util.List evictCancelled(String bucketKey) { + LinkedHashSet bucket = byBucket.get(bucketKey); + if (bucket == null) { + return java.util.Collections.emptyList(); + } + java.util.List evicted = new java.util.ArrayList<>(); + for (java.util.Iterator it = bucket.iterator(); it.hasNext();) { + QueuedRequest req = it.next(); + if (req.task().isCancelled()) { + it.remove(); + depth.decrementAndGet(); + evicted.add(req); + } + } + if (bucket.isEmpty()) { + byBucket.remove(bucketKey); + } + return evicted; + } + + /** + * Removes a specific parked request from its bucket (used on cancellation/eviction). Must be called while holding + * the per-bucket lock. Returns {@code true} if it was present and removed (in which case depth is decremented). + * O(1) average — this is the hot path under a cancellation storm, so the bucket is a {@link LinkedHashSet} rather + * than a deque (whose {@code remove(Object)} would be O(n), making a mass cancel on one bucket O(n^2)). + */ + boolean remove(QueuedRequest req) { + LinkedHashSet bucket = byBucket.get(req.bucketKey); + if (bucket == null) { + return false; + } + boolean removed = bucket.remove(req); + if (removed) { + depth.decrementAndGet(); + } + if (bucket.isEmpty()) { + byBucket.remove(req.bucketKey); + } + return removed; + } + + /** Snapshot of the bucket keys that currently have at least one parked request. */ + java.util.Set bucketKeys() { + return byBucket.keySet(); + } + + int currentDepth() { + return depth.get(); + } + + long peakDepth() { + return peak.get(); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupQueueService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupQueueService.java new file mode 100644 index 0000000000000..139bf07eafe63 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupQueueService.java @@ -0,0 +1,396 @@ +/* + * 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 org.opensearch.common.lease.Releasable; +import org.opensearch.common.util.concurrent.KeyedLock; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.tasks.TaskCancelledException; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.stats.WorkloadGroupState; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; + +/** + * Coordinator-local owner of the per-workload-group request queues. When {@link WorkloadGroupService} would reject a + * search at a throttle limit, it parks the request here instead (if queueing is enabled and the group's queue has + * room). A parked request holds no thread — only its {@link ActionListener} and open connection — and is admitted + * later when a permit frees: + *

+ * Concurrency: each bucket's request set is guarded by a per-bucket lock ({@link KeyedLock}); the group-total depth + * counter is an atomic inside {@link WorkloadGroupQueue}. A parked listener is never completed while a bucket lock is held — every + * admission dispatches the listener completion onto an executor to avoid running the downstream search inline on the + * releasing/grant/sweep thread (and to avoid deep recursion when many drain at once). + */ +@ExperimentalApi +public class WorkloadGroupQueueService { + + private final ThreadPool threadPool; + private final WorkloadGroupsStateAccessor stateAccessor; + // One queue per group id, created lazily on first enqueue. The queue holds only the parked requests + depth; the + // per-bucket capacity (queue.size_per_bucket) is NOT baked into it — it is passed in per enqueue from live config, + // so a dynamic queue.size_per_bucket change takes effect immediately + // (WorkloadGroupQueue.offer(req, sizePerBucket)). The per-group ceiling is the fixed MAX_GROUP_QUEUE_DEPTH. + private final Map queuesByGroup = new ConcurrentHashMap<>(); + private final KeyedLock bucketLocks = new KeyedLock<>(); + + public WorkloadGroupQueueService(ThreadPool threadPool, WorkloadGroupsStateAccessor stateAccessor) { + this.threadPool = threadPool; + this.stateAccessor = stateAccessor; + } + + /** + * Parks a throttle-denied request. Returns {@code false} if queueing is disabled ({@code sizePerBucket <= 0}), the + * request's own bucket is full, or the group's fixed total ceiling + * ({@link WorkloadGroupQueueSettings#MAX_GROUP_QUEUE_DEPTH}) is reached — the caller then rejects with a 429. On + * success the request holds no thread; it is admitted later by a drain, or evicted+failed on task cancellation + * (client disconnect / {@code cancel_after_time_interval}). + *

+ * There is no queue timeout: a parked request has no wall-clock deadline. Legitimate queue wait is unbounded (it + * grows with backlog depth over throughput), so a fixed cap would eventually cancel healthy, still-connected + * requests under a large slow burst. Clients bound their own wait via task cancellation instead. + * + * @param groupId the workload group id (queue identity + stat key) + * @param bucketKey the throttle bucket the request is waiting on + * @param task the search task (observed for cancellation) + * @param sizePerBucket the group's configured {@code queue.size_per_bucket} + * @param listener the parked admission listener (already context-preserving) + * @return {@code true} if enqueued, {@code false} if rejected (disabled / bucket full / group full) + */ + public boolean tryEnqueue( + String groupId, + String bucketKey, + WorkloadGroupTask task, + int sizePerBucket, + ActionListener listener + ) { + if (sizePerBucket <= 0) { + return false; // queueing disabled: not a queue rejection, the caller rejects with the normal throttle 429 + } + // Capture the enqueue instant once for the queue-wait metric. The request is never re-parked (the sweep + // re-attempts in place), so this instant is stable for the request's whole queue lifetime and waitNanos at + // admit reflects the true total time parked. + final long enqueueNanos = threadPool.relativeTimeInNanos(); + WorkloadGroupQueue.QueuedRequest req = new WorkloadGroupQueue.QueuedRequest(listener, bucketKey, task, enqueueNanos); + if (enqueue(groupId, req, sizePerBucket) == false) { + incrementQueueRejection(groupId); // bucket or group queue full + return false; + } + incrementQueued(groupId); + return true; + } + + // Parks a request (fresh or re-parked) and registers its cancellation callback. Returns false if the queue is full. + // Cancellation is registered AFTER a successful offer (it references the enqueued request) but STILL UNDER the + // bucket lock: a concurrent drain that admits this request must take the same lock, so installing the handle inside + // the lock guarantees a racing admit cannot poll the request before its cancellation handle exists (which would + // leave the callback registered on an already-admitted task forever). addOnCancelledCallback runs the callback + // immediately if the task is already cancelled, closing the race where cancellation lands during enqueue. + private boolean enqueue(String groupId, WorkloadGroupQueue.QueuedRequest req, int sizePerBucket) { + WorkloadGroupQueue queue = queuesByGroup.computeIfAbsent(groupId, k -> new WorkloadGroupQueue()); + try (Releasable ignored = bucketLocks.acquire(lockKey(groupId, req.bucketKey()))) { + if (queue.offer(req, sizePerBucket) == false) { + return false; + } + req.cancellationHandle = req.task().addOnCancelledCallback(() -> evictCancelled(groupId, req)); + return true; + } + } + + /** + * Node-tier drain: a node permit for {@code bucketKey} just freed on this coordinator. Admit the oldest waiter for + * that bucket if a node permit can be re-acquired. Admits at most one (one freed permit == one slot). The + * {@code nodeAcquire} function re-acquires a node-tier permit for the bucket (or returns {@code null} if a racing + * arrival took the freed slot). Guarded so it is only called when the queue is non-empty (see caller). + */ + public void drainNode(String groupId, String bucketKey, Function nodeAcquire) { + WorkloadGroupQueue queue = queuesByGroup.get(groupId); + if (queue == null) { + return; + } + WorkloadGroupQueue.QueuedRequest req; + Releasable permit; + try (Releasable ignored = bucketLocks.acquire(lockKey(groupId, bucketKey))) { + // Peek-then-acquire-then-remove under the lock: only remove a waiter once we hold a permit for it, so a + // failed re-acquire never drops a request from the queue. + if (hasWaiter(queue, bucketKey) == false) { + return; + } + permit = nodeAcquire.apply(bucketKey); + if (permit == null) { + return; // slot taken by a racing arrival; leave the waiter queued + } + req = queue.pollOldest(bucketKey); + if (req == null) { + // No waiter after all (shouldn't happen under the lock, but be safe): release the permit we took. + permit.close(); + return; + } + } + admit(req, permit); + } + + /** + * Cluster-tier drain ({@link WorkloadGroupSharedThrottleService.GrantConsumer}): the bucket owner pushed a grant + * carrying a reserved shared permit. Hand it to the oldest waiter for {@code bucketKey}. If there is no drainable + * waiter, returns {@code false} so the shared service returns the reserved permit to the next waiter. + * + * @return {@code true} if a waiter was admitted with the permit, {@code false} if there was none (caller releases) + */ + public boolean admitWithPermit(String bucketKey, Releasable permit) { + String groupId = groupIdOf(bucketKey); + WorkloadGroupQueue queue = queuesByGroup.get(groupId); + if (queue == null) { + return false; + } + WorkloadGroupQueue.QueuedRequest req; + try (Releasable ignored = bucketLocks.acquire(lockKey(groupId, bucketKey))) { + req = queue.pollOldest(bucketKey); + } + if (req == null) { + return false; + } + admit(req, permit); + return true; + } + + /** Group ids that currently hold a queue object. Small (groups are few); used to find backlogs needing a release. */ + public Set queuedGroupIds() { + return queuesByGroup.keySet(); + } + + /** + * Releases a group's ENTIRE backlog across all of its buckets, untracked. See + * {@link #admitAllUntracked(String, String)} for why a no-op permit is the right thing to hand out here. + * + * @return the number of requests admitted + */ + public int admitAllUntracked(String groupId) { + WorkloadGroupQueue queue = queuesByGroup.get(groupId); + if (queue == null) { + return 0; + } + int released = 0; + // Snapshot the bucket keys: admitAllUntracked mutates the bucket map (an emptied bucket is pruned). + for (String bucketKey : new ArrayList<>(queue.bucketKeys())) { + released += admitAllUntracked(groupId, bucketKey); + } + return released; + } + + /** + * Admits every request parked for {@code bucketKey} with an untracked (no-op) permit, i.e. releases the + * backlog without holding any throttle slot. Called only when the group has no throttle limit left to enforce, so + * there is nothing for these requests to wait for: continuing to hold them would be an indefinite stall, since a + * parked request has no deadline and an unthrottled group produces no permit completions to drive a drain. + *

+ * A no-op permit matches the established fail-open semantics ({@code admitWithPermit(bucketKey, () -> {})}): the + * request runs untracked and its {@code close()} does nothing. Requests are collected under the bucket lock but + * completed by {@link #admit} outside it, which also re-checks cancellation per request. + * + * @return the number of requests admitted + */ + private int admitAllUntracked(String groupId, String bucketKey) { + WorkloadGroupQueue queue = queuesByGroup.get(groupId); + if (queue == null) { + return 0; + } + final List drained = new ArrayList<>(); + try (Releasable ignored = bucketLocks.acquire(lockKey(groupId, bucketKey))) { + WorkloadGroupQueue.QueuedRequest req; + while ((req = queue.pollOldest(bucketKey)) != null) { + drained.add(req); + } + } + for (WorkloadGroupQueue.QueuedRequest req : drained) { + admit(req, () -> {}); + } + return drained.size(); + } + + // The group id is the bucketKey prefix before the first ':' (":group" or "::", see + // WorkloadGroupService.buildBucketKey). Group ids are base64 UUIDs with no ':', so the first ':' is unambiguous. + private static String groupIdOf(String bucketKey) { + int idx = bucketKey.indexOf(':'); + return idx < 0 ? bucketKey : bucketKey.substring(0, idx); + } + + /** + * Backstop sweep, run on the {@link WorkloadGroupService} scheduled loop. Two jobs, both done in place so + * a still-waiting request keeps its original identity: + *

    + *
  1. Cancelled-entry cleanup — remove and fail (with {@link TaskCancelledException}) every parked request + * whose task was cancelled. This is defense-in-depth for the per-request cancellation callback, which normally + * evicts a cancelled entry immediately; the sweep catches any it missed. There is NO time-based eviction: a + * still-live parked request is never removed regardless of how long it has waited (queue wait is unbounded by + * design; a client bounds its own wait via {@code cancel_after_time_interval} / disconnect).
    2. + *
  2. Node-tier backstop drain — for each bucket that still has a waiter, try {@code drain} (a node-tier + * local re-acquire) once, admitting the head if a node permit is free. This recovers a request that the + * node-completion chain missed. It does NOT re-contact the shared owner: the owner recovers its own lost + * grants (grant-failure reclaim) and crashed reservations (lease TTL), and re-registering a shared waiter + * every tick is exactly what caused unbounded owner-side waiter-count growth.
    3. + *
+ * + * @param drain admits the oldest waiter for a (groupId, bucketKey) against a freshly re-acquired node permit if one + * is available; a no-op if the bucket is empty or no permit is free (typically + * {@link WorkloadGroupService#sweepDrainNode}) + */ + public void sweep(SweepDrain drain) { + for (Map.Entry entry : queuesByGroup.entrySet()) { + String groupId = entry.getKey(); + WorkloadGroupQueue queue = entry.getValue(); + // Snapshot bucket keys so we don't iterate a map being mutated by concurrent drains. + for (String bucketKey : new ArrayList<>(queue.bucketKeys())) { + List evicted; + try (Releasable ignored = bucketLocks.acquire(lockKey(groupId, bucketKey))) { + evicted = queue.evictCancelled(bucketKey); + } + for (WorkloadGroupQueue.QueuedRequest req : evicted) { + failCancelled(req); + } + // Node-tier backstop: try to admit the head if a local permit is free. drainNode holds the bucket lock + // internally and admits at most one; safe to call even if the bucket is now empty. + drain.drain(groupId, bucketKey); + } + } + // Note: an emptied group queue's map entry is intentionally NOT pruned here. A concurrent tryEnqueue could + // offer to the same WorkloadGroupQueue between an "is it empty" check and its removal, which would orphan that + // freshly-parked request (present in the queue object but no longer reachable from queuesByGroup). The leak is + // one small empty object per group ever used — negligible (groups are few and long-lived) — and not worth a + // race to reclaim. A parked request whose task is cancelled (e.g. for a DELETED group, whose in-flight searches + // are cancelled) is still failed here by the cancelled-entry cleanup above. + } + + /** Total parked requests across all groups on this coordinator. Cheap; used for the node-drain fast-out. */ + public int totalDepth() { + int total = 0; + for (WorkloadGroupQueue queue : queuesByGroup.values()) { + total += queue.currentDepth(); + } + return total; + } + + /** Current parked depth for a group (0 if none). Package-private for stats. */ + public int currentDepth(String groupId) { + WorkloadGroupQueue queue = queuesByGroup.get(groupId); + return queue == null ? 0 : queue.currentDepth(); + } + + /** Peak parked depth for a group since its queue was created (0 if none). Package-private for stats. */ + public long peakDepth(String groupId) { + WorkloadGroupQueue queue = queuesByGroup.get(groupId); + return queue == null ? 0L : queue.peakDepth(); + } + + // --- internals --- + + // Completes a parked listener with an acquired permit, off the caller's thread. Deregisters the cancellation + // callback first (the request is leaving the queue). + // + // The ENTIRE body — including the cancelled-task branch that closes the permit — runs on the GENERIC executor, never + // inline on the caller's (draining/completion) thread. This is load-bearing for recursion safety: for a node-tier + // permit, permit.close() is the wrapNodePermit wrapper whose close() re-enters drainNode -> admit. Running close() + // on the caller thread would let a run of consecutively-cancelled waiters recurse close -> drainNode -> admit -> + // close ... one synchronous frame per waiter (StackOverflow under a cancel storm). Dispatching first means each such + // hop is a fresh executor task, so the chain unwinds across tasks rather than down one stack. + private void admit(WorkloadGroupQueue.QueuedRequest req, Releasable permit) { + req.releaseCancellationHandle(); + // Record the queue wait at the admission instant (this thread), before the executor hop, so the metric is the + // true time parked and not inflated by dispatch latency. Only requests that actually parked reach admit(), so + // this counts wait among queued requests. Recorded even for a task cancelled-while-queued: it still waited. + recordQueueWait(groupIdOf(req.bucketKey()), TimeUnit.NANOSECONDS.toMillis(req.waitNanos(threadPool.relativeTimeInNanos()))); + threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> { + if (req.task().isCancelled()) { + permit.close(); + req.listener().onFailure(new TaskCancelledException("task cancelled while queued")); + return; + } + req.listener().onResponse(permit); + }); + } + + // Cancellation callback: remove the entry (if still queued) and fail it. Runs on whatever thread cancelled the task + // — including, for a task already cancelled at enqueue time, inline under the enqueue bucket lock (KeyedLock is + // reentrant, so re-acquiring below does not deadlock). The failure is dispatched on the GENERIC executor rather + // than completed inline so the listener is never completed while a bucket lock is held (consistent with admit()). + private void evictCancelled(String groupId, WorkloadGroupQueue.QueuedRequest req) { + WorkloadGroupQueue queue = queuesByGroup.get(groupId); + if (queue == null) { + return; + } + final boolean removed; + try (Releasable ignored = bucketLocks.acquire(lockKey(groupId, req.bucketKey()))) { + removed = queue.remove(req); + } + if (removed) { + threadPool.executor(ThreadPool.Names.GENERIC) + .execute(() -> req.listener().onFailure(new TaskCancelledException("task cancelled while queued"))); + } + } + + // Fail a cancelled request the sweep evicted. The entry is already removed from the queue. + private void failCancelled(WorkloadGroupQueue.QueuedRequest req) { + req.releaseCancellationHandle(); + req.listener().onFailure(new TaskCancelledException("task cancelled while queued")); + } + + private static boolean hasWaiter(WorkloadGroupQueue queue, String bucketKey) { + return queue.bucketKeys().contains(bucketKey); + } + + private static String lockKey(String groupId, String bucketKey) { + return groupId + '\0' + bucketKey; + } + + private void incrementQueued(String groupId) { + WorkloadGroupState state = stateAccessor.getWorkloadGroupStateMap().get(groupId); + if (state != null) { + state.totalQueued.inc(); + } + } + + private void recordQueueWait(String groupId, long waitMillis) { + WorkloadGroupState state = stateAccessor.getWorkloadGroupStateMap().get(groupId); + if (state != null) { + state.recordQueueWaitMillis(waitMillis); + } + } + + private void incrementQueueRejection(String groupId) { + WorkloadGroupState state = stateAccessor.getWorkloadGroupStateMap().get(groupId); + if (state != null) { + state.totalQueueRejections.inc(); + } + } + + /** + * Node-tier backstop drain for the sweep: admit the oldest waiter for {@code (groupId, bucketKey)} against a + * freshly re-acquired node permit, if one is free; otherwise a no-op. Implemented by {@link WorkloadGroupService} + * (which owns the node permit tracker) so the queue service holds no throttle logic. + */ + @ExperimentalApi + @FunctionalInterface + public interface SweepDrain { + void drain(String groupId, String bucketKey); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupQueueSettings.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupQueueSettings.java new file mode 100644 index 0000000000000..d553b5f302d76 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupQueueSettings.java @@ -0,0 +1,131 @@ +/* + * 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 org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; + +import java.util.Map; + +/** + * Registry of valid workload group request-queue settings with their validators. Queue config is a nested + * {@code queue} object (a {@link Settings} bag) alongside {@code throttling}, so per-key null clears a field and an + * absent key keeps the existing value with no extra bookkeeping. + *

+ * When a request is denied by a throttle limit, instead of an immediate 429 the coordinator may park it in a bounded + * queue and admit it once a permit frees. The single user-facing knob is {@code size_per_bucket}: the maximum parked + * requests per throttle bucket (per coordinator), mirroring how the throttle limits ({@code node_limit}, + * {@code shared_limit}) are themselves per-bucket. {@code 0} disables queueing (immediate reject preserved). Keying the + * cap per bucket gives fairness for {@code attribute=username}/{@code role}: one principal's flood cannot consume + * another principal's per-bucket allowance. That fairness is bounded rather than absolute — see the group ceiling below, + * at which admission reverts to first-come-first-served across buckets. For {@code attribute=group} there is a single + * bucket, so it is simply the group's queue depth. + *

+ * Above the per-bucket cap sits a fixed, non-configurable per-group ceiling ({@link #MAX_GROUP_QUEUE_DEPTH}) on the + * total parked requests across all of a group's buckets on one coordinator. It is a safety backstop, not a + * fairness knob: because bucket keys for {@code username}/{@code role} are attacker-controlled (derived from the request + * principal), a purely per-bucket cap would let unbounded distinct principals each allocate {@code size_per_bucket} + * slots, so the group ceiling bounds the coordinator's parked footprint (heap + open connections) regardless of bucket + * cardinality. {@link #MAX_SIZE_PER_BUCKET} is pinned to the same value, so validation never accepts a per-bucket depth + * the ceiling could not honour: a single-bucket group ({@code attribute=group}) may queue the whole group budget in its + * one bucket, while a many-bucket group reaches the ceiling first — the intended shed point, where a 429 is the correct + * response. See {@code WorkloadGroupQueue} for enforcement (a request is admitted only if both its bucket is under + * {@code size_per_bucket} AND the group total is under this ceiling). + *

+ * There is deliberately no user-facing queue timeout, and no timeout of any kind. Legitimate queue wait is + * unbounded — it grows with backlog depth over drain throughput — so any fixed wall-clock cap would eventually cancel + * healthy, still-connected requests under a large slow burst. A client bounds its own wait with + * {@code cancel_after_time_interval} (per request, or the {@code search.cancel_after_time_interval} cluster setting): + * its cancellation timer is armed before throttle admission, so it fires while the request is parked and evicts the + * entry promptly. A parked request is therefore bounded only by task cancellation (client disconnect or + * {@code cancel_after_time_interval}); the queue itself never expires an entry by time — the backstop sweep only + * removes entries whose task is already cancelled. + */ +@ExperimentalApi +public class WorkloadGroupQueueSettings { + + /** Default per-bucket queue depth: {@code 0} disables queueing (over-limit requests are rejected immediately). */ + public static final int DEFAULT_SIZE_PER_BUCKET = 0; + + /** + * Fixed, non-configurable ceiling on the TOTAL parked requests for one group across all its buckets on a single + * coordinator. A pure OOM/footprint backstop against attacker-controlled bucket cardinality (username/role buckets + * come from the request principal), NOT a latency or fairness knob. Grounded in the retained heap of a parked + * request (its held {@code SearchRequest} + open channel + task ~ tens of KB): 10,000 parked ~= a few hundred MB of + * pinned request memory, a high-but-acceptable backstop on a small heap and negligible on a large one. + */ + public static final int MAX_GROUP_QUEUE_DEPTH = 10_000; + + /** + * Maximum configurable per-bucket queue depth. Pinned to {@link #MAX_GROUP_QUEUE_DEPTH}: the group ceiling caps the + * total across all of a group's buckets, so a larger per-bucket value could never be honoured and accepting one + * would silently mislead. A single-bucket group ({@code attribute=group}) can therefore still queue the entire group + * budget in its one bucket. + */ + public static final int MAX_SIZE_PER_BUCKET = MAX_GROUP_QUEUE_DEPTH; + + /** Per-coordinator bounded queue depth per throttle bucket. {@code 0} = queueing disabled. */ + public static final Setting SIZE_PER_BUCKET = Setting.intSetting( + "size_per_bucket", + DEFAULT_SIZE_PER_BUCKET, + DEFAULT_SIZE_PER_BUCKET + ); + + private static final Map> REGISTERED_SETTINGS = Map.of(SIZE_PER_BUCKET.getKey(), SIZE_PER_BUCKET); + + private WorkloadGroupQueueSettings() { + throw new UnsupportedOperationException("Utility class"); + } + + /** + * Per-key validation: every key must be registered and {@code size_per_bucket} must be a non-negative integer no + * greater than {@link #MAX_SIZE_PER_BUCKET}. Safe to run on a partial fragment from an update request. + * {@code size_per_bucket} is currently the only queue key, so there are no cross-field rules and + * therefore no merged-config validator (unlike {@link WorkloadGroupThrottleSettings}, which has real cross-field rules). + * + * @param queue the queue settings to validate + * @throws IllegalArgumentException if any key is unknown or {@code size_per_bucket} is invalid + */ + public static void validate(Settings queue) { + if (queue == null) { + return; + } + for (String key : queue.keySet()) { + String value = queue.get(key); + if (REGISTERED_SETTINGS.containsKey(key) == false) { + throw new IllegalArgumentException("Unknown queue setting: " + key); + } + // null value means "clear this key" — skip value validation + if (value == null) { + continue; + } + if (SIZE_PER_BUCKET.getKey().equals(key)) { + validateSizePerBucket(value); + } + } + } + + // Rejects a size_per_bucket that is not a non-negative integer in [0, MAX_SIZE_PER_BUCKET]. + private static void validateSizePerBucket(String value) { + final int parsed; + try { + parsed = Integer.parseInt(value); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("queue.size_per_bucket must be an integer but was '" + value + "'"); + } + if (parsed < 0) { + throw new IllegalArgumentException("queue.size_per_bucket must be non-negative but was " + parsed); + } + if (parsed > MAX_SIZE_PER_BUCKET) { + throw new IllegalArgumentException("queue.size_per_bucket must not exceed " + MAX_SIZE_PER_BUCKET + " but was " + parsed); + } + } + +} diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java index 173b32706d603..9f853d2109890 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java @@ -66,6 +66,8 @@ public class WorkloadGroupService extends AbstractLifecycleComponent 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; + // Coordinator-local request queues; late-bound. Null => no queueing (throttle denial rejects immediately, as before). + private volatile WorkloadGroupQueueService queueService; public WorkloadGroupService( WorkloadGroupTaskCancellationService taskCancellationService, @@ -135,6 +137,13 @@ void doRun() { } taskCancellationService.cancelTasks(nodeDuressTrackers::isNodeInDuress, activeWorkloadGroups, deletedWorkloadGroups); taskCancellationService.pruneDeletedWorkloadGroups(deletedWorkloadGroups); + // Backstop sweep: reap cancelled queued requests (defense-in-depth for the per-request cancel callback) and, as + // a node-tier backstop, admit a waiter if a local permit is free. Owner-push and node-completion are the primary + // drains; task cancellation (client disconnect / cancel_after_time_interval) is the primary parked-request exit. + final WorkloadGroupQueueService qs = queueService; + if (qs != null) { + qs.sweep(this::sweepDrainNode); + } } /** @@ -192,6 +201,59 @@ public void clusterChanged(ClusterChangedEvent event) { } } this.activeWorkloadGroups = new HashSet<>(currentMetadata.workloadGroups().values()); + releaseBacklogForUnthrottledGroups(currentWorkloadGroups); + } + + /** + * Immediately releases the parked backlog of any group that no longer has a throttle limit (throttling was disabled, + * or the group was deleted). Nothing remains for those requests to wait for, and an unthrottled group takes the + * no-permit fast path — so it generates no permit completions to drive a drain chain, and a parked request has no + * deadline. Reacting to the config change is therefore the only thing that frees these requests: there is no + * periodic backstop for this case, so if this listener does not reach a node holding a backlog (e.g. its queue service + * is not wired yet), those requests wait for client cancellation. + *

+ * Note disabling throttling necessarily disables queueing in the same update ({@code WorkloadGroup} rejects a queue + * with no throttle limit, and {@code validateMergedConfig} rejects a throttling block whose limits are all unset), so + * this is the only reachable way to end up with a backlog and nothing to drain it. + *

+ * Runs on the cluster-applier thread, so the actual release is dispatched to {@code GENERIC}: emptying a deep queue + * polls under each bucket's lock and completes one listener per request, which must not sit on the applier thread. + */ + private void releaseBacklogForUnthrottledGroups(Map currentWorkloadGroups) { + final WorkloadGroupQueueService qs = queueService; + if (qs == null) { + return; + } + for (String groupId : qs.queuedGroupIds()) { + if (qs.currentDepth(groupId) == 0) { + continue; // no backlog to release + } + WorkloadGroup workloadGroup = currentWorkloadGroups.get(groupId); + if (workloadGroup != null) { + Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling(); + if (WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling) != WorkloadGroupThrottleSettings.UNSET_LIMIT + || WorkloadGroupThrottleSettings.SHARED_LIMIT.get(throttling) != WorkloadGroupThrottleSettings.UNSET_LIMIT) { + continue; // still throttled: the normal drain paths own this backlog + } + } + // Group is gone, or has no limit left. (A deleted group's tasks are also being cancelled; admit() re-checks + // cancellation per request, so releasing here is safe either way and self-heals a missed cancellation.) + threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> { + try { + int released = qs.admitAllUntracked(groupId); + if (released > 0) { + logger.info( + "Released {} queued request(s) for workload group [{}]: throttling is no longer configured, " + + "so there is nothing left to wait for.", + released, + groupId + ); + } + } catch (Exception e) { + logger.warn("Failed to release the queued backlog for workload group [" + groupId + "]", e); + } + }); + } } /** @@ -222,12 +284,15 @@ public WorkloadGroupStats nodeStats(Set workloadGroupIds, Boolean reques } } } + final WorkloadGroupQueueService qs = queueService; if (existingStateMap != null) { existingStateMap.forEach((workloadGroupId, currentState) -> { boolean shouldInclude = workloadGroupIds.contains("_all") || workloadGroupIds.contains(workloadGroupId); if (shouldInclude) { if (requestedBreached == null || requestedBreached == resourceLimitBreached(workloadGroupId, currentState)) { - statsHolderMap.put(workloadGroupId, WorkloadGroupStatsHolder.from(currentState)); + long queuedCurrent = qs == null ? 0L : qs.currentDepth(workloadGroupId); + long queuePeak = qs == null ? 0L : qs.peakDepth(workloadGroupId); + statsHolderMap.put(workloadGroupId, WorkloadGroupStatsHolder.from(currentState, queuedCurrent, queuePeak)); } } }); @@ -328,6 +393,15 @@ public void setSharedThrottleService(WorkloadGroupSharedThrottleService sharedTh this.sharedThrottleService = sharedThrottleService; } + /** + * Late-binds the coordinator-local request-queue service. When unset, a throttle denial rejects immediately (the + * pre-queueing behavior); when set, a denial may park the request instead (if the group's + * {@code queue.size_per_bucket} > 0). + */ + public void setQueueService(WorkloadGroupQueueService queueService) { + this.queueService = queueService; + } + /** * Two-tier throttle admission for a search request. Notifies {@code listener} with: *

    @@ -342,11 +416,12 @@ public void setSharedThrottleService(WorkloadGroupSharedThrottleService sharedTh * 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 task the search task (carries the workload group id; observed for cancellation while queued) * @param principal the raw {@code WORKLOAD_GROUP_PRINCIPAL_HEADER} value, or {@code null} * @param listener receives the permit / null / 429 */ - public void acquireThrottlePermit(String workloadGroupId, String principal, ActionListener listener) { + public void acquireThrottlePermit(WorkloadGroupTask task, String principal, ActionListener listener) { + final String workloadGroupId = task.getWorkloadGroupId(); final ThrottlePlan plan; try { plan = resolveThrottlePlan(workloadGroupId, principal); @@ -363,7 +438,7 @@ public void acquireThrottlePermit(String workloadGroupId, String principal, Acti // 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); + Releasable localPermit = acquireNodePermit(plan); if (localPermit != null) { listener.onResponse(localPermit); return; @@ -373,11 +448,63 @@ public void acquireThrottlePermit(String workloadGroupId, String principal, Acti // 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 -> { + final WorkloadGroupQueueService qs = queueService; + // Queueing is active only when the group has a non-zero queue AND is not in MONITOR mode (monitor observes + // and always admits — it must never park). When active, take the ENQUEUE-FIRST path; otherwise the request + // runs directly on grant and is admitted (monitor) or 429'd (no queue) on denial. + final boolean queueingActive = qs != null && plan.queueSizePerBucket > 0 && plan.monitorMode == false; + + if (queueingActive) { + // ENQUEUE-FIRST: park the request BEFORE contacting the owner. The owner registers this coordinator as a + // waiter synchronously when it denies an acquire (in handleAcquire, before the reply is even sent), so if + // we enqueued only after the denial reply there would be a window in which the owner believes we are + // waiting while our queue is still empty — a racing grant would then find nothing to admit, return the + // slot "unused", and deregister us, stranding the request (worst on a shared-only group with no node-tier + // drain). Parking first closes that window by construction: whenever a grant or owner-push arrives, the + // request is already in the queue. The acquire below is now a pure "supply" signal — a granted slot + // drains the OLDEST queued request (FIFO), which may differ from this one; that is fine, since supply is + // matched to demand by count and shared_limit is still gated solely by the owner's tryAcquire. + if (qs.tryEnqueue(plan.workloadGroupId, plan.bucketKey, task, plan.queueSizePerBucket, listener) == false) { + // Queue full (bucket or group ceiling): reject with the throttle 429 (tryEnqueue already counted the + // queue rejection). + incrementThrottled(plan.workloadGroupId); + listener.onFailure(new OpenSearchRejectedExecutionException("Request throttled: " + plan.describeBreach(false) + ".")); + return; + } + // Request is parked (its listener is held by the queue). Ask the owner for a shared slot; wantsQueue=true + // so a denial registers this coordinator for owner-push. This callback NEVER completes the request's + // listener directly — it only supplies a permit to the queue. + sharedThrottleService.acquireAsync(plan.bucketKey, plan.sharedLimit, true, ActionListener.wrap(permit -> { + if (permit != null) { + // Granted a shared slot: hand it to the oldest queued request. If a concurrent drain (owner-push + // or node-tier completion) already emptied the bucket, release the slot rather than hold it. + if (qs.admitWithPermit(plan.bucketKey, permit) == false) { + permit.close(); + } + } else { + // Fail-open (owner unreachable / empty ring): admit one queued request with no shared permit + // (untracked), matching the shipped fail-open semantics. No-op if the bucket already drained. + qs.admitWithPermit(plan.bucketKey, () -> {}); + } + }, e -> { + // acquireAsync only ever fails with the message-less denial marker (transport errors fail open via + // onResponse(null)). Denied: the request stays parked and the owner has registered this coordinator + // (wantsQueue), so owner-push drains it when a slot frees. Nothing to complete here. + if (e instanceof OpenSearchRejectedExecutionException == false) { + logger.warn( + "Unexpected shared-acquire failure for a queued request in workload group [" + workloadGroupId + "]", + e + ); + } + })); + return; + } + + // Non-queueing shared path (queue disabled or MONITOR mode): unchanged — acquire with the request's own + // listener; a denial admits (monitor) or 429s via onThrottleBreach. + sharedThrottleService.acquireAsync(plan.bucketKey, plan.sharedLimit, false, ActionListener.wrap(listener::onResponse, e -> { if (e instanceof OpenSearchRejectedExecutionException) { - // At the shared limit. In MONITOR mode observe only (log, admit, no stat); otherwise reject with the - // recomposed message (the shared tier only has the opaque bucket key) and count it. - onThrottleBreach(plan, false, listener); + onThrottleBreach(plan, false, task, listener); } else { listener.onFailure(e); } @@ -387,7 +514,7 @@ public void acquireThrottlePermit(String workloadGroupId, String principal, Acti 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. - onThrottleBreach(plan, true, listener); + onThrottleBreach(plan, true, task, listener); } 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. @@ -395,10 +522,65 @@ public void acquireThrottlePermit(String workloadGroupId, String principal, Acti } } - // Terminal handling when a request would be throttled at a tier's limit. In MONITOR mode the group only observes: - // log that the request WOULD have been rejected, then admit it (onResponse(null)) without touching total_throttled. - // In any other mode, count the rejection and fail with the user-facing 429. - private void onThrottleBreach(ThrottlePlan plan, boolean nodeTier, ActionListener listener) { + // Acquires a node-local permit for admission, wrapped so its close() drains the bucket's queue. + private Releasable acquireNodePermit(ThrottlePlan plan) { + return wrapNodePermit(throttleTracker.tryAcquire(plan.bucketKey, plan.nodeLimit), plan.workloadGroupId, plan.bucketKey); + } + + /** + * The group's current {@code node_limit} from cluster state, or {@link WorkloadGroupThrottleSettings#UNSET_LIMIT} + * if the group is gone or the node tier is not configured. Read fresh (not captured) everywhere a drain re-acquires, + * so a live {@code node_limit} update takes effect on the very next drain. + */ + private int currentNodeLimit(String groupId) { + WorkloadGroup workloadGroup = getWorkloadGroupById(groupId); + if (workloadGroup == null) { + return WorkloadGroupThrottleSettings.UNSET_LIMIT; + } + return WorkloadGroupThrottleSettings.NODE_LIMIT.get(workloadGroup.getMutableWorkloadGroupFragment().getThrottling()); + } + + // Wraps a raw node permit so its close() releases the slot AND drains one waiter for the bucket — a freed node + // permit creates room for one waiting request on this coordinator. Returns null if the raw permit is null (limit + // reached). The wrapping is applied to BOTH the initial admission permit and the permit handed to a drained + // waiter, so the node-completion drain chains continuously instead of dying after one hop (each hop is a separate, + // asynchronous request completion, so there is no synchronous recursion). Guarded by a cheap per-group depth check + // so the common unthrottled path pays only a single map lookup past the shipped decrement. + // + // node_limit is re-read from cluster state on each drain rather than captured when the chain started: a busy bucket's + // drain chain can run for a long time (one hop per request completion), so a captured value would keep admitting + // against a stale limit across a live node_limit update — over-admitting above a lowered limit until the chain broke. + // The re-read sits INSIDE the depth guard, so it costs nothing unless this group actually has a backlog to drain. + private Releasable wrapNodePermit(Releasable raw, String groupId, String bucketKey) { + if (raw == null) { + return null; + } + final WorkloadGroupQueueService qs = queueService; + if (qs == null) { + return raw; + } + return () -> { + raw.close(); + // Fast-out scoped to THIS group: drainNode only ever admits a waiter for (groupId, bucketKey), so guard on + // this group's depth, not the cluster-wide total. A cheap single-map lookup, and it avoids a spurious + // lock-acquire + permit re-acquire/release when some *other* group is the one that is backlogged. + if (qs.currentDepth(groupId) > 0) { + final int liveNodeLimit = currentNodeLimit(groupId); + if (liveNodeLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) { + // Node tier no longer configured (or group deleted): this tier must not admit — a still-configured + // shared_limit would be bypassed. If NO limit remains, the backlog has nothing to wait for, but it is + // not stranded: clusterChanged releases it as soon as throttling is disabled. + return; + } + qs.drainNode(groupId, bucketKey, key -> wrapNodePermit(throttleTracker.tryAcquire(key, liveNodeLimit), groupId, key)); + } + }; + } + + // Terminal handling when a request would be throttled at a tier's limit. Order: MONITOR observes only (log + admit, + // no stat); else try to park the request in the queue (if queueing is enabled and has room); else count the + // rejection and fail with the user-facing 429. + private void onThrottleBreach(ThrottlePlan plan, boolean nodeTier, WorkloadGroupTask task, ActionListener listener) { if (plan.monitorMode) { // DEBUG, not INFO: this fires once per would-be-throttled request, so INFO would spam a hot bucket under // load. The message names the throttle attribute value (username/role), but that is the caller's own @@ -407,10 +589,43 @@ private void onThrottleBreach(ThrottlePlan plan, boolean nodeTier, ActionListene listener.onResponse(null); return; } + // Queue-then-reject: hold the request instead of rejecting, if queueing is enabled and both the request's bucket + // and the group have room. A parked request holds no thread — only its listener + open connection — and is + // admitted later by a node-completion drain or an owner grant, or evicted by task cancellation (client + // disconnect / cancel_after_time_interval). There is no queue timeout: a parked request has no wall-clock + // deadline. + final WorkloadGroupQueueService qs = queueService; + if (qs != null + && plan.queueSizePerBucket > 0 + && qs.tryEnqueue(plan.workloadGroupId, plan.bucketKey, task, plan.queueSizePerBucket, listener)) { + return; // parked; listener completed later + } incrementThrottled(plan.workloadGroupId); listener.onFailure(new OpenSearchRejectedExecutionException("Request throttled: " + plan.describeBreach(nodeTier) + ".")); } + /** + * Node-tier backstop drain for the sweep: admit the oldest waiter for {@code bucketKey} against a freshly acquired + * node permit if one is free; a no-op otherwise. Wired into {@link WorkloadGroupQueueService#sweep}. This recovers + * a request the node-completion chain missed without re-running full admission or re-contacting the shared owner + * (the owner recovers its own lost grants and reservations), and — crucially — without dequeuing-and-re-parking, so + * a still-waiting request stays parked in place (there is no queue timeout; its only deadline is task cancellation + * via {@code cancel_after_time_interval} / client disconnect). + */ + private void sweepDrainNode(String groupId, String bucketKey) { + // node_limit is not carried per-bucket here; the drain lambda re-derives the permit for the bucket. We only + // engage the node tier as a backstop — the shared tier is drained by owner-push. Read the group's CURRENT + // node_limit (same live read the completion drain uses); if the node tier isn't configured, nothing to drain. + final int nodeLimit = currentNodeLimit(groupId); + if (nodeLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) { + // Group deleted, or shared-only config: the node tier can't admit; owner-push handles the shared tier. A + // group with NO limit left is not this path's problem either — clusterChanged releases that backlog the + // moment throttling is disabled (see releaseBacklogForUnthrottledGroups). + return; + } + queueService.drainNode(groupId, bucketKey, key -> wrapNodePermit(throttleTracker.tryAcquire(key, nodeLimit), groupId, key)); + } + // 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) { @@ -455,7 +670,23 @@ private ThrottlePlan resolveThrottlePlan(String workloadGroupId, String principa // MONITOR mode observes only: the limit is still evaluated so a breach can be logged, but the request is never // rejected and no stat is updated (consistent with how MONITOR is dormant on the resource cancellation path). boolean monitorMode = workloadGroup.getResiliencyMode() == MutableWorkloadGroupFragment.ResiliencyMode.MONITOR; - return new ThrottlePlan(workloadGroupId, bucketKey, nodeLimit, sharedLimit, workloadGroup.getName(), attribute, value, monitorMode); + // Queue config (used only if a throttle limit is breached). Absent => 0 => no queueing (reject immediately). + // queue.size_per_bucket is the only queue knob (the per-group total is the fixed MAX_GROUP_QUEUE_DEPTH ceiling); + // there is no queue timeout — a client's own deadline comes from cancel_after_time_interval, and a parked + // request has no wall-clock deadline. + Settings queue = workloadGroup.getMutableWorkloadGroupFragment().getQueue(); + int queueSizePerBucket = WorkloadGroupQueueSettings.SIZE_PER_BUCKET.get(queue); + return new ThrottlePlan( + workloadGroupId, + bucketKey, + nodeLimit, + sharedLimit, + workloadGroup.getName(), + attribute, + value, + monitorMode, + queueSizePerBucket + ); } // The resolved throttle configuration for a single request. Carries the human-readable group name and the throttle @@ -470,6 +701,7 @@ private static class ThrottlePlan { final String attribute; // "group" | "username" | "role" final String value; // the resolved principal value for username/role; null for whole-group throttling final boolean monitorMode; // group is in MONITOR resiliency mode -> observe (log), never reject or count + final int queueSizePerBucket; // queue.size_per_bucket for this group (0 => queueing disabled) ThrottlePlan( String workloadGroupId, @@ -479,7 +711,8 @@ private static class ThrottlePlan { String groupName, String attribute, String value, - boolean monitorMode + boolean monitorMode, + int queueSizePerBucket ) { this.workloadGroupId = workloadGroupId; this.bucketKey = bucketKey; @@ -489,6 +722,7 @@ private static class ThrottlePlan { this.attribute = attribute; this.value = value; this.monitorMode = monitorMode; + this.queueSizePerBucket = queueSizePerBucket; } // "workload group [analytics]" or "workload group [analytics] for username [alice]". diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java index ecaba972446d0..928055f2233cb 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java @@ -12,6 +12,7 @@ import org.apache.logging.log4j.Logger; import org.opensearch.cluster.ClusterChangedEvent; import org.opensearch.cluster.ClusterStateListener; +import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.cluster.node.DiscoveryNode; import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; @@ -33,8 +34,10 @@ import org.opensearch.transport.TransportService; import java.io.IOException; +import java.util.LinkedHashSet; import java.util.Map; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; @@ -59,6 +62,9 @@ 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"; + // Owner -> coordinator: a shared slot freed and this coordinator has a registered waiter for the bucket; here is a + // reserved permit to admit one queued request against (queueing / owner-push tier). + public static final String GRANT_ACTION_NAME = "internal:wlm/throttle/shared/grant"; // 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 @@ -75,11 +81,14 @@ public class WorkloadGroupSharedThrottleService implements ClusterStateListener // is uncommon (no default search timeout, but heavy aggs/scripts can exceed it); if it becomes a problem the fix is // permit renewal or deriving the TTL from a request deadline, not a larger constant. static final long PERMIT_TTL_NANOS = TimeValue.timeValueMinutes(5).nanos(); - // How often the owner sweeps expired permits. This is a pure memory-hygiene backstop, NOT a correctness mechanism: - // tryAcquire() already prunes a bucket's expired permits on every acquire, so an active bucket self-heals and can - // never over-reject due to a stuck permit. The sweep only reclaims the map entry for a bucket that abandoned a permit - // 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. + // How often the owner sweeps expired permits. Originally pure memory hygiene, on the reasoning that tryAcquire() + // prunes a bucket's expired permits on every acquire, so an active bucket self-heals and can never over-reject. + // QUEUEING MADE THIS A CORRECTNESS PATH TOO: a bucket can now hold parked requests while receiving no new acquires + // (the clients already in the queue are waiting, not arriving), so there may be no tryAcquire to do that pruning. For + // such a bucket this sweep is the only thing that reclaims a crashed holder's permit AND the only thing that then + // drives owner-push for the freed slot (see start()). Consequence to be aware of: a permit frees at its TTL but is + // only discovered on the next sweep, so worst-case discovery latency is TTL + SWEEP_INTERVAL. Parked requests have + // no deadline, so they wait rather than fail — but they wait that long. static final TimeValue SWEEP_INTERVAL = TimeValue.timeValueMinutes(5); private static final Logger logger = LogManager.getLogger(WorkloadGroupSharedThrottleService.class); @@ -93,6 +102,25 @@ public class WorkloadGroupSharedThrottleService implements ClusterStateListener private final AtomicReference ring = new AtomicReference<>(); private volatile Scheduler.Cancellable sweepTask; + // OWNER-SIDE waiter registry for owner-push queue draining: for each bucket this node owns, the ordered SET of + // coordinators that have at least one request parked for the bucket. A set (not a count) so registration is + // idempotent and the registry is bounded at <= N coordinators per bucket — a coordinator that re-registers (e.g. + // still parked across several acquires) does not inflate it. INSERTION-ORDERED (LinkedHashSet) so a grant rotates + // the chosen coordinator to the tail (see pickAndRotate): membership persists across a successful hand-off (the + // coordinator may still have more queued requests), and successive freed slots round-robin fairly across + // coordinators instead of repeatedly serving whichever one hashes first. A coordinator is removed only when a grant + // to it comes back unused (no more queued requests) or it disconnects, so the registry self-reconciles with real + // demand. No request identity is held (the requests live on their coordinators). LinkedHashSet is NOT thread-safe, + // so every access — read, size, add, remove, rotate — MUST go through compute/computeIfPresent on this map, whose + // per-key exclusive remapping is the sole lock guarding the inner set. + private final Map> waitersByBucket = new ConcurrentHashMap<>(); + + // COORDINATOR-SIDE: how a pushed grant is turned into an admitted queued request. Late-bound (the queue service is + // constructed after this service); returns true if a queued request was admitted with the reserved permit, false + // if there was none (this service then returns the reserved slot to the owner). Null before wiring => no grant is + // ever consumed (a received grant is returned), which is safe. + private volatile GrantConsumer grantConsumer; + public WorkloadGroupSharedThrottleService(ClusterService clusterService, ThreadPool threadPool, TransportService transportService) { this(clusterService, threadPool, transportService, new SharedThrottleTracker()); } @@ -126,6 +154,23 @@ public WorkloadGroupSharedThrottleService( ReleasePermitRequest::new, (request, channel, task) -> { tracker.release(request.bucketKey, request.permitId); + // If this release is a coordinator returning an UNUSED grant (it had no queued request for the bucket), + // drop it from the waiter set so it stops drawing wasted grants (remote analog of the local removeWaiter). + if (request.queueEmptyOnNodeId.isEmpty() == false) { + removeWaiterByNodeId(request.bucketKey, request.queueEmptyOnNodeId); + } + // A slot just freed on the owner. If a coordinator is waiting on this bucket, hand the freed slot to one + // of them (reserve-then-grant) so its queued request can drain — owner-push. No-op if no waiters. + onSharedSlotFreed(request.bucketKey, request.sharedLimit); + channel.sendResponse(TransportResponse.Empty.INSTANCE); + } + ); + transportService.registerRequestHandler( + GRANT_ACTION_NAME, + ThreadPool.Names.SAME, + GrantPermitRequest::new, + (request, channel, task) -> { + handleGrant(request); channel.sendResponse(TransportResponse.Empty.INSTANCE); } ); @@ -138,7 +183,7 @@ public void start() { // set yet". This method only schedules the memory-hygiene sweep. sweepTask = threadPool.scheduleWithFixedDelay(() -> { try { - tracker.sweepExpired(); + sweepExpiredAndDrive(); } catch (Exception e) { logger.warn("Shared throttle TTL sweep failed", e); } @@ -185,6 +230,15 @@ public void clusterChanged(ClusterChangedEvent event) { * The listener may be invoked inline (local-owner short-circuit) or on a transport thread. */ public void acquireAsync(String bucketKey, int sharedLimit, ActionListener listener) { + acquireAsync(bucketKey, sharedLimit, false, listener); + } + + /** + * As {@link #acquireAsync(String, int, ActionListener)}, but {@code wantsQueue} tells the owner to register this + * coordinator as a waiter for the bucket if the acquire is denied, so a later freed slot is pushed back here as a + * grant (owner-push queue draining). Pass {@code true} only when this group has queueing enabled. + */ + public void acquireAsync(String bucketKey, int sharedLimit, boolean wantsQueue, ActionListener listener) { final ThrottleOwnerSelector currentRing = ring.get(); final DiscoveryNode owner = currentRing.ownerFor(bucketKey).orElse(null); if (owner == null) { @@ -198,8 +252,11 @@ public void acquireAsync(String bucketKey, int sharedLimit, ActionListener() { @Override @@ -229,7 +286,7 @@ public AcquirePermitResponse read(StreamInput in) throws IOException { @Override public void handleResponse(AcquirePermitResponse response) { if (response.granted) { - listener.onResponse(releaseRemote(owner, bucketKey, permitId)); + listener.onResponse(releaseRemote(owner, bucketKey, sharedLimit, permitId)); } else { listener.onFailure(deniedMarker()); } @@ -243,7 +300,9 @@ public void handleException(TransportException exp) { // best-effort release for this permitId to reclaim it immediately; release-by-id is idempotent, so // if no permit 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, permitId); + // Reclaim only (no owner-push): this permit likely never existed; UNSET_LIMIT tells the owner to + // release without driving a grant. + sendRelease(owner, bucketKey, WorkloadGroupThrottleSettings.UNSET_LIMIT, permitId, ""); listener.onResponse(null); } @@ -258,26 +317,57 @@ public String executor() { // Owner-side admission. Package-private for tests. AcquirePermitResponse handleAcquire(AcquirePermitRequest request) { boolean granted = tracker.tryAcquire(request.bucketKey, request.sharedLimit, request.permitId, request.ttlNanos); + if (granted == false && request.wantsQueue && request.requestingNodeId.isEmpty() == false) { + // Denied and the coordinator will park the request -> remember it so a freed slot is pushed back as a grant. + // The RPC carries only the coordinator's persistent node id, so resolve it to the live DiscoveryNode here: + // the waiter set must hold the real node because sendGrant needs its transport address. A node that has since + // left the cluster resolves to null and is simply not registered, which is correct — there is nothing to push + // a grant to. + final DiscoveryNode requestingNode = clusterService.state().nodes().get(request.requestingNodeId); + if (requestingNode != null) { + registerWaiter(request.bucketKey, requestingNode); + } else { + logger.debug( + "Not registering waiter for bucket [{}]: node [{}] is no longer in the cluster", + request.bucketKey, + request.requestingNodeId + ); + } + } return new AcquirePermitResponse(granted); } - private Releasable releaseLocal(String bucketKey, String permitId) { - return releaseOnce(() -> tracker.release(bucketKey, permitId)); + // A local release path must also drive owner-push when this node owns the bucket: releasing frees a slot that a + // registered waiter should get. sharedLimit is threaded so the reserve step can re-check the limit. + private Releasable releaseLocal(String bucketKey, int sharedLimit, String permitId) { + return releaseOnce(() -> { + tracker.release(bucketKey, permitId); + onSharedSlotFreed(bucketKey, sharedLimit); + }); } - private Releasable releaseRemote(DiscoveryNode owner, String bucketKey, String permitId) { - return releaseOnce(() -> sendRelease(owner, bucketKey, permitId)); + private Releasable releaseRemote(DiscoveryNode owner, String bucketKey, int sharedLimit, String permitId) { + // The remote RELEASE RPC carries sharedLimit so the owner can drive owner-push after freeing the slot. + return releaseOnce(() -> sendRelease(owner, bucketKey, sharedLimit, permitId, "")); + } + + // Returns an UNUSED remote grant: a RELEASE tagged with this coordinator's node id so the owner also deregisters it + // from the bucket's waiter set before re-driving owner-push to the next waiter. + private void sendReleaseUnusedGrant(DiscoveryNode owner, String bucketKey, int sharedLimit, String permitId, DiscoveryNode self) { + sendRelease(owner, bucketKey, sharedLimit, permitId, self.getId()); } // 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 permit — so failures are logged at debug only. - private void sendRelease(DiscoveryNode owner, String bucketKey, String permitId) { + // fatal — the owner's TTL sweep reclaims the permit — so failures are logged at debug only. Carries sharedLimit so + // the owner can drive owner-push (grant a waiter the freed slot) after releasing. {@code queueEmptyOnNodeId} is set + // only when returning an unused grant, so the owner deregisters that coordinator; empty for a normal release. + private void sendRelease(DiscoveryNode owner, String bucketKey, int sharedLimit, String permitId, String queueEmptyOnNodeId) { final TransportRequestOptions options = TransportRequestOptions.builder().withTimeout(ACQUIRE_TIMEOUT).build(); transportService.sendRequest( owner, RELEASE_ACTION_NAME, - new ReleasePermitRequest(bucketKey, permitId), + new ReleasePermitRequest(bucketKey, permitId, sharedLimit, queueEmptyOnNodeId), options, new TransportResponseHandler() { @Override @@ -319,6 +409,317 @@ private static String permitId() { return UUIDs.base64UUID(); } + /** + * Late-binds the coordinator-side grant consumer (the queue service). A received grant admits one queued request + * via this; before it is set, a grant is returned to the owner (safe: the reserved slot re-enters the pool). + */ + public void setGrantConsumer(GrantConsumer grantConsumer) { + this.grantConsumer = grantConsumer; + } + + // OWNER-SIDE: record that a coordinator is waiting on a bucket this node owns. Idempotent (a set), so re-registering + // an already-known waiter is a no-op — it keeps its existing queue position (LinkedHashSet.add does not reorder an + // element already present), so re-registration can't let a coordinator jump the rotation. Registry stays bounded at + // <= N coordinators per bucket. + private void registerWaiter(String bucketKey, DiscoveryNode coordinator) { + waitersByBucket.compute(bucketKey, (k, nodes) -> { + if (nodes == null) { + nodes = new LinkedHashSet<>(); + } + nodes.add(coordinator); + return nodes; + }); + } + + // OWNER-SIDE: drop a coordinator from a bucket's waiter set (it reported no more queued requests for the bucket, + // via an unused grant). Prunes the bucket entry when the last waiter leaves. + private void removeWaiter(String bucketKey, DiscoveryNode coordinator) { + removeWaiterByNodeId(bucketKey, coordinator.getId()); + } + + // As removeWaiter, but keyed on the coordinator's PERSISTENT node id alone — which is all the RELEASE RPC carries, + // since deregistration needs identity and nothing else. One removal implementation, and one identity basis for the + // whole registry: registration also resolves from the persistent id (see handleAcquire). The scan is O(n) where + // remove(node) was O(1), which is irrelevant — n is bounded by the number of coordinators in the cluster. Stays + // inside computeIfPresent because that per-key remapping is the sole lock guarding the non-thread-safe + // LinkedHashSet, and removeIf preserves insertion order so pickAndRotate's round-robin survives. + // + // Note this also clears a stale entry left by a restarted coordinator: DiscoveryNode identity is ephemeralId, so a + // restart leaves a second entry under the same persistent id, and matching on that id removes both. That is what we + // want — the old incarnation's parked requests died with its JVM. + private void removeWaiterByNodeId(String bucketKey, String nodeId) { + waitersByBucket.computeIfPresent(bucketKey, (k, nodes) -> { + nodes.removeIf(n -> n.getId().equals(nodeId)); + return nodes.isEmpty() ? null : nodes; + }); + } + + /** + * One TTL-sweep pass: reclaim expired permits, then drive owner-push for every bucket that gained free capacity. + * Package-private so tests can run a sweep deterministically instead of waiting on the scheduler. + *

    + * Reclaiming an expired permit frees a shared slot with NO release RPC behind it — the holder crashed, or its release + * was lost — so this is the only place that free slot is ever observed. Without the owner-push call a coordinator + * with a parked request stays registered as a waiter while capacity sits idle and, because parked requests have no + * deadline, strands until some unrelated release happens to re-drive the bucket. + */ + void sweepExpiredAndDrive() { + for (String bucketKey : tracker.sweepExpired()) { + onSharedSlotFreed(bucketKey, currentSharedLimit(bucketKey)); + } + } + + /** + * The current {@code shared_limit} for a bucket, resolved from cluster state, or + * {@link WorkloadGroupThrottleSettings#UNSET_LIMIT} if the group is gone or the shared tier is not configured. Used + * by the TTL sweep, which knows only a bucket key: unlike the release path, no {@code sharedLimit} travels with an + * expiry. {@link #onSharedSlotFreed} treats {@code UNSET_LIMIT} as "do not grant", so an unresolvable bucket simply + * skips owner-push rather than guessing a limit. + */ + private int currentSharedLimit(String bucketKey) { + // The group id is the bucketKey prefix before the first ':' (see WorkloadGroupService.buildBucketKey); group ids + // are base64 UUIDs with no ':', so the first ':' is unambiguous. + final int idx = bucketKey.indexOf(':'); + final String groupId = idx < 0 ? bucketKey : bucketKey.substring(0, idx); + final WorkloadGroup workloadGroup = clusterService.state().metadata().workloadGroups().get(groupId); + if (workloadGroup == null) { + return WorkloadGroupThrottleSettings.UNSET_LIMIT; + } + return WorkloadGroupThrottleSettings.SHARED_LIMIT.get(workloadGroup.getMutableWorkloadGroupFragment().getThrottling()); + } + + // OWNER-SIDE: a shared slot for this bucket just freed. Reserve it and push a grant to one waiting coordinator. + // Iterative, not recursive: if a chosen waiter turns out to be gone (disconnected), we reclaim and try the next + // one in a bounded loop (at most one pass over the <= N waiters), so a burst of stale/undeliverable waiters can + // never blow the stack. The unused-grant case (coordinator connected but has no queued request) is handled off + // this thread by the grant round-trip, which removes the waiter and re-drives once — not looped here. + private void onSharedSlotFreed(String bucketKey, int sharedLimit) { + if (sharedLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) { + return; // reclaim-only release (e.g. failed-acquire cleanup): never drive a grant + } + // Absolute safety cap on total iterations to bound work and rule out livelock, while still allowing the loop to + // react to waiters that REGISTER during it. A fixed snapshot of waiterCount is not enough: a stale local waiter + // or a disconnected waiter frees the reserved slot without handing it off, and a coordinator can registerWaiter + // (on a concurrent denied acquire) after the snapshot — leaving a free slot with an un-granted waiter, which + // would otherwise strand until the next unrelated release (a spurious queue.timeout despite free capacity). + // The cap is generous (waiters at entry, doubled, plus a constant); each iteration makes progress (serves, + // drops a stale/disconnected waiter, or stops), so the loop terminates well within it in practice. + final int maxAttempts = Math.max(1, waiterCount(bucketKey) * 2 + 8); + for (int attempt = 0; attempt < maxAttempts; attempt++) { + final DiscoveryNode target = pickAndRotate(bucketKey); + if (target == null) { + return; // no waiters currently registered + } + final String reservedPermitId = permitId(); + // Reserve the freed slot so a concurrent acquire can't take it before the grant lands. If the bucket is at + // its limit again (a racing acquire beat us), don't over-grant: stop. The waiter stays registered (peek did + // not remove it), so the next release re-drives owner-push to it. + if (tracker.tryAcquire(bucketKey, sharedLimit, reservedPermitId, PERMIT_TTL_NANOS) == false) { + return; + } + if (target.getId().equals(clusterService.localNode().getId())) { + // Local waiter: consume in-process. Returns true (STOP) if it admitted a queued request or admission + // failed unexpectedly; false only when this coordinator had no queued request — it then already + // released the reserved slot AND removed itself from the set, so the loop advances to another waiter + // with the re-freed slot (including any registered concurrently with this loop). + if (consumeGrantLocal(bucketKey, sharedLimit, reservedPermitId, target)) { + return; + } + // fall through: slot freed but not handed off -> re-check for a (possibly newly-registered) waiter + } else if (transportService.nodeConnected(target) == false) { + // Disconnected waiter: reclaim the reserved slot, drop it, and loop to the next waiter — handled here + // in the bounded loop rather than by recursing through sendGrant. + tracker.release(bucketKey, reservedPermitId); + removeWaiter(bucketKey, target); + // fall through: slot freed but not handed off -> re-check + } else { + // Remote, connected waiter: fire the grant and stop. Delivery failure and unused-grant return are + // handled asynchronously by sendGrant's response handler + the grant handler (which re-drive once). + sendGrant(target, bucketKey, sharedLimit, reservedPermitId); + return; + } + } + } + + // OWNER-SIDE: consume a grant for a LOCAL waiter without a network hop. Returns true if the caller's drain loop + // should STOP for this freed slot — either a queued request was admitted (slot handed off), or admission failed + // unexpectedly and retrying is unsafe. Returns false ONLY when this coordinator had no queued request (reserved slot + // released, waiter removed), so the caller's loop can try the next waiter with the re-freed slot. + // + // Recursion safety: on ADMIT we hand a self-re-driving reservedPermit — but its close() fires later, at request + // completion (async), so re-driving then is fine and not re-entrant. On the NO-request path we release the reserved + // permit DIRECTLY (not by closing a re-driving permit) and return false so the caller's bounded loop advances, + // rather than recursing through onSharedSlotFreed. + private boolean consumeGrantLocal(String bucketKey, int sharedLimit, String reservedPermitId, DiscoveryNode self) { + final GrantConsumer consumer = grantConsumer; + if (consumer == null) { + tracker.release(bucketKey, reservedPermitId); + removeWaiter(bucketKey, self); + return false; + } + final DiscoveryNode owner = clusterService.localNode(); // local path: this node owns the bucket + boolean admitted; + try { + admitted = consumer.admit(bucketKey, reservedPermit(owner, bucketKey, sharedLimit, reservedPermitId)); + } catch (Exception e) { + // admit() dispatches on GENERIC (which never rejects); a throw here means the executor is shutting down. + // Release the reserved slot and STOP the loop — do NOT retry. admitWithPermit polls the head request before + // admit() can throw, so a retry against the still-registered waiter would poll-and-drop a further request on + // every iteration during a dispatch outage. The waiter stays registered; a later release re-drives it. + logger.debug("Queue grant admit failed for bucket [" + bucketKey + "]", e); + tracker.release(bucketKey, reservedPermitId); + return true; + } + if (admitted == false) { + // admitWithPermit does not close the permit on a false return; release the reserved permit directly so the + // slot is reused by this loop. (Not via the permit's close(), which would re-drive owner-push inline.) + tracker.release(bucketKey, reservedPermitId); + removeWaiter(bucketKey, self); // this coordinator has nothing queued for the bucket + return false; + } + return true; + } + + // Current number of coordinators waiting on a bucket (0 if none). Reads the size under the map's per-key remapping + // lock (returning the set unchanged), since the LinkedHashSet is not safe to size concurrently with a rotate/add. + private int waiterCount(String bucketKey) { + final int[] count = new int[1]; + waitersByBucket.computeIfPresent(bucketKey, (k, nodes) -> { + count[0] = nodes.size(); + return nodes; + }); + return count[0]; + } + + // OWNER-SIDE: pick the head waiting coordinator for the bucket and ROTATE it to the tail, WITHOUT removing it. + // Returns null if none. Membership means "this coordinator has at least one request parked for the bucket", so a + // waiter stays registered across a successful grant hand-off — it may still have more queued requests, and each + // subsequent release re-drives owner-push. Rotating the picked coordinator to the tail gives round-robin fairness: + // successive freed slots serve different coordinators in turn instead of repeatedly serving whichever one is first, + // so one coordinator's deep backlog can't starve another's single request. A waiter is dropped only by the explicit + // self-reconciling signals: an unused grant returned (remote via ReleaseRequest.unusedGrantFrom, local via + // consumeGrantLocal) or a disconnect. Removing it here on pick would sever the drain chain after the first request, + // stranding the rest until queue.timeout despite free capacity. Runs under the map's per-key remapping lock (the + // sole guard for the non-thread-safe LinkedHashSet). + private DiscoveryNode pickAndRotate(String bucketKey) { + final DiscoveryNode[] picked = new DiscoveryNode[1]; + waitersByBucket.computeIfPresent(bucketKey, (k, nodes) -> { + final java.util.Iterator it = nodes.iterator(); + if (it.hasNext()) { + final DiscoveryNode head = it.next(); + picked[0] = head; + it.remove(); // detach from the head... + nodes.add(head); // ...and re-append at the tail (round-robin), keeping it registered + } + return nodes.isEmpty() ? null : nodes; + }); + return picked[0]; + } + + // OWNER-SIDE: fire-and-forget GRANT RPC pushing a reserved permit to a waiting coordinator. If it can't be + // delivered, reclaim the reserved slot (release by id) so it isn't lost until TTL. + private void sendGrant(DiscoveryNode coordinator, String bucketKey, int sharedLimit, String reservedPermitId) { + if (transportService.nodeConnected(coordinator) == false) { + tracker.release(bucketKey, reservedPermitId); // waiter gone; reclaim immediately + removeWaiter(bucketKey, coordinator); // it's disconnected; drop it from the set + onSharedSlotFreed(bucketKey, sharedLimit); // try the next waiter (iterative; this call is not re-entrant here) + return; + } + final TransportRequestOptions options = TransportRequestOptions.builder().withTimeout(ACQUIRE_TIMEOUT).build(); + transportService.sendRequest( + coordinator, + GRANT_ACTION_NAME, + new GrantPermitRequest(bucketKey, sharedLimit, reservedPermitId), + 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) { + // Grant undeliverable: reclaim the reserved slot now (TTL is the ultimate backstop) and re-drive + // owner-push so a still-waiting coordinator gets the slot instead of it idling until TTL. + logger.debug("Shared throttle grant to [{}] for bucket [{}] failed; reclaiming", coordinator.getId(), bucketKey); + tracker.release(bucketKey, reservedPermitId); + onSharedSlotFreed(bucketKey, sharedLimit); + } + + @Override + public String executor() { + return ThreadPool.Names.SAME; + } + } + ); + } + + // COORDINATOR-SIDE: a grant arrived. Hand the reserved permit to the queue service to admit one queued request; if + // there is none (or no consumer wired yet), return the reserved slot — closing the permit releases the permit AND + // re-drives owner-push so the owner tries the next waiter (this is how a stale waiter count self-drains). + private void handleGrant(GrantPermitRequest request) { + consumeGrant(request.bucketKey, request.sharedLimit, request.permitId); + } + + private void consumeGrant(String bucketKey, int sharedLimit, String reservedPermitId) { + final GrantConsumer consumer = grantConsumer; + if (consumer == null) { + returnUnusedGrant(bucketKey, sharedLimit, reservedPermitId); // not wired yet -> return the slot + deregister + return; + } + // The permit handed to a successfully-admitted request releases only the reserved permit on completion (which + // re-drives owner-push at the owner). It does NOT carry the unused-grant deregister signal — an admitted + // coordinator is a legitimate ongoing waiter if it has more queued requests. + final DiscoveryNode owner = ring.get().ownerFor(bucketKey).orElse(null); + boolean admitted; + try { + admitted = consumer.admit(bucketKey, reservedPermit(owner, bucketKey, sharedLimit, reservedPermitId)); + } catch (Exception e) { + logger.warn("Queue grant admit failed for bucket [" + bucketKey + "]", e); + returnUnusedGrant(bucketKey, sharedLimit, reservedPermitId); + return; + } + if (admitted == false) { + // No queued request on this coordinator for the bucket: return the reserved slot AND tell the owner to + // deregister this coordinator so it stops drawing wasted grants. + returnUnusedGrant(bucketKey, sharedLimit, reservedPermitId); + } + } + + // Returns an unused reserved slot to its owner, tagging the release so the owner deregisters this coordinator from + // the bucket's waiter set (it has no queued request for the bucket) and then re-drives owner-push to the next + // waiter. If this node is the owner, does it in-process. + private void returnUnusedGrant(String bucketKey, int sharedLimit, String reservedPermitId) { + final DiscoveryNode owner = ring.get().ownerFor(bucketKey).orElse(null); + final DiscoveryNode self = clusterService.localNode(); + if (owner == null) { + return; // ring empty; the reserved permit (if any) is reclaimed by TTL + } + if (owner.getId().equals(self.getId())) { + tracker.release(bucketKey, reservedPermitId); + removeWaiter(bucketKey, self); + onSharedSlotFreed(bucketKey, sharedLimit); + } else { + sendReleaseUnusedGrant(owner, bucketKey, sharedLimit, reservedPermitId, self); + } + } + + // A Releasable for a reserved permit, releasing the same way a normal acquired permit does: locally if this node + // owns the bucket, else via a RELEASE RPC. Both paths re-drive owner-push (releaseLocal directly; releaseRemote via + // the owner's RELEASE handler), so returning an unused grant flows the slot to the next waiter. + private Releasable reservedPermit(DiscoveryNode owner, String bucketKey, int sharedLimit, String permitId) { + if (owner == null) { + return releaseOnce(() -> {}); // ring empty; nothing to release remotely + } + if (owner.getId().equals(clusterService.localNode().getId())) { + return releaseLocal(bucketKey, sharedLimit, permitId); + } + return releaseRemote(owner, bucketKey, sharedLimit, permitId); + } + // 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 @@ -336,6 +737,11 @@ ThrottleOwnerSelector ring() { return ring.get(); } + // Package-private accessor for tests: number of coordinators currently registered as waiters on a bucket. + int waiterCountForTest(String bucketKey) { + return waiterCount(bucketKey); + } + /** * Shared body helpers for the RPC types below. All three serialize their body as ONE count-prefixed, name-keyed map * rather than as positional fields, so a peer built from a different commit of the same release stays interoperable. @@ -398,8 +804,44 @@ private static boolean requireBoolean(Map body, String key) { throw new IllegalStateException("wlm shared-throttle: key [" + key + "] missing or not a Boolean [" + value + "]"); } - // Adding a field? Read it strictly only if a version gate guarantees the sender has it; if it can be absent from a - // same-version build, read it with a safe default instead. See the rule above. + // --- fields added after the map format shipped: optional, with a safe default. See the rule above. --- + // + // Queueing added four such fields (requesting_node_id / wants_queue on acquire, shared_limit / + // queue_empty_on_node_id on release). A peer predating queueing sends none of them, so reading them strictly would + // throw on every message from every such node. Each default reproduces the pre-queueing behaviour exactly. + + private static String optionalString(Map body, String key, String fallback) { + final Object value = body.get(key); + if (value instanceof String s) { + return s; + } + if (value != null) { + logger.debug("wlm shared-throttle: key [{}] is not a String ([{}]); using [{}]", key, value, fallback); + } + return fallback; + } + + private static long optionalLong(Map body, String key, long fallback) { + final Object value = body.get(key); + if (value instanceof Number n) { + return n.longValue(); + } + if (value != null) { + logger.debug("wlm shared-throttle: key [{}] is not a Number ([{}]); using [{}]", key, value, fallback); + } + return fallback; + } + + private static boolean optionalBoolean(Map body, String key, boolean fallback) { + final Object value = body.get(key); + if (value instanceof Boolean b) { + return b; + } + if (value != null) { + logger.debug("wlm shared-throttle: key [{}] is not a Boolean ([{}]); using [{}]", key, value, fallback); + } + return fallback; + } /** * {@code coord -> owner}. Acquire RPC: a coordinator asks the bucket's owner to admit one request under @@ -410,17 +852,46 @@ public static class AcquirePermitRequest extends TransportRequest { static final String KEY_SHARED_LIMIT = "shared_limit"; static final String KEY_PERMIT_ID = "permit_id"; static final String KEY_TTL_NANOS = "ttl_nanos"; + static final String KEY_REQUESTING_NODE_ID = "requesting_node_id"; + static final String KEY_WANTS_QUEUE = "wants_queue"; final String bucketKey; final int sharedLimit; final String permitId; final long ttlNanos; + /** + * Owner-push: who is asking, and whether they will park the request on denial (so the owner should register a + * waiter and later push a grant). Both are ADDED fields — read with {@code optional*} and a safe default, never + * {@code require*}, because a peer predating queueing does not send them (see the read policy above). + *

    + * Only the coordinator's PERSISTENT node id travels, not the {@link DiscoveryNode}: a DiscoveryNode is not in the + * {@code writeGenericValue} registry so it cannot ride in the map, and the owner does not need the object on the + * wire — it resolves the live node from cluster state via {@code DiscoveryNodes.get(nodeId)}, which is O(1) and + * indexed on exactly this id. The owner needs the resolved node (not just an identity) because {@code sendGrant} + * addresses it. Empty means "not an owner-push acquire". + */ + final String requestingNodeId; + final boolean wantsQueue; + // Convenience for callers/tests that don't use owner-push (no waiter registration on denial). AcquirePermitRequest(String bucketKey, int sharedLimit, String permitId, long ttlNanos) { + this(bucketKey, sharedLimit, permitId, ttlNanos, "", false); + } + + AcquirePermitRequest( + String bucketKey, + int sharedLimit, + String permitId, + long ttlNanos, + String requestingNodeId, + boolean wantsQueue + ) { this.bucketKey = bucketKey; this.sharedLimit = sharedLimit; this.permitId = permitId; this.ttlNanos = ttlNanos; + this.requestingNodeId = requestingNodeId == null ? "" : requestingNodeId; + this.wantsQueue = wantsQueue; } AcquirePermitRequest(StreamInput in) throws IOException { @@ -430,6 +901,10 @@ public static class AcquirePermitRequest extends TransportRequest { this.sharedLimit = requireNumber(body, KEY_SHARED_LIMIT).intValue(); this.permitId = requireString(body, KEY_PERMIT_ID); this.ttlNanos = requireNumber(body, KEY_TTL_NANOS).longValue(); + // Added by queueing, so optional: a peer predating it sends neither key. The defaults are the pre-queueing + // behaviour — no owner-push, no waiter registered — so an older peer's acquire is handled exactly as before. + this.requestingNodeId = optionalString(body, KEY_REQUESTING_NODE_ID, ""); + this.wantsQueue = optionalBoolean(body, KEY_WANTS_QUEUE, false); // Any other key is a field this build does not know about: ignored on purpose. That is the tolerance. } @@ -437,7 +912,20 @@ public static class AcquirePermitRequest extends TransportRequest { public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeMap( - Map.of(KEY_BUCKET, bucketKey, KEY_SHARED_LIMIT, sharedLimit, KEY_PERMIT_ID, permitId, KEY_TTL_NANOS, ttlNanos), + Map.of( + KEY_BUCKET, + bucketKey, + KEY_SHARED_LIMIT, + sharedLimit, + KEY_PERMIT_ID, + permitId, + KEY_TTL_NANOS, + ttlNanos, + KEY_REQUESTING_NODE_ID, + requestingNodeId, + KEY_WANTS_QUEUE, + wantsQueue + ), StreamOutput::writeString, StreamOutput::writeGenericValue ); @@ -473,13 +961,37 @@ public void writeTo(StreamOutput out) throws IOException { public static class ReleasePermitRequest extends TransportRequest { static final String KEY_BUCKET = "bucket_key"; static final String KEY_PERMIT_ID = "permit_id"; + static final String KEY_SHARED_LIMIT = "shared_limit"; + static final String KEY_QUEUE_EMPTY_ON_NODE_ID = "queue_empty_on_node_id"; final String bucketKey; final String permitId; + /** + * The bucket's shared limit, so the owner can drive owner-push (grant a waiter) after freeing the slot. + * {@code UNSET_LIMIT} means "release only, do not drive push" — which is also the default when an older peer + * omits the key, so such a release behaves exactly as it did before queueing existed. + */ + final int sharedLimit; + /** + * Set to the coordinator's PERSISTENT node id only when this release is returning an UNUSED grant: the owner then + * deregisters that coordinator from the bucket's waiter set, since it has no queued request. Empty for a normal + * release, and empty is also the default when an older peer omits the key. + *

    + * Persistent node id rather than a {@link DiscoveryNode} for the same reason as {@code requestingNodeId} above, + * and rather than an ephemeralId so that registration and deregistration key the waiter set on ONE identity. + */ + final String queueEmptyOnNodeId; + // Convenience for a normal release that does not drive owner-push. ReleasePermitRequest(String bucketKey, String permitId) { + this(bucketKey, permitId, WorkloadGroupThrottleSettings.UNSET_LIMIT, ""); + } + + ReleasePermitRequest(String bucketKey, String permitId, int sharedLimit, String queueEmptyOnNodeId) { this.bucketKey = bucketKey; this.permitId = permitId; + this.sharedLimit = sharedLimit; + this.queueEmptyOnNodeId = queueEmptyOnNodeId == null ? "" : queueEmptyOnNodeId; } ReleasePermitRequest(StreamInput in) throws IOException { @@ -487,6 +999,10 @@ public static class ReleasePermitRequest extends TransportRequest { final Map body = readBody(in); this.bucketKey = requireString(body, KEY_BUCKET); this.permitId = requireString(body, KEY_PERMIT_ID); + // Added by queueing, so optional. The defaults are the pre-queueing behaviour: no owner-push, no waiter + // deregistration. + this.sharedLimit = (int) optionalLong(body, KEY_SHARED_LIMIT, WorkloadGroupThrottleSettings.UNSET_LIMIT); + this.queueEmptyOnNodeId = optionalString(body, KEY_QUEUE_EMPTY_ON_NODE_ID, ""); // Any other key is a field this build does not know about: ignored on purpose. That is the tolerance. } @@ -494,10 +1010,71 @@ public static class ReleasePermitRequest extends TransportRequest { public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeMap( - Map.of(KEY_BUCKET, bucketKey, KEY_PERMIT_ID, permitId), + Map.of( + KEY_BUCKET, + bucketKey, + KEY_PERMIT_ID, + permitId, + KEY_SHARED_LIMIT, + sharedLimit, + KEY_QUEUE_EMPTY_ON_NODE_ID, + queueEmptyOnNodeId + ), StreamOutput::writeString, StreamOutput::writeGenericValue ); } } + + /** + * {@code owner -> coord}. Grant: a reserved shared permit for a bucket the coordinator is waiting on. Carries the + * bucket's shared limit so a returned (unused) grant can re-drive owner-push toward the next waiter. + */ + public static class GrantPermitRequest extends TransportRequest { + // Born with the map format, so all three keys are baseline and read strictly. + static final String KEY_BUCKET = "bucket_key"; + static final String KEY_SHARED_LIMIT = "shared_limit"; + static final String KEY_PERMIT_ID = "permit_id"; + + final String bucketKey; + final int sharedLimit; + final String permitId; + + GrantPermitRequest(String bucketKey, int sharedLimit, String permitId) { + this.bucketKey = bucketKey; + this.sharedLimit = sharedLimit; + this.permitId = permitId; + } + + GrantPermitRequest(StreamInput in) throws IOException { + super(in); + final Map body = readBody(in); + this.bucketKey = requireString(body, KEY_BUCKET); + this.sharedLimit = requireNumber(body, KEY_SHARED_LIMIT).intValue(); + this.permitId = requireString(body, KEY_PERMIT_ID); + // Any other key is a field this build does not know about: ignored on purpose. That is the tolerance. + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeMap( + Map.of(KEY_BUCKET, bucketKey, KEY_SHARED_LIMIT, sharedLimit, KEY_PERMIT_ID, permitId), + StreamOutput::writeString, + StreamOutput::writeGenericValue + ); + } + } + + /** + * COORDINATOR-SIDE seam: consumes a pushed grant by admitting one queued request against the reserved permit. + * Returns {@code true} if a queued request was admitted (it now owns the permit), {@code false} if there was none + * (the caller returns the reserved slot). Implemented by the queue service; late-bound via + * {@link #setGrantConsumer}. + */ + @ExperimentalApi + @FunctionalInterface + public interface GrantConsumer { + boolean admit(String bucketKey, Releasable reservedPermit); + } } diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java index e07d9981434a3..15950b11ab3e9 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java @@ -11,11 +11,14 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.common.annotation.PublicApi; +import org.opensearch.common.lease.Releasable; import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.tasks.TaskId; import org.opensearch.tasks.CancellableTask; +import java.util.ArrayList; +import java.util.List; import java.util.Map; import java.util.function.LongSupplier; import java.util.function.Supplier; @@ -43,6 +46,14 @@ public class WorkloadGroupTask extends CancellableTask { private String workloadGroupId; private boolean isWorkloadGroupSet = false; + // Cancellation callbacks for a request parked in the WLM request queue. A queued request holds no thread; when the + // task is cancelled (client disconnect or cancel_after timer, both of which cancel the task rather than completing + // the parked listener), the queue must evict the entry and fail the listener. Guarded by its own lock and fired at + // most once: once cancelled, a later register runs immediately so there is no lost-cancellation race. + private final Object cancelCallbackLock = new Object(); + private List onCancelledCallbacks = new ArrayList<>(); + private boolean cancellationNotified = false; + public WorkloadGroupTask(long id, String type, String action, String description, TaskId parentTaskId, Map headers) { this(id, type, action, description, parentTaskId, headers, NO_TIMEOUT, System::nanoTime); } @@ -110,4 +121,53 @@ public boolean isWorkloadGroupSet() { public boolean shouldCancelChildrenOnCancellation() { return false; } + + /** + * Registers a callback invoked once if this task is cancelled, and returns a {@link Releasable} that deregisters it. + * If the task is already cancelled the callback runs immediately (so there is no window where a cancellation between + * the cancel check and registration is lost). Used by the WLM request queue to evict and fail a parked request whose + * client disconnected or whose {@code cancel_after} timer fired — both cancel the task rather than completing the + * parked listener, so the queue observes cancellation here. + * + * @param callback run at most once on cancellation + * @return a {@link Releasable} that removes the callback if the request is admitted/drained before any cancellation + */ + public final Releasable addOnCancelledCallback(Runnable callback) { + synchronized (cancelCallbackLock) { + if (cancellationNotified || isCancelled()) { + // Already cancelled: run now rather than register, so a cancellation that landed before registration + // is never dropped. Nothing to deregister. + callback.run(); + return () -> {}; + } + onCancelledCallbacks.add(callback); + } + return () -> { + synchronized (cancelCallbackLock) { + if (onCancelledCallbacks != null) { + onCancelledCallbacks.remove(callback); + } + } + }; + } + + @Override + protected void onCancelled() { + final List callbacks; + synchronized (cancelCallbackLock) { + if (cancellationNotified) { + return; + } + cancellationNotified = true; + callbacks = onCancelledCallbacks; + onCancelledCallbacks = null; + } + for (Runnable callback : callbacks) { + try { + callback.run(); + } catch (Exception e) { + logger.warn("WorkloadGroupTask onCancelled callback failed", e); + } + } + } } diff --git a/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupState.java b/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupState.java index 9efc2dbf7a5d6..4ad37a3ee502b 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupState.java +++ b/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupState.java @@ -13,6 +13,7 @@ import java.util.EnumMap; import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; /** * This class will keep the point in time view of the workload group stats @@ -43,6 +44,35 @@ public class WorkloadGroupState { */ public final CounterMetric totalThrottled = new CounterMetric(); + /** + * Cumulative requests parked in the workload group's request queue (admitted into the queue after a throttle + * denial), since the OpenSearch start time. + */ + public final CounterMetric totalQueued = new CounterMetric(); + + /** + * Cumulative requests rejected because the workload group's request queue was full, since the OpenSearch start time. + */ + public final CounterMetric totalQueueRejections = new CounterMetric(); + + /** + * Cumulative time (in millis) that admitted requests spent parked in the queue, summed across all requests that + * were queued and then admitted. Paired with {@link #queueWaitCount} this yields the mean wait; use + * {@link #maxQueueWaitMillis} for the tail. Recorded only for requests that actually parked (never-queued requests + * do not contribute), so the mean reflects wait among queued requests, not all requests. + */ + public final CounterMetric totalQueueWaitMillis = new CounterMetric(); + + /** + * Number of admitted requests that had been parked in the queue (the denominator for mean queue wait). + */ + public final CounterMetric queueWaitCount = new CounterMetric(); + + /** + * High-water mark (in millis) of any single admitted request's queue wait, since the OpenSearch start time. + */ + private final AtomicLong maxQueueWaitMillis = new AtomicLong(0); + /** * This is used to store the resource type state both for CPU and MEMORY */ @@ -93,6 +123,59 @@ public long getTotalThrottled() { return totalThrottled.count(); } + /** + * + * @return requests parked in the workload group's request queue + */ + public long getTotalQueued() { + return totalQueued.count(); + } + + /** + * + * @return requests rejected because the workload group's request queue was full + */ + public long getTotalQueueRejections() { + return totalQueueRejections.count(); + } + + /** + * Records the queue wait of one request that was parked and then admitted: adds to the cumulative sum and count + * and advances the high-water mark. Called once per admitted-from-queue request. + * + * @param waitMillis how long the request was parked, in millis (non-negative) + */ + public void recordQueueWaitMillis(long waitMillis) { + if (waitMillis < 0) { + waitMillis = 0; + } + totalQueueWaitMillis.inc(waitMillis); + queueWaitCount.inc(); + final long w = waitMillis; + maxQueueWaitMillis.accumulateAndGet(w, Math::max); + } + + /** + * @return cumulative parked time (millis) summed over admitted-from-queue requests + */ + public long getTotalQueueWaitMillis() { + return totalQueueWaitMillis.count(); + } + + /** + * @return number of admitted requests that had been parked (denominator for mean wait) + */ + public long getQueueWaitCount() { + return queueWaitCount.count(); + } + + /** + * @return the longest single queue wait observed (millis) + */ + public long getMaxQueueWaitMillis() { + return maxQueueWaitMillis.get(); + } + /** * getter for workload group resource state * @return the workload group resource state diff --git a/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupStats.java b/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupStats.java index fada855f4eefe..d514093bf182c 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupStats.java +++ b/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupStats.java @@ -97,11 +97,27 @@ public static class WorkloadGroupStatsHolder implements ToXContentObject, Writea public static final String TOTAL_CANCELLATIONS = "total_cancellations"; public static final String FAILURES = "failures"; public static final String THROTTLED = "total_throttled"; + public static final String QUEUED = "total_queued"; + public static final String QUEUE_REJECTIONS = "total_queue_rejections"; + public static final String QUEUED_CURRENT = "queued_current"; + public static final String QUEUE_PEAK = "queue_peak"; + public static final String TOTAL_QUEUE_WAIT_MILLIS = "total_queue_wait_millis"; + public static final String QUEUE_WAIT_COUNT = "queue_wait_count"; + public static final String MAX_QUEUE_WAIT_MILLIS = "max_queue_wait_millis"; private long completions; private long rejections; private long failures; private long cancellations; private long throttled; + private long queued; + private long queueRejections; + private long queuedCurrent; + private long queuePeak; + // Cumulative parked time + count (mean = sum/count) and the single-request high-water mark, for queued-then- + // admitted requests. Populated from WorkloadGroupState in from(...); 0 via the plain constructors. + private long totalQueueWaitMillis; + private long queueWaitCount; + private long maxQueueWaitMillis; private Map resourceStats; // this is needed to support the factory method @@ -114,12 +130,31 @@ public WorkloadGroupStatsHolder( long cancellations, long throttled, Map resourceStats + ) { + this(completions, rejections, failures, cancellations, throttled, 0, 0, 0, 0, resourceStats); + } + + public WorkloadGroupStatsHolder( + long completions, + long rejections, + long failures, + long cancellations, + long throttled, + long queued, + long queueRejections, + long queuedCurrent, + long queuePeak, + Map resourceStats ) { this.completions = completions; this.rejections = rejections; this.failures = failures; this.cancellations = cancellations; this.throttled = throttled; + this.queued = queued; + this.queueRejections = queueRejections; + this.queuedCurrent = queuedCurrent; + this.queuePeak = queuePeak; this.resourceStats = resourceStats; } @@ -128,9 +163,16 @@ public WorkloadGroupStatsHolder(StreamInput in) throws IOException { this.rejections = in.readVLong(); this.failures = in.readVLong(); this.cancellations = in.readVLong(); - // total_throttled is version-gated so a pre-throttling node's stats stream stays readable. + // total_throttled and the queue stats are version-gated so a pre-throttling node's stats stream stays readable. if (in.getVersion().onOrAfter(Version.V_3_7_0)) { this.throttled = in.readVLong(); + this.queued = in.readVLong(); + this.queueRejections = in.readVLong(); + this.queuedCurrent = in.readVLong(); + this.queuePeak = in.readVLong(); + this.totalQueueWaitMillis = in.readVLong(); + this.queueWaitCount = in.readVLong(); + this.maxQueueWaitMillis = in.readVLong(); } this.resourceStats = in.readMap((i) -> ResourceType.fromName(i.readString()), ResourceStats::new); } @@ -151,16 +193,57 @@ public long getThrottled() { return throttled; } + public long getQueued() { + return queued; + } + + public long getQueueRejections() { + return queueRejections; + } + + public long getQueuedCurrent() { + return queuedCurrent; + } + + public long getQueuePeak() { + return queuePeak; + } + + public long getTotalQueueWaitMillis() { + return totalQueueWaitMillis; + } + + public long getQueueWaitCount() { + return queueWaitCount; + } + + public long getMaxQueueWaitMillis() { + return maxQueueWaitMillis; + } + public Map getResourceStats() { return resourceStats; } /** - * static factory method to convert {@link WorkloadGroupState} into {@link WorkloadGroupStatsHolder} + * static factory method to convert {@link WorkloadGroupState} into {@link WorkloadGroupStatsHolder}, with no + * live queue depth (used where a queue service is not available, e.g. tests). * @param workloadGroupState which needs to be converted * @return WorkloadGroupStatsHolder object */ public static WorkloadGroupStatsHolder from(WorkloadGroupState workloadGroupState) { + return from(workloadGroupState, 0L, 0L); + } + + /** + * static factory method to convert {@link WorkloadGroupState} into {@link WorkloadGroupStatsHolder}, including + * the point-in-time queue depth gauges (which live in the queue service, not the state). + * @param workloadGroupState which needs to be converted + * @param queuedCurrent current queued depth for this group + * @param queuePeak peak queued depth for this group + * @return WorkloadGroupStatsHolder object + */ + public static WorkloadGroupStatsHolder from(WorkloadGroupState workloadGroupState, long queuedCurrent, long queuePeak) { final WorkloadGroupStatsHolder statsHolder = new WorkloadGroupStatsHolder(); Map resourceStatsMap = new HashMap<>(); @@ -174,6 +257,13 @@ public static WorkloadGroupStatsHolder from(WorkloadGroupState workloadGroupStat statsHolder.failures = workloadGroupState.getFailures(); statsHolder.cancellations = workloadGroupState.getTotalCancellations(); statsHolder.throttled = workloadGroupState.getTotalThrottled(); + statsHolder.queued = workloadGroupState.getTotalQueued(); + statsHolder.queueRejections = workloadGroupState.getTotalQueueRejections(); + statsHolder.queuedCurrent = queuedCurrent; + statsHolder.queuePeak = queuePeak; + statsHolder.totalQueueWaitMillis = workloadGroupState.getTotalQueueWaitMillis(); + statsHolder.queueWaitCount = workloadGroupState.getQueueWaitCount(); + statsHolder.maxQueueWaitMillis = workloadGroupState.getMaxQueueWaitMillis(); statsHolder.resourceStats = resourceStatsMap; return statsHolder; } @@ -192,6 +282,13 @@ public static void writeTo(StreamOutput out, WorkloadGroupStatsHolder statsHolde // version-gated to match the StreamInput ctor; read/write order must stay in sync. if (out.getVersion().onOrAfter(Version.V_3_7_0)) { out.writeVLong(statsHolder.throttled); + out.writeVLong(statsHolder.queued); + out.writeVLong(statsHolder.queueRejections); + out.writeVLong(statsHolder.queuedCurrent); + out.writeVLong(statsHolder.queuePeak); + out.writeVLong(statsHolder.totalQueueWaitMillis); + out.writeVLong(statsHolder.queueWaitCount); + out.writeVLong(statsHolder.maxQueueWaitMillis); } out.writeMap(statsHolder.resourceStats, (o, val) -> o.writeString(val.getName()), ResourceStats::writeTo); } @@ -209,6 +306,13 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws // builder.field(FAILURES, failures); builder.field(TOTAL_CANCELLATIONS, cancellations); builder.field(THROTTLED, throttled); + builder.field(QUEUED, queued); + builder.field(QUEUE_REJECTIONS, queueRejections); + builder.field(QUEUED_CURRENT, queuedCurrent); + builder.field(QUEUE_PEAK, queuePeak); + builder.field(TOTAL_QUEUE_WAIT_MILLIS, totalQueueWaitMillis); + builder.field(QUEUE_WAIT_COUNT, queueWaitCount); + builder.field(MAX_QUEUE_WAIT_MILLIS, maxQueueWaitMillis); for (ResourceType resourceType : ResourceType.getSortedValues()) { ResourceStats resourceStats1 = resourceStats.get(resourceType); @@ -230,12 +334,33 @@ public boolean equals(Object o) { && Objects.equals(resourceStats, that.resourceStats) && failures == that.failures && cancellations == that.cancellations - && throttled == that.throttled; + && throttled == that.throttled + && queued == that.queued + && queueRejections == that.queueRejections + && queuedCurrent == that.queuedCurrent + && queuePeak == that.queuePeak + && totalQueueWaitMillis == that.totalQueueWaitMillis + && queueWaitCount == that.queueWaitCount + && maxQueueWaitMillis == that.maxQueueWaitMillis; } @Override public int hashCode() { - return Objects.hash(completions, rejections, cancellations, failures, throttled, resourceStats); + return Objects.hash( + completions, + rejections, + cancellations, + failures, + throttled, + queued, + queueRejections, + queuedCurrent, + queuePeak, + totalQueueWaitMillis, + queueWaitCount, + maxQueueWaitMillis, + resourceStats + ); } } diff --git a/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java b/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java index e3c02197c7cc9..998242ac3de06 100644 --- a/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java +++ b/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java @@ -82,6 +82,13 @@ public void testToString() { + " \"total_rejections\" : 0,\n" + " \"total_cancellations\" : 0,\n" + " \"total_throttled\" : 0,\n" + + " \"total_queued\" : 0,\n" + + " \"total_queue_rejections\" : 0,\n" + + " \"queued_current\" : 0,\n" + + " \"queue_peak\" : 0,\n" + + " \"total_queue_wait_millis\" : 0,\n" + + " \"queue_wait_count\" : 0,\n" + + " \"max_queue_wait_millis\" : 0,\n" + " \"cpu\" : {\n" + " \"current_usage\" : 0.0,\n" + " \"cancellations\" : 0,\n" diff --git a/server/src/test/java/org/opensearch/cluster/metadata/WorkloadGroupTests.java b/server/src/test/java/org/opensearch/cluster/metadata/WorkloadGroupTests.java index 26117b9d2223c..21f44d5dcdc92 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/WorkloadGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/WorkloadGroupTests.java @@ -41,7 +41,9 @@ static WorkloadGroup createRandomWorkloadGroup(String _id) { // Generate a valid throttling config: either disabled (empty), or enabled with a required attribute plus // at least one positive limit (so the effective ceiling is >= 1). Settings.Builder throttling = Settings.builder(); + boolean throttlingConfigured = false; if (randomBoolean()) { + throttlingConfigured = true; throttling.put("attribute", randomFrom("group", "username", "role")); if (randomBoolean()) { throttling.put("node_limit", randomIntBetween(1, 100)); @@ -55,10 +57,16 @@ static WorkloadGroup createRandomWorkloadGroup(String _id) { } } } + // Queue config is only valid alongside a throttle limit (a queue with nothing to queue is rejected). + // size_per_bucket is the only queue setting. + Settings.Builder queue = Settings.builder(); + if (throttlingConfigured && randomBoolean()) { + queue.put("size_per_bucket", randomIntBetween(1, 1000)); + } return new WorkloadGroup( name, _id, - new MutableWorkloadGroupFragment(randomMode(), resourceLimit, Settings.EMPTY, throttling.build()), + new MutableWorkloadGroupFragment(randomMode(), resourceLimit, Settings.EMPTY, throttling.build(), queue.build()), Instant.now().getMillis() ); } @@ -753,4 +761,97 @@ public void testSettingsNullFromXContentClearsSettings() throws IOException { // Settings should be empty (cleared) assertTrue(fragment.getSettings().isEmpty()); } + + public void testQueueRequiresThrottleLimit() { + // A queue with no throttle limit has nothing to queue -> rejected at the workload-group level. + IllegalArgumentException e = expectThrows( + IllegalArgumentException.class, + () -> new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.EMPTY, // no throttling + Settings.builder().put("size_per_bucket", 100).build() + ), + System.currentTimeMillis() + ) + ); + assertTrue(e.getMessage(), e.getMessage().contains("queue requires a throttle limit")); + } + + public void testQueueTimeoutIsRejectedAsUnknownSetting() { + // queue.timeout was removed: a client bounds its wait via cancel_after_time_interval, and the queue has no + // wall-clock deadline. A stale queue.timeout config must be rejected as an unknown key, not silently accepted. + Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 10).build(); + IllegalArgumentException e = expectThrows( + IllegalArgumentException.class, + () -> new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + throttling, + Settings.builder().put("size_per_bucket", 10).put("timeout", "30s").build() + ), + System.currentTimeMillis() + ) + ); + assertTrue(e.getMessage(), e.getMessage().contains("Unknown queue setting")); + } + + public void testToXContentEmitsQueue() throws IOException { + long currentTimeInMillis = Instant.now().getMillis(); + String workloadGroupId = UUIDs.randomBase64UUID(); + Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 10).build(); + Settings queue = Settings.builder().put("size_per_bucket", 200).build(); + WorkloadGroup workloadGroup = new WorkloadGroup( + "TestWorkloadGroup", + workloadGroupId, + new MutableWorkloadGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.30), Settings.EMPTY, throttling, queue), + currentTimeInMillis + ); + XContentBuilder builder = JsonXContent.contentBuilder(); + workloadGroup.toXContent(builder, ToXContent.EMPTY_PARAMS); + String expected = String.format( + Locale.ROOT, + "{\"_id\":\"%s\",\"name\":\"TestWorkloadGroup\",\"resiliency_mode\":\"enforced\"," + + "\"resource_limits\":{\"cpu\":0.3}," + + "\"settings\":{}," + + "\"throttling\":{\"attribute\":\"username\",\"node_limit\":10}," + + "\"queue\":{\"size_per_bucket\":200}," + + "\"updated_at\":%d}", + workloadGroupId, + currentTimeInMillis + ); + assertEquals(expected, builder.toString()); + } + + public void testToXContentOmitsUnsetQueue() throws IOException { + Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 10).build(); + WorkloadGroup workloadGroup = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5), Settings.EMPTY, throttling), + System.currentTimeMillis() + ); + XContentBuilder builder = JsonXContent.contentBuilder(); + workloadGroup.toXContent(builder, ToXContent.EMPTY_PARAMS); + assertFalse(builder.toString().contains("queue")); + } + + public void testQueueNullFromXContentClearsQueue() throws IOException { + String json = "{\"_id\":\"test_id\",\"name\":\"test\",\"resiliency_mode\":\"enforced\"," + + "\"resource_limits\":{\"memory\":0.5}," + + "\"queue\":null," + + "\"updated_at\":1720047207}"; + XContentParser parser = createParser(JsonXContent.jsonXContent, json); + WorkloadGroup.Builder builder = WorkloadGroup.Builder.fromXContent(parser); + MutableWorkloadGroupFragment fragment = builder.getMutableWorkloadGroupFragment(); + assertTrue(fragment.getQueue().isEmpty()); + } } diff --git a/server/src/test/java/org/opensearch/wlm/SharedThrottleTrackerTests.java b/server/src/test/java/org/opensearch/wlm/SharedThrottleTrackerTests.java index 8e1373bf6e14d..3b3a19ea607a8 100644 --- a/server/src/test/java/org/opensearch/wlm/SharedThrottleTrackerTests.java +++ b/server/src/test/java/org/opensearch/wlm/SharedThrottleTrackerTests.java @@ -10,6 +10,7 @@ import org.opensearch.test.OpenSearchTestCase; +import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -215,4 +216,27 @@ public void testSweepRecomputesMinExpiryEnablingLaterSkip() { } assertEquals("no further scans while survivors are fresh", 1, tracker.pruneScanCount()); } + + public void testSweepExpiredReportsOnlyBucketsThatActuallyFreedCapacity() { + // The owning service drives owner-push from this return value, so it must name exactly the buckets that gained + // free capacity: a miss strands a waiting coordinator (an expiry has no release RPC to re-drive the bucket), and + // a false positive would make the owner reserve-and-grant a slot that does not exist. + AtomicLong now = new AtomicLong(0L); + SharedThrottleTracker tracker = new SharedThrottleTracker(now::get); + assertTrue(tracker.tryAcquire("expiring", 5, "lease-a", 1000L)); + assertTrue(tracker.tryAcquire("surviving", 5, "lease-b", 100_000L)); + + // Nothing has expired yet. + assertTrue("no expiry yet -> no bucket reported", tracker.sweepExpired().isEmpty()); + + // Past only the first lease's TTL. + now.set(1001L); + List freed = tracker.sweepExpired(); + assertEquals("exactly the bucket whose lease was reclaimed", List.of("expiring"), freed); + assertEquals(0, tracker.inFlight("expiring")); + assertEquals(1, tracker.inFlight("surviving")); + + // Idempotent: a second pass has nothing left to reclaim for that bucket. + assertTrue("a repeat sweep must not re-report an already-reclaimed bucket", tracker.sweepExpired().isEmpty()); + } } diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueServiceTests.java new file mode 100644 index 0000000000000..0fa5f62e18fed --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueServiceTests.java @@ -0,0 +1,136 @@ +/* + * 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.action.search.SearchTask; +import org.opensearch.common.lease.Releasable; +import org.opensearch.core.action.ActionListener; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; + +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +public class WorkloadGroupQueueServiceTests extends OpenSearchTestCase { + + private ThreadPool threadPool; + private WorkloadGroupsStateAccessor stateAccessor; + private WorkloadGroupQueueService service; + + @Override + public void setUp() throws Exception { + super.setUp(); + threadPool = new TestThreadPool(getTestName()); + stateAccessor = new WorkloadGroupsStateAccessor(); + stateAccessor.addNewWorkloadGroup("g1"); + service = new WorkloadGroupQueueService(threadPool, stateAccessor); + } + + @Override + public void tearDown() throws Exception { + threadPool.shutdown(); + super.tearDown(); + } + + private static SearchTask task() { + return new SearchTask(randomNonNegativeLong(), "", "", () -> "", null, null); + } + + public void testTryEnqueueRejectsWhenDisabled() { + boolean parked = service.tryEnqueue("g1", "g1:group", task(), 0, ActionListener.wrap(r -> {}, e -> {})); + assertFalse(parked); + assertEquals(0, service.currentDepth("g1")); + } + + public void testTryEnqueueParksAndCountsDepth() { + boolean parked = service.tryEnqueue("g1", "g1:group", task(), 5, ActionListener.wrap(r -> {}, e -> {})); + assertTrue(parked); + assertEquals(1, service.currentDepth("g1")); + assertEquals(1, service.totalDepth()); + } + + public void testTryEnqueueCapacityIsPerBucket() { + // The size_per_bucket cap is threaded through tryEnqueue per bucket, not as one budget for the whole group: + // saturating alice's bucket must not deny bob. + assertTrue(service.tryEnqueue("g1", "g1:username:alice", task(), 1, ActionListener.wrap(r -> {}, e -> {}))); + assertFalse(service.tryEnqueue("g1", "g1:username:alice", task(), 1, ActionListener.wrap(r -> {}, e -> {}))); + assertTrue(service.tryEnqueue("g1", "g1:username:bob", task(), 1, ActionListener.wrap(r -> {}, e -> {}))); + assertEquals(2, service.currentDepth("g1")); + } + + public void testDrainNodeAdmitsOneWaiterWithPermit() throws Exception { + AtomicReference admittedPermit = new AtomicReference<>(); + AtomicInteger admitted = new AtomicInteger(); + assertTrue(service.tryEnqueue("g1", "g1:group", task(), 5, ActionListener.wrap(p -> { + admittedPermit.set(p); + admitted.incrementAndGet(); + }, e -> {}))); + + Releasable permit = () -> {}; + service.drainNode("g1", "g1:group", key -> permit); + + assertBusy(() -> assertEquals(1, admitted.get())); + assertSame(permit, admittedPermit.get()); + assertEquals(0, service.currentDepth("g1")); // drained + } + + public void testDrainNodeAdmitsAtMostOnePerFreedPermit() throws Exception { + AtomicInteger admitted = new AtomicInteger(); + for (int i = 0; i < 3; i++) { + assertTrue(service.tryEnqueue("g1", "g1:group", task(), 5, ActionListener.wrap(p -> admitted.incrementAndGet(), e -> {}))); + } + assertEquals(3, service.currentDepth("g1")); + + // A single freed node permit admits exactly one waiter. + service.drainNode("g1", "g1:group", key -> (Releasable) () -> {}); + assertBusy(() -> assertEquals(1, admitted.get())); + assertEquals(2, service.currentDepth("g1")); // two still parked + } + + public void testDrainNodeLeavesWaiterQueuedWhenNoPermit() { + AtomicInteger admitted = new AtomicInteger(); + assertTrue(service.tryEnqueue("g1", "g1:group", task(), 5, ActionListener.wrap(p -> admitted.incrementAndGet(), e -> {}))); + // nodeAcquire returns null (limit reached) -> nothing admitted, request stays queued. + service.drainNode("g1", "g1:group", key -> null); + assertEquals(0, admitted.get()); + assertEquals(1, service.currentDepth("g1")); + } + + // Regression for the recursion fix: admit() must complete the parked listener on the executor, never inline on the + // caller's (drain/completion) thread. This is what breaks the node-tier close()->drainNode->admit->close() chain + // into separate executor tasks (a synchronous inline completion allowed StackOverflow under a cancel storm). We + // assert the common admit path hands off to a different thread than the caller — the same dispatch that also + // governs the cancelled branch's permit.close(). + public void testAdmitCompletesOffCallerThread() throws Exception { + Thread callerThread = Thread.currentThread(); + AtomicReference respondedOn = new AtomicReference<>(); + assertTrue( + service.tryEnqueue("g1", "g1:group", task(), 5, ActionListener.wrap(p -> respondedOn.set(Thread.currentThread()), e -> {})) + ); + service.drainNode("g1", "g1:group", key -> (Releasable) () -> {}); + assertBusy(() -> assertNotNull(respondedOn.get())); + assertNotSame("admit must complete the listener off the caller thread (recursion-safety)", callerThread, respondedOn.get()); + } + + // Regression for the exactly-once contract on cancellation-during-enqueue: an already-cancelled task must not be + // left parked, and its listener is failed exactly once. + public void testEnqueueOfAlreadyCancelledTaskFailsExactlyOnce() throws Exception { + SearchTask t = task(); + t.cancel("already gone"); + AtomicInteger failures = new AtomicInteger(); + AtomicBoolean admittedFlag = new AtomicBoolean(false); + // tryEnqueue may report parked=true, but the immediate cancellation callback evicts + fails it exactly once. + service.tryEnqueue("g1", "g1:group", t, 5, ActionListener.wrap(p -> admittedFlag.set(true), e -> failures.incrementAndGet())); + assertBusy(() -> assertEquals(1, failures.get())); + assertFalse(admittedFlag.get()); + assertEquals(0, service.currentDepth("g1")); + } +} diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueSettingsTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueSettingsTests.java new file mode 100644 index 0000000000000..28cc2d928591c --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueSettingsTests.java @@ -0,0 +1,81 @@ +/* + * 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.settings.Settings; +import org.opensearch.test.OpenSearchTestCase; + +public class WorkloadGroupQueueSettingsTests extends OpenSearchTestCase { + + public void testDefaults() { + assertEquals(0, WorkloadGroupQueueSettings.SIZE_PER_BUCKET.get(Settings.EMPTY).intValue()); + } + + public void testValidAcceptsSizePerBucket() { + Settings queue = Settings.builder().put("size_per_bucket", 200).build(); + WorkloadGroupQueueSettings.validate(queue); // no throw + assertEquals(200, WorkloadGroupQueueSettings.SIZE_PER_BUCKET.get(queue).intValue()); + } + + public void testValidRejectsTimeoutKey() { + // queue.timeout is no longer a setting: it must be rejected as an unknown key so a stale config surfaces clearly. + Settings queue = Settings.builder().put("timeout", "30s").build(); + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> WorkloadGroupQueueSettings.validate(queue)); + assertTrue(e.getMessage(), e.getMessage().contains("Unknown queue setting")); + } + + public void testValidRejectsOldSizeKey() { + // queue.size was renamed to queue.size_per_bucket (the cap is per throttle bucket, not per group). The old key + // must be rejected as unknown rather than silently ignored, so a stale config is not read as "queueing enabled". + Settings queue = Settings.builder().put("size", 100).build(); + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> WorkloadGroupQueueSettings.validate(queue)); + assertTrue(e.getMessage(), e.getMessage().contains("Unknown queue setting")); + } + + public void testValidNullIsNoOp() { + WorkloadGroupQueueSettings.validate(null); + } + + public void testValidRejectsUnknownKey() { + Settings queue = Settings.builder().put("bogus", 1).build(); + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> WorkloadGroupQueueSettings.validate(queue)); + assertTrue(e.getMessage(), e.getMessage().contains("Unknown queue setting")); + } + + public void testValidRejectsNegativeSizePerBucket() { + Settings queue = Settings.builder().put("size_per_bucket", -5).build(); + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> WorkloadGroupQueueSettings.validate(queue)); + assertTrue(e.getMessage(), e.getMessage().contains("non-negative")); + } + + public void testValidRejectsNonIntegerSizePerBucket() { + Settings queue = Settings.builder().put("size_per_bucket", "abc").build(); + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> WorkloadGroupQueueSettings.validate(queue)); + assertTrue(e.getMessage(), e.getMessage().contains("must be an integer")); + } + + public void testValidRejectsSizePerBucketAboveMax() { + Settings queue = Settings.builder().put("size_per_bucket", WorkloadGroupQueueSettings.MAX_SIZE_PER_BUCKET + 1).build(); + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> WorkloadGroupQueueSettings.validate(queue)); + assertTrue(e.getMessage(), e.getMessage().contains("must not exceed")); + } + + public void testValidAcceptsSizePerBucketAtMax() { + // The boundary itself is legal: a single-bucket group (attribute=group) may queue the entire group budget. + Settings queue = Settings.builder().put("size_per_bucket", WorkloadGroupQueueSettings.MAX_SIZE_PER_BUCKET).build(); + WorkloadGroupQueueSettings.validate(queue); // no throw + } + + public void testMaxSizePerBucketIsPinnedToGroupCeiling() { + // Invariant: the configurable per-bucket cap can never exceed the fixed per-group ceiling, otherwise validation + // would accept a depth the group total could never honour. + assertEquals(WorkloadGroupQueueSettings.MAX_GROUP_QUEUE_DEPTH, WorkloadGroupQueueSettings.MAX_SIZE_PER_BUCKET); + } + +} diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueTests.java new file mode 100644 index 0000000000000..423443fe5def3 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupQueueTests.java @@ -0,0 +1,228 @@ +/* + * 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.action.search.SearchTask; +import org.opensearch.common.lease.Releasable; +import org.opensearch.core.action.ActionListener; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.List; +import java.util.concurrent.atomic.AtomicReference; + +public class WorkloadGroupQueueTests extends OpenSearchTestCase { + + private static WorkloadGroupTask task() { + return new SearchTask(randomNonNegativeLong(), "", "", () -> "", null, null); + } + + private static WorkloadGroupQueue.QueuedRequest req(String bucketKey) { + return new WorkloadGroupQueue.QueuedRequest(ActionListener.wrap(r -> {}, e -> {}), bucketKey, task(), 0L); + } + + private static WorkloadGroupQueue.QueuedRequest req(String bucketKey, WorkloadGroupTask task) { + return new WorkloadGroupQueue.QueuedRequest(ActionListener.wrap(r -> {}, e -> {}), bucketKey, task, 0L); + } + + public void testOfferRejectedWhenDisabled() { + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + assertFalse(queue.offer(req("g:group"), 0)); + assertEquals(0, queue.currentDepth()); + } + + public void testOfferUpToBucketCapacityThenReject() { + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + assertTrue(queue.offer(req("g:group"), 2)); + assertTrue(queue.offer(req("g:group"), 2)); + assertFalse(queue.offer(req("g:group"), 2)); // this bucket is full + assertEquals(2, queue.currentDepth()); + assertEquals(2L, queue.peakDepth()); + } + + public void testBucketCapacityReflectsCurrentSizePerCall() { + // The cap is re-read per offer (dynamic queue.size_per_bucket), not frozen at construction. + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + assertTrue(queue.offer(req("g:group"), 5)); // bucket depth 1, cap 5 + assertTrue(queue.offer(req("g:group"), 5)); // bucket depth 2, cap 5 + // Cap lowered to 2 mid-flight: the bucket already holds 2, so the next offer is rejected immediately. Requests + // already parked are not evicted. + assertFalse(queue.offer(req("g:group"), 2)); + assertEquals(2, queue.currentDepth()); + // Cap raised to 3: a further offer now succeeds. + assertTrue(queue.offer(req("g:group"), 3)); + assertEquals(3, queue.currentDepth()); + } + + public void testCapacityIsPerBucketNotSharedAcrossBuckets() { + // Cross-principal fairness: each bucket gets its own size_per_bucket budget, so one principal filling its queue + // cannot deny another principal capacity (the pre-rename behavior, where one shared budget was consumed + // first-come-first-served across buckets). + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + assertTrue(queue.offer(req("g:username:alice"), 2)); + assertTrue(queue.offer(req("g:username:alice"), 2)); + assertFalse(queue.offer(req("g:username:alice"), 2)); // alice's own bucket is full + // bob and carol are unaffected by alice saturating hers. + assertTrue(queue.offer(req("g:username:bob"), 2)); + assertTrue(queue.offer(req("g:username:bob"), 2)); + assertTrue(queue.offer(req("g:username:carol"), 2)); + assertEquals(5, queue.currentDepth()); + } + + public void testGroupCeilingRejectsOnceTotalDepthIsReached() { + // The fixed per-group ceiling bounds the coordinator's parked footprint regardless of bucket cardinality: + // username/role bucket keys come from the request principal, so a purely per-bucket cap would let unbounded + // distinct principals each allocate size_per_bucket slots. One request per bucket isolates the ceiling from the + // per-bucket cap. + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + for (int i = 0; i < WorkloadGroupQueueSettings.MAX_GROUP_QUEUE_DEPTH; i++) { + assertTrue(queue.offer(req("g:username:u" + i), 1)); + } + assertEquals(WorkloadGroupQueueSettings.MAX_GROUP_QUEUE_DEPTH, queue.currentDepth()); + // A brand-new bucket is under its own cap but the group total is at the ceiling -> rejected. + assertFalse(queue.offer(req("g:username:overflow"), 1)); + assertEquals(WorkloadGroupQueueSettings.MAX_GROUP_QUEUE_DEPTH, queue.currentDepth()); + } + + public void testGroupCeilingRejectionLeavesNoEmptyBucket() { + // offer() reads the per-bucket depth before reserving the group counter, so a group-ceiling rejection must not + // leave an empty bucket set behind — a present bucket key means "has a live waiter", which the drain paths rely + // on (hasWaiter is a bucketKeys() membership check). + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + for (int i = 0; i < WorkloadGroupQueueSettings.MAX_GROUP_QUEUE_DEPTH; i++) { + assertTrue(queue.offer(req("g:username:u" + i), 1)); + } + assertFalse(queue.offer(req("g:username:overflow"), 1)); + assertFalse(queue.bucketKeys().contains("g:username:overflow")); + } + + public void testBucketRejectionLeavesNoEmptyBucketAndDoesNotConsumeGroupBudget() { + // A per-bucket rejection happens before the group counter is touched, so it neither inflates depth/peak nor + // creates a bucket entry for a request that was never parked. + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + assertTrue(queue.offer(req("g:username:alice"), 1)); + assertFalse(queue.offer(req("g:username:alice"), 1)); // alice's bucket full + assertEquals(1, queue.currentDepth()); + assertEquals(1L, queue.peakDepth()); // the rejected offer never reserved a slot + // A disabled queue (size_per_bucket == 0) likewise creates nothing. + assertFalse(queue.offer(req("g:username:bob"), 0)); + assertFalse(queue.bucketKeys().contains("g:username:bob")); + assertEquals(1, queue.currentDepth()); + } + + public void testPollOldestIsPerBucketFifo() { + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + WorkloadGroupQueue.QueuedRequest a1 = req("g:username:alice"); + WorkloadGroupQueue.QueuedRequest a2 = req("g:username:alice"); + assertTrue(queue.offer(a1, 10)); + assertTrue(queue.offer(a2, 10)); + assertSame(a1, queue.peekOldest("g:username:alice")); // peek does not remove + assertSame(a1, queue.pollOldest("g:username:alice")); // oldest first + assertSame(a2, queue.pollOldest("g:username:alice")); + assertNull(queue.pollOldest("g:username:alice")); + assertEquals(0, queue.currentDepth()); + } + + public void testNoHeadOfLineBlockAcrossBuckets() { + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + WorkloadGroupQueue.QueuedRequest alice = req("g:username:alice"); + WorkloadGroupQueue.QueuedRequest bob = req("g:username:bob"); + assertTrue(queue.offer(alice, 10)); + assertTrue(queue.offer(bob, 10)); + // A drain for bob's bucket returns bob even though alice was enqueued first (different bucket, no HoL blocking). + assertSame(bob, queue.pollOldest("g:username:bob")); + assertEquals(1, queue.currentDepth()); + assertSame(alice, queue.pollOldest("g:username:alice")); + } + + public void testRemoveDecrementsDepthAndPrunes() { + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + WorkloadGroupQueue.QueuedRequest r = req("g:group"); + assertTrue(queue.offer(r, 10)); + assertTrue(queue.bucketKeys().contains("g:group")); + assertTrue(queue.remove(r)); + assertEquals(0, queue.currentDepth()); + assertFalse(queue.bucketKeys().contains("g:group")); // empty bucket pruned + assertFalse(queue.remove(r)); // idempotent: already gone + } + + public void testPeakTracksHighWaterMark() { + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + assertTrue(queue.offer(req("g:group"), 10)); + assertTrue(queue.offer(req("g:group"), 10)); + assertEquals(2L, queue.peakDepth()); + queue.pollOldest("g:group"); + assertEquals(1, queue.currentDepth()); + assertEquals(2L, queue.peakDepth()); // peak does not decrease + } + + public void testEvictCancelledRemovesOnlyCancelledPreservingSurvivors() { + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + WorkloadGroupQueue.QueuedRequest survivorA = req("g:group"); // live + WorkloadGroupTask cancelledTask = task(); + WorkloadGroupQueue.QueuedRequest cancelled = req("g:group", cancelledTask); + WorkloadGroupQueue.QueuedRequest survivorB = req("g:group"); // live + assertTrue(queue.offer(survivorA, 10)); + assertTrue(queue.offer(cancelled, 10)); + assertTrue(queue.offer(survivorB, 10)); + + cancelledTask.cancel("client disconnect"); + List evicted = queue.evictCancelled("g:group"); + assertEquals(1, evicted.size()); + assertSame(cancelled, evicted.get(0)); + assertEquals(2, queue.currentDepth()); // both live survivors remain + // Survivors keep their place and identity; no time-based eviction ever removes a live request. + assertSame(survivorA, queue.pollOldest("g:group")); + assertSame(survivorB, queue.pollOldest("g:group")); + } + + public void testEvictCancelledLeavesLiveRequestsRegardlessOfAge() { + // There is no wall-clock deadline: a live parked request is never evicted by the sweep no matter how long it + // has waited. Only a cancelled task is removed. + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + WorkloadGroupQueue.QueuedRequest live = req("g:group"); + assertTrue(queue.offer(live, 10)); + assertTrue(queue.evictCancelled("g:group").isEmpty()); + assertEquals(1, queue.currentDepth()); + assertSame(live, queue.pollOldest("g:group")); + } + + public void testEvictCancelledRemovesCancelledTaskAndPrunes() { + WorkloadGroupQueue queue = new WorkloadGroupQueue(); + WorkloadGroupTask t = task(); + WorkloadGroupQueue.QueuedRequest r = req("g:group", t); + assertTrue(queue.offer(r, 10)); + t.cancel("client disconnect"); + List evicted = queue.evictCancelled("g:group"); + assertEquals(1, evicted.size()); + assertSame(r, evicted.get(0)); + assertEquals(0, queue.currentDepth()); + assertFalse(queue.bucketKeys().contains("g:group")); // pruned + } + + public void testWaitNanos() { + long enqueue = 5_000_000L; + WorkloadGroupQueue.QueuedRequest r = new WorkloadGroupQueue.QueuedRequest( + ActionListener.wrap(x -> {}, e -> {}), + "g:group", + task(), + enqueue + ); + assertEquals(2000L, r.waitNanos(enqueue + 2000)); + assertEquals(0L, r.waitNanos(enqueue)); // admitted instantly + assertEquals(0L, r.waitNanos(enqueue - 100)); // clock skew guard: never negative + } + + public void testCapturesListenerAndBucket() { + AtomicReference got = new AtomicReference<>(); + ActionListener listener = ActionListener.wrap(got::set, e -> {}); + WorkloadGroupQueue.QueuedRequest r = new WorkloadGroupQueue.QueuedRequest(listener, "g:username:alice", task(), 0L); + assertEquals("g:username:alice", r.bucketKey()); + assertSame(listener, r.listener()); + } +} diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java index 6ecf956e817ed..938d27d2e2233 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java @@ -24,6 +24,7 @@ import org.opensearch.common.logging.Loggers; import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.util.concurrent.OpenSearchExecutors; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; @@ -507,7 +508,16 @@ private void stubClusterStateWithGroup(WorkloadGroup wg) { 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)); + // acquireThrottlePermit takes the task (it carries the workload group id and is observed for cancellation while + // queued). Build a SearchTask whose workload group id is the requested one via a real thread-context header + // (mockThreadPool is a Mockito mock, so use a self-contained ThreadContext here). + WorkloadGroupTask task = new SearchTask(1, "", "", () -> "", null, null); + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + if (workloadGroupId != null) { + threadContext.putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, workloadGroupId); + } + task.setWorkloadGroupId(threadContext); + service.acquireThrottlePermit(task, principal, ActionListener.wrap(permit::set, failure::set)); if (failure.get() != null) { if (failure.get() instanceof RuntimeException) { throw (RuntimeException) failure.get(); @@ -530,6 +540,130 @@ private WorkloadGroup throttledGroup(String id, Settings throttling, MutableWork ); } + // Same as throttledGroup but with queueing enabled, so a throttle denial can park instead of rejecting. + private WorkloadGroup queueingGroup(String id, Settings throttling, Settings queue, MutableWorkloadGroupFragment.ResiliencyMode mode) { + return new WorkloadGroup( + id + "-name", + id, + new MutableWorkloadGroupFragment(mode, Map.of(ResourceType.MEMORY, 0.5), Settings.EMPTY, throttling, queue), + 1L + ); + } + + // Delivers a clusterChanged event whose CURRENT state holds exactly `currentGroups` (previous state holds `previous`). + private void deliverWorkloadGroupsChanged(Map previous, Map currentGroups) { + ClusterChangedEvent event = Mockito.mock(ClusterChangedEvent.class); + ClusterState previousState = Mockito.mock(ClusterState.class); + ClusterState currentState = Mockito.mock(ClusterState.class); + Metadata previousMetadata = Mockito.mock(Metadata.class); + Metadata currentMetadata = Mockito.mock(Metadata.class); + when(event.previousState()).thenReturn(previousState); + when(event.state()).thenReturn(currentState); + when(previousState.metadata()).thenReturn(previousMetadata); + when(currentState.metadata()).thenReturn(currentMetadata); + when(previousMetadata.workloadGroups()).thenReturn(previous); + when(currentMetadata.workloadGroups()).thenReturn(currentGroups); + workloadGroupService.clusterChanged(event); + } + + public void testDisablingThrottlingImmediatelyReleasesTheQueuedBacklog() { + // Disabling throttling leaves parked requests with nothing to wait for, and an unthrottled group takes the + // no-permit fast path — so it produces no permit completions to drive a drain chain, and parked requests have no + // deadline. Without an explicit release they wait forever. Reacting to the config change must free them at once + // (the sweep is only the backstop). Note disabling throttling necessarily disables queueing in the same update: + // WorkloadGroup rejects a queue with no throttle limit, and validateMergedConfig rejects an all-unset throttling + // block, so "throttling off, queueing on" is unreachable and this is the only way to strand a backlog. + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + when(mockThreadPool.executor(ThreadPool.Names.GENERIC)).thenReturn(OpenSearchExecutors.newDirectExecutorService()); + + Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); + Settings queue = Settings.builder().put("size_per_bucket", 5).build(); + WorkloadGroup throttled = queueingGroup("wg-1", throttling, queue, MutableWorkloadGroupFragment.ResiliencyMode.ENFORCED); + stubClusterStateWithGroup(throttled); + WorkloadGroupQueueService queueService = new WorkloadGroupQueueService(mockThreadPool, mockWorkloadGroupsStateAccessor); + workloadGroupService.setQueueService(queueService); + + // One request holds the single node slot; two more park behind it. + assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); + assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); + assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); + assertEquals("two requests should be parked", 2, queueService.currentDepth("wg-1")); + + // Operator disables throttling (and therefore queueing) on the group. + WorkloadGroup unthrottled = new WorkloadGroup( + "wg-1-name", + "wg-1", + new MutableWorkloadGroupFragment(MutableWorkloadGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5)), + 2L + ); + deliverWorkloadGroupsChanged(Map.of("wg-1", throttled), Map.of("wg-1", unthrottled)); + + assertEquals("the whole backlog must be released as soon as throttling is disabled", 0, queueService.currentDepth("wg-1")); + } + + public void testCompletionDrainHonoursALiveNodeLimitDecrease() { + // Regression: the node-tier drain chain used to capture node_limit when the chain started and reuse it for every + // subsequent hop. A busy bucket's chain runs one hop per request completion, so it can outlive a live node_limit + // update — and reusing the captured (higher) value kept admitting above the NEW lower limit until the chain broke. + // The drain must re-read node_limit from cluster state instead. + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings queue = Settings.builder().put("size_per_bucket", 5).build(); + + // Start at node_limit=2 and fill both slots, then park a third request. + Settings limit2 = Settings.builder().put("attribute", "group").put("node_limit", 2).build(); + stubClusterStateWithGroup(queueingGroup("wg-1", limit2, queue, MutableWorkloadGroupFragment.ResiliencyMode.ENFORCED)); + // This is the only test here that actually admits FROM the queue, so it needs a real executor: admit() + // dispatches the parked listener off the caller thread. Direct (same-thread) keeps the assertions synchronous — + // depth is decremented under the bucket lock before admit() runs, so inline dispatch does not affect what we assert. + when(mockThreadPool.executor(ThreadPool.Names.GENERIC)).thenReturn(OpenSearchExecutors.newDirectExecutorService()); + WorkloadGroupQueueService queueService = new WorkloadGroupQueueService(mockThreadPool, mockWorkloadGroupsStateAccessor); + workloadGroupService.setQueueService(queueService); + + Releasable first = acquireThrottlePermitSync(workloadGroupService, "wg-1", null); + Releasable second = acquireThrottlePermitSync(workloadGroupService, "wg-1", null); + assertNotNull(first); + assertNotNull(second); + // The 3rd breaches node_limit=2 with no shared tier configured, so it parks. + assertNull("third request should be parked, not admitted", acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); + assertEquals(1, queueService.currentDepth("wg-1")); + + // Operator lowers node_limit to 1 while the bucket is busy with a backlog. + Settings limit1 = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); + stubClusterStateWithGroup(queueingGroup("wg-1", limit1, queue, MutableWorkloadGroupFragment.ResiliencyMode.ENFORCED)); + + // Completing one request drops in-flight to 1, which is already AT the new limit — so the drain must not admit + // the parked request. With the old captured limit of 2 it would have, over-admitting above the configured 1. + first.close(); + assertEquals("drain must respect the lowered node_limit, leaving the request parked", 1, queueService.currentDepth("wg-1")); + + // Completing the second drops in-flight to 0, leaving room under the new limit -> the parked request drains. + second.close(); + assertEquals("a slot under the new limit must still drain the backlog", 0, queueService.currentDepth("wg-1")); + } + + public void testMonitorModeNeverParksEvenWhenQueueingIsEnabled() { + // MONITOR observes and always admits — it must never park a request. Both queueing entry points are gated on + // monitorMode == false; without that guard a would-be-throttled monitor request would sit in the queue holding + // its listener and never be answered (monitor mode is supposed to be a dry run, so this would be a hang, not a + // rejection). Note a parked request and an admitted monitor request BOTH surface as a null permit here, so the + // queue depth is what actually distinguishes them. + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); + Settings queue = Settings.builder().put("size_per_bucket", 5).build(); + stubClusterStateWithGroup(queueingGroup("wg-1", throttling, queue, MutableWorkloadGroupFragment.ResiliencyMode.MONITOR)); + WorkloadGroupQueueService queueService = new WorkloadGroupQueueService(mockThreadPool, mockWorkloadGroupsStateAccessor); + workloadGroupService.setQueueService(queueService); + + assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); // fills node_limit=1 + // The 2nd request breaches node_limit. Under MONITOR it is admitted untracked (null permit), never queued. + assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); + assertEquals("a monitor-mode request must never be parked", 0, queueService.currentDepth("wg-1")); + assertEquals(0, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); + } + public void testAcquireThrottleReturnsNullWhenNodeLimitUnset() { when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); @@ -888,6 +1022,57 @@ public void testMonitorModeObservesButDoesNotThrottleOnSharedTier() { assertEquals(0, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); } + public void testMonitorModeNeverParksOnSharedTierWithQueueingEnabled() { + // The shared tier takes the ENQUEUE-FIRST path, which parks the request BEFORE asking the owner for a slot. That + // path is gated on `queueSizePerBucket > 0 && monitorMode == false`; this pins the monitorMode half. Without it a + // monitor-mode request would be parked before the acquire even happens — i.e. MONITOR would silently start + // holding requests instead of being a dry run. The node-tier equivalent is guarded separately inside + // onThrottleBreach, so this is the case that actually covers the enqueue-first gate. + 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(); + Settings queue = Settings.builder().put("size_per_bucket", 5).build(); + WorkloadGroup group = queueingGroup("wg-1", throttling, queue, MutableWorkloadGroupFragment.ResiliencyMode.MONITOR); + + 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) + ); + ClusterState previous = Mockito.mock(ClusterState.class); + when(previous.nodes()).thenReturn(DiscoveryNodes.EMPTY_NODES); + sharedService.clusterChanged(new ClusterChangedEvent("test", clusterState, previous)); + workloadGroupService.setSharedThrottleService(sharedService); + WorkloadGroupQueueService queueService = new WorkloadGroupQueueService(mockThreadPool, mockWorkloadGroupsStateAccessor); + workloadGroupService.setQueueService(queueService); + + // 1st fills the local slot, 2nd fills the shared slot. + assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); + assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); + // 3rd overflows to the shared tier at its limit. Under MONITOR it must be admitted untracked, never parked. + assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); + assertEquals("a monitor-mode request must never be parked on the shared tier", 0, queueService.currentDepth("wg-1")); + assertEquals(0, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); + } + 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 index 06d7a6abd5521..c2bb8e149902a 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java @@ -10,12 +10,15 @@ import org.opensearch.cluster.ClusterChangedEvent; 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.io.stream.BytesStreamOutput; import org.opensearch.common.lease.Releasable; +import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; @@ -26,11 +29,14 @@ import org.opensearch.transport.TransportService; import java.io.IOException; +import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import org.mockito.Mockito; @@ -43,6 +49,7 @@ public class WorkloadGroupSharedThrottleServiceTests extends OpenSearchTestCase private ThreadPool threadPool; private TransportService transportService; private DiscoveryNode localNode; + private Metadata metadata; @Override public void setUp() throws Exception { @@ -62,6 +69,11 @@ public void setUp() throws Exception { ClusterState state = Mockito.mock(ClusterState.class); DiscoveryNodes singleDataNode = DiscoveryNodes.builder().add(localNode).localNodeId("local").build(); when(state.nodes()).thenReturn(singleDataNode); + // The TTL sweep resolves a bucket's live shared_limit from cluster-state metadata, so every test needs a + // metadata stub. Default to "no workload groups"; tests that need a real limit re-stub workloadGroups(). + metadata = Mockito.mock(Metadata.class); + when(metadata.workloadGroups()).thenReturn(Map.of()); + when(state.metadata()).thenReturn(metadata); when(clusterService.state()).thenReturn(state); when(clusterService.localNode()).thenReturn(localNode); } @@ -94,6 +106,65 @@ private static Releasable awaitGrant(WorkloadGroupSharedThrottleService service, return permit.get(); } + public void testTtlSweepDrivesOwnerPushForALeaseExpiredSlot() { + // Regression: a permit reclaimed by the TTL sweep frees a shared slot with NO release RPC behind it (the holder + // crashed, or its release was lost), so the sweep is the ONLY observer of that free slot. It must drive + // owner-push. Previously sweepExpired() returned void and the sweep ignored the freed capacity, so a coordinator + // with a parked request stayed registered as a waiter while the slot sat idle — and since parked requests have no + // deadline, it stranded until some unrelated release happened to re-drive the bucket. + final int sharedLimit = 1; + final String groupId = "g1"; + final String bucket = groupId + ":group"; + + // Controllable clock so the permit can be expired deterministically instead of sleeping past the 5-minute TTL. + final AtomicLong nanos = new AtomicLong(0L); + WorkloadGroupSharedThrottleService service = new WorkloadGroupSharedThrottleService( + clusterService, + threadPool, + transportService, + new SharedThrottleTracker(nanos::get) + ); + deliverNodesChanged(service, clusterService.state().nodes()); + + // The sweep resolves shared_limit from cluster state (an expiry carries no limit, unlike a release). + Settings throttling = Settings.builder().put("attribute", "group").put("shared_limit", sharedLimit).build(); + WorkloadGroup group = new WorkloadGroup( + "g1-name", + groupId, + new MutableWorkloadGroupFragment( + MutableWorkloadGroupFragment.ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + throttling + ), + 1L + ); + when(metadata.workloadGroups()).thenReturn(Map.of(groupId, group)); + + final AtomicInteger admits = new AtomicInteger(0); + service.setGrantConsumer((bucketKey, reservedPermit) -> { + admits.incrementAndGet(); + return true; // stands in for the queue service admitting one parked request + }); + + // Take the only shared slot, then deny an acquire with wantsQueue=true so this coordinator is a registered waiter + // with a parked request. Both happen at t=0, while the holder's permit is still live. + assertNotNull("first acquire takes the only shared slot", awaitGrant(service, bucket, sharedLimit)); + AtomicReference denial = new AtomicReference<>(); + service.acquireAsync(bucket, sharedLimit, true, ActionListener.wrap(p -> fail("must be denied while at limit"), denial::set)); + assertTrue("acquire at limit must be denied", denial.get() instanceof OpenSearchRejectedExecutionException); + assertEquals("coordinator is registered as a waiter", 1, service.waiterCountForTest(bucket)); + + // Simulate the holder vanishing: advance past the permit TTL WITHOUT any release RPC. Nothing has pruned the + // permit yet (no acquire has touched the bucket), so the slot is expired-but-unreclaimed. + nanos.set(WorkloadGroupSharedThrottleService.PERMIT_TTL_NANOS + 1); + assertEquals("nothing can have been admitted before the sweep runs", 0, admits.get()); + + // One sweep pass must reclaim the expired permit AND hand the freed slot to the waiting coordinator. + service.sweepExpiredAndDrive(); + assertEquals("the sweep must drive owner-push for the slot it freed", 1, admits.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 @@ -157,6 +228,63 @@ public void testLocalOwnerGrantsThenDeniesAtLimit() { assertNotNull(awaitGrant(service, "b", 1)); } + public void testOwnerPushDrainsEveryParkedRequestOnOneCoordinator() { + // Regression: one coordinator (here the local owner) parks SEVERAL requests for a shared bucket, but the owner + // waiter registry holds one Set membership per coordinator. A grant must NOT deregister the coordinator on a + // successful admit — it may still have more queued requests — so each successive freed slot drains the next + // parked request. The bug drained only the first and stranded the rest until queue.timeout despite free capacity. + final int sharedLimit = 1; + final String bucket = "b"; + WorkloadGroupSharedThrottleService service = newService(); + + // A stubbed coordinator-side consumer standing in for the queue service: it holds `parked` requests and admits + // one per grant, capturing the reserved permit so the test can "complete" that request by closing it (which + // re-drives owner-push, exactly like a real request finishing). + final AtomicInteger parked = new AtomicInteger(3); + final AtomicInteger admits = new AtomicInteger(0); + final List heldPermits = new ArrayList<>(); + service.setGrantConsumer((bucketKey, reservedPermit) -> { + if (parked.get() <= 0) { + return false; // nothing left to admit -> caller returns the unused grant and deregisters this waiter + } + parked.decrementAndGet(); + admits.incrementAndGet(); + heldPermits.add(reservedPermit); + return true; + }); + + // Fill the single shared slot, then issue 3 denied acquires with wantsQueue=true. Each denial registers this + // (local) coordinator as a waiter — idempotently, so the Set holds exactly ONE membership for 3 parked requests. + Releasable slotHolder = awaitGrant(service, bucket, sharedLimit); + assertNotNull("first acquire takes the only shared slot", slotHolder); + for (int i = 0; i < 3; i++) { + AtomicReference denial = new AtomicReference<>(); + service.acquireAsync(bucket, sharedLimit, true, ActionListener.wrap(p -> fail("must be denied while at limit"), denial::set)); + assertTrue("acquire at limit must be denied (429)", denial.get() instanceof OpenSearchRejectedExecutionException); + } + assertEquals("one Set membership for the coordinator regardless of parked count", 1, service.waiterCountForTest(bucket)); + + // Release the in-flight slot -> owner-push admits the FIRST parked request and the coordinator stays registered. + slotHolder.close(); + assertEquals("first freed slot drains exactly one parked request", 1, admits.get()); + assertEquals("coordinator must remain registered while it still has parked requests", 1, service.waiterCountForTest(bucket)); + + // Each admitted request completing frees the slot again and must drain the NEXT parked request. + heldPermits.remove(0).close(); + assertEquals("second freed slot drains the second parked request", 2, admits.get()); + assertEquals(1, service.waiterCountForTest(bucket)); + + heldPermits.remove(0).close(); + assertEquals("third freed slot drains the third (last) parked request", 3, admits.get()); + + // The last request completes with nothing left queued: the next grant comes back unused, so the coordinator + // self-deregisters and the freed slot returns to the pool. No stranding, no leaked permit. + heldPermits.remove(0).close(); + assertEquals("no further admits once the queue is empty", 3, admits.get()); + assertEquals("coordinator self-reconciles out of the registry when it has nothing queued", 0, service.waiterCountForTest(bucket)); + assertEquals("no shared permit leaked after the burst fully drains", 0, service.tracker().inFlight(bucket)); + } + public void testDoubleCloseReleasesOnce() { WorkloadGroupSharedThrottleService service = newService(); Releasable p = awaitGrant(service, "b", 2); diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupTaskTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupTaskTests.java index 341f31993f800..643b86d6e883a 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupTaskTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupTaskTests.java @@ -8,11 +8,13 @@ package org.opensearch.wlm; +import org.opensearch.common.lease.Releasable; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; import static org.opensearch.wlm.WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER; import static org.opensearch.wlm.WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER; @@ -41,4 +43,31 @@ public void testSuccessfulSetWorkloadGroupId() { sut.setWorkloadGroupId(threadPool.getThreadContext()); assertEquals("akfanglkaglknag2332", sut.getWorkloadGroupId()); } + + public void testOnCancelledCallbackFiresOnceOnCancel() { + AtomicInteger fired = new AtomicInteger(0); + sut.addOnCancelledCallback(fired::incrementAndGet); + assertEquals(0, fired.get()); + sut.cancel("test"); + assertEquals(1, fired.get()); + // A second cancel must not re-fire the callback. + sut.cancel("again"); + assertEquals(1, fired.get()); + } + + public void testOnCancelledCallbackRunsImmediatelyIfAlreadyCancelled() { + sut.cancel("test"); + AtomicInteger fired = new AtomicInteger(0); + // Registering after cancellation runs the callback immediately (no lost-cancellation race). + sut.addOnCancelledCallback(fired::incrementAndGet); + assertEquals(1, fired.get()); + } + + public void testDeregisteredCallbackDoesNotFire() { + AtomicInteger fired = new AtomicInteger(0); + Releasable handle = sut.addOnCancelledCallback(fired::incrementAndGet); + handle.close(); // deregister (request admitted/drained before any cancellation) + sut.cancel("test"); + assertEquals(0, fired.get()); + } } diff --git a/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java b/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java index 9e1bcfedd6999..86651c19d3b2f 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java @@ -50,7 +50,7 @@ public void testToXContent() throws IOException { wlmStats.toXContent(builder, ToXContent.EMPTY_PARAMS); builder.endObject(); assertEquals( - "{\"workload_groups\":{\"afakjklaj304041-afaka\":{\"total_completions\":123456789,\"total_rejections\":13,\"total_cancellations\":0,\"total_throttled\":5,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", + "{\"workload_groups\":{\"afakjklaj304041-afaka\":{\"total_completions\":123456789,\"total_rejections\":13,\"total_cancellations\":0,\"total_throttled\":5,\"total_queued\":0,\"total_queue_rejections\":0,\"queued_current\":0,\"queue_peak\":0,\"total_queue_wait_millis\":0,\"queue_wait_count\":0,\"max_queue_wait_millis\":0,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", builder.toString() ); } diff --git a/server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStateTests.java b/server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStateTests.java index 4aa70d17064e0..80cfc7114d5f3 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStateTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStateTests.java @@ -70,4 +70,27 @@ public void testRandomWorkloadGroupsStateUpdates() { assertEquals(5, workloadGroupState.getResourceState().get(ResourceType.MEMORY).cancellations.count()); } + public void testRecordQueueWaitAggregates() { + WorkloadGroupState state = new WorkloadGroupState(); + assertEquals(0, state.getQueueWaitCount()); + assertEquals(0, state.getTotalQueueWaitMillis()); + assertEquals(0, state.getMaxQueueWaitMillis()); + + state.recordQueueWaitMillis(100); + state.recordQueueWaitMillis(300); + state.recordQueueWaitMillis(50); + + assertEquals(3, state.getQueueWaitCount()); + assertEquals(450, state.getTotalQueueWaitMillis()); // sum -> mean = 150 + assertEquals(300, state.getMaxQueueWaitMillis()); // high-water mark, does not decrease + } + + public void testRecordQueueWaitClampsNegative() { + WorkloadGroupState state = new WorkloadGroupState(); + state.recordQueueWaitMillis(-5); // clock skew guard + assertEquals(1, state.getQueueWaitCount()); + assertEquals(0, state.getTotalQueueWaitMillis()); + assertEquals(0, state.getMaxQueueWaitMillis()); + } + } diff --git a/server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStatsTests.java b/server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStatsTests.java index 356d90edce041..2c359b5821181 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStatsTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStatsTests.java @@ -8,24 +8,17 @@ package org.opensearch.wlm.stats; -import org.opensearch.Version; -import org.opensearch.cluster.node.DiscoveryNode; -import org.opensearch.cluster.node.DiscoveryNodeRole; import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.test.AbstractWireSerializingTestCase; -import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.test.VersionUtils; import org.opensearch.wlm.ResourceType; import java.io.IOException; import java.util.HashMap; import java.util.Map; -import static java.util.Collections.emptyMap; - public class WorkloadGroupStatsTests extends AbstractWireSerializingTestCase { public void testToXContent() throws IOException { @@ -48,11 +41,34 @@ public void testToXContent() throws IOException { workloadGroupStats.toXContent(builder, ToXContent.EMPTY_PARAMS); builder.endObject(); assertEquals( - "{\"workload_groups\":{\"afakjklaj304041-afaka\":{\"total_completions\":123456789,\"total_rejections\":13,\"total_cancellations\":0,\"total_throttled\":5,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", + "{\"workload_groups\":{\"afakjklaj304041-afaka\":{\"total_completions\":123456789,\"total_rejections\":13,\"total_cancellations\":0,\"total_throttled\":5,\"total_queued\":0,\"total_queue_rejections\":0,\"queued_current\":0,\"queue_peak\":0,\"total_queue_wait_millis\":0,\"queue_wait_count\":0,\"max_queue_wait_millis\":0,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", builder.toString() ); } + // The randomized createTestInstance() builds holders via the constructor, which leaves the three queue-wait fields + // at 0 (they are only populated from WorkloadGroupState via from(...)). This explicit round-trip drives them + // non-zero through the real production path so a write/read order or mapping bug among them is caught. + public void testQueueWaitFieldsSurviveWireRoundTrip() throws IOException { + WorkloadGroupState state = new WorkloadGroupState(); + state.recordQueueWaitMillis(100); + state.recordQueueWaitMillis(300); // sum=400, count=2, max=300 + WorkloadGroupStats.WorkloadGroupStatsHolder holder = WorkloadGroupStats.WorkloadGroupStatsHolder.from(state, 7L, 9L); + assertEquals(400L, holder.getTotalQueueWaitMillis()); + assertEquals(2L, holder.getQueueWaitCount()); + assertEquals(300L, holder.getMaxQueueWaitMillis()); + + WorkloadGroupStats original = new WorkloadGroupStats(Map.of("g", holder)); + WorkloadGroupStats roundTripped = copyWriteable(original, writableRegistry(), WorkloadGroupStats::new); + WorkloadGroupStats.WorkloadGroupStatsHolder rt = roundTripped.getStats().get("g"); + assertEquals(400L, rt.getTotalQueueWaitMillis()); + assertEquals(2L, rt.getQueueWaitCount()); + assertEquals(300L, rt.getMaxQueueWaitMillis()); + assertEquals(7L, rt.getQueuedCurrent()); + assertEquals(9L, rt.getQueuePeak()); + assertEquals(original, roundTripped); + } + @Override protected Writeable.Reader instanceReader() { return WorkloadGroupStats::new; @@ -60,31 +76,57 @@ protected Writeable.Reader instanceReader() { @Override protected WorkloadGroupStats createTestInstance() { - Map stats = new HashMap<>(); - stats.put( - randomAlphaOfLength(10), - new WorkloadGroupStats.WorkloadGroupStatsHolder( - randomNonNegativeLong(), - randomNonNegativeLong(), - randomNonNegativeLong(), - randomNonNegativeLong(), - randomNonNegativeLong(), - Map.of( - ResourceType.CPU, - new WorkloadGroupStats.ResourceStats( - randomDoubleBetween(0.0, 0.90, false), - randomNonNegativeLong(), - randomNonNegativeLong() - ) + return new WorkloadGroupStats(Map.of(randomAlphaOfLength(10), randomStatsHolder())); + } + + // Uses the full constructor with a random value for EVERY field — including all four queue stats — so the wire + // round-trip (testSerialization) and equals/hashCode (testEqualsAndHashcode) actually exercise the new fields. + // A previous version used the 6-arg ctor, leaving the queue fields hard-zeroed, so a write/read order swap among + // them would have been invisible. + private static WorkloadGroupStats.WorkloadGroupStatsHolder randomStatsHolder() { + return new WorkloadGroupStats.WorkloadGroupStatsHolder( + randomNonNegativeLong(), // completions + randomNonNegativeLong(), // rejections + randomNonNegativeLong(), // failures + randomNonNegativeLong(), // cancellations + randomNonNegativeLong(), // throttled + randomNonNegativeLong(), // queued + randomNonNegativeLong(), // queueRejections + randomNonNegativeLong(), // queuedCurrent + randomNonNegativeLong(), // queuePeak + Map.of( + ResourceType.CPU, + new WorkloadGroupStats.ResourceStats( + randomDoubleBetween(0.0, 0.90, false), + randomNonNegativeLong(), + randomNonNegativeLong() ) ) ); - DiscoveryNode discoveryNode = new DiscoveryNode( - "node", - OpenSearchTestCase.buildNewFakeTransportAddress(), - emptyMap(), - DiscoveryNodeRole.BUILT_IN_ROLES, - VersionUtils.randomCompatibleVersion(random(), Version.CURRENT) + } + + @Override + protected WorkloadGroupStats mutateInstance(WorkloadGroupStats instance) { + // Flip exactly one queue field on one holder so the mutation-inequality check actually asserts that queue + // fields participate in equals/hashCode and, via the round-trip, that their wire order is honored. + Map stats = new HashMap<>(instance.getStats()); + String key = stats.isEmpty() ? randomAlphaOfLength(10) : stats.keySet().iterator().next(); + WorkloadGroupStats.WorkloadGroupStatsHolder h = stats.get(key); + long bump = h == null ? 1 : h.getQueuedCurrent() + 1; + stats.put( + key, + new WorkloadGroupStats.WorkloadGroupStatsHolder( + h == null ? 0 : h.getCompletions(), + h == null ? 0 : h.getRejections(), + 0, + h == null ? 0 : h.getCancellations(), + h == null ? 0 : h.getThrottled(), + h == null ? 0 : h.getQueued(), + h == null ? 0 : h.getQueueRejections(), + bump, // mutated field: queuedCurrent + h == null ? 0 : h.getQueuePeak(), + h == null ? Map.of() : h.getResourceStats() + ) ); return new WorkloadGroupStats(stats); } From b278658eae74132dba8cb431429b44c55b52a2bd Mon Sep 17 00:00:00 2001 From: David Zane Date: Fri, 28 Aug 2026 16:54:06 -0700 Subject: [PATCH 2/2] Let the permit owner decide when to drive owner-push Signed-off-by: David Zane --- .../opensearch/wlm/SharedThrottleTracker.java | 12 +- .../WorkloadGroupSharedThrottleService.java | 112 +++++++++--------- ...rkloadGroupSharedThrottleServiceTests.java | 22 ++++ 3 files changed, 90 insertions(+), 56 deletions(-) diff --git a/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java b/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java index 054bd93acf07f..a021b8f693ff2 100644 --- a/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java +++ b/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java @@ -126,15 +126,23 @@ public boolean tryAcquire(String bucketKey, int sharedLimit, String permitId, lo * * @param bucketKey the throttle bucket identifier * @param permitId the permit id returned to the coordinator at acquire time + * @return {@code true} if a live permit was actually removed, i.e. a slot genuinely just freed. This is the + * authoritative answer to "did capacity become available?", and the owner is the only node that can give it: + * a coordinator sending a speculative release (e.g. after a lost acquire reply) cannot know whether its + * permit was ever recorded. Callers drive owner-push on {@code true} only, so a no-op release never produces + * a phantom grant and — the case a coordinator-supplied hint used to get wrong — a release that DID free a + * slot always drives one. */ - public void release(String bucketKey, String permitId) { + public boolean release(String bucketKey, String permitId) { + final boolean[] removed = new boolean[1]; permitsByBucket.computeIfPresent(bucketKey, (k, bucket) -> { - bucket.permits.remove(permitId); + removed[0] = bucket.permits.remove(permitId) != null; // minExpiry is intentionally left unchanged: removing a permit can only raise the true earliest expiry, so // the cached value stays a valid (possibly loose) lower bound. Recomputing here would add an O(size) scan // to the release hot path for no correctness benefit. return bucket.permits.isEmpty() ? null : bucket; }); + return removed[0]; } /** diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java index 928055f2233cb..21b99a0aabf5f 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java @@ -153,15 +153,7 @@ public WorkloadGroupSharedThrottleService( ThreadPool.Names.SAME, ReleasePermitRequest::new, (request, channel, task) -> { - tracker.release(request.bucketKey, request.permitId); - // If this release is a coordinator returning an UNUSED grant (it had no queued request for the bucket), - // drop it from the waiter set so it stops drawing wasted grants (remote analog of the local removeWaiter). - if (request.queueEmptyOnNodeId.isEmpty() == false) { - removeWaiterByNodeId(request.bucketKey, request.queueEmptyOnNodeId); - } - // A slot just freed on the owner. If a coordinator is waiting on this bucket, hand the freed slot to one - // of them (reserve-then-grant) so its queued request can drain — owner-push. No-op if no waiters. - onSharedSlotFreed(request.bucketKey, request.sharedLimit); + handleRelease(request); channel.sendResponse(TransportResponse.Empty.INSTANCE); } ); @@ -252,7 +244,7 @@ public void acquireAsync(String bucketKey, int sharedLimit, boolean wantsQueue, // 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, permitId, ttlNanos)) { - listener.onResponse(releaseLocal(bucketKey, sharedLimit, permitId)); + listener.onResponse(releaseLocal(bucketKey, permitId)); } else { if (wantsQueue) { registerWaiter(bucketKey, clusterService.localNode()); @@ -286,7 +278,7 @@ public AcquirePermitResponse read(StreamInput in) throws IOException { @Override public void handleResponse(AcquirePermitResponse response) { if (response.granted) { - listener.onResponse(releaseRemote(owner, bucketKey, sharedLimit, permitId)); + listener.onResponse(releaseRemote(owner, bucketKey, permitId)); } else { listener.onFailure(deniedMarker()); } @@ -302,7 +294,7 @@ public void handleException(TransportException exp) { logger.debug("Shared throttle acquire to owner [{}] for bucket [{}] failed; failing open", owner.getId(), bucketKey); // Reclaim only (no owner-push): this permit likely never existed; UNSET_LIMIT tells the owner to // release without driving a grant. - sendRelease(owner, bucketKey, WorkloadGroupThrottleSettings.UNSET_LIMIT, permitId, ""); + sendRelease(owner, bucketKey, permitId, ""); listener.onResponse(null); } @@ -339,22 +331,24 @@ AcquirePermitResponse handleAcquire(AcquirePermitRequest request) { // A local release path must also drive owner-push when this node owns the bucket: releasing frees a slot that a // registered waiter should get. sharedLimit is threaded so the reserve step can re-check the limit. - private Releasable releaseLocal(String bucketKey, int sharedLimit, String permitId) { + private Releasable releaseLocal(String bucketKey, String permitId) { return releaseOnce(() -> { - tracker.release(bucketKey, permitId); - onSharedSlotFreed(bucketKey, sharedLimit); + // Same rule as the remote path's handleRelease, so local-owner and remote-owner behave identically. + if (tracker.release(bucketKey, permitId)) { + onSharedSlotFreed(bucketKey, currentSharedLimit(bucketKey)); + } }); } - private Releasable releaseRemote(DiscoveryNode owner, String bucketKey, int sharedLimit, String permitId) { + private Releasable releaseRemote(DiscoveryNode owner, String bucketKey, String permitId) { // The remote RELEASE RPC carries sharedLimit so the owner can drive owner-push after freeing the slot. - return releaseOnce(() -> sendRelease(owner, bucketKey, sharedLimit, permitId, "")); + return releaseOnce(() -> sendRelease(owner, bucketKey, permitId, "")); } // Returns an UNUSED remote grant: a RELEASE tagged with this coordinator's node id so the owner also deregisters it // from the bucket's waiter set before re-driving owner-push to the next waiter. - private void sendReleaseUnusedGrant(DiscoveryNode owner, String bucketKey, int sharedLimit, String permitId, DiscoveryNode self) { - sendRelease(owner, bucketKey, sharedLimit, permitId, self.getId()); + private void sendReleaseUnusedGrant(DiscoveryNode owner, String bucketKey, String permitId, DiscoveryNode self) { + sendRelease(owner, bucketKey, permitId, self.getId()); } // Fire-and-forget RELEASE RPC to the bucket owner. Bounded by the same timeout as acquire so a half-open @@ -362,12 +356,12 @@ private void sendReleaseUnusedGrant(DiscoveryNode owner, String bucketKey, int s // fatal — the owner's TTL sweep reclaims the permit — so failures are logged at debug only. Carries sharedLimit so // the owner can drive owner-push (grant a waiter the freed slot) after releasing. {@code queueEmptyOnNodeId} is set // only when returning an unused grant, so the owner deregisters that coordinator; empty for a normal release. - private void sendRelease(DiscoveryNode owner, String bucketKey, int sharedLimit, String permitId, String queueEmptyOnNodeId) { + private void sendRelease(DiscoveryNode owner, String bucketKey, String permitId, String queueEmptyOnNodeId) { final TransportRequestOptions options = TransportRequestOptions.builder().withTimeout(ACQUIRE_TIMEOUT).build(); transportService.sendRequest( owner, RELEASE_ACTION_NAME, - new ReleasePermitRequest(bucketKey, permitId, sharedLimit, queueEmptyOnNodeId), + new ReleasePermitRequest(bucketKey, permitId, queueEmptyOnNodeId), options, new TransportResponseHandler() { @Override @@ -431,6 +425,29 @@ private void registerWaiter(String bucketKey, DiscoveryNode coordinator) { }); } + /** + * OWNER-SIDE release: free the permit, deregister the coordinator if it reported its queue for the bucket is empty, + * and drive owner-push if a slot genuinely freed. Package-private and shared with the transport handler rather than + * duplicated, so a test can drive it without the two copies drifting apart. + *

    + * Both decisions are the OWNER's, not the coordinator's. Whether a slot freed comes from whether the remove actually + * hit a live permit — a coordinator's release may be speculative (see the lost acquire-reply path), so it cannot + * know. The ceiling is resolved from this node's cluster state, whose view is the one being enforced. Deciding here + * means a no-op release never produces a phantom grant and — the case a coordinator-supplied hint got wrong — a + * release that DID free a slot always drives one. + */ + void handleRelease(ReleasePermitRequest request) { + final boolean freed = tracker.release(request.bucketKey, request.permitId); + // If this release is a coordinator returning an UNUSED grant (it had no queued request for the bucket), drop it + // from the waiter set so it stops drawing wasted grants (remote analog of the local removeWaiter). + if (request.queueEmptyOnNodeId.isEmpty() == false) { + removeWaiterByNodeId(request.bucketKey, request.queueEmptyOnNodeId); + } + if (freed) { + onSharedSlotFreed(request.bucketKey, currentSharedLimit(request.bucketKey)); + } + } + // OWNER-SIDE: drop a coordinator from a bucket's waiter set (it reported no more queued requests for the bucket, // via an unused grant). Prunes the bucket entry when the last waiter leaves. private void removeWaiter(String bucketKey, DiscoveryNode coordinator) { @@ -667,7 +684,7 @@ private void handleGrant(GrantPermitRequest request) { private void consumeGrant(String bucketKey, int sharedLimit, String reservedPermitId) { final GrantConsumer consumer = grantConsumer; if (consumer == null) { - returnUnusedGrant(bucketKey, sharedLimit, reservedPermitId); // not wired yet -> return the slot + deregister + returnUnusedGrant(bucketKey, reservedPermitId); // not wired yet -> return the slot + deregister return; } // The permit handed to a successfully-admitted request releases only the reserved permit on completion (which @@ -679,31 +696,35 @@ private void consumeGrant(String bucketKey, int sharedLimit, String reservedPerm admitted = consumer.admit(bucketKey, reservedPermit(owner, bucketKey, sharedLimit, reservedPermitId)); } catch (Exception e) { logger.warn("Queue grant admit failed for bucket [" + bucketKey + "]", e); - returnUnusedGrant(bucketKey, sharedLimit, reservedPermitId); + returnUnusedGrant(bucketKey, reservedPermitId); return; } if (admitted == false) { // No queued request on this coordinator for the bucket: return the reserved slot AND tell the owner to // deregister this coordinator so it stops drawing wasted grants. - returnUnusedGrant(bucketKey, sharedLimit, reservedPermitId); + returnUnusedGrant(bucketKey, reservedPermitId); } } // Returns an unused reserved slot to its owner, tagging the release so the owner deregisters this coordinator from // the bucket's waiter set (it has no queued request for the bucket) and then re-drives owner-push to the next // waiter. If this node is the owner, does it in-process. - private void returnUnusedGrant(String bucketKey, int sharedLimit, String reservedPermitId) { + private void returnUnusedGrant(String bucketKey, String reservedPermitId) { final DiscoveryNode owner = ring.get().ownerFor(bucketKey).orElse(null); final DiscoveryNode self = clusterService.localNode(); if (owner == null) { return; // ring empty; the reserved permit (if any) is reclaimed by TTL } if (owner.getId().equals(self.getId())) { - tracker.release(bucketKey, reservedPermitId); + // Same rules as the remote path's handleRelease, so local-owner and remote-owner behave identically: push is + // driven only if a permit really went away, against this node's own view of the ceiling. + final boolean freed = tracker.release(bucketKey, reservedPermitId); removeWaiter(bucketKey, self); - onSharedSlotFreed(bucketKey, sharedLimit); + if (freed) { + onSharedSlotFreed(bucketKey, currentSharedLimit(bucketKey)); + } } else { - sendReleaseUnusedGrant(owner, bucketKey, sharedLimit, reservedPermitId, self); + sendReleaseUnusedGrant(owner, bucketKey, reservedPermitId, self); } } @@ -715,9 +736,9 @@ private Releasable reservedPermit(DiscoveryNode owner, String bucketKey, int sha return releaseOnce(() -> {}); // ring empty; nothing to release remotely } if (owner.getId().equals(clusterService.localNode().getId())) { - return releaseLocal(bucketKey, sharedLimit, permitId); + return releaseLocal(bucketKey, permitId); } - return releaseRemote(owner, bucketKey, sharedLimit, permitId); + return releaseRemote(owner, bucketKey, permitId); } // Message-less "denied" marker: the bucket is at its shared limit. Carries no user text because this service has @@ -961,17 +982,10 @@ public void writeTo(StreamOutput out) throws IOException { public static class ReleasePermitRequest extends TransportRequest { static final String KEY_BUCKET = "bucket_key"; static final String KEY_PERMIT_ID = "permit_id"; - static final String KEY_SHARED_LIMIT = "shared_limit"; static final String KEY_QUEUE_EMPTY_ON_NODE_ID = "queue_empty_on_node_id"; final String bucketKey; final String permitId; - /** - * The bucket's shared limit, so the owner can drive owner-push (grant a waiter) after freeing the slot. - * {@code UNSET_LIMIT} means "release only, do not drive push" — which is also the default when an older peer - * omits the key, so such a release behaves exactly as it did before queueing existed. - */ - final int sharedLimit; /** * Set to the coordinator's PERSISTENT node id only when this release is returning an UNUSED grant: the owner then * deregisters that coordinator from the bucket's waiter set, since it has no queued request. Empty for a normal @@ -982,15 +996,14 @@ public static class ReleasePermitRequest extends TransportRequest { */ final String queueEmptyOnNodeId; - // Convenience for a normal release that does not drive owner-push. + // Convenience for a normal release (not returning an unused grant). ReleasePermitRequest(String bucketKey, String permitId) { - this(bucketKey, permitId, WorkloadGroupThrottleSettings.UNSET_LIMIT, ""); + this(bucketKey, permitId, ""); } - ReleasePermitRequest(String bucketKey, String permitId, int sharedLimit, String queueEmptyOnNodeId) { + ReleasePermitRequest(String bucketKey, String permitId, String queueEmptyOnNodeId) { this.bucketKey = bucketKey; this.permitId = permitId; - this.sharedLimit = sharedLimit; this.queueEmptyOnNodeId = queueEmptyOnNodeId == null ? "" : queueEmptyOnNodeId; } @@ -999,9 +1012,9 @@ public static class ReleasePermitRequest extends TransportRequest { final Map body = readBody(in); this.bucketKey = requireString(body, KEY_BUCKET); this.permitId = requireString(body, KEY_PERMIT_ID); - // Added by queueing, so optional. The defaults are the pre-queueing behaviour: no owner-push, no waiter - // deregistration. - this.sharedLimit = (int) optionalLong(body, KEY_SHARED_LIMIT, WorkloadGroupThrottleSettings.UNSET_LIMIT); + // Added by queueing, so optional; empty means "not an unused-grant return", which is the pre-queueing + // behaviour. Note the bucket's shared limit is deliberately NOT on the wire: the owner resolves it from its + // own cluster state, and decides whether to drive owner-push from whether a permit was really removed. this.queueEmptyOnNodeId = optionalString(body, KEY_QUEUE_EMPTY_ON_NODE_ID, ""); // Any other key is a field this build does not know about: ignored on purpose. That is the tolerance. } @@ -1010,16 +1023,7 @@ public static class ReleasePermitRequest extends TransportRequest { public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeMap( - Map.of( - KEY_BUCKET, - bucketKey, - KEY_PERMIT_ID, - permitId, - KEY_SHARED_LIMIT, - sharedLimit, - KEY_QUEUE_EMPTY_ON_NODE_ID, - queueEmptyOnNodeId - ), + Map.of(KEY_BUCKET, bucketKey, KEY_PERMIT_ID, permitId, KEY_QUEUE_EMPTY_ON_NODE_ID, queueEmptyOnNodeId), StreamOutput::writeString, StreamOutput::writeGenericValue ); diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java index c2bb8e149902a..252e739b34c00 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java @@ -78,6 +78,27 @@ public void setUp() throws Exception { when(clusterService.localNode()).thenReturn(localNode); } + // The owner resolves a bucket's shared_limit from its OWN cluster state (it no longer arrives on the RELEASE RPC), so + // any test that expects owner-push to be driven has to model the group. Without this the limit reads as UNSET and + // onSharedSlotFreed correctly declines to grant. + private void stubGroupFor(String bucketKey, int sharedLimit) { + final int idx = bucketKey.indexOf(':'); + final String groupId = idx < 0 ? bucketKey : bucketKey.substring(0, idx); + Settings throttling = Settings.builder().put("attribute", "group").put("shared_limit", sharedLimit).build(); + WorkloadGroup group = new WorkloadGroup( + groupId + "-name", + groupId, + new MutableWorkloadGroupFragment( + MutableWorkloadGroupFragment.ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + throttling + ), + 1L + ); + when(metadata.workloadGroups()).thenReturn(Map.of(groupId, group)); + } + 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); @@ -236,6 +257,7 @@ public void testOwnerPushDrainsEveryParkedRequestOnOneCoordinator() { final int sharedLimit = 1; final String bucket = "b"; WorkloadGroupSharedThrottleService service = newService(); + stubGroupFor(bucket, sharedLimit); // A stubbed coordinator-side consumer standing in for the queue service: it holds `parked` requests and admits // one per grant, capturing the reserved permit so the test can "complete" that request by closing it (which