Do not block initialization of the build server when build server is unresponsive in returning the list of test files

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -204,17 +204,20 @@ package final class TestSourceKitLSPClient: MessageHandler, Sendable {
     }
     preInitialization?(self)
     if initialize {
-      _ = try await self.send(
-        InitializeRequest(
-          processId: nil,
-          rootPath: nil,
-          rootURI: nil,
-          initializationOptions: initializationOptions,
-          capabilities: capabilities,
-          trace: .off,
-          workspaceFolders: workspaceFolders
+      let capabilities = capabilities
+      try await withTimeout(.seconds(defaultTimeout)) {
+        _ = try await self.send(
+          InitializeRequest(
+            processId: nil,
+            rootPath: nil,
+            rootURI: nil,
+            initializationOptions: initializationOptions,
+            capabilities: capabilities,
+            trace: .off,
+            workspaceFolders: workspaceFolders
+          )
         )
-      )
+      }
     }
   }
 

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -339,13 +339,7 @@ package final actor SemanticIndexManager {
   }
 
   private func waitForBuildGraphGenerationTasks() async {
-    await withTaskGroup(of: Void.self) { taskGroup in
-      for generateBuildGraphTask in scheduleIndexingTasks.values {
-        taskGroup.addTask {
-          await generateBuildGraphTask.value
-        }
-      }
-    }
+    await scheduleIndexingTasks.values.concurrentForEach { await $0.value }
   }
 
   /// Wait for all in-progress index tasks to finish.
@@ -355,18 +349,13 @@ package final actor SemanticIndexManager {
     // can await the index tasks below.
     await waitForBuildGraphGenerationTasks()
 
-    await withTaskGroup(of: Void.self) { taskGroup in
-      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 inProgressIndexTasks.concurrentForEach { (_, status) in
+      switch status {
+      case .waitingForPreparation(preparationTaskID: _, indexTask: let indexTask),
+        .preparing(preparationTaskID: _, indexTask: let indexTask),
+        .updatingIndexStore(updateIndexStoreTask: _, indexTask: let indexTask):
+        await indexTask.value
       }
-      await taskGroup.waitForAll()
     }
     index.pollForUnitChangesAndWait()
     logger.debug("Done waiting for up-to-date index")
@@ -786,17 +775,9 @@ package final actor SemanticIndexManager {
       }
       indexTasksWereScheduled(newIndexTasks)
     }
-    let indexTasksImmutable = indexTasks
 
     return Task(priority: priority) {
-      await withTaskGroup(of: Void.self) { taskGroup in
-        for indexTask in indexTasksImmutable {
-          taskGroup.addTask {
-            await indexTask.value
-          }
-        }
-        await taskGroup.waitForAll()
-      }
+      await indexTasks.concurrentForEach { await $0.value }
     }
   }
 }

Sources/SourceKitLSP/Swift/SyntacticTestIndex.swift

