From ce813f11c2387f59f460e1b70a6b1736dea9aba6 Mon Sep 17 00:00:00 2001
From: Kaushal Kumar <ravi.kaushal97@gmail.com>
Date: Sun, 25 Aug 2024 18:14:45 -0700
Subject: [PATCH] add total cancellations

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>
---
 .../opensearch/wlm/stats/QueryGroupState.java    | 16 ++++++++++++++++
 .../wlm/stats/QueryGroupStateTests.java          |  5 +++++
 2 files changed, 21 insertions(+)

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 f2f3c1ff1f2ce..b53cdc11ea2ea 100644
--- a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java
+++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java
@@ -33,6 +33,11 @@ public class QueryGroupState {
      */
     private final AtomicLong failures = new AtomicLong();
 
+    /**
+     * This will track total number of cancellations in the query group due to all resource type breaches
+     */
+    private final AtomicLong totalCancellations = new AtomicLong();
+
     /**
      * This is used to store the resource type state both for CPU and MEMORY
      */
@@ -68,6 +73,10 @@ public long getFailures() {
         return failures.get();
     }
 
+    public long getTotalCancellations() {
+        return totalCancellations.get();
+    }
+
     /**
      * getter for query group resource state
      * @return the query group resource state
@@ -97,6 +106,13 @@ public void incrementFailures() {
         failures.incrementAndGet();
     }
 
+    /**
+     * this is a call back to increment total cancellations for a query group
+     */
+    public void incrementTotalCancellations() {
+        totalCancellations.incrementAndGet();
+    }
+
     /**
      * This class holds the resource level stats for the query group
      */
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 00d308beeed72..722a8bf9235fe 100644
--- a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java
+++ b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java
@@ -33,6 +33,10 @@ public void testRandomQueryGroupsStateUpdates() {
             } else {
                 updaterThreads.add(new Thread(() -> queryGroupState.getResourceState().get(ResourceType.MEMORY).incrementCancellations()));
             }
+
+            if (i%5 == 3 || i%5 == 4) {
+                updaterThreads.add(new Thread(() -> queryGroupState.incrementTotalCancellations()));
+            }
         }
 
         // trigger the updates
@@ -49,6 +53,7 @@ public void testRandomQueryGroupsStateUpdates() {
         assertEquals(5, queryGroupState.getCompletions());
         assertEquals(5, queryGroupState.getRejections());
         assertEquals(5, queryGroupState.getFailures());
+        assertEquals(10, queryGroupState.getTotalCancellations());
         assertEquals(5, queryGroupState.getResourceState().get(ResourceType.CPU).getCancellations());
         assertEquals(5, queryGroupState.getResourceState().get(ResourceType.MEMORY).getCancellations());
     }