[Tiered Caching] Fixing flaky tiered cache test (opensearch-project#1…

…2650)

* Fixing flaky tiered cache test

Signed-off-by: Sagar Upadhyaya <sagar.upadhyaya.121@gmail.com>

* Removing unnecessary comment

Signed-off-by: Sagar Upadhyaya <sagar.upadhyaya.121@gmail.com>

* Removing unused variable

Signed-off-by: Sagar Upadhyaya <sagar.upadhyaya.121@gmail.com>

---------

Signed-off-by: Sagar Upadhyaya <sagar.upadhyaya.121@gmail.com>

modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java

@@ -65,7 +65,6 @@ public void onRemoval(RemovalNotification<K, V> notification) {
  try (ReleasableLock ignore = writeLock.acquire()) {
  diskCache.put(notification.getKey(), notification.getValue());
  }
- removalListener.onRemoval(notification);
  }
  })
  .setKeyType(builder.cacheConfig.getKeyType())

modules/cache-common/src/test/java/org/opensearch/cache/common/tier/MockDiskCache.java

@@ -11,6 +11,9 @@
 import org.opensearch.common.cache.CacheType;
 import org.opensearch.common.cache.ICache;
 import org.opensearch.common.cache.LoadAwareCacheLoader;
+import org.opensearch.common.cache.RemovalListener;
+import org.opensearch.common.cache.RemovalNotification;
+import org.opensearch.common.cache.RemovalReason;
 import org.opensearch.common.cache.store.builders.ICacheBuilder;
 import org.opensearch.common.cache.store.config.CacheConfig;
 
