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

When the user opens documents from three targets A, B, and C in quick succession, then we don’t want to schedule preparation of wait until A *and* B are finished preparing before preparing C.

Instead, we want to
- Finish for preparation of A to finish if it has already started by the time the file in C is opened. This is done so we always make progress during preparation and don’t get into a scenario where preparation is always cancelled if a user switches between two targets more quickly than it takes to prepare those targets.
- Not prepare B because it is no longer relevant and we haven’t started any progress here. Essentially, we pretend that the hop to B never happened.

Author: ahoppen

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,14 @@ public final actor SemanticIndexManager {
   /// After the file is indexed, it is removed from this dictionary.
   private var inProgressIndexTasks: [DocumentURI: InProgressIndexStore] = [:]
 
+  /// The last 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.
+  private var lastPrepareForEditorTask: (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 +295,35 @@ 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 lastPrepareForEditorTask?.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 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)
       }
     }
+    lastPrepareForEditorTask?.task.cancel()
+    lastPrepareForEditorTask = (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 +346,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 +363,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)