From a5dfc2326c69adea662cc0578f1e76778c445936 Mon Sep 17 00:00:00 2001
From: Gareth Pearce <Tilps@users.noreply.github.com>
Date: Mon, 16 Nov 2020 17:20:00 +1100
Subject: [PATCH] Fix cancellation support for Mutex when lock becomes
 available between tryLock and suspending.

---
 .../common/src/sync/Mutex.kt                  | 12 +++++-
 .../common/test/sync/MutexTest.kt             | 37 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/kotlinx-coroutines-core/common/src/sync/Mutex.kt b/kotlinx-coroutines-core/common/src/sync/Mutex.kt
index 73aaab5fbf..36f62acd3b 100644
--- a/kotlinx-coroutines-core/common/src/sync/Mutex.kt
+++ b/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
                         }
                     }
diff --git a/kotlinx-coroutines-core/common/test/sync/MutexTest.kt b/kotlinx-coroutines-core/common/test/sync/MutexTest.kt
index c5d0ccf187..0cfd48031a 100644
--- a/kotlinx-coroutines-core/common/test/sync/MutexTest.kt
+++ b/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))
     }
-}
\ No newline at end of file
+
+    @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)
+    }
+}