From d6820a555377d9509d7546ace586f738385f8ae3 Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Mon, 10 Aug 2026 11:47:58 +0700 Subject: [PATCH] fix(coordinator): give each tab its own row count task so one cannot orphan another --- CHANGELOG.md | 1 + .../Coordinators/PaginationCoordinator.swift | 8 +- .../QueryExecutionCoordinator+Helpers.swift | 5 +- ...MainContentCoordinator+RowCountTasks.swift | 32 ++++ .../Views/Main/MainContentCoordinator.swift | 7 +- .../MainContentCoordinatorRefreshTests.swift | 10 +- .../Main/RowCountTaskLifecycleTests.swift | 153 ++++++++++++++++++ 7 files changed, 201 insertions(+), 15 deletions(-) create mode 100644 TablePro/Views/Main/Extensions/MainContentCoordinator+RowCountTasks.swift create mode 100644 TableProTests/Views/Main/RowCountTaskLifecycleTests.swift diff --git a/CHANGELOG.md b/CHANGELOG.md index 54bea0cb5..ef0891ef8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Closing a window now stops an exact row count that is still running, and pointing a tab at another table stops the count for the table you left. Both used to run to completion against a server nobody was waiting on. (#2059) - Clicking a table while another one is still loading now loads the table you clicked. The tab used to fill with the previous table's rows under the new table's name, and the table you clicked never loaded at all. - Stop no longer aborts a save, a discard, or a schema change. It cancels reads only, so a write can no longer be cut off half way. - Ending a session while its window stayed open could lose that window's tabs. Tabs are now written to disk before the session goes away, which also covers the disconnect tool used by AI clients and Reset Sample Database. diff --git a/TablePro/Core/Coordinators/PaginationCoordinator.swift b/TablePro/Core/Coordinators/PaginationCoordinator.swift index a96127a55..0dfda890a 100644 --- a/TablePro/Core/Coordinators/PaginationCoordinator.swift +++ b/TablePro/Core/Coordinators/PaginationCoordinator.swift @@ -130,8 +130,7 @@ final class PaginationCoordinator { func cancelCurrentQuery() { parent.cancelInFlightQueryTask() - parent.currentRowCountTask?.cancel() - parent.currentRowCountTask = nil + parent.cancelAllRowCountTasks() parent.tabExecution.invalidateAll() parent.toolbarState.setExecuting(false) for idx in parent.tabManager.tabs.indices @@ -167,7 +166,7 @@ final class PaginationCoordinator { parent.tabManager.mutate(at: index) { $0.pagination.isCountingExact = true } let contentEpoch = parent.tabExecution.contentEpoch(for: tabId) - parent.currentRowCountTask = Task(priority: .userInitiated) { [parent] in + let task = Task(priority: .userInitiated) { [parent] in let count = await Self.exactRowCount( scope: scope, tableName: tableName, @@ -178,7 +177,7 @@ final class PaginationCoordinator { guard !Task.isCancelled else { return } guard parent.tabExecution.isSameContent(contentEpoch, for: tabId) else { return } - parent.currentRowCountTask = nil + parent.clearRowCountTask(for: tabId) parent.tabManager.mutate(tabId: tabId) { tab in tab.pagination.isCountingExact = false guard let count, count >= 0 else { return } @@ -186,6 +185,7 @@ final class PaginationCoordinator { tab.pagination.isApproximateRowCount = false } } + parent.setRowCountTask(task, for: tabId) } private static func exactRowCount( diff --git a/TablePro/Core/Coordinators/QueryExecutionCoordinator+Helpers.swift b/TablePro/Core/Coordinators/QueryExecutionCoordinator+Helpers.swift index 024b54529..ff562cf9a 100644 --- a/TablePro/Core/Coordinators/QueryExecutionCoordinator+Helpers.swift +++ b/TablePro/Core/Coordinators/QueryExecutionCoordinator+Helpers.swift @@ -415,7 +415,7 @@ extension QueryExecutionCoordinator { ) { let isNonSQL = PluginManager.shared.editorLanguage(for: connectionType) != .sql - parent.currentRowCountTask = Task(priority: .utility) { [weak self, parent] in + let task = Task(priority: .utility) { [weak self, parent] in guard let self else { return } guard !parent.isTearingDown else { return } @@ -488,7 +488,7 @@ extension QueryExecutionCoordinator { await MainActor.run { guard parent.tabExecution.isCurrent(claim) else { return } - parent.currentRowCountTask = nil + parent.clearRowCountTask(for: tabId) guard let outcome else { return } parent.tabManager.mutate(tabId: tabId) { tab in let applied = outcome.appliedTotal @@ -497,6 +497,7 @@ extension QueryExecutionCoordinator { } } } + parent.setRowCountTask(task, for: tabId) } static func rowCountPlan( diff --git a/TablePro/Views/Main/Extensions/MainContentCoordinator+RowCountTasks.swift b/TablePro/Views/Main/Extensions/MainContentCoordinator+RowCountTasks.swift new file mode 100644 index 000000000..104026a5e --- /dev/null +++ b/TablePro/Views/Main/Extensions/MainContentCoordinator+RowCountTasks.swift @@ -0,0 +1,32 @@ +// +// MainContentCoordinator+RowCountTasks.swift +// TablePro +// +// A row count belongs to the tab that asked for it. One shared handle per window meant the +// second tab to ask dropped the only reference to the first tab's task, leaving it running with +// nothing able to cancel it, not even teardown, while it held the coordinator alive. +// + +import Foundation + +extension MainContentCoordinator { + /// Cancels whatever this tab had running first, because a tab only ever wants its newest count. + internal func setRowCountTask(_ task: Task, for tabId: UUID) { + rowCountTasks[tabId]?.cancel() + rowCountTasks[tabId] = task + } + + /// Drops a finished task's handle. The task is already done, so there is nothing to cancel. + internal func clearRowCountTask(for tabId: UUID) { + rowCountTasks.removeValue(forKey: tabId) + } + + internal func cancelRowCountTask(for tabId: UUID) { + rowCountTasks.removeValue(forKey: tabId)?.cancel() + } + + internal func cancelAllRowCountTasks() { + for task in rowCountTasks.values { task.cancel() } + rowCountTasks.removeAll() + } +} diff --git a/TablePro/Views/Main/MainContentCoordinator.swift b/TablePro/Views/Main/MainContentCoordinator.swift index 8d3111ffb..ded4b15b3 100644 --- a/TablePro/Views/Main/MainContentCoordinator.swift +++ b/TablePro/Views/Main/MainContentCoordinator.swift @@ -210,10 +210,9 @@ final class MainContentCoordinator { /// property and invalidates its readers. internal var tabExecution = TabExecutionRegistry() @ObservationIgnored internal var currentQueryTask: Task? - @ObservationIgnored internal var currentRowCountTask: Task? + @ObservationIgnored internal var rowCountTasks: [UUID: Task] = [:] @ObservationIgnored internal var tableLoadTasks: [UUID: (token: UUID, task: Task)] = [:] @ObservationIgnored internal var redisDatabaseSwitchTask: Task? - @ObservationIgnored private var changeManagerUpdateTask: Task? @ObservationIgnored private var periodicSaveTask: Task? @ObservationIgnored private var draftSaveTask: Task? @ObservationIgnored private var terminationObserver: NSObjectProtocol? @@ -758,8 +757,7 @@ final class MainContentCoordinator { refreshCoalesceTask = nil for entry in tableLoadTasks.values { entry.task.cancel() } tableLoadTasks.removeAll() - changeManagerUpdateTask?.cancel() - changeManagerUpdateTask = nil + cancelAllRowCountTasks() periodicSaveTask?.cancel() periodicSaveTask = nil draftSaveTask?.cancel() @@ -1325,6 +1323,7 @@ final class MainContentCoordinator { internal func supersedeExecution(for tabId: UUID) { tabExecution.invalidate(tabId) cancelTableLoad(for: tabId) + cancelRowCountTask(for: tabId) cancelInFlightQueryTask(reach: .supersededNavigation) } diff --git a/TableProTests/Views/Main/MainContentCoordinatorRefreshTests.swift b/TableProTests/Views/Main/MainContentCoordinatorRefreshTests.swift index 4885384e2..4b69341b6 100644 --- a/TableProTests/Views/Main/MainContentCoordinatorRefreshTests.swift +++ b/TableProTests/Views/Main/MainContentCoordinatorRefreshTests.swift @@ -197,14 +197,14 @@ struct MainContentCoordinatorRefreshTests { @Test("A finished row count leaves no handle that fakes an in-flight query") func cancelWithStaleRowCountHandleDoesNotTouchDriver() { withInjectedDriver { connection, driver in - let (coordinator, _) = makeCoordinator(connection: connection) - let finishedRowCount = Task {} - coordinator.currentRowCountTask = finishedRowCount + let (coordinator, tabManager) = makeCoordinator(connection: connection) + let tabId = addTableTab(to: tabManager) + coordinator.setRowCountTask(Task {}, for: tabId) coordinator.cancelCurrentQuery() #expect(driver.cancelQueryCallCount == 0) - #expect(coordinator.currentRowCountTask == nil) + #expect(coordinator.rowCountTasks.isEmpty) } } @@ -220,7 +220,7 @@ struct MainContentCoordinatorRefreshTests { tabManager.tabs[idx].execution.lastExecutedAt = Date() for _ in 0..<4 { - coordinator.currentRowCountTask = Task {} + coordinator.setRowCountTask(Task {}, for: tabId) coordinator.handleRefresh(hasPendingTableOps: false, onDiscard: {}) coordinator.currentQueryTask?.cancel() coordinator.currentQueryTask = nil diff --git a/TableProTests/Views/Main/RowCountTaskLifecycleTests.swift b/TableProTests/Views/Main/RowCountTaskLifecycleTests.swift new file mode 100644 index 000000000..bc299270b --- /dev/null +++ b/TableProTests/Views/Main/RowCountTaskLifecycleTests.swift @@ -0,0 +1,153 @@ +// +// RowCountTaskLifecycleTests.swift +// TableProTests +// +// A row count is per tab, so its handle is too. One shared slot per window let the second tab to +// ask drop the first tab's task with nothing left able to cancel it, not even teardown (#2059). +// + +import Foundation +@testable import TablePro +import Testing + +@Suite("Row count task lifecycle") +@MainActor +struct RowCountTaskLifecycleTests { + @Test("A tab's second row count cancels its first") + func secondCountForTheSameTabCancelsTheFirst() { + let (coordinator, tabManager) = Self.makeCoordinator() + let tabId = Self.addTableTab(to: tabManager) + let first = Self.neverEndingTask() + defer { first.cancel() } + + coordinator.setRowCountTask(first, for: tabId) + coordinator.setRowCountTask(Self.neverEndingTask(), for: tabId) + + #expect(first.isCancelled) + #expect(coordinator.rowCountTasks.count == 1) + } + + /// The bug: one slot for the whole window meant tab B's count evicted tab A's handle without + /// cancelling it, so tab A's task ran on with no reference left to stop it. + @Test("A count on another tab leaves the first tab's count alone") + func countOnAnotherTabDoesNotDisturbTheFirst() { + let (coordinator, tabManager) = Self.makeCoordinator() + let tabA = Self.addTableTab(to: tabManager, tableName: "orders") + let tabB = Self.addTableTab(to: tabManager, tableName: "customers") + let countA = Self.neverEndingTask() + let countB = Self.neverEndingTask() + defer { + countA.cancel() + countB.cancel() + } + + coordinator.setRowCountTask(countA, for: tabA) + coordinator.setRowCountTask(countB, for: tabB) + + #expect(countA.isCancelled == false) + #expect(coordinator.rowCountTasks.count == 2) + } + + @Test("Retargeting a tab cancels only that tab's row count") + func supersedingATabCancelsOnlyItsOwnCount() { + let (coordinator, tabManager) = Self.makeCoordinator() + let tabA = Self.addTableTab(to: tabManager, tableName: "orders") + let tabB = Self.addTableTab(to: tabManager, tableName: "customers") + let countA = Self.neverEndingTask() + let countB = Self.neverEndingTask() + defer { + countA.cancel() + countB.cancel() + } + coordinator.setRowCountTask(countA, for: tabA) + coordinator.setRowCountTask(countB, for: tabB) + + coordinator.supersedeExecution(for: tabA) + + #expect(countA.isCancelled) + #expect(countB.isCancelled == false) + #expect(coordinator.rowCountTasks[tabA] == nil) + } + + /// Closing a window used to leave every in-flight count running, each holding the coordinator + /// alive through a strong capture until its own driver round trip returned. + @Test("Teardown cancels every tab's row count") + func teardownCancelsEveryRowCount() { + let (coordinator, tabManager) = Self.makeCoordinator() + let tabA = Self.addTableTab(to: tabManager, tableName: "orders") + let tabB = Self.addTableTab(to: tabManager, tableName: "customers") + let countA = Self.neverEndingTask() + let countB = Self.neverEndingTask() + coordinator.setRowCountTask(countA, for: tabA) + coordinator.setRowCountTask(countB, for: tabB) + + coordinator.teardown() + + #expect(countA.isCancelled) + #expect(countB.isCancelled) + #expect(coordinator.rowCountTasks.isEmpty) + } + + @Test("Stop cancels every tab's row count") + func stopCancelsEveryRowCount() { + let (coordinator, tabManager) = Self.makeCoordinator() + let tabA = Self.addTableTab(to: tabManager, tableName: "orders") + let tabB = Self.addTableTab(to: tabManager, tableName: "customers") + let countA = Self.neverEndingTask() + let countB = Self.neverEndingTask() + coordinator.setRowCountTask(countA, for: tabA) + coordinator.setRowCountTask(countB, for: tabB) + + coordinator.cancelCurrentQuery() + + #expect(countA.isCancelled) + #expect(countB.isCancelled) + #expect(coordinator.rowCountTasks.isEmpty) + } + + /// A task that finished on its own drops its handle without cancelling a successor that may + /// already have taken the slot. + @Test("Clearing a finished count never cancels anything") + func clearingAFinishedCountCancelsNothing() { + let (coordinator, tabManager) = Self.makeCoordinator() + let tabId = Self.addTableTab(to: tabManager) + let count = Self.neverEndingTask() + defer { count.cancel() } + coordinator.setRowCountTask(count, for: tabId) + + coordinator.clearRowCountTask(for: tabId) + + #expect(count.isCancelled == false) + #expect(coordinator.rowCountTasks.isEmpty) + } + + private static func makeCoordinator() -> (MainContentCoordinator, QueryTabManager) { + let tabManager = QueryTabManager() + let coordinator = MainContentCoordinator( + connection: TestFixtures.makeConnection(), + tabManager: tabManager, + changeManager: DataChangeManager(), + toolbarState: ConnectionToolbarState() + ) + return (coordinator, tabManager) + } + + private static func addTableTab( + to tabManager: QueryTabManager, + tableName: String = "users" + ) -> UUID { + let tab = QueryTab( + title: tableName, + query: "SELECT * FROM \(tableName)", + tabType: .table, + tableName: tableName + ) + tabManager.tabs.append(tab) + tabManager.selectedTabId = tab.id + return tab.id + } + + private static func neverEndingTask() -> Task { + Task { _ = try? await Task.sleep(for: .seconds(60)) } + } +}