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..f3c3c92599c3e --- /dev/null +++ b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmNodeThrottlingIT.java @@ -0,0 +1,472 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm; + +import org.apache.logging.log4j.LogManager; +import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest; +import org.opensearch.action.index.IndexResponse; +import org.opensearch.action.search.SearchRequestBuilder; +import org.opensearch.action.support.WriteRequest; +import org.opensearch.cluster.metadata.WorkloadGroup; +import org.opensearch.common.action.ActionFuture; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType; +import org.opensearch.plugins.Plugin; +import org.opensearch.plugins.PluginsService; +import org.opensearch.rule.RuleAttribute; +import org.opensearch.rule.RuleFrameworkPlugin; +import org.opensearch.rule.RulePersistenceServiceRegistry; +import org.opensearch.rule.RuleRoutingServiceRegistry; +import org.opensearch.rule.action.CreateRuleAction; +import org.opensearch.rule.action.CreateRuleRequest; +import org.opensearch.rule.autotagging.AutoTaggingRegistry; +import org.opensearch.rule.autotagging.FeatureType; +import org.opensearch.rule.autotagging.Rule; +import org.opensearch.script.MockScriptPlugin; +import org.opensearch.script.Script; +import org.opensearch.script.ScriptType; +import org.opensearch.search.lookup.LeafFieldsLookup; +import org.opensearch.test.OpenSearchIntegTestCase; +import org.opensearch.wlm.MutableWorkloadGroupFragment; +import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.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); + 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); + + // Second search while the first is still in-flight: must be rejected (429). + Throwable rejection = expectThrows(Throwable.class, () -> blockingSearch(indexName).get()); + 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, so the in-flight search + // gauge reflects only the one still-blocked search (not two). This guards against the gauge leak where a + // throttle rejection increments 'current' via onRequestStart but never reaches onRequestEnd/onRequestFailure. + assertEquals("in-flight search gauge must exclude the throttle-rejected request", 1L, currentInFlightSearches()); + + // 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 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).get()); + assertTrue( + "Expected an OpenSearchRejectedExecutionException in the cause chain but was: " + rejection, + hasRejectedExecutionCause(rejection) + ); + assertEquals("total_throttled should increment by exactly one", throttledBefore + 1, getThrottled(workloadGroupId)); + + // 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 { + 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(); + return extractStatField(response.toString(), groupId, "total_completions"); + } + + private int getThrottled(String groupId) 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(); + return extractStatField(response.toString(), groupId, "total_throttled"); + } + + /** + * 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; + } + + private int extractStatField(String responseBody, String workloadGroupId, String fieldName) { + int total = 0; + String groupKey = "\"" + workloadGroupId + "\""; + String field = "\"" + fieldName + "\""; + int index = 0; + while ((index = responseBody.indexOf(groupKey, index)) != -1) { + int groupStart = responseBody.indexOf("{", index); + int fieldIndex = responseBody.indexOf(field, groupStart); + if (fieldIndex == -1) break; + int colonIndex = responseBody.indexOf(":", fieldIndex); + int commaIndex = responseBody.indexOf(",", colonIndex); + String numberStr = responseBody.substring(colonIndex + 1, commaIndex).trim(); + total += Integer.parseInt(numberStr); + index = commaIndex; + } + return total; + } + + private SearchRequestBuilder blockingSearch(String indexName) { + return client().prepareSearch(indexName) + .setQuery(scriptQuery(new Script(ScriptType.INLINE, "mockscript", ScriptedBlockPlugin.SCRIPT_NAME, Collections.emptyMap()))); + } + + // Injects the resolved principal directly into the thread-context header (WORKLOAD_GROUP_PRINCIPAL_HEADER). + // In production the WLM auto-tagging filter sets this header from the security plugin's principal extractor; here we + // stand in for that so the core username/role bucket-keying can be exercised without a real security plugin. + private org.opensearch.transport.client.Client clientAs(String username) { + return client().filterWithHeader( + Map.of(org.opensearch.wlm.WorkloadGroupTask.WORKLOAD_GROUP_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("wlm.workload_group.mode", mode); + ClusterUpdateSettingsRequest request = new ClusterUpdateSettingsRequest().persistentSettings(settings); + 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()); + } + } + + /** + * 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 { + assertBusy(() -> assertFalse(shouldBlock.get())); + } 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..d8603ef072df7 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,32 @@ public LogicalOperator getLogicalOperator() { } } + AttributeExtractor principalExtractor = 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()); + principalExtractor = attributeExtensions.get(attribute).getAttributeExtractor(); + attributeExtractors.add(principalExtractor); } Optional label = ruleProcessingService.evaluateLabel(attributeExtractors); label.ifPresent(s -> threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, s)); + // Propagate the principal so core-side throttling can build per-username / per-role buckets. + if (principalExtractor != null) { + stashPrincipal(principalExtractor); + } chain.proceed(task, action, request, listener); } + + /** + * Stashes the extractor's {@code subfield|value} principal tokens (joined by + * {@link WorkloadGroupTask#WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER}) into the thread context for core-side + * throttling. No header is set when the extractor yields nothing, leaving username/role throttling to fail open. + */ + private void stashPrincipal(AttributeExtractor principalExtractor) { + String principal = String.join(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER, principalExtractor.extract()); + if (principal.isEmpty() == false) { + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER, principal); + } + } } 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..e456718737d1b 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,7 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.core.action.ActionResponse; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.plugin.wlm.spi.AttributeExtractorExtension; import org.opensearch.rule.InMemoryRuleProcessingService; import org.opensearch.rule.RuleAttribute; import org.opensearch.rule.attribute_extractor.AttributeExtractor; @@ -38,6 +39,7 @@ import java.util.Map; import java.util.Optional; +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 +102,92 @@ public void testApplyForInValidRequest() { verify(ruleProcessingService, times(0)).evaluateLabel(anyList()); } + public void testApplyStashesPrincipalHeaderWhenExtractorPresent() { + // 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); + try (ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext()) { + when(svc.evaluateLabel(anyList())).thenReturn(Optional.of("QG")); + filter.apply(mock(Task.class), "Test", request, ActionRequestMetadata.empty(), null, chain); + + // Both principal tokens are joined (by WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER) into the header for + // core-side throttling. + assertEquals( + "username|alice" + WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER + "role|admin", + threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER) + ); + } + } + + public void testApplyDoesNotStashPrincipalHeaderWhenNoExtractor() { + // Default filter from setUp has no principal attribute/extractor -> header must be absent. + 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")); + autoTaggingActionFilter.apply(mock(Task.class), "Test", request, ActionRequestMetadata.empty(), null, chain); + assertNull(threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER)); + } + } + public void testApplyForScrollRequestWithOriginalIndices() { SearchScrollRequest request = mock(SearchScrollRequest.class); ActionFilterChain chain = mock(TestActionFilterChain.class); 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 b69c8578e5703..5785688fa244c 100644 --- a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java +++ b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java @@ -57,6 +57,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; @@ -67,6 +68,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; @@ -107,6 +109,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; @@ -188,6 +191,7 @@ public class TransportSearchAction extends HandledTransportAction) SearchRequest::new); this.client = client; @@ -231,6 +236,7 @@ public TransportSearchAction( this.tracer = tracer; this.taskResourceTrackingService = taskResourceTrackingService; this.indicesService = indicesService; + this.workloadGroupService = workloadGroupService; } private Map buildPerIndexAliasFilter( @@ -463,7 +469,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, @@ -474,14 +480,31 @@ 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). Principal header + // is null unless the security plugin's extractor set it. + try { + String principal = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER); + Releasable throttlePermit = workloadGroupService.acquireThrottleOrReject( + ((WorkloadGroupTask) task).getWorkloadGroupId(), + principal + ); + if (throttlePermit != null) { + updatedListener = ActionListener.runBefore(updatedListener, throttlePermit::close); + } + } catch (OpenSearchRejectedExecutionException e) { + updatedListener.onFailure(e); + return; + } } + searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext); + PipelinedRequest searchRequest; ActionListener listener; try { @@ -512,7 +535,12 @@ void executeRequest( ); } }, listener::onFailure); - searchRequest.transformRequest(requestTransformListener); + try { + searchRequest.transformRequest(requestTransformListener); + } catch (Exception e) { + // Route through updatedListener so a chained throttle permit is released rather than leaked. + updatedListener.onFailure(e); + } } } diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java index 64c398e1d5e90..b45f3fb81bddc 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; @@ -59,6 +61,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 +316,91 @@ public void rejectIfNeeded(String workloadGroupId) { }); } + /** + * Acquires one node-level throttle permit for the request, or returns {@code null} (nothing to release) when the + * request is not throttled: WLM disabled, default/unknown group, no {@code node_limit}, or no resolvable bucket + * (see {@link #resolveThrottleBucketKey}). The bucket depends on the group's throttle {@code attribute}. + * + * @param workloadGroupId the workload group the request is assigned to + * @param principal the raw {@code WORKLOAD_GROUP_PRINCIPAL_HEADER} value, or {@code null} (see resolver) + * @return a permit to close on request completion, or {@code null} if not throttled + * @throws OpenSearchRejectedExecutionException if the bucket is already at its node limit + */ + public Releasable acquireThrottleOrReject(String workloadGroupId, String principal) { + if (workloadManagementSettings.getWlmMode() != WlmMode.ENABLED) { + return null; + } + if (workloadGroupId == null || workloadGroupId.equals(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get())) { + return null; + } + try { + WorkloadGroup workloadGroup = getWorkloadGroupById(workloadGroupId); + if (workloadGroup == null) { + return null; + } + Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling(); + int nodeLimit = WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling); + if (nodeLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) { + return null; + } + // A null bucket means the request can't be attributed (e.g. username/role with no principal) -> fail open. + String bucketKey = resolveThrottleBucketKey( + workloadGroupId, + WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling), + principal + ); + if (bucketKey == null) { + return null; + } + try { + return throttleTracker.acquire(bucketKey, nodeLimit); + } catch (OpenSearchRejectedExecutionException e) { + // Record the rejection without ever letting a stats failure swallow the 429. Use the raw state map, + // not the DEFAULT-fallback accessor, so a not-yet-registered group isn't misattributed to DEFAULT. + try { + WorkloadGroupState workloadGroupState = workloadGroupsStateAccessor.getWorkloadGroupStateMap().get(workloadGroupId); + if (workloadGroupState != null) { + workloadGroupState.totalThrottled.inc(); + } + } catch (Exception statsException) { + logger.warn("Failed to record throttle stat for workload group [" + workloadGroupId + "]", statsException); + } + throw e; + } + } catch (OpenSearchRejectedExecutionException e) { + throw e; // the intended 429 + } catch (Exception e) { + // A bug in the throttle path must never fail an otherwise-valid search, so fail open. + logger.warn("Skipping node-level throttle for workload group [" + workloadGroupId + "] due to an error", e); + return null; + } + } + + /** + * Builds the throttle bucket key: {@code :group} for whole-group throttling, or + * {@code ::} keyed by the matching principal subfield for {@code username}/{@code role}. + * Returns {@code null} (fail open, not throttled) when the principal is absent or lacks a token for the subfield. + */ + private String resolveThrottleBucketKey(String workloadGroupId, String attribute, String principal) { + if ("group".equals(attribute)) { + return workloadGroupId + ":group"; + } + if (principal == null || principal.isEmpty()) { + return null; + } + String subfieldPrefix = attribute + "|"; + 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) { + return workloadGroupId + ":" + attribute + ":" + value; + } + } + } + return null; + } + 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..e07d9981434a3 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java @@ -30,6 +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"; + /** + * Carries the request's principal for username/role throttling: {@code subfield|value} tokens + * (e.g. {@code username|alice}) joined by {@link #WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER}. Only set when the + * security plugin's principal extractor is installed; absent otherwise (username/role throttling then fails open). + * Consumed only on the origin coordinator (synchronously, before shard fan-out), so it is not propagated cross-node. + */ + public static final String WORKLOAD_GROUP_PRINCIPAL_HEADER = "workloadGroupPrincipal"; + 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; 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..1c8ededc7f5dd --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java @@ -0,0 +1,87 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.lease.Releasable; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Tracks the number of in-flight requests per throttle bucket on a single node and enforces a per-node cap. + *

