Remove CB transport action

Signed-off-by: John Mazanec <jmazane@amazon.com>

Author: jmazanec15

qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/StatsIT.java

@@ -7,7 +7,6 @@
 
 import org.apache.hc.core5.http.io.entity.EntityUtils;
 import org.junit.Before;
-import org.opensearch.Version;
 import org.opensearch.client.Response;
 import org.opensearch.client.ResponseException;
 import org.opensearch.knn.plugin.stats.KNNStats;
@@ -22,7 +21,7 @@ public class StatsIT extends AbstractRollingUpgradeTestCase {
     @Before
     public void setUp() throws Exception {
         super.setUp();
-        this.knnStats = new KNNStats(null, () -> Version.CURRENT);
+        this.knnStats = new KNNStats();
     }
 
     // Validate if all the KNN Stats metrics from old version are present in new version

src/main/java/org/opensearch/knn/plugin/KNNPlugin.java

@@ -58,8 +58,6 @@
 import org.opensearch.knn.plugin.transport.DeleteModelTransportAction;
 import org.opensearch.knn.plugin.transport.GetModelAction;
 import org.opensearch.knn.plugin.transport.GetModelTransportAction;
-import org.opensearch.knn.plugin.transport.KNNCircuitBreakerTrippedAction;
-import org.opensearch.knn.plugin.transport.KNNCircuitBreakerTrippedTransportAction;
 import org.opensearch.knn.plugin.transport.KNNStatsAction;
 import org.opensearch.knn.plugin.transport.KNNStatsTransportAction;
 import org.opensearch.knn.plugin.transport.KNNWarmupAction;
@@ -215,7 +213,7 @@ public Collection<Object> createComponents(
 
         clusterService.addListener(TrainingJobClusterStateListener.getInstance());
 
-        KNNStats knnStats = new KNNStats(client, () -> clusterService.getClusterManagerService().getMinNodeVersion());
+        KNNStats knnStats = new KNNStats();
         return ImmutableList.of(knnStats);
     }
 
@@ -275,8 +273,7 @@ public List<RestHandler> getRestHandlers(
             new ActionHandler<>(RemoveModelFromCacheAction.INSTANCE, RemoveModelFromCacheTransportAction.class),
             new ActionHandler<>(SearchModelAction.INSTANCE, SearchModelTransportAction.class),
             new ActionHandler<>(UpdateModelGraveyardAction.INSTANCE, UpdateModelGraveyardTransportAction.class),
-            new ActionHandler<>(ClearCacheAction.INSTANCE, ClearCacheTransportAction.class),
-            new ActionHandler<>(KNNCircuitBreakerTrippedAction.INSTANCE, KNNCircuitBreakerTrippedTransportAction.class)
+            new ActionHandler<>(ClearCacheAction.INSTANCE, ClearCacheTransportAction.class)
         );
     }
 

src/main/java/org/opensearch/knn/plugin/stats/CircuitBreakerStat.java

@@ -5,17 +5,7 @@
 
 package org.opensearch.knn.plugin.stats;
 
