feat(wlm): add search.max_buckets to workload group settings

plugins/workload-management/src/javaRestTest/java/org/opensearch/rest/WorkloadManagementRestIT.java

@@ -17,6 +17,8 @@
 
 import java.io.IOException;
 import java.util.Locale;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 public class WorkloadManagementRestIT extends OpenSearchRestTestCase {
 
@@ -158,7 +160,8 @@ public void testSearchSettings() throws Exception {
                     "search.default_search_timeout": "30s",
                     "search.cancel_after_time_interval": "1m",
                     "search.max_concurrent_shard_requests": "5",
-                    "search.batched_reduce_size": "512"
+                    "search.batched_reduce_size": "512",
+                    "search.max_buckets": "1000"
                 }
             }""";
         Response response = performOperation("PUT", "_wlm/workload_group", createJson);
@@ -173,6 +176,7 @@ public void testSearchSettings() throws Exception {
         assertTrue(responseBody.contains("\"search.cancel_after_time_interval\":\"1m\""));
         assertTrue(responseBody.contains("\"search.max_concurrent_shard_requests\":\"5\""));
         assertTrue(responseBody.contains("\"search.batched_reduce_size\":\"512\""));
+        assertTrue(responseBody.contains("\"search.max_buckets\":\"1000\""));
 
         // Update search settings
         String updateJson = """
