Clean Up Child Processes (#1885)

CodeEditApp · Sep 25, 2024 · e4f03fe · e4f03fe
1 parent 9251301
commit e4f03fe
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 5 deletions.
diff --git a/CodeEdit/AppDelegate.swift b/CodeEdit/AppDelegate.swift
@@ -8,13 +8,17 @@
 import SwiftUI
 import CodeEditSymbols
 import CodeEditSourceEditor
+import OSLog
 
 final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
+    private let logger = Logger(subsystem: Bundle.main.bundleIdentifier ?? "", category: "AppDelegate")
     private let updater = SoftwareUpdater()
 
     @Environment(\.openWindow)
     var openWindow
 
+    @LazyService var lspService: LSPService
+
     func applicationDidFinishLaunching(_ notification: Notification) {
         setupServiceContainer()
         enableWindowSizeSaveOnQuit()
@@ -115,6 +119,20 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
         }
     }
 
+    /// Defers the application terminate message until we've finished cleanup.
+    ///
+    /// All paths _must_ call `NSApplication.shared.reply(toApplicationShouldTerminate: true)` as soon as possible.
+    ///
+    /// The two things needing deferring are:
+    /// - Language server cancellation
+    /// - Outstanding document changes.
+    ///
+    /// Things that don't need deferring (happen immediately):
+    /// - Task termination.
+    /// These are called immediately if no documents need closing, and are called by
+    /// ``documentController(_:didCloseAll:contextInfo:)`` if there are documents we need to defer for.
+    ///
+    /// See ``terminateLanguageServers()`` and ``documentController(_:didCloseAll:contextInfo:)`` for deferring tasks.
     func applicationShouldTerminate(_ sender: NSApplication) -> NSApplication.TerminateReply {
         let projects: [String] = CodeEditDocumentController.shared.documents
             .compactMap { ($0 as? WorkspaceDocument)?.fileURL?.path }
@@ -128,10 +146,13 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
                 didCloseAllSelector: #selector(documentController(_:didCloseAll:contextInfo:)),
                 contextInfo: nil
             )
+            // `documentController(_:didCloseAll:contextInfo:)` will call `terminateLanguageServers()`
             return .terminateLater
         }
 
-        return .terminateNow
+        terminateTasks()
+        terminateLanguageServers()
+        return .terminateLater
     }
 
     func applicationShouldTerminateAfterLastWindowClosed(_ sender: NSApplication) -> Bool {
@@ -224,7 +245,28 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
 
     @objc
     func documentController(_ docController: NSDocumentController, didCloseAll: Bool, contextInfo: Any) {
-        NSApplication.shared.reply(toApplicationShouldTerminate: didCloseAll)
+        if didCloseAll {
+            terminateTasks()
+            terminateLanguageServers()
+        }
+    }
+
+    /// Terminates running language servers. Used during app termination to ensure resources are freed.
+    private func terminateLanguageServers() {
+        Task {
+            await lspService.stopAllServers()
+            await MainActor.run {
+                NSApplication.shared.reply(toApplicationShouldTerminate: true)
+            }
+        }
+    }
+
+    /// Terminates all running tasks. Used during app termination to ensure resources are freed.
+    private func terminateTasks() {
+        let documents = CodeEditDocumentController.shared.documents.compactMap({ $0 as? WorkspaceDocument })
+        documents.forEach { workspace in
+            workspace.taskManager?.stopAllTasks()
+        }
     }
 
     /// Setup all the services into a ServiceContainer for the application to use.

diff --git a/CodeEdit/Features/LSP/Service/LSPService.swift b/CodeEdit/Features/LSP/Service/LSPService.swift
@@ -269,10 +269,24 @@ final class LSPService: ObservableObject {
     }
 
     /// Goes through all active language servers and attempts to shut them down.
-    func stopAllServers() async throws {
-        for key in languageClients.keys {
-            try await stopServer(forLanguage: key.languageId, workspacePath: key.workspacePath)
+    func stopAllServers() async {
+        await withThrowingTaskGroup(of: Void.self) { group in
+            for (key, server) in languageClients {
+                group.addTask {
+                    do {
+                        try await server.shutdown()
+                    } catch {
+                        self.logger.error("Shutting down \(key.languageId.rawValue): Error \(error)")
+                        throw error
+                    }
+                }
+            }
+        }
+        languageClients.removeAll()
+        eventListeningTasks.forEach { (_, value) in
+            value.cancel()
         }
+        eventListeningTasks.removeAll()
     }
 }