Adapt ktor tcp transport implementation to the latest change in transport API and simplify some parts

* drop the ability to change dispatcher - always use Default
* improve a bit usage of dispatchers and close sequences

Author: whyoleg

rsocket-transports/ktor-tcp/api/rsocket-transport-ktor-tcp.api

@@ -8,10 +8,7 @@ public final class io/rsocket/kotlin/transport/ktor/tcp/KtorTcpClientTransport$F
 }
 
 public abstract interface class io/rsocket/kotlin/transport/ktor/tcp/KtorTcpClientTransportBuilder : io/rsocket/kotlin/transport/RSocketTransportBuilder {
-	public abstract fun dispatcher (Lkotlin/coroutines/CoroutineContext;)V
-	public fun inheritDispatcher ()V
 	public abstract fun selectorManager (Lio/ktor/network/selector/SelectorManager;Z)V
-	public abstract fun selectorManagerDispatcher (Lkotlin/coroutines/CoroutineContext;)V
 	public abstract fun socketOptions (Lkotlin/jvm/functions/Function1;)V
 }
 
@@ -31,10 +28,7 @@ public final class io/rsocket/kotlin/transport/ktor/tcp/KtorTcpServerTransport$F
 }
 
 public abstract interface class io/rsocket/kotlin/transport/ktor/tcp/KtorTcpServerTransportBuilder : io/rsocket/kotlin/transport/RSocketTransportBuilder {
-	public abstract fun dispatcher (Lkotlin/coroutines/CoroutineContext;)V
-	public fun inheritDispatcher ()V
 	public abstract fun selectorManager (Lio/ktor/network/selector/SelectorManager;Z)V
-	public abstract fun selectorManagerDispatcher (Lkotlin/coroutines/CoroutineContext;)V
 	public abstract fun socketOptions (Lkotlin/jvm/functions/Function1;)V
 }
 

rsocket-transports/ktor-tcp/src/commonMain/kotlin/io/rsocket/kotlin/transport/ktor/tcp/KtorTcpClientTransport.kt

