tcp optimization and graceful close

src/jvmMain/kotlin/net/sergeych/kiloparsec/adapter/asyncSocketToDevice.kt

@@ -3,13 +3,14 @@ package net.sergeych.kiloparsec.adapter
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.channels.ClosedReceiveChannelException
-import net.sergeych.bintools.toDump
+import kotlinx.coroutines.flow.MutableStateFlow
 import net.sergeych.crypto.encodeVarUnsigned
 import net.sergeych.crypto.readVarUnsigned
 import net.sergeych.kiloparsec.Transport
 import net.sergeych.mp_logger.LogTag
 import net.sergeych.mp_logger.warning
 import net.sergeych.mp_tools.globalLaunch
+import net.sergeych.tools.waitFor
 import java.net.InetSocketAddress
 import java.net.StandardSocketOptions.TCP_NODELAY
 import java.nio.ByteBuffer
@@ -22,8 +23,7 @@ private val log = LogTag("ASTD")
 /**
  * Prepend block with its size, varint-encoded
  */
-private fun encode(block: UByteArray): ByteArray
-    = (encodeVarUnsigned(block.size.toUInt()) + block).toByteArray()
+private fun encode(block: UByteArray): ByteArray = (encodeVarUnsigned(block.size.toUInt()) + block).toByteArray()
 
 /**
  * Convert asynchronous socket to a [Transport.Device] using non-blocking nio,
@@ -35,14 +35,19 @@ suspend fun asyncSocketToDevice(socket: AsynchronousSocketChannel): InetTranspor
     val deferredDevice = CompletableDeferred<InetTransportDevice>()
     globalLaunch {
         coroutineScope {
+            val sendQueueEmpty = MutableStateFlow(true)
+            val receiving = MutableStateFlow(false)
             fun stop() {
                 cancel()
             }
+            // We're in block mode, every block we send worth immediate sending, we do not
+            // send partial blocks, so:
             socket.setOption(TCP_NODELAY, true)
+
             // socket input is to be parsed for blocks, so we receive bytes
             // and decode them to blocks
             val input = Channel<UByte>(1024)
-            // copy from socket to input:
+            // copy incoming data from the socket to input channel:
             launch {
                 val data = ByteArray(1024)
                 val inb = ByteBuffer.wrap(data)
@@ -53,32 +58,41 @@ suspend fun asyncSocketToDevice(socket: AsynchronousSocketChannel): InetTranspor
                     }
                     if (size < 0) stop()
                     else {
-                        println("recvd:\n${data.sliceArray(0..<size).toDump()}\n------------------")
+//                        println("recvd:\n${data.sliceArray(0..<size).toDump()}\n------------------")
                         for (i in 0..<size) input.send(data[i].toUByte())
                     }
                 }
             }
-            // output is blocks, so we sent blocks:
+
+            // output is blocks, so we sent transformed, framed blocks:
             val outputBlocks = Channel<UByteArray>()
             // copy from output to socket:
             launch {
                 try {
                     while (isActive) {
                         // wait for the first block to send
+                        sendQueueEmpty.value = outputBlocks.isEmpty
                         var data = encode(outputBlocks.receive())
+
+                        // now we're sending, so queue state is sending:
+                        sendQueueEmpty.value = false
+
                         // if there are more, take them all (NO_DELAY optimization)
                         while (!outputBlocks.isEmpty)
                             data += encode(outputBlocks.receive())
-                        // now send the aggregate:
+
+                        // now send it all together:
                         val outBuff = ByteBuffer.wrap(data)
                         val cnt = suspendCoroutine { continuation ->
                             socket.write(outBuff, continuation, IntCompletionHandler)
                         }
                         // be sure it was all sent
-                        if( outBuff.position() != data.size || cnt != data.size) {
-                            throw RuntimeException("PArtial write!")
+                        if (outBuff.position() != data.size || cnt != data.size) {
+                            throw RuntimeException("unexpected partial write")
                         }
                     }
+                    // in the case of just breaking out of the loop:
+                    sendQueueEmpty.value = true
                 } catch (_: ClosedReceiveChannelException) {
                     stop()
                 }
@@ -89,7 +103,9 @@ suspend fun asyncSocketToDevice(socket: AsynchronousSocketChannel): InetTranspor
             launch {
                 try {
                     while (isActive) {
+                        receiving.value = !input.isEmpty
                         val size = readVarUnsigned(input)
+                        receiving.value = true
                         if (size == 0u) log.warning { "zero size block is ignored!" }
                         else {
                             val block = UByteArray(size.toInt())
@@ -105,11 +121,33 @@ suspend fun asyncSocketToDevice(socket: AsynchronousSocketChannel): InetTranspor
                     inputBlocks.send(null)
                     stop()
                 }
+                receiving.value = false
             }
-//            SocketAddress.
+
+            // wait until send queue is empty
+            suspend fun flush() {
+                yield()
+                // do not slow down with collect if it is ok by now:
+                if (!sendQueueEmpty.value || !outputBlocks.isEmpty)
+                    // wait until all output is sent
+                    sendQueueEmpty.waitFor { it && outputBlocks.isEmpty }
+            }
+
             val addr = socket.remoteAddress as InetSocketAddress
             deferredDevice.complete(
-                InetTransportDevice(inputBlocks, outputBlocks, JvmNetworkAddress(addr.address,addr.port)) { stop() }
+                InetTransportDevice(inputBlocks, outputBlocks, JvmNetworkAddress(addr.address, addr.port),
+                    { flush() }
+                ) {
+                    yield()
+                    // wait until all received data are parsed, but not too long
+                    withTimeoutOrNull( 1000 ) {
+                        receiving.waitFor { !it }
+                    }
+                    // graceful close: flush output
+                    flush()
+                    // then stop it
+                    stop()
+                }
             )
         }
         globalLaunch { socket.close() }