[Tiered caching] Supporting removal function on EhcacheDiskCache iter…

…ator (opensearch-project#12653) * [Tiered caching] Supporting removal function on EhcacheDiskCache iterator Signed-off-by: Sagar Upadhyaya <[email protected]> * Minor refactoring in unit test Signed-off-by: Sagar Upadhyaya <[email protected]> --------- Signed-off-by: Sagar Upadhyaya <[email protected]>
amkhar · Mar 15, 2024 · ec61cbc · ec61cbc
1 parent ccdf3ff
commit ec61cbc
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 0 deletions.
diff --git a/plugins/cache-ehcache/src/main/java/org/opensearch/cache/store/disk/EhcacheDiskCache.java b/plugins/cache-ehcache/src/main/java/org/opensearch/cache/store/disk/EhcacheDiskCache.java
@@ -407,6 +407,11 @@ public K next() {
             }
             return iterator.next().getKey();
         }
+
+        @Override
+        public void remove() {
+            iterator.remove(); // Calls underlying ehcache iterator.remove()
+        }
     }
 
     /**

diff --git a/...ns/cache-ehcache/src/test/java/org/opensearch/cache/store/disk/EhCacheDiskCacheTests.java b/...ns/cache-ehcache/src/test/java/org/opensearch/cache/store/disk/EhCacheDiskCacheTests.java
@@ -22,6 +22,7 @@
 import org.opensearch.test.OpenSearchSingleNodeTestCase;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -481,6 +482,49 @@ public String load(String key) throws Exception {
         }
     }
 
+    public void testEhcacheKeyIteratorWithRemove() throws IOException {
+        Settings settings = Settings.builder().build();
+        try (NodeEnvironment env = newNodeEnvironment(settings)) {
+            ICache<String, String> ehcacheTest = new EhcacheDiskCache.Builder<String, String>().setDiskCacheAlias("test1")
+                .setThreadPoolAlias("ehcacheTest")
+                .setStoragePath(env.nodePaths()[0].indicesPath.toString() + "/request_cache")
+                .setIsEventListenerModeSync(true)
+                .setKeyType(String.class)
+                .setValueType(String.class)
+                .setCacheType(CacheType.INDICES_REQUEST_CACHE)
+                .setSettings(settings)
+                .setExpireAfterAccess(TimeValue.MAX_VALUE)
+                .setMaximumWeightInBytes(CACHE_SIZE_IN_BYTES)
+                .setRemovalListener(new MockRemovalListener<>())
+                .build();
+
+            int randomKeys = randomIntBetween(2, 100);
+            for (int i = 0; i < randomKeys; i++) {
+                ehcacheTest.put(UUID.randomUUID().toString(), UUID.randomUUID().toString());
+            }
+            long originalSize = ehcacheTest.count();
+            assertEquals(randomKeys, originalSize);
+
+            // Now try removing subset of keys and verify
+            List<String> removedKeyList = new ArrayList<>();
+            for (Iterator<String> iterator = ehcacheTest.keys().iterator(); iterator.hasNext();) {
+                String key = iterator.next();
+                if (randomBoolean()) {
+                    removedKeyList.add(key);
+                    iterator.remove();
+                }
+            }
+            // Verify the removed key doesn't exist anymore.
+            for (String ehcacheKey : removedKeyList) {
+                assertNull(ehcacheTest.get(ehcacheKey));
+            }
+            // Verify ehcache entry size again.
+            assertEquals(originalSize - removedKeyList.size(), ehcacheTest.count());
+            ehcacheTest.close();
+        }
+
+    }
+
     private static String generateRandomString(int length) {
         String characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
         StringBuilder randomString = new StringBuilder(length);