From f33d424ad6ebf1fcb9914adfb09bf9c539f6b960 Mon Sep 17 00:00:00 2001 From: David Zane Date: Wed, 26 Aug 2026 12:25:25 -0700 Subject: [PATCH] Add per-node request throttling to workload groups A workload group can now cap the number of search requests a single node admits for it concurrently. Overflow is rejected with a 429 instead of queueing up behind the resource-based cancellation path. Throttle config is a new nested "throttling" object on a workload group, modeled as a Settings bag like the existing "settings" field, so a per-key null clears one field and an absent key keeps its current value. "attribute" selects the dimension the limit is keyed by: group, or username/role for one bucket per principal, resolved from the security plugin's principal.* attributes at enforcement time. "node_limit" is the per-node concurrent allowance for each bucket. The caller's principal for username/role keying is carried on the coordinator's task object rather than in a ThreadContext request header, so a client cannot supply it, it is never serialized onto shard requests or to remote clusters in a cross-cluster search, and the concurrent _msearch sub-requests that share one thread context each get their own value. Admission runs in TransportSearchAction before onRequestStart, so a rejection never leaves the request-operations gauges incremented. The permit is chained onto the response listener and released exactly once, including on the transformRequest failure path. Rejections are counted in a new total_throttled field on _wlm/stats. A group in monitor resiliency mode observes only: a breach is logged at DEBUG and the request is admitted, matching how MONITOR is dormant on the cancellation path, and it is not counted. The throttle path fails open, so an unattributable request, an unregistered group, or an unexpected error skips throttling rather than failing a valid search. Validation rejects a limit with no attribute, an unknown throttling key, a negative or int-overflowing limit, and a config whose effective ceiling is 0. Setting "throttling" to null disables throttling. The throttling field is gated behind 3.9.0 on the wire, separately from the already-released settings gate, so mixed-version clusters stay wire-compatible. The cluster-wide shared tier is not included here; shared_limit is not an accepted key, so it cannot be set with no effect. Scroll continuations draw on the same budget: exempting them would make node_limit evadable by appending ?scroll= to a query. Each _msearch sub-search takes its own permit, so a throttled sub-search reports 429 inside the enclosing 200 response. Query rewriting can issue a nested coordinator search on the same node -- a terms lookup with a subquery does -- while the outer request already holds its bucket's permit, and the nested request inherits the same workload group id, so it resolves to the same bucket. Charging it a second permit made the request compete with itself: with node_limit=N, N such requests all got a 429 at exactly the configured concurrency. Admission now skips a request whose bucket an ancestor task already holds. To make that visible, searches issued by the rewrite phase are parented on the task that triggered the rewrite, via new overloads of SearchService and IndicesService getRewriteContext. That parenting is applied only when the request actually holds a permit, so a search in a group without throttling issues rewrite requests exactly as before; a throttled request's rewrite searches do now report a parent action to system-generated search pipeline selection, where before they reported none. Creating or updating a throttling config is rejected when the cluster still has a pre-3.9 node, instead of returning 200 and silently dropping the field on the wire, and when the attribute keys on a principal but no principal attribute provider is installed, which could never enforce. Both are checked in the cluster-manager transport actions rather than in a state applier, because throwing while applying state wedges the cluster-manager. Signed-off-by: Emily Guo <35637792+LilyCaroline17@users.noreply.github.com> Signed-off-by: David Zane Co-authored-by: Emily Guo <35637792+LilyCaroline17@users.noreply.github.com> --- .../plugin/wlm/WlmNodeThrottlingIT.java | 672 ++++++++++++++++++ .../plugin/wlm/AutoTaggingActionFilter.java | 33 +- .../TransportCreateWorkloadGroupAction.java | 9 + .../TransportUpdateWorkloadGroupAction.java | 9 + .../WorkloadGroupPersistenceService.java | 61 ++ .../wlm/AutoTaggingActionFilterTests.java | 183 +++++ .../WorkloadGroupPersistenceServiceTests.java | 63 ++ .../test/wlm/10_workload_group.yml | 124 ++++ .../search/StreamTransportSearchAction.java | 7 +- .../action/search/TransportSearchAction.java | 122 +++- .../search/TransportSearchScrollAction.java | 25 +- .../cluster/metadata/WorkloadGroup.java | 133 +++- .../opensearch/indices/IndicesService.java | 17 + .../admin/cluster/RestWlmStatsAction.java | 6 +- .../org/opensearch/search/SearchService.java | 10 + .../wlm/MutableWorkloadGroupFragment.java | 113 ++- .../opensearch/wlm/WorkloadGroupService.java | 166 +++++ .../org/opensearch/wlm/WorkloadGroupTask.java | 51 ++ .../wlm/WorkloadGroupThrottleSettings.java | 158 ++++ .../wlm/WorkloadGroupThrottleTracker.java | 107 +++ .../wlm/WorkloadGroupsStateAccessor.java | 15 +- .../wlm/stats/WorkloadGroupState.java | 13 + .../wlm/stats/WorkloadGroupStats.java | 26 +- .../cluster/wlm/WlmStatsResponseTests.java | 2 + .../WlmPaginationStrategyTests.java | 2 +- .../search/TransportSearchActionTests.java | 4 +- .../cluster/metadata/WorkloadGroupTests.java | 388 +++++++++- .../cluster/RestWlmStatsActionTests.java | 7 +- .../snapshots/SnapshotResiliencyTests.java | 4 +- .../wlm/WorkloadGroupServiceTests.java | 325 +++++++++ .../wlm/WorkloadGroupTaskTests.java | 18 + .../WorkloadGroupThrottleTrackerTests.java | 166 +++++ ...oadGroupRequestOperationListenerTests.java | 6 + .../opensearch/wlm/stats/WlmStatsTests.java | 3 +- .../wlm/stats/WorkloadGroupStatsTests.java | 37 +- 35 files changed, 3012 insertions(+), 73 deletions(-) create mode 100644 plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmNodeThrottlingIT.java create mode 100644 server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleSettings.java create mode 100644 server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java create mode 100644 server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java diff --git a/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmNodeThrottlingIT.java b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmNodeThrottlingIT.java new file mode 100644 index 0000000000000..408e60ee2abf5 --- /dev/null +++ b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmNodeThrottlingIT.java @@ -0,0 +1,672 @@ +/* + * 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.ActionRequest; +import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest; +import org.opensearch.action.index.IndexResponse; +import org.opensearch.action.search.SearchRequestBuilder; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.support.ActionFilter; +import org.opensearch.action.support.ActionFilterChain; +import org.opensearch.action.support.ActionRequestMetadata; +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.action.ActionListener; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.indices.TermsLookup; +import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType; +import org.opensearch.plugins.ActionPlugin; +import org.opensearch.plugins.Plugin; +import org.opensearch.plugins.PluginsService; +import org.opensearch.rest.RestHeaderDefinition; +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.tasks.Task; +import org.opensearch.test.OpenSearchIntegTestCase; +import org.opensearch.wlm.MutableWorkloadGroupFragment; +import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadGroupTask; +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.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; + +/** + * End-to-end integration test for per-node WLM request throttling ({@code node_limit}, {@code attribute=group}). + *

