Add additional WLM search settings

CHANGELOG.md

@@ -23,6 +23,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))

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

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

@@ -31,6 +31,12 @@ private WorkloadGroupSearchSettings() {
      */
     public enum WlmSearchSetting {
         // Query parameters (applied to SearchRequest)
+        /** Setting for batched reduce size */
+        BATCHED_REDUCE_SIZE("batched_reduce_size", WorkloadGroupSearchSettings::validateBatchedReduceSize),
+        /** Setting for canceling search requests after a time interval */
+        CANCEL_AFTER_TIME_INTERVAL("cancel_after_time_interval", WorkloadGroupSearchSettings::validateTimeValue),
+        /** Setting for maximum concurrent shard requests */
+        MAX_CONCURRENT_SHARD_REQUESTS("max_concurrent_shard_requests", WorkloadGroupSearchSettings::validatePositiveInt),
         /** Setting for search request timeout */
         TIMEOUT("timeout", WorkloadGroupSearchSettings::validateTimeValue);
 
@@ -114,4 +120,38 @@ private static String validateTimeValue(String value) {
             return e.getMessage();
         }
     }
+
+    /**
+     * Validates a positive integer string.
+     * @param value the string to validate
+     * @return null if valid, error message if invalid
+     */
+    private static String validatePositiveInt(String value) {
+        try {
+            int intValue = Integer.parseInt(value);
+            if (intValue < 1) {
+                return "must be positive";
+            }
+            return null;
+        } catch (NumberFormatException e) {
+            return "must be a valid integer";
+        }
+    }
+
+    /**
+     * Validates batched reduce size (must be >= 2).
+     * @param value the string to validate
+     * @return null if valid, error message if invalid
+     */
+    private static String validateBatchedReduceSize(String value) {
+        try {
+            int intValue = Integer.parseInt(value);
+            if (intValue < 2) {
+                return "must be >= 2";
+            }
+            return null;
+        } catch (NumberFormatException e) {
+            return "must be a valid integer";
+        }
+    }
 }

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

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 positive"));
+
+        exception = expectThrows(
+            IllegalArgumentException.class,
+            () -> WorkloadGroupSearchSettings.WlmSearchSetting.MAX_CONCURRENT_SHARD_REQUESTS.validate("-1")
+        );
+        assertTrue(exception.getMessage().contains("must be positive"));
+
+        exception = expectThrows(
+            IllegalArgumentException.class,
+            () -> WorkloadGroupSearchSettings.WlmSearchSetting.MAX_CONCURRENT_SHARD_REQUESTS.validate("abc")
+        );
+        assertTrue(exception.getMessage().contains("must be a valid integer"));
+    }
+
     public void testValidateSearchSettingsValid() {
         Map<String, String> 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("must be a valid integer"));
+    }
 }

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<String, String> searchSettings) {
         return new WorkloadGroup(
             "test-name",