IO: introduce simple ByteReadChannel.discard implementation

Author: Sergey Mashkov

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

@@ -1477,6 +1477,43 @@ internal class ByteBufferChannel(
         }
     }
 
+    override suspend fun discard(max: Long): Long {
+        require(max >= 0) { "max shouldn't be negative: $max" }
+
+        var discarded = 0L
+
+        reading {
+            val n = it.tryReadAtMost(minOf(Int.MAX_VALUE.toLong(), max).toInt())
+            bytesRead(it, n)
+            discarded += n
+            true
+        }
+
+        if (discarded == max || isClosedForRead) return discarded
+
+        return discardSuspend(discarded, max)
+    }
+
+    private suspend fun discardSuspend(discarded0: Long, max: Long): Long {
+        var discarded = discarded0
+
+        while (discarded < max) {
+            val rc = reading {
+                val n = it.tryReadAtMost(minOf(Int.MAX_VALUE.toLong(), max - discarded).toInt())
+                bytesRead(it, n)
+                discarded += n
+
+                true
+            }
+
+            if (!rc) {
+                if (isClosedForRead || !readSuspend(1)) break
+            }
+        }
+
+        return discarded
+    }
+
     private suspend fun readBlockSuspend(min: Int, block: (ByteBuffer) -> Unit) {
         if (!readSuspend(min.coerceAtLeast(1))) {
             if (min > 0)

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

@@ -156,6 +156,13 @@ public interface ByteReadChannel {
      * @see ByteWriteChannel.close
      */
     fun cancel(cause: Throwable? = null): Boolean
+
+    /**
+     * Discard up to [max] bytes
+     *
+     * @return number of bytes were discarded
+     */
+    suspend fun discard(max: Long = Long.MAX_VALUE): Long
 }
 
 suspend fun ByteReadChannel.joinTo(dst: ByteWriteChannel, closeOnEnd: Boolean) {

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

@@ -469,4 +469,88 @@ class ByteBufferChannelScenarioTest : TestBase() {
         finish(6)
     }
 
+    @Test
+    fun testDiscardExisting() = runBlocking {
+        launch(coroutineContext) {
+            expect(1)
+            ch.writeInt(1)
+            ch.writeInt(2)
+            expect(2)
+        }
+
+        yield()
+        expect(3)
+
+        assertEquals(4, ch.discard(4))
+        assertEquals(2, ch.readInt())
+
+        finish(4)
+    }
+
+    @Test
+    fun testDiscardPartiallyExisting() = runBlocking {
+        ch.writeInt(1)
+
+        launch(coroutineContext) {
+            expect(1)
+            assertEquals(8, ch.discard(8))
+            expect(3)
+        }
+
+        yield()
+        expect(2)
+
+        ch.writeInt(2)
+        yield()
+
+        expect(4)
+        assertEquals(0, ch.availableForRead)
+        finish(5)
+    }
+
+    @Test
+    fun testDiscardPartiallyExisting2() = runBlocking {
+        launch(coroutineContext) {
+            expect(1)
+            assertEquals(8, ch.discard(8))
+            expect(4)
+        }
+
+        yield()
+
+        expect(2)
+        ch.writeInt(1)
+        yield()
+        expect(3)
+        assertEquals(0, ch.availableForRead)
+
+        ch.writeInt(2)
+        yield()
+        expect(5)
+        assertEquals(0, ch.availableForRead)
+        finish(6)
+    }
+
+    @Test
+    fun testDiscardClose() = runBlocking {
+        launch(coroutineContext) {
+            expect(1)
+            assertEquals(8, ch.discard())
+            expect(4)
+        }
+
+        yield()
+
+        expect(2)
+        ch.writeInt(1)
+        yield()
+        ch.writeInt(2)
+        yield()
+
+        expect(3)
+        ch.close()
+        yield()
+
+        finish(5)
+    }
 }