Fix comments

Signed-off-by: David Zane <davizane@amazon.com>

Author: dzane17

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

@@ -91,7 +91,7 @@ public void testUpdate() throws Exception {
                     "cpu" : 1.1,
                     "memory" : -0.1
                 }
-            }'""";
+            }""";
         assertThrows(ResponseException.class, () -> performOperation("PUT", "_wlm/workload_group/analytics5", json));
         assertThrows(
             ResponseException.class,

server/src/main/java/org/opensearch/action/search/SearchRequest.java

@@ -649,20 +649,15 @@ public int getMaxConcurrentShardRequests() {
         return maxConcurrentShardRequests == 0 ? 5 : maxConcurrentShardRequests;
     }
 
-    /**
-     * Returns the raw internal value of maxConcurrentShardRequests without applying the default.
-     */
-    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
      * cluster can be throttled with this number to reduce the cluster load. The default is {@code 5}
      */
     public void setMaxConcurrentShardRequests(int maxConcurrentShardRequests) {
-        validatePositiveInteger(maxConcurrentShardRequests, "maxConcurrentShardRequests");
+        if (maxConcurrentShardRequests < 1) {
+            throw new IllegalArgumentException("maxConcurrentShardRequests must be >= 1");
+        }
         this.maxConcurrentShardRequests = maxConcurrentShardRequests;
     }
 
@@ -680,7 +675,9 @@ public void setMaxConcurrentShardRequests(int maxConcurrentShardRequests) {
      * </ul>
      */
     public void setPreFilterShardSize(int preFilterShardSize) {
-        validatePositiveInteger(preFilterShardSize, "preFilterShardSize");
+        if (preFilterShardSize < 1) {
+            throw new IllegalArgumentException("preFilterShardSize must be >= 1");
+        }
         this.preFilterShardSize = preFilterShardSize;
     }
 
@@ -738,18 +735,6 @@ public static int resolveTrackTotalHitsUpTo(Scroll scroll, SearchSourceBuilder s
             : source.trackTotalHitsUpTo();
     }
 
-    /**
-     * Validates that an integer value is positive (>= 1).
-     * @param value the value to validate
-     * @param parameterName the name of the parameter for error messages
-     * @throws IllegalArgumentException if the value is not positive
-     */
-    public static void validatePositiveInteger(int value, String parameterName) {
-        if (value < 1) {
-            throw new IllegalArgumentException(parameterName + " must be >= 1");
-        }
-    }
-
     public void setCancelAfterTimeInterval(TimeValue cancelAfterTimeInterval) {
         this.cancelAfterTimeInterval = cancelAfterTimeInterval;
     }

server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java

@@ -25,7 +25,6 @@
 
 import java.io.IOException;
 import java.util.HashMap;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;

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

@@ -1140,8 +1140,7 @@ protected Node(final Environment initialEnvironment, Collection<PluginInfo> clas
 
             final WorkloadGroupRequestOperationListener workloadGroupRequestOperationListener = new WorkloadGroupRequestOperationListener(
                 workloadGroupService,
-                threadPool,
-                clusterService
+                threadPool
             );
 
             // register all standard SearchRequestOperationsCompositeListenerFactory to the SearchRequestOperationsCompositeListenerFactory

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

@@ -73,7 +73,7 @@ public MutableWorkloadGroupFragment(StreamInput in) throws IOException {
         if (in.getVersion().onOrAfter(Version.V_3_5_0)) {
             // Read null marker: true means searchSettings is null (not specified)
             boolean isNull = in.readBoolean();
-            searchSettings = isNull ? null : in.readMap(StreamInput::readOptionalString, StreamInput::readOptionalString);
+            searchSettings = isNull ? null : in.readMap(StreamInput::readString, StreamInput::readString);
         } else {
             searchSettings = new HashMap<>();
         }
@@ -195,7 +195,7 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getVersion().onOrAfter(Version.V_3_5_0)) {
             out.writeBoolean(searchSettings == null);
             if (searchSettings != null) {
-                out.writeMap(searchSettings, StreamOutput::writeOptionalString, StreamOutput::writeOptionalString);
+                out.writeMap(searchSettings, StreamOutput::writeString, StreamOutput::writeString);
             }
         }
     }

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

@@ -56,7 +56,7 @@ public String getSettingName() {
         void validate(String value) {
             String error = validator.apply(value);
             if (error != null) {
-                throw new IllegalArgumentException("Invalid value for " + settingName + ": " + error);
+                throw new IllegalArgumentException("Invalid value '" + value + "' for " + settingName + ": " + error);
             }
         }
 
@@ -85,6 +85,12 @@ public static void validateSearchSettings(Map<String, String> searchSettings) {
             return;
         }
         for (Map.Entry<String, String> entry : searchSettings.entrySet()) {
+            if (entry.getKey() == null) {
+                throw new IllegalArgumentException("Search setting key cannot be null");
+            }
+            if (entry.getValue() == null) {
+                throw new IllegalArgumentException("Search setting value cannot be null for key: " + entry.getKey());
+            }
             WlmSearchSetting setting = WlmSearchSetting.fromKey(entry.getKey());
             if (setting == null) {
                 throw new IllegalArgumentException("Unknown search setting: " + entry.getKey());

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

@@ -325,6 +325,15 @@ public Set<WorkloadGroup> getActiveWorkloadGroups() {
         return activeWorkloadGroups;
     }
 
+    /**
+     * Returns the workload group with the given ID, or null if not found.
+     * @param workloadGroupId the workload group identifier
+     * @return the WorkloadGroup or null
+     */
+    public WorkloadGroup getWorkloadGroupById(String workloadGroupId) {
+        return clusterService.state().metadata().workloadGroups().get(workloadGroupId);
+    }
+
     public Set<WorkloadGroup> getDeletedWorkloadGroups() {
         return deletedWorkloadGroups;
     }

server/src/main/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListener.java

@@ -15,8 +15,6 @@
 import org.opensearch.action.search.SearchRequestContext;
 import org.opensearch.action.search.SearchRequestOperationsListener;
 import org.opensearch.cluster.metadata.WorkloadGroup;
-import org.opensearch.cluster.metadata.WorkloadGroupMetadata;
-import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.wlm.WorkloadGroupSearchSettings;
 import org.opensearch.wlm.WorkloadGroupService;
@@ -32,16 +30,10 @@ public class WorkloadGroupRequestOperationListener extends SearchRequestOperatio
     private static final Logger logger = LogManager.getLogger(WorkloadGroupRequestOperationListener.class);
     private final WorkloadGroupService workloadGroupService;
     private final ThreadPool threadPool;
-    private final ClusterService clusterService;
 
-    public WorkloadGroupRequestOperationListener(
-        WorkloadGroupService workloadGroupService,
-        ThreadPool threadPool,
-        ClusterService clusterService
-    ) {
+    public WorkloadGroupRequestOperationListener(WorkloadGroupService workloadGroupService, ThreadPool threadPool) {
         this.workloadGroupService = workloadGroupService;
         this.threadPool = threadPool;
-        this.clusterService = clusterService;
     }
 
     /**
@@ -69,24 +61,18 @@ protected void onRequestFailure(SearchPhaseContext context, SearchRequestContext
      * @param searchRequest the search request to modify
      */
     private void applyWorkloadGroupSearchSettings(String workloadGroupId, SearchRequest searchRequest) {
-        // Skip if no workload group ID (default group is added later)
         if (workloadGroupId == null) {
+            // Return if request contains no WLM group assignment (default group is added later)
             return;
         }
 
-        WorkloadGroupMetadata metadata = clusterService.state().metadata().custom(WorkloadGroupMetadata.TYPE);
-        if (metadata == null) {
-            return;
-        }
+        WorkloadGroup workloadGroup = workloadGroupService.getWorkloadGroupById(workloadGroupId);
 
-        // Get the workload group object by ID
-        WorkloadGroup workloadGroup = metadata.workloadGroups().get(workloadGroupId);
         if (workloadGroup == null) {
             return;
         }
 
         // Loop through WLM group search settings and apply them as needed
-        // WLM settings are applied only if the corresponding setting is not already set in the request
         for (Map.Entry<String, String> entry : workloadGroup.getSearchSettings().entrySet()) {
             try {
                 WorkloadGroupSearchSettings.WlmSearchSetting settingKey = WorkloadGroupSearchSettings.WlmSearchSetting.fromKey(

server/src/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java

@@ -13,7 +13,6 @@
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.metadata.Metadata;
 import org.opensearch.cluster.metadata.WorkloadGroup;
-import org.opensearch.cluster.metadata.WorkloadGroupMetadata;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
@@ -69,7 +68,7 @@ public void setUp() throws Exception {
         Metadata mockMetaData = mock(Metadata.class);
         when(mockClusterState.metadata()).thenReturn(mockMetaData);
         workloadGroupService = mock(WorkloadGroupService.class);
-        sut = new WorkloadGroupRequestOperationListener(workloadGroupService, testThreadPool, mockClusterService);
+        sut = new WorkloadGroupRequestOperationListener(workloadGroupService, testThreadPool);
         mockSearchRequest = new SearchRequest();
         mockSearchRequestContext = mock(SearchRequestContext.class);
         when(mockSearchRequestContext.getRequest()).thenReturn(mockSearchRequest);
@@ -147,7 +146,7 @@ public void testMultiThreadedValidWorkloadGroupRequestFailures() {
             Collections.emptySet()
         );
 
-        sut = new WorkloadGroupRequestOperationListener(workloadGroupService, testThreadPool, mockClusterService);
+        sut = new WorkloadGroupRequestOperationListener(workloadGroupService, testThreadPool);
 
         List<Thread> threads = new ArrayList<>();
         for (int i = 0; i < ITERATIONS; i++) {
@@ -265,7 +264,7 @@ private void assertSuccess(
                 Collections.emptySet(),
                 Collections.emptySet()
             );
-            sut = new WorkloadGroupRequestOperationListener(workloadGroupService, testThreadPool, mockClusterService);
+            sut = new WorkloadGroupRequestOperationListener(workloadGroupService, testThreadPool);
             sut.onRequestFailure(null, null);
 
             HashSet<String> set = new HashSet<>();
@@ -296,11 +295,7 @@ public void testApplySearchSettings_NullWorkloadGroupId() {
 
     public void testApplySearchSettings_NullMetadata() {
         testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "test-wg-id");
-        ClusterState state = mock(ClusterState.class);
-        Metadata metadata = mock(Metadata.class);
-        when(mockClusterService.state()).thenReturn(state);
-        when(state.metadata()).thenReturn(metadata);
-        when(metadata.custom(WorkloadGroupMetadata.TYPE)).thenReturn(null);
+        when(workloadGroupService.getWorkloadGroupById("test-wg-id")).thenReturn(null);
 
         sut.onRequestStart(mockSearchRequestContext);
 
@@ -309,7 +304,7 @@ public void testApplySearchSettings_NullMetadata() {
 
     public void testApplySearchSettings_WorkloadGroupNotFound() {
         testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "non-existent-id");
-        setupWorkloadGroupMetadata(Map.of());
+        when(workloadGroupService.getWorkloadGroupById("non-existent-id")).thenReturn(null);
 
         sut.onRequestStart(mockSearchRequestContext);
 
@@ -321,24 +316,14 @@ public void testApplySearchSettings_PhaseTook() {
 
         String wgId = "test-wg";
         WorkloadGroup wg = createWorkloadGroup(wgId, Map.of("phase_took", "true"));
-        setupWorkloadGroupMetadata(Map.of(wgId, wg));
+        when(workloadGroupService.getWorkloadGroupById(wgId)).thenReturn(wg);
         testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, wgId);
 
         sut.onRequestStart(mockSearchRequestContext);
 
         assertTrue(mockSearchRequest.isPhaseTook());
     }
 
-    private void setupWorkloadGroupMetadata(Map<String, WorkloadGroup> workloadGroups) {
-        ClusterState state = mock(ClusterState.class);
-        Metadata metadata = mock(Metadata.class);
-        WorkloadGroupMetadata wgMetadata = new WorkloadGroupMetadata(workloadGroups);
-        when(mockClusterService.state()).thenReturn(state);
-        when(state.metadata()).thenReturn(metadata);
-        when(metadata.custom(WorkloadGroupMetadata.TYPE)).thenReturn(wgMetadata);
-        when(metadata.workloadGroups()).thenReturn(workloadGroups);
-    }
-
     private WorkloadGroup createWorkloadGroup(String id, Map<String, String> searchSettings) {
         return new WorkloadGroup(
             "test-name",