Tidy up and reduce copying of byte buffers

CHANGELOG

@@ -1,5 +1,8 @@
 vNEXT (in development)
 
+ * (Internal) improved the way byte buffers are used to
+   reduce array copying and clean up code
+
 v1.0.1
 
  * Fixed issue with very long packets not fitting in buffers

src/main/kotlin/com/dmdirc/ktirc/io/Buffers.kt

@@ -0,0 +1,14 @@
+package com.dmdirc.ktirc.io
+
+import kotlinx.io.pool.DefaultPool
+import java.nio.ByteBuffer
+
+private const val BUFFER_SIZE = 32768
+private const val MAXIMUM_POOL_SIZE = 1024
+
+internal val byteBufferPool = ByteBufferPool(MAXIMUM_POOL_SIZE, BUFFER_SIZE)
+
+internal class ByteBufferPool(maximumPoolSize: Int, private val bufferSize: Int) : DefaultPool<ByteBuffer>(maximumPoolSize) {
+    override fun produceInstance(): ByteBuffer = ByteBuffer.allocate(bufferSize)
+    override fun clearInstance(instance: ByteBuffer): ByteBuffer = instance.apply { clear() }
+}

src/main/kotlin/com/dmdirc/ktirc/io/LineBufferedSocket.kt

@@ -8,6 +8,7 @@ import kotlinx.coroutines.channels.SendChannel
 import kotlinx.coroutines.channels.produce
 import kotlinx.coroutines.io.ByteWriteChannel
 import kotlinx.io.core.String
+import kotlinx.io.pool.useInstance
 import java.net.InetSocketAddress
 import java.nio.ByteBuffer
 import java.security.SecureRandom
