move ResourceType enum to wlm

Signed-off-by: Kaushal Kumar <[email protected]>
kaushalmahi12 · Aug 28, 2024 · 353b1c8 · 353b1c8
1 parent 1df409f
commit 353b1c8
Show file tree

Hide file tree

Showing 20 changed files with 34 additions and 28 deletions.
diff --git a/...agement/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/...agement/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java
@@ -29,7 +29,7 @@
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse;
 import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest;
-import org.opensearch.search.ResourceType;
+import org.opensearch.wlm.ResourceType;
 
 import java.util.Collection;
 import java.util.EnumMap;

diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java
@@ -31,7 +31,7 @@
 import java.util.Set;
 
 import static org.opensearch.cluster.metadata.QueryGroup.builder;
-import static org.opensearch.search.ResourceType.fromName;
+import static org.opensearch.wlm.ResourceType.fromName;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;

diff --git a/...nt/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java b/...nt/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java
@@ -24,9 +24,9 @@
 import org.opensearch.plugin.wlm.QueryGroupTestUtils;
 import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse;
 import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest;
-import org.opensearch.search.ResourceType;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.wlm.ResourceType;
 
 import java.util.ArrayList;
 import java.util.Collection;

diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java
@@ -17,7 +17,7 @@
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
-import org.opensearch.search.ResourceType;
+import org.opensearch.wlm.ResourceType;
 import org.joda.time.Instant;
 
 import java.io.IOException;

diff --git a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java b/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java
@@ -18,7 +18,6 @@
 import org.opensearch.common.settings.Setting;
 import org.opensearch.monitor.jvm.JvmStats;
 import org.opensearch.monitor.process.ProcessProbe;
-import org.opensearch.search.ResourceType;
 import org.opensearch.search.backpressure.settings.SearchBackpressureMode;
 import org.opensearch.search.backpressure.settings.SearchBackpressureSettings;
 import org.opensearch.search.backpressure.settings.SearchShardTaskSettings;
@@ -43,6 +42,7 @@
 import org.opensearch.tasks.TaskResourceTrackingService.TaskCompletionListener;
 import org.opensearch.threadpool.Scheduler;
 import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.wlm.ResourceType;
 
 import java.io.IOException;
 import java.util.ArrayList;

diff --git a/server/src/main/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackers.java b/server/src/main/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackers.java
@@ -9,7 +9,7 @@
 package org.opensearch.search.backpressure.trackers;
 
 import org.opensearch.common.util.Streak;
-import org.opensearch.search.ResourceType;
+import org.opensearch.wlm.ResourceType;
 
 import java.util.Map;
 import java.util.function.BooleanSupplier;

diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java b/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java
@@ -8,7 +8,6 @@
 
 package org.opensearch.wlm;
 
-import org.opensearch.search.ResourceType;
 import org.opensearch.tasks.Task;
 
 import java.util.List;

diff --git a/...a/org/opensearch/search/ResourceType.java → ...java/org/opensearch/wlm/ResourceType.java b/...a/org/opensearch/search/ResourceType.java → ...java/org/opensearch/wlm/ResourceType.java
@@ -6,7 +6,7 @@
  * compatible open source license.
  */
 
-package org.opensearch.search;
+package org.opensearch.wlm;
 
 import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.core.common.io.stream.StreamOutput;
