Preserve mutex invariant: throw ISE when incorrect access with owner is detected (#3620)

* Behaviour compatible with the previous implementation is restored
* For everything else, there is #3626

Signed-off-by: Nikita Koval <[email protected]>
Co-authored-by: Nikita Koval <[email protected]>

Author: qwwdfsad

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

@@ -170,14 +170,36 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
         acquire(contWithOwner)
     }
 
-    override fun tryLock(owner: Any?): Boolean =
-        if (tryAcquire()) {
-            assert { this.owner.value === NO_OWNER }
-            this.owner.value = owner
-            true
-        } else {
-            false
+    override fun tryLock(owner: Any?): Boolean = when (tryLockImpl(owner)) {
+        TRY_LOCK_SUCCESS -> true
+        TRY_LOCK_FAILED -> false
+        TRY_LOCK_ALREADY_LOCKED_BY_OWNER -> error("This mutex is already locked by the specified owner: $owner")
+        else -> error("unexpected")
+    }
+
+    private fun tryLockImpl(owner: Any?): Int {
+        while (true) {
+            if (tryAcquire()) {
+                assert { this.owner.value === NO_OWNER }
+                this.owner.value = owner
+                return TRY_LOCK_SUCCESS
+            } else {
+                // The semaphore permit acquisition has failed.
+                // However, we need to check that this mutex is not
+                // locked by our owner.
+                if (owner != null) {
+                    // Is this mutex locked by our owner?
+                    if (holdsLock(owner)) return TRY_LOCK_ALREADY_LOCKED_BY_OWNER
+                    // This mutex is either locked by another owner or unlocked.
+                    // In the latter case, it is possible that it WAS locked by
+                    // our owner when the semaphore permit acquisition has failed.
+                    // To preserve linearizability, the operation restarts in this case.
+                    if (!isLocked) continue
+                }
+                return TRY_LOCK_FAILED
+            }
         }
+    }
 
     override fun unlock(owner: Any?) {
         while (true) {
@@ -205,10 +227,17 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
     )
 
     protected open fun onLockRegFunction(select: SelectInstance<*>, owner: Any?) {
-        onAcquireRegFunction(SelectInstanceWithOwner(select, owner), owner)
+        if (owner != null && holdsLock(owner)) {
+            select.selectInRegistrationPhase(ON_LOCK_ALREADY_LOCKED_BY_OWNER)
+        } else {
+            onAcquireRegFunction(SelectInstanceWithOwner(select, owner), owner)
+        }
     }
 
     protected open fun onLockProcessResult(owner: Any?, result: Any?): Any? {
+        if (result == ON_LOCK_ALREADY_LOCKED_BY_OWNER) {
+            error("This mutex is already locked by the specified owner: $owner")
+        }
         return this
     }
 
@@ -263,3 +292,8 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
 }
 
 private val NO_OWNER = Symbol("NO_OWNER")
+private val ON_LOCK_ALREADY_LOCKED_BY_OWNER = Symbol("ALREADY_LOCKED_BY_OWNER")
+
+private const val TRY_LOCK_SUCCESS = 0
+private const val TRY_LOCK_FAILED = 1
+private const val TRY_LOCK_ALREADY_LOCKED_BY_OWNER = 2

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

@@ -4,8 +4,8 @@
 
 package kotlinx.coroutines.sync
 
-import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
+import kotlinx.coroutines.selects.*
 import kotlin.test.*
 
 class MutexTest : TestBase() {
@@ -138,4 +138,14 @@ class MutexTest : TestBase() {
         assertTrue(mutex.holdsLock(owner2))
         finish(4)
     }
+
+    @Test
+    fun testIllegalStateInvariant() = runTest {
+        val mutex = Mutex()
+        val owner = Any()
+        assertTrue(mutex.tryLock(owner))
+        assertFailsWith<IllegalStateException> { mutex.tryLock(owner) }
+        assertFailsWith<IllegalStateException> { mutex.lock(owner) }
+        assertFailsWith<IllegalStateException> { select { mutex.onLock(owner) {} } }
+    }
 }

kotlinx-coroutines-core/jvm/test/lincheck/MutexLincheckTest.kt

@@ -16,15 +16,17 @@ import org.jetbrains.kotlinx.lincheck.paramgen.*
 class MutexLincheckTest : AbstractLincheckTest() {
     private val mutex = Mutex()
 
-    @Operation
+    @Operation(handleExceptionsAsResult = [IllegalStateException::class])
     fun tryLock(@Param(name = "owner") owner: Int) = mutex.tryLock(owner.asOwnerOrNull)
 
+    // TODO: `lock()` with non-null owner is non-linearizable
     @Operation(promptCancellation = true)
-    suspend fun lock(@Param(name = "owner") owner: Int) = mutex.lock(owner.asOwnerOrNull)
+    suspend fun lock() = mutex.lock(null)
 
+    // TODO: `onLock` with non-null owner is non-linearizable
     // onLock may suspend in case of clause re-registration.
     @Operation(allowExtraSuspension = true, promptCancellation = true)
-    suspend fun onLock(@Param(name = "owner") owner: Int) = select<Unit> { mutex.onLock(owner.asOwnerOrNull) {} }
+    suspend fun onLock() = select<Unit> { mutex.onLock(null) {} }
 
     @Operation(handleExceptionsAsResult = [IllegalStateException::class])
     fun unlock(@Param(name = "owner") owner: Int) = mutex.unlock(owner.asOwnerOrNull)