@@ -55,7 +56,7 @@ internal class LineBufferedSocketImpl(coroutineScope: CoroutineScope, private va
 
         runBlocking {
             if (tls) {
-                with (SSLContext.getInstance("TLSv1.2")) {
+                with(SSLContext.getInstance("TLSv1.2")) {
                     init(null, tlsTrustManager?.let { arrayOf(it) }, SecureRandom.getInstanceStrong())
                     socket = TlsSocket(this@LineBufferedSocketImpl, socket, this, host)
                 }
@@ -74,35 +75,31 @@ internal class LineBufferedSocketImpl(coroutineScope: CoroutineScope, private va
 
     override val receiveChannel
         get() = produce {
-            defaultPool.borrow { lineBuffer ->
+            byteBufferPool.useInstance { lineBuffer ->
                 while (socket.isOpen) {
-                    defaultPool.borrow { buffer ->
-                        val bytesRead = socket.read(buffer)
-                        if (bytesRead == -1) {
-                            return@produce
-                        }
-                        var lastLine = 0
-                        for (i in 0 until bytesRead) {
-                            if (buffer[i] == CARRIAGE_RETURN || buffer[i] == LINE_FEED) {
-                                val length = i - lastLine + lineBuffer.position()
-
-                                if (length > 1) {
-                                    val output = ByteBuffer.allocate(length)
-
-                                    lineBuffer.flip()
-                                    output.put(lineBuffer)
-                                    lineBuffer.clear()
-
-                                    output.put(buffer.array(), lastLine, i - lastLine)
-                                    log.fine { "<<< ${String(output.array())}" }
-                                    send(output.array())
-                                }
-
-                                lastLine = i + 1
+                    val buffer = socket.read() ?: return@produce
+                    var lastLine = 0
+                    for (i in 0 until buffer.limit()) {
+                        if (buffer[i] == CARRIAGE_RETURN || buffer[i] == LINE_FEED) {
+                            val length = i - lastLine + lineBuffer.position()
+
+                            if (length > 1) {
+                                val output = ByteBuffer.allocate(length)
+
+                                lineBuffer.flip()
+                                output.put(lineBuffer)
+                                lineBuffer.clear()
+
+                                output.put(buffer.array(), lastLine, i - lastLine)
+                                log.fine { "<<< ${String(output.array())}" }
+                                send(output.array())
                             }
+
+                            lastLine = i + 1
                         }
-                        lineBuffer.put(buffer.array(), lastLine, bytesRead - lastLine)
                     }
+                    lineBuffer.put(buffer.array(), lastLine, buffer.limit() - lastLine)
+                    byteBufferPool.recycle(buffer)
                 }
             }
         }

src/main/kotlin/com/dmdirc/ktirc/io/Sockets.kt

@@ -7,36 +7,17 @@ import kotlinx.coroutines.io.ByteWriteChannel
 import kotlinx.coroutines.io.close
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.suspendCancellableCoroutine
-import kotlinx.io.pool.DefaultPool
+import kotlinx.io.pool.useInstance
 import java.net.SocketAddress
 import java.nio.ByteBuffer
 import java.nio.channels.*
 import kotlin.coroutines.resume
 import kotlin.coroutines.resumeWithException
 
-internal const val BUFFER_SIZE = 32768
-internal const val POOL_SIZE = 128
-
-internal val defaultPool = ByteBufferPool()
-
-internal class ByteBufferPool : DefaultPool<ByteBuffer>(POOL_SIZE) {
-    override fun produceInstance(): ByteBuffer = ByteBuffer.allocate(BUFFER_SIZE)
-    override fun clearInstance(instance: ByteBuffer): ByteBuffer = instance.apply { clear() }
-
-    inline fun <T> borrow(block: (ByteBuffer) -> T): T {
-        val buffer = borrow()
-        try {
-            return block(buffer)
-        } finally {
-            recycle(buffer)
-        }
-    }
-}
-
 internal interface Socket {
     fun bind(socketAddress: SocketAddress)
     suspend fun connect(socketAddress: SocketAddress)
-    suspend fun read(buffer: ByteBuffer): Int
+    suspend fun read(): ByteBuffer?
     fun close()
     val write: ByteWriteChannel
     val isOpen: Boolean
@@ -73,9 +54,9 @@ internal class PlainTextSocket(private val scope: CoroutineScope) : Socket {
         client.close()
     }
 
-    override suspend fun read(buffer: ByteBuffer) = try {
+    override suspend fun read() = try {
+        val buffer = byteBufferPool.borrow()
         val bytes = suspendCancellableCoroutine<Int> { continuation ->
-
             client.closeOnCancel(continuation)
             client.read(buffer, continuation, asyncIOHandler())
         }
@@ -83,15 +64,17 @@ internal class PlainTextSocket(private val scope: CoroutineScope) : Socket {
         if (bytes == -1) {
             close()
         }
-        bytes
+
+        buffer.flip()
+        buffer
     } catch (_: ClosedChannelException) {
         // Ignore
-        0
+        null
     }
 
     private suspend fun writeLoop() {
         while (client.isOpen) {
-            defaultPool.borrow { buffer ->
+            byteBufferPool.useInstance { buffer ->
                 writeChannel.readAvailable(buffer)
                 buffer.flip()
                 try {

src/main/kotlin/com/dmdirc/ktirc/io/Tls.kt

@@ -7,6 +7,7 @@ import kotlinx.coroutines.channels.ClosedReceiveChannelException
 import kotlinx.coroutines.io.ByteChannel
 import kotlinx.coroutines.io.ByteWriteChannel
 import kotlinx.coroutines.launch
+import kotlinx.io.pool.useInstance
 import java.net.SocketAddress
 import java.nio.ByteBuffer
 import java.security.cert.CertificateException
@@ -97,19 +98,15 @@ internal class TlsSocket(
         }
     }
 
-    override suspend fun read(buffer: ByteBuffer) = try {
-        val nextBuffer = outgoingAppBuffers.receive()
-        val bytes = nextBuffer.limit()
-        buffer.put(nextBuffer)
-        defaultPool.recycle(nextBuffer)
-        bytes
+    override suspend fun read() = try {
+        outgoingAppBuffers.receive()
     } catch (_: ClosedReceiveChannelException) {
-        -1
+        null
     }
 
     private suspend fun wrap(): SSLEngineResult? {
         var result: SSLEngineResult? = null
-        defaultPool.borrow { netBuffer ->
+        byteBufferPool.useInstance { netBuffer ->
             if (engine.handshakeStatus <= SSLEngineResult.HandshakeStatus.FINISHED) {
                 writeChannel.readAvailable(incomingAppBuffer)
             }
@@ -125,24 +122,25 @@ internal class TlsSocket(
 
     private suspend fun unwrap(networkRead: Boolean = incomingNetBuffer.position() == 0): SSLEngineResult? {
         if (networkRead) {
-            val bytes = socket.read(incomingNetBuffer.slice())
-            if (bytes == -1) {
+            val buffer = socket.read()
+            if (buffer == null) {
                 close()
                 return null
             }
-            incomingNetBuffer.position(incomingNetBuffer.position() + bytes)
+            incomingNetBuffer.put(buffer)
+            byteBufferPool.recycle(buffer)
         }
 
         incomingNetBuffer.flip()
 
-        val buffer = defaultPool.borrow()
+        val buffer = byteBufferPool.borrow()
         val result = engine.unwrap(incomingNetBuffer, buffer)
         incomingNetBuffer.compact()
         if (buffer.position() > 0) {
             buffer.flip()
             outgoingAppBuffers.send(buffer)
         } else {
-            defaultPool.recycle(buffer)
+            byteBufferPool.recycle(buffer)
         }
 
         return if (result?.status == SSLEngineResult.Status.BUFFER_UNDERFLOW && !networkRead) {
@@ -158,9 +156,9 @@ internal class TlsSocket(
         socket.close()
 
         // Release any buffers we've got queued up
-        while(true) {
+        while (true) {
             outgoingAppBuffers.poll()?.let {
-                defaultPool.recycle(it)
+                byteBufferPool.recycle(it)
             } ?: break
         }
 

src/test/kotlin/com/dmdirc/ktirc/io/ByteBufferPoolTest.kt

@@ -0,0 +1,62 @@
+package com.dmdirc.ktirc.io
+
+import kotlinx.io.pool.useInstance
+import org.junit.jupiter.api.Assertions.*
+import org.junit.jupiter.api.Test
+import java.nio.ByteBuffer
+
+internal class ByteBufferPoolTest {
+
+    @Test
+    fun `it allows borrowing of multiple unique bytebuffers`() {
+        val pool = ByteBufferPool(5, 10)
+        val buffer1 = pool.borrow()
+        val buffer2 = pool.borrow()
+        val buffer3 = pool.borrow()
+
+        assertFalse(buffer1 === buffer2)
+        assertFalse(buffer2 === buffer3)
+        assertFalse(buffer1 === buffer3)
+    }
+
+    @Test
+    fun `it produces buffers of the correct size`() {
+        val pool = ByteBufferPool(5, 12)
+        val buffer = pool.borrow()
+        assertEquals(12, buffer.limit())
+    }
+
+    @Test
+    fun `it reuses recycled buffers`() {
+        val pool = ByteBufferPool(1, 10)
+
+        val buffer1 = pool.borrow()
+        pool.recycle(buffer1)
+
+        val buffer2 = pool.borrow()
+        assertTrue(buffer1 === buffer2)
+    }
+
+    @Test
+    fun `it resets buffers when reborrowing`() {
+        val pool = ByteBufferPool(1, 10)
+        val buffer1 = pool.borrow()
+        buffer1.put("31137".toByteArray())
+        pool.recycle(buffer1)
+
+        val buffer2 = pool.borrow()
+        assertTrue(buffer1 === buffer2)
+        assertEquals(0, buffer2.position())
+        assertEquals(10, buffer2.limit())
+    }
+
+    @Test
+    fun `borrow with block automatically returns`() {
+        val pool = ByteBufferPool(1, 10)
+        var buffer1: ByteBuffer? = null
+        pool.useInstance {  buffer1 = it }
+        val buffer2 = pool.borrow()
+        assertTrue(buffer1 === buffer2)
+    }
+
+}