Implemented/tested counter+removal list setup to allow more removals

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

Author: Peter Alfonsi

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

@@ -34,6 +34,8 @@
 
 import org.opensearch.common.metrics.CounterMetric;
 
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
@@ -73,7 +75,8 @@ public int getValue() {
     protected final int modulo;
     KeyStoreStats stats;
     protected RoaringBitmap rbm;
-    private HashSet<Integer> collidedInts;
+    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();
@@ -89,7 +92,8 @@ public RBMIntKeyLookupStore(KeystoreModuloValue moduloValue, long memSizeCapInBy
         sizeEstimator = new RBMSizeEstimator(modulo);
         this.stats = new KeyStoreStats(memSizeCapInBytes, calculateMaxNumEntries(memSizeCapInBytes));
         this.rbm = new RoaringBitmap();
-        collidedInts = new HashSet<>();
+        this.collidedIntCounters = new HashMap<>();
+        this.removalSets = new HashMap<>();
     }
 
     protected int calculateMaxNumEntries(long memSizeCapInBytes) {
@@ -105,7 +109,14 @@ private final int transform(int value) {
 
     private void handleCollisions(int transformedValue) {
         stats.numCollisions.inc();
-        collidedInts.add(transformedValue);
+        CounterMetric numCollisions = collidedIntCounters.get(transformedValue);
+        if (numCollisions == null) { // First time the transformedValue has had a collision
+            numCollisions = new CounterMetric();
+            numCollisions.inc(2);
+            collidedIntCounters.put(transformedValue, numCollisions); // Initialize the number of colliding keys to 2
+        } else {
+            numCollisions.inc();
+        }
     }
 
     @Override
@@ -130,6 +141,16 @@ public boolean add(Integer value) throws Exception {
                 stats.size.inc();
                 return true;
             }
+            // If the value is already pending removal, take it out of the removalList
+            HashSet<Integer> removalSet = removalSets.get(transformedValue);
+            if (removalSet != null) {
+                removalSet.remove(value);
+                // Don't increment the counter - this is handled by handleCollisions() later
+                if (removalSet.isEmpty()) {
+                    removalSets.remove(transformedValue);
+                }
+            }
+
             handleCollisions(transformedValue);
             return false;
         } finally {
@@ -159,6 +180,13 @@ public Integer getInternalRepresentation(Integer value) {
         return Integer.valueOf(transform(value));
     }
 
+    /**
+     * Attempts to remove a value from the keystore. WARNING: Removing keys which have not been added to the keystore
+     * may cause undefined behavior, including future false negatives!!
+     * @param value The value to attempt to remove.
+     * @return true if the value was removed, false otherwise
+     * @throws Exception
+     */
     @Override
     public boolean remove(Integer value) throws Exception {
         if (value == null) {
@@ -170,24 +198,54 @@ public boolean remove(Integer value) throws Exception {
             if (!rbm.contains(transformedValue)) { // saves additional transform() call
                 return false;
             }
+            // move below into write lock
             stats.numRemovalAttempts.inc();
-            if (collidedInts.contains(transformedValue)) {
-                return false;
-            }
         } finally {
             readLock.unlock();
         }
         writeLock.lock();
         try {
-            rbm.remove(transformedValue);
-            stats.size.dec();
-            stats.numSuccessfulRemovals.inc();
+            CounterMetric numCollisions = collidedIntCounters.get(transformedValue);
+            if (numCollisions != null) {
+                // This transformed value has had a collision before
+                HashSet<Integer> removalSet = removalSets.get(transformedValue);
+                if (removalSet == null) {
+                    // First time a removal has been attempted for this transformed value
+                    HashSet<Integer> newRemovalSet = new HashSet<>();
+                    newRemovalSet.add(value); // Add the key value, not the transformed value, to the list of attempted removals for this transformedValue
+                    removalSets.put(transformedValue, newRemovalSet);
+                    numCollisions.dec();
+                } else {
+                    if (removalSet.contains(value)) {
+                        return false; // We have already attempted to remove this value. Do nothing
+                    }
+                    removalSet.add(value);
+                    numCollisions.dec();
+                    // If numCollisions has reached zero, we can safely remove all values in removalList
+                    if (numCollisions.count() == 0) {
+                        removeFromRBM(transformedValue);
+                        collidedIntCounters.remove(transformedValue);
+                        removalSets.remove(transformedValue);
+                        return true;
+                    }
+                }
+                return false;
+            }
+            // Otherwise, there's not been a collision for this transformedValue, so we can safely remove
+            removeFromRBM(transformedValue);
             return true;
         } finally {
             writeLock.unlock();
         }
     }
 
+    // Helper fn for remove()
+    private void removeFromRBM(int transformedValue) {
+        rbm.remove(transformedValue);
+        stats.size.dec();
+        stats.numSuccessfulRemovals.inc();
+    }
+
     @Override
     public int getSize() {
         readLock.lock();
@@ -218,7 +276,7 @@ public boolean isCollision(Integer value1, Integer value2) {
 
     @Override
     public long getMemorySizeInBytes() {
-        return sizeEstimator.getSizeInBytes((int) stats.size.count()) + RBMSizeEstimator.getHashsetMemSizeInBytes(collidedInts.size());
+        return sizeEstimator.getSizeInBytes((int) stats.size.count()); // + RBMSizeEstimator.getHashsetMemSizeInBytes(collidedInts.size());
     }
 
     @Override
@@ -234,7 +292,8 @@ public boolean isFull() {
     @Override
     public void regenerateStore(Integer[] newValues) throws Exception {
         rbm.clear();
-        collidedInts = new HashSet<>();
+        collidedIntCounters = new HashMap<>();
+        removalSets = new HashMap<>();
         stats.size = new CounterMetric();
         stats.numAddAttempts = new CounterMetric();
         stats.numCollisions = new CounterMetric();
@@ -265,6 +324,14 @@ public boolean valueHasHadCollision(Integer value) {
         if (value == null) {
             return false;
         }
-        return collidedInts.contains(transform(value));
+        return collidedIntCounters.containsKey(transform(value));
+    }
+
+    CounterMetric getNumCollisionsForValue(int value) { // package private for testing
+        return collidedIntCounters.get(transform(value));
+    }
+
+    HashSet<Integer> getRemovalSetForValue(int value) {
+        return removalSets.get(transform(value));
     }
 }

server/src/test/java/org/opensearch/common/cache/tier/keystore/RBMIntKeyLookupStoreTests.java

@@ -34,9 +34,11 @@
 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 java.util.ArrayList;
+import java.util.HashSet;
 import java.util.Random;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -277,12 +279,10 @@ public void testNullInputs() throws Exception {
     }
 
     public void testMemoryCapValueInitialization() {
-        // double[] logModulos = new double[] { 0.0, 31.2, 30, 29, 28, 13 };
-        // 0, 31, 29, 28, 26
         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 }; // check the numbers closer later
+        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++) {
@@ -295,4 +295,106 @@ public void testMemoryCapValueInitialization() {
         }
 
     }
+
+    public void testRemovalLogic() throws Exception {
+        RBMIntKeyLookupStore.KeystoreModuloValue moduloValue = RBMIntKeyLookupStore.KeystoreModuloValue.TWO_TO_TWENTY_SIX;
+        int modulo = moduloValue.getValue();
+        RBMIntKeyLookupStore kls = new RBMIntKeyLookupStore(moduloValue, 0L);
+
+        // Test standard sequence: add K1, K2, K3 which all transform to C, then:
+        // Remove K3
+        // Remove K2, re-add it, re-remove it twice (duplicate should do nothing)
+        // Remove K1, which should finally actually remove everything
+        int c = -42;
+        int k1 = c + modulo;
+        int k2 = c + 2 * modulo;
+        int k3 = c + 3 * modulo;
+        kls.add(k1);
+        assertTrue(kls.contains(k1));
+        assertTrue(kls.contains(k3));
+        kls.add(k2);
+        CounterMetric numCollisions = kls.getNumCollisionsForValue(k2);
+        assertNotNull(numCollisions);
+        assertEquals(2, numCollisions.count());
+        kls.add(k3);
+        assertEquals(3, numCollisions.count());
+        assertEquals(1, kls.getSize());
+
+        boolean removed = kls.remove(k3);
+        assertFalse(removed);
+        HashSet<Integer> removalSet = kls.getRemovalSetForValue(k3);
+        assertEquals(1, removalSet.size());
+        assertTrue(removalSet.contains(k3));
+        assertEquals(2, numCollisions.count());
+        assertEquals(1, kls.getSize());
+
+        removed = kls.remove(k2);
+        assertFalse(removed);
+        assertEquals(2, removalSet.size());
+        assertTrue(removalSet.contains(k2));
+        assertEquals(1, numCollisions.count());
+        assertEquals(1, kls.getSize());
+
+        kls.add(k2);
+        assertEquals(1, removalSet.size());
+        assertFalse(removalSet.contains(k2));
+        assertEquals(2, numCollisions.count());
+        assertEquals(1, kls.getSize());
+
+        removed = kls.remove(k2);
+        assertFalse(removed);
+        assertEquals(2, removalSet.size());
+        assertTrue(removalSet.contains(k2));
+        assertEquals(1, numCollisions.count());
+        assertEquals(1, kls.getSize());
+
+        removed = kls.remove(k2);
+        assertFalse(removed);
+        assertEquals(2, removalSet.size());
+        assertTrue(removalSet.contains(k2));
+        assertEquals(1, numCollisions.count());
+        assertEquals(1, kls.getSize());
+
+        removed = kls.remove(k1);
+        assertTrue(removed);
+        assertNull(kls.getRemovalSetForValue(k1));
+        assertNull(kls.getNumCollisionsForValue(k1));
+        assertEquals(0, kls.getSize());
+    }
+
+    public void testRemovalLogicWithHashCollision() throws Exception {
+        RBMIntKeyLookupStore.KeystoreModuloValue moduloValue = RBMIntKeyLookupStore.KeystoreModuloValue.TWO_TO_TWENTY_SIX;
+        int modulo = moduloValue.getValue();
+        RBMIntKeyLookupStore kls = new RBMIntKeyLookupStore(moduloValue, 0L);
+
+        // Test adding K1 twice (maybe two keys hash to K1), then removing it twice.
+        // We expect it to be unable to remove the last one, but there should be no false negatives.
+        int c = 77;
+        int k1 = c + modulo;
+        int k2 = c + 2 * modulo;
+        kls.add(k1);
+        kls.add(k2);
+        CounterMetric numCollisions = kls.getNumCollisionsForValue(k1);
+        assertEquals(2, numCollisions.count());
+        kls.add(k1);
+        assertEquals(3, numCollisions.count());
+
+        boolean removed = kls.remove(k1);
+        assertFalse(removed);
+        HashSet<Integer> removalSet = kls.getRemovalSetForValue(k1);
+        assertTrue(removalSet.contains(k1));
+        assertEquals(2, numCollisions.count());
+
+        removed = kls.remove(k2);
+        assertFalse(removed);
+        assertTrue(removalSet.contains(k2));
+        assertEquals(1, numCollisions.count());
+
+        removed = kls.remove(k1);
+        assertFalse(removed);
+        assertTrue(removalSet.contains(k1));
+        assertEquals(1, numCollisions.count());
+        assertTrue(kls.contains(k1));
+        assertTrue(kls.contains(k2)); 
+    }
 }