Don’t re-index file if we are waiting for its preparation

This fixes an issue that caused the index progress to not make progress at the beginning when opening a new project because of the following:

When opening the project, we would cause package resolution and schedule indexing of all source files after the package has been resolved (and dependencies been downloaded). But dependency resolution will add the dependency files to be written to disk and the editor will inform us of these modified files, which will cause us to schedule a new index operation for these files. Thus, when the initial indexing of the files finishes, we won’t increase the number of tasks that we have already indexed.

To fix this, don’t schedule new indexing of these files if we are still waiting for their preparation to start.

This should not be a big performance win because even before this change, when we were trying to index the file the second time, we would have realized that the index store is up to date and thus do a no-op. The main benefit of this change is to have the index counter be more accurate.

Author: ahoppen

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -51,7 +51,8 @@ private struct OpaqueQueuedIndexTask: Equatable {
 }
 
 private enum InProgressIndexStore {
-  /// We are waiting for preparation of the file's target to finish before we can index it.
+  /// We are waiting for preparation of the file's target to be scheduled. The next step is that we wait for
+  /// prepration to finish before we can update the index store for this file.
   ///
   /// `preparationTaskID` identifies the preparation task so that we can transition a file's index state to
   /// `updatingIndexStore` when its preparation task has finished.
@@ -60,6 +61,16 @@ private enum InProgressIndexStore {
   /// task is still the sole owner of it and responsible for its cancellation.
   case waitingForPreparation(preparationTaskID: UUID, indexTask: Task<Void, Never>)
 
+  /// We have started preparing this file and are waiting for preparation to finish before we can update the index
+  /// store for this file.
+  ///
+  /// `preparationTaskID` identifies the preparation task so that we can transition a file's index state to
+  /// `updatingIndexStore` when its preparation task has finished.
+  ///
+  /// `indexTask` is a task that finishes after both preparation and index store update are done. Whoever owns the index
+  /// task is still the sole owner of it and responsible for its cancellation.
+  case preparing(preparationTaskID: UUID, indexTask: Task<Void, Never>)
+
   /// The file's target has been prepared and we are updating the file's index store.
   ///
   /// `updateIndexStoreTask` is the task that updates the index store itself.
@@ -210,7 +221,7 @@ package final actor SemanticIndexManager {
     }
     let indexTasks = inProgressIndexTasks.mapValues { status in
       switch status {
-      case .waitingForPreparation:
+      case .waitingForPreparation, .preparing:
         return IndexTaskStatus.scheduled
       case .updatingIndexStore(updateIndexStoreTask: let updateIndexStoreTask, indexTask: _):
         return updateIndexStoreTask.isExecuting ? IndexTaskStatus.executing : IndexTaskStatus.scheduled
@@ -284,7 +295,17 @@ package final actor SemanticIndexManager {
           }
         if !indexFilesWithUpToDateUnit {
           let index = index.checked(for: .modifiedFiles)
-          filesToIndex = filesToIndex.filter { !index.hasUpToDateUnit(for: $0) }
+          filesToIndex = filesToIndex.filter {
+            if index.hasUpToDateUnit(for: $0) {
+              return false
+            }
+            if case .waitingForPreparation = inProgressIndexTasks[$0] {
+              // We haven't started preparing the file yet. Scheduling a new index operation for it won't produce any
+              // more recent results.
+              return false
+            }
+            return true
+          }
         }
         await scheduleBackgroundIndex(files: filesToIndex, indexFilesWithUpToDateUnit: indexFilesWithUpToDateUnit)
         generateBuildGraphTask = nil
@@ -312,11 +333,11 @@ package final actor SemanticIndexManager {
       for (_, status) in inProgressIndexTasks {
         switch status {
         case .waitingForPreparation(preparationTaskID: _, indexTask: let indexTask),
+          .preparing(preparationTaskID: _, indexTask: let indexTask),
           .updatingIndexStore(updateIndexStoreTask: _, indexTask: let indexTask):
           taskGroup.addTask {
             await indexTask.value
           }
-
         }
       }
       await taskGroup.waitForAll()
@@ -462,7 +483,9 @@ package final actor SemanticIndexManager {
   private func prepare(
     targets: [BuildTargetIdentifier],
     purpose: TargetPreparationPurpose,
-    priority: TaskPriority?
+    priority: TaskPriority?,
+    executionStatusChangedCallback: @escaping (QueuedTask<AnyIndexTaskDescription>, TaskExecutionState) async -> Void =
+      { _, _ in }
   ) 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.
@@ -490,6 +513,7 @@ package final actor SemanticIndexManager {
       return
     }
     let preparationTask = await indexTaskScheduler.schedule(priority: priority, taskDescription) { task, newState in
+      await executionStatusChangedCallback(task, newState)
       guard case .finished = newState else {
         self.indexProgressStatusDidChange()
         return
@@ -547,28 +571,36 @@ package final actor SemanticIndexManager {
         testHooks: testHooks
       )
     )
+
     let updateIndexTask = await indexTaskScheduler.schedule(priority: priority, taskDescription) { task, newState in
       guard case .finished = newState else {
         self.indexProgressStatusDidChange()
         return
       }
       for fileAndTarget in filesAndTargets {
-        if case .updatingIndexStore(OpaqueQueuedIndexTask(task), _) = self.inProgressIndexTasks[
-          fileAndTarget.file.sourceFile
-        ] {
-          self.inProgressIndexTasks[fileAndTarget.file.sourceFile] = nil
+        switch self.inProgressIndexTasks[fileAndTarget.file.sourceFile] {
+        case .updatingIndexStore(let registeredTask, _):
+          if registeredTask == OpaqueQueuedIndexTask(task) {
+            self.inProgressIndexTasks[fileAndTarget.file.sourceFile] = nil
+          }
+        case .waitingForPreparation(let registeredTask, _), .preparing(let registeredTask, _):
+          if registeredTask == preparationTaskID {
+            self.inProgressIndexTasks[fileAndTarget.file.sourceFile] = nil
+          }
+        case nil:
+          break
         }
       }
       self.indexProgressStatusDidChange()
     }
     for fileAndTarget in filesAndTargets {
-      if case .waitingForPreparation(preparationTaskID, let indexTask) = inProgressIndexTasks[
-        fileAndTarget.file.sourceFile
-      ] {
+      switch inProgressIndexTasks[fileAndTarget.file.sourceFile] {
+      case .waitingForPreparation(preparationTaskID, let indexTask), .preparing(preparationTaskID, let indexTask):
         inProgressIndexTasks[fileAndTarget.file.sourceFile] = .updatingIndexStore(
           updateIndexStoreTask: OpaqueQueuedIndexTask(updateIndexTask),
           indexTask: indexTask
         )
+      default: break
       }
     }
     return await updateIndexTask.waitToFinishPropagatingCancellation()
@@ -639,9 +671,24 @@ package final actor SemanticIndexManager {
     // (https://github.com/swiftlang/sourcekit-lsp/issues/1262)
     for targetsBatch in sortedTargets.partition(intoBatchesOfSize: 1) {
       let preparationTaskID = UUID()
+      let filesToIndex = targetsBatch.flatMap({ filesByTarget[$0]! })
+
       let indexTask = Task(priority: priority) {
         // First prepare the targets.
-        await prepare(targets: targetsBatch, purpose: .forIndexing, priority: priority)
+        await prepare(targets: targetsBatch, purpose: .forIndexing, priority: priority) { task, newState in
+          if case .executing = newState {
+            for file in filesToIndex {
+              if case .waitingForPreparation(preparationTaskID: preparationTaskID, indexTask: let indexTask) =
+                self.inProgressIndexTasks[file.sourceFile]
+              {
+                self.inProgressIndexTasks[file.sourceFile] = .preparing(
+                  preparationTaskID: preparationTaskID,
+                  indexTask: indexTask
+                )
+              }
+            }
+          }
+        }
 
         // And after preparation is done, index the files in the targets.
         await withTaskGroup(of: Void.self) { taskGroup in
@@ -665,7 +712,6 @@ package final actor SemanticIndexManager {
       }
       indexTasks.append(indexTask)
 
-      let filesToIndex = targetsBatch.flatMap({ filesByTarget[$0]! })
       // The number of index tasks that don't currently have an in-progress task associated with it.
       // The denominator in the index progress should get incremented by this amount.
       // We don't want to increment the denominator for tasks that already have an index in progress.