[Tiered Caching] Expose new cache stats API (opensearch-project#13237)

Step 3 out of 4 --------- Signed-off-by: Peter Alfonsi <[email protected]> Co-authored-by: Peter Alfonsi <[email protected]>
peteralfonsi · Sep 3, 2024 · 88718a2 · 88718a2
1 parent 5e97082
commit 88718a2
Show file tree

Hide file tree

Showing 37 changed files with 1,355 additions and 170 deletions.
diff --git a/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamInput.java b/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamInput.java
@@ -79,6 +79,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.Date;
 import java.util.EnumSet;
 import java.util.HashMap;
@@ -89,6 +90,8 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
+import java.util.SortedMap;
+import java.util.TreeMap;
 import java.util.concurrent.TimeUnit;
 import java.util.function.IntFunction;
 
@@ -641,12 +644,47 @@ public <K, V> Map<K, V> readMap(Writeable.Reader<K> keyReader, Writeable.Reader<
             return Collections.emptyMap();
         }
         Map<K, V> map = new HashMap<>(size);
+        readIntoMap(keyReader, valueReader, map, size);
+        return map;
+    }
+
+    /**
+     * Read a serialized map into a SortedMap using the default ordering for the keys. If the result is empty it might be immutable.
+     */
+    public <K extends Comparable<K>, V> SortedMap<K, V> readOrderedMap(Writeable.Reader<K> keyReader, Writeable.Reader<V> valueReader)
+        throws IOException {
+        return readOrderedMap(keyReader, valueReader, null);
+    }
+
+    /**
+     * Read a serialized map into a SortedMap, specifying a Comparator for the keys. If the result is empty it might be immutable.
+     */
+    public <K extends Comparable<K>, V> SortedMap<K, V> readOrderedMap(
+        Writeable.Reader<K> keyReader,
+        Writeable.Reader<V> valueReader,
+        @Nullable Comparator<K> keyComparator
+    ) throws IOException {
+        int size = readArraySize();
+        if (size == 0) {
+            return Collections.emptySortedMap();
+        }
+        SortedMap<K, V> sortedMap;
+        if (keyComparator == null) {
+            sortedMap = new TreeMap<>();
+        } else {
+            sortedMap = new TreeMap<>(keyComparator);
+        }
+        readIntoMap(keyReader, valueReader, sortedMap, size);
+        return sortedMap;
+    }
+
+    private <K, V> void readIntoMap(Writeable.Reader<K> keyReader, Writeable.Reader<V> valueReader, Map<K, V> map, int size)
+        throws IOException {
         for (int i = 0; i < size; i++) {
             K key = keyReader.read(this);
             V value = valueReader.read(this);
             map.put(key, value);
         }
-        return map;
     }
 
     /**

diff --git a/...les/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java b/...les/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java
@@ -230,7 +230,7 @@ public void close() throws IOException {
     }
 
     @Override
-    public ImmutableCacheStatsHolder stats() {
+    public ImmutableCacheStatsHolder stats(String[] levels) {
         return null; // TODO: in TSC stats PR
     }
 

diff --git a/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/MockDiskCache.java b/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/MockDiskCache.java
@@ -99,6 +99,11 @@ public ImmutableCacheStatsHolder stats() {
         return null;
     }
 
+    @Override
+    public ImmutableCacheStatsHolder stats(String[] levels) {
+        return null;
+    }
+
     @Override
     public void close() {
 

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
@@ -164,7 +164,7 @@ private EhcacheDiskCache(Builder<K, V> builder) {
         this.cache = buildCache(Duration.ofMillis(expireAfterAccess.getMillis()), builder);
         List<String> dimensionNames = Objects.requireNonNull(builder.dimensionNames, "Dimension names can't be null");
         // If this cache is being used, FeatureFlags.PLUGGABLE_CACHE is already on, so we can always use the DefaultCacheStatsHolder.
-        this.cacheStatsHolder = new DefaultCacheStatsHolder(dimensionNames);
+        this.cacheStatsHolder = new DefaultCacheStatsHolder(dimensionNames, EhcacheDiskCacheFactory.EHCACHE_DISK_CACHE_NAME);
     }
 
     @SuppressWarnings({ "rawtypes" })
@@ -446,12 +446,13 @@ public void close() {
     }
 
     /**
-     * Relevant stats for this cache.
-     * @return CacheStats
+     * Relevant stats for this cache, aggregated by levels.
+     * @param levels The levels to aggregate by.
+     * @return ImmutableCacheStatsHolder
      */
     @Override
