Fix max size for shouldCache check

Signed-off-by: Peter Alfonsi <[email protected]>
peteralfonsi · Dec 12, 2024 · c802a55 · c802a55
1 parent 96a6eed
commit c802a55
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 1 deletion.
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
@@ -426,6 +426,11 @@ public ImmutableCacheStatsHolder stats(String[] levels) {
             return null;
         }
 
+        @Override
+        public long getMaxHeapBytes() {
+            return onHeapCache.getMaxHeapBytes();
+        }
+
         @Override
         public void close() throws IOException {
             for (Map.Entry<ICache<K, V>, TierInfo> cacheEntry : caches.entrySet()) {
@@ -654,6 +659,11 @@ public ImmutableCacheStatsHolder stats(String[] levels) {
         return statsHolder.getImmutableCacheStatsHolder(levels);
     }
 
+    @Override
+    public long getMaxHeapBytes() {
+        return numberOfSegments * tieredSpilloverCacheSegments[0].getMaxHeapBytes();
+    }
+
     // Package private for testing.
     @SuppressWarnings({ "unchecked" })
     Iterable<ICacheKey<K>> getOnHeapCacheKeys() {

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
@@ -138,6 +138,11 @@ public ImmutableCacheStatsHolder stats(String[] levels) {
         return null;
     }
 
+    @Override
+    public long getMaxHeapBytes() {
+        return 0;
+    }
+
     @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
@@ -496,6 +496,11 @@ public ImmutableCacheStatsHolder stats(String[] levels) {
         return cacheStatsHolder.getImmutableCacheStatsHolder(levels);
     }
 
+    @Override
+    public long getMaxHeapBytes() {
+        return 0;
+    }
+
     /**
      * This iterator wraps ehCache iterator and only iterates over its keys.
      * @param <K> Type of key

diff --git a/server/src/main/java/org/opensearch/common/cache/ICache.java b/server/src/main/java/org/opensearch/common/cache/ICache.java
@@ -53,6 +53,9 @@ default ImmutableCacheStatsHolder stats() {
     // Return stats aggregated by the provided levels. If levels is null or an empty array, return total stats only.
     ImmutableCacheStatsHolder stats(String[] levels);
 
+    // Return the bytes in memory this cache can consume.
+    long getMaxHeapBytes();
+
     /**
      * Factory to create objects.
      */

diff --git a/server/src/main/java/org/opensearch/common/cache/store/OpenSearchOnHeapCache.java b/server/src/main/java/org/opensearch/common/cache/store/OpenSearchOnHeapCache.java
@@ -155,6 +155,11 @@ public ImmutableCacheStatsHolder stats(String[] levels) {
         return cacheStatsHolder.getImmutableCacheStatsHolder(levels);
     }
 
+    @Override
+    public long getMaxHeapBytes() {
+        return maximumWeight;
+    }
+
     @Override
     public void onRemoval(RemovalNotification<ICacheKey<K>, V> notification) {
         removalListener.onRemoval(notification);

diff --git a/server/src/main/java/org/opensearch/indices/query/PluggableQueryCache.java b/server/src/main/java/org/opensearch/indices/query/PluggableQueryCache.java
@@ -196,7 +196,8 @@ public PluggableQueryCache(
                 .build(),
             CacheType.INDICES_QUERY_CACHE
         );
-        this.maxRamBytesUsed = 1_000_000; // TODO: i think this just controls the max size going in rn. Didnt want to deal w the settings rn
+        this.maxRamBytesUsed = innerCache.getMaxHeapBytes();
+        int k = 0;
         // In theory it may be useful to test this with non-TSC internal caches for regression purposes.
         // For example, if we just do an OpenSearchOnHeapCache of the same size, how much worse does it perform?
     }