Merge pull request #2081 from ahoppen/cancel-in-progress-indexing

Terminate pending background indexing and preparation tasks when shutting down SourceKit-LSP

Author: ahoppen

Sources/BuildSystemIntegration/BuildSystemManager.swift

@@ -559,8 +559,18 @@ package actor BuildSystemManager: QueueBasedMessageHandler {
       return
     }
     await orLog("Sending shutdown request to build server") {
-      _ = try await buildSystemAdapter.send(BuildShutdownRequest())
-      await buildSystemAdapter.send(OnBuildExitNotification())
+      // Give the build server 2 seconds to shut down by itself. If it doesn't shut down within that time, terminate it.
+      try await withTimeout(.seconds(2)) {
+        _ = try await buildSystemAdapter.send(BuildShutdownRequest())
+        await buildSystemAdapter.send(OnBuildExitNotification())
+      }
+    }
+    if case .external(let externalBuildSystemAdapter) = buildSystemAdapter {
+      await orLog("Terminating external build server") {
+        // Give the build server 1 second to exit after receiving the `build/exit` notification. If it doesn't exit
+        // within that time, terminate it.
+        try await externalBuildSystemAdapter.terminateIfRunning(after: .seconds(1))
+      }
     }
     self.buildSystemAdapter = nil
   }

Sources/BuildSystemIntegration/ExternalBuildSystemAdapter.swift

@@ -102,6 +102,9 @@ actor ExternalBuildSystemAdapter {
   /// server crashes.
   private var initializeRequest: InitializeBuildRequest?
 
+  /// The process that runs the external BSP server.
+  private var process: Process?
+
   /// The date at which `clangd` was last restarted.
   /// Used to delay restarting in case of a crash loop.
   private var lastRestart: Date?
@@ -175,7 +178,7 @@ actor ExternalBuildSystemAdapter {
       serverPath = interpreterPath
     }
 
-    return try JSONRPCConnection.start(
+    let (connection, process) = try JSONRPCConnection.start(
       executable: serverPath,
       arguments: serverArgs,
       name: "BSP-Server",
@@ -194,7 +197,9 @@ actor ExternalBuildSystemAdapter {
           }
         }
       }
-    ).connection
+    )
+    self.process = process
+    return connection
   }
 
   private static func getConfigPath(for workspaceFolder: URL? = nil, onlyConsiderRoot: Bool = false) -> URL? {
@@ -295,6 +300,11 @@ actor ExternalBuildSystemAdapter {
     // SourceKit-LSP to discard cached information.
     self.messagesToSourceKitLSPHandler.handle(OnBuildTargetDidChangeNotification(changes: nil))
   }
+
+  /// If the build server is still running after `duration`, terminate it. Otherwise a no-op.
+  package func terminateIfRunning(after duration: Duration) async throws {
+    try await process?.terminateIfRunning(after: duration)
+  }
 }
 
 fileprivate extension URL {

Sources/SKTestSupport/CustomBuildServerTestProject.swift

@@ -208,7 +208,7 @@ package extension CustomBuildServer {
     )
   }
 
