IO: propagate channel cancellation to controller coroutine

Author: Sergey Mashkov

core/kotlinx-coroutines-io/src/main/kotlin/kotlinx/coroutines/experimental/io/ByteBufferChannel.kt

@@ -5,6 +5,7 @@ package kotlinx.coroutines.experimental.io
 import kotlinx.atomicfu.*
 import kotlinx.coroutines.experimental.*
 import kotlinx.coroutines.experimental.channels.*
+import kotlinx.coroutines.experimental.internal.*
 import kotlinx.coroutines.experimental.io.internal.*
 import kotlinx.coroutines.experimental.io.packet.*
 import kotlinx.io.core.*
@@ -50,6 +51,17 @@ internal class ByteBufferChannel(
     private var readPosition = 0
     private var writePosition = 0
 
+    @Volatile
+    private var attachedJob: Job? = null
+
+    internal fun attachJob(job: Job) {
+        attachedJob?.cancel()
+        attachedJob = job
+        job.invokeOnCompletion {
+            attachedJob = null
+        }
+    }
+
     override var readByteOrder: ByteOrder = ByteOrder.BIG_ENDIAN
     override var writeByteOrder: ByteOrder = ByteOrder.BIG_ENDIAN
         set(newOrder) {
@@ -92,10 +104,12 @@ internal class ByteBufferChannel(
             joining?.let { ensureClosedJoined(it) }
         }
 
+        if (cause != null) attachedJob?.cancel(cause)
+
         return true
     }
 
-    suspend override fun cancel(cause: Throwable?): Boolean {
+    override fun cancel(cause: Throwable?): Boolean {
         return close(cause ?: CancellationException("Channel has been cancelled"))
     }
 

core/kotlinx-coroutines-io/src/main/kotlin/kotlinx/coroutines/experimental/io/ByteReadChannel.kt

@@ -145,7 +145,7 @@ public interface ByteReadChannel {
      */
     suspend fun read(min: Int = 1, block: (ByteBuffer) -> Unit)
 
-    suspend fun cancel(cause: Throwable?): Boolean
+    fun cancel(cause: Throwable? = null): Boolean
 }
 
 suspend fun ByteReadChannel.joinTo(dst: ByteWriteChannel, closeOnEnd: Boolean) {

core/kotlinx-coroutines-io/src/main/kotlin/kotlinx/coroutines/experimental/io/ReaderJob.kt

@@ -34,7 +34,12 @@ fun reader(coroutineContext: CoroutineContext,
 fun reader(coroutineContext: CoroutineContext,
            autoFlush: Boolean = false,
            parent: Job? = null,
-           block: suspend ReaderScope.() -> Unit): ReaderJob = reader(coroutineContext, ByteChannel(autoFlush), parent, block)
+           block: suspend ReaderScope.() -> Unit): ReaderJob {
+    val channel = ByteChannel(autoFlush) as ByteBufferChannel
+    val job = reader(coroutineContext, channel, parent, block)
+    channel.attachJob(job)
+    return job
+}
 
 private class ReaderCoroutine(context: CoroutineContext, channel: ByteChannel)
     : ByteChannelCoroutine(context, channel), ReaderJob, ReaderScope

core/kotlinx-coroutines-io/src/main/kotlin/kotlinx/coroutines/experimental/io/WriterJob.kt

@@ -34,7 +34,12 @@ fun writer(coroutineContext: CoroutineContext,
 fun writer(coroutineContext: CoroutineContext,
            autoFlush: Boolean = false,
            parent: Job? = null,
-           block: suspend WriterScope.() -> Unit): WriterJob = writer(coroutineContext, ByteChannel(autoFlush), parent, block)
+           block: suspend WriterScope.() -> Unit): WriterJob {
+    val channel = ByteChannel(autoFlush) as ByteBufferChannel
+    val job = writer(coroutineContext, channel, parent, block)
+    channel.attachJob(job)
+    return job
+}
 
 private class WriterCoroutine(ctx: CoroutineContext, channel: ByteChannel)
     : ByteChannelCoroutine(ctx, channel), WriterScope, WriterJob

core/kotlinx-coroutines-io/src/test/kotlin/kotlinx/coroutines/experimental/io/ByteBufferChannelTest.kt

@@ -1088,6 +1088,25 @@ class ByteBufferChannelTest {
         }
     }
 
+    @Test
+    fun testCancelWriter() = runBlocking {
+        val sub = writer(DefaultDispatcher) {
+            delay(1000000L)
+        }
+
+        sub.channel.cancel()
+        sub.join()
+    }
+
+    @Test
+    fun testCancelReader() = runBlocking {
+        val sub = reader(DefaultDispatcher) {
+            delay(10000000L)
+        }
+
+        sub.channel.close(CancellationException())
+        sub.join()
+    }
 
     private inline fun buildPacket(block: ByteWritePacket.() -> Unit): ByteReadPacket {
         val builder = BytePacketBuilder(0, pktPool)