diff --git a/CHANGELOG.md b/CHANGELOG.md index 01b431aa258ea..e87ca2522c23c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add intra segment support for single-value metric aggregations ([#20503](https://github.com/opensearch-project/OpenSearch/pull/20503)) - Add new setting property 'Sensitive' for tiering dynamic settings ([#20901](https://github.com/opensearch-project/OpenSearch/pull/20901)) - Add ref_path support for package-based hunspell dictionary loading ([#20840](https://github.com/opensearch-project/OpenSearch/pull/20840)) +- Add additional WLM search settings ([#20830](https://github.com/opensearch-project/OpenSearch/issues/20830)) - Add support for enabling pluggable data formats, starting with phase-1 of decoupling shard from engine, and introducing basic abstractions ([#20675](https://github.com/opensearch-project/OpenSearch/pull/20675)) - Add warmup phase to wait for lag to catch up in pull-based ingestion before serving ([#20526](https://github.com/opensearch-project/OpenSearch/pull/20526)) 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 47f9c8eb63dad..d5640e846c248 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 @@ -148,14 +148,17 @@ public void testOperationWhenWlmDisabled() throws Exception { } public void testSearchSettings() throws Exception { - // Create with search_settings + // Create with all search_settings String createJson = """ { "name": "search_test", "resiliency_mode": "enforced", "resource_limits": {"cpu": 0.3, "memory": 0.3}, "search_settings": { - "timeout": "30s" + "timeout": "30s", + "cancel_after_time_interval": "1m", + "max_concurrent_shard_requests": "5", + "batched_reduce_size": "512" } }"""; Response response = performOperation("PUT", "_wlm/workload_group", createJson); @@ -166,12 +169,18 @@ public void testSearchSettings() throws Exception { String responseBody = EntityUtils.toString(getResponse.getEntity()); assertTrue(responseBody.contains("\"search_settings\"")); assertTrue(responseBody.contains("\"timeout\":\"30s\"")); + assertTrue(responseBody.contains("\"cancel_after_time_interval\":\"1m\"")); + assertTrue(responseBody.contains("\"max_concurrent_shard_requests\":\"5\"")); + assertTrue(responseBody.contains("\"batched_reduce_size\":\"512\"")); // Update search_settings String updateJson = """ { "search_settings": { - "timeout": "1m" + "timeout": "1m", + "cancel_after_time_interval": "5m", + "max_concurrent_shard_requests": "10", + "batched_reduce_size": "256" } }"""; Response updateResponse = performOperation("PUT", "_wlm/workload_group/search_test", updateJson); @@ -181,6 +190,9 @@ public void testSearchSettings() throws Exception { Response getResponse2 = performOperation("GET", "_wlm/workload_group/search_test", null); String responseBody2 = EntityUtils.toString(getResponse2.getEntity()); assertTrue(responseBody2.contains("\"timeout\":\"1m\"")); + assertTrue(responseBody2.contains("\"cancel_after_time_interval\":\"5m\"")); + assertTrue(responseBody2.contains("\"max_concurrent_shard_requests\":\"10\"")); + assertTrue(responseBody2.contains("\"batched_reduce_size\":\"256\"")); performOperation("DELETE", "_wlm/workload_group/search_test", null); } diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequest.java b/server/src/main/java/org/opensearch/action/search/SearchRequest.java index a1e6e7605cbdb..0adcd25c3a168 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchRequest.java +++ b/server/src/main/java/org/opensearch/action/search/SearchRequest.java @@ -39,6 +39,7 @@ import org.opensearch.action.IndicesRequest; import org.opensearch.action.support.IndicesOptions; import org.opensearch.common.Nullable; +import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.annotation.PublicApi; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.unit.TimeValue; @@ -649,6 +650,15 @@ public int getMaxConcurrentShardRequests() { return maxConcurrentShardRequests == 0 ? 5 : maxConcurrentShardRequests; } + /** + * Returns the raw value of maxConcurrentShardRequests without applying the default. + * A value of {@code 0} means the user has not explicitly set this parameter. + */ + @ExperimentalApi + public int getMaxConcurrentShardRequestsRaw() { + return maxConcurrentShardRequests; + } + /** * Sets the number of shard requests that should be executed concurrently on a single node. This value should be used as a * protection mechanism to reduce the number of shard requests fired per high level search request. Searches that hit the entire diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupSearchSettings.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSearchSettings.java index 5dc66ab8babd0..2c3d36c65c91a 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadGroupSearchSettings.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupSearchSettings.java @@ -8,10 +8,11 @@ package org.opensearch.wlm; +import org.opensearch.common.settings.Setting; import org.opensearch.common.unit.TimeValue; import java.util.Map; -import java.util.function.Function; +import java.util.function.Consumer; /** * Registry of valid workload group search settings with their validators @@ -31,13 +32,19 @@ private WorkloadGroupSearchSettings() { */ public enum WlmSearchSetting { // Query parameters (applied to SearchRequest) + /** Setting for batched reduce size */ + BATCHED_REDUCE_SIZE("batched_reduce_size", v -> Setting.parseInt(v, 2, "batched_reduce_size")), + /** Setting for canceling search requests after a time interval */ + CANCEL_AFTER_TIME_INTERVAL("cancel_after_time_interval", v -> TimeValue.parseTimeValue(v, "cancel_after_time_interval")), + /** Setting for maximum concurrent shard requests */ + MAX_CONCURRENT_SHARD_REQUESTS("max_concurrent_shard_requests", v -> Setting.parseInt(v, 1, "max_concurrent_shard_requests")), /** Setting for search request timeout */ - TIMEOUT("timeout", WorkloadGroupSearchSettings::validateTimeValue); + TIMEOUT("timeout", v -> TimeValue.parseTimeValue(v, "timeout")); private final String settingName; - private final Function validator; + private final Consumer validator; - WlmSearchSetting(String settingName, Function validator) { + WlmSearchSetting(String settingName, Consumer validator) { this.settingName = settingName; this.validator = validator; } @@ -56,9 +63,10 @@ public String getSettingName() { * @throws IllegalArgumentException if the value is invalid */ void validate(String value) { - String error = validator.apply(value); - if (error != null) { - throw new IllegalArgumentException("Invalid value '" + value + "' for " + settingName + ": " + error); + try { + validator.accept(value); + } catch (Exception e) { + throw new IllegalArgumentException("Invalid value '" + value + "' for " + settingName + ": " + e.getMessage()); } } @@ -101,17 +109,4 @@ public static void validateSearchSettings(Map searchSettings) { } } - /** - * Validates a time value string. - * @param value the string to validate - * @return null if valid, error message if invalid - */ - private static String validateTimeValue(String value) { - try { - TimeValue.parseTimeValue(value, "validation"); - return null; - } catch (Exception e) { - return e.getMessage(); - } - } } 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 8c0010f539550..138dd58f33e3f 100644 --- a/server/src/main/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListener.java +++ b/server/src/main/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListener.java @@ -94,6 +94,34 @@ private void applyWorkloadGroupSearchSettings(String workloadGroupId, SearchRequ ); } break; + case BATCHED_REDUCE_SIZE: + // Only apply WLM batched reduce size when the request uses the default value + // TODO: batchedReduceSize is a primitive int with no sentinel value, so we cannot + // distinguish between "not set" and "explicitly set to 512 (the default)". If a user + // explicitly sets batched_reduce_size=512, WLM will still override it. Consider adding + // a raw accessor or tracking boolean similar to maxConcurrentShardRequests. + int wlmBatchedReduceSize = Integer.parseInt(entry.getValue()); + if (searchRequest.getBatchedReduceSize() == SearchRequest.DEFAULT_BATCHED_REDUCE_SIZE) { + searchRequest.setBatchedReduceSize(wlmBatchedReduceSize); + } + break; + case CANCEL_AFTER_TIME_INTERVAL: + // Only apply WLM cancel_after_time_interval when the request has none set + if (searchRequest.getCancelAfterTimeInterval() == null) { + searchRequest.setCancelAfterTimeInterval( + TimeValue.parseTimeValue( + entry.getValue(), + WorkloadGroupSearchSettings.WlmSearchSetting.CANCEL_AFTER_TIME_INTERVAL.getSettingName() + ) + ); + } + break; + case MAX_CONCURRENT_SHARD_REQUESTS: + // Raw value 0 means not explicitly set; only apply WLM when not explicitly set + if (searchRequest.getMaxConcurrentShardRequestsRaw() == 0) { + searchRequest.setMaxConcurrentShardRequests(Integer.parseInt(entry.getValue())); + } + break; } } catch (Exception e) { logger.error("Failed to apply workload group setting [{}={}]: {}", entry.getKey(), entry.getValue(), e); diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSearchSettingsTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSearchSettingsTests.java index c2212662064d5..bbe8930d3156e 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadGroupSearchSettingsTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadGroupSearchSettingsTests.java @@ -16,10 +16,31 @@ public class WorkloadGroupSearchSettingsTests extends OpenSearchTestCase { public void testEnumSettingNames() { + assertEquals("batched_reduce_size", WorkloadGroupSearchSettings.WlmSearchSetting.BATCHED_REDUCE_SIZE.getSettingName()); + assertEquals( + "cancel_after_time_interval", + WorkloadGroupSearchSettings.WlmSearchSetting.CANCEL_AFTER_TIME_INTERVAL.getSettingName() + ); + assertEquals( + "max_concurrent_shard_requests", + WorkloadGroupSearchSettings.WlmSearchSetting.MAX_CONCURRENT_SHARD_REQUESTS.getSettingName() + ); assertEquals("timeout", WorkloadGroupSearchSettings.WlmSearchSetting.TIMEOUT.getSettingName()); } public void testFromKeyValidSettings() { + assertEquals( + WorkloadGroupSearchSettings.WlmSearchSetting.BATCHED_REDUCE_SIZE, + WorkloadGroupSearchSettings.WlmSearchSetting.fromKey("batched_reduce_size") + ); + assertEquals( + WorkloadGroupSearchSettings.WlmSearchSetting.CANCEL_AFTER_TIME_INTERVAL, + WorkloadGroupSearchSettings.WlmSearchSetting.fromKey("cancel_after_time_interval") + ); + assertEquals( + WorkloadGroupSearchSettings.WlmSearchSetting.MAX_CONCURRENT_SHARD_REQUESTS, + WorkloadGroupSearchSettings.WlmSearchSetting.fromKey("max_concurrent_shard_requests") + ); assertEquals(WorkloadGroupSearchSettings.WlmSearchSetting.TIMEOUT, WorkloadGroupSearchSettings.WlmSearchSetting.fromKey("timeout")); } @@ -33,6 +54,7 @@ public void testValidateTimeValue() { WorkloadGroupSearchSettings.WlmSearchSetting.TIMEOUT.validate("30s"); WorkloadGroupSearchSettings.WlmSearchSetting.TIMEOUT.validate("5m"); WorkloadGroupSearchSettings.WlmSearchSetting.TIMEOUT.validate("1h"); + WorkloadGroupSearchSettings.WlmSearchSetting.CANCEL_AFTER_TIME_INTERVAL.validate("1h"); } public void testValidateInvalidTimeValue() { @@ -43,9 +65,35 @@ public void testValidateInvalidTimeValue() { assertTrue(exception.getMessage().contains("Invalid value")); } + public void testValidatePositiveInt() { + WorkloadGroupSearchSettings.WlmSearchSetting.MAX_CONCURRENT_SHARD_REQUESTS.validate("1"); + WorkloadGroupSearchSettings.WlmSearchSetting.MAX_CONCURRENT_SHARD_REQUESTS.validate("100"); + } + + public void testValidateInvalidPositiveInt() { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> WorkloadGroupSearchSettings.WlmSearchSetting.MAX_CONCURRENT_SHARD_REQUESTS.validate("0") + ); + assertTrue(exception.getMessage().contains("must be >= 1")); + + exception = expectThrows( + IllegalArgumentException.class, + () -> WorkloadGroupSearchSettings.WlmSearchSetting.MAX_CONCURRENT_SHARD_REQUESTS.validate("-1") + ); + assertTrue(exception.getMessage().contains("must be >= 1")); + + exception = expectThrows( + IllegalArgumentException.class, + () -> WorkloadGroupSearchSettings.WlmSearchSetting.MAX_CONCURRENT_SHARD_REQUESTS.validate("abc") + ); + assertTrue(exception.getMessage().contains("Invalid value")); + } + public void testValidateSearchSettingsValid() { Map settings = new HashMap<>(); settings.put("timeout", "30s"); + settings.put("max_concurrent_shard_requests", "5"); // Should not throw exception WorkloadGroupSearchSettings.validateSearchSettings(settings); @@ -106,4 +154,29 @@ public void testValidateSearchSettingsEmpty() { // Should not throw exception for empty map WorkloadGroupSearchSettings.validateSearchSettings(settings); } + + public void testValidateBatchedReduceSize() { + WorkloadGroupSearchSettings.WlmSearchSetting.BATCHED_REDUCE_SIZE.validate("2"); + WorkloadGroupSearchSettings.WlmSearchSetting.BATCHED_REDUCE_SIZE.validate("512"); + } + + public void testValidateInvalidBatchedReduceSize() { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> WorkloadGroupSearchSettings.WlmSearchSetting.BATCHED_REDUCE_SIZE.validate("1") + ); + assertTrue(exception.getMessage().contains("must be >= 2")); + + exception = expectThrows( + IllegalArgumentException.class, + () -> WorkloadGroupSearchSettings.WlmSearchSetting.BATCHED_REDUCE_SIZE.validate("0") + ); + assertTrue(exception.getMessage().contains("must be >= 2")); + + exception = expectThrows( + IllegalArgumentException.class, + () -> WorkloadGroupSearchSettings.WlmSearchSetting.BATCHED_REDUCE_SIZE.validate("abc") + ); + assertTrue(exception.getMessage().contains("Invalid value")); + } } 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 7863045de2dea..5eea4b03151dd 100644 --- a/server/src/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java +++ b/server/src/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java @@ -288,33 +288,104 @@ private void setupMockedWorkloadGroupsFromClusterState() { // Tests for applyWorkloadGroupSearchSettings public void testApplySearchSettings_NullWorkloadGroupId() { + mockSearchRequest.setBatchedReduceSize(256); + // No workload group ID in thread context sut.onRequestStart(mockSearchRequestContext); - // Request should remain unchanged - source is null, no timeout set - assertNull(mockSearchRequest.source()); + // Request should remain unchanged + assertEquals(256, mockSearchRequest.getBatchedReduceSize()); } public void testApplySearchSettings_WorkloadGroupNotFound() { + mockSearchRequest.setBatchedReduceSize(256); + testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "non-existent-id"); when(workloadGroupService.getWorkloadGroupById("non-existent-id")).thenReturn(null); sut.onRequestStart(mockSearchRequestContext); - assertNull(mockSearchRequest.source()); + assertEquals(256, mockSearchRequest.getBatchedReduceSize()); } - public void testApplySearchSettings_EmptySearchSettings() { - mockSearchRequest.source(new SearchSourceBuilder()); + public void testApplySearchSettings_BatchedReduceSize_WlmAppliedWhenDefault() { + // Request uses default value (512) + assertEquals(SearchRequest.DEFAULT_BATCHED_REDUCE_SIZE, mockSearchRequest.getBatchedReduceSize()); String wgId = "test-wg"; - WorkloadGroup wg = createWorkloadGroup(wgId, Map.of()); + WorkloadGroup wg = createWorkloadGroup(wgId, Map.of("batched_reduce_size", "100")); when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); sut.onRequestStart(mockSearchRequestContext); - assertNull(mockSearchRequest.source().timeout()); // No settings applied + assertEquals(100, mockSearchRequest.getBatchedReduceSize()); // WLM applied + } + + public void testApplySearchSettings_BatchedReduceSize_RequestAlreadySet() { + mockSearchRequest.setBatchedReduceSize(50); // explicitly set by user + + String wgId = "test-wg"; + WorkloadGroup wg = createWorkloadGroup(wgId, Map.of("batched_reduce_size", "100")); + when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); + + sut.onRequestStart(mockSearchRequestContext); + + assertEquals(50, mockSearchRequest.getBatchedReduceSize()); // Request value preserved + } + + public void testApplySearchSettings_MaxConcurrentShardRequests_WlmAppliedWhenDefault() { + // Request has default value (not explicitly set), raw field is 0 + assertEquals(0, mockSearchRequest.getMaxConcurrentShardRequestsRaw()); + + String wgId = "test-wg"; + WorkloadGroup wg = createWorkloadGroup(wgId, Map.of("max_concurrent_shard_requests", "10")); + when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); + + sut.onRequestStart(mockSearchRequestContext); + + assertEquals(10, mockSearchRequest.getMaxConcurrentShardRequests()); // WLM applied + } + + public void testApplySearchSettings_MaxConcurrentShardRequests_RequestAlreadySet() { + mockSearchRequest.setMaxConcurrentShardRequests(20); // explicitly set by user + + String wgId = "test-wg"; + WorkloadGroup wg = createWorkloadGroup(wgId, Map.of("max_concurrent_shard_requests", "5")); + when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); + + sut.onRequestStart(mockSearchRequestContext); + + assertEquals(20, mockSearchRequest.getMaxConcurrentShardRequests()); // Request value preserved + } + + public void testApplySearchSettings_CancelAfterTimeInterval_WlmAppliedWhenNull() { + assertNull(mockSearchRequest.getCancelAfterTimeInterval()); + + String wgId = "test-wg"; + WorkloadGroup wg = createWorkloadGroup(wgId, Map.of("cancel_after_time_interval", "30s")); + when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); + + sut.onRequestStart(mockSearchRequestContext); + + assertEquals(TimeValue.timeValueSeconds(30), mockSearchRequest.getCancelAfterTimeInterval()); + } + + public void testApplySearchSettings_CancelAfterTimeInterval_RequestAlreadySet() { + mockSearchRequest.setCancelAfterTimeInterval(TimeValue.timeValueSeconds(10)); + + String wgId = "test-wg"; + WorkloadGroup wg = createWorkloadGroup(wgId, Map.of("cancel_after_time_interval", "30s")); + when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); + + sut.onRequestStart(mockSearchRequestContext); + + assertEquals(TimeValue.timeValueSeconds(10), mockSearchRequest.getCancelAfterTimeInterval()); // Request value preserved } public void testApplySearchSettings_Timeout_WlmAppliedWhenNull() { @@ -357,6 +428,37 @@ public void testApplySearchSettings_Timeout_NullSource() { assertNull(mockSearchRequest.source()); // Should not throw, source remains null } + public void testApplySearchSettings_EmptySearchSettings() { + mockSearchRequest.source(new SearchSourceBuilder()); + + String wgId = "test-wg"; + WorkloadGroup wg = createWorkloadGroup(wgId, Map.of()); + when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); + + sut.onRequestStart(mockSearchRequestContext); + + assertNull(mockSearchRequest.source().timeout()); // No settings applied + } + + public void testApplySearchSettings_MultipleSettings() { + mockSearchRequest.source(new SearchSourceBuilder()); + + String wgId = "test-wg"; + WorkloadGroup wg = createWorkloadGroup( + wgId, + Map.of("batched_reduce_size", "100", "max_concurrent_shard_requests", "5", "timeout", "30s") + ); + when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg); + testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId); + + sut.onRequestStart(mockSearchRequestContext); + + assertEquals(100, mockSearchRequest.getBatchedReduceSize()); + assertEquals(5, mockSearchRequest.getMaxConcurrentShardRequests()); + assertEquals(TimeValue.timeValueSeconds(30), mockSearchRequest.source().timeout()); + } + private WorkloadGroup createWorkloadGroup(String id, Map searchSettings) { return new WorkloadGroup( "test-name",