@@ -23,9 +26,12 @@ public class MockDiskCache<K, V> implements ICache<K, V> {
  int maxSize;
  long delay;
 
- public MockDiskCache(int maxSize, long delay) {
+ private final RemovalListener<K, V> removalListener;
+
+ public MockDiskCache(int maxSize, long delay, RemovalListener<K, V> removalListener) {
  this.maxSize = maxSize;
  this.delay = delay;
+ this.removalListener = removalListener;
  this.cache = new ConcurrentHashMap<K, V>();
  }
 
@@ -38,7 +44,7 @@ public V get(K key) {
  @Override
  public void put(K key, V value) {
  if (this.cache.size() >= maxSize) { // For simplification
- return;
+ this.removalListener.onRemoval(new RemovalNotification<>(key, value, RemovalReason.EVICTED));
  }
  try {
  Thread.sleep(delay);
@@ -101,7 +107,10 @@ public MockDiskCacheFactory(long delay, int maxSize) {
 
  @Override
  public <K, V> ICache<K, V> create(CacheConfig<K, V> config, CacheType cacheType, Map<String, Factory> cacheFactories) {
- return new Builder<K, V>().setMaxSize(maxSize).setDeliberateDelay(delay).build();
+ return new Builder<K, V>().setMaxSize(maxSize)
+ .setDeliberateDelay(delay)
+ .setRemovalListener(config.getRemovalListener())
+ .build();
  }
 
  @Override
@@ -117,7 +126,7 @@ public static class Builder<K, V> extends ICacheBuilder<K, V> {
 
  @Override
  public ICache<K, V> build() {
- return new MockDiskCache<K, V>(this.maxSize, this.delay);
+ return new MockDiskCache<K, V>(this.maxSize, this.delay, this.getRemovalListener());
  }
 
  public Builder<K, V> setMaxSize(int maxSize) {

modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java

@@ -42,7 +42,7 @@ public void testComputeIfAbsentWithoutAnyOnHeapCacheEviction() throws Exception
 
  MockCacheRemovalListener<String, String> removalListener = new MockCacheRemovalListener<>();
  TieredSpilloverCache<String, String> tieredSpilloverCache = intializeTieredSpilloverCache(
- onHeapCacheSize,
+ keyValueSize,
  randomIntBetween(1, 4),
  removalListener,
  Settings.builder()
@@ -142,10 +142,6 @@ public void testComputeIfAbsentWithFactoryBasedCacheCreation() throws Exception
  LoadAwareCacheLoader<String, String> tieredCacheLoader = getLoadAwareCacheLoader();
  tieredSpilloverCache.computeIfAbsent(key, tieredCacheLoader);
  }
- long actualDiskCacheSize = tieredSpilloverCache.getDiskCache().count();
- assertEquals(actualDiskCacheSize, removalListener.evictionsMetric.count()); // Evictions from onHeap equal to
- // disk cache size.
-
  tieredSpilloverCache.getOnHeapCache().keys().forEach(onHeapKeys::add);
  tieredSpilloverCache.getDiskCache().keys().forEach(diskTierKeys::add);
 
@@ -290,9 +286,6 @@ public void testComputeIfAbsentWithEvictionsFromOnHeapCache() throws Exception {
  LoadAwareCacheLoader<String, String> tieredCacheLoader = getLoadAwareCacheLoader();
  tieredSpilloverCache.computeIfAbsent(key, tieredCacheLoader);
  }
- long actualDiskCacheSize = tieredSpilloverCache.getDiskCache().count();
- assertEquals(actualDiskCacheSize, removalListener.evictionsMetric.count()); // Evictions from onHeap equal to
- // disk cache size.
 
  tieredSpilloverCache.getOnHeapCache().keys().forEach(onHeapKeys::add);
  tieredSpilloverCache.getDiskCache().keys().forEach(diskTierKeys::add);
@@ -328,15 +321,15 @@ public void testComputeIfAbsentWithEvictionsFromOnHeapCache() throws Exception {
  }
  }
 
- public void testComputeIfAbsentWithEvictionsFromBothTier() throws Exception {
+ public void testComputeIfAbsentWithEvictionsFromTieredCache() throws Exception {
  int onHeapCacheSize = randomIntBetween(10, 30);
  int diskCacheSize = randomIntBetween(onHeapCacheSize + 1, 100);
  int totalSize = onHeapCacheSize + diskCacheSize;
  int keyValueSize = 50;
 
  MockCacheRemovalListener<String, String> removalListener = new MockCacheRemovalListener<>();
  TieredSpilloverCache<String, String> tieredSpilloverCache = intializeTieredSpilloverCache(
- onHeapCacheSize,
+ keyValueSize,
  diskCacheSize,
  removalListener,
  Settings.builder()
@@ -349,13 +342,13 @@ public void testComputeIfAbsentWithEvictionsFromBothTier() throws Exception {
  .build(),
  0
  );
-
  int numOfItems = randomIntBetween(totalSize + 1, totalSize * 3);
  for (int iter = 0; iter < numOfItems; iter++) {
  LoadAwareCacheLoader<String, String> tieredCacheLoader = getLoadAwareCacheLoader();
  tieredSpilloverCache.computeIfAbsent(UUID.randomUUID().toString(), tieredCacheLoader);
  }
- assertTrue(removalListener.evictionsMetric.count() > 0);
+ int evictions = numOfItems - (totalSize);
+ assertEquals(evictions, removalListener.evictionsMetric.count());
  }
 
  public void testGetAndCount() throws Exception {
@@ -366,7 +359,7 @@ public void testGetAndCount() throws Exception {
 
  MockCacheRemovalListener<String, String> removalListener = new MockCacheRemovalListener<>();
  TieredSpilloverCache<String, String> tieredSpilloverCache = intializeTieredSpilloverCache(
- onHeapCacheSize,
+ keyValueSize,
  diskCacheSize,
  removalListener,
  Settings.builder()
@@ -418,7 +411,7 @@ public void testPut() {
 
  MockCacheRemovalListener<String, String> removalListener = new MockCacheRemovalListener<>();
  TieredSpilloverCache<String, String> tieredSpilloverCache = intializeTieredSpilloverCache(
- onHeapCacheSize,
+ keyValueSize,
  diskCacheSize,
  removalListener,
  Settings.builder()
@@ -519,7 +512,7 @@ public void testInvalidate() {
 
  MockCacheRemovalListener<String, String> removalListener = new MockCacheRemovalListener<>();
  TieredSpilloverCache<String, String> tieredSpilloverCache = intializeTieredSpilloverCache(
- onHeapCacheSize,
+ keyValueSize,
  diskCacheSize,
  removalListener,
  Settings.builder()
@@ -744,7 +737,7 @@ public String load(String key) {
  assertEquals(1, numberOfTimesKeyLoaded); // It should be loaded only once.
  }
 
- public void testConcurrencyForEvictionFlow() throws Exception {
+ public void testConcurrencyForEvictionFlowFromOnHeapToDiskTier() throws Exception {
  int diskCacheSize = randomIntBetween(450, 800);
 
  MockCacheRemovalListener<String, String> removalListener = new MockCacheRemovalListener<>();
@@ -828,7 +821,6 @@ public String load(String key) {
  countDownLatch.await();
  assertNotNull(actualValue.get());
  countDownLatch1.await();
- assertEquals(1, removalListener.evictionsMetric.count());
  assertEquals(1, tieredSpilloverCache.getOnHeapCache().count());
  assertEquals(1, onDiskCache.count());
  assertNotNull(onDiskCache.get(keyToBeEvicted));
@@ -883,7 +875,6 @@ private TieredSpilloverCache<String, String> intializeTieredSpilloverCache(
  .build()
  )
  .build();
-
  ICache.Factory mockDiskCacheFactory = new MockDiskCache.MockDiskCacheFactory(diskDeliberateDelay, diskCacheSize);
 
  return new TieredSpilloverCache.Builder<String, String>().setCacheType(CacheType.INDICES_REQUEST_CACHE)