+ * 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. + * + * @param bucketKey the throttle bucket identifier + * @param nodeLimit the maximum concurrent in-flight requests this node may admit for the bucket + * @return a {@link Releasable} that decrements the bucket's in-flight count exactly once when closed + * @throws OpenSearchRejectedExecutionException (HTTP 429) if the bucket is already at the limit + */ + public Releasable acquire(String bucketKey, int nodeLimit) { + // Create-and-increment atomically with respect to the decrement-and-remove in release(), so a concurrent + // release draining a bucket to 0 can never orphan the counter this acquire is about to use. + AtomicInteger counter = inFlightByBucket.compute(bucketKey, (k, existing) -> { + AtomicInteger c = existing != null ? existing : new AtomicInteger(0); + c.incrementAndGet(); + return c; + }); + if (counter.get() > nodeLimit) { + // Over the cap: roll back this increment and reject. (Reading get() after compute may over-reject under + // concurrent acquires on the same bucket, but never admits over the limit — the safe direction.) + release(bucketKey, counter); + throw new OpenSearchRejectedExecutionException( + "Node-level workload group throttle limit reached: " + nodeLimit + " concurrent in-flight requests" + ); + } + return releaseOnce(bucketKey, counter); + } + + /** + * Current in-flight count for a bucket, or 0 if the bucket has no active requests. Package-private for tests. + */ + int inFlight(String bucketKey) { + AtomicInteger counter = inFlightByBucket.get(bucketKey); + return counter == null ? 0 : counter.get(); + } + + // 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 compute() makes the + // decrement-and-maybe-remove atomic per key, so a concurrent acquire can't be orphaned by a remove. + 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..e3476bb05c2fc 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) { 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..fada855f4eefe 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,10 @@ public WorkloadGroupStatsHolder(StreamInput in) throws IOException { this.rejections = in.readVLong(); this.failures = in.readVLong(); this.cancellations = in.readVLong(); + // total_throttled is version-gated so a pre-throttling node's stats stream stays readable. + if (in.getVersion().onOrAfter(Version.V_3_7_0)) { + this.throttled = in.readVLong(); + } this.resourceStats = in.readMap((i) -> ResourceType.fromName(i.readString()), ResourceStats::new); } @@ -138,6 +147,10 @@ public long getCancellations() { return cancellations; } + public long getThrottled() { + return throttled; + } + public Map getResourceStats() { return resourceStats; } @@ -160,6 +173,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 +189,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 order must stay in sync. + if (out.getVersion().onOrAfter(Version.V_3_7_0)) { + out.writeVLong(statsHolder.throttled); + } out.writeMap(statsHolder.resourceStats, (o, val) -> o.writeString(val.getName()), ResourceStats::writeTo); } @@ -190,6 +208,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 +229,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 74ac72aedbf2d..2070fa24e5d03 100644 --- a/server/src/test/java/org/opensearch/action/search/TransportSearchActionTests.java +++ b/server/src/test/java/org/opensearch/action/search/TransportSearchActionTests.java @@ -112,6 +112,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; @@ -1247,7 +1248,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/rest/action/admin/cluster/RestWlmStatsActionTests.java b/server/src/test/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsActionTests.java index 2f30181c263a7..25ead3cdd0829 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,7 +102,7 @@ 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()); } @@ -122,7 +122,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 11a481fdebdbd..a7535984ceca0 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; @@ -2409,7 +2410,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..18ea170d86e95 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java @@ -14,6 +14,7 @@ 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; @@ -479,6 +480,199 @@ 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 new WorkloadGroup( + id + "-name", + id, + new MutableWorkloadGroupFragment( + MutableWorkloadGroupFragment.ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + throttling + ), + 1L + ); + } + + 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 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 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/WorkloadGroupThrottleTrackerTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java new file mode 100644 index 0000000000000..475760834af5a --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java @@ -0,0 +1,81 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.common.lease.Releasable; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.test.OpenSearchTestCase; + +public class WorkloadGroupThrottleTrackerTests extends OpenSearchTestCase { + + public void testAcquireUnderLimitSucceeds() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + Releasable p1 = tracker.acquire("bucket", 2); + Releasable p2 = tracker.acquire("bucket", 2); + assertEquals(2, tracker.inFlight("bucket")); + p1.close(); + p2.close(); + } + + public void testAcquireAtLimitRejects() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + tracker.acquire("bucket", 1); + OpenSearchRejectedExecutionException e = expectThrows( + OpenSearchRejectedExecutionException.class, + () -> tracker.acquire("bucket", 1) + ); + assertTrue(e.getMessage().contains("throttle limit reached")); + assertFalse(e.getMessage().contains("bucket")); + // rejected acquire must not leave the count inflated + assertEquals(1, tracker.inFlight("bucket")); + } + + public void testReleaseFreesAPermit() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + Releasable p = tracker.acquire("bucket", 1); + // at the limit + expectThrows(OpenSearchRejectedExecutionException.class, () -> tracker.acquire("bucket", 1)); + p.close(); + // permit freed -> a fresh acquire now succeeds + Releasable p2 = tracker.acquire("bucket", 1); + assertEquals(1, tracker.inFlight("bucket")); + p2.close(); + } + + public void testDrainToZeroRemovesBucketThenReacquire() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + Releasable p = tracker.acquire("bucket", 5); + assertEquals(1, tracker.inFlight("bucket")); + p.close(); + assertEquals(0, tracker.inFlight("bucket")); + // re-acquiring after the bucket drained (and was removed) works and starts from 1 + Releasable p2 = tracker.acquire("bucket", 5); + assertEquals(1, tracker.inFlight("bucket")); + p2.close(); + } + + public void testReleaseIsIdempotent() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + Releasable p = tracker.acquire("bucket", 5); + p.close(); + p.close(); // double close must not decrement twice + assertEquals(0, tracker.inFlight("bucket")); + } + + public void testBucketsAreIndependent() { + WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); + tracker.acquire("a", 1); + // "a" is full but "b" is a separate bucket + expectThrows(OpenSearchRejectedExecutionException.class, () -> tracker.acquire("a", 1)); + Releasable pb = tracker.acquire("b", 1); + assertEquals(1, tracker.inFlight("a")); + assertEquals(1, tracker.inFlight("b")); + pb.close(); + } +} 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..356d90edce041 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,7 +48,7 @@ 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() ); } @@ -67,6 +68,7 @@ protected WorkloadGroupStats createTestInstance() { randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong(), + randomNonNegativeLong(), Map.of( ResourceType.CPU, new WorkloadGroupStats.ResourceStats(