From 658eb309b1b80a8a0b53171a1a33a54375ffb2f6 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 27 Jun 2024 21:54:27 +0000 Subject: [PATCH] Fix flaky test TieredSpilloverCacheTests.testComputeIfAbsentConcurrently (#14550) * Fix flaky test TieredSpilloverCacheTests.testComputeIfAbsentConcurrently Signed-off-by: Sagar Upadhyaya * Addressing comment Signed-off-by: Sagar Upadhyaya --------- Signed-off-by: Sagar Upadhyaya Signed-off-by: Sagar Upadhyaya Co-authored-by: Sagar Upadhyaya (cherry picked from commit 729276f7e8f799ca21ace12bb2767869322ff253) Signed-off-by: github-actions[bot] --- .../common/tier/TieredSpilloverCache.java | 21 ++++++++++++------- .../tier/TieredSpilloverCacheTests.java | 4 ++-- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java b/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java index b6d6913a9f8d4..f69c56808b2a1 100644 --- a/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java +++ b/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java @@ -195,7 +195,16 @@ public V computeIfAbsent(ICacheKey key, LoadAwareCacheLoader, V> // and it only has to be loaded one time, we should report one miss and the rest hits. But, if we do stats in // getValueFromTieredCache(), // we will see all misses. Instead, handle stats in computeIfAbsent(). - Tuple cacheValueTuple = getValueFromTieredCache(false).apply(key); + Tuple cacheValueTuple; + CompletableFuture, V>> future = null; + try (ReleasableLock ignore = readLock.acquire()) { + cacheValueTuple = getValueFromTieredCache(false).apply(key); + if (cacheValueTuple == null) { + // Only one of the threads will succeed putting a future into map for the same key. + // Rest will fetch existing future and wait on that to complete. + future = completableFutureMap.putIfAbsent(key, new CompletableFuture<>()); + } + } List heapDimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, TIER_DIMENSION_VALUE_ON_HEAP); List diskDimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, TIER_DIMENSION_VALUE_DISK); @@ -203,7 +212,7 @@ public V computeIfAbsent(ICacheKey key, LoadAwareCacheLoader, V> // Add the value to the onHeap cache. We are calling computeIfAbsent which does another get inside. // This is needed as there can be many requests for the same key at the same time and we only want to load // the value once. - V value = compute(key, loader); + V value = compute(key, loader, future); // Handle stats if (loader.isLoaded()) { // The value was just computed and added to the cache by this thread. Register a miss for the heap cache, and the disk cache @@ -232,10 +241,8 @@ public V computeIfAbsent(ICacheKey key, LoadAwareCacheLoader, V> return cacheValueTuple.v1(); } - private V compute(ICacheKey key, LoadAwareCacheLoader, V> loader) throws Exception { - // Only one of the threads will succeed putting a future into map for the same key. - // Rest will fetch existing future and wait on that to complete. - CompletableFuture, V>> future = completableFutureMap.putIfAbsent(key, new CompletableFuture<>()); + private V compute(ICacheKey key, LoadAwareCacheLoader, V> loader, CompletableFuture, V>> future) + throws Exception { // Handler to handle results post processing. Takes a tuple or exception as an input and returns // the value. Also before returning value, puts the value in cache. BiFunction, V>, Throwable, Void> handler = (pair, ex) -> { @@ -253,7 +260,7 @@ private V compute(ICacheKey key, LoadAwareCacheLoader, V> loader logger.warn("Exception occurred while trying to compute the value", ex); } } - completableFutureMap.remove(key); // Remove key from map as not needed anymore. + completableFutureMap.remove(key);// Remove key from map as not needed anymore. return null; }; V value = null; diff --git a/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java b/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java index b9c7bbdb77d3d..c6440a1e1797f 100644 --- a/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java +++ b/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java @@ -760,7 +760,7 @@ public void testInvalidateAll() throws Exception { } public void testComputeIfAbsentConcurrently() throws Exception { - int onHeapCacheSize = randomIntBetween(100, 300); + int onHeapCacheSize = randomIntBetween(500, 700); int diskCacheSize = randomIntBetween(200, 400); int keyValueSize = 50; @@ -782,7 +782,7 @@ public void testComputeIfAbsentConcurrently() throws Exception { 0 ); - int numberOfSameKeys = randomIntBetween(10, onHeapCacheSize - 1); + int numberOfSameKeys = randomIntBetween(400, onHeapCacheSize - 1); ICacheKey key = getICacheKey(UUID.randomUUID().toString()); String value = UUID.randomUUID().toString();