Introduce test hooks that can be used to monitor when preparation and index tasks finish

Author: ahoppen

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

@@ -35,6 +35,9 @@ 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
 
+  /// 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)"
         )
       }
+      testHooks.preparationTaskDidFinish?(self)
       logger.log(
         "Finished preparation in \(Date().timeIntervalSince(startDate) * 1000, privacy: .public)ms: \(targetsToPrepareDescription)"
       )

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -60,6 +60,8 @@ 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>?
@@ -121,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
@@ -296,7 +300,8 @@ public final actor SemanticIndexManager {
     let taskDescription = AnyIndexTaskDescription(
       PreparationTaskDescription(
         targetsToPrepare: targetsToPrepare,
-        buildSystemManager: self.buildSystemManager
+        buildSystemManager: self.buildSystemManager,
+        testHooks: testHooks
       )
     )
     if !targetsToPrepare.isEmpty {
@@ -355,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) -> Void)?
+
+  /// A callback that is called when an index task finishes.
+  public var updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) -> Void)?
+
+  public init(
+    preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) -> Void)? = nil,
+    updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) -> 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)
       }
+      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,16 +23,6 @@ extension SourceKitLSPServer {
 
   /// Configuration options for the SourceKitServer.
   public struct Options: Sendable {
-    /// Callbacks that allow inspection of internal state modifications during testing.
-    public struct TestHooks: Sendable {
-      /// A callback that is called when an index task finishes.
-      public var indexTaskDidFinish: (@Sendable () -> Void)?
-
-      public init(indexTaskDidFinish: (@Sendable () -> Void)? = nil) {
-        self.indexTaskDidFinish = indexTaskDidFinish
-      }
-    }
-
     /// Additional compiler flags (e.g. `-Xswiftc` for SwiftPM projects) and other build-related
     /// configuration.
     public var buildSetup: BuildSetup
@@ -58,7 +49,7 @@ extension SourceKitLSPServer {
     /// notification when running unit tests.
     public var swiftPublishDiagnosticsDebounceDuration: TimeInterval
 
-    public var testHooks: TestHooks
+    public var indexTestHooks: IndexTestHooks
 
     public init(
       buildSetup: BuildSetup = .default,
@@ -68,7 +59,7 @@ extension SourceKitLSPServer {
       completionOptions: SKCompletionOptions = .init(),
       generatedInterfacesPath: AbsolutePath = defaultDirectoryForGeneratedInterfaces,
       swiftPublishDiagnosticsDebounceDuration: TimeInterval = 2, /* 2s */
-      testHooks: TestHooks = TestHooks()
+      indexTestHooks: IndexTestHooks = IndexTestHooks()
     ) {
       self.buildSetup = buildSetup
       self.clangdOptions = clangdOptions
@@ -77,7 +68,7 @@ extension SourceKitLSPServer {
       self.completionOptions = completionOptions
       self.generatedInterfacesPath = generatedInterfacesPath
       self.swiftPublishDiagnosticsDebounceDuration = swiftPublishDiagnosticsDebounceDuration
-      self.testHooks = testHooks
+      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.testHooks.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.testHooks.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

Tests/SourceKitLSPTests/BackgroundIndexingTests.swift

@@ -165,11 +165,11 @@ final class BackgroundIndexingTests: XCTestCase {
 
   func testBackgroundIndexingHappensWithLowPriority() async throws {
     var serverOptions = backgroundIndexingOptions
-    serverOptions.testHooks.indexTaskDidFinish = {
-      XCTAssert(
-        Task.currentPriority == .low,
-        "An index task ran with priority \(Task.currentPriority)"
-      )
+    serverOptions.indexTestHooks.preparationTaskDidFinish = { taskDescription in
+      XCTAssert(Task.currentPriority == .low, "\(taskDescription) ran with priority \(Task.currentPriority)")
+    }
+    serverOptions.indexTestHooks.updateIndexStoreTaskDidFinish = { taskDescription in
+      XCTAssert(Task.currentPriority == .low, "\(taskDescription) ran with priority \(Task.currentPriority)")
     }
     let project = try await SwiftPMTestProject(
       files: [