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

[dzane17]

Description

Onboards search.max_buckets as a Workload Management search setting. Today, search.max_buckets is a single cluster-wide setting, leaving operators with no way to apply per-tenant bucket limits. This feature gives users the flexibility to assign more restrictive bucket limits to protect the system, or more permissive limits when their use case calls for it.

Unlike the four settings added in #21523, search.max_buckets does not exist as a per-request parameter — it only exists as a cluster setting. The value is enforced by the MultiBucketConsumerService class at two sites: (1) during the query phase on each data node, and (2) during the final reduce on the coordinator. This PR includes the necessary piping so that MultiBucketConsumerService can resolve the WLM max_buckets value from the workload group ID passed in thread context (already propagated to shards).

Because there is no per-request value to override, override_request_values is not relevant for this setting. The WLM value, when defined, always wins over the cluster setting — workload groups are admin-configured and more granular than the cluster-wide setting.

Related Issues

Resolves #20556
Part of #20555

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit b3de5eb)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Null Pointer Risk When workloadGroupService is null, resolveMaxBuckets() returns the cluster default. However, the constructor accepts a nullable WorkloadGroupService parameter without documenting this contract. If callers pass null unintentionally (e.g., during partial initialization), the code silently falls back to cluster defaults instead of failing fast. This could mask configuration errors in production where WLM is expected to be active. 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); } Flaky Test Risk The test testSearchMaxBucketsEnforcedAtRequestPath assumes that the cluster default permits >1 bucket and that the untagged request will succeed. If the cluster default is changed (e.g., via test configuration or a future default change), the untagged request at line 414 could fail, causing the test to break. The test does not explicitly set or verify the cluster default before asserting success. // 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());

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to b3de5eb

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Use JSON parser instead of regex Parsing JSON responses with regex is fragile and error-prone. Use a proper JSON parser (e.g., ObjectMapper or XContentParser) to extract the _id field reliably, especially since the response format may include escaped characters or whitespace variations that break the regex. plugins/workload-management/src/javaRestTest/java/org/opensearch/rest/WorkloadManagementRestIT.java [504-511] -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); + Map<String, Object> map = XContentHelper.convertToMap( + XContentType.JSON.xContent(), + body, + false + ); + Object id = map.get("_id"); + assertTrue("could not find _id in response: " + body, id != null); + return id.toString(); } Suggestion importance[1-10]: 6 __ Why: Using a proper JSON parser instead of regex is more robust and maintainable. However, for test code with controlled response formats, the regex approach is acceptable. The suggestion offers a moderate improvement in code quality and reliability.	Low
	Narrow exception handling scope The catch-all exception handler may mask critical errors during workload group resolution. Consider catching only specific exceptions (e.g., IllegalStateException, NullPointerException) that are expected during normal operation, and let unexpected exceptions propagate to surface configuration or system issues. server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java [105-123] int resolveMaxBuckets() { + if (workloadGroupService == null) { + return maxBucket; + } 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) { + } catch (IllegalStateException | IllegalArgumentException e) { logger.warn("Failed to resolve workload group [search.max_buckets]; falling back to cluster default", e); return maxBucket; } } Suggestion importance[1-10]: 4 __ Why: While narrowing exception types can improve error visibility, the current catch-all approach is reasonable for a fallback mechanism. The suggestion provides marginal improvement in error handling but doesn't address a critical issue.	Low

---

Previous suggestions

Suggestions up to commit 055c663

Category	Suggestion	Impact
General	Narrow exception handling scope The catch-all exception handler may mask critical errors during workload group resolution. Consider catching only specific exceptions (e.g., IllegalStateException, NullPointerException) that are expected during normal operation, and let unexpected exceptions propagate to surface potential bugs in the workload group service. server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java [105-123] int resolveMaxBuckets() { + if (workloadGroupService == null) { + return maxBucket; + } try { - if (workloadGroupService == null) { - return maxBucket; - } WorkloadGroup workloadGroup = workloadGroupService.resolveFromThreadContext(); 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) { + } catch (IllegalStateException | IllegalArgumentException e) { logger.warn("Failed to resolve workload group [search.max_buckets]; falling back to cluster default", e); return maxBucket; } } Suggestion importance[1-10]: 5 __ Why: The suggestion to narrow exception handling is reasonable for better error visibility, but the impact is moderate. The current catch-all approach provides resilience by ensuring resolveMaxBuckets() always returns a valid value, which is appropriate for a fallback mechanism. The suggested specific exceptions (IllegalStateException, IllegalArgumentException) may not cover all failure scenarios from workloadGroupService.resolveFromThreadContext().	Low

