Let the server control the batch size instead

Author: rockbruno

Contributor Documentation/BSP Extensions.md

@@ -22,6 +22,9 @@ export interface SourceKitInitializeBuildResponseData {
    * for `swiftc` or `clang` invocations **/
   indexStorePath?: string;
 
+  /** Options to control how many targets should be prepared simultaneously by SourceKit-LSP. */
+  multiTargetPreparation?: MultiTargetPreparationSupport
+
   /** Whether the server set the `outputPath` property in the `buildTarget/sources` request */
   outputPathsProvider?: bool;
 
@@ -34,7 +37,7 @@ export interface SourceKitInitializeBuildResponseData {
   /** The files to watch for changes.
    * Changes to these files are sent to the BSP server using `workspace/didChangeWatchedFiles`.
    * `FileSystemWatcher` is the same as in LSP. */
-  watchers: [FileSystemWatcher]?
+  watchers?: [FileSystemWatcher]
 }
 ```
 

Sources/BuildServerIntegration/BuiltInBuildServer.swift

@@ -35,6 +35,9 @@ package protocol BuiltInBuildServer: AnyObject, Sendable {
   /// The path to put the index database, if any.
   var indexDatabasePath: URL? { get async }
 
+  /// Whether the build server can prepare multiple targets in parallel.
+  var supportsMultiTargetPreparation: Bool { get }
+
   /// Whether the build server is capable of preparing a target for indexing and determining the output paths for the
   /// target, ie. whether the `prepare` method has been implemented and this build server populates the `outputPath`
   /// property in the `buildTarget/sources` request.

Sources/BuildServerIntegration/BuiltInBuildServerAdapter.swift

@@ -94,6 +94,7 @@ actor BuiltInBuildServerAdapter: QueueBasedMessageHandler {
         indexStorePath: await orLog("getting index store file path") {
           try await underlyingBuildServer.indexStorePath?.filePath
         },
+        multiTargetPreparation: MultiTargetPreparationSupport(supported: underlyingBuildServer.supportsMultiTargetPreparation),
         outputPathsProvider: underlyingBuildServer.supportsPreparationAndOutputPaths,
         prepareProvider: underlyingBuildServer.supportsPreparationAndOutputPaths,
         sourceKitOptionsProvider: true,

Sources/BuildServerIntegration/FixedCompilationDatabaseBuildServer.swift

@@ -59,6 +59,8 @@ package actor FixedCompilationDatabaseBuildServer: BuiltInBuildServer {
     indexStorePath?.deletingLastPathComponent().appendingPathComponent("IndexDatabase")
   }
 
+  package nonisolated var supportsMultiTargetPreparation: Bool { true }
+
   package nonisolated var supportsPreparationAndOutputPaths: Bool { false }
 
   private static func parseCompileFlags(at configPath: URL) throws -> [String] {

Sources/BuildServerIntegration/JSONCompilationDatabaseBuildServer.swift

@@ -96,6 +96,8 @@ package actor JSONCompilationDatabaseBuildServer: BuiltInBuildServer {
     indexStorePath?.deletingLastPathComponent().appendingPathComponent("IndexDatabase")
   }
 
+  package nonisolated var supportsMultiTargetPreparation: Bool { true }
+
   package nonisolated var supportsPreparationAndOutputPaths: Bool { false }
 
   package init(

Sources/BuildServerIntegration/LegacyBuildServer.swift

@@ -138,6 +138,8 @@ actor LegacyBuildServer: MessageHandler, BuiltInBuildServer {
     connectionToSourceKitLSP.send(OnBuildTargetDidChangeNotification(changes: nil))
   }
 
+  package nonisolated var supportsMultiTargetPreparation: Bool { true }
+
   package nonisolated var supportsPreparationAndOutputPaths: Bool { false }
 
   package func buildTargets(request: WorkspaceBuildTargetsRequest) async throws -> WorkspaceBuildTargetsResponse {

Sources/BuildServerIntegration/SwiftPMBuildServer.swift

@@ -459,6 +459,8 @@ package actor SwiftPMBuildServer: BuiltInBuildServer {
     signposter.endInterval("Reloading package", state)
   }
 
+  package nonisolated var supportsMultiTargetPreparation: Bool { false }
+
   package nonisolated var supportsPreparationAndOutputPaths: Bool { options.backgroundIndexingOrDefault }
 
   package var buildPath: URL {

Sources/BuildServerProtocol/Messages/InitializeBuildRequest.swift

@@ -269,6 +269,9 @@ public struct SourceKitInitializeBuildResponseData: LSPAnyCodable, Codable, Send
   /// The path at which SourceKit-LSP can store its index database, aggregating data from `indexStorePath`
   public var indexStorePath: String?
 
+  /// Options to control how many targets should be prepared simultaneously by SourceKit-LSP.
+  public var multiTargetPreparation: MultiTargetPreparationSupport?
+
   /// Whether the server implements the `buildTarget/outputPaths` request.
   public var outputPathsProvider: Bool?
 
@@ -278,9 +281,6 @@ 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]?
 
@@ -289,14 +289,14 @@ public struct SourceKitInitializeBuildResponseData: LSPAnyCodable, Codable, Send
   public init(
     indexDatabasePath: String? = nil,
     indexStorePath: String? = nil,
-    indexTaskBatchSize: Int? = nil,
+    multiTargetPreparation: MultiTargetPreparationSupport? = nil,
     watchers: [FileSystemWatcher]? = nil,
     prepareProvider: Bool? = nil,
     sourceKitOptionsProvider: Bool? = nil
   ) {
     self.indexDatabasePath = indexDatabasePath
     self.indexStorePath = indexStorePath
-    self.indexTaskBatchSize = indexTaskBatchSize
+    self.multiTargetPreparation = multiTargetPreparation
     self.watchers = watchers
     self.prepareProvider = prepareProvider
     self.sourceKitOptionsProvider = sourceKitOptionsProvider
@@ -305,15 +305,15 @@ public struct SourceKitInitializeBuildResponseData: LSPAnyCodable, Codable, Send
   public init(
     indexDatabasePath: String? = nil,
     indexStorePath: String? = nil,
-    indexTaskBatchSize: Int? = nil,
+    multiTargetPreparation: MultiTargetPreparationSupport? = nil,
     outputPathsProvider: Bool? = nil,
     prepareProvider: Bool? = nil,
     sourceKitOptionsProvider: Bool? = nil,
     watchers: [FileSystemWatcher]? = nil
   ) {
     self.indexDatabasePath = indexDatabasePath
     self.indexStorePath = indexStorePath
-    self.indexTaskBatchSize = indexTaskBatchSize
+    self.multiTargetPreparation = multiTargetPreparation
     self.outputPathsProvider = outputPathsProvider
     self.prepareProvider = prepareProvider
     self.sourceKitOptionsProvider = sourceKitOptionsProvider
@@ -327,8 +327,8 @@ 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 .dictionary(let multiTargetPreparation) = dictionary[CodingKeys.multiTargetPreparation.stringValue] {
+      self.multiTargetPreparation = MultiTargetPreparationSupport(fromLSPDictionary: multiTargetPreparation)
     }
     if case .bool(let outputPathsProvider) = dictionary[CodingKeys.outputPathsProvider.stringValue] {
       self.outputPathsProvider = outputPathsProvider
@@ -352,8 +352,8 @@ 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 multiTargetPreparation {
+      result[CodingKeys.multiTargetPreparation.stringValue] = multiTargetPreparation.encodeToLSPAny()
     }
     if let outputPathsProvider {
       result[CodingKeys.outputPathsProvider.stringValue] = .bool(outputPathsProvider)
@@ -370,3 +370,26 @@ public struct SourceKitInitializeBuildResponseData: LSPAnyCodable, Codable, Send
     return .dictionary(result)
   }
 }
+
+public struct MultiTargetPreparationSupport: LSPAnyCodable, Codable, Sendable {
+  /// Whether the build server can prepare multiple targets in parallel. If this value is omitted, it is assumed to be `true`.
+  public var supported: Bool?
+
+  public init(supported: Bool? = nil) {
+    self.supported = supported
+  }
+
+  public init?(fromLSPDictionary dictionary: [String: LanguageServerProtocol.LSPAny]) {
+    if case .bool(let supported) = dictionary[CodingKeys.supported.stringValue] {
+      self.supported = supported
+    }
+  }
+
+  public func encodeToLSPAny() -> LanguageServerProtocol.LSPAny {
+    var result: [String: LSPAny] = [:]
+    if let supported {
+      result[CodingKeys.supported.stringValue] = .bool(supported)
+    }
+    return .dictionary(result)
+  }
+}

Sources/SKTestSupport/CustomBuildServerTestProject.swift

@@ -175,17 +175,15 @@ package extension CustomBuildServer {
 
   func initializationResponseSupportingBackgroundIndexing(
     projectRoot: URL,
-    outputPathsProvider: Bool,
-    indexTaskBatchSize: Int? = nil
+    outputPathsProvider: Bool
   ) 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
+        sourceKitOptionsProvider: true,
       )
     )
   }

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -222,8 +222,8 @@ 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
+  /// Determines whether or not the `SemanticIndexManager` should dispatch preparation tasks in batches.
+  private let shouldIndexInParallel: Bool
 
   /// Callback that is called when `progressStatus` might have changed.
   private let indexProgressStatusDidChange: @Sendable () -> Void
@@ -264,7 +264,7 @@ package final actor SemanticIndexManager {
     updateIndexStoreTimeout: Duration,
     hooks: IndexHooks,
     indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>,
-    indexTaskBatchSize: Int,
+    shouldIndexInParallel: Bool,
     logMessageToIndexLog:
       @escaping @Sendable (
         _ message: String, _ type: WindowMessageType, _ structure: StructuredLogKind
@@ -277,7 +277,7 @@ package final actor SemanticIndexManager {
     self.updateIndexStoreTimeout = updateIndexStoreTimeout
     self.hooks = hooks
     self.indexTaskScheduler = indexTaskScheduler
-    self.indexTaskBatchSize = indexTaskBatchSize
+    self.shouldIndexInParallel = shouldIndexInParallel
     self.logMessageToIndexLog = logMessageToIndexLog
     self.indexTasksWereScheduled = indexTasksWereScheduled
     self.indexProgressStatusDidChange = indexProgressStatusDidChange
@@ -654,7 +654,6 @@ package final actor SemanticIndexManager {
     guard !targetsToPrepare.isEmpty else {
       return
     }
-
     let taskDescription = AnyIndexTaskDescription(
       PreparationTaskDescription(
         targetsToPrepare: targetsToPrepare,
@@ -882,7 +881,14 @@ package final actor SemanticIndexManager {
 
     var indexTasks: [Task<Void, Never>] = []
 
-    for targetsBatch in sortedTargets.partition(intoBatchesOfSize: indexTaskBatchSize) {
+    let batchSize: Int
+    if shouldIndexInParallel {
+      let processorCount = ProcessInfo.processInfo.activeProcessorCount
+      batchSize = max(1, processorCount * 5)
+    } else {
+      batchSize = 1
+    }
+    for targetsBatch in sortedTargets.partition(intoBatchesOfSize: batchSize) {
       let preparationTaskID = UUID()
       let filesToIndex = targetsBatch.flatMap({ filesByTarget[$0]! })
 
@@ -906,10 +912,7 @@ package final actor SemanticIndexManager {
         // And after preparation is done, index the files in the targets.
         await withTaskGroup(of: Void.self) { taskGroup in
           for target in targetsBatch {
-            // TODO: Once swiftc supports indexing of multiple files in a single invocation, increase the batch size to
-            // allow it to share AST builds between multiple files within a target.
-            // (https://github.com/swiftlang/sourcekit-lsp/issues/1268)
-            for fileBatch in filesByTarget[target]!.partition(intoBatchesOfSize: 1) {
+            for fileBatch in filesByTarget[target]!.partition(intoBatchesOfSize: batchSize) {
               taskGroup.addTask {
                 await self.updateIndexStore(
                   for: fileBatch,