⚡ Harden event tap lifecycle during shutdown

Author: mrkai77

Loop/App/AppDelegate.swift

@@ -139,8 +139,17 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
             return .terminateLater
         }
 
+        // LoopManager and WindowDragManager are explicitly shut down so that their
+        // event monitors are stopped immediately (in case they are active)
+        LoopManager.shared.shutdown()
+        WindowDragManager.shared.shutdown()
+
         shutdownTask = Task { @MainActor in
-            await StashManager.shared.shutdown()
+            let didFinishStashShutdown = await runStashShutdownWithTimeout(.seconds(3))
+            if !didFinishStashShutdown {
+                log.warn("Timed out while restoring stashed windows during termination. Continuing shutdown.")
+            }
+
             self.shutdownTask = nil
             sender.reply(toApplicationShouldTerminate: true)
         }
@@ -153,4 +162,40 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
             urlCommandHandler.handle(url)
         }
     }
+
+    private func runStashShutdownWithTimeout(_ duration: Duration) async -> Bool {
+        return await withCheckedContinuation { continuation in
+            let reply = OneShotContinuation(continuation)
+
+            let shutdownTask = Task { @MainActor in
+                await StashManager.shared.shutdown()
+                reply.resume(returning: true)
+            }
+
+            Task {
+                try? await Task.sleep(for: duration)
+                shutdownTask.cancel()
+                reply.resume(returning: false)
+            }
+        }
+    }
+}
+
+private final class OneShotContinuation<T>: @unchecked Sendable {
+    private let lock = NSLock()
+    private var didResume = false
+    private let continuation: CheckedContinuation<T, Never>
+
+    init(_ continuation: CheckedContinuation<T, Never>) {
+        self.continuation = continuation
+    }
+
+    func resume(returning result: T) {
+        lock.lock()
+        defer { lock.unlock() }
+
+        guard !didResume else { return }
+        didResume = true
+        continuation.resume(returning: result)
+    }
 }

Loop/Core/LoopManager.swift

@@ -6,6 +6,7 @@
 //
 
 import Defaults
+import os
 import Scribe
 import SwiftUI
 
@@ -25,7 +26,22 @@ final class LoopManager {
 
     private var accessibilityCheckerTask: Task<(), Never>?
 
-    private(set) var isLoopActive: Bool = false
+    private(set) var isLoopActive: Bool = false {
+        didSet {
+            let value = isLoopActive
+            isLoopActiveMirror.withLock { $0 = value }
+        }
+    }
+
+    private let isLoopActiveMirror = OSAllocatedUnfairLock<Bool>(initialState: false)
+    nonisolated var isLoopActiveAtomic: Bool {
+        isLoopActiveMirror.withLock { $0 }
+    }
+
+    private let hasParentCycleActionMirror = OSAllocatedUnfairLock<Bool>(initialState: false)
+    nonisolated var hasParentCycleActionAtomic: Bool {
+        hasParentCycleActionMirror.withLock { $0 }
+    }
 
     private lazy var triggerKeyTimeoutTimer = TriggerKeyTimeoutTimer(
         closeCallback: { [weak self] forceClose in
@@ -46,7 +62,7 @@ final class LoopManager {
             }
         },
         checkIfLoopOpen: { [weak self] in
-            self?.isLoopActive ?? false
+            self?.isLoopActiveAtomic ?? false
         }
     )
 
@@ -61,7 +77,7 @@ final class LoopManager {
                 await self?.closeLoop(forceClose: forceClose)
             }
         },
-        checkIfLoopOpen: { [weak self] in self?.isLoopActive ?? false }
+        checkIfLoopOpen: { [weak self] in self?.isLoopActiveAtomic ?? false }
     )
 
     private(set) lazy var mouseInteractionObserver = MouseInteractionObserver(
@@ -81,9 +97,9 @@ final class LoopManager {
             }
         },
         canSelectNextCycleitem: { [weak self] in
-            self?.resizeContext.parentAction != nil
+            self?.hasParentCycleActionAtomic ?? false
         },
-        checkIfLoopOpen: { [weak self] in self?.isLoopActive ?? false }
+        checkIfLoopOpen: { [weak self] in self?.isLoopActiveAtomic ?? false }
     )
 
     func start() {
@@ -103,6 +119,19 @@ final class LoopManager {
             }
         }
     }
