IO: eliminate stealing if packets use different pools

Author: Sergey Mashkov

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

@@ -8,11 +8,11 @@ import java.nio.*
 import java.nio.charset.*
 import java.util.*
 
-internal class ByteReadPacketImpl(internal val packets: ArrayDeque<ByteBuffer>, internal val pool: ObjectPool<ByteBuffer>) : ByteReadPacket {
+internal class ByteReadPacketImpl(internal val buffers: ArrayDeque<ByteBuffer>, internal val pool: ObjectPool<ByteBuffer>) : ByteReadPacket {
     override val remaining: Int
-        get() = if (packets.isNotEmpty()) packets.sumBy { it.remaining() } else 0
+        get() = if (buffers.isNotEmpty()) buffers.sumBy { it.remaining() } else 0
 
-    internal fun steal(): ByteBuffer = packets.pollFirst() ?: throw IllegalStateException("EOF")
+    internal fun steal(): ByteBuffer = buffers.pollFirst() ?: throw IllegalStateException("EOF")
 
     override fun readAvailable(dst: ByteArray, offset: Int, length: Int): Int {
         var copied = 0
@@ -228,27 +228,27 @@ internal class ByteReadPacketImpl(internal val packets: ArrayDeque<ByteBuffer>,
     }
 
     override fun release() {
-        while (packets.isNotEmpty()) {
-            recycle(packets.remove())
+        while (buffers.isNotEmpty()) {
+            recycle(buffers.remove())
         }
     }
 
     override fun copy(): ByteReadPacket {
-        if (packets.isEmpty()) return ByteReadPacketEmpty
-        val copyDeque = ArrayDeque<ByteBuffer>(packets.size)
+        if (buffers.isEmpty()) return ByteReadPacketEmpty
+        val copyDeque = ArrayDeque<ByteBuffer>(buffers.size)
 
-        for (p in packets) {
+        for (p in buffers) {
             copyDeque.add(pool.borrow().also { it.put(p.duplicate()); it.flip() })
         }
 
         return ByteReadPacketImpl(copyDeque, pool)
     }
 
     private inline fun reading(size: Int, block: (ByteBuffer) -> Boolean): Boolean {
-        if (packets.isEmpty()) return false
+        if (buffers.isEmpty()) return false
 
         var visited = false
-        var buffer = packets.peekFirst()
+        var buffer = buffers.peekFirst()
         var stop = false
 
         while (!stop) {
@@ -262,26 +262,26 @@ internal class ByteReadPacketImpl(internal val packets: ArrayDeque<ByteBuffer>,
             }
 
             if (!buffer.hasRemaining()) {
-                packets.removeFirst()
+                buffers.removeFirst()
                 recycle(buffer)
 
-                if (packets.isEmpty()) break
-                buffer = packets.peekFirst()
+                if (buffers.isEmpty()) break
+                buffer = buffers.peekFirst()
             }
         }
 
         return visited
     }
 
     private fun tryStealBytesFromNextBuffer(size: Int, buffer: ByteBuffer): Boolean {
-        if (packets.size == 1) {
+        if (buffers.size == 1) {
             return false
         }
 
-        packets.removeFirst()
+        buffers.removeFirst()
 
         val extraBytes = size - buffer.remaining()
-        val next = packets.peekFirst()
+        val next = buffers.peekFirst()
 
         if (extraBytes > next.remaining()) return  false
 
@@ -291,7 +291,7 @@ internal class ByteReadPacketImpl(internal val packets: ArrayDeque<ByteBuffer>,
         }
         buffer.flip()
 
-        packets.addFirst(buffer)
+        buffers.addFirst(buffer)
         return true
     }
 

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

@@ -116,27 +116,30 @@ internal class ByteWritePacketImpl(private var headerSizeHint: Int, private val
 
     private fun writePacketSingle(p: ByteReadPacketSingle) {
         val initialRemaining = p.remaining
+        val samePool = p.pool === this.pool
         if (initialRemaining > 0) {
             size += initialRemaining
-            writePacketBuffer(p.steal())
+            writePacketBuffer(p.steal(), samePool, p.pool)
         }
     }
 
     private fun writePacketMultiple(p: ByteReadPacketImpl) {
         val initialRemaining = p.remaining
         if (initialRemaining > 0) {
             size += initialRemaining
+            val samePool = p.pool === this.pool
+            val packetPool = p.pool
 
             do {
-                writePacketBuffer(p.steal())
+                writePacketBuffer(p.steal(), samePool, packetPool)
             } while (p.remaining > 0)
         }
     }
 
-    private fun writePacketBuffer(buffer: ByteBuffer) {
+    private fun writePacketBuffer(buffer: ByteBuffer, samePool: Boolean, packetPool: ObjectPool<ByteBuffer>) {
         val last = last()
 
-        if (buffer.position() > 0 && last != null) {
+        if (samePool && buffer.position() > 0 && last != null) {
             if (last === buffers || buffersCount() == 1) {
                 val count = last.position() - headerSizeHint
                 if (count < PACKET_MAX_COPY_SIZE && count <= buffer.position()) {
@@ -152,11 +155,16 @@ internal class ByteWritePacketImpl(private var headerSizeHint: Int, private val
             }
         } else if (last != null && last.remaining() <= buffer.remaining() && buffer.remaining() < PACKET_MAX_COPY_SIZE) {
             last.put(buffer)
-            recycle(buffer)
+            packetPool.recycle(buffer)
             return
         }
 
-        last(buffer.also { it.compact() })
+        if (samePool) {
+            last(buffer.also { it.compact() })
+        } else {
+            writeFully(buffer)
+            packetPool.recycle(buffer)
+        }
     }
 
     override fun writePacketUnconsumed(p: ByteReadPacket) {
@@ -166,7 +174,7 @@ internal class ByteWritePacketImpl(private var headerSizeHint: Int, private val
                 p.buffer?.duplicate()?.let { writeFully(it) }
             }
             is ByteReadPacketImpl -> {
-                for (buffer in p.packets) {
+                for (buffer in p.buffers) {
                     writeFully(buffer.duplicate())
                 }
             }