Simplified RBM size estimator

Signed-off-by: Peter Alfonsi <[email protected]>

Author: Peter Alfonsi

server/src/main/java/org/opensearch/common/cache/tier/keystore/KeyLookupStore.java

@@ -103,7 +103,7 @@ public interface KeyLookupStore<T> {
 
     /**
      * Returns an estimate of the store's memory usage.
-     * @return The memory usage, in MB
+     * @return The memory usage
      */
     long getMemorySizeInBytes();
 

server/src/main/java/org/opensearch/common/cache/tier/keystore/KeyStoreStats.java

@@ -22,17 +22,15 @@ public class KeyStoreStats {
     protected CounterMetric numAddAttempts;
     protected CounterMetric numCollisions;
     protected boolean guaranteesNoFalseNegatives;
-    protected final int maxNumEntries;
     protected AtomicBoolean atCapacity;
     protected CounterMetric numRemovalAttempts;
     protected CounterMetric numSuccessfulRemovals;
 
-    protected KeyStoreStats(long memSizeCapInBytes, int maxNumEntries) {
+    protected KeyStoreStats(long memSizeCapInBytes) {
         this.size = new CounterMetric();
         this.numAddAttempts = new CounterMetric();
         this.numCollisions = new CounterMetric();
         this.memSizeCapInBytes = memSizeCapInBytes;
-        this.maxNumEntries = maxNumEntries;
         this.atCapacity = new AtomicBoolean(false);
         this.numRemovalAttempts = new CounterMetric();
         this.numSuccessfulRemovals = new CounterMetric();

server/src/main/java/org/opensearch/common/cache/tier/keystore/RBMIntKeyLookupStore.java

@@ -34,7 +34,6 @@
 
 import org.opensearch.common.metrics.CounterMetric;
 
-import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.concurrent.locks.Lock;
@@ -77,10 +76,14 @@ public int getValue() {
     protected RoaringBitmap rbm;
     private HashMap<Integer, CounterMetric> collidedIntCounters;
     private HashMap<Integer, HashSet<Integer>> removalSets;
-    protected RBMSizeEstimator sizeEstimator;
     protected final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
     protected final Lock readLock = lock.readLock();
     protected final Lock writeLock = lock.writeLock();
+    private long mostRecentByteEstimate;
+    private final int REFRESH_SIZE_EST_INTERVAL = 10000;
+    // Refresh size estimate every X new elements. Refreshes use the RBM's internal size estimator, which takes ~0.01 ms,
+    // so we don't want to do it on every get(), and it doesn't matter much if there are +- 10000 keys in this store
+    // in terms of storage impact
 
     // Default constructor sets modulo = 2^28
     public RBMIntKeyLookupStore(long memSizeCapInBytes) {
@@ -89,18 +92,11 @@ public RBMIntKeyLookupStore(long memSizeCapInBytes) {
 
     public RBMIntKeyLookupStore(KeystoreModuloValue moduloValue, long memSizeCapInBytes) {
         this.modulo = moduloValue.getValue();
-        sizeEstimator = new RBMSizeEstimator(modulo);
-        this.stats = new KeyStoreStats(memSizeCapInBytes, calculateMaxNumEntries(memSizeCapInBytes));
+        this.stats = new KeyStoreStats(memSizeCapInBytes);
         this.rbm = new RoaringBitmap();
         this.collidedIntCounters = new HashMap<>();
         this.removalSets = new HashMap<>();
-    }
-
-    protected int calculateMaxNumEntries(long memSizeCapInBytes) {
-        if (memSizeCapInBytes == 0) {
-            return Integer.MAX_VALUE;
-        }
-        return sizeEstimator.getNumEntriesFromSizeInBytes(memSizeCapInBytes);
+        this.mostRecentByteEstimate = 0L;
     }
 
     private final int transform(int value) {
@@ -125,7 +121,11 @@ public boolean add(Integer value) {
             return false;
         }
         stats.numAddAttempts.inc();
-        if (stats.size.count() == stats.maxNumEntries) {
+
+        if (getSize() % REFRESH_SIZE_EST_INTERVAL == 0) {
+            mostRecentByteEstimate = getMemorySizeInBytes();
+        }
+        if (getMemorySizeCapInBytes() > 0 && mostRecentByteEstimate > getMemorySizeCapInBytes()) {
             stats.atCapacity.set(true);
             return false;
         }
@@ -273,9 +273,24 @@ public boolean isCollision(Integer value1, Integer value2) {
         return transform(value1) == transform(value2);
     }
 
+     static double getRBMSizeMultiplier(int numEntries, int modulo) {
+        double x = Math.log10((double) numEntries / modulo);
+        if (x < -5) {
+            return 7.0;
+        }
+        if (x < -2.75) {
+            return -2.5 * x - 5.5;
+        }
+        if (x <= 0) {
+            return -3.0 / 22.0 * x + 1;
+        }
+        return 1;
+    }
+
     @Override
     public long getMemorySizeInBytes() {
-        return sizeEstimator.getSizeInBytes((int) stats.size.count()); // + RBMSizeEstimator.getHashsetMemSizeInBytes(collidedInts.size());
+        double multiplier = getRBMSizeMultiplier((int) stats.size.count(), modulo);
+        return (long) (rbm.getSizeInBytes() * multiplier);
     }
 
     @Override

server/src/main/java/org/opensearch/common/cache/tier/keystore/RBMSizeEstimator.java

@@ -32,10 +32,9 @@
 package org.opensearch.common.cache.tier.keystore;
 
 import org.opensearch.common.Randomness;
-import org.opensearch.common.cache.tier.keystore.RBMIntKeyLookupStore;
-import org.opensearch.common.cache.tier.keystore.RBMSizeEstimator;
 import org.opensearch.common.metrics.CounterMetric;
 import org.opensearch.test.OpenSearchTestCase;
+import org.roaringbitmap.RoaringBitmap;
 
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -45,8 +44,10 @@
 import java.util.concurrent.ThreadPoolExecutor;
 
 public class RBMIntKeyLookupStoreTests extends OpenSearchTestCase {
+
+    final int BYTES_IN_MB = 1048576;
     public void testInit() {
-        long memCap = 100 * RBMSizeEstimator.BYTES_IN_MB;
+        long memCap = 100 * BYTES_IN_MB;
         RBMIntKeyLookupStore kls = new RBMIntKeyLookupStore(memCap);
         assertEquals(0, kls.getSize());
         assertEquals(RBMIntKeyLookupStore.KeystoreModuloValue.TWO_TO_TWENTY_EIGHT.getValue(), kls.modulo);
@@ -129,16 +130,35 @@ public void testAddingDuplicates() throws Exception {
     }
 
     public void testMemoryCapBlocksAdd() throws Exception {
+        // Now that we're using a modified version of rbm.getSizeInBytes(), which doesn't provide an inverse function,
+        // we have to test filling just an RBM with random test values first so that we can get the resulting memory cap limit
+        // to use with our modified size estimate.
+        // This is much noisier so the precision is lower.
+
+        // It is necessary to use randomly distributed integers for both parts of this test, as we would do with hashes in the cache,
+        // as that's what our size estimator is designed for.
+        // If we add a run of integers, our size estimator is not valid, especially for small RBMs.
+
+        int[] maxEntriesArr = new int[] { 1342000, 100000, 3000000};
+        long[] rbmReportedSizes = new long[4];
+        Random rand = Randomness.get();
+        for (int j = 0; j < maxEntriesArr.length; j++)  {
+            RoaringBitmap rbm = new RoaringBitmap();
+            for (int i = 0; i < maxEntriesArr[j]; i++) {
+                rbm.add(rand.nextInt());
+            }
+            rbmReportedSizes[j] = rbm.getSizeInBytes();
+        }
         RBMIntKeyLookupStore.KeystoreModuloValue moduloValue = RBMIntKeyLookupStore.KeystoreModuloValue.TWO_TO_TWENTY_NINE;
-        for (int maxEntries : new int[] { 2342000, 1000, 100000 }) {
-            long memSizeCapInBytes = RBMSizeEstimator.getSizeInBytesWithModuloValue(maxEntries, moduloValue);
+        for (int i = 0; i < maxEntriesArr.length; i++) {
+            double multiplier = RBMIntKeyLookupStore.getRBMSizeMultiplier(maxEntriesArr[i], moduloValue.getValue());
+            long memSizeCapInBytes = (long) (rbmReportedSizes[i] * multiplier);
+            //long memSizeCapInBytes = RBMSizeEstimator.getSizeInBytesWithModuloValue(maxEntries, moduloValue);
             RBMIntKeyLookupStore kls = new RBMIntKeyLookupStore(moduloValue, memSizeCapInBytes);
-            for (int j = 0; j < maxEntries + 1000; j++) {
-                kls.add(j);
+            for (int j = 0; j < maxEntriesArr[i] + 5000; j++) {
+                kls.add(rand.nextInt());
             }
-            assertTrue(Math.abs(maxEntries - kls.getSize()) < (double) maxEntries / 25);
-            // exact cap varies a small amount bc of floating point, especially when we use bytes instead of MB for calculations
-            // precision gets much worse when we compose the two functions, as we do here, but this wouldn't happen in an actual use case
+            assertTrue(Math.abs(maxEntriesArr[i] - kls.getSize()) < (double) maxEntriesArr[i] / 10);
         }
     }
 
@@ -205,7 +225,7 @@ public void testConcurrency() throws Exception {
     }
 
     public void testRemoveNoCollisions() throws Exception {
-        long memCap = 100L * RBMSizeEstimator.BYTES_IN_MB;
+        long memCap = 100L * BYTES_IN_MB;
         int numToAdd = 195000;
         RBMIntKeyLookupStore kls = new RBMIntKeyLookupStore(RBMIntKeyLookupStore.KeystoreModuloValue.NONE, memCap);
         // there should be no collisions for sequential positive numbers up to modulo
@@ -222,7 +242,7 @@ public void testRemoveNoCollisions() throws Exception {
 
     public void testRemoveWithCollisions() throws Exception {
         int modulo = (int) Math.pow(2, 26);
-        long memCap = 100L * RBMSizeEstimator.BYTES_IN_MB;
+        long memCap = 100L * BYTES_IN_MB;
         RBMIntKeyLookupStore kls = new RBMIntKeyLookupStore(RBMIntKeyLookupStore.KeystoreModuloValue.TWO_TO_TWENTY_SIX, memCap);
         for (int i = 0; i < 10; i++) {
             kls.add(i);
@@ -278,24 +298,6 @@ public void testNullInputs() throws Exception {
         assertEquals(4, kls.getSize());
     }
 
-    public void testMemoryCapValueInitialization() {
-        RBMIntKeyLookupStore.KeystoreModuloValue[] mods = RBMIntKeyLookupStore.KeystoreModuloValue.values();
-        double[] expectedMultipliers = new double[] { 1.2, 1.2, 1, 1, 1 };
-        double[] expectedSlopes = new double[] { 0.637, 0.637, 0.619, 0.614, 0.629 };
-        double[] expectedIntercepts = new double[] { 3.091, 3.091, 2.993, 2.905, 2.603 };
-        long memSizeCapInBytes = (long) 100.0 * RBMSizeEstimator.BYTES_IN_MB;
-        double delta = 0.01;
-        for (int i = 0; i < mods.length; i++) {
-            RBMIntKeyLookupStore.KeystoreModuloValue moduloValue = mods[i];
-            RBMIntKeyLookupStore kls = new RBMIntKeyLookupStore(moduloValue, memSizeCapInBytes);
-            assertEquals(kls.stats.memSizeCapInBytes, kls.getMemorySizeCapInBytes(), 1.0);
-            assertEquals(expectedMultipliers[i], kls.sizeEstimator.bufferMultiplier, delta);
-            assertEquals(expectedSlopes[i], kls.sizeEstimator.slope, delta);
-            assertEquals(expectedIntercepts[i], kls.sizeEstimator.intercept, delta);
-        }
-
-    }
-
     public void testRemovalLogic() throws Exception {
         RBMIntKeyLookupStore.KeystoreModuloValue moduloValue = RBMIntKeyLookupStore.KeystoreModuloValue.TWO_TO_TWENTY_SIX;
         int modulo = moduloValue.getValue();
@@ -395,6 +397,6 @@ public void testRemovalLogicWithHashCollision() throws Exception {
         assertTrue(removalSet.contains(k1));
         assertEquals(1, numCollisions.count());
         assertTrue(kls.contains(k1));
-        assertTrue(kls.contains(k2)); 
+        assertTrue(kls.contains(k2));
     }
 }