Mutex and Semaphore review

Author: ndkoval

kotlinx-coroutines-core/common/src/JobSupport.kt

@@ -1205,7 +1205,7 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
     /**
      * @suppress **This is unstable API and it is subject to change.**
      */
-    internal suspend fun awaitInternal(): Any? {
+    protected suspend fun awaitInternal(): Any? {
         // fast-path -- check state (avoid extra object creation)
         while (true) { // lock-free loop on state
             val state = this.state
@@ -1235,9 +1235,8 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
         cont.getResult()
     }
 
-
     @Suppress("UNCHECKED_CAST")
-    internal val onAwaitInternal: SelectClause1<*> get() = SelectClause1Impl<Any?>(
+    protected val onAwaitInternal: SelectClause1<*> get() = SelectClause1Impl<Any?>(
         clauseObject = this@JobSupport,
         regFunc = JobSupport::onAwaitInternalRegFunc as RegistrationFunction,
         processResFunc = JobSupport::onAwaitInternalProcessResFunc as ProcessResultFunction

kotlinx-coroutines-core/common/src/internal/LocalAtomics.common.kt

@@ -14,6 +14,7 @@ internal expect class LocalAtomicInt(value: Int) {
     fun get(): Int
     fun set(value: Int)
     fun decrementAndGet(): Int
+    fun incrementAndGet(): Int
 }
 
 internal inline var LocalAtomicInt.value

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

@@ -9,6 +9,7 @@ import kotlinx.coroutines.*
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.selects.*
 import kotlin.contracts.*
+import kotlin.jvm.*
 import kotlin.native.concurrent.*
 
 /**
@@ -20,7 +21,7 @@ import kotlin.native.concurrent.*
  *
  * JVM API note:
  * Memory semantic of the [Mutex] is similar to `synchronized` block on JVM:
- * An unlock on a [Mutex] happens-before every subsequent successful lock on that [Mutex].
+ * An unlock operation on a [Mutex] happens-before every subsequent successful lock on that [Mutex].
  * Unsuccessful call to [tryLock] do not have any memory effects.
  */
 public interface Mutex {
@@ -171,9 +172,12 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
 
     override fun tryLock(owner: Any?): Boolean =
         if (tryAcquire()) {
+            assert { this.owner.value === NO_OWNER }
             this.owner.value = owner
             true
-        } else false
+        } else {
+            false
+        }
 
     override fun unlock(owner: Any?) {
         while (true) {
@@ -192,7 +196,7 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
         }
     }
 
-    @Suppress("UNCHECKED_CAST")
+    @Suppress("UNCHECKED_CAST", "OverridingDeprecatedMember")
     override val onLock: SelectClause2<Any?, Mutex> get() = SelectClause2Impl(
         clauseObject = this,
         regFunc = MutexImpl::onLockRegFunction as RegistrationFunction,
@@ -201,52 +205,67 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
     )
 
     protected open fun onLockRegFunction(select: SelectInstance<*>, owner: Any?) {
-        onAcquire.regFunc(this, SelectInstanceWithOwner(select, owner), owner)
+        onAcquireRegFunction(SelectInstanceWithOwner(select, owner), owner)
     }
 
     protected open fun onLockProcessResult(owner: Any?, result: Any?): Any? {
-        onAcquire.processResFunc(this, null, result)
         return this
     }
 
-    protected val onCancellationUnlockConstructor: OnCancellationConstructor =
-        { select: SelectInstance<*>, owner: Any?, ignoredResult: Any? ->
+    private val onCancellationUnlockConstructor: OnCancellationConstructor =
+        { _: SelectInstance<*>, owner: Any?, _: Any? ->
             { unlock(owner) }
         }
 
     private inner class CancellableContinuationWithOwner(
+        @JvmField
         val cont: CancellableContinuation<Unit>,
+        @JvmField
         val owner: Any?
     ) : CancellableContinuation<Unit> by cont {
         override fun tryResume(value: Unit, idempotent: Any?, onCancellation: ((cause: Throwable) -> Unit)?): Any? {
-            val token = cont.tryResume(value, idempotent, onCancellation)
-            if (token !== null) this@MutexImpl.owner.value = owner
+            assert { this@MutexImpl.owner.value === NO_OWNER }
+            val token = cont.tryResume(value, idempotent) {
+                assert { this@MutexImpl.owner.value.let { it === NO_OWNER ||it === owner } }
+                this@MutexImpl.owner.value = owner
+                unlock(owner)
+            }
+            if (token != null) {
+                assert { this@MutexImpl.owner.value === NO_OWNER }
+                this@MutexImpl.owner.value = owner
+            }
             return token
         }
 
         override fun resume(value: Unit, onCancellation: ((cause: Throwable) -> Unit)?) {
+            assert { this@MutexImpl.owner.value === NO_OWNER }
             this@MutexImpl.owner.value = owner
-            cont.resume(value, onCancellation)
+            cont.resume(value) { unlock(owner) }
         }
     }
 
     private inner class SelectInstanceWithOwner<Q>(
+        @JvmField
         val select: SelectInstance<Q>,
+        @JvmField
         val owner: Any?
     ) : SelectInstance<Q> by select {
         override fun trySelect(clauseObject: Any, result: Any?): Boolean {
+            assert { this@MutexImpl.owner.value === NO_OWNER }
+            this@MutexImpl.owner.value = owner
             return select.trySelect(clauseObject, result).also { success ->
-                if (success) this@MutexImpl.owner.value = owner
+                if (!success) this@MutexImpl.owner.value = NO_OWNER
             }
         }
 
         override fun selectInRegistrationPhase(internalResult: Any?) {
-            select.selectInRegistrationPhase(internalResult)
+            assert { this@MutexImpl.owner.value === NO_OWNER }
             this@MutexImpl.owner.value = owner
+            select.selectInRegistrationPhase(internalResult)
         }
     }
 
-    override fun toString() = "Mutex@${hexAddress}"
+    override fun toString() = "Mutex@${hexAddress}[isLocked=$isLocked,owner=${owner.value}]"
 }
 
 @SharedImmutable

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

@@ -215,25 +215,15 @@ internal open class SemaphoreImpl(private val permits: Int, acquiredPermits: Int
         }
     }
 
-    val onAcquire: SelectClause1<Semaphore> get() = SelectClause1Impl(
-        clauseObject = this,
-        regFunc = SemaphoreImpl::onAcquireRegFunction as RegistrationFunction,
-        processResFunc = SemaphoreImpl::onAcquireProcessResultFunction as ProcessResultFunction
-    )
-
+    // We do not fully support `onAcquire` as it is needed only for `Mutex.onLock`.
     @Suppress("UNUSED_PARAMETER")
-    private fun onAcquireRegFunction(select: SelectInstance<*>, ignoredParam: Any?) =
+    protected fun onAcquireRegFunction(select: SelectInstance<*>, ignoredParam: Any?) =
         acquire(
             waiter = select,
             suspend = { s -> addAcquireToQueue(s) },
             onAcquired = { s -> s.selectInRegistrationPhase(Unit) }
         )
 
-    @Suppress("UNUSED_PARAMETER", "RedundantNullableReturnType")
-    private fun onAcquireProcessResultFunction(param: Any?, result: Any?): Any? {
-        return this
-    }
-
     /**
      * Decrements the number of available permits
      * and ensures that it is not greater than [permits]

kotlinx-coroutines-core/js/src/internal/LocalAtomics.kt

@@ -12,4 +12,5 @@ internal actual class LocalAtomicInt actual constructor(private var value: Int)
     actual fun get(): Int = value
 
     actual fun decrementAndGet(): Int = --value
+    actual fun incrementAndGet(): Int = ++value
 }

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

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2016-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.internal.*
+import kotlinx.coroutines.selects.*
+import kotlinx.coroutines.sync.*
+import org.junit.*
+import java.util.concurrent.*
+
+class MutexStressTest : TestBase() {
+    @Test
+    fun testStressCancellationDoesNotBreakMutex() = runTest {
+        val mutex = Mutex()
+        val mutexJobNumber = 3
+        val mutexOwners = Array(mutexJobNumber) { "$it" }
+        val dispatcher = Executors.newFixedThreadPool(mutexJobNumber + 2).asCoroutineDispatcher()
+        var counter = 0
+        val counterLocal = Array(mutexJobNumber) { LocalAtomicInt(0) }
+        val completed = LocalAtomicInt(0)
+        val mutexJobLauncher: (jobNumber: Int) -> Job = { jobId ->
+            val coroutineName = "MutexJob-$jobId"
+            launch(dispatcher + CoroutineName(coroutineName)) {
+                while (completed.value == 0) {
+                    mutex.holdsLock(mutexOwners[(jobId + 1) % mutexJobNumber])
+                    if (mutex.tryLock(mutexOwners[jobId])) {
+                        counterLocal[jobId].incrementAndGet()
+                        counter++
+                        mutex.unlock(mutexOwners[jobId])
+                    }
+                    mutex.withLock(mutexOwners[jobId]) {
+                        counterLocal[jobId].incrementAndGet()
+                        counter++
+                    }
+                    select<Unit> {
+                        mutex.onLock(mutexOwners[jobId]) {
+                            counterLocal[jobId].incrementAndGet()
+                            counter++
+                            mutex.unlock(mutexOwners[jobId])
+                        }
+                    }
+                }
+            }
+        }
+        val mutexJobs = (0 until mutexJobNumber).map { mutexJobLauncher(it) }.toMutableList()
+        val checkProgressJob = launch(dispatcher + CoroutineName("checkProgressJob")) {
+            var lastCounterLocalSnapshot = (0 until mutexJobNumber).map { 0 }
+            while (completed.value == 0) {
+                delay(500)
+                val c = counterLocal.map { it.value }
+                for (i in 0 until mutexJobNumber) {
+                    assert(c[i] > lastCounterLocalSnapshot[i]) { "No progress in MutexJob-$i" }
+                }
+                lastCounterLocalSnapshot = c
+            }
+        }
+        val cancellationJob = launch(dispatcher + CoroutineName("cancellationJob")) {
+            var cancellingJobId = 0
+            while (completed.value == 0) {
+                val jobToCancel = mutexJobs.removeFirst()
+                jobToCancel.cancelAndJoin()
+                mutexJobs += mutexJobLauncher(cancellingJobId)
+                cancellingJobId = (cancellingJobId + 1) % mutexJobNumber
+            }
+        }
+        delay(2000L * stressTestMultiplier)
+        completed.value = 1
+        cancellationJob.join()
+        mutexJobs.forEach { it.join() }
+        checkProgressJob.join()
+        check(counter == counterLocal.sumOf { it.value })
+        dispatcher.close()
+    }
+}

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

@@ -5,7 +5,6 @@
 package kotlinx.coroutines.lincheck
 
 import kotlinx.coroutines.*
-import kotlinx.coroutines.selects.*
 import kotlinx.coroutines.sync.*
 import org.jetbrains.kotlinx.lincheck.*
 import org.jetbrains.kotlinx.lincheck.annotations.Operation
@@ -20,9 +19,6 @@ abstract class SemaphoreLincheckTestBase(permits: Int) : AbstractLincheckTest()
     @Operation(promptCancellation = true, allowExtraSuspension = true)
     suspend fun acquire() = semaphore.acquire()
 
-    @Operation(promptCancellation = true, allowExtraSuspension = true)
-    suspend fun onAcquire(): Unit = select { semaphore.onAcquire {} }
-
     @Operation(handleExceptionsAsResult = [IllegalStateException::class])
     fun release() = semaphore.release()
 

kotlinx-coroutines-core/native/src/internal/LocalAtomics.kt

@@ -17,4 +17,5 @@ internal actual class LocalAtomicInt actual constructor(value: Int) {
     actual fun get(): Int = iRef.value
 
     actual fun decrementAndGet(): Int = iRef.decrementAndGet()
+    actual fun incrementAndGet(): Int = iRef.incrementAndGet()
 }