perf(coroutines): reduce smart lock overhead

CHANGELOG.md

@@ -2,6 +2,16 @@
 
 ## 0.32.1
 
+* `Coroutines`:
+  * `SmartSemaphore`:
+    * Optimize permit checks, acquisition, uncontended releases and already satisfied waits to reduce allocations while preserving cancellation handling.
+  * `SmartMutex`:
+    * Optimize waits for unlocked mutexes and uncontended unlocks to reduce allocations while preserving cancellation handling.
+  * `SmartRWLocker`:
+    * Remove redundant coroutine contexts during read releases, write unlocks and cancellation cleanup.
+  * `SmartKeyRWLocker`:
+    * Reduce allocations during keyed operations by inlining the locker allocation helper and using a hash map without insertion order tracking.
+
 ## 0.32.0
 
 * `Versions`:

coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/SmartKeyRWLocker.kt

@@ -37,7 +37,7 @@ class SmartKeyRWLocker<T>(
         readPermits = globalLockerReadPermits,
         writeIsLocked = globalLockerWriteIsLocked
     )
-    private val lockers = mutableMapOf<T, SmartRWLocker>()
+    private val lockers = hashMapOf<T, SmartRWLocker>()
     private val lockersMutex = Mutex()
     private val lockersWritingLocker = SmartSemaphore.Mutable(Int.MAX_VALUE)
     private val globalWritingLocker = SmartSemaphore.Mutable(Int.MAX_VALUE)
