Show a work done progress that shows the index progress

This makes it a lot easier to work on background indexing because you can easily see how background indexing is making progress.

Resolves #1257
rdar://127474057

Author: ahoppen

Sources/SKCore/TaskScheduler.swift

@@ -76,6 +76,19 @@ public protocol TaskDescriptionProtocol: Identifiable, Sendable, CustomLogString
   var estimatedCPUCoreCount: Int { get }
 }
 
+/// Parameter that's passed to `executionStateChangedCallback` to indicate the new state of a scheduled task.
+public enum TaskExecutionState {
+  /// The task started executing.
+  case executing
+
+  /// The task was cancelled and will be re-scheduled for execution later. Will be followed by another call with
+  /// `executing`.
+  case cancelledToBeRescheduled
+
+  /// The task has finished executing. Now more state updates will come after this one.
+  case finished
+}
+
 fileprivate actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
   /// Result of `executionTask` / the tasks in `executionTaskCreatedContinuation`.
   /// See doc comment on `executionTask`.
@@ -136,9 +149,18 @@ fileprivate actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
   /// Gets reset every time `executionTask` finishes.
   nonisolated(unsafe) private var cancelledToBeRescheduled: AtomicBool = .init(initialValue: false)
 
-  init(priority: TaskPriority? = nil, description: TaskDescription) async {
+  /// A callback that will be called when the task starts executing, is cancelled to be rescheduled, or when it finishes
+  /// execution.
+  private let executionStateChangedCallback: (@Sendable (TaskExecutionState) async -> Void)?
+
+  init(
+    priority: TaskPriority? = nil,
+    description: TaskDescription,
+    executionStateChangedCallback: (@Sendable (TaskExecutionState) async -> Void)?
+  ) async {
     self._priority = .init(initialValue: priority?.rawValue ?? Task.currentPriority.rawValue)
     self.description = description
+    self.executionStateChangedCallback = executionStateChangedCallback
 
     var updatePriorityContinuation: AsyncStream<Void>.Continuation!
     let updatePriorityStream = AsyncStream {
@@ -194,16 +216,19 @@ fileprivate actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
     }
     executionTask = task
     executionTaskCreatedContinuation.yield(task)
+    await executionStateChangedCallback?(.executing)
     return await task.value
   }
 
   /// Implementation detail of `execute` that is called after `self.description.execute()` finishes.
-  private func finalizeExecution() -> ExecutionTaskFinishStatus {
+  private func finalizeExecution() async -> ExecutionTaskFinishStatus {
     self.executionTask = nil
     if Task.isCancelled && self.cancelledToBeRescheduled.value {
+      await executionStateChangedCallback?(.cancelledToBeRescheduled)
       self.cancelledToBeRescheduled.value = false
       return ExecutionTaskFinishStatus.cancelledToBeRescheduled
     } else {
+      await executionStateChangedCallback?(.finished)
       return ExecutionTaskFinishStatus.terminated
     }
   }
@@ -308,12 +333,17 @@ public actor TaskScheduler<TaskDescription: TaskDescriptionProtocol> {
   @discardableResult
   public func schedule(
     priority: TaskPriority? = nil,
-    _ taskDescription: TaskDescription
+    _ taskDescription: TaskDescription,
+    @_inheritActorContext executionStateChangedCallback: (@Sendable (TaskExecutionState) async -> Void)? = nil
   ) async -> Task<Void, Never> {
     withLoggingSubsystemAndScope(subsystem: taskSchedulerSubsystem, scope: nil) {
       logger.debug("Scheduling \(taskDescription.forLogging)")
     }
-    let queuedTask = await QueuedTask(priority: priority, description: taskDescription)
+    let queuedTask = await QueuedTask(
+      priority: priority,
+      description: taskDescription,
+      executionStateChangedCallback: executionStateChangedCallback
+    )
     pendingTasks.append(queuedTask)
     Task.detached(priority: priority ?? Task.currentPriority) {
       // Poke the `TaskScheduler` to execute a new task. If the `TaskScheduler` is already working at its capacity

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -19,8 +19,22 @@ import SKCore
 private enum FileIndexStatus {
   /// The index is up-to-date.
   case upToDate
-  /// The file is being indexed by the given task.
-  case inProgress(Task<Void, Never>)
+  /// The file is not up to date and we have scheduled a task to index it but that index operation hasn't been started
+  /// yet.
+  case scheduled(Task<Void, Never>)
+  /// We are currently actively indexing this file, ie. we are running a subprocess that indexes the file.
+  case executing(Task<Void, Never>)
+
+  var description: String {
+    switch self {
+    case .upToDate:
+      return "upToDate"
+    case .scheduled:
+      return "scheduled"
+    case .executing:
+      return "executing"
+    }
+  }
 }
 
 /// Schedules index tasks and keeps track of the index status of files.
@@ -46,22 +60,51 @@ public final actor SemanticIndexManager {
   /// workspaces.
   private let indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>
 
+  /// Called when files are scheduled to be indexed.
+  ///
+  /// The parameter is the number of files that were scheduled to be indexed.
+  private let indexTasksWereScheduled: @Sendable (_ numberOfFileScheduled: Int) -> Void
+
   /// Callback that is called when an index task has finished.
   ///
-  /// Currently only used for testing.
-  private let indexTaskDidFinish: (@Sendable () -> Void)?
+  /// An object observing this property probably wants to check `inProgressIndexTasks` when the callback is called to
+  /// get the current list of in-progress index tasks.
+  ///
+  /// The number of `indexTaskDidFinish` calls does not have to relate to the number of `indexTasksWereScheduled` calls.
+  private let indexTaskDidFinish: @Sendable () -> Void
 
   // MARK: - Public API
 
+  /// The files that still need to be indexed.
+  ///
+  /// See `FileIndexStatus` for the distinction between `scheduled` and `executing`.
+  public var inProgressIndexTasks: (scheduled: [DocumentURI], executing: [DocumentURI]) {
+    let scheduled = indexStatus.compactMap { (uri: DocumentURI, status: FileIndexStatus) in
+      if case .scheduled = status {
+        return uri
+      }
+      return nil
+    }
+    let inProgress = indexStatus.compactMap { (uri: DocumentURI, status: FileIndexStatus) in
+      if case .executing = status {
+        return uri
+      }
+      return nil
+    }
+    return (scheduled, inProgress)
+  }
+
   public init(
     index: UncheckedIndex,
     buildSystemManager: BuildSystemManager,
     indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>,
-    indexTaskDidFinish: (@Sendable () -> Void)?
+    indexTasksWereScheduled: @escaping @Sendable (Int) -> Void,
+    indexTaskDidFinish: @escaping @Sendable () -> Void
   ) {
     self.index = index.checked(for: .modifiedFiles)
     self.buildSystemManager = buildSystemManager
     self.indexTaskScheduler = indexTaskScheduler
+    self.indexTasksWereScheduled = indexTasksWereScheduled
     self.indexTaskDidFinish = indexTaskDidFinish
   }
 
@@ -93,7 +136,7 @@ public final actor SemanticIndexManager {
     await withTaskGroup(of: Void.self) { taskGroup in
       for (_, status) in indexStatus {
         switch status {
-        case .inProgress(let task):
+        case .scheduled(let task), .executing(let task):
           taskGroup.addTask {
             await task.value
           }
@@ -138,7 +181,7 @@ public final actor SemanticIndexManager {
       )
     )
     await self.indexTaskScheduler.schedule(priority: priority, taskDescription).value
-    self.indexTaskDidFinish?()
+    self.indexTaskDidFinish()
   }
 
   /// Update the index store for the given files, assuming that their targets have already been prepared.
@@ -150,11 +193,44 @@ public final actor SemanticIndexManager {
         index: self.index.unchecked
       )
     )
-    await self.indexTaskScheduler.schedule(priority: priority, taskDescription).value
-    for file in files {
-      self.indexStatus[file] = .upToDate
+    let updateIndexStoreTask = await self.indexTaskScheduler.schedule(priority: priority, taskDescription) { newState in
+      switch newState {
+      case .executing:
+        for file in files {
+          if case .scheduled(let task) = self.indexStatus[file] {
+            self.indexStatus[file] = .executing(task)
+          } else {
+            logger.fault(
+              """
+              Index status of \(file) is in an unexpected state \
+              '\(self.indexStatus[file]?.description ?? "<nil>", privacy: .public)' when update index store task \
+              started executing
+              """
+            )
+          }
+        }
+      case .cancelledToBeRescheduled:
+        for file in files {
+          if case .executing(let task) = self.indexStatus[file] {
+            self.indexStatus[file] = .scheduled(task)
+          } else {
+            logger.fault(
+              """
+              Index status of \(file) is in an unexpected state \
+              '\(self.indexStatus[file]?.description ?? "<nil>", privacy: .public)' when update index store task \
+              is cancelled to be rescheduled.
+              """
+            )
+          }
+        }
+      case .finished:
+        for file in files {
+          self.indexStatus[file] = .upToDate
+        }
+        self.indexTaskDidFinish()
+      }
     }
-    self.indexTaskDidFinish?()
+    await updateIndexStoreTask.value
   }
 
   /// Index the given set of files at the given priority.
@@ -226,9 +302,15 @@ public final actor SemanticIndexManager {
       }
       indexTasks.append(indexTask)
 
-      for file in targetsBatch.flatMap({ filesByTarget[$0]! }) {
-        indexStatus[file] = .inProgress(indexTask)
+      let filesToIndex = targetsBatch.flatMap({ filesByTarget[$0]! })
+      for file in filesToIndex {
+        // indexStatus will get set to `.upToDate` by `updateIndexStore`. Setting it to `.upToDate` cannot race with
+        // setting it to `.scheduled` because we don't have an `await` call between the creation of `indexTask` and
+        // this loop, so we still have exclusive access to the `SemanticIndexManager` actor and hence `updateIndexStore`
+        // can't execute until we have set all index statuses to `.scheduled`.
+        indexStatus[file] = .scheduled(indexTask)
       }
+      indexTasksWereScheduled(filesToIndex.count)
     }
     let indexTasksImmutable = indexTasks
 

Sources/SourceKitLSP/CMakeLists.txt

@@ -3,6 +3,7 @@ add_library(SourceKitLSP STATIC
   CapabilityRegistry.swift
   DocumentManager.swift
   DocumentSnapshot+FromFileContents.swift
+  IndexProgressManager.swift
   IndexStoreDB+MainFilesProvider.swift
   LanguageService.swift
   Rename.swift

Sources/SourceKitLSP/IndexProgressManager.swift

@@ -0,0 +1,103 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 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
+//
+//===----------------------------------------------------------------------===//
+
+import LanguageServerProtocol
+import SKSupport
+import SemanticIndex
+
+/// Listens for index status updates from `SemanticIndexManagers`. From that information, it manages a
+/// `WorkDoneProgress` that communicates the index progress to the editor.
+actor IndexProgressManager {
+  /// A queue on which `indexTaskWasQueued` and `indexStatusDidChange` are handled.
+  ///
+  /// This allows the two functions two be `nonisolated` (and eg. the caller of `indexStatusDidChange` doesn't have to
+  /// wait for the work done progress to be updated) while still guaranteeing that we handle them in the order they
+  /// were called.
+  private let queue = AsyncQueue<Serial>()
+
+  /// The `SourceKitLSPServer` for which this manages the index progress. It gathers all `SemanticIndexManagers` from
+  /// the workspaces in the `SourceKitLSPServer`.
+  private weak var sourceKitLSPServer: SourceKitLSPServer?
+
+  /// This is the target number of index tasks (eg. the `3` in `1/3 done`).
+  ///
+  /// Every time a new index task is scheduled, this number gets incremented, so that it only ever increases.
+  /// When indexing of one session is done (ie. when there are no more `scheduled` or `executing` tasks in any
+  /// `SemanticIndexManager`), `queuedIndexTasks` gets reset to 0 and the work done progress gets ended.
+  /// This way, when the next work done progress is started, it starts at zero again.
+  ///
+  /// The number of outstanding tasks is determined from the `scheduled` and `executing` tasks in all the
+  /// `SemanticIndexManager`s.
+  private var queuedIndexTasks = 0
+
+  /// While there are ongoing index tasks, a `WorkDoneProgressManager` that displays the work done progress.
+  private var workDoneProgress: WorkDoneProgressManager?
+
+  init(sourceKitLSPServer: SourceKitLSPServer) {
+    self.sourceKitLSPServer = sourceKitLSPServer
+  }
+
+  /// Called when a new file is scheduled to be indexed. Increments the target index count, eg. the 3 in `1/3`.
+  nonisolated func indexTaskWasQueued(count: Int) {
+    queue.async {
+      await self.indexTaskWasQueuedImpl(count: count)
+    }
+  }
+
+  private func indexTaskWasQueuedImpl(count: Int) async {
+    queuedIndexTasks += count
+    await indexStatusDidChangeImpl()
+  }
+
+  /// Called when a `SemanticIndexManager` finishes indexing a file. Adjusts the done index count, eg. the 1 in `1/3`.
+  nonisolated func indexStatusDidChange() {
+    queue.async {
+      await self.indexStatusDidChangeImpl()
+    }
+  }
+
+  private func indexStatusDidChangeImpl() async {
+    guard let sourceKitLSPServer else {
+      workDoneProgress = nil
+      return
+    }
+    var scheduled: [DocumentURI] = []
+    var executing: [DocumentURI] = []
+    for indexManager in await sourceKitLSPServer.workspaces.compactMap({ $0.semanticIndexManager }) {
+      let inProgress = await indexManager.inProgressIndexTasks
+      scheduled += inProgress.scheduled
+      executing += inProgress.executing
+    }
+
+    if scheduled.isEmpty && executing.isEmpty {
+      // Nothing left to index. Reset the target count and dismiss the work done progress.
+      queuedIndexTasks = 0
+      workDoneProgress = nil
+      return
+    }
+
+    let finishedTasks = queuedIndexTasks - scheduled.count - executing.count
+    let message = "\(finishedTasks) / \(queuedIndexTasks)"
+
+    let percentage = Int(Double(finishedTasks) / Double(queuedIndexTasks) * 100)
+    if let workDoneProgress {
+      workDoneProgress.update(message: message, percentage: percentage)
+    } else {
+      workDoneProgress = await WorkDoneProgressManager(
+        server: sourceKitLSPServer,
+        title: "Indexing",
+        message: message,
+        percentage: percentage
+      )
+    }
+  }
+}

Sources/SourceKitLSP/SourceKitLSPServer+Options.swift

@@ -49,6 +49,11 @@ extension SourceKitLSPServer {
     /// notification when running unit tests.
     public var swiftPublishDiagnosticsDebounceDuration: TimeInterval
 
+    /// A callback that is called when an index task finishes.
+    ///
+    /// Intended for testing purposes.
+    public var indexTaskDidFinish: (@Sendable () -> Void)?
+
     public init(
       buildSetup: BuildSetup = .default,
       clangdOptions: [String] = [],