From c71ae2a7ac153f24641ed5bf6e2f5613d5680cdc Mon Sep 17 00:00:00 2001 From: Kaushal Kumar Date: Wed, 28 Aug 2024 09:27:14 -0700 Subject: [PATCH] Add query group stats constructs (#15343) * add query group stats constructs Signed-off-by: Kaushal Kumar * add changelog entry Signed-off-by: Kaushal Kumar * add packageinfo for stats Signed-off-by: Kaushal Kumar * add total cancellations Signed-off-by: Kaushal Kumar * add more granular level rejections Signed-off-by: Kaushal Kumar * add toXContent test cases Signed-off-by: Kaushal Kumar * move ResourceType enum to wlm Signed-off-by: Kaushal Kumar * update the comment for query group stats Signed-off-by: Kaushal Kumar --------- Signed-off-by: Kaushal Kumar --- CHANGELOG.md | 1 + .../service/QueryGroupPersistenceService.java | 2 +- .../plugin/wlm/QueryGroupTestUtils.java | 2 +- .../QueryGroupPersistenceServiceTests.java | 2 +- .../cluster/metadata/QueryGroup.java | 2 +- .../SearchBackpressureService.java | 2 +- .../trackers/NodeDuressTrackers.java | 2 +- .../wlm/QueryGroupLevelResourceUsageView.java | 1 - .../{search => wlm}/ResourceType.java | 14 +- .../opensearch/wlm/stats/QueryGroupState.java | 103 ++++++++ .../opensearch/wlm/stats/QueryGroupStats.java | 228 ++++++++++++++++++ .../opensearch/wlm/stats/package-info.java | 12 + ...QueryGroupResourceUsageTrackerService.java | 2 +- .../metadata/QueryGroupMetadataTests.java | 2 +- .../cluster/metadata/QueryGroupTests.java | 2 +- .../SearchBackpressureServiceTests.java | 6 +- .../trackers/NodeDuressTrackersTests.java | 2 +- ...QueryGroupLevelResourceUsageViewTests.java | 1 - .../{search => wlm}/ResourceTypeTests.java | 3 +- .../wlm/stats/QueryGroupStateTests.java | 73 ++++++ .../wlm/stats/QueryGroupStatsTests.java | 75 ++++++ ...GroupResourceUsageTrackerServiceTests.java | 2 +- 22 files changed, 517 insertions(+), 22 deletions(-) rename server/src/main/java/org/opensearch/{search => wlm}/ResourceType.java (87%) create mode 100644 server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java create mode 100644 server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java create mode 100644 server/src/main/java/org/opensearch/wlm/stats/package-info.java rename server/src/test/java/org/opensearch/{search => wlm}/ResourceTypeTests.java (96%) create mode 100644 server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java create mode 100644 server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b95e9ec57733..f3d7525fec7f2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - [Streaming Indexing] Introduce bulk HTTP API streaming flavor ([#15381](https://github.com/opensearch-project/OpenSearch/pull/15381)) - Add support for centralize snapshot creation with pinned timestamp ([#15124](https://github.com/opensearch-project/OpenSearch/pull/15124)) - Add concurrent search support for Derived Fields ([#15326](https://github.com/opensearch-project/OpenSearch/pull/15326)) +- [Workload Management] Add query group stats constructs ([#15343](https://github.com/opensearch-project/OpenSearch/pull/15343))) ### Dependencies - Bump `netty` from 4.1.111.Final to 4.1.112.Final ([#15081](https://github.com/opensearch-project/OpenSearch/pull/15081)) 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 new file mode 100644 index 0000000000000..93cfcea697c43 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java @@ -0,0 +1,103 @@ +/* + * 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.stats; + +import org.opensearch.common.metrics.CounterMetric; +import org.opensearch.wlm.ResourceType; + +import java.util.EnumMap; +import java.util.Map; + +/** + * This class will keep the point in time view of the query group stats + */ +public class QueryGroupState { + /** + * completions at the query group level, this is a cumulative counter since the Opensearch start time + */ + final CounterMetric completions = new CounterMetric(); + + /** + * rejections at the query group level, this is a cumulative counter since the OpenSearch start time + */ + final CounterMetric totalRejections = new CounterMetric(); + + /** + * this will track the cumulative failures in a query group + */ + final CounterMetric failures = new CounterMetric(); + + /** + * This will track total number of cancellations in the query group due to all resource type breaches + */ + final CounterMetric totalCancellations = new CounterMetric(); + + /** + * This is used to store the resource type state both for CPU and MEMORY + */ + private final Map resourceState; + + public QueryGroupState() { + resourceState = new EnumMap<>(ResourceType.class); + for (ResourceType resourceType : ResourceType.values()) { + if (resourceType.hasStatsEnabled()) { + resourceState.put(resourceType, new ResourceTypeState(resourceType)); + } + } + } + + /** + * + * @return completions in the query group + */ + public long getCompletions() { + return completions.count(); + } + + /** + * + * @return rejections in the query group + */ + public long getTotalRejections() { + return totalRejections.count(); + } + + /** + * + * @return failures in the query group + */ + public long getFailures() { + return failures.count(); + } + + public long getTotalCancellations() { + return totalCancellations.count(); + } + + /** + * getter for query group resource state + * @return the query group resource state + */ + public Map getResourceState() { + return resourceState; + } + + /** + * This class holds the resource level stats for the query group + */ + public static class ResourceTypeState { + final ResourceType resourceType; + final CounterMetric cancellations = new CounterMetric(); + final CounterMetric rejections = new CounterMetric(); + + public ResourceTypeState(ResourceType resourceType) { + this.resourceType = 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 new file mode 100644 index 0000000000000..d39bf104332da --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java @@ -0,0 +1,228 @@ +/* + * 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.stats; + +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.wlm.ResourceType; + +import java.io.IOException; +import java.util.Map; +import java.util.Objects; + +/** + * { + * "queryGroupID": { + * "completions": 1233234234, + * "rejections": 12, + * "failures": 97, + * "total_cancellations": 474, + * "CPU": { "current_usage": 49.6, "cancellation": 432, "rejections": 8 }, + * "MEMORY": { "current_usage": 39.6, "cancellation": 42, "rejections": 4 } + * }, + * ... + * ... + * } + */ +public class QueryGroupStats implements ToXContentObject, Writeable { + private final Map stats; + + public QueryGroupStats(Map stats) { + this.stats = stats; + } + + public QueryGroupStats(StreamInput in) throws IOException { + stats = in.readMap(StreamInput::readString, QueryGroupStatsHolder::new); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeMap(stats, StreamOutput::writeString, QueryGroupStatsHolder::writeTo); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject("query_groups"); + for (Map.Entry queryGroupStats : stats.entrySet()) { + builder.startObject(queryGroupStats.getKey()); + queryGroupStats.getValue().toXContent(builder, params); + builder.endObject(); + } + builder.endObject(); + return builder; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + QueryGroupStats that = (QueryGroupStats) o; + return Objects.equals(stats, that.stats); + } + + @Override + public int hashCode() { + return Objects.hash(stats); + } + + /** + * This is a stats holder object which will hold the data for a query group at a point in time + * the instance will only be created on demand through stats api + */ + public static class QueryGroupStatsHolder implements ToXContentObject, Writeable { + public static final String COMPLETIONS = "completions"; + public static final String REJECTIONS = "rejections"; + public static final String TOTAL_CANCELLATIONS = "total_cancellations"; + public static final String FAILURES = "failures"; + private final long completions; + private final long rejections; + private final long failures; + private final long totalCancellations; + private final Map resourceStats; + + public QueryGroupStatsHolder( + long completions, + long rejections, + long failures, + long totalCancellations, + Map resourceStats + ) { + this.completions = completions; + this.rejections = rejections; + this.failures = failures; + this.totalCancellations = totalCancellations; + this.resourceStats = resourceStats; + } + + public QueryGroupStatsHolder(StreamInput in) throws IOException { + this.completions = in.readVLong(); + this.rejections = in.readVLong(); + this.failures = in.readVLong(); + this.totalCancellations = in.readVLong(); + this.resourceStats = in.readMap((i) -> ResourceType.fromName(i.readString()), ResourceStats::new); + } + + /** + * Writes the {@param statsHolder} to {@param out} + * @param out StreamOutput + * @param statsHolder QueryGroupStatsHolder + * @throws IOException exception + */ + public static void writeTo(StreamOutput out, QueryGroupStatsHolder statsHolder) throws IOException { + out.writeVLong(statsHolder.completions); + out.writeVLong(statsHolder.rejections); + out.writeVLong(statsHolder.failures); + out.writeVLong(statsHolder.totalCancellations); + out.writeMap(statsHolder.resourceStats, (o, val) -> o.writeString(val.getName()), ResourceStats::writeTo); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + QueryGroupStatsHolder.writeTo(out, this); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.field(COMPLETIONS, completions); + builder.field(REJECTIONS, rejections); + builder.field(FAILURES, failures); + builder.field(TOTAL_CANCELLATIONS, totalCancellations); + for (Map.Entry resourceStat : resourceStats.entrySet()) { + ResourceType resourceType = resourceStat.getKey(); + ResourceStats resourceStats1 = resourceStat.getValue(); + builder.startObject(resourceType.getName()); + resourceStats1.toXContent(builder, params); + builder.endObject(); + } + return builder; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + QueryGroupStatsHolder that = (QueryGroupStatsHolder) o; + return completions == that.completions + && rejections == that.rejections + && Objects.equals(resourceStats, that.resourceStats) + && failures == that.failures + && totalCancellations == that.totalCancellations; + } + + @Override + public int hashCode() { + return Objects.hash(completions, rejections, totalCancellations, failures, resourceStats); + } + } + + /** + * point in time resource level stats holder + */ + public static class ResourceStats implements ToXContentObject, Writeable { + public static final String CURRENT_USAGE = "current_usage"; + public static final String CANCELLATIONS = "cancellations"; + public static final double PRECISION = 1e-9; + private final double currentUsage; + private final long cancellations; + private final long rejections; + + public ResourceStats(double currentUsage, long cancellations, long rejections) { + this.currentUsage = currentUsage; + this.cancellations = cancellations; + this.rejections = rejections; + } + + public ResourceStats(StreamInput in) throws IOException { + this.currentUsage = in.readDouble(); + this.cancellations = in.readVLong(); + this.rejections = in.readVLong(); + } + + /** + * Writes the {@param stats} to {@param out} + * @param out StreamOutput + * @param stats QueryGroupStatsHolder + * @throws IOException exception + */ + public static void writeTo(StreamOutput out, ResourceStats stats) throws IOException { + out.writeDouble(stats.currentUsage); + out.writeVLong(stats.cancellations); + out.writeVLong(stats.rejections); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + ResourceStats.writeTo(out, this); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.field(CURRENT_USAGE, currentUsage); + builder.field(CANCELLATIONS, cancellations); + builder.field(QueryGroupStatsHolder.REJECTIONS, rejections); + return builder; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + ResourceStats that = (ResourceStats) o; + return (currentUsage - that.currentUsage) < PRECISION && cancellations == that.cancellations && rejections == that.rejections; + } + + @Override + public int hashCode() { + return Objects.hash(currentUsage, cancellations, rejections); + } + } +} diff --git a/server/src/main/java/org/opensearch/wlm/stats/package-info.java b/server/src/main/java/org/opensearch/wlm/stats/package-info.java new file mode 100644 index 0000000000000..2facf8d16df22 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/stats/package-info.java @@ -0,0 +1,12 @@ +/* + * 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. + */ + +/** + * Query group stats related artifacts + */ +package org.opensearch.wlm.stats; 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 new file mode 100644 index 0000000000000..576eec7be1888 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java @@ -0,0 +1,73 @@ +/* + * 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.stats; + +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.wlm.ResourceType; + +import java.util.ArrayList; +import java.util.List; + +public class QueryGroupStateTests extends OpenSearchTestCase { + QueryGroupState queryGroupState; + + public void testRandomQueryGroupsStateUpdates() { + queryGroupState = new QueryGroupState(); + List updaterThreads = new ArrayList<>(); + + for (int i = 0; i < 25; i++) { + if (i % 5 == 0) { + updaterThreads.add(new Thread(() -> queryGroupState.completions.inc())); + } else if (i % 5 == 1) { + updaterThreads.add(new Thread(() -> { + queryGroupState.totalRejections.inc(); + if (randomBoolean()) { + queryGroupState.getResourceState().get(ResourceType.CPU).rejections.inc(); + } else { + queryGroupState.getResourceState().get(ResourceType.MEMORY).rejections.inc(); + } + })); + } else if (i % 5 == 2) { + updaterThreads.add(new Thread(() -> queryGroupState.failures.inc())); + } else if (i % 5 == 3) { + updaterThreads.add(new Thread(() -> queryGroupState.getResourceState().get(ResourceType.CPU).cancellations.inc())); + } else { + updaterThreads.add(new Thread(() -> queryGroupState.getResourceState().get(ResourceType.MEMORY).cancellations.inc())); + } + + if (i % 5 == 3 || i % 5 == 4) { + updaterThreads.add(new Thread(() -> queryGroupState.totalCancellations.inc())); + } + } + + // trigger the updates + updaterThreads.forEach(Thread::start); + // wait for updates to be finished + updaterThreads.forEach(thread -> { + try { + thread.join(); + } catch (InterruptedException ignored) { + + } + }); + + assertEquals(5, queryGroupState.getCompletions()); + assertEquals(5, queryGroupState.getTotalRejections()); + + final long sumOfRejectionsDueToResourceTypes = queryGroupState.getResourceState().get(ResourceType.CPU).rejections.count() + + queryGroupState.getResourceState().get(ResourceType.MEMORY).rejections.count(); + assertEquals(sumOfRejectionsDueToResourceTypes, queryGroupState.getTotalRejections()); + + assertEquals(5, queryGroupState.getFailures()); + assertEquals(10, queryGroupState.getTotalCancellations()); + assertEquals(5, queryGroupState.getResourceState().get(ResourceType.CPU).cancellations.count()); + assertEquals(5, queryGroupState.getResourceState().get(ResourceType.MEMORY).cancellations.count()); + } + +} diff --git a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java new file mode 100644 index 0000000000000..661c3a7beae40 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java @@ -0,0 +1,75 @@ +/* + * 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.stats; + +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.test.AbstractWireSerializingTestCase; +import org.opensearch.wlm.ResourceType; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +public class QueryGroupStatsTests extends AbstractWireSerializingTestCase { + + public void testToXContent() throws IOException { + final Map stats = new HashMap<>(); + final String queryGroupId = "afakjklaj304041-afaka"; + stats.put( + queryGroupId, + new QueryGroupStats.QueryGroupStatsHolder( + 123456789, + 2, + 0, + 13, + Map.of(ResourceType.CPU, new QueryGroupStats.ResourceStats(0.3, 13, 2)) + ) + ); + XContentBuilder builder = JsonXContent.contentBuilder(); + QueryGroupStats queryGroupStats = new QueryGroupStats(stats); + builder.startObject(); + queryGroupStats.toXContent(builder, ToXContent.EMPTY_PARAMS); + builder.endObject(); + assertEquals( + "{\"query_groups\":{\"afakjklaj304041-afaka\":{\"completions\":123456789,\"rejections\":2,\"failures\":0,\"total_cancellations\":13,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", + builder.toString() + ); + } + + @Override + protected Writeable.Reader instanceReader() { + return QueryGroupStats::new; + } + + @Override + protected QueryGroupStats createTestInstance() { + Map stats = new HashMap<>(); + stats.put( + randomAlphaOfLength(10), + new QueryGroupStats.QueryGroupStatsHolder( + randomNonNegativeLong(), + randomNonNegativeLong(), + randomNonNegativeLong(), + randomNonNegativeLong(), + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats( + randomDoubleBetween(0.0, 0.90, false), + randomNonNegativeLong(), + randomNonNegativeLong() + ) + ) + ) + ); + return new QueryGroupStats(stats); + } +} 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;