Optimize CancellableContinuationImpl.invokeOnCancellation(..) for Segments

ndkoval · ndkoval · commit 34b7b987bd6a · 2023-02-10T19:29:53.000+01:00
diff --git a/benchmarks/src/jmh/kotlin/benchmarks/SequentialSemaphoreBenchmark.kt b/benchmarks/src/jmh/kotlin/benchmarks/SequentialSemaphoreBenchmark.kt
@@ -10,8 +10,8 @@ import org.openjdk.jmh.annotations.*
 import java.util.concurrent.TimeUnit
 import kotlin.test.*
 
-@Warmup(iterations = 5, time = 100)
-@Measurement(iterations = 10, time = 100)
+@Warmup(iterations = 5, time = 1)
+@Measurement(iterations = 10, time = 1)
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(TimeUnit.MILLISECONDS)
 @State(Scope.Benchmark)
diff --git a/kotlinx-coroutines-core/common/src/CancellableContinuationImpl.kt b/kotlinx-coroutines-core/common/src/CancellableContinuationImpl.kt
@@ -11,10 +11,18 @@ import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 import kotlin.jvm.*
 
+private const val DECISION_SHIFT = 29
+private const val INDEX_MASK = (1 shl DECISION_SHIFT) - 1
+private const val NO_INDEX = INDEX_MASK
 private const val UNDECIDED = 0
 private const val SUSPENDED = 1
 private const val RESUMED = 2
 
+private inline val Int.decision get() = this shr DECISION_SHIFT
+private inline val Int.index get() = this and INDEX_MASK
+@Suppress("NOTHING_TO_INLINE")
+private inline fun construct(decision: Int, index: Int) = (decision shl DECISION_SHIFT) + index
+
 @JvmField
 internal val RESUME_TOKEN = Symbol("RESUME_TOKEN")
 
