Join gc and new window detection into a single request to MacApp & make it parallel

Reduce thread context switching

Author: nikitabobko

Sources/AppBundle/command/impl/CloseCommand.swift

@@ -12,8 +12,13 @@ struct CloseCommand: Command {
             }
             // Access ax directly. Not cool :(
             if try await args.quitIfLastWindow.andAsyncMainActor(try await window.macAppUnsafe.getAxWindowsCount() == 1) {
-                if window.macAppUnsafe.nsApp.terminate() {
-                    window.asMacWindow().macApp.destroy(skipClosedWindowsCache: true)
+                let app = window.macAppUnsafe
+                if app.nsApp.terminate() {
+                    for workspace in Workspace.all {
+                        for window in workspace.allLeafWindowsRecursive where window.app.pid == app.pid {
+                            (window as! MacWindow).garbageCollect(skipClosedWindowsCache: true)
+                        }
+                    }
                     return true
                 } else {
                     return io.err("Failed to quit '\(window.app.name ?? "Unknown app")'")

Sources/AppBundle/getNativeFocusedWindow.swift

@@ -11,7 +11,11 @@ private var focusedApp: (any AbstractApp)? {
             return appForTests
         } else {
             check(appForTests == nil)
-            return try await NSWorkspace.shared.frontmostApplication?.macApp
+            if let frontmostApplication = NSWorkspace.shared.frontmostApplication {
+                return try await MacApp.getOrRegister(frontmostApplication)
+            } else {
+                return nil
+            }
         }
     }
 }

Sources/AppBundle/layout/refresh.swift

@@ -18,9 +18,8 @@ func runRefreshSession(
 func runRefreshSessionBlocking(_ event: RefreshSessionEvent) async throws {
     if !TrayMenuModel.shared.isEnabled { return }
     try await $refreshSessionEventForDebug.withValue(event) {
-        try await gc()
+        try await refresh()
         gcMonitors()
-        try await detectNewAppsAndWindows()
 
         let nativeFocused = try await getNativeFocusedWindow()
         if let nativeFocused { try await debugWindowsIfRecording(nativeFocused) }
@@ -87,18 +86,19 @@ func refreshModel() async throws {
 }
 
 @MainActor
-private func gc() async throws {
+private func refresh() async throws {
     // Garbage collect terminated apps and windows before working with all windows
-    MacApp.gcTerminatedApps()
+    let mapping = try await MacApp.refreshAllAndGetAliveWindowIds(frontmostAppBundleId: NSWorkspace.shared.frontmostApplication?.bundleIdentifier)
+    let aliveWindowIds = mapping.values.flatMap { $0 }
 
-    let frontmostAppBundleId = NSWorkspace.shared.frontmostApplication?.bundleIdentifier
-    var aliveIds: Set<UInt32> = []
-    for (_, app) in MacApp.allAppsMap {
-        aliveIds.formUnion(try await app.gcDeadWindowsAndGetAliveIds(frontmostAppBundleId: frontmostAppBundleId))
-    }
     for window in MacWindow.allWindows {
-        if !aliveIds.contains(window.windowId) {
-            window.garbageCollect(skipClosedWindowsCache: false, unregisterAxWindow: false)
+        if !aliveWindowIds.contains(window.windowId) {
+            window.garbageCollect(skipClosedWindowsCache: false)
+        }
+    }
+    for (app, windowIds) in mapping {
+        for windowId in windowIds {
+            try await MacWindow.getOrRegister(windowId: windowId, macApp: app)
         }
     }
 
@@ -171,12 +171,3 @@ private func normalizeContainers() {
         workspace.normalizeContainers()
     }
 }
-
-@MainActor
-private func detectNewAppsAndWindows() async throws {
-    for app in try await detectNewApps() { // todo parallelize
-        for id in try await app.detectNewWindowsAndGetIds() {
-            _ = try await MacWindow.getOrRegister(windowId: id, macApp: app as! MacApp)
-        }
-    }
-}

Sources/AppBundle/tree/AbstractApp.swift

@@ -8,7 +8,6 @@ protocol AbstractApp: AnyObject, Hashable, AeroAny {
     var name: String? { get }
     var execPath: String? { get }
     var bundlePath: String? { get }
-    @MainActor func detectNewWindowsAndGetIds() async throws -> [UInt32]
 }
 
 extension AbstractApp {

Sources/AppBundle/tree/MacApp.swift

@@ -8,8 +8,8 @@ final class MacApp: AbstractApp {
     /*conforms*/ let pid: Int32
     /*conforms*/ let bundleId: String?
     let nsApp: NSRunningApplication
-    private let axApp: ThreadGuardedValue<AXUIElement>
     let isZoom: Bool
+    private let axApp: ThreadGuardedValue<AXUIElement>
     private let appAxSubscriptions: ThreadGuardedValue<[AxSubscription]> // keep subscriptions in memory
     private let windows: ThreadGuardedValue<[UInt32: AxWindow]> = .init([:])
     private var thread: Thread?
@@ -34,7 +34,8 @@ final class MacApp: AbstractApp {
     }
 
     @MainActor
-    static func get(_ nsApp: NSRunningApplication) async throws -> MacApp? {
+    @discardableResult
+    static func getOrRegister(_ nsApp: NSRunningApplication) async throws -> MacApp? {
         // Don't perceive any of the lock screen windows as real windows
         // Otherwise, false positive ax notifications might trigger that lead to gcWindows
         if nsApp.bundleIdentifier == lockScreenAppBundleId {
@@ -236,68 +237,67 @@ final class MacApp: AbstractApp {
         } ?? ""
     }
 
-    @MainActor func detectNewWindowsAndGetIds() async throws -> [UInt32] {
-        try await thread?.runInLoop { [axApp, windows, nsApp] job in
-            guard let newWindows = axApp.threadGuarded.get(Ax.windowsAttr, signpostEvent: nsApp.idForDebug) else { return Array(windows.threadGuarded.keys) }
-            var result: [UInt32] = []
-            for window in newWindows {
-                try job.checkCancellation()
-                if let windowId = windows.threadGuarded.getOrRegisterAxWindow(window, nsApp)?.windowId {
-                    result.append(windowId)
+    @MainActor
+    static func refreshAllAndGetAliveWindowIds(frontmostAppBundleId: String?) async throws -> [MacApp: [UInt32]] {
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            // Register new apps
+            for nsApp in NSWorkspace.shared.runningApplications where nsApp.activationPolicy == .regular {
+                try group.addTaskOrCancelAll { @Sendable @MainActor in
+                    _ = try await getOrRegister(nsApp)
+                }
+            }
+            try await group.waitForAll()
+        }
+        return try await withThrowingTaskGroup(of: (pid_t, [UInt32]).self, returning: [MacApp: [UInt32]].self) { group in
+            // gc dead apps. refresh underlying windows
+            for (_, app) in MacApp.allAppsMap {
+                try group.addTaskOrCancelAll { @Sendable @MainActor in
+                    (app.pid, try await app.refreshAndGetAliveWindowIds(frontmostAppBundleId: frontmostAppBundleId))
+                }
+            }
+            var result: [MacApp: [UInt32]] = [:]
+            for try await (pid, windowIds) in group {
+                if let app = allAppsMap[pid] {
+                    result[app] = windowIds
                 }
             }
             return result
-        } ?? []
+        }
     }
 
     @MainActor
-    func gcDeadWindowsAndGetAliveIds(frontmostAppBundleId: String?) async throws -> Set<UInt32> {
-        try await thread?.runInLoop { [nsApp, windows] (job) -> Set<UInt32> in
+    private func refreshAndGetAliveWindowIds(frontmostAppBundleId: String?) async throws -> [UInt32] {
+        if nsApp.isTerminated {
+            MacApp.allAppsMap.removeValue(forKey: pid)
+            thread?.runInLoopAsync { [windows, appAxSubscriptions, axApp] job in
+                axApp.destroy()
+                appAxSubscriptions.destroy()
+                windows.destroy()
+                CFRunLoopStop(CFRunLoopGetCurrent())
+            }
+            thread = nil // Disallow all future job submissions
+            return []
+        }
+        guard let thread else { return [] }
+        return try await thread.runInLoop { [nsApp, windows, axApp] (job) -> [UInt32] in
+            var result: [UInt32: AxWindow] = windows.threadGuarded
             // Second line of defence against lock screen. See the first line of defence: closedWindowsCache
             // Second and third lines of defence are technically needed only to avoid potential flickering
-            let _windows: [UInt32: AxWindow] = windows.threadGuarded
-            if frontmostAppBundleId == lockScreenAppBundleId { return Set(_windows.keys) }
-            let toKeepAlive: [UInt32: AxWindow] = try _windows.filter {
-                try job.checkCancellation()
-                return $0.value.ax.containingWindowId(signpostEvent: nsApp.idForDebug) != nil
+            if frontmostAppBundleId != lockScreenAppBundleId {
+                result = try result.filter {
+                    try job.checkCancellation()
+                    return $0.value.ax.containingWindowId(signpostEvent: nsApp.idForDebug) != nil
+                }
             }
-            windows.threadGuarded = toKeepAlive
-            return Set(toKeepAlive.keys)
-        } ?? []
-    }
 
-    @MainActor
-    func unregisterWindow(_ windowId: UInt32) {
-        thread?.runInLoopAsync { [windows] job in
-            windows.threadGuarded.removeValue(forKey: windowId)
-        }
-    }
-
-    @MainActor
-    static func gcTerminatedApps() {
-        for app in allAppsMap.values where app.nsApp.isTerminated {
-            app.destroy(skipClosedWindowsCache: true)
-        }
-    }
+            for window in axApp.threadGuarded.get(Ax.windowsAttr, signpostEvent: nsApp.idForDebug) ?? [] {
+                try job.checkCancellation()
+                result.getOrRegisterAxWindow(window, nsApp)
+            }
 
-    @MainActor
-    func destroy(skipClosedWindowsCache: Bool) {
-        MacApp.allAppsMap.removeValue(forKey: nsApp.processIdentifier)
-        for (_, window) in MacWindow.allWindowsMap where window.app.pid == self.pid {
-            window.garbageCollect(skipClosedWindowsCache: skipClosedWindowsCache, unregisterAxWindow: false)
+            windows.threadGuarded = result
+            return Array(result.keys)
         }
-        thread?.runInLoopAsync { [windows, appAxSubscriptions, axApp] job in
-            axApp.destroy()
-            appAxSubscriptions.destroy()
-            windows.destroy()
-            CFRunLoopStop(CFRunLoopGetCurrent())
-        }
-        thread = nil // Disallow all future job submissions
-    }
-
-    private func getThreadOrCancel() throws -> Thread { // todo convert untyped throws to throws across the whole app
-        if let thread { return thread }
-        throw CancellationError()
     }
 
     @MainActor // todo swift is stupid
@@ -342,6 +342,7 @@ private class AxWindow {
 }
 
 extension [UInt32: AxWindow] {
+    @discardableResult
     fileprivate mutating func getOrRegisterAxWindow(_ axWindow: AXUIElement, _ nsApp: NSRunningApplication) -> AxWindow? {
         guard let id = axWindow.containingWindowId() else { return nil }
         if let existing = self[id] {
@@ -429,7 +430,4 @@ extension NSRunningApplication {
     var idForDebug: String {
         "PID: \(processIdentifier) ID: \(bundleIdentifier ?? executableURL?.description ?? "")"
     }
-
-    @MainActor
-    var macApp: MacApp? { get async throws { try await MacApp.get(self) } }
 }

Sources/AppBundle/tree/MacWindow.swift

@@ -16,6 +16,7 @@ final class MacWindow: Window {
     @MainActor static var allWindows: [MacWindow] { Array(allWindowsMap.values) }
 
     @MainActor
+    @discardableResult
     static func getOrRegister(windowId: UInt32, macApp: MacApp) async throws -> MacWindow {
         if let existing = allWindowsMap[windowId] { return existing }
         let rect = try await macApp.getAxRect(windowId)
@@ -79,7 +80,7 @@ final class MacWindow: Window {
     // skipClosedWindowsCache is an optimization when it's definitely not necessary to cache closed window.
     //                        If you are unsure, it's better to pass `false`
     @MainActor
-    func garbageCollect(skipClosedWindowsCache: Bool, unregisterAxWindow: Bool = true) {
+    func garbageCollect(skipClosedWindowsCache: Bool) {
         if MacWindow.allWindowsMap.removeValue(forKey: windowId) == nil {
             return
         }
@@ -104,9 +105,6 @@ final class MacWindow: Window {
                     break // Don't switch back on popup destruction
             }
         }
-        if unregisterAxWindow {
-            macApp.unregisterWindow(windowId)
-        }
     }
 
     @MainActor override var title: String { get async throws { try await macApp.getAxTitle(windowId) ?? "" } }
@@ -119,7 +117,7 @@ final class MacWindow: Window {
     }
 
     override func closeAxWindow() {
-        garbageCollect(skipClosedWindowsCache: true, unregisterAxWindow: false)
+        garbageCollect(skipClosedWindowsCache: true)
         macApp.closeAndUnregisterAxWindow(windowId)
     }
 

Sources/AppBundle/util/ThrowingTaskGroupEx.swift

@@ -0,0 +1,9 @@
+public extension ThrowingTaskGroup {
+    mutating func addTaskOrCancelAll(priority: TaskPriority? = nil, operation: sending @escaping @isolated(any) () async throws -> ChildTaskResult) throws {
+        let succ = addTaskUnlessCancelled(priority: priority, operation: operation)
+        if !succ {
+            cancelAll()
+            throw CancellationError()
+        }
+    }
+}

Sources/AppBundle/util/appBundleUtil.swift

@@ -59,20 +59,6 @@ extension String? {
     var isNilOrEmpty: Bool { self == nil || self?.isEmpty == true }
 }
 
-@MainActor
-func detectNewApps() async throws -> [any AbstractApp] {
-    if isUnitTest {
-        return appForTests.asList()
-    }
-    var result = [any AbstractApp]()
-    for _app in NSWorkspace.shared.runningApplications where _app.activationPolicy == .regular {
-        if let app = try await _app.macApp {
-            result.append(app)
-        }
-    }
-    return result
-}
-
 @MainActor
 func terminateApp() -> Never {
     NSApplication.shared.terminate(nil)