Suggestions up to commit 20e59c7

Category	Suggestion	Impact
General	Narrow exception handling scope The catch-all exception handler silently swallows all exceptions, which could mask critical errors like NullPointerException or IllegalStateException. Consider catching only specific expected exceptions (e.g., IllegalArgumentException from settings parsing) and allowing unexpected exceptions to propagate or be logged at error level. server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java [109-127] int resolveMaxBuckets() { + if (workloadGroupService == null) { + return maxBucket; + } try { - if (workloadGroupService == null) { - return maxBucket; - } WorkloadGroup workloadGroup = workloadGroupService.resolveFromThreadContext(threadPool.getThreadContext()); 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); + } catch (IllegalArgumentException e) { + logger.warn("Invalid workload group [search.max_buckets] setting; falling back to cluster default", e); return maxBucket; } } Suggestion importance[1-10]: 5 __ Why: The suggestion correctly identifies that the catch-all exception handler could mask unexpected errors. However, the improved code only catches IllegalArgumentException, which may not cover all expected exceptions from resolveFromThreadContext() or getSettings(). The suggestion is valid but the implementation could be more comprehensive in identifying all expected exception types.	Low

Suggestions up to commit f28c859

Category	Suggestion	Impact
Possible issue	Validate resolved max_buckets value The method performs multiple null checks but doesn't validate the resolved max_buckets value before returning it. If the WLM setting contains an invalid or corrupted value, it could bypass validation. Add a bounds check to ensure the resolved value is non-negative before returning it. server/src/main/java/org/opensearch/search/aggregations/MultiBucketConsumerService.java [112-137] 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); + int resolvedValue = WorkloadGroupSearchSettings.WLM_MAX_BUCKETS.get(wlmSettings); + if (resolvedValue < 0) { + logger.warn("Invalid workload group [search.max_buckets] value: {}; falling back to cluster default", resolvedValue); + return maxBucket; + } + return resolvedValue; } catch (Exception e) { logger.warn("Failed to resolve workload group [search.max_buckets]; falling back to cluster default", e); return maxBucket; } } Suggestion importance[1-10]: 3 __ Why: The suggestion adds a bounds check for the resolved max_buckets value. However, the WLM_MAX_BUCKETS setting is already defined with validation (Setting.intSetting("search.max_buckets", DEFAULT_MAX_BUCKETS, 0)) that ensures values are >= 0, making this additional check redundant. The suggestion is technically correct but offers minimal value since the validation already exists at the setting definition level.	Low

[github-actions]

✅ Gradle check result for f28c859: SUCCESS

[codecov]

Codecov Report

❌ Patch coverage is 96.00000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 73.57%. Comparing base (6d41554) to head (b3de5eb).
⚠️ Report is 16 commits behind head on main.

Files with missing lines	Patch %	Lines
...earch/aggregations/MultiBucketConsumerService.java	93.75%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21721      +/-   ##
============================================
+ Coverage     73.41%   73.57%   +0.16%     
- Complexity    75312    75413     +101     
============================================
  Files          6028     6028              
  Lines        341999   342016      +17     
  Branches      49185    49187       +2     
============================================
+ Hits         251094   251654     +560     
+ Misses        70951    70401     -550     
- Partials      19954    19961       +7     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Persistent review updated to latest commit 20e59c7

[github-actions]

❌ Gradle check result for 20e59c7: null

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit 055c663

[github-actions]

❌ Gradle check result for 055c663: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit b3de5eb

[github-actions]

❌ Gradle check result for b3de5eb: null

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

✅ Gradle check result for b3de5eb: SUCCESS