Minor tweaks

Signed-off-by: John Mazanec <[email protected]>
opensearch-project · Feb 27, 2025 · c91a084 · c91a084
1 parent 3fc15f7
commit c91a084
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 72 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManager.java b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManager.java
@@ -54,7 +54,7 @@ public class NativeMemoryCacheManager implements Closeable {
     private static NativeMemoryCacheManager INSTANCE;
     @Setter
     private static ThreadPool threadPool;
-    public static int CB_TIME_INTERVAL = 2 * 60; // seconds
+    public static int CB_TIME_INTERVAL = 20; // seconds
 
     private Cache<String, NativeMemoryAllocation> cache;
     private Deque<String> accessRecencyQueue;

diff --git a/src/test/java/org/opensearch/knn/index/KNNCircuitBreakerIT.java b/src/test/java/org/opensearch/knn/index/KNNCircuitBreakerIT.java
@@ -144,13 +144,6 @@ private double getGraphMemoryPercentage() throws Exception {
         return Double.parseDouble(nodeStatsResponse.getFirst().get(StatNames.GRAPH_MEMORY_USAGE_PERCENTAGE.getName()).toString());
     }
 
-    public boolean isCbTripped() throws Exception {
-        Response response = getKnnStats(Collections.emptyList(), Collections.singletonList("circuit_breaker_triggered"));
-        String responseBody = EntityUtils.toString(response.getEntity());
-        Map<String, Object> clusterStats = parseClusterStatsResponse(responseBody);
-        return Boolean.parseBoolean(clusterStats.get("circuit_breaker_triggered").toString());
-    }
-
     public void testCbTripped() throws Exception {
         setupIndices();
         testClusterLevelCircuitBreaker();

diff --git a/src/test/java/org/opensearch/knn/index/OpenSearchIT.java b/src/test/java/org/opensearch/knn/index/OpenSearchIT.java
@@ -185,70 +185,6 @@ public void testEndToEnd() throws Exception {
         fail("Graphs are not getting evicted");
     }
 
-    public void testAddDoc_blockedWhenCbTrips() throws Exception {
-        createKnnIndex(INDEX_NAME, createKnnIndexMapping(FIELD_NAME, 2));
-        updateClusterSettings("knn.circuit_breaker.triggered", "true");
-
-        Float[] vector = { 6.0f, 6.0f };
-        ResponseException ex = expectThrows(ResponseException.class, () -> addKnnDoc(INDEX_NAME, "1", FIELD_NAME, vector));
-        String expMessage =
-            "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.";
-        assertThat(EntityUtils.toString(ex.getResponse().getEntity()), containsString(expMessage));
-
-        // reset
-        updateClusterSettings("knn.circuit_breaker.triggered", "false");
-        addKnnDoc(INDEX_NAME, "1", FIELD_NAME, vector);
-    }
-
-    public void testUpdateDoc_blockedWhenCbTrips() throws Exception {
-        createKnnIndex(INDEX_NAME, createKnnIndexMapping(FIELD_NAME, 2));
-        Float[] vector = { 6.0f, 6.0f };
-        addKnnDoc(INDEX_NAME, "1", FIELD_NAME, vector);
-
-        // update
-        updateClusterSettings("knn.circuit_breaker.triggered", "true");
-        Float[] updatedVector = { 8.0f, 8.0f };
-        ResponseException ex = expectThrows(ResponseException.class, () -> updateKnnDoc(INDEX_NAME, "1", FIELD_NAME, vector));
-        String expMessage =
-            "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.";
-        assertThat(EntityUtils.toString(ex.getResponse().getEntity()), containsString(expMessage));
-
-        // reset
-        updateClusterSettings("knn.circuit_breaker.triggered", "false");
-        updateKnnDoc(INDEX_NAME, "1", FIELD_NAME, vector);
-    }
-
-    public void testAddAndSearchIndex_whenCBTrips() throws Exception {
-        createKnnIndex(INDEX_NAME, createKnnIndexMapping(FIELD_NAME, 2));
-        for (int i = 1; i <= 4; i++) {
-            Float[] vector = { (float) i, (float) (i + 1) };
-            addKnnDoc(INDEX_NAME, Integer.toString(i), FIELD_NAME, vector);
-        }
-
-        float[] queryVector = { 1.0f, 1.0f }; // vector to be queried
-        int k = 10; // nearest 10 neighbor
-        KNNQueryBuilder knnQueryBuilder = new KNNQueryBuilder(FIELD_NAME, queryVector, k);
-        Response response = searchKNNIndex(INDEX_NAME, knnQueryBuilder, k);
-        List<KNNResult> results = parseSearchResponse(EntityUtils.toString(response.getEntity()), FIELD_NAME);
-        assertEquals(4, results.size());
-
-        updateClusterSettings("knn.circuit_breaker.triggered", "true");
-        // Try add another doc
-        Float[] vector = { 1.0f, 2.0f };
-        ResponseException ex = expectThrows(ResponseException.class, () -> addKnnDoc(INDEX_NAME, "5", FIELD_NAME, vector));
-
-        // Still get 4 docs
-        response = searchKNNIndex(INDEX_NAME, knnQueryBuilder, k);
-        results = parseSearchResponse(EntityUtils.toString(response.getEntity()), FIELD_NAME);
-        assertEquals(4, results.size());
-
-        updateClusterSettings("knn.circuit_breaker.triggered", "false");
-        addKnnDoc(INDEX_NAME, "5", FIELD_NAME, vector);
-        response = searchKNNIndex(INDEX_NAME, knnQueryBuilder, k);
-        results = parseSearchResponse(EntityUtils.toString(response.getEntity()), FIELD_NAME);
-        assertEquals(5, results.size());
-    }
-
     public void testIndexingVectorValidation_differentSizes() throws Exception {
         Settings settings = Settings.builder().put(getKNNDefaultIndexSettings()).build();
 

diff --git a/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java b/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java
@@ -2401,4 +2401,15 @@ protected void setupSnapshotRestore(String index, String snapshot, String reposi
         // create snapshot
         createSnapshot(repository, snapshot, true);
     }
+
+    protected boolean isCbTripped() throws Exception {
+        Response response = getKnnStats(Collections.emptyList(), Collections.singletonList("circuit_breaker_triggered"));
+        String responseBody = EntityUtils.toString(response.getEntity());
+        Map<String, Object> clusterStats = parseClusterStatsResponse(responseBody);
+        return Boolean.parseBoolean(clusterStats.get("circuit_breaker_triggered").toString());
+    }
+
+    protected void tripCB(String index, String snapshot) {
+        setupSnapshotRestore(index, snapshot, "repo-" + randomLowerCaseString());
+    }
 }