@@ -21,15 +21,17 @@
  */
 @PublicApi(since = "2.x")
 public enum ResourceType {
-    CPU("cpu", task -> task.getTotalResourceUtilization(ResourceStats.CPU)),
-    MEMORY("memory", task -> task.getTotalResourceUtilization(ResourceStats.MEMORY));
+    CPU("cpu", task -> task.getTotalResourceUtilization(ResourceStats.CPU), true),
+    MEMORY("memory", task -> task.getTotalResourceUtilization(ResourceStats.MEMORY), true);
 
     private final String name;
     private final Function<Task, Long> getResourceUsage;
+    private final boolean statsEnabled;
 
-    ResourceType(String name, Function<Task, Long> getResourceUsage) {
+    ResourceType(String name, Function<Task, Long> getResourceUsage, boolean statsEnabled) {
         this.name = name;
         this.getResourceUsage = getResourceUsage;
+        this.statsEnabled = statsEnabled;
     }
 
     /**
@@ -63,4 +65,8 @@ public String getName() {
     public long getResourceUsage(Task task) {
         return getResourceUsage.apply(task);
     }
+
+    public boolean hasStatsEnabled() {
+        return statsEnabled;
+    }
 }
diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java
@@ -9,7 +9,7 @@
 package org.opensearch.wlm.stats;
 
 import org.opensearch.common.metrics.CounterMetric;
-import org.opensearch.search.ResourceType;
+import org.opensearch.wlm.ResourceType;
 
 import java.util.EnumMap;
 import java.util.Map;
@@ -45,8 +45,11 @@ public class QueryGroupState {
 
     public QueryGroupState() {
         resourceState = new EnumMap<>(ResourceType.class);
-        resourceState.put(ResourceType.CPU, new ResourceTypeState(ResourceType.CPU));
-        resourceState.put(ResourceType.MEMORY, new ResourceTypeState(ResourceType.MEMORY));
+        for (ResourceType resourceType : ResourceType.values()) {
+            if (resourceType.hasStatsEnabled()) {
+                resourceState.put(resourceType, new ResourceTypeState(resourceType));
+            }
+        }
     }
 
     /**

diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java
@@ -13,7 +13,7 @@
 import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
-import org.opensearch.search.ResourceType;
+import org.opensearch.wlm.ResourceType;
 
 import java.io.IOException;
 import java.util.Map;

diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java
@@ -8,11 +8,11 @@
 
 package org.opensearch.wlm.tracker;
 
-import org.opensearch.search.ResourceType;
 import org.opensearch.tasks.Task;
 import org.opensearch.tasks.TaskResourceTrackingService;
 import org.opensearch.wlm.QueryGroupLevelResourceUsageView;
 import org.opensearch.wlm.QueryGroupTask;
+import org.opensearch.wlm.ResourceType;
 
 import java.util.EnumMap;
 import java.util.EnumSet;

diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java
@@ -14,8 +14,8 @@
 import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
-import org.opensearch.search.ResourceType;
 import org.opensearch.test.AbstractDiffableSerializationTestCase;
+import org.opensearch.wlm.ResourceType;
 
 import java.io.IOException;
 import java.util.Collections;

diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java
@@ -14,8 +14,8 @@
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
-import org.opensearch.search.ResourceType;
 import org.opensearch.test.AbstractSerializingTestCase;
+import org.opensearch.wlm.ResourceType;
 import org.joda.time.Instant;
 
 import java.io.IOException;

diff --git a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java
@@ -16,7 +16,6 @@
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.xcontent.XContentBuilder;
-import org.opensearch.search.ResourceType;
 import org.opensearch.search.backpressure.settings.SearchBackpressureMode;
 import org.opensearch.search.backpressure.settings.SearchBackpressureSettings;
 import org.opensearch.search.backpressure.settings.SearchShardTaskSettings;
@@ -40,6 +39,7 @@
 import org.opensearch.test.transport.MockTransportService;
 import org.opensearch.threadpool.TestThreadPool;
 import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.wlm.ResourceType;
 import org.junit.After;
 import org.junit.Before;
 
@@ -56,9 +56,9 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.LongSupplier;
 
-import static org.opensearch.search.ResourceType.CPU;
-import static org.opensearch.search.ResourceType.MEMORY;
 import static org.opensearch.search.backpressure.SearchBackpressureTestHelpers.createMockTaskWithResourceStats;
+import static org.opensearch.wlm.ResourceType.CPU;
+import static org.opensearch.wlm.ResourceType.MEMORY;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.any;

diff --git a/...er/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java b/...er/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java
@@ -8,9 +8,9 @@
 
 package org.opensearch.search.backpressure.trackers;
 
-import org.opensearch.search.ResourceType;
 import org.opensearch.search.backpressure.trackers.NodeDuressTrackers.NodeDuressTracker;
 import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.wlm.ResourceType;
 
 import java.util.EnumMap;
 

diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java
@@ -10,7 +10,6 @@
 
 import org.opensearch.action.search.SearchAction;
 import org.opensearch.core.tasks.TaskId;
-import org.opensearch.search.ResourceType;
 import org.opensearch.tasks.Task;
 import org.opensearch.test.OpenSearchTestCase;
 

diff --git a/.../opensearch/search/ResourceTypeTests.java → ...org/opensearch/wlm/ResourceTypeTests.java b/.../opensearch/search/ResourceTypeTests.java → ...org/opensearch/wlm/ResourceTypeTests.java
@@ -6,14 +6,13 @@
  * compatible open source license.
  */
 
-package org.opensearch.search;
+package org.opensearch.wlm;
 
 import org.opensearch.action.search.SearchShardTask;
 import org.opensearch.core.tasks.resourcetracker.ResourceStats;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.test.OpenSearchTestCase;
 
-import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 

diff --git a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java
@@ -8,8 +8,8 @@
 
 package org.opensearch.wlm.stats;
 
-import org.opensearch.search.ResourceType;
 import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.wlm.ResourceType;
 
 import java.util.ArrayList;
 import java.util.List;

diff --git a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java
@@ -12,8 +12,8 @@
 import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
-import org.opensearch.search.ResourceType;
 import org.opensearch.test.AbstractWireSerializingTestCase;
+import org.opensearch.wlm.ResourceType;
 
 import java.io.IOException;
 import java.util.HashMap;

diff --git a/.../src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java b/.../src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java
@@ -12,7 +12,6 @@
 import org.opensearch.action.search.SearchTask;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.core.tasks.resourcetracker.ResourceStats;
-import org.opensearch.search.ResourceType;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.Task;
 import org.opensearch.tasks.TaskResourceTrackingService;
@@ -21,6 +20,7 @@
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.wlm.QueryGroupLevelResourceUsageView;
 import org.opensearch.wlm.QueryGroupTask;
+import org.opensearch.wlm.ResourceType;
 import org.junit.After;
 import org.junit.Before;