Split up-to-date status tracking and index progress tracking

We were mixing the up-to-date status and in-progress status of an index task in `SemanticIndexManager`. This meant that a single `QueuedTask` in the task scheduler could be needed for eg. both preparation for editor functionality in a file of that target and to re-index a file in that target. This dual ownership made it unclear, which caller would be entitled to cancel the task. Furthermore, we needed to duplicate some logic from the preparation task dependencies in `SemanticIndexManager.prepare`.

To simplify things:
- Split the up-to-date status and the in-progress status into two different data structures
- Make the caller of `prepare` and `scheduleIndex` responsible for cancellation of the task it has scheduled. `TaskScheduler` might receive more scheduled tasks this way but the additional tasks should all be no-ops because the status is known to be up-to-date when they execute.

Author: ahoppen

Sources/SKCore/TaskScheduler.swift

@@ -87,7 +87,7 @@ public enum TaskExecutionState {
   case finished
 }
 
-fileprivate actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
+public actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
   /// Result of `executionTask` / the tasks in `executionTaskCreatedContinuation`.
   /// See doc comment on `executionTask`.
   enum ExecutionTaskFinishStatus {
@@ -147,14 +147,38 @@ fileprivate actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
   /// Gets reset every time `executionTask` finishes.
   nonisolated(unsafe) private var cancelledToBeRescheduled: 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
+  }
+
+  /// Wait for the task to finish.
+  ///
+  /// If the tasks that waits for this queued task to finished is cancelled, the QueuedTask will still continue
+  /// executing.
+  public func waitToFinish() async {
+    return await resultTask.value
+  }
+
+  /// Wait for the task to finish.
+  ///
+  /// If the tasks that waits for this queued task to finished is cancelled, the QueuedTask will also be cancelled.
+  /// This assumes that the caller of this method has unique control over the task and is the only one interested in its
+  /// value.
+  public func waitToFinishPropagatingCancellation() async {
+    return await resultTask.valuePropagatingCancellation
+  }
+
   /// 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)?
+  private let executionStateChangedCallback: (@Sendable (QueuedTask, TaskExecutionState) async -> Void)?
 
   init(
     priority: TaskPriority? = nil,
     description: TaskDescription,
-    executionStateChangedCallback: (@Sendable (TaskExecutionState) async -> Void)?
+    executionStateChangedCallback: (@Sendable (QueuedTask, TaskExecutionState) async -> Void)?
   ) async {
     self._priority = .init(initialValue: priority?.rawValue ?? Task.currentPriority.rawValue)
     self.description = description
@@ -214,19 +238,21 @@ fileprivate actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
     }
     executionTask = task
     executionTaskCreatedContinuation.yield(task)
-    await executionStateChangedCallback?(.executing)
+    _isExecuting.value = true
+    await executionStateChangedCallback?(self, .executing)
     return await task.value
   }
 
   /// Implementation detail of `execute` that is called after `self.description.execute()` finishes.
   private func finalizeExecution() async -> ExecutionTaskFinishStatus {
     self.executionTask = nil
+    _isExecuting.value = false
     if Task.isCancelled && self.cancelledToBeRescheduled.value {
-      await executionStateChangedCallback?(.cancelledToBeRescheduled)
+      await executionStateChangedCallback?(self, .cancelledToBeRescheduled)
       self.cancelledToBeRescheduled.value = false
       return ExecutionTaskFinishStatus.cancelledToBeRescheduled
     } else {
-      await executionStateChangedCallback?(.finished)
+      await executionStateChangedCallback?(self, .finished)
       return ExecutionTaskFinishStatus.terminated
     }
   }
