From b0e7ee08a5052ecd8d892ba3c7eea209cffceeb2 Mon Sep 17 00:00:00 2001
From: owenhalpert <ohalpert@gmail.com>
Date: Wed, 4 Dec 2024 14:10:31 -0800
Subject: [PATCH] Update CacheMaintainer to be a generic ScheduledExecutor
 class

Signed-off-by: owenhalpert <ohalpert@gmail.com>
---
 .../knn/index/memory/NativeMemoryCacheManager.java |  2 +-
 .../knn/index/{ => util}/ScheduledExecutor.java    | 14 ++++----------
 .../quantizationState/QuantizationStateCache.java  |  6 ++++--
 .../opensearch/knn/index/CacheMaintainerTests.java |  1 +
 4 files changed, 10 insertions(+), 13 deletions(-)
 rename src/main/java/org/opensearch/knn/index/{ => util}/ScheduledExecutor.java (73%)

diff --git a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManager.java b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManager.java
index f0d932cd5..4babde422 100644
--- 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/src/main/java/org/opensearch/knn/index/ScheduledExecutor.java b/src/main/java/org/opensearch/knn/index/util/ScheduledExecutor.java
similarity index 73%
rename from src/main/java/org/opensearch/knn/index/ScheduledExecutor.java
rename to src/main/java/org/opensearch/knn/index/util/ScheduledExecutor.java
index 74c1c4a8b..82f3f2fd4 100644
--- a/src/main/java/org/opensearch/knn/index/ScheduledExecutor.java
+++ b/src/main/java/org/opensearch/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/src/main/java/org/opensearch/knn/quantization/models/quantizationState/QuantizationStateCache.java b/src/main/java/org/opensearch/knn/quantization/models/quantizationState/QuantizationStateCache.java
index 06a4dbbbf..c0e61feb3 100644
--- a/src/main/java/org/opensearch/knn/quantization/models/quantizationState/QuantizationStateCache.java
+++ b/src/main/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
index 0281593e6..f4f3291bf 100644
--- 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;