diff --git a/plugins/workload-management/src/javaRestTest/java/org/opensearch/rest/WorkloadManagementRestIT.java b/plugins/workload-management/src/javaRestTest/java/org/opensearch/rest/WorkloadManagementRestIT.java index aa72dbfaacc4b..d9ac4d10c6e47 100644 --- a/plugins/workload-management/src/javaRestTest/java/org/opensearch/rest/WorkloadManagementRestIT.java +++ b/plugins/workload-management/src/javaRestTest/java/org/opensearch/rest/WorkloadManagementRestIT.java @@ -17,6 +17,8 @@ import java.io.IOException; import java.util.Locale; +import java.util.regex.Matcher; +import java.util.regex.Pattern; public class WorkloadManagementRestIT extends OpenSearchRestTestCase { @@ -158,7 +160,8 @@ public void testSearchSettings() throws Exception { "search.default_search_timeout": "30s", "search.cancel_after_time_interval": "1m", "search.max_concurrent_shard_requests": "5", - "search.batched_reduce_size": "512" + "search.batched_reduce_size": "512", + "search.max_buckets": "1000" } }"""; Response response = performOperation("PUT", "_wlm/workload_group", createJson); @@ -173,6 +176,7 @@ public void testSearchSettings() throws Exception { assertTrue(responseBody.contains("\"search.cancel_after_time_interval\":\"1m\"")); assertTrue(responseBody.contains("\"search.max_concurrent_shard_requests\":\"5\"")); assertTrue(responseBody.contains("\"search.batched_reduce_size\":\"512\"")); + assertTrue(responseBody.contains("\"search.max_buckets\":\"1000\"")); // Update search settings String updateJson = """ @@ -181,7 +185,8 @@ public void testSearchSettings() throws Exception { "search.default_search_timeout": "1m", "search.cancel_after_time_interval": "5m", "search.max_concurrent_shard_requests": "10", - "search.batched_reduce_size": "256" + "search.batched_reduce_size": "256", + "search.max_buckets": "500" } }"""; Response updateResponse = performOperation("PUT", "_wlm/workload_group/search_test", updateJson); @@ -194,6 +199,7 @@ public void testSearchSettings() throws Exception { assertTrue(responseBody2.contains("\"search.cancel_after_time_interval\":\"5m\"")); assertTrue(responseBody2.contains("\"search.max_concurrent_shard_requests\":\"10\"")); assertTrue(responseBody2.contains("\"search.batched_reduce_size\":\"256\"")); + assertTrue(responseBody2.contains("\"search.max_buckets\":\"500\"")); performOperation("DELETE", "_wlm/workload_group/search_test", null); } @@ -315,6 +321,109 @@ public void testSearchSettingsInvalidSettingsRejected() throws Exception { String invalidTimeBody = EntityUtils.toString(invalidTimeException.getResponse().getEntity()); assertTrue(invalidTimeBody.contains("search.cancel_after_time_interval")); assertTrue(invalidTimeBody.contains("Invalid value")); + + // Invalid value for max_buckets (must be >= 0) + String invalidMaxBucketsJson = """ + { + "name": "invalid_test", + "resiliency_mode": "enforced", + "resource_limits": {"cpu": 0.3, "memory": 0.3}, + "settings": { + "search.max_buckets": "-1" + } + }"""; + ResponseException invalidMaxBucketsException = expectThrows( + ResponseException.class, + () -> performOperation("PUT", "_wlm/workload_group", invalidMaxBucketsJson) + ); + String invalidMaxBucketsBody = EntityUtils.toString(invalidMaxBucketsException.getResponse().getEntity()); + assertTrue(invalidMaxBucketsBody.contains("search.max_buckets")); + assertTrue(invalidMaxBucketsBody.contains("Invalid value")); + } + + public void testSearchMaxBucketsCreateAndUpdate() throws Exception { + // Create a WLM group with a small max_buckets value + String createJson = """ + { + "name": "max_buckets_test", + "resiliency_mode": "enforced", + "resource_limits": {"cpu": 0.3, "memory": 0.3}, + "settings": { + "search.max_buckets": "100" + } + }"""; + Response response = performOperation("PUT", "_wlm/workload_group", createJson); + assertEquals(200, response.getStatusLine().getStatusCode()); + + Response getResponse = performOperation("GET", "_wlm/workload_group/max_buckets_test", null); + assertTrue(EntityUtils.toString(getResponse.getEntity()).contains("\"search.max_buckets\":\"100\"")); + + // Update to a larger value + String updateJson = """ + {"settings": {"search.max_buckets": "5000"}}"""; + Response updateResponse = performOperation("PUT", "_wlm/workload_group/max_buckets_test", updateJson); + assertEquals(200, updateResponse.getStatusLine().getStatusCode()); + + Response getResponse2 = performOperation("GET", "_wlm/workload_group/max_buckets_test", null); + assertTrue(EntityUtils.toString(getResponse2.getEntity()).contains("\"search.max_buckets\":\"5000\"")); + + // Exercise the request path with an aggregation — confirms the resolver is wired + // through MultiBucketConsumerService without errors. Resolution semantics are + // verified in MultiBucketConsumerServiceTests. + performOperation("PUT", "wlm-buckets-idx", "{\"settings\":{\"number_of_shards\":1,\"number_of_replicas\":0}}"); + performOperation("POST", "wlm-buckets-idx/_doc", "{\"k\":\"v1\"}"); + performOperation("POST", "wlm-buckets-idx/_refresh", null); + + Request searchRequest = new Request("POST", "wlm-buckets-idx/_search"); + searchRequest.setJsonEntity("{\"size\":0,\"aggs\":{\"by_k\":{\"terms\":{\"field\":\"k.keyword\"}}}}"); + searchRequest.setOptions(searchRequest.getOptions().toBuilder().addHeader("X-opaque-id", "wlm=max_buckets_test")); + Response searchResponse = client().performRequest(searchRequest); + assertEquals(200, searchResponse.getStatusLine().getStatusCode()); + + performOperation("DELETE", "wlm-buckets-idx", null); + performOperation("DELETE", "_wlm/workload_group/max_buckets_test", null); + } + + public void testSearchMaxBucketsEnforcedAtRequestPath() throws Exception { + // Cluster default permits the aggregation; WLM cap is smaller and must win. + String wgName = "max_buckets_enforced_test"; + String createJson = String.format(Locale.ROOT, """ + { + "name": "%s", + "resiliency_mode": "enforced", + "resource_limits": {"cpu": 0.3, "memory": 0.3}, + "settings": { + "search.max_buckets": "1" + } + }""", wgName); + Response response = performOperation("PUT", "_wlm/workload_group", createJson); + assertEquals(200, response.getStatusLine().getStatusCode()); + + String wgId = extractWorkloadGroupId(performOperation("GET", "_wlm/workload_group/" + wgName, null)); + + performOperation("PUT", "wlm-buckets-enforce-idx", "{\"settings\":{\"number_of_shards\":1,\"number_of_replicas\":0}}"); + performOperation("POST", "wlm-buckets-enforce-idx/_doc", "{\"k\":\"v1\"}"); + performOperation("POST", "wlm-buckets-enforce-idx/_doc", "{\"k\":\"v2\"}"); + performOperation("POST", "wlm-buckets-enforce-idx/_refresh", null); + + String body = "{\"size\":0,\"aggs\":{\"by_k\":{\"terms\":{\"field\":\"k.keyword\"}}}}"; + + // Same query without the WLM tag passes — cluster default allows >1 bucket. + Request unTagged = new Request("POST", "wlm-buckets-enforce-idx/_search"); + unTagged.setJsonEntity(body); + assertEquals(200, client().performRequest(unTagged).getStatusLine().getStatusCode()); + + // With the workload group attached, the WLM cap of 1 is enforced. + Request tagged = new Request("POST", "wlm-buckets-enforce-idx/_search"); + tagged.setJsonEntity(body); + tagged.setOptions(tagged.getOptions().toBuilder().addHeader("workloadGroupId", wgId)); + ResponseException rejected = expectThrows(ResponseException.class, () -> client().performRequest(tagged)); + String rejectedBody = EntityUtils.toString(rejected.getResponse().getEntity()); + assertTrue("expected too_many_buckets error, got: " + rejectedBody, rejectedBody.contains("too_many_buckets")); + assertTrue("expected limit of 1 in error, got: " + rejectedBody, rejectedBody.contains("\"max_buckets\":1")); + + performOperation("DELETE", "wlm-buckets-enforce-idx", null); + performOperation("DELETE", "_wlm/workload_group/" + wgName, null); } public void testSearchSettingsMergeSemantics() throws Exception { @@ -392,6 +501,15 @@ static String getUpdateJson(String resiliencyMode, double cpu, double memory) { }""", resiliencyMode, cpu, memory); } + private static final Pattern WORKLOAD_GROUP_ID_PATTERN = Pattern.compile("\"_id\"\\s*:\\s*\"([^\"]+)\""); + + private String extractWorkloadGroupId(Response response) throws Exception { + String body = EntityUtils.toString(response.getEntity()); + Matcher m = WORKLOAD_GROUP_ID_PATTERN.matcher(body); + assertTrue("could not find _id in response: " + body, m.find()); + return m.group(1); + } + Response performOperation(String method, String uriPath, String json) throws IOException { Request request = new Request(method, uriPath); if (json != null) { diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index 5844a83128636..a99ed50338b1b 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -1631,7 +1631,8 @@ protected Node(final Environment initialEnvironment, Collection clas searchModule.getIndexSearcherExecutor(threadPool), taskResourceTrackingService, searchModule.getConcurrentSearchRequestDeciderFactories(), - searchModule.getPluginProfileMetricsProviders() + searchModule.getPluginProfileMetricsProviders(), + workloadGroupService ); final List> tasksExecutors = pluginsService.filterPlugins(PersistentTaskPlugin.class) @@ -2389,7 +2390,8 @@ protected SearchService newSearchService( Executor indexSearcherExecutor, TaskResourceTrackingService taskResourceTrackingService, Collection concurrentSearchDeciderFactories, - List pluginProfilers + List pluginProfilers, + WorkloadGroupService workloadGroupService ) { return new SearchService( clusterService, @@ -2404,7 +2406,8 @@ protected SearchService newSearchService( indexSearcherExecutor, taskResourceTrackingService, concurrentSearchDeciderFactories, - pluginProfilers + pluginProfilers, + workloadGroupService ); } diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java index 4a74170f71271..bc83688983a70 100644 --- a/server/src/main/java/org/opensearch/search/SearchService.java +++ b/server/src/main/java/org/opensearch/search/SearchService.java @@ -158,6 +158,7 @@ import org.opensearch.threadpool.ThreadPool; import org.opensearch.threadpool.ThreadPool.Names; import org.opensearch.transport.TransportRequest; +import org.opensearch.wlm.WorkloadGroupService; import java.io.IOException; import java.util.ArrayList; @@ -529,7 +530,8 @@ public SearchService( Executor indexSearcherExecutor, TaskResourceTrackingService taskResourceTrackingService, Collection concurrentSearchDeciderFactories, - List pluginProfilers + List pluginProfilers, + WorkloadGroupService workloadGroupService ) { Settings settings = clusterService.getSettings(); this.threadPool = threadPool; @@ -543,7 +545,8 @@ public SearchService( this.multiBucketConsumerService = new MultiBucketConsumerService( clusterService, settings, - circuitBreakerService.getBreaker(CircuitBreaker.REQUEST) + circuitBreakerService.getBreaker(CircuitBreaker.REQUEST), + workloadGroupService ); this.indexSearcherExecutor = indexSearcherExecutor; this.taskResourceTrackingService = taskResourceTrackingService; diff --git a/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java b/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java index 35186422fceaa..34a9135d7e736 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java +++ b/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java @@ -31,6 +31,9 @@ package org.opensearch.search.aggregations; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.annotation.PublicApi; import org.opensearch.common.settings.Setting; @@ -42,6 +45,8 @@ import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.search.aggregations.bucket.BucketsAggregator; +import org.opensearch.wlm.WorkloadGroupSearchSettings; +import org.opensearch.wlm.WorkloadGroupService; import java.io.IOException; import java.util.concurrent.atomic.LongAdder; @@ -56,6 +61,8 @@ * @opensearch.internal */ public class MultiBucketConsumerService { + private static final Logger logger = LogManager.getLogger(MultiBucketConsumerService.class); + public static final int DEFAULT_MAX_BUCKETS = 65535; public static final Setting MAX_BUCKET_SETTING = Setting.intSetting( "search.max_buckets", @@ -66,11 +73,18 @@ public class MultiBucketConsumerService { ); private final CircuitBreaker breaker; + private final WorkloadGroupService workloadGroupService; private volatile int maxBucket; - public MultiBucketConsumerService(ClusterService clusterService, Settings settings, CircuitBreaker breaker) { + public MultiBucketConsumerService( + ClusterService clusterService, + Settings settings, + CircuitBreaker breaker, + WorkloadGroupService workloadGroupService + ) { this.breaker = breaker; + this.workloadGroupService = workloadGroupService; this.maxBucket = MAX_BUCKET_SETTING.get(settings); clusterService.getClusterSettings().addSettingsUpdateConsumer(MAX_BUCKET_SETTING, this::setMaxBucket); } @@ -79,6 +93,35 @@ private void setMaxBucket(int maxBucket) { this.maxBucket = maxBucket; } + /** + * Resolves the effective max-buckets limit for the current request by consulting the + * workload group (if any) attached to the calling thread context. If the request has no + * workload group, the group is unknown, or the group does not define + * {@code search.max_buckets}, the cluster-level default is returned. + *

+ * The WLM-set value, when present, always wins — {@code override_request_values} is not + * relevant because {@code search.max_buckets} is not a per-request parameter. + */ + int resolveMaxBuckets() { + try { + if (workloadGroupService == null) { + return maxBucket; + } + WorkloadGroup workloadGroup = workloadGroupService.getCurrentWorkloadGroup(); + if (workloadGroup == null) { + return maxBucket; + } + Settings wlmSettings = workloadGroup.getSettings(); + if (wlmSettings == null || wlmSettings.hasValue(WorkloadGroupSearchSettings.WLM_MAX_BUCKETS.getKey()) == false) { + return maxBucket; + } + return WorkloadGroupSearchSettings.WLM_MAX_BUCKETS.get(wlmSettings); + } catch (Exception e) { + logger.warn("Failed to resolve workload group [search.max_buckets]; falling back to cluster default", e); + return maxBucket; + } + } + /** * Thrown when there are too many buckets * @@ -216,6 +259,6 @@ public int getLimit() { } public MultiBucketConsumer create() { - return new MultiBucketConsumer(maxBucket, breaker); + return new MultiBucketConsumer(resolveMaxBuckets(), breaker); } } diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupSearchSettings.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSearchSettings.java index 4ad771c5d3152..a9e97c41979e6 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupSearchSettings.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSearchSettings.java @@ -12,6 +12,7 @@ import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; +import org.opensearch.search.aggregations.MultiBucketConsumerService; import java.util.Map; @@ -62,6 +63,19 @@ private WorkloadGroupSearchSettings() { */ public static final Setting WLM_BATCHED_REDUCE_SIZE = Setting.intSetting("search.batched_reduce_size", 512, 2); + /** + * The WLM max buckets setting. Caps the number of aggregation buckets a request in this + * workload group may produce. Mirrors the cluster-level {@code search.max_buckets}; when + * set on a workload group, this value always takes precedence over the cluster default for + * requests assigned to the group. {@code override_request_values} is not relevant for this + * setting because {@code search.max_buckets} is not a per-request parameter. + */ + public static final Setting WLM_MAX_BUCKETS = Setting.intSetting( + "search.max_buckets", + MultiBucketConsumerService.DEFAULT_MAX_BUCKETS, + 0 + ); + /** * Controls whether WLM search settings should override values explicitly set in the * search request query parameters. When {@code false} (default), WLM settings are only @@ -82,6 +96,8 @@ private WorkloadGroupSearchSettings() { WLM_MAX_CONCURRENT_SHARD_REQUESTS, "search.batched_reduce_size", WLM_BATCHED_REDUCE_SIZE, + "search.max_buckets", + WLM_MAX_BUCKETS, "override_request_values", WLM_OVERRIDE_REQUEST_VALUES ); diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java index 5c87ebbb13fea..64c398e1d5e90 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java @@ -334,6 +334,19 @@ public WorkloadGroup getWorkloadGroupById(String workloadGroupId) { return clusterService.state().metadata().workloadGroups().get(workloadGroupId); } + /** + * Returns the workload group attached to the calling thread context, or null if the current + * request does not map to a workload group (no header set, or the referenced group does not + * exist). + */ + public WorkloadGroup getCurrentWorkloadGroup() { + String workloadGroupId = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER); + if (workloadGroupId == null) { + return null; + } + return getWorkloadGroupById(workloadGroupId); + } + public Set getDeletedWorkloadGroups() { return deletedWorkloadGroups; } diff --git a/server/src/main/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListener.java b/server/src/main/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListener.java index 73f29b09776e1..2ee503c52246b 100644 --- a/server/src/main/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListener.java +++ b/server/src/main/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListener.java @@ -44,7 +44,8 @@ public WorkloadGroupRequestOperationListener(WorkloadGroupService workloadGroupS protected void onRequestStart(SearchRequestContext searchRequestContext) { final String workloadGroupId = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER); workloadGroupService.rejectIfNeeded(workloadGroupId); - applyWorkloadGroupSearchSettings(workloadGroupId, searchRequestContext.getRequest()); + WorkloadGroup workloadGroup = workloadGroupService.getCurrentWorkloadGroup(); + applyWorkloadGroupSearchSettings(workloadGroup, searchRequestContext.getRequest()); } @Override @@ -61,15 +62,10 @@ protected void onRequestFailure(SearchPhaseContext context, SearchRequestContext * applied when the request does not already have an explicit value set. * When {@code true}, WLM settings always take precedence over request-level values. * - * @param workloadGroupId the workload group identifier from thread context + * @param workloadGroup the resolved workload group, or null if none is attached to the request * @param searchRequest the search request to modify */ - private void applyWorkloadGroupSearchSettings(String workloadGroupId, SearchRequest searchRequest) { - if (workloadGroupId == null) { - return; - } - - WorkloadGroup workloadGroup = workloadGroupService.getWorkloadGroupById(workloadGroupId); + private void applyWorkloadGroupSearchSettings(WorkloadGroup workloadGroup, SearchRequest searchRequest) { if (workloadGroup == null) { return; } diff --git a/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java b/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java new file mode 100644 index 0000000000000..9a7baf37a2492 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java @@ -0,0 +1,202 @@ +/* + * 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.search.aggregations; + +import org.opensearch.cluster.metadata.WorkloadGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.common.breaker.CircuitBreaker; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.MutableWorkloadGroupFragment; +import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadGroupService; +import org.opensearch.wlm.WorkloadGroupTask; + +import java.util.Map; +import java.util.Set; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class MultiBucketConsumerServiceTests extends OpenSearchTestCase { + + private ThreadPool threadPool; + private ClusterService clusterService; + private CircuitBreaker breaker; + private WorkloadGroupService workloadGroupService; + + @Override + public void setUp() throws Exception { + super.setUp(); + threadPool = new TestThreadPool(getTestName()); + breaker = mock(CircuitBreaker.class); + clusterService = mock(ClusterService.class); + ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, Set.of(MultiBucketConsumerService.MAX_BUCKET_SETTING)); + when(clusterService.getClusterSettings()).thenReturn(clusterSettings); + workloadGroupService = mock(WorkloadGroupService.class); + } + + @Override + public void tearDown() throws Exception { + terminate(threadPool); + super.tearDown(); + } + + public void testResolveFallsBackToClusterDefaultWhenNoHeader() { + MultiBucketConsumerService svc = new MultiBucketConsumerService( + clusterService, + Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), + breaker, + workloadGroupService + ); + // No WORKLOAD_GROUP_ID_HEADER in thread context + assertEquals(100, svc.resolveMaxBuckets()); + } + + public void testResolveFallsBackWhenWorkloadGroupNotFound() { + MultiBucketConsumerService svc = new MultiBucketConsumerService( + clusterService, + Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), + breaker, + workloadGroupService + ); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "missing-id"); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(null); + assertEquals(100, svc.resolveMaxBuckets()); + } + + public void testResolveFallsBackWhenWorkloadGroupHasNoMaxBucketsSetting() { + MultiBucketConsumerService svc = new MultiBucketConsumerService( + clusterService, + Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), + breaker, + workloadGroupService + ); + WorkloadGroup wg = createWorkloadGroup("wg-id", Settings.builder().put("search.default_search_timeout", "30s").build()); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); + assertEquals(100, svc.resolveMaxBuckets()); + } + + public void testResolveUsesWlmValueWhenSet() { + MultiBucketConsumerService svc = new MultiBucketConsumerService( + clusterService, + Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), + breaker, + workloadGroupService + ); + WorkloadGroup wg = createWorkloadGroup("wg-id", Settings.builder().put("search.max_buckets", "42").build()); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); + assertEquals(42, svc.resolveMaxBuckets()); + } + + public void testResolveWlmValueOverridesClusterEvenWhenLarger() { + MultiBucketConsumerService svc = new MultiBucketConsumerService( + clusterService, + Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), + breaker, + workloadGroupService + ); + WorkloadGroup wg = createWorkloadGroup("wg-id", Settings.builder().put("search.max_buckets", "10000").build()); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); + assertEquals(10000, svc.resolveMaxBuckets()); + } + + public void testResolveFallsBackWhenWorkloadGroupServiceNull() { + MultiBucketConsumerService svc = new MultiBucketConsumerService( + clusterService, + Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), + breaker, + null + ); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); + assertEquals(100, svc.resolveMaxBuckets()); + } + + public void testResolveFallsBackWhenLookupThrows() { + MultiBucketConsumerService svc = new MultiBucketConsumerService( + clusterService, + Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), + breaker, + workloadGroupService + ); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); + when(workloadGroupService.getCurrentWorkloadGroup()).thenThrow(new RuntimeException("boom")); + assertEquals(100, svc.resolveMaxBuckets()); + } + + public void testOverrideRequestValuesIsIrrelevantWhenFalse() { + // override_request_values=false should NOT prevent the WLM max_buckets from applying, + // because max_buckets is not a per-request value. The WLM-set value always wins. + MultiBucketConsumerService svc = new MultiBucketConsumerService( + clusterService, + Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), + breaker, + workloadGroupService + ); + WorkloadGroup wg = createWorkloadGroup( + "wg-id", + Settings.builder().put("search.max_buckets", "42").put("override_request_values", "false").build() + ); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); + assertEquals(42, svc.resolveMaxBuckets()); + } + + public void testOverrideRequestValuesIsIrrelevantWhenTrue() { + // override_request_values=true behaves identically to false for max_buckets — both + // result in the WLM value being applied. This pins the no-op contract. + MultiBucketConsumerService svc = new MultiBucketConsumerService( + clusterService, + Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), + breaker, + workloadGroupService + ); + WorkloadGroup wg = createWorkloadGroup( + "wg-id", + Settings.builder().put("search.max_buckets", "42").put("override_request_values", "true").build() + ); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); + assertEquals(42, svc.resolveMaxBuckets()); + } + + public void testCreateUsesResolvedLimit() { + MultiBucketConsumerService svc = new MultiBucketConsumerService( + clusterService, + Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), + breaker, + workloadGroupService + ); + WorkloadGroup wg = createWorkloadGroup("wg-id", Settings.builder().put("search.max_buckets", "7").build()); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); + + MultiBucketConsumerService.MultiBucketConsumer consumer = svc.create(); + assertEquals(7, consumer.getLimit()); + } + + private WorkloadGroup createWorkloadGroup(String id, Settings searchSettings) { + return new WorkloadGroup( + "test-name", + id, + new MutableWorkloadGroupFragment( + MutableWorkloadGroupFragment.ResiliencyMode.SOFT, + Map.of(ResourceType.MEMORY, 0.5), + searchSettings + ), + System.currentTimeMillis() + ); + } +} diff --git a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java index ae6c5ba87a8ce..5d90890c05c14 100644 --- a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java +++ b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java @@ -2370,7 +2370,8 @@ public void onFailure(final Exception e) { null, new TaskResourceTrackingService(settings, clusterSettings, threadPool), Collections.emptyList(), - Collections.emptyList() + Collections.emptyList(), + null ); SearchPhaseController searchPhaseController = new SearchPhaseController( writableRegistry(), diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSearchSettingsTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSearchSettingsTests.java index fe9efc2615766..aef9e3070dec0 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSearchSettingsTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSearchSettingsTests.java @@ -33,6 +33,11 @@ public void testWlmBatchedReduceSizeSettingExists() { assertEquals("search.batched_reduce_size", WorkloadGroupSearchSettings.WLM_BATCHED_REDUCE_SIZE.getKey()); } + public void testWlmMaxBucketsSettingExists() { + assertNotNull(WorkloadGroupSearchSettings.WLM_MAX_BUCKETS); + assertEquals("search.max_buckets", WorkloadGroupSearchSettings.WLM_MAX_BUCKETS.getKey()); + } + public void testWlmOverrideRequestValuesSettingExists() { assertNotNull(WorkloadGroupSearchSettings.WLM_OVERRIDE_REQUEST_VALUES); assertEquals("override_request_values", WorkloadGroupSearchSettings.WLM_OVERRIDE_REQUEST_VALUES.getKey()); @@ -104,6 +109,32 @@ public void testValidateBatchedReduceSizeInvalid() { assertTrue(exception.getMessage().contains("search.batched_reduce_size")); } + public void testValidateMaxBuckets() { + Settings settings = Settings.builder().put("search.max_buckets", "0").build(); + WorkloadGroupSearchSettings.validate(settings); + + settings = Settings.builder().put("search.max_buckets", "65535").build(); + WorkloadGroupSearchSettings.validate(settings); + + settings = Settings.builder().put("search.max_buckets", "1000000").build(); + WorkloadGroupSearchSettings.validate(settings); + } + + public void testValidateMaxBucketsInvalid() { + Settings settings = Settings.builder().put("search.max_buckets", "-1").build(); + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> WorkloadGroupSearchSettings.validate(settings) + ); + assertTrue(exception.getMessage().contains("Invalid value")); + assertTrue(exception.getMessage().contains("search.max_buckets")); + + Settings settings2 = Settings.builder().put("search.max_buckets", "abc").build(); + exception = expectThrows(IllegalArgumentException.class, () -> WorkloadGroupSearchSettings.validate(settings2)); + assertTrue(exception.getMessage().contains("Invalid value")); + assertTrue(exception.getMessage().contains("search.max_buckets")); + } + public void testValidateOverrideRequestValues() { Settings settings = Settings.builder().put("override_request_values", "true").build(); WorkloadGroupSearchSettings.validate(settings); @@ -118,6 +149,7 @@ public void testValidateMultipleSettings() { .put("search.cancel_after_time_interval", "1m") .put("search.max_concurrent_shard_requests", "5") .put("search.batched_reduce_size", "256") + .put("search.max_buckets", "1000") .put("override_request_values", "true") .build(); WorkloadGroupSearchSettings.validate(settings); diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java index 989d390853cc2..8da689ab2bc89 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java @@ -14,7 +14,9 @@ import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import org.opensearch.search.backpressure.trackers.NodeDuressTrackers; import org.opensearch.tasks.Task; @@ -441,6 +443,42 @@ public void testOnTaskCompleted() { mockThreadPool.shutdown(); } + public void testGetCurrentWorkloadGroupReturnsNullWhenHeaderMissing() { + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + when(mockThreadPool.getThreadContext()).thenReturn(threadContext); + assertNull(workloadGroupService.getCurrentWorkloadGroup()); + } + + public void testGetCurrentWorkloadGroupReturnsGroupWhenPresent() { + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + threadContext.putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-1"); + when(mockThreadPool.getThreadContext()).thenReturn(threadContext); + WorkloadGroup wg = new WorkloadGroup( + "wg-1-name", + "wg-1", + new MutableWorkloadGroupFragment(MutableWorkloadGroupFragment.ResiliencyMode.SOFT, Map.of(ResourceType.MEMORY, 0.5)), + 1L + ); + ClusterState clusterState = Mockito.mock(ClusterState.class); + Metadata metadata = Mockito.mock(Metadata.class); + when(mockClusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + when(metadata.workloadGroups()).thenReturn(Map.of("wg-1", wg)); + assertSame(wg, workloadGroupService.getCurrentWorkloadGroup()); + } + + public void testGetCurrentWorkloadGroupReturnsNullWhenGroupMissing() { + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + threadContext.putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "missing-id"); + when(mockThreadPool.getThreadContext()).thenReturn(threadContext); + ClusterState clusterState = Mockito.mock(ClusterState.class); + Metadata metadata = Mockito.mock(Metadata.class); + when(mockClusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + when(metadata.workloadGroups()).thenReturn(Collections.emptyMap()); + assertNull(workloadGroupService.getCurrentWorkloadGroup()); + } + 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/listeners/WorkloadGroupRequestOperationListenerTests.java b/server/src/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java index 5594c77fd5d58..31071d7acf1c3 100644 --- a/server/src/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java +++ b/server/src/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java @@ -298,7 +298,7 @@ public void testApplySearchSettings_NullWorkloadGroupId() { public void testApplySearchSettings_WorkloadGroupNotFound() { testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "non-existent-id"); - when(workloadGroupService.getWorkloadGroupById("non-existent-id")).thenReturn(null); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(null); sut.onRequestStart(mockSearchRequestContext); @@ -310,7 +310,7 @@ public void testApplySearchSettings_EmptySearchSettings() { String wgId = "test-wg"; WorkloadGroup wg = createWorkloadGroup(wgId, Settings.EMPTY); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -324,7 +324,7 @@ public void testApplySearchSettings_Timeout_WlmAppliedWhenNull() { String wgId = "test-wg"; WorkloadGroup wg = createWorkloadGroup(wgId, Settings.builder().put("search.default_search_timeout", "1m").build()); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -337,7 +337,7 @@ public void testApplySearchSettings_Timeout_RequestAlreadySet() { String wgId = "test-wg"; WorkloadGroup wg = createWorkloadGroup(wgId, Settings.builder().put("search.default_search_timeout", "10s").build()); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -350,7 +350,7 @@ public void testApplySearchSettings_Timeout_NullSource() { String wgId = "test-wg"; WorkloadGroup wg = createWorkloadGroup(wgId, Settings.builder().put("search.default_search_timeout", "30s").build()); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -363,7 +363,7 @@ public void testApplySearchSettings_CancelAfterTimeInterval_WlmAppliedWhenNull() String wgId = "test-wg"; WorkloadGroup wg = createWorkloadGroup(wgId, Settings.builder().put("search.cancel_after_time_interval", "30s").build()); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -376,7 +376,7 @@ public void testApplySearchSettings_CancelAfterTimeInterval_RequestAlreadySet() String wgId = "test-wg"; WorkloadGroup wg = createWorkloadGroup(wgId, Settings.builder().put("search.cancel_after_time_interval", "30s").build()); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -390,7 +390,7 @@ public void testApplySearchSettings_MaxConcurrentShardRequests_WlmAppliedWhenDef String wgId = "test-wg"; WorkloadGroup wg = createWorkloadGroup(wgId, Settings.builder().put("search.max_concurrent_shard_requests", "10").build()); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -403,7 +403,7 @@ public void testApplySearchSettings_MaxConcurrentShardRequests_RequestAlreadySet String wgId = "test-wg"; WorkloadGroup wg = createWorkloadGroup(wgId, Settings.builder().put("search.max_concurrent_shard_requests", "5").build()); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -417,7 +417,7 @@ public void testApplySearchSettings_BatchedReduceSize_WlmAppliedWhenDefault() { String wgId = "test-wg"; WorkloadGroup wg = createWorkloadGroup(wgId, Settings.builder().put("search.batched_reduce_size", "100").build()); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -430,7 +430,7 @@ public void testApplySearchSettings_BatchedReduceSize_RequestAlreadySet() { String wgId = "test-wg"; WorkloadGroup wg = createWorkloadGroup(wgId, Settings.builder().put("search.batched_reduce_size", "100").build()); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -456,7 +456,7 @@ public void testApplySearchSettings_OverrideRequestValues_True() { .put("override_request_values", "true") .build() ); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -486,7 +486,7 @@ public void testApplySearchSettings_OverrideRequestValues_False() { .put("override_request_values", "false") .build() ); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -511,7 +511,7 @@ public void testApplySearchSettings_MultipleSettings() { .put("search.batched_reduce_size", "100") .build() ); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -533,7 +533,7 @@ public void testApplySearchSettings_OverrideRequestValues_DefaultsToFalseWhenAbs wgId, Settings.builder().put("search.default_search_timeout", "1m").put("search.cancel_after_time_interval", "2m").build() ); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); @@ -562,7 +562,7 @@ public void testApplySearchSettings_OverrideRequestValues_TrueWithRequestUnset() .put("override_request_values", "true") .build() ); - when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); diff --git a/test/framework/src/main/java/org/opensearch/node/MockNode.java b/test/framework/src/main/java/org/opensearch/node/MockNode.java index 23ef62dcaf02e..99ddae8c89c4a 100644 --- a/test/framework/src/main/java/org/opensearch/node/MockNode.java +++ b/test/framework/src/main/java/org/opensearch/node/MockNode.java @@ -73,6 +73,7 @@ import org.opensearch.transport.TransportInterceptor; import org.opensearch.transport.TransportService; import org.opensearch.transport.client.node.NodeClient; +import org.opensearch.wlm.WorkloadGroupService; import java.nio.file.Path; import java.util.Collection; @@ -176,7 +177,8 @@ protected SearchService newSearchService( Executor indexSearcherExecutor, TaskResourceTrackingService taskResourceTrackingService, Collection concurrentSearchDeciderFactories, - List pluginProfilers + List pluginProfilers, + WorkloadGroupService workloadGroupService ) { if (getPluginsService().filterPlugins(MockSearchService.TestPlugin.class).isEmpty()) { return super.newSearchService( @@ -192,7 +194,8 @@ protected SearchService newSearchService( indexSearcherExecutor, taskResourceTrackingService, concurrentSearchDeciderFactories, - pluginProfilers + pluginProfilers, + workloadGroupService ); } return new MockSearchService( @@ -205,7 +208,8 @@ protected SearchService newSearchService( fetchPhase, circuitBreakerService, indexSearcherExecutor, - taskResourceTrackingService + taskResourceTrackingService, + workloadGroupService ); } diff --git a/test/framework/src/main/java/org/opensearch/search/MockSearchService.java b/test/framework/src/main/java/org/opensearch/search/MockSearchService.java index e3bc166e56d6b..992ab584aaa46 100644 --- a/test/framework/src/main/java/org/opensearch/search/MockSearchService.java +++ b/test/framework/src/main/java/org/opensearch/search/MockSearchService.java @@ -44,6 +44,7 @@ import org.opensearch.search.query.QueryPhase; import org.opensearch.tasks.TaskResourceTrackingService; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.WorkloadGroupService; import java.util.Collections; import java.util.HashMap; @@ -99,7 +100,8 @@ public MockSearchService( FetchPhase fetchPhase, CircuitBreakerService circuitBreakerService, Executor indexSearcherExecutor, - TaskResourceTrackingService taskResourceTrackingService + TaskResourceTrackingService taskResourceTrackingService, + WorkloadGroupService workloadGroupService ) { super( clusterService, @@ -114,7 +116,8 @@ public MockSearchService( indexSearcherExecutor, taskResourceTrackingService, Collections.emptyList(), - Collections.emptyList() + Collections.emptyList(), + workloadGroupService ); }