fix(coordinator): give each tab its own row count task so one cannot orphan another

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 quickly through the sidebar no longer keeps loading column details for tables you have already left. Each abandoned table held the connection for another 75-90ms, so every table clicked after it waited that much longer. (#2058)
 - 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.

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,14 +177,15 @@ 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 }
                 tab.pagination.totalRowCount = count
                 tab.pagination.isApproximateRowCount = false
             }
         }
+        parent.setRowCountTask(task, for: tabId)
     }
 
     private static func exactRowCount(

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(

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<Void, Never>, 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()
+    }
+}

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<Void, Never>?
-    @ObservationIgnored internal var currentRowCountTask: Task<Void, Never>?
+    @ObservationIgnored internal var rowCountTasks: [UUID: Task<Void, Never>] = [:]
     @ObservationIgnored internal var tableLoadTasks: [UUID: (token: UUID, task: Task<Void, Never>)] = [:]
     @ObservationIgnored internal var redisDatabaseSwitchTask: Task<Void, Never>?
-    @ObservationIgnored private var changeManagerUpdateTask: Task<Void, Never>?
     @ObservationIgnored private var periodicSaveTask: Task<Void, Never>?
     @ObservationIgnored private var draftSaveTask: Task<Void, Never>?
     @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)
     }
 

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<Void, Never> {}
-            coordinator.currentRowCountTask = finishedRowCount
+            let (coordinator, tabManager) = makeCoordinator(connection: connection)
+            let tabId = addTableTab(to: tabManager)
+            coordinator.setRowCountTask(Task<Void, Never> {}, 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<Void, Never> {}
+                coordinator.setRowCountTask(Task<Void, Never> {}, for: tabId)
                 coordinator.handleRefresh(hasPendingTableOps: false, onDiscard: {})
                 coordinator.currentQueryTask?.cancel()
                 coordinator.currentQueryTask = nil

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<Void, Never> {
+        Task { _ = try? await Task.sleep(for: .seconds(60)) }
+    }
+}