add circuit breaker trigger count stat (#274)

Signed-off-by: Yaliang Wu <ylwu@amazon.com>
opensearch-project · Apr 8, 2022 · 18f9065 · 18f9065
1 parent 08ca3ab
commit 18f9065
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 6 deletions.
diff --git a/plugin/build.gradle b/plugin/build.gradle
@@ -207,8 +207,7 @@ jacocoTestReport {
 List<String> jacocoExclusions = [
         // TODO: add more unit test to meet the minimal test coverage.
         'org.opensearch.ml.constant.CommonValue',
-        'org.opensearch.ml.plugin.*',
-        'org.opensearch.ml.task.MLPredictTaskRunner',
+        'org.opensearch.ml.plugin.MachineLearningPlugin*',
         'org.opensearch.ml.rest.AbstractMLSearchAction*',
         'org.opensearch.ml.rest.RestMLExecuteAction' //0.3
 ]

diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
@@ -162,6 +162,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;
@@ -78,6 +79,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");
         }
         try {

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;
@@ -29,15 +30,17 @@
 import org.opensearch.ml.common.MLTaskType;
 import org.opensearch.ml.common.breaker.MLCircuitBreakerService;
 import org.opensearch.ml.common.exception.MLLimitExceededException;
+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.TransportService;
 
 public class TaskRunnerTests extends OpenSearchTestCase {
 
     @Mock
     MLTaskManager mlTaskManager;
-    @Mock
     MLStats mlStats;
     @Mock
     MLTaskDispatcher mlTaskDispatcher;
@@ -52,6 +55,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) {
             @Override
@@ -98,11 +109,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());
     }
 }