add check for WLM mode to control whether CRUD requests proceed

Signed-off-by: Ruirui Zhang <[email protected]>

Author: ruai0511

CHANGELOG.md

@@ -10,7 +10,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add NodeResourceUsageStats to ClusterInfo ([#18480](https://github.com/opensearch-project/OpenSearch/issues/18472))
 - Introduce SecureHttpTransportParameters experimental API (to complement SecureTransportParameters counterpart) ([#18572](https://github.com/opensearch-project/OpenSearch/issues/18572))
 - Create equivalents of JSM's AccessController in the java agent ([#18346](https://github.com/opensearch-project/OpenSearch/issues/18346))
-- Introduced a new cluster-level API to fetch remote store metadata (segments and translogs) for each shard of an index. ([#18257](https://github.com/opensearch-project/OpenSearch/pull/18257))
+- [WLM] Add WLM mode validation for workload group CRUD requests ([#18346](https://github.com/opensearch-project/OpenSearch/issues/18346))
+- Introduced a new cluster-level API to fetch remote store metadata (segments and translogs) for each shard of an index. ([#18652](https://github.com/opensearch-project/OpenSearch/pull/18652))
 - Add last index request timestamp columns to the `_cat/indices` API. ([10766](https://github.com/opensearch-project/OpenSearch/issues/10766))
 - Introduce a new pull-based ingestion plugin for file-based indexing (for local testing) ([#18591](https://github.com/opensearch-project/OpenSearch/pull/18591))
 - Add support for search pipeline in search and msearch template ([#18564](https://github.com/opensearch-project/OpenSearch/pull/18564))

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java

@@ -68,6 +68,7 @@
 import org.opensearch.transport.TransportService;
 import org.opensearch.transport.client.Client;
 import org.opensearch.watcher.ResourceWatcherService;
+import org.opensearch.wlm.WorkloadManagementSettings;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -91,6 +92,8 @@ public class WorkloadManagementPlugin extends Plugin implements ActionPlugin, Sy
     private static FeatureType featureType;
     private static RulePersistenceService rulePersistenceService;
     private static RuleRoutingService ruleRoutingService;
+    private WorkloadManagementSettings workloadManagementSettings;
+
     private AutoTaggingActionFilter autoTaggingActionFilter;
 
     /**
@@ -112,6 +115,7 @@ public Collection<Object> createComponents(
         IndexNameExpressionResolver indexNameExpressionResolver,
         Supplier<RepositoriesService> repositoriesServiceSupplier
     ) {
+        workloadManagementSettings = new WorkloadManagementSettings(environment.settings(), clusterService.getClusterSettings());
         featureType = new WorkloadGroupFeatureType(new WorkloadGroupFeatureValueValidator(clusterService));
         RuleEntityParser parser = new XContentRuleParser(featureType);
         AttributeValueStoreFactory attributeValueStoreFactory = new AttributeValueStoreFactory(
@@ -141,7 +145,7 @@ public Collection<Object> createComponents(
         );
 
         autoTaggingActionFilter = new AutoTaggingActionFilter(ruleProcessingService, threadPool);
-        return List.of(refreshMechanism);
+        return List.of(refreshMechanism, workloadManagementSettings);
     }
 
     @Override
@@ -181,10 +185,10 @@ public List<RestHandler> getRestHandlers(
         Supplier<DiscoveryNodes> nodesInCluster
     ) {
         return List.of(
-            new RestCreateWorkloadGroupAction(),
+            new RestCreateWorkloadGroupAction(workloadManagementSettings),
             new RestGetWorkloadGroupAction(),
-            new RestDeleteWorkloadGroupAction(),
-            new RestUpdateWorkloadGroupAction()
+            new RestDeleteWorkloadGroupAction(workloadManagementSettings),
+            new RestUpdateWorkloadGroupAction(workloadManagementSettings)
         );
     }
 

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateWorkloadGroupAction.java

@@ -21,6 +21,8 @@
 import org.opensearch.rest.RestResponse;
 import org.opensearch.rest.action.RestResponseListener;
 import org.opensearch.transport.client.node.NodeClient;
+import org.opensearch.wlm.WlmMode;
+import org.opensearch.wlm.WorkloadManagementSettings;
 
 import java.io.IOException;
 import java.util.List;
@@ -35,10 +37,15 @@
  */
 public class RestCreateWorkloadGroupAction extends BaseRestHandler {
 
+    private final WorkloadManagementSettings workloadManagementSettings;
+
     /**
      * Constructor for RestCreateWorkloadGroupAction
+     * @param workloadManagementSettings the WorkloadManagementSettings instance to access the current WLM mode
      */
-    public RestCreateWorkloadGroupAction() {}
+    public RestCreateWorkloadGroupAction(WorkloadManagementSettings workloadManagementSettings) {
+        this.workloadManagementSettings = workloadManagementSettings;
+    }
 
     @Override
     public String getName() {
@@ -55,6 +62,9 @@ public List<Route> routes() {
 
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        if (workloadManagementSettings.getWlmMode() == WlmMode.DISABLED) {
+            throw new IllegalStateException("Workload management mode is DISABLED. Cannot create workload group.");
+        }
         try (XContentParser parser = request.contentParser()) {
             CreateWorkloadGroupRequest createWorkloadGroupRequest = CreateWorkloadGroupRequest.fromXContent(parser);
             return channel -> client.execute(

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupAction.java

@@ -14,6 +14,8 @@
 import org.opensearch.rest.RestRequest;
 import org.opensearch.rest.action.RestToXContentListener;
 import org.opensearch.transport.client.node.NodeClient;
+import org.opensearch.wlm.WlmMode;
+import org.opensearch.wlm.WorkloadManagementSettings;
 
 import java.io.IOException;
 import java.util.List;
@@ -27,10 +29,15 @@
  */
 public class RestDeleteWorkloadGroupAction extends BaseRestHandler {
 
+    private final WorkloadManagementSettings workloadManagementSettings;
+
     /**
      * Constructor for RestDeleteWorkloadGroupAction
+     * @param workloadManagementSettings the WorkloadManagementSettings instance to access the current WLM mode
      */
-    public RestDeleteWorkloadGroupAction() {}
+    public RestDeleteWorkloadGroupAction(WorkloadManagementSettings workloadManagementSettings) {
+        this.workloadManagementSettings = workloadManagementSettings;
+    }
 
     @Override
     public String getName() {
@@ -47,6 +54,10 @@ public List<Route> routes() {
 
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        if (workloadManagementSettings.getWlmMode() == WlmMode.DISABLED) {
+            throw new IllegalStateException("Workload management mode is DISABLED. Cannot delete workload group.");
+        }
+
         DeleteWorkloadGroupRequest deleteWorkloadGroupRequest = new DeleteWorkloadGroupRequest(request.param("name"));
         deleteWorkloadGroupRequest.clusterManagerNodeTimeout(
             request.paramAsTime("cluster_manager_timeout", deleteWorkloadGroupRequest.clusterManagerNodeTimeout())

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateWorkloadGroupAction.java

@@ -21,6 +21,8 @@
 import org.opensearch.rest.RestResponse;
 import org.opensearch.rest.action.RestResponseListener;
 import org.opensearch.transport.client.node.NodeClient;
+import org.opensearch.wlm.WlmMode;
+import org.opensearch.wlm.WorkloadManagementSettings;
 
 import java.io.IOException;
 import java.util.List;
@@ -35,10 +37,15 @@
  */
 public class RestUpdateWorkloadGroupAction extends BaseRestHandler {
 
+    private final WorkloadManagementSettings workloadManagementSettings;
+
     /**
      * Constructor for RestUpdateWorkloadGroupAction
+     * @param workloadManagementSettings the WorkloadManagementSettings instance to access the current WLM mode
      */
-    public RestUpdateWorkloadGroupAction() {}
+    public RestUpdateWorkloadGroupAction(WorkloadManagementSettings workloadManagementSettings) {
+        this.workloadManagementSettings = workloadManagementSettings;
+    }
 
     @Override
     public String getName() {
@@ -55,6 +62,9 @@ public List<Route> routes() {
 
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        if (workloadManagementSettings.getWlmMode() == WlmMode.DISABLED) {
+            throw new IllegalStateException("Workload management mode is DISABLED. Cannot update workload group.");
+        }
         try (XContentParser parser = request.contentParser()) {
             UpdateWorkloadGroupRequest updateWorkloadGroupRequest = UpdateWorkloadGroupRequest.fromXContent(parser, request.param("name"));
             return channel -> client.execute(

plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rest/RestCreateWorkloadGroupActionTests.java

@@ -0,0 +1,68 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.wlm.rest;
+
+import org.opensearch.rest.RestRequest;
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.transport.client.node.NodeClient;
+import org.opensearch.wlm.WlmMode;
+import org.opensearch.wlm.WorkloadManagementSettings;
+import org.junit.Before;
+
+import java.io.IOException;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class RestCreateWorkloadGroupActionTests extends OpenSearchTestCase {
+    private WorkloadManagementSettings workloadManagementSettings;
+    private RestCreateWorkloadGroupAction restAction;
+    private NodeClient nodeClient;
+    private RestRequest restRequest;
+
+    @Before
+    public void setup() {
+        workloadManagementSettings = mock(WorkloadManagementSettings.class);
+        restAction = new RestCreateWorkloadGroupAction(workloadManagementSettings);
+        nodeClient = mock(NodeClient.class);
+        restRequest = mock(RestRequest.class);
+    }
+
+    public void testPrepareRequestThrowsWhenWlmModeDisabled() throws IOException {
+        when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED);
+        try {
+            restAction.prepareRequest(restRequest, nodeClient);
+            fail("Expected OpenSearchException when WLM mode is DISABLED");
+        } catch (Exception e) {
+            assertTrue(e.getMessage().contains("Cannot create workload group."));
+        }
+    }
+
+    public void testPrepareRequestSucceedsWhenWlmModeEnabled() throws IOException {
+        when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED);
+        try {
+            restAction.prepareRequest(restRequest, nodeClient);
+        } catch (Exception e) {
+            if (e.getMessage().contains("Cannot create workload group.")) {
+                fail("Expected no exception when WLM mode is ENABLED");
+            }
+        }
+    }
+
+    public void testPrepareRequestSucceedsWhenWlmModeMonitorOnly() throws IOException {
+        when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.MONITOR_ONLY);
+        try {
+            restAction.prepareRequest(restRequest, nodeClient);
+        } catch (Exception e) {
+            if (e.getMessage().contains("Cannot create workload group.")) {
+                fail("Expected no exception when WLM mode is MONITOR_ONLY");
+            }
+        }
+    }
+}

plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupActionTests.java

@@ -20,6 +20,8 @@
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.test.rest.FakeRestRequest;
 import org.opensearch.transport.client.node.NodeClient;
+import org.opensearch.wlm.WlmMode;
+import org.opensearch.wlm.WorkloadManagementSettings;
 
 import java.util.List;
 
@@ -34,13 +36,14 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class RestDeleteWorkloadGroupActionTests extends OpenSearchTestCase {
     /**
      * Test case to validate the construction for RestDeleteWorkloadGroupAction
      */
     public void testConstruction() {
-        RestDeleteWorkloadGroupAction action = new RestDeleteWorkloadGroupAction();
+        RestDeleteWorkloadGroupAction action = new RestDeleteWorkloadGroupAction(mock(WorkloadManagementSettings.class));
         assertNotNull(action);
         assertEquals("delete_workload_group", action.getName());
         List<RestHandler.Route> routes = action.routes();
@@ -55,7 +58,9 @@ public void testConstruction() {
      */
     @SuppressWarnings("unchecked")
     public void testPrepareRequest() throws Exception {
-        RestDeleteWorkloadGroupAction restDeleteWorkloadGroupAction = new RestDeleteWorkloadGroupAction();
+        RestDeleteWorkloadGroupAction restDeleteWorkloadGroupAction = new RestDeleteWorkloadGroupAction(
+            mock(WorkloadManagementSettings.class)
+        );
         NodeClient nodeClient = mock(NodeClient.class);
         RestRequest realRequest = new FakeRestRequest();
         realRequest.params().put("name", NAME_ONE);
@@ -82,4 +87,16 @@ public void testPrepareRequest() throws Exception {
             any(RestToXContentListener.class)
         );
     }
+
+    public void testPrepareRequestThrowsWhenWlmModeDisabled() {
+        WorkloadManagementSettings workloadManagementSettings = mock(WorkloadManagementSettings.class);
+        RestDeleteWorkloadGroupAction restAction = new RestDeleteWorkloadGroupAction(workloadManagementSettings);
+        when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED);
+        try {
+            restAction.prepareRequest(mock(RestRequest.class), mock(NodeClient.class));
+            fail("Expected OpenSearchException when WLM mode is DISABLED");
+        } catch (Exception e) {
+            assertTrue(e.getMessage().contains("Cannot delete workload group."));
+        }
+    }
 }

plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rest/RestUpdateWorkloadGroupActionTests.java

@@ -0,0 +1,66 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.wlm.rest;
+
+import org.opensearch.rest.RestRequest;
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.transport.client.node.NodeClient;
+import org.opensearch.wlm.WlmMode;
+import org.opensearch.wlm.WorkloadManagementSettings;
+import org.junit.Before;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class RestUpdateWorkloadGroupActionTests extends OpenSearchTestCase {
+    private WorkloadManagementSettings workloadManagementSettings;
+    private RestUpdateWorkloadGroupAction restAction;
+    private NodeClient nodeClient;
+    private RestRequest restRequest;
+
+    @Before
+    public void setup() {
+        workloadManagementSettings = mock(WorkloadManagementSettings.class);
+        restAction = new RestUpdateWorkloadGroupAction(workloadManagementSettings);
+        nodeClient = mock(NodeClient.class);
+        restRequest = mock(RestRequest.class);
+    }
+
+    public void testPrepareRequestThrowsWhenWlmModeDisabled() {
+        when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED);
+        try {
+            restAction.prepareRequest(restRequest, nodeClient);
+            fail("Expected OpenSearchException when WLM mode is DISABLED");
+        } catch (Exception e) {
+            assertTrue(e.getMessage().contains("Cannot update workload group."));
+        }
+    }
+
+    public void testPrepareRequestSucceedsWhenWlmModeEnabled() {
+        when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED);
+        try {
+            restAction.prepareRequest(restRequest, nodeClient);
+        } catch (Exception e) {
+            if (e.getMessage().contains("Cannot update workload group.")) {
+                fail("Expected no exception when WLM mode is ENABLED");
+            }
+        }
+    }
+
+    public void testPrepareRequestSucceedsWhenWlmModeMonitorOnly() {
+        when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.MONITOR_ONLY);
+        try {
+            restAction.prepareRequest(restRequest, nodeClient);
+        } catch (Exception e) {
+            if (e.getMessage().contains("Cannot update workload group.")) {
+                fail("Expected no exception when WLM mode is MONITOR_ONLY");
+            }
+        }
+    }
+}