@@ -181,7 +185,8 @@ public void testSearchSettings() throws Exception {
                     "search.default_search_timeout": "1m",
                     "search.cancel_after_time_interval": "5m",
                     "search.max_concurrent_shard_requests": "10",
-                    "search.batched_reduce_size": "256"
+                    "search.batched_reduce_size": "256",
+                    "search.max_buckets": "500"
                 }
             }""";
         Response updateResponse = performOperation("PUT", "_wlm/workload_group/search_test", updateJson);
@@ -194,6 +199,7 @@ public void testSearchSettings() throws Exception {
         assertTrue(responseBody2.contains("\"search.cancel_after_time_interval\":\"5m\""));
         assertTrue(responseBody2.contains("\"search.max_concurrent_shard_requests\":\"10\""));
         assertTrue(responseBody2.contains("\"search.batched_reduce_size\":\"256\""));
+        assertTrue(responseBody2.contains("\"search.max_buckets\":\"500\""));
 
         performOperation("DELETE", "_wlm/workload_group/search_test", null);
     }
@@ -315,6 +321,109 @@ public void testSearchSettingsInvalidSettingsRejected() throws Exception {
         String invalidTimeBody = EntityUtils.toString(invalidTimeException.getResponse().getEntity());
         assertTrue(invalidTimeBody.contains("search.cancel_after_time_interval"));
         assertTrue(invalidTimeBody.contains("Invalid value"));
+
+        // Invalid value for max_buckets (must be >= 0)
+        String invalidMaxBucketsJson = """
+            {
+                "name": "invalid_test",
+                "resiliency_mode": "enforced",
+                "resource_limits": {"cpu": 0.3, "memory": 0.3},
+                "settings": {
+                    "search.max_buckets": "-1"
+                }
+            }""";
+        ResponseException invalidMaxBucketsException = expectThrows(
+            ResponseException.class,
+            () -> performOperation("PUT", "_wlm/workload_group", invalidMaxBucketsJson)
+        );
+        String invalidMaxBucketsBody = EntityUtils.toString(invalidMaxBucketsException.getResponse().getEntity());
+        assertTrue(invalidMaxBucketsBody.contains("search.max_buckets"));
+        assertTrue(invalidMaxBucketsBody.contains("Invalid value"));
+    }
+
+    public void testSearchMaxBucketsCreateAndUpdate() throws Exception {
+        // Create a WLM group with a small max_buckets value
+        String createJson = """
+            {
+                "name": "max_buckets_test",
+                "resiliency_mode": "enforced",
+                "resource_limits": {"cpu": 0.3, "memory": 0.3},
+                "settings": {
+                    "search.max_buckets": "100"
+                }
+            }""";
+        Response response = performOperation("PUT", "_wlm/workload_group", createJson);
+        assertEquals(200, response.getStatusLine().getStatusCode());
+
+        Response getResponse = performOperation("GET", "_wlm/workload_group/max_buckets_test", null);
+        assertTrue(EntityUtils.toString(getResponse.getEntity()).contains("\"search.max_buckets\":\"100\""));
+
+        // Update to a larger value
+        String updateJson = """
+            {"settings": {"search.max_buckets": "5000"}}""";
+        Response updateResponse = performOperation("PUT", "_wlm/workload_group/max_buckets_test", updateJson);
+        assertEquals(200, updateResponse.getStatusLine().getStatusCode());
+
+        Response getResponse2 = performOperation("GET", "_wlm/workload_group/max_buckets_test", null);
+        assertTrue(EntityUtils.toString(getResponse2.getEntity()).contains("\"search.max_buckets\":\"5000\""));
+
+        // Exercise the request path with an aggregation — confirms the resolver is wired
+        // through MultiBucketConsumerService without errors. Resolution semantics are
+        // verified in MultiBucketConsumerServiceTests.
+        performOperation("PUT", "wlm-buckets-idx", "{\"settings\":{\"number_of_shards\":1,\"number_of_replicas\":0}}");
+        performOperation("POST", "wlm-buckets-idx/_doc", "{\"k\":\"v1\"}");
+        performOperation("POST", "wlm-buckets-idx/_refresh", null);
+
+        Request searchRequest = new Request("POST", "wlm-buckets-idx/_search");
+        searchRequest.setJsonEntity("{\"size\":0,\"aggs\":{\"by_k\":{\"terms\":{\"field\":\"k.keyword\"}}}}");
+        searchRequest.setOptions(searchRequest.getOptions().toBuilder().addHeader("X-opaque-id", "wlm=max_buckets_test"));
+        Response searchResponse = client().performRequest(searchRequest);
+        assertEquals(200, searchResponse.getStatusLine().getStatusCode());
+
+        performOperation("DELETE", "wlm-buckets-idx", null);
+        performOperation("DELETE", "_wlm/workload_group/max_buckets_test", null);
+    }
+
+    public void testSearchMaxBucketsEnforcedAtRequestPath() throws Exception {
+        // Cluster default permits the aggregation; WLM cap is smaller and must win.
+        String wgName = "max_buckets_enforced_test";
+        String createJson = String.format(Locale.ROOT, """
+            {
+                "name": "%s",
+                "resiliency_mode": "enforced",
+                "resource_limits": {"cpu": 0.3, "memory": 0.3},
+                "settings": {
+                    "search.max_buckets": "1"
+                }
+            }""", wgName);
+        Response response = performOperation("PUT", "_wlm/workload_group", createJson);
+        assertEquals(200, response.getStatusLine().getStatusCode());
+
+        String wgId = extractWorkloadGroupId(performOperation("GET", "_wlm/workload_group/" + wgName, null));
+
+        performOperation("PUT", "wlm-buckets-enforce-idx", "{\"settings\":{\"number_of_shards\":1,\"number_of_replicas\":0}}");
+        performOperation("POST", "wlm-buckets-enforce-idx/_doc", "{\"k\":\"v1\"}");
+        performOperation("POST", "wlm-buckets-enforce-idx/_doc", "{\"k\":\"v2\"}");
+        performOperation("POST", "wlm-buckets-enforce-idx/_refresh", null);
+
+        String body = "{\"size\":0,\"aggs\":{\"by_k\":{\"terms\":{\"field\":\"k.keyword\"}}}}";
+
+        // Same query without the WLM tag passes — cluster default allows >1 bucket.
+        Request unTagged = new Request("POST", "wlm-buckets-enforce-idx/_search");
+        unTagged.setJsonEntity(body);
+        assertEquals(200, client().performRequest(unTagged).getStatusLine().getStatusCode());
+
+        // With the workload group attached, the WLM cap of 1 is enforced.
+        Request tagged = new Request("POST", "wlm-buckets-enforce-idx/_search");
+        tagged.setJsonEntity(body);
+        tagged.setOptions(tagged.getOptions().toBuilder().addHeader("workloadGroupId", wgId));
+        ResponseException rejected = expectThrows(ResponseException.class, () -> client().performRequest(tagged));
+        String rejectedBody = EntityUtils.toString(rejected.getResponse().getEntity());
+        assertTrue("expected too_many_buckets error, got: " + rejectedBody, rejectedBody.contains("too_many_buckets"));
+        assertTrue("expected limit of 1 in error, got: " + rejectedBody, rejectedBody.contains("\"max_buckets\":1"));
+
+        performOperation("DELETE", "wlm-buckets-enforce-idx", null);
+        performOperation("DELETE", "_wlm/workload_group/" + wgName, null);
     }
 
     public void testSearchSettingsMergeSemantics() throws Exception {
@@ -392,6 +501,15 @@ static String getUpdateJson(String resiliencyMode, double cpu, double memory) {
             }""", resiliencyMode, cpu, memory);
     }
 
