Add infrastructure to test which targets are being prepared

Sources/SKTestSupport/MultiFileTestProject.swift

@@ -84,6 +84,7 @@ public class MultiFileTestProject {
     serverOptions: SourceKitLSPServer.Options = .testDefault,
     usePullDiagnostics: Bool = true,
     preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
+    cleanUp: (() -> Void)? = nil,
     testName: String = #function
   ) async throws {
     scratchDirectory = try testScratchDir(testName: testName)
@@ -122,6 +123,7 @@ public class MultiFileTestProject {
         if cleanScratchDirectories {
           try? FileManager.default.removeItem(at: scratchDirectory)
         }
+        cleanUp?()
       }
     )
   }

Sources/SKTestSupport/SwiftPMTestProject.swift

@@ -47,6 +47,7 @@ public class SwiftPMTestProject: MultiFileTestProject {
     pollIndex: Bool = true,
     preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     usePullDiagnostics: Bool = true,
+    cleanUp: (() -> Void)? = nil,
     testName: String = #function
   ) async throws {
     var filesByPath: [RelativeFileLocation: String] = [:]
@@ -72,6 +73,7 @@ public class SwiftPMTestProject: MultiFileTestProject {
       serverOptions: serverOptions,
       usePullDiagnostics: usePullDiagnostics,
       preInitialization: preInitialization,
+      cleanUp: cleanUp,
       testName: testName
     )
 

Sources/SemanticIndex/CMakeLists.txt

@@ -4,6 +4,7 @@ add_library(SemanticIndex STATIC
   IndexTaskDescription.swift
   PreparationTaskDescription.swift
   SemanticIndexManager.swift
+  TestHooks.swift
   UpdateIndexStoreTaskDescription.swift
 )
 set_target_properties(SemanticIndex PROPERTIES

Sources/SemanticIndex/PreparationTaskDescription.swift

@@ -30,11 +30,14 @@ public struct PreparationTaskDescription: IndexTaskDescription {
   public let id = preparationIDForLogging.fetchAndIncrement()
 
   /// The targets that should be prepared.
-  private let targetsToPrepare: [ConfiguredTarget]
+  public let targetsToPrepare: [ConfiguredTarget]
 
   /// The build system manager that is used to get the toolchain and build settings for the files to index.
   private let buildSystemManager: BuildSystemManager
 
+  /// Test hooks that should be called when the preparation task finishes.
+  private let testHooks: IndexTestHooks
+
   /// The task is idempotent because preparing the same target twice produces the same result as preparing it once.
   public var isIdempotent: Bool { true }
 
@@ -50,10 +53,12 @@ public struct PreparationTaskDescription: IndexTaskDescription {
 
   init(
     targetsToPrepare: [ConfiguredTarget],
-    buildSystemManager: BuildSystemManager
+    buildSystemManager: BuildSystemManager,
+    testHooks: IndexTestHooks
   ) {
     self.targetsToPrepare = targetsToPrepare
     self.buildSystemManager = buildSystemManager
+    self.testHooks = testHooks
   }
 
   public func execute() async {
@@ -79,6 +84,7 @@ public struct PreparationTaskDescription: IndexTaskDescription {
           "Preparation failed: \(error.forLogging)"
         )
       }
+      await testHooks.preparationTaskDidFinish?(self)
       logger.log(
         "Finished preparation in \(Date().timeIntervalSince(startDate) * 1000, privacy: .public)ms: \(targetsToPrepareDescription)"
       )

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -38,6 +38,19 @@ private enum IndexStatus<T> {
   }
 }
 
+/// See `SemanticIndexManager.preparationStatus`
+private struct PreparationTaskStatusData {
+  /// A UUID to track the task. This is used to ensure that status updates from this task don't update
+  /// `preparationStatus` for targets that are tracked by a different task.
+  let taskID: UUID
+
+  /// The list of targets that are being prepared in a joint preparation operation.
+  let targets: [ConfiguredTarget]
+
+  /// The task that prepares the target
+  let task: Task<Void, Never>
+}
+
 /// Schedules index tasks and keeps track of the index status of files.
 public final actor SemanticIndexManager {
   /// The underlying index. This is used to check if the index of a file is already up-to-date, in which case it doesn't
@@ -47,20 +60,16 @@ public final actor SemanticIndexManager {
   /// The build system manager that is used to get compiler arguments for a file.
   private let buildSystemManager: BuildSystemManager
 
+  private let testHooks: IndexTestHooks
+
   /// The task to generate the build graph (resolving package dependencies, generating the build description,
   /// ...). `nil` if no build graph is currently being generated.
   private var generateBuildGraphTask: Task<Void, Never>?
 
   /// The preparation status of the targets that the `SemanticIndexManager` has started preparation for.
   ///
   /// Targets will be removed from this dictionary when they are no longer known to be up-to-date.
-  ///
-  /// The associated values of the `IndexStatus` are:
-  ///  - A UUID to track the task. This is used to ensure that status updates from this task don't update
-  ///    `preparationStatus` for targets that are tracked by a different task.
-  ///  - The list of targets that are being prepared in a joint preparation operation
-  ///  - The task that prepares the target
-  private var preparationStatus: [ConfiguredTarget: IndexStatus<(UUID, [ConfiguredTarget], Task<Void, Never>)>] = [:]
+  private var preparationStatus: [ConfiguredTarget: IndexStatus<PreparationTaskStatusData>] = [:]
 
   /// The index status of the source files that the `SemanticIndexManager` knows about.
   ///
@@ -114,12 +123,14 @@ public final actor SemanticIndexManager {
   public init(
     index: UncheckedIndex,
     buildSystemManager: BuildSystemManager,
+    testHooks: IndexTestHooks,
     indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>,
     indexTasksWereScheduled: @escaping @Sendable (Int) -> Void,
     indexTaskDidFinish: @escaping @Sendable () -> Void
   ) {
     self.index = index
     self.buildSystemManager = buildSystemManager
+    self.testHooks = testHooks
     self.indexTaskScheduler = indexTaskScheduler
     self.indexTasksWereScheduled = indexTasksWereScheduled
     self.indexTaskDidFinish = indexTaskDidFinish
@@ -268,16 +279,16 @@ public final actor SemanticIndexManager {
       switch preparationStatus[target] {
       case .upToDate:
         break
-      case .scheduled((_, let existingTaskTargets, let task)), .executing((_, let existingTaskTargets, let task)):
+      case .scheduled(let existingTaskData), .executing(let existingTaskData):
         // If we already have a task scheduled that prepares fewer targets, await that instead of overriding the
         // target's preparation status with a longer-running task. The key benefit here is that when we get many
         // preparation requests for the same target (eg. one for every text document request sent to a file), we don't
         // re-create new `PreparationTaskDescription`s for every preparation request. Instead, all the preparation
         // requests await the same task. At the same time, if we have a multi-file preparation request and then get a
         // single-file preparation request, we will override the preparation of that target with the single-file
         // preparation task, ensuring that the task gets prepared as quickly as possible.
-        if existingTaskTargets.count <= targets.count {
-          preparationTasksToAwait.append(task)
+        if existingTaskData.targets.count <= targets.count {
+          preparationTasksToAwait.append(existingTaskData.task)
         } else {
           targetsToPrepare.append(target)
         }
@@ -289,7 +300,8 @@ public final actor SemanticIndexManager {
     let taskDescription = AnyIndexTaskDescription(
       PreparationTaskDescription(
         targetsToPrepare: targetsToPrepare,
-        buildSystemManager: self.buildSystemManager
+        buildSystemManager: self.buildSystemManager,
+        testHooks: testHooks
       )
     )
     if !targetsToPrepare.isEmpty {
@@ -301,20 +313,22 @@ public final actor SemanticIndexManager {
         switch newState {
         case .executing:
           for target in targetsToPrepare {
-            if case .scheduled((taskID, let targets, let task)) = self.preparationStatus[target] {
-              self.preparationStatus[target] = .executing((taskID, targets, task))
+            if case .scheduled(let existingTaskData) = self.preparationStatus[target], existingTaskData.taskID == taskID
+            {
+              self.preparationStatus[target] = .executing(existingTaskData)
             }
           }
         case .cancelledToBeRescheduled:
           for target in targetsToPrepare {
-            if case .executing((taskID, let targets, let task)) = self.preparationStatus[target] {
-              self.preparationStatus[target] = .scheduled((taskID, targets, task))
+            if case .executing(let existingTaskData) = self.preparationStatus[target], existingTaskData.taskID == taskID
+            {
+              self.preparationStatus[target] = .scheduled(existingTaskData)
             }
           }
         case .finished:
           for target in targetsToPrepare {
             switch self.preparationStatus[target] {
-            case .executing((taskID, _, _)):
+            case .executing(let existingTaskData) where existingTaskData.taskID == taskID:
               self.preparationStatus[target] = .upToDate
             default:
               break
@@ -324,14 +338,16 @@ public final actor SemanticIndexManager {
         }
       }
       for target in targetsToPrepare {
-        preparationStatus[target] = .scheduled((taskID, targetsToPrepare, preparationTask))
+        preparationStatus[target] = .scheduled(
+          PreparationTaskStatusData(taskID: taskID, targets: targetsToPrepare, task: preparationTask)
+        )
       }
       preparationTasksToAwait.append(preparationTask)
     }
     await withTaskGroup(of: Void.self) { taskGroup in
       for task in preparationTasksToAwait {
         taskGroup.addTask {
-          await task.value
+          await task.valuePropagatingCancellation
         }
       }
       await taskGroup.waitForAll()
@@ -344,7 +360,8 @@ public final actor SemanticIndexManager {
       UpdateIndexStoreTaskDescription(
         filesToIndex: Set(files),
         buildSystemManager: self.buildSystemManager,
-        index: index
+        index: index,
+        testHooks: testHooks
       )
     )
     let updateIndexStoreTask = await self.indexTaskScheduler.schedule(priority: priority, taskDescription) { newState in

Sources/SemanticIndex/TestHooks.swift

@@ -0,0 +1,27 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+/// Callbacks that allow inspection of internal state modifications during testing.
+public struct IndexTestHooks: Sendable {
+  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(
+    preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) async -> Void)? = nil,
+    updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)? = nil
+  ) {
+    self.preparationTaskDidFinish = preparationTaskDidFinish
+    self.updateIndexStoreTaskDidFinish = updateIndexStoreTaskDidFinish
+  }
+}

Sources/SemanticIndex/UpdateIndexStoreTaskDescription.swift

@@ -49,6 +49,9 @@ public struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
   /// case we don't need to index it again.
   private let index: UncheckedIndex
 
+  /// Test hooks that should be called when the index task finishes.
+  private let testHooks: IndexTestHooks
+
   /// The task is idempotent because indexing the same file twice produces the same result as indexing it once.
   public var isIdempotent: Bool { true }
 
@@ -65,11 +68,13 @@ public struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
   init(
     filesToIndex: Set<FileToIndex>,
     buildSystemManager: BuildSystemManager,
-    index: UncheckedIndex
+    index: UncheckedIndex,
+    testHooks: IndexTestHooks
   ) {
     self.filesToIndex = filesToIndex
     self.buildSystemManager = buildSystemManager
     self.index = index
+    self.testHooks = testHooks
   }
 
   public func execute() async {
@@ -94,6 +99,7 @@ public struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
       for file in filesToIndex {
         await updateIndexStoreForSingleFile(file)
       }
+      await testHooks.updateIndexStoreTaskDidFinish?(self)
       logger.log(
         "Finished updating index store in \(Date().timeIntervalSince(startDate) * 1000, privacy: .public)ms: \(filesToIndexDescription)"
       )

Sources/SourceKitLSP/SourceKitLSPServer+Options.swift

@@ -14,6 +14,7 @@ import Foundation
 import LanguageServerProtocol
 import SKCore
 import SKSupport
+import SemanticIndex
 
 import struct TSCBasic.AbsolutePath
 import struct TSCBasic.RelativePath
@@ -22,7 +23,6 @@ extension SourceKitLSPServer {
 
   /// Configuration options for the SourceKitServer.
   public struct Options: Sendable {
-
     /// Additional compiler flags (e.g. `-Xswiftc` for SwiftPM projects) and other build-related
     /// configuration.
     public var buildSetup: BuildSetup
@@ -49,10 +49,7 @@ 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 var indexTestHooks: IndexTestHooks
 
     public init(
       buildSetup: BuildSetup = .default,
@@ -61,7 +58,8 @@ extension SourceKitLSPServer {
       indexOptions: IndexOptions = .init(),
       completionOptions: SKCompletionOptions = .init(),
       generatedInterfacesPath: AbsolutePath = defaultDirectoryForGeneratedInterfaces,
-      swiftPublishDiagnosticsDebounceDuration: TimeInterval = 2 /* 2s */
+      swiftPublishDiagnosticsDebounceDuration: TimeInterval = 2, /* 2s */
+      indexTestHooks: IndexTestHooks = IndexTestHooks()
     ) {
       self.buildSetup = buildSetup
       self.clangdOptions = clangdOptions
@@ -70,6 +68,7 @@ extension SourceKitLSPServer {
       self.completionOptions = completionOptions
       self.generatedInterfacesPath = generatedInterfacesPath
       self.swiftPublishDiagnosticsDebounceDuration = swiftPublishDiagnosticsDebounceDuration
+      self.indexTestHooks = indexTestHooks
     }
   }
 }

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -1211,7 +1211,6 @@ extension SourceKitLSPServer {
       logger.log("Cannot open workspace before server is initialized")
       return nil
     }
-    let indexTaskDidFinishCallback = options.indexTaskDidFinish
     var options = self.options
     options.buildSetup = self.options.buildSetup.merging(buildSetup(for: workspaceFolder))
     return try? await Workspace(
@@ -1237,7 +1236,6 @@ extension SourceKitLSPServer {
       },
       indexTaskDidFinish: { [weak self] in
         self?.indexProgressManager.indexStatusDidChange()
-        indexTaskDidFinishCallback?()
       }
     )
   }
@@ -1286,7 +1284,6 @@ extension SourceKitLSPServer {
       if self.workspaces.isEmpty {
         logger.error("no workspace found")
 
-        let indexTaskDidFinishCallback = self.options.indexTaskDidFinish
         let workspace = await Workspace(
           documentManager: self.documentManager,
           rootUri: req.rootURI,
@@ -1302,7 +1299,6 @@ extension SourceKitLSPServer {
           },
           indexTaskDidFinish: { [weak self] in
             self?.indexProgressManager.indexStatusDidChange()
-            indexTaskDidFinishCallback?()
           }
         )
 

Sources/SourceKitLSP/Workspace.swift

@@ -112,6 +112,7 @@ public final class Workspace: Sendable {
       self.semanticIndexManager = SemanticIndexManager(
         index: uncheckedIndex,
         buildSystemManager: buildSystemManager,
+        testHooks: options.indexTestHooks,
         indexTaskScheduler: indexTaskScheduler,
         indexTasksWereScheduled: indexTasksWereScheduled,
         indexTaskDidFinish: indexTaskDidFinish