Refine MergeInterceptor to not cancel current job when awaiter exists (#1950)

* Refine MergeInterceptor to not cancel current job when awaiter exists

* Refine

* Cleanup abstraction

* Cleanup abstraction

* Refine

* Refine

* Fix MutableSharedFlow skip f when no awaiters

---------

Co-authored-by: Dude so hot <djohn@fbi.gov>

app/src/main/kotlin/com/hippo/ehviewer/coil/MergeInterceptor.kt

@@ -1,22 +1,61 @@
 package com.hippo.ehviewer.coil
 
+import arrow.core.merge
+import arrow.fx.coroutines.raceN
 import coil3.decode.DataSource
+import coil3.imageLoader
 import coil3.intercept.Interceptor
+import coil3.memory.MemoryCache
 import coil3.request.ImageResult
 import coil3.request.SuccessResult
-import moe.tarsin.coroutines.NamedMutex
-import moe.tarsin.coroutines.withLockNeedSuspend
+import kotlinx.coroutines.CancellableContinuation
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.buffer
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.receiveAsFlow
+import kotlinx.coroutines.flow.shareIn
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.suspendCancellableCoroutine
+import moe.tarsin.coroutines.Counter
+import moe.tarsin.coroutines.Tracker
+import moe.tarsin.coroutines.counter
+import moe.tarsin.coroutines.runSuspendCatching
+import moe.tarsin.coroutines.use
+
+typealias F = suspend () -> Unit
+
+suspend inline fun <T> CancellableContinuation<T>.evalAndResume(f: suspend () -> T) = runSuspendCatching { f() }.takeIf { isActive }?.let(::resumeWith)
+
+class TightRope : CoroutineScope, Counter by counter() {
+    override val coroutineContext = Dispatchers.IO + Job()
+    val actions = Channel<F>()
+    val flow = actions.receiveAsFlow().map(F::invoke).buffer(0).shareIn(this, SharingStarted.WhileSubscribed())
+    suspend inline fun <R> sendAndAwait(crossinline block: suspend () -> R) = raceN(
+        { flow.collect {} },
+        { suspendCancellableCoroutine { cont -> launch { actions.send { cont.evalAndResume(block) } } } },
+    ).merge()
+}
+
+object TightRopeTracker : Tracker<TightRope, String>() {
+    override fun new() = TightRope()
+    override fun free(e: TightRope) = e.cancel()
+}
 
 object MergeInterceptor : Interceptor {
-    private val mutex = NamedMutex<String>()
-
     override suspend fun intercept(chain: Interceptor.Chain): ImageResult {
         val req = chain.request
         val key = req.memoryCacheKey
         return if (key != null) {
-            val (result, suspended) = mutex.withLockNeedSuspend(key) { chain.proceed() }
+            val cacheKey = MemoryCache.Key(key, req.memoryCacheKeyExtras)
+            val cached = req.context.imageLoader.memoryCache!![cacheKey] != null
+            val result = TightRopeTracker.use(key) { sendAndAwait { chain.proceed() } }
             when (result) {
-                is SuccessResult if (suspended) -> result.copy(dataSource = DataSource.MEMORY)
+                is SuccessResult if cached -> result.copy(dataSource = DataSource.MEMORY)
                 else -> result
             }
         } else {

app/src/main/kotlin/moe/tarsin/coroutines/LockPool.kt

@@ -1,47 +0,0 @@
-package moe.tarsin.coroutines
-
-import kotlin.contracts.InvocationKind
-import kotlin.contracts.contract
-
-interface LockPool<Lock, K> {
-    fun acquire(key: K): Lock
-    fun release(key: K, lock: Lock)
-    suspend fun Lock.lock()
-    fun Lock.tryLock(): Boolean
-    fun Lock.unlock()
-}
-
-suspend inline fun <Lock, K, R> LockPool<Lock, K>.withLock(key: K, action: () -> R): R {
-    contract {
-        callsInPlace(action, InvocationKind.EXACTLY_ONCE)
-    }
-    val lock = acquire(key)
-    return try {
-        lock.lock()
-        return try {
-            action()
-        } finally {
-            lock.unlock()
-        }
-    } finally {
-        release(key, lock)
-    }
-}
-
-suspend inline fun <Lock, K, R> LockPool<Lock, K>.withLockNeedSuspend(key: K, action: () -> R): Pair<R, Boolean> {
-    contract {
-        callsInPlace(action, InvocationKind.EXACTLY_ONCE)
-    }
-    val lock = acquire(key)
-    return try {
-        val mustSuspend = !lock.tryLock()
-        if (mustSuspend) lock.lock()
-        return try {
-            action() to mustSuspend
-        } finally {
-            lock.unlock()
-        }
-    } finally {
-        release(key, lock)
-    }
-}

app/src/main/kotlin/moe/tarsin/coroutines/LockTracker.kt

@@ -0,0 +1,59 @@
+package moe.tarsin.coroutines
+
+import androidx.collection.mutableScatterMapOf
+import kotlin.contracts.InvocationKind
+import kotlin.contracts.contract
+
+interface Counter {
+    var ref: Int
+    val isFree get() = ref == 0
+}
+
+fun counter() = object : Counter {
+    override var ref = 0
+}
+
+abstract class Tracker<T : Counter, K> {
+    private val active = mutableScatterMapOf<K, T>()
+    abstract fun new(): T
+    abstract fun free(e: T)
+    fun acquire(key: K) = synchronized(active) { active.getOrPut(key, ::new).apply { ref++ } }
+    fun release(key: K, e: T) = synchronized(active) {
+        e.ref--
+        if (e.isFree) {
+            free(e)
+            active.remove(key)
+        }
+    }
+}
+
+inline fun <T : Counter, K, R> Tracker<T, K>.use(key: K, block: T.() -> R): R {
+    contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }
+    val inst = acquire(key)
+    try {
+        return block(inst)
+    } finally {
+        release(key, inst)
+    }
+}
+
+abstract class LockTracker<Lock : Counter, K> : Tracker<Lock, K>() {
+    abstract suspend fun Lock.lock()
+    abstract fun Lock.unlock()
+}
+
+suspend inline fun <Lock : Counter, K, R> LockTracker<Lock, K>.withLock(key: K, action: () -> R): R {
+    contract {
+        callsInPlace(action, InvocationKind.EXACTLY_ONCE)
+    }
+    use(key) {
+        lock()
+        try {
+            return action()
+        } finally {
+            unlock()
+        }
+    }
+}

app/src/main/kotlin/moe/tarsin/coroutines/NamedMutex.kt

@@ -1,35 +1,21 @@
 package moe.tarsin.coroutines
 
-import androidx.collection.MutableScatterMap
-import androidx.collection.mutableScatterMapOf
 import io.ktor.utils.io.pool.DefaultPool
 import kotlinx.coroutines.sync.Mutex
 
-class MutexTracker(mutex: Mutex = Mutex(), private var count: Int = 0) : Mutex by mutex {
-    operator fun inc() = apply { count++ }
-    operator fun dec() = apply { count-- }
-    val isFree
-        get() = count == 0
-}
+class MutexWithCounter() : Mutex by Mutex(), Counter by counter()
 
-object MutexPool : DefaultPool<MutexTracker>(capacity = 32) {
-    override fun produceInstance() = MutexTracker()
-    override fun validateInstance(mutex: MutexTracker) {
+object MutexPool : DefaultPool<MutexWithCounter>(capacity = 32) {
+    override fun produceInstance() = MutexWithCounter()
+    override fun validateInstance(mutex: MutexWithCounter) {
         check(!mutex.isLocked)
         check(mutex.isFree)
     }
 }
 
-class NamedMutex<K>(val active: MutableScatterMap<K, MutexTracker> = mutableScatterMapOf()) : LockPool<MutexTracker, K> {
-    override fun acquire(key: K) = synchronized(active) { active.getOrPut(key) { MutexPool.borrow() }.inc() }
-    override fun release(key: K, lock: MutexTracker) = synchronized(active) {
-        lock.dec()
-        if (lock.isFree) {
-            active.remove(key)
-            MutexPool.recycle(lock)
-        }
-    }
-    override suspend fun MutexTracker.lock() = lock()
-    override fun MutexTracker.tryLock() = tryLock()
-    override fun MutexTracker.unlock() = unlock()
+class NamedMutex<K>() : LockTracker<MutexWithCounter, K>() {
+    override suspend fun MutexWithCounter.lock() = lock()
+    override fun MutexWithCounter.unlock() = unlock()
+    override fun new() = MutexPool.borrow()
+    override fun free(e: MutexWithCounter) = MutexPool.recycle(e)
 }

app/src/main/kotlin/moe/tarsin/coroutines/NamedSemaphore.kt

@@ -1,36 +1,22 @@
 package moe.tarsin.coroutines
 
-import androidx.collection.mutableScatterMapOf
 import io.ktor.utils.io.pool.DefaultPool
 import kotlinx.coroutines.sync.Semaphore
 
-class SemaphoreTracker(semaphore: Semaphore, private var count: Int = 0) : Semaphore by semaphore {
-    operator fun inc() = apply { count++ }
-    operator fun dec() = apply { count-- }
-    val isFree
-        get() = count == 0
-}
+class SemaphoreWithCounter(semaphore: Semaphore) : Semaphore by semaphore, Counter by counter()
 
-class SemaphorePool(val permits: Int) : DefaultPool<SemaphoreTracker>(capacity = 32) {
-    override fun produceInstance() = SemaphoreTracker(semaphore = Semaphore(permits = permits))
-    override fun validateInstance(semaphore: SemaphoreTracker) {
+class SemaphorePool(val permits: Int) : DefaultPool<SemaphoreWithCounter>(capacity = 32) {
+    override fun produceInstance() = SemaphoreWithCounter(semaphore = Semaphore(permits = permits))
+    override fun validateInstance(semaphore: SemaphoreWithCounter) {
         check(semaphore.availablePermits == permits)
         check(semaphore.isFree)
     }
 }
 
-class NamedSemaphore<K>(val permits: Int) : LockPool<SemaphoreTracker, K> {
+class NamedSemaphore<K>(val permits: Int) : LockTracker<SemaphoreWithCounter, K>() {
     val pool = SemaphorePool(permits = permits)
-    val active = mutableScatterMapOf<K, SemaphoreTracker>()
-    override fun acquire(key: K) = synchronized(active) { active.getOrPut(key) { pool.borrow() }.inc() }
-    override fun release(key: K, lock: SemaphoreTracker) = synchronized(active) {
-        lock.dec()
-        if (lock.isFree) {
-            active.remove(key)
-            pool.recycle(lock)
-        }
-    }
-    override suspend fun SemaphoreTracker.lock() = acquire()
-    override fun SemaphoreTracker.tryLock() = tryAcquire()
-    override fun SemaphoreTracker.unlock() = release()
+    override suspend fun SemaphoreWithCounter.lock() = acquire()
+    override fun SemaphoreWithCounter.unlock() = release()
+    override fun new() = pool.borrow()
+    override fun free(e: SemaphoreWithCounter) = pool.recycle(e)
 }