From 32a04f1d9f46286e231bdcd5a65ca57cc6aa12c1 Mon Sep 17 00:00:00 2001 From: NathanFallet Date: Thu, 23 Jul 2026 13:22:40 +0200 Subject: [PATCH] fix: restore correctness after reconnect --- .../kdriver/core/browser/DefaultBrowser.kt | 5 + .../dev/kdriver/core/connection/Connection.kt | 23 ++++ .../core/connection/DefaultConnection.kt | 85 ++++++++++-- .../core/network/BaseFetchInterception.kt | 20 +-- .../core/network/BaseRequestExpectation.kt | 6 +- .../kotlin/dev/kdriver/core/tab/DefaultTab.kt | 9 ++ .../connection/ReconnectionIntegrationTest.kt | 51 +++++++- .../core/connection/ReconnectionTest.kt | 123 +++++++++++++++++- .../kdriver/opentelemetry/OpenTelemetryTab.kt | 5 + 9 files changed, 299 insertions(+), 28 deletions(-) diff --git a/core/src/commonMain/kotlin/dev/kdriver/core/browser/DefaultBrowser.kt b/core/src/commonMain/kotlin/dev/kdriver/core/browser/DefaultBrowser.kt index a91c01bf7..45179ad3f 100644 --- a/core/src/commonMain/kotlin/dev/kdriver/core/browser/DefaultBrowser.kt +++ b/core/src/commonMain/kotlin/dev/kdriver/core/browser/DefaultBrowser.kt @@ -295,6 +295,11 @@ open class DefaultBrowser( } connection.target.setDiscoverTargets(discover = true) + // The event collectors above survive a reconnect (they read the connection's shared flow), + // but the new CDP session won't emit target events until discovery is re-enabled on it. + connection.registerReconnectRestore("targetDiscovery") { + connection.target.setDiscoverTargets(discover = true) + } } updateTargets() diff --git a/core/src/commonMain/kotlin/dev/kdriver/core/connection/Connection.kt b/core/src/commonMain/kotlin/dev/kdriver/core/connection/Connection.kt index 569fdee95..e0ec8f45e 100644 --- a/core/src/commonMain/kotlin/dev/kdriver/core/connection/Connection.kt +++ b/core/src/commonMain/kotlin/dev/kdriver/core/connection/Connection.kt @@ -25,6 +25,29 @@ interface Connection : BrowserTarget, CDP { @InternalCdpApi suspend fun close() + /** + * Registers a [restore] action to be re-run automatically after the connection transparently + * reconnects. + * + * A reconnect opens a fresh CDP session that has lost all session-scoped state (enabled domains, + * overrides such as the user agent, scripts injected with `addScriptToEvaluateOnNewDocument`, …). + * Features that establish such state register a restorer here — keyed by [key] so it can be + * replaced or removed — and remove it with [unregisterReconnectRestore] when they tear that state + * down. Each restorer runs at most once per reconnect, before the command that triggered the + * restore proceeds; a concurrent command may still race ahead of an in-progress restore. A + * restorer that throws is logged and skipped rather than failing the command. + * + * A restorer must only re-issue CDP commands; it must not itself call + * [registerReconnectRestore]/[unregisterReconnectRestore]. + */ + suspend fun registerReconnectRestore(key: Any, restore: suspend () -> Unit) + + /** + * Removes a restorer previously registered with [registerReconnectRestore]. No-op if [key] is + * not registered. + */ + suspend fun unregisterReconnectRestore(key: Any) + /** * Updates the target information by fetching it from the CDP. * 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 ceb343f7d..389c4c971 100644 --- a/core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt +++ b/core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt @@ -77,6 +77,14 @@ open class DefaultConnection( private var prepareHeadlessDone = false private var prepareExpertDone = false + // Session state to re-apply after a reconnect (a new CDP session loses enabled domains, overrides, + // injected scripts, …). Guarded by [reconnectRestoreMutex]; [needsRestore] flips true on reconnect. + private val reconnectRestoreMutex = Mutex() + private val reconnectRestores = LinkedHashMap Unit>() + + @Volatile + private var needsRestore = false + private val allMessages = MutableSharedFlow(extraBufferCapacity = Channel.UNLIMITED) @InternalCdpApi @@ -101,15 +109,21 @@ open class DefaultConnection( // 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() } } + // Clear `transport` and cancel the old receive loop *before* closing the old transport, + // so the loop ending on that close sees `transport !== old` and no-ops instead of + // spuriously re-flagging `disconnected` (which would cascade into extra reconnects). + val old = transport transport = null + socketSubscription?.cancel() + socketSubscription = null + old?.let { runCatching { it.close() } } 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. + // A reconnect opens a *new* CDP session, so all session state applied on the old one + // is gone. Re-run the built-in preparation, and flag the registered restorers (enabled + // domains, overrides, …) to be replayed before the next command proceeds. prepareHeadlessDone = false prepareExpertDone = false + needsRestore = true } val t = transport ?: createTransport().also { transport = it } if (!t.isActive) { @@ -175,19 +189,54 @@ open class DefaultConnection( pending.forEach { it.completeExceptionally(cause) } } - @InternalCdpApi - override suspend fun callCommand(method: String, parameter: JsonElement?, mode: CommandMode): JsonElement? { - connect() + override suspend fun registerReconnectRestore(key: Any, restore: suspend () -> Unit) { + reconnectRestoreMutex.withLock { reconnectRestores[key] = restore } + } + + override suspend fun unregisterReconnectRestore(key: Any) { + reconnectRestoreMutex.withLock { reconnectRestores.remove(key) } + } + + /** + * Re-applies the registered session state once after a reconnect. Snapshots the restorers and + * clears [needsRestore] under the lock, then runs them *outside* it: a restorer only re-issues + * CDP commands (whose [callCommand] sees [needsRestore] already false and so won't recurse here), + * and running off-lock avoids a deadlock if one of those commands itself triggers a reconnect. + */ + private suspend fun runReconnectRestores() { + val restores = reconnectRestoreMutex.withLock { + if (!needsRestore) return + needsRestore = false + reconnectRestores.values.toList() + } + for (restore in restores) { + try { + restore() + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { + logger.error("Reconnect restore failed: {}", e) + } + } + } - if (mode == CommandMode.DEFAULT) owner?.let { browser -> - // Serialize preparation so concurrent first commands run it once, not N times (ISSUE-4). - // prepare* issue ONE_SHOT commands, which skip this block, so prepareMutex isn't re-entered. + /** + * Runs the built-in browser-side preparation for a DEFAULT command, on the current session. + * prepare* issue ONE_SHOT commands, which skip this step, so [prepareMutex] isn't re-entered. + * Serialized so concurrent first commands run it once, not N times (ISSUE-4). + */ + private suspend fun ensurePrepared(mode: CommandMode) { + if (mode != CommandMode.DEFAULT) return + owner?.let { browser -> prepareMutex.withLock { if (browser.config.expert) prepareExpert() if (browser.config.headless) prepareHeadless() } } + } + @InternalCdpApi + override suspend fun callCommand(method: String, parameter: JsonElement?, mode: CommandMode): JsonElement? { val requestId = currentIdMutex.withLock { currentId++ } val jsonString = Serialization.json.encodeToString(Request(requestId, method, parameter)) val timeout = owner?.config?.commandTimeout ?: Defaults.COMMAND_TIMEOUT @@ -199,6 +248,20 @@ open class DefaultConnection( // the command may already have run (that case surfaces as a ConnectionClosedException). var sendError: Throwable? = null repeat(SEND_ATTEMPTS) { + connect() + // Restore + prepare the (possibly freshly reconnected) session *before* sending on it, so + // the command always runs against a fully re-established session — this covers a reconnect + // that happens here in the retry loop, not just one detected before it. + // + // Restore MUST run before (and outside of) ensurePrepared: prepare* issue their own inner + // commands, which would otherwise trigger the restore while prepareMutex is held, and a + // restorer's DEFAULT command re-entering the non-reentrant prepareMutex would deadlock. + // Running restore first is also correct precedence-wise — a restorer's own DEFAULT command + // re-runs prepare internally, so e.g. a user user-agent override still lands last. + if (needsRestore) runReconnectRestores() + ensurePrepared(mode) + // Fetch the transport only now: prepare/restore issue their own commands, which may have + // reconnected, so a transport captured before them could be the stale, disposed one. 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 diff --git a/core/src/commonMain/kotlin/dev/kdriver/core/network/BaseFetchInterception.kt b/core/src/commonMain/kotlin/dev/kdriver/core/network/BaseFetchInterception.kt index 0ce0bfee7..84f86432a 100644 --- a/core/src/commonMain/kotlin/dev/kdriver/core/network/BaseFetchInterception.kt +++ b/core/src/commonMain/kotlin/dev/kdriver/core/network/BaseFetchInterception.kt @@ -19,6 +19,12 @@ open class BaseFetchInterception( private var responseDeferred = CompletableDeferred() private var job: Job? = null + private val requestPattern = Fetch.RequestPattern( + urlPattern = urlPattern, + resourceType = resourceType, + requestStage = requestStage, + ) + private val handler: suspend (Fetch.RequestPausedParameter) -> Unit = handler@{ event -> responseDeferred.complete(event) job?.cancel() @@ -32,18 +38,14 @@ open class BaseFetchInterception( job = coroutineScope.launch(start = CoroutineStart.UNDISPATCHED) { tab.fetch.requestPaused.collect { handler(it) } } - tab.fetch.enable( - listOf( - Fetch.RequestPattern( - urlPattern = urlPattern, - resourceType = resourceType, - requestStage = requestStage - ) - ) - ) + tab.fetch.enable(listOf(requestPattern)) + // The requestPaused collector survives a reconnect (it reads the connection's shared flow), + // but Fetch must be re-enabled on the new CDP session for the browser to keep pausing requests. + tab.registerReconnectRestore(this) { tab.fetch.enable(listOf(requestPattern)) } } private suspend fun teardown() { + tab.unregisterReconnectRestore(this) job?.cancel() tab.fetch.disable() } diff --git a/core/src/commonMain/kotlin/dev/kdriver/core/network/BaseRequestExpectation.kt b/core/src/commonMain/kotlin/dev/kdriver/core/network/BaseRequestExpectation.kt index de244a18f..b6d8e4e34 100644 --- a/core/src/commonMain/kotlin/dev/kdriver/core/network/BaseRequestExpectation.kt +++ b/core/src/commonMain/kotlin/dev/kdriver/core/network/BaseRequestExpectation.kt @@ -63,9 +63,13 @@ open class BaseRequestExpectation( tab.network.loadingFinished.collect { loadingFinishedHandler(it) } } tab.network.enable() + // The event collectors survive a reconnect; re-enable Network on the new session so the + // browser keeps emitting the request/response events this expectation waits for. + tab.registerReconnectRestore(this) { tab.network.enable() } } - private fun teardown() { + private suspend fun teardown() { + tab.unregisterReconnectRestore(this) requestJob?.cancel() responseJob?.cancel() loadingFinishedJob?.cancel() diff --git a/core/src/commonMain/kotlin/dev/kdriver/core/tab/DefaultTab.kt b/core/src/commonMain/kotlin/dev/kdriver/core/tab/DefaultTab.kt index 0229ff779..c030c95c0 100644 --- a/core/src/commonMain/kotlin/dev/kdriver/core/tab/DefaultTab.kt +++ b/core/src/commonMain/kotlin/dev/kdriver/core/tab/DefaultTab.kt @@ -118,6 +118,15 @@ open class DefaultTab( acceptLanguage = acceptLanguage, platform = platform ) + // A reconnect opens a new CDP session that would drop this override, so re-apply it then. + // Keyed so a later setUserAgent replaces (rather than stacks) the restorer. + registerReconnectRestore("userAgent") { + network.setUserAgentOverride( + userAgent = ua, + acceptLanguage = acceptLanguage, + platform = platform + ) + } } override suspend fun getWindow(): dev.kdriver.cdp.domain.Browser.GetWindowForTargetReturn { diff --git a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionIntegrationTest.kt b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionIntegrationTest.kt index 3c87e5472..2fae6af94 100644 --- a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionIntegrationTest.kt +++ b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionIntegrationTest.kt @@ -7,6 +7,7 @@ import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.runBlocking import kotlinx.coroutines.withTimeout import kotlinx.serialization.json.* +import java.util.concurrent.atomic.AtomicInteger import kotlin.test.* /** @@ -36,6 +37,21 @@ class ReconnectionIntegrationTest { return result!!.jsonObject["result"]!!.jsonObject["value"]!!.jsonPrimitive.int } + /** + * Evaluates, retrying through a transient reconnect. The single command *in flight* when the + * socket is cut may itself be the casualty (its send lands on the dying socket, then the reply + * never comes) — by design we never resend an awaited command. The connection self-heals, so the + * next attempt reconnects and succeeds, exactly as kdriver's high-level retry loops rely on. + */ + private suspend fun DefaultConnection.evaluateRecovering(expression: String, timeoutMs: Long = 15_000): Int = + withTimeout(timeoutMs) { + while (true) { + val result = runCatching { evaluate(expression) } + if (result.isSuccess) return@withTimeout result.getOrThrow() + } + error("unreachable") + } + @Test fun reconnects_afterHardCut_throughProxy() = runBlocking { val host = createBrowser(this, headless = true, sandbox = false) @@ -51,9 +67,38 @@ class ReconnectionIntegrationTest { // 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") }) + // The connection reconnects through the proxy on its own and commands succeed again. + assertEquals(4, conn.evaluateRecovering("2 + 2")) + } finally { + conn.close() + proxy.stop() + host.stop() + } + } + + @Test + fun runsRegisteredRestore_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) + val restoreCount = AtomicInteger(0) + + try { + assertEquals(2, conn.evaluate("1 + 1")) + + // A feature (e.g. FetchInterception) registers a restorer to re-establish its session + // state after a reconnect. Here it just counts its runs. + conn.registerReconnectRestore("probe") { restoreCount.incrementAndGet() } + + proxy.severAll() + + // The connection reconnects once; the restorer must run exactly once as part of that + // single reconnect, no matter how many commands drive the recovery. + assertEquals(4, conn.evaluateRecovering("2 + 2")) + assertEquals(9, conn.evaluateRecovering("3 + 3 + 3")) + assertEquals(1, restoreCount.get(), "the registered restorer must run once per reconnect") } finally { conn.close() proxy.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 index 5ed9c7b97..3a4db7036 100644 --- a/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionTest.kt +++ b/core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionTest.kt @@ -17,10 +17,7 @@ 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 +import kotlin.test.* /** * Behaviour of [DefaultConnection] when the underlying transport disconnects and a later command @@ -279,6 +276,124 @@ class ReconnectionTest { connection.close() } + @Test + fun runsRegisteredRestore_oncePerReconnect() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport() + val connection = TestConnection(this, transport) + var restoreCount = 0 + connection.registerReconnectRestore("k") { restoreCount++ } + + // A command that just connects (no reconnect) must not run the restore. + val c1 = async { connection.callCommand("A", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + c1.await() + assertEquals(0, restoreCount, "restore must not run without a reconnect") + + // A command after a drop reconnects and runs the restore exactly once. + transport.drop() + val c2 = async { connection.callCommand("B", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + c2.await() + assertEquals(1, restoreCount) + + // A further command while still connected does not run it again. + val c3 = async { connection.callCommand("C", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + c3.await() + assertEquals(1, restoreCount, "restore runs once per reconnect, not once per command") + + // A second drop reconnects again and re-runs the restore. + transport.drop() + val c4 = async { connection.callCommand("D", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + c4.await() + assertEquals(2, restoreCount) + + connection.close() + } + + @Test + fun restoreThatIssuesCommands_runsAfterReconnect() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport().apply { autoRespond = true } + val connection = TestConnection(this, transport) + connection.registerReconnectRestore("k") { + connection.callCommand("Restore.enable", null, CommandMode.ONE_SHOT) + } + + connection.callCommand("A", null, CommandMode.ONE_SHOT) + assertEquals(0, transport.sent.count { it.contains("Restore.enable") }) + + transport.drop() + connection.callCommand("B", null, CommandMode.ONE_SHOT) + assertEquals( + 1, + transport.sent.count { it.contains("Restore.enable") }, + "a restorer that re-issues a CDP command must run it once on the reconnected session", + ) + + connection.close() + } + + @Test + fun doesNotRun_unregisteredRestore_afterReconnect() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport() + val connection = TestConnection(this, transport) + var restoreCount = 0 + connection.registerReconnectRestore("k") { restoreCount++ } + + val c1 = async { connection.callCommand("A", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + c1.await() + + connection.unregisterReconnectRestore("k") + + transport.drop() + val c2 = async { connection.callCommand("B", null, CommandMode.ONE_SHOT) } + transport.respondToLast() + c2.await() + assertEquals(0, restoreCount, "an unregistered restore must not run") + + connection.close() + } + + @Test + fun replaysPreparationAndRestore_afterReconnect_withoutDeadlock() = runTest(UnconfinedTestDispatcher()) { + val transport = ControllableTransport().apply { autoRespond = true } + // Headless browser => prepareHeadless runs (a Runtime.evaluate under prepareMutex). + 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 } + + var restoreCount = 0 + // A restorer that issues a DEFAULT command: this is the combination that used to deadlock, as + // prepare's own (ONE_SHOT) command would trigger the restore while prepareMutex was held, and + // the restorer's DEFAULT command then re-entered the non-reentrant prepareMutex. + connection.registerReconnectRestore("k") { + connection.callCommand("Restore.enable", null, CommandMode.DEFAULT) + restoreCount++ + } + + connection.callCommand("Some.method", null, CommandMode.DEFAULT) + assertEquals(1, transport.sent.count { it.contains("Runtime.evaluate") }) + assertEquals(0, restoreCount) + + transport.drop() + // Must complete (no deadlock): re-prepares the new session AND runs the restorer. + withTimeout(5_000) { connection.callCommand("Some.method", null, CommandMode.DEFAULT) } + assertEquals( + 2, + transport.sent.count { it.contains("Runtime.evaluate") }, + "preparation replayed after reconnect" + ) + assertEquals(1, restoreCount, "restorer ran after reconnect") + assertTrue(transport.sent.any { it.contains("Restore.enable") }, "the restorer's DEFAULT command was sent") + + connection.close() + } + @Test fun reusesConnection_afterHalfOpenTimeout() = runTest(UnconfinedTestDispatcher()) { val transport = ControllableTransport() diff --git a/opentelemetry/src/commonMain/kotlin/dev/kdriver/opentelemetry/OpenTelemetryTab.kt b/opentelemetry/src/commonMain/kotlin/dev/kdriver/opentelemetry/OpenTelemetryTab.kt index 26ee93402..cfa8937a4 100644 --- a/opentelemetry/src/commonMain/kotlin/dev/kdriver/opentelemetry/OpenTelemetryTab.kt +++ b/opentelemetry/src/commonMain/kotlin/dev/kdriver/opentelemetry/OpenTelemetryTab.kt @@ -407,6 +407,11 @@ class OpenTelemetryTab( Span.current().addEvent("kdriver.tab.close") } + override suspend fun registerReconnectRestore(key: Any, restore: suspend () -> Unit) = + tab.registerReconnectRestore(key, restore) + + override suspend fun unregisterReconnectRestore(key: Any) = tab.unregisterReconnectRestore(key) + override suspend fun updateTarget() = tab.updateTarget() override suspend fun wait(t: Long?) = tab.wait(t)