-  func dummyTargetSourcesResponse(_ files: some Sequence<DocumentURI>) -> BuildTargetSourcesResponse {
+  func dummyTargetSourcesResponse(files: some Sequence<DocumentURI>) -> BuildTargetSourcesResponse {
     return BuildTargetSourcesResponse(items: [
       SourcesItem(target: .dummy, sources: files.map { SourceItem(uri: $0, kind: .file, generated: false) })
     ])

Sources/SemanticIndex/TaskScheduler.swift

@@ -324,6 +324,9 @@ package actor TaskScheduler<TaskDescription: TaskDescriptionProtocol> {
   /// The queue of pending tasks that haven't been scheduled for execution yet.
   private var pendingTasks: [QueuedTask<TaskDescription>] = []
 
+  /// Whether `shutDown` has been called on the `TaskScheduler` and it should thus schedule any new tasks.
+  private var isShutDown: Bool = false
+
   /// An ordered list of task priorities to the number of tasks that might execute concurrently at that or a lower
   /// priority. As the task priority is decreased, the number of current tasks becomes more restricted.
   ///
@@ -440,6 +443,18 @@ package actor TaskScheduler<TaskDescription: TaskDescriptionProtocol> {
     return queuedTask
   }
 
+  /// Cancel all in-progress tasks and wait for them to finish, either by honoring the cancellation or finishing their
+  /// work.
+  ///
+  /// After `shutDown` has been called, no more tasks will be executed on this `TaskScheduler`.
+  package func shutDown() async {
+    self.isShutDown = true
+    await self.currentlyExecutingTasks.concurrentForEach { task in
+      task.cancel()
+      await task.waitToFinish()
+    }
+  }
+
   private static func normalize(
     maxConcurrentTasksByPriority: [(priority: TaskPriority, maxConcurrentTasks: Int)]
   ) -> [(priority: TaskPriority, maxConcurrentTasks: Int)] {

Sources/SourceKitLSP/Clang/ClangLanguageService.swift

@@ -100,12 +100,8 @@ actor ClangLanguageService: LanguageService, MessageHandler {
   /// Type to map `clangd`'s semantic token legend to SourceKit-LSP's.
   private var semanticTokensTranslator: SemanticTokensLegendTranslator? = nil
 
-  /// While `clangd` is running, its PID.
-  #if os(Windows)
-  private var hClangd: HANDLE = INVALID_HANDLE_VALUE
-  #else
-  private var clangdPid: Int32?
-  #endif
+  /// While `clangd` is running, its `Process` object.
+  private var clangdProcess: Process?
 
   /// Creates a language server for the given client referencing the clang binary specified in `toolchain`.
   /// Returns `nil` if `clangd` can't be found.
@@ -159,11 +155,7 @@ actor ClangLanguageService: LanguageService, MessageHandler {
   ///
   /// - Parameter terminationStatus: The exit code of `clangd`.
   private func handleClangdTermination(terminationStatus: Int32) {
-    #if os(Windows)
-    self.hClangd = INVALID_HANDLE_VALUE
-    #else
-    self.clangdPid = nil
-    #endif
+    self.clangdProcess = nil
     if terminationStatus != 0 {
       self.state = .connectionInterrupted
       self.restartClangd()
@@ -196,11 +188,7 @@ actor ClangLanguageService: LanguageService, MessageHandler {
     )
     self.clangd = connectionToClangd
 
-    #if os(Windows)
-    self.hClangd = process.processHandle
-    #else
-    self.clangdPid = process.processIdentifier
-    #endif
+    self.clangdProcess = process
   }
 
   /// Restart `clangd` after it has crashed.
@@ -318,21 +306,7 @@ actor ClangLanguageService: LanguageService, MessageHandler {
   }
 
   func crash() {
-    // Since `clangd` doesn't have a method to crash it, terminate it.
-    #if os(Windows)
-    if self.hClangd != INVALID_HANDLE_VALUE {
-      // We can potentially deadlock the process if a kobject is a pending state.
-      // Unfortunately, the `OpenProcess(PROCESS_TERMINATE, ...)`, `CreateRemoteThread`, `ExitProcess` dance, while
-      // safer, can also indefinitely spin as `CreateRemoteThread` may not be serviced depending on the state of
-      // the process. This just attempts to terminate the process, risking a deadlock and resource leaks, which is fine
-      // since we only use `crash` from tests.
-      _ = TerminateProcess(self.hClangd, 255)
-    }
-    #else
-    if let pid = self.clangdPid {
-      kill(pid, SIGKILL)  // ignore-unacceptable-language
-    }
-    #endif
+    clangdProcess?.terminateImmediately()
   }
 }
 
@@ -401,15 +375,19 @@ extension ClangLanguageService {
   }
 
   package func shutdown() async {
-    let clangd = clangd!
-    await withCheckedContinuation { continuation in
-      _ = clangd.send(ShutdownRequest()) { _ in
-        Task {
-          clangd.send(ExitNotification())
-          continuation.resume()
-        }
+    _ = await orLog("Shutting down clangd") {
+      guard let clangd else { return }
+      // Give clangd 2 seconds to shut down by itself. If it doesn't shut down within that time, terminate it.
+      try await withTimeout(.seconds(2)) {
+        _ = try await clangd.send(ShutdownRequest())
+        clangd.send(ExitNotification())
       }
     }
+    await orLog("Terminating clangd") {
+      // Give clangd 1 second to exit after receiving the `exit` notification. If it doesn't exit within that time,
+      // terminate it.
+      try await clangdProcess?.terminateIfRunning(after: .seconds(1))
+    }
   }
 
   // MARK: - Text synchronization

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -1189,26 +1189,27 @@ extension SourceKitLSPServer {
   /// The server shall not be used to handle more requests (other than possibly
   /// `shutdown` and `exit`) and should attempt to flush any buffered state
   /// immediately, such as sending index changes to disk.
+  ///
+  ///  - Note: this method should be safe to call multiple times, since we want to be resilient against multiple
+  //     possible shutdown sequences, including pipe failure.
   package func prepareForExit() async {
-    // Note: this method should be safe to call multiple times, since we want to
-    // be resilient against multiple possible shutdown sequences, including
-    // pipe failure.
-
-    // Theoretically, new workspaces could be added while we are awaiting inside
-    // the loop. But since we are currently exiting, it doesn't make sense for
-    // the client to open new workspaces.
-    for workspace in self.workspaces {
-      await workspace.buildSystemManager.setMainFilesProvider(nil)
-      workspace.closeIndex()
-    }
-    // Make sure we emit all pending log messages. When we're not using `NonDarwinLogger` this is a no-op.
-    await NonDarwinLogger.flush()
-  }
-
-  func shutdown(_ request: ShutdownRequest) async throws -> ShutdownRequest.Response {
-    await prepareForExit()
+    // We are shutting down / closing all workspaces and language services, so clear the arrays caching them.
+    let languageServices = self.languageServices
+    self.languageServices = [:]
+
+    let workspaces = await self.workspaceQueue.async {
+      let workspaces = self.workspaces
+      self.workspacesAndIsImplicit = []
+      return workspaces
+    }.valuePropagatingCancellation
 
+    // Concurrently shut all things down.
     await withTaskGroup(of: Void.self) { taskGroup in
+      taskGroup.addTask {
+        await orLog("Shutting down index scheduler") {
+          await self.indexTaskScheduler.shutDown()
+        }
+      }
       for service in languageServices.values.flatMap({ $0 }) {
         taskGroup.addTask {
           await service.shutdown()
@@ -1217,20 +1218,20 @@ extension SourceKitLSPServer {
       for workspace in workspaces {
         taskGroup.addTask {
           await orLog("Shutting down build server") {
-            // If the build server doesn't shut down in 1 second, don't delay SourceKit-LSP's shutdown because of it.
-            try await withTimeout(.seconds(2)) {
-              await workspace.buildSystemManager.shutdown()
-            }
+            await workspace.buildSystemManager.shutdown()
+            await workspace.buildSystemManager.setMainFilesProvider(nil)
+            workspace.closeIndex()
           }
         }
       }
     }
 
-    // We have a semantic guarantee that no request or notification should be
-    // sent to an LSP server after the shutdown request. Thus, there's no chance
-    // that a new language service has been started during the above 'await'
-    // call.
-    languageServices = [:]
+    // Make sure we emit all pending log messages. When we're not using `NonDarwinLogger` this is a no-op.
+    await NonDarwinLogger.flush()
+  }
+
+  func shutdown(_ request: ShutdownRequest) async throws -> ShutdownRequest.Response {
+    await prepareForExit()
 
     // Wait for all services to shut down before sending the shutdown response.
     // Otherwise we might terminate sourcekit-lsp while it still has open

Sources/SwiftExtensions/CMakeLists.txt

@@ -13,6 +13,7 @@ set(sources
   NSLock+WithLock.swift
   PipeAsStringHandler.swift
   Platform.swift
+  Process+terminate.swift
   ResultExtensions.swift
   Sequence+AsyncMap.swift
   Sequence+ContainsAnyIn.swift

Sources/SwiftExtensions/Process+terminate.swift

@@ -0,0 +1,38 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2025 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+package import Foundation
+
+extension Foundation.Process {
+  /// If the process has not exited after `duration`, terminate it.
+  package func terminateIfRunning(after duration: Duration, pollInterval: Duration = .milliseconds(5)) async throws {
+    for _ in 0..<Int(duration.seconds / pollInterval.seconds) {
+      if !self.isRunning {
+        break
+      }
+      try await Task.sleep(for: pollInterval)
+    }
+    if self.isRunning {
+      self.terminate()
+    }
+  }
+
+  /// On Posix platforms, send a SIGKILL to the process. On Windows, terminate the process.
+  package func terminateImmediately() {
+    // TODO: We should also terminate all child processes (https://github.com/swiftlang/sourcekit-lsp/issues/2080)
+    #if os(Windows)
+    self.terminate()
+    #else
+    Foundation.kill(processIdentifier, SIGKILL)  // ignore-unacceptable-language
+    #endif
+  }
+}

Sources/TSCExtensions/Process+Run.swift

@@ -43,6 +43,7 @@ extension Process {
         // Give the process 2 seconds to react to a SIGINT. If that doesn't work, terminate the process.
         try await Task.sleep(for: .seconds(2))
         if !hasExited.value {
+          // TODO: We should also terminate all child processes (https://github.com/swiftlang/sourcekit-lsp/issues/2080)
           #if os(Windows)
           // Windows does not define SIGKILL. Process.signal sends a `terminate` to the underlying Foundation process
           // for any signal that is not SIGINT. Use `SIGABRT` to terminate the process.