Addressed Sagar's comments besides new counter/arraylist setup for removing keys

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

Author: Peter Alfonsi

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

@@ -30,7 +30,7 @@
  * GitHub history for details.
  */
 
-package org.opensearch.indices;
+package org.opensearch.common.cache.tier.keystore;
 
 /**
  * An interface for objects that hold an in-memory record of hashes of keys in the disk cache.
@@ -85,15 +85,14 @@ public interface KeyLookupStore<T> {
      * Returns the number of times add() has been run, including unsuccessful attempts.
      * @return The number of adding attempts.
      */
-    int getTotalAdds();
+    int getAddAttempts();
 
     /**
      * Returns the number of times add() has returned false due to a collision.
      * @return The number of collisions.
      */
     int getCollisions();
 
-
     /**
      * Checks if two values would collide after being transformed by this store's transformation.
      * @param value1 The first value to compare.

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

@@ -0,0 +1,40 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.common.cache.tier.keystore;
+
+import org.opensearch.common.metrics.CounterMetric;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * A stats holder for use in KeyLookupStore implementations.
+ * Getters should be exposed by the KeyLookupStore which uses it.
+ */
+public class KeyStoreStats {
+    protected CounterMetric size;
+    protected long memSizeCapInBytes;
+    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) {
+        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/indices/RBMIntKeyLookupStore.java renamed to server/src/main/java/org/opensearch/common/cache/tier/keystore/RBMIntKeyLookupStore.java

@@ -30,15 +30,16 @@
  * GitHub history for details.
  */
 
-package org.opensearch.indices;
+package org.opensearch.common.cache.tier.keystore;
 
 import org.opensearch.common.metrics.CounterMetric;
-import org.roaringbitmap.RoaringBitmap;
 
 import java.util.HashSet;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
+import org.roaringbitmap.RoaringBitmap;
+
 /**
  * This class implements KeyLookupStore<Integer> using a roaring bitmap with a modulo applied to values.
  * The modulo increases the density of values, which makes RBMs more memory-efficient. The recommended modulo is ~2^28.
@@ -48,39 +49,43 @@
  * The store estimates its memory footprint and will stop adding more values once it reaches its memory cap.
  */
 public class RBMIntKeyLookupStore implements KeyLookupStore<Integer> {
-    protected final int modulo;
-    protected class KeyStoreStats {
-        protected int size;
-        protected long memSizeCapInBytes;
-        protected CounterMetric numAddAttempts;
-        protected CounterMetric numCollisions;
-        protected boolean guaranteesNoFalseNegatives;
-        protected int maxNumEntries;
-        protected boolean atCapacity;
-        protected CounterMetric numRemovalAttempts;
-        protected CounterMetric numSuccessfulRemovals;
-        protected KeyStoreStats(long memSizeCapInBytes, int maxNumEntries) {
-            this.size = 0;
-            this.numAddAttempts = new CounterMetric();
-            this.numCollisions = new CounterMetric();
-            this.memSizeCapInBytes = memSizeCapInBytes;
-            this.maxNumEntries = maxNumEntries;
-            this.atCapacity = false;
-            this.numRemovalAttempts = new CounterMetric();
-            this.numSuccessfulRemovals = new CounterMetric();
+    /**
+     * An enum representing modulo values for use in the keystore
+     */
+    public enum KeystoreModuloValue {
+        NONE(0), // No modulo applied
+        TWO_TO_THIRTY_ONE((int) Math.pow(2, 31)),
+        TWO_TO_TWENTY_NINE((int) Math.pow(2, 29)), // recommended value
+        TWO_TO_TWENTY_EIGHT((int) Math.pow(2, 28)),
+        TWO_TO_TWENTY_SIX((int) Math.pow(2, 26));
+
+        private final int value;
+
+        private KeystoreModuloValue(int value) {
+            this.value = value;
+        }
+
+        public int getValue() {
+            return this.value;
         }
     }
 
-    protected KeyStoreStats stats;
+    protected final int modulo;
+    KeyStoreStats stats;
     protected RoaringBitmap rbm;
     private HashSet<Integer> collidedInts;
     protected RBMSizeEstimator sizeEstimator;
     protected final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
     protected final Lock readLock = lock.readLock();
     protected final Lock writeLock = lock.writeLock();
 
-    RBMIntKeyLookupStore(int modulo, long memSizeCapInBytes) {
-        this.modulo = modulo;
+    // Default constructor sets modulo = 2^28
+    public RBMIntKeyLookupStore(long memSizeCapInBytes) {
+        this(KeystoreModuloValue.TWO_TO_TWENTY_EIGHT, memSizeCapInBytes);
+    }
+
+    public RBMIntKeyLookupStore(KeystoreModuloValue moduloValue, long memSizeCapInBytes) {
+        this.modulo = moduloValue.getValue();
         sizeEstimator = new RBMSizeEstimator(modulo);
         this.stats = new KeyStoreStats(memSizeCapInBytes, calculateMaxNumEntries(memSizeCapInBytes));
         this.rbm = new RoaringBitmap();
@@ -94,11 +99,11 @@ protected int calculateMaxNumEntries(long memSizeCapInBytes) {
         return sizeEstimator.getNumEntriesFromSizeInBytes(memSizeCapInBytes);
     }
 
-    protected final int transform(int value) {
+    private final int transform(int value) {
         return modulo == 0 ? value : value % modulo;
     }
 
-    protected void handleCollisions(int transformedValue) {
+    private void handleCollisions(int transformedValue) {
         stats.numCollisions.inc();
         collidedInts.add(transformedValue);
     }
@@ -108,18 +113,21 @@ public boolean add(Integer value) throws Exception {
         if (value == null) {
             return false;
         }
-        writeLock.lock();
         stats.numAddAttempts.inc();
+        if (stats.size.count() == stats.maxNumEntries) {
+            stats.atCapacity.set(true);
+            return false;
+        }
+        int transformedValue = transform(value);
+
+        writeLock.lock();
         try {
-            if (stats.size == stats.maxNumEntries) {
-                stats.atCapacity = true;
-                return false;
-            }
-            int transformedValue = transform(value);
-            boolean alreadyContained = contains(value);
+            boolean alreadyContained;
+            // saves calling transform() an additional time
+            alreadyContained = rbm.contains(transformedValue);
             if (!alreadyContained) {
                 rbm.add(transformedValue);
-                stats.size++;
+                stats.size.inc();
                 return true;
             }
             handleCollisions(transformedValue);
@@ -159,7 +167,7 @@ public boolean remove(Integer value) throws Exception {
         int transformedValue = transform(value);
         readLock.lock();
         try {
-            if (!contains(value)) {
+            if (!rbm.contains(transformedValue)) { // saves additional transform() call
                 return false;
             }
             stats.numRemovalAttempts.inc();
@@ -172,7 +180,7 @@ public boolean remove(Integer value) throws Exception {
         writeLock.lock();
         try {
             rbm.remove(transformedValue);
-            stats.size--;
+            stats.size.dec();
             stats.numSuccessfulRemovals.inc();
             return true;
         } finally {
@@ -184,14 +192,14 @@ public boolean remove(Integer value) throws Exception {
     public int getSize() {
         readLock.lock();
         try {
-            return stats.size;
+            return (int) stats.size.count();
         } finally {
             readLock.unlock();
         }
     }
 
     @Override
-    public int getTotalAdds() {
+    public int getAddAttempts() {
         return (int) stats.numAddAttempts.count();
     }
 
@@ -200,7 +208,6 @@ public int getCollisions() {
         return (int) stats.numCollisions.count();
     }
 
-
     @Override
     public boolean isCollision(Integer value1, Integer value2) {
         if (value1 == null || value2 == null) {
@@ -211,7 +218,7 @@ public boolean isCollision(Integer value1, Integer value2) {
 
     @Override
     public long getMemorySizeInBytes() {
-        return sizeEstimator.getSizeInBytes(stats.size) + RBMSizeEstimator.getHashsetMemSizeInBytes(collidedInts.size());
+        return sizeEstimator.getSizeInBytes((int) stats.size.count()) + RBMSizeEstimator.getHashsetMemSizeInBytes(collidedInts.size());
     }
 
     @Override
@@ -221,14 +228,14 @@ public long getMemorySizeCapInBytes() {
 
     @Override
     public boolean isFull() {
-        return stats.atCapacity;
+        return stats.atCapacity.get();
     }
 
     @Override
     public void regenerateStore(Integer[] newValues) throws Exception {
         rbm.clear();
         collidedInts = new HashSet<>();
-        stats.size = 0;
+        stats.size = new CounterMetric();
         stats.numAddAttempts = new CounterMetric();
         stats.numCollisions = new CounterMetric();
         stats.guaranteesNoFalseNegatives = true;
@@ -241,12 +248,11 @@ public void regenerateStore(Integer[] newValues) throws Exception {
         }
     }
 
-
-
     @Override
     public void clear() throws Exception {
-        regenerateStore(new Integer[]{});
+        regenerateStore(new Integer[] {});
     }
+
     public int getNumRemovalAttempts() {
         return (int) stats.numRemovalAttempts.count();
     }

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

@@ -30,7 +30,9 @@
  * GitHub history for details.
  */
 
-package org.opensearch.indices;
+package org.opensearch.common.cache.tier.keystore;
+
+import org.opensearch.common.cache.tier.keystore.RBMIntKeyLookupStore;
 
 /**
  * A class used to estimate roaring bitmap memory sizes (and hash set sizes).
@@ -111,12 +113,16 @@ public static long getSizeInBytesWithModulo(int numEntries, int modulo) {
         return (long) ((long) Math.pow(numEntries, memValues[1]) * (long) Math.pow(10, memValues[2]) * memValues[0]);
     }
 
+    public static long getSizeInBytesWithModuloValue(int numEntries, RBMIntKeyLookupStore.KeystoreModuloValue moduloValue) {
+        double[] memValues = calculateMemoryCoefficients(moduloValue.getValue());
+        return (long) ((long) Math.pow(numEntries, memValues[1]) * (long) Math.pow(10, memValues[2]) * memValues[0]);
+    }
+
     public static int getNumEntriesFromSizeInBytesWithModulo(long sizeInBytes, int modulo) {
         double[] memValues = calculateMemoryCoefficients(modulo);
         return (int) Math.pow(sizeInBytes / (memValues[0] * Math.pow(10, memValues[2])), 1 / memValues[1]);
     }
 
-
     protected static long convertMBToBytes(double valMB) {
         return (long) (valMB * BYTES_IN_MB);
     }