@@ -327,8 +353,10 @@ public actor TaskScheduler<TaskDescription: TaskDescriptionProtocol> {
   public func schedule(
     priority: TaskPriority? = nil,
     _ taskDescription: TaskDescription,
-    @_inheritActorContext executionStateChangedCallback: (@Sendable (TaskExecutionState) async -> Void)? = nil
-  ) async -> Task<Void, Never> {
+    @_inheritActorContext executionStateChangedCallback: (
+      @Sendable (QueuedTask<TaskDescription>, TaskExecutionState) async -> Void
+    )? = nil
+  ) async -> QueuedTask<TaskDescription> {
     let queuedTask = await QueuedTask(
       priority: priority,
       description: taskDescription,
@@ -341,7 +369,7 @@ public actor TaskScheduler<TaskDescription: TaskDescriptionProtocol> {
       // queued task.
       await self.poke()
     }
-    return queuedTask.resultTask
+    return queuedTask
   }
 
   /// Trigger all queued tasks to update their priority.

Sources/SKSupport/Sequence+AsyncMap.swift

@@ -39,4 +39,19 @@ extension Sequence {
 
     return result
   }
+
+  /// Just like `Sequence.map` but allows an `async` transform function.
+  public func asyncFilter(
+    @_inheritActorContext _ predicate: @Sendable (Element) async throws -> Bool
+  ) async rethrows -> [Element] {
+    var result: [Element] = []
+
+    for element in self {
+      if try await predicate(element) {
+        result.append(element)
+      }
+    }
+
+    return result
+  }
 }

Sources/SemanticIndex/CMakeLists.txt

@@ -2,6 +2,7 @@
 add_library(SemanticIndex STATIC
   CheckedIndex.swift
   CompilerCommandLineOption.swift
+  IndexStatusManager.swift
   IndexTaskDescription.swift
   PreparationTaskDescription.swift
   SemanticIndexManager.swift

Sources/SemanticIndex/IndexStatusManager.swift

@@ -0,0 +1,70 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 Foundation
+import SKCore
+
+/// Keeps track of whether an item (a target or file to index) is up-to-date.
+actor IndexUpToDateStatusManager<Item: Hashable> {
+  private enum Status {
+    /// The item is up-to-date.
+    case upToDate
+
+    /// The target or file has been marked out-of-date at the given date.
+    ///
+    /// Keeping track of the date is necessary so that we don't mark a target as up-to-date if we have the following
+    /// ordering of events:
+    ///  - Preparation started
+    ///  - Target marked out of date
+    ///  - Preparation finished
+    case outOfDate(Date)
+  }
+
+  private var status: [Item: Status] = [:]
+
+  /// Mark the target or file as up-to-date from a preparation/update-indexstore operation started at
+  /// `updateOperationStartDate`.
+  ///
+  /// See comment on `Status.outOfDate` why `updateOperationStartDate` needs to be passed.
+  func markUpToDate(_ items: [Item], updateOperationStartDate: Date) {
+    for item in items {
+      switch status[item] {
+      case .upToDate:
+        break
+      case .outOfDate(let markedOutOfDate):
+        if markedOutOfDate < updateOperationStartDate {
+          status[item] = .upToDate
+        }
+      case nil:
+        status[item] = .upToDate
+      }
+    }
+  }
+
+  func markOutOfDate(_ items: some Collection<Item>) {
+    let date = Date()
+    for item in items {
+      status[item] = .outOfDate(date)
+    }
+  }
+
+  func markAllOutOfDate() {
+    markOutOfDate(status.keys)
+  }
+
+  func isUpToDate(_ item: Item) -> Bool {
+    if case .upToDate = status[item] {
+      return true
+    }
+    return false
+  }
+}

Sources/SemanticIndex/PreparationTaskDescription.swift

@@ -35,6 +35,8 @@ public struct PreparationTaskDescription: IndexTaskDescription {
   /// The build system manager that is used to get the toolchain and build settings for the files to index.
   private let buildSystemManager: BuildSystemManager
 
+  private let preparationUpToDateStatus: IndexUpToDateStatusManager<ConfiguredTarget>
+
   /// Test hooks that should be called when the preparation task finishes.
   private let testHooks: IndexTestHooks
 
@@ -54,10 +56,12 @@ public struct PreparationTaskDescription: IndexTaskDescription {
   init(
     targetsToPrepare: [ConfiguredTarget],
     buildSystemManager: BuildSystemManager,
+    preparationUpToDateStatus: IndexUpToDateStatusManager<ConfiguredTarget>,
     testHooks: IndexTestHooks
   ) {
     self.targetsToPrepare = targetsToPrepare
     self.buildSystemManager = buildSystemManager
+    self.preparationUpToDateStatus = preparationUpToDateStatus
     self.testHooks = testHooks
   }
 
@@ -66,17 +70,23 @@ 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)") {
-      let startDate = Date()
-      let targetsToPrepare = targetsToPrepare.sorted(by: {
+      let targetsToPrepare = await targetsToPrepare.asyncFilter {
+        await !preparationUpToDateStatus.isUpToDate($0)
+      }.sorted(by: {
         ($0.targetID, $0.runDestinationID) < ($1.targetID, $1.runDestinationID)
       })
+      if targetsToPrepare.isEmpty {
+        return
+      }
+
       let targetsToPrepareDescription =
         targetsToPrepare
         .map { "\($0.targetID)-\($0.runDestinationID)" }
         .joined(separator: ", ")
       logger.log(
         "Starting preparation with priority \(Task.currentPriority.rawValue, privacy: .public): \(targetsToPrepareDescription)"
       )
+      let startDate = Date()
       do {
         try await buildSystemManager.prepare(targets: targetsToPrepare)
       } catch {
@@ -85,6 +95,9 @@ public struct PreparationTaskDescription: IndexTaskDescription {
         )
       }
       await testHooks.preparationTaskDidFinish?(self)
+      if !Task.isCancelled {
+        await preparationUpToDateStatus.markUpToDate(targetsToPrepare, updateOperationStartDate: startDate)
+      }
       logger.log(
         "Finished preparation in \(Date().timeIntervalSince(startDate) * 1000, privacy: .public)ms: \(targetsToPrepareDescription)"
       )