@@ -34,56 +34,44 @@ public sealed interface KtorTcpClientTransport : RSocketTransport {
 
 @OptIn(RSocketTransportApi::class)
 public sealed interface KtorTcpClientTransportBuilder : RSocketTransportBuilder<KtorTcpClientTransport> {
-    public fun dispatcher(context: CoroutineContext)
-    public fun inheritDispatcher(): Unit = dispatcher(EmptyCoroutineContext)
-
-    public fun selectorManagerDispatcher(context: CoroutineContext)
     public fun selectorManager(manager: SelectorManager, manage: Boolean)
-
     public fun socketOptions(block: SocketOptions.TCPClientSocketOptions.() -> Unit)
-
     //TODO: TLS support
 }
 
 private class KtorTcpClientTransportBuilderImpl : KtorTcpClientTransportBuilder {
-    private var dispatcher: CoroutineContext = Dispatchers.Default
-    private var selector: KtorTcpSelector = KtorTcpSelector.FromContext(Dispatchers.IoCompatible)
+    private var selectorManager: SelectorManager? = null
+    private var manageSelectorManager: Boolean = true
     private var socketOptions: SocketOptions.TCPClientSocketOptions.() -> Unit = {}
 
-    override fun dispatcher(context: CoroutineContext) {
-        check(context[Job] == null) { "Dispatcher shouldn't contain job" }
-        this.dispatcher = context
-    }
-
     override fun socketOptions(block: SocketOptions.TCPClientSocketOptions.() -> Unit) {
         this.socketOptions = block
     }
 
-    override fun selectorManagerDispatcher(context: CoroutineContext) {
-        check(context[Job] == null) { "Dispatcher shouldn't contain job" }
-        this.selector = KtorTcpSelector.FromContext(context)
-    }
-
     override fun selectorManager(manager: SelectorManager, manage: Boolean) {
-        this.selector = KtorTcpSelector.FromInstance(manager, manage)
+        this.selectorManager = manager
+        this.manageSelectorManager = manage
     }
 
     @RSocketTransportApi
-    override fun buildTransport(context: CoroutineContext): KtorTcpClientTransport {
-        val transportContext = context.supervisorContext() + dispatcher
-        return KtorTcpClientTransportImpl(
-            coroutineContext = transportContext,
-            socketOptions = socketOptions,
-            selectorManager = selector.createFor(transportContext)
-        )
-    }
+    override fun buildTransport(context: CoroutineContext): KtorTcpClientTransport = KtorTcpClientTransportImpl(
+        coroutineContext = context.supervisorContext() + Dispatchers.Default,
+        socketOptions = socketOptions,
+        selectorManager = selectorManager ?: SelectorManager(Dispatchers.IoCompatible),
+        manageSelectorManager = manageSelectorManager
+    )
 }
 
 private class KtorTcpClientTransportImpl(
     override val coroutineContext: CoroutineContext,
     private val socketOptions: SocketOptions.TCPClientSocketOptions.() -> Unit,
     private val selectorManager: SelectorManager,
+    manageSelectorManager: Boolean,
 ) : KtorTcpClientTransport {
+    init {
+        if (manageSelectorManager) coroutineContext.job.invokeOnCompletion { selectorManager.close() }
+    }
+
     override fun target(remoteAddress: SocketAddress): RSocketClientTarget = KtorTcpClientTargetImpl(
         coroutineContext = coroutineContext.supervisorContext(),
         socketOptions = socketOptions,
@@ -101,10 +89,17 @@ private class KtorTcpClientTargetImpl(
     private val selectorManager: SelectorManager,
     private val remoteAddress: SocketAddress,
 ) : RSocketClientTarget {
-
     @RSocketTransportApi
-    override fun connectClient(handler: RSocketConnectionHandler): Job = launch {
-        val socket = aSocket(selectorManager).tcp().connect(remoteAddress, socketOptions)
-        handler.handleKtorTcpConnection(socket)
+    override suspend fun connectClient(): RSocketConnection {
+        currentCoroutineContext().ensureActive()
+        coroutineContext.ensureActive()
+
+        return withContext(Dispatchers.IoCompatible) {
+            val socket = aSocket(selectorManager).tcp().connect(remoteAddress, socketOptions)
+            KtorTcpConnection(
+                parentContext = this@KtorTcpClientTargetImpl.coroutineContext,
+                socket = socket
+            )
+        }
     }
 }

rsocket-transports/ktor-tcp/src/commonMain/kotlin/io/rsocket/kotlin/transport/ktor/tcp/KtorTcpConnection.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2015-2024 the original author or authors.
+ * Copyright 2015-2025 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -24,86 +24,95 @@ import io.rsocket.kotlin.transport.internal.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
 import kotlinx.io.*
+import kotlin.coroutines.*
 
 @RSocketTransportApi
-internal suspend fun RSocketConnectionHandler.handleKtorTcpConnection(socket: Socket): Unit = coroutineScope {
-    val outboundQueue = PrioritizationFrameQueue(Channel.BUFFERED)
-    val inbound = bufferChannel(Channel.BUFFERED)
+internal class KtorTcpConnection(
+    parentContext: CoroutineContext,
+    private val socket: Socket,
+) : RSocketSequentialConnection {
+    private val outboundQueue = PrioritizationFrameQueue(Channel.BUFFERED)
+    private val inbound = bufferChannel(Channel.BUFFERED)
 
-    val readerJob = launch {
-        val input = socket.openReadChannel()
-        try {
-            while (true) inbound.send(input.readFrame() ?: break)
-            input.cancel(null)
-        } catch (cause: Throwable) {
-            input.cancel(cause)
-            throw cause
-        }
-    }.onCompletion { inbound.cancel() }
+    override val coroutineContext: CoroutineContext = parentContext.childContext()
 
-    val writerJob = launch {
-        val output = socket.openWriteChannel()
-        try {
-            while (true) {
-                // we write all available frames here, and only after it flush
-                // in this case, if there are several buffered frames we can send them in one go
-                // avoiding unnecessary flushes
-                output.writeFrame(outboundQueue.dequeueFrame() ?: break)
-                while (true) output.writeFrame(outboundQueue.tryDequeueFrame() ?: break)
-                output.flush()
+    init {
+        @OptIn(DelicateCoroutinesApi::class)
+        launch(start = CoroutineStart.ATOMIC) {
+            val outboundJob = launch {
+                nonCancellable {
+                    val output = socket.openWriteChannel()
+                    try {
+                        while (true) {
+                            // we write all available frames here, and only after it flush
+                            // in this case, if there are several buffered frames we can send them in one go
+                            // avoiding unnecessary flushes
+                            output.writeFrame(outboundQueue.dequeueFrame() ?: break)
+                            while (true) output.writeFrame(outboundQueue.tryDequeueFrame() ?: break)
+                            output.flush()
+                        }
+                        output.flushAndClose()
+                    } catch (cause: Throwable) {
+                        output.cancel(cause)
+                        throw cause
+                    }
+                }
+            }.onCompletion {
+                outboundQueue.cancel()
             }
-            output.close(null)
-        } catch (cause: Throwable) {
-            output.close(cause)
-            throw cause
-        }
-    }.onCompletion { outboundQueue.cancel() }
 
-    try {
-        handleConnection(KtorTcpConnection(outboundQueue, inbound))
-    } finally {
-        readerJob.cancel()
-        outboundQueue.close() // will cause `writerJob` completion
-        // even if it was cancelled, we still need to close socket and await it closure
-        withContext(NonCancellable) {
-            // await completion of read/write and then close socket
-            readerJob.join()
-            writerJob.join()
-            // close socket
-            socket.close()
-            socket.socketContext.join()
+            val inboundJob = launch {
+                val input = socket.openReadChannel()
+                try {
+                    while (true) {
+                        inbound.send(input.readFrame() ?: break)
+                    }
+                    input.cancel(null)
+                } catch (cause: Throwable) {
+                    input.cancel(cause)
+                    throw cause
+                }
+            }.onCompletion {
+                inbound.cancel()
+            }
+
+            try {
+                awaitCancellation()
+            } finally {
+                nonCancellable {
+                    outboundQueue.close()
+                    outboundJob.join()
+                    inboundJob.join()
+                    // await socket completion
+                    socket.close()
+                    socket.socketContext.join()
+                }
+            }
         }
     }
-}
 
-@RSocketTransportApi
-private class KtorTcpConnection(
-    private val outboundQueue: PrioritizationFrameQueue,
-    private val inbound: ReceiveChannel<Buffer>,
-) : RSocketSequentialConnection {
-    override val isClosedForSend: Boolean get() = outboundQueue.isClosedForSend
     override suspend fun sendFrame(streamId: Int, frame: Buffer) {
         return outboundQueue.enqueueFrame(streamId, frame)
     }
 
     override suspend fun receiveFrame(): Buffer? {
         return inbound.receiveCatching().getOrNull()
     }
-}
 
-@OptIn(InternalAPI::class)
-private fun ByteWriteChannel.writeFrame(frame: Buffer) {
-    writeBuffer.writeInt24(frame.size.toInt())
-    writeBuffer.transferFrom(frame)
-}
+    @OptIn(InternalAPI::class)
+    private fun ByteWriteChannel.writeFrame(frame: Buffer) {
+        writeBuffer.writeInt24(frame.size.toInt())
+        writeBuffer.transferFrom(frame)
+    }
 
-@OptIn(InternalAPI::class)
-private suspend fun ByteReadChannel.readFrame(): Buffer? {
-    while (availableForRead < 3 && awaitContent(3)) yield()
-    if (availableForRead == 0) return null
+    @OptIn(InternalAPI::class)
+    private suspend fun ByteReadChannel.readFrame(): Buffer? {
+        while (availableForRead < 3 && awaitContent(3)) yield()
+        if (availableForRead == 0) return null
 
-    val length = readBuffer.readInt24()
-    return readBuffer(length).also {
-        it.require(length.toLong())
+        val length = readBuffer.readInt24()
+        return readBuffer(length).also {
+            it.require(length.toLong())
+        }
     }
 }