From 4ee84b1ef02359239c1a82b155685fce145bf240 Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Sun, 13 Sep 2026 22:32:40 +0700 Subject: [PATCH 1/4] fix(connections): fail queued driver work on disconnect and gate Redis database selection --- CHANGELOG.md | 2 + .../Core/Concurrency/SessionDriverGate.swift | 29 ++- .../Database/DatabaseManager+Sessions.swift | 16 +- .../ViewModels/RedisKeyTreeViewModel.swift | 8 +- .../MainContentCoordinator+Navigation.swift | 8 +- .../Concurrency/SessionDriverGateTests.swift | 110 ++++++++++ .../DatabaseManagerDisconnectTests.swift | 69 ++++++ .../DatabaseSwitchLeaseOrderingTests.swift | 47 +++- .../RedisDatabaseSelectionGateTests.swift | 200 ++++++++++++++++++ 9 files changed, 464 insertions(+), 25 deletions(-) create mode 100644 TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift diff --git a/CHANGELOG.md b/CHANGELOG.md index df3b69ac07..53f7b93a0e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Dropped tables left staged after a save that failed part way, and drops staged in another database unstaged on refresh. - Table tab showing another database's rows after a database switch on PostgreSQL, Redshift and CockroachDB. - 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. ## [0.74.0] - 2026-09-13 diff --git a/TablePro/Core/Concurrency/SessionDriverGate.swift b/TablePro/Core/Concurrency/SessionDriverGate.swift index 6244d659b2..7da213e1e7 100644 --- a/TablePro/Core/Concurrency/SessionDriverGate.swift +++ b/TablePro/Core/Concurrency/SessionDriverGate.swift @@ -13,6 +13,10 @@ import Foundation /// /// The body runs inline in the caller's own task rather than in a detached one, so /// cancellation still reaches the work. +/// +/// A turn is owned by the ticket that took it, not by the connection. A drain ends the turn of a +/// holder that may still be stuck in a driver call, and when that holder finally returns, its +/// release must not free or hand off a turn a later session has taken since. @MainActor final class SessionDriverGate { private struct Waiter { @@ -20,15 +24,15 @@ final class SessionDriverGate { let continuation: CheckedContinuation } - private var holders: Set = [] + private var owners: [UUID: UUID] = [:] private var waiters: [UUID: [Waiter]] = [:] func withExclusiveAccess( _ connectionId: UUID, _ body: () async throws -> T ) async throws -> T { - try await acquire(connectionId) - defer { release(connectionId) } + let ticket = try await acquire(connectionId) + defer { release(connectionId, ticket: ticket) } return try await body() } @@ -42,19 +46,19 @@ final class SessionDriverGate { /// Releases a connection that is going away, failing everyone still queued for it. func drain(connectionId: UUID) { - holders.remove(connectionId) + owners.removeValue(forKey: connectionId) let pending = waiters.removeValue(forKey: connectionId) ?? [] for waiter in pending { waiter.continuation.resume(throwing: CancellationError()) } } - private func acquire(_ connectionId: UUID) async throws { - guard holders.contains(connectionId) else { - holders.insert(connectionId) - return - } + private func acquire(_ connectionId: UUID) async throws -> UUID { let ticket = UUID() + guard owners[connectionId] != nil else { + owners[connectionId] = ticket + return ticket + } try await withTaskCancellationHandler( operation: { try await enqueue(ticket: ticket, connectionId: connectionId) }, onCancel: { [weak self] in @@ -63,6 +67,7 @@ final class SessionDriverGate { } } ) + return ticket } private func enqueue(ticket: UUID, connectionId: UUID) async throws { @@ -90,14 +95,16 @@ final class SessionDriverGate { waiter.continuation.resume(throwing: CancellationError()) } - private func release(_ connectionId: UUID) { + private func release(_ connectionId: UUID, ticket: UUID) { + guard owners[connectionId] == ticket else { return } guard var pending = waiters[connectionId], !pending.isEmpty else { - holders.remove(connectionId) + owners.removeValue(forKey: connectionId) waiters.removeValue(forKey: connectionId) return } let next = pending.removeFirst() waiters[connectionId] = pending.isEmpty ? nil : pending + owners[connectionId] = next.ticket next.continuation.resume() } } diff --git a/TablePro/Core/Database/DatabaseManager+Sessions.swift b/TablePro/Core/Database/DatabaseManager+Sessions.swift index d36a02edf4..d080c7127b 100644 --- a/TablePro/Core/Database/DatabaseManager+Sessions.swift +++ b/TablePro/Core/Database/DatabaseManager+Sessions.swift @@ -335,7 +335,12 @@ extension DatabaseManager { try await reconnectOntoDatabase(database, for: connectionId) } else if let adapter = driver as? PluginDriverAdapter { let grouping = pm?.schema.databaseGroupingStrategy ?? .byDatabase + let sessionStartedAt = session(for: connectionId)?.connectedAt try await sessionDriverGate.withExclusiveAccess(connectionId) { + try Task.checkCancellation() + guard session(for: connectionId)?.connectedAt == sessionStartedAt else { + throw CancellationError() + } try await adapter.switchDatabase(to: database) if grouping == .bySchema { await resetSchema(on: adapter, to: pm?.schema.defaultSchemaName) @@ -376,10 +381,10 @@ extension DatabaseManager { /// one the reconnect was about to disconnect. Counting it as an operation keeps the monitor's /// ping and a waiting lease's verification off the driver while it is being replaced. /// - /// A switch can now wait for its turn, and a disconnect does not fail what is waiting, so the - /// session it was asked on is checked once the turn comes. A connection closed and opened again - /// in between is a new session, and moving it would switch, or disconnect, a session nobody - /// asked this of. + /// A switch can now wait for its turn. A disconnect fails what is still queued, but a turn + /// handed over just before the session went away has already left the queue, so the session it + /// was asked on is checked once the turn comes. A connection closed and opened again in between + /// is a new session, and moving it would switch, or disconnect, a session nobody asked this of. private func reconnectOntoDatabase(_ database: String, for connectionId: UUID) async throws { let sessionStartedAt = session(for: connectionId)?.connectedAt try await sessionDriverGate.withExclusiveAccess(connectionId) { @@ -756,8 +761,11 @@ extension DatabaseManager { userRequestedDisconnects.contains(connectionId) } + /// Drains the driver gate in the same step the entry goes, so nothing still queued for this + /// session wakes to find a reopened one under the same id and runs there. internal func removeSessionEntry(for connectionId: UUID) { activeSessions.removeValue(forKey: connectionId) + sessionDriverGate.drain(connectionId: connectionId) connectionStatusVersions.removeValue(forKey: connectionId) forgetVerification(for: connectionId) AppEvents.shared.connectionStatusChanged.send( diff --git a/TablePro/ViewModels/RedisKeyTreeViewModel.swift b/TablePro/ViewModels/RedisKeyTreeViewModel.swift index 9c00898b61..e0af9c2b6a 100644 --- a/TablePro/ViewModels/RedisKeyTreeViewModel.swift +++ b/TablePro/ViewModels/RedisKeyTreeViewModel.swift @@ -32,13 +32,17 @@ internal final class RedisKeyTreeViewModel { isTruncated = false defer { isLoading = false } - guard let driver = DatabaseManager.shared.driver(for: connectionId) else { + guard DatabaseManager.shared.driver(for: connectionId) != nil else { clear() return } + let scope = DatabaseScope(connectionId: connectionId, database: database, schema: nil) + let limit = Self.maxKeys do { - let result = try await driver.execute(query: "KEYTREE LIMIT \(Self.maxKeys)") + let result = try await DatabaseManager.shared.withMetadataDriver(scope: scope) { driver in + try await driver.execute(query: "KEYTREE LIMIT \(limit)") + } let keyColumnIndex = result.columns.firstIndex(of: "Key") ?? 0 let typeColumnIndex = result.columns.firstIndex(of: "Type") ?? 1 diff --git a/TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift b/TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift index 6688d1500c..b3e77b30c3 100644 --- a/TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift +++ b/TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift @@ -722,10 +722,9 @@ extension MainContentCoordinator { redisDatabaseSwitchTask = Task { [weak self] in guard let self else { return } do { - if let adapter = DatabaseManager.shared.driver(for: connId) as? PluginDriverAdapter { - try await adapter.switchDatabase(to: String(dbIndex)) - } + try await DatabaseManager.shared.switchDatabase(to: database, for: connId, persist: false) } catch { + guard !DatabaseCancellationDiagnosis.isCancellation(error) else { return } if !Task.isCancelled { navigationLogger.error("Failed to SELECT Redis db\(dbIndex): \(error.localizedDescription, privacy: .public)") } @@ -735,9 +734,6 @@ extension MainContentCoordinator { return } guard !Task.isCancelled else { return } - DatabaseManager.shared.updateSession(connId) { session in - session.browseDatabase = database - } toolbarState.currentDatabase = database executeTableTabQueryDirectly() diff --git a/TableProTests/Core/Concurrency/SessionDriverGateTests.swift b/TableProTests/Core/Concurrency/SessionDriverGateTests.swift index 4bbae3d9a3..75934649ae 100644 --- a/TableProTests/Core/Concurrency/SessionDriverGateTests.swift +++ b/TableProTests/Core/Concurrency/SessionDriverGateTests.swift @@ -63,6 +63,15 @@ private func drainMainActor(_ times: Int = 8) async { } } +/// Waits for callers to queue. The bound is there so a caller that never queues fails the +/// assertions after it rather than hanging the suite. +@MainActor +private func waitForWaiters(_ count: Int, on gate: SessionDriverGate, _ connectionId: UUID) async { + for _ in 0..<10_000 where gate.waiterCount(for: connectionId) < count { + await Task.yield() + } +} + @Suite("SessionDriverGate") @MainActor struct SessionDriverGateTests { @@ -219,6 +228,107 @@ struct SessionDriverGateTests { try await holder.value } + /// A drain ends the turn of a holder that can still be stuck in a driver call. When that call + /// finally returns, its release belongs to a turn that is already over. + @Test("A drained holder's late release does not hand a later holder's turn to the next caller") + func drainedReleaseDoesNotHandOffALaterTurn() async throws { + let gate = SessionDriverGate() + let connectionId = UUID() + let log = EventLog() + let drainedEntered = TestSignal() + let releaseDrained = TestSignal() + + let drained = Task { @MainActor in + try await gate.withExclusiveAccess(connectionId) { + drainedEntered.signal() + await releaseDrained.wait() + } + } + await drainedEntered.wait() + gate.drain(connectionId: connectionId) + + let laterEntered = TestSignal() + let releaseLater = TestSignal() + let later = Task { @MainActor in + try await gate.withExclusiveAccess(connectionId) { + log.record("later-start") + laterEntered.signal() + await releaseLater.wait() + log.record("later-end") + } + } + await laterEntered.wait() + + let queued = Task { @MainActor in + try await gate.withExclusiveAccess(connectionId) { + log.record("queued-start") + } + } + await waitForWaiters(1, on: gate, connectionId) + #expect(gate.waiterCount(for: connectionId) == 1) + + releaseDrained.signal() + try await drained.value + + #expect(gate.waiterCount(for: connectionId) == 1) + #expect(log.events == ["later-start"]) + + releaseLater.signal() + try await later.value + try await queued.value + + #expect(log.events == ["later-start", "later-end", "queued-start"]) + } + + @Test("A drained holder's late release does not free a later holder's turn") + func drainedReleaseDoesNotFreeALaterTurn() async throws { + let gate = SessionDriverGate() + let connectionId = UUID() + let log = EventLog() + let drainedEntered = TestSignal() + let releaseDrained = TestSignal() + + let drained = Task { @MainActor in + try await gate.withExclusiveAccess(connectionId) { + drainedEntered.signal() + await releaseDrained.wait() + } + } + await drainedEntered.wait() + gate.drain(connectionId: connectionId) + + let laterEntered = TestSignal() + let releaseLater = TestSignal() + let later = Task { @MainActor in + try await gate.withExclusiveAccess(connectionId) { + log.record("later-start") + laterEntered.signal() + await releaseLater.wait() + log.record("later-end") + } + } + await laterEntered.wait() + + releaseDrained.signal() + try await drained.value + + let arriving = Task { @MainActor in + try await gate.withExclusiveAccess(connectionId) { + log.record("arriving-start") + } + } + await waitForWaiters(1, on: gate, connectionId) + + #expect(gate.waiterCount(for: connectionId) == 1) + #expect(log.events == ["later-start"]) + + releaseLater.signal() + try await later.value + try await arriving.value + + #expect(log.events == ["later-start", "later-end", "arriving-start"]) + } + @Test("The body observes its own task's cancellation") func bodyObservesCallerCancellation() async throws { let gate = SessionDriverGate() diff --git a/TableProTests/Core/Database/DatabaseManagerDisconnectTests.swift b/TableProTests/Core/Database/DatabaseManagerDisconnectTests.swift index 01ed814862..87485b5202 100644 --- a/TableProTests/Core/Database/DatabaseManagerDisconnectTests.swift +++ b/TableProTests/Core/Database/DatabaseManagerDisconnectTests.swift @@ -103,4 +103,73 @@ struct DatabaseManagerDisconnectTests { #expect(DatabaseManager.shared.activeSessions[id] == nil) } + + /// A driver stuck in a call keeps its turn across a disconnect. Work queued behind it has to end + /// with the session, or it wakes when that call returns and runs on whatever holds the id then. + @Test("Disconnecting fails the work still queued for the session's driver") + func disconnectFailsQueuedDriverWork() async throws { + let connection = TestFixtures.makeConnection(name: "Queued") + var session = ConnectionSession(connection: connection, driver: MockDatabaseDriver(connection: connection)) + session.status = .connected + DatabaseManager.shared.injectSession(session, for: connection.id) + defer { DatabaseManager.shared.removeSession(for: connection.id) } + let gate = DatabaseManager.shared.sessionDriverGate + + let acquired = Latch() + let release = Latch() + let holder = Task { @MainActor in + try await gate.withExclusiveAccess(connection.id) { + acquired.open() + await release.wait() + } + } + await acquired.wait() + + let scope = DatabaseScope(connectionId: connection.id, database: connection.database, schema: nil) + let lease = Task { @MainActor in + try await DatabaseManager.shared.withScopedDriver( + scope: scope, + route: .sessionDriver, + cancellation: .cancellableRead + ) { driver in + driver.connection.database + } + } + for _ in 0..<10_000 where gate.waiterCount(for: connection.id) < 1 { + await Task.yield() + } + #expect(gate.waiterCount(for: connection.id) == 1) + + await DatabaseManager.shared.disconnectSession(connection.id) + + #expect(gate.waiterCount(for: connection.id) == 0) + + release.open() + try await holder.value + + await #expect(throws: CancellationError.self) { + try await lease.value + } + } +} + +@MainActor +private final class Latch { + private var waiters: [CheckedContinuation] = [] + private var isOpen = false + + func open() { + guard !isOpen else { return } + isOpen = true + let pending = waiters + waiters = [] + for waiter in pending { + waiter.resume() + } + } + + func wait() async { + guard !isOpen else { return } + await withCheckedContinuation { waiters.append($0) } + } } diff --git a/TableProTests/Core/Database/DatabaseSwitchLeaseOrderingTests.swift b/TableProTests/Core/Database/DatabaseSwitchLeaseOrderingTests.swift index d0420a53db..f10e342797 100644 --- a/TableProTests/Core/Database/DatabaseSwitchLeaseOrderingTests.swift +++ b/TableProTests/Core/Database/DatabaseSwitchLeaseOrderingTests.swift @@ -233,8 +233,8 @@ struct DatabaseSwitchLeaseOrderingTests { #expect(original.switchSchemaCallCount == 0) } - /// A disconnect does not fail what is waiting for the driver, so a switch still queued when the - /// connection is closed and opened again would otherwise move the session that replaced it. + /// A switch still queued when the connection is closed and opened again would otherwise move the + /// session that replaced it. @Test("A switch queued behind the driver is dropped when the session it was asked on has gone") func queuedSwitchIsDroppedForAReplacedSession() async throws { let connection = makeSession() @@ -272,4 +272,47 @@ struct DatabaseSwitchLeaseOrderingTests { #expect(reopened.disconnectCallCount == 0) #expect(reopened.switchSchemaCallCount == 0) } + + /// A lease that waited behind a holder stuck on the old session used to wake when that holder + /// returned, find a live session under the same id, and run the old tab's work there. + @Test("A lease queued when its session ends never runs on the session opened after it") + func queuedLeaseNeverRunsOnAReopenedSession() async throws { + let connection = makeSession() + defer { cleanUp(connection.id) } + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + let ran = LeaseRecord() + let app = DatabaseScope(connectionId: connection.id, database: "app", schema: nil) + let lease = Task { @MainActor in + try await DatabaseManager.shared.withScopedDriver( + scope: app, + route: .sessionDriver, + 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.finalizeConnectionFailure(for: connection.id, cancelled: false) + var session = ConnectionSession(connection: connection, driver: MockDatabaseDriver(connection: connection)) + session.status = .connected + session.browseDatabase = "app" + DatabaseManager.shared.injectSession(session, for: connection.id) + + release.open() + try await holder.value + + await #expect(throws: CancellationError.self) { + try await lease.value + } + #expect(!ran.didRun) + } +} + +@MainActor +private final class LeaseRecord { + var didRun = false } diff --git a/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift b/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift new file mode 100644 index 0000000000..8c0fad8341 --- /dev/null +++ b/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift @@ -0,0 +1,200 @@ +// +// RedisDatabaseSelectionGateTests.swift +// TableProTests +// +// Selecting a Redis database moves the connection's one shared driver. It used to send its SELECT +// around the driver gate, so the SELECT could land between a browse's SCAN and its TYPE/TTL +// pipeline, or in the middle of a key tree load, and both then read another database's keys. +// + +import Foundation +@testable import TablePro +import TableProPluginKit +import Testing + +@MainActor +private final class Latch { + private var waiters: [CheckedContinuation] = [] + private var isOpen = false + + func open() { + guard !isOpen else { return } + isOpen = true + let pending = waiters + waiters = [] + for waiter in pending { + waiter.resume() + } + } + + func wait() async { + guard !isOpen else { return } + await withCheckedContinuation { waiters.append($0) } + } +} + +@Suite("Redis database selection and the session driver gate", .serialized) +@MainActor +struct RedisDatabaseSelectionGateTests { + private func makeSession() -> (connection: DatabaseConnection, recorder: RecordingRedisPluginDriver) { + let connection = TestFixtures.makeConnection(database: "0", type: .redis) + let recorder = RecordingRedisPluginDriver() + var session = ConnectionSession( + connection: connection, + driver: PluginDriverAdapter(connection: connection, pluginDriver: recorder) + ) + session.status = .connected + session.browseDatabase = "0" + DatabaseManager.shared.injectSession(session, for: connection.id) + return (connection, recorder) + } + + private func makeCoordinator(for connection: DatabaseConnection) -> MainContentCoordinator { + MainContentCoordinator( + connection: connection, + tabManager: QueryTabManager(), + changeManager: DataChangeManager(), + toolbarState: ConnectionToolbarState() + ) + } + + private func cleanUp(_ connectionId: UUID) { + DatabaseManager.shared.removeSession(for: connectionId) + SharedSidebarState.removeConnection(connectionId) + } + + /// Holds the connection's driver until `release` opens, and returns only once it holds it. + private func holdDriver(_ connectionId: UUID, until release: Latch) async -> Task { + let acquired = Latch() + let holder = Task { @MainActor in + try await DatabaseManager.shared.sessionDriverGate.withExclusiveAccess(connectionId) { + acquired.open() + await release.wait() + } + } + await acquired.wait() + return holder + } + + /// The bound is there so a caller that never queues, which is the regression these tests guard, + /// fails the assertions after it rather than hanging the suite. + private func waitForQueuedCallers(_ count: Int, on connectionId: UUID) async { + for _ in 0..<10_000 where DatabaseManager.shared.sessionDriverGate.waiterCount(for: connectionId) < count { + await Task.yield() + } + } + + @Test("Selecting a database waits for the driver and switches once its turn comes") + func selectionWaitsForTheDriver() async throws { + let (connection, recorder) = makeSession() + defer { cleanUp(connection.id) } + let coordinator = makeCoordinator(for: connection) + defer { coordinator.teardown() } + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + coordinator.openTableTab("db2") + await waitForQueuedCallers(1, on: connection.id) + + #expect(DatabaseManager.shared.sessionDriverGate.waiterCount(for: connection.id) == 1) + #expect(recorder.switchedDatabases.isEmpty) + + release.open() + try await holder.value + await coordinator.redisDatabaseSwitchTask?.value + + #expect(recorder.switchedDatabases == ["2"]) + #expect(DatabaseManager.shared.session(for: connection.id)?.browseDatabase == "2") + } + + @Test("A selection superseded while it waits never moves the connection") + func supersededSelectionNeverSwitches() async throws { + let (connection, recorder) = makeSession() + defer { cleanUp(connection.id) } + let coordinator = makeCoordinator(for: connection) + defer { coordinator.teardown() } + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + coordinator.openTableTab("db2") + await waitForQueuedCallers(1, on: connection.id) + let superseded = coordinator.redisDatabaseSwitchTask + + coordinator.openTableTab("db3") + await superseded?.value + await waitForQueuedCallers(1, on: connection.id) + + release.open() + try await holder.value + await coordinator.redisDatabaseSwitchTask?.value + + #expect(recorder.switchedDatabases == ["3"]) + #expect(DatabaseManager.shared.session(for: connection.id)?.browseDatabase == "3") + } + + @Test("Loading the key tree waits for the driver") + func keyTreeLoadWaitsForTheDriver() async throws { + let (connection, recorder) = makeSession() + defer { cleanUp(connection.id) } + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + let viewModel = RedisKeyTreeViewModel() + let load = Task { @MainActor in + await viewModel.loadKeys(connectionId: connection.id, database: "2", separator: ":") + } + await waitForQueuedCallers(1, on: connection.id) + + #expect(DatabaseManager.shared.sessionDriverGate.waiterCount(for: connection.id) == 1) + #expect(recorder.executedQueries.isEmpty) + + release.open() + try await holder.value + await load.value + + #expect(recorder.executedQueries == ["KEYTREE LIMIT \(RedisKeyTreeViewModel.maxKeys)"]) + } +} + +/// Records the calls that move or read the connection. The ping answers, because Redis declares a +/// health monitor and a check before use would otherwise fail the session. +private final class RecordingRedisPluginDriver: PluginDatabaseDriver, @unchecked Sendable { + private let lock = NSLock() + private var switched: [String] = [] + private var executed: [String] = [] + + var switchedDatabases: [String] { + lock.withLock { switched } + } + + var executedQueries: [String] { + lock.withLock { executed } + } + + func ping() async throws {} + + func switchDatabase(to database: String) async throws { + lock.withLock { switched.append(database) } + } + + func execute(query: String) async throws -> PluginQueryResult { + lock.withLock { executed.append(query) } + return PluginQueryResult(columns: [], columnTypeNames: [], rows: [], rowsAffected: 0, executionTime: 0) + } + + func connect() async throws {} + func disconnect() {} + func fetchTables(schema: String?) async throws -> [PluginTableInfo] { [] } + func fetchColumns(table: String, schema: String?) async throws -> [PluginColumnInfo] { [] } + func fetchIndexes(table: String, schema: String?) async throws -> [PluginIndexInfo] { [] } + func fetchForeignKeys(table: String, schema: String?) async throws -> [PluginForeignKeyInfo] { [] } + func fetchTableDDL(table: String, schema: String?) async throws -> String { "" } + func fetchViewDefinition(view: String, schema: String?) async throws -> String { "" } + func fetchTableMetadata(table: String, schema: String?) async throws -> PluginTableMetadata { + PluginTableMetadata(tableName: table) + } + func fetchDatabases() async throws -> [String] { [] } + func fetchDatabaseMetadata(_ database: String) async throws -> PluginDatabaseMetadata { + PluginDatabaseMetadata(name: database) + } +} From 7b8a4338fd2838d68dfa7575b53a60bffe008f8a Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Sun, 13 Sep 2026 23:11:59 +0700 Subject: [PATCH 2/4] fix(connections): fail a handed-off gate turn after a drain and switch the driver installed when the turn comes --- .../Core/Concurrency/SessionDriverGate.swift | 5 +++ .../Database/DatabaseManager+Sessions.swift | 16 ++++--- .../Concurrency/SessionDriverGateTests.swift | 42 +++++++++++++++++++ .../RedisDatabaseSelectionGateTests.swift | 27 ++++++++++++ 4 files changed, 84 insertions(+), 6 deletions(-) diff --git a/TablePro/Core/Concurrency/SessionDriverGate.swift b/TablePro/Core/Concurrency/SessionDriverGate.swift index 7da213e1e7..35aec640ab 100644 --- a/TablePro/Core/Concurrency/SessionDriverGate.swift +++ b/TablePro/Core/Concurrency/SessionDriverGate.swift @@ -67,6 +67,11 @@ final class SessionDriverGate { } } ) + /// A hand-off resumes this caller before it runs, so a drain can land in between, and the + /// turn it was handed ended with that drain. + guard owners[connectionId] == ticket else { + throw CancellationError() + } return ticket } diff --git a/TablePro/Core/Database/DatabaseManager+Sessions.swift b/TablePro/Core/Database/DatabaseManager+Sessions.swift index d080c7127b..16a181d885 100644 --- a/TablePro/Core/Database/DatabaseManager+Sessions.swift +++ b/TablePro/Core/Database/DatabaseManager+Sessions.swift @@ -333,18 +333,22 @@ extension DatabaseManager { if pm?.capabilities.requiresReconnectForDatabaseSwitch == true { try await reconnectOntoDatabase(database, for: connectionId) - } else if let adapter = driver as? PluginDriverAdapter { + } else if driver is PluginDriverAdapter { let grouping = pm?.schema.databaseGroupingStrategy ?? .byDatabase let sessionStartedAt = session(for: connectionId)?.connectedAt - try await sessionDriverGate.withExclusiveAccess(connectionId) { + let adapter = try await sessionDriverGate.withExclusiveAccess(connectionId) { try Task.checkCancellation() guard session(for: connectionId)?.connectedAt == sessionStartedAt else { throw CancellationError() } + guard let adapter = self.driver(for: connectionId) as? PluginDriverAdapter else { + throw DatabaseError.notConnected + } try await adapter.switchDatabase(to: database) if grouping == .bySchema { await resetSchema(on: adapter, to: pm?.schema.defaultSchemaName) } + return adapter } updateSession(connectionId) { session in session.browseDatabase = database @@ -381,10 +385,10 @@ extension DatabaseManager { /// one the reconnect was about to disconnect. Counting it as an operation keeps the monitor's /// ping and a waiting lease's verification off the driver while it is being replaced. /// - /// A switch can now wait for its turn. A disconnect fails what is still queued, but a turn - /// handed over just before the session went away has already left the queue, so the session it - /// was asked on is checked once the turn comes. A connection closed and opened again in between - /// is a new session, and moving it would switch, or disconnect, a session nobody asked this of. + /// A switch can now wait for its turn. A disconnect fails every caller still waiting for one, + /// and the session the switch was asked on is checked again once the turn comes: a connection + /// closed and opened again is a new session, and moving it would switch, or disconnect, a + /// session nobody asked this of. private func reconnectOntoDatabase(_ database: String, for connectionId: UUID) async throws { let sessionStartedAt = session(for: connectionId)?.connectedAt try await sessionDriverGate.withExclusiveAccess(connectionId) { diff --git a/TableProTests/Core/Concurrency/SessionDriverGateTests.swift b/TableProTests/Core/Concurrency/SessionDriverGateTests.swift index 75934649ae..02e4548816 100644 --- a/TableProTests/Core/Concurrency/SessionDriverGateTests.swift +++ b/TableProTests/Core/Concurrency/SessionDriverGateTests.swift @@ -329,6 +329,48 @@ struct SessionDriverGateTests { #expect(log.events == ["later-start", "later-end", "arriving-start"]) } + /// The holder drains in the same synchronous step as its own release, which is after the hand-off + /// resumed the next caller and before that caller has run. + @Test("A drain fails a caller that was handed the gate but has not started") + func drainFailsAHandedOffCallerThatHasNotStarted() async throws { + let gate = SessionDriverGate() + let connectionId = UUID() + let log = EventLog() + let holderEntered = TestSignal() + let releaseHolder = TestSignal() + + let holder = Task { @MainActor in + try await gate.withExclusiveAccess(connectionId) { + holderEntered.signal() + await releaseHolder.wait() + } + gate.drain(connectionId: connectionId) + } + await holderEntered.wait() + + let handedOff = Task { @MainActor in + try await gate.withExclusiveAccess(connectionId) { + log.record("handed-off-ran") + } + } + await waitForWaiters(1, on: gate, connectionId) + #expect(gate.waiterCount(for: connectionId) == 1) + + releaseHolder.signal() + try await holder.value + + await #expect(throws: CancellationError.self) { + try await handedOff.value + } + #expect(log.events.isEmpty) + + let ran = BoolBox() + try await gate.withExclusiveAccess(connectionId) { + ran.value = true + } + #expect(ran.value) + } + @Test("The body observes its own task's cancellation") func bodyObservesCallerCancellation() async throws { let gate = SessionDriverGate() diff --git a/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift b/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift index 8c0fad8341..2c842618e9 100644 --- a/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift +++ b/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift @@ -132,6 +132,33 @@ struct RedisDatabaseSelectionGateTests { #expect(DatabaseManager.shared.session(for: connection.id)?.browseDatabase == "3") } + /// A reconnect replaces the driver on the same session, so the session check alone cannot tell + /// that the handle the selection was asked on is gone. + @Test("A selection queued behind the driver switches the driver installed when its turn comes") + func selectionSwitchesTheDriverInstalledWhenItRuns() async throws { + let (connection, original) = makeSession() + defer { cleanUp(connection.id) } + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + let selection = Task { @MainActor in + try await DatabaseManager.shared.switchDatabase(to: "2", for: connection.id, persist: false) + } + await waitForQueuedCallers(1, on: connection.id) + #expect(DatabaseManager.shared.sessionDriverGate.waiterCount(for: connection.id) == 1) + + let replacement = RecordingRedisPluginDriver() + DatabaseManager.shared.updateSession(connection.id) { session in + session.driver = PluginDriverAdapter(connection: connection, pluginDriver: replacement) + } + release.open() + try await holder.value + try await selection.value + + #expect(replacement.switchedDatabases == ["2"]) + #expect(original.switchedDatabases.isEmpty) + } + @Test("Loading the key tree waits for the driver") func keyTreeLoadWaitsForTheDriver() async throws { let (connection, recorder) = makeSession() From 18db581b02fdc159e32b3183de5feb466e24265a Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Mon, 14 Sep 2026 00:12:27 +0700 Subject: [PATCH 3/4] fix(connections): decline a Redis table load when its selection fails without being superseded --- .../MainContentCoordinator+Navigation.swift | 6 ++-- .../RedisDatabaseSelectionGateTests.swift | 29 +++++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift b/TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift index b3e77b30c3..9f9f9e8513 100644 --- a/TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift +++ b/TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift @@ -724,10 +724,8 @@ extension MainContentCoordinator { do { try await DatabaseManager.shared.switchDatabase(to: database, for: connId, persist: false) } catch { - guard !DatabaseCancellationDiagnosis.isCancellation(error) else { return } - if !Task.isCancelled { - navigationLogger.error("Failed to SELECT Redis db\(dbIndex): \(error.localizedDescription, privacy: .public)") - } + guard !Task.isCancelled else { return } + navigationLogger.error("Failed to SELECT Redis db\(dbIndex): \(error.localizedDescription, privacy: .public)") if let tabId = tabManager.selectedTab?.id { declineTableLoad(for: tabId) } diff --git a/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift b/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift index 2c842618e9..b3165eaa08 100644 --- a/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift +++ b/TableProTests/Views/Main/RedisDatabaseSelectionGateTests.swift @@ -132,6 +132,35 @@ struct RedisDatabaseSelectionGateTests { #expect(DatabaseManager.shared.session(for: connection.id)?.browseDatabase == "3") } + /// Only a newer selection owns the tab's load. A selection that fails for any other reason, + /// including one drained by a disconnect, has to give the load back or the tab keeps its spinner. + @Test("A selection drained while it waits gives the tab's load back") + func drainedSelectionDeclinesTheLoad() async throws { + let (connection, recorder) = makeSession() + defer { cleanUp(connection.id) } + let coordinator = makeCoordinator(for: connection) + defer { coordinator.teardown() } + let release = Latch() + let holder = await holdDriver(connection.id, until: release) + + coordinator.openTableTab("db2") + await waitForQueuedCallers(1, on: connection.id) + let superseded = coordinator.redisDatabaseSwitchTask + coordinator.openTableTab("db3") + await superseded?.value + await waitForQueuedCallers(1, on: connection.id) + #expect(coordinator.tabManager.selectedTab?.pagination.isLoading == true) + + DatabaseManager.shared.removeSession(for: connection.id) + await coordinator.redisDatabaseSwitchTask?.value + + #expect(coordinator.tabManager.selectedTab?.pagination.isLoading == false) + #expect(recorder.switchedDatabases.isEmpty) + + release.open() + try await holder.value + } + /// A reconnect replaces the driver on the same session, so the session check alone cannot tell /// that the handle the selection was asked on is gone. @Test("A selection queued behind the driver switches the driver installed when its turn comes") From ccc47dfd82a2d303f1cf852a1aa147f6187d828d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ng=C3=B4=20Qu=E1=BB=91c=20=C4=90=E1=BA=A1t?= Date: Mon, 14 Sep 2026 06:20:47 +0700 Subject: [PATCH 4/4] fix(connections): keep table loads working while a database switch reconnects (#2826) * fix(connections): keep table loads working while a database switch reconnects * fix(connections): hold the pool only while a tunnel recovery connects and fail parked leases on close * fix(connections): leave tunnel recovery's pooled connections as they were --- CHANGELOG.md | 2 + .../Coordinators/PaginationCoordinator.swift | 9 +- .../Database/DatabaseManager+Health.swift | 18 ++ .../DatabaseManager+ScopedDriver.swift | 130 +++++++-- .../Query/DatabaseTreeMetadataService.swift | 6 +- .../Query/MetadataConnectionPool.swift | 261 ++++++++++++++---- .../MainContentCoordinator+QueryHelpers.swift | 24 ++ .../Views/Main/MainContentCoordinator.swift | 6 +- .../DatabaseSwitchLeaseOrderingTests.swift | 221 ++++++++++++++- .../SwitchDatabasePooledConnectionTests.swift | 148 ++++++++++ .../Query/MetadataConnectionPoolTests.swift | 225 +++++++++++++++ 11 files changed, 952 insertions(+), 98 deletions(-) create mode 100644 TableProTests/Core/Database/SwitchDatabasePooledConnectionTests.swift diff --git a/CHANGELOG.md b/CHANGELOG.md index 53f7b93a0e..238fd68808 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 6ed9207832..d17c6b363b 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 b52e07fffe..0447558932 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 005599a705..62a76b679e 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/Services/Query/DatabaseTreeMetadataService.swift b/TablePro/Core/Services/Query/DatabaseTreeMetadataService.swift index 5a305cd36d..658c0b16e0 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 6b3435ab2e..8d52bd1bb3 100644 --- a/TablePro/Core/Services/Query/MetadataConnectionPool.swift +++ b/TablePro/Core/Services/Query/MetadataConnectionPool.swift @@ -56,11 +56,42 @@ 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 scope: DatabaseScope + 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 +111,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 +133,41 @@ 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 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 } + } + + /// 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 +203,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 +236,83 @@ 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() + } + 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 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(for: scope) + 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 +320,77 @@ 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(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, scope: scope) }, + onCancel: { [weak self] in + Task { @MainActor in + self?.failTransportWaiter(ticket: ticket, connectionId: connectionId) + } + } + ) + } + } + + 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[scope.connectionId, default: []].append( + TransportWaiter(ticket: ticket, scope: scope, 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 +415,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 55110842cd..a920f28fb3 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 fa9b110f19..5f9986e539 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 f10e342797..9d3e85ccc3 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 0000000000..552b3dc535 --- /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 3ab20df493..dc0436fb46 100644 --- a/TableProTests/Core/Services/Query/MetadataConnectionPoolTests.swift +++ b/TableProTests/Core/Services/Query/MetadataConnectionPoolTests.swift @@ -151,6 +151,231 @@ 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) + } + + /// 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) @MainActor struct MetadataConnectionPoolIdleEvictionTests {