Cancel preparation tasks for editor functionality if the preparation task hasn't been started yet and the document is no longer active

Sources/SKCore/TaskScheduler.swift

@@ -147,13 +147,20 @@ public actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
   /// Gets reset every time `executionTask` finishes.
   nonisolated(unsafe) private var cancelledToBeRescheduled: AtomicBool = .init(initialValue: false)
 
+  /// Whether `resultTask` has been cancelled.
+  private nonisolated(unsafe) var resultTaskCancelled: AtomicBool = .init(initialValue: false)
+
   private nonisolated(unsafe) var _isExecuting: AtomicBool = .init(initialValue: false)
 
   /// Whether the task is currently executing or still queued to be executed later.
   public nonisolated var isExecuting: Bool {
     return _isExecuting.value
   }
 
+  public nonisolated func cancel() {
+    resultTask.cancel()
+  }
+
   /// Wait for the task to finish.
   ///
   /// If the tasks that waits for this queued task to finished is cancelled, the QueuedTask will still continue
@@ -197,31 +204,35 @@ public actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
     self.executionTaskCreatedContinuation = executionTaskCreatedContinuation
 
     self.resultTask = Task.detached(priority: priority) {
-      await withTaskGroup(of: Void.self) { taskGroup in
-        taskGroup.addTask {
-          for await _ in updatePriorityStream {
-            self.priority = Task.currentPriority
+      await withTaskCancellationHandler {
+        await withTaskGroup(of: Void.self) { taskGroup in
+          taskGroup.addTask {
+            for await _ in updatePriorityStream {
+              self.priority = Task.currentPriority
+            }
           }
-        }
-        taskGroup.addTask {
-          for await task in executionTaskCreatedStream {
-            switch await task.value {
-            case .cancelledToBeRescheduled:
-              // Break the switch and wait for a new `executionTask` to be placed into `executionTaskCreatedStream`.
-              break
-            case .terminated:
-              // The task finished. We are done with this `QueuedTask`
-              return
+          taskGroup.addTask {
+            for await task in executionTaskCreatedStream {
+              switch await task.valuePropagatingCancellation {
+              case .cancelledToBeRescheduled:
+                // Break the switch and wait for a new `executionTask` to be placed into `executionTaskCreatedStream`.
+                break
+              case .terminated:
+                // The task finished. We are done with this `QueuedTask`
+                return
+              }
             }
           }
+          // The first (update priority) task never finishes, so this waits for the second (wait for execution) task
+          // to terminate.
+          // Afterwards we also cancel the update priority task.
+          for await _ in taskGroup {
+            taskGroup.cancelAll()
+            return
+          }
         }
-        // The first (update priority) task never finishes, so this waits for the second (wait for execution) task
-        // to terminate.
-        // Afterwards we also cancel the update priority task.
-        for await _ in taskGroup {
-          taskGroup.cancelAll()
-          return
-        }
+      } onCancel: {
+        self.resultTaskCancelled.value = true
       }
     }
   }
@@ -233,7 +244,9 @@ public actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
   func execute() async -> ExecutionTaskFinishStatus {
     precondition(executionTask == nil, "Task started twice")
     let task = Task.detached(priority: self.priority) {
-      await self.description.execute()
+      if !Task.isCancelled && !self.resultTaskCancelled.value {
+        await self.description.execute()
+      }
       return await self.finalizeExecution()
     }
     executionTask = task

Sources/SKSupport/LineTable.swift

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 import LSPLogging
+
 #if canImport(os)
 import os
 #endif

Sources/SKSwiftPMWorkspace/SwiftPMBuildSystem.swift

@@ -489,6 +489,9 @@ extension SwiftPMBuildSystem: SKCore.BuildSystem {
       "--disable-index-store",
       "--target", target.targetID,
     ]
+    if Task.isCancelled {
+      return
+    }
     let process = try Process.launch(arguments: arguments, workingDirectory: nil)
     let result = try await process.waitUntilExitSendingSigIntOnTaskCancellation()
     switch result.exitStatus.exhaustivelySwitchable {

Sources/SemanticIndex/PreparationTaskDescription.swift

@@ -70,6 +70,7 @@ public struct PreparationTaskDescription: IndexTaskDescription {
     // See comment in `withLoggingScope`.
     // The last 2 digits should be sufficient to differentiate between multiple concurrently running preparation operations
     await withLoggingScope("preparation-\(id % 100)") {
+      await testHooks.preparationTaskDidStart?(self)
       let targetsToPrepare = await targetsToPrepare.asyncFilter {
         await !preparationUpToDateStatus.isUpToDate($0)
       }.sorted(by: {

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -86,6 +86,17 @@ public final actor SemanticIndexManager {
   /// After the file is indexed, it is removed from this dictionary.
   private var inProgressIndexTasks: [DocumentURI: InProgressIndexStore] = [:]
 
+  /// The currently running task that prepares a document for editor functionality.
+  ///
+  /// This is used so we can cancel preparation tasks for documents that the user is no longer interacting with and
+  /// avoid the following scenario: The user browses through documents from targets A, B, and C in quick succession. We
+  /// don't want stack preparation of A, B, and C. Instead we want to only prepare target C - and also finish
+  /// preparation of A if it has already started when the user opens C.
+  ///
+  /// `id` is a unique ID that identifies the preparation task and is used to set `inProgressPrepareForEditorTask` to
+  /// `nil` when the current in progress task finishes.
+  private var inProgressPrepareForEditorTask: (id: UUID, document: DocumentURI, task: Task<Void, Never>)? = nil
+
   /// The `TaskScheduler` that manages the scheduling of index tasks. This is shared among all `SemanticIndexManager`s
   /// in the process, to ensure that we don't schedule more index operations than processor cores from multiple
   /// workspaces.
@@ -287,20 +298,39 @@ public final actor SemanticIndexManager {
   /// Schedule preparation of the target that contains the given URI, building all modules that the file depends on.
   ///
   /// This is intended to be called when the user is interacting with the document at the given URI.
-  public func schedulePreparation(of uri: DocumentURI, priority: TaskPriority? = nil) {
-    Task(priority: priority) {
+  public func schedulePreparationForEditorFunctionality(
+    of uri: DocumentURI,
+    priority: TaskPriority? = nil
+  ) {
+    if inProgressPrepareForEditorTask?.document == uri {
+      // We are already preparing this document, so nothing to do. This is necessary to avoid the following scenario:
+      // Determining the canonical configured target for a document takes 1s and we get a new document request for the
+      // document ever 0.5s, which would cancel the previous in-progress preparation task, cancelling the canonical
+      // configured target configuration, never actually getting to the actual preparation.
+      return
+    }
+    let id = UUID()
+    let task = Task(priority: priority) {
       await withLoggingScope("preparation") {
         guard let target = await buildSystemManager.canonicalConfiguredTarget(for: uri) else {
           return
         }
+        if Task.isCancelled {
+          return
+        }
         await self.prepare(targets: [target], priority: priority)
+        if inProgressPrepareForEditorTask?.id == id {
+          inProgressPrepareForEditorTask = nil
+        }
       }
     }
+    inProgressPrepareForEditorTask?.task.cancel()
+    inProgressPrepareForEditorTask = (id, uri, task)
   }
 
   // MARK: - Helper functions
 
-  /// Prepare the given targets for indexing
+  /// Prepare the given targets for indexing.
   private func prepare(targets: [ConfiguredTarget], priority: TaskPriority?) async {
     // Perform a quick initial check whether the target is up-to-date, in which case we don't need to schedule a
     // preparation operation at all.
@@ -323,6 +353,9 @@ public final actor SemanticIndexManager {
         testHooks: testHooks
       )
     )
+    if Task.isCancelled {
+      return
+    }
     let preparationTask = await indexTaskScheduler.schedule(priority: priority, taskDescription) { task, newState in
       guard case .finished = newState else {
         return
@@ -337,7 +370,17 @@ public final actor SemanticIndexManager {
     for target in targetsToPrepare {
       inProgressPreparationTasks[target] = OpaqueQueuedIndexTask(preparationTask)
     }
-    return await preparationTask.waitToFinishPropagatingCancellation()
+    await withTaskCancellationHandler {
+      return await preparationTask.waitToFinish()
+    } onCancel: {
+      // Only cancel the preparation task if it hasn't started executing yet. This ensures that we always make progress
+      // during preparation and can't get into the following scenario: The user has two target A and B that both take
+      // 10s to prepare. The user is now switching between the files every 5 seconds, which would always cause
+      // preparation for one target to get cancelled, never resulting in an up-to-date preparation status.
+      if !preparationTask.isExecuting {
+        preparationTask.cancel()
+      }
+    }
   }
 
   /// Update the index store for the given files, assuming that their targets have already been prepared.

Sources/SemanticIndex/TestHooks.swift

@@ -12,15 +12,19 @@
 
 /// Callbacks that allow inspection of internal state modifications during testing.
 public struct IndexTestHooks: Sendable {
+  public var preparationTaskDidStart: (@Sendable (PreparationTaskDescription) async -> Void)?
+
   public var preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) async -> Void)?
 
   /// A callback that is called when an index task finishes.
   public var updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)?
 
   public init(
+    preparationTaskDidStart: (@Sendable (PreparationTaskDescription) async -> Void)? = nil,
     preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) async -> Void)? = nil,
     updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)? = nil
   ) {
+    self.preparationTaskDidStart = preparationTaskDidStart
     self.preparationTaskDidFinish = preparationTaskDidFinish
     self.updateIndexStoreTaskDidFinish = updateIndexStoreTaskDidFinish
   }

Sources/SemanticIndex/UpdateIndexStoreTaskDescription.swift

@@ -301,6 +301,9 @@ public struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
     processArguments: [String],
     workingDirectory: AbsolutePath?
   ) async throws {
+    if Task.isCancelled {
+      return
+    }
     let process = try Process.launch(
       arguments: processArguments,
       workingDirectory: workingDirectory

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -1000,7 +1000,9 @@ extension SourceKitLSPServer: MessageHandler {
       // which prepares the files. For files that are open but aren't being worked on (eg. a different tab), we don't
       // get requests, ensuring that we don't unnecessarily prepare them.
       let workspace = await self.workspaceForDocument(uri: textDocumentRequest.textDocument.uri)
-      await workspace?.semanticIndexManager?.schedulePreparation(of: textDocumentRequest.textDocument.uri)
+      await workspace?.semanticIndexManager?.schedulePreparationForEditorFunctionality(
+        of: textDocumentRequest.textDocument.uri
+      )
     }
 
     switch request {
@@ -1543,7 +1545,7 @@ extension SourceKitLSPServer {
       )
       return
     }
-    await workspace.semanticIndexManager?.schedulePreparation(of: uri)
+    await workspace.semanticIndexManager?.schedulePreparationForEditorFunctionality(of: uri)
     await openDocument(notification, workspace: workspace)
   }
 
@@ -1599,7 +1601,7 @@ extension SourceKitLSPServer {
       )
       return
     }
-    await workspace.semanticIndexManager?.schedulePreparation(of: uri)
+    await workspace.semanticIndexManager?.schedulePreparationForEditorFunctionality(of: uri)
 
     // If the document is ready, we can handle the change right now.
     documentManager.edit(notification)