diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index ba5161a2c855e..7561a2f6f99c3 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/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 index 5ba1ad5334712..e165645775d5c 100644 --- 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/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java index a516ffdde839e..5cb3d8fc6d11f 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/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 index 9b5c6bc2369a6..a971aa58940ba 100644 --- 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 index c26c5d63a3573..a85bc69b766cb 100644 --- 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 index ae60a82fc2816..c27c50ac12c0f 100644 --- 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 index 2fd743dc3f83f..7577c8573ec10 100644 --- 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/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/wlm/ResourceType.java similarity index 87% rename from server/src/main/java/org/opensearch/search/ResourceType.java rename to server/src/main/java/org/opensearch/wlm/ResourceType.java index 0cba2222a6e20..adf384995c91d 100644 --- a/server/src/main/java/org/opensearch/search/ResourceType.java +++ b/server/src/main/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 getResourceUsage; + private final boolean statsEnabled; - ResourceType(String name, Function getResourceUsage) { + ResourceType(String name, Function 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 index 44d2c13b58388..93cfcea697c43 100644 --- 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 index 6b9fef979c53e..6a7da53f4a986 100644 --- 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 index bfbf5d8a452d1..15852b5bbe6a8 100644 --- 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 index 06734b8e0bac2..f5e667de73d93 100644 --- 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 index 884b364fb26b8..f4d3e5ceb1784 100644 --- 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 index 15d0fcd10d701..a444eb42eac2e 100644 --- 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/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java b/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java index 801576bdf89d4..7c52840c099d4 100644 --- a/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java +++ b/server/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 index 7f6419505fec2..532bf3de95bd6 100644 --- 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/server/src/test/java/org/opensearch/search/ResourceTypeTests.java b/server/src/test/java/org/opensearch/wlm/ResourceTypeTests.java similarity index 96% rename from server/src/test/java/org/opensearch/search/ResourceTypeTests.java rename to server/src/test/java/org/opensearch/wlm/ResourceTypeTests.java index 78827b8b1bdad..737cbb37b554c 100644 --- a/server/src/test/java/org/opensearch/search/ResourceTypeTests.java +++ b/server/src/test/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 index 76f4b787a34f4..576eec7be1888 100644 --- 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 index c114c811bbe07..661c3a7beae40 100644 --- 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/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java index 967119583c25f..ca2891cb532f2 100644 --- a/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java +++ b/server/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;