Introduced ReceiveChannel.consumeAll; all operators on ReceiveChannel

fully consume the original channel (using a helper consume extension)

Author: elizarov

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/AbstractChannel.kt

@@ -398,6 +398,11 @@ public abstract class AbstractSendChannel<E> : SendChannel<E> {
             remove()
         }
 
+        override fun resumeSendClosed(closed: Closed<*>) {
+            if (select.trySelect(null))
+                select.resumeSelectCancellableWithException(closed.sendException)
+        }
+
         override fun toString(): String = "SendSelect($pollResult)[$channel, $select]"
     }
 
@@ -407,6 +412,7 @@ public abstract class AbstractSendChannel<E> : SendChannel<E> {
         override val pollResult: Any? get() = element
         override fun tryResumeSend(idempotent: Any?): Any? = SEND_RESUMED
         override fun completeResumeSend(token: Any) { check(token === SEND_RESUMED) }
+        override fun resumeSendClosed(closed: Closed<*>) {}
     }
 }
 
@@ -567,6 +573,25 @@ public abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E>
         return if (result === POLL_FAILED) null else receiveOrNullResult(result)
     }
 
+    override fun consumeAll() {
+        close(cause = null)
+        // now cleanup send queue
+        consumeAllInternal()
+    }
+
+    // Note: this function is invoked when channel is already closed
+    protected open fun consumeAllInternal() {
+        val closed = closedForSend ?: error("Cannot happen")
+        while (true) {
+            val send = takeFirstSendOrPeekClosed() ?: error("Cannot happen")
+            if (send is Closed<*>) {
+                check(send === closed)
+                return // cleaned
+            }
+            send.resumeSendClosed(closed)
+        }
+    }
+
     public final override fun iterator(): ChannelIterator<E> = Itr(this)
 
     // ------ registerSelectReceive ------
@@ -909,6 +934,7 @@ public interface Send {
     val pollResult: Any? // E | Closed
     fun tryResumeSend(idempotent: Any?): Any?
     fun completeResumeSend(token: Any)
+    fun resumeSendClosed(closed: Closed<*>)
 }
 
 /**
@@ -922,7 +948,7 @@ public interface ReceiveOrClosed<in E> {
 }
 
 /**
- * Represents closed channel.
+ * Represents sender for a specific element.
  * @suppress **This is unstable API and it is subject to change.**
  */
 @Suppress("UNCHECKED_CAST")
@@ -932,6 +958,7 @@ public class SendElement(
 ) : LockFreeLinkedListNode(), Send {
     override fun tryResumeSend(idempotent: Any?): Any? = cont.tryResume(Unit, idempotent)
     override fun completeResumeSend(token: Any) = cont.completeResume(token)
+    override fun resumeSendClosed(closed: Closed<*>) = cont.resumeWithException(closed.sendException)
     override fun toString(): String = "SendElement($pollResult)[$cont]"
 }
 
@@ -951,6 +978,7 @@ public class Closed<in E>(
     override fun completeResumeSend(token: Any) { check(token === CLOSE_RESUMED) }
     override fun tryResumeReceive(value: E, idempotent: Any?): Any? = CLOSE_RESUMED
     override fun completeResumeReceive(token: Any) { check(token === CLOSE_RESUMED) }
+    override fun resumeSendClosed(closed: Closed<*>) = error("Should be never invoked")
     override fun toString(): String = "Closed[$closeCause]"
 }
 

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/ArrayBroadcastChannel.kt

@@ -209,7 +209,7 @@ class ArrayBroadcastChannel<E>(
         override val isBufferAlwaysFull: Boolean get() = error("Should not be used")
         override val isBufferFull: Boolean get() = error("Should not be used")
 
-        override fun close() {
+        override fun consumeAll() {
             if (close(cause = null))
                 broadcastChannel.updateHead(removeSub = this)
         }

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/ArrayChannel.kt

@@ -231,4 +231,18 @@ public open class ArrayChannel<E>(
             send!!.completeResumeSend(token!!)
         return result
     }
+
+    // Note: this function is invoked when channel is already closed
+    override fun consumeAllInternal() {
+        // clear buffer first
+        lock.withLock {
+            repeat(size) {
+                buffer[head] = 0
+                head = (head + 1) % capacity
+            }
+            size = 0
+        }
+        // then clean all queued senders
+        super.consumeAllInternal()
+    }
 }

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/BroadcastChannel.kt

@@ -76,10 +76,12 @@ public fun <E> BroadcastChannel(capacity: Int): BroadcastChannel<E> =
 /**
  * Return type for [BroadcastChannel.openSubscription] that can be used to [receive] elements from the
  * open subscription and to [close] it to unsubscribe.
+ *
+ * Note, that invocation of [consumeAll] also closes subscription.
  */
 public interface SubscriptionReceiveChannel<out T> : ReceiveChannel<T>, Closeable {
     /**
-     * Closes this subscription.
+     * Closes this subscription. This is a synonym for [consumeAll].
      */
-    public override fun close()
+    public override fun close() = consumeAll()
 }

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/Channel.kt

@@ -87,7 +87,9 @@ public interface SendChannel<in E> {
     /**
      * Closes this channel with an optional exceptional [cause].
      * This is an idempotent operation -- repeated invocations of this function have no effect and return `false`.
-     * Conceptually, its sends a special "close token" over this channel. Immediately after invocation of this function
+     * Conceptually, its sends a special "close token" over this channel.
+     *
+     * Immediately after invocation of this function
      * [isClosedForSend] starts returning `true`. However, [isClosedForReceive][ReceiveChannel.isClosedForReceive]
      * on the side of [ReceiveChannel] starts returning `true` only after all previously sent elements
      * are received.
@@ -192,6 +194,18 @@ public interface ReceiveChannel<out E> {
      * throws the original [close][SendChannel.close] cause exception if the channel has _failed_.
      */
     public operator fun iterator(): ChannelIterator<E>
+
+    /**
+     * Consumes all remaining elements from this channel. This function marks the channel as _consumed_
+     * by closing it normally (unless it was already closed) and removes all buffered sent elements from it.
+     *
+     * Immediately after invocation of this function [isClosedForReceive] and
+     * [isClosedForSend][SendChannel.isClosedForSend]
+     * on the side of [SendChannel] start returning `true`, so all attempts to send to this channel
+     * afterwards will throw [ClosedSendChannelException], while attempts to receive will throw
+     * [ClosedReceiveChannelException].
+     */
+    public fun consumeAll()
 }
 
 /**