Fix cache regression in reflectIfNewer and cancellation hang in loadIfNewer

Author: rossdanderson

util/src/main/kotlin/com/caplin/integration/datasourcex/util/store/FlowStoreCache.kt

@@ -6,6 +6,7 @@ import com.github.benmanes.caffeine.cache.AsyncCache
 import com.github.benmanes.caffeine.cache.Caffeine
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.ConcurrentHashMap
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.future.await
 import kotlinx.coroutines.launch
@@ -34,7 +35,7 @@ class FlowStoreCache<K : Any, V : Any>(
     inFlight.putIfAbsent(key, mine)?.let {
       return it.await()
     }
-    loaderScope.launch {
+    val job = loaderScope.launch {
       try {
         val loaded = load(key)
         mine.complete(
@@ -46,6 +47,13 @@ class FlowStoreCache<K : Any, V : Any>(
         inFlight.remove(key, mine)
       }
     }
+    job.invokeOnCompletion { throwable ->
+      if (throwable != null) {
+        mine.completeExceptionally(throwable)
+      } else if (!mine.isDone) {
+        mine.completeExceptionally(CancellationException("Job completed without completing future"))
+      }
+    }
     return mine.await()
   }
 
@@ -63,14 +71,21 @@ class FlowStoreCache<K : Any, V : Any>(
   }
 
   /**
-   * Applies [event] to a resident entry only, gated on version; absent keys await the next read.
+   * Applies [event] to a resident or in-flight loading entry only, gated on version; other absent
+   * keys await the next read.
    */
   fun reflectIfNewer(event: VersionedMapEvent<K, V>) {
-    asyncCache().asMap().computeIfPresent(event.key) { _, old ->
-      val current = old.getNow(null)
-      if (current != null && event.version > current.version)
-          CompletableFuture.completedFuture(event.toEntry())
-      else old
+    asyncCache().asMap().compute(event.key) { _, old ->
+      val current = old?.getNow(null)
+      if (current != null) {
+        if (event.version > current.version) CompletableFuture.completedFuture(event.toEntry()) else old
+      } else if (old != null) {
+        old
+      } else if (inFlight.containsKey(event.key)) {
+        CompletableFuture.completedFuture(event.toEntry())
+      } else {
+        null
+      }
     }
   }
 }

util/src/test/kotlin/com/caplin/integration/datasourcex/util/store/FlowStoreCacheTest.kt

@@ -1,10 +1,13 @@
 package com.caplin.integration.datasourcex.util.store
 
+import com.caplin.integration.datasourcex.util.flow.VersionedMapEvent
 import com.github.benmanes.caffeine.cache.Caffeine
+import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.core.spec.style.FunSpec
 import io.kotest.matchers.nulls.shouldBeNull
 import io.kotest.matchers.shouldBe
 import java.util.concurrent.atomic.AtomicInteger
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -53,4 +56,46 @@ class FlowStoreCacheTest :
         cache.putIfNewer("k", Live("v5b", 5L)) shouldBe Live("v5", 5L) // equal rejected
         cache.putIfNewer("k", Live("v9", 9L)) shouldBe Live("v9", 9L)
       }
+
+      test("reflectIfNewer updates the cache during an in-flight load, preventing regression").config(
+          coroutineTestScope = false
+      ) {
+        val cache = Caffeine.newBuilder().buildFlowStoreCache<String, String>(scope)
+        val loadGate = CompletableDeferred<Unit>()
+        val resultGate = CompletableDeferred<Unit>()
+
+        val loadJob = async {
+          cache.loadIfNewer("k") {
+            loadGate.complete(Unit)
+            resultGate.await()
+            Versioned("v5", 5L)
+          }
+        }
+
+        // Wait until the load is executing and in-flight
+        loadGate.await()
+
+        // Call reflectIfNewer with a newer version (Version 6) while load is in-flight
+        cache.reflectIfNewer(VersionedMapEvent.Upsert("k", "v6", 6L))
+
+        // Complete the load (which yields Version 5)
+        resultGate.complete(Unit)
+        loadJob.await()
+
+        // Verify that the cache keeps the newer Version 6 and did not regress to Version 5
+        cache.getIfPresent("k") shouldBe Live("v6", 6L)
+      }
+
+      test("loadIfNewer throws CancellationException if the scope is cancelled").config(
+          coroutineTestScope = false
+      ) {
+        val deadScope = CoroutineScope(Dispatchers.Default)
+        deadScope.cancel()
+
+        val cache = Caffeine.newBuilder().buildFlowStoreCache<String, String>(deadScope)
+
+        shouldThrow<CancellationException> {
+          cache.loadIfNewer("k") { Versioned("v", 1L) }
+        }
+      }
     })