From 38e23e57a7b63d91e5873523b103c9f3f8f74109 Mon Sep 17 00:00:00 2001 From: David Zane Date: Mon, 18 May 2026 10:48:20 -0700 Subject: [PATCH 1/4] feat(wlm): add search.max_buckets to workload group settings Signed-off-by: David Zane --- .../rest/WorkloadManagementRestIT.java | 69 +++++- .../main/java/org/opensearch/node/Node.java | 9 +- .../org/opensearch/search/SearchService.java | 8 +- .../MultiBucketConsumerService.java | 61 ++++- .../wlm/WorkloadGroupSearchSettings.java | 16 ++ .../MultiBucketConsumerServiceTests.java | 212 ++++++++++++++++++ .../snapshots/SnapshotResiliencyTests.java | 3 +- .../wlm/WorkloadGroupSearchSettingsTests.java | 32 +++ .../java/org/opensearch/node/MockNode.java | 10 +- .../opensearch/search/MockSearchService.java | 7 +- 10 files changed, 412 insertions(+), 15 deletions(-) create mode 100644 server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java 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..1dc927e686c82 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 @@ -158,7 +158,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 +174,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 +183,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 +197,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 +319,67 @@ 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 testSearchSettingsMergeSemantics() throws Exception { 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..7266f870ea7f7 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,9 @@ public SearchService( this.multiBucketConsumerService = new MultiBucketConsumerService( clusterService, settings, - circuitBreakerService.getBreaker(CircuitBreaker.REQUEST) + circuitBreakerService.getBreaker(CircuitBreaker.REQUEST), + threadPool, + 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..bc0274abd7cc7 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java +++ b/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java @@ -31,10 +31,14 @@ 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; import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.common.breaker.CircuitBreaker; import org.opensearch.core.common.breaker.CircuitBreakingException; import org.opensearch.core.common.io.stream.StreamInput; @@ -42,6 +46,10 @@ import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.search.aggregations.bucket.BucketsAggregator; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.WorkloadGroupSearchSettings; +import org.opensearch.wlm.WorkloadGroupService; +import org.opensearch.wlm.WorkloadGroupTask; import java.io.IOException; import java.util.concurrent.atomic.LongAdder; @@ -56,6 +64,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 +76,21 @@ public class MultiBucketConsumerService { ); private final CircuitBreaker breaker; + private final ThreadPool threadPool; + private final WorkloadGroupService workloadGroupService; private volatile int maxBucket; - public MultiBucketConsumerService(ClusterService clusterService, Settings settings, CircuitBreaker breaker) { + public MultiBucketConsumerService( + ClusterService clusterService, + Settings settings, + CircuitBreaker breaker, + ThreadPool threadPool, + WorkloadGroupService workloadGroupService + ) { this.breaker = breaker; + this.threadPool = threadPool; + this.workloadGroupService = workloadGroupService; this.maxBucket = MAX_BUCKET_SETTING.get(settings); clusterService.getClusterSettings().addSettingsUpdateConsumer(MAX_BUCKET_SETTING, this::setMaxBucket); } @@ -79,6 +99,43 @@ 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 (threadPool == null || workloadGroupService == null) { + return maxBucket; + } + ThreadContext threadContext = threadPool.getThreadContext(); + if (threadContext == null) { + return maxBucket; + } + String workloadGroupId = threadContext.getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER); + if (workloadGroupId == null) { + return maxBucket; + } + WorkloadGroup workloadGroup = workloadGroupService.getWorkloadGroupById(workloadGroupId); + 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 +273,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/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..76481b9963ad6 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java @@ -0,0 +1,212 @@ +/* + * 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, + threadPool, + 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, + threadPool, + workloadGroupService + ); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "missing-id"); + when(workloadGroupService.getWorkloadGroupById("missing-id")).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, + threadPool, + 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.getWorkloadGroupById("wg-id")).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, + threadPool, + 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.getWorkloadGroupById("wg-id")).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, + threadPool, + 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.getWorkloadGroupById("wg-id")).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, + threadPool, + 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, + threadPool, + workloadGroupService + ); + threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); + when(workloadGroupService.getWorkloadGroupById("wg-id")).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, + threadPool, + 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.getWorkloadGroupById("wg-id")).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, + threadPool, + 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.getWorkloadGroupById("wg-id")).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, + threadPool, + 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.getWorkloadGroupById("wg-id")).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/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 ); } From 7d869756871f630af4c05218d0da153da536da3d Mon Sep 17 00:00:00 2001 From: David Zane Date: Thu, 21 May 2026 14:14:21 -0600 Subject: [PATCH 2/4] Review comment fixes Signed-off-by: David Zane --- .../rest/WorkloadManagementRestIT.java | 53 +++++++++++++++++++ .../MultiBucketConsumerService.java | 14 +---- .../opensearch/wlm/WorkloadGroupService.java | 13 +++++ ...WorkloadGroupRequestOperationListener.java | 12 ++--- .../MultiBucketConsumerServiceTests.java | 16 +++--- .../wlm/WorkloadGroupServiceTests.java | 35 ++++++++++++ ...oadGroupRequestOperationListenerTests.java | 32 +++++------ 7 files changed, 131 insertions(+), 44 deletions(-) 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 1dc927e686c82..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 { @@ -382,6 +384,48 @@ public void testSearchMaxBucketsCreateAndUpdate() throws Exception { 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 { // Create with multiple settings String createJson = """ @@ -457,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/search/aggregations/MultiBucketConsumerService.java b/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java index bc0274abd7cc7..2911e2db80141 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java +++ b/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java @@ -38,7 +38,6 @@ import org.opensearch.common.annotation.PublicApi; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; -import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.common.breaker.CircuitBreaker; import org.opensearch.core.common.breaker.CircuitBreakingException; import org.opensearch.core.common.io.stream.StreamInput; @@ -49,7 +48,6 @@ import org.opensearch.threadpool.ThreadPool; import org.opensearch.wlm.WorkloadGroupSearchSettings; import org.opensearch.wlm.WorkloadGroupService; -import org.opensearch.wlm.WorkloadGroupTask; import java.io.IOException; import java.util.concurrent.atomic.LongAdder; @@ -110,18 +108,10 @@ private void setMaxBucket(int maxBucket) { */ int resolveMaxBuckets() { try { - if (threadPool == null || workloadGroupService == null) { + if (workloadGroupService == null) { return maxBucket; } - ThreadContext threadContext = threadPool.getThreadContext(); - if (threadContext == null) { - return maxBucket; - } - String workloadGroupId = threadContext.getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER); - if (workloadGroupId == null) { - return maxBucket; - } - WorkloadGroup workloadGroup = workloadGroupService.getWorkloadGroupById(workloadGroupId); + WorkloadGroup workloadGroup = workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext()); if (workloadGroup == null) { return maxBucket; } diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java index 5c87ebbb13fea..a7b6d3729f8af 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java @@ -17,6 +17,7 @@ import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.lifecycle.AbstractLifecycleComponent; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import org.opensearch.monitor.jvm.JvmStats; import org.opensearch.monitor.process.ProcessProbe; @@ -334,6 +335,18 @@ public WorkloadGroup getWorkloadGroupById(String workloadGroupId) { return clusterService.state().metadata().workloadGroups().get(workloadGroupId); } + /** + * Resolves the workload group attached to the calling thread context, or null if there is + * no workload group ID header set or the referenced group does not exist. + */ + public WorkloadGroup resolveFromThreadContext(ThreadContext threadContext) { + String workloadGroupId = threadContext.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..728f26a4ae743 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.resolveFromThreadContext(threadPool.getThreadContext()); + 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 index 76481b9963ad6..976f7ea6e5c38 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java @@ -72,7 +72,7 @@ public void testResolveFallsBackWhenWorkloadGroupNotFound() { workloadGroupService ); threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "missing-id"); - when(workloadGroupService.getWorkloadGroupById("missing-id")).thenReturn(null); + when(workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(null); assertEquals(100, svc.resolveMaxBuckets()); } @@ -86,7 +86,7 @@ public void testResolveFallsBackWhenWorkloadGroupHasNoMaxBucketsSetting() { ); 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.getWorkloadGroupById("wg-id")).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); assertEquals(100, svc.resolveMaxBuckets()); } @@ -100,7 +100,7 @@ public void testResolveUsesWlmValueWhenSet() { ); 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.getWorkloadGroupById("wg-id")).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); assertEquals(42, svc.resolveMaxBuckets()); } @@ -114,7 +114,7 @@ public void testResolveWlmValueOverridesClusterEvenWhenLarger() { ); 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.getWorkloadGroupById("wg-id")).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); assertEquals(10000, svc.resolveMaxBuckets()); } @@ -139,7 +139,7 @@ public void testResolveFallsBackWhenLookupThrows() { workloadGroupService ); threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); - when(workloadGroupService.getWorkloadGroupById("wg-id")).thenThrow(new RuntimeException("boom")); + when(workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext())).thenThrow(new RuntimeException("boom")); assertEquals(100, svc.resolveMaxBuckets()); } @@ -158,7 +158,7 @@ public void testOverrideRequestValuesIsIrrelevantWhenFalse() { 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.getWorkloadGroupById("wg-id")).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); assertEquals(42, svc.resolveMaxBuckets()); } @@ -177,7 +177,7 @@ public void testOverrideRequestValuesIsIrrelevantWhenTrue() { 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.getWorkloadGroupById("wg-id")).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); assertEquals(42, svc.resolveMaxBuckets()); } @@ -191,7 +191,7 @@ public void testCreateUsesResolvedLimit() { ); 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.getWorkloadGroupById("wg-id")).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); MultiBucketConsumerService.MultiBucketConsumer consumer = svc.create(); assertEquals(7, consumer.getLimit()); diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java index 989d390853cc2..df52df60b49f9 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,39 @@ public void testOnTaskCompleted() { mockThreadPool.shutdown(); } + public void testResolveFromThreadContextReturnsNullWhenHeaderMissing() { + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + assertNull(workloadGroupService.resolveFromThreadContext(threadContext)); + } + + public void testResolveFromThreadContextReturnsGroupWhenPresent() { + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + threadContext.putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-1"); + 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.resolveFromThreadContext(threadContext)); + } + + public void testResolveFromThreadContextReturnsNullWhenGroupMissing() { + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + threadContext.putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "missing-id"); + 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.resolveFromThreadContext(threadContext)); + } + 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..acbc0d33b93da 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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); From 06eb76b78f0f113ef52328f8b5a0a5463caa0e26 Mon Sep 17 00:00:00 2001 From: David Zane Date: Fri, 22 May 2026 10:49:01 -0600 Subject: [PATCH 3/4] Use threadpool in WorkloadGroupRequestOperationListener Signed-off-by: David Zane --- .../org/opensearch/search/SearchService.java | 1 - .../MultiBucketConsumerService.java | 6 +--- .../opensearch/wlm/WorkloadGroupService.java | 5 ++- ...WorkloadGroupRequestOperationListener.java | 2 +- .../MultiBucketConsumerServiceTests.java | 26 +++++---------- .../wlm/WorkloadGroupServiceTests.java | 9 ++++-- ...oadGroupRequestOperationListenerTests.java | 32 +++++++++---------- 7 files changed, 34 insertions(+), 47 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java index 7266f870ea7f7..bc83688983a70 100644 --- a/server/src/main/java/org/opensearch/search/SearchService.java +++ b/server/src/main/java/org/opensearch/search/SearchService.java @@ -546,7 +546,6 @@ public SearchService( clusterService, settings, circuitBreakerService.getBreaker(CircuitBreaker.REQUEST), - threadPool, workloadGroupService ); this.indexSearcherExecutor = indexSearcherExecutor; 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 2911e2db80141..3eb5986db21c0 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java +++ b/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java @@ -45,7 +45,6 @@ import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.search.aggregations.bucket.BucketsAggregator; -import org.opensearch.threadpool.ThreadPool; import org.opensearch.wlm.WorkloadGroupSearchSettings; import org.opensearch.wlm.WorkloadGroupService; @@ -74,7 +73,6 @@ public class MultiBucketConsumerService { ); private final CircuitBreaker breaker; - private final ThreadPool threadPool; private final WorkloadGroupService workloadGroupService; private volatile int maxBucket; @@ -83,11 +81,9 @@ public MultiBucketConsumerService( ClusterService clusterService, Settings settings, CircuitBreaker breaker, - ThreadPool threadPool, WorkloadGroupService workloadGroupService ) { this.breaker = breaker; - this.threadPool = threadPool; this.workloadGroupService = workloadGroupService; this.maxBucket = MAX_BUCKET_SETTING.get(settings); clusterService.getClusterSettings().addSettingsUpdateConsumer(MAX_BUCKET_SETTING, this::setMaxBucket); @@ -111,7 +107,7 @@ int resolveMaxBuckets() { if (workloadGroupService == null) { return maxBucket; } - WorkloadGroup workloadGroup = workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext()); + WorkloadGroup workloadGroup = workloadGroupService.resolveFromThreadContext(); if (workloadGroup == null) { return maxBucket; } diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java index a7b6d3729f8af..e83251433e0d7 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java @@ -17,7 +17,6 @@ import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.lifecycle.AbstractLifecycleComponent; -import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import org.opensearch.monitor.jvm.JvmStats; import org.opensearch.monitor.process.ProcessProbe; @@ -339,8 +338,8 @@ public WorkloadGroup getWorkloadGroupById(String workloadGroupId) { * Resolves the workload group attached to the calling thread context, or null if there is * no workload group ID header set or the referenced group does not exist. */ - public WorkloadGroup resolveFromThreadContext(ThreadContext threadContext) { - String workloadGroupId = threadContext.getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER); + public WorkloadGroup resolveFromThreadContext() { + String workloadGroupId = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER); if (workloadGroupId == null) { return null; } 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 728f26a4ae743..fe90d611152a1 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,7 @@ public WorkloadGroupRequestOperationListener(WorkloadGroupService workloadGroupS protected void onRequestStart(SearchRequestContext searchRequestContext) { final String workloadGroupId = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER); workloadGroupService.rejectIfNeeded(workloadGroupId); - WorkloadGroup workloadGroup = workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext()); + WorkloadGroup workloadGroup = workloadGroupService.resolveFromThreadContext(); applyWorkloadGroupSearchSettings(workloadGroup, searchRequestContext.getRequest()); } diff --git a/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java b/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java index 976f7ea6e5c38..7f04b54a88750 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java @@ -56,7 +56,6 @@ public void testResolveFallsBackToClusterDefaultWhenNoHeader() { clusterService, Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), breaker, - threadPool, workloadGroupService ); // No WORKLOAD_GROUP_ID_HEADER in thread context @@ -68,11 +67,10 @@ public void testResolveFallsBackWhenWorkloadGroupNotFound() { clusterService, Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), breaker, - threadPool, workloadGroupService ); threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "missing-id"); - when(workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(null); + when(workloadGroupService.resolveFromThreadContext()).thenReturn(null); assertEquals(100, svc.resolveMaxBuckets()); } @@ -81,12 +79,11 @@ public void testResolveFallsBackWhenWorkloadGroupHasNoMaxBucketsSetting() { clusterService, Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), breaker, - threadPool, 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.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).thenReturn(wg); assertEquals(100, svc.resolveMaxBuckets()); } @@ -95,12 +92,11 @@ public void testResolveUsesWlmValueWhenSet() { clusterService, Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), breaker, - threadPool, 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.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).thenReturn(wg); assertEquals(42, svc.resolveMaxBuckets()); } @@ -109,12 +105,11 @@ public void testResolveWlmValueOverridesClusterEvenWhenLarger() { clusterService, Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), breaker, - threadPool, 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.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).thenReturn(wg); assertEquals(10000, svc.resolveMaxBuckets()); } @@ -123,7 +118,6 @@ public void testResolveFallsBackWhenWorkloadGroupServiceNull() { clusterService, Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), breaker, - threadPool, null ); threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); @@ -135,11 +129,10 @@ public void testResolveFallsBackWhenLookupThrows() { clusterService, Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), breaker, - threadPool, workloadGroupService ); threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); - when(workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext())).thenThrow(new RuntimeException("boom")); + when(workloadGroupService.resolveFromThreadContext()).thenThrow(new RuntimeException("boom")); assertEquals(100, svc.resolveMaxBuckets()); } @@ -150,7 +143,6 @@ public void testOverrideRequestValuesIsIrrelevantWhenFalse() { clusterService, Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), breaker, - threadPool, workloadGroupService ); WorkloadGroup wg = createWorkloadGroup( @@ -158,7 +150,7 @@ public void testOverrideRequestValuesIsIrrelevantWhenFalse() { 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.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).thenReturn(wg); assertEquals(42, svc.resolveMaxBuckets()); } @@ -169,7 +161,6 @@ public void testOverrideRequestValuesIsIrrelevantWhenTrue() { clusterService, Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), breaker, - threadPool, workloadGroupService ); WorkloadGroup wg = createWorkloadGroup( @@ -177,7 +168,7 @@ public void testOverrideRequestValuesIsIrrelevantWhenTrue() { 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.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).thenReturn(wg); assertEquals(42, svc.resolveMaxBuckets()); } @@ -186,12 +177,11 @@ public void testCreateUsesResolvedLimit() { clusterService, Settings.builder().put(MultiBucketConsumerService.MAX_BUCKET_SETTING.getKey(), 100).build(), breaker, - threadPool, 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.resolveFromThreadContext(threadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).thenReturn(wg); MultiBucketConsumerService.MultiBucketConsumer consumer = svc.create(); assertEquals(7, consumer.getLimit()); diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java index df52df60b49f9..b17effbb86e9b 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java @@ -445,12 +445,14 @@ public void testOnTaskCompleted() { public void testResolveFromThreadContextReturnsNullWhenHeaderMissing() { ThreadContext threadContext = new ThreadContext(Settings.EMPTY); - assertNull(workloadGroupService.resolveFromThreadContext(threadContext)); + when(mockThreadPool.getThreadContext()).thenReturn(threadContext); + assertNull(workloadGroupService.resolveFromThreadContext()); } public void testResolveFromThreadContextReturnsGroupWhenPresent() { 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", @@ -462,18 +464,19 @@ public void testResolveFromThreadContextReturnsGroupWhenPresent() { when(mockClusterService.state()).thenReturn(clusterState); when(clusterState.metadata()).thenReturn(metadata); when(metadata.workloadGroups()).thenReturn(Map.of("wg-1", wg)); - assertSame(wg, workloadGroupService.resolveFromThreadContext(threadContext)); + assertSame(wg, workloadGroupService.resolveFromThreadContext()); } public void testResolveFromThreadContextReturnsNullWhenGroupMissing() { 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.resolveFromThreadContext(threadContext)); + assertNull(workloadGroupService.resolveFromThreadContext()); } public void testShouldSBPHandle() { 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 acbc0d33b93da..b0f02e32db9fb 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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(null); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).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.resolveFromThreadContext(testThreadPool.getThreadContext())).thenReturn(wg); + when(workloadGroupService.resolveFromThreadContext()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); From b3de5eb1b28eef03daf19ab612f8b41f6ddc9291 Mon Sep 17 00:00:00 2001 From: David Zane Date: Fri, 22 May 2026 13:05:54 -0600 Subject: [PATCH 4/4] Change method to getCurrentWorkloadGroup Signed-off-by: David Zane --- .../MultiBucketConsumerService.java | 2 +- .../opensearch/wlm/WorkloadGroupService.java | 7 ++-- ...WorkloadGroupRequestOperationListener.java | 2 +- .../MultiBucketConsumerServiceTests.java | 16 +++++----- .../wlm/WorkloadGroupServiceTests.java | 12 +++---- ...oadGroupRequestOperationListenerTests.java | 32 +++++++++---------- 6 files changed, 36 insertions(+), 35 deletions(-) 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 3eb5986db21c0..34a9135d7e736 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java +++ b/server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java @@ -107,7 +107,7 @@ int resolveMaxBuckets() { if (workloadGroupService == null) { return maxBucket; } - WorkloadGroup workloadGroup = workloadGroupService.resolveFromThreadContext(); + WorkloadGroup workloadGroup = workloadGroupService.getCurrentWorkloadGroup(); if (workloadGroup == null) { return maxBucket; } diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java index e83251433e0d7..64c398e1d5e90 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java @@ -335,10 +335,11 @@ public WorkloadGroup getWorkloadGroupById(String workloadGroupId) { } /** - * Resolves the workload group attached to the calling thread context, or null if there is - * no workload group ID header set or the referenced group does not exist. + * 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 resolveFromThreadContext() { + public WorkloadGroup getCurrentWorkloadGroup() { String workloadGroupId = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER); if (workloadGroupId == null) { return null; 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 fe90d611152a1..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,7 @@ public WorkloadGroupRequestOperationListener(WorkloadGroupService workloadGroupS protected void onRequestStart(SearchRequestContext searchRequestContext) { final String workloadGroupId = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER); workloadGroupService.rejectIfNeeded(workloadGroupId); - WorkloadGroup workloadGroup = workloadGroupService.resolveFromThreadContext(); + WorkloadGroup workloadGroup = workloadGroupService.getCurrentWorkloadGroup(); applyWorkloadGroupSearchSettings(workloadGroup, searchRequestContext.getRequest()); } diff --git a/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java b/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java index 7f04b54a88750..9a7baf37a2492 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/MultiBucketConsumerServiceTests.java @@ -70,7 +70,7 @@ public void testResolveFallsBackWhenWorkloadGroupNotFound() { workloadGroupService ); threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "missing-id"); - when(workloadGroupService.resolveFromThreadContext()).thenReturn(null); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(null); assertEquals(100, svc.resolveMaxBuckets()); } @@ -83,7 +83,7 @@ public void testResolveFallsBackWhenWorkloadGroupHasNoMaxBucketsSetting() { ); 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.resolveFromThreadContext()).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); assertEquals(100, svc.resolveMaxBuckets()); } @@ -96,7 +96,7 @@ public void testResolveUsesWlmValueWhenSet() { ); 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.resolveFromThreadContext()).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); assertEquals(42, svc.resolveMaxBuckets()); } @@ -109,7 +109,7 @@ public void testResolveWlmValueOverridesClusterEvenWhenLarger() { ); 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.resolveFromThreadContext()).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); assertEquals(10000, svc.resolveMaxBuckets()); } @@ -132,7 +132,7 @@ public void testResolveFallsBackWhenLookupThrows() { workloadGroupService ); threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-id"); - when(workloadGroupService.resolveFromThreadContext()).thenThrow(new RuntimeException("boom")); + when(workloadGroupService.getCurrentWorkloadGroup()).thenThrow(new RuntimeException("boom")); assertEquals(100, svc.resolveMaxBuckets()); } @@ -150,7 +150,7 @@ public void testOverrideRequestValuesIsIrrelevantWhenFalse() { 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.resolveFromThreadContext()).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); assertEquals(42, svc.resolveMaxBuckets()); } @@ -168,7 +168,7 @@ public void testOverrideRequestValuesIsIrrelevantWhenTrue() { 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.resolveFromThreadContext()).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); assertEquals(42, svc.resolveMaxBuckets()); } @@ -181,7 +181,7 @@ public void testCreateUsesResolvedLimit() { ); 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.resolveFromThreadContext()).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); MultiBucketConsumerService.MultiBucketConsumer consumer = svc.create(); assertEquals(7, consumer.getLimit()); diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java index b17effbb86e9b..8da689ab2bc89 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupServiceTests.java @@ -443,13 +443,13 @@ public void testOnTaskCompleted() { mockThreadPool.shutdown(); } - public void testResolveFromThreadContextReturnsNullWhenHeaderMissing() { + public void testGetCurrentWorkloadGroupReturnsNullWhenHeaderMissing() { ThreadContext threadContext = new ThreadContext(Settings.EMPTY); when(mockThreadPool.getThreadContext()).thenReturn(threadContext); - assertNull(workloadGroupService.resolveFromThreadContext()); + assertNull(workloadGroupService.getCurrentWorkloadGroup()); } - public void testResolveFromThreadContextReturnsGroupWhenPresent() { + public void testGetCurrentWorkloadGroupReturnsGroupWhenPresent() { ThreadContext threadContext = new ThreadContext(Settings.EMPTY); threadContext.putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "wg-1"); when(mockThreadPool.getThreadContext()).thenReturn(threadContext); @@ -464,10 +464,10 @@ public void testResolveFromThreadContextReturnsGroupWhenPresent() { when(mockClusterService.state()).thenReturn(clusterState); when(clusterState.metadata()).thenReturn(metadata); when(metadata.workloadGroups()).thenReturn(Map.of("wg-1", wg)); - assertSame(wg, workloadGroupService.resolveFromThreadContext()); + assertSame(wg, workloadGroupService.getCurrentWorkloadGroup()); } - public void testResolveFromThreadContextReturnsNullWhenGroupMissing() { + public void testGetCurrentWorkloadGroupReturnsNullWhenGroupMissing() { ThreadContext threadContext = new ThreadContext(Settings.EMPTY); threadContext.putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "missing-id"); when(mockThreadPool.getThreadContext()).thenReturn(threadContext); @@ -476,7 +476,7 @@ public void testResolveFromThreadContextReturnsNullWhenGroupMissing() { when(mockClusterService.state()).thenReturn(clusterState); when(clusterState.metadata()).thenReturn(metadata); when(metadata.workloadGroups()).thenReturn(Collections.emptyMap()); - assertNull(workloadGroupService.resolveFromThreadContext()); + assertNull(workloadGroupService.getCurrentWorkloadGroup()); } public void testShouldSBPHandle() { 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 b0f02e32db9fb..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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).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.resolveFromThreadContext()).thenReturn(wg); + when(workloadGroupService.getCurrentWorkloadGroup()).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext);