Add query group stats constructs (#15343)

* add query group stats constructs

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

* add changelog entry

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

* add packageinfo for stats

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

* add total cancellations

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

* add more granular level rejections

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

* add toXContent test cases

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

* move ResourceType enum to wlm

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

* update the comment for query group stats

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

---------

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

Author: kaushalmahi12

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))

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;

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;

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;

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;

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;

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;

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;

server/src/main/java/org/opensearch/search/ResourceType.java renamed to 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<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;
+    }
 }

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<ResourceType, ResourceTypeState> 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<ResourceType, ResourceTypeState> 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;
+        }
+    }
+}