revert some changes

Author: whyoleg

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/connection/Connection.kt

@@ -38,7 +38,7 @@ internal abstract class Connection2(
 
     // connection establishment part
 
-    abstract suspend fun establishConnection(handler: ConnectionEstablishmentHandler<*>): ConnectionConfig
+    abstract suspend fun establishConnection(handler: ConnectionEstablishmentHandler): ConnectionConfig
 
     // setup completed, start handling requests
     abstract suspend fun handleConnection(inbound: ConnectionInbound)

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/connection/ConnectionEstablishmentHandler.kt

@@ -26,14 +26,14 @@ import kotlinx.coroutines.*
 import kotlin.coroutines.*
 
 @RSocketTransportApi
-internal abstract class ConnectionEstablishmentHandler<T>(
+internal abstract class ConnectionEstablishmentHandler(
     private val isClient: Boolean,
     private val frameCodec: FrameCodec,
     private val connectionAcceptor: ConnectionAcceptor,
     private val interceptors: Interceptors,
-) : RSocketConnectionInitializer<T> {
+    private val requesterDeferred: CompletableDeferred<RSocket>?,
+) : RSocketConnectionInitializer<Unit> {
     abstract suspend fun establishConnection(context: ConnectionEstablishmentContext): ConnectionConfig
-    protected abstract fun transform(requester: RSocket): T
 
     private suspend fun wrapConnection(
         connection: RSocketConnection,
@@ -54,8 +54,8 @@ internal abstract class ConnectionEstablishmentHandler<T>(
     }
 
     @Suppress("SuspendFunctionOnCoroutineScope")
-    private suspend fun CoroutineScope.handleConnection(connection: Connection2): T {
-        return try {
+    private suspend fun CoroutineScope.handleConnection(connection: Connection2) {
+        try {
             val connectionConfig = connection.establishConnection(this@ConnectionEstablishmentHandler)
             try {
                 val requester = interceptors.wrapRequester(connection)
@@ -77,14 +77,12 @@ internal abstract class ConnectionEstablishmentHandler<T>(
                     responder.cancel("Connection closed", cause)
                 }
 
-                val keepAliveHandler = KeepAliveHandler(connectionConfig.keepAlive, connection, this)
-                connection.launch {
-                    connection.handleConnection(
-                        ConnectionInbound(connection.coroutineContext, responder, keepAliveHandler)
-                    )
-                }
+                requesterDeferred?.complete(requester)
 
-                transform(requester)
+                val keepAliveHandler = KeepAliveHandler(connectionConfig.keepAlive, connection, this)
+                connection.handleConnection(
+                    ConnectionInbound(connection.coroutineContext, responder, keepAliveHandler)
+                )
             } catch (cause: Throwable) {
                 connectionConfig.setupPayload.close()
                 throw cause
@@ -103,7 +101,7 @@ internal abstract class ConnectionEstablishmentHandler<T>(
         }
     }
 
-    override suspend fun RSocketConnection.initialize(): T {
-        return handleConnection(wrapConnection(this, coroutineContext.supervisorContext()))
+    override suspend fun RSocketConnection.initialize() {
+        handleConnection(wrapConnection(this, coroutineContext.supervisorContext()))
     }
 }

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/connection/MultiplexedConnection.kt

@@ -40,7 +40,7 @@ internal class MultiplexedConnection(
         storage.clear()
     }
 
-    override suspend fun establishConnection(handler: ConnectionEstablishmentHandler<*>): ConnectionConfig {
+    override suspend fun establishConnection(handler: ConnectionEstablishmentHandler): ConnectionConfig {
         return handler.establishConnection(EstablishmentContext())
     }
 

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/connection/OldConnection.kt

@@ -39,7 +39,7 @@ internal class OldConnection(
     init {
         @OptIn(DelicateCoroutinesApi::class)
         launch(start = CoroutineStart.ATOMIC) {
-            val outboundJob = launch {
+            launch {
                 nonCancellable {
                     while (true) {
                         connection.send(outboundQueue.dequeueFrame() ?: break)
@@ -52,10 +52,7 @@ internal class OldConnection(
             try {
                 awaitCancellation()
             } finally {
-                nonCancellable {
-                    outboundQueue.close()
-                    outboundJob.join()
-                }
+                outboundQueue.close()
             }
         }
     }

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/connection/SequentialConnection.kt

@@ -38,7 +38,7 @@ internal class SequentialConnection(
         storage.clear().forEach { it.close() }
     }
 
-    override suspend fun establishConnection(handler: ConnectionEstablishmentHandler<*>): ConnectionConfig {
+    override suspend fun establishConnection(handler: ConnectionEstablishmentHandler): ConnectionConfig {
         return handler.establishConnection(EstablishmentContext())
     }
 

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/core/RSocketConnector.kt

@@ -20,8 +20,10 @@ import io.rsocket.kotlin.*
 import io.rsocket.kotlin.connection.*
 import io.rsocket.kotlin.frame.*
 import io.rsocket.kotlin.frame.io.*
+import io.rsocket.kotlin.internal.io.*
 import io.rsocket.kotlin.logging.*
 import io.rsocket.kotlin.transport.*
+import kotlinx.coroutines.*
 import kotlin.coroutines.*
 
 @OptIn(RSocketTransportApi::class, RSocketLoggingApi::class)
@@ -58,16 +60,27 @@ public class RSocketConnector internal constructor(
     }
 
     private suspend fun connectOnce(transport: RSocketClientTarget): RSocket {
-        return SetupConnection().runInitializer(transport.connectClient().logging(frameLogger))
+        val requesterDeferred = CompletableDeferred<RSocket>()
+        val connectJob = SetupConnection(requesterDeferred).launchInitializer(
+            transport.connectClient().logging(frameLogger)
+        ).onCompletion {
+            if (it != null) requesterDeferred.completeExceptionally(it)
+        }
+        return try {
+            requesterDeferred.await()
+        } catch (cause: Throwable) {
+            connectJob.cancel("RSocketConnector.connect was cancelled", cause)
+            throw cause
+        }
     }
 
-    private inner class SetupConnection : ConnectionEstablishmentHandler<RSocket>(
+    private inner class SetupConnection(requesterDeferred: CompletableDeferred<RSocket>) : ConnectionEstablishmentHandler(
         isClient = true,
         frameCodec = FrameCodec(maxFragmentSize),
         connectionAcceptor = acceptor,
-        interceptors = interceptors
+        interceptors = interceptors,
+        requesterDeferred = requesterDeferred
     ) {
-        override fun transform(requester: RSocket): RSocket = requester
         override suspend fun establishConnection(context: ConnectionEstablishmentContext): ConnectionConfig {
             val connectionConfig = connectionConfigProvider()
             try {
@@ -86,6 +99,5 @@ public class RSocketConnector internal constructor(
             }
             return connectionConfig
         }
-
     }
 }

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/core/RSocketServer.kt

@@ -63,14 +63,13 @@ public class RSocketServer internal constructor(
         AcceptConnection(acceptor).launchInitializer(connection.logging(frameLogger))
     }
 
-    private inner class AcceptConnection(acceptor: ConnectionAcceptor) : ConnectionEstablishmentHandler<Unit>(
+    private inner class AcceptConnection(acceptor: ConnectionAcceptor) : ConnectionEstablishmentHandler(
         isClient = false,
         frameCodec = FrameCodec(maxFragmentSize),
         connectionAcceptor = acceptor,
         interceptors = interceptors,
+        requesterDeferred = null
     ) {
-        override fun transform(requester: RSocket) = Unit
-
         override suspend fun establishConnection(context: ConnectionEstablishmentContext): ConnectionConfig {
             val setupFrame = context.receiveFrame()
             return try {