@@ -45,7 +45,7 @@ class SmartKeyRWLocker<T>(
     private fun allocateLockerWithoutLock(key: T) = lockers.getOrPut(key) {
         SmartRWLocker(perKeyReadPermits)
     }
-    private suspend fun allocateLocker(key: T) = lockersMutex.withLock {
+    private suspend inline fun allocateLocker(key: T) = lockersMutex.withLock {
         lockers.getOrPut(key) {
             SmartRWLocker(perKeyReadPermits)
         }

coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/SmartMutex.kt

@@ -2,6 +2,7 @@ package dev.inmo.micro_utils.coroutines
 
 import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.first
@@ -94,19 +95,32 @@ sealed interface SmartMutex {
          * If [isLocked] == true - will change it to false and return true. If current call will not unlock this
          * [SmartMutex] - false
          */
-        suspend fun unlock(): Boolean = withContext(NonCancellable) {
-            if (_lockStateFlow.value) {
-                internalChangesMutex.withLock {
-                    if (_lockStateFlow.value) {
-                        _lockStateFlow.value = false
-                        true
-                    } else {
-                        false
-                    }
-                }
-            } else {
-                false
+        suspend fun unlock(): Boolean {
+            if (!_lockStateFlow.value) {
+                return false
             }
+            // The uncontended path cannot suspend, including during cancelled cleanup.
+            if (internalChangesMutex.tryLock()) {
+                try {
+                    return unlockState()
+                } finally {
+                    internalChangesMutex.unlock()
+                }
+            }
+            return withContext(NonCancellable) {
+                internalChangesMutex.withLock {
+                    unlockState()
+                }
+            }
+        }
+
+        // The caller must hold internalChangesMutex while updating the state.
+        private fun unlockState(): Boolean {
+            if (!_lockStateFlow.value) {
+                return false
+            }
+            _lockStateFlow.value = false
+            return true
         }
     }
 }
@@ -134,4 +148,11 @@ suspend inline fun <T> SmartMutex.Mutable.withLock(action: () -> T): T {
  * Anyway, after the end of this block there are no any guaranties that [SmartMutex.isLocked] == false due to the fact
  * that some other parties may lock it again
  */
-suspend fun SmartMutex.waitUnlock() = lockStateFlow.first { !it }
+suspend fun SmartMutex.waitUnlock(): Boolean {
+    currentCoroutineContext().ensureActive()
+    return if (!lockStateFlow.value) {
+        false
+    } else {
+        lockStateFlow.first { !it }
+    }
+}

coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/SmartRWLocker.kt

@@ -1,8 +1,6 @@
 package dev.inmo.micro_utils.coroutines
 
 import kotlinx.coroutines.CancellationException
-import kotlinx.coroutines.NonCancellable
-import kotlinx.coroutines.withContext
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.InvocationKind
 import kotlin.contracts.contract
@@ -39,9 +37,7 @@ class SmartRWLocker(private val readPermits: Int = Int.MAX_VALUE, writeIsLocked:
     /**
      * Release one read permit in [readSemaphore]
      */
-    suspend fun releaseRead(): Boolean = withContext(NonCancellable) {
-        _readSemaphore.release()
-    }
+    suspend fun releaseRead(): Boolean = _readSemaphore.release()
 
     /**
      * Locking [writeMutex] and wait while all [readSemaphore] permits will be freed
@@ -51,9 +47,7 @@ class SmartRWLocker(private val readPermits: Int = Int.MAX_VALUE, writeIsLocked:
         try {
             _readSemaphore.acquire(readPermits)
         } catch (e: CancellationException) {
-            withContext(NonCancellable) {
-                _writeMutex.unlock()
-            }
+            _writeMutex.unlock()
             throw e
         }
     }
@@ -61,13 +55,12 @@ class SmartRWLocker(private val readPermits: Int = Int.MAX_VALUE, writeIsLocked:
     /**
      * Unlock [writeMutex]
      */
-    suspend fun unlockWrite(): Boolean = withContext(NonCancellable) {
+    suspend fun unlockWrite(): Boolean =
         _writeMutex.unlock().also { unlocked ->
             if (unlocked) {
                 _readSemaphore.release(readPermits)
             }
         }
-    }
 }
 
 /**

coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/SmartSemaphore.kt

@@ -2,6 +2,7 @@ package dev.inmo.micro_utils.coroutines
 
 import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.first
@@ -63,7 +64,7 @@ sealed interface SmartSemaphore {
 
         fun immutable() = Immutable(permitsStateFlow, maxPermits)
 
-        private fun checkedPermits(permits: Int) = permits.coerceIn(1 .. this.maxPermits)
+        private fun checkedPermits(permits: Int) = permits.coerceIn(1, maxPermits)
 
         /**
          * Holds call until this [SmartSemaphore] will be re-locked. That means that current method will
@@ -75,7 +76,8 @@ sealed interface SmartSemaphore {
                 do {
                     val shouldContinue = internalChangesMutex.withLock {
                         val requiredPermits = checkedPermits - acquiredPermits
-                        val acquiring = minOf(freePermits, requiredPermits).takeIf { it > 0 } ?: return@withLock true
+                        val acquiring = minOf(freePermits, requiredPermits)
+                        if (acquiring <= 0) return@withLock true
                         acquiredPermits += acquiring
                         _freePermitsStateFlow.value -= acquiring
 
@@ -137,20 +139,31 @@ sealed interface SmartSemaphore {
          * If [freePermits] == true - will change it to false and return true. If current call will not unlock this
          * [SmartSemaphore] - false
          */
-        suspend fun release(permits: Int = 1): Boolean = withContext(NonCancellable) {
+        suspend fun release(permits: Int = 1): Boolean {
             val checkedPermits = checkedPermits(permits)
-            if (_freePermitsStateFlow.value < maxPermits) {
-                internalChangesMutex.withLock {
-                    if (_freePermitsStateFlow.value < maxPermits) {
-                        _freePermitsStateFlow.value = minOf(_freePermitsStateFlow.value + checkedPermits, maxPermits)
-                        true
-                    } else {
-                        false
-                    }
+            if (_freePermitsStateFlow.value >= maxPermits) return false
+
+            // Updating permits without suspension remains safe in a cancelled coroutine.
+            if (internalChangesMutex.tryLock()) {
+                try {
+                    return releaseLocked(checkedPermits)
+                } finally {
+                    internalChangesMutex.unlock()
                 }
-            } else {
-                false
             }
+
+            return withContext(NonCancellable) {
+                internalChangesMutex.withLock {
+                    releaseLocked(checkedPermits)
+                }
+            }
+        }
+
+        private fun releaseLocked(permits: Int): Boolean {
+            val freePermits = _freePermitsStateFlow.value
+            if (freePermits >= maxPermits) return false
+            _freePermitsStateFlow.value = minOf(freePermits + permits, maxPermits)
+            return true
         }
     }
 }
@@ -178,5 +191,14 @@ suspend inline fun <T> SmartSemaphore.Mutable.withAcquire(permits: Int = 1, acti
  * Anyway, after the end of this block there are no any guaranties that [SmartSemaphore.freePermits] >= [permits] due to
  * the fact that some other parties may lock it again
  */
-suspend fun SmartSemaphore.waitRelease(permits: Int = 1) = permitsStateFlow.first { it >= permits }
-suspend fun SmartSemaphore.waitReleaseAll() = permitsStateFlow.first { it == maxPermits }
+suspend fun SmartSemaphore.waitRelease(permits: Int = 1): Int {
+    currentCoroutineContext().ensureActive()
+    val availablePermits = freePermits
+    return if (availablePermits >= permits) availablePermits else permitsStateFlow.first { it >= permits }
+}
+
+suspend fun SmartSemaphore.waitReleaseAll(): Int {
+    currentCoroutineContext().ensureActive()
+    val availablePermits = freePermits
+    return if (availablePermits == maxPermits) availablePermits else permitsStateFlow.first { it == maxPermits }
+}