+    private static final Pattern WORKLOAD_GROUP_ID_PATTERN = Pattern.compile("\"_id\"\\s*:\\s*\"([^\"]+)\"");
+
+    private String extractWorkloadGroupId(Response response) throws Exception {
+        String body = EntityUtils.toString(response.getEntity());
+        Matcher m = WORKLOAD_GROUP_ID_PATTERN.matcher(body);
+        assertTrue("could not find _id in response: " + body, m.find());
+        return m.group(1);
+    }
+
     Response performOperation(String method, String uriPath, String json) throws IOException {
         Request request = new Request(method, uriPath);
         if (json != null) {

server/src/main/java/org/opensearch/node/Node.java

@@ -1631,7 +1631,8 @@ protected Node(final Environment initialEnvironment, Collection<PluginInfo> clas
                 searchModule.getIndexSearcherExecutor(threadPool),
                 taskResourceTrackingService,
                 searchModule.getConcurrentSearchRequestDeciderFactories(),
-                searchModule.getPluginProfileMetricsProviders()
+                searchModule.getPluginProfileMetricsProviders(),
+                workloadGroupService
             );
 
             final List<PersistentTasksExecutor<?>> tasksExecutors = pluginsService.filterPlugins(PersistentTaskPlugin.class)
@@ -2389,7 +2390,8 @@ protected SearchService newSearchService(
         Executor indexSearcherExecutor,
         TaskResourceTrackingService taskResourceTrackingService,
         Collection<ConcurrentSearchRequestDecider.Factory> concurrentSearchDeciderFactories,
-        List<SearchPlugin.ProfileMetricsProvider> pluginProfilers
+        List<SearchPlugin.ProfileMetricsProvider> pluginProfilers,
+        WorkloadGroupService workloadGroupService
     ) {
         return new SearchService(
             clusterService,
@@ -2404,7 +2406,8 @@ protected SearchService newSearchService(
             indexSearcherExecutor,
             taskResourceTrackingService,
             concurrentSearchDeciderFactories,
-            pluginProfilers
+            pluginProfilers,
+            workloadGroupService
         );
     }
 

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<ConcurrentSearchRequestDecider.Factory> concurrentSearchDeciderFactories,
-        List<SearchPlugin.ProfileMetricsProvider> pluginProfilers
+        List<SearchPlugin.ProfileMetricsProvider> pluginProfilers,
+        WorkloadGroupService workloadGroupService
     ) {
         Settings settings = clusterService.getSettings();
         this.threadPool = threadPool;
@@ -543,7 +545,8 @@ public SearchService(
         this.multiBucketConsumerService = new MultiBucketConsumerService(
             clusterService,
             settings,
-            circuitBreakerService.getBreaker(CircuitBreaker.REQUEST)
+            circuitBreakerService.getBreaker(CircuitBreaker.REQUEST),
+            workloadGroupService
         );
         this.indexSearcherExecutor = indexSearcherExecutor;
         this.taskResourceTrackingService = taskResourceTrackingService;

server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java

@@ -31,6 +31,9 @@
 
 package org.opensearch.search.aggregations;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.cluster.metadata.WorkloadGroup;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.common.settings.Setting;
@@ -42,6 +45,8 @@
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.search.aggregations.bucket.BucketsAggregator;
+import org.opensearch.wlm.WorkloadGroupSearchSettings;
+import org.opensearch.wlm.WorkloadGroupService;
 
 import java.io.IOException;
 import java.util.concurrent.atomic.LongAdder;
@@ -56,6 +61,8 @@
  * @opensearch.internal
  */
 public class MultiBucketConsumerService {
+    private static final Logger logger = LogManager.getLogger(MultiBucketConsumerService.class);
+
     public static final int DEFAULT_MAX_BUCKETS = 65535;
     public static final Setting<Integer> MAX_BUCKET_SETTING = Setting.intSetting(
         "search.max_buckets",
@@ -66,11 +73,18 @@ public class MultiBucketConsumerService {
     );
 
     private final CircuitBreaker breaker;
+    private final WorkloadGroupService workloadGroupService;
 
     private volatile int maxBucket;
 
-    public MultiBucketConsumerService(ClusterService clusterService, Settings settings, CircuitBreaker breaker) {
+    public MultiBucketConsumerService(
+        ClusterService clusterService,
+        Settings settings,
+        CircuitBreaker breaker,
+        WorkloadGroupService workloadGroupService
+    ) {
         this.breaker = breaker;
+        this.workloadGroupService = workloadGroupService;
         this.maxBucket = MAX_BUCKET_SETTING.get(settings);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(MAX_BUCKET_SETTING, this::setMaxBucket);
     }
@@ -79,6 +93,35 @@ private void setMaxBucket(int maxBucket) {
         this.maxBucket = maxBucket;
     }
 
+    /**
+     * Resolves the effective max-buckets limit for the current request by consulting the
+     * workload group (if any) attached to the calling thread context. If the request has no
+     * workload group, the group is unknown, or the group does not define
+     * {@code search.max_buckets}, the cluster-level default is returned.
+     * <p>
+     * The WLM-set value, when present, always wins — {@code override_request_values} is not
+     * relevant because {@code search.max_buckets} is not a per-request parameter.
+     */
+    int resolveMaxBuckets() {
+        try {
+            if (workloadGroupService == null) {
+                return maxBucket;
+            }
+            WorkloadGroup workloadGroup = workloadGroupService.getCurrentWorkloadGroup();
+            if (workloadGroup == null) {
+                return maxBucket;
+            }
+            Settings wlmSettings = workloadGroup.getSettings();
+            if (wlmSettings == null || wlmSettings.hasValue(WorkloadGroupSearchSettings.WLM_MAX_BUCKETS.getKey()) == false) {
+                return maxBucket;
+            }
+            return WorkloadGroupSearchSettings.WLM_MAX_BUCKETS.get(wlmSettings);
+        } catch (Exception e) {
+            logger.warn("Failed to resolve workload group [search.max_buckets]; falling back to cluster default", e);
+            return maxBucket;
+        }
+    }
+
     /**
      * Thrown when there are too many buckets
      *
@@ -216,6 +259,6 @@ public int getLimit() {
     }
 
     public MultiBucketConsumer create() {
-        return new MultiBucketConsumer(maxBucket, breaker);
+        return new MultiBucketConsumer(resolveMaxBuckets(), breaker);
     }
 }

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<Integer> 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<Integer> 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
     );

server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java

@@ -334,6 +334,19 @@ public WorkloadGroup getWorkloadGroupById(String workloadGroupId) {
         return clusterService.state().metadata().workloadGroups().get(workloadGroupId);
     }
 
+    /**
+     * Returns the workload group attached to the calling thread context, or null if the current
+     * request does not map to a workload group (no header set, or the referenced group does not
+     * exist).
+     */
+    public WorkloadGroup getCurrentWorkloadGroup() {
+        String workloadGroupId = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER);
+        if (workloadGroupId == null) {
+            return null;
+        }
+        return getWorkloadGroupById(workloadGroupId);
+    }
+
     public Set<WorkloadGroup> getDeletedWorkloadGroups() {
         return deletedWorkloadGroups;
     }