Remove usage of cluster level setting for circuit breaker

Simplification of circuit breaker management. Before, we were
periodically updating the circuit breaker cluster setting to have a
global circuit breaker. However, this is inconvenient and not actually
useful.

This change removes that logic and only trips based on node level
circuit breaker. For knn stats, in order to fetch the cluster level
circuit breaker, a transport call is made to check all of the nodes.
This isnt super efficient but its made for stats calls so its not on the
critical path.

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

Author: jmazanec15

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

@@ -21,7 +21,7 @@ public class StatsIT extends AbstractRollingUpgradeTestCase {
     @Before
     public void setUp() throws Exception {
         super.setUp();
-        this.knnStats = new KNNStats();
+        this.knnStats = new KNNStats(null);
     }
 
     // Validate if all the KNN Stats metrics from old version are present in new version

src/main/java/org/opensearch/knn/index/KNNCircuitBreaker.java

@@ -5,34 +5,22 @@
 
 package org.opensearch.knn.index;
 
+import lombok.Getter;
 import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
-import org.opensearch.knn.plugin.stats.StatNames;
-import org.opensearch.knn.plugin.transport.KNNStatsAction;
-import org.opensearch.knn.plugin.transport.KNNStatsNodeResponse;
-import org.opensearch.knn.plugin.transport.KNNStatsRequest;
-import org.opensearch.knn.plugin.transport.KNNStatsResponse;
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
-import org.opensearch.transport.client.Client;
-import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.threadpool.ThreadPool;
 
-import java.util.ArrayList;
-import java.util.List;
-
 /**
  * Runs the circuit breaker logic and updates the settings
  */
+@Getter
 public class KNNCircuitBreaker {
-    private static Logger logger = LogManager.getLogger(KNNCircuitBreaker.class);
     public static final String KNN_CIRCUIT_BREAKER_TIER = "knn_cb_tier";
     public static int CB_TIME_INTERVAL = 2 * 60; // seconds
 
     private static KNNCircuitBreaker INSTANCE;
-    private ThreadPool threadPool;
-    private ClusterService clusterService;
-    private Client client;
+
+    private boolean isTripped = false;
 
     private KNNCircuitBreaker() {}
 
@@ -44,68 +32,34 @@ public static synchronized KNNCircuitBreaker getInstance() {
     }
 
     /**
-     * SetInstance of Circuit Breaker
+     * Initialize the circuit breaker
      *
-     * @param instance KNNCircuitBreaker instance
+     * @param threadPool ThreadPool instance
      */
-    public static synchronized void setInstance(KNNCircuitBreaker instance) {
-        INSTANCE = instance;
-    }
-
-    public void initialize(ThreadPool threadPool, ClusterService clusterService, Client client) {
-        this.threadPool = threadPool;
-        this.clusterService = clusterService;
-        this.client = client;
+    public void initialize(ThreadPool threadPool) {
         NativeMemoryCacheManager nativeMemoryCacheManager = NativeMemoryCacheManager.getInstance();
         Runnable runnable = () -> {
-            if (nativeMemoryCacheManager.isCacheCapacityReached() && clusterService.localNode().isDataNode()) {
+            if (isTripped) {
                 long currentSizeKiloBytes = nativeMemoryCacheManager.getCacheSizeInKilobytes();
                 long circuitBreakerLimitSizeKiloBytes = KNNSettings.state().getCircuitBreakerLimit().getKb();
                 long circuitBreakerUnsetSizeKiloBytes = (long) ((KNNSettings.getCircuitBreakerUnsetPercentage() / 100)
                     * circuitBreakerLimitSizeKiloBytes);
-                /**
-                 * Unset capacityReached flag if currentSizeBytes is less than circuitBreakerUnsetSizeBytes
-                 */
-                if (currentSizeKiloBytes <= circuitBreakerUnsetSizeKiloBytes) {
-                    nativeMemoryCacheManager.setCacheCapacityReached(false);
-                }
-            }
-
-            // Leader node untriggers CB if all nodes have not reached their max capacity
-            if (KNNSettings.isCircuitBreakerTriggered() && clusterService.state().nodes().isLocalNodeElectedClusterManager()) {
-                KNNStatsRequest knnStatsRequest = new KNNStatsRequest();
-                knnStatsRequest.addStat(StatNames.CACHE_CAPACITY_REACHED.getName());
-                knnStatsRequest.timeout(new TimeValue(1000 * 10)); // 10 second timeout
 
-                try {
-                    KNNStatsResponse knnStatsResponse = client.execute(KNNStatsAction.INSTANCE, knnStatsRequest).get();
-                    List<KNNStatsNodeResponse> nodeResponses = knnStatsResponse.getNodes();
-
-                    List<String> nodesAtMaxCapacity = new ArrayList<>();
-                    for (KNNStatsNodeResponse nodeResponse : nodeResponses) {
-                        if ((Boolean) nodeResponse.getStatsMap().get(StatNames.CACHE_CAPACITY_REACHED.getName())) {
-                            nodesAtMaxCapacity.add(nodeResponse.getNode().getId());
-                        }
-                    }
-
-                    if (!nodesAtMaxCapacity.isEmpty()) {
-                        logger.info(
-                            "[KNN] knn.circuit_breaker.triggered stays set. Nodes at max cache capacity: "
-                                + String.join(",", nodesAtMaxCapacity)
-                                + "."
-                        );
-                    } else {
-                        logger.info(
-                            "[KNN] Cache capacity below 75% of the circuit breaker limit for all nodes."
-                                + " Unsetting knn.circuit_breaker.triggered flag."
-                        );
-                        KNNSettings.state().updateCircuitBreakerSettings(false);
-                    }
-                } catch (Exception e) {
-                    logger.error("[KNN] Exception getting stats: " + e);
+                // Unset capacityReached flag if currentSizeBytes is less than circuitBreakerUnsetSizeBytes
+                if (currentSizeKiloBytes <= circuitBreakerUnsetSizeKiloBytes) {
+                    setTripped(false);
                 }
             }
         };
-        this.threadPool.scheduleWithFixedDelay(runnable, TimeValue.timeValueSeconds(CB_TIME_INTERVAL), ThreadPool.Names.GENERIC);
+        threadPool.scheduleWithFixedDelay(runnable, TimeValue.timeValueSeconds(CB_TIME_INTERVAL), ThreadPool.Names.GENERIC);
+    }
+
+    /**
+     * Set the circuit breaker flag
+     *
+     * @param isTripped true if circuit breaker is tripped, false otherwise
+     */
+    public synchronized void setTripped(boolean isTripped) {
+        this.isTripped = isTripped;
     }
 }

src/main/java/org/opensearch/knn/index/KNNSettings.java

@@ -10,15 +10,12 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.OpenSearchParseException;
-import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest;
-import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsResponse;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.Booleans;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
-import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.unit.ByteSizeUnit;
 import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.index.IndexModule;
@@ -28,7 +25,6 @@
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCacheManager;
 import org.opensearch.monitor.jvm.JvmInfo;
 import org.opensearch.monitor.os.OsProbe;
-import org.opensearch.transport.client.Client;
 
 import java.security.InvalidParameterException;
 import java.util.Arrays;
@@ -317,7 +313,8 @@ public class KNNSettings {
         KNN_CIRCUIT_BREAKER_TRIGGERED,
         false,
         NodeScope,
-        Dynamic
+        Dynamic,
+        Setting.Property.Deprecated
     );
 
     public static final Setting<Double> KNN_CIRCUIT_BREAKER_UNSET_PERCENTAGE_SETTING = Setting.doubleSetting(
@@ -473,8 +470,8 @@ public class KNNSettings {
     private final static Map<String, Setting<?>> FEATURE_FLAGS = getFeatureFlags().stream()
         .collect(toUnmodifiableMap(Setting::getKey, Function.identity()));
 
+    @Setter
     private ClusterService clusterService;
-    private Client client;
     @Setter
     private Optional<String> nodeCbAttribute;
 
@@ -638,10 +635,6 @@ public static boolean isKNNPluginEnabled() {
         return KNNSettings.state().getSettingValue(KNNSettings.KNN_PLUGIN_ENABLED);
     }
 
-    public static boolean isCircuitBreakerTriggered() {
-        return KNNSettings.state().getSettingValue(KNNSettings.KNN_CIRCUIT_BREAKER_TRIGGERED);
-    }
-
     /**
      * Retrieves the node-specific circuit breaker limit based on the existing settings.
      *
@@ -806,8 +799,7 @@ public static boolean isShardLevelRescoringDisabledForDiskBasedVector(String ind
             .getAsBoolean(KNN_DISK_VECTOR_SHARD_LEVEL_RESCORING_DISABLED, false);
     }
 
-    public void initialize(Client client, ClusterService clusterService) {
-        this.client = client;
+    public void initialize(ClusterService clusterService) {
         this.clusterService = clusterService;
         this.nodeCbAttribute = Optional.empty();
         setSettingsUpdateConsumers();
@@ -841,35 +833,6 @@ public static ByteSizeValue parseknnMemoryCircuitBreakerValue(String sValue, Byt
         }
     }
 
-    /**
-     * Updates knn.circuit_breaker.triggered setting to true/false
-     * @param flag true/false
-     */
-    public synchronized void updateCircuitBreakerSettings(boolean flag) {
-        ClusterUpdateSettingsRequest clusterUpdateSettingsRequest = new ClusterUpdateSettingsRequest();
-        Settings circuitBreakerSettings = Settings.builder().put(KNNSettings.KNN_CIRCUIT_BREAKER_TRIGGERED, flag).build();
-        clusterUpdateSettingsRequest.persistentSettings(circuitBreakerSettings);
-        client.admin().cluster().updateSettings(clusterUpdateSettingsRequest, new ActionListener<ClusterUpdateSettingsResponse>() {
-            @Override
-            public void onResponse(ClusterUpdateSettingsResponse clusterUpdateSettingsResponse) {
-                logger.debug(
-                    "Cluster setting {}, acknowledged: {} ",
-                    clusterUpdateSettingsRequest.persistentSettings(),
-                    clusterUpdateSettingsResponse.isAcknowledged()
-                );
-            }
-
-            @Override
-            public void onFailure(Exception e) {
-                logger.info(
-                    "Exception while updating circuit breaker setting {} to {}",
-                    clusterUpdateSettingsRequest.persistentSettings(),
-                    e.getMessage()
-                );
-            }
-        });
-    }
-
     public static ByteSizeValue getVectorStreamingMemoryLimit() {
         return KNNSettings.state().getSettingValue(KNN_VECTOR_STREAMING_MEMORY_LIMIT_IN_MB);
     }
@@ -888,10 +851,6 @@ public static int getEfSearchParam(String index) {
             );
     }
 
-    public void setClusterService(ClusterService clusterService) {
-        this.clusterService = clusterService;
-    }
-
     static class SpaceTypeValidator implements Setting.Validator<String> {
 
         @Override

src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperUtil.java

@@ -20,6 +20,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.opensearch.Version;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.knn.index.KNNCircuitBreaker;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.KnnCircuitBreakerException;
 import org.opensearch.knn.index.SpaceType;
@@ -114,9 +115,9 @@ public static int getExpectedVectorLength(final KNNVectorFieldType knnVectorFiel
      * Validate if the circuit breaker is triggered
      */
     static void validateIfCircuitBreakerIsNotTriggered() {
-        if (KNNSettings.isCircuitBreakerTriggered()) {
+        if (KNNCircuitBreaker.getInstance().isTripped()) {
             throw new KnnCircuitBreakerException(
-                "Parsing the created knn vector fields prior to indexing has failed as the circuit breaker triggered.  This indicates that the cluster is low on memory resources and cannot index more documents at the moment. Check _plugins/_knn/stats for the circuit breaker status."
+                "Parsing the created knn vector fields prior to indexing has failed as the circuit breaker triggered.  This indicates that the node is low on memory resources and cannot index more documents at the moment. Check _plugins/_knn/stats for the circuit breaker status."
             );
         }
     }

src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManager.java

@@ -24,6 +24,7 @@
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.knn.common.exception.OutOfNativeMemoryException;
 import org.opensearch.knn.common.featureflags.KNNFeatureFlags;
+import org.opensearch.knn.index.KNNCircuitBreaker;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.plugin.stats.StatNames;
 import org.opensearch.threadpool.ThreadPool;
@@ -413,24 +414,6 @@ public void invalidateAll() {
         cache.invalidateAll();
     }
 
-    /**
-     * Returns whether or not the capacity of the cache has been reached
-     *
-     * @return Boolean of whether cache limit has been reached
-     */
-    public Boolean isCacheCapacityReached() {
-        return cacheCapacityReached.get();
-    }
-
-    /**
-     * Sets cache capacity reached
-     *
-     * @param value Boolean value to set cache Capacity Reached to
-     */
-    public void setCacheCapacityReached(Boolean value) {
-        cacheCapacityReached.set(value);
-    }
-
     /**
      * Get the stats of all of the OpenSearch indices currently loaded into the cache
      *
@@ -461,8 +444,7 @@ private void onRemoval(RemovalNotification<String, NativeMemoryAllocation> remov
         nativeMemoryAllocation.close();
 
         if (RemovalCause.SIZE == removalNotification.getCause()) {
-            KNNSettings.state().updateCircuitBreakerSettings(true);
-            setCacheCapacityReached(true);
+            KNNCircuitBreaker.getInstance().setTripped(true);
         }
 
         logger.debug("[KNN] Cache evicted. Key {}, Reason: {}", removalNotification.getKey(), removalNotification.getCause());

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

@@ -60,6 +60,8 @@
 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;
@@ -202,22 +204,22 @@ public Collection<Object> createComponents(
         VectorReader vectorReader = new VectorReader(client);
         NativeMemoryLoadStrategy.TrainingLoadStrategy.initialize(vectorReader);
 
-        KNNSettings.state().initialize(client, clusterService);
+        KNNSettings.state().initialize(clusterService);
         KNNClusterUtil.instance().initialize(clusterService);
         ModelDao.OpenSearchKNNModelDao.initialize(client, clusterService, environment.settings());
         ModelCache.initialize(ModelDao.OpenSearchKNNModelDao.getInstance(), clusterService);
         TrainingJobRunner.initialize(threadPool, ModelDao.OpenSearchKNNModelDao.getInstance());
         TrainingJobClusterStateListener.initialize(threadPool, ModelDao.OpenSearchKNNModelDao.getInstance(), clusterService);
         QuantizationStateCache.setThreadPool(threadPool);
         NativeMemoryCacheManager.setThreadPool(threadPool);
-        KNNCircuitBreaker.getInstance().initialize(threadPool, clusterService, client);
+        KNNCircuitBreaker.getInstance().initialize(threadPool);
         KNNQueryBuilder.initialize(ModelDao.OpenSearchKNNModelDao.getInstance());
         KNNWeight.initialize(ModelDao.OpenSearchKNNModelDao.getInstance());
         TrainingModelRequest.initialize(ModelDao.OpenSearchKNNModelDao.getInstance(), clusterService);
 
         clusterService.addListener(TrainingJobClusterStateListener.getInstance());
 
-        knnStats = new KNNStats();
+        knnStats = new KNNStats(client);
         return ImmutableList.of(knnStats);
     }
 
@@ -277,7 +279,8 @@ 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<>(ClearCacheAction.INSTANCE, ClearCacheTransportAction.class),
+            new ActionHandler<>(KNNCircuitBreakerTrippedAction.INSTANCE, KNNCircuitBreakerTrippedTransportAction.class)
         );
     }