+
+    func shutdown() {
+        accessibilityCheckerTask?.cancel()
+        accessibilityCheckerTask = nil
+
+        keybindTrigger.stop()
+        middleClickTrigger.stop()
+        mouseInteractionObserver.stop()
+        triggerKeyTimeoutTimer.cancel()
+
+        isLoopActive = false
+        hasParentCycleActionMirror.withLock { $0 = false }
+    }
 }
 
 // MARK: - Opening/Closing Loop
@@ -134,6 +163,9 @@ extension LoopManager {
             return
         }
 
+        isLoopActive = true
+        hasParentCycleActionMirror.withLock { $0 = false }
+
         log.info("Opening Loop with starting action: \(startingAction.description) and target window: \(window?.description ?? "(none)")")
 
         // Refresh accent colors in case user has enabled the wallpaper processor
@@ -163,7 +195,6 @@ extension LoopManager {
 
         indicatorService.openAndUpdate(context: resizeContext)
 
-        isLoopActive = true
         await changeAction(startingAction, disableHapticFeedback: true)
 
         triggerKeyTimeoutTimer.start()
@@ -175,6 +206,7 @@ extension LoopManager {
 
         indicatorService.closeAll()
         isLoopActive = false
+        hasParentCycleActionMirror.withLock { $0 = false }
 
         triggerKeyTimeoutTimer.cancel()
         mouseInteractionObserver.stop()
@@ -312,7 +344,7 @@ extension LoopManager {
                     if let lastAction = await WindowRecords.shared.getCurrentAction(for: targetWindow),
                        lastAction.getName() != screenSwitchingCustomActionName,
                        !lastAction.forceProportionalFrameOnScreenChange {
-                        resizeContext.setAction(to: lastAction, parent: nil)
+                        setResizeAction(to: lastAction, parent: nil)
                     } else {
                         let currentFrame = targetWindow.frame
 
@@ -327,7 +359,7 @@ extension LoopManager {
                             height: currentFrame.height / adjustedBounds.height
                         )
 
-                        resizeContext.setAction(
+                        setResizeAction(
                             to: .init(
                                 .custom,
                                 keybind: [],
@@ -344,15 +376,15 @@ extension LoopManager {
                         )
                     }
                 } else {
-                    resizeContext.setAction(to: .init(.center), parent: nil)
+                    setResizeAction(to: .init(.center), parent: nil)
                 }
             }
 
             resizeContext.setScreen(to: newScreen)
             indicatorService.openAndUpdate(context: resizeContext)
 
             if let parent = newParentAction {
-                resizeContext.setAction(to: newAction, parent: newParentAction)
+                setResizeAction(to: newAction, parent: newParentAction)
                 await changeAction(parent, triggeredFromScreenChange: true)
             } else {
                 if !Defaults[.previewVisibility] {
@@ -377,7 +409,7 @@ extension LoopManager {
 
         if newAction != resizeContext.action || newAction.canRepeat {
             let previousActionWasNoOp = resizeContext.action.direction.isNoOp
-            resizeContext.setAction(to: newAction, parent: newParentAction)
+            setResizeAction(to: newAction, parent: newParentAction)
             if !Defaults[.previewVisibility], !previousActionWasNoOp {
                 await resizeContext.refreshResolvedState()
             }
@@ -459,6 +491,11 @@ extension LoopManager {
         }
     }
 
+    private func setResizeAction(to newAction: WindowAction, parent newParentAction: WindowAction?) {
+        resizeContext.setAction(to: newAction, parent: newParentAction)
+        hasParentCycleActionMirror.withLock { $0 = newParentAction != nil }
+    }
+
     /// Resolves the target screen for `screenToResizeOn`.
     ///
     /// By default, this uses the user's `useScreenWithCursor` setting.

Loop/Core/Observers/MouseInteractionObserver.swift

@@ -54,6 +54,8 @@ final class MouseInteractionObserver {
     }
 
     func start(initialMousePosition: CGPoint) {
+        stop()
+
         screenBounds = NSScreen.screens.first(where: { $0.frame.contains(initialMousePosition) })?.frame
 
         if let screenBounds {

Loop/Core/WindowDragManager.swift

@@ -56,7 +56,17 @@ final class WindowDragManager {
         }
     }
 
+    func shutdown() {
+        accessibilityCheckerTask?.cancel()
+        accessibilityCheckerTask = nil
+        removeListeners()
+        resetDragState()
+        previewController.close()
+    }
+
     private func setupListeners() {
+        removeListeners()
+
         let leftMouseDraggedMonitor = PassiveEventMonitor(
             "snapping_left_mouse_dragged_monitor",
             events: [.leftMouseDragged],

Loop/Utilities/Event Monitoring/BaseEventTapMonitor.swift

@@ -21,32 +21,19 @@ class BaseEventTapMonitor: EventMonitorProtocol, Identifiable, Equatable {
     private(set) var isEnabled: Bool = false
 
     deinit {
-        if isEnabled {
-            stop()
-        }
-
-        // Clean up run loop source and event tap
-        if let runLoop, let runLoopSource {
-            CFRunLoopRemoveSource(runLoop, runLoopSource, .commonModes)
-            self.runLoopSource = nil
-        }
-
-        if let eventTap {
-            CFMachPortInvalidate(eventTap)
-            self.eventTap = nil
-        }
+        tearDownEventTap()
     }
 
     func setupRunLoopSource(eventTap: CFMachPort, readableIdentifier: String) {
-        // Runloop is already running here. In the future, we can investigate running the mach port on another thread.
-        let runLoop = CFRunLoopGetMain()
+        let runLoop = EventTapThread.shared.runLoop
         self.readableIdentifier = readableIdentifier
 
         if let runLoopSource = CFMachPortCreateRunLoopSource(kCFAllocatorDefault, eventTap, 0) {
             self.eventTap = eventTap
             self.runLoop = runLoop
             self.runLoopSource = runLoopSource
             CFRunLoopAddSource(runLoop, runLoopSource, .commonModes)
+            CFRunLoopWakeUp(runLoop)
         }
     }
 
@@ -64,19 +51,66 @@ class BaseEventTapMonitor: EventMonitorProtocol, Identifiable, Equatable {
     }
 
     func stop() {
-        guard let eventTap else { return }
+        guard eventTap != nil else { return }
 
         if let readableIdentifier {
             log.info("Stopping BaseEventTapMonitor '\(readableIdentifier)'")
         } else {
             log.info("Stopping BaseEventTapMonitor with ID \(id)")
         }
 
-        CGEvent.tapEnable(tap: eventTap, enable: false)
-        isEnabled = false
+        tearDownEventTap()
     }
 
     static func == (lhs: BaseEventTapMonitor, rhs: BaseEventTapMonitor) -> Bool {
         lhs.id == rhs.id
     }
+
+    private func tearDownEventTap() {
+        guard eventTap != nil || runLoopSource != nil else { return }
+
+        let eventTap = eventTap
+        let runLoop = runLoop
+        let runLoopSource = runLoopSource
+
+        self.eventTap = nil
+        self.runLoop = nil
+        self.runLoopSource = nil
+        isEnabled = false
+
+        guard let runLoop else {
+            if let eventTap {
+                CGEvent.tapEnable(tap: eventTap, enable: false)
+                CFMachPortInvalidate(eventTap)
+            }
+            return
+        }
+
+        let cleanup = {
+            if let eventTap {
+                CGEvent.tapEnable(tap: eventTap, enable: false)
+            }
+
+            if let runLoopSource {
+                CFRunLoopRemoveSource(runLoop, runLoopSource, .commonModes)
+            }
+
+            if let eventTap {
+                CFMachPortInvalidate(eventTap)
+            }
+        }
+
+        if CFRunLoopGetCurrent() == runLoop {
+            cleanup()
+            return
+        }
+
+        let finished = DispatchSemaphore(value: 0)
+        CFRunLoopPerformBlock(runLoop, CFRunLoopMode.commonModes.rawValue) {
+            cleanup()
+            finished.signal()
+        }
+        CFRunLoopWakeUp(runLoop)
+        finished.wait()
+    }
 }