[Tiered caching] Framework changes

CHANGELOG.md

@@ -98,6 +98,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Remote cluster state] Make index and global metadata upload timeout dynamic cluster settings ([#10814](https://github.com/opensearch-project/OpenSearch/pull/10814))
 - Added cluster setting cluster.restrict.index.replication_type to restrict setting of index setting replication type ([#10866](https://github.com/opensearch-project/OpenSearch/pull/10866))
 - Add cluster state stats ([#10670](https://github.com/opensearch-project/OpenSearch/pull/10670))
+- [Tiered caching] Framework changes ([#10753](https://github.com/opensearch-project/OpenSearch/pull/10753)
 
 ### Dependencies
 - Bump `com.google.api.grpc:proto-google-common-protos` from 2.10.0 to 2.25.1 ([#10208](https://github.com/opensearch-project/OpenSearch/pull/10208), [#10298](https://github.com/opensearch-project/OpenSearch/pull/10298))

server/src/main/java/org/opensearch/common/cache/Cache.java

@@ -422,68 +422,74 @@
             }
         });
         if (value == null) {
-            // we need to synchronize loading of a value for a given key; however, holding the segment lock while
-            // invoking load can lead to deadlock against another thread due to dependent key loading; therefore, we
-            // need a mechanism to ensure that load is invoked at most once, but we are not invoking load while holding
-            // the segment lock; to do this, we atomically put a future in the map that can load the value, and then
-            // get the value from this future on the thread that won the race to place the future into the segment map
-            CacheSegment<K, V> segment = getCacheSegment(key);
-            CompletableFuture<Entry<K, V>> future;
-            CompletableFuture<Entry<K, V>> completableFuture = new CompletableFuture<>();
+            value = compute(key, loader);
+        }
+        return value;
+    }
 
-            try (ReleasableLock ignored = segment.writeLock.acquire()) {
-                future = segment.map.putIfAbsent(key, completableFuture);
-            }
+    public V compute(K key, CacheLoader<K, V> loader) throws ExecutionException {
+        long now = now();
+        // we need to synchronize loading of a value for a given key; however, holding the segment lock while
+        // invoking load can lead to deadlock against another thread due to dependent key loading; therefore, we
+        // need a mechanism to ensure that load is invoked at most once, but we are not invoking load while holding
+        // the segment lock; to do this, we atomically put a future in the map that can load the value, and then
+        // get the value from this future on the thread that won the race to place the future into the segment map
+        CacheSegment<K, V> segment = getCacheSegment(key);
+        CompletableFuture<Entry<K, V>> future;
+        CompletableFuture<Entry<K, V>> completableFuture = new CompletableFuture<>();
 
-            BiFunction<? super Entry<K, V>, Throwable, ? extends V> handler = (ok, ex) -> {
-                if (ok != null) {
-                    try (ReleasableLock ignored = lruLock.acquire()) {
-                        promote(ok, now);
-                    }
-                    return ok.value;
-                } else {
-                    try (ReleasableLock ignored = segment.writeLock.acquire()) {
-                        CompletableFuture<Entry<K, V>> sanity = segment.map.get(key);
-                        if (sanity != null && sanity.isCompletedExceptionally()) {
-                            segment.map.remove(key);
-                        }
-                    }
-                    return null;
-                }
-            };
+        try (ReleasableLock ignored = segment.writeLock.acquire()) {
+            future = segment.map.putIfAbsent(key, completableFuture);
+        }
 
-            CompletableFuture<V> completableValue;
-            if (future == null) {
-                future = completableFuture;
-                completableValue = future.handle(handler);
-                V loaded;
-                try {
-                    loaded = loader.load(key);
-                } catch (Exception e) {
-                    future.completeExceptionally(e);
-                    throw new ExecutionException(e);
-                }
-                if (loaded == null) {
-                    NullPointerException npe = new NullPointerException("loader returned a null value");
-                    future.completeExceptionally(npe);
-                    throw new ExecutionException(npe);
-                } else {
-                    future.complete(new Entry<>(key, loaded, now));
+        BiFunction<? super Entry<K, V>, Throwable, ? extends V> handler = (ok, ex) -> {
+            if (ok != null) {
+                try (ReleasableLock ignored = lruLock.acquire()) {
+                    promote(ok, now);
                 }
+                return ok.value;
             } else {
-                completableValue = future.handle(handler);
+                try (ReleasableLock ignored = segment.writeLock.acquire()) {
+                    CompletableFuture<Entry<K, V>> sanity = segment.map.get(key);
+                    if (sanity != null && sanity.isCompletedExceptionally()) {
+                        segment.map.remove(key);
+                    }
+                }
+                return null;
             }
+        };
 
+        CompletableFuture<V> completableValue;
+        if (future == null) {
+            future = completableFuture;
+            completableValue = future.handle(handler);
+            V loaded;
             try {
-                value = completableValue.get();
-                // check to ensure the future hasn't been completed with an exception
-                if (future.isCompletedExceptionally()) {
-                    future.get(); // call get to force the exception to be thrown for other concurrent callers
-                    throw new IllegalStateException("the future was completed exceptionally but no exception was thrown");
-                }
-            } catch (InterruptedException e) {
-                throw new IllegalStateException(e);
+                loaded = loader.load(key);
+            } catch (Exception e) {
+                future.completeExceptionally(e);
+                throw new ExecutionException(e);
             }
+            if (loaded == null) {
+                NullPointerException npe = new NullPointerException("loader returned a null value");
+                future.completeExceptionally(npe);
+                throw new ExecutionException(npe);
+            } else {
+                future.complete(new Entry<>(key, loaded, now));
+            }
+        } else {
+            completableValue = future.handle(handler);
+        }
+        V value;
+        try {
+            value = completableValue.get();
+            // check to ensure the future hasn't been completed with an exception
+            if (future.isCompletedExceptionally()) {
+                future.get(); // call get to force the exception to be thrown for other concurrent callers
+                throw new IllegalStateException("the future was completed exceptionally but no exception was thrown");
+            }
+        } catch (InterruptedException e) {
+            throw new IllegalStateException(e);
         }
         return value;
     }

server/src/main/java/org/opensearch/common/cache/RemovalNotification.java

@@ -32,6 +32,8 @@
 
 package org.opensearch.common.cache;
 
+import org.opensearch.common.cache.tier.TierType;
+
 /**
  * Notification when an element is removed from the cache
  *
@@ -42,11 +44,17 @@ public class RemovalNotification<K, V> {
     private final K key;
     private final V value;
     private final RemovalReason removalReason;
+    private final TierType tierType;
 
     public RemovalNotification(K key, V value, RemovalReason removalReason) {
+        this(key, value, removalReason, TierType.ON_HEAP);
+    }
+
+    public RemovalNotification(K key, V value, RemovalReason removalReason, TierType tierType) {
         this.key = key;
         this.value = value;
         this.removalReason = removalReason;
+        this.tierType = tierType;
     }
 
     public K getKey() {
@@ -60,4 +68,8 @@ public V getValue() {
     public RemovalReason getRemovalReason() {
         return removalReason;
     }
+
+    public TierType getTierType() {
+        return tierType;
+    }
 }

server/src/main/java/org/opensearch/common/cache/tier/CachingTier.java

@@ -0,0 +1,45 @@
+/*
+ * 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;
+
+import org.opensearch.common.cache.RemovalListener;
+
+/**
+ * Caching tier interface. Can be implemented/extended by concrete classes to provide different flavors of cache like
+ * onHeap, disk etc.
+ * @param <K> Type of key
+ * @param <V> Type of value
+ */
+public interface CachingTier<K, V> {
+
+    V get(K key);
+
+    void put(K key, V value);
+
+    V computeIfAbsent(K key, TieredCacheLoader<K, V> loader) throws Exception;
+
+    void invalidate(K key);
+
+    V compute(K key, TieredCacheLoader<K, V> loader) throws Exception;
+
+    void setRemovalListener(RemovalListener<K, V> removalListener);
+
+    void invalidateAll();
+
+    Iterable<K> keys();
+
+    int count();
+
+    TierType getTierType();
+
+    /**
+     * Force any outstanding size-based and time-based evictions to occur
+     */
+    default void refresh() {}
+}

server/src/main/java/org/opensearch/common/cache/tier/DiskCachingTier.java

@@ -0,0 +1,18 @@
+/*
+ * 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;
+
+/**
+ * This is specific to disk caching tier and can be used to add methods which are specific to disk tier.
+ * @param <K> Type of key
+ * @param <V> Type of value
+ */
+public interface DiskCachingTier<K, V> extends CachingTier<K, V> {
+
+}

server/src/main/java/org/opensearch/common/cache/tier/OnHeapCachingTier.java

@@ -0,0 +1,16 @@
+/*
+ * 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;
+
+/**
+ * This is specific to onHeap caching tier and can be used to add methods which are specific to this tier.
+ * @param <K> Type of key
+ * @param <V> Type of value
+ */
+public interface OnHeapCachingTier<K, V> extends CachingTier<K, V> {}

server/src/main/java/org/opensearch/common/cache/tier/OpenSearchOnHeapCache.java

@@ -0,0 +1,136 @@
+/*
+ * 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;
+
+import org.opensearch.common.cache.Cache;
+import org.opensearch.common.cache.CacheBuilder;
+import org.opensearch.common.cache.RemovalListener;
+import org.opensearch.common.cache.RemovalNotification;
+import org.opensearch.common.unit.TimeValue;
+
+import java.util.Objects;
+import java.util.concurrent.ExecutionException;
+import java.util.function.ToLongBiFunction;
+
+/**
+ * This variant of on-heap cache uses OpenSearch custom cache implementation.
+ * @param <K> Type of key
+ * @param <V> Type of value
+ */
+public class OpenSearchOnHeapCache<K, V> implements OnHeapCachingTier<K, V>, RemovalListener<K, V> {
+
+    private final Cache<K, V> cache;
+    private RemovalListener<K, V> removalListener;
+
+    private OpenSearchOnHeapCache(Builder<K, V> builder) {
+        Objects.requireNonNull(builder.weigher);
+        CacheBuilder<K, V> cacheBuilder = CacheBuilder.<K, V>builder()
+            .setMaximumWeight(builder.maxWeightInBytes)
+            .weigher(builder.weigher)
+            .removalListener(this);
+        if (builder.expireAfterAcess != null) {
+            cacheBuilder.setExpireAfterAccess(builder.expireAfterAcess);
+        }
+        cache = cacheBuilder.build();
+    }
+
+    @Override
+    public void setRemovalListener(RemovalListener<K, V> removalListener) {
+        this.removalListener = removalListener;
+    }
+
+    @Override
+    public void invalidateAll() {
+        cache.invalidateAll();
+    }
+
+    @Override
+    public Iterable<K> keys() {
+        return this.cache.keys();
+    }
+
+    @Override
+    public int count() {
+        return cache.count();
+    }
+
+    @Override
+    public TierType getTierType() {
+        return TierType.ON_HEAP;
+    }
+
+    @Override
+    public V get(K key) {
+        return cache.get(key);
+    }
+
+    @Override
+    public void put(K key, V value) {
+        cache.put(key, value);
+    }
+
+    @Override
+    public V computeIfAbsent(K key, TieredCacheLoader<K, V> loader) throws ExecutionException {
+        return cache.computeIfAbsent(key, key1 -> loader.load(key));
+    }
+
+    @Override
+    public void invalidate(K key) {
+        cache.invalidate(key);
+    }
+
+    @Override
+    public V compute(K key, TieredCacheLoader<K, V> loader) throws Exception {
+        return cache.compute(key, key1 -> loader.load(key));
+    }
+
+    @Override
+    public void refresh() {
+        cache.refresh();
+    }
+
+    @Override
+    public void onRemoval(RemovalNotification<K, V> notification) {
+        removalListener.onRemoval(notification);
+    }
+
+    /**
+     * Builder object
+     * @param <K> Type of key
+     * @param <V> Type of value
+     */
+    public static class Builder<K, V> {
+        private long maxWeightInBytes;
+
+        private ToLongBiFunction<K, V> weigher;
+
+        private TimeValue expireAfterAcess;
+
+        public Builder() {}
+
+        public Builder<K, V> setMaximumWeight(long sizeInBytes) {
+            this.maxWeightInBytes = sizeInBytes;
+            return this;
+        }
+
+        public Builder<K, V> setWeigher(ToLongBiFunction<K, V> weigher) {
+            this.weigher = weigher;
+            return this;
+        }
+
+        public Builder<K, V> setExpireAfterAccess(TimeValue expireAfterAcess) {
+            this.expireAfterAcess = expireAfterAcess;
+            return this;
+        }
+
+        public OpenSearchOnHeapCache<K, V> build() {
+            return new OpenSearchOnHeapCache<K, V>(this);
+        }
+    }
+}