Fix cancellation support for Mutex when lock becomes available between tryLock and suspending.

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

@@ -201,7 +201,17 @@ 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)
+                            val token = cont.tryResume(Unit, idempotent = null) {
+                                // if this continuation gets cancelled during dispatch to the caller, then release
+                                // the lock
+                                unlock(owner)
+                            }
+                            if (token != null) {
+                                cont.completeResume(token)
+                            } else {
+                                // failure to get token implies already cancelled
+                                unlock(owner)
+                            }
                             return@sc
                         }
                     }

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

@@ -4,10 +4,14 @@
 
 package kotlinx.coroutines.sync
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlin.test.*
 
 class MutexTest : TestBase() {
+    private val enterCount = atomic(0)
+    private val releasedCount = atomic(0)
+
     @Test
     fun testSimple() = runTest {
         val mutex = Mutex()
@@ -106,4 +110,35 @@ class MutexTest : TestBase() {
         assertFalse(mutex.holdsLock(firstOwner))
         assertFalse(mutex.holdsLock(secondOwner))
     }
-}
+
+    @Test
+    fun cancelLock() = runTest() {
+        val mutex = Mutex()
+        enterCount.value = 0
+        releasedCount.value = 0
+        repeat(1000) {
+            val job = launch(Dispatchers.Default) {
+                val owner = Any()
+                try {
+                    enterCount.incrementAndGet()
+                    mutex.withLock(owner) {}
+                    // repeat to give an increase in race probability
+                    mutex.withLock(owner) {}
+                } finally {
+                    // should be no way lock is still held by owner here
+                    if (mutex.holdsLock(owner)) {
+                        // if it is held, ensure test case doesn't lockup
+                        mutex.unlock(owner)
+                    } else {
+                        releasedCount.incrementAndGet()
+                    }
+                }
+            }
+            mutex.withLock {
+                job.cancel()
+            }
+            job.join()
+        }
+        assertEquals(enterCount.value, releasedCount.value)
+    }
+}