RendezvousChannel is now an open class with a afterClose and onEnqueuedReceive/onCancelledReceive extension points

elizarov · elizarov · commit 0a7883910513 · 2017-02-15T17:55:29.000+03:00
More efficient cancellable suspending functions for Channel.send/receive and Mutex.lock
diff --git a/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/CancellableContinuation.kt b/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/CancellableContinuation.kt
@@ -16,6 +16,7 @@
 
 package kotlinx.coroutines.experimental
 
+import kotlinx.coroutines.experimental.internal.LockFreeLinkedListNode
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater
 import kotlin.coroutines.experimental.Continuation
 import kotlin.coroutines.experimental.ContinuationInterceptor
@@ -118,8 +119,27 @@ public inline suspend fun <T> suspendCancellableCoroutine(
         safe.getResult()
     }
 
+
+/**
+ * Removes a given node on cancellation.
+ * @suppress **This is unstable API and it is subject to change.**
+ */
+public fun CancellableContinuation<*>.removeOnCancel(node: LockFreeLinkedListNode): Job.Registration =
+    onCompletion(RemoveOnCancel(this, node))
+
 // --------------- implementation details ---------------
 
+private class RemoveOnCancel(
+    cont: CancellableContinuation<*>,
+    val node: LockFreeLinkedListNode
+) : JobNode<CancellableContinuation<*>>(cont)  {
+    override fun invoke(reason: Throwable?) {
+        if (job.isCancelled)
+            node.remove()
+    }
+    override fun toString() = "RemoveOnCancel[$node]"
+}
+
 @PublishedApi
 internal fun getParentJobOrAbort(cont: Continuation<*>): Job? {
     val job = cont.context[Job]
diff --git a/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/Job.kt b/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/Job.kt
@@ -185,13 +185,6 @@ public fun Job.unregisterOnCompletion(registration: Job.Registration): Job.Regis
 public fun Job.cancelFutureOnCompletion(future: Future<*>): Job.Registration =
     onCompletion(CancelFutureOnCompletion(this, future))
 
-/**
- * Removes a given node on completion.
- * @suppress **This is unstable API and it is subject to change.**
- */
-public fun Job.removeOnCompletion(node: LockFreeLinkedListNode): Job.Registration =
-    onCompletion(RemoveOnCompletion(this, node))
-
 /**
  * @suppress **Deprecated**: `join` is now a member function of `Job`.
  */
@@ -326,13 +319,13 @@ public open class JobSupport(active: Boolean) : AbstractCoroutineContextElement(
         var completionException: Throwable? = null
         when (expect) {
             // SINGLE/SINGLE+ state -- one completion handler (common case)
-            is JobNode -> try {
+            is JobNode<*> -> try {
                 expect.invoke(cause)
             } catch (ex: Throwable) {
                 completionException = ex
             }
             // LIST state -- a list of completion handlers
-            is NodeList -> expect.forEach<JobNode> { node ->
+            is NodeList -> expect.forEach<JobNode<*>> { node ->
                 try {
                     node.invoke(cause)
                 } catch (ex: Throwable) {
@@ -401,7 +394,7 @@ public open class JobSupport(active: Boolean) : AbstractCoroutineContextElement(
     }
 
     final override fun onCompletion(handler: CompletionHandler): Job.Registration {
-        var nodeCache: JobNode? = null
+        var nodeCache: JobNode<*>? = null
         while (true) { // lock-free loop on state
             val state = this.state
             when {
@@ -417,7 +410,7 @@ public open class JobSupport(active: Boolean) : AbstractCoroutineContextElement(
                     STATE.compareAndSet(this, state, NodeList(active = 0))
                 }
                 // SINGLE/SINGLE+ state -- one completion handler
-                state is JobNode -> {
+                state is JobNode<*> -> {
                     // try to promote it to list (SINGLE+ state)
                     state.addFirstIfEmpty(NodeList(active = 1))
                     // it must be in SINGLE+ state or state has changed (node could have need removed from state)
@@ -451,13 +444,13 @@ public open class JobSupport(active: Boolean) : AbstractCoroutineContextElement(
         cont.unregisterOnCompletion(onCompletion(ResumeOnCompletion(this, cont)))
     }
 
-    internal fun removeNode(node: JobNode) {
+    internal fun removeNode(node: JobNode<*>) {
         // remove logic depends on the state of the job
         while (true) { // lock-free loop on job state
             val state = this.state
             when (state) {
                 // SINGE/SINGLE+ state -- one completion handler
-                is JobNode -> {
+                is JobNode<*> -> {
                     if (state !== this) return // a different job node --> we were already removed
                     // try remove and revert back to empty state
                     if (STATE.compareAndSet(this, state, EmptyActive)) return
@@ -493,8 +486,8 @@ public open class JobSupport(active: Boolean) : AbstractCoroutineContextElement(
      */
     protected open fun afterCompletion(state: Any?) {}
 
-    private fun makeNode(handler: CompletionHandler): JobNode =
-            (handler as? JobNode)?.also { require(it.job === this) }
+    private fun makeNode(handler: CompletionHandler): JobNode<*> =
+            (handler as? JobNode<*>)?.also { require(it.job === this) }
                     ?: InvokeOnCompletion(this, handler)
 
     // for nicer debugging
@@ -524,7 +517,7 @@ public open class JobSupport(active: Boolean) : AbstractCoroutineContextElement(
             append(if (isActive) "{Active}" else "{New}")
             append("[")
             var first = true
-            this@NodeList.forEach<JobNode> { node ->
+            this@NodeList.forEach<JobNode<*>> { node ->
                 if (first) first = false else append(", ")
                 append(node)
             }
@@ -565,8 +558,8 @@ private class Empty(override val isActive: Boolean) : JobSupport.Incomplete {
     override fun toString(): String = "Empty{${if (isActive) "Active" else "New" }}"
 }
 
-internal abstract class JobNode(
-    val job: Job
+internal abstract class JobNode<out J : Job>(
+    val job: J
 ) : LockFreeLinkedListNode(), Job.Registration, CompletionHandler, JobSupport.Incomplete {
     final override val isActive: Boolean get() = true
     // if unregister is called on this instance, then Job was an instance of JobSupport that added this node it itself
@@ -578,39 +571,39 @@ internal abstract class JobNode(
 private class InvokeOnCompletion(
     job: Job,
     val handler: CompletionHandler
-) : JobNode(job)  {
+) : JobNode<Job>(job)  {
     override fun invoke(reason: Throwable?) = handler.invoke(reason)
     override fun toString() = "InvokeOnCompletion[${handler::class.java.name}@${Integer.toHexString(System.identityHashCode(handler))}]"
 }
 
 private class ResumeOnCompletion(
     job: Job,
     val continuation: Continuation<Unit>
-) : JobNode(job)  {
+) : JobNode<Job>(job)  {
     override fun invoke(reason: Throwable?) = continuation.resume(Unit)
     override fun toString() = "ResumeOnCompletion[$continuation]"
 }
 
 private class UnregisterOnCompletion(
     job: Job,
     val registration: Job.Registration
-) : JobNode(job) {
+) : JobNode<Job>(job) {
     override fun invoke(reason: Throwable?) = registration.unregister()
     override fun toString(): String = "UnregisterOnCompletion[$registration]"
 }
 
 private class CancelOnCompletion(
     parentJob: Job,
     val subordinateJob: Job
-) : JobNode(parentJob) {
+) : JobNode<Job>(parentJob) {
     override fun invoke(reason: Throwable?) { subordinateJob.cancel(reason) }
     override fun toString(): String = "CancelOnCompletion[$subordinateJob]"
 }
 
 private class CancelFutureOnCompletion(
     job: Job,
     val future: Future<*>
-) : JobNode(job)  {
+) : JobNode<Job>(job)  {
     override fun invoke(reason: Throwable?) {
         // Don't interrupt when cancelling future on completion, because no one is going to reset this
         // interruption flag and it will cause spurious failures elsewhere
@@ -619,14 +612,6 @@ private class CancelFutureOnCompletion(
     override fun toString() = "CancelFutureOnCompletion[$future]"
 }
 
-private class RemoveOnCompletion(
-    job: Job,
-    val node: LockFreeLinkedListNode
-) : JobNode(job)  {
-    override fun invoke(reason: Throwable?) { node.remove() }
-    override fun toString() = "RemoveOnCompletion[$node]"
-}
-
 private class JobImpl(parent: Job? = null) : JobSupport(true) {
     init { initParentJob(parent) }
 }
diff --git a/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/Mutex.kt b/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/Mutex.kt
@@ -118,7 +118,7 @@ public class Mutex(locked: Boolean = false) {
                     if (state.addLastIf(waiter, { this.state === state })) {
                         // added to waiter list!
                         cont.initCancellability()
-                        cont.removeOnCompletion(waiter)
+                        cont.removeOnCancel(waiter)
                         return@sc
                     }
                 }
diff --git a/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/AbstractChannel.kt b/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/AbstractChannel.kt
@@ -19,7 +19,7 @@ package kotlinx.coroutines.experimental.channels
 import kotlinx.coroutines.experimental.CancellableContinuation
 import kotlinx.coroutines.experimental.internal.LockFreeLinkedListHead
 import kotlinx.coroutines.experimental.internal.LockFreeLinkedListNode
-import kotlinx.coroutines.experimental.removeOnCompletion
+import kotlinx.coroutines.experimental.removeOnCancel
 import kotlinx.coroutines.experimental.suspendCancellableCoroutine
 
 /**
@@ -72,17 +72,17 @@ public abstract class AbstractChannel<E> : Channel<E> {
 
     // ------ SendChannel ------
 
-    override val isClosedForSend: Boolean get() = closedForSend != null
-    override val isFull: Boolean get() = queue.next() !is ReceiveOrClosed<*> && isBufferFull
+    public final override val isClosedForSend: Boolean get() = closedForSend != null
+    public final override val isFull: Boolean get() = queue.next() !is ReceiveOrClosed<*> && isBufferFull
 
-    suspend override fun send(element: E) {
+    public final override suspend fun send(element: E) {
         // fast path -- try offer non-blocking
         if (offer(element)) return
         // slow-path does suspend
         return sendSuspend(element)
     }
 
-    override fun offer(element: E): Boolean {
+    public final override fun offer(element: E): Boolean {
         val result = offerInternal(element)
         return when {
             result === OFFER_SUCCESS -> true
@@ -96,7 +96,7 @@ public abstract class AbstractChannel<E> : Channel<E> {
         loop@ while (true) {
             if (enqueueSend(send)) {
                 cont.initCancellability() // make it properly cancellable
-                cont.removeOnCompletion(send)
+                cont.removeOnCancel(send)
                 return@sc
             }
             // hm... something is not right. try to offer
@@ -120,13 +120,16 @@ public abstract class AbstractChannel<E> : Channel<E> {
         else
             queue.addLastIfPrev(send, { it !is ReceiveOrClosed<*> })
 
-    override fun close(cause: Throwable?): Boolean {
+    public final override fun close(cause: Throwable?): Boolean {
         val closed = Closed<E>(cause)
         while (true) {
             val receive = takeFirstReceiveOrPeekClosed()
             if (receive == null) {
                 // queue empty or has only senders -- try add last "Closed" item to the queue
-                if (queue.addLastIfPrev(closed, { it !is ReceiveOrClosed<*> })) return true
+                if (queue.addLastIfPrev(closed, { it !is ReceiveOrClosed<*> })) {
+                    afterClose(cause)
+                    return true
+                }
                 continue // retry on failure
             }
             if (receive is Closed<*>) return false // already marked as closed -- nothing to do
@@ -135,6 +138,11 @@ public abstract class AbstractChannel<E> : Channel<E> {
         }
     }
 
+    /**
+     * Invoked after successful [close].
+     */
+    protected open fun afterClose(cause: Throwable?) {}
+
     /**
      * Retrieves first receiving waiter from the queue or returns closed token.
      */
@@ -143,11 +151,11 @@ public abstract class AbstractChannel<E> : Channel<E> {
 
     // ------ ReceiveChannel ------
 
-    override val isClosedForReceive: Boolean get() = closedForReceive != null && isBufferEmpty
-    override val isEmpty: Boolean get() = queue.next() !is Send && isBufferEmpty
+    public final override val isClosedForReceive: Boolean get() = closedForReceive != null && isBufferEmpty
+    public final override val isEmpty: Boolean get() = queue.next() !is Send && isBufferEmpty
 
     @Suppress("UNCHECKED_CAST")
-    suspend override fun receive(): E {
+    public final override suspend fun receive(): E {
         // fast path -- try poll non-blocking
         val result = pollInternal()
         if (result !== POLL_EMPTY) return receiveResult(result)
@@ -167,7 +175,7 @@ public abstract class AbstractChannel<E> : Channel<E> {
         while (true) {
             if (enqueueReceive(receive)) {
                 cont.initCancellability() // make it properly cancellable
-                cont.removeOnCompletion(receive)
+                removeReceiveOnCancel(cont, receive)
                 return@sc
             }
             // hm... something is not right. try to poll
@@ -183,14 +191,16 @@ public abstract class AbstractChannel<E> : Channel<E> {
         }
     }
 
-    private fun enqueueReceive(receive: Receive<E>) =
-        if (hasBuffer)
-            queue.addLastIfPrevAndIf(receive, { it !is Send }, { isBufferEmpty })
-        else
+    private fun enqueueReceive(receive: Receive<E>): Boolean {
+        val result = if (hasBuffer)
+            queue.addLastIfPrevAndIf(receive, { it !is Send }, { isBufferEmpty }) else
             queue.addLastIfPrev(receive, { it !is Send })
+        if (result) onEnqueuedReceive()
+        return result
+    }
 
     @Suppress("UNCHECKED_CAST")
-    suspend override fun receiveOrNull(): E? {
+    public final override suspend fun receiveOrNull(): E? {
         // fast path -- try poll non-blocking
         val result = pollInternal()
         if (result !== POLL_EMPTY) return receiveOrNullResult(result)
@@ -213,7 +223,7 @@ public abstract class AbstractChannel<E> : Channel<E> {
         while (true) {
             if (enqueueReceive(receive)) {
                 cont.initCancellability() // make it properly cancellable
-                cont.removeOnCompletion(receive)
+                removeReceiveOnCancel(cont, receive)
                 return@sc
             }
             // hm... something is not right. try to poll
@@ -233,12 +243,12 @@ public abstract class AbstractChannel<E> : Channel<E> {
     }
 
     @Suppress("UNCHECKED_CAST")
-    override fun poll(): E? {
+    public final override fun poll(): E? {
         val result = pollInternal()
         return if (result === POLL_EMPTY) null else receiveOrNullResult(result)
     }
 
-    override fun iterator(): ChannelIterator<E> = Iterator(this)
+    public final override fun iterator(): ChannelIterator<E> = Iterator(this)
 
     /**
      * Retrieves first sending waiter from the queue or returns closed token.
@@ -262,6 +272,23 @@ public abstract class AbstractChannel<E> : Channel<E> {
         override fun toString(): String = string
     }
 
+    private fun removeReceiveOnCancel(cont: CancellableContinuation<*>, receive: Receive<*>) {
+        cont.onCompletion {
+            if (cont.isCancelled && receive.remove())
+                onCancelledReceive()
+        }
+    }
+
+    /**
+     * Invoked when receiver is successfully enqueued to the queue of waiting receivers.
+     */
+    protected open fun onEnqueuedReceive() {}
+
+    /**
+     * Invoked when enqueued receiver was successfully cancelled.
+     */
+    protected open fun onCancelledReceive() {}
+
     private class Iterator<E>(val channel: AbstractChannel<E>) : ChannelIterator<E> {
         var result: Any? = POLL_EMPTY // E | POLL_CLOSED | POLL_EMPTY
 
@@ -288,7 +315,7 @@ public abstract class AbstractChannel<E> : Channel<E> {
             while (true) {
                 if (channel.enqueueReceive(receive)) {
                     cont.initCancellability() // make it properly cancellable
-                    cont.removeOnCompletion(receive)
+                    channel.removeReceiveOnCancel(cont, receive)
                     return@sc
                 }
                 // hm... something is not right. try to poll
diff --git a/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/RendezvousChannel.kt b/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/RendezvousChannel.kt

Original file line number	Diff line number	Diff line change
`@@ -118,7 +118,7 @@ public class Mutex(locked: Boolean = false) {`
`118`	`118`	`if (state.addLastIf(waiter, { this.state === state })) {`
`119`	`119`	`// added to waiter list!`
`120`	`120`	`cont.initCancellability()`
`121`		`- cont.removeOnCompletion(waiter)`
	`121`	`+ cont.removeOnCancel(waiter)`
`122`	`122`	`return@sc`
`123`	`123`	`}`
`124`	`124`	`}`