In progress

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

Author: ndkoval

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

@@ -33,7 +33,7 @@ internal val RESUME_TOKEN = Symbol("RESUME_TOKEN")
 internal open class CancellableContinuationImpl<in T>(
     final override val delegate: Continuation<T>,
     resumeMode: Int
-) : DispatchedTask<T>(resumeMode), CancellableContinuation<T>, CoroutineStackFrame, Waiter {
+) : DispatchedTask<T>(resumeMode), CancellableContinuation<T>, CoroutineStackFrame, Waiter, SegmentDisposable {
     init {
         assert { resumeMode != MODE_UNINITIALIZED } // invalid mode for CancellableContinuationImpl
     }
@@ -246,7 +246,7 @@ internal open class CancellableContinuationImpl<in T>(
     private fun callSegmentOnCancellation(segment: Segment<*>, cause: Throwable?) {
         val index = _decisionAndIndex.value.index
         check(index != NO_INDEX) { "The index for segment.invokeOnCancellation(..) is broken" }
-        callCancelHandlerSafely { segment.invokeOnCancellation(index, cause) }
+        callCancelHandlerSafely { segment.onCancellation(index, cause) }
     }
 
 
@@ -372,11 +372,11 @@ internal open class CancellableContinuationImpl<in T>(
      * [segment] and [index] in this [CancellableContinuationImpl].
      * ```
      * invokeOnCancellation { cause ->
-     *   segment.invokeOnCancellation(index, cause)
+     *   segment.onCancellation(index, cause)
      * }
      * ```
      */
-    internal fun invokeOnCancellation(segment: Segment<*>, index: Int) {
+    override fun disposeOnCancellation(segment: Segment<*>, index: Int) {
         _decisionAndIndex.update {
             check(it.index == NO_INDEX) {
                 "invokeOnCancellation should be invoked at most once"

kotlinx-coroutines-core/common/src/channels/BufferedChannel.kt

@@ -175,24 +175,13 @@ internal open class BufferedChannel<E>(
         index: Int
     ) {
         if (onUndeliveredElement == null) {
-            invokeOnCancellation(SenderOrReceiverCancellationHandler(segment, index).asHandler)
+            this as CancellableContinuationImpl<*>
+            disposeOnCancellation(segment, index)
         } else {
             invokeOnCancellation(SenderWithOnUndeliveredElementCancellationHandler(segment, index, context).asHandler)
         }
     }
 
-    // TODO: Replace with a more efficient cancellation mechanism for segments when #3084 is finished.
-    private inner class SenderOrReceiverCancellationHandler(
-        private val segment: ChannelSegment<E>,
-        private val index: Int
-    ) : BeforeResumeCancelHandler(), DisposableHandle {
-        override fun dispose() {
-            segment.onCancellation(index)
-        }
-
-        override fun invoke(cause: Throwable?) = dispose()
-    }
-
     private inner class SenderWithOnUndeliveredElementCancellationHandler(
         private val segment: ChannelSegment<E>,
         private val index: Int,
@@ -743,7 +732,8 @@ internal open class BufferedChannel<E>(
 
     private fun CancellableContinuation<*>.prepareReceiverForSuspension(segment: ChannelSegment<E>, index: Int) {
         onReceiveEnqueued()
-        invokeOnCancellation(SenderOrReceiverCancellationHandler(segment, index).asHandler)
+        this as CancellableContinuationImpl<*>
+        disposeOnCancellation(segment, index)
     }
 
     private fun onClosedReceiveOnNoWaiterSuspend(cont: CancellableContinuation<E>) {
@@ -1508,7 +1498,8 @@ internal open class BufferedChannel<E>(
         index: Int
     ) {
         if (onUndeliveredElement == null) {
-            disposeOnCompletion(SenderOrReceiverCancellationHandler(segment, index))
+            this as SelectImplementation<*>
+            disposeOnCancellation(segment, index)
         } else {
             disposeOnCompletion(SenderWithOnUndeliveredElementCancellationHandler(segment, index, context))
         }
@@ -1567,7 +1558,8 @@ internal open class BufferedChannel<E>(
         index: Int
     ) {
         onReceiveEnqueued()
-        disposeOnCompletion(SenderOrReceiverCancellationHandler(segment, index))
+        this as SelectImplementation<*>
+        disposeOnCancellation(segment, index)
     }
 
     private fun onClosedSelectOnReceive(select: SelectInstance<*>) {
@@ -2871,6 +2863,10 @@ internal class ChannelSegment<E>(id: Long, prev: ChannelSegment<E>?, channel: Bu
         }
     }
 
+    override fun onCancellation(index: Int, cause: Throwable?) {
+        onCancellation(index)
+    }
+
     /**
      *  Returns `true` if the request is successfully cancelled,
      *  and no rendezvous has happened. We need this knowledge

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

@@ -223,9 +223,9 @@ internal abstract class Segment<S : Segment<S>>(val id: Long, prev: S?, pointers
     /**
      * This function is invoked on continuation cancellation when this segment
      * with the specified [index] are installed as cancellation handler via
-     * `CancellableContinuationImpl.invokeOnCancellation(Segment, Int)`.
+     * `SegmentDisposable.disposeOnCancellation(Segment, Int)`.
      */
-    internal open fun invokeOnCancellation(index: Int, cause: Throwable?) {}
+    internal open fun onCancellation(index: Int, cause: Throwable?) {}
 
     /**
      * Invoked on each slot clean-up; should not be invoked twice for the same slot.
@@ -235,6 +235,10 @@ internal abstract class Segment<S : Segment<S>>(val id: Long, prev: S?, pointers
     }
 }
 
+internal interface SegmentDisposable {
+    fun disposeOnCancellation(segment: Segment<*>, index: Int)
+}
+
 private inline fun AtomicInt.addConditionally(delta: Int, condition: (cur: Int) -> Boolean): Boolean {
     while (true) {
         val cur = this.value

kotlinx-coroutines-core/common/src/selects/Select.kt

@@ -238,7 +238,7 @@ public sealed interface SelectInstance<in R> {
      */
     public fun selectInRegistrationPhase(internalResult: Any?)
 }
-internal interface SelectInstanceInternal<R>: SelectInstance<R>, Waiter
+internal interface SelectInstanceInternal<R>: SelectInstance<R>, Waiter, SegmentDisposable
 
 @PublishedApi
 internal open class SelectImplementation<R> constructor(
@@ -380,7 +380,12 @@ internal open class SelectImplementation<R> constructor(
      * After that, if the clause is successfully registered (so, it has not completed immediately),
      * this [DisposableHandle] is stored into the corresponding [ClauseData] instance.
      */
-    private var disposableHandle: DisposableHandle? = null
+    private var disposableHandleOrSegment: Any? = null
+
+    /**
+     * TODO
+     */
+    private var indexInSegment: Int = -1
 
     /**
      * Stores the result passed via [selectInRegistrationPhase] during clause registration
@@ -469,8 +474,10 @@ internal open class SelectImplementation<R> constructor(
             // This also guarantees that the list of clauses cannot be cleared
             // in the registration phase, so it is safe to read it with "!!".
             if (!reregister) clauses!! += this
-            disposableHandle = this@SelectImplementation.disposableHandle
-            this@SelectImplementation.disposableHandle = null
+            disposableHandleOrSegment = this@SelectImplementation.disposableHandleOrSegment
+            indexInSegment = this@SelectImplementation.indexInSegment
+            this@SelectImplementation.disposableHandleOrSegment = null
+            this@SelectImplementation.indexInSegment = -1
         } else {
             // This clause has been selected!
             // Update the state correspondingly.
@@ -493,7 +500,12 @@ internal open class SelectImplementation<R> constructor(
     }
 
     override fun disposeOnCompletion(disposableHandle: DisposableHandle) {
-        this.disposableHandle = disposableHandle
+        this.disposableHandleOrSegment = disposableHandle
+    }
+
+    override fun disposeOnCancellation(segment: Segment<*>, index: Int) {
+        this.disposableHandleOrSegment = segment
+        this.indexInSegment = index
     }
 
     override fun selectInRegistrationPhase(internalResult: Any?) {
@@ -556,7 +568,8 @@ internal open class SelectImplementation<R> constructor(
      */
     private fun reregisterClause(clauseObject: Any) {
         val clause = findClause(clauseObject)!! // it is guaranteed that the corresponding clause is presented
-        clause.disposableHandle = null
+        clause.disposableHandleOrSegment = null
+        clause.indexInSegment = -1
         clause.register(reregister = true)
     }
 
@@ -692,7 +705,7 @@ internal open class SelectImplementation<R> constructor(
         // Invoke all cancellation handlers except for the
         // one related to the selected clause, if specified.
         clauses.forEach { clause ->
-            if (clause !== selectedClause) clause.disposableHandle?.dispose()
+            if (clause !== selectedClause) clause.dispose()
         }
         // We do need to clean all the data to avoid memory leaks.
         this.state.value = STATE_COMPLETED
@@ -716,7 +729,7 @@ internal open class SelectImplementation<R> constructor(
         // a concurrent clean-up procedure has already completed, and it is safe to finish.
         val clauses = this.clauses ?: return
         // Remove this `select` instance from all the clause object (channels, mutexes, etc.).
-        clauses.forEach { it.disposableHandle?.dispose() }
+        clauses.forEach { it.dispose() }
         // We do need to clean all the data to avoid memory leaks.
         this.internalResult = NO_RESULT
         this.clauses = null
@@ -731,9 +744,11 @@ internal open class SelectImplementation<R> constructor(
         private val processResFunc: ProcessResultFunction,
         private val param: Any?, // the user-specified param
         private val block: Any, // the user-specified block, which should be called if this clause becomes selected
-        @JvmField val onCancellationConstructor: OnCancellationConstructor?,
-        @JvmField var disposableHandle: DisposableHandle? = null
+        @JvmField val onCancellationConstructor: OnCancellationConstructor?
     ) {
+        @JvmField var disposableHandleOrSegment: Any? = null
+        @JvmField var indexInSegment: Int = -1
+
         /**
          * Tries to register the specified [select] instance in [clauseObject] and check
          * whether the registration succeeded or a rendezvous has happened during the registration.
@@ -788,6 +803,17 @@ internal open class SelectImplementation<R> constructor(
             }
         }
 
+        fun dispose() {
+            with(disposableHandleOrSegment) {
+                if (this is Segment<*>) {
+                    this.onCancellation(indexInSegment, null)
+                } else {
+                    this as DisposableHandle
+                    this.dispose()
+                }
+            }
+        }
+
         fun createOnCancellationAction(select: SelectInstance<*>, internalResult: Any?) =
             onCancellationConstructor?.invoke(select, param, internalResult)
     }

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

@@ -165,7 +165,8 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
         lockSuspend(owner)
     }
 
-    private suspend fun lockSuspend(owner: Any?) = suspendCancellableCoroutineReusable { cont ->
+    private suspend fun lockSuspend(owner: Any?) = suspendCancellableCoroutineReusable<Unit> { cont ->
+        cont as CancellableContinuationImpl<Unit>
         val contWithOwner = CancellableContinuationWithOwner(cont, owner)
         acquire(contWithOwner)
     }
@@ -205,7 +206,7 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
     )
 
     protected open fun onLockRegFunction(select: SelectInstance<*>, owner: Any?) {
-        onAcquireRegFunction(SelectInstanceWithOwner(select, owner), owner)
+        onAcquireRegFunction(SelectInstanceWithOwner(select as SelectInstanceInternal<*>, owner), owner)
     }
 
     protected open fun onLockProcessResult(owner: Any?, result: Any?): Any? {
@@ -214,10 +215,10 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
 
     private inner class CancellableContinuationWithOwner(
         @JvmField
-        val cont: CancellableContinuation<Unit>,
+        val cont: CancellableContinuationImpl<Unit>,
         @JvmField
         val owner: Any?
-    ) : CancellableContinuation<Unit> by cont {
+    ) : CancellableContinuation<Unit> by cont, SegmentDisposable 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) {
@@ -241,10 +242,10 @@ internal open class MutexImpl(locked: Boolean) : SemaphoreImpl(1, if (locked) 1
 
     private inner class SelectInstanceWithOwner<Q>(
         @JvmField
-        val select: SelectInstance<Q>,
+        val select: SelectInstanceInternal<Q>,
         @JvmField
         val owner: Any?
-    ) : SelectInstanceInternal<Q> by select as SelectInstanceInternal<Q> {
+    ) : SelectInstanceInternal<Q> by select {
         override fun trySelect(clauseObject: Any, result: Any?): Boolean {
             assert { this@MutexImpl.owner.value === NO_OWNER }
             return select.trySelect(clauseObject, result).also { success ->

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

@@ -185,7 +185,7 @@ internal open class SemaphoreImpl(private val permits: Int, acquiredPermits: Int
 
     private suspend fun acquireSlowPath() = suspendCancellableCoroutineReusable<Unit> sc@ { cont ->
         // Try to suspend.
-        if (addAcquireToQueue(cont)) return@sc
+        if (addAcquireToQueue(cont as SegmentDisposable)) return@sc
         // The suspension has been failed
         // due to the synchronous resumption mode.
         // Restart the whole `acquire`.
@@ -195,7 +195,7 @@ internal open class SemaphoreImpl(private val permits: Int, acquiredPermits: Int
     @JsName("acquireCont")
     protected fun acquire(waiter: CancellableContinuation<Unit>) = acquire(
         waiter = waiter,
-        suspend = { cont -> addAcquireToQueue(cont) },
+        suspend = { cont -> addAcquireToQueue(cont as SegmentDisposable) },
         onAcquired = { cont -> cont.resume(Unit, onCancellationRelease) }
     )
 
@@ -219,7 +219,7 @@ internal open class SemaphoreImpl(private val permits: Int, acquiredPermits: Int
     protected fun onAcquireRegFunction(select: SelectInstance<*>, ignoredParam: Any?) =
         acquire(
             waiter = select,
-            suspend = { s -> addAcquireToQueue(s) },
+            suspend = { s -> addAcquireToQueue(s as SegmentDisposable) },
             onAcquired = { s -> s.selectInRegistrationPhase(Unit) }
         )
 
@@ -281,24 +281,15 @@ internal open class SemaphoreImpl(private val permits: Int, acquiredPermits: Int
     /**
      * Returns `false` if the received permit cannot be used and the calling operation should restart.
      */
-    private fun addAcquireToQueue(waiter: Any): Boolean {
+    private fun addAcquireToQueue(waiter: SegmentDisposable): Boolean {
         val curTail = this.tail.value
         val enqIdx = enqIdx.getAndIncrement()
         val segment = this.tail.findSegmentAndMoveForward(id = enqIdx / SEGMENT_SIZE, startFrom = curTail,
             createNewSegment = ::createSegment).segment // cannot be closed
         val i = (enqIdx % SEGMENT_SIZE).toInt()
         // the regular (fast) path -- if the cell is empty, try to install continuation
         if (segment.cas(i, null, waiter)) { // installed continuation successfully
-            when (waiter) {
-                is CancellableContinuation<*> -> {
-                    waiter as CancellableContinuationImpl<*>
-                    waiter.invokeOnCancellation(segment, i)
-                }
-                is SelectInstance<*> -> {
-                    waiter.disposeOnCompletion(CancelSemaphoreAcquisitionHandler(segment, i))
-                }
-                else -> error("unexpected: $waiter")
-            }
+            waiter.disposeOnCancellation(segment, i)
             return true
         }
         // On CAS failure -- the cell must be either PERMIT or BROKEN
@@ -385,7 +376,7 @@ private class SemaphoreSegment(id: Long, prev: SemaphoreSegment?, pointers: Int)
 
     // Cleans the acquirer slot located by the specified index
     // and removes this segment physically if all slots are cleaned.
-    override fun invokeOnCancellation(index: Int, cause: Throwable?) {
+    override fun onCancellation(index: Int, cause: Throwable?) {
         // Clean the slot
         set(index, CANCELLED)
         // Remove this segment if needed

kotlinx-coroutines-core/common/test/CancellableContinuationHandlersTest.kt

@@ -164,10 +164,10 @@ class CancellableContinuationHandlersTest : TestBase() {
     @Test
     fun testSegmentAsHandler() = runTest {
         class MySegment : Segment<MySegment>(0, null, 0) {
-            override val maxSlots: Int get() = 0
+            override val numberOfSlots: Int get() = 0
 
             var invokeOnCancellationCalled = false
-            override fun invokeOnCancellation(index: Int, cause: Throwable?) {
+            override fun onCancellation(index: Int, cause: Throwable?) {
                 invokeOnCancellationCalled = true
             }
         }
@@ -177,7 +177,7 @@ class CancellableContinuationHandlersTest : TestBase() {
             suspendCancellableCoroutine<Unit> { c ->
                 expect(2)
                 c as CancellableContinuationImpl<*>
-                c.invokeOnCancellation(s, 0)
+                c.disposeOnCancellation(s, 0)
                 c.cancel()
             }
         } catch (e: CancellationException) {