KTOR-6883 Add callbacks before waiting for connection establishment t…

…o prevent frame receiving delays (#4012) (cherry picked from commit 0837f12)
ktorio · Apr 5, 2024 · 28f1d00 · 28f1d00
1 parent df04b16
commit 28f1d00
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/ktor-client/ktor-client-core/js/src/io/ktor/client/engine/js/JsClientEngine.kt b/ktor-client/ktor-client-core/js/src/io/ktor/client/engine/js/JsClientEngine.kt
@@ -90,15 +90,15 @@ internal class JsClientEngine(
         val urlString = request.url.toString()
         val socket: WebSocket = createWebSocket(urlString, request.headers)
 
+        val session = JsWebSocketSession(callContext, socket)
+
         try {
             socket.awaitConnection()
         } catch (cause: Throwable) {
             callContext.cancel(CancellationException("Failed to connect to $urlString", cause))
             throw cause
         }
 
-        val session = JsWebSocketSession(callContext, socket)
-
         return HttpResponseData(
             HttpStatusCode.SwitchingProtocols,
             requestTime,

diff --git a/ktor-client/ktor-client-tests/common/test/io/ktor/client/tests/WebSocketTest.kt b/ktor-client/ktor-client-tests/common/test/io/ktor/client/tests/WebSocketTest.kt
@@ -17,6 +17,7 @@ import io.ktor.utils.io.charsets.*
 import io.ktor.websocket.*
 import kotlinx.coroutines.*
 import kotlin.test.*
+import kotlin.time.Duration.Companion.seconds
 
 internal val ENGINES_WITHOUT_WS = listOf("Android", "Apache", "Apache5", "Curl", "DarwinLegacy")
 
@@ -322,4 +323,28 @@ class WebSocketTest : ClientLoader() {
             }
         }
     }
+
+    @Test
+    fun testImmediateReceiveAfterConnect() = clientTests(ENGINES_WITHOUT_WS) {
+        config {
+            install(WebSockets)
+        }
+
+        test { client ->
+            withTimeout(10_000) {
+                coroutineScope {
+                    val defs = (1..100).map {
+                        async {
+                            client.webSocket("$TEST_WEBSOCKET_SERVER/websockets/headers") {
+                                val frame = withTimeoutOrNull(1.seconds) { incoming.receive() }
+                                assertNotNull(frame)
+                                assertIs<Frame.Text>(frame)
+                            }
+                        }
+                    }
+                    defs.awaitAll()
+                }
+            }
+        }
+    }
 }