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(