Update CacheMaintainer to be a generic ScheduledExecutor class

Signed-off-by: owenhalpert <[email protected]>
opensearch-project · Dec 19, 2024 · b0e7ee0 · b0e7ee0
1 parent 601a24a
commit b0e7ee0
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 13 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
@@ -23,7 +23,7 @@
 import org.opensearch.knn.common.exception.OutOfNativeMemoryException;
 import org.opensearch.knn.common.featureflags.KNNFeatureFlags;
 import org.opensearch.knn.index.KNNSettings;
-import org.opensearch.knn.index.ScheduledExecutor;
+import org.opensearch.knn.index.util.ScheduledExecutor;
 import org.opensearch.knn.plugin.stats.StatNames;
 
 import java.io.Closeable;

diff --git a/...ensearch/knn/index/ScheduledExecutor.java → ...rch/knn/index/util/ScheduledExecutor.java b/...ensearch/knn/index/ScheduledExecutor.java → ...rch/knn/index/util/ScheduledExecutor.java
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.knn.index;
+package org.opensearch.knn.index.util;
 
 import java.io.Closeable;
 import java.util.concurrent.Executors;
@@ -12,11 +12,10 @@
 
 /**
  * Executes a task periodically
-
  */
 public class ScheduledExecutor implements Closeable {
-    private final ScheduledExecutorService executor;
-    final Runnable task;
+    final ScheduledExecutorService executor;
+    public final Runnable task;
 
     /**
      * @param task task to be completed
@@ -25,12 +24,7 @@ public class ScheduledExecutor implements Closeable {
     public ScheduledExecutor(Runnable task, long scheduleMillis) {
         this.task = task;
         this.executor = Executors.newSingleThreadScheduledExecutor();
-        executor.scheduleAtFixedRate(
-            task,
-            0,
-            scheduleMillis,
-            TimeUnit.MILLISECONDS
-        );
+        executor.scheduleAtFixedRate(task, 0, scheduleMillis, TimeUnit.MILLISECONDS);
     }
 
     @Override

diff --git a/...java/org/opensearch/knn/quantization/models/quantizationState/QuantizationStateCache.java b/...java/org/opensearch/knn/quantization/models/quantizationState/QuantizationStateCache.java
@@ -15,7 +15,7 @@
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.knn.index.KNNSettings;
-import org.opensearch.knn.index.ScheduledExecutor;
+import org.opensearch.knn.index.util.ScheduledExecutor;
 
 import java.io.Closeable;
 import java.io.IOException;
@@ -141,6 +141,8 @@ public void clear() {
 
     @Override
     public void close() throws IOException {
-        cacheMaintainer.close();
+        if (cacheMaintainer != null) {
+            cacheMaintainer.close();
+        }
     }
 }
diff --git a/src/test/java/org/opensearch/knn/index/CacheMaintainerTests.java b/src/test/java/org/opensearch/knn/index/CacheMaintainerTests.java
@@ -8,6 +8,7 @@
 import com.google.common.cache.Cache;
 import com.google.common.cache.CacheBuilder;
 import org.junit.Test;
+import org.opensearch.knn.index.util.ScheduledExecutor;
 
 import java.util.concurrent.TimeUnit;