+ * The scripted-block plugin holds a search in-flight (occupying a throttle permit) so that a second concurrent + * search deterministically exceeds the node limit and must be rejected with a 429 + * ({@link OpenSearchRejectedExecutionException}). This exercises the real coordinator admission hook in + * {@code TransportSearchAction} through auto-tagging, not a mocked service. + */ +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 1, numClientNodes = 0, supportsDedicatedMasters = false) +public class WlmNodeThrottlingIT 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); + plugins.add(TestPrincipalPlugin.class); + return plugins; + } + + @Before + public void registerFeatureTypeIfMissingOnAllNodes() { + // AutoTaggingRegistry is a JVM-static singleton, but each test (Scope.TEST) restarts the cluster and rebuilds + // the feature type — including its WorkloadGroupFeatureValueValidator, which is bound to that cluster's live + // ClusterService. Always refresh the registry to the current cluster's feature type; otherwise a later test + // would validate rules against a previous (dead) cluster's state and fail with "not a valid workload group id". + 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 testSecondConcurrentSearchRejectedWhenNodeLimitReached() throws Exception { + String workloadGroupId = "wlm_throttle_group"; + String ruleId = "wlm_throttle_rule"; + String indexName = "throttle_index"; + + setWlmMode("enabled"); + + // Workload group throttled to a single in-flight request per node. + WorkloadGroup workloadGroup = createThrottledWorkloadGroup("throttle_test_group", workloadGroupId, 1); + updateWorkloadGroupInClusterState(PUT, workloadGroup); + + FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME); + createRule(ruleId, "throttle rule", indexName, featureType, workloadGroupId); + + indexDocument(indexName); + + // Rule propagation to the in-memory processing service is asynchronous. Wait until a + // (non-blocking) search is actually tagged to the throttled group before exercising + // the concurrency scenario, otherwise the requests are untagged and never throttled. + assertBusy(() -> { + int before = getCompletions(workloadGroupId); + client().prepareSearch(indexName).setQuery(org.opensearch.index.query.QueryBuilders.matchAllQuery()).get(); + int after = getCompletions(workloadGroupId); + assertTrue("Expected search to be tagged to the throttled workload group", after > before); + }, 30, TimeUnit.SECONDS); + + List plugins = initBlockFactory(); + + // First search: blocks in the query phase, holding the single permit. + ActionFuture blockedSearch = blockingSearch(indexName).execute(); + awaitForBlock(plugins); + + int throttledBefore = getThrottled(workloadGroupId); + long inFlightBefore = currentInFlightSearches(); + + // Second search while the first is still in-flight: must be rejected (429). + Throwable rejection = expectThrows(Throwable.class, () -> blockingSearch(indexName).execute().actionGet(TIMEOUT)); + assertTrue( + "Expected an OpenSearchRejectedExecutionException in the cause chain but was: " + rejection, + hasRejectedExecutionCause(rejection) + ); + + // The rejection must be counted in total_throttled. + assertEquals("total_throttled should increment by exactly one", throttledBefore + 1, getThrottled(workloadGroupId)); + + // The rejected request must NOT have entered the request-operations start path. This guards against the gauge + // leak where a throttle rejection increments 'current' via onRequestStart but never reaches + // onRequestEnd/onRequestFailure. + // + // Poll for the gauge to settle rather than comparing two instantaneous samples: the gauge is node-global and + // the WLM rule-sync job issues its own search every few seconds, so any single pair of samples can differ by + // that traffic in either direction. The steady state is well defined here -- the first search is still blocked + // and nothing else in this test is running -- so the gauge must come back to inFlightBefore. A real leak is a + // permanent +1 and never settles, so the poll still fails on the regression it is guarding. + assertBusy( + () -> assertEquals( + "in-flight search gauge must exclude the throttle-rejected request", + inFlightBefore, + currentInFlightSearches() + ), + 30, + TimeUnit.SECONDS + ); + + // Release the block; the first search should complete successfully. + disableBlocks(plugins); + assertNotNull(blockedSearch.actionGet(TIMEOUT)); + + // Once the blocked search finishes, the gauge must drain back to zero (no leaked in-flight count). + assertBusy(() -> assertEquals("in-flight search gauge must drain to zero", 0L, currentInFlightSearches()), 30, TimeUnit.SECONDS); + } + + public void testScrollContinuationIsThrottled() throws Exception { + String workloadGroupId = "wlm_scroll_throttle_group"; + String ruleId = "wlm_scroll_throttle_rule"; + String indexName = "scroll_throttle_index"; + + setWlmMode("enabled"); + WorkloadGroup workloadGroup = createThrottledWorkloadGroup("scroll_throttle_test_group", workloadGroupId, 1); + updateWorkloadGroupInClusterState(PUT, workloadGroup); + + FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME); + createRule(ruleId, "scroll throttle rule", indexName, featureType, workloadGroupId); + + indexDocument(indexName); + + assertBusy(() -> { + int before = getCompletions(workloadGroupId); + client().prepareSearch(indexName).setQuery(org.opensearch.index.query.QueryBuilders.matchAllQuery()).get(); + int after = getCompletions(workloadGroupId); + assertTrue("Expected search to be tagged to the throttled workload group", after > before); + }, 30, TimeUnit.SECONDS); + + // Open a scroll context with a cheap query so the initial search releases its permit immediately. + String scrollId = client().prepareSearch(indexName) + .setQuery(org.opensearch.index.query.QueryBuilders.matchAllQuery()) + .setSize(1) + .setScroll(TIMEOUT) + .get() + .getScrollId(); + try { + List plugins = initBlockFactory(); + ActionFuture blockedSearch = blockingSearch(indexName).execute(); + awaitForBlock(plugins); + + int throttledBefore = getThrottled(workloadGroupId); + + // The group's only permit is held. A scroll continuation must be rejected like any other search -- if it is + // admitted, node_limit is evadable simply by adding ?scroll= to a query. + final String sid = scrollId; + Throwable rejection = expectThrows( + Throwable.class, + () -> client().prepareSearchScroll(sid).setScroll(TIMEOUT).execute().actionGet(TIMEOUT) + ); + assertTrue("Expected a scroll continuation to be throttled but was: " + rejection, hasRejectedExecutionCause(rejection)); + assertEquals("a throttled scroll must be counted", throttledBefore + 1, getThrottled(workloadGroupId)); + + disableBlocks(plugins); + assertNotNull(blockedSearch.actionGet(TIMEOUT)); + + // With the permit released the same scroll continues normally, proving the rejection was the throttle and + // not a broken scroll context. + assertNotNull(client().prepareSearchScroll(sid).setScroll(TIMEOUT).get()); + } finally { + client().prepareClearScroll().addScrollId(scrollId).get(); + } + } + + public void testUsernameThrottlingKeepsPerUserBuckets() throws Exception { + String workloadGroupId = "wlm_user_throttle_group"; + String ruleId = "wlm_user_throttle_rule"; + String indexName = "user_throttle_index"; + + setWlmMode("enabled"); + + // Group throttled per-username to a single in-flight request per node (attribute = username). + WorkloadGroup workloadGroup = createThrottledWorkloadGroup("user_throttle_test_group", workloadGroupId, 1, "username"); + updateWorkloadGroupInClusterState(PUT, workloadGroup); + + FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME); + // The rule's feature value (the workload group id) is validated against applied cluster state, which the group + // update above populates asynchronously. Wait until the group is visible in cluster state before creating the + // rule, otherwise rule creation races the update and fails validation. + 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); + createRule(ruleId, "user throttle rule", indexName, featureType, workloadGroupId); + + indexDocument(indexName); + + // Wait for rule propagation: a search tagged as alice must reach the group before the concurrency scenario. + assertBusy(() -> { + int before = getCompletions(workloadGroupId); + searchAs("alice", indexName).setQuery(org.opensearch.index.query.QueryBuilders.matchAllQuery()).get(); + int after = getCompletions(workloadGroupId); + assertTrue("Expected search to be tagged to the throttled workload group", after > before); + }, 30, TimeUnit.SECONDS); + + List plugins = initBlockFactory(); + + // alice's first search blocks in the query phase, holding her single per-user permit. + ActionFuture aliceBlocked = blockingSearchAs("alice", indexName).execute(); + awaitForBlock(plugins); + + int throttledBefore = getThrottled(workloadGroupId); + + // alice's second concurrent search hits her per-user node_limit -> 429. + Throwable rejection = expectThrows(Throwable.class, () -> blockingSearchAs("alice", indexName).execute().actionGet(TIMEOUT)); + assertTrue( + "Expected an OpenSearchRejectedExecutionException in the cause chain but was: " + rejection, + hasRejectedExecutionCause(rejection) + ); + assertEquals("total_throttled should increment by exactly one", throttledBefore + 1, getThrottled(workloadGroupId)); + + // bob is a different principal -> a different bucket -> admitted even while alice is at her limit. + // (bob's search also blocks; we just need it to get past admission, so run it async and then release.) + ActionFuture bobBlocked = blockingSearchAs("bob", indexName).execute(); + assertBusy(() -> { + int blocked = 0; + for (ScriptedBlockPlugin plugin : plugins) { + blocked += plugin.hits.get(); + } + assertThat("bob's search should have been admitted and reached the blocking script", blocked, greaterThan(1)); + }, 30, TimeUnit.SECONDS); + // bob was admitted, so no additional throttle beyond alice's one rejection. + assertEquals("bob must not be throttled by alice's bucket", throttledBefore + 1, getThrottled(workloadGroupId)); + + // Release the blocks; both alice's and bob's blocked searches complete successfully. + disableBlocks(plugins); + assertNotNull(aliceBlocked.actionGet(TIMEOUT)); + assertNotNull(bobBlocked.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 getCompletions(String groupId) throws Exception { + return sumGroupStat(groupId, org.opensearch.wlm.stats.WorkloadGroupStats.WorkloadGroupStatsHolder::getCompletions); + } + + private int getThrottled(String groupId) throws Exception { + return sumGroupStat(groupId, org.opensearch.wlm.stats.WorkloadGroupStats.WorkloadGroupStatsHolder::getThrottled); + } + + /** + * Sums one stat for a workload group across every node's WLM stats, read from the response objects directly. The + * group may be absent from a node that has not registered it yet, which contributes nothing. + */ + private int sumGroupStat( + String groupId, + java.util.function.ToLongFunction extractor + ) throws Exception { + org.opensearch.action.admin.cluster.wlm.WlmStatsRequest request = new org.opensearch.action.admin.cluster.wlm.WlmStatsRequest( + null, + new java.util.HashSet<>(Collections.singletonList(groupId)), + null + ); + org.opensearch.action.admin.cluster.wlm.WlmStatsResponse response = client().execute( + org.opensearch.action.admin.cluster.wlm.WlmStatsAction.INSTANCE, + request + ).get(); + long total = 0; + for (org.opensearch.wlm.stats.WlmStats nodeStats : response.getNodes()) { + org.opensearch.wlm.stats.WorkloadGroupStats.WorkloadGroupStatsHolder holder = nodeStats.getWorkloadGroupStats() + .getStats() + .get(groupId); + if (holder != null) { + total += extractor.applyAsLong(holder); + } + } + return Math.toIntExact(total); + } + + /** + * Sums the current in-flight search gauge ({@link org.opensearch.action.search.SearchRequestStats#getTookCurrent()}) + * across all data nodes. This is the counter incremented in {@code onRequestStart} and decremented in + * {@code onRequestEnd}/{@code onRequestFailure}; a throttle rejection must never touch it. + */ + private long currentInFlightSearches() { + long total = 0; + for (org.opensearch.action.search.SearchRequestStats stats : internalCluster().getDataNodeInstances( + org.opensearch.action.search.SearchRequestStats.class + )) { + total += stats.getTookCurrent(); + } + return total; + } + + public void testNestedRewriteSearchIsNotChargedASecondPermit() throws Exception { + String workloadGroupId = "wlm_nested_group"; + String ruleId = "wlm_nested_rule"; + String indexName = "orders"; + String lookupIndex = "lookupidx"; + + setWlmMode("enabled"); + // node_limit=1 is the case that exposes re-entrancy: the outer search holds the group's only permit while its + // rewrite phase issues a nested coordinator search that resolves to the same bucket. + WorkloadGroup workloadGroup = createThrottledWorkloadGroup("nested_test_group", workloadGroupId, 1); + updateWorkloadGroupInClusterState(PUT, workloadGroup); + assertBusy( + () -> assertTrue( + "workload group not yet applied in cluster state", + client().admin().cluster().prepareState().get().getState().metadata().workloadGroups().containsKey(workloadGroupId) + ), + 30, + TimeUnit.SECONDS + ); + + FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME); + createRule(ruleId, "nested rule", indexName, featureType, workloadGroupId); + + indexDocument(indexName); + // The lookup index deliberately matches no rule, so the nested search inherits the outer request's workload + // group id from the thread context -- the same bucket the outer request already holds a permit for. + assertAcked( + client().admin() + .indices() + .prepareCreate(lookupIndex) + .setSettings(Settings.builder().put("index.number_of_shards", 1).put("index.number_of_replicas", 0)) + ); + client().prepareIndex(lookupIndex) + .setId("1") + .setSource(Map.of("uid", "value")) + .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE) + .get(); + + assertBusy(() -> { + int before = getCompletions(workloadGroupId); + client().prepareSearch(indexName).setQuery(org.opensearch.index.query.QueryBuilders.matchAllQuery()).get(); + assertTrue("Expected search to be tagged to the throttled workload group", getCompletions(workloadGroupId) > before); + }, 30, TimeUnit.SECONDS); + + int throttledBefore = getThrottled(workloadGroupId); + + // A terms lookup with a subquery issues a full nested coordinator search during the rewrite phase. With zero + // other load this must succeed: the outer request already paid for the bucket. + TermsLookup lookup = new TermsLookup(lookupIndex, null, "uid", org.opensearch.index.query.QueryBuilders.matchAllQuery()); + SearchResponse response = client().prepareSearch(indexName) + .setQuery(org.opensearch.index.query.QueryBuilders.termsLookupQuery("field", lookup)) + .execute() + .actionGet(TIMEOUT); + assertEquals(RestStatus.OK, response.status()); + assertEquals("a nested rewrite search must not be counted as throttled", throttledBefore, getThrottled(workloadGroupId)); + + // The exemption must be scoped to nesting only -- a genuinely concurrent second request still gets a 429, + // otherwise the fix would have silently disabled throttling for this group. + List plugins = initBlockFactory(); + ActionFuture blocked = blockingSearch(indexName).execute(); + awaitForBlock(plugins); + try { + Throwable rejection = expectThrows(Throwable.class, () -> blockingSearch(indexName).execute().actionGet(TIMEOUT)); + assertTrue( + "an independent concurrent request must still be throttled but was: " + rejection, + hasRejectedExecutionCause(rejection) + ); + } finally { + disableBlocks(plugins); + assertNotNull(blocked.actionGet(TIMEOUT)); + } + } + + private SearchRequestBuilder blockingSearch(String indexName) { + return client().prepareSearch(indexName) + .setQuery(scriptQuery(new Script(ScriptType.INLINE, "mockscript", ScriptedBlockPlugin.SCRIPT_NAME, Collections.emptyMap()))); + } + + // In production the WLM auto-tagging filter sets the task's throttle principal from the security plugin's principal + // extractor. There is no such extractor here, so TestPrincipalPlugin below stands in for it, reading the username + // from a test-only header and setting it on the task exactly as the real filter does. This exercises the real + // plumbing (task field -> throttle admission) rather than simulating it. + private org.opensearch.transport.client.Client clientAs(String username) { + return client().filterWithHeader(Map.of(TestPrincipalPlugin.TEST_PRINCIPAL_HEADER, "username|" + username)); + } + + private SearchRequestBuilder searchAs(String username, String indexName) { + return clientAs(username).prepareSearch(indexName); + } + + private SearchRequestBuilder blockingSearchAs(String username, String indexName) { + return clientAs(username).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(WorkloadManagementSettings.WLM_MODE_SETTING.getKey(), mode); + ClusterUpdateSettingsRequest request = new ClusterUpdateSettingsRequest().persistentSettings(settings); + assertAcked(client().admin().cluster().updateSettings(request).get()); + } + + private WorkloadGroup createThrottledWorkloadGroup(String name, String id, int nodeLimit) { + return createThrottledWorkloadGroup(name, id, nodeLimit, "group"); + } + + private WorkloadGroup createThrottledWorkloadGroup(String name, String id, int nodeLimit, String attribute) { + Settings throttling = Settings.builder() + .put(WorkloadGroupThrottleSettings.ATTRIBUTE.getKey(), attribute) + .put(WorkloadGroupThrottleSettings.NODE_LIMIT.getKey(), nodeLimit) + .build(); + return new WorkloadGroup( + name, + id, + new MutableWorkloadGroupFragment( + MutableWorkloadGroupFragment.ResiliencyMode.SOFT, + Map.of(ResourceType.CPU, 0.9, ResourceType.MEMORY, 0.9), + Settings.EMPTY, + throttling + ), + Instant.now().getMillis() + ); + } + + private void indexDocument(String indexName) { + assertAcked( + client().admin() + .indices() + .prepareCreate(indexName) + .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()); + } + } + + /** + * Stands in for the security plugin's principal extractor. Registers a test-only header as both a REST header and a + * task header, then copies it onto the task's throttle principal exactly as the real auto-tagging filter does, so the + * IT exercises the production plumbing (task field -> throttle admission) instead of simulating it. + */ + public static class TestPrincipalPlugin extends Plugin implements ActionPlugin { + static final String TEST_PRINCIPAL_HEADER = "test_throttle_principal"; + + @Override + public Collection getRestHeaders() { + return List.of(new RestHeaderDefinition(TEST_PRINCIPAL_HEADER, false)); + } + + @Override + public Collection getTaskHeaders() { + return List.of(TEST_PRINCIPAL_HEADER); + } + + @Override + public List getActionFilters() { + return List.of(new ActionFilter() { + @Override + public int order() { + return 0; + } + + @Override + public void apply( + Task task, + String action, + Req request, + ActionRequestMetadata metadata, + ActionListener listener, + ActionFilterChain chain + ) { + if (task instanceof WorkloadGroupTask) { + String principal = task.getHeader(TEST_PRINCIPAL_HEADER); + if (principal != null) { + ((WorkloadGroupTask) task).setThrottlePrincipal(principal); + } + } + chain.proceed(task, action, request, listener); + } + }); + } + } + + /** + * 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(WlmNodeThrottlingIT.class).info("Blocking on the document {}", fieldsLookup.get("_id")); + hits.incrementAndGet(); + try { + // Explicit, generous budget: the default overload is 10s, but callers hold a search here while + // running their own 30s assertBusy waits, so the default would expire first and surface as a + // baffling "expected false but was true" failure inside an unrelated assertion. + assertBusy(() -> assertFalse(shouldBlock.get()), 120, TimeUnit.SECONDS); + } catch (Exception e) { + throw new RuntimeException(e); + } + return true; + }); + } + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/AutoTaggingActionFilter.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/AutoTaggingActionFilter.java index bea19e17073ec..48ef87452c7f4 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/AutoTaggingActionFilter.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/AutoTaggingActionFilter.java @@ -117,14 +117,45 @@ public LogicalOperator getLogicalOperator() { } } + List principalValues = null; if (featureType.getAllowedAttributesRegistry().containsKey(PRINCIPAL_ATTRIBUTE_NAME)) { Attribute attribute = featureType.getAllowedAttributesRegistry().get(PRINCIPAL_ATTRIBUTE_NAME); assert attributeExtensions.containsKey(attribute); - attributeExtractors.add(attributeExtensions.get(attribute).getAttributeExtractor()); + final AttributeExtractor extractor = attributeExtensions.get(attribute).getAttributeExtractor(); + // Materialize once. The value is needed both for label evaluation and for the principal header, and + // AttributeExtractor.extract() carries no re-iterability contract -- a stream-backed implementation would + // yield nothing the second time and silently disable username/role throttling. + final List values = new ArrayList<>(); + extractor.extract().forEach(values::add); + principalValues = values; + attributeExtractors.add(new AttributeExtractor<>() { + @Override + public Attribute getAttribute() { + return extractor.getAttribute(); + } + + @Override + public Iterable extract() { + return values; + } + + @Override + public LogicalOperator getLogicalOperator() { + return extractor.getLogicalOperator(); + } + }); } Optional label = ruleProcessingService.evaluateLabel(attributeExtractors); label.ifPresent(s -> threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, s)); + // Hand the principal to core-side throttling so it can build per-username / per-role buckets. It goes on the + // task, not into the thread context: see WorkloadGroupTask#setThrottlePrincipal. + if (principalValues != null && task instanceof WorkloadGroupTask) { + String principal = String.join(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER, principalValues); + if (principal.isEmpty() == false) { + ((WorkloadGroupTask) task).setThrottlePrincipal(principal); + } + } chain.proceed(task, action, request, listener); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java index 2039f1cb590ff..69519cc13872e 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java @@ -71,6 +71,15 @@ protected void clusterManagerOperation( ClusterState clusterState, ActionListener listener ) { + try { + WorkloadGroupPersistenceService.validateThrottlingIsEnforceable( + request.getWorkloadGroup().getMutableWorkloadGroupFragment().getThrottling(), + clusterState + ); + } catch (Exception e) { + listener.onFailure(e); + return; + } workloadGroupPersistenceService.persistInClusterStateMetadata(request.getWorkloadGroup(), listener); } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java index ef639d44b4155..4b651f1f8e67e 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java @@ -71,6 +71,15 @@ protected void clusterManagerOperation( ClusterState clusterState, ActionListener listener ) { + try { + WorkloadGroupPersistenceService.validateThrottlingIsEnforceable( + request.getmMutableWorkloadGroupFragment().getThrottling(), + clusterState + ); + } catch (Exception e) { + listener.onFailure(e); + return; + } workloadGroupPersistenceService.updateInClusterStateMetadata(request, listener); } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java index ecf5c7f68fa70..5831161b9ae0c 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java @@ -11,6 +11,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.ResourceNotFoundException; +import org.opensearch.Version; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; import org.opensearch.cluster.AckedClusterStateUpdateTask; import org.opensearch.cluster.ClusterState; @@ -27,12 +28,17 @@ import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; +import org.opensearch.plugin.wlm.WorkloadManagementPlugin; import org.opensearch.plugin.wlm.action.CreateWorkloadGroupResponse; import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupResponse; +import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType; +import org.opensearch.rule.autotagging.AutoTaggingRegistry; +import org.opensearch.rule.autotagging.FeatureType; import org.opensearch.wlm.MutableWorkloadGroupFragment; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadGroupThrottleSettings; import java.util.Collection; import java.util.EnumMap; @@ -366,4 +372,59 @@ public int getMaxWorkloadGroupCount() { public ClusterService getClusterService() { return clusterService; } + + /** + * Rejects a throttling config the cluster cannot actually honour. Both cases below would otherwise return a 200 for + * a config that silently never takes effect: + *

+ * Called from the cluster-manager transport actions rather than from a cluster-state applier or settings update + * consumer on purpose: throwing while applying cluster state wedges the cluster-manager. + * + * @param throttling the incoming throttling fragment, may be {@code null} or empty (both fine: nothing to honour) + * @param clusterState state used to read the oldest node version in the cluster + * @throws IllegalArgumentException if the config cannot be enforced + */ + public static void validateThrottlingIsEnforceable(Settings throttling, ClusterState clusterState) { + if (throttling == null || throttling.isEmpty()) { + return; + } + Version minNodeVersion = clusterState.nodes().getMinNodeVersion(); + if (minNodeVersion.before(Version.V_3_9_0)) { + throw new IllegalArgumentException( + "workload group throttling requires every node to be on " + + Version.V_3_9_0 + + " or later, but the oldest node in the cluster is on " + + minNodeVersion + + ". The throttling config would be silently dropped; complete the upgrade first." + ); + } + // ATTRIBUTE.get returns "" (its default), not null, when the key is absent -- which is the normal shape of a + // partial update that only changes the limit. Only an explicitly principal-keyed attribute is checked here; the + // merged config is validated separately. + String attribute = WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling); + if (attribute == null || attribute.isEmpty() || "group".equals(attribute)) { + return; + } + try { + FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME); + if (featureType.getAllowedAttributesRegistry().containsKey(WorkloadManagementPlugin.PRINCIPAL_ATTRIBUTE_NAME) == false) { + throw new IllegalArgumentException( + "throttling attribute [" + + attribute + + "] needs a principal attribute provider (the security plugin) to be installed, otherwise the " + + "limit can never be enforced. Use attribute [group] instead." + ); + } + } catch (ResourceNotFoundException e) { + // Feature type not registered on this node yet. Skip rather than reject a config that is probably fine -- + // the throttle path fails open anyway, so a false rejection here is worse than a missed warning. + logger.debug("WLM feature type not registered; skipping principal-attribute check for throttling config", e); + } + } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/AutoTaggingActionFilterTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/AutoTaggingActionFilterTests.java index 40995d70c0848..cf0fc9911c9b9 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/AutoTaggingActionFilterTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/AutoTaggingActionFilterTests.java @@ -18,6 +18,8 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.core.action.ActionResponse; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.plugin.wlm.spi.AttributeExtractorExtension; import org.opensearch.rule.InMemoryRuleProcessingService; import org.opensearch.rule.RuleAttribute; import org.opensearch.rule.attribute_extractor.AttributeExtractor; @@ -37,7 +39,9 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; +import static org.opensearch.plugin.wlm.WorkloadManagementPlugin.PRINCIPAL_ATTRIBUTE_NAME; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.anyList; import static org.mockito.Mockito.doAnswer; @@ -100,6 +104,181 @@ public void testApplyForInValidRequest() { verify(ruleProcessingService, times(0)).evaluateLabel(anyList()); } + public void testApplySetsThrottlePrincipalOnTaskWhenExtractorPresent() { + // A feature type that includes a "principal" attribute + a matching extractor extension in the map. + Attribute principalAttr = new Attribute() { + @Override + public String getName() { + return PRINCIPAL_ATTRIBUTE_NAME; + } + + @Override + public void validateAttribute() {} + + @Override + public void writeTo(StreamOutput out) throws IOException {} + }; + FeatureType featureTypeWithPrincipal = new FeatureType() { + @Override + public String getName() { + return "wlm"; + } + + @Override + public Map getOrderedAttributes() { + return Map.of(principalAttr, 1); + } + }; + AttributeExtractor principalExtractor = new AttributeExtractor<>() { + @Override + public Attribute getAttribute() { + return principalAttr; + } + + @Override + public Iterable extract() { + return List.of("username|alice", "role|admin"); + } + + @Override + public LogicalOperator getLogicalOperator() { + return LogicalOperator.OR; + } + }; + AttributeExtractorExtension extension = () -> principalExtractor; + Map extensions = Map.of(principalAttr, extension); + + InMemoryRuleProcessingService svc = spy( + new InMemoryRuleProcessingService( + new AttributeValueStoreFactory(featureTypeWithPrincipal, DefaultAttributeValueStore::new), + null + ) + ); + AutoTaggingActionFilter filter = new AutoTaggingActionFilter( + svc, + threadPool, + extensions, + mock(WlmClusterSettingValuesProvider.class), + featureTypeWithPrincipal + ); + + SearchRequest request = mock(SearchRequest.class); + when(request.indices()).thenReturn(new String[] { "foo" }); + ActionFilterChain chain = mock(TestActionFilterChain.class); + WorkloadGroupTask task = newWorkloadGroupTask(); + try (ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext()) { + when(svc.evaluateLabel(anyList())).thenReturn(Optional.of("QG")); + filter.apply(task, "Test", request, ActionRequestMetadata.empty(), null, chain); + + // Both principal tokens are joined (by WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER) onto the task for + // core-side throttling. + assertEquals( + "username|alice" + WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER + "role|admin", + task.getThrottlePrincipal() + ); + // The principal must NOT land in the thread context: request headers are serialized onto every outgoing + // transport request, which would ship the caller's identity to every shard and to remote clusters. + assertNull(threadPool.getThreadContext().getHeader("workloadGroupPrincipal")); + } + } + + public void testApplyTwiceOnOneThreadContextIsTolerated() { + // _msearch dispatches each sub-search through the filter chain on a single ThreadContext with no + // stashContext(), and ThreadContext.putHeader throws when the key is already present. The filter must therefore + // tolerate running more than once, otherwise every sub-request after the first fails. Carrying the principal on + // the task instead of in the thread context is what makes that safe, and gives each sub-request its own value. + Attribute principalAttr = new Attribute() { + @Override + public String getName() { + return PRINCIPAL_ATTRIBUTE_NAME; + } + + @Override + public void validateAttribute() {} + + @Override + public void writeTo(StreamOutput out) throws IOException {} + }; + FeatureType featureTypeWithPrincipal = new FeatureType() { + @Override + public String getName() { + return "wlm"; + } + + @Override + public Map getOrderedAttributes() { + return Map.of(principalAttr, 1); + } + }; + AtomicInteger extractCalls = new AtomicInteger(); + AttributeExtractor principalExtractor = new AttributeExtractor<>() { + @Override + public Attribute getAttribute() { + return principalAttr; + } + + @Override + public Iterable extract() { + extractCalls.incrementAndGet(); + return List.of("username|alice"); + } + + @Override + public LogicalOperator getLogicalOperator() { + return LogicalOperator.OR; + } + }; + AttributeExtractorExtension extension = () -> principalExtractor; + InMemoryRuleProcessingService svc = spy( + new InMemoryRuleProcessingService( + new AttributeValueStoreFactory(featureTypeWithPrincipal, DefaultAttributeValueStore::new), + null + ) + ); + AutoTaggingActionFilter filter = new AutoTaggingActionFilter( + svc, + threadPool, + Map.of(principalAttr, extension), + mock(WlmClusterSettingValuesProvider.class), + featureTypeWithPrincipal + ); + + SearchRequest request = mock(SearchRequest.class); + when(request.indices()).thenReturn(new String[] { "foo" }); + ActionFilterChain chain = mock(TestActionFilterChain.class); + WorkloadGroupTask first = newWorkloadGroupTask(); + WorkloadGroupTask second = newWorkloadGroupTask(); + try (ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext()) { + // No label, so the (separate, pre-existing) workload-group-id header is not set and this test isolates the + // principal. + when(svc.evaluateLabel(anyList())).thenReturn(Optional.empty()); + filter.apply(first, "Test", request, ActionRequestMetadata.empty(), null, chain); + filter.apply(second, "Test", request, ActionRequestMetadata.empty(), null, chain); + + // Each sub-request carries its own principal, and neither run threw on a duplicate key. + assertEquals("username|alice", first.getThrottlePrincipal()); + assertEquals("username|alice", second.getThrottlePrincipal()); + assertNull(threadPool.getThreadContext().getHeader("workloadGroupPrincipal")); + // The principal is materialized once per request and reused for both label evaluation and the task field; + // extract() carries no re-iterability contract, so calling it twice per request risks yielding nothing. + assertEquals("extract() must be invoked once per request", 2, extractCalls.get()); + } + } + + public void testApplyLeavesThrottlePrincipalUnsetWhenNoExtractor() { + // Default filter from setUp has no principal attribute/extractor -> the task's principal stays null, which is + // what makes username/role throttling fail open rather than bucket everyone together. + SearchRequest request = mock(SearchRequest.class); + when(request.indices()).thenReturn(new String[] { "foo" }); + ActionFilterChain chain = mock(TestActionFilterChain.class); + try (ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext()) { + when(ruleProcessingService.evaluateLabel(anyList())).thenReturn(Optional.of("QG")); + WorkloadGroupTask task = newWorkloadGroupTask(); + autoTaggingActionFilter.apply(task, "Test", request, ActionRequestMetadata.empty(), null, chain); + assertNull(task.getThrottlePrincipal()); + } + } + public void testApplyForScrollRequestWithOriginalIndices() { SearchScrollRequest request = mock(SearchScrollRequest.class); ActionFilterChain chain = mock(TestActionFilterChain.class); @@ -169,6 +348,10 @@ public void validateAttribute() {} public void writeTo(StreamOutput out) throws IOException {} } + private static WorkloadGroupTask newWorkloadGroupTask() { + return new WorkloadGroupTask(1L, "transport", "Test", "test task", TaskId.EMPTY_TASK_ID, Map.of()); + } + private static class TestActionFilterChain implements ActionFilterChain { @Override public void proceed(Task task, String action, ActionRequest request, ActionListener listener) { diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceServiceTests.java index 51911b2b67df8..006773aa2df43 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceServiceTests.java @@ -9,6 +9,7 @@ package org.opensearch.plugin.wlm.service; import org.opensearch.ResourceNotFoundException; +import org.opensearch.Version; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; import org.opensearch.cluster.AckedClusterStateUpdateTask; import org.opensearch.cluster.ClusterName; @@ -16,11 +17,14 @@ import org.opensearch.cluster.ClusterStateUpdateTask; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.WorkloadGroup; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.collect.Tuple; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.transport.TransportAddress; import org.opensearch.plugin.wlm.WorkloadManagementTestUtils; import org.opensearch.plugin.wlm.action.CreateWorkloadGroupResponse; import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupRequest; @@ -530,4 +534,63 @@ public void testUpdateInClusterStateMetadataFailure() { workloadGroupPersistenceService.updateInClusterStateMetadata(updateWorkloadGroupRequest, listener); verify(listener).onFailure(any(RuntimeException.class)); } + + private static ClusterState clusterStateWithOldestNode(Version version) { + DiscoveryNode node = new DiscoveryNode( + "node-1", + new TransportAddress(TransportAddress.META_ADDRESS, 9300), + Map.of(), + Set.of(), + version + ); + return ClusterState.builder(new ClusterName("test")) + .nodes(DiscoveryNodes.builder().add(node).localNodeId("node-1").clusterManagerNodeId("node-1").build()) + .build(); + } + + private static Settings throttling(String attribute, Integer nodeLimit) { + Settings.Builder builder = Settings.builder(); + if (attribute != null) { + builder.put("attribute", attribute); + } + if (nodeLimit != null) { + builder.put("node_limit", nodeLimit); + } + return builder.build(); + } + + public void testValidateThrottlingRejectsClusterWithAPreThrottlingNode() { + // The throttling field is gated on the wire, so a pre-3.9 node in the cluster means the config is dropped in + // transit and the group silently comes back without it. Reject rather than return 200 for a no-op. + IllegalArgumentException e = expectThrows( + IllegalArgumentException.class, + () -> WorkloadGroupPersistenceService.validateThrottlingIsEnforceable( + throttling("group", 5), + clusterStateWithOldestNode(Version.V_3_8_0) + ) + ); + assertTrue(e.getMessage(), e.getMessage().contains("requires every node to be on")); + assertTrue("the message must name the version actually found", e.getMessage().contains(Version.V_3_8_0.toString())); + } + + public void testValidateThrottlingAcceptsWhenEveryNodeSupportsIt() { + WorkloadGroupPersistenceService.validateThrottlingIsEnforceable( + throttling("group", 5), + clusterStateWithOldestNode(Version.V_3_9_0) + ); + } + + public void testValidateThrottlingIgnoresAbsentAndEmptyConfig() { + // Nothing to honour, so an old node in the cluster is not a problem: this is the shape of an update that does + // not touch throttling at all, and of "throttling": null / {}. + ClusterState oldCluster = clusterStateWithOldestNode(Version.V_3_8_0); + WorkloadGroupPersistenceService.validateThrottlingIsEnforceable(null, oldCluster); + WorkloadGroupPersistenceService.validateThrottlingIsEnforceable(Settings.EMPTY, oldCluster); + } + + public void testValidateThrottlingAllowsPartialUpdateWithoutAnAttribute() { + // A limit-only update carries no attribute. ATTRIBUTE.get returns "" rather than null for an absent key, so a + // validator that only null-checks would wrongly reject this. + WorkloadGroupPersistenceService.validateThrottlingIsEnforceable(throttling(null, 9), clusterStateWithOldestNode(Version.V_3_9_0)); + } } diff --git a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_workload_group.yml b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_workload_group.yml index a9ba5d300c9fa..66bed7268bb7f 100644 --- a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_workload_group.yml +++ b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_workload_group.yml @@ -161,3 +161,127 @@ name: "analytics2" - match: { acknowledged: true } + +--- +"test throttling field on the WorkloadGroup API": + - skip: + version: " - 3.8.99" + reason: "throttling was added to workload groups in 3.9" + + - do: + cluster.put_settings: + flat_settings: true + body: + transient: + wlm.workload_group.mode: "enabled" + + # create with throttling: the limit round-trips as a JSON number, not a string + - do: + create_workload_group_context: + body: + { + "name": "throttled", + "resiliency_mode": "enforced", + "resource_limits": { + "cpu": 0.1 + }, + "throttling": { + "attribute": "group", + "node_limit": 5 + } + } + + - match: { name: "throttled" } + - match: { throttling.attribute: "group" } + - match: { throttling.node_limit: 5 } + + - do: + get_workload_group_context: + name: "throttled" + + - match: { workload_groups.0.throttling.attribute: "group" } + - match: { workload_groups.0.throttling.node_limit: 5 } + + # a partial update keeps the key it does not mention + - do: + update_workload_group_context: + name: "throttled" + body: + { + "throttling": { + "node_limit": 9 + } + } + + - match: { throttling.attribute: "group" } + - match: { throttling.node_limit: 9 } + + # an unknown throttle key is rejected + - do: + catch: /Unknown throttle setting/ + update_workload_group_context: + name: "throttled" + body: + { + "throttling": { + "not_a_real_limit": 1 + } + } + + # an invalid attribute is rejected + - do: + catch: /throttling.attribute must be one of/ + update_workload_group_context: + name: "throttled" + body: + { + "throttling": { + "attribute": "index" + } + } + + # username/role keying needs a principal attribute provider; without one the limit could never be enforced, so the + # config is rejected rather than silently accepted and ignored + - do: + catch: /needs a principal attribute provider/ + update_workload_group_context: + name: "throttled" + body: + { + "throttling": { + "attribute": "username" + } + } + + # a negative limit is rejected + - do: + catch: /must be non-negative/ + update_workload_group_context: + name: "throttled" + body: + { + "throttling": { + "node_limit": -1 + } + } + + # "throttling": null disables throttling entirely, and the field is then omitted + - do: + update_workload_group_context: + name: "throttled" + body: + { + "throttling": null + } + + - do: + get_workload_group_context: + name: "throttled" + + - is_false: workload_groups.0.throttling + + - do: + delete_workload_group_context: + name: "throttled" + + - match: { acknowledged: true } diff --git a/server/src/main/java/org/opensearch/action/search/StreamTransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/StreamTransportSearchAction.java index 8474121115222..ae47458e4e993 100644 --- a/server/src/main/java/org/opensearch/action/search/StreamTransportSearchAction.java +++ b/server/src/main/java/org/opensearch/action/search/StreamTransportSearchAction.java @@ -30,6 +30,7 @@ import org.opensearch.transport.StreamTransportService; import org.opensearch.transport.Transport; import org.opensearch.transport.client.node.NodeClient; +import org.opensearch.wlm.WorkloadGroupService; import java.util.Map; import java.util.Set; @@ -59,7 +60,8 @@ public StreamTransportSearchAction( SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory, Tracer tracer, TaskResourceTrackingService taskResourceTrackingService, - IndicesService indicesService + IndicesService indicesService, + WorkloadGroupService workloadGroupService ) { super( client, @@ -78,7 +80,8 @@ public StreamTransportSearchAction( searchRequestOperationsCompositeListenerFactory, tracer, taskResourceTrackingService, - indicesService + indicesService, + workloadGroupService ); } 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 9efdb2c4cd1d9..a268b8f0ea1d5 100644 --- a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java +++ b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java @@ -59,6 +59,7 @@ import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.Nullable; import org.opensearch.common.inject.Inject; +import org.opensearch.common.lease.Releasable; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Setting.Property; import org.opensearch.common.unit.TimeValue; @@ -69,6 +70,7 @@ import org.opensearch.core.common.breaker.CircuitBreaker; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import org.opensearch.core.index.Index; import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.indices.breaker.CircuitBreakerService; @@ -110,6 +112,7 @@ import org.opensearch.transport.client.Client; import org.opensearch.transport.client.OriginSettingClient; import org.opensearch.transport.client.node.NodeClient; +import org.opensearch.wlm.WorkloadGroupService; import org.opensearch.wlm.WorkloadGroupTask; import java.util.ArrayList; @@ -191,6 +194,7 @@ public class TransportSearchAction extends HandledTransportAction) SearchRequest::new); this.client = client; @@ -240,6 +245,7 @@ public TransportSearchAction( clusterService.getClusterSettings(), new ClusterStateFieldDomainProvider() ); + this.workloadGroupService = workloadGroupService; } private Map buildPerIndexAliasFilter( @@ -472,7 +478,7 @@ void executeRequest( final Span requestSpan = tracer.startSpan(SpanBuilder.from(task, actionName)); try (final SpanScope spanScope = tracer.withSpanInScope(requestSpan)) { SearchRequestOperationsListener.CompositeListener requestOperationsListeners; - final ActionListener updatedListener = TraceableActionListener.create(originalListener, requestSpan, tracer); + ActionListener updatedListener = TraceableActionListener.create(originalListener, requestSpan, tracer); requestOperationsListeners = searchRequestOperationsCompositeListenerFactory.buildCompositeListener( originalSearchRequest, logger, @@ -483,14 +489,33 @@ void executeRequest( originalSearchRequest, taskResourceTrackingService::getTaskResourceUsageFromThreadContext ); - searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext); // At this point either the QUERY_GROUP_ID header will be present in ThreadContext either via ActionFilter // or HTTP header (HTTP header will be deprecated once ActionFilter is implemented) if (task instanceof WorkloadGroupTask) { ((WorkloadGroupTask) task).setWorkloadGroupId(threadPool.getThreadContext()); + // Node-level throttle admission. Runs before onRequestStart so a rejection doesn't leak the request + // gauges (decremented only on request end/failure, which the early return skips). The principal is null + // unless the WLM auto-tagging filter set it from the security plugin's extractor. + try { + Releasable throttlePermit = workloadGroupService.acquireThrottleOrReject( + (WorkloadGroupTask) task, + bucketsHeldByAncestors(task) + ); + if (throttlePermit != null) { + // runAfter, not runBefore: it releases in a finally, so a failure to release cannot turn a + // successful search into a client-visible error, and the slot is held until the response has + // actually been handed downstream rather than freed just before it. + updatedListener = ActionListener.runAfter(updatedListener, throttlePermit::close); + } + } catch (OpenSearchRejectedExecutionException e) { + updatedListener.onFailure(e); + return; + } } + searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext); + PipelinedRequest searchRequest; ActionListener listener; try { @@ -517,13 +542,100 @@ void executeRequest( } else { Rewriteable.rewriteAndFetch( sr.source(), - searchService.getRewriteContext(timeProvider::getAbsoluteStartMillis, searchRequest), + // Parent the rewrite phase's searches (a terms lookup with a subquery issues one) on this task, + // so throttle admission can recognise them as nested and not charge the request twice for its + // own bucket. See bucketsHeldByAncestors. + // + // Only when this request actually holds a permit. Otherwise there is no bucket to inherit, and + // EMPTY_TASK_ID leaves the rewrite client unwrapped -- so a search in a group without throttling + // behaves exactly as before, rather than every search in the cluster gaining a parent task it + // never had. + searchService.getRewriteContext( + timeProvider::getAbsoluteStartMillis, + searchRequest, + holdsThrottlePermit(task) ? localTaskId(task) : TaskId.EMPTY_TASK_ID + ), rewriteListener ); } }, listener::onFailure); - searchRequest.transformRequest(requestTransformListener); + try { + searchRequest.transformRequest(requestTransformListener); + } catch (Exception e) { + // Same listener the asynchronous failure path uses above, so a synchronous throw and an async failure + // are reported identically; it wraps updatedListener, so the throttle permit is still released. + listener.onFailure(e); + } + } + } + + /** + * Throttle buckets that ancestor tasks of this request already hold a permit for. + *

+ * A coordinator search can issue a nested coordinator search on the same node while holding a permit -- a terms + * lookup with a subquery does this during the rewrite phase -- and the nested request resolves to the same bucket. + * Admission uses this set to admit such a request without charging it a second permit; see + * {@link WorkloadGroupService#acquireThrottleOrReject(WorkloadGroupTask, Set)}. + *

+ * Only local ancestors are visible, which is exactly the right scope: the throttle is per node, so an ancestor on + * another node holds a permit against that node's budget rather than this one's. + */ + private Set bucketsHeldByAncestors(final Task task) { + TaskId parentTaskId = task.getParentTaskId(); + if (parentTaskId == null || parentTaskId.isSet() == false) { + return Set.of(); + } + final String localNodeId = localNodeId(); + if (localNodeId == null) { + return Set.of(); } + Set held = null; + // Real depth is 1 (a nested rewrite search under a coordinator search); the bound only guards against a + // pathological or cyclic parent chain. + for (int depth = 0; depth < 10 && parentTaskId.isSet() && localNodeId.equals(parentTaskId.getNodeId()); depth++) { + Task ancestor = taskManager.getTask(parentTaskId.getId()); + if (ancestor == null) { + break; + } + if (ancestor instanceof WorkloadGroupTask) { + String bucketKey = ((WorkloadGroupTask) ancestor).getHeldThrottleBucket(); + if (bucketKey != null) { + if (held == null) { + held = new HashSet<>(); + } + held.add(bucketKey); + } + } + parentTaskId = ancestor.getParentTaskId(); + if (parentTaskId == null) { + break; + } + } + return held == null ? Set.of() : held; + } + + /** Whether this request took a throttle permit, and therefore has a bucket a nested search could inherit. */ + private static boolean holdsThrottlePermit(final Task task) { + return task instanceof WorkloadGroupTask && ((WorkloadGroupTask) task).getHeldThrottleBucket() != null; + } + + /** + * The local node's id, or {@code null} if this node is not in the applied cluster state yet. Read via + * {@code state().nodes()} rather than {@link ClusterService#localNode()} because that throws when the node is not + * started, and neither parenting the rewrite phase nor the re-entrancy check is worth failing a search over. + */ + private String localNodeId() { + DiscoveryNode localNode = clusterService.state().nodes().getLocalNode(); + return localNode == null ? null : localNode.getId(); + } + + /** + * A {@link TaskId} addressing {@code task} on this node, or {@link TaskId#EMPTY_TASK_ID} if the local node id is + * unavailable -- in which case the rewrite phase simply issues unparented requests, as it did before. + */ + private TaskId localTaskId(final Task task) { + String localNodeId = localNodeId(); + return localNodeId == null ? TaskId.EMPTY_TASK_ID : new TaskId(localNodeId, task.getId()); } private Task extractParentTask(final SearchRequest searchRequest) { diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchScrollAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchScrollAction.java index c6383acb3d767..ff9a60519a52a 100644 --- a/server/src/main/java/org/opensearch/action/search/TransportSearchScrollAction.java +++ b/server/src/main/java/org/opensearch/action/search/TransportSearchScrollAction.java @@ -36,13 +36,17 @@ import org.opensearch.action.support.HandledTransportAction; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; +import org.opensearch.common.lease.Releasable; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.tasks.Task; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; +import org.opensearch.wlm.WorkloadGroupService; import org.opensearch.wlm.WorkloadGroupTask; +import java.util.Set; + /** * Perform the search scroll * @@ -54,6 +58,7 @@ public class TransportSearchScrollAction extends HandledTransportAction) SearchScrollRequest::new); this.clusterService = clusterService; this.searchTransportService = searchTransportService; this.searchPhaseController = searchPhaseController; this.threadPool = threadPool; + this.workloadGroupService = workloadGroupService; } @Override protected void doExecute(Task task, SearchScrollRequest request, ActionListener listener) { + // Holds the throttle permit release once one is acquired, so every exit below (including the catch) frees it. + ActionListener throttledListener = listener; try { if (task instanceof WorkloadGroupTask) { ((WorkloadGroupTask) task).setWorkloadGroupId(threadPool.getThreadContext()); + // A scroll continuation occupies the node like any other search, so it draws on the same node-level + // budget. Exempting it would make node_limit evadable by appending ?scroll= to a query. + // A scroll continuation issues no nested coordinator search of its own, so there is no ancestor bucket + // to inherit; see TransportSearchAction#bucketsHeldByAncestors. + Releasable throttlePermit = workloadGroupService.acquireThrottleOrReject((WorkloadGroupTask) task, Set.of()); + if (throttlePermit != null) { + throttledListener = ActionListener.runAfter(throttledListener, throttlePermit::close); + } } ParsedScrollId scrollId = request.parseScrollId(); @@ -91,7 +108,7 @@ protected void doExecute(Task task, SearchScrollRequest request, ActionListener< request, (SearchTask) task, scrollId, - listener + throttledListener ); break; case ParsedScrollId.QUERY_AND_FETCH_TYPE: // TODO can we get rid of this? @@ -103,7 +120,7 @@ protected void doExecute(Task task, SearchScrollRequest request, ActionListener< request, (SearchTask) task, scrollId, - listener + throttledListener ); break; default: @@ -111,7 +128,7 @@ protected void doExecute(Task task, SearchScrollRequest request, ActionListener< } action.run(); } catch (Exception e) { - listener.onFailure(e); + throttledListener.onFailure(e); } } } 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 d2e14948c297e..d09d6b2dbf9e6 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java @@ -8,6 +8,8 @@ package org.opensearch.cluster.metadata; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.opensearch.cluster.AbstractDiffable; import org.opensearch.cluster.Diff; import org.opensearch.common.UUIDs; @@ -22,6 +24,7 @@ import org.opensearch.wlm.MutableWorkloadGroupFragment; import org.opensearch.wlm.MutableWorkloadGroupFragment.ResiliencyMode; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadGroupThrottleSettings; import org.joda.time.Instant; import java.io.IOException; @@ -46,6 +49,8 @@ @PublicApi(since = "2.18.0") public class WorkloadGroup extends AbstractDiffable implements ToXContentObject { + private static final Logger logger = LogManager.getLogger(WorkloadGroup.class); + public static final String _ID_STRING = "_id"; public static final String NAME_STRING = "name"; public static final String UPDATED_AT_STRING = "updated_at"; @@ -61,6 +66,16 @@ public WorkloadGroup(String name, MutableWorkloadGroupFragment mutableWorkloadGr } public WorkloadGroup(String name, String _id, MutableWorkloadGroupFragment mutableWorkloadGroupFragment, long updatedAt) { + this(name, _id, mutableWorkloadGroupFragment, updatedAt, false); + } + + private WorkloadGroup( + String name, + String _id, + MutableWorkloadGroupFragment mutableWorkloadGroupFragment, + long updatedAt, + boolean deserializing + ) { Objects.requireNonNull(name, "WorkloadGroup.name can't be null"); Objects.requireNonNull(mutableWorkloadGroupFragment.getResourceLimits(), "WorkloadGroup.resourceLimits can't be null"); Objects.requireNonNull(mutableWorkloadGroupFragment.getResiliencyMode(), "WorkloadGroup.resiliencyMode can't be null"); @@ -74,15 +89,38 @@ public WorkloadGroup(String name, String _id, MutableWorkloadGroupFragment mutab throw new IllegalArgumentException("WorkloadGroup.updatedAtInMillis is not a valid epoch"); } - // Normalize null settings to empty Settings for storage - if (mutableWorkloadGroupFragment.getSettings() == null) { + // 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()); + if (normalizedSettings.equals(mutableWorkloadGroupFragment.getSettings()) == false + || normalizedThrottling.equals(mutableWorkloadGroupFragment.getThrottling()) == false) { mutableWorkloadGroupFragment = new MutableWorkloadGroupFragment( mutableWorkloadGroupFragment.getResiliencyMode(), mutableWorkloadGroupFragment.getResourceLimits(), - Settings.EMPTY + normalizedSettings, + normalizedThrottling ); } + // Cross-field checks on the merged throttling config (attribute required with a limit; ceiling must be >= 1). + // On the deserialization path these are advisory: a newer node may legitimately relax them (e.g. by adding a + // second limit key), and throwing while applying published cluster state would wedge this node out of the + // cluster rather than reject one API call. Enforcement fails open on config it cannot interpret. + if (deserializing) { + try { + WorkloadGroupThrottleSettings.validateMergedConfig(mutableWorkloadGroupFragment.getThrottling()); + } catch (IllegalArgumentException e) { + logger.warn( + "Accepting workload group [{}] with a throttling config this node considers invalid ({}); " + + "throttling will not be enforced for it here", + name, + e.getMessage() + ); + } + } else { + WorkloadGroupThrottleSettings.validateMergedConfig(mutableWorkloadGroupFragment.getThrottling()); + } + this.name = name; this._id = _id; this.mutableWorkloadGroupFragment = mutableWorkloadGroupFragment; @@ -100,7 +138,7 @@ public static boolean isValid(long updatedAt) { } public WorkloadGroup(StreamInput in) throws IOException { - this(in.readString(), in.readString(), new MutableWorkloadGroupFragment(in), in.readLong()); + this(in.readString(), in.readString(), new MutableWorkloadGroupFragment(in), in.readLong(), true); } public static WorkloadGroup updateExistingWorkloadGroup( @@ -114,40 +152,69 @@ public static WorkloadGroup updateExistingWorkloadGroup( } final ResiliencyMode mode = Optional.ofNullable(mutableWorkloadGroupFragment.getResiliencyMode()) .orElse(existingGroup.getResiliencyMode()); - // Handle settings update with merge semantics: - // null settings = not specified in request (keep existing) - // empty Settings = clear all settings - // non-empty Settings = merge with existing; keys with null values are removed - final Settings mutableFragmentSettings = mutableWorkloadGroupFragment.getSettings(); - final Settings updatedSettings; - if (mutableFragmentSettings == null) { - // Not specified - keep existing - updatedSettings = Settings.builder().put(existingGroup.getSettings()).build(); - } else if (mutableFragmentSettings.isEmpty()) { - // Explicitly empty - clear all settings - updatedSettings = Settings.EMPTY; - } else { - // Merge: start with existing settings, overlay new values, remove null-valued keys - Settings.Builder builder = Settings.builder().put(existingGroup.getSettings()); - for (String key : mutableFragmentSettings.keySet()) { - String value = mutableFragmentSettings.get(key); - if (value == null) { - // null value means "clear this setting" - builder.remove(key); - } else { - builder.put(key, value); - } - } - updatedSettings = builder.build(); - } + final Settings updatedSettings = mergeSettings(existingGroup.getSettings(), mutableWorkloadGroupFragment.getSettings()); + final Settings updatedThrottling = mergeSettings( + existingGroup.getMutableWorkloadGroupFragment().getThrottling(), + mutableWorkloadGroupFragment.getThrottling() + ); return new WorkloadGroup( existingGroup.getName(), existingGroup.get_id(), - new MutableWorkloadGroupFragment(mode, updatedResourceLimits, updatedSettings), + new MutableWorkloadGroupFragment(mode, updatedResourceLimits, updatedSettings, updatedThrottling), Instant.now().getMillis() ); } + /** + * Drops null-valued keys from a settings bag before storage. A null value is the API gesture for "clear this key", + * which only carries meaning during an update merge (which consumes it); any that reach a persisted group, e.g. a + * null sent on create where there is nothing to clear, are dropped so stored config never contains a null value. + * + * @param s the settings to normalize (may be null) + * @return the settings with all null-valued keys removed, or empty if {@code s} is null + */ + private static Settings stripClearMarkers(Settings s) { + if (s == null) { + return Settings.EMPTY; + } + Settings.Builder builder = Settings.builder(); + for (String key : s.keySet()) { + String value = s.get(key); + if (value != null) { + builder.put(key, value); + } + } + return builder.build(); + } + + /** + * Merges an incoming settings bag from an update request onto the existing one: + * a null incoming bag (field absent) keeps existing, an empty incoming bag clears all, and a non-empty bag overlays + * its values with a per-key null value clearing that key. + * + * @param existing the currently stored settings + * @param incoming the settings from the update request (may be null) + * @return the merged settings + */ + private static Settings mergeSettings(Settings existing, Settings incoming) { + if (incoming == null) { + return Settings.builder().put(existing).build(); + } + if (incoming.isEmpty()) { + return Settings.EMPTY; + } + Settings.Builder builder = Settings.builder().put(existing); + for (String key : incoming.keySet()) { + String value = incoming.get(key); + if (value == null) { + builder.remove(key); + } else { + builder.put(key, value); + } + } + return builder.build(); + } + @Override public void writeTo(StreamOutput out) throws IOException { out.writeString(name); @@ -298,8 +365,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)) { - // "settings": null means clear all settings + if (fieldName.equals(MutableWorkloadGroupFragment.SETTINGS_STRING) + || fieldName.equals(MutableWorkloadGroupFragment.THROTTLING_STRING)) { mutableWorkloadGroupFragment1.parseField(parser, fieldName); } } diff --git a/server/src/main/java/org/opensearch/indices/IndicesService.java b/server/src/main/java/org/opensearch/indices/IndicesService.java index 2d36ec3cd686f..e34c2368b9f06 100644 --- a/server/src/main/java/org/opensearch/indices/IndicesService.java +++ b/server/src/main/java/org/opensearch/indices/IndicesService.java @@ -96,6 +96,7 @@ import org.opensearch.core.index.Index; import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.indices.breaker.CircuitBreakerService; +import org.opensearch.core.tasks.TaskId; import org.opensearch.core.util.FileSystemUtils; import org.opensearch.core.xcontent.MediaTypeRegistry; import org.opensearch.core.xcontent.NamedXContentRegistry; @@ -190,6 +191,7 @@ import org.opensearch.storage.slowlogs.TieredStorageSearchSlowLog; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.Client; +import org.opensearch.transport.client.ParentTaskAssigningClient; import java.io.Closeable; import java.io.IOException; @@ -2284,6 +2286,21 @@ public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis) { return getRewriteContext(nowInMillis, false); } + /** + * Returns a new {@link QueryRewriteContext} whose async actions issue their requests as children of + * {@code parentTaskId}. Rewriting can issue real requests -- a terms lookup with a subquery runs a search -- and + * without a parent those look like fresh top-level requests to anything that tracks the task tree. + */ + public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis, TaskId parentTaskId) { + return new BaseQueryRewriteContext( + xContentRegistry, + namedWriteableRegistry, + parentTaskId != null && parentTaskId.isSet() ? new ParentTaskAssigningClient(client, parentTaskId) : client, + nowInMillis, + false + ); + } + /** * Returns a new {@link QueryRewriteContext} for query validation with the given {@code now} provider */ diff --git a/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsAction.java b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsAction.java index 43d55613768b0..95bc28f491729 100644 --- a/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsAction.java +++ b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsAction.java @@ -192,6 +192,8 @@ protected Table createTableWithHeaders(PageToken pageToken, boolean verbose) { table.addCell("|"); table.addCell("TOTAL_CANCELLATIONS", verbose ? "desc:Total Canceled Queries" : ""); table.addCell("|"); + table.addCell("TOTAL_THROTTLED", verbose ? "desc:Total Throttled Queries" : ""); + table.addCell("|"); table.addCell("CPU_USAGE", verbose ? "desc:CPU Usage" : ""); table.addCell("|"); table.addCell("MEMORY_USAGE", verbose ? "desc:Memory Usage" : ""); @@ -213,6 +215,8 @@ protected void addRow(Table table, String nodeId, String workloadGroupId, Worklo table.addCell("|"); table.addCell(statsHolder.getCancellations()); table.addCell("|"); + table.addCell(statsHolder.getThrottled()); + table.addCell("|"); WorkloadGroupStats.ResourceStats cpuStats = statsHolder.getResourceStats().get(ResourceType.CPU); WorkloadGroupStats.ResourceStats memoryStats = statsHolder.getResourceStats().get(ResourceType.MEMORY); @@ -236,7 +240,7 @@ protected void addFooterRow(Table table, int COLUMN_COUNT) { * Builds a tabular response with '|' column separators. */ protected void buildTable(Table table, List paginatedStats, WlmPaginationStrategy paginationStrategy) { - final int COLUMN_COUNT = 13; + final int COLUMN_COUNT = 15; for (WlmStats wlmStats : paginatedStats) { String nodeId = wlmStats.getNode().getId(); diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java index 4729517917d22..e64e0d8fe59cd 100644 --- a/server/src/main/java/org/opensearch/search/SearchService.java +++ b/server/src/main/java/org/opensearch/search/SearchService.java @@ -79,6 +79,7 @@ import org.opensearch.core.index.Index; import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.indices.breaker.CircuitBreakerService; +import org.opensearch.core.tasks.TaskId; import org.opensearch.index.IndexNotFoundException; import org.opensearch.index.IndexService; import org.opensearch.index.IndexSettings; @@ -2020,6 +2021,15 @@ public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis, IndicesRe return new QueryCoordinatorContext(indicesService.getRewriteContext(nowInMillis), searchRequest); } + /** + * Returns a new {@link QueryCoordinatorContext} whose async rewrite actions issue their requests as children of + * {@code parentTaskId}. Query rewriting can issue real requests (a terms lookup with a subquery runs a search), and + * parenting them lets per-request admission control tell a nested request from a fresh one. + */ + public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis, IndicesRequest searchRequest, TaskId parentTaskId) { + return new QueryCoordinatorContext(indicesService.getRewriteContext(nowInMillis, parentTaskId), searchRequest); + } + /** * Returns a new {@link QueryCoordinatorContext} with the given {@code now} provider and {@link IndicesRequest searchRequest} */ diff --git a/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java b/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java index 0d767a5b9fa7b..54b83f1768b0f 100644 --- a/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java +++ b/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java @@ -36,11 +36,18 @@ public class MutableWorkloadGroupFragment extends AbstractDiffable resourceLimits; private Settings settings; + private Settings throttling; - public static final List acceptedFieldNames = List.of(RESILIENCY_MODE_STRING, RESOURCE_LIMITS_STRING, SETTINGS_STRING); + public static final List acceptedFieldNames = List.of( + RESILIENCY_MODE_STRING, + RESOURCE_LIMITS_STRING, + SETTINGS_STRING, + THROTTLING_STRING + ); public MutableWorkloadGroupFragment() {} @@ -52,11 +59,22 @@ public MutableWorkloadGroupFragment(ResiliencyMode resiliencyMode, Map resourceLimits, Settings settings) { + this(resiliencyMode, resourceLimits, settings, Settings.EMPTY); + } + + public MutableWorkloadGroupFragment( + ResiliencyMode resiliencyMode, + Map resourceLimits, + Settings settings, + Settings throttling + ) { validateResourceLimits(resourceLimits); WorkloadGroupSearchSettings.validate(settings); + WorkloadGroupThrottleSettings.validate(throttling); this.resiliencyMode = resiliencyMode; this.resourceLimits = resourceLimits; this.settings = settings != null ? settings : Settings.EMPTY; + this.throttling = throttling != null ? throttling : Settings.EMPTY; } public MutableWorkloadGroupFragment(StreamInput in) throws IOException { @@ -79,6 +97,17 @@ public MutableWorkloadGroupFragment(StreamInput in) throws IOException { } else { settings = Settings.EMPTY; } + // throttling is newer than settings, so it needs its own gate: a 3.7/3.8 peer writes only settings, and + // reading a throttling bag that was never written would desync the stream for every field after it. + // Decode "not on the wire" as null, not Settings.EMPTY: this class doubles as the partial update fragment, where + // an empty bag is the explicit "clear all throttling" gesture, so EMPTY here would make any update routed + // through a pre-3.9 node silently wipe the group's throttling config. Null means "field absent, keep existing"; + // WorkloadGroup's constructor normalizes it to EMPTY for a full object. + if (in.getVersion().onOrAfter(Version.V_3_9_0)) { + throttling = Settings.readOptionalSettingsFromStream(in); + } else { + throttling = null; + } } interface FieldParser { @@ -119,12 +148,25 @@ public Settings parseField(XContentParser parser) throws IOException { } } + static class ThrottlingParser implements FieldParser { + public Settings parseField(XContentParser parser) throws IOException { + // "throttling": null means clear all throttling (disable) + if (parser.currentToken() == XContentParser.Token.VALUE_NULL) { + return Settings.EMPTY; + } + Settings throttling = Settings.fromXContent(parser); + WorkloadGroupThrottleSettings.validate(throttling); + return throttling; + } + } + static class FieldParserFactory { static Optional> fieldParserFor(String fieldName) { return switch (fieldName) { case RESILIENCY_MODE_STRING -> Optional.of(new ResiliencyModeParser()); case RESOURCE_LIMITS_STRING -> Optional.of(new ResourceLimitsParser()); case SETTINGS_STRING -> Optional.of(new SearchSettingsParser()); + case THROTTLING_STRING -> Optional.of(new ThrottlingParser()); default -> Optional.empty(); }; } @@ -153,21 +195,55 @@ static Optional> fieldParserFor(String fieldName) { }, SETTINGS_STRING, (builder) -> { try { builder.startObject(SETTINGS_STRING); - Settings s = settings != null ? settings : Settings.EMPTY; - Map sortedSettingsMap = new TreeMap<>(); - for (String key : s.keySet()) { - sortedSettingsMap.put(key, s.get(key)); - } - for (Map.Entry e : sortedSettingsMap.entrySet()) { - builder.field(e.getKey(), e.getValue()); - } + writeSettingsFields(builder, settings); builder.endObject(); return null; } catch (IOException e) { throw new IllegalStateException("writing error encountered for the field " + SETTINGS_STRING); } + }, THROTTLING_STRING, (builder) -> { + try { + // Unlike settings (always emitted as {}), throttling is omitted entirely when unset. + Settings t = throttling != null ? throttling : Settings.EMPTY; + if (t.isEmpty() == false) { + builder.startObject(THROTTLING_STRING); + writeThrottlingFields(builder, t); + builder.endObject(); + } + return null; + } catch (IOException e) { + throw new IllegalStateException("writing error encountered for the field " + THROTTLING_STRING); + } }); + // Emits every stored throttling key, limits as JSON numbers. Iterating the bag rather than a hardcoded allowlist + // matters because this xContent is also the on-disk cluster-state format: a key that is accepted and stored but not + // emitted here would survive in a running cluster and then be silently lost on a full-cluster restart. + private static void writeThrottlingFields(XContentBuilder builder, Settings t) throws IOException { + Map sorted = new TreeMap<>(); + for (String key : t.keySet()) { + sorted.put(key, t.get(key)); + } + for (Map.Entry e : sorted.entrySet()) { + if (WorkloadGroupThrottleSettings.isLimitKey(e.getKey())) { + builder.field(e.getKey(), Integer.parseInt(e.getValue())); + } else { + builder.field(e.getKey(), e.getValue()); + } + } + } + + private static void writeSettingsFields(XContentBuilder builder, Settings s) throws IOException { + Settings source = s != null ? s : Settings.EMPTY; + Map sorted = new TreeMap<>(); + for (String key : source.keySet()) { + sorted.put(key, source.get(key)); + } + for (Map.Entry e : sorted.entrySet()) { + builder.field(e.getKey(), e.getValue()); + } + } + public static boolean shouldParse(String field) { return FieldParserFactory.fieldParserFor(field).isPresent(); } @@ -181,6 +257,7 @@ public void parseField(XContentParser parser, String field) { case RESILIENCY_MODE_STRING -> setResiliencyMode((ResiliencyMode) value); case RESOURCE_LIMITS_STRING -> setResourceLimits((Map) value); case SETTINGS_STRING -> setSettings((Settings) value); + case THROTTLING_STRING -> setThrottling((Settings) value); } } catch (IllegalArgumentException e) { throw e; @@ -210,6 +287,10 @@ public void writeTo(StreamOutput out) throws IOException { out.writeBoolean(false); out.writeMap(Map.of(), StreamOutput::writeString, StreamOutput::writeString); } + // Mirrors the read path: only 3.9+ peers expect a throttling bag on the wire. + if (out.getVersion().onOrAfter(Version.V_3_9_0)) { + Settings.writeOptionalSettingsToStream(throttling, out); + } } public static void validateResourceLimits(Map resourceLimits) { @@ -234,12 +315,13 @@ public boolean equals(Object o) { MutableWorkloadGroupFragment that = (MutableWorkloadGroupFragment) o; return Objects.equals(resiliencyMode, that.resiliencyMode) && Objects.equals(resourceLimits, that.resourceLimits) - && Objects.equals(settings, that.settings); + && Objects.equals(settings, that.settings) + && Objects.equals(throttling, that.throttling); } @Override public int hashCode() { - return Objects.hash(resiliencyMode, resourceLimits, settings); + return Objects.hash(resiliencyMode, resourceLimits, settings, throttling); } public ResiliencyMode getResiliencyMode() { @@ -254,6 +336,10 @@ public Settings getSettings() { return settings; } + public Settings getThrottling() { + return throttling; + } + /** * 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 @@ -299,4 +385,9 @@ void setSettings(Settings settings) { this.settings = settings != null ? settings : Settings.EMPTY; } + void setThrottling(Settings throttling) { + WorkloadGroupThrottleSettings.validate(throttling); + this.throttling = throttling != null ? throttling : Settings.EMPTY; + } + } diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java index 64c398e1d5e90..4abdb9c16b3ef 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java @@ -16,7 +16,9 @@ import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.lease.Releasable; import org.opensearch.common.lifecycle.AbstractLifecycleComponent; +import org.opensearch.common.settings.Settings; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import org.opensearch.monitor.jvm.JvmStats; import org.opensearch.monitor.process.ProcessProbe; @@ -37,6 +39,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.function.Consumer; import static org.opensearch.wlm.tracker.WorkloadGroupResourceUsageTrackerService.TRACKED_RESOURCES; @@ -59,6 +62,8 @@ public class WorkloadGroupService extends AbstractLifecycleComponent private final Set deletedWorkloadGroups; private final NodeDuressTrackers nodeDuressTrackers; private final WorkloadGroupsStateAccessor workloadGroupsStateAccessor; + // Node-local in-flight throttle counters, keyed by throttle bucket. No cross-node coordination in this tier. + private final WorkloadGroupThrottleTracker throttleTracker = new WorkloadGroupThrottleTracker(); public WorkloadGroupService( WorkloadGroupTaskCancellationService taskCancellationService, @@ -312,6 +317,167 @@ public void rejectIfNeeded(String workloadGroupId) { }); } + /** + * Group-and-principal seam over {@link #acquireThrottleOrReject(WorkloadGroupTask, Set)} for tests that want to + * exercise bucket resolution and the limit directly, without building a task and a thread context to carry the + * workload group id. Package-private on purpose: production callers go through the task-aware variant so the + * acquired bucket is recorded and re-entrancy is handled. + * + * @param workloadGroupId the workload group the request is assigned to + * @param principal the caller's joined principal tokens, or {@code null} (see resolver) + * @return a permit to close on request completion, or {@code null} if not throttled + * @throws OpenSearchRejectedExecutionException if the bucket is already at its node limit + */ + Releasable acquireThrottleOrReject(String workloadGroupId, String principal) { + return acquireThrottleOrReject(workloadGroupId, principal, Set.of(), bucketKey -> {}); + } + + /** + * Acquires one node-level throttle permit for the request, or returns {@code null} (nothing to release) when the + * request is not throttled: WLM disabled, default/unknown group, no {@code node_limit}, no resolvable bucket (see + * {@link #resolveThrottleAttributeValue}), or a bucket an ancestor task already holds. The bucket depends on the + * group's throttle {@code attribute}. + * + * @param task the request's task; its held bucket is recorded on a successful acquire + * @param bucketsHeldByAncestors buckets that ancestor tasks of this request already hold a permit for, so a nested + * coordinator search is not charged twice for its own request's bucket + * @return a permit to close on request completion, or {@code null} if not throttled + * @throws OpenSearchRejectedExecutionException if the bucket is already at its node limit + */ + public Releasable acquireThrottleOrReject(WorkloadGroupTask task, Set bucketsHeldByAncestors) { + return acquireThrottleOrReject( + task.getWorkloadGroupId(), + task.getThrottlePrincipal(), + bucketsHeldByAncestors, + task::setHeldThrottleBucket + ); + } + + private Releasable acquireThrottleOrReject( + String workloadGroupId, + String principal, + Set bucketsHeldByAncestors, + Consumer onAcquired + ) { + if (workloadManagementSettings.getWlmMode() != WlmMode.ENABLED) { + return null; + } + if (workloadGroupId == null || workloadGroupId.equals(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get())) { + return null; + } + try { + WorkloadGroup workloadGroup = getWorkloadGroupById(workloadGroupId); + if (workloadGroup == null) { + return null; + } + Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling(); + // Cheap early-out so a group that never configured throttling does not pay for parsing an absent limit on + // every search request. + if (throttling == null || throttling.isEmpty()) { + return null; + } + int nodeLimit = WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling); + if (nodeLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) { + return null; + } + String attribute = WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling); + // A null value means the request can't be attributed (e.g. username/role with no principal) -> fail open. + String attributeValue = resolveThrottleAttributeValue(attribute, principal); + if (attributeValue == null) { + return null; + } + String bucketKey = workloadGroupId + ":" + attribute + ":" + attributeValue; + + // Re-entrancy. A coordinator search can issue a nested coordinator search on this same node while holding + // this bucket's permit: a terms lookup with a subquery does exactly that during the rewrite phase, and the + // nested request inherits the same workloadGroupId (the thread context is not stashed) so it resolves to + // the same bucket. Charging it a second permit makes the request compete with itself -- with node_limit=N, + // N such requests would all be rejected at precisely the configured concurrency. The ancestor already paid + // for this bucket, so admit the nested request without a second permit. + if (bucketsHeldByAncestors.contains(bucketKey)) { + return null; + } + + Releasable permit = throttleTracker.tryAcquire(bucketKey, nodeLimit); + if (permit != null) { + onAcquired.accept(bucketKey); + return permit; + } + + // Over the limit. Name the group and the throttle dimension so both the log line and the 429 identify who + // was throttled -- the bucket key alone is opaque to an operator. + String target = "workload group [" + workloadGroup.getName() + "]"; + if ("group".equals(attribute) == false) { + target += " for " + attribute + " [" + attributeValue + "]"; + } + if (workloadGroup.getResiliencyMode() == MutableWorkloadGroupFragment.ResiliencyMode.MONITOR) { + // MONITOR observes only: log that the request WOULD have been rejected, then admit it without touching + // total_throttled, consistent with MONITOR being dormant on the cancellation path. DEBUG, not INFO: + // this fires once per would-be-throttled request, so INFO would spam a hot bucket under load. + logger.debug( + "Request would be throttled (monitor mode, not rejected): {} reached its per-node limit of {} concurrent requests.", + target, + nodeLimit + ); + return null; + } + // Record the rejection without ever letting a stats failure swallow the 429. Use the raw state map, not the + // DEFAULT-fallback accessor, so a not-yet-registered group isn't misattributed to DEFAULT. + try { + WorkloadGroupState workloadGroupState = workloadGroupsStateAccessor.getWorkloadGroupStateMap().get(workloadGroupId); + if (workloadGroupState != null) { + workloadGroupState.totalThrottled.inc(); + } + } catch (Exception statsException) { + logger.warn("Failed to record throttle stat for workload group [" + workloadGroupId + "]", statsException); + } + throw new OpenSearchRejectedExecutionException( + "Request throttled: " + target + " reached its per-node limit of " + nodeLimit + " concurrent requests." + ); + } catch (OpenSearchRejectedExecutionException e) { + throw e; // the intended 429 + } catch (Exception e) { + // A bug in the throttle path must never fail an otherwise-valid search, so fail open. DEBUG, not WARN: a + // deterministic failure in here would otherwise emit a stack trace at the full query rate. + logger.debug(() -> "Skipping node-level throttle for workload group [" + workloadGroupId + "] due to an error", e); + return null; + } + } + + /** + * Resolves the value the throttle bucket is keyed by: the literal {@code "group"} for whole-group throttling, or + * the principal's {@code username} / {@code role} subfield value. + *

+ * A principal may carry several values for one subfield (a user in many roles). The request is charged to exactly + * one of them, chosen as the lexicographically smallest so the bucket is stable: picking whichever value the + * extractor happened to emit first would let the same user land in different buckets on different requests, and so + * draw more than one allowance. + * + * @return the attribute value, or {@code null} to fail open (not throttled) when the principal is absent or has no + * usable value for the subfield + */ + private String resolveThrottleAttributeValue(String attribute, String principal) { + if ("group".equals(attribute)) { + return "group"; + } + if (principal == null || principal.isEmpty()) { + return null; + } + // Trim the token, not the value: trimming past the delimiter would fold "username|alice " into alice's bucket. + String subfieldPrefix = attribute + "|"; + String selected = null; + for (String token : principal.split(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER)) { + String trimmed = token.trim(); + if (trimmed.startsWith(subfieldPrefix)) { + String value = trimmed.substring(subfieldPrefix.length()); + if (value.isEmpty() == false && (selected == null || value.compareTo(selected) < 0)) { + selected = value; + } + } + } + return selected; + } + private double getNormalisedRejectionThreshold(double limit, ResourceType resourceType) { if (resourceType == ResourceType.CPU) { return limit * workloadManagementSettings.getNodeLevelCpuRejectionThreshold(); diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java index 636e9178775f9..0bccbc366665b 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java @@ -30,10 +30,14 @@ public class WorkloadGroupTask extends CancellableTask { private static final Logger logger = LogManager.getLogger(WorkloadGroupTask.class); public static final String WORKLOAD_GROUP_ID_HEADER = "workloadGroupId"; + /** Separator between the {@code subfield|value} principal tokens carried by {@link #getThrottlePrincipal()}. */ + public static final String WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER = "\u001F"; public static final Supplier DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER = () -> "DEFAULT_WORKLOAD_GROUP"; private final LongSupplier nanoTimeSupplier; private String workloadGroupId; private boolean isWorkloadGroupSet = false; + private volatile String throttlePrincipal; + private volatile String heldThrottleBucket; 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); @@ -90,6 +94,53 @@ public final void setWorkloadGroupId(final ThreadContext threadContext) { } } + /** + * Records the caller's principal for {@code username}/{@code role} throttling: {@code subfield|value} tokens + * (e.g. {@code username|alice}) joined by {@link #WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER}. Set on the coordinator + * by the WLM auto-tagging action filter, from the security plugin's principal extractor, before the action executes. + *

+ * Deliberately held on the task rather than in the {@link ThreadContext}: a ThreadContext request header is + * serialized onto every outgoing transport request, which would ship the caller's identity to every shard and to + * remote clusters in a cross-cluster search even though only the coordinator reads it. A task field is also not + * something a client can supply, and is naturally per-request, so concurrent sub-requests sharing one thread context + * (an {@code _msearch}) cannot collide. + * + * @param throttlePrincipal the joined principal tokens, or {@code null} when no extractor is installed + */ + public void setThrottlePrincipal(final String throttlePrincipal) { + this.throttlePrincipal = throttlePrincipal; + } + + /** + * The caller's principal for throttle bucket resolution, or {@code null} when unknown, in which case + * {@code username}/{@code role} throttling fails open. + */ + public String getThrottlePrincipal() { + return throttlePrincipal; + } + + /** + * Records the throttle bucket this task successfully took a permit for, so a nested coordinator search issued + * while this one is in flight can recognise that its bucket is already paid for and skip admission. Set by + * {@code WorkloadGroupService#acquireThrottleOrReject} on a successful acquire only. + *

+ * Deliberately not cleared on release: the value is scoped to the task, which is unregistered when the request + * finishes, and a rewrite round that issues a nested search after the outer permit has been released must still + * be recognised as nested rather than charged a fresh permit. + * + * @param heldThrottleBucket the bucket key a permit is held for + */ + public void setHeldThrottleBucket(final String heldThrottleBucket) { + this.heldThrottleBucket = heldThrottleBucket; + } + + /** + * The throttle bucket this task holds (or held) a permit for, or {@code null} if it was never throttled. + */ + public String getHeldThrottleBucket() { + return heldThrottleBucket; + } + public long getElapsedTime() { return nanoTimeSupplier.getAsLong() - getStartTimeNanos(); } diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleSettings.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleSettings.java new file mode 100644 index 0000000000000..6c506c2337428 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleSettings.java @@ -0,0 +1,158 @@ +/* + * 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.math.BigInteger; +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Registry of valid workload group throttle settings with their validators. Throttle config is a nested + * {@code throttling} object (a {@link Settings} bag) like {@code settings}, so per-key null clears a field and an + * absent key keeps the existing value with no extra bookkeeping. + */ +@ExperimentalApi +public class WorkloadGroupThrottleSettings { + + /** Sentinel for an unset limit, matching the {@code -1 = not set} convention of {@code WLM_SEARCH_TIMEOUT}. */ + public static final int UNSET_LIMIT = -1; + + /** Upper bound for a limit. Limits are stored in an int-backed setting, so a larger value would overflow the {@code int}. */ + public static final int MAX_LIMIT = Integer.MAX_VALUE; + + /** Dimension the limit is keyed by: {@code group} (whole group) or per {@code username} / {@code role}. No default: unset when absent. */ + public static final Setting ATTRIBUTE = Setting.simpleString("attribute"); + + /** Per-node in-flight allowance admitted locally with no coordination. {@code -1} means unset. */ + public static final Setting NODE_LIMIT = Setting.intSetting("node_limit", UNSET_LIMIT, UNSET_LIMIT); + + /** + * Allowed attribute values; {@code username} / {@code role} map to the security {@code principal.*} attributes at + * enforcement. Ordered so validation errors enumerate them the same way every time ({@code Set.of} iteration order + * varies between JVM runs). + */ + public static final Set ALLOWED_ATTRIBUTES = Collections.unmodifiableSet( + new LinkedHashSet<>(List.of("group", "username", "role")) + ); + + private static final Map> REGISTERED_SETTINGS = Map.of( + ATTRIBUTE.getKey(), + ATTRIBUTE, + NODE_LIMIT.getKey(), + NODE_LIMIT + ); + + private WorkloadGroupThrottleSettings() { + throw new UnsupportedOperationException("Utility class"); + } + + /** True for keys holding an integer limit, which xContent must emit as a JSON number rather than a string. */ + static boolean isLimitKey(String key) { + return NODE_LIMIT.getKey().equals(key); + } + + /** + * Per-key validation: every key must be registered, {@code attribute} must be an allowed value, and each limit + * must be a non-negative integer no greater than {@link #MAX_LIMIT} ({@code -1} is the internal "unset" sentinel and + * is not user-settable). Safe to run on a partial fragment from an update request; the cross-field checks live in + * {@link #validateMergedConfig(Settings)}. + * + * @param throttling the throttling settings to validate + * @throws IllegalArgumentException if any key is unknown or any value is invalid + */ + public static void validate(Settings throttling) { + if (throttling == null) { + return; + } + for (String key : throttling.keySet()) { + String value = throttling.get(key); + if (REGISTERED_SETTINGS.containsKey(key) == false) { + throw new IllegalArgumentException("Unknown throttle setting: " + key); + } + // null value means "clear this key" — skip value validation + if (value == null) { + continue; + } + // Limits are stored with -1 as the internal "unset" sentinel, but a user may only send a non-negative integer. + if (NODE_LIMIT.getKey().equals(key)) { + validateUserLimit(key, value); + } + } + String attribute = throttling.get(ATTRIBUTE.getKey()); + if (attribute != null && ALLOWED_ATTRIBUTES.contains(attribute) == false) { + throw new IllegalArgumentException( + "throttling.attribute must be one of " + ALLOWED_ATTRIBUTES + " but was '" + attribute + "'" + ); + } + } + + // Rejects a user-supplied limit that is not a non-negative integer in [0, MAX_LIMIT]; -1 is reserved as the internal + // "unset" sentinel. Parsed as a BigInteger so a value that is a well-formed integer but too large for the int-backed + // setting (e.g. Integer.MAX_VALUE + 1) reports an overflow instead of being mislabelled as "not an integer". + private static void validateUserLimit(String key, String value) { + final BigInteger parsed; + try { + parsed = new BigInteger(value); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("throttling." + key + " must be an integer but was '" + value + "'"); + } + if (parsed.signum() < 0) { + throw new IllegalArgumentException("throttling." + key + " must be non-negative but was " + parsed); + } + if (parsed.compareTo(BigInteger.valueOf(MAX_LIMIT)) > 0) { + throw new IllegalArgumentException("throttling." + key + " must not exceed " + MAX_LIMIT + " but was " + parsed); + } + } + + /** + * Cross-field validation on a fully-merged throttling config. A limit may only be set alongside an attribute + * (a limit with no attribute is meaningless), and when throttling is configured the effective ceiling + * {@code max(0, node_limit)} must be at least 1, since a ceiling of 0 rejects every request. Must be called on + * the merged result, not a partial update fragment. + * + * @param throttling the merged throttling settings + * @throws IllegalArgumentException if a limit is set without an attribute, or the effective ceiling is 0 + */ + public static void validateMergedConfig(Settings throttling) { + if (throttling == null || throttling.isEmpty()) { + return; + } + boolean hasAttribute = throttling.hasValue(ATTRIBUTE.getKey()); + boolean hasNode = throttling.hasValue(NODE_LIMIT.getKey()); + + if (hasNode && hasAttribute == false) { + throw new IllegalArgumentException("throttling.attribute is required when a throttle limit is set"); + } + + // An attribute on its own configures nothing, so say that rather than reporting a zero ceiling: no limit was + // ever set, so nothing "would reject all requests". + if (hasNode == false) { + throw new IllegalArgumentException( + "throttling.node_limit is required when throttling.attribute is set; " + + "set throttling as null to disable throttling instead" + ); + } + int node = NODE_LIMIT.get(throttling); + if (node < 1) { + throw new IllegalArgumentException( + "Effective throttle ceiling is 0 (node_limit=" + + node + + "); this would reject all requests. " + + "Set node_limit to a positive value, or set throttling as null to disable throttling" + ); + } + } +} diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java new file mode 100644 index 0000000000000..4a173471732df --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java @@ -0,0 +1,107 @@ +/* + * 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 java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Tracks the number of in-flight requests per throttle bucket on a single node and enforces a per-node cap. + *

+ * A bucket is identified by an opaque key (see {@code WorkloadGroupService} for how the key is built from a + * workload group and its throttle attribute). A counter exists only while a bucket has at least one in-flight + * request: it is created on first acquire and removed when it drains back to zero, so memory scales with the + * number of concurrently active buckets rather than the total population of users/roles. + *

+ * This tier is fully local — no cross-node coordination — mirroring the acquire/rollback + {@link Releasable} + * release shape of {@link org.opensearch.index.IndexingPressure}. + */ +@ExperimentalApi +public class WorkloadGroupThrottleTracker { + + private final Map inFlightByBucket = new ConcurrentHashMap<>(); + + /** + * Attempts to admit one request into the given bucket under the per-node limit. + *

+ * Returns {@code null} rather than throwing when the bucket is full: the caller decides what a breach means + * (reject, or observe-only in MONITOR mode), and building a rejection exception here would mean allocating and + * filling in a stack trace on the search hot path only to discard it in the observe-only case. + * + * @param bucketKey the throttle bucket identifier + * @param nodeLimit the maximum concurrent in-flight requests this node may admit for the bucket + * @return a {@link Releasable} that decrements the bucket's in-flight count exactly once when closed, or + * {@code null} if the bucket is already at the limit + */ + public Releasable tryAcquire(String bucketKey, int nodeLimit) { + // Create-and-increment inside compute() so the counter this acquire is about to use cannot be removed by a + // concurrent release between lookup and increment (see release() for why removal is safe). + // observed holds this thread's own post-increment count, captured under the map's per-key lock. Checking that + // instead of a later counter.get() keeps the decision exact: a get() could see an unrelated concurrent + // acquire's increment and reject a request that was actually within the limit. + final int[] observed = new int[1]; + AtomicInteger counter = inFlightByBucket.compute(bucketKey, (k, existing) -> { + AtomicInteger c = existing != null ? existing : new AtomicInteger(0); + observed[0] = c.incrementAndGet(); + return c; + }); + if (observed[0] > nodeLimit) { + // Over the cap: roll back this increment and report the breach to the caller. + release(bucketKey, counter); + return null; + } + return releaseOnce(bucketKey, counter); + } + + /** + * Current in-flight count for a bucket, or 0 if the bucket has no active requests. Package-private for tests. + * Note this returns 0 both for an absent bucket and for one that is present with a zero count; use + * {@link #bucketCount()} to distinguish them. + */ + int inFlight(String bucketKey) { + AtomicInteger counter = inFlightByBucket.get(bucketKey); + return counter == null ? 0 : counter.get(); + } + + /** + * Number of buckets currently holding a counter. Package-private for tests, which use it to assert that a bucket + * is actually evicted once it drains rather than merely reading back as zero. + */ + int bucketCount() { + return inFlightByBucket.size(); + } + + // Wraps release in a one-shot guard so a double close (e.g. onRequestEnd and onRequestFailure) decrements once. + private Releasable releaseOnce(String bucketKey, AtomicInteger counter) { + AtomicBoolean released = new AtomicBoolean(false); + return () -> { + if (released.compareAndSet(false, true)) { + release(bucketKey, counter); + } + }; + } + + // Decrements the bucket and removes the map entry once it drains to 0. + // + // The decrement deliberately sits outside the compute(), so decrement-and-remove is NOT atomic as a pair. What + // makes removal safe is an invariant instead: decrements are 1:1 with prior increments (each permit closes at most + // once, and the over-limit path rolls back its own increment), so the counter equals outstanding permits plus + // pending rollbacks and is therefore >= 1 while any permit is held. Removal requires <= 0, so it can only happen + // when no permit is outstanding and no acquire can be orphaned. A refactor that adds a second decrement site, or + // that keys removal on identity rather than the <= 0 check, breaks this silently. + private void release(String bucketKey, AtomicInteger counter) { + counter.decrementAndGet(); + inFlightByBucket.compute(bucketKey, (k, existing) -> (existing != null && existing.get() <= 0) ? null : existing); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupsStateAccessor.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupsStateAccessor.java index 582730bbf0b33..b77ddbc40d3cb 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupsStateAccessor.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupsStateAccessor.java @@ -10,20 +10,20 @@ import org.opensearch.wlm.stats.WorkloadGroupState; -import java.util.HashMap; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; /** * This class is used to decouple {@link WorkloadGroupService} and {@link org.opensearch.wlm.cancellation.WorkloadGroupTaskCancellationService} to share the * {@link WorkloadGroupState}s */ public class WorkloadGroupsStateAccessor { - // This map does not need to be concurrent since we will process the cluster state change serially and update - // this map with new additions and deletions of entries. WorkloadGroupState is thread safe + // Concurrent: structural updates happen on the cluster-applier thread while request threads read concurrently + // (throttle admission, stat updates, cancellation). WorkloadGroupState is itself thread safe. private final Map workloadGroupStateMap; public WorkloadGroupsStateAccessor() { - this(new HashMap<>()); + this(new ConcurrentHashMap<>()); } public WorkloadGroupsStateAccessor(Map workloadGroupStateMap) { @@ -39,10 +39,15 @@ public Map getWorkloadGroupStateMap() { /** * return WorkloadGroupState for the given workloadGroupId - * @param workloadGroupId + * @param workloadGroupId may be null when a request carried no workload group header * @return WorkloadGroupState for the given workloadGroupId, if id is invalid return default workload group state */ public WorkloadGroupState getWorkloadGroupState(String workloadGroupId) { + // The backing map is a ConcurrentHashMap, which rejects a null key, and an untagged request legitimately has + // no id (e.g. the failure listener reads the header unconditionally). Fall back to DEFAULT instead of throwing. + if (workloadGroupId == null) { + return workloadGroupStateMap.get(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get()); + } return workloadGroupStateMap.getOrDefault( workloadGroupId, workloadGroupStateMap.get(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get()) 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 a3715eb72f385..9efc2dbf7a5d6 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupState.java +++ b/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupState.java @@ -38,6 +38,11 @@ public class WorkloadGroupState { */ public final CounterMetric totalCancellations = new CounterMetric(); + /** + * This will track the cumulative requests throttled (rejected by the node-level in-flight throttle) in the workload group since the OpenSearch start time + */ + public final CounterMetric totalThrottled = new CounterMetric(); + /** * This is used to store the resource type state both for CPU and MEMORY */ @@ -80,6 +85,14 @@ public long getTotalCancellations() { return totalCancellations.count(); } + /** + * + * @return requests throttled in the workload group + */ + public long getTotalThrottled() { + return totalThrottled.count(); + } + /** * 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 1174424ed398e..0bad5eb8d6ad6 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupStats.java +++ b/server/src/main/java/org/opensearch/wlm/stats/WorkloadGroupStats.java @@ -8,6 +8,7 @@ package org.opensearch.wlm.stats; +import org.opensearch.Version; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.common.io.stream.Writeable; @@ -95,10 +96,12 @@ public static class WorkloadGroupStatsHolder implements ToXContentObject, Writea public static final String REJECTIONS = "total_rejections"; public static final String TOTAL_CANCELLATIONS = "total_cancellations"; public static final String FAILURES = "failures"; + public static final String THROTTLED = "total_throttled"; private long completions; private long rejections; private long failures; private long cancellations; + private long throttled; private Map resourceStats; // this is needed to support the factory method @@ -109,12 +112,14 @@ public WorkloadGroupStatsHolder( long rejections, long failures, long cancellations, + long throttled, Map resourceStats ) { this.completions = completions; this.rejections = rejections; this.failures = failures; this.cancellations = cancellations; + this.throttled = throttled; this.resourceStats = resourceStats; } @@ -123,6 +128,12 @@ public WorkloadGroupStatsHolder(StreamInput in) throws IOException { this.rejections = in.readVLong(); this.failures = in.readVLong(); this.cancellations = in.readVLong(); + // total_throttled arrives with throttling in 3.9, so it must be gated on that version and not on the + // older gate used by fields that already shipped: a 3.7/3.8 peer never writes it, and reading it anyway + // would consume the resourceStats map header and desync everything after it. + if (in.getVersion().onOrAfter(Version.V_3_9_0)) { + this.throttled = in.readVLong(); + } this.resourceStats = in.readMap((i) -> ResourceType.fromName(i.readString()), ResourceStats::new); } @@ -138,6 +149,10 @@ public long getCancellations() { return cancellations; } + public long getThrottled() { + return throttled; + } + public Map getResourceStats() { return resourceStats; } @@ -160,6 +175,7 @@ public static WorkloadGroupStatsHolder from(WorkloadGroupState workloadGroupStat statsHolder.rejections = workloadGroupState.getTotalRejections(); statsHolder.failures = workloadGroupState.getFailures(); statsHolder.cancellations = workloadGroupState.getTotalCancellations(); + statsHolder.throttled = workloadGroupState.getTotalThrottled(); statsHolder.resourceStats = resourceStatsMap; return statsHolder; } @@ -175,6 +191,10 @@ public static void writeTo(StreamOutput out, WorkloadGroupStatsHolder statsHolde out.writeVLong(statsHolder.rejections); out.writeVLong(statsHolder.failures); out.writeVLong(statsHolder.cancellations); + // version-gated to match the StreamInput ctor; read/write gates and order must stay in sync. + if (out.getVersion().onOrAfter(Version.V_3_9_0)) { + out.writeVLong(statsHolder.throttled); + } out.writeMap(statsHolder.resourceStats, (o, val) -> o.writeString(val.getName()), ResourceStats::writeTo); } @@ -190,6 +210,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.field(REJECTIONS, rejections); // builder.field(FAILURES, failures); builder.field(TOTAL_CANCELLATIONS, cancellations); + builder.field(THROTTLED, throttled); for (ResourceType resourceType : ResourceType.getSortedValues()) { ResourceStats resourceStats1 = resourceStats.get(resourceType); @@ -210,12 +231,13 @@ public boolean equals(Object o) { && rejections == that.rejections && Objects.equals(resourceStats, that.resourceStats) && failures == that.failures - && cancellations == that.cancellations; + && cancellations == that.cancellations + && throttled == that.throttled; } @Override public int hashCode() { - return Objects.hash(completions, rejections, cancellations, failures, resourceStats); + return Objects.hash(completions, rejections, cancellations, failures, throttled, 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 d1e13546935b2..e3c02197c7cc9 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 @@ -46,6 +46,7 @@ public class WlmStatsResponseTests extends OpenSearchTestCase { 0, 1, 0, + 0, Map.of( ResourceType.CPU, new WorkloadGroupStats.ResourceStats(0, 0, 0), @@ -80,6 +81,7 @@ public void testToString() { + " \"total_completions\" : 0,\n" + " \"total_rejections\" : 0,\n" + " \"total_cancellations\" : 0,\n" + + " \"total_throttled\" : 0,\n" + " \"cpu\" : {\n" + " \"current_usage\" : 0.0,\n" + " \"cancellations\" : 0,\n" diff --git a/server/src/test/java/org/opensearch/action/pagination/WlmPaginationStrategyTests.java b/server/src/test/java/org/opensearch/action/pagination/WlmPaginationStrategyTests.java index 9b7d3347664e8..69ff8eaee7c5f 100644 --- a/server/src/test/java/org/opensearch/action/pagination/WlmPaginationStrategyTests.java +++ b/server/src/test/java/org/opensearch/action/pagination/WlmPaginationStrategyTests.java @@ -240,7 +240,7 @@ public void testFindIndex_found() { WorkloadGroupStats.ResourceStats dummyStats = new WorkloadGroupStats.ResourceStats(0.1, 2, 3); Map resourceMap = Map.of(ResourceType.CPU, dummyStats); - WorkloadGroupStats.WorkloadGroupStatsHolder holder = new WorkloadGroupStats.WorkloadGroupStatsHolder(1, 2, 3, 4, resourceMap); + WorkloadGroupStats.WorkloadGroupStatsHolder holder = new WorkloadGroupStats.WorkloadGroupStatsHolder(1, 2, 3, 4, 5, resourceMap); Map groupStats = new HashMap<>(); groupStats.put("group-1", holder); diff --git a/server/src/test/java/org/opensearch/action/search/TransportSearchActionTests.java b/server/src/test/java/org/opensearch/action/search/TransportSearchActionTests.java index e1f515d9f5c19..6ce224b5a0911 100644 --- a/server/src/test/java/org/opensearch/action/search/TransportSearchActionTests.java +++ b/server/src/test/java/org/opensearch/action/search/TransportSearchActionTests.java @@ -113,6 +113,7 @@ import org.opensearch.transport.TransportRequestOptions; import org.opensearch.transport.TransportService; import org.opensearch.transport.client.node.NodeClient; +import org.opensearch.wlm.WorkloadGroupService; import java.util.ArrayList; import java.util.Arrays; @@ -1251,7 +1252,8 @@ public void testResolveIndices() { new SearchRequestOperationsCompositeListenerFactory(), mock(Tracer.class), mock(TaskResourceTrackingService.class), - mock(IndicesService.class) + mock(IndicesService.class), + mock(WorkloadGroupService.class) ); // Actual test cases start here: 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 0e57b739a5cd2..c883d2271101d 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/WorkloadGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/WorkloadGroupTests.java @@ -8,9 +8,13 @@ package org.opensearch.cluster.metadata; +import org.opensearch.Version; import org.opensearch.common.UUIDs; +import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.settings.Settings; import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; +import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; @@ -19,6 +23,7 @@ import org.opensearch.wlm.MutableWorkloadGroupFragment; import org.opensearch.wlm.MutableWorkloadGroupFragment.ResiliencyMode; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadGroupThrottleSettings; import org.joda.time.Instant; import java.io.IOException; @@ -37,7 +42,19 @@ static WorkloadGroup createRandomWorkloadGroup(String _id) { String name = randomAlphaOfLength(10); Map resourceLimit = new HashMap<>(); resourceLimit.put(ResourceType.MEMORY, randomDoubleBetween(0.0, 0.80, false)); - return new WorkloadGroup(name, _id, new MutableWorkloadGroupFragment(randomMode(), resourceLimit), Instant.now().getMillis()); + // Generate a valid throttling config: either disabled (empty), or enabled with a required attribute plus + // a positive node_limit (so the effective ceiling is >= 1). + Settings.Builder throttling = Settings.builder(); + if (randomBoolean()) { + throttling.put("attribute", randomFrom("group", "username", "role")); + throttling.put("node_limit", randomIntBetween(1, 100)); + } + return new WorkloadGroup( + name, + _id, + new MutableWorkloadGroupFragment(randomMode(), resourceLimit, Settings.EMPTY, throttling.build()), + Instant.now().getMillis() + ); } private static ResiliencyMode randomMode() { @@ -376,6 +393,375 @@ public void testUpdateOverrideRequestValuesPersistsThroughMerge() { assertEquals("true", updated.getSettings().get("override_request_values")); } + public void testToXContentOmitsUnsetThrottling() throws IOException { + WorkloadGroup workloadGroup = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5), Settings.EMPTY), + System.currentTimeMillis() + ); + XContentBuilder builder = JsonXContent.contentBuilder(); + workloadGroup.toXContent(builder, ToXContent.EMPTY_PARAMS); + assertFalse(builder.toString().contains("throttling")); + } + + public void testToXContentEmitsThrottling() throws IOException { + long currentTimeInMillis = Instant.now().getMillis(); + String workloadGroupId = UUIDs.randomBase64UUID(); + Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 10).build(); + WorkloadGroup workloadGroup = new WorkloadGroup( + "TestWorkloadGroup", + workloadGroupId, + new MutableWorkloadGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.30), Settings.EMPTY, throttling), + 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}," + + "\"updated_at\":%d}", + workloadGroupId, + currentTimeInMillis + ); + assertEquals(expected, builder.toString()); + } + + public void testNegativeThrottleLimitRejected() { + // -1 is the internal "unset" sentinel and, like any negative value, is not user-settable. + for (int badLimit : new int[] { -1, -2 }) { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("node_limit", badLimit).build() + ) + ); + assertTrue(exception.getMessage().contains("node_limit must be non-negative")); + } + } + + public void testThrottleLimitExceedingMaxRejected() { + // Integer.MAX_VALUE + 1: a well-formed non-negative integer, but too large for the int-backed setting. The error + // must call out the overflow rather than falsely claiming it is "not an integer". + String tooLarge = Long.toString((long) Integer.MAX_VALUE + 1); + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", tooLarge).build() + ) + ); + assertTrue(exception.getMessage().contains("node_limit must not exceed " + Integer.MAX_VALUE)); + assertTrue(exception.getMessage().contains(tooLarge)); + } + + public void testThrottleLimitAtMaxAccepted() { + // Integer.MAX_VALUE is the largest limit the int-backed setting can hold and must be accepted. + WorkloadGroup workloadGroup = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", Integer.MAX_VALUE).build() + ), + System.currentTimeMillis() + ); + Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling(); + assertEquals(Integer.valueOf(Integer.MAX_VALUE), WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling)); + } + + public void testNonNumericThrottleLimitRejected() { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", "not_a_number").build() + ) + ); + assertTrue(exception.getMessage().contains("node_limit must be an integer")); + } + + public void testInvalidThrottleAttributeRejected() { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "index").put("node_limit", 5).build() + ) + ); + assertTrue(exception.getMessage().contains("throttling.attribute must be one of")); + } + + public void testUnknownThrottleKeyRejected() { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("bogus_limit", 5).build() + ) + ); + assertTrue(exception.getMessage().contains("Unknown throttle setting")); + } + + public void testZeroEffectiveCeilingRejected() { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", 0).build() + ), + System.currentTimeMillis() + ) + ); + assertTrue(exception.getMessage().contains("Effective throttle ceiling is 0")); + } + + public void testAttributeWithoutLimitRejected() { + // An attribute alone configures nothing, so the error must say a limit is missing rather than report a zero + // ceiling (nothing was set, so nothing "would reject all requests"), and must never leak the -1 sentinel. + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").build() + ), + System.currentTimeMillis() + ) + ); + assertTrue(exception.getMessage().contains("throttling.node_limit is required when throttling.attribute is set")); + assertFalse(exception.getMessage().contains("-1")); + assertFalse(exception.getMessage().contains("ceiling")); + } + + public void testLimitWithoutAttributeRejected() { + // A throttle limit requires an attribute; a limit with no attribute is rejected. + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("node_limit", 5).build() + ), + System.currentTimeMillis() + ) + ); + assertTrue(exception.getMessage().contains("throttling.attribute is required")); + } + + public void testWholeGroupThrottleWithExplicitAttribute() { + // attribute has no default; whole-group throttling must be requested explicitly with attribute=group. + WorkloadGroup workloadGroup = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "group").put("node_limit", 5).build() + ), + System.currentTimeMillis() + ); + Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling(); + assertEquals("group", WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling)); + assertEquals(Integer.valueOf(5), WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling)); + } + + public void testUpdateMergesThrottling() { + WorkloadGroup existing = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", 10).build() + ), + System.currentTimeMillis() + ); + + // Update only node_limit — the absent attribute key should keep its existing value + MutableWorkloadGroupFragment updateFragment = new MutableWorkloadGroupFragment( + null, + Map.of(), + Settings.EMPTY, + Settings.builder().put("node_limit", 50).build() + ); + + WorkloadGroup updated = WorkloadGroup.updateExistingWorkloadGroup(existing, updateFragment); + Settings throttling = updated.getMutableWorkloadGroupFragment().getThrottling(); + assertEquals("username", WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling)); + assertEquals(Integer.valueOf(50), WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling)); + } + + public void testUpdateWithNullClearsThrottleKeys() throws IOException { + // Clearing every throttle key individually is equivalent to disabling throttling: the merge consumes each null + // and the bag collapses to empty. Clearing only node_limit is rejected instead, because that would leave an + // attribute with no limit, which configures nothing; "throttling": null is the way to disable one key at a time. + String json = "{\"resource_limits\":{\"memory\":0.5},\"throttling\":{\"attribute\":null,\"node_limit\":null}}"; + XContentParser parser = createParser(JsonXContent.jsonXContent, json); + MutableWorkloadGroupFragment clearAll = WorkloadGroup.Builder.fromXContent(parser).getMutableWorkloadGroupFragment(); + + WorkloadGroup updated = WorkloadGroup.updateExistingWorkloadGroup(throttledGroup(), clearAll); + assertTrue(updated.getMutableWorkloadGroupFragment().getThrottling().isEmpty()); + + String clearLimitOnly = "{\"resource_limits\":{\"memory\":0.5},\"throttling\":{\"node_limit\":null}}"; + XContentParser limitParser = createParser(JsonXContent.jsonXContent, clearLimitOnly); + MutableWorkloadGroupFragment clearLimit = WorkloadGroup.Builder.fromXContent(limitParser).getMutableWorkloadGroupFragment(); + + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> WorkloadGroup.updateExistingWorkloadGroup(throttledGroup(), clearLimit) + ); + assertTrue(exception.getMessage().contains("throttling.node_limit is required when throttling.attribute is set")); + } + + public void testUpdateFromPreThrottlingPeerPreservesThrottling() throws IOException { + // A pre-3.9 node has no throttling field, so it writes none. Decoding "absent" as an empty bag would make + // mergeSettings treat it as the explicit "clear all" gesture and silently delete the group's throttling on an + // update that never mentioned throttling. + MutableWorkloadGroupFragment update = new MutableWorkloadGroupFragment( + ResiliencyMode.SOFT, + Map.of(), + Settings.EMPTY, + Settings.EMPTY + ); + MutableWorkloadGroupFragment asSeenByCurrentNode = copyWriteable( + update, + new NamedWriteableRegistry(Collections.emptyList()), + MutableWorkloadGroupFragment::new, + Version.V_3_8_0 + ); + + WorkloadGroup updated = WorkloadGroup.updateExistingWorkloadGroup(throttledGroup(), asSeenByCurrentNode); + Settings throttling = updated.getMutableWorkloadGroupFragment().getThrottling(); + assertEquals("username", WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling)); + assertEquals(Integer.valueOf(10), WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling)); + assertEquals(ResiliencyMode.SOFT, updated.getResiliencyMode()); + } + + public void testDeserializationAcceptsThrottlingThisNodeConsidersInvalid() throws IOException { + // Cluster state published by a newer node may use throttling rules this node does not know (e.g. a second limit + // key, making node_limit optional). Rejecting it here would wedge the node out of the cluster instead of + // failing one API call, so the deserialization path must accept it. + WorkloadGroup valid = throttledGroup(); + BytesStreamOutput out = new BytesStreamOutput(); + out.writeString(valid.getName()); + out.writeString(valid.get_id()); + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + // attribute with no limit: rejected on the API path, must be tolerated on the wire + Settings.builder().put("attribute", "username").build() + ).writeTo(out); + out.writeLong(System.currentTimeMillis()); + + StreamInput in = out.bytes().streamInput(); + WorkloadGroup deserialized = new WorkloadGroup(in); + assertEquals( + "username", + WorkloadGroupThrottleSettings.ATTRIBUTE.get(deserialized.getMutableWorkloadGroupFragment().getThrottling()) + ); + + // The same config through the API path is still rejected. + expectThrows( + IllegalArgumentException.class, + () -> new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").build() + ), + System.currentTimeMillis() + ) + ); + } + + private static WorkloadGroup throttledGroup() { + return new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", 10).build() + ), + System.currentTimeMillis() + ); + } + + public void testUpdateWithNullThrottlingObjectDisables() throws IOException { + WorkloadGroup existing = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", 10).build() + ), + System.currentTimeMillis() + ); + + // "throttling": null disables throttling entirely + String json = "{\"resource_limits\":{\"memory\":0.5},\"throttling\":null}"; + XContentParser parser = createParser(JsonXContent.jsonXContent, json); + MutableWorkloadGroupFragment updateFragment = WorkloadGroup.Builder.fromXContent(parser).getMutableWorkloadGroupFragment(); + + WorkloadGroup updated = WorkloadGroup.updateExistingWorkloadGroup(existing, updateFragment); + assertTrue(updated.getMutableWorkloadGroupFragment().getThrottling().isEmpty()); + } + + public void testCreateDropsNullThrottleValues() throws IOException { + // On create there is nothing to clear, so null-valued keys are dropped rather than persisted; an + // all-null throttling object therefore collapses to empty (disabled) instead of hitting a ceiling error. + WorkloadGroup allNull = parseCreate( + "{\"resiliency_mode\":\"enforced\",\"resource_limits\":{\"memory\":0.5}," + "\"throttling\":{\"node_limit\":null}}" + ); + Settings throttling = allNull.getMutableWorkloadGroupFragment().getThrottling(); + assertTrue(throttling.isEmpty()); + assertFalse(throttling.keySet().contains("node_limit")); // raw check: null-valued key was dropped, not persisted + } + + private WorkloadGroup parseCreate(String json) throws IOException { + XContentParser parser = createParser(JsonXContent.jsonXContent, json); + return WorkloadGroup.Builder.fromXContent(parser).name("test")._id("test_id").updatedAt(System.currentTimeMillis()).build(); + } + public void testSettingsNullFromXContentClearsSettings() throws IOException { // Simulate parsing {"settings": null} via XContent String json = "{\"_id\":\"test_id\",\"name\":\"test\",\"resiliency_mode\":\"enforced\"," diff --git a/server/src/test/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsActionTests.java b/server/src/test/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsActionTests.java index 2f30181c263a7..a4cd8cf86aeb8 100644 --- a/server/src/test/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsActionTests.java +++ b/server/src/test/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsActionTests.java @@ -102,14 +102,15 @@ public void testCreateTableWithHeaders() { public void testAddRow() { Table table = action.createTableWithHeaders(null, true); - WorkloadGroupStats.WorkloadGroupStatsHolder stats = new WorkloadGroupStats.WorkloadGroupStatsHolder(1, 2, 3, 4, new HashMap<>()); + WorkloadGroupStats.WorkloadGroupStatsHolder stats = new WorkloadGroupStats.WorkloadGroupStatsHolder(1, 2, 3, 4, 5, new HashMap<>()); action.addRow(table, "node1", "group1", stats); assertEquals(1, table.getRows().size()); } public void testAddFooterRow() { Table table = action.createTableWithHeaders(null, true); - action.addFooterRow(table, 13); + // Derive the width from the headers rather than hardcoding it, so adding a column cannot silently desync. + action.addFooterRow(table, table.getHeaders().size()); assertEquals(1, table.getRows().size()); } @@ -122,7 +123,7 @@ public void testBuildTable() { statsMap.put(ResourceType.CPU, cpuStats); statsMap.put(ResourceType.MEMORY, memoryStats); - WorkloadGroupStats.WorkloadGroupStatsHolder statsHolder = new WorkloadGroupStats.WorkloadGroupStatsHolder(1, 2, 3, 4, statsMap); + WorkloadGroupStats.WorkloadGroupStatsHolder statsHolder = new WorkloadGroupStats.WorkloadGroupStatsHolder(1, 2, 3, 4, 5, statsMap); Map groupStats = new HashMap<>(); groupStats.put("groupA", statsHolder); WorkloadGroupStats stats = new WorkloadGroupStats(groupStats); diff --git a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java index 95e7ae5384f49..870f8b70f50f2 100644 --- a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java +++ b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java @@ -251,6 +251,7 @@ import org.opensearch.transport.TransportService; import org.opensearch.transport.client.AdminClient; import org.opensearch.transport.client.node.NodeClient; +import org.opensearch.wlm.WorkloadGroupService; import org.junit.After; import org.junit.Before; @@ -2410,7 +2411,8 @@ public void onFailure(final Exception e) { searchRequestOperationsCompositeListenerFactory, NoopTracer.INSTANCE, new TaskResourceTrackingService(settings, clusterSettings, threadPool), - mockIndicesService + mockIndicesService, + mock(WorkloadGroupService.class) ) ); actions.put( diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java index 8da689ab2bc89..47f619152462e 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java @@ -14,10 +14,12 @@ import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.lease.Releasable; import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.core.tasks.TaskId; import org.opensearch.search.backpressure.trackers.NodeDuressTrackers; import org.opensearch.tasks.Task; import org.opensearch.test.OpenSearchTestCase; @@ -479,6 +481,329 @@ public void testGetCurrentWorkloadGroupReturnsNullWhenGroupMissing() { assertNull(workloadGroupService.getCurrentWorkloadGroup()); } + private void stubClusterStateWithGroup(WorkloadGroup wg) { + ClusterState clusterState = Mockito.mock(ClusterState.class); + Metadata metadata = Mockito.mock(Metadata.class); + when(mockClusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + when(metadata.workloadGroups()).thenReturn(Map.of(wg.get_id(), wg)); + } + + private WorkloadGroup throttledGroup(String id, Settings throttling) { + return throttledGroup(id, throttling, MutableWorkloadGroupFragment.ResiliencyMode.ENFORCED); + } + + private WorkloadGroup throttledGroup(String id, Settings throttling, MutableWorkloadGroupFragment.ResiliencyMode mode) { + return new WorkloadGroup( + id + "-name", + id, + new MutableWorkloadGroupFragment(mode, Map.of(ResourceType.MEMORY, 0.5), Settings.EMPTY, throttling), + 1L + ); + } + + public void testAcquireThrottleAdmitsNestedRequestWithoutASecondPermit() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + // The outer request takes the group's only permit and records the bucket it holds. + WorkloadGroupTask outer = throttleTask("wg-1"); + Releasable outerPermit = workloadGroupService.acquireThrottleOrReject(outer, Set.of()); + assertNotNull(outerPermit); + String heldBucket = outer.getHeldThrottleBucket(); + assertNotNull("a successful acquire must record the bucket on the task", heldBucket); + + // A nested coordinator search (a terms lookup with a subquery runs one during the rewrite phase) resolves to + // the same bucket. At node_limit=1 charging it again would 429 the request that spawned it, so it is admitted + // with no permit of its own -- null means "nothing to release", not "throttled". + WorkloadGroupTask nested = throttleTask("wg-1"); + assertNull(workloadGroupService.acquireThrottleOrReject(nested, Set.of(heldBucket))); + assertNull("an exempt request must not record a bucket it never took", nested.getHeldThrottleBucket()); + assertEquals( + "an exemption is not a throttle", + 0, + mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled() + ); + + // The exemption is scoped to the inherited bucket only: an independent request still hits the limit, so this + // cannot silently disable throttling for the group. + WorkloadGroupTask independent = throttleTask("wg-1"); + expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject(independent, Set.of())); + // An ancestor holding some *other* bucket must not exempt this request either. + WorkloadGroupTask unrelatedAncestor = throttleTask("wg-1"); + expectThrows( + OpenSearchRejectedExecutionException.class, + () -> workloadGroupService.acquireThrottleOrReject(unrelatedAncestor, Set.of("wg-1:group:something-else")) + ); + + outerPermit.close(); + } + + private WorkloadGroupTask throttleTask(String workloadGroupId) { + WorkloadGroupTask task = new WorkloadGroupTask(1L, "transport", "Search", "test task", TaskId.EMPTY_TASK_ID, Map.of()); + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + threadContext.putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, workloadGroupId); + task.setWorkloadGroupId(threadContext); + return task; + } + + public void testAcquireThrottleReturnsNullWhenNodeLimitUnset() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + stubClusterStateWithGroup(throttledGroup("wg-1", Settings.EMPTY)); // throttling not configured + assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); + } + + public void testAcquireThrottleReturnsNullWhenWlmDisabled() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED); + assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); + } + + public void testAcquireThrottleRejectsAtLimitAndIncrementsStat() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + Releasable permit = workloadGroupService.acquireThrottleOrReject("wg-1", null); // first admit succeeds + assertNotNull(permit); + // second admit hits node_limit of 1 -> 429 + total_throttled incremented + expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject("wg-1", null)); + assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); + + // releasing the first permit frees the slot so a subsequent acquire succeeds + permit.close(); + assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); + } + + public void testAcquireThrottleMonitorModeObservesWithoutRejecting() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling, MutableWorkloadGroupFragment.ResiliencyMode.MONITOR)); + + assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); // first admit takes the only slot + // A MONITOR group observes only: an over-limit request is admitted (null permit, nothing to release) rather + // than rejected, and the would-be rejection is not counted. + assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); + assertEquals(0, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); + } + + public void testAcquireThrottleSoftModeStillRejects() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling, MutableWorkloadGroupFragment.ResiliencyMode.SOFT)); + + assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); + // Only MONITOR is observe-only; SOFT enforces the throttle like ENFORCED does. + expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject("wg-1", null)); + assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); + } + + public void testAcquireThrottleUsernameKeepsPerUserBuckets() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + // alice takes her single slot; a second alice request is rejected. + Releasable alice = workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice"); + assertNotNull(alice); + expectThrows( + OpenSearchRejectedExecutionException.class, + () -> workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice") + ); + assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); + + // bob is a different bucket, so he is admitted even while alice is at her limit. + Releasable bob = workloadGroupService.acquireThrottleOrReject("wg-1", "username|bob"); + assertNotNull(bob); + + // releasing alice frees her bucket + alice.close(); + assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice")); + } + + public void testAcquireThrottleUsernameWithCommaDoesNotCollide() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + String delim = WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER; + // principal for user "a,b" with a role token appended + String userAB = "username|a,b" + delim + "role|admin"; + // user "a" is a genuinely different principal + String userA = "username|a"; + + Releasable ab = workloadGroupService.acquireThrottleOrReject("wg-1", userAB); // fills "a,b" bucket + assertNotNull(ab); + // user "a" must NOT be treated as the same bucket as "a,b" -> still admitted + assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", userA)); + // a second "a,b" request hits the "a,b" bucket limit -> rejected + expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject("wg-1", userAB)); + } + + public void testAcquireThrottleRolePicksMatchingSubfieldFromMultiTokenPrincipal() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "role").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + // A principal header may carry both subfields; the role bucket must key off the role token only. + String delim = WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER; + assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice" + delim + "role|admin")); + expectThrows( + OpenSearchRejectedExecutionException.class, + () -> workloadGroupService.acquireThrottleOrReject("wg-1", "username|bob" + delim + "role|admin") + ); + assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); + } + + public void testAcquireThrottleRoleBucketIsStableAcrossTokenOrder() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "role").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + // A user in several roles must land in one deterministic bucket. If the resolver took whichever role token came + // first, the same user would draw a fresh allowance whenever the extractor changed its ordering. + String delim = WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER; + assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", "role|admin" + delim + "role|analyst")); + expectThrows( + OpenSearchRejectedExecutionException.class, + () -> workloadGroupService.acquireThrottleOrReject("wg-1", "role|analyst" + delim + "role|admin") + ); + assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); + } + + public void testThrottleRejectionNamesGroupAndAttribute() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice")); + OpenSearchRejectedExecutionException e = expectThrows( + OpenSearchRejectedExecutionException.class, + () -> workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice") + ); + // The operator (and the caller) must be able to tell which group and which principal was throttled. + assertTrue(e.getMessage(), e.getMessage().contains("workload group [wg-1-name]")); + assertTrue(e.getMessage(), e.getMessage().contains("username [alice]")); + assertTrue(e.getMessage(), e.getMessage().contains("per-node limit of 1")); + } + + public void testThrottleRejectionForWholeGroupOmitsAttributeClause() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); + OpenSearchRejectedExecutionException e = expectThrows( + OpenSearchRejectedExecutionException.class, + () -> workloadGroupService.acquireThrottleOrReject("wg-1", null) + ); + assertTrue(e.getMessage(), e.getMessage().contains("workload group [wg-1-name]")); + assertFalse(e.getMessage(), e.getMessage().contains(" for group ")); + } + + public void testIncrementFailuresForUntaggedRequestDoesNotThrow() { + // The search failure listener reads the workload group header unconditionally, so it legitimately passes null + // for an untagged request. The state map is a ConcurrentHashMap, which rejects a null key. + workloadGroupService.incrementFailuresFor(null); + assertEquals( + 1, + mockWorkloadGroupsStateAccessor.getWorkloadGroupState(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get()).getFailures() + ); + } + + public void testAcquireThrottleFailsOpenWhenPrincipalMissingForUsername() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + // No principal (e.g. security plugin not installed) or no matching subfield -> not throttled (fail open). + assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); + assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", "")); + assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", "role|admin")); // no username token + assertEquals(0, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); + } + + /** + * A failure while recording the total_throttled stat must NOT swallow the rejection and admit the over-limit + * request. Whether the state map lookup returns null (group not yet registered / just deleted) or throws, the + * 429 must still propagate. + */ + public void testAcquireThrottleStillRejectsWhenStatUpdateFails() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + // state map with no entry for wg-1 (as during the state-registration lag) -> raw get(id) returns null + WorkloadGroupsStateAccessor emptyMapAccessor = Mockito.mock(WorkloadGroupsStateAccessor.class); + when(emptyMapAccessor.getWorkloadGroupStateMap()).thenReturn(new HashMap<>()); + WorkloadGroupService serviceWithNullState = new WorkloadGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + emptyMapAccessor, + new HashSet<>(), + new HashSet<>() + ); + + assertNotNull(serviceWithNullState.acquireThrottleOrReject("wg-1", null)); // first admit fills the single slot + // second acquire is over the limit; a null state must not let the stat update swallow the 429 + expectThrows(OpenSearchRejectedExecutionException.class, () -> serviceWithNullState.acquireThrottleOrReject("wg-1", null)); + + // accessor whose state-map lookup throws must also still propagate the 429 + WorkloadGroupsStateAccessor throwingStateAccessor = Mockito.mock(WorkloadGroupsStateAccessor.class); + when(throwingStateAccessor.getWorkloadGroupStateMap()).thenThrow(new RuntimeException("state map race")); + WorkloadGroupService serviceWithThrowingState = new WorkloadGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + throwingStateAccessor, + new HashSet<>(), + new HashSet<>() + ); + + assertNotNull(serviceWithThrowingState.acquireThrottleOrReject("wg-1", null)); // fills the single slot + expectThrows(OpenSearchRejectedExecutionException.class, () -> serviceWithThrowingState.acquireThrottleOrReject("wg-1", null)); + } + + /** + * During the state-registration lag a node can enforce a new group's limit before its clusterChanged() registers + * the state. The rejection stat must not be misattributed to the DEFAULT group in that window. + */ + public void testAcquireThrottleDoesNotMisattributeToDefaultDuringRegistrationLag() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); + stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); + + // DEFAULT group state exists, but wg-1 is NOT yet registered (registration lag). + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get()); + + assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); // fills the single slot + expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject("wg-1", null)); + + // the rejection must NOT have landed on the DEFAULT group + assertEquals( + 0, + mockWorkloadGroupsStateAccessor.getWorkloadGroupState(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get()) + .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/WorkloadGroupTaskTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupTaskTests.java index 341f31993f800..0085e0f8da19d 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupTaskTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupTaskTests.java @@ -41,4 +41,22 @@ public void testSuccessfulSetWorkloadGroupId() { sut.setWorkloadGroupId(threadPool.getThreadContext()); assertEquals("akfanglkaglknag2332", sut.getWorkloadGroupId()); } + + public void testThrottlePrincipalIsPerTaskAndNotAHeader() { + assertNull("an absent principal must read as null so username/role throttling fails open", sut.getThrottlePrincipal()); + + sut.setThrottlePrincipal("username|alice"); + assertEquals("username|alice", sut.getThrottlePrincipal()); + + // An _msearch runs every sub-request through the filter chain on one thread context, so holding the principal per + // task is what keeps one sub-request's caller from being billed to another's throttle bucket. + WorkloadGroupTask other = new WorkloadGroupTask(124, "transport", "Search", "test task", null, Collections.emptyMap()); + assertNull(other.getThrottlePrincipal()); + + // The principal must stay out of the header maps: a ThreadContext request header is serialized onto every + // outgoing transport request, which would ship the caller's identity to every shard and to remote clusters in a + // cross-cluster search even though only the coordinator reads it. + assertNull(sut.getHeader("workloadGroupPrincipal")); + assertNull(threadPool.getThreadContext().getHeader("workloadGroupPrincipal")); + } } diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java new file mode 100644 index 0000000000000..214394ceb552b --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java @@ -0,0 +1,166 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.common.lease.Releasable; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +public class WorkloadGroupThrottleTrackerTests extends OpenSearchTestCase { + + public void testAcquireUnderLimitSucceeds() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + Releasable p1 = tracker.tryAcquire("bucket", 2); + Releasable p2 = tracker.tryAcquire("bucket", 2); + assertNotNull(p1); + assertNotNull(p2); + assertEquals(2, tracker.inFlight("bucket")); + p1.close(); + p2.close(); + } + + public void testAcquireAtLimitIsRefused() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + assertNotNull(tracker.tryAcquire("bucket", 1)); + // Over the cap the tracker reports the breach by returning null; building the 429 is the caller's job. + assertNull(tracker.tryAcquire("bucket", 1)); + // a refused acquire must not leave the count inflated + assertEquals(1, tracker.inFlight("bucket")); + } + + public void testExactlyNAdmittedForLimitAboveOne() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + int limit = randomIntBetween(2, 5); + List permits = new ArrayList<>(); + for (int i = 0; i < limit; i++) { + Releasable p = tracker.tryAcquire("bucket", limit); + assertNotNull("acquire " + i + " of " + limit + " must be admitted", p); + permits.add(p); + } + assertEquals(limit, tracker.inFlight("bucket")); + assertNull("the limit+1'th acquire must be refused", tracker.tryAcquire("bucket", limit)); + assertEquals(limit, tracker.inFlight("bucket")); + permits.forEach(Releasable::close); + assertEquals(0, tracker.bucketCount()); + } + + public void testReleaseFreesAPermit() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + Releasable p = tracker.tryAcquire("bucket", 1); + assertNull(tracker.tryAcquire("bucket", 1)); // at the limit + p.close(); + // permit freed -> a fresh acquire now succeeds + Releasable p2 = tracker.tryAcquire("bucket", 1); + assertNotNull(p2); + assertEquals(1, tracker.inFlight("bucket")); + p2.close(); + } + + public void testDrainToZeroRemovesBucketThenReacquire() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + Releasable p = tracker.tryAcquire("bucket", 5); + assertEquals(1, tracker.inFlight("bucket")); + assertEquals(1, tracker.bucketCount()); + p.close(); + // bucketCount, not inFlight: inFlight returns 0 for an absent bucket AND for one still present at zero, so only + // bucketCount actually proves the entry was evicted. Without this the memory bound is untested. + assertEquals(0, tracker.bucketCount()); + // re-acquiring after the bucket drained (and was removed) works and starts from 1 + Releasable p2 = tracker.tryAcquire("bucket", 5); + assertEquals(1, tracker.inFlight("bucket")); + p2.close(); + assertEquals(0, tracker.bucketCount()); + } + + public void testRefusedAcquireDoesNotLeaveAnEmptyBucketBehind() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + assertNull(tracker.tryAcquire("bucket", 0)); + // The rollback must remove the entry it created, otherwise a stream of refused requests for distinct buckets + // would accumulate zero-valued entries forever. + assertEquals(0, tracker.bucketCount()); + } + + public void testReleaseIsIdempotent() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + // Two permits so the bucket survives the first close: with only one, the entry is evicted and a buggy second + // decrement would land on an orphaned counter that inFlight() can no longer see, making the test vacuous. + Releasable p1 = tracker.tryAcquire("bucket", 5); + Releasable p2 = tracker.tryAcquire("bucket", 5); + p1.close(); + p1.close(); // double close must not decrement twice + assertEquals(1, tracker.inFlight("bucket")); + p2.close(); + assertEquals(0, tracker.bucketCount()); + } + + public void testBucketsAreIndependent() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + assertNotNull(tracker.tryAcquire("a", 1)); + assertNull(tracker.tryAcquire("a", 1)); // "a" is full but "b" is a separate bucket + Releasable pb = tracker.tryAcquire("b", 1); + assertNotNull(pb); + assertEquals(1, tracker.inFlight("a")); + assertEquals(1, tracker.inFlight("b")); + pb.close(); + } + + /** + * Hammers one bucket from many threads. Asserts the two properties the whole design rests on and that no + * single-threaded test can show: the limit is never exceeded, and every entry is evicted once the dust settles. + */ + public void testConcurrentAcquireNeverExceedsLimitAndFullyDrains() throws Exception { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + final int limit = 4; + final int threads = 8; + final int iterations = 200; + final AtomicInteger concurrentlyHeld = new AtomicInteger(); + final AtomicInteger maxObserved = new AtomicInteger(); + final CountDownLatch start = new CountDownLatch(1); + final CountDownLatch done = new CountDownLatch(threads); + final List failures = new ArrayList<>(); + + for (int t = 0; t < threads; t++) { + Thread thread = new Thread(() -> { + try { + start.await(); + for (int i = 0; i < iterations; i++) { + Releasable p = tracker.tryAcquire("hot", limit); + if (p != null) { + int held = concurrentlyHeld.incrementAndGet(); + maxObserved.accumulateAndGet(held, Math::max); + concurrentlyHeld.decrementAndGet(); + p.close(); + } + } + } catch (Throwable e) { + synchronized (failures) { + failures.add(e); + } + } finally { + done.countDown(); + } + }); + thread.start(); + } + start.countDown(); + assertTrue("threads did not finish in time", done.await(60, TimeUnit.SECONDS)); + + synchronized (failures) { + assertTrue("worker threads threw: " + failures, failures.isEmpty()); + } + assertTrue("admitted " + maxObserved.get() + " concurrently for a limit of " + limit, maxObserved.get() <= limit); + assertEquals("counter did not return to zero", 0, tracker.inFlight("hot")); + assertEquals("bucket entry was not evicted", 0, tracker.bucketCount()); + } +} diff --git a/server/src/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java b/server/src/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java index 31071d7acf1c3..ac61ed88e77b6 100644 --- a/server/src/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java +++ b/server/src/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java @@ -107,6 +107,7 @@ public void testValidWorkloadGroupRequestFailure() throws IOException { 0, 1, 0, + 0, Map.of( ResourceType.CPU, new WorkloadGroupStats.ResourceStats(0, 0, 0), @@ -120,6 +121,7 @@ public void testValidWorkloadGroupRequestFailure() throws IOException { 0, 0, 0, + 0, Map.of( ResourceType.CPU, new WorkloadGroupStats.ResourceStats(0, 0, 0), @@ -182,6 +184,7 @@ public void testMultiThreadedValidWorkloadGroupRequestFailures() { 0, ITERATIONS, 0, + 0, Map.of( ResourceType.CPU, new WorkloadGroupStats.ResourceStats(0, 0, 0), @@ -195,6 +198,7 @@ public void testMultiThreadedValidWorkloadGroupRequestFailures() { 0, 0, 0, + 0, Map.of( ResourceType.CPU, new WorkloadGroupStats.ResourceStats(0, 0, 0), @@ -217,6 +221,7 @@ public void testInvalidWorkloadGroupFailure() throws IOException { 0, 0, 0, + 0, Map.of( ResourceType.CPU, new WorkloadGroupStats.ResourceStats(0, 0, 0), @@ -230,6 +235,7 @@ public void testInvalidWorkloadGroupFailure() throws IOException { 0, 1, 0, + 0, Map.of( ResourceType.CPU, new WorkloadGroupStats.ResourceStats(0, 0, 0), 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 5589db7c0c20d..9e1bcfedd6999 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java @@ -39,6 +39,7 @@ public void testToXContent() throws IOException { 13, 2, 0, + 5, Map.of(ResourceType.CPU, new WorkloadGroupStats.ResourceStats(0.3, 13, 2)) ) ); @@ -49,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,\"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,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", builder.toString() ); } 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 d7d77761aa9fa..ea4f69c821185 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStatsTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/WorkloadGroupStatsTests.java @@ -38,6 +38,7 @@ public void testToXContent() throws IOException { 13, 2, 0, + 5, Map.of(ResourceType.CPU, new WorkloadGroupStats.ResourceStats(0.3, 13, 2)) ) ); @@ -47,11 +48,44 @@ 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,\"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,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", builder.toString() ); } + public void testThrottledIsVersionGatedAndKeepsOlderStreamsAligned() throws IOException { + WorkloadGroupStats original = new WorkloadGroupStats( + Map.of( + "group-1", + new WorkloadGroupStats.WorkloadGroupStatsHolder( + 100, + 13, + 2, + 7, + 5, + Map.of(ResourceType.CPU, new WorkloadGroupStats.ResourceStats(0.3, 11, 4)) + ) + ) + ); + + // A 3.9 peer exchanges total_throttled, and everything after it on the wire stays aligned. + WorkloadGroupStats.WorkloadGroupStatsHolder current = copyInstance(original, Version.V_3_9_0).getStats().get("group-1"); + assertEquals(5, current.getThrottled()); + assertEquals(100, current.getCompletions()); + assertEquals(0.3, current.getResourceStats().get(ResourceType.CPU).getCurrentUsage(), 0.0); + + // A pre-throttling peer never writes total_throttled, so it must read back as 0 and -- the actual hazard -- + // the resourceStats map that follows it must still deserialize instead of being consumed as the throttled slot. + WorkloadGroupStats.WorkloadGroupStatsHolder legacy = copyInstance(original, Version.V_3_8_0).getStats().get("group-1"); + assertEquals(0, legacy.getThrottled()); + assertEquals(100, legacy.getCompletions()); + assertEquals(13, legacy.getRejections()); + assertEquals(7, legacy.getCancellations()); + assertEquals(0.3, legacy.getResourceStats().get(ResourceType.CPU).getCurrentUsage(), 0.0); + assertEquals(11, legacy.getResourceStats().get(ResourceType.CPU).getCancellations()); + assertEquals(4, legacy.getResourceStats().get(ResourceType.CPU).getRejections()); + } + @Override protected Writeable.Reader instanceReader() { return WorkloadGroupStats::new; @@ -67,6 +101,7 @@ protected WorkloadGroupStats createTestInstance() { randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong(), + randomNonNegativeLong(), Map.of( ResourceType.CPU, new WorkloadGroupStats.ResourceStats(