-    public ImmutableCacheStatsHolder stats() {
-        return cacheStatsHolder.getImmutableCacheStatsHolder();
+    public ImmutableCacheStatsHolder stats(String[] levels) {
+        return cacheStatsHolder.getImmutableCacheStatsHolder(levels);
     }
 
     /**
@@ -510,15 +511,15 @@ private long getNewValuePairSize(CacheEvent<? extends ICacheKey<K>, ? extends By
         public void onEvent(CacheEvent<? extends ICacheKey<K>, ? extends ByteArrayWrapper> event) {
             switch (event.getType()) {
                 case CREATED:
-                    cacheStatsHolder.incrementEntries(event.getKey().dimensions);
+                    cacheStatsHolder.incrementItems(event.getKey().dimensions);
                     cacheStatsHolder.incrementSizeInBytes(event.getKey().dimensions, getNewValuePairSize(event));
                     assert event.getOldValue() == null;
                     break;
                 case EVICTED:
                     this.removalListener.onRemoval(
                         new RemovalNotification<>(event.getKey(), deserializeValue(event.getOldValue()), RemovalReason.EVICTED)
                     );
-                    cacheStatsHolder.decrementEntries(event.getKey().dimensions);
+                    cacheStatsHolder.decrementItems(event.getKey().dimensions);
                     cacheStatsHolder.decrementSizeInBytes(event.getKey().dimensions, getOldValuePairSize(event));
                     cacheStatsHolder.incrementEvictions(event.getKey().dimensions);
                     assert event.getNewValue() == null;
@@ -527,15 +528,15 @@ public void onEvent(CacheEvent<? extends ICacheKey<K>, ? extends ByteArrayWrappe
                     this.removalListener.onRemoval(
                         new RemovalNotification<>(event.getKey(), deserializeValue(event.getOldValue()), RemovalReason.EXPLICIT)
                     );
-                    cacheStatsHolder.decrementEntries(event.getKey().dimensions);
+                    cacheStatsHolder.decrementItems(event.getKey().dimensions);
                     cacheStatsHolder.decrementSizeInBytes(event.getKey().dimensions, getOldValuePairSize(event));
                     assert event.getNewValue() == null;
                     break;
                 case EXPIRED:
                     this.removalListener.onRemoval(
                         new RemovalNotification<>(event.getKey(), deserializeValue(event.getOldValue()), RemovalReason.INVALIDATED)
                     );
-                    cacheStatsHolder.decrementEntries(event.getKey().dimensions);
+                    cacheStatsHolder.decrementItems(event.getKey().dimensions);
                     cacheStatsHolder.decrementSizeInBytes(event.getKey().dimensions, getOldValuePairSize(event));
                     assert event.getNewValue() == null;
                     break;

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
@@ -93,7 +93,7 @@ public void testBasicGetAndPut() throws IOException {
                 String value = ehcacheTest.get(getICacheKey(entry.getKey()));
                 assertEquals(entry.getValue(), value);
             }
-            assertEquals(randomKeys, ehcacheTest.stats().getTotalEntries());
+            assertEquals(randomKeys, ehcacheTest.stats().getTotalItems());
             assertEquals(randomKeys, ehcacheTest.stats().getTotalHits());
             assertEquals(expectedSize, ehcacheTest.stats().getTotalSizeInBytes());
             assertEquals(randomKeys, ehcacheTest.count());
@@ -217,7 +217,7 @@ public void testConcurrentPut() throws Exception {
                 assertEquals(entry.getValue(), value);
             }
             assertEquals(randomKeys, ehcacheTest.count());
-            assertEquals(randomKeys, ehcacheTest.stats().getTotalEntries());
+            assertEquals(randomKeys, ehcacheTest.stats().getTotalItems());
             ehcacheTest.close();
         }
     }
@@ -416,7 +416,7 @@ public String load(ICacheKey<String> key) {
             assertEquals(1, numberOfTimesValueLoaded);
             assertEquals(0, ((EhcacheDiskCache) ehcacheTest).getCompletableFutureMap().size());
             assertEquals(1, ehcacheTest.stats().getTotalMisses());
-            assertEquals(1, ehcacheTest.stats().getTotalEntries());
+            assertEquals(1, ehcacheTest.stats().getTotalItems());
             assertEquals(numberOfRequest - 1, ehcacheTest.stats().getTotalHits());
             assertEquals(1, ehcacheTest.count());
             ehcacheTest.close();