Unlock Mutex and release Semaphore during cancellation on a fast branch of slow-path in Mutex/Semaphore (#2396)

Fixes #2390

Co-authored-by: Gareth Pearce <Tilps@users.noreply.github.com>

Author: qwwdfsad

kotlinx-coroutines-core/common/src/sync/Mutex.kt

@@ -201,7 +201,8 @@ internal class MutexImpl(locked: Boolean) : Mutex, SelectClause2<Any?, Mutex> {
                         // try lock
                         val update = if (owner == null) EMPTY_LOCKED else Empty(owner)
                         if (_state.compareAndSet(state, update)) { // locked
-                            cont.resume(Unit)
+                            // TODO implement functional type in LockCont as soon as we get rid of legacy JS
+                            cont.resume(Unit) { unlock(owner) }
                             return@sc
                         }
                     }

kotlinx-coroutines-core/common/src/sync/Semaphore.kt

@@ -172,7 +172,7 @@ private class SemaphoreImpl(private val permits: Int, acquiredPermits: Int) : Se
             if (addAcquireToQueue(cont)) return@sc
             val p = _availablePermits.getAndDecrement()
             if (p > 0) { // permit acquired
-                cont.resume(Unit)
+                cont.resume(Unit, onCancellationRelease)
                 return@sc
             }
         }
@@ -206,9 +206,8 @@ private class SemaphoreImpl(private val permits: Int, acquiredPermits: Int) : Se
         // On CAS failure -- the cell must be either PERMIT or BROKEN
         // If the cell already has PERMIT from tryResumeNextFromQueue, try to grab it
         if (segment.cas(i, PERMIT, TAKEN)) { // took permit thus eliminating acquire/release pair
-            // The following resume must always succeed, since continuation was not published yet and we don't have
-            // to pass onCancellationRelease handle, since the coroutine did not suspend yet and cannot be cancelled
-            cont.resume(Unit)
+            /// This continuation is not yet published, but still can be cancelled via outer job
+            cont.resume(Unit, onCancellationRelease)
             return true
         }
         assert { segment.get(i) === BROKEN } // it must be broken in this case, no other way around it

kotlinx-coroutines-core/common/test/sync/MutexTest.kt

@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines.sync
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlin.test.*
 
@@ -106,4 +107,4 @@ class MutexTest : TestBase() {
         assertFalse(mutex.holdsLock(firstOwner))
         assertFalse(mutex.holdsLock(secondOwner))
     }
-}
+}

kotlinx-coroutines-core/jvm/test/sync/MutexStressTest.kt

@@ -90,4 +90,21 @@ class MutexStressTest : TestBase() {
             }
         }
     }
-}
+
+    @Test
+    fun testShouldBeUnlockedOnCancellation() = runTest {
+        val mutex = Mutex()
+        val n = 1000 * stressTestMultiplier
+        repeat(n) {
+            val job = launch(Dispatchers.Default) {
+                mutex.lock()
+                mutex.unlock()
+            }
+            mutex.withLock {
+                job.cancel()
+            }
+            job.join()
+            assertFalse { mutex.isLocked }
+        }
+    }
+}

kotlinx-coroutines-core/jvm/test/sync/SemaphoreStressTest.kt

@@ -2,7 +2,7 @@ package kotlinx.coroutines.sync
 
 import kotlinx.coroutines.*
 import org.junit.Test
-import kotlin.test.assertEquals
+import kotlin.test.*
 
 class SemaphoreStressTest : TestBase() {
     @Test
@@ -90,4 +90,21 @@ class SemaphoreStressTest : TestBase() {
             }
         }
     }
+
+    @Test
+    fun testShouldBeUnlockedOnCancellation() = runTest {
+        val semaphore = Semaphore(1)
+        val n = 1000 * stressTestMultiplier
+        repeat(n) {
+            val job = launch(Dispatchers.Default) {
+                semaphore.acquire()
+                semaphore.release()
+            }
+            semaphore.withPermit {
+                job.cancel()
+            }
+            job.join()
+            assertTrue { semaphore.availablePermits == 1 }
+        }
+    }
 }