IO: add BufferView internal, add ByteChannel read and write with lambda

Author: Sergey Mashkov

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

@@ -4,9 +4,11 @@ package kotlinx.coroutines.experimental.io
 
 import kotlinx.coroutines.experimental.CancellableContinuation
 import kotlinx.coroutines.experimental.channels.ClosedReceiveChannelException
+import kotlinx.coroutines.experimental.io.buffers.*
 import kotlinx.coroutines.experimental.io.internal.*
 import kotlinx.coroutines.experimental.io.packet.*
 import kotlinx.coroutines.experimental.suspendCancellableCoroutine
+import java.io.*
 import java.nio.BufferOverflowException
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater
 
@@ -370,26 +372,60 @@ internal class ByteBufferChannel(
     suspend override fun readPacket(size: Int, headerSizeHint: Int): ByteReadPacket {
         closed?.cause?.let { throw it }
 
-        if (size == 0) return ByteReadPacketEmpty
+        if (size == 0) return ByteReadPacketViewBased.Empty
 
-        val builder = ByteWritePacketImpl(headerSizeHint, BufferPool)
+        val builder = ByteWritePacketImpl(headerSizeHint, BufferView.Pool)
         val buffer = BufferPool.borrow()
+        var remaining = size
+
+        try {
+            while (remaining > 0) {
+                buffer.clear()
+                if (buffer.remaining() > remaining) {
+                    buffer.limit(remaining)
+                }
+
+                val rc = readAsMuchAsPossible(buffer)
+                if (rc == 0) break
+
+                buffer.flip()
+                builder.writeFully(buffer)
+
+                remaining -= rc
+            }
+        } catch (t: Throwable) {
+            BufferPool.recycle(buffer)
+            builder.release()
+            throw t
+        }
+
+        return if (remaining == 0) {
+            BufferPool.recycle(buffer)
+            builder.build()
+        } else {
+            readPacketSuspend(remaining, builder, buffer)
+        }
+    }
+
+    private suspend fun readPacketSuspend(size: Int, builder: ByteWritePacketImpl, buffer: ByteBuffer): ByteReadPacket {
+        var remaining = size
 
         try {
-            var remaining = size
             while (remaining > 0) {
                 buffer.clear()
                 if (buffer.remaining() > remaining) {
                     buffer.limit(remaining)
                 }
 
                 val rc = readFully(buffer)
+
                 buffer.flip()
                 builder.writeFully(buffer)
 
                 remaining -= rc
             }
 
+
             return builder.build()
         } catch (t: Throwable) {
             builder.release()
@@ -837,82 +873,125 @@ internal class ByteBufferChannel(
         closed?.sendException?.let { packet.release(); throw it }
 
         when (packet) {
-            is ByteReadPacketEmpty -> return
-            is ByteReadPacketSingle -> writeSingleBufferPacket(packet)
-            is ByteReadPacketImpl -> writeMultipleBufferPacket(packet)
+            ByteReadPacketEmpty -> return
+            is ByteReadPacketViewBased -> writeViewBasedPacket(packet)
             else -> writeExternalPacket(packet)
         }
     }
 
-    private suspend fun writeSingleBufferPacket(packet: ByteReadPacketSingle) {
-        val buffer = packet.steal()
-        val t = try {
-            writeAsMuchAsPossible(buffer)
-            null
-        } catch (t: Throwable) {
-            t
-        }
+    override suspend fun write(min: Int, block: (ByteBuffer) -> Unit) {
+        require(min > 0) { "min should be positive"}
 
-        if (t != null) {
-            packet.pool.recycle(buffer)
-            throw t
+        var written = false
+
+        writing {
+            if (it.availableForWrite >= min) {
+                val position = this.position()
+                val l = this.limit()
+                block(this)
+                if (l != this.limit()) throw IllegalStateException("buffer limit modified")
+                val delta = position() - position
+                if (delta < 0) throw IllegalStateException("position has been moved backward: pushback is not supported")
+
+                if (!it.tryWriteExact(delta)) throw IllegalStateException()
+                bytesWritten(it, delta)
+                written = true
+            }
         }
 
-        if (buffer.hasRemaining()) {
-            return writeSingleBufferPacketSuspend(buffer, packet)
+        if (!written) {
+            return writeBlockSuspend(min, block)
         }
+    }
 
-        packet.pool.recycle(buffer)
+    private suspend fun writeBlockSuspend(min: Int, block: (ByteBuffer) -> Unit) {
+        writeSuspend(min)
+        write(min, block)
     }
 
-    private suspend fun writeMultipleBufferPacket(packet: ByteReadPacketImpl) {
-        var buffer: ByteBuffer? = null
+    override suspend fun read(min: Int, block: (ByteBuffer) -> Unit) {
+        require(min > 0) { "min should be positive" }
 
-        val t = try {
-            while (packet.remaining > 0) {
-                buffer = packet.steal()
-                writeAsMuchAsPossible(buffer)
-                if (buffer.hasRemaining()) break
-                packet.pool.recycle(buffer)
-                buffer = null
-            }
-            null
-        } catch (t: Throwable) { t }
+        val read = reading {
+            if (it.availableForRead >= min) {
+                val position = this.position()
+                val l = this.limit()
+                block(this)
+                if (l != this.limit()) throw IllegalStateException("buffer limit modified")
+                val delta = position() - position
+                if (delta < 0) throw IllegalStateException("position has been moved backward: pushback is not supported")
 
-        if (t != null) {
-            buffer?.let { packet.pool.recycle(it) }
-            packet.release()
-            throw t
+                if (!it.tryReadExact(delta)) throw IllegalStateException()
+                bytesRead(it, delta)
+                true
+            }
+            else false
         }
 
-        if (buffer != null) {
-            return writeMultipleBufferPacketSuspend(buffer, packet)
+        if (!read) {
+            readBlockSuspend(min, block)
         }
+    }
 
-        packet.release()
+    private suspend fun readBlockSuspend(min: Int, block: (ByteBuffer) -> Unit) {
+        if (!readSuspend(min)) throw EOFException("Got EOF but at least $min bytes were expected")
+        read(min, block)
     }
 
-    private suspend fun writeSingleBufferPacketSuspend(buffer: ByteBuffer, packet: ByteReadPacketSingle) {
-        try {
-            writeFully(buffer)
-        } finally {
-            packet.pool.recycle(buffer)
+    private suspend fun writeViewBasedPacket(packet: ByteReadPacketViewBased) {
+        while (!packet.isEmpty) {
+            val node = packet.steal() ?: break
+
+            var written: Int
+
+            while (node.readRemaining > 0) {
+                try {
+                    written = 0
+                    writing {
+                        val size = minOf(remaining(), node.readRemaining)
+                        if (!it.tryWriteExact(size)) throw IllegalStateException()
+                        written = node.read(this, size)
+                        bytesWritten(it, size)
+                    }
+                } catch (t: Throwable) {
+                    node.release()
+                    packet.release()
+                    throw t
+                }
+
+                if (written == 0 && node.readRemaining > 0) {
+                    return writeViewBasedPacketSuspend(packet, node)
+                }
+            }
+
+            node.release()
         }
     }
 
-    private suspend fun writeMultipleBufferPacketSuspend(rem: ByteBuffer, packet: ByteReadPacketImpl) {
-        var buffer = rem
+    private suspend fun writeViewBasedPacketSuspend(packet: ByteReadPacketViewBased, node0: BufferView) {
+        var node: BufferView? = node0
 
         try {
-            do {
-                writeFully(buffer)
-                if (packet.remaining == 0) break
-                packet.pool.recycle(buffer)
-                buffer = packet.steal()
-            } while (true)
-        } finally {
-            packet.pool.recycle(buffer)
+            while ((node != null && node.readRemaining > 0) || !packet.isEmpty) {
+                while (node != null && node.readRemaining > 0) {
+                    writeSuspend(1)
+                    writing {
+                        val size = minOf(remaining(), node!!.readRemaining)
+                        if (!it.tryWriteExact(size)) throw IllegalStateException()
+                        node!!.read(this, size)
+                        bytesWritten(it, size)
+                    }
+
+                    if (node.readRemaining == 0) {
+                        node.release()
+                        node = packet.steal()
+                    }
+                }
+            }
+        } catch (t: Throwable) {
+            node?.release()
             packet.release()
+            throw t
         }
     }
 
@@ -927,6 +1006,8 @@ internal class ByteBufferChannel(
                 if (buffer.hasRemaining()) {
                     buffer.compact()
                     break
+                } else {
+                    buffer.compact()
                 }
             }
 
@@ -935,15 +1016,19 @@ internal class ByteBufferChannel(
             t
         }
 
+        if (t != null) {
+            BufferPool.recycle(buffer)
+            packet.release()
+            throw t
+        }
+
         buffer.flip()
         if (buffer.hasRemaining()) {
             return writeExternalPacketSuspend(buffer, packet)
         }
 
         BufferPool.recycle(buffer)
         packet.release()
-
-        if (t != null) throw t
     }
 
     private suspend fun writeExternalPacketSuspend(buffer: ByteBuffer, packet: ByteReadPacket) {

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

@@ -117,6 +117,21 @@ public interface ByteReadChannel {
      * and no characters were read.
      */
     suspend fun <A : Appendable> readUTF8LineTo(out: A, limit: Int = Int.MAX_VALUE): Boolean
+
+    /**
+     * Invokes [block] when it will be possible to read at least [min] bytes
+     * providing byte buffer to it so lambda can read from the buffer
+     * up to [ByteBuffer.remaining] bytes. If there are no [min] bytes available then the invocation could
+     * suspend until the requirement will be met.
+     *
+     * Warning: it is not guaranteed that all of remaining bytes will be represented as a single byte buffer
+     * eg: it could be 4 bytes available for read but the provided byte buffer could have only 2 remaining bytes:
+     * in this case you have to invoke read again (with decreased [min] accordingly).
+     *
+     * @param min amount of bytes available for read, should be positive
+     * @param block to be invoked when at least [min] bytes available for read
+     */
+    suspend fun read(min: Int = 1, block: (ByteBuffer) -> Unit)
 }
 
 /**

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

@@ -59,6 +59,21 @@ public interface ByteWriteChannel {
     suspend fun writeFully(src: ByteArray) = writeFully(src, 0, src.size)
     suspend fun writeFully(src: ByteBuffer)
 
+    /**
+     * Invokes [block] when it will be possible to write at least [min] bytes
+     * providing byte buffer to it so lambda can write to the buffer
+     * up to [ByteBuffer.remaining] bytes. If there are no [min] bytes spaces available then the invocation could
+     * suspend until the requirement will be met.
+     *
+     * Warning: it is not guaranteed that all of remaining bytes will be represented as a single byte buffer
+     * eg: it could be 4 bytes available for write but the provided byte buffer could have only 2 remaining bytes:
+     * in this case you have to invoke write again (with decreased [min] accordingly).
+     *
+     * @param min amount of bytes available for write, should be positive
+     * @param block to be invoked when at least [min] bytes free capacity available
+     */
+    suspend fun write(min: Int = 1, block: (ByteBuffer) -> Unit)
+
     /**
      * Writes a [packet] fully or fails if channel get closed before the whole packet has been written
      */