Update BoundedCache

Author: davidh44

utils/src/main/java/software/amazon/awssdk/utils/cache/bounded/BoundedCache.java

@@ -17,108 +17,128 @@
 
 import java.util.Iterator;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
-import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Validate;
 
 /**
- * A thread-safe cache implementation that returns the value for a specified key,
- * retrieving it by either getting the stored value from the cache or using a supplied function to calculate that value
- * and add it to the cache.
+ * A thread-safe cache implementation that returns the value for a specified key, retrieving it by either getting the stored
+ * value from the cache or using a supplied function to calculate that value and add it to the cache.
  * <p>
- * When the cache is full, a new value will push out an unspecified value.
+ * When the cache is full, batch eviction of random values will be performed, with a default evictionBatchSize of 10.
  * <p>
- * The user can configure the maximum size of the cache, which is set to a default of 100.
+ * The user can configure the maximum size of the cache, which is set to a default of 150.
  * <p>
- * Null values are not cached.
+ * Keys must not be null, otherwise an error will be thrown. Null values are not cached.
  */
 @SdkProtectedApi
 @ThreadSafe
 public final class BoundedCache<K, V>  {
-
-    private static final Logger log = Logger.loggerFor(BoundedCache.class);
-
-    private static final int DEFAULT_SIZE = 100;
+    private static final int DEFAULT_CACHE_SIZE = 150;
+    private static final int DEFAULT_EVICTION_BATCH_SIZE = 10;
 
     private final ConcurrentHashMap<K, V> cache;
-    private final Function<K, V> valueSupplier;
+    private final Function<K, V> valueMappingFunction;
     private final int maxCacheSize;
+    private final int evictionBatchSize;
     private final Object cacheLock;
-
-    private BoundedCache(Builder<K, V> builder) {
-        this.valueSupplier = builder.supplier;
-        this.maxCacheSize = builder.maxSize != null ?
-                            Validate.isPositive(builder.maxSize, "maxSize")
-                                                    : DEFAULT_SIZE;
+    private final AtomicInteger cacheSize;
+
+    private BoundedCache(Builder<K, V> b) {
+        this.valueMappingFunction = b.mappingFunction;
+        this.maxCacheSize = b.maxSize != null ? Validate.isPositive(b.maxSize, "maxSize") : DEFAULT_CACHE_SIZE;
+        this.evictionBatchSize = b.evictionBatchSize != null ?
+                                 Validate.isPositive(b.evictionBatchSize, "evictionBatchSize") :
+                                 DEFAULT_EVICTION_BATCH_SIZE;
         this.cache = new ConcurrentHashMap<>();
         this.cacheLock = new Object();
+        this.cacheSize = new AtomicInteger();
     }
 
     /**
-     * Get a value based on the key. If the value exists in the cache, it's returned.
+     * Get a value based on the key. The key must not be null, otherwise an error is thrown.
+     * If the value exists in the cache, it's returned.
      * Otherwise, the value is calculated based on the supplied function {@link Builder#builder(Function)}.
      */
     public V get(K key) {
+        Validate.paramNotNull(key, "key");
         V value = cache.get(key);
         if (value != null) {
             return value;
         }
 
-        V newValue = valueSupplier.apply(key);
+        V newValue = valueMappingFunction.apply(key);
+
+        // If the value is null, just return it without caching
         if (newValue == null) {
             return null;
         }
 
         synchronized (cacheLock) {
+            // Check again inside the synchronized block in case another thread added the value
             value = cache.get(key);
             if (value != null) {
                 return value;
             }
 
-            if (cache.size() >= maxCacheSize) {
+            if (cacheSize.get() >= maxCacheSize) {
                 cleanup();
             }
 
             cache.put(key, newValue);
+            cacheSize.incrementAndGet();
             return newValue;
         }
     }
 
     /**
-     * Clean up the cache by removing an unspecified entry
+     * Clean up the cache by batch removing random entries of evictionBatchSize
      */
     private void cleanup() {
         Iterator<K> iterator = cache.keySet().iterator();
-        if (iterator.hasNext()) {
-            K key = iterator.next();
-            cache.remove(key);
+        int count = 0;
+        while (iterator.hasNext() && count < evictionBatchSize) {
+            iterator.next();
+            iterator.remove();
+            count++;
+            cacheSize.decrementAndGet();
         }
     }
 
     public int size() {
-        return cache.size();
+        return cacheSize.get();
     }
 
-    public static <K, V> BoundedCache.Builder<K, V> builder(Function<K, V> supplier) {
-        return new Builder<>(supplier);
+    public boolean containsKey(K key) {
+        return cache.containsKey(key);
+    }
+
+    public static <K, V> BoundedCache.Builder<K, V> builder(Function<K, V> mappingFunction) {
+        return new Builder<>(mappingFunction);
     }
 
     public static final class Builder<K, V> {
 
-        private final Function<K, V> supplier;
+        private final Function<K, V> mappingFunction;
         private Integer maxSize;
+        private Integer evictionBatchSize;
 
-        private Builder(Function<K, V> supplier) {
-            this.supplier = supplier;
+        private Builder(Function<K, V> mappingFunction) {
+            this.mappingFunction = mappingFunction;
         }
 
         public Builder<K, V> maxSize(Integer maxSize) {
             this.maxSize = maxSize;
             return this;
         }
 
+        public Builder<K, V> evictionBatchSize(Integer evictionBatchSize) {
+            this.evictionBatchSize = evictionBatchSize;
+            return this;
+        }
+
         public BoundedCache<K, V> build() {
             return new BoundedCache<>(this);
         }

utils/src/main/java/software/amazon/awssdk/utils/uri/SdkUri.java

@@ -37,8 +37,8 @@ public final class SdkUri {
     private static final int MAX_INT_DIGITS_BASE_10 = 10;
 
     /*
-     * The default BoundedCache size is 100, but for a single service call we cache at least 3 different URIs so the cache size is
-     * increased a bit to account for the different URIs.
+     * Same value as default BoundedCache size. This contrasts to the default LruCache size of 100, since for a single service
+     * call we cache at least 3 different URIs, so the cache size is increased a bit to account for the different URIs.
      */
     private static final int CACHE_SIZE = 150;
 

utils/src/test/java/software/amazon/awssdk/utils/SdkUriTest.java

@@ -41,6 +41,7 @@ void resetCache() throws IllegalAccessException {
         cacheField.setAccessible(true);
         cacheField.set(SdkUri.getInstance(), BoundedCache.builder(UriConstructorArgs::newInstance)
                                                          .maxSize(100)
+                                                         .evictionBatchSize(5)
                                                          .build());
     }