From bf3e4ee179a071cd2e981feb2d59d88f5b42a148 Mon Sep 17 00:00:00 2001 From: NathanFallet Date: Mon, 20 Jul 2026 16:10:29 +0200 Subject: [PATCH 1/2] fix: auto reconnect --- .../core/connection/DefaultConnection.kt | 132 +++++++++---- .../core/connection/CommandTimeoutTest.kt | 2 +- .../connection/ReconnectionIntegrationTest.kt | 100 ++++++++++ .../core/connection/ReconnectionTest.kt | 180 ++++++++++++++++++ .../dev/kdriver/core/connection/TcpProxy.kt | 107 +++++++++++ 5 files changed, 488 insertions(+), 33 deletions(-) create mode 100644 core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionIntegrationTest.kt create mode 100644 core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionTest.kt create mode 100644 core/src/jvmTest/kotlin/dev/kdriver/core/connection/TcpProxy.kt diff --git a/core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt b/core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt index 9863b60f0..9c58b6cc2 100644 --- a/core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt +++ b/core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt @@ -17,6 +17,7 @@ import kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.withLock import kotlinx.serialization.json.JsonElement import kotlinx.serialization.json.jsonPrimitive +import kotlin.concurrent.Volatile import kotlin.reflect.KClass import kotlin.time.Clock import kotlin.time.Duration.Companion.milliseconds @@ -31,9 +32,21 @@ open class DefaultConnection( override var owner: Browser? = null, ) : OwnedConnection { + companion object { + // One initial send plus one resend after a transparent reconnect. The resend is only ever reached + // when a send fails (bytes never left the client), so it cannot double-execute a command. + private const val SEND_ATTEMPTS = 2 + } + private val logger = KtorSimpleLogger("Connection") - private val transport: WebSocketTransport by lazy { createTransport() } + // Recreated on every (re)connect. A dropped socket can keep reporting a stale `isActive` for a + // short window, so once the receive loop has observed a disconnect we throw the whole transport + // away and build a fresh one rather than trusting it (see [connect]). + private var transport: WebSocketTransport? = null + + @Volatile + private var disconnected = false private var socketSubscription: Job? = null @@ -67,22 +80,35 @@ open class DefaultConnection( @InternalCdpApi override val generatedDomains: MutableMap, Domain> = mutableMapOf() - private suspend fun connect() { - if (transport.isActive) return + private suspend fun connect(): WebSocketTransport { + transport?.let { if (it.isActive && !disconnected) return it } // Guard so concurrent first commands don't each open a session (which would leak the // duplicate sockets/listeners). Double-checked: skip the lock once connected (ISSUE-4). - connectMutex.withLock { - if (transport.isActive) return@withLock - transport.connect() - startListening() + return connectMutex.withLock { + transport?.let { if (it.isActive && !disconnected) return@withLock it } + // Once the receive loop has reported a disconnect we don't trust the old transport's + // `isActive` (it can lag the real socket close), so dispose it and build a fresh one. This + // makes a command issued after a drop reconnect deterministically instead of reusing a + // dead socket and failing (or hanging until the command timeout). + if (disconnected) { + transport?.let { old -> runCatching { old.close() } } + transport = null + disconnected = false + } + val t = transport ?: createTransport().also { transport = it } + if (!t.isActive) { + t.connect() + startListening(t) + } + t } } - private fun startListening() { + private fun startListening(t: WebSocketTransport) { socketSubscription?.cancel() socketSubscription = messageListeningScope.launch { try { - transport.incoming().collect { text -> + t.incoming().collect { text -> try { logger.debug("WS < CDP: ${text.take(owner?.config?.debugStringLimit ?: Defaults.DEBUG_STRING_LIMIT)}") val received = Serialization.json.decodeFromString(text) @@ -99,16 +125,27 @@ open class DefaultConnection( } // incoming() completed without error => the socket was closed. Fail any in-flight // commands so their callers observe the disconnect instead of hanging (ISSUE-3). - failPendingRequests(ConnectionClosedException()) + onTransportGone(t, ConnectionClosedException()) } catch (e: CancellationException) { throw e } catch (e: Exception) { logger.error("WebSocket receive loop terminated: {}", e) - failPendingRequests(ConnectionClosedException(cause = e)) + onTransportGone(t, ConnectionClosedException(cause = e)) } } } + /** + * Reacts to a receive loop ending because its socket went away: marks the connection for a + * reconnect and fails any in-flight commands. Ignored if [t] is no longer the current transport + * (a stale loop from a transport we've already replaced must not disturb the new one). + */ + private suspend fun onTransportGone(t: WebSocketTransport, cause: Throwable) { + if (transport !== t) return + disconnected = true + failPendingRequests(cause) + } + /** * Completes every in-flight request waiter exceptionally and clears the registry, so callers * parked in [callCommand] observe a failure rather than hanging when the connection goes away. @@ -136,33 +173,64 @@ open class DefaultConnection( } val requestId = currentIdMutex.withLock { currentId++ } - // Register the response waiter *before* sending, so a reply that arrives before we start - // awaiting is still captured (the receive loop completes this deferred). Awaiting the - // response via a replay-0 shared flow after sending could miss it and hang (ISSUE-1). - val deferred = CompletableDeferred() - pendingRequestsMutex.withLock { pendingRequests[requestId] = deferred } - try { - val jsonString = Serialization.json.encodeToString(Request(requestId, method, parameter)) - transport.send(jsonString) - logger.debug("WS > CDP: ${jsonString.take(owner?.config?.debugStringLimit ?: Defaults.DEBUG_STRING_LIMIT)}") + val jsonString = Serialization.json.encodeToString(Request(requestId, method, parameter)) + val timeout = owner?.config?.commandTimeout ?: Defaults.COMMAND_TIMEOUT + + // A dropped socket can still report itself active for a short window, so the send below can + // fail (or the next command can land on a dead transport). A failed send means the bytes + // never left the client, so reconnecting and resending is safe — no risk of executing the + // command twice. We only retry the *send*: once a reply is awaited we never resend, because + // the command may already have run (that case surfaces as a ConnectionClosedException). + var sendError: Throwable? = null + repeat(SEND_ATTEMPTS) { + val transport = connect() + // Register the response waiter *before* sending, so a reply that arrives before we start + // awaiting is still captured (the receive loop completes this deferred). Awaiting the + // response via a replay-0 shared flow after sending could miss it and hang (ISSUE-1). + val deferred = CompletableDeferred() + pendingRequestsMutex.withLock { pendingRequests[requestId] = deferred } + + val sent = try { + transport.send(jsonString) + true + } catch (e: CancellationException) { + // A send on a dead Ktor session surfaces as a channel CancellationException even + // though *this* coroutine isn't cancelled; only the latter must propagate. + if (!currentCoroutineContext().isActive) throw e + sendError = e + disconnected = true + false + } catch (e: Exception) { + sendError = e + disconnected = true + false + } + if (!sent) { + pendingRequestsMutex.withLock { pendingRequests.remove(requestId) } + return@repeat // reconnect and resend on the next iteration + } - val timeout = owner?.config?.commandTimeout ?: Defaults.COMMAND_TIMEOUT - // A non-null Message.Response means success; null can only come from the timeout, so - // there's no ambiguity with a legitimate value. A timeout <= 0 waits indefinitely. - val result = - if (timeout > 0) withTimeoutOrNull(timeout.milliseconds) { deferred.await() } - ?: throw CommandTimeoutException(method, requestId, timeout) - else deferred.await() - result.error?.throwAsException(method) - return result.result - } finally { - pendingRequestsMutex.withLock { pendingRequests.remove(requestId) } + logger.debug("WS > CDP: ${jsonString.take(owner?.config?.debugStringLimit ?: Defaults.DEBUG_STRING_LIMIT)}") + try { + // A non-null Message.Response means success; null can only come from the timeout, so + // there's no ambiguity with a legitimate value. A timeout <= 0 waits indefinitely. + val result = + if (timeout > 0) withTimeoutOrNull(timeout.milliseconds) { deferred.await() } + ?: throw CommandTimeoutException(method, requestId, timeout) + else deferred.await() + result.error?.throwAsException(method) + return result.result + } finally { + pendingRequestsMutex.withLock { pendingRequests.remove(requestId) } + } } + throw sendError ?: ConnectionClosedException() } @InternalCdpApi override suspend fun close() { - transport.close() + transport?.close() + transport = null socketSubscription?.cancel() socketSubscription = null // Fail any commands still awaiting a reply that will now never come (ISSUE-3). diff --git a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/CommandTimeoutTest.kt b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/CommandTimeoutTest.kt index a4c26f6ee..52b25d57e 100644 --- a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/CommandTimeoutTest.kt +++ b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/CommandTimeoutTest.kt @@ -43,7 +43,7 @@ class CommandTimeoutTest { @Test fun callCommand_failsWithCommandTimeout_whenNoResponse() = runTest(UnconfinedTestDispatcher()) { val transport = SilentTransport() - val connection = object : DefaultConnection("ws://stub/devtools/page/stub", this) { + val connection = object : DefaultConnection("ws://stub/devtools/page/stub", this@runTest) { override fun createTransport(): WebSocketTransport = transport } diff --git a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionIntegrationTest.kt b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionIntegrationTest.kt new file mode 100644 index 000000000..3c87e5472 --- /dev/null +++ b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionIntegrationTest.kt @@ -0,0 +1,100 @@ +package dev.kdriver.core.connection + +import dev.kdriver.cdp.CommandMode +import dev.kdriver.core.browser.createBrowser +import dev.kdriver.core.exceptions.CommandTimeoutException +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.withTimeout +import kotlinx.serialization.json.* +import kotlin.test.* + +/** + * End-to-end reconnection behaviour against a real Chrome, with a [TcpProxy] in the middle so the + * test can cut or freeze the socket at will. + * + * A real browser is launched (the "host"), and the test opens its own [DefaultConnection] to one of + * its page targets *through the proxy* — the per-target WebSocket URL is built from host:port, so it + * always goes through the proxy regardless of what Chrome advertises for the browser endpoint. + */ +class ReconnectionIntegrationTest { + + private class ProxiedConnection( + private val wsUrl: String, + scope: CoroutineScope, + ) : DefaultConnection(wsUrl, scope) { + override fun createTransport(): WebSocketTransport = KtorWebSocketTransport(wsUrl) + } + + private fun evalParam(expression: String): JsonObject = buildJsonObject { + put("expression", expression) + put("returnByValue", true) + } + + private suspend fun DefaultConnection.evaluate(expression: String): Int { + val result = callCommand("Runtime.evaluate", evalParam(expression), CommandMode.ONE_SHOT) + return result!!.jsonObject["result"]!!.jsonObject["value"]!!.jsonPrimitive.int + } + + @Test + fun reconnects_afterHardCut_throughProxy() = runBlocking { + val host = createBrowser(this, headless = true, sandbox = false) + val proxy = TcpProxy("127.0.0.1", host.config.port!!, this).also { it.start() } + val targetId = assertNotNull(host.mainTab?.targetId, "host browser should have a page target") + val wsUrl = "ws://127.0.0.1:${proxy.port}/devtools/page/$targetId" + val conn = ProxiedConnection(wsUrl, this) + + try { + // Sanity: the proxied connection works end-to-end. + assertEquals(2, conn.evaluate("1 + 1")) + + // Hard cut: the socket is closed under the connection. + proxy.severAll() + + // The next command must transparently reconnect through the proxy and succeed, without + // the caller having to do anything. + assertEquals(4, withTimeout(15_000) { conn.evaluate("2 + 2") }) + } finally { + conn.close() + proxy.stop() + host.stop() + } + } + + @Test + fun halfOpenSocket_isBoundedByCommandTimeout_andRecovers() = runBlocking { + // Short command timeout so a half-open socket fails fast instead of hanging for 30s. + val host = createBrowser(this) { + headless = true + sandbox = false + commandTimeout = 3_000 + } + val proxy = TcpProxy("127.0.0.1", host.config.port!!, this).also { it.start() } + val targetId = assertNotNull(host.mainTab?.targetId, "host browser should have a page target") + val wsUrl = "ws://127.0.0.1:${proxy.port}/devtools/page/$targetId" + // owner = host so the connection picks up the short commandTimeout from its config. + val conn = ProxiedConnection(wsUrl, this).also { it.owner = host } + + try { + assertEquals(2, conn.evaluate("1 + 1")) + + // Half-open: sockets stay open but no application bytes flow, so the reply never comes + // and the socket never closes. This is the case keep-alive pings would be needed to + // detect; for now it is merely *bounded* by the command timeout rather than hanging. + proxy.freeze() + val start = System.currentTimeMillis() + val error = runCatching { conn.evaluate("2 + 2") }.exceptionOrNull() + val elapsed = System.currentTimeMillis() - start + assertIs(error, "a half-open socket must fail with a command timeout, not hang") + assertTrue(elapsed in 3_000..15_000, "should fail around the 3s command timeout (was ${elapsed}ms)") + + // Once traffic flows again the same session is healthy, so commands resume working. + proxy.unfreeze() + assertEquals(9, withTimeout(15_000) { conn.evaluate("3 + 3 + 3") }) + } finally { + conn.close() + proxy.stop() + host.stop() + } + } +} diff --git a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionTest.kt b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionTest.kt new file mode 100644 index 000000000..5fea85dac --- /dev/null +++ b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionTest.kt @@ -0,0 +1,180 @@ +package dev.kdriver.core.connection + +import dev.kdriver.cdp.CommandMode +import dev.kdriver.cdp.Serialization +import dev.kdriver.core.exceptions.CommandTimeoutException +import kotlinx.coroutines.* +import kotlinx.coroutines.channels.Channel +import kotlinx.coroutines.flow.Flow +import kotlinx.coroutines.flow.receiveAsFlow +import kotlinx.coroutines.test.UnconfinedTestDispatcher +import kotlinx.coroutines.test.runTest +import kotlinx.serialization.json.jsonObject +import kotlinx.serialization.json.jsonPrimitive +import kotlinx.serialization.json.long +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertIs +import kotlin.test.assertNotNull + +/** + * Behaviour of [DefaultConnection] when the underlying transport disconnects and a later command + * must transparently re-establish the connection. + * + * The production symptom (commands timing out at 30s across a session) points at a connection that + * has gone away without the caller being able to recover. These tests pin down what happens at the + * transport seam so we can reason about reconnection independently of a real browser. + */ +class ReconnectionTest { + + /** + * A transport the test fully drives: + * - [respondTo] / [respond] inject a CDP reply for an in-flight request, + * - [drop] simulates a clean socket close (FIN): `incoming()` completes, [isActive] flips false, + * - going silent (never responding, never dropping) simulates a half-open socket. + * + * [connectCount] lets a test assert whether a *new* connection was opened (i.e. a reconnect). + */ + private class ControllableTransport : WebSocketTransport { + private var channel = Channel(Channel.UNLIMITED) + + @Volatile + override var isActive: Boolean = false + private set + + var connectCount: Int = 0 + private set + + val sent = mutableListOf() + + /** When true, [send] fails as a dead Ktor session would; cleared on the next [connect]. */ + @Volatile + var failSend: Boolean = false + + override suspend fun connect() { + connectCount++ + // A reconnect reuses this instance, so hand out a fresh channel for the new receive loop. + if (channel.isClosedForSend) channel = Channel(Channel.UNLIMITED) + failSend = false + isActive = true + } + + override suspend fun send(message: String) { + // A send on a dropped Ktor session throws a channel CancellationException. + if (failSend) throw CancellationException("Channel was cancelled") + sent.add(message) + } + + override fun incoming(): Flow = channel.receiveAsFlow() + + override suspend fun close() { + isActive = false + channel.close() + } + + /** Simulate the socket dropping cleanly: the receive loop sees [incoming] complete. */ + fun drop() { + isActive = false + channel.close() + } + + /** Deliver a successful CDP reply for the request with the given [id]. */ + fun respond(id: Long, result: String = """{}""") { + channel.trySend("""{"id":$id,"result":$result}""") + } + + /** Deliver a successful CDP reply for the most recently sent request. */ + fun respondToLast(result: String = """{}""") { + val id = Serialization.json.parseToJsonElement(sent.last()) + .jsonObject["id"]!!.jsonPrimitive.long + respond(id, result) + } + } + + private class TestConnection( + scope: CoroutineScope, + private val transport: ControllableTransport, + ) : DefaultConnection("ws://stub/devtools/page/stub", scope) { + override fun createTransport(): WebSocketTransport = transport + } + + @Test + fun reconnects_onNextCommand_afterCleanDrop() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport() + val connection = TestConnection(this, transport) + + // First command opens the connection and succeeds. + val first = async { connection.callCommand("A.first", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + assertNotNull(withTimeout(2_000) { first.await() }) + assertEquals(1, transport.connectCount) + + // The socket drops cleanly. + transport.drop() + + // The next command must transparently re-open the transport and succeed. + val second = async { connection.callCommand("B.second", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + assertNotNull(withTimeout(2_000) { second.await() }) + assertEquals(2, transport.connectCount, "a clean drop must lead to a reconnect on next command") + + connection.close() + } + + @Test + fun reconnectsAndResends_whenSendFailsOnDeadSocket() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport() + val connection = TestConnection(this, transport) + + val first = async { connection.callCommand("A.first", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + assertNotNull(withTimeout(2_000) { first.await() }) + + // The socket is dead but still reports active (the race window): the send fails. The command + // must reconnect and resend transparently rather than surfacing the failure. + transport.failSend = true + val second = async { connection.callCommand("B.second", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + assertNotNull(withTimeout(2_000) { second.await() }) + assertEquals(2, transport.connectCount, "a failed send must trigger a reconnect and resend") + + connection.close() + } + + @Test + fun reusesConnection_afterHalfOpenTimeout() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport() + val connection = TestConnection(this, transport) + + // Half-open socket: the command is sent but no reply ever comes and the socket never drops. + // With no owner attached, the connection falls back to Config.Defaults.COMMAND_TIMEOUT; the + // virtual clock advances it instantly. + val first = CompletableDeferred() + launch { + try { + connection.callCommand("A.first", null, CommandMode.ONE_SHOT) + } catch (e: Throwable) { + first.complete(e) + } + } + assertIs(withTimeout(60_000) { first.await() }) + + // Characterise the seam: on a half-open socket the transport still *looks* active, so the + // next command reuses the same (dead) socket and times out again. DefaultConnection alone + // cannot recover here -- detection has to come from the transport reporting a disconnect. + // That is what the WebSocket keep-alive ping does in production (KtorWebSocketTransport): a + // missing pong closes the session, turning this case into reconnects_onNextCommand_afterCleanDrop. + val second = CompletableDeferred() + launch { + try { + connection.callCommand("B.second", null, CommandMode.ONE_SHOT) + } catch (e: Throwable) { + second.complete(e) + } + } + assertIs(withTimeout(60_000) { second.await() }) + assertEquals(1, transport.connectCount, "current behaviour: no reconnect after a timeout") + + connection.close() + } +} diff --git a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/TcpProxy.kt b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/TcpProxy.kt new file mode 100644 index 000000000..6245c9a23 --- /dev/null +++ b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/TcpProxy.kt @@ -0,0 +1,107 @@ +package dev.kdriver.core.connection + +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.launch +import java.net.ServerSocket +import java.net.Socket +import java.util.* + +/** + * A dumb TCP relay used by reconnection integration tests. + * + * It forwards raw bytes between a kdriver client and a real Chrome DevTools endpoint, which makes it + * transparent for both the HTTP (`/json/version`) and WebSocket traffic. Because it sits on the TCP + * stream it can reproduce, on demand, the two failure modes we care about: + * - [severAll] closes the sockets (RST/FIN) -> a clean disconnect, + * - [freeze]/[unfreeze] keep the sockets open but stop forwarding application bytes -> a half-open + * socket (the peer looks alive at the TCP level but never answers), which only WebSocket keep-alive + * pings can detect. + * + * @param upstreamHost Host of the real Chrome DevTools endpoint. + * @param upstreamPort Port of the real Chrome DevTools endpoint. + * @param scope Scope the relay coroutines run in. + */ +class TcpProxy( + private val upstreamHost: String, + private val upstreamPort: Int, + private val scope: CoroutineScope, +) { + + private val server = ServerSocket(0) + private val sockets = Collections.synchronizedList(mutableListOf()) + + @Volatile + private var forwarding = true + + /** The local port clients should connect to. */ + val port: Int get() = server.localPort + + fun start() { + scope.launch(Dispatchers.IO) { + while (!server.isClosed) { + val client = try { + server.accept() + } catch (_: Exception) { + break + } + val upstream = try { + Socket(upstreamHost, upstreamPort) + } catch (_: Exception) { + runCatching { client.close() } + continue + } + sockets.add(client) + sockets.add(upstream) + pump(client, upstream) + pump(upstream, client) + } + } + } + + private fun pump(from: Socket, to: Socket) { + scope.launch(Dispatchers.IO) { + val buffer = ByteArray(16 * 1024) + try { + val input = from.getInputStream() + val output = to.getOutputStream() + while (true) { + val read = input.read(buffer) + if (read < 0) break + // While frozen we keep draining the socket (so TCP stays healthy and no RST is + // sent) but drop the bytes: the peer sees silence at the application level. + if (forwarding) { + output.write(buffer, 0, read) + output.flush() + } + } + } catch (_: Exception) { + // Socket closed (severed or torn down): just end this pump. + } + } + } + + /** Simulate a clean disconnect: close every live socket so both ends see the stream end. */ + fun severAll() { + synchronized(sockets) { + sockets.forEach { runCatching { it.close() } } + sockets.clear() + } + } + + /** Simulate a half-open socket: keep connections open but stop forwarding application bytes. */ + fun freeze() { + forwarding = false + } + + /** Resume forwarding after a [freeze]. */ + fun unfreeze() { + forwarding = true + } + + fun stop() { + severAll() + runCatching { server.close() } + } + +} From 1159707aea781ab04a87e0141f24de79e537a6ec Mon Sep 17 00:00:00 2001 From: NathanFallet Date: Mon, 20 Jul 2026 18:56:14 +0200 Subject: [PATCH 2/2] fixes --- .../core/connection/DefaultConnection.kt | 45 ++++- .../core/connection/KtorWebSocketTransport.kt | 7 +- .../core/connection/ReconnectionTest.kt | 188 +++++++++++++++--- .../dev/kdriver/core/connection/TcpProxy.kt | 9 +- 4 files changed, 205 insertions(+), 44 deletions(-) diff --git a/core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt b/core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt index 9c58b6cc2..ceb343f7d 100644 --- a/core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt +++ b/core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt @@ -42,12 +42,20 @@ open class DefaultConnection( // Recreated on every (re)connect. A dropped socket can keep reporting a stale `isActive` for a // short window, so once the receive loop has observed a disconnect we throw the whole transport - // away and build a fresh one rather than trusting it (see [connect]). + // away and build a fresh one rather than trusting it (see [connect]). @Volatile because it is + // read off-lock on the fast path of [connect], from the receive-loop coroutine in [onTransportGone], + // and in [close]. + @Volatile private var transport: WebSocketTransport? = null @Volatile private var disconnected = false + // Set once [close] has run; terminal. [connect] refuses to re-open a closed connection so a + // command issued after close fails instead of silently resurrecting the socket. + @Volatile + private var closed = false + private var socketSubscription: Job? = null private val connectMutex = Mutex() @@ -81,10 +89,12 @@ open class DefaultConnection( override val generatedDomains: MutableMap, Domain> = mutableMapOf() private suspend fun connect(): WebSocketTransport { + if (closed) throw ConnectionClosedException("Connection closed") transport?.let { if (it.isActive && !disconnected) return it } // Guard so concurrent first commands don't each open a session (which would leak the // duplicate sockets/listeners). Double-checked: skip the lock once connected (ISSUE-4). return connectMutex.withLock { + if (closed) throw ConnectionClosedException("Connection closed") transport?.let { if (it.isActive && !disconnected) return@withLock it } // Once the receive loop has reported a disconnect we don't trust the old transport's // `isActive` (it can lag the real socket close), so dispose it and build a fresh one. This @@ -94,6 +104,12 @@ open class DefaultConnection( transport?.let { old -> runCatching { old.close() } } transport = null disconnected = false + // A reconnect opens a *new* CDP session, so the browser-side preparation applied on + // the old one is gone: re-run it on the next command. NOTE: domains the caller enabled + // themselves (Network.enable, Page.enable, event subscriptions, …) are NOT restored — + // restoring arbitrary caller state on reconnect is out of scope here. + prepareHeadlessDone = false + prepareExpertDone = false } val t = transport ?: createTransport().also { transport = it } if (!t.isActive) { @@ -198,14 +214,18 @@ open class DefaultConnection( // though *this* coroutine isn't cancelled; only the latter must propagate. if (!currentCoroutineContext().isActive) throw e sendError = e - disconnected = true false } catch (e: Exception) { sendError = e - disconnected = true false } if (!sent) { + // Treat the failed send as this transport going away: mark it for reconnect and fail + // any *other* commands in flight on it (ISSUE-3), but only if it is still the current + // transport — a send that fails late on an already-replaced transport must not tear + // down the healthy one that superseded it. onTransportGone clears our own just-added + // waiter too; we resend on the next iteration. + onTransportGone(transport, ConnectionClosedException(cause = sendError)) pendingRequestsMutex.withLock { pendingRequests.remove(requestId) } return@repeat // reconnect and resend on the next iteration } @@ -224,15 +244,24 @@ open class DefaultConnection( pendingRequestsMutex.withLock { pendingRequests.remove(requestId) } } } - throw sendError ?: ConnectionClosedException() + // Every send attempt failed: the connection is gone. Surface it as a ConnectionClosedException + // (never the raw send failure, which may be a CancellationException a caller would mistake for + // its own cancellation), keeping the original send error as the cause. + throw ConnectionClosedException(cause = sendError) } @InternalCdpApi override suspend fun close() { - transport?.close() - transport = null - socketSubscription?.cancel() - socketSubscription = null + // Take connectMutex so close is ordered against an in-progress (re)connect: without it, a + // concurrent connect() could reassign `transport`/`socketSubscription` right after we cleared + // them, resurrecting a live socket and receive loop that outlive close(). + connectMutex.withLock { + closed = true + transport?.close() + transport = null + socketSubscription?.cancel() + socketSubscription = null + } // Fail any commands still awaiting a reply that will now never come (ISSUE-3). failPendingRequests(ConnectionClosedException("Connection closed")) } diff --git a/core/src/commonMain/kotlin/dev/kdriver/core/connection/KtorWebSocketTransport.kt b/core/src/commonMain/kotlin/dev/kdriver/core/connection/KtorWebSocketTransport.kt index 5670f145b..5fd806d3f 100644 --- a/core/src/commonMain/kotlin/dev/kdriver/core/connection/KtorWebSocketTransport.kt +++ b/core/src/commonMain/kotlin/dev/kdriver/core/connection/KtorWebSocketTransport.kt @@ -5,9 +5,9 @@ import io.ktor.client.* import io.ktor.client.plugins.websocket.* import io.ktor.http.* import io.ktor.websocket.* -import kotlinx.coroutines.isActive import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.flow +import kotlinx.coroutines.isActive /** * Default [WebSocketTransport] backed by a Ktor WebSocket session. @@ -39,7 +39,10 @@ class KtorWebSocketTransport( } override suspend fun send(message: String) { - session?.send(message) + // Fail loudly instead of silently succeeding when there is no session: a false success would + // make the caller wait for a reply to a frame that was never sent. + val session = session ?: throw IllegalStateException("WebSocket session is not connected") + session.send(message) } override fun incoming(): Flow = flow { diff --git a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionTest.kt b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionTest.kt index 5fea85dac..5ed9c7b97 100644 --- a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionTest.kt +++ b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionTest.kt @@ -2,7 +2,12 @@ package dev.kdriver.core.connection import dev.kdriver.cdp.CommandMode import dev.kdriver.cdp.Serialization +import dev.kdriver.core.browser.Browser +import dev.kdriver.core.browser.Config import dev.kdriver.core.exceptions.CommandTimeoutException +import dev.kdriver.core.exceptions.ConnectionClosedException +import io.mockk.every +import io.mockk.mockk import kotlinx.coroutines.* import kotlinx.coroutines.channels.Channel import kotlinx.coroutines.flow.Flow @@ -29,9 +34,12 @@ class ReconnectionTest { /** * A transport the test fully drives: - * - [respondTo] / [respond] inject a CDP reply for an in-flight request, - * - [drop] simulates a clean socket close (FIN): `incoming()` completes, [isActive] flips false, - * - going silent (never responding, never dropping) simulates a half-open socket. + * - [respond] / [respondToLast] inject a CDP reply for an in-flight request, + * - [drop] simulates a dropped socket that still *reports* active (the production race): only the + * receive loop ending signals the disconnect, + * - [failSend] fails the next send like a dead Ktor session (cleared on reconnect); [keepFailingSend] + * keeps failing across reconnects, + * - [autoRespond] echoes a canned reply for every sent frame (used to drive the prepare* path). * * [connectCount] lets a test assert whether a *new* connection was opened (i.e. a reconnect). */ @@ -47,10 +55,18 @@ class ReconnectionTest { val sent = mutableListOf() - /** When true, [send] fails as a dead Ktor session would; cleared on the next [connect]. */ + /** Fails the next send as a dropped Ktor session would; cleared on the next [connect]. */ @Volatile var failSend: Boolean = false + /** Like [failSend] but persists across reconnects, to exercise exhausted retries. */ + @Volatile + var keepFailingSend: Boolean = false + + /** When true, every sent frame gets an immediate canned reply (drives the prepare* path). */ + @Volatile + var autoRespond: Boolean = false + override suspend fun connect() { connectCount++ // A reconnect reuses this instance, so hand out a fresh channel for the new receive loop. @@ -61,8 +77,19 @@ class ReconnectionTest { override suspend fun send(message: String) { // A send on a dropped Ktor session throws a channel CancellationException. - if (failSend) throw CancellationException("Channel was cancelled") + if (failSend || keepFailingSend) throw CancellationException("Channel was cancelled") sent.add(message) + if (autoRespond) { + val request = Serialization.json.parseToJsonElement(message).jsonObject + val id = request["id"]!!.jsonPrimitive.long + val method = request["method"]!!.jsonPrimitive.content + // Runtime.evaluate must parse into an EvaluateReturn with a value so prepareHeadless + // completes (and sets its done flag); anything else just needs a well-formed reply. + val result = if (method == "Runtime.evaluate") + """{"result":{"type":"string","value":"HeadlessChrome/1.0"}}""" + else "{}" + channel.trySend("""{"id":$id,"result":$result}""") + } } override fun incoming(): Flow = channel.receiveAsFlow() @@ -72,9 +99,8 @@ class ReconnectionTest { channel.close() } - /** Simulate the socket dropping cleanly: the receive loop sees [incoming] complete. */ + /** Simulate a dropped socket that still reports active: only the receive loop ends. */ fun drop() { - isActive = false channel.close() } @@ -98,6 +124,19 @@ class ReconnectionTest { override fun createTransport(): WebSocketTransport = transport } + /** Runs [block] and captures its thrown exception without failing the test scope. */ + private fun CoroutineScope.captureError(block: suspend () -> Unit): CompletableDeferred { + val error = CompletableDeferred() + launch { + try { + block() + } catch (e: Throwable) { + error.complete(e) + } + } + return error + } + @Test fun reconnects_onNextCommand_afterCleanDrop() = runTest(UnconfinedTestDispatcher()) { val transport = ControllableTransport() @@ -109,14 +148,14 @@ class ReconnectionTest { assertNotNull(withTimeout(2_000) { first.await() }) assertEquals(1, transport.connectCount) - // The socket drops cleanly. + // The socket drops but keeps reporting active (the production race): only the receive loop + // ending flags the disconnect. The next command must reconnect via that signal. transport.drop() - // The next command must transparently re-open the transport and succeed. val second = async { connection.callCommand("B.second", null, CommandMode.ONE_SHOT) } transport.respondToLast() assertNotNull(withTimeout(2_000) { second.await() }) - assertEquals(2, transport.connectCount, "a clean drop must lead to a reconnect on next command") + assertEquals(2, transport.connectCount, "a drop must lead to a reconnect on the next command") connection.close() } @@ -141,6 +180,105 @@ class ReconnectionTest { connection.close() } + @Test + fun failsOtherInFlightCommand_whenASendTriggersReconnect() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport() + val connection = TestConnection(this, transport) + + // Command A is sent and parks awaiting its reply. + val aError = captureError { connection.callCommand("A.first", null, CommandMode.ONE_SHOT) } + + // Command B's send fails on the (now dead) socket. The reconnect it triggers must also fail A, + // which is in flight on that same dead transport, instead of leaving it to hang to the timeout. + transport.failSend = true + val second = async { connection.callCommand("B.second", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + + assertNotNull(withTimeout(2_000) { second.await() }) + assertIs(withTimeout(2_000) { aError.await() }) + assertEquals(2, transport.connectCount) + + connection.close() + } + + @Test + fun doesNotResend_afterAnAwaitedReplyDies() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport() + val connection = TestConnection(this, transport) + + // Command is sent (successfully) and parks awaiting its reply — the resend-unsafe window. + val error = captureError { connection.callCommand("A.first", null, CommandMode.ONE_SHOT) } + + // The socket dies while awaiting. The command must surface the failure, NOT resend (it may + // already have executed). + transport.drop() + + assertIs(withTimeout(2_000) { error.await() }) + assertEquals(1, transport.connectCount, "a command must not resend once it is awaiting a reply") + + connection.close() + } + + @Test + fun throwsConnectionClosed_whenAllSendAttemptsFail() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport().apply { keepFailingSend = true } + val connection = TestConnection(this, transport) + + val error = captureError { connection.callCommand("A.first", null, CommandMode.ONE_SHOT) } + + // Bounded retries: initial connect + one reconnect attempt, then give up with a clear error + // (never the raw CancellationException from the failed send). + assertIs(withTimeout(2_000) { error.await() }) + assertEquals(2, transport.connectCount) + + connection.close() + } + + @Test + fun callCommandAfterClose_throwsAndDoesNotReconnect() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport() + val connection = TestConnection(this, transport) + + val first = async { connection.callCommand("A.first", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + assertNotNull(withTimeout(2_000) { first.await() }) + + connection.close() + + val error = captureError { connection.callCommand("B.second", null, CommandMode.ONE_SHOT) } + assertIs(withTimeout(2_000) { error.await() }) + assertEquals(1, transport.connectCount, "close is terminal: no silent reconnect afterwards") + } + + @Test + fun replaysPreparation_afterReconnect() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport().apply { autoRespond = true } + // A headless browser triggers prepareHeadless (a Runtime.evaluate + Network.setUserAgentOverride). + val config = mockk(relaxed = true) + every { config.headless } returns true + every { config.expert } returns false + val browser = mockk(relaxed = true) + every { browser.config } returns config + val connection = TestConnection(this, transport).also { it.owner = browser } + + fun evaluateCount() = transport.sent.count { it.contains("Runtime.evaluate") } + + // First DEFAULT command runs the preparation once. + connection.callCommand("Some.method", null, CommandMode.DEFAULT) + assertEquals(1, evaluateCount()) + + // While still connected, preparation is not repeated. + connection.callCommand("Some.method", null, CommandMode.DEFAULT) + assertEquals(1, evaluateCount(), "preparation must not repeat while the session is alive") + + // After a reconnect the new CDP session has lost the preparation, so it must be re-applied. + transport.drop() + connection.callCommand("Some.method", null, CommandMode.DEFAULT) + assertEquals(2, evaluateCount(), "preparation must be replayed on the new session after a reconnect") + + connection.close() + } + @Test fun reusesConnection_afterHalfOpenTimeout() = runTest(UnconfinedTestDispatcher()) { val transport = ControllableTransport() @@ -149,31 +287,17 @@ class ReconnectionTest { // Half-open socket: the command is sent but no reply ever comes and the socket never drops. // With no owner attached, the connection falls back to Config.Defaults.COMMAND_TIMEOUT; the // virtual clock advances it instantly. - val first = CompletableDeferred() - launch { - try { - connection.callCommand("A.first", null, CommandMode.ONE_SHOT) - } catch (e: Throwable) { - first.complete(e) - } - } + val first = captureError { connection.callCommand("A.first", null, CommandMode.ONE_SHOT) } assertIs(withTimeout(60_000) { first.await() }) - // Characterise the seam: on a half-open socket the transport still *looks* active, so the - // next command reuses the same (dead) socket and times out again. DefaultConnection alone - // cannot recover here -- detection has to come from the transport reporting a disconnect. - // That is what the WebSocket keep-alive ping does in production (KtorWebSocketTransport): a - // missing pong closes the session, turning this case into reconnects_onNextCommand_afterCleanDrop. - val second = CompletableDeferred() - launch { - try { - connection.callCommand("B.second", null, CommandMode.ONE_SHOT) - } catch (e: Throwable) { - second.complete(e) - } - } + // Characterise the seam: on a half-open socket the transport still *looks* active and its + // receive loop never ends, so nothing flags a disconnect — the next command reuses the same + // socket and times out again. DefaultConnection alone cannot recover here; detection would + // require the transport itself to report the disconnect (e.g. a keep-alive ping), which the + // current build does not do, so this case is only bounded by the command timeout. + val second = captureError { connection.callCommand("B.second", null, CommandMode.ONE_SHOT) } assertIs(withTimeout(60_000) { second.await() }) - assertEquals(1, transport.connectCount, "current behaviour: no reconnect after a timeout") + assertEquals(1, transport.connectCount, "no reconnect after a half-open timeout") connection.close() } diff --git a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/TcpProxy.kt b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/TcpProxy.kt index 6245c9a23..c0b29c41e 100644 --- a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/TcpProxy.kt +++ b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/TcpProxy.kt @@ -15,8 +15,8 @@ import java.util.* * stream it can reproduce, on demand, the two failure modes we care about: * - [severAll] closes the sockets (RST/FIN) -> a clean disconnect, * - [freeze]/[unfreeze] keep the sockets open but stop forwarding application bytes -> a half-open - * socket (the peer looks alive at the TCP level but never answers), which only WebSocket keep-alive - * pings can detect. + * socket (the peer looks alive at the TCP level but never answers), which requires an application- + * level liveness check (e.g. keep-alive pings) to detect. * * @param upstreamHost Host of the real Chrome DevTools endpoint. * @param upstreamPort Port of the real Chrome DevTools endpoint. @@ -77,6 +77,11 @@ class TcpProxy( } } catch (_: Exception) { // Socket closed (severed or torn down): just end this pump. + } finally { + // Propagate the stream end to the peer so a one-sided close (e.g. the client closing + // its connection) tears down the paired upstream socket instead of leaking it. + runCatching { from.close() } + runCatching { to.close() } } } }