IO: add Packet.writePacket optimizations, add object pool verifier

Author: Sergey Mashkov

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

@@ -367,12 +367,12 @@ internal class ByteBufferChannel(
         return readAvailable(dst)
     }
 
-    suspend override fun readPacket(size: Int): ByteReadPacket {
+    suspend override fun readPacket(size: Int, headerSizeHint: Int): ByteReadPacket {
         closed?.cause?.let { throw it }
 
         if (size == 0) return ByteReadPacketEmpty
 
-        val builder = ByteWritePacketImpl(BufferPool)
+        val builder = ByteWritePacketImpl(headerSizeHint, BufferPool)
         val buffer = BufferPool.borrow()
 
         try {
@@ -874,6 +874,7 @@ internal class ByteBufferChannel(
                 writeAsMuchAsPossible(buffer)
                 if (buffer.hasRemaining()) break
                 packet.pool.recycle(buffer)
+                buffer = null
             }
             null
         } catch (t: Throwable) { t }

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

@@ -54,9 +54,9 @@ public interface ByteReadChannel {
 
     /**
      * Reads the specified amount of bytes and makes a byte packet from them. Fails if channel has been closed
-     * and not enough bytes available.
+     * and not enough bytes available. Accepts [headerSizeHint] to be provided, see [WritePacket].
      */
-    suspend fun readPacket(size: Int): ByteReadPacket
+    suspend fun readPacket(size: Int, headerSizeHint: Int = 0): ByteReadPacket
 
     /**
      * Reads a long number (suspending if not enough bytes available) or fails if channel has been closed

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

@@ -168,8 +168,8 @@ suspend fun ByteWriteChannel.writeChar(ch: Char) {
     writeShort(ch.toInt())
 }
 
-inline suspend fun ByteWriteChannel.writePacket(builder: ByteWritePacket.() -> Unit) {
-    writePacket(buildPacket(builder))
+inline suspend fun ByteWriteChannel.writePacket(headerSizeHint: Int = 0, builder: ByteWritePacket.() -> Unit) {
+    writePacket(buildPacket(headerSizeHint, builder))
 }
 
 suspend fun ByteWriteChannel.writePacketSuspend(builder: suspend ByteWritePacket.() -> Unit) {

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

@@ -10,7 +10,7 @@ import java.util.*
 
 internal class ByteReadPacketImpl(internal val packets: ArrayDeque<ByteBuffer>, internal val pool: ObjectPool<ByteBuffer>) : ByteReadPacket {
     override val remaining: Int
-        get() = packets.sumBy { it.remaining() }
+        get() = if (packets.isNotEmpty()) packets.sumBy { it.remaining() } else 0
 
     internal fun steal(): ByteBuffer = packets.pollFirst() ?: throw IllegalStateException("EOF")
 

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

@@ -24,11 +24,13 @@ interface ByteWritePacket : Appendable {
 
     fun release()
     fun build(): ByteReadPacket
+    fun reset()
 
     fun outputStream(): OutputStream
     fun writerUTF8(): Writer
 
     fun writePacket(p: ByteReadPacket)
+    @Deprecated("")
     fun writePacketUnconsumed(p: ByteReadPacket)
 
     override fun append(csq: CharSequence): ByteWritePacket {

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

@@ -5,8 +5,13 @@ import java.io.*
 import java.nio.ByteBuffer
 import java.nio.CharBuffer
 import java.util.*
+import kotlin.NoSuchElementException
+
+internal class ByteWritePacketImpl(private var headerSizeHint: Int, private val pool: ObjectPool<ByteBuffer>) : ByteWritePacket {
+    init {
+        require(headerSizeHint >= 0) { "shouldn't be negative: headerSizeHint = $headerSizeHint" }
+    }
 
-internal class ByteWritePacketImpl(private val pool: ObjectPool<ByteBuffer>) : ByteWritePacket {
     override var size: Int = 0
         private set
     private var buffers: Any? = null
@@ -86,22 +91,72 @@ internal class ByteWritePacketImpl(private val pool: ObjectPool<ByteBuffer>) : B
     override fun writePacket(p: ByteReadPacket) {
         when (p) {
             is ByteReadPacketEmpty -> {}
-            is ByteReadPacketSingle -> {
-                if (p.remaining > 0) {
-                    size += p.remaining
-                    last(p.steal().also { it.compact() })
+            is ByteReadPacketSingle -> writePacketSingle(p)
+            is ByteReadPacketImpl -> writePacketMultiple(p)
+            else -> writeFully(p.readBytes())
+        }
+    }
+
+    private fun reverseCopyToForeignBuffer(count: Int, last: ByteBuffer, buffer: ByteBuffer) {
+        val startOffset = buffer.position() - count
+        val l = buffer.limit()
+        buffer.position(startOffset)
+        buffer.limit(startOffset + count)
+
+        last.flip()
+        last.position(last.limit() - count)
+        buffer.put(last)
+
+        recycleLast()
+        headerSizeHint = startOffset
+        buffer.limit(buffer.capacity())
+        buffer.position(l)
+        last(buffer)
+    }
+
+    private fun writePacketSingle(p: ByteReadPacketSingle) {
+        val initialRemaining = p.remaining
+        if (initialRemaining > 0) {
+            size += initialRemaining
+            writePacketBuffer(p.steal())
+        }
+    }
+
+    private fun writePacketMultiple(p: ByteReadPacketImpl) {
+        val initialRemaining = p.remaining
+        if (initialRemaining > 0) {
+            size += initialRemaining
+
+            do {
+                writePacketBuffer(p.steal())
+            } while (p.remaining > 0)
+        }
+    }
+
+    private fun writePacketBuffer(buffer: ByteBuffer) {
+        val last = last()
+
+        if (buffer.position() > 0 && last != null) {
+            if (last === buffers || buffersCount() == 1) {
+                val count = last.position() - headerSizeHint
+                if (count < PACKET_MAX_COPY_SIZE && count <= buffer.position()) {
+                    reverseCopyToForeignBuffer(count, last, buffer)
+                    return
                 }
-            }
-            is ByteReadPacketImpl -> {
-                size += p.remaining
-                while (p.remaining > 0) {
-                    last(p.steal().also { it.compact() })
+            } else {
+                val count = last.position()
+                if (count < PACKET_MAX_COPY_SIZE && count == buffer.position()) {
+                    reverseCopyToForeignBuffer(count, last, buffer)
+                    return
                 }
             }
-            else -> {
-                writeFully(p.readBytes())
-            }
+        } else if (last != null && last.remaining() <= buffer.remaining() && buffer.remaining() < PACKET_MAX_COPY_SIZE) {
+            last.put(buffer)
+            recycle(buffer)
+            return
         }
+
+        last(buffer.also { it.compact() })
     }
 
     override fun writePacketUnconsumed(p: ByteReadPacket) {
@@ -264,22 +319,43 @@ internal class ByteWritePacketImpl(private val pool: ObjectPool<ByteBuffer>) : B
     }
 
     override fun build(): ByteReadPacket {
-        val bs = buffers ?: return ByteReadPacketEmpty
+        val bs = buffers
         buffers = null
+        size = 0
 
-        return if (bs is ArrayDeque<*>) {
-            @Suppress("UNCHECKED_CAST")
-            when {
-                bs.isEmpty() -> ByteReadPacketEmpty
-                bs.size == 1 -> ByteReadPacketSingle((bs.first as ByteBuffer).also { it.flip() }, pool)
-                else -> ByteReadPacketImpl((bs as ArrayDeque<ByteBuffer>).also {
-                    for (b in bs) {
-                        b.flip()
-                    }
-                }, pool)
+        return when (bs) {
+            null -> ByteReadPacketEmpty
+            is ArrayDeque<*> -> {
+                @Suppress("UNCHECKED_CAST") buildMultiBufferPacket(bs as ArrayDeque<ByteBuffer>)
+            }
+            else -> ByteReadPacketSingle((bs as ByteBuffer).also { switchBufferToRead(true, it) }, pool)
+        }
+    }
+
+    private fun buildMultiBufferPacket(buffers: ArrayDeque<ByteBuffer>): ByteReadPacket {
+        return when (buffers.size) {
+            0 -> ByteReadPacketEmpty
+            1 -> ByteReadPacketSingle(buffers.first.also { switchBufferToRead(true, it) }, pool)
+            else -> {
+                val it = buffers.iterator()
+                switchBufferToRead(true, it.next())
+                do {
+                    switchBufferToRead(false, it.next())
+                } while (it.hasNext())
+
+                ByteReadPacketImpl(buffers, pool)
+            }
+        }
+    }
+
+    private fun switchBufferToRead(first: Boolean, bb: ByteBuffer) {
+        bb.flip()
+
+        if (first) {
+            val skip = headerSizeHint
+            if (skip > 0) {
+                bb.position(bb.position() + skip)
             }
-        } else {
-            ByteReadPacketSingle((bs as ByteBuffer).also { it.flip() }, pool)
         }
     }
 
@@ -297,6 +373,10 @@ internal class ByteWritePacketImpl(private val pool: ObjectPool<ByteBuffer>) : B
         }
     }
 
+    override fun reset() {
+        release()
+    }
+
     private inline fun write(size: Int, block: (ByteBuffer) -> Unit) {
         val buffer = last()?.takeIf { it.remaining() >= size }
 
@@ -313,6 +393,9 @@ internal class ByteWritePacketImpl(private val pool: ObjectPool<ByteBuffer>) : B
         val new = pool.borrow()
         new.clear()
         last(new)
+        if (buffers === new) {
+            new.position(headerSizeHint)
+        }
         return new
     }
 
@@ -328,15 +411,44 @@ internal class ByteWritePacketImpl(private val pool: ObjectPool<ByteBuffer>) : B
     private fun last(new: ByteBuffer) {
         @Suppress("UNCHECKED_CAST")
         if (buffers is ArrayDeque<*>) (buffers as ArrayDeque<ByteBuffer>).addLast(new)
-        else if (buffers == null) buffers = new
-        else {
+        else if (buffers == null) {
+            buffers = new
+        } else {
             val dq = ArrayDeque<ByteBuffer>()
             dq.addFirst(buffers as ByteBuffer)
             dq.addLast(new)
             buffers = dq
         }
     }
 
+    private fun recycleLast() {
+        val b = buffers
+        when (b) {
+            is ArrayDeque<*> -> {
+                recycle(b.pollLast() as ByteBuffer)
+                if (b.isEmpty()) {
+                    buffers = null
+                }
+            }
+            is ByteBuffer -> {
+                buffers = null
+                recycle(b)
+            }
+            else -> throw NoSuchElementException("Unable to recycle last buffer: buffers chain is empty")
+        }
+    }
+
+    private fun buffersCount(): Int {
+        val bb = buffers
+
+        return when (bb) {
+            null -> 0
+            is ByteBuffer -> 1
+            is ArrayDeque<*> -> bb.size
+            else -> 0
+        }
+    }
+
     private fun recycle(buffer: ByteBuffer) {
         pool.recycle(buffer)
     }

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

@@ -5,17 +5,37 @@ import java.nio.ByteBuffer
 
 internal val PACKET_BUFFER_SIZE = getIOIntProperty("PacketBufferSize", 4096)
 private val PACKET_BUFFER_POOL_SIZE = getIOIntProperty("PacketBufferPoolSize", 128)
+internal val PACKET_MAX_COPY_SIZE = getIOIntProperty("PacketMaxCopySize", 500)
 
-private val PacketBufferPool: ObjectPool<ByteBuffer> =
+internal val PacketBufferPool: ObjectPool<ByteBuffer> =
     object : ObjectPoolImpl<ByteBuffer>(PACKET_BUFFER_POOL_SIZE) {
         override fun produceInstance(): ByteBuffer = ByteBuffer.allocateDirect(PACKET_BUFFER_SIZE)
         override fun clearInstance(instance: ByteBuffer) = instance.apply { clear() }
     }
 
-inline fun buildPacket(block: ByteWritePacket.() -> Unit): ByteReadPacket =
-        WritePacket().apply(block).build()
+internal inline fun buildPacket(pool: ObjectPool<ByteBuffer> = PacketBufferPool, headerSizeHint: Int, block: ByteWritePacket.() -> Unit): ByteReadPacket {
+    val w = ByteWritePacketImpl(headerSizeHint, pool)
+    try {
+        block(w)
+        return w.build()
+    } catch (t: Throwable) {
+        w.release()
+        throw t
+    }
+}
+
+inline fun buildPacket(headerSizeHint: Int = 0, block: ByteWritePacket.() -> Unit): ByteReadPacket {
+    val w = WritePacket(headerSizeHint)
+    try {
+        block(w)
+        return w.build()
+    } catch (t: Throwable) {
+        w.release()
+        throw t
+    }
+}
 
-fun WritePacket(): ByteWritePacket = ByteWritePacketImpl(PacketBufferPool)
+fun WritePacket(headerSizeHint: Int = 0): ByteWritePacket = ByteWritePacketImpl(headerSizeHint, PacketBufferPool)
 
 fun ByteReadPacket.readUTF8Line(estimate: Int = 16, limit: Int = Int.MAX_VALUE): String? {
     val sb = StringBuilder(estimate)