From 9ccddae60a0e0887babd2943acf9a82e19fff5a1 Mon Sep 17 00:00:00 2001
From: Dude so hot <djohn@fbi.gov>
Date: Tue, 15 Oct 2024 19:58:08 +0800
Subject: [PATCH] Introduce Better NamedMutex

---
 .../moe/tarsin/coroutines/NamedMutex.kt       | 49 ++++++++++++++-----
 1 file changed, 36 insertions(+), 13 deletions(-)

diff --git a/app/src/main/kotlin/moe/tarsin/coroutines/NamedMutex.kt b/app/src/main/kotlin/moe/tarsin/coroutines/NamedMutex.kt
index 35ce5e8ac..e55d9403d 100644
--- a/app/src/main/kotlin/moe/tarsin/coroutines/NamedMutex.kt
+++ b/app/src/main/kotlin/moe/tarsin/coroutines/NamedMutex.kt
@@ -1,23 +1,46 @@
 package moe.tarsin.coroutines
 
-import arrow.core.memoize
-import kotlin.contracts.InvocationKind
-import kotlin.contracts.contract
+import androidx.collection.mutableScatterMapOf
+import io.ktor.utils.io.pool.DefaultPool
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 
-class NamedMutex<T> {
-    val innerMutexGetter = { _: T -> Mutex() }.memoize()
+class MutexTracker(
+    val mutex: Mutex = Mutex(),
+    var count: Int = 0,
+)
+
+val MutexTracker.free
+    get() = count == 0
+
+operator fun MutexTracker.inc() = apply { count++ }
+operator fun MutexTracker.dec() = apply { count-- }
+
+class MutexPool(capacity: Int = 16) : DefaultPool<MutexTracker>(capacity) {
+    override fun produceInstance() = MutexTracker()
+    override fun validateInstance(mutex: MutexTracker) {
+        check(!mutex.mutex.isLocked)
+        check(mutex.free)
+    }
 }
 
-suspend inline fun <T, K> NamedMutex<K>.withLock(key: K, owner: Any? = null, action: () -> T): T {
-    contract {
-        callsInPlace(action, InvocationKind.EXACTLY_ONCE)
-    }
+class NamedMutex<K> {
+    val pool = MutexPool()
+    val active = mutableScatterMapOf<K, MutexTracker>()
+    val lock = Any()
+}
 
-    val mutex = synchronized(innerMutexGetter) {
-        innerMutexGetter(key)
+suspend inline fun <K, R> NamedMutex<K>.withLock(key: K, owner: Any? = null, action: () -> R): R {
+    val mutex = synchronized(lock) { active.getOrPut(key) { pool.borrow() }.inc() }
+    return try {
+        mutex.mutex.withLock(owner, action)
+    } finally {
+        synchronized(lock) {
+            mutex.dec()
+            if (mutex.free) {
+                active.remove(key)
+                pool.recycle(mutex)
+            }
+        }
     }
-
-    return mutex.withLock(owner, action)
 }