@@ -56,9 +64,9 @@ internal open class CancellableContinuationImpl<in T>(
         |  RESUMED  |
         +-----------+
 
-        Note: both tryResume and trySuspend can be invoked at most once, first invocation wins
+        Note: both tryResume and trySuspend can be invoked at most once, first invocation wins.
      */
-    private val _decision = atomic(UNDECIDED)
+    private val _decisionAndIndex = atomic(construct(UNDECIDED, NO_INDEX))
 
     /*
        === Internal states ===
@@ -144,7 +152,7 @@ internal open class CancellableContinuationImpl<in T>(
             detachChild()
             return false
         }
-        _decision.value = UNDECIDED
+        _decisionAndIndex.value = construct(UNDECIDED, NO_INDEX)
         _state.value = Active
         return true
     }
@@ -194,10 +202,11 @@ internal open class CancellableContinuationImpl<in T>(
         _state.loop { state ->
             if (state !is NotCompleted) return false // false if already complete or cancelling
             // Active -- update to final state
-            val update = CancelledContinuation(this, cause, handled = state is CancelHandler)
+            val update = CancelledContinuation(this, cause, handled = state is CancelHandler || state is Segment<*>)
             if (!_state.compareAndSet(state, update)) return@loop // retry on cas failure
             // Invoke cancel handler if it was present
             (state as? CancelHandler)?.let { callCancelHandler(it, cause) }
+            (state as? Segment<*>)?.let { callSegmentOnCancellation(it, cause) }
             // Complete state update
             detachChildIfNonResuable()
             dispatchResume(resumeMode) // no need for additional cancellation checks
@@ -234,6 +243,13 @@ internal open class CancellableContinuationImpl<in T>(
     fun callCancelHandler(handler: CancelHandler, cause: Throwable?) =
         callCancelHandlerSafely { handler.invoke(cause) }
 
+    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) }
+    }
+
+
     fun callOnCancellation(onCancellation: (cause: Throwable) -> Unit, cause: Throwable) {
         try {
             onCancellation.invoke(cause)
@@ -253,19 +269,19 @@ internal open class CancellableContinuationImpl<in T>(
         parent.getCancellationException()
 
     private fun trySuspend(): Boolean {
-        _decision.loop { decision ->
-            when (decision) {
-                UNDECIDED -> if (this._decision.compareAndSet(UNDECIDED, SUSPENDED)) return true
+        _decisionAndIndex.loop { cur ->
+            when (cur.decision) {
+                UNDECIDED -> if (this._decisionAndIndex.compareAndSet(cur, construct(SUSPENDED, cur.index))) return true
                 RESUMED -> return false
                 else -> error("Already suspended")
             }
         }
     }
 
     private fun tryResume(): Boolean {
-        _decision.loop { decision ->
-            when (decision) {
-                UNDECIDED -> if (this._decision.compareAndSet(UNDECIDED, RESUMED)) return true
+        _decisionAndIndex.loop { cur ->
+            when (cur.decision) {
+                UNDECIDED -> if (this._decisionAndIndex.compareAndSet(cur, construct(RESUMED, cur.index))) return true
                 SUSPENDED -> return false
                 else -> error("Already resumed")
             }
@@ -350,14 +366,39 @@ internal open class CancellableContinuationImpl<in T>(
     override fun resume(value: T, onCancellation: ((cause: Throwable) -> Unit)?) =
         resumeImpl(value, resumeMode, onCancellation)
 
+    /**
+     * An optimized version for the code below that does not allocate
+     * a cancellation handler object and efficiently stores the specified
+     * [segment] and [index] in this [CancellableContinuationImpl].
+     * ```
+     * invokeOnCancellation { cause ->
+     *   segment.invokeOnCancellation(index, cause)
+     * }
+     * ```
+     */
+    internal fun invokeOnCancellation(segment: Segment<*>, index: Int) {
+        _decisionAndIndex.update {
+            check(it.index == NO_INDEX) {
+                "invokeOnCancellation should be invoked at most once"
+            }
+            construct(it.decision, index)
+        }
+        invokeOnCancellationImpl(segment)
+    }
+
     public override fun invokeOnCancellation(handler: CompletionHandler) {
         val cancelHandler = makeCancelHandler(handler)
+        invokeOnCancellationImpl(cancelHandler)
+    }
+
+    private fun invokeOnCancellationImpl(handler: Any) {
+        assert { handler is CancelHandler || handler is Segment<*> }
         _state.loop { state ->
             when (state) {
                 is Active -> {
-                    if (_state.compareAndSet(state, cancelHandler)) return // quit on cas success
+                    if (_state.compareAndSet(state, handler)) return // quit on cas success
                 }
-                is CancelHandler -> multipleHandlersError(handler, state)
+                is CancelHandler, is Segment<*> -> multipleHandlersError(handler, state)
                 is CompletedExceptionally -> {
                     /*
                      * Continuation was already cancelled or completed exceptionally.
@@ -371,7 +412,13 @@ internal open class CancellableContinuationImpl<in T>(
                      * because we play type tricks on Kotlin/JS and handler is not necessarily a function there
                      */
                     if (state is CancelledContinuation) {
-                        callCancelHandler(handler, (state as? CompletedExceptionally)?.cause)
+                        val cause: Throwable? = (state as? CompletedExceptionally)?.cause
+                        if (handler is CancelHandler) {
+                            callCancelHandler(handler, cause)
+                        } else {
+                            val segment = handler as Segment<*>
+                            callSegmentOnCancellation(segment, cause)
+                        }
                     }
                     return
                 }
@@ -380,14 +427,16 @@ internal open class CancellableContinuationImpl<in T>(
                      * Continuation was already completed, and might already have cancel handler.
                      */
                     if (state.cancelHandler != null) multipleHandlersError(handler, state)
-                    // BeforeResumeCancelHandler does not need to be called on a completed continuation
-                    if (cancelHandler is BeforeResumeCancelHandler) return
+                    // BeforeResumeCancelHandler and Segment.invokeOnCancellation(..)
+                    // do NOT need to be called on completed continuation.
+                    if (handler is BeforeResumeCancelHandler || handler is Segment<*>) return
+                    handler as CancelHandler
                     if (state.cancelled) {
                         // Was already cancelled while being dispatched -- invoke the handler directly
                         callCancelHandler(handler, state.cancelCause)
                         return
                     }
-                    val update = state.copy(cancelHandler = cancelHandler)
+                    val update = state.copy(cancelHandler = handler)
                     if (_state.compareAndSet(state, update)) return // quit on cas success
                 }
                 else -> {
@@ -396,15 +445,16 @@ internal open class CancellableContinuationImpl<in T>(
                      * Change its state to CompletedContinuation, unless we have BeforeResumeCancelHandler which
                      * does not need to be called in this case.
                      */
-                    if (cancelHandler is BeforeResumeCancelHandler) return
-                    val update = CompletedContinuation(state, cancelHandler = cancelHandler)
+                    if (handler is BeforeResumeCancelHandler || handler is Segment<*>) return
+                    handler as CancelHandler
+                    val update = CompletedContinuation(state, cancelHandler = handler)
                     if (_state.compareAndSet(state, update)) return // quit on cas success
                 }
             }
         }
     }
 
-    private fun multipleHandlersError(handler: CompletionHandler, state: Any?) {
+    private fun multipleHandlersError(handler: Any, state: Any?) {
         error("It's prohibited to register multiple handlers, tried to register $handler, already has $state")
     }
 
diff --git a/kotlinx-coroutines-core/common/src/internal/ConcurrentLinkedList.kt b/kotlinx-coroutines-core/common/src/internal/ConcurrentLinkedList.kt
@@ -188,7 +188,15 @@ internal abstract class ConcurrentLinkedListNode<N : ConcurrentLinkedListNode<N>
  * Essentially, this is a node in the Michael-Scott queue algorithm,
  * but with maintaining [prev] pointer for efficient [remove] implementation.
  */
-internal abstract class Segment<S : Segment<S>>(val id: Long, prev: S?, pointers: Int): ConcurrentLinkedListNode<S>(prev) {
+internal abstract class Segment<S : Segment<S>>(val id: Long, prev: S?, pointers: Int) :
+    ConcurrentLinkedListNode<S>(prev),
+    // Segments typically store waiting continuations. Thus, on cancellation, the corresponding
+    // slot should be cleaned and the segment should be removed if it becomes full of cancelled cells.
+    // To install such a handler efficiently, without creating an extra object, we allow storing
+    // segments as cancellation handlers in [CancellableContinuationImpl] state, putting the slot
+    // index in another field. The details are here: https://github.com/Kotlin/kotlinx.coroutines/pull/3084.
+    NotCompleted
+{
     /**
      * This property should return the number of slots in this segment,
      * it is used to define whether the segment is logically removed.
@@ -212,6 +220,13 @@ internal abstract class Segment<S : Segment<S>>(val id: Long, prev: S?, pointers
     // returns `true` if this segment is logically removed after the decrement.
     internal fun decPointers() = cleanedAndPointers.addAndGet(-(1 shl POINTERS_SHIFT)) == numberOfSlots && !isTail
 
+    /**
+     * This function is invoked on continuation cancellation when this segment
+     * with the specified [index] are installed as cancellation handler via
+     * `CancellableContinuationImpl.invokeOnCancellation(Segment, Int)`.
+     */
+    internal open fun invokeOnCancellation(index: Int, cause: Throwable?) {}
+
     /**
      * Invoked on each slot clean-up; should not be invoked twice for the same slot.
      */
diff --git a/kotlinx-coroutines-core/common/src/sync/Semaphore.kt b/kotlinx-coroutines-core/common/src/sync/Semaphore.kt
@@ -291,7 +291,8 @@ internal open class SemaphoreImpl(private val permits: Int, acquiredPermits: Int
         if (segment.cas(i, null, waiter)) { // installed continuation successfully
             when (waiter) {
                 is CancellableContinuation<*> -> {
-                    waiter.invokeOnCancellation(CancelSemaphoreAcquisitionHandler(segment, i).asHandler)
+                    waiter as CancellableContinuationImpl<*>
+                    waiter.invokeOnCancellation(segment, i)
                 }
                 is SelectInstance<*> -> {
                     waiter.disposeOnCompletion(CancelSemaphoreAcquisitionHandler(segment, i))
@@ -362,20 +363,7 @@ internal open class SemaphoreImpl(private val permits: Int, acquiredPermits: Int
     }
 }
 
-private class CancelSemaphoreAcquisitionHandler(
-    private val segment: SemaphoreSegment,
-    private val index: Int
-) : CancelHandler(), DisposableHandle {
-    override fun invoke(cause: Throwable?) = dispose()
-
-    override fun dispose() {
-        segment.cancel(index)
-    }
-
-    override fun toString() = "CancelSemaphoreAcquisitionHandler[$segment, $index]"
-}
-
-private fun createSegment(id: Long, prev: SemaphoreSegment) = SemaphoreSegment(id, prev, 0)
+private fun createSegment(id: Long, prev: SemaphoreSegment?) = SemaphoreSegment(id, prev, 0)
 
 private class SemaphoreSegment(id: Long, prev: SemaphoreSegment?, pointers: Int) : Segment<SemaphoreSegment>(id, prev, pointers) {
     val acquirers = atomicArrayOfNulls<Any?>(SEGMENT_SIZE)
@@ -397,7 +385,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.
-    fun cancel(index: Int) {
+    override fun invokeOnCancellation(index: Int, cause: Throwable?) {
         // Clean the slot
         set(index, CANCELLED)
         // Remove this segment if needed
diff --git a/kotlinx-coroutines-core/common/test/CancellableContinuationHandlersTest.kt b/kotlinx-coroutines-core/common/test/CancellableContinuationHandlersTest.kt
@@ -6,6 +6,7 @@
 
 package kotlinx.coroutines
 
+import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
 import kotlin.test.*
 
@@ -159,4 +160,31 @@ class CancellableContinuationHandlersTest : TestBase() {
         }
         finish(3)
     }
+
+    @Test
+    fun testSegmentAsHandler() = runTest {
+        class MySegment : Segment<MySegment>(0, null, 0) {
+            override val maxSlots: Int get() = 0
+
+            var invokeOnCancellationCalled = false
+            override fun invokeOnCancellation(index: Int, cause: Throwable?) {
+                invokeOnCancellationCalled = true
+            }
+        }
+        val s = MySegment()
+        expect(1)
+        try {
+            suspendCancellableCoroutine<Unit> { c ->
+                expect(2)
+                c as CancellableContinuationImpl<*>
+                c.invokeOnCancellation(s, 0)
+                c.cancel()
+            }
+        } catch (e: CancellationException) {
+            expect(3)
+        }
+        expect(4)
+        check(s.invokeOnCancellationCalled)
+        finish(5)
+    }
 }
