Add unit tests for maintenance scheduling

Signed-off-by: owenhalpert <[email protected]>

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

@@ -16,6 +16,7 @@
 import com.google.common.cache.CacheStats;
 import com.google.common.cache.RemovalCause;
 import com.google.common.cache.RemovalNotification;
+import lombok.Getter;
 import lombok.Setter;
 import org.apache.commons.lang.Validate;
 import org.apache.logging.log4j.LogManager;
@@ -58,6 +59,7 @@ public class NativeMemoryCacheManager implements Closeable {
     private final ExecutorService executor;
     private AtomicBoolean cacheCapacityReached;
     private long maxWeight;
+    @Getter
     private Cancellable maintenanceTask;
 
     NativeMemoryCacheManager() {

src/main/java/org/opensearch/knn/quantization/models/quantizationState/QuantizationStateCache.java

@@ -41,6 +41,7 @@ public class QuantizationStateCache implements Closeable {
     private long maxCacheSizeInKB;
     @Getter
     private Instant evictedDueToSizeAt;
+    @Getter
     private Cancellable maintenanceTask;
 
     @VisibleForTesting

src/test/java/org/opensearch/knn/index/memory/NativeMemoryCacheManagerTests.java

@@ -22,6 +22,7 @@
 import org.opensearch.knn.plugin.KNNPlugin;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.test.OpenSearchSingleNodeTestCase;
+import org.opensearch.threadpool.Scheduler.Cancellable;
 import org.opensearch.threadpool.ThreadPool;
 
 import java.io.IOException;
@@ -70,6 +71,8 @@ protected Collection<Class<? extends Plugin>> getPlugins() {
 
     public void testRebuildCache() throws ExecutionException, InterruptedException {
         NativeMemoryCacheManager nativeMemoryCacheManager = new NativeMemoryCacheManager();
+        Cancellable task1 = nativeMemoryCacheManager.getMaintenanceTask();
+        assertNotNull(task1);
 
         // Put entry in cache and check that the weight matches
         int size = 10;
@@ -84,6 +87,9 @@ public void testRebuildCache() throws ExecutionException, InterruptedException {
         // Sleep for a second or two so that the executor can invalidate all entries
         Thread.sleep(2000);
 
+        assertTrue(task1.isCancelled());
+        assertNotNull(nativeMemoryCacheManager.getMaintenanceTask());
+
         assertEquals(0, nativeMemoryCacheManager.getCacheSizeInKilobytes());
         nativeMemoryCacheManager.close();
     }
@@ -484,6 +490,16 @@ public void testGetIndicesCacheStats() throws IOException, ExecutionException {
         nativeMemoryCacheManager.close();
     }
 
+    public void testMaintenanceScheduled() {
+        NativeMemoryCacheManager nativeMemoryCacheManager = new NativeMemoryCacheManager();
+        Cancellable maintenanceTask = nativeMemoryCacheManager.getMaintenanceTask();
+
+        assertNotNull(maintenanceTask);
+
+        nativeMemoryCacheManager.close();
+        assertTrue(maintenanceTask.isCancelled());
+    }
+
     private static class TestNativeMemoryAllocation implements NativeMemoryAllocation {
 
         int size;

src/test/java/org/opensearch/knn/quantization/models/quantizationState/QuantizationStateCacheTests.java

@@ -17,6 +17,7 @@
 import org.opensearch.knn.KNNTestCase;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.quantization.models.quantizationParams.ScalarQuantizationParams;
+import org.opensearch.threadpool.Scheduler;
 import org.opensearch.threadpool.ThreadPool;
 
 import java.io.IOException;
@@ -467,4 +468,27 @@ public void testCacheEvictionDueToSize() throws IOException {
         cache.close();
         assertNotNull(cache.getEvictedDueToSizeAt());
     }
+
+    public void testMaintenanceScheduled() throws Exception {
+        QuantizationStateCache quantizationStateCache= new QuantizationStateCache();
+        Scheduler.Cancellable maintenanceTask = quantizationStateCache.getMaintenanceTask();
+
+        assertNotNull(maintenanceTask);
+
+        quantizationStateCache.close();
+        assertTrue(maintenanceTask.isCancelled());
+    }
+
+    public void testMaintenanceWithRebuild() throws Exception {
+        QuantizationStateCache quantizationStateCache= new QuantizationStateCache();
+        Scheduler.Cancellable task1 = quantizationStateCache.getMaintenanceTask();
+        assertNotNull(task1);
+
+        quantizationStateCache.rebuildCache();
+
+        Scheduler.Cancellable task2 = quantizationStateCache.getMaintenanceTask();
+        assertTrue(task1.isCancelled());
+        assertNotNull(task2);
+        quantizationStateCache.close();
+    }
 }