add circuit breaker trigger count stat (#274) (#321)

Signed-off-by: Yaliang Wu <[email protected]>
opensearch-project · May 13, 2022 · 14dcf76 · 14dcf76
1 parent 15d8ae3
commit 14dcf76
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 4 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
@@ -160,6 +160,7 @@ public Collection<Object> createComponents(
         stats.put(StatNames.ML_TOTAL_REQUEST_COUNT, new MLStat<>(false, new CounterSupplier()));
         stats.put(StatNames.ML_TOTAL_FAILURE_COUNT, new MLStat<>(false, new CounterSupplier()));
         stats.put(StatNames.ML_TOTAL_MODEL_COUNT, new MLStat<>(false, new CounterSupplier()));
+        stats.put(StatNames.ML_TOTAL_CIRCUIT_BREAKER_TRIGGER_COUNT, new MLStat<>(false, new CounterSupplier()));
         this.mlStats = new MLStats(stats);
 
         mlIndicesHandler = new MLIndicesHandler(clusterService, client);

diff --git a/plugin/src/main/java/org/opensearch/ml/stats/StatNames.java b/plugin/src/main/java/org/opensearch/ml/stats/StatNames.java
@@ -17,6 +17,7 @@ public class StatNames {
     public static String ML_TOTAL_REQUEST_COUNT = "ml_total_request_count";
     public static String ML_TOTAL_FAILURE_COUNT = "ml_total_failure_count";
     public static String ML_TOTAL_MODEL_COUNT = "ml_total_model_count";
+    public static String ML_TOTAL_CIRCUIT_BREAKER_TRIGGER_COUNT = "ml_total_circuit_breaker_trigger_count";
 
     public static String requestCountStat(FunctionName functionName, ActionName actionName) {
         return String.format(Locale.ROOT, "ml_%s_%s_request_count", functionName, actionName).toLowerCase(Locale.ROOT);

diff --git a/plugin/src/main/java/org/opensearch/ml/task/MLTaskRunner.java b/plugin/src/main/java/org/opensearch/ml/task/MLTaskRunner.java
@@ -6,6 +6,7 @@
 package org.opensearch.ml.task;
 
 import static org.opensearch.ml.stats.StatNames.ML_EXECUTING_TASK_COUNT;
+import static org.opensearch.ml.stats.StatNames.ML_TOTAL_CIRCUIT_BREAKER_TRIGGER_COUNT;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -88,6 +89,7 @@ protected void handleAsyncMLTaskComplete(MLTask mlTask) {
 
     public void run(Request request, TransportService transportService, ActionListener<Response> listener) {
         if (mlCircuitBreakerService.isOpen()) {
+            mlStats.getStat(ML_TOTAL_CIRCUIT_BREAKER_TRIGGER_COUNT).increment();
             throw new MLLimitExceededException("Circuit breaker is open");
         }
         if (!request.isDispatchTask()) {

diff --git a/plugin/src/test/java/org/opensearch/ml/task/TaskRunnerTests.java b/plugin/src/test/java/org/opensearch/ml/task/TaskRunnerTests.java
@@ -16,6 +16,7 @@
 
 import java.time.Instant;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -31,7 +32,10 @@
 import org.opensearch.ml.common.parameter.MLTaskState;
 import org.opensearch.ml.common.parameter.MLTaskType;
 import org.opensearch.ml.common.transport.MLTaskRequest;
+import org.opensearch.ml.stats.MLStat;
 import org.opensearch.ml.stats.MLStats;
+import org.opensearch.ml.stats.StatNames;
+import org.opensearch.ml.stats.suppliers.CounterSupplier;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.transport.TransportResponseHandler;
 import org.opensearch.transport.TransportService;
@@ -40,7 +44,6 @@ public class TaskRunnerTests extends OpenSearchTestCase {
 
     @Mock
     MLTaskManager mlTaskManager;
-    @Mock
     MLStats mlStats;
     @Mock
     MLTaskDispatcher mlTaskDispatcher;
@@ -57,6 +60,14 @@ public class TaskRunnerTests extends OpenSearchTestCase {
 
     @Before
     public void setup() {
+        Map<String, MLStat<?>> stats = new ConcurrentHashMap<>();
+        stats.put(StatNames.ML_EXECUTING_TASK_COUNT, new MLStat<>(false, new CounterSupplier()));
+        stats.put(StatNames.ML_TOTAL_REQUEST_COUNT, new MLStat<>(false, new CounterSupplier()));
+        stats.put(StatNames.ML_TOTAL_FAILURE_COUNT, new MLStat<>(false, new CounterSupplier()));
+        stats.put(StatNames.ML_TOTAL_MODEL_COUNT, new MLStat<>(false, new CounterSupplier()));
+        stats.put(StatNames.ML_TOTAL_CIRCUIT_BREAKER_TRIGGER_COUNT, new MLStat<>(false, new CounterSupplier()));
+        mlStats = new MLStats(stats);
+
         MockitoAnnotations.openMocks(this);
         mlTaskRunner = new MLTaskRunner(mlTaskManager, mlStats, mlTaskDispatcher, mlCircuitBreakerService, clusterService) {
             @Override
@@ -113,11 +124,11 @@ public void testHandleAsyncMLTaskComplete_SyncTask() {
     }
 
     public void testRun_CircuitBreakerOpen() {
-        exceptionRule.expect(MLLimitExceededException.class);
-        exceptionRule.expectMessage("Circuit breaker is open");
         when(mlCircuitBreakerService.isOpen()).thenReturn(true);
         TransportService transportService = mock(TransportService.class);
         ActionListener listener = mock(ActionListener.class);
-        mlTaskRunner.run(null, transportService, listener);
+        expectThrows(MLLimitExceededException.class, () -> mlTaskRunner.run(null, transportService, listener));
+        Long value = (Long) mlStats.getStat(StatNames.ML_TOTAL_CIRCUIT_BREAKER_TRIGGER_COUNT).getValue();
+        assertEquals(1L, value.longValue());
     }
 }