Skip to content

Commit 729276f

Browse files
sgup432Sagar Upadhyaya
and
Sagar Upadhyaya
authored
Fix flaky test TieredSpilloverCacheTests.testComputeIfAbsentConcurrently (#14550)
* Fix flaky test TieredSpilloverCacheTests.testComputeIfAbsentConcurrently Signed-off-by: Sagar Upadhyaya <sagar.upadhyaya.121@gmail.com> * Addressing comment Signed-off-by: Sagar Upadhyaya <upasagar@amazon.com> --------- Signed-off-by: Sagar Upadhyaya <sagar.upadhyaya.121@gmail.com> Signed-off-by: Sagar Upadhyaya <upasagar@amazon.com> Co-authored-by: Sagar Upadhyaya <upasagar@amazon.com>
1 parent 8839904 commit 729276f

File tree

2 files changed

+16
-9
lines changed

2 files changed

+16
-9
lines changed

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

+14-7
Original file line numberDiff line numberDiff line change
@@ -195,15 +195,24 @@ public V computeIfAbsent(ICacheKey<K> key, LoadAwareCacheLoader<ICacheKey<K>, V>
195195
// and it only has to be loaded one time, we should report one miss and the rest hits. But, if we do stats in
196196
// getValueFromTieredCache(),
197197
// we will see all misses. Instead, handle stats in computeIfAbsent().
198-
Tuple<V, String> cacheValueTuple = getValueFromTieredCache(false).apply(key);
198+
Tuple<V, String> cacheValueTuple;
199+
CompletableFuture<Tuple<ICacheKey<K>, V>> future = null;
200+
try (ReleasableLock ignore = readLock.acquire()) {
201+
cacheValueTuple = getValueFromTieredCache(false).apply(key);
202+
if (cacheValueTuple == null) {
203+
// Only one of the threads will succeed putting a future into map for the same key.
204+
// Rest will fetch existing future and wait on that to complete.
205+
future = completableFutureMap.putIfAbsent(key, new CompletableFuture<>());
206+
}
207+
}
199208
List<String> heapDimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, TIER_DIMENSION_VALUE_ON_HEAP);
200209
List<String> diskDimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, TIER_DIMENSION_VALUE_DISK);
201210

202211
if (cacheValueTuple == null) {
203212
// Add the value to the onHeap cache. We are calling computeIfAbsent which does another get inside.
204213
// This is needed as there can be many requests for the same key at the same time and we only want to load
205214
// the value once.
206-
V value = compute(key, loader);
215+
V value = compute(key, loader, future);
207216
// Handle stats
208217
if (loader.isLoaded()) {
209218
// 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<K> key, LoadAwareCacheLoader<ICacheKey<K>, V>
232241
return cacheValueTuple.v1();
233242
}
234243

235-
private V compute(ICacheKey<K> key, LoadAwareCacheLoader<ICacheKey<K>, V> loader) throws Exception {
236-
// Only one of the threads will succeed putting a future into map for the same key.
237-
// Rest will fetch existing future and wait on that to complete.
238-
CompletableFuture<Tuple<ICacheKey<K>, V>> future = completableFutureMap.putIfAbsent(key, new CompletableFuture<>());
244+
private V compute(ICacheKey<K> key, LoadAwareCacheLoader<ICacheKey<K>, V> loader, CompletableFuture<Tuple<ICacheKey<K>, V>> future)
245+
throws Exception {
239246
// Handler to handle results post processing. Takes a tuple<key, value> or exception as an input and returns
240247
// the value. Also before returning value, puts the value in cache.
241248
BiFunction<Tuple<ICacheKey<K>, V>, Throwable, Void> handler = (pair, ex) -> {
@@ -253,7 +260,7 @@ private V compute(ICacheKey<K> key, LoadAwareCacheLoader<ICacheKey<K>, V> loader
253260
logger.warn("Exception occurred while trying to compute the value", ex);
254261
}
255262
}
256-
completableFutureMap.remove(key); // Remove key from map as not needed anymore.
263+
completableFutureMap.remove(key);// Remove key from map as not needed anymore.
257264
return null;
258265
};
259266
V value = null;

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -760,7 +760,7 @@ public void testInvalidateAll() throws Exception {
760760
}
761761

762762
public void testComputeIfAbsentConcurrently() throws Exception {
763-
int onHeapCacheSize = randomIntBetween(100, 300);
763+
int onHeapCacheSize = randomIntBetween(500, 700);
764764
int diskCacheSize = randomIntBetween(200, 400);
765765
int keyValueSize = 50;
766766

@@ -782,7 +782,7 @@ public void testComputeIfAbsentConcurrently() throws Exception {
782782
0
783783
);
784784

785-
int numberOfSameKeys = randomIntBetween(10, onHeapCacheSize - 1);
785+
int numberOfSameKeys = randomIntBetween(400, onHeapCacheSize - 1);
786786
ICacheKey<String> key = getICacheKey(UUID.randomUUID().toString());
787787
String value = UUID.randomUUID().toString();
788788

0 commit comments

Comments
 (0)