diff --git a/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmClusterThrottlingIT.java b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmClusterThrottlingIT.java new file mode 100644 index 0000000000000..84da62a996194 --- /dev/null +++ b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmClusterThrottlingIT.java @@ -0,0 +1,386 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm; + +import org.apache.logging.log4j.LogManager; +import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest; +import org.opensearch.action.admin.cluster.wlm.WlmStatsAction; +import org.opensearch.action.admin.cluster.wlm.WlmStatsRequest; +import org.opensearch.action.admin.cluster.wlm.WlmStatsResponse; +import org.opensearch.action.index.IndexResponse; +import org.opensearch.action.search.SearchRequestBuilder; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.support.WriteRequest; +import org.opensearch.cluster.metadata.WorkloadGroup; +import org.opensearch.common.action.ActionFuture; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType; +import org.opensearch.plugins.Plugin; +import org.opensearch.plugins.PluginsService; +import org.opensearch.rule.RuleAttribute; +import org.opensearch.rule.RuleFrameworkPlugin; +import org.opensearch.rule.RulePersistenceServiceRegistry; +import org.opensearch.rule.RuleRoutingServiceRegistry; +import org.opensearch.rule.action.CreateRuleAction; +import org.opensearch.rule.action.CreateRuleRequest; +import org.opensearch.rule.autotagging.AutoTaggingRegistry; +import org.opensearch.rule.autotagging.FeatureType; +import org.opensearch.rule.autotagging.Rule; +import org.opensearch.script.MockScriptPlugin; +import org.opensearch.script.Script; +import org.opensearch.script.ScriptType; +import org.opensearch.search.lookup.LeafFieldsLookup; +import org.opensearch.test.OpenSearchIntegTestCase; +import org.opensearch.wlm.MutableWorkloadGroupFragment; +import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadGroupThrottleSettings; +import org.opensearch.wlm.WorkloadManagementSettings; +import org.opensearch.wlm.stats.WlmStats; +import org.opensearch.wlm.stats.WorkloadGroupStats.WorkloadGroupStatsHolder; +import org.joda.time.Instant; +import org.junit.After; +import org.junit.Before; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; +import java.util.function.ToLongFunction; + +import static org.opensearch.index.query.QueryBuilders.scriptQuery; +import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked; + +/** + * End-to-end integration test for cluster-level ({@code shared_limit}) WLM request throttling across a multi-node + * cluster. Whereas {@link WlmNodeThrottlingIT} proves the per-node tier, this verifies the distributed shared tier: + * a bucket's cluster-wide in-flight count is enforced by a single owner node regardless of which coordinator a + * request lands on, so concurrent searches spread across coordinators still cap at {@code shared_limit}. + *

+ * The group sets only {@code shared_limit} (no {@code node_limit}), so every admitted request must consult the + * bucket owner — exercising both the local-owner short-circuit and the cross-node acquire RPC depending on where + * the coordinator sits relative to the owner. + */ +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 3, numClientNodes = 0, supportsDedicatedMasters = false) +public class WlmClusterThrottlingIT extends OpenSearchIntegTestCase { + + private static final TimeValue TIMEOUT = new TimeValue(30, TimeUnit.SECONDS); + private static final String PUT = "PUT"; + + @Override + protected Collection> nodePlugins() { + List> plugins = new ArrayList<>(super.nodePlugins()); + plugins.add(WlmAutoTaggingIT.TestWorkloadManagementPlugin.class); + plugins.add(RuleFrameworkPlugin.class); + plugins.add(ClusterScriptedBlockPlugin.class); + return plugins; + } + + @Before + public void registerFeatureTypeIfMissingOnAllNodes() { + AutoTaggingRegistry.featureTypesRegistryMap.remove(WorkloadGroupFeatureType.NAME); + FeatureType featureType = WlmAutoTaggingIT.TestWorkloadManagementPlugin.featureType; + AutoTaggingRegistry.registerFeatureType(featureType); + + for (String node : internalCluster().getNodeNames()) { + RulePersistenceServiceRegistry persistenceRegistry = internalCluster().getInstance(RulePersistenceServiceRegistry.class, node); + RuleRoutingServiceRegistry routingRegistry = internalCluster().getInstance(RuleRoutingServiceRegistry.class, node); + try { + routingRegistry.getRuleRoutingService(featureType); + } catch (IllegalArgumentException ex) { + persistenceRegistry.register(featureType, WlmAutoTaggingIT.TestWorkloadManagementPlugin.rulePersistenceService); + routingRegistry.register(featureType, WlmAutoTaggingIT.TestWorkloadManagementPlugin.ruleRoutingService); + } + } + } + + @After + public void clearWlmModeSetting() throws Exception { + Settings.Builder builder = Settings.builder().putNull(WorkloadManagementSettings.WLM_MODE_SETTING.getKey()); + assertAcked(client().admin().cluster().prepareUpdateSettings().setPersistentSettings(builder).get()); + } + + public void testClusterWideCeilingHoldsAcrossCoordinators() throws Exception { + String workloadGroupId = "wlm_shared_throttle_group"; + String ruleId = "wlm_shared_throttle_rule"; + String indexName = "shared_throttle_index"; + + setWlmMode("enabled"); + + // shared_limit = 2, node_limit unset -> a cluster-wide ceiling of 2 in-flight for the whole group, + // enforced by the bucket owner no matter which coordinator receives the request. + WorkloadGroup workloadGroup = createSharedThrottledGroup("shared_test_group", workloadGroupId, 2); + updateWorkloadGroupInClusterState(PUT, workloadGroup); + + assertBusy(() -> { + boolean present = client().admin() + .cluster() + .prepareState() + .get() + .getState() + .metadata() + .workloadGroups() + .containsKey(workloadGroupId); + assertTrue("workload group not yet applied in cluster state", present); + }, 30, TimeUnit.SECONDS); + + FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME); + createRule(ruleId, "shared throttle rule", indexName, featureType, workloadGroupId); + indexDocument(indexName); + + // Wait for rule propagation on EVERY coordinator this test drives, not just a random one. The rule refresh is + // asynchronous and can reach nodes at different times; a search on an un-refreshed coordinator stays untagged + // and bypasses admit(), which would break the ceiling assertions below. Poll each named node's own client + // until a search through it is tagged to the group (its completions advance). + for (String node : internalCluster().getNodeNames()) { + assertBusy(() -> { + long before = getCompletions(workloadGroupId); + try { + client(node).prepareSearch(indexName).setQuery(QueryBuilders.matchAllQuery()).get(); + } catch (Exception e) { + // Each probe consults the shared tier (shared_limit=2, node_limit unset) and briefly holds a permit; + // remote releases are fire-and-forget, so a burst of probes can transiently exhaust the budget and + // throttle this one. That is not a propagation failure — turn the 429 into an AssertionError so + // assertBusy retries it (assertBusy only re-runs on AssertionError; a raw 429 would escape). Any + // other exception is a real failure and propagates. + assertFalse("transient throttle during propagation probe — retry: " + e, hasRejectedExecutionCause(e)); + throw e; + } + long after = getCompletions(workloadGroupId); + assertTrue("search via [" + node + "] not yet tagged to the throttled workload group", after > before); + }, 30, TimeUnit.SECONDS); + } + + List plugins = initBlockFactory(); + + List coordinators = new ArrayList<>(List.of(internalCluster().getNodeNames())); + // Fill the cluster-wide budget of 2 using two DIFFERENT coordinators. + ActionFuture first = blockingSearchVia(coordinators.get(0), indexName).execute(); + ActionFuture second = blockingSearchVia(coordinators.get(1 % coordinators.size()), indexName).execute(); + awaitBlockedCount(plugins, 2); + + long throttledBefore = getThrottled(workloadGroupId); + + // A third concurrent search on yet another coordinator must be rejected: the cluster-wide ceiling of 2 is + // reached, and the owner (a single node) sees the count regardless of coordinator. + String thirdCoordinator = coordinators.get(2 % coordinators.size()); + Throwable rejection = expectThrows(Throwable.class, () -> blockingSearchVia(thirdCoordinator, indexName).get()); + assertTrue( + "Expected an OpenSearchRejectedExecutionException in the cause chain but was: " + rejection, + hasRejectedExecutionCause(rejection) + ); + assertEquals("total_throttled should increment by exactly one", throttledBefore + 1, getThrottled(workloadGroupId)); + + // Release the blocks; the two admitted searches complete successfully. + disableBlocks(plugins); + assertNotNull(first.actionGet(TIMEOUT)); + assertNotNull(second.actionGet(TIMEOUT)); + + // Once the budget drains, a new search is admitted again. The two fills may have released REMOTELY (when the + // bucket owner is a third node, release is a fire-and-forget RPC with no happens-before to this re-admission), + // so the owner's in-flight count may not have dropped the instant actionGet() above returned. Poll: retry a + // non-blocking search until it is admitted (tolerating a transient 429 during the drain window), rather than + // assuming the drain is immediately visible. + assertBusy(() -> { + try { + client(thirdCoordinator).prepareSearch(indexName).setQuery(QueryBuilders.matchAllQuery()).get(); + } catch (Exception e) { + assertFalse("still draining (transient 429) — retry", hasRejectedExecutionCause(e)); + throw e; + } + }, 30, TimeUnit.SECONDS); + } + + // Helpers + + private static boolean hasRejectedExecutionCause(Throwable t) { + for (Throwable cur = t; cur != null; cur = cur.getCause()) { + if (cur instanceof OpenSearchRejectedExecutionException) { + return true; + } + if (cur.getCause() == cur) { + break; + } + } + return false; + } + + private long getCompletions(String groupId) throws Exception { + return sumAcrossNodes(groupId, WorkloadGroupStatsHolder::getCompletions); + } + + private long getThrottled(String groupId) throws Exception { + return sumAcrossNodes(groupId, WorkloadGroupStatsHolder::getThrottled); + } + + // Sums a per-group stat across all nodes using the typed WlmStats response (no brittle string parsing). + private long sumAcrossNodes(String groupId, ToLongFunction extractor) throws Exception { + WlmStatsRequest request = new WlmStatsRequest(null, new HashSet<>(Collections.singletonList(groupId)), null); + WlmStatsResponse response = client().execute(WlmStatsAction.INSTANCE, request).get(); + long total = 0; + for (WlmStats nodeStats : response.getNodes()) { + WorkloadGroupStatsHolder holder = nodeStats.getWorkloadGroupStats().getStats().get(groupId); + if (holder != null) { + total += extractor.applyAsLong(holder); + } + } + return total; + } + + private SearchRequestBuilder blockingSearchVia(String nodeName, String indexName) { + return client(nodeName).prepareSearch(indexName) + .setQuery( + scriptQuery(new Script(ScriptType.INLINE, "mockscript", ClusterScriptedBlockPlugin.SCRIPT_NAME, Collections.emptyMap())) + ); + } + + private List initBlockFactory() { + List plugins = new ArrayList<>(); + for (PluginsService pluginsService : internalCluster().getDataNodeInstances(PluginsService.class)) { + plugins.addAll(pluginsService.filterPlugins(ClusterScriptedBlockPlugin.class)); + } + for (ClusterScriptedBlockPlugin plugin : plugins) { + plugin.reset(); + plugin.enableBlock(); + } + return plugins; + } + + private void awaitBlockedCount(List plugins, int expected) throws Exception { + assertBusy(() -> { + int blocked = 0; + for (ClusterScriptedBlockPlugin plugin : plugins) { + blocked += plugin.hits.get(); + } + assertEquals("expected exactly " + expected + " searches blocked in-flight", expected, blocked); + }, 30, TimeUnit.SECONDS); + } + + private void disableBlocks(List plugins) { + for (ClusterScriptedBlockPlugin plugin : plugins) { + plugin.disableBlock(); + } + } + + private void createRule(String ruleId, String ruleName, String indexPattern, FeatureType featureType, String workloadGroupId) + throws Exception { + Rule rule = new Rule( + ruleId, + ruleName, + Map.of(RuleAttribute.INDEX_PATTERN, Set.of(indexPattern)), + featureType, + workloadGroupId, + Instant.now().toString() + ); + client().execute(CreateRuleAction.INSTANCE, new CreateRuleRequest(rule)).get(); + } + + private void setWlmMode(String mode) throws Exception { + Settings.Builder settings = Settings.builder().put("wlm.workload_group.mode", mode); + ClusterUpdateSettingsRequest request = new ClusterUpdateSettingsRequest().persistentSettings(settings); + client().admin().cluster().updateSettings(request).get(); + } + + private WorkloadGroup createSharedThrottledGroup(String name, String id, int sharedLimit) { + Settings throttling = Settings.builder() + .put(WorkloadGroupThrottleSettings.ATTRIBUTE.getKey(), "group") + .put(WorkloadGroupThrottleSettings.SHARED_LIMIT.getKey(), sharedLimit) + .build(); + return new WorkloadGroup( + name, + id, + new MutableWorkloadGroupFragment( + MutableWorkloadGroupFragment.ResiliencyMode.SOFT, + Map.of(ResourceType.CPU, 0.9, ResourceType.MEMORY, 0.9), + Settings.EMPTY, + throttling + ), + Instant.now().getMillis() + ); + } + + private void indexDocument(String indexName) { + assertAcked( + client().admin() + .indices() + .prepareCreate(indexName) + // One shard so each search produces exactly one blocking script hit, making the in-flight count exact. + // Throttle admission happens on the coordinator that receives the request, independent of where the + // shard lives, so a single shard is sufficient to exercise the cross-coordinator ceiling. + .setSettings(Settings.builder().put("index.number_of_shards", 1).put("index.number_of_replicas", 0)) + ); + IndexResponse response = client().prepareIndex(indexName) + .setId("1") + .setSource(Map.of("field", "value")) + .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE) + .get(); + assertEquals(org.opensearch.action.DocWriteResponse.Result.CREATED, response.getResult()); + } + + private void updateWorkloadGroupInClusterState(String method, WorkloadGroup workloadGroup) throws InterruptedException { + WlmAutoTaggingIT.ExceptionCatchingListener listener = new WlmAutoTaggingIT.ExceptionCatchingListener(); + client().execute( + WlmAutoTaggingIT.TestClusterUpdateTransportAction.ACTION, + new WlmAutoTaggingIT.TestClusterUpdateRequest(workloadGroup, method), + listener + ); + boolean completed = listener.getLatch().await(TIMEOUT.getSeconds(), TimeUnit.SECONDS); + assertTrue("cluster-state update did not complete in time", completed); + if (listener.getException() != null) { + throw new AssertionError("cluster-state update failed", listener.getException()); + } + } + + /** + * Test script plugin that blocks during the query phase until released, keeping a search in-flight on whichever + * coordinator dispatched it. + */ + public static class ClusterScriptedBlockPlugin extends MockScriptPlugin { + static final String SCRIPT_NAME = "cluster_search_block"; + + private final AtomicInteger hits = new AtomicInteger(); + private final AtomicBoolean shouldBlock = new AtomicBoolean(true); + + public void reset() { + hits.set(0); + } + + public void disableBlock() { + shouldBlock.set(false); + } + + public void enableBlock() { + shouldBlock.set(true); + } + + @Override + public Map, Object>> pluginScripts() { + return Collections.singletonMap(SCRIPT_NAME, params -> { + LeafFieldsLookup fieldsLookup = (LeafFieldsLookup) params.get("_fields"); + LogManager.getLogger(WlmClusterThrottlingIT.class).info("Blocking on the document {}", fieldsLookup.get("_id")); + hits.incrementAndGet(); + try { + assertBusy(() -> assertFalse(shouldBlock.get())); + } catch (Exception e) { + throw new RuntimeException(e); + } + return true; + }); + } + } +} diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java index 5785688fa244c..0d781c597224e 100644 --- a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java +++ b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java @@ -39,6 +39,7 @@ import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsRequest; import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsResponse; import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.ContextPreservingActionListener; import org.opensearch.action.support.HandledTransportAction; import org.opensearch.action.support.IndicesOptions; import org.opensearch.action.support.TimeoutTaskCancellationUtility; @@ -68,7 +69,6 @@ import org.opensearch.core.common.breaker.CircuitBreaker; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.common.io.stream.Writeable; -import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import org.opensearch.core.index.Index; import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.indices.breaker.CircuitBreakerService; @@ -485,25 +485,71 @@ void executeRequest( // or HTTP header (HTTP header will be deprecated once ActionFilter is implemented) if (task instanceof WorkloadGroupTask) { ((WorkloadGroupTask) task).setWorkloadGroupId(threadPool.getThreadContext()); - // Node-level throttle admission. Runs before onRequestStart so a rejection doesn't leak the request - // gauges (decremented only on request end/failure, which the early return skips). Principal header - // is null unless the security plugin's extractor set it. - try { - String principal = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER); - Releasable throttlePermit = workloadGroupService.acquireThrottleOrReject( - ((WorkloadGroupTask) task).getWorkloadGroupId(), - principal - ); - if (throttlePermit != null) { - updatedListener = ActionListener.runBefore(updatedListener, throttlePermit::close); - } - } catch (OpenSearchRejectedExecutionException e) { - updatedListener.onFailure(e); - return; - } + // Two-tier throttle admission (node-local + cluster-level shared). Runs before onRequestStart so a + // rejection doesn't leak the request gauges (decremented only on request end/failure, which the early + // return skips). The node-local tier is granted synchronously here (zero added latency); only an + // overflow to the shared tier does an asynchronous owner round-trip, so this call may complete inline + // or on a transport thread. Principal header is null unless the security plugin's extractor set it. + final String principal = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER); + final ActionListener outerListener = updatedListener; + // Preserve the request's thread context (workload group + principal headers, etc.) so that when the + // shared-tier admission callback runs on a transport thread, downstream search setup sees them. + final ActionListener admissionListener = ContextPreservingActionListener.wrapPreservingContext( + ActionListener.wrap(throttlePermit -> { + ActionListener proceedListener = throttlePermit == null + ? outerListener + : ActionListener.runBefore(outerListener, throttlePermit::close); + proceedWithSearch( + task, + originalSearchRequest, + searchAsyncActionProvider, + proceedListener, + searchRequestContext, + timeProvider, + requestSpan + ); + }, outerListener::onFailure), + threadPool.getThreadContext() + ); + workloadGroupService.acquireThrottlePermit(((WorkloadGroupTask) task).getWorkloadGroupId(), principal, admissionListener); + } else { + proceedWithSearch( + task, + originalSearchRequest, + searchAsyncActionProvider, + updatedListener, + searchRequestContext, + timeProvider, + requestSpan + ); } + } + } - searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext); + /** + * Runs the search once throttle admission has granted (or been skipped). Extracted from {@link #executeRequest} + * so it can be invoked either inline (node-local grant / not throttled) or asynchronously from the shared-tier + * acquire callback. Re-enters the request span scope because the async callback may run on a transport thread + * where the caller's span scope is no longer active. + */ + private void proceedWithSearch( + Task task, + SearchRequest originalSearchRequest, + SearchAsyncActionProvider searchAsyncActionProvider, + ActionListener updatedListener, + SearchRequestContext searchRequestContext, + SearchTimeProvider timeProvider, + Span requestSpan + ) { + try (final SpanScope ignore = tracer.withSpanInScope(requestSpan)) { + try { + searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext); + } catch (Exception e) { + // Route through updatedListener (which may hold a throttle permit via runBefore) so a failure here + // releases the permit rather than leaking it. A node-local permit has no TTL to reclaim it. + updatedListener.onFailure(e); + return; + } PipelinedRequest searchRequest; ActionListener listener; diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index 832a9721ce7e5..a7104e5a45862 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -309,6 +309,7 @@ import org.opensearch.usage.UsageService; import org.opensearch.watcher.ResourceWatcherService; import org.opensearch.wlm.WorkloadGroupService; +import org.opensearch.wlm.WorkloadGroupSharedThrottleService; import org.opensearch.wlm.WorkloadGroupsStateAccessor; import org.opensearch.wlm.WorkloadManagementSettings; import org.opensearch.wlm.WorkloadManagementTransportInterceptor; @@ -1423,6 +1424,16 @@ protected Node(final Environment initialEnvironment, Collection clas ) : Optional.empty(); + // Cluster-level (shared_limit) WLM throttle tier. Built here because it needs the transport service for + // the owner round-trip; late-bound into the WorkloadGroupService (constructed earlier) so the two-tier + // admit() path can reach it. + final WorkloadGroupSharedThrottleService workloadGroupSharedThrottleService = new WorkloadGroupSharedThrottleService( + clusterService, + threadPool, + transportService + ); + workloadGroupService.setSharedThrottleService(workloadGroupSharedThrottleService); + TopNSearchTasksLogger taskConsumer = new TopNSearchTasksLogger(settings, settingsModule.getClusterSettings()); transportService.getTaskManager().registerTaskResourceConsumer(taskConsumer); streamTransportService.ifPresent(service -> service.getTaskManager().registerTaskResourceConsumer(taskConsumer)); @@ -1709,6 +1720,7 @@ protected Node(final Environment initialEnvironment, Collection clas b.bind(TaskResourceTrackingService.class).toInstance(taskResourceTrackingService); b.bind(SearchBackpressureService.class).toInstance(searchBackpressureService); b.bind(WorkloadGroupService.class).toInstance(workloadGroupService); + b.bind(WorkloadGroupSharedThrottleService.class).toInstance(workloadGroupSharedThrottleService); b.bind(AdmissionControlService.class).toInstance(admissionControlService); b.bind(UsageService.class).toInstance(usageService); b.bind(AggregationUsageService.class).toInstance(searchModule.getValuesSourceRegistry().getUsageService()); @@ -1963,6 +1975,7 @@ public Node start() throws NodeValidationException { nodeService.getSearchBackpressureService().start(); nodeService.getTaskCancellationMonitoringService().start(); injector.getInstance(WorkloadGroupService.class).start(); + injector.getInstance(WorkloadGroupSharedThrottleService.class).start(); final ClusterService clusterService = injector.getInstance(ClusterService.class); @@ -2153,6 +2166,7 @@ private Node stop() { injector.getInstance(NodeResourceUsageTracker.class).stop(); injector.getInstance(ResourceUsageCollectorService.class).stop(); injector.getInstance(WorkloadGroupService.class).stop(); + injector.getInstance(WorkloadGroupSharedThrottleService.class).stop(); nodeService.getMonitorService().stop(); nodeService.getSearchBackpressureService().stop(); injector.getInstance(GatewayService.class).stop(); diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java b/server/src/main/java/org/opensearch/wlm/NodeThrottleTracker.java similarity index 61% rename from server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java rename to server/src/main/java/org/opensearch/wlm/NodeThrottleTracker.java index 1c8ededc7f5dd..428098dd30e74 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleTracker.java +++ b/server/src/main/java/org/opensearch/wlm/NodeThrottleTracker.java @@ -10,7 +10,6 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.lease.Releasable; -import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -18,7 +17,9 @@ import java.util.concurrent.atomic.AtomicInteger; /** - * Tracks the number of in-flight requests per throttle bucket on a single node and enforces a per-node cap. + * Node-local tier of workload-group request throttling: tracks in-flight requests per throttle bucket on a single + * node and enforces the bucket's {@code node_limit}. Its cluster-wide sibling is {@link SharedThrottleTracker}, and + * {@code WorkloadGroupService.acquireThrottlePermit} composes the two (local first, then shared on overflow). *

* A bucket is identified by an opaque key (see {@code WorkloadGroupService} for how the key is built from a * workload group and its throttle attribute). A counter exists only while a bucket has at least one in-flight @@ -26,36 +27,43 @@ * number of concurrently active buckets rather than the total population of users/roles. *

* This tier is fully local — no cross-node coordination — mirroring the acquire/rollback + {@link Releasable} - * release shape of {@link org.opensearch.index.IndexingPressure}. + * release shape of {@link org.opensearch.index.IndexingPressure}. Because acquire and release happen in the same + * process, a permit is a self-releasing {@link Releasable} and needs no lease id or TTL (contrast + * {@link SharedThrottleTracker}, whose distributed acquire/release requires both). When a two-tier throttle is + * configured, an exhausted local tier falls through to the cluster-level {@link SharedThrottleTracker} rather than + * rejecting outright; see {@link #tryAcquire}, which returns {@code null} on exhaustion instead of throwing. */ @ExperimentalApi -public class WorkloadGroupThrottleTracker { +public class NodeThrottleTracker { private final Map inFlightByBucket = new ConcurrentHashMap<>(); /** - * Attempts to admit one request into the given bucket under the per-node limit. + * Attempts to admit one request into the given bucket under the per-node limit, returning {@code null} instead + * of throwing when the bucket is already at the limit. This is the fall-through-friendly form: a caller that + * also has a cluster-level shared pool can try the shared pool on a {@code null} result rather than rejecting. * * @param bucketKey the throttle bucket identifier * @param nodeLimit the maximum concurrent in-flight requests this node may admit for the bucket - * @return a {@link Releasable} that decrements the bucket's in-flight count exactly once when closed - * @throws OpenSearchRejectedExecutionException (HTTP 429) if the bucket is already at the limit + * @return a {@link Releasable} that decrements the bucket's in-flight count exactly once when closed, or + * {@code null} if the per-node limit is already reached (nothing was acquired, nothing to release) */ - public Releasable acquire(String bucketKey, int nodeLimit) { + public Releasable tryAcquire(String bucketKey, int nodeLimit) { // Create-and-increment atomically with respect to the decrement-and-remove in release(), so a concurrent - // release draining a bucket to 0 can never orphan the counter this acquire is about to use. + // release draining a bucket to 0 can never orphan the counter this acquire is about to use. Capture THIS + // caller's post-increment value from inside compute() and decide on it — reading counter.get() afterwards + // would observe other concurrent acquirers' increments too, causing spurious over-declines (all racers seeing + // an inflated count and rolling back, admitting zero when a slot was free). + final int[] countAfterIncrement = new int[1]; AtomicInteger counter = inFlightByBucket.compute(bucketKey, (k, existing) -> { AtomicInteger c = existing != null ? existing : new AtomicInteger(0); - c.incrementAndGet(); + countAfterIncrement[0] = c.incrementAndGet(); return c; }); - if (counter.get() > nodeLimit) { - // Over the cap: roll back this increment and reject. (Reading get() after compute may over-reject under - // concurrent acquires on the same bucket, but never admits over the limit — the safe direction.) + if (countAfterIncrement[0] > nodeLimit) { + // Over the cap: roll back this increment and decline. Never admits over the limit. release(bucketKey, counter); - throw new OpenSearchRejectedExecutionException( - "Node-level workload group throttle limit reached: " + nodeLimit + " concurrent in-flight requests" - ); + return null; } return releaseOnce(bucketKey, counter); } diff --git a/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java b/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java new file mode 100644 index 0000000000000..8a005c1cfddf8 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/SharedThrottleTracker.java @@ -0,0 +1,139 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.common.annotation.ExperimentalApi; + +import java.util.Iterator; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.LongSupplier; + +/** + * Cluster-level tier of workload-group request throttling: the authoritative, cluster-wide in-flight counter for a + * throttle bucket, enforcing the bucket's {@code shared_limit}. Its node-local sibling is {@link NodeThrottleTracker}, + * and {@code WorkloadGroupService.acquireThrottlePermit} composes the two (local first, then shared on overflow). On the node that + * owns a bucket (chosen by {@link ThrottleOwnerSelector}), this holds one live-lease set per active bucket. + *

+ * Unlike the node tier, acquire and release happen on different nodes (a coordinator acquires from the owner, then + * releases via a separate RPC), so each permit is a lease with an id and a TTL rather than a self-releasing + * {@link org.opensearch.common.lease.Releasable}. + *

+ * Correctness under concurrency and node churn: + *

+ * A bucket entry exists only while it has at least one live lease, so memory scales with concurrently-active + * buckets, not the total user/role population. + */ +@ExperimentalApi +public class SharedThrottleTracker { + + private final Map> leasesByBucket = new ConcurrentHashMap<>(); + private final LongSupplier nanoTimeSupplier; + + public SharedThrottleTracker() { + this(System::nanoTime); + } + + public SharedThrottleTracker(LongSupplier nanoTimeSupplier) { + this.nanoTimeSupplier = nanoTimeSupplier; + } + + /** + * Attempts to admit one request into a bucket under its cluster-wide shared limit. + * + * @param bucketKey the throttle bucket identifier + * @param sharedLimit the maximum concurrent in-flight requests allowed across the cluster for this bucket + * @param leaseId a globally-unique lease id minted by the requesting coordinator + * @param ttlNanos how long the lease may live before {@link #sweepExpired} may reclaim it + * @return {@code true} if the lease was granted (live count was below the limit), {@code false} if the bucket + * is already at the limit + */ + public boolean tryAcquire(String bucketKey, int sharedLimit, String leaseId, long ttlNanos) { + final long expiresAt = nanoTimeSupplier.getAsLong() + ttlNanos; + final boolean[] granted = new boolean[1]; + leasesByBucket.compute(bucketKey, (k, leases) -> { + if (leases == null) { + leases = new ConcurrentHashMap<>(); + } + // Only prune when the bucket looks full. An expired lease can change the outcome only when we would + // otherwise reject; below the limit there is a free slot regardless, so we skip the O(size) scan and keep + // the common under-limit acquire O(1). A saturated bucket still self-heals its stuck leases right here, at + // the one moment it matters. (Leases in an idle bucket are reclaimed by the periodic sweep instead.) + if (leases.size() >= sharedLimit) { + pruneExpired(leases); + } + if (leases.size() < sharedLimit) { + leases.put(leaseId, expiresAt); + granted[0] = true; + } + // Never leave an empty map behind (keeps the live set == active buckets). + return leases.isEmpty() ? null : leases; + }); + return granted[0]; + } + + /** + * Releases a previously-granted lease. Removing by lease id makes this idempotent and safe: an unknown id + * (already swept, double release, or a lease that belonged to a previous owner of a since-remapped bucket) is + * a no-op and can never decrement a live lease that belongs to some other request. + * + * @param bucketKey the throttle bucket identifier + * @param leaseId the lease id returned to the coordinator at acquire time + */ + public void release(String bucketKey, String leaseId) { + leasesByBucket.computeIfPresent(bucketKey, (k, leases) -> { + leases.remove(leaseId); + return leases.isEmpty() ? null : leases; + }); + } + + /** + * Reclaims all expired leases across every bucket. Intended to be called periodically by the owning service. + * Safe to run concurrently with {@link #tryAcquire}/{@link #release} thanks to per-key {@code compute}. + */ + public void sweepExpired() { + for (String bucketKey : leasesByBucket.keySet()) { + leasesByBucket.computeIfPresent(bucketKey, (k, leases) -> { + pruneExpired(leases); + return leases.isEmpty() ? null : leases; + }); + } + } + + // Current live (non-expired-at-read-time) count for a bucket. Package-private for tests. + int inFlight(String bucketKey) { + Map leases = leasesByBucket.get(bucketKey); + return leases == null ? 0 : leases.size(); + } + + // Number of buckets currently holding at least one lease. Package-private for tests. + int activeBuckets() { + return leasesByBucket.size(); + } + + private void pruneExpired(Map leases) { + final long now = nanoTimeSupplier.getAsLong(); + for (Iterator> it = leases.entrySet().iterator(); it.hasNext();) { + if (it.next().getValue() <= now) { + it.remove(); + } + } + } +} diff --git a/server/src/main/java/org/opensearch/wlm/ThrottleOwnerSelector.java b/server/src/main/java/org/opensearch/wlm/ThrottleOwnerSelector.java new file mode 100644 index 0000000000000..c311e6ec34bce --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/ThrottleOwnerSelector.java @@ -0,0 +1,141 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.Version; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.cluster.routing.Murmur3HashFunction; +import org.opensearch.common.annotation.ExperimentalApi; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.SortedMap; +import java.util.TreeMap; + +/** + * Maps a throttle bucket to the single node that authoritatively owns its cluster-level counter, using a + * consistent-hash ring built from cluster state. Every node computes the same ring from the same + * {@link DiscoveryNodes}, so any coordinator can locate a bucket's owner with no election or lookup, and node + * join/leave only remaps roughly {@code 1/N} of buckets. + *

+ * Design choices that matter: + *

    + *
  • Data nodes only. Cluster-manager nodes are deliberately excluded — a domain may run zero or one + * dedicated cluster-manager, so they are not a dependable ownership pool.
  • + *
  • Version-filtered. During a rolling upgrade, nodes older than {@link #MIN_OWNER_VERSION} (which do + * not run the shared-throttle handler) are excluded from the ring, so every bucket lands on a node that can + * actually enforce it rather than silently going unthrottled. This trades extra remapping during the + * upgrade window for real enforcement.
  • + *
  • Virtual nodes. Each physical node is placed at {@link #VIRTUAL_NODES} points on the ring so buckets + * spread evenly and a departing node's share redistributes across the survivors rather than piling onto one + * neighbour.
  • + *
  • Empty ring ⇒ no owner. When no eligible node exists (bootstrap, coordinating-only topologies, or a + * fleet with no upgraded data nodes yet), {@link #ownerFor} returns empty and the caller fails open.
  • + *
+ * Instances are immutable snapshots; {@link WorkloadGroupSharedThrottleService} rebuilds one whenever the + * discovery-node set changes. + *

+ * Transient ceiling breach on rebalance. When a bucket remaps to a new owner (node join/leave), the previous + * owner still holds leases for that bucket's in-flight requests while the new owner starts counting from zero, so the + * bucket's cluster-wide in-flight count can briefly exceed {@code shared_limit} (up to ~2x) until the old leases drain + * or TTL out. This is inherent to stateless consistent hashing and is accepted for this experimental feature. + */ +@ExperimentalApi +public class ThrottleOwnerSelector { + + /** + * First version whose nodes run the shared-throttle transport handler and may therefore own buckets. Nodes below + * this are excluded from the ring so a bucket never maps to an owner that would answer the acquire RPC with an + * unknown-action failure (which fails open, silently disabling {@code shared_limit} for that bucket). + *

+ * TODO(merge): set this to the actual release that first ships these handlers before upstreaming. On this feature + * branch it is a placeholder; if the handlers land in a build after {@code V_3_7_0}, this must be bumped to that + * version, otherwise a rolling upgrade with older-but-same-major data nodes would map buckets to nodes lacking the + * handler. (Tracked alongside the throttle-settings version-gate TODO.) + */ + public static final Version MIN_OWNER_VERSION = Version.V_3_7_0; + + private static final int VIRTUAL_NODES = 128; + + // Ring position -> node id. TreeMap gives the ceiling/first lookup that consistent hashing needs. + private final SortedMap ring; + private final Map nodesById; + + private ThrottleOwnerSelector(SortedMap ring, Map nodesById) { + this.ring = ring; + this.nodesById = nodesById; + } + + /** + * Builds a ring from the current discovery nodes, including only data nodes at or above + * {@link #MIN_OWNER_VERSION}. + */ + public static ThrottleOwnerSelector fromDiscoveryNodes(DiscoveryNodes discoveryNodes) { + final SortedMap ring = new TreeMap<>(); + final Map nodesById = new HashMap<>(); + for (DiscoveryNode node : discoveryNodes.getDataNodes().values()) { + if (node.getVersion().onOrAfter(MIN_OWNER_VERSION) == false) { + continue; + } + nodesById.put(node.getId(), node); + for (int i = 0; i < VIRTUAL_NODES; i++) { + // Hash the node id with a replica suffix to scatter its virtual points around the ring. + ring.put(Murmur3HashFunction.hash(node.getId() + "#" + i), node.getId()); + } + } + return new ThrottleOwnerSelector(ring, nodesById); + } + + /** + * @return the node that owns the given bucket, or empty if the ring has no eligible node (fail open). + */ + public Optional ownerFor(String bucketKey) { + if (ring.isEmpty()) { + return Optional.empty(); + } + final int hash = Murmur3HashFunction.hash(bucketKey); + // First ring point at or after the bucket's hash, wrapping around to the first point. + SortedMap tail = ring.tailMap(hash); + Integer point = tail.isEmpty() ? ring.firstKey() : tail.firstKey(); + return Optional.ofNullable(nodesById.get(ring.get(point))); + } + + /** + * The eligible owner nodes in this snapshot as a set of full {@link DiscoveryNode}s. Used to decide whether a + * rebuild is needed: comparing whole nodes (not just persistent ids) detects a same-id restart, where the node + * keeps its persistent id but gets a new ephemeral id/address — otherwise the ring would keep the stale node and + * {@code nodeConnected} would fail open for its buckets indefinitely. + */ + Set eligibleNodeSet() { + return new HashSet<>(nodesById.values()); + } + + /** + * Computes the eligible owner set for a candidate {@link DiscoveryNodes} without building the full ring + * (no virtual-node hashing). Lets a caller cheaply check whether membership actually changed before paying to + * rebuild the ring on every cluster-state update. + */ + static Set eligibleNodeSet(DiscoveryNodes discoveryNodes) { + Set eligible = new HashSet<>(); + for (DiscoveryNode node : discoveryNodes.getDataNodes().values()) { + if (node.getVersion().onOrAfter(MIN_OWNER_VERSION)) { + eligible.add(node); + } + } + return eligible; + } + + boolean isEmpty() { + return ring.isEmpty(); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java index b45f3fb81bddc..a7e385f74ae00 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java @@ -19,6 +19,7 @@ import org.opensearch.common.lease.Releasable; import org.opensearch.common.lifecycle.AbstractLifecycleComponent; import org.opensearch.common.settings.Settings; +import org.opensearch.core.action.ActionListener; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import org.opensearch.monitor.jvm.JvmStats; import org.opensearch.monitor.process.ProcessProbe; @@ -62,7 +63,9 @@ public class WorkloadGroupService extends AbstractLifecycleComponent private final NodeDuressTrackers nodeDuressTrackers; private final WorkloadGroupsStateAccessor workloadGroupsStateAccessor; // Node-local in-flight throttle counters, keyed by throttle bucket. No cross-node coordination in this tier. - private final WorkloadGroupThrottleTracker throttleTracker = new WorkloadGroupThrottleTracker(); + private final NodeThrottleTracker throttleTracker = new NodeThrottleTracker(); + // Cluster-level (shared_limit) tier; late-bound after the transport service exists. Null => only the local tier. + private volatile WorkloadGroupSharedThrottleService sharedThrottleService; public WorkloadGroupService( WorkloadGroupTaskCancellationService taskCancellationService, @@ -317,73 +320,184 @@ public void rejectIfNeeded(String workloadGroupId) { } /** - * Acquires one node-level throttle permit for the request, or returns {@code null} (nothing to release) when the - * request is not throttled: WLM disabled, default/unknown group, no {@code node_limit}, or no resolvable bucket - * (see {@link #resolveThrottleBucketKey}). The bucket depends on the group's throttle {@code attribute}. + * Late-binds the cluster-level ({@code shared_limit}) throttle tier. It is constructed after the transport + * service (which it needs for the owner round-trip), so it cannot be a constructor argument. When unset, only + * the node-local tier applies and a shared-only config fails open. + */ + public void setSharedThrottleService(WorkloadGroupSharedThrottleService sharedThrottleService) { + this.sharedThrottleService = sharedThrottleService; + } + + /** + * Two-tier throttle admission for a search request. Notifies {@code listener} with: + *

    + *
  • a non-null {@link Releasable} — admitted; close it exactly once on request completion to release the + * permit (works for both a node-local and a cluster-level shared permit);
  • + *
  • {@code null} — admitted but not throttle-tracked (throttling disabled/not configured for this request, + * request not attributable to a bucket, or a fail-open path); nothing to release;
  • + *
  • {@link ActionListener#onFailure} with an {@link OpenSearchRejectedExecutionException} (HTTP 429) — the + * bucket is at its limit.
  • + *
+ * The node-local tier is checked synchronously (zero added latency on the common path); only an overflow to the + * shared tier does the asynchronous owner round-trip, so the calling thread is never blocked. The listener may + * therefore be invoked inline or, for the shared-tier overflow, on a transport thread. * * @param workloadGroupId the workload group the request is assigned to - * @param principal the raw {@code WORKLOAD_GROUP_PRINCIPAL_HEADER} value, or {@code null} (see resolver) - * @return a permit to close on request completion, or {@code null} if not throttled - * @throws OpenSearchRejectedExecutionException if the bucket is already at its node limit + * @param principal the raw {@code WORKLOAD_GROUP_PRINCIPAL_HEADER} value, or {@code null} + * @param listener receives the permit / null / 429 */ - public Releasable acquireThrottleOrReject(String workloadGroupId, String principal) { + public void acquireThrottlePermit(String workloadGroupId, String principal, ActionListener listener) { + final ThrottlePlan plan; + try { + plan = resolveThrottlePlan(workloadGroupId, principal); + } catch (Exception e) { + // A bug in the throttle-resolution path must never fail an otherwise-valid search, so fail open. + logger.warn("Skipping throttle for workload group [" + workloadGroupId + "] due to an error", e); + listener.onResponse(null); + return; + } + if (plan == null) { + listener.onResponse(null); // not throttled / not attributable -> fail open + return; + } + + // Node-local tier: synchronous, no cross-node coordination. Granting here is the zero-latency common path. + if (plan.nodeLimit != WorkloadGroupThrottleSettings.UNSET_LIMIT) { + Releasable localPermit = throttleTracker.tryAcquire(plan.bucketKey, plan.nodeLimit); + if (localPermit != null) { + listener.onResponse(localPermit); + return; + } + // Local allowance exhausted -> fall through to the shared pool if one is configured. + } + + // Cluster-level shared tier (asynchronous owner round-trip). + if (plan.sharedLimit != WorkloadGroupThrottleSettings.UNSET_LIMIT && sharedThrottleService != null) { + sharedThrottleService.acquireAsync(plan.bucketKey, plan.sharedLimit, ActionListener.wrap(listener::onResponse, e -> { + if (e instanceof OpenSearchRejectedExecutionException) { + // Recompose the denial with the human-readable group name/attribute (the shared tier only has the + // opaque bucket key), and count it. + incrementThrottled(workloadGroupId); + listener.onFailure(sharedTierRejection(plan)); + } else { + listener.onFailure(e); + } + })); + return; + } + + if (plan.nodeLimit != WorkloadGroupThrottleSettings.UNSET_LIMIT && plan.sharedLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) { + // Node-only config with the local allowance exhausted and no shared tier to overflow to -> reject. + incrementThrottled(workloadGroupId); + listener.onFailure(nodeTierRejection(plan)); + } else { + // Either a shared tier was configured but is unavailable (not yet wired), or a shared-only config with no + // wired tier. Fail open rather than reject, consistent with every other shared-tier-unavailable path. + listener.onResponse(null); + } + } + + // User-facing 429 for the node-local tier. Names the group (and attribute, if keyed) and the per-node limit knob. + private static OpenSearchRejectedExecutionException nodeTierRejection(ThrottlePlan plan) { + return new OpenSearchRejectedExecutionException( + "Request throttled: " + plan.describeTarget() + " reached its per-node limit of " + plan.nodeLimit + " concurrent requests." + ); + } + + // User-facing 429 for the cluster-level shared tier. Same shape as the node message, naming the cluster-wide limit. + private static OpenSearchRejectedExecutionException sharedTierRejection(ThrottlePlan plan) { + return new OpenSearchRejectedExecutionException( + "Request throttled: " + + plan.describeTarget() + + " reached its cluster-wide limit of " + + plan.sharedLimit + + " concurrent requests." + ); + } + + // Records a throttle rejection. Uses the raw state map, not the DEFAULT-fallback accessor, so a not-yet-registered + // group isn't misattributed to DEFAULT, and never lets a stats failure swallow the 429. + private void incrementThrottled(String workloadGroupId) { + try { + WorkloadGroupState workloadGroupState = workloadGroupsStateAccessor.getWorkloadGroupStateMap().get(workloadGroupId); + if (workloadGroupState != null) { + workloadGroupState.totalThrottled.inc(); + } + } catch (Exception statsException) { + logger.warn("Failed to record throttle stat for workload group [" + workloadGroupId + "]", statsException); + } + } + + /** + * Resolves the throttle configuration for a request into a bucket key plus the node and shared limits, or + * {@code null} when the request is not throttled: WLM disabled, default/unknown group, throttling not configured + * (both limits unset), or the request can't be attributed to a bucket (e.g. username/role with no principal). + */ + private ThrottlePlan resolveThrottlePlan(String workloadGroupId, String principal) { if (workloadManagementSettings.getWlmMode() != WlmMode.ENABLED) { return null; } if (workloadGroupId == null || workloadGroupId.equals(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get())) { return null; } - try { - WorkloadGroup workloadGroup = getWorkloadGroupById(workloadGroupId); - if (workloadGroup == null) { - return null; - } - Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling(); - int nodeLimit = WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling); - if (nodeLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) { - return null; - } - // A null bucket means the request can't be attributed (e.g. username/role with no principal) -> fail open. - String bucketKey = resolveThrottleBucketKey( - workloadGroupId, - WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling), - principal - ); - if (bucketKey == null) { - return null; - } - try { - return throttleTracker.acquire(bucketKey, nodeLimit); - } catch (OpenSearchRejectedExecutionException e) { - // Record the rejection without ever letting a stats failure swallow the 429. Use the raw state map, - // not the DEFAULT-fallback accessor, so a not-yet-registered group isn't misattributed to DEFAULT. - try { - WorkloadGroupState workloadGroupState = workloadGroupsStateAccessor.getWorkloadGroupStateMap().get(workloadGroupId); - if (workloadGroupState != null) { - workloadGroupState.totalThrottled.inc(); - } - } catch (Exception statsException) { - logger.warn("Failed to record throttle stat for workload group [" + workloadGroupId + "]", statsException); - } - throw e; - } - } catch (OpenSearchRejectedExecutionException e) { - throw e; // the intended 429 - } catch (Exception e) { - // A bug in the throttle path must never fail an otherwise-valid search, so fail open. - logger.warn("Skipping node-level throttle for workload group [" + workloadGroupId + "] due to an error", e); + WorkloadGroup workloadGroup = getWorkloadGroupById(workloadGroupId); + if (workloadGroup == null) { return null; } + Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling(); + int nodeLimit = WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling); + int sharedLimit = WorkloadGroupThrottleSettings.SHARED_LIMIT.get(throttling); + if (nodeLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT && sharedLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) { + return null; // throttling not configured + } + String attribute = WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling); + String value = resolveAttributeValue(attribute, principal); // null for whole-group; the principal value otherwise + String bucketKey = buildBucketKey(workloadGroupId, attribute, value); + if (bucketKey == null) { + return null; // can't attribute the request to a bucket -> fail open + } + return new ThrottlePlan(bucketKey, nodeLimit, sharedLimit, workloadGroup.getName(), attribute, value); + } + + // The resolved throttle configuration for a single request. Carries the human-readable group name and the throttle + // dimension (attribute + resolved value) purely so a rejection message can name them; the opaque bucketKey remains + // the identity used by both tracker tiers. + private static class ThrottlePlan { + final String bucketKey; + final int nodeLimit; + final int sharedLimit; + final String groupName; + final String attribute; // "group" | "username" | "role" + final String value; // the resolved principal value for username/role; null for whole-group throttling + + ThrottlePlan(String bucketKey, int nodeLimit, int sharedLimit, String groupName, String attribute, String value) { + this.bucketKey = bucketKey; + this.nodeLimit = nodeLimit; + this.sharedLimit = sharedLimit; + this.groupName = groupName; + this.attribute = attribute; + this.value = value; + } + + // "workload group [analytics]" or "workload group [analytics] for username [alice]". + String describeTarget() { + String base = "workload group [" + groupName + "]"; + if ("group".equals(attribute) || value == null) { + return base; + } + return base + " for " + attribute + " [" + value + "]"; + } } /** - * Builds the throttle bucket key: {@code :group} for whole-group throttling, or - * {@code ::} keyed by the matching principal subfield for {@code username}/{@code role}. - * Returns {@code null} (fail open, not throttled) when the principal is absent or lacks a token for the subfield. + * Resolves the throttle attribute's value for a request: {@code null} for whole-group throttling + * ({@code attribute == "group"}), otherwise the value of the matching {@code username}/{@code role} subfield + * parsed out of the principal header. Returns {@code null} (fail open, not throttled) when a keyed attribute has + * no usable principal token. */ - private String resolveThrottleBucketKey(String workloadGroupId, String attribute, String principal) { + private String resolveAttributeValue(String attribute, String principal) { if ("group".equals(attribute)) { - return workloadGroupId + ":group"; + return null; } if (principal == null || principal.isEmpty()) { return null; @@ -394,13 +508,28 @@ private String resolveThrottleBucketKey(String workloadGroupId, String attribute if (trimmed.startsWith(subfieldPrefix)) { String value = trimmed.substring(subfieldPrefix.length()); if (value.isEmpty() == false) { - return workloadGroupId + ":" + attribute + ":" + value; + return value; } } } return null; } + /** + * Assembles the opaque bucket key used by both tracker tiers: {@code :group} for whole-group throttling, + * or {@code ::} for a {@code username}/{@code role} bucket. Returns {@code null} when a + * keyed attribute has no resolved value (fail open). + */ + private String buildBucketKey(String workloadGroupId, String attribute, String value) { + if ("group".equals(attribute)) { + return workloadGroupId + ":group"; + } + if (value == null) { + return null; + } + return workloadGroupId + ":" + attribute + ":" + value; + } + private double getNormalisedRejectionThreshold(double limit, ResourceType resourceType) { if (resourceType == ResourceType.CPU) { return limit * workloadManagementSettings.getNodeLevelCpuRejectionThreshold(); diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java new file mode 100644 index 0000000000000..7ee5d10a5f782 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSharedThrottleService.java @@ -0,0 +1,417 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.cluster.ClusterChangedEvent; +import org.opensearch.cluster.ClusterStateListener; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.UUIDs; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.lease.Releasable; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.core.transport.TransportResponse; +import org.opensearch.threadpool.Scheduler; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportException; +import org.opensearch.transport.TransportRequest; +import org.opensearch.transport.TransportRequestOptions; +import org.opensearch.transport.TransportResponseHandler; +import org.opensearch.transport.TransportService; + +import java.io.IOException; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Cluster-level ({@code shared_limit}) throttle tier. Each throttle bucket has one authoritative in-flight counter + * living on the node that owns it per a consistent-hash ring ({@link ThrottleOwnerSelector}); this service is both + * the coordinator-side client that acquires/releases a shared permit for an incoming request and the + * owner-side host of the {@link SharedThrottleTracker} that answers those requests for the buckets it owns. + *

+ * Latency: the acquire is asynchronous so the calling (transport) thread never blocks on the round-trip — critical + * under the very bursts throttling defends against. When the coordinator itself owns the bucket, the acquire + * short-circuits to a direct in-memory tracker call with no network hop at all. + *

+ * Availability: any failure to reach the owner (timeout, disconnect, missing handler on an old node, or an empty + * ring) fails open — the request is admitted with no shared permit. This is a deliberate, un-toggleable + * choice: a single unreachable owner must not turn a network blip into a cluster-wide rejection storm for its + * share of buckets. A stuck lease is instead reclaimed by the owner's TTL sweep. + */ +@ExperimentalApi +public class WorkloadGroupSharedThrottleService implements ClusterStateListener { + + /** Internal transport action names (not client-facing REST). */ + public static final String ACQUIRE_ACTION_NAME = "internal:wlm/throttle/shared/acquire"; + public static final String RELEASE_ACTION_NAME = "internal:wlm/throttle/shared/release"; + + // Fixed operational constants. Deliberately not cluster settings: they are internal-coordination knobs an operator + // would never need to tune, and fail-open + a generous TTL make them safe as constants. Promote to a setting later + // (non-breaking) only if a real deployment need emerges. + // + // Timeout for the acquire round-trip to a bucket's owner. Small so a slow/unreachable owner fails open quickly + // rather than adding latency to the request path. + static final TimeValue ACQUIRE_TIMEOUT = TimeValue.timeValueMillis(200); + // Time-to-live for a lease on the owner. Reclaims a lease left behind by a crashed coordinator or a lost release + // RPC. Set generously above the typical request duration so a still-running request's lease is not reclaimed early. + // Accepted tradeoff: leases are NOT renewed, so a search that runs longer than this TTL has its lease reclaimed + // while still executing, which can transiently admit one-or-more requests beyond shared_limit for that bucket + // (self-correcting, fail-open direction — same flavor as the ring-rebalance transient breach). Long-running search + // is uncommon (no default search timeout, but heavy aggs/scripts can exceed it); if it becomes a problem the fix is + // lease renewal or deriving the TTL from a request deadline, not a larger constant. + static final long LEASE_TTL_NANOS = TimeValue.timeValueMinutes(5).nanos(); + // How often the owner sweeps expired leases. This is a pure memory-hygiene backstop, NOT a correctness mechanism: + // tryAcquire() already prunes a bucket's expired leases on every acquire, so an active bucket self-heals and can + // never over-reject due to a stuck lease. The sweep only reclaims the map entry for a bucket that abandoned a lease + // and then went completely idle (rejects nothing). An entry can't be reclaimed before its TTL elapses anyway, so + // sweeping faster than the TTL is pointless; run it infrequently. + static final TimeValue SWEEP_INTERVAL = TimeValue.timeValueMinutes(5); + + private static final Logger logger = LogManager.getLogger(WorkloadGroupSharedThrottleService.class); + + private final ClusterService clusterService; + private final ThreadPool threadPool; + private final TransportService transportService; + private final SharedThrottleTracker tracker; + + // Immutable ring snapshot, swapped wholesale on discovery-node changes so readers never see a half-built ring. + private final AtomicReference ring = new AtomicReference<>(); + private volatile Scheduler.Cancellable sweepTask; + + public WorkloadGroupSharedThrottleService(ClusterService clusterService, ThreadPool threadPool, TransportService transportService) { + this(clusterService, threadPool, transportService, new SharedThrottleTracker()); + } + + public WorkloadGroupSharedThrottleService( + ClusterService clusterService, + ThreadPool threadPool, + TransportService transportService, + SharedThrottleTracker tracker + ) { + this.clusterService = clusterService; + this.threadPool = threadPool; + this.transportService = transportService; + this.tracker = tracker; + // Start with an empty ring (fail-open) rather than reading cluster state here: during node construction the + // ClusterApplierService has no state yet. The ring is populated by the first clusterChanged whose eligible + // owner set differs from this empty ring — which includes a single-node cluster's very first applied state + // (its node set does not "change" relative to the coordinator-seeded initial state, so we must NOT gate on + // nodesChanged()). + this.ring.set(ThrottleOwnerSelector.fromDiscoveryNodes(DiscoveryNodes.EMPTY_NODES)); + clusterService.addListener(this); + transportService.registerRequestHandler( + ACQUIRE_ACTION_NAME, + ThreadPool.Names.SAME, + AcquireRequest::new, + (request, channel, task) -> channel.sendResponse(handleAcquire(request)) + ); + transportService.registerRequestHandler( + RELEASE_ACTION_NAME, + ThreadPool.Names.SAME, + ReleaseRequest::new, + (request, channel, task) -> { + tracker.release(request.bucketKey, request.leaseId); + channel.sendResponse(TransportResponse.Empty.INSTANCE); + } + ); + } + + /** Starts the periodic TTL sweep. Idempotent-safe to call once at node start. */ + public void start() { + // The ring is populated by clusterChanged (see the constructor), not here: at node start the initial cluster + // state may not be applied yet, so reading clusterService.state() here can throw "initial cluster state not + // set yet". This method only schedules the memory-hygiene sweep. + sweepTask = threadPool.scheduleWithFixedDelay(() -> { + try { + tracker.sweepExpired(); + } catch (Exception e) { + logger.warn("Shared throttle TTL sweep failed", e); + } + }, SWEEP_INTERVAL, ThreadPool.Names.GENERIC); + } + + public void stop() { + if (sweepTask != null) { + sweepTask.cancel(); + } + } + + @Override + public void clusterChanged(ClusterChangedEvent event) { + // Rebuild whenever the eligible-owner set differs from the current ring, rather than gating on + // event.nodesChanged(). nodesChanged() is a delta vs the previous applied state, and the coordinator seeds the + // initial applied state already containing the local node — so on a single-node cluster (or any cluster whose + // membership is stable after this service starts) nodesChanged() is never true and the ring would stay empty, + // silently disabling shared throttling. + // + // Compare the eligible-node set FIRST, without building the ring (cheap: a version-filtered data-node scan, no + // virtual-node hashing), so an unchanged membership — the common case, including every cluster-state update + // while WLM is disabled — pays nothing. Compare whole DiscoveryNodes, not just persistent ids, so a same-id + // restart (new ephemeral id/address) is treated as a change; otherwise the ring would keep the stale node and + // its buckets would fail open forever. + Set candidate = ThrottleOwnerSelector.eligibleNodeSet(event.state().nodes()); + if (ring.get().eligibleNodeSet().equals(candidate) == false) { + ring.set(ThrottleOwnerSelector.fromDiscoveryNodes(event.state().nodes())); + } + } + + /** + * Asynchronously acquires one cluster-level shared permit for {@code bucketKey}, notifying {@code listener} with: + *

    + *
  • a non-null {@link Releasable} — granted; close it on request completion to release the shared permit;
  • + *
  • {@code null} — fail open: admit the request without a shared permit (owner unreachable, empty + * ring, or any error). Nothing to release.
  • + *
+ * A denial (bucket at its shared limit) is signalled via {@link ActionListener#onFailure} with a + * message-less {@link OpenSearchRejectedExecutionException} — a pure "denied" marker. This service holds + * only the opaque bucket key, not the human-readable group/attribute, so {@link WorkloadGroupService} recomposes + * the user-facing 429 message; the exception's type (not its text) is what distinguishes a denial from a + * transport error (which also arrives via {@code onFailure} but must pass through as fail-open, not a 429). + * The listener may be invoked inline (local-owner short-circuit) or on a transport thread. + */ + public void acquireAsync(String bucketKey, int sharedLimit, ActionListener listener) { + final ThrottleOwnerSelector currentRing = ring.get(); + final DiscoveryNode owner = currentRing.ownerFor(bucketKey).orElse(null); + if (owner == null) { + listener.onResponse(null); // empty ring -> fail open + return; + } + + final String leaseId = leaseId(); + final long ttlNanos = LEASE_TTL_NANOS; + + // Local-owner short-circuit: this coordinator owns the bucket, so hit the tracker directly with no network hop. + if (owner.getId().equals(clusterService.localNode().getId())) { + if (tracker.tryAcquire(bucketKey, sharedLimit, leaseId, ttlNanos)) { + listener.onResponse(releaseLocal(bucketKey, leaseId)); + } else { + listener.onFailure(deniedMarker()); + } + return; + } + + // Owner known-disconnected: fail open immediately rather than sending into a dead socket and waiting out + // ACQUIRE_TIMEOUT. This is an in-memory connection-map lookup (ConcurrentHashMap.containsKey), not a network + // call, so it costs nothing. It collapses the bulk of the post-node-drop detection window; the tiny remaining + // sub-window (socket dead but not yet flagged) still hits the timeout and is handled by handleException below. + if (transportService.nodeConnected(owner) == false) { + listener.onResponse(null); + return; + } + + final TransportRequestOptions options = TransportRequestOptions.builder().withTimeout(ACQUIRE_TIMEOUT).build(); + transportService.sendRequest( + owner, + ACQUIRE_ACTION_NAME, + new AcquireRequest(bucketKey, sharedLimit, leaseId, ttlNanos), + options, + new TransportResponseHandler() { + @Override + public AcquireResponse read(StreamInput in) throws IOException { + return new AcquireResponse(in); + } + + @Override + public void handleResponse(AcquireResponse response) { + if (response.granted) { + listener.onResponse(releaseRemote(owner, bucketKey, leaseId)); + } else { + listener.onFailure(deniedMarker()); + } + } + + @Override + public void handleException(TransportException exp) { + // Owner unreachable / timed out -> fail open. The owner may have GRANTED this lease before the + // response was lost (e.g. the acquire arrived but the reply timed out), which would otherwise + // occupy a shared slot until its TTL and cause false 429s once connectivity recovers. Send a + // best-effort release for this leaseId to reclaim it immediately; release-by-id is idempotent, so + // if no lease was created it is a harmless no-op. Then admit (fail open). + logger.debug("Shared throttle acquire to owner [{}] for bucket [{}] failed; failing open", owner.getId(), bucketKey); + sendRelease(owner, bucketKey, leaseId); + listener.onResponse(null); + } + + @Override + public String executor() { + return ThreadPool.Names.SAME; + } + } + ); + } + + // Owner-side admission. Package-private for tests. + AcquireResponse handleAcquire(AcquireRequest request) { + boolean granted = tracker.tryAcquire(request.bucketKey, request.sharedLimit, request.leaseId, request.ttlNanos); + return new AcquireResponse(granted); + } + + private Releasable releaseLocal(String bucketKey, String leaseId) { + return releaseOnce(() -> tracker.release(bucketKey, leaseId)); + } + + private Releasable releaseRemote(DiscoveryNode owner, String bucketKey, String leaseId) { + return releaseOnce(() -> sendRelease(owner, bucketKey, leaseId)); + } + + // Fire-and-forget RELEASE RPC to the bucket owner. Bounded by the same timeout as acquire so a half-open + // connection can't leave the response handler pending until the connection is torn down. A lost release is not + // fatal — the owner's TTL sweep reclaims the lease — so failures are logged at debug only. + private void sendRelease(DiscoveryNode owner, String bucketKey, String leaseId) { + final TransportRequestOptions options = TransportRequestOptions.builder().withTimeout(ACQUIRE_TIMEOUT).build(); + transportService.sendRequest( + owner, + RELEASE_ACTION_NAME, + new ReleaseRequest(bucketKey, leaseId), + options, + new TransportResponseHandler() { + @Override + public TransportResponse.Empty read(StreamInput in) { + return TransportResponse.Empty.INSTANCE; + } + + @Override + public void handleResponse(TransportResponse.Empty response) {} + + @Override + public void handleException(TransportException exp) { + logger.debug( + "Shared throttle release to owner [{}] for bucket [{}] failed; TTL will reclaim", + owner.getId(), + bucketKey + ); + } + + @Override + public String executor() { + return ThreadPool.Names.SAME; + } + } + ); + } + + // Guards against a double close (e.g. onRequestEnd and onRequestFailure) releasing twice. + private static Releasable releaseOnce(Runnable release) { + final AtomicBoolean released = new AtomicBoolean(false); + return () -> { + if (released.compareAndSet(false, true)) { + release.run(); + } + }; + } + + private static String leaseId() { + return UUIDs.base64UUID(); + } + + // Message-less "denied" marker: the bucket is at its shared limit. Carries no user text because this service has + // only the opaque bucket key; WorkloadGroupService recomposes the user-facing 429 with the group name/attribute. + // The exception TYPE (OpenSearchRejectedExecutionException) is the signal — the orchestrator uses it to tell a + // denial (recompose as a 429) apart from a transport error (pass through as fail-open). + private static OpenSearchRejectedExecutionException deniedMarker() { + return new OpenSearchRejectedExecutionException(); + } + + // Package-private accessor for tests. + SharedThrottleTracker tracker() { + return tracker; + } + + ThrottleOwnerSelector ring() { + return ring.get(); + } + + /** + * Acquire RPC: a coordinator asks the bucket's owner to admit one request under {@code sharedLimit}. + */ + public static class AcquireRequest extends TransportRequest { + final String bucketKey; + final int sharedLimit; + final String leaseId; + final long ttlNanos; + + AcquireRequest(String bucketKey, int sharedLimit, String leaseId, long ttlNanos) { + this.bucketKey = bucketKey; + this.sharedLimit = sharedLimit; + this.leaseId = leaseId; + this.ttlNanos = ttlNanos; + } + + AcquireRequest(StreamInput in) throws IOException { + super(in); + this.bucketKey = in.readString(); + this.sharedLimit = in.readVInt(); + this.leaseId = in.readString(); + this.ttlNanos = in.readVLong(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(bucketKey); + out.writeVInt(sharedLimit); + out.writeString(leaseId); + out.writeVLong(ttlNanos); + } + } + + /** + * Acquire RPC response: whether the owner granted a shared permit. + */ + public static class AcquireResponse extends TransportResponse { + final boolean granted; + + AcquireResponse(boolean granted) { + this.granted = granted; + } + + AcquireResponse(StreamInput in) throws IOException { + this.granted = in.readBoolean(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeBoolean(granted); + } + } + + /** + * Release RPC: fire-and-forget, tells the owner to drop a previously-granted lease. + */ + public static class ReleaseRequest extends TransportRequest { + final String bucketKey; + final String leaseId; + + ReleaseRequest(String bucketKey, String leaseId) { + this.bucketKey = bucketKey; + this.leaseId = leaseId; + } + + ReleaseRequest(StreamInput in) throws IOException { + super(in); + this.bucketKey = in.readString(); + this.leaseId = in.readString(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(bucketKey); + out.writeString(leaseId); + } + } +} diff --git a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java index a7535984ceca0..66b5ea3d7bf82 100644 --- a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java +++ b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java @@ -164,6 +164,7 @@ import org.opensearch.common.Nullable; import org.opensearch.common.SetOnce; import org.opensearch.common.cache.module.CacheModule; +import org.opensearch.common.lease.Releasable; import org.opensearch.common.network.NetworkModule; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.IndexScopedSettings; @@ -297,6 +298,8 @@ import static org.hamcrest.Matchers.iterableWithSize; import static org.hamcrest.Matchers.lessThanOrEqualTo; import static org.hamcrest.Matchers.notNullValue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -2381,6 +2384,13 @@ public void onFailure(final Exception e) { ); SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory = new SearchRequestOperationsCompositeListenerFactory(); + // admit() is void+async: a default mock would never invoke the listener, leaving every search hung. + // Stub it to admit immediately with no throttle permit (the not-throttled path), matching production. + WorkloadGroupService workloadGroupService = mock(WorkloadGroupService.class); + doAnswer(invocation -> { + invocation.>getArgument(2).onResponse(null); + return null; + }).when(workloadGroupService).acquireThrottlePermit(any(), any(), any()); actions.put( SearchAction.INSTANCE, new TransportSearchAction( @@ -2411,7 +2421,7 @@ public void onFailure(final Exception e) { NoopTracer.INSTANCE, new TaskResourceTrackingService(settings, clusterSettings, threadPool), mockIndicesService, - mock(WorkloadGroupService.class) + workloadGroupService ) ); actions.put( diff --git a/server/src/test/java/org/opensearch/wlm/NodeThrottleTrackerTests.java b/server/src/test/java/org/opensearch/wlm/NodeThrottleTrackerTests.java new file mode 100644 index 0000000000000..80d1bbaa90a9a --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/NodeThrottleTrackerTests.java @@ -0,0 +1,120 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.common.lease.Releasable; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +public class NodeThrottleTrackerTests extends OpenSearchTestCase { + + public void testAcquireUnderLimitSucceeds() { + NodeThrottleTracker tracker = new NodeThrottleTracker(); + Releasable p1 = tracker.tryAcquire("bucket", 2); + Releasable p2 = tracker.tryAcquire("bucket", 2); + assertNotNull(p1); + assertNotNull(p2); + assertEquals(2, tracker.inFlight("bucket")); + p1.close(); + p2.close(); + } + + public void testAcquireAtLimitReturnsNull() { + NodeThrottleTracker tracker = new NodeThrottleTracker(); + assertNotNull(tracker.tryAcquire("bucket", 1)); + // at the limit -> tryAcquire returns null (production then falls through to the shared tier or rejects) + assertNull(tracker.tryAcquire("bucket", 1)); + // declined acquire must not leave the count inflated + assertEquals(1, tracker.inFlight("bucket")); + } + + public void testReleaseFreesAPermit() { + NodeThrottleTracker tracker = new NodeThrottleTracker(); + Releasable p = tracker.tryAcquire("bucket", 1); + assertNotNull(p); + // at the limit + assertNull(tracker.tryAcquire("bucket", 1)); + p.close(); + // permit freed -> a fresh acquire now succeeds + Releasable p2 = tracker.tryAcquire("bucket", 1); + assertNotNull(p2); + assertEquals(1, tracker.inFlight("bucket")); + p2.close(); + } + + public void testDrainToZeroRemovesBucketThenReacquire() { + NodeThrottleTracker tracker = new NodeThrottleTracker(); + Releasable p = tracker.tryAcquire("bucket", 5); + assertNotNull(p); + assertEquals(1, tracker.inFlight("bucket")); + p.close(); + assertEquals(0, tracker.inFlight("bucket")); + // re-acquiring after the bucket drained (and was removed) works and starts from 1 + Releasable p2 = tracker.tryAcquire("bucket", 5); + assertNotNull(p2); + assertEquals(1, tracker.inFlight("bucket")); + p2.close(); + } + + public void testReleaseIsIdempotent() { + NodeThrottleTracker tracker = new NodeThrottleTracker(); + Releasable p = tracker.tryAcquire("bucket", 5); + assertNotNull(p); + p.close(); + p.close(); // double close must not decrement twice + assertEquals(0, tracker.inFlight("bucket")); + } + + public void testBucketsAreIndependent() { + NodeThrottleTracker tracker = new NodeThrottleTracker(); + assertNotNull(tracker.tryAcquire("a", 1)); + // "a" is full but "b" is a separate bucket + assertNull(tracker.tryAcquire("a", 1)); + Releasable pb = tracker.tryAcquire("b", 1); + assertNotNull(pb); + assertEquals(1, tracker.inFlight("a")); + assertEquals(1, tracker.inFlight("b")); + pb.close(); + } + + public void testConcurrentAcquireAdmitsExactlyLimit() throws Exception { + // Many threads race to acquire the same bucket at once. tryAcquire must admit exactly nodeLimit of them — + // never fewer (the over-decline race where all racers observe an inflated shared count and roll back) and + // never more (over-admit). Each caller decides on its own captured post-increment value. + final NodeThrottleTracker tracker = new NodeThrottleTracker(); + final int limit = 10; + final int threads = 64; + final CountDownLatch start = new CountDownLatch(1); + final AtomicInteger granted = new AtomicInteger(0); + Thread[] workers = new Thread[threads]; + for (int i = 0; i < threads; i++) { + workers[i] = new Thread(() -> { + try { + start.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + if (tracker.tryAcquire("hot", limit) != null) { + granted.incrementAndGet(); + } + }); + workers[i].start(); + } + start.countDown(); + for (Thread w : workers) { + w.join(TimeUnit.SECONDS.toMillis(30)); + } + assertEquals("must admit exactly the limit under contention", limit, granted.get()); + assertEquals(limit, tracker.inFlight("hot")); + } +} diff --git a/server/src/test/java/org/opensearch/wlm/SharedThrottleTrackerTests.java b/server/src/test/java/org/opensearch/wlm/SharedThrottleTrackerTests.java new file mode 100644 index 0000000000000..bcc12b663e61a --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/SharedThrottleTrackerTests.java @@ -0,0 +1,152 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.test.OpenSearchTestCase; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +public class SharedThrottleTrackerTests extends OpenSearchTestCase { + + private static final long TTL = TimeUnit.MINUTES.toNanos(5); + + public void testGrantsUpToLimitThenDenies() { + SharedThrottleTracker tracker = new SharedThrottleTracker(); + assertTrue(tracker.tryAcquire("b", 2, "l1", TTL)); + assertTrue(tracker.tryAcquire("b", 2, "l2", TTL)); + assertFalse("third acquire must be denied at limit 2", tracker.tryAcquire("b", 2, "l3", TTL)); + assertEquals(2, tracker.inFlight("b")); + } + + public void testReleaseFreesASlot() { + SharedThrottleTracker tracker = new SharedThrottleTracker(); + assertTrue(tracker.tryAcquire("b", 1, "l1", TTL)); + assertFalse(tracker.tryAcquire("b", 1, "l2", TTL)); + tracker.release("b", "l1"); + assertEquals(0, tracker.inFlight("b")); + assertTrue("slot freed after release", tracker.tryAcquire("b", 1, "l3", TTL)); + } + + public void testReleaseByUnknownLeaseIsNoOp() { + SharedThrottleTracker tracker = new SharedThrottleTracker(); + assertTrue(tracker.tryAcquire("b", 2, "l1", TTL)); + // Unknown lease (already swept, double release, or belonged to a previous owner of a remapped bucket). + tracker.release("b", "does-not-exist"); + tracker.release("unknown-bucket", "l1"); + assertEquals("a stray release must not decrement a live lease", 1, tracker.inFlight("b")); + } + + public void testDoubleReleaseDecrementsOnce() { + SharedThrottleTracker tracker = new SharedThrottleTracker(); + assertTrue(tracker.tryAcquire("b", 2, "l1", TTL)); + assertTrue(tracker.tryAcquire("b", 2, "l2", TTL)); + tracker.release("b", "l1"); + tracker.release("b", "l1"); // second release of the same lease id is a no-op + assertEquals(1, tracker.inFlight("b")); + } + + public void testDrainedBucketIsRemoved() { + SharedThrottleTracker tracker = new SharedThrottleTracker(); + assertTrue(tracker.tryAcquire("b", 5, "l1", TTL)); + assertEquals(1, tracker.activeBuckets()); + tracker.release("b", "l1"); + assertEquals("bucket entry removed when it drains to zero", 0, tracker.activeBuckets()); + } + + public void testExpiredLeaseIsReclaimedBySweep() { + AtomicLong clock = new AtomicLong(0); + SharedThrottleTracker tracker = new SharedThrottleTracker(clock::get); + assertTrue(tracker.tryAcquire("b", 1, "l1", 100)); + assertFalse(tracker.tryAcquire("b", 1, "l2", 100)); // at limit + clock.set(101); // lease l1 has now expired + tracker.sweepExpired(); + assertEquals(0, tracker.inFlight("b")); + assertTrue("slot reclaimed after TTL sweep", tracker.tryAcquire("b", 1, "l3", 100)); + } + + public void testExpiredLeaseReclaimedLazilyOnAcquire() { + AtomicLong clock = new AtomicLong(0); + SharedThrottleTracker tracker = new SharedThrottleTracker(clock::get); + assertTrue(tracker.tryAcquire("b", 1, "l1", 100)); + clock.set(101); // l1 expired; a fresh acquire should prune it and be granted without an explicit sweep + assertTrue(tracker.tryAcquire("b", 1, "l2", 100)); + assertEquals(1, tracker.inFlight("b")); + } + + public void testSaturatedBucketWithExpiredLeaseIsReclaimedOnAcquire() { + // Prune now runs only when the bucket is at/over its limit; verify a fully-saturated bucket with one expired + // lease still admits the next acquire (the case the "prune only when full" optimization must not break). + AtomicLong clock = new AtomicLong(0); + SharedThrottleTracker tracker = new SharedThrottleTracker(clock::get); + assertTrue(tracker.tryAcquire("b", 2, "l1", 100)); // expires at 100 + assertTrue(tracker.tryAcquire("b", 2, "l2", 500)); // expires at 500 + assertFalse("bucket is full", tracker.tryAcquire("b", 2, "l3", 100)); + clock.set(200); // l1 expired, l2 still live + // Bucket is at the limit -> acquire prunes l1 and grants; l2 remains. + assertTrue(tracker.tryAcquire("b", 2, "l4", 500)); + assertEquals(2, tracker.inFlight("b")); + } + + public void testSweepRacingLateReleaseDoesNotUnderCount() { + AtomicLong clock = new AtomicLong(0); + SharedThrottleTracker tracker = new SharedThrottleTracker(clock::get); + assertTrue(tracker.tryAcquire("b", 5, "l1", 100)); + assertTrue(tracker.tryAcquire("b", 5, "l2", 100)); + clock.set(101); // both expired + tracker.sweepExpired(); // removes l1 and l2 + // A late release for an already-swept lease must be a no-op, not push the count negative. + tracker.release("b", "l1"); + assertEquals(0, tracker.inFlight("b")); + // fresh lease unaffected + clock.set(150); + assertTrue(tracker.tryAcquire("b", 5, "l3", 100)); + assertEquals(1, tracker.inFlight("b")); + } + + public void testConcurrentAcquireNeverExceedsLimit() throws Exception { + final SharedThrottleTracker tracker = new SharedThrottleTracker(); + final int limit = 10; + final int threads = 64; + final CountDownLatch start = new CountDownLatch(1); + final AtomicInteger granted = new AtomicInteger(0); + Thread[] workers = new Thread[threads]; + for (int i = 0; i < threads; i++) { + final int id = i; + workers[i] = new Thread(() -> { + try { + start.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + if (tracker.tryAcquire("hot", limit, "lease-" + id, TTL)) { + granted.incrementAndGet(); + } + }); + workers[i].start(); + } + start.countDown(); + for (Thread w : workers) { + w.join(); + } + assertEquals("exactly limit grants under contention", limit, granted.get()); + assertEquals(limit, tracker.inFlight("hot")); + } + + public void testBucketsAreIndependent() { + SharedThrottleTracker tracker = new SharedThrottleTracker(); + assertTrue(tracker.tryAcquire("a", 1, "la", TTL)); + assertTrue("different bucket has its own budget", tracker.tryAcquire("b", 1, "lb", TTL)); + assertFalse(tracker.tryAcquire("a", 1, "la2", TTL)); + assertEquals(2, tracker.activeBuckets()); + } +} diff --git a/server/src/test/java/org/opensearch/wlm/ThrottleOwnerSelectorTests.java b/server/src/test/java/org/opensearch/wlm/ThrottleOwnerSelectorTests.java new file mode 100644 index 0000000000000..815f577f35cdd --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/ThrottleOwnerSelectorTests.java @@ -0,0 +1,141 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.Version; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodeRole; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; + +public class ThrottleOwnerSelectorTests extends OpenSearchTestCase { + + private DiscoveryNode dataNode(String id, Version version) { + return new DiscoveryNode( + "name_" + id, + id, + buildNewFakeTransportAddress(), + Collections.emptyMap(), + Set.of(DiscoveryNodeRole.DATA_ROLE), + version + ); + } + + private DiscoveryNode managerOnlyNode(String id) { + return new DiscoveryNode( + "name_" + id, + id, + buildNewFakeTransportAddress(), + Collections.emptyMap(), + Set.of(DiscoveryNodeRole.CLUSTER_MANAGER_ROLE), + Version.CURRENT + ); + } + + private DiscoveryNodes nodesOf(DiscoveryNode... nodes) { + DiscoveryNodes.Builder builder = DiscoveryNodes.builder(); + for (DiscoveryNode node : nodes) { + builder.add(node); + } + builder.localNodeId(nodes[0].getId()); + return builder.build(); + } + + public void testEmptyRingWhenNoDataNodes() { + ThrottleOwnerSelector selector = ThrottleOwnerSelector.fromDiscoveryNodes(nodesOf(managerOnlyNode("m1"))); + assertTrue(selector.isEmpty()); + assertFalse("empty ring must yield no owner (fail open)", selector.ownerFor("group-a:group").isPresent()); + } + + public void testClusterManagerOnlyNodesExcluded() { + ThrottleOwnerSelector selector = ThrottleOwnerSelector.fromDiscoveryNodes( + nodesOf(dataNode("d1", Version.CURRENT), managerOnlyNode("m1")) + ); + // Only the data node is eligible; the manager-only node must never be an owner. + for (int i = 0; i < 200; i++) { + assertEquals("d1", selector.ownerFor("bucket-" + i).orElseThrow().getId()); + } + } + + public void testOldVersionNodesExcludedFromRing() { + DiscoveryNode newNode = dataNode("new", Version.CURRENT); + DiscoveryNode oldNode = dataNode("old", Version.V_3_6_0); // below MIN_OWNER_VERSION + ThrottleOwnerSelector selector = ThrottleOwnerSelector.fromDiscoveryNodes(nodesOf(newNode, oldNode)); + assertEquals(1, selector.eligibleNodeSet().size()); + for (int i = 0; i < 200; i++) { + assertEquals("only the upgraded node may own buckets", "new", selector.ownerFor("bucket-" + i).orElseThrow().getId()); + } + } + + public void testDeterministicForSameNodeSet() { + DiscoveryNodes nodes = nodesOf(dataNode("a", Version.CURRENT), dataNode("b", Version.CURRENT), dataNode("c", Version.CURRENT)); + ThrottleOwnerSelector s1 = ThrottleOwnerSelector.fromDiscoveryNodes(nodes); + ThrottleOwnerSelector s2 = ThrottleOwnerSelector.fromDiscoveryNodes(nodes); + for (int i = 0; i < 500; i++) { + String key = "bucket-" + i; + assertEquals(s1.ownerFor(key).orElseThrow().getId(), s2.ownerFor(key).orElseThrow().getId()); + } + } + + public void testDistributionAcrossNodesIsReasonablyBalanced() { + ThrottleOwnerSelector selector = ThrottleOwnerSelector.fromDiscoveryNodes( + nodesOf(dataNode("a", Version.CURRENT), dataNode("b", Version.CURRENT), dataNode("c", Version.CURRENT)) + ); + Map counts = new HashMap<>(); + int n = 3000; + for (int i = 0; i < n; i++) { + String owner = selector.ownerFor("bucket-" + i).orElseThrow().getId(); + counts.merge(owner, 1, Integer::sum); + } + // With 128 virtual nodes each, every physical node should get a non-trivial share (well above 10%). + assertEquals(3, counts.size()); + for (int c : counts.values()) { + assertTrue("each node should own a meaningful share, saw " + counts, c > n / 10); + } + } + + public void testRebalanceOnlyRemapsSmallFraction() { + DiscoveryNode a = dataNode("a", Version.CURRENT); + DiscoveryNode b = dataNode("b", Version.CURRENT); + DiscoveryNode c = dataNode("c", Version.CURRENT); + DiscoveryNode d = dataNode("d", Version.CURRENT); + ThrottleOwnerSelector before = ThrottleOwnerSelector.fromDiscoveryNodes(nodesOf(a, b, c, d)); + ThrottleOwnerSelector after = ThrottleOwnerSelector.fromDiscoveryNodes(nodesOf(a, b, c)); // d leaves + + int n = 5000; + int remapped = 0; + for (int i = 0; i < n; i++) { + String key = "bucket-" + i; + String ownerBefore = before.ownerFor(key).orElseThrow().getId(); + String ownerAfter = after.ownerFor(key).orElseThrow().getId(); + if (ownerBefore.equals(ownerAfter) == false) { + remapped++; + } + } + // Removing 1 of 4 nodes should remap roughly 1/4 of keys. Allow generous slack, but it must be far from "all". + double fraction = (double) remapped / n; + assertTrue("remap fraction should be well under half, was " + fraction, fraction < 0.45); + // Keys not owned by the departed node 'd' should mostly keep their owner. + for (int i = 0; i < n; i++) { + String key = "bucket-" + i; + if (before.ownerFor(key).orElseThrow().getId().equals("d") == false) { + assertEquals( + "survivor-owned keys must not move", + before.ownerFor(key).orElseThrow().getId(), + after.ownerFor(key).orElseThrow().getId() + ); + } + } + } +} diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java index 18ea170d86e95..f8fabb8ff0b85 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java @@ -13,11 +13,15 @@ import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.WorkloadGroup; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodeRole; +import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.lease.Releasable; import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.action.ActionListener; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import org.opensearch.search.backpressure.trackers.NodeDuressTrackers; import org.opensearch.tasks.Task; @@ -36,6 +40,7 @@ import java.util.HashSet; import java.util.Map; import java.util.Set; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BooleanSupplier; import org.mockito.ArgumentCaptor; @@ -488,6 +493,25 @@ private void stubClusterStateWithGroup(WorkloadGroup wg) { when(metadata.workloadGroups()).thenReturn(Map.of(wg.get_id(), wg)); } + /** + * Drives the async {@link WorkloadGroupService#acquireThrottlePermit} back into the synchronous "permit or 429" + * contract these node-tier tests assert on. With no shared throttle service wired, it completes inline: responds + * with a permit (or null when not throttled) or fails with the 429. Rethrows the failure so {@code expectThrows} + * works. + */ + private Releasable acquireThrottlePermitSync(WorkloadGroupService service, String workloadGroupId, String principal) { + AtomicReference permit = new AtomicReference<>(); + AtomicReference failure = new AtomicReference<>(); + service.acquireThrottlePermit(workloadGroupId, principal, ActionListener.wrap(permit::set, failure::set)); + if (failure.get() != null) { + if (failure.get() instanceof RuntimeException) { + throw (RuntimeException) failure.get(); + } + throw new RuntimeException(failure.get()); + } + return permit.get(); + } + private WorkloadGroup throttledGroup(String id, Settings throttling) { return new WorkloadGroup( id + "-name", @@ -506,12 +530,12 @@ public void testAcquireThrottleReturnsNullWhenNodeLimitUnset() { when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); stubClusterStateWithGroup(throttledGroup("wg-1", Settings.EMPTY)); // throttling not configured - assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); + assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); } public void testAcquireThrottleReturnsNullWhenWlmDisabled() { when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED); - assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); + assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); } public void testAcquireThrottleRejectsAtLimitAndIncrementsStat() { @@ -520,15 +544,19 @@ public void testAcquireThrottleRejectsAtLimitAndIncrementsStat() { Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).build(); stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); - Releasable permit = workloadGroupService.acquireThrottleOrReject("wg-1", null); // first admit succeeds + Releasable permit = acquireThrottlePermitSync(workloadGroupService, "wg-1", null); // first admit succeeds assertNotNull(permit); // second admit hits node_limit of 1 -> 429 + total_throttled incremented - expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject("wg-1", null)); + OpenSearchRejectedExecutionException e = expectThrows( + OpenSearchRejectedExecutionException.class, + () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", null) + ); + assertEquals("Request throttled: workload group [wg-1-name] reached its per-node limit of 1 concurrent requests.", e.getMessage()); assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); // releasing the first permit frees the slot so a subsequent acquire succeeds permit.close(); - assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); + assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); } public void testAcquireThrottleUsernameKeepsPerUserBuckets() { @@ -537,22 +565,26 @@ public void testAcquireThrottleUsernameKeepsPerUserBuckets() { Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 1).build(); stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); - // alice takes her single slot; a second alice request is rejected. - Releasable alice = workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice"); + // alice takes her single slot; a second alice request is rejected — and the message names the username. + Releasable alice = acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|alice"); assertNotNull(alice); - expectThrows( + OpenSearchRejectedExecutionException e = expectThrows( OpenSearchRejectedExecutionException.class, - () -> workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice") + () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|alice") + ); + assertEquals( + "Request throttled: workload group [wg-1-name] for username [alice] reached its per-node limit of 1 concurrent requests.", + e.getMessage() ); assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); // bob is a different bucket, so he is admitted even while alice is at her limit. - Releasable bob = workloadGroupService.acquireThrottleOrReject("wg-1", "username|bob"); + Releasable bob = acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|bob"); assertNotNull(bob); // releasing alice frees her bucket alice.close(); - assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice")); + assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|alice")); } public void testAcquireThrottleUsernameWithCommaDoesNotCollide() { @@ -567,12 +599,12 @@ public void testAcquireThrottleUsernameWithCommaDoesNotCollide() { // user "a" is a genuinely different principal String userA = "username|a"; - Releasable ab = workloadGroupService.acquireThrottleOrReject("wg-1", userAB); // fills "a,b" bucket + Releasable ab = acquireThrottlePermitSync(workloadGroupService, "wg-1", userAB); // fills "a,b" bucket assertNotNull(ab); // user "a" must NOT be treated as the same bucket as "a,b" -> still admitted - assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", userA)); + assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", userA)); // a second "a,b" request hits the "a,b" bucket limit -> rejected - expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject("wg-1", userAB)); + expectThrows(OpenSearchRejectedExecutionException.class, () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", userAB)); } public void testAcquireThrottleRolePicksMatchingSubfieldFromMultiTokenPrincipal() { @@ -583,10 +615,10 @@ public void testAcquireThrottleRolePicksMatchingSubfieldFromMultiTokenPrincipal( // A principal header may carry both subfields; the role bucket must key off the role token only. String delim = WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER; - assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", "username|alice" + delim + "role|admin")); + assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|alice" + delim + "role|admin")); expectThrows( OpenSearchRejectedExecutionException.class, - () -> workloadGroupService.acquireThrottleOrReject("wg-1", "username|bob" + delim + "role|admin") + () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", "username|bob" + delim + "role|admin") ); assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); } @@ -598,9 +630,9 @@ public void testAcquireThrottleFailsOpenWhenPrincipalMissingForUsername() { stubClusterStateWithGroup(throttledGroup("wg-1", throttling)); // No principal (e.g. security plugin not installed) or no matching subfield -> not throttled (fail open). - assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); - assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", "")); - assertNull(workloadGroupService.acquireThrottleOrReject("wg-1", "role|admin")); // no username token + assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); + assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", "")); + assertNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", "role|admin")); // no username token assertEquals(0, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); } @@ -628,9 +660,9 @@ public void testAcquireThrottleStillRejectsWhenStatUpdateFails() { new HashSet<>() ); - assertNotNull(serviceWithNullState.acquireThrottleOrReject("wg-1", null)); // first admit fills the single slot + assertNotNull(acquireThrottlePermitSync(serviceWithNullState, "wg-1", null)); // first admit fills the single slot // second acquire is over the limit; a null state must not let the stat update swallow the 429 - expectThrows(OpenSearchRejectedExecutionException.class, () -> serviceWithNullState.acquireThrottleOrReject("wg-1", null)); + expectThrows(OpenSearchRejectedExecutionException.class, () -> acquireThrottlePermitSync(serviceWithNullState, "wg-1", null)); // accessor whose state-map lookup throws must also still propagate the 429 WorkloadGroupsStateAccessor throwingStateAccessor = Mockito.mock(WorkloadGroupsStateAccessor.class); @@ -646,8 +678,8 @@ public void testAcquireThrottleStillRejectsWhenStatUpdateFails() { new HashSet<>() ); - assertNotNull(serviceWithThrowingState.acquireThrottleOrReject("wg-1", null)); // fills the single slot - expectThrows(OpenSearchRejectedExecutionException.class, () -> serviceWithThrowingState.acquireThrottleOrReject("wg-1", null)); + assertNotNull(acquireThrottlePermitSync(serviceWithThrowingState, "wg-1", null)); // fills the single slot + expectThrows(OpenSearchRejectedExecutionException.class, () -> acquireThrottlePermitSync(serviceWithThrowingState, "wg-1", null)); } /** @@ -662,8 +694,8 @@ public void testAcquireThrottleDoesNotMisattributeToDefaultDuringRegistrationLag // DEFAULT group state exists, but wg-1 is NOT yet registered (registration lag). mockWorkloadGroupsStateAccessor.addNewWorkloadGroup(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get()); - assertNotNull(workloadGroupService.acquireThrottleOrReject("wg-1", null)); // fills the single slot - expectThrows(OpenSearchRejectedExecutionException.class, () -> workloadGroupService.acquireThrottleOrReject("wg-1", null)); + assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); // fills the single slot + expectThrows(OpenSearchRejectedExecutionException.class, () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); // the rejection must NOT have landed on the DEFAULT group assertEquals( @@ -673,6 +705,70 @@ public void testAcquireThrottleDoesNotMisattributeToDefaultDuringRegistrationLag ); } + /** + * The headline two-tier path: node_limit=1 + shared_limit=1 with a shared throttle service wired in. The first + * request takes the local slot, the second overflows and takes the shared slot, and the third is rejected (both + * tiers full) with total_throttled incremented. Uses a single-data-node ring so the shared acquire resolves to + * this node and runs via the local-owner short-circuit (no real network). + */ + public void testAdmitFallsThroughNodeTierToSharedTier() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + mockWorkloadGroupsStateAccessor.addNewWorkloadGroup("wg-1"); + + Settings throttling = Settings.builder().put("attribute", "group").put("node_limit", 1).put("shared_limit", 1).build(); + WorkloadGroup group = throttledGroup("wg-1", throttling); + + // Cluster state must expose both the group metadata and a single data node (so this node owns every bucket). + DiscoveryNode localNode = new DiscoveryNode( + "local", + "local", + buildNewFakeTransportAddress(), + Collections.emptyMap(), + Set.of(DiscoveryNodeRole.DATA_ROLE), + org.opensearch.Version.CURRENT + ); + DiscoveryNodes nodes = DiscoveryNodes.builder().add(localNode).localNodeId("local").build(); + ClusterState clusterState = Mockito.mock(ClusterState.class); + Metadata metadata = Mockito.mock(Metadata.class); + when(mockClusterService.state()).thenReturn(clusterState); + when(mockClusterService.localNode()).thenReturn(localNode); + when(clusterState.metadata()).thenReturn(metadata); + when(clusterState.nodes()).thenReturn(nodes); + when(metadata.workloadGroups()).thenReturn(Map.of(group.get_id(), group)); + + WorkloadGroupSharedThrottleService sharedService = new WorkloadGroupSharedThrottleService( + mockClusterService, + mockThreadPool, + Mockito.mock(org.opensearch.transport.TransportService.class) + ); + // Populate the ring (the constructor starts empty; a nodes-changed event builds it — matches production). + ClusterState previous = Mockito.mock(ClusterState.class); + when(previous.nodes()).thenReturn(DiscoveryNodes.EMPTY_NODES); + sharedService.clusterChanged(new ClusterChangedEvent("test", clusterState, previous)); + workloadGroupService.setSharedThrottleService(sharedService); + + // 1st: local node_limit slot. + Releasable local = acquireThrottlePermitSync(workloadGroupService, "wg-1", null); + assertNotNull(local); + // 2nd: local exhausted -> falls through to shared_limit slot. + Releasable shared = acquireThrottlePermitSync(workloadGroupService, "wg-1", null); + assertNotNull(shared); + // 3rd: both tiers full -> 429 from the shared tier (with the cluster-wide message), and total_throttled bumped. + OpenSearchRejectedExecutionException e = expectThrows( + OpenSearchRejectedExecutionException.class, + () -> acquireThrottlePermitSync(workloadGroupService, "wg-1", null) + ); + assertEquals( + "Request throttled: workload group [wg-1-name] reached its cluster-wide limit of 1 concurrent requests.", + e.getMessage() + ); + assertEquals(1, mockWorkloadGroupsStateAccessor.getWorkloadGroupState("wg-1").getTotalThrottled()); + + // Releasing the shared permit frees the shared slot so a subsequent request is admitted again. + shared.close(); + assertNotNull(acquireThrottlePermitSync(workloadGroupService, "wg-1", null)); + } + public void testShouldSBPHandle() { SearchTask task = createMockTaskWithResourceStats(SearchTask.class, 100, 200, 0, 12); WorkloadGroupState workloadGroupState = new WorkloadGroupState(); diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java new file mode 100644 index 0000000000000..792b94a58f750 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTests.java @@ -0,0 +1,282 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.cluster.ClusterChangedEvent; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodeRole; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.lease.Releasable; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.util.Collections; +import java.util.Set; +import java.util.concurrent.atomic.AtomicReference; + +import org.mockito.Mockito; + +import static org.mockito.Mockito.when; + +public class WorkloadGroupSharedThrottleServiceTests extends OpenSearchTestCase { + + private ClusterService clusterService; + private ThreadPool threadPool; + private TransportService transportService; + private DiscoveryNode localNode; + + @Override + public void setUp() throws Exception { + super.setUp(); + localNode = new DiscoveryNode( + "local", + "local", + buildNewFakeTransportAddress(), + Collections.emptyMap(), + Set.of(DiscoveryNodeRole.DATA_ROLE), + org.opensearch.Version.CURRENT + ); + clusterService = Mockito.mock(ClusterService.class); + threadPool = Mockito.mock(ThreadPool.class); + transportService = Mockito.mock(TransportService.class); + + ClusterState state = Mockito.mock(ClusterState.class); + DiscoveryNodes singleDataNode = DiscoveryNodes.builder().add(localNode).localNodeId("local").build(); + when(state.nodes()).thenReturn(singleDataNode); + when(clusterService.state()).thenReturn(state); + when(clusterService.localNode()).thenReturn(localNode); + } + + private WorkloadGroupSharedThrottleService newService() { + // Single data node => this node owns every bucket => acquire uses the local short-circuit (no real network). + WorkloadGroupSharedThrottleService service = new WorkloadGroupSharedThrottleService(clusterService, threadPool, transportService); + // The ring is empty until a cluster-state change populates it (matches production: the constructor does not + // read cluster state). Deliver a nodesChanged event with the current nodes. + deliverNodesChanged(service, clusterService.state().nodes()); + return service; + } + + private void deliverNodesChanged(WorkloadGroupSharedThrottleService service, DiscoveryNodes nodes) { + ClusterState previous = Mockito.mock(ClusterState.class); + when(previous.nodes()).thenReturn(DiscoveryNodes.EMPTY_NODES); + ClusterState current = Mockito.mock(ClusterState.class); + when(current.nodes()).thenReturn(nodes); + ClusterChangedEvent event = new ClusterChangedEvent("test", current, previous); + service.clusterChanged(event); + } + + private static Releasable awaitGrant(WorkloadGroupSharedThrottleService service, String bucket, int limit) { + AtomicReference permit = new AtomicReference<>(); + AtomicReference failure = new AtomicReference<>(); + service.acquireAsync(bucket, limit, ActionListener.wrap(permit::set, failure::set)); + if (failure.get() != null) { + throw failure.get() instanceof RuntimeException re ? re : new RuntimeException(failure.get()); + } + return permit.get(); + } + + public void testRingPopulatesWhenNodeSetUnchangedVsPreviousState() { + // Regression for the single-node no-op: the coordinator seeds the initial applied state already containing the + // local node, so the first real clusterChanged has previous.nodes() == current.nodes() and nodesChanged() is + // false. The ring must still populate (we compare eligible owner sets, not the delta), otherwise shared + // throttling silently does nothing on a single-node cluster. + WorkloadGroupSharedThrottleService service = new WorkloadGroupSharedThrottleService(clusterService, threadPool, transportService); + DiscoveryNodes nodes = DiscoveryNodes.builder().add(localNode).localNodeId("local").build(); + ClusterState previous = Mockito.mock(ClusterState.class); + when(previous.nodes()).thenReturn(nodes); // same node set as current -> nodesChanged() == false + ClusterState current = Mockito.mock(ClusterState.class); + when(current.nodes()).thenReturn(nodes); + ClusterChangedEvent event = new ClusterChangedEvent("test", current, previous); + assertFalse("precondition: this event reports no node change", event.nodesChanged()); + + service.clusterChanged(event); + + // shared_limit=1 must now actually enforce (grant then reject), not fail open. + assertNotNull(awaitGrant(service, "b", 1)); + expectThrows(OpenSearchRejectedExecutionException.class, () -> awaitGrant(service, "b", 1)); + } + + public void testSameIdRestartRebuildsRingWithFreshNode() { + // A node can restart keeping its persistent id but with a new ephemeral id/address. Comparing only persistent + // ids would treat this as "no change" and keep the stale DiscoveryNode in the ring, so nodeConnected(stale) + // would fail open forever. The ring must rebuild and hold the FRESH node object. + WorkloadGroupSharedThrottleService service = new WorkloadGroupSharedThrottleService(clusterService, threadPool, transportService); + + DiscoveryNode first = new DiscoveryNode( + "n1", + "n1", + buildNewFakeTransportAddress(), + Collections.emptyMap(), + Set.of(DiscoveryNodeRole.DATA_ROLE), + org.opensearch.Version.CURRENT + ); + deliverNodesChanged(service, DiscoveryNodes.builder().add(first).localNodeId("n1").build()); + assertSame("ring should hold the first node instance", first, service.ring().ownerFor("b").orElseThrow()); + + // Same persistent id "n1", but a brand-new DiscoveryNode instance (fresh ephemeral id + transport address). + DiscoveryNode restarted = new DiscoveryNode( + "n1", + "n1", + buildNewFakeTransportAddress(), + Collections.emptyMap(), + Set.of(DiscoveryNodeRole.DATA_ROLE), + org.opensearch.Version.CURRENT + ); + assertNotEquals("restarted node must not equal the old instance", first, restarted); + deliverNodesChanged(service, DiscoveryNodes.builder().add(restarted).localNodeId("n1").build()); + + assertSame("ring must have rebuilt with the restarted node instance", restarted, service.ring().ownerFor("b").orElseThrow()); + } + + public void testLocalOwnerGrantsThenDeniesAtLimit() { + WorkloadGroupSharedThrottleService service = newService(); + Releasable p1 = awaitGrant(service, "b", 1); + assertNotNull(p1); + expectThrows(OpenSearchRejectedExecutionException.class, () -> awaitGrant(service, "b", 1)); + // release frees the shared slot + p1.close(); + assertNotNull(awaitGrant(service, "b", 1)); + } + + public void testDoubleCloseReleasesOnce() { + WorkloadGroupSharedThrottleService service = newService(); + Releasable p = awaitGrant(service, "b", 2); + assertNotNull(awaitGrant(service, "b", 2)); // second slot + p.close(); + p.close(); // must not double-release + // one slot still held, so only one more grant is available + assertNotNull(awaitGrant(service, "b", 2)); + expectThrows(OpenSearchRejectedExecutionException.class, () -> awaitGrant(service, "b", 2)); + } + + public void testEmptyRingFailsOpen() { + // Cluster with no eligible data node (coordinating/manager-only) => empty ring => fail open (null permit). + DiscoveryNode managerOnly = new DiscoveryNode( + "m", + "m", + buildNewFakeTransportAddress(), + Collections.emptyMap(), + Set.of(DiscoveryNodeRole.CLUSTER_MANAGER_ROLE), + org.opensearch.Version.CURRENT + ); + ClusterState state = Mockito.mock(ClusterState.class); + when(state.nodes()).thenReturn(DiscoveryNodes.builder().add(managerOnly).localNodeId("m").build()); + when(clusterService.state()).thenReturn(state); + when(clusterService.localNode()).thenReturn(managerOnly); + + WorkloadGroupSharedThrottleService service = newService(); + AtomicReference permit = new AtomicReference<>(); + AtomicReference called = new AtomicReference<>(false); + service.acquireAsync("b", 1, ActionListener.wrap(p -> { + called.set(true); + permit.set(p); + }, e -> fail("must not fail: " + e))); + assertTrue("listener must be invoked inline", called.get()); + assertNull("fail-open yields a null permit (admit, untracked)", permit.get()); + } + + public void testDisconnectedOwnerFailsOpenWithoutSendingRequest() { + // Ring owner is a remote data node (not the local node), so this is NOT the local short-circuit path. + DiscoveryNode remote = new DiscoveryNode( + "remote", + "remote", + buildNewFakeTransportAddress(), + Collections.emptyMap(), + Set.of(DiscoveryNodeRole.DATA_ROLE), + org.opensearch.Version.CURRENT + ); + ClusterState state = Mockito.mock(ClusterState.class); + // Only the remote node is a data node, so every bucket hashes to it; local is a coordinating-only node. + when(state.nodes()).thenReturn(DiscoveryNodes.builder().add(remote).localNodeId("local").build()); + when(clusterService.state()).thenReturn(state); + when(clusterService.localNode()).thenReturn(localNode); + // The owner is known-disconnected. + when(transportService.nodeConnected(remote)).thenReturn(false); + + WorkloadGroupSharedThrottleService service = newService(); + AtomicReference permit = new AtomicReference<>(); + AtomicReference called = new AtomicReference<>(false); + service.acquireAsync("b", 1, ActionListener.wrap(p -> { + called.set(true); + permit.set(p); + }, e -> fail("must not fail: " + e))); + + // The listener completed inline with a null permit. This proves the disconnected pre-check fired: on the + // transport path the listener is only invoked from a sendRequest callback, but transportService is a plain + // mock whose sendRequest is a no-op that never calls back — so an inline null response can only come from the + // nodeConnected()==false fast-fail branch, not from a request sent into the dead socket. + assertTrue("listener must be invoked inline (no RTT to a dead node)", called.get()); + assertNull("disconnected owner -> fail open with a null permit", permit.get()); + } + + public void testHandleAcquireIsOwnerSideAdmission() { + WorkloadGroupSharedThrottleService service = newService(); + // Directly exercise the owner-side handler (what a remote coordinator's RPC would hit). + assertTrue( + service.handleAcquire( + new WorkloadGroupSharedThrottleService.AcquireRequest("b", 1, "lease-1", WorkloadGroupSharedThrottleService.LEASE_TTL_NANOS) + ).granted + ); + assertFalse( + service.handleAcquire( + new WorkloadGroupSharedThrottleService.AcquireRequest("b", 1, "lease-2", WorkloadGroupSharedThrottleService.LEASE_TTL_NANOS) + ).granted + ); + assertEquals(1, service.tracker().inFlight("b")); + } + + public void testAcquireRequestSerializationRoundTrip() throws Exception { + WorkloadGroupSharedThrottleService.AcquireRequest original = new WorkloadGroupSharedThrottleService.AcquireRequest( + "grp1:username:alice", + 42, + "lease-xyz", + 123_456_789L + ); + WorkloadGroupSharedThrottleService.AcquireRequest copy = copyWriteable( + original, + writableRegistry(), + WorkloadGroupSharedThrottleService.AcquireRequest::new + ); + assertEquals(original.bucketKey, copy.bucketKey); + assertEquals(original.sharedLimit, copy.sharedLimit); + assertEquals(original.leaseId, copy.leaseId); + assertEquals(original.ttlNanos, copy.ttlNanos); + } + + public void testAcquireResponseSerializationRoundTrip() throws Exception { + for (boolean granted : new boolean[] { true, false }) { + WorkloadGroupSharedThrottleService.AcquireResponse original = new WorkloadGroupSharedThrottleService.AcquireResponse(granted); + WorkloadGroupSharedThrottleService.AcquireResponse copy = copyWriteable( + original, + writableRegistry(), + WorkloadGroupSharedThrottleService.AcquireResponse::new + ); + assertEquals(granted, copy.granted); + } + } + + public void testReleaseRequestSerializationRoundTrip() throws Exception { + WorkloadGroupSharedThrottleService.ReleaseRequest original = new WorkloadGroupSharedThrottleService.ReleaseRequest( + "grp1:group", + "lease-abc" + ); + WorkloadGroupSharedThrottleService.ReleaseRequest copy = copyWriteable( + original, + writableRegistry(), + WorkloadGroupSharedThrottleService.ReleaseRequest::new + ); + assertEquals(original.bucketKey, copy.bucketKey); + assertEquals(original.leaseId, copy.leaseId); + } +} diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTransportTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTransportTests.java new file mode 100644 index 0000000000000..0fec6160e6b65 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSharedThrottleServiceTransportTests.java @@ -0,0 +1,227 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.Version; +import org.opensearch.cluster.ClusterChangedEvent; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.lease.Releasable; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.telemetry.tracing.noop.NoopTracer; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.transport.MockTransportService; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.ConnectTransportException; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Exercises the cross-node (remote owner) acquire/release path of {@link WorkloadGroupSharedThrottleService} + * over a real two-node {@link MockTransportService} harness. Mockito cannot be used to stub the acquire round-trip + * because {@code TransportService#sendRequest(node, action, request, options, handler)} is {@code final}; instead we + * stand up two genuine transports so the {@code internal:wlm/throttle/shared/acquire} and {@code .../release} + * RPCs travel over the wire and hit the owner-side handlers for real. + */ +public class WorkloadGroupSharedThrottleServiceTransportTests extends OpenSearchTestCase { + + private ThreadPool threadPool; + private MockTransportService coordinatorTransport; + private MockTransportService ownerTransport; + + private ClusterService coordinatorClusterService; + private ClusterService ownerClusterService; + + private WorkloadGroupSharedThrottleService coordinatorService; + private WorkloadGroupSharedThrottleService ownerService; + + private DiscoveryNode coordinatorNode; + private DiscoveryNode ownerNode; + + // A bucket key whose ring owner is the remote (owner) node, forcing the coordinator down the transport path. + private String remoteKey; + + @Override + public void setUp() throws Exception { + super.setUp(); + threadPool = new TestThreadPool(getClass().getName()); + + // Two real transports. Default node roles include DATA at Version.CURRENT (>= MIN_OWNER_VERSION), so both + // nodes are eligible ring owners. + coordinatorTransport = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool, NoopTracer.INSTANCE); + ownerTransport = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool, NoopTracer.INSTANCE); + coordinatorTransport.start(); + coordinatorTransport.acceptIncomingRequests(); + ownerTransport.start(); + ownerTransport.acceptIncomingRequests(); + // The coordinator must be able to reach the owner for the acquire/release RPCs. + coordinatorTransport.connectToNode(ownerTransport.getLocalDiscoNode()); + + coordinatorNode = coordinatorTransport.getLocalDiscoNode(); + ownerNode = ownerTransport.getLocalDiscoNode(); + + // Each service sees itself as the local node. + coordinatorClusterService = mock(ClusterService.class); + when(coordinatorClusterService.localNode()).thenReturn(coordinatorNode); + ownerClusterService = mock(ClusterService.class); + when(ownerClusterService.localNode()).thenReturn(ownerNode); + + // Build one service per transport so BOTH register their acquire/release handlers on their own transport. + coordinatorService = new WorkloadGroupSharedThrottleService(coordinatorClusterService, threadPool, coordinatorTransport); + ownerService = new WorkloadGroupSharedThrottleService(ownerClusterService, threadPool, ownerTransport); + + // Drive an identical membership (both nodes as data nodes) into both services so they build the same ring and + // agree on a single deterministic owner per bucket. + DiscoveryNodes bothNodes = DiscoveryNodes.builder() + .add(coordinatorNode) + .add(ownerNode) + .localNodeId(coordinatorNode.getId()) + .build(); + deliverNodes(coordinatorService, bothNodes); + deliverNodes(ownerService, bothNodes); + + // Find a bucket key that the coordinator's ring maps to the REMOTE owner node, so acquireAsync takes the + // transport path rather than the local short-circuit. Guard the whole harness against a mis-built ring. + for (int i = 0; i < 10000 && remoteKey == null; i++) { + String candidate = "bucket-" + i; + DiscoveryNode owner = coordinatorService.ring().ownerFor(candidate).orElse(null); + if (owner != null && owner.getId().equals(ownerNode.getId())) { + remoteKey = candidate; + } + } + assertNotNull("could not find a bucket owned by the remote owner node", remoteKey); + assertEquals( + "chosen bucket must be owned by the remote node (transport path), not the coordinator", + ownerNode.getId(), + coordinatorService.ring().ownerFor(remoteKey).orElseThrow().getId() + ); + } + + @Override + public void tearDown() throws Exception { + coordinatorTransport.close(); + ownerTransport.close(); + ThreadPool.terminate(threadPool, 10, TimeUnit.SECONDS); + super.tearDown(); + } + + private void deliverNodes(WorkloadGroupSharedThrottleService service, DiscoveryNodes nodes) { + ClusterState previous = mock(ClusterState.class); + when(previous.nodes()).thenReturn(DiscoveryNodes.EMPTY_NODES); + ClusterState current = mock(ClusterState.class); + when(current.nodes()).thenReturn(nodes); + service.clusterChanged(new ClusterChangedEvent("test", current, previous)); + } + + /** + * Happy path across nodes: the coordinator asks the remote owner for a permit, the owner grants it over the wire, + * and closing the returned {@link Releasable} sends the fire-and-forget RELEASE RPC that drains the owner's tracker. + */ + public void testRemoteAcquireGrantedThenReleaseRoundTrips() throws Exception { + CapturingListener listener = new CapturingListener(); + coordinatorService.acquireAsync(remoteKey, 1, listener); + listener.await(); + + assertNull("granted acquire must not fail", listener.failure.get()); + assertTrue("listener must have been notified via onResponse", listener.responded.get()); + Releasable permit = listener.response.get(); + assertNotNull("remote owner granted a permit, so a non-null Releasable is expected", permit); + // The grant was recorded on the OWNER node's tracker (the acquire handler ran there before responding). + assertEquals("owner tracker must hold exactly one in-flight lease", 1, ownerService.tracker().inFlight(remoteKey)); + + // Closing the permit fires the RELEASE RPC (fire-and-forget), which the owner applies asynchronously. + permit.close(); + assertBusy(() -> assertEquals("release RPC must drain the owner's in-flight count", 0, ownerService.tracker().inFlight(remoteKey))); + } + + /** + * When the remote owner is already at its shared limit for the bucket, the owner denies the acquire over the wire + * and the coordinator surfaces an {@link OpenSearchRejectedExecutionException} via {@code onFailure} (a 429). + */ + public void testRemoteAcquireDeniedReturns429() throws Exception { + // Pre-fill the owner's tracker to the limit directly so the next acquire is denied at the source. + assertTrue(ownerService.tracker().tryAcquire(remoteKey, 1, "pre", WorkloadGroupSharedThrottleService.LEASE_TTL_NANOS)); + assertEquals(1, ownerService.tracker().inFlight(remoteKey)); + + CapturingListener listener = new CapturingListener(); + coordinatorService.acquireAsync(remoteKey, 1, listener); + listener.await(); + + assertFalse("denied acquire must not invoke onResponse", listener.responded.get()); + assertNotNull("denied acquire must invoke onFailure", listener.failure.get()); + assertTrue( + "denial must be an OpenSearchRejectedExecutionException (429), was: " + listener.failure.get(), + listener.failure.get() instanceof OpenSearchRejectedExecutionException + ); + } + + /** + * If the acquire RPC to the owner fails at send time, the coordinator must FAIL OPEN — admit the request with a + * {@code null} permit via {@code onResponse}, never {@code onFailure}. We keep the owner connected (so the + * in-memory {@code nodeConnected} pre-check passes) and inject a send behavior that throws, forcing the failure + * through {@code handleException} rather than the disconnected pre-check. + */ + public void testFailOpenWhenOwnerTransportUnavailable() throws Exception { + assertTrue( + "precondition: owner must be connected so we exercise the send path, not the pre-check", + coordinatorTransport.nodeConnected(ownerNode) + ); + coordinatorTransport.addSendBehavior(ownerTransport, (connection, requestId, action, request, options) -> { + if (WorkloadGroupSharedThrottleService.ACQUIRE_ACTION_NAME.equals(action)) { + // Simulate the send blowing up; TransportService routes this to the handler's handleException. + throw new ConnectTransportException(connection.getNode(), "simulated acquire send failure"); + } + connection.sendRequest(requestId, action, request, options); + }); + + CapturingListener listener = new CapturingListener(); + coordinatorService.acquireAsync(remoteKey, 1, listener); + listener.await(); + + assertTrue("fail-open must invoke onResponse, not onFailure", listener.responded.get()); + assertNull("fail-open on transport error must yield a null permit (admit, untracked)", listener.response.get()); + assertNull("fail-open must never propagate a failure to the listener", listener.failure.get()); + } + + /** Captures the outcome of an async acquire, distinguishing onResponse(null) from an unfired listener via a latch. */ + private static final class CapturingListener implements ActionListener { + final CountDownLatch latch = new CountDownLatch(1); + final AtomicReference response = new AtomicReference<>(); + final AtomicReference failure = new AtomicReference<>(); + final AtomicBoolean responded = new AtomicBoolean(false); + + @Override + public void onResponse(Releasable releasable) { + response.set(releasable); + responded.set(true); + latch.countDown(); + } + + @Override + public void onFailure(Exception e) { + failure.set(e); + latch.countDown(); + } + + void await() throws InterruptedException { + assertTrue("listener was not invoked within the timeout", latch.await(10, TimeUnit.SECONDS)); + } + } +} diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java deleted file mode 100644 index 475760834af5a..0000000000000 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupThrottleTrackerTests.java +++ /dev/null @@ -1,81 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.wlm; - -import org.opensearch.common.lease.Releasable; -import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; -import org.opensearch.test.OpenSearchTestCase; - -public class WorkloadGroupThrottleTrackerTests extends OpenSearchTestCase { - - public void testAcquireUnderLimitSucceeds() { - WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); - Releasable p1 = tracker.acquire("bucket", 2); - Releasable p2 = tracker.acquire("bucket", 2); - assertEquals(2, tracker.inFlight("bucket")); - p1.close(); - p2.close(); - } - - public void testAcquireAtLimitRejects() { - WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); - tracker.acquire("bucket", 1); - OpenSearchRejectedExecutionException e = expectThrows( - OpenSearchRejectedExecutionException.class, - () -> tracker.acquire("bucket", 1) - ); - assertTrue(e.getMessage().contains("throttle limit reached")); - assertFalse(e.getMessage().contains("bucket")); - // rejected acquire must not leave the count inflated - assertEquals(1, tracker.inFlight("bucket")); - } - - public void testReleaseFreesAPermit() { - WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); - Releasable p = tracker.acquire("bucket", 1); - // at the limit - expectThrows(OpenSearchRejectedExecutionException.class, () -> tracker.acquire("bucket", 1)); - p.close(); - // permit freed -> a fresh acquire now succeeds - Releasable p2 = tracker.acquire("bucket", 1); - assertEquals(1, tracker.inFlight("bucket")); - p2.close(); - } - - public void testDrainToZeroRemovesBucketThenReacquire() { - WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); - Releasable p = tracker.acquire("bucket", 5); - assertEquals(1, tracker.inFlight("bucket")); - p.close(); - assertEquals(0, tracker.inFlight("bucket")); - // re-acquiring after the bucket drained (and was removed) works and starts from 1 - Releasable p2 = tracker.acquire("bucket", 5); - assertEquals(1, tracker.inFlight("bucket")); - p2.close(); - } - - public void testReleaseIsIdempotent() { - WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); - Releasable p = tracker.acquire("bucket", 5); - p.close(); - p.close(); // double close must not decrement twice - assertEquals(0, tracker.inFlight("bucket")); - } - - public void testBucketsAreIndependent() { - WorkloadGroupThrottleTracker tracker = new WorkloadGroupThrottleTracker(); - tracker.acquire("a", 1); - // "a" is full but "b" is a separate bucket - expectThrows(OpenSearchRejectedExecutionException.class, () -> tracker.acquire("a", 1)); - Releasable pb = tracker.acquire("b", 1); - assertEquals(1, tracker.inFlight("a")); - assertEquals(1, tracker.inFlight("b")); - pb.close(); - } -}