Preserve mutex invariant: throw ISE when incorrect access with owenr is detected

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

Author: ndkoval

Diff for: kotlinx-coroutines-core/common/src/sync/Mutex.kt

@@ -165,20 +165,41 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
         lockSuspend(owner)
     }
 
-    private suspend fun lockSuspend(owner: Any?) = suspendCancellableCoroutineReusable<Unit> { cont ->
-        cont as CancellableContinuationImpl<Unit>
+    private suspend fun lockSuspend(owner: Any?) = suspendCancellableCoroutineReusable { cont ->
         val contWithOwner = CancellableContinuationWithOwner(cont, owner)
         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) {
@@ -206,19 +227,26 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
     )
 
     protected open fun onLockRegFunction(select: SelectInstance<*>, owner: Any?) {
-        onAcquireRegFunction(SelectInstanceWithOwner(select as SelectInstanceInternal<*>, 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
     }
 
     private inner class CancellableContinuationWithOwner(
         @JvmField
-        val cont: CancellableContinuationImpl<Unit>,
+        val cont: CancellableContinuation<Unit>,
         @JvmField
         val owner: Any?
-    ) : CancellableContinuation<Unit> by cont, Waiter by cont {
+    ) : CancellableContinuation<Unit> by cont {
         override fun tryResume(value: Unit, idempotent: Any?, onCancellation: ((cause: Throwable) -> Unit)?): Any? {
             assert { this@MutexImpl.owner.value === NO_OWNER }
             val token = cont.tryResume(value, idempotent) {
@@ -242,10 +270,10 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
 
     private inner class SelectInstanceWithOwner<Q>(
         @JvmField
-        val select: SelectInstanceInternal<Q>,
+        val select: SelectInstance<Q>,
         @JvmField
         val owner: Any?
-    ) : SelectInstanceInternal<Q> by select {
+    ) : SelectInstanceInternal<Q> by select as SelectInstanceInternal<Q> {
         override fun trySelect(clauseObject: Any, result: Any?): Boolean {
             assert { this@MutexImpl.owner.value === NO_OWNER }
             return select.trySelect(clauseObject, result).also { success ->
@@ -264,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

Diff for: 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) {} } }
+    }
 }

Diff for: 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)