@@ -18,13 +18,28 @@ import SwiftExtensions
 
 /// Task metadata for `SyntacticTestIndexer.indexingQueue`
 fileprivate enum TaskMetadata: DependencyTracker, Equatable {
-  case read
+  /// Determine the list of test files from the build system and scan them for tests. Only created when the
+  /// `SyntacticTestIndex` is created
+  case initialPopulation
+
+  /// Index the files in the given set for tests
   case index(Set<DocumentURI>)
 
+  /// Retrieve information about syntactically discovered tests from the index.
+  case read
+
   /// Reads can be concurrent and files can be indexed concurrently. But we need to wait for all files to finish
   /// indexing before reading the index.
   func isDependency(of other: TaskMetadata) -> Bool {
     switch (self, other) {
+    case (.initialPopulation, _):
+      // The initial population need to finish before we can do anything with the task.
+      return true
+    case (_, .initialPopulation):
+      // Should never happen because the initial population should only be scheduled once before any other operations
+      // on the test index. But be conservative in case we do get an `initialPopulation` somewhere in between and use it
+      // as a full blocker on the queue.
+      return true
     case (.read, .read):
       // We allow concurrent reads
       return false
@@ -109,7 +124,23 @@ actor SyntacticTestIndex {
   /// indexing tasks to finish.
   private let indexingQueue = AsyncQueue<TaskMetadata>()
 
-  init() {}
+  init(determineTestFiles: @Sendable @escaping () async -> [DocumentURI]) {
+    indexingQueue.async(priority: .low, metadata: .initialPopulation) {
+      let testFiles = await determineTestFiles()
+
+      // Divide the files into multiple batches. This is more efficient than spawning a new task for every file, mostly
+      // because it keeps the number of pending items in `indexingQueue` low and adding a new task to `indexingQueue` is
+      // in O(number of pending tasks), since we need to scan for dependency edges to add, which would make scanning files
+      // be O(number of files).
+      // Over-subscribe the processor count in case one batch finishes more quickly than another.
+      let batches = testFiles.partition(intoNumberOfBatches: ProcessInfo.processInfo.processorCount * 4)
+      await batches.concurrentForEach { filesInBatch in
+        for uri in filesInBatch {
+          await self.rescanFileAssumingOnQueue(uri)
+        }
+      }
+    }
+  }
 
   private func removeFilesFromIndex(_ removedFiles: Set<DocumentURI>) {
     self.removedFiles.formUnion(removedFiles)

Sources/SourceKitLSP/Workspace.swift

@@ -127,7 +127,7 @@ package final class Workspace: Sendable, BuildSystemManagerDelegate {
   }
 
   /// The index that syntactically scans the workspace for tests.
-  let syntacticTestIndex = SyntacticTestIndex()
+  let syntacticTestIndex: SyntacticTestIndex
 
   /// Language service for an open document, if available.
   private let documentService: ThreadSafeBox<[DocumentURI: LanguageService]> = ThreadSafeBox(initialValue: [:])
@@ -180,13 +180,15 @@ package final class Workspace: Sendable, BuildSystemManagerDelegate {
     } else {
       self.semanticIndexManager = nil
     }
+    // Trigger an initial population of `syntacticTestIndex`.
+    self.syntacticTestIndex = SyntacticTestIndex(determineTestFiles: {
+      await orLog("Getting list of test files for initial syntactic index population") {
+        try await buildSystemManager.testFiles()
+      } ?? []
+    })
     await indexDelegate?.addMainFileChangedCallback { [weak self] in
       await self?.buildSystemManager.mainFilesChanged()
     }
-    // Trigger an initial population of `syntacticTestIndex`.
-    if let testFiles = await orLog("Getting initial test files", { try await self.buildSystemManager.testFiles() }) {
-      await syntacticTestIndex.listOfTestFilesDidChange(testFiles)
-    }
     if let semanticIndexManager {
       await semanticIndexManager.scheduleBuildGraphGenerationAndBackgroundIndexAllFiles(
         filesToIndex: nil,
@@ -224,7 +226,7 @@ package final class Workspace: Sendable, BuildSystemManagerDelegate {
           // `SourceKitLSPServer` has been destructed. We are tearing down the
           // language server. Nothing left to do.
           logger.error(
-            "Ignoring notificaiton \(type(of: notification).method) because connection to editor has been closed"
+            "Ignoring notification \(type(of: notification).method) because connection to editor has been closed"
           )
           return
         }

Sources/SwiftExtensions/AsyncUtils.swift

@@ -165,6 +165,17 @@ extension Collection where Element: Sendable {
       count = indexedResults.count
     }
   }
+
+  /// Invoke `body` for every element in the collection and wait for all calls of `body` to finish
+  package func concurrentForEach(_ body: @escaping @Sendable (Element) async -> Void) async {
+    await withTaskGroup(of: Void.self) { taskGroup in
+      for element in self {
+        taskGroup.addTask {
+          await body(element)
+        }
+      }
+    }
+  }
 }
 
 package struct TimeoutError: Error, CustomStringConvertible {

Tests/BuildSystemIntegrationTests/BuildServerBuildSystemTests.swift

@@ -17,9 +17,15 @@ import LanguageServerProtocol
 import LanguageServerProtocolExtensions
 import SKOptions
 import SKTestSupport
+import SourceKitLSP
+import SwiftExtensions
 import TSCBasic
 import XCTest
 
+#if os(Windows)
+import WinSDK
+#endif
+
 final class BuildServerBuildSystemTests: XCTestCase {
   func testBuildSettingsFromBuildServer() async throws {
     let project = try await BuildServerTestProject(
@@ -517,4 +523,51 @@ final class BuildServerBuildSystemTests: XCTestCase {
     // We currently pick the canonical target alphabetically, which means that `bsp://first` wins over `bsp://second`
     XCTAssert(try XCTUnwrap(optionsWithoutTarget).compilerArguments.contains("-DFIRST"))
   }
+
+  func testDontBlockBuildServerInitializationIfBuildSystemIsUnresponsive() async throws {
+    // A build server that responds to the initialize request but not to any other requests.
+    final class UnresponsiveBuildServer: MessageHandler {
+      func handle(_ notification: some LanguageServerProtocol.NotificationType) {}
+
+      func handle<Request: RequestType>(
+        _ request: Request,
+        id: RequestID,
+        reply: @escaping @Sendable (LSPResult<Request.Response>) -> Void
+      ) {
+        switch request {
+        case is InitializeBuildRequest:
+          reply(
+            .success(
+              InitializeBuildResponse(
+                displayName: "UnresponsiveBuildServer",
+                version: "",
+                bspVersion: "2.2.0",
+                capabilities: BuildServerCapabilities()
+              ) as! Request.Response
+            )
+          )
+        default:
+          #if os(Windows)
+          Sleep(60 * 60 * 1000 /*ms*/)
+          #else
+          sleep(60 * 60 /*s*/)
+          #endif
+          XCTFail("Build server should be terminated before finishing the timeout")
+        }
+      }
+    }
+
+    // Creating the `MultiFileTestProject` waits for the initialize response and times out if it doesn't receive one.
+    // Make sure that we get that response back.
+    _ = try await MultiFileTestProject(
+      files: ["Test.swift": ""],
+      hooks: Hooks(
+        buildSystemHooks: BuildSystemHooks(injectBuildServer: { _, _ in
+          let connection = LocalConnection(receiverName: "Unresponsive build system")
+          connection.start(handler: UnresponsiveBuildServer())
+          return connection
+        })
+      )
+    )
+  }
 }