WorkloadManagementIT

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

server/src/internalClusterTest/java/org/opensearch/wlm/WorkloadManagementIT.java

@@ -0,0 +1,85 @@
+/*
+ * 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.wlm;
+
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+import org.junit.After;
+import org.junit.Before;
+import org.opensearch.client.Request;
+import org.opensearch.client.RequestOptions;
+import org.opensearch.client.Response;
+import org.opensearch.client.ResponseException;
+import org.opensearch.cluster.metadata.QueryGroup;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.plugins.Plugin;
+import org.opensearch.test.OpenSearchIntegTestCase;
+import org.opensearch.test.ParameterizedStaticSettingsOpenSearchIntegTestCase;
+
+import java.io.IOException;
+import java.util.*;
+
+import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
+import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
+import static org.opensearch.wlm.QueryGroupTask.QUERY_GROUP_ID_HEADER;
+
+public class WorkloadManagementIT extends ParameterizedStaticSettingsOpenSearchIntegTestCase {
+
+    private static String queryGroupName = "test";
+
+    public WorkloadManagementIT(Settings nodeSettings) {
+        super(nodeSettings);
+    }
+    @ParametersFactory
+    public static Collection<Object[]> parameters() {
+        return Arrays.asList(
+            new Object[] { Settings.builder().put(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey(), false).build() },
+            new Object[] { Settings.builder().put(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey(), true).build() }
+        );
+    }
+
+    @Before
+    public final void setupNodeSettings() {
+        Settings request = Settings.builder()
+            .put(WorkloadManagementSettings.NODE_LEVEL_MEMORY_REJECTION_THRESHOLD.getKey(), 0.8)
+            .put(WorkloadManagementSettings.NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD.getKey(), 0.9)
+            .put(WorkloadManagementSettings.NODE_LEVEL_CPU_REJECTION_THRESHOLD.getKey(), 0.8)
+            .put(WorkloadManagementSettings.NODE_LEVEL_CPU_CANCELLATION_THRESHOLD.getKey(), 0.9)
+            .build();
+        assertAcked(client().admin().cluster().prepareUpdateSettings().setPersistentSettings(request).get());
+    }
+
+    @After
+    public final void cleanupNodeSettings() {
+        assertAcked(
+            client().admin()
+                .cluster()
+                .prepareUpdateSettings()
+                .setPersistentSettings(Settings.builder().putNull("*"))
+                .setTransientSettings(Settings.builder().putNull("*"))
+        );
+    }
+
+    public void testMemoryThresholdTaskRejection() throws IOException, IOException {
+        QueryGroup queryGroup = new QueryGroup(queryGroupName, new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.3)));
+
+        QueryGroup queryGroup2 = new QueryGroup(queryGroupName + "2", new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.3)));
+        String queryGroupId = queryGroup.get_id();
+        String queryGroupId2 = queryGroup2.get_id();
+
+        client().admin().cluster().prepareState().get().getState().metadata().queryGroups().put(queryGroupId, queryGroup);
+        client().admin().cluster().prepareState().get().getState().metadata().queryGroups().put(queryGroupId2, queryGroup2);
+
+        Request request = new Request("GET", "/");
+        RequestOptions.Builder options = request.getOptions().toBuilder();
+        options.addHeader(QUERY_GROUP_ID_HEADER, queryGroupId);
+        request.setOptions(options);
+        Response response = getRestClient().performRequest(request);
+
+    }
+}

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

@@ -856,7 +856,7 @@ public Map<String, View> views() {
     public Map<String, QueryGroup> queryGroups() {
         return Optional.ofNullable((QueryGroupMetadata) this.custom(QueryGroupMetadata.TYPE))
             .map(QueryGroupMetadata::queryGroups)
-            .orElse(Collections.emptyMap());
+            .orElse(new HashMap<>());
     }
 
     public DecommissionAttributeMetadata decommissionAttributeMetadata() {