Address comments

Signed-off-by: John Mazanec <[email protected]>
opensearch-project · Mar 18, 2024 · c1c39d2 · c1c39d2
1 parent 730381e
commit c1c39d2
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 10 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/IndexUtil.java b/src/main/java/org/opensearch/knn/index/IndexUtil.java
@@ -274,7 +274,7 @@ public static boolean isVersionOnOrAfterMinRequiredVersion(Version version, Stri
      * @return true if state can be shared; false otherwise
      */
     public static boolean isSharedIndexStateRequired(KNNEngine knnEngine, String modelId, long indexAddr) {
-        if (StringUtils.isEmpty(modelId) || KNNEngine.FAISS != knnEngine) {
+        if (StringUtils.isEmpty(modelId)) {
             return false;
         }
         return JNIService.isSharedIndexStateRequired(indexAddr, knnEngine);

diff --git a/src/main/java/org/opensearch/knn/index/memory/SharedIndexStateManager.java b/src/main/java/org/opensearch/knn/index/memory/SharedIndexStateManager.java
@@ -97,15 +97,15 @@ public void release(SharedIndexState sharedIndexState) {
         this.readWriteLock.writeLock().lock();
 
         try {
-            if (!sharedIndexStateCache.containsKey(sharedIndexState.getModelId())) {
+            SharedIndexStateEntry sharedIndexStateEntry;
+            if ((sharedIndexStateEntry = sharedIndexStateCache.get(sharedIndexState.getModelId())) == null) {
                 // This should not happen. Will log the error and return to prevent crash
                 log.error("Attempting to evict model from cache but it is not present: {}", sharedIndexState.getModelId());
                 this.readWriteLock.writeLock().unlock();
                 return;
             }
 
-            long refCount = sharedIndexStateCache.get(sharedIndexState.getModelId()).decRef();
-            if (refCount <= 0) {
+            if (sharedIndexStateEntry.decRef() <= 0) {
                 log.info("Evicting entry from shared index state cache for key {}", sharedIndexState.getModelId());
                 sharedIndexStateCache.remove(sharedIndexState.getModelId());
                 JNIService.freeSharedIndexState(sharedIndexState.getSharedIndexStateAddress(), sharedIndexState.getKnnEngine());

diff --git a/src/test/java/org/opensearch/knn/index/IndexUtilTests.java b/src/test/java/org/opensearch/knn/index/IndexUtilTests.java
@@ -228,12 +228,6 @@ public void testIsShareableStateContainedInIndex_whenIndexNotModelBased_thenRetu
         assertFalse(IndexUtil.isSharedIndexStateRequired(knnEngine, modelId, TEST_INDEX_ADDRESS));
     }
 
-    public void testIsShareableStateContainedInIndex_whenEngineIsNotFaiss_thenReturnFalse() {
-        String modelId = "test-model";
-        KNNEngine knnEngine = KNNEngine.NMSLIB;
-        assertFalse(IndexUtil.isSharedIndexStateRequired(knnEngine, modelId, TEST_INDEX_ADDRESS));
-    }
-
     public void testIsShareableStateContainedInIndex_whenFaissHNSWIsUsed_thenReturnFalse() {
         jniServiceMockedStatic.when(() -> JNIService.isSharedIndexStateRequired(anyLong(), any())).thenReturn(false);
         String modelId = "test-model";