Allow alternate build systems to have larger indexing batch sizes

Author: rockbruno

Sources/BuildServerIntegration/BuildServerManager.swift

@@ -322,9 +322,13 @@ package actor BuildServerManager: QueueBasedMessageHandler {
 
   private var connectionToClient: BuildServerManagerConnectionToClient
 
-  /// The build serer adapter that is used to answer build server queries.
+  /// The build server adapter that is used to answer build server queries.
   private var buildServerAdapter: BuildServerAdapter?
 
+  /// The kind of underlying build server adapter that is being managed.
+  /// `nil` if the `BuildServerManager` does not have an underlying build system.
+  private var kind: BuildServerSpec.Kind?
+
   /// The build server adapter after initialization finishes. When sending messages to the BSP server, this should be
   /// preferred over `buildServerAdapter` because no messages must be sent to the build server before initialization
   /// finishes.
@@ -451,6 +455,7 @@ package actor BuildServerManager: QueueBasedMessageHandler {
     self.toolchainRegistry = toolchainRegistry
     self.options = options
     self.connectionToClient = connectionToClient
+    self.kind = buildServerSpec?.kind
     self.configPath = buildServerSpec?.configPath
     self.buildServerAdapter = await buildServerSpec?.createBuildServerAdapter(
       toolchainRegistry: toolchainRegistry,

Sources/BuildServerProtocol/Messages/InitializeBuildRequest.swift

@@ -278,6 +278,9 @@ public struct SourceKitInitializeBuildResponseData: LSPAnyCodable, Codable, Send
   /// Whether the server implements the `textDocument/sourceKitOptions` request.
   public var sourceKitOptionsProvider: Bool?
 
+  /// The number of targets to prepare concurrently, when an index request is scheduled.
+  public var indexTaskBatchSize: Int?
+
   /// The files to watch for changes.
   public var watchers: [FileSystemWatcher]?
 
@@ -286,12 +289,14 @@ public struct SourceKitInitializeBuildResponseData: LSPAnyCodable, Codable, Send
   public init(
     indexDatabasePath: String? = nil,
     indexStorePath: String? = nil,
+    indexTaskBatchSize: Int? = nil,
     watchers: [FileSystemWatcher]? = nil,
     prepareProvider: Bool? = nil,
     sourceKitOptionsProvider: Bool? = nil
   ) {
     self.indexDatabasePath = indexDatabasePath
     self.indexStorePath = indexStorePath
+    self.indexTaskBatchSize = indexTaskBatchSize
     self.watchers = watchers
     self.prepareProvider = prepareProvider
     self.sourceKitOptionsProvider = sourceKitOptionsProvider
@@ -300,13 +305,15 @@ public struct SourceKitInitializeBuildResponseData: LSPAnyCodable, Codable, Send
   public init(
     indexDatabasePath: String? = nil,
     indexStorePath: String? = nil,
+    indexTaskBatchSize: Int? = nil,
     outputPathsProvider: Bool? = nil,
     prepareProvider: Bool? = nil,
     sourceKitOptionsProvider: Bool? = nil,
     watchers: [FileSystemWatcher]? = nil
   ) {
     self.indexDatabasePath = indexDatabasePath
     self.indexStorePath = indexStorePath
+    self.indexTaskBatchSize = indexTaskBatchSize
     self.outputPathsProvider = outputPathsProvider
     self.prepareProvider = prepareProvider
     self.sourceKitOptionsProvider = sourceKitOptionsProvider
@@ -320,6 +327,9 @@ public struct SourceKitInitializeBuildResponseData: LSPAnyCodable, Codable, Send
     if case .string(let indexStorePath) = dictionary[CodingKeys.indexStorePath.stringValue] {
       self.indexStorePath = indexStorePath
     }
+    if case .int(let indexTaskBatchSize) = dictionary[CodingKeys.indexTaskBatchSize.stringValue] {
+      self.indexTaskBatchSize = indexTaskBatchSize
+    }
     if case .bool(let outputPathsProvider) = dictionary[CodingKeys.outputPathsProvider.stringValue] {
       self.outputPathsProvider = outputPathsProvider
     }
@@ -342,6 +352,9 @@ public struct SourceKitInitializeBuildResponseData: LSPAnyCodable, Codable, Send
     if let indexStorePath {
       result[CodingKeys.indexStorePath.stringValue] = .string(indexStorePath)
     }
+    if let indexTaskBatchSize {
+      result[CodingKeys.indexTaskBatchSize.stringValue] = .int(indexTaskBatchSize)
+    }
     if let outputPathsProvider {
       result[CodingKeys.outputPathsProvider.stringValue] = .bool(outputPathsProvider)
     }

Sources/SKTestSupport/CustomBuildServerTestProject.swift

@@ -175,12 +175,14 @@ package extension CustomBuildServer {
 
   func initializationResponseSupportingBackgroundIndexing(
     projectRoot: URL,
-    outputPathsProvider: Bool
+    outputPathsProvider: Bool,
+    indexTaskBatchSize: Int? = nil
   ) throws -> InitializeBuildResponse {
     return initializationResponse(
       initializeData: SourceKitInitializeBuildResponseData(
         indexDatabasePath: try projectRoot.appendingPathComponent("index-db").filePath,
         indexStorePath: try projectRoot.appendingPathComponent("index-store").filePath,
+        indexTaskBatchSize: indexTaskBatchSize,
         outputPathsProvider: outputPathsProvider,
         prepareProvider: true,
         sourceKitOptionsProvider: true

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -222,6 +222,9 @@ package final actor SemanticIndexManager {
   /// The parameter is the number of files that were scheduled to be indexed.
   private let indexTasksWereScheduled: @Sendable (_ numberOfFileScheduled: Int) -> Void
 
+  /// The number of targets to prepare concurrently, whenever a index request is scheduled.
+  private let indexTaskBatchSize: Int
+
   /// Callback that is called when `progressStatus` might have changed.
   private let indexProgressStatusDidChange: @Sendable () -> Void
 
@@ -261,6 +264,7 @@ package final actor SemanticIndexManager {
     updateIndexStoreTimeout: Duration,
     hooks: IndexHooks,
     indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>,
+    indexTaskBatchSize: Int,
     logMessageToIndexLog:
       @escaping @Sendable (
         _ message: String, _ type: WindowMessageType, _ structure: StructuredLogKind
@@ -273,6 +277,7 @@ package final actor SemanticIndexManager {
     self.updateIndexStoreTimeout = updateIndexStoreTimeout
     self.hooks = hooks
     self.indexTaskScheduler = indexTaskScheduler
+    self.indexTaskBatchSize = indexTaskBatchSize
     self.logMessageToIndexLog = logMessageToIndexLog
     self.indexTasksWereScheduled = indexTasksWereScheduled
     self.indexProgressStatusDidChange = indexProgressStatusDidChange
@@ -877,10 +882,7 @@ package final actor SemanticIndexManager {
 
     var indexTasks: [Task<Void, Never>] = []
 
-    // TODO: When we can index multiple targets concurrently in SwiftPM, increase the batch size to half the
-    // processor count, so we can get parallelism during preparation.
-    // (https://github.com/swiftlang/sourcekit-lsp/issues/1262)
-    for targetsBatch in sortedTargets.partition(intoBatchesOfSize: 1) {
+    for targetsBatch in sortedTargets.partition(intoBatchesOfSize: indexTaskBatchSize) {
       let preparationTaskID = UUID()
       let filesToIndex = targetsBatch.flatMap({ filesByTarget[$0]! })
 

Sources/SourceKitLSP/Workspace.swift

@@ -162,12 +162,17 @@ package final class Workspace: Sendable, BuildServerManagerDelegate {
     if options.backgroundIndexingOrDefault, let uncheckedIndex,
       await buildServerManager.initializationData?.prepareProvider ?? false
     {
+      // TODO: When we can index multiple targets concurrently in SwiftPM, we may want to default
+      // to something else other than 1.
+      // (https://github.com/swiftlang/sourcekit-lsp/issues/1262)
+      let batchSize = await buildServerManager.initializationData?.indexTaskBatchSize ?? 1
       self.semanticIndexManager = SemanticIndexManager(
         index: uncheckedIndex,
         buildServerManager: buildServerManager,
         updateIndexStoreTimeout: options.indexOrDefault.updateIndexStoreTimeoutOrDefault,
         hooks: hooks.indexHooks,
         indexTaskScheduler: indexTaskScheduler,
+        indexTaskBatchSize: batchSize,
         logMessageToIndexLog: { [weak sourceKitLSPServer] in
           sourceKitLSPServer?.logMessageToIndexLog(message: $0, type: $1, structure: $2)
         },

Tests/SourceKitLSPTests/BackgroundIndexingTests.swift

@@ -2624,6 +2624,81 @@ final class BackgroundIndexingTests: XCTestCase {
     let symbols = try await project.testClient.send(WorkspaceSymbolsRequest(query: "myTestFu"))
     XCTAssertEqual(symbols?.count, 1)
   }
+
+  func testBuildServerUsesCustomTaskBatchSize() async throws {
+    final class BuildServer: CustomBuildServer {
+      let inProgressRequestsTracker = CustomBuildServerInProgressRequestTracker()
+      private let projectRoot: URL
+      private var testFileURL: URL { projectRoot.appendingPathComponent("test.swift").standardized }
+
+      required init(projectRoot: URL, connectionToSourceKitLSP: any LanguageServerProtocol.Connection) {
+        self.projectRoot = projectRoot
+      }
+
+      func initializeBuildRequest(_ request: InitializeBuildRequest) async throws -> InitializeBuildResponse {
+        return try initializationResponseSupportingBackgroundIndexing(
+          projectRoot: projectRoot,
+          outputPathsProvider: false,
+          indexTaskBatchSize: 3
+        )
+      }
+
+      func buildTargetSourcesRequest(_ request: BuildTargetSourcesRequest) async throws -> BuildTargetSourcesResponse {
+        var dummyTargets = [BuildTargetIdentifier]()
+        for i in 0..<10 {
+          dummyTargets.append(BuildTargetIdentifier(uri: try! URI(string: "dummy://dummy-\(i)")))
+        }
+        return BuildTargetSourcesResponse(items: dummyTargets.map {
+          SourcesItem(target: $0, sources: [SourceItem(uri: URI(testFileURL), kind: .file, generated: false)])
+        })
+      }
+
+      func textDocumentSourceKitOptionsRequest(
+        _ request: TextDocumentSourceKitOptionsRequest
+      ) async throws -> TextDocumentSourceKitOptionsResponse? {
+        return TextDocumentSourceKitOptionsResponse(compilerArguments: [request.textDocument.uri.pseudoPath])
+      }
+    }
+
+    let preparationTaskSemaphore = WrappedSemaphore(name: "Received a preparation task")
+    let preparationTasks = ThreadSafeBox<[PreparationTaskDescription]>(initialValue: [])
+    let project = try await CustomBuildServerTestProject(
+      files: [
+        "test.swift": """
+        func testFunction() {}
+        """
+      ],
+      buildServer: BuildServer.self,
+      hooks: Hooks(
+        indexHooks: IndexHooks(
+          preparationTaskDidStart: { task in
+            preparationTasks.withLock { preparationTasks in
+              // Ignore everything after the 4th batch to avoid flakiness.
+              if preparationTasks.count < 4 {
+                preparationTasks.append(task)
+              }
+              if preparationTasks.count == 4 {
+                preparationTaskSemaphore.signal()
+              }
+            }
+          }
+        )
+      ),
+      enableBackgroundIndexing: true
+    )
+
+    _ = try await project.testClient.send(SynchronizeRequest(index: true))
+
+    preparationTaskSemaphore.waitOrXCTFail()
+
+    // The test project has 10 targets, and we should have received them in batches of 3.
+    XCTAssertEqual(preparationTasks.value.count, 4)
+    let preparedTargetBatches = preparationTasks.value.map(\.targetsToPrepare).sorted { $0.count > $1.count }
+    XCTAssertEqual(preparedTargetBatches[0].count, 3)
+    XCTAssertEqual(preparedTargetBatches[1].count, 3)
+    XCTAssertEqual(preparedTargetBatches[2].count, 3)
+    XCTAssertEqual(preparedTargetBatches[3].count, 1)
+  }
 }
 
 extension HoverResponseContents {