-import lombok.Getter;
-import org.opensearch.Version;
-import org.opensearch.core.action.ActionListener;
-import org.opensearch.knn.index.KNNSettings;
-import org.opensearch.knn.plugin.transport.KNNCircuitBreakerTrippedAction;
-import org.opensearch.knn.plugin.transport.KNNCircuitBreakerTrippedRequest;
-import org.opensearch.transport.client.Client;
-
-import java.util.function.Supplier;
-
-import static org.opensearch.knn.index.KNNSettings.KNN_CIRCUIT_BREAKER_TRIGGERED;
+import java.util.List;
 
 /**
  * Cluster stat that checks if the circuit breaker is enabled. For clusters on or after version 3.0, this stat will
@@ -24,49 +14,25 @@
  */
 public class CircuitBreakerStat extends KNNStat<Boolean> {
 
-    private final Client client;
-    private final Supplier<Version> minVersionSupplier;
-    @Getter
-    private Boolean cbTripped;
-
-    /**
-     *
-     * @param client Client used to execute transport call to get CB values of nodes
-     * @param minVersionSupplier Minimum version supplier to provide minimum node version in the cluster
-     */
-    public CircuitBreakerStat(Client client, Supplier<Version> minVersionSupplier) {
+    public CircuitBreakerStat() {
         super(true, null);
-        this.client = client;
-        this.minVersionSupplier = minVersionSupplier;
-        this.cbTripped = null;
     }
 
     @Override
-    public ActionListener<Void> setupContext(ActionListener<Void> actionListener) {
-        // If there are any nodes in the cluster before 3.0, then we need to fall back to checking the CB via cluster
-        // setting
-        if (minVersionSupplier.get().before(Version.V_3_0_0)) {
-            return ActionListener.wrap(voidResponse -> {
-                cbTripped = KNNSettings.state().getSettingValue(KNN_CIRCUIT_BREAKER_TRIGGERED);
-                actionListener.onResponse(voidResponse);
-            }, actionListener::onFailure);
-        }
-        return ActionListener.wrap(
-            voidResponse -> client.execute(
-                KNNCircuitBreakerTrippedAction.INSTANCE,
-                new KNNCircuitBreakerTrippedRequest(),
-                ActionListener.wrap(knnCircuitBreakerTrippedResponse -> {
-                    cbTripped = knnCircuitBreakerTrippedResponse.isTripped();
-                    actionListener.onResponse(voidResponse);
-                }, actionListener::onFailure)
-            ),
-            actionListener::onFailure
-        );
+    public List<String> dependentNodeStats() {
+        return List.of(StatNames.CACHE_CAPACITY_REACHED.getName());
     }
 
     @Override
     public Boolean getValue() {
-        assert cbTripped != null;
-        return cbTripped;
+        return false;
+    }
+
+    @Override
+    public Boolean getValue(KNNNodeStatAggregation aggregation) {
+        if (aggregation == null) {
+            return false;
+        }
+        return aggregation.isClusterLevelCircuitBreakerTripped();
     }
 }

src/main/java/org/opensearch/knn/plugin/stats/KNNNodeStatAggregation.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.plugin.stats;
+
+import lombok.Getter;
+import org.opensearch.knn.plugin.transport.KNNStatsNodeResponse;
+
+import java.util.List;
+
+/**
+ * Class contains aggregations of node stats that can be used as cluster stats. For instance, for the circuit breaker
+ * stat, we want to check if any nodes in the cluster have their circuit breaker tripped. This is reported as the
+ * cache_capacity_reached stat. We need to perform an aggregation on this stat to check if the circuit breaker is
+ * tripped for the cluster. We cannot rely on the node stats returned for this because some nodes may be filtered out
+ */
+public class KNNNodeStatAggregation {
+    @Getter
+    private final boolean isClusterLevelCircuitBreakerTripped;
+
+    /**
+     * From a set of node responses, create the aggregate stats
+     *
+     * @param nodeResponses
+     */
+    public KNNNodeStatAggregation(List<KNNStatsNodeResponse> nodeResponses) {
+        this.isClusterLevelCircuitBreakerTripped = nodeResponses.stream().anyMatch(r -> {
+            if (r == null) {
+                return false;
+            }
+
+            if (r.getStatsMap() == null) {
+                return false;
+            }
+
+            if (r.getStatsMap().containsKey(StatNames.CACHE_CAPACITY_REACHED.getName()) == false) {
+                return false;
+            }
+
+            return (boolean) r.getStatsMap().get(StatNames.CACHE_CAPACITY_REACHED.getName());
+        });
+    }
+}

src/main/java/org/opensearch/knn/plugin/stats/KNNStat.java

@@ -6,8 +6,9 @@
 package org.opensearch.knn.plugin.stats;
 
 import lombok.Getter;
-import org.opensearch.core.action.ActionListener;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.function.Supplier;
 
 /**
@@ -30,13 +31,13 @@ public KNNStat(Boolean isClusterLevel, Supplier<T> supplier) {
     }
 
     /**
-     * Allows stats to set context via asynchronous calls. This should only be used for cluster level stats.
+     * Allows a cluster stat to depend on node stats. This should only be set for cluster stats and should only return
+     * node stats.
      *
-     * @param actionListener listener to call once the context is setup. If no async calls are made, then do nothing.
-     * @return listener that will execute setup on response.
+     * @return list of dependent node stat names. Null if none
      */
-    public ActionListener<Void> setupContext(ActionListener<Void> actionListener) {
-        return actionListener;
+    public List<String> dependentNodeStats() {
+        return Collections.emptyList();
     }
 
     /**
@@ -47,4 +48,14 @@ public ActionListener<Void> setupContext(ActionListener<Void> actionListener) {
     public T getValue() {
         return supplier.get();
     }
+
+    /**
+     * Get the value of the statistic potentially using the {@link KNNNodeStatAggregation}
+     *
+     * @param aggregation that can be used for cluster stats
+     * @return value of the stat
+     */
+    public T getValue(KNNNodeStatAggregation aggregation) {
+        return supplier.get();
+    }
 }

src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java

@@ -7,7 +7,6 @@
 
 import com.google.common.cache.CacheStats;
 import com.google.common.collect.ImmutableMap;
-import org.opensearch.Version;
 import org.opensearch.knn.common.KNNConstants;
 import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
 import org.opensearch.knn.index.engine.KNNEngine;
@@ -21,7 +20,6 @@
 import org.opensearch.knn.plugin.stats.suppliers.ModelIndexStatusSupplier;
 import org.opensearch.knn.plugin.stats.suppliers.ModelIndexingDegradingSupplier;
 import org.opensearch.knn.plugin.stats.suppliers.NativeMemoryCacheManagerSupplier;
-import org.opensearch.transport.client.Client;
 
 import java.time.temporal.ChronoUnit;
 import java.util.HashMap;
@@ -33,16 +31,12 @@
  */
 public class KNNStats {
 
-    private final Client client;
-    private final Supplier<Version> minVersionSupplier;
     private final Map<String, KNNStat<?>> knnStats;
 
     /**
      * Constructor
      */
-    public KNNStats(Client client, Supplier<Version> minVersionSupplier) {
-        this.client = client;
-        this.minVersionSupplier = minVersionSupplier;
+    public KNNStats() {
         this.knnStats = buildStatsMap();
     }
 
@@ -151,7 +145,7 @@ private void addNativeMemoryStats(ImmutableMap.Builder<String, KNNStat<?>> build
             .put(StatNames.GRAPH_QUERY_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_QUERY_REQUESTS)))
             .put(StatNames.GRAPH_INDEX_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_INDEX_ERRORS)))
             .put(StatNames.GRAPH_INDEX_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_INDEX_REQUESTS)))
-            .put(StatNames.CIRCUIT_BREAKER_TRIGGERED.getName(), new CircuitBreakerStat(client, minVersionSupplier));
+            .put(StatNames.CIRCUIT_BREAKER_TRIGGERED.getName(), new CircuitBreakerStat());
     }
 
     private void addEngineStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {