From 1375da90715cfda99d4ffd276ef984f9599cdbb5 Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Mon, 14 Sep 2026 00:40:09 +0700 Subject: [PATCH 1/3] fix(connections): keep table loads working while a database switch reconnects --- CHANGELOG.md | 2 + .../Coordinators/PaginationCoordinator.swift | 9 +- .../Database/DatabaseManager+Health.swift | 18 ++ .../DatabaseManager+ScopedDriver.swift | 130 +++++++--- .../Database/DatabaseManager+Tunnel.swift | 4 + .../Query/DatabaseTreeMetadataService.swift | 6 +- .../Query/MetadataConnectionPool.swift | 240 ++++++++++++++---- .../MainContentCoordinator+QueryHelpers.swift | 24 ++ .../Views/Main/MainContentCoordinator.swift | 6 +- .../DatabaseSwitchLeaseOrderingTests.swift | 221 +++++++++++++++- .../SwitchDatabasePooledConnectionTests.swift | 148 +++++++++++ .../Query/MetadataConnectionPoolTests.swift | 168 ++++++++++++ 12 files changed, 878 insertions(+), 98 deletions(-) create mode 100644 TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift diff --git a/CHANGELOG.md b/CHANGELOG.md index 53f7b93a0..238fd6880 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Wrong approximate row count for a PostgreSQL or PGlite table outside the current schema. - Queued queries running on a reopened connection, or stalling it, after a disconnect. - Redis keys read from the wrong database when switching databases while a load was running. +- "This tab is on" error on a table tab for the previous database when it reloads during a database switch. +- Table on another database opening empty with no error while the connection switches database. ## [0.74.0] - 2026-09-13 diff --git a/TablePro/Core/Coordinators/PaginationCoordinator.swift b/TablePro/Core/Coordinators/PaginationCoordinator.swift index 6ed920783..d17c6b363 100644 --- a/TablePro/Core/Coordinators/PaginationCoordinator.swift +++ b/TablePro/Core/Coordinators/PaginationCoordinator.swift @@ -303,8 +303,7 @@ final class PaginationCoordinator { /// would discard its own rows. It registers as unclaimed work instead, which is what keeps /// the titlebar reporting it, and releases that on every exit including cancellation. let workToken = parent.tabExecution.beginUnclaimedWork(for: tabId) - - let route = DatabaseManager.shared.executionRoute(for: scope) + let isTableTab = parent.tabManager.tabs[idx].tabType == .table let startedAt = ContinuousClock.Instant.now let fetchAllTask = Task { [weak self, parent] in @@ -314,11 +313,7 @@ final class PaginationCoordinator { do { let start = CFAbsoluteTimeGetCurrent() progressLog.info("[fetchAll] executing full query: \(baseQuery.prefix(100), privacy: .public)") - let result = try await DatabaseManager.shared.withScopedDriver( - scope: scope, - route: route, - cancellation: .cancellableRead - ) { driver in + let result = try await parent.withExecutionDriver(scope: scope, isTableTab: isTableTab) { driver in try await driver.executeUserQuery( query: baseQuery, rowCap: nil, diff --git a/TablePro/Core/Database/DatabaseManager+Health.swift b/TablePro/Core/Database/DatabaseManager+Health.swift index b52e07fff..044755893 100644 --- a/TablePro/Core/Database/DatabaseManager+Health.swift +++ b/TablePro/Core/Database/DatabaseManager+Health.swift @@ -138,6 +138,11 @@ extension DatabaseManager { let attemptedDriver = session.driver await SchemaService.shared.prepareForReload(connectionId: connectionId) await DatabaseTreeMetadataService.shared.handleReconnect(connectionId: connectionId) + /// A connection that stopped answering has most likely taken its pooled connections with it, + /// and a rebuilt tunnel moves every one of them to a new port, so pooled work waits for the + /// replacement rather than dialing what is being torn down. + MetadataConnectionPool.shared.beginTransportReplacement(connectionId: connectionId) + defer { MetadataConnectionPool.shared.endTransportReplacement(connectionId: connectionId) } do { guard let result = try await trackOperation(sessionId: connectionId, operation: { @@ -355,6 +360,19 @@ extension DatabaseManager { await SchemaService.shared.prepareForReload(connectionId: sessionId) await DatabaseTreeMetadataService.shared.handleReconnect(connectionId: sessionId) + /// A pooled connection stands on the effective connection and its own database, so the pool is + /// only held back when this rebuilds a tunnel on a new port or recovers a session that had + /// stopped answering. A database switch over a live, direct connection changes neither, and + /// closing the pool there withdrew an open that another database's table load was waiting on. + let replacesPooledTransport = session.connection.activeTunnelKind != nil || session.liveness != .live + if replacesPooledTransport { + MetadataConnectionPool.shared.beginTransportReplacement(connectionId: sessionId) + } + defer { + if replacesPooledTransport { + MetadataConnectionPool.shared.endTransportReplacement(connectionId: sessionId) + } + } await stopHealthMonitor(for: sessionId) diff --git a/TablePro/Core/Database/DatabaseManager+ScopedDriver.swift b/TablePro/Core/Database/DatabaseManager+ScopedDriver.swift index 005599a70..62a76b679 100644 --- a/TablePro/Core/Database/DatabaseManager+ScopedDriver.swift +++ b/TablePro/Core/Database/DatabaseManager+ScopedDriver.swift @@ -15,6 +15,13 @@ enum ScopedDriverRoute: Equatable { case unavailable(String) } +/// What a table read's turn on the session driver came to. +private enum TableReadTurn: Sendable { + case ran(T) + /// The route decided once the turn came, which is never the session driver. + case moved(ScopedDriverRoute) +} + extension DatabaseManager { /// A metadata read needs no transaction, no temp tables and no cancellation handle, /// so it takes a pooled connection and leaves the shared driver where it is. @@ -76,29 +83,7 @@ extension DatabaseManager { cancellation: DriverCancellationPolicy, _ body: @Sendable @escaping (DatabaseDriver) async throws -> T ) async throws -> T { - let leased: @Sendable (DatabaseDriver) async throws -> T - if cancellation.isTracked { - let connectionId = scope.connectionId - let token = UUID() - let entry = RunningDriver(driver: nil, policy: cancellation) - leased = { driver in - await MainActor.run { - DatabaseManager.shared.runningDrivers[connectionId, default: [:]][token] = - entry.adopting(driver) - } - do { - let value = try await body(driver) - await MainActor.run { DatabaseManager.shared.releaseRunningDriver(token, for: connectionId) } - return value - } catch { - await MainActor.run { DatabaseManager.shared.releaseRunningDriver(token, for: connectionId) } - throw error - } - } - } else { - leased = body - } - + let leased = trackedLease(for: scope.connectionId, cancellation: cancellation, body) switch route { case .unavailable(let message): throw DatabaseError.queryFailed(message) @@ -111,6 +96,67 @@ extension DatabaseManager { } } + /// A table tab's read is a SELECT the app built from the tab's own table, so it depends on + /// nothing the session holds: no transaction, no temp table, no variable. That makes it the one + /// kind of work that can follow a route change it waited through. A database switch on an engine + /// that reconnects to perform one moves the browsed database while holding the gate, so a read + /// queued behind it for the database being left is re-routed once its turn comes, instead of + /// being refused by `pin`. User SQL never comes here: the session it was written against is gone, + /// and a refusal is the right answer for it. + /// + /// The gate is left before the read is dispatched again, so a pooled read never holds it. The + /// loop ends: a turn only ever moves the read off the session driver, and neither the pool nor an + /// unavailable route queues on the gate again. + func withTableReadDriver( + scope: DatabaseScope, + cancellation: DriverCancellationPolicy, + _ body: @Sendable @escaping (DatabaseDriver) async throws -> T + ) async throws -> T { + let leased = trackedLease(for: scope.connectionId, cancellation: cancellation, body) + var route = executionRoute(for: scope) + while true { + switch route { + case .unavailable(let message): + throw DatabaseError.queryFailed(message) + case .pooled: + return try await MetadataConnectionPool.shared.withDriver(scope: scope, leased) + case .sessionDriver: + switch try await withTableReadSessionDriver(scope: scope, leased) { + case .ran(let value): + return value + case .moved(let decided): + route = decided + } + } + } + } + + /// Registers the driver a tracked lease runs on for the length of its body, whichever route it + /// took, so Stop reaches the handle the work is actually on. + private func trackedLease( + for connectionId: UUID, + cancellation: DriverCancellationPolicy, + _ body: @Sendable @escaping (DatabaseDriver) async throws -> T + ) -> @Sendable (DatabaseDriver) async throws -> T { + guard cancellation.isTracked else { return body } + let token = UUID() + let entry = RunningDriver(driver: nil, policy: cancellation) + return { driver in + await MainActor.run { + DatabaseManager.shared.runningDrivers[connectionId, default: [:]][token] = + entry.adopting(driver) + } + do { + let value = try await body(driver) + await MainActor.run { DatabaseManager.shared.releaseRunningDriver(token, for: connectionId) } + return value + } catch { + await MainActor.run { DatabaseManager.shared.releaseRunningDriver(token, for: connectionId) } + throw error + } + } + } + internal func releaseRunningDriver(_ token: UUID, for connectionId: UUID) { runningDrivers[connectionId]?.removeValue(forKey: token) if runningDrivers[connectionId]?.isEmpty == true { @@ -183,26 +229,50 @@ extension DatabaseManager { scope: DatabaseScope, _ body: @Sendable @escaping (DatabaseDriver) async throws -> T ) async throws -> T { + try await withSessionDriverTurn(connectionId: scope.connectionId) { driver in + try await pin(driver, to: scope) + return try await body(driver) + } + } + + /// The route is asked again through `executionRoute` itself, so it reads the same inputs the + /// caller's first answer did, the browsed database among them. The driver's own connection is no + /// substitute: it carries the database name the connection resolved to, not the one browsed. + private func withTableReadSessionDriver( + scope: DatabaseScope, + _ body: @Sendable @escaping (DatabaseDriver) async throws -> T + ) async throws -> TableReadTurn { + try await withSessionDriverTurn(connectionId: scope.connectionId) { driver in + let route = executionRoute(for: scope) + guard route == .sessionDriver else { return .moved(route) } + try await pin(driver, to: scope) + return .ran(try await body(driver)) + } + } + + private func withSessionDriverTurn( + connectionId: UUID, + _ turn: (DatabaseDriver) async throws -> R + ) async throws -> R { /// Outside the gate on purpose. A verification that has to reconnect runs the whole /// reconnect, which restores the schema and the database on the new driver, and doing that /// while holding the gate would deadlock the very thing waiting to be pinned. - await verifyBeforeUse(scope.connectionId) + await verifyBeforeUse(connectionId) /// A check that failed and could not recover left the driver installed and disconnected, /// so the presence of a driver below is not enough. Refusing here is the point of checking /// at all: without it the user's own work runs on a handle the app already knows is dead. - guard isUsable(scope.connectionId) else { + guard isUsable(connectionId) else { throw DatabaseError.notConnected } - return try await sessionDriverGate.withExclusiveAccess(scope.connectionId) { - try await trackOperation(sessionId: scope.connectionId) { + return try await sessionDriverGate.withExclusiveAccess(connectionId) { + try await trackOperation(sessionId: connectionId) { try Task.checkCancellation() /// Asked again once the lease has its turn, because a database switch that held the /// gate can have left the driver the same way while this lease waited. - guard isUsable(scope.connectionId), let driver = driver(for: scope.connectionId) else { + guard isUsable(connectionId), let driver = driver(for: connectionId) else { throw DatabaseError.notConnected } - try await pin(driver, to: scope) - return try await body(driver) + return try await turn(driver) } } } diff --git a/TablePro/Core/Database/DatabaseManager+Tunnel.swift b/TablePro/Core/Database/DatabaseManager+Tunnel.swift index 9a2b54e91..88d72c0f6 100644 --- a/TablePro/Core/Database/DatabaseManager+Tunnel.swift +++ b/TablePro/Core/Database/DatabaseManager+Tunnel.swift @@ -140,6 +140,10 @@ extension DatabaseManager { recoveringConnectionIds.insert(connectionId) defer { recoveringConnectionIds.remove(connectionId) } + /// The tunnel every pooled connection dialed is gone, and recovering binds a new one on another + /// port, so pooled work waits for it instead of dialing a port another tunnel may since own. + MetadataConnectionPool.shared.beginTransportReplacement(connectionId: connectionId) + defer { MetadataConnectionPool.shared.endTransportReplacement(connectionId: connectionId) } Self.logger.warning("\(kind, privacy: .public) tunnel died for connection: \(session.connection.name)") diff --git a/TablePro/Core/Services/Query/DatabaseTreeMetadataService.swift b/TablePro/Core/Services/Query/DatabaseTreeMetadataService.swift index 5a305cd36..658c0b16e 100644 --- a/TablePro/Core/Services/Query/DatabaseTreeMetadataService.swift +++ b/TablePro/Core/Services/Query/DatabaseTreeMetadataService.swift @@ -623,9 +623,13 @@ final class DatabaseTreeMetadataService: CatalogChangeTarget { /// A fetch still running on the driver the reconnect replaced answers for the old session, so /// every key of the connection is superseded before anything suspends and none of those fetches /// may commit. What is already on screen stays until the new session's own loads replace it. + /// + /// The pooled connections are not this service's to close. They stand on the transport rather + /// than on the session driver, and `DatabaseManager`, which rebuilds the transport, holds them + /// back while it does. Closing them on every reconnect withdrew an open a table tab on another + /// database was waiting on, and that tab then showed nothing and no error. func handleReconnect(connectionId: UUID) async { supersedeEveryKey(of: connectionId) - MetadataConnectionPool.shared.closeAll(connectionId: connectionId) SchemaForeignKeyStore.shared.invalidate(connectionId: connectionId) await resetPending(connectionId: connectionId) } diff --git a/TablePro/Core/Services/Query/MetadataConnectionPool.swift b/TablePro/Core/Services/Query/MetadataConnectionPool.swift index 6b3435ab2..d0693af37 100644 --- a/TablePro/Core/Services/Query/MetadataConnectionPool.swift +++ b/TablePro/Core/Services/Query/MetadataConnectionPool.swift @@ -56,11 +56,41 @@ final class MetadataConnectionPool { } } + typealias DriverOpener = @MainActor (DatabaseScope) async throws -> DatabaseDriver + + /// Why an open still in progress was taken away from the callers waiting on it. + private enum Withdrawal { + /// The transport it was dialing is being rebuilt, so its callers take a connection again + /// once the replacement is in place. + case transportReplaced + /// What it was opened for is going away, so nothing may open it again. + case closed + } + + /// One open in progress, shared by every caller that asks for its key while it runs. + @MainActor + private final class PendingOpen { + let task: Task + var withdrawal: Withdrawal? + + init(task: Task) { + self.task = task + } + } + + private struct TransportWaiter { + let ticket: UUID + let continuation: CheckedContinuation + } + private var entries: [Key: Entry] = [:] - private var pending: [Key: Task] = [:] + private var pending: [Key: PendingOpen] = [:] + private var transportReplacements: [UUID: Int] = [:] + private var transportWaiters: [UUID: [TransportWaiter]] = [:] + private let openDriver: DriverOpener private let maxPerConnection = 6 - private let operationTimeoutSeconds: Double = 15 - private let preparationTimeoutSeconds: Double = 60 + private static let operationTimeoutSeconds: Double = 15 + private static let preparationTimeoutSeconds: Double = 60 private var sweeper: Task? /// How long a pooled connection may sit unused before it is handed back. @@ -80,7 +110,11 @@ final class MetadataConnectionPool { /// day waking up. private static let sweepInterval: Duration = .seconds(60) - private init() {} + private init(openDriver: DriverOpener? = nil) { + self.openDriver = openDriver ?? { scope in + try await MetadataConnectionPool.openSessionDriver(for: scope) + } + } func withDriver( scope: DatabaseScope, @@ -98,27 +132,40 @@ final class MetadataConnectionPool { /// needs: PostgreSQL refuses `ALTER DATABASE ... RENAME` while any backend is connected to it, /// and an expanded row or a tab that ran a query there leaves one here. func closeAll(connectionId: UUID, database: String) { - for key in pending.keys - where key.scope.connectionId == connectionId && key.scope.database == database { - pending[key]?.cancel() - pending.removeValue(forKey: key) - } - for key in entries.keys - where key.scope.connectionId == connectionId && key.scope.database == database { - closeOrDeferEntry(forKey: key) + closeEntries(withdrawingOpensAs: .closed) { scope in + scope.connectionId == connectionId && scope.database == database } - stopSweeperIfEmpty() } func closeAll(connectionId: UUID) { - for key in pending.keys where key.scope.connectionId == connectionId { - pending[key]?.cancel() - pending.removeValue(forKey: key) + closeEntries(withdrawingOpensAs: .closed) { $0.connectionId == connectionId } + } + + /// Holds a connection's pooled work while its transport is rebuilt. + /// + /// A pooled entry stands on the session's effective connection plus its own database, so it only + /// has to go when that endpoint does: a tunnel rebuilt on a new local port, or a reconnect + /// recovering a connection that stopped answering. An open already dialing is withdrawn, and the + /// callers waiting on it take a connection again after `endTransportReplacement` rather than + /// failing a load nobody cancelled. A new caller waits for the same moment instead of dialing an + /// endpoint about to close. Replacements nest, and pooled work resumes when the last one ends. + func beginTransportReplacement(connectionId: UUID) { + transportReplacements[connectionId, default: 0] += 1 + closeEntries(withdrawingOpensAs: .transportReplaced) { $0.connectionId == connectionId } + } + + /// Has to follow every `beginTransportReplacement` on every exit, cancellation included, or the + /// connection's pooled work waits for good. + func endTransportReplacement(connectionId: UUID) { + guard let depth = transportReplacements[connectionId] else { return } + guard depth == 1 else { + transportReplacements[connectionId] = depth - 1 + return } - for key in entries.keys where key.scope.connectionId == connectionId { - closeOrDeferEntry(forKey: key) + transportReplacements.removeValue(forKey: connectionId) + for waiter in transportWaiters.removeValue(forKey: connectionId) ?? [] { + waiter.continuation.resume() } - stopSweeperIfEmpty() } #if DEBUG @@ -154,9 +201,20 @@ final class MetadataConnectionPool { } /// A pool of its own, so a test that moves the clock or empties the pool cannot close the - /// entries another test injected into the shared one. - internal static func isolatedForTesting() -> MetadataConnectionPool { - MetadataConnectionPool() + /// entries another test injected into the shared one. `openDriver` stands in for opening a real + /// connection, which a test with no plugin cannot do. + internal static func isolatedForTesting(openDriver: DriverOpener? = nil) -> MetadataConnectionPool { + MetadataConnectionPool(openDriver: openDriver) + } + + /// How many callers are waiting for a transport replacement to end, so a test can wait for one + /// to park instead of guessing how many scheduler turns that takes. + internal func transportWaiterCount(for connectionId: UUID) -> Int { + transportWaiters[connectionId]?.count ?? 0 + } + + internal func isReplacingTransport(for connectionId: UUID) -> Bool { + transportReplacements[connectionId] != nil } #endif @@ -176,35 +234,65 @@ final class MetadataConnectionPool { } } + private func closeEntries( + withdrawingOpensAs withdrawal: Withdrawal, + where matches: (DatabaseScope) -> Bool + ) { + for key in pending.keys where matches(key.scope) { + guard let open = pending.removeValue(forKey: key) else { continue } + open.withdrawal = withdrawal + open.task.cancel() + } + for key in entries.keys where matches(key.scope) { + closeOrDeferEntry(forKey: key) + } + stopSweeperIfEmpty() + } + private func acquireEntry(scope: DatabaseScope, workload: Workload) async throws -> Entry { - let connectionId = scope.connectionId let key = Key(scope: scope, workload: workload) - /// A cached entry is only worth reusing while it is both connected and recent. The - /// staleness half matters even with the sweeper running, because a Mac that slept comes - /// back with entries the sweeper never got to and sockets the server has long since - /// closed. The old entry is closed rather than left behind: overwriting `entries[key]` - /// with a fresh one used to leak the driver it replaced. - if let entry = entries[key] { - if entry.driver.status == .connected, !entry.driver.hasLostConnection, !Self.isStale(entry.lastUsed) { + while true { + try await waitForTransport(of: scope.connectionId) + if let entry = reusableEntry(forKey: key) { return entry } - closeOrDeferEntry(forKey: key) + let open = try pending[key] ?? startOpen(forKey: key) + let failure = await completion(of: open, forKey: key) + /// A withdrawn open can end either way: a driver whose connect honours cancellation + /// throws, and one that finished first returns without keeping its entry. + if open.withdrawal == .transportReplaced { + try Task.checkCancellation() + continue + } + if let failure { + throw failure + } + guard let entry = entries[key] else { throw DatabaseError.notConnected } + return entry } + } - if let inFlight = pending[key] { - try await inFlight.value - guard let entry = entries[key] else { throw DatabaseError.notConnected } + /// A cached entry is only worth reusing while it is both connected and recent. The + /// staleness half matters even with the sweeper running, because a Mac that slept comes + /// back with entries the sweeper never got to and sockets the server has long since + /// closed. The old entry is closed rather than left behind: overwriting `entries[key]` + /// with a fresh one used to leak the driver it replaced. + private func reusableEntry(forKey key: Key) -> Entry? { + guard let entry = entries[key] else { return nil } + if entry.driver.status == .connected, !entry.driver.hasLostConnection, !Self.isStale(entry.lastUsed) { return entry } + closeOrDeferEntry(forKey: key) + return nil + } - guard DatabaseManager.shared.session(for: connectionId) != nil else { + private func startOpen(forKey key: Key) throws -> PendingOpen { + guard DatabaseManager.shared.session(for: key.scope.connectionId) != nil else { throw DatabaseError.notConnected } - - evictIdleIfNeeded(for: connectionId) - + evictIdleIfNeeded(for: key.scope.connectionId) let task = Task { [self] in - let entry = try await openEntry(key: key) + let entry = Entry(driver: try await openDriver(key.scope)) if Task.isCancelled { entry.driver.disconnect() return @@ -212,22 +300,76 @@ final class MetadataConnectionPool { entries[key] = entry startSweeperIfNeeded() } - pending[key] = task - defer { if pending[key] == task { pending.removeValue(forKey: key) } } - try await task.value + let open = PendingOpen(task: task) + pending[key] = open + return open + } - guard let entry = entries[key] else { throw DatabaseError.notConnected } - return entry + /// Waits for an open to finish, takes it off the pending list if nothing replaced it there, and + /// returns how it failed. + private func completion(of open: PendingOpen, forKey key: Key) async -> Error? { + defer { + if pending[key] === open { + pending.removeValue(forKey: key) + } + } + do { + try await open.task.value + return nil + } catch { + return error + } + } + + /// Parks a caller for as long as the connection's transport is being replaced. A caller that is + /// cancelled while parked stops waiting at once rather than when the replacement ends. + private func waitForTransport(of connectionId: UUID) async throws { + while transportReplacements[connectionId] != nil { + let ticket = UUID() + try await withTaskCancellationHandler( + operation: { try await parkForTransport(ticket: ticket, connectionId: connectionId) }, + onCancel: { [weak self] in + Task { @MainActor in + self?.failTransportWaiter(ticket: ticket, connectionId: connectionId) + } + } + ) + } + } + + private func parkForTransport(ticket: UUID, connectionId: UUID) async throws { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + guard !Task.isCancelled else { + continuation.resume(throwing: CancellationError()) + return + } + transportWaiters[connectionId, default: []].append( + TransportWaiter(ticket: ticket, continuation: continuation) + ) + } + } + + /// Removes the ticket before resuming it, so a cancellation racing the end of a replacement can + /// only ever find one of them. + private func failTransportWaiter(ticket: UUID, connectionId: UUID) { + guard var waiters = transportWaiters[connectionId], + let index = waiters.firstIndex(where: { $0.ticket == ticket }) + else { + return + } + let waiter = waiters.remove(at: index) + transportWaiters[connectionId] = waiters.isEmpty ? nil : waiters + waiter.continuation.resume(throwing: CancellationError()) } - private func openEntry(key: Key) async throws -> Entry { - guard let session = DatabaseManager.shared.session(for: key.scope.connectionId) else { + private static func openSessionDriver(for scope: DatabaseScope) async throws -> DatabaseDriver { + guard let session = DatabaseManager.shared.session(for: scope.connectionId) else { throw DatabaseError.notConnected } var connection = session.effectiveConnection ?? session.connection let plan = Self.planConnection( configuredDatabase: connection.database, - targetDatabase: key.scope.database, + targetDatabase: scope.database, authenticationIsDatabaseScoped: connection.type.authenticationIsDatabaseScoped, runsStartupCommands: DatabaseManager.hasStartupCommands(session.connection.startupCommands), switchesDatabaseWithoutReconnecting: PluginManager.shared @@ -252,14 +394,14 @@ final class MetadataConnectionPool { if let database = plan.switchDatabase { try await Self.switchDatabase(driver, to: database, timeoutSeconds: operationTimeoutSeconds) } - if let schema = key.scope.schema { + if let schema = scope.schema { try await Self.switchSchema(driver, to: schema, timeoutSeconds: operationTimeoutSeconds) } } catch { driver.disconnect() throw error } - return Entry(driver: driver) + return driver } static func connect(_ driver: DatabaseDriver, database: String, timeoutSeconds: Double) async throws { diff --git a/TablePro/Views/Main/Extensions/MainContentCoordinator+QueryHelpers.swift b/TablePro/Views/Main/Extensions/MainContentCoordinator+QueryHelpers.swift index 55110842c..a920f28fb 100644 --- a/TablePro/Views/Main/Extensions/MainContentCoordinator+QueryHelpers.swift +++ b/TablePro/Views/Main/Extensions/MainContentCoordinator+QueryHelpers.swift @@ -25,6 +25,30 @@ extension MainContentCoordinator { ).post() } + /// A table tab's SELECT is the app's own, so it may follow a database switch it waited through + /// onto a pooled connection. The choice belongs to the tab and never to the statement: an editor + /// SELECT can read a temp table or sit inside the user's open transaction, and moving it to + /// another connection would lose both. + func withExecutionDriver( + scope: DatabaseScope, + isTableTab: Bool, + _ body: @Sendable @escaping (DatabaseDriver) async throws -> T + ) async throws -> T { + guard isTableTab else { + return try await services.databaseManager.withScopedDriver( + scope: scope, + route: services.databaseManager.executionRoute(for: scope), + cancellation: .cancellableRead, + body + ) + } + return try await services.databaseManager.withTableReadDriver( + scope: scope, + cancellation: .cancellableRead, + body + ) + } + func finishFailedQuery( _ error: Error, tabId: UUID, diff --git a/TablePro/Views/Main/MainContentCoordinator.swift b/TablePro/Views/Main/MainContentCoordinator.swift index fa9b110f1..5f9986e53 100644 --- a/TablePro/Views/Main/MainContentCoordinator.swift +++ b/TablePro/Views/Main/MainContentCoordinator.swift @@ -1320,6 +1320,7 @@ final class MainContentCoordinator { } return } + let isTableTab = tab.tabType == .table let queryTask = Task { [weak self] in guard let self else { return } @@ -1348,10 +1349,9 @@ final class MainContentCoordinator { let fetchBeganAt = ContinuousClock.now do { - let fetchResult = try await services.databaseManager.withScopedDriver( + let fetchResult = try await withExecutionDriver( scope: scope, - route: services.databaseManager.executionRoute(for: scope), - cancellation: .cancellableRead + isTableTab: isTableTab ) { [queryExecutor] driver in try await queryExecutor.executeQuery( driver: driver, diff --git a/TableProTests/Core/Database/DatabaseSwitchLeaseOrderingTests.swift b/TableProTests/Core/Database/DatabaseSwitchLeaseOrderingTests.swift index f10e34279..9d3e85ccc 100644 --- a/TableProTests/Core/Database/DatabaseSwitchLeaseOrderingTests.swift +++ b/TableProTests/Core/Database/DatabaseSwitchLeaseOrderingTests.swift @@ -37,13 +37,15 @@ private final class Latch { @MainActor struct DatabaseSwitchLeaseOrderingTests { private static let typeId = "LeaseOrderingReconnectFake" + private static let unpooledTypeId = "LeaseOrderingUnpooledReconnectFake" /// A type that reopens its connection to change database, and whose driver plugin is not - /// registered, so every switch reaches the reconnect and fails there. - private func registerTypeIfNeeded() { - guard PluginMetadataRegistry.shared.snapshot(forRegisteredTypeId: Self.typeId) == nil else { return } + /// registered, so every switch reaches the reconnect and fails there. `pools` says whether it can + /// open a second connection for a database the session has left. + private func registerTypeIfNeeded(_ typeId: String = Self.typeId, pools: Bool = true) { + guard PluginMetadataRegistry.shared.snapshot(forRegisteredTypeId: typeId) == nil else { return } let defaults = PluginMetadataSnapshot.CapabilityFlags.defaults - let capabilities = PluginMetadataSnapshot.CapabilityFlags( + var capabilities = PluginMetadataSnapshot.CapabilityFlags( supportsSchemaSwitching: true, supportsImport: defaults.supportsImport, supportsExport: defaults.supportsExport, @@ -56,17 +58,18 @@ struct DatabaseSwitchLeaseOrderingTests { requiresReconnectForDatabaseSwitch: true, supportsDropDatabase: defaults.supportsDropDatabase ) + capabilities.supportsConnectionPooling = pools let snapshot = PluginMetadataSnapshot( - displayName: Self.typeId, iconName: "cylinder", defaultPort: 1_234, + displayName: typeId, iconName: "cylinder", defaultPort: 1_234, requiresAuthentication: true, supportsForeignKeys: true, supportsSchemaEditing: true, - isDownloadable: false, primaryUrlScheme: "leaseorderingfake", parameterStyle: .questionMark, + isDownloadable: false, primaryUrlScheme: typeId.lowercased(), parameterStyle: .questionMark, navigationModel: .standard, explainVariants: [], pathFieldRole: .database, - supportsHealthMonitor: false, urlSchemes: ["leaseorderingfake"], postConnectActions: [], + supportsHealthMonitor: false, urlSchemes: [typeId.lowercased()], postConnectActions: [], brandColorHex: "#000000", queryLanguageName: "SQL", editorLanguage: .sql, connectionMode: .network, supportsDatabaseSwitching: true, capabilities: capabilities, schema: .defaults, editor: .defaults, connection: .defaults ) - PluginMetadataRegistry.shared.register(snapshot: snapshot, forTypeId: Self.typeId) + PluginMetadataRegistry.shared.register(snapshot: snapshot, forTypeId: typeId) } private func makeSession() -> DatabaseConnection { @@ -109,6 +112,36 @@ struct DatabaseSwitchLeaseOrderingTests { } } + /// Waits for work running off the main actor to reach a point it reports. Bounded for the same + /// reason as `waitForQueuedCallers`. + private func waitUntil(_ condition: () -> Bool) async { + for _ in 0..<2_000 where !condition() { + try? await Task.sleep(for: .milliseconds(5)) + } + } + + /// A session browsing `app`, with the connection's pooled connection to `app` seeded, so a read + /// the pool serves is told apart from one the session driver serves. + private func makeTableReadSession(type: DatabaseType) -> (DatabaseConnection, MockDatabaseDriver) { + let connection = TestFixtures.makeConnection(database: "app", type: type) + var session = ConnectionSession(connection: connection, driver: MockDatabaseDriver(connection: connection)) + session.status = .connected + session.browseDatabase = "app" + DatabaseManager.shared.injectSession(session, for: connection.id) + let pooled = MockDatabaseDriver(connection: connection) + MetadataConnectionPool.shared.injectEntry(pooled, scope: appScope(connection)) + return (connection, pooled) + } + + private func appScope(_ connection: DatabaseConnection) -> DatabaseScope { + DatabaseScope(connectionId: connection.id, database: "app", schema: nil) + } + + private func cleanUpTableRead(_ connectionId: UUID) { + MetadataConnectionPool.shared.closeAll(connectionId: connectionId) + DatabaseManager.shared.removeSession(for: connectionId) + } + /// The reported schedule, with its order fixed by the gate rather than by timing: the switch /// queues for the driver first, then a table load already routed onto the switch's target queues /// behind it. This switch cannot reconnect, so it rolls back and leaves the driver dead. The @@ -310,6 +343,178 @@ struct DatabaseSwitchLeaseOrderingTests { } #expect(!ran.didRun) } + + // MARK: - A table read queued through a switch + + /// The reported schedule: a table load for `app` queues behind a switch to `orders` on an engine + /// that reconnects to switch. When its turn came it was refused with "This tab is on app", although + /// a pooled connection to `app` would have served it and running it again did. + @Test("A table read that waited through a switch runs on a pooled connection to its own database") + func tableReadFollowsTheSwitchOntoThePool() async throws { + let (connection, pooled) = makeTableReadSession(type: .postgresql) + defer { cleanUpTableRead(connection.id) } + let app = appScope(connection) + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + let read = Task { @MainActor in + try await DatabaseManager.shared.withTableReadDriver(scope: app, cancellation: .cancellableRead) { driver in + driver === pooled + } + } + await waitForQueuedCallers(1, on: connection.id) + #expect(DatabaseManager.shared.sessionDriverGate.waiterCount(for: connection.id) == 1) + + DatabaseManager.shared.updateSession(connection.id) { $0.browseDatabase = "orders" } + release.open() + try await holder.value + + #expect(try await read.value) + } + + /// The re-route belongs to table reads alone. A COMMIT or an editor statement queued the same way + /// was written against the session that is gone, so it is still refused and never reaches the pool. + @Test("Other work queued through the same switch is still refused on the session driver") + func otherWorkQueuedThroughASwitchIsStillRefused() async throws { + let (connection, _) = makeTableReadSession(type: .postgresql) + defer { cleanUpTableRead(connection.id) } + let app = appScope(connection) + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + let ran = LeaseRecord() + let lease = Task { @MainActor in + try await DatabaseManager.shared.withScopedDriver( + scope: app, + route: DatabaseManager.shared.executionRoute(for: app), + cancellation: .cancellableRead + ) { _ in + await MainActor.run { ran.didRun = true } + } + } + await waitForQueuedCallers(1, on: connection.id) + #expect(DatabaseManager.shared.sessionDriverGate.waiterCount(for: connection.id) == 1) + + DatabaseManager.shared.updateSession(connection.id) { $0.browseDatabase = "orders" } + release.open() + try await holder.value + + await #expect { + try await lease.value + } throws: { error in + guard case .queryFailed(let message) = error as? DatabaseError else { return false } + return message.contains("app") + } + #expect(!ran.didRun) + } + + @Test("A re-routed table read leaves the session driver free while it runs on the pool") + func reroutedTableReadDoesNotHoldTheGate() async throws { + let (connection, pooled) = makeTableReadSession(type: .postgresql) + defer { cleanUpTableRead(connection.id) } + let app = appScope(connection) + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + let running = LeaseRecord() + let finish = Latch() + let read = Task { @MainActor in + try await DatabaseManager.shared.withTableReadDriver(scope: app, cancellation: .cancellableRead) { driver in + await MainActor.run { running.didRun = true } + await finish.wait() + return driver === pooled + } + } + await waitForQueuedCallers(1, on: connection.id) + DatabaseManager.shared.updateSession(connection.id) { $0.browseDatabase = "orders" } + release.open() + try await holder.value + await waitUntil { running.didRun } + #expect(running.didRun) + + let probe = LeaseRecord() + let prober = Task { @MainActor in + try await DatabaseManager.shared.sessionDriverGate.withExclusiveAccess(connection.id) { + probe.didRun = true + } + } + await waitUntil { probe.didRun } + + #expect(probe.didRun) + #expect(DatabaseManager.shared.sessionDriverGate.waiterCount(for: connection.id) == 0) + finish.open() + try await prober.value + #expect(try await read.value) + } + + /// The driver is judged before the route, so a read that waited behind a driver that stopped + /// answering is refused like any other lease rather than slipping out to a pooled connection. + @Test("A table read queued behind a driver that stopped answering is refused and never reaches the pool") + func tableReadBehindADeadDriverIsRefused() async throws { + let (connection, _) = makeTableReadSession(type: .postgresql) + defer { cleanUpTableRead(connection.id) } + let app = appScope(connection) + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + let ran = LeaseRecord() + let read = Task { @MainActor in + try await DatabaseManager.shared.withTableReadDriver(scope: app, cancellation: .cancellableRead) { _ in + await MainActor.run { ran.didRun = true } + } + } + await waitForQueuedCallers(1, on: connection.id) + #expect(DatabaseManager.shared.sessionDriverGate.waiterCount(for: connection.id) == 1) + + DatabaseManager.shared.updateSession(connection.id) { session in + session.browseDatabase = "orders" + session.liveness = .unreachable(nil) + } + release.open() + try await holder.value + + await #expect { + try await read.value + } throws: { error in + guard case .notConnected = error as? DatabaseError else { return false } + return true + } + #expect(!ran.didRun) + } + + @Test("A table read left behind by a switch on an engine that cannot pool names its database") + func tableReadWithNowhereToGoNamesItsDatabase() async throws { + registerTypeIfNeeded(Self.unpooledTypeId, pools: false) + let (connection, _) = makeTableReadSession(type: DatabaseType(rawValue: Self.unpooledTypeId)) + defer { + cleanUpTableRead(connection.id) + PluginMetadataRegistry.shared.unregister(typeId: Self.unpooledTypeId) + } + let app = appScope(connection) + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + let ran = LeaseRecord() + let read = Task { @MainActor in + try await DatabaseManager.shared.withTableReadDriver(scope: app, cancellation: .cancellableRead) { _ in + await MainActor.run { ran.didRun = true } + } + } + await waitForQueuedCallers(1, on: connection.id) + #expect(DatabaseManager.shared.sessionDriverGate.waiterCount(for: connection.id) == 1) + + DatabaseManager.shared.updateSession(connection.id) { $0.browseDatabase = "orders" } + release.open() + try await holder.value + + await #expect { + try await read.value + } throws: { error in + guard case .queryFailed(let message) = error as? DatabaseError else { return false } + return message.contains("app") + } + #expect(!ran.didRun) + } } @MainActor diff --git a/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift b/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift new file mode 100644 index 000000000..552b3dc53 --- /dev/null +++ b/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift @@ -0,0 +1,148 @@ +// +// SwitchDatabasePooledConnectionTests.swift +// TableProTests +// +// A database switch on an engine that reconnects to perform one used to close every pooled +// connection the connection held. A table tab on another database that was dialing one lost its +// open and showed nothing, and idle pooled connections to other databases were closed for a switch +// that never touched them. +// + +import Foundation +@testable import TablePro +import TableProPluginKit +import Testing + +@Suite("Switch database and pooled connections", .serialized) +@MainActor +struct SwitchDatabasePooledConnectionTests { + /// Reopens its connection to change database and has no driver plugin registered, so every switch + /// reaches the reconnect and fails there with no network and no waiting. + private static let typeId = "PooledSwitchReconnectFake" + + private func registerTypeIfNeeded() { + guard PluginMetadataRegistry.shared.snapshot(forRegisteredTypeId: Self.typeId) == nil else { return } + let defaults = PluginMetadataSnapshot.CapabilityFlags.defaults + let capabilities = PluginMetadataSnapshot.CapabilityFlags( + supportsSchemaSwitching: true, + supportsImport: defaults.supportsImport, + supportsExport: defaults.supportsExport, + supportsSSH: defaults.supportsSSH, + supportsSSL: defaults.supportsSSL, + supportsCascadeDrop: defaults.supportsCascadeDrop, + supportsForeignKeyDisable: defaults.supportsForeignKeyDisable, + supportsReadOnlyMode: defaults.supportsReadOnlyMode, + supportsQueryProgress: defaults.supportsQueryProgress, + requiresReconnectForDatabaseSwitch: true, + supportsDropDatabase: defaults.supportsDropDatabase + ) + let snapshot = PluginMetadataSnapshot( + displayName: Self.typeId, iconName: "cylinder", defaultPort: 1_234, + requiresAuthentication: true, supportsForeignKeys: true, supportsSchemaEditing: true, + isDownloadable: false, primaryUrlScheme: "pooledswitchfake", parameterStyle: .questionMark, + navigationModel: .standard, explainVariants: [], pathFieldRole: .database, + supportsHealthMonitor: false, urlSchemes: ["pooledswitchfake"], postConnectActions: [], + brandColorHex: "#000000", queryLanguageName: "SQL", editorLanguage: .sql, + connectionMode: .network, supportsDatabaseSwitching: true, + capabilities: capabilities, schema: .defaults, editor: .defaults, connection: .defaults + ) + PluginMetadataRegistry.shared.register(snapshot: snapshot, forTypeId: Self.typeId) + } + + /// A direct connection, with no tunnel of any kind, browsing `app`. + private func makeSession(stoppedAnswering: Bool = false) -> DatabaseConnection { + registerTypeIfNeeded() + var connection = TestFixtures.makeConnection(database: "app") + connection.type = DatabaseType(rawValue: Self.typeId) + var session = ConnectionSession(connection: connection, driver: MockDatabaseDriver(connection: connection)) + session.status = .connected + session.browseDatabase = "app" + if stoppedAnswering { + session.liveness = .unreachable(nil) + } + DatabaseManager.shared.injectSession(session, for: connection.id) + return connection + } + + /// Stands in for the pooled connection a table tab on `reports` opened. + private func seedPooledConnection(for connectionId: UUID) -> MockDatabaseDriver { + let pooled = MockDatabaseDriver() + MetadataConnectionPool.shared.injectEntry( + pooled, + scope: DatabaseScope(connectionId: connectionId, database: "reports", schema: nil) + ) + return pooled + } + + private func cleanUp(_ connectionId: UUID) { + MetadataConnectionPool.shared.closeAll(connectionId: connectionId) + DatabaseManager.shared.removeSession(for: connectionId) + AppSettingsStorage.shared.saveLastDatabase(nil, for: connectionId) + AppSettingsStorage.shared.saveLastSchema(nil, for: connectionId) + PluginMetadataRegistry.shared.unregister(typeId: Self.typeId) + } + + @Test("A switch over a direct connection leaves the pooled connection to another database open") + func switchLeavesOtherDatabasesPooled() async { + let connection = makeSession() + defer { cleanUp(connection.id) } + let pooled = seedPooledConnection(for: connection.id) + + try? await DatabaseManager.shared.switchDatabase(to: "unreachable", for: connection.id, persist: false) + + /// Only the reconnect branch moves the session off `.connected`, so this proves it ran. + #expect(DatabaseManager.shared.session(for: connection.id)?.status != .connected) + #expect(MetadataConnectionPool.shared.pooledDriverCount(for: connection.id) == 1) + #expect(pooled.disconnectCallCount == 0) + #expect(!MetadataConnectionPool.shared.isReplacingTransport(for: connection.id)) + } + + /// A session that had stopped answering has most likely lost its pooled connections too, so a + /// switch that reconnects it is a recovery and still clears them, and releases the pool again when + /// the reconnect fails. + @Test("A switch that reconnects a connection that stopped answering closes its pooled connections") + func recoveringSwitchClosesThePool() async { + let connection = makeSession(stoppedAnswering: true) + defer { cleanUp(connection.id) } + let pooled = seedPooledConnection(for: connection.id) + + try? await DatabaseManager.shared.switchDatabase(to: "unreachable", for: connection.id, persist: false) + + #expect(DatabaseManager.shared.session(for: connection.id)?.status != .connected) + #expect(MetadataConnectionPool.shared.pooledDriverCount(for: connection.id) == 0) + #expect(pooled.disconnectCallCount == 1) + #expect(!MetadataConnectionPool.shared.isReplacingTransport(for: connection.id)) + } + + @Test("A background reconnect closes the pooled connections and releases the pool when it is done") + func healthReconnectClosesThePool() async { + FakeMSSQLPluginRegistration.registerIfNeeded() + var connection = TestFixtures.makeConnection(name: "Prod") + connection.type = DatabaseType(rawValue: FakeMSSQLPlugin.databaseTypeId) + var session = ConnectionSession(connection: connection) + session.status = .connected + DatabaseManager.shared.injectSession(session, for: connection.id) + let pooled = seedPooledConnection(for: connection.id) + + _ = await DatabaseManager.shared.performHealthMonitorReconnect(connectionId: connection.id) + + #expect(MetadataConnectionPool.shared.pooledDriverCount(for: connection.id) == 0) + #expect(pooled.disconnectCallCount == 1) + #expect(!MetadataConnectionPool.shared.isReplacingTransport(for: connection.id)) + MetadataConnectionPool.shared.closeAll(connectionId: connection.id) + DatabaseManager.shared.removeSession(for: connection.id) + await SchemaService.shared.invalidate(connectionId: connection.id) + } + + @Test("The sidebar's reconnect handling leaves pooled connections to the transport's owner") + func treeReconnectLeavesThePool() async { + let connectionId = UUID() + let pooled = seedPooledConnection(for: connectionId) + defer { MetadataConnectionPool.shared.closeAll(connectionId: connectionId) } + + await DatabaseTreeMetadataService.shared.handleReconnect(connectionId: connectionId) + + #expect(MetadataConnectionPool.shared.pooledDriverCount(for: connectionId) == 1) + #expect(pooled.disconnectCallCount == 0) + } +} diff --git a/TableProTests/Core/Services/Query/MetadataConnectionPoolTests.swift b/TableProTests/Core/Services/Query/MetadataConnectionPoolTests.swift index 3ab20df49..2dcbd0f7a 100644 --- a/TableProTests/Core/Services/Query/MetadataConnectionPoolTests.swift +++ b/TableProTests/Core/Services/Query/MetadataConnectionPoolTests.swift @@ -151,6 +151,174 @@ private final class PoolBodyFlag: @unchecked Sendable { var value = false } +/// Stands in for opening a pooled connection. Each open takes the next connect delay, and a delayed +/// connect throws when its open is cancelled, the way libpq's cooperative connect does. +@MainActor +private final class RecordingOpener { + private var connectDelays: [Double] + private(set) var opened: [MockDatabaseDriver] = [] + + init(connectDelays: [Double] = []) { + self.connectDelays = connectDelays + } + + func open(_ scope: DatabaseScope) async throws -> DatabaseDriver { + let driver = MockDatabaseDriver() + driver.connectDelaySeconds = connectDelays.isEmpty ? 0 : connectDelays.removeFirst() + opened.append(driver) + try await driver.connect() + return driver + } +} + +/// A reconnect that rebuilds a connection's transport used to close every pooled entry and cancel +/// every open in progress, so a table load dialing a pooled connection failed as a user cancel and +/// showed nothing. +@Suite("MetadataConnectionPool transport replacement", .serialized) +@MainActor +struct MetadataConnectionPoolTransportReplacementTests { + private func makeSession() -> (DatabaseConnection, DatabaseScope) { + let connection = TestFixtures.makeConnection(database: "shop") + var session = ConnectionSession(connection: connection, driver: MockDatabaseDriver(connection: connection)) + session.status = .connected + DatabaseManager.shared.injectSession(session, for: connection.id) + return (connection, DatabaseScope(connectionId: connection.id, database: "shop", schema: nil)) + } + + /// Bounded, so a caller that never reaches the point being waited for fails the assertions after + /// it rather than hanging the suite. + private func waitUntil(_ condition: () -> Bool) async { + for _ in 0..<2_000 where !condition() { + try? await Task.sleep(for: .milliseconds(5)) + } + } + + @Test("A lease waiting on an open the replacement withdrew takes a new connection once it ends") + func withdrawnOpenIsRetriedAfterTheReplacement() async throws { + let (connection, scope) = makeSession() + defer { DatabaseManager.shared.removeSession(for: connection.id) } + let opener = RecordingOpener(connectDelays: [5, 0]) + let pool = MetadataConnectionPool.isolatedForTesting(openDriver: { try await opener.open($0) }) + defer { pool.closeAll(connectionId: connection.id) } + + let lease = Task { @MainActor in + try await pool.withDriver(scope: scope) { driver in ObjectIdentifier(driver) } + } + await waitUntil { opener.opened.count == 1 } + #expect(opener.opened.count == 1) + + pool.beginTransportReplacement(connectionId: connection.id) + await waitUntil { pool.transportWaiterCount(for: connection.id) == 1 } + #expect(pool.transportWaiterCount(for: connection.id) == 1) + #expect(opener.opened.count == 1) + + pool.endTransportReplacement(connectionId: connection.id) + let ranOn = try await lease.value + + #expect(opener.opened.count == 2) + #expect(ranOn == opener.opened.last.map { ObjectIdentifier($0) }) + } + + /// A new open during the replacement would read the effective connection before the reconnect + /// replaced it, and dial a tunnel port that is about to close. + @Test("A lease that arrives during a replacement opens nothing until it ends") + func leaseDuringAReplacementWaitsToOpen() async throws { + let (connection, scope) = makeSession() + defer { DatabaseManager.shared.removeSession(for: connection.id) } + let opener = RecordingOpener() + let pool = MetadataConnectionPool.isolatedForTesting(openDriver: { try await opener.open($0) }) + defer { pool.closeAll(connectionId: connection.id) } + + pool.beginTransportReplacement(connectionId: connection.id) + let lease = Task { @MainActor in + try await pool.withDriver(scope: scope) { _ in } + } + await waitUntil { pool.transportWaiterCount(for: connection.id) == 1 } + #expect(pool.transportWaiterCount(for: connection.id) == 1) + #expect(opener.opened.isEmpty) + + pool.endTransportReplacement(connectionId: connection.id) + try await lease.value + + #expect(opener.opened.count == 1) + } + + @Test("Overlapping replacements hold pooled work until the last one ends") + func overlappingReplacementsHoldUntilTheLastEnds() async throws { + let (connection, scope) = makeSession() + defer { DatabaseManager.shared.removeSession(for: connection.id) } + let opener = RecordingOpener() + let pool = MetadataConnectionPool.isolatedForTesting(openDriver: { try await opener.open($0) }) + defer { pool.closeAll(connectionId: connection.id) } + + pool.beginTransportReplacement(connectionId: connection.id) + pool.beginTransportReplacement(connectionId: connection.id) + let lease = Task { @MainActor in + try await pool.withDriver(scope: scope) { _ in } + } + await waitUntil { pool.transportWaiterCount(for: connection.id) == 1 } + + pool.endTransportReplacement(connectionId: connection.id) + #expect(pool.transportWaiterCount(for: connection.id) == 1) + #expect(opener.opened.isEmpty) + + pool.endTransportReplacement(connectionId: connection.id) + try await lease.value + + #expect(opener.opened.count == 1) + #expect(!pool.isReplacingTransport(for: connection.id)) + } + + @Test("A lease cancelled while it waits for a replacement stops waiting at once") + func cancelledWaiterStopsWaiting() async throws { + let (connection, scope) = makeSession() + defer { DatabaseManager.shared.removeSession(for: connection.id) } + let opener = RecordingOpener() + let pool = MetadataConnectionPool.isolatedForTesting(openDriver: { try await opener.open($0) }) + defer { pool.closeAll(connectionId: connection.id) } + + pool.beginTransportReplacement(connectionId: connection.id) + defer { pool.endTransportReplacement(connectionId: connection.id) } + let lease = Task { @MainActor in + try await pool.withDriver(scope: scope) { _ in } + } + await waitUntil { pool.transportWaiterCount(for: connection.id) == 1 } + #expect(pool.transportWaiterCount(for: connection.id) == 1) + + lease.cancel() + + await #expect(throws: CancellationError.self) { + try await lease.value + } + #expect(pool.transportWaiterCount(for: connection.id) == 0) + #expect(opener.opened.isEmpty) + } + + /// A rename needs every backend on the database gone, and PostgreSQL refuses it while one is + /// attached, so an open withdrawn for it must never be dialed again. + @Test("Closing a database withdraws its open for good and fails the lease waiting on it") + func closingADatabaseFailsItsWaiter() async throws { + let (connection, scope) = makeSession() + defer { DatabaseManager.shared.removeSession(for: connection.id) } + let opener = RecordingOpener(connectDelays: [5]) + let pool = MetadataConnectionPool.isolatedForTesting(openDriver: { try await opener.open($0) }) + defer { pool.closeAll(connectionId: connection.id) } + + let lease = Task { @MainActor in + try await pool.withDriver(scope: scope) { _ in } + } + await waitUntil { opener.opened.count == 1 } + #expect(opener.opened.count == 1) + + pool.closeAll(connectionId: connection.id, database: "shop") + + await #expect(throws: (any Error).self) { + try await lease.value + } + #expect(opener.opened.count == 1) + } +} + @Suite("MetadataConnectionPool idle eviction", .serialized) @MainActor struct MetadataConnectionPoolIdleEvictionTests { From b2e1759e101febfcf14a669c2d7bba33d1571e27 Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Mon, 14 Sep 2026 00:59:44 +0700 Subject: [PATCH 2/3] fix(connections): hold the pool only while a tunnel recovery connects and fail parked leases on close --- .../Database/DatabaseManager+Tunnel.swift | 10 ++-- .../Query/MetadataConnectionPool.swift | 41 +++++++++---- .../SwitchDatabasePooledConnectionTests.swift | 43 ++++++++++++++ .../Query/MetadataConnectionPoolTests.swift | 57 +++++++++++++++++++ 4 files changed, 137 insertions(+), 14 deletions(-) diff --git a/TablePro/Core/Database/DatabaseManager+Tunnel.swift b/TablePro/Core/Database/DatabaseManager+Tunnel.swift index 88d72c0f6..0f028d601 100644 --- a/TablePro/Core/Database/DatabaseManager+Tunnel.swift +++ b/TablePro/Core/Database/DatabaseManager+Tunnel.swift @@ -140,10 +140,6 @@ extension DatabaseManager { recoveringConnectionIds.insert(connectionId) defer { recoveringConnectionIds.remove(connectionId) } - /// The tunnel every pooled connection dialed is gone, and recovering binds a new one on another - /// port, so pooled work waits for it instead of dialing a port another tunnel may since own. - MetadataConnectionPool.shared.beginTransportReplacement(connectionId: connectionId) - defer { MetadataConnectionPool.shared.endTransportReplacement(connectionId: connectionId) } Self.logger.warning("\(kind, privacy: .public) tunnel died for connection: \(session.connection.name)") @@ -161,6 +157,12 @@ extension DatabaseManager { Self.logger.info("\(kind, privacy: .public) reconnect attempt \(retryCount + 1)/\(maxRetries) in \(delay)s for: \(session.connection.name)") try? await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) + /// Held for the attempt, never across the wait before it. The tunnel every pooled connection + /// dialed is gone and the attempt binds a new one on another port, so pooled work waits for + /// it rather than dialing a port another tunnel may since own. Held across the backoff, which + /// reaches two minutes, a sidebar read would spin with no error for the whole recovery. + MetadataConnectionPool.shared.beginTransportReplacement(connectionId: connectionId) + defer { MetadataConnectionPool.shared.endTransportReplacement(connectionId: connectionId) } do { try await connectToSession(session.connection) Self.logger.info("Successfully reconnected \(kind, privacy: .public) tunnel for: \(session.connection.name)") diff --git a/TablePro/Core/Services/Query/MetadataConnectionPool.swift b/TablePro/Core/Services/Query/MetadataConnectionPool.swift index d0693af37..8d52bd1bb 100644 --- a/TablePro/Core/Services/Query/MetadataConnectionPool.swift +++ b/TablePro/Core/Services/Query/MetadataConnectionPool.swift @@ -80,6 +80,7 @@ final class MetadataConnectionPool { private struct TransportWaiter { let ticket: UUID + let scope: DatabaseScope let continuation: CheckedContinuation } @@ -145,10 +146,11 @@ final class MetadataConnectionPool { /// /// A pooled entry stands on the session's effective connection plus its own database, so it only /// has to go when that endpoint does: a tunnel rebuilt on a new local port, or a reconnect - /// recovering a connection that stopped answering. An open already dialing is withdrawn, and the - /// callers waiting on it take a connection again after `endTransportReplacement` rather than - /// failing a load nobody cancelled. A new caller waits for the same moment instead of dialing an - /// endpoint about to close. Replacements nest, and pooled work resumes when the last one ends. + /// recovering a connection that stopped answering. An open already dialing is withdrawn, and once + /// it has returned the callers waiting on it take a connection again after `endTransportReplacement` + /// rather than failing a load nobody cancelled. A new caller waits for the same moment instead of + /// dialing an endpoint about to close. Replacements nest, and pooled work resumes when the last one + /// ends. Closing the connection, or one of its databases, fails the callers waiting for it. func beginTransportReplacement(connectionId: UUID) { transportReplacements[connectionId, default: 0] += 1 closeEntries(withdrawingOpensAs: .transportReplaced) { $0.connectionId == connectionId } @@ -243,16 +245,34 @@ final class MetadataConnectionPool { open.withdrawal = withdrawal open.task.cancel() } + if withdrawal == .closed { + failTransportWaiters(where: matches) + } for key in entries.keys where matches(key.scope) { closeOrDeferEntry(forKey: key) } stopSweeperIfEmpty() } + /// A caller parked for a replacement is waiting on its connection as much as one waiting on a + /// pending open, so closing what it waits for fails it too. Left parked, it would wake when the + /// replacement ends and open a connection for a session, or a database, that has gone. + private func failTransportWaiters(where matches: (DatabaseScope) -> Bool) { + for (connectionId, waiters) in transportWaiters { + let failing = waiters.filter { matches($0.scope) } + guard !failing.isEmpty else { continue } + let remaining = waiters.filter { !matches($0.scope) } + transportWaiters[connectionId] = remaining.isEmpty ? nil : remaining + for waiter in failing { + waiter.continuation.resume(throwing: CancellationError()) + } + } + } + private func acquireEntry(scope: DatabaseScope, workload: Workload) async throws -> Entry { let key = Key(scope: scope, workload: workload) while true { - try await waitForTransport(of: scope.connectionId) + try await waitForTransport(for: scope) if let entry = reusableEntry(forKey: key) { return entry } @@ -323,11 +343,12 @@ final class MetadataConnectionPool { /// Parks a caller for as long as the connection's transport is being replaced. A caller that is /// cancelled while parked stops waiting at once rather than when the replacement ends. - private func waitForTransport(of connectionId: UUID) async throws { + private func waitForTransport(for scope: DatabaseScope) async throws { + let connectionId = scope.connectionId while transportReplacements[connectionId] != nil { let ticket = UUID() try await withTaskCancellationHandler( - operation: { try await parkForTransport(ticket: ticket, connectionId: connectionId) }, + operation: { try await parkForTransport(ticket: ticket, scope: scope) }, onCancel: { [weak self] in Task { @MainActor in self?.failTransportWaiter(ticket: ticket, connectionId: connectionId) @@ -337,14 +358,14 @@ final class MetadataConnectionPool { } } - private func parkForTransport(ticket: UUID, connectionId: UUID) async throws { + private func parkForTransport(ticket: UUID, scope: DatabaseScope) async throws { try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in guard !Task.isCancelled else { continuation.resume(throwing: CancellationError()) return } - transportWaiters[connectionId, default: []].append( - TransportWaiter(ticket: ticket, continuation: continuation) + transportWaiters[scope.connectionId, default: []].append( + TransportWaiter(ticket: ticket, scope: scope, continuation: continuation) ) } } diff --git a/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift b/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift index 552b3dc53..b7a308644 100644 --- a/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift +++ b/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift @@ -82,6 +82,14 @@ struct SwitchDatabasePooledConnectionTests { PluginMetadataRegistry.shared.unregister(typeId: Self.typeId) } + /// Bounded, so a recovery that never reaches the point being waited for fails the assertions + /// after it rather than hanging the suite. + private func waitUntil(_ condition: () -> Bool) async { + for _ in 0..<2_000 where !condition() { + try? await Task.sleep(for: .milliseconds(5)) + } + } + @Test("A switch over a direct connection leaves the pooled connection to another database open") func switchLeavesOtherDatabasesPooled() async { let connection = makeSession() @@ -145,4 +153,39 @@ struct SwitchDatabasePooledConnectionTests { #expect(MetadataConnectionPool.shared.pooledDriverCount(for: connectionId) == 1) #expect(pooled.disconnectCallCount == 0) } + + /// A tunnel recovery waits out a backoff before every attempt, reaching two minutes. Holding the + /// pool across those waits left every pooled read on the connection spinning with no error for the + /// whole recovery. The recovery clears the driver immediately before its first two-second wait, + /// which is the handshake this reads. + @Test("A tunnel recovery holds the pool only while an attempt connects") + func tunnelRecoveryHoldsThePoolOnlyWhileItConnects() async { + FakeMSSQLPluginRegistration.registerIfNeeded() + var connection = TestFixtures.makeConnection(name: "Tunneled") + connection.type = DatabaseType(rawValue: FakeMSSQLPlugin.databaseTypeId) + var session = ConnectionSession(connection: connection, driver: MockDatabaseDriver(connection: connection)) + session.status = .connected + DatabaseManager.shared.injectSession(session, for: connection.id) + let pooled = seedPooledConnection(for: connection.id) + + let recovery = Task { @MainActor in + await DatabaseManager.shared.recoverDeadTunnel( + connectionId: connection.id, kind: "SSH", disconnectedMessage: "The tunnel closed." + ) + } + await waitUntil { DatabaseManager.shared.session(for: connection.id)?.driver == nil } + #expect(DatabaseManager.shared.session(for: connection.id)?.driver == nil) + #expect(!MetadataConnectionPool.shared.isReplacingTransport(for: connection.id)) + + await recovery.value + + #expect(DatabaseManager.shared.session(for: connection.id)?.driver != nil) + #expect(MetadataConnectionPool.shared.pooledDriverCount(for: connection.id) == 0) + #expect(pooled.disconnectCallCount == 1) + #expect(!MetadataConnectionPool.shared.isReplacingTransport(for: connection.id)) + await DatabaseManager.shared.stopHealthMonitor(for: connection.id) + MetadataConnectionPool.shared.closeAll(connectionId: connection.id) + DatabaseManager.shared.removeSession(for: connection.id) + await SchemaService.shared.invalidate(connectionId: connection.id) + } } diff --git a/TableProTests/Core/Services/Query/MetadataConnectionPoolTests.swift b/TableProTests/Core/Services/Query/MetadataConnectionPoolTests.swift index 2dcbd0f7a..dc0436fb4 100644 --- a/TableProTests/Core/Services/Query/MetadataConnectionPoolTests.swift +++ b/TableProTests/Core/Services/Query/MetadataConnectionPoolTests.swift @@ -317,6 +317,63 @@ struct MetadataConnectionPoolTransportReplacementTests { } #expect(opener.opened.count == 1) } + + /// Ordered so that a parked lease the close failed to reach opens a connection once the + /// replacement ends, and fails the expectations, rather than hanging the suite. + @Test("Closing a connection fails the leases parked for its replacement at once") + func closingAConnectionFailsParkedLeases() async throws { + let (connection, scope) = makeSession() + defer { DatabaseManager.shared.removeSession(for: connection.id) } + let opener = RecordingOpener() + let pool = MetadataConnectionPool.isolatedForTesting(openDriver: { try await opener.open($0) }) + defer { pool.closeAll(connectionId: connection.id) } + + pool.beginTransportReplacement(connectionId: connection.id) + let lease = Task { @MainActor in + try await pool.withDriver(scope: scope) { _ in } + } + await waitUntil { pool.transportWaiterCount(for: connection.id) == 1 } + #expect(pool.transportWaiterCount(for: connection.id) == 1) + + pool.closeAll(connectionId: connection.id) + #expect(pool.transportWaiterCount(for: connection.id) == 0) + pool.endTransportReplacement(connectionId: connection.id) + + await #expect(throws: CancellationError.self) { + try await lease.value + } + #expect(opener.opened.isEmpty) + } + + @Test("Closing one database fails only the leases parked for that database") + func closingADatabaseFailsOnlyItsParkedLeases() async throws { + let (connection, shop) = makeSession() + defer { DatabaseManager.shared.removeSession(for: connection.id) } + let reports = DatabaseScope(connectionId: connection.id, database: "reports", schema: nil) + let opener = RecordingOpener() + let pool = MetadataConnectionPool.isolatedForTesting(openDriver: { try await opener.open($0) }) + defer { pool.closeAll(connectionId: connection.id) } + + pool.beginTransportReplacement(connectionId: connection.id) + let shopLease = Task { @MainActor in + try await pool.withDriver(scope: shop) { _ in } + } + let reportsLease = Task { @MainActor in + try await pool.withDriver(scope: reports) { _ in } + } + await waitUntil { pool.transportWaiterCount(for: connection.id) == 2 } + #expect(pool.transportWaiterCount(for: connection.id) == 2) + + pool.closeAll(connectionId: connection.id, database: "shop") + #expect(pool.transportWaiterCount(for: connection.id) == 1) + pool.endTransportReplacement(connectionId: connection.id) + + await #expect(throws: CancellationError.self) { + try await shopLease.value + } + try await reportsLease.value + #expect(opener.opened.count == 1) + } } @Suite("MetadataConnectionPool idle eviction", .serialized) From a241591015aa8e9cdf93211ad4031cb6ef58ecb0 Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Mon, 14 Sep 2026 01:12:44 +0700 Subject: [PATCH 3/3] fix(connections): leave tunnel recovery's pooled connections as they were --- .../Database/DatabaseManager+Tunnel.swift | 6 --- .../SwitchDatabasePooledConnectionTests.swift | 43 ------------------- 2 files changed, 49 deletions(-) diff --git a/TablePro/Core/Database/DatabaseManager+Tunnel.swift b/TablePro/Core/Database/DatabaseManager+Tunnel.swift index 0f028d601..9a2b54e91 100644 --- a/TablePro/Core/Database/DatabaseManager+Tunnel.swift +++ b/TablePro/Core/Database/DatabaseManager+Tunnel.swift @@ -157,12 +157,6 @@ extension DatabaseManager { Self.logger.info("\(kind, privacy: .public) reconnect attempt \(retryCount + 1)/\(maxRetries) in \(delay)s for: \(session.connection.name)") try? await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) - /// Held for the attempt, never across the wait before it. The tunnel every pooled connection - /// dialed is gone and the attempt binds a new one on another port, so pooled work waits for - /// it rather than dialing a port another tunnel may since own. Held across the backoff, which - /// reaches two minutes, a sidebar read would spin with no error for the whole recovery. - MetadataConnectionPool.shared.beginTransportReplacement(connectionId: connectionId) - defer { MetadataConnectionPool.shared.endTransportReplacement(connectionId: connectionId) } do { try await connectToSession(session.connection) Self.logger.info("Successfully reconnected \(kind, privacy: .public) tunnel for: \(session.connection.name)") diff --git a/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift b/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift index b7a308644..552b3dc53 100644 --- a/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift +++ b/TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift @@ -82,14 +82,6 @@ struct SwitchDatabasePooledConnectionTests { PluginMetadataRegistry.shared.unregister(typeId: Self.typeId) } - /// Bounded, so a recovery that never reaches the point being waited for fails the assertions - /// after it rather than hanging the suite. - private func waitUntil(_ condition: () -> Bool) async { - for _ in 0..<2_000 where !condition() { - try? await Task.sleep(for: .milliseconds(5)) - } - } - @Test("A switch over a direct connection leaves the pooled connection to another database open") func switchLeavesOtherDatabasesPooled() async { let connection = makeSession() @@ -153,39 +145,4 @@ struct SwitchDatabasePooledConnectionTests { #expect(MetadataConnectionPool.shared.pooledDriverCount(for: connectionId) == 1) #expect(pooled.disconnectCallCount == 0) } - - /// A tunnel recovery waits out a backoff before every attempt, reaching two minutes. Holding the - /// pool across those waits left every pooled read on the connection spinning with no error for the - /// whole recovery. The recovery clears the driver immediately before its first two-second wait, - /// which is the handshake this reads. - @Test("A tunnel recovery holds the pool only while an attempt connects") - func tunnelRecoveryHoldsThePoolOnlyWhileItConnects() async { - FakeMSSQLPluginRegistration.registerIfNeeded() - var connection = TestFixtures.makeConnection(name: "Tunneled") - connection.type = DatabaseType(rawValue: FakeMSSQLPlugin.databaseTypeId) - var session = ConnectionSession(connection: connection, driver: MockDatabaseDriver(connection: connection)) - session.status = .connected - DatabaseManager.shared.injectSession(session, for: connection.id) - let pooled = seedPooledConnection(for: connection.id) - - let recovery = Task { @MainActor in - await DatabaseManager.shared.recoverDeadTunnel( - connectionId: connection.id, kind: "SSH", disconnectedMessage: "The tunnel closed." - ) - } - await waitUntil { DatabaseManager.shared.session(for: connection.id)?.driver == nil } - #expect(DatabaseManager.shared.session(for: connection.id)?.driver == nil) - #expect(!MetadataConnectionPool.shared.isReplacingTransport(for: connection.id)) - - await recovery.value - - #expect(DatabaseManager.shared.session(for: connection.id)?.driver != nil) - #expect(MetadataConnectionPool.shared.pooledDriverCount(for: connection.id) == 0) - #expect(pooled.disconnectCallCount == 1) - #expect(!MetadataConnectionPool.shared.isReplacingTransport(for: connection.id)) - await DatabaseManager.shared.stopHealthMonitor(for: connection.id) - MetadataConnectionPool.shared.closeAll(connectionId: connection.id) - DatabaseManager.shared.removeSession(for: connection.id) - await SchemaService.shared.invalidate(connectionId: connection.id) - } }