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

We were blocking the initialization response on `self.buildSystemManager.testFiles`, which requires the list of test files to be determined. Make that operation asynchronous so that a slow build server can’t take down all of SourceKit-LSP.

Author: ahoppen

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/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,
@@ -223,7 +225,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
         }

Tests/BuildSystemIntegrationTests/BuildServerBuildSystemTests.swift

@@ -16,6 +16,8 @@ import Foundation
 import LanguageServerProtocol
 import LanguageServerProtocolExtensions
 import SKTestSupport
+import SourceKitLSP
+import SwiftExtensions
 import TSCBasic
 import XCTest
 
@@ -369,4 +371,47 @@ final class BuildServerBuildSystemTests: XCTestCase {
       return diags.fullReport?.items.map(\.message) == ["DEBUG SET"]
     }
   }
+
+  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:
+          sleep(UInt32(defaultTimeout * 10))
+          break
+        }
+      }
+    }
+
+    // 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
+        })
+      )
+    )
+  }
 }