Allow non-SwiftPM build systems to have larger indexing batch sizes

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,17 @@ 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]
+}
+
+export interface MultiTargetPreparationSupport {
+  /** Whether the build server can prepare multiple targets in parallel.
+   * If this value is not provided, it will be assumed to be `true`. */
+  supported?: bool;
+
+  /** The number of targets to prepare in parallel.
+   * If not provided, SourceKit-LSP will calculate an appropriate value based on the environment. */
+  batchSize?: int;
 }
 ```
 

Sources/BuildServerIntegration/BuildServerManager.swift

@@ -322,7 +322,7 @@ 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 build server adapter after initialization finishes. When sending messages to the BSP server, this should be

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 {
     connectionToSourceKitLSP.send(OnBuildTargetDidChangeNotification(changes: nil))
   }
 
+  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?
 
@@ -286,12 +289,14 @@ public struct SourceKitInitializeBuildResponseData: LSPAnyCodable, Codable, Send
   public init(
     indexDatabasePath: String? = nil,
     indexStorePath: String? = nil,
+    multiTargetPreparation: MultiTargetPreparationSupport? = nil,
     watchers: [FileSystemWatcher]? = nil,
     prepareProvider: Bool? = nil,
     sourceKitOptionsProvider: Bool? = nil
   ) {
     self.indexDatabasePath = indexDatabasePath
     self.indexStorePath = indexStorePath
+    self.multiTargetPreparation = multiTargetPreparation
     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,
+    multiTargetPreparation: MultiTargetPreparationSupport? = nil,
     outputPathsProvider: Bool? = nil,
     prepareProvider: Bool? = nil,
     sourceKitOptionsProvider: Bool? = nil,
     watchers: [FileSystemWatcher]? = nil
   ) {
     self.indexDatabasePath = indexDatabasePath
     self.indexStorePath = indexStorePath
+    self.multiTargetPreparation = multiTargetPreparation
     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 .dictionary(let multiTargetPreparation) = dictionary[CodingKeys.multiTargetPreparation.stringValue] {
+      self.multiTargetPreparation = MultiTargetPreparationSupport(fromLSPDictionary: multiTargetPreparation)
+    }
     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 multiTargetPreparation {
+      result[CodingKeys.multiTargetPreparation.stringValue] = multiTargetPreparation.encodeToLSPAny()
+    }
     if let outputPathsProvider {
       result[CodingKeys.outputPathsProvider.stringValue] = .bool(outputPathsProvider)
     }
@@ -357,3 +370,38 @@ 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 not provided, it will be assumed to be `true`.
+  public var supported: Bool?
+
+  /// The number of targets to prepare in parallel.
+  /// If not provided, SourceKit-LSP will calculate an appropriate value based on the environment.
+  public var batchSize: Int?
+
+  public init(supported: Bool? = nil, batchSize: Int? = nil) {
+    self.supported = supported
+    self.batchSize = batchSize
+  }
+
+  public init?(fromLSPDictionary dictionary: [String: LanguageServerProtocol.LSPAny]) {
+    if case .bool(let supported) = dictionary[CodingKeys.supported.stringValue] {
+      self.supported = supported
+    }
+    if case .int(let batchSize) = dictionary[CodingKeys.batchSize.stringValue] {
+      self.batchSize = batchSize
+    }
+  }
+
+  public func encodeToLSPAny() -> LanguageServerProtocol.LSPAny {
+    var result: [String: LSPAny] = [:]
+    if let supported {
+      result[CodingKeys.supported.stringValue] = .bool(supported)
+    }
+    if let batchSize {
+      result[CodingKeys.batchSize.stringValue] = .int(batchSize)
+    }
+    return .dictionary(result)
+  }
+}

Sources/SKTestSupport/CustomBuildServerTestProject.swift

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

Sources/SemanticIndex/PreparationTaskDescription.swift

@@ -39,6 +39,9 @@ package struct PreparationTaskDescription: IndexTaskDescription {
 
   private let preparationUpToDateTracker: UpToDateTracker<BuildTargetIdentifier, DummySecondaryKey>
 
+  /// The purpose of the preparation task.
+  private let purpose: TargetPreparationPurpose
+
   /// See `SemanticIndexManager.logMessageToIndexLog`.
   private let logMessageToIndexLog:
     @Sendable (_ message: String, _ type: WindowMessageType, _ structure: StructuredLogKind) -> Void
@@ -63,6 +66,7 @@ package struct PreparationTaskDescription: IndexTaskDescription {
     targetsToPrepare: [BuildTargetIdentifier],
     buildServerManager: BuildServerManager,
     preparationUpToDateTracker: UpToDateTracker<BuildTargetIdentifier, DummySecondaryKey>,
+    purpose: TargetPreparationPurpose,
     logMessageToIndexLog:
       @escaping @Sendable (
         _ message: String, _ type: WindowMessageType, _ structure: StructuredLogKind
@@ -72,6 +76,7 @@ package struct PreparationTaskDescription: IndexTaskDescription {
     self.targetsToPrepare = targetsToPrepare
     self.buildServerManager = buildServerManager
     self.preparationUpToDateTracker = preparationUpToDateTracker
+    self.purpose = purpose
     self.logMessageToIndexLog = logMessageToIndexLog
     self.hooks = hooks
   }
@@ -119,14 +124,22 @@ package struct PreparationTaskDescription: IndexTaskDescription {
     to currentlyExecutingTasks: [PreparationTaskDescription]
   ) -> [TaskDependencyAction<PreparationTaskDescription>] {
     return currentlyExecutingTasks.compactMap { (other) -> TaskDependencyAction<PreparationTaskDescription>? in
-      if other.targetsToPrepare.count > self.targetsToPrepare.count {
-        // If there is an prepare operation with more targets already running, suspend it.
-        // The most common use case for this is if we prepare all targets simultaneously during the initial preparation
-        // when a project is opened and need a single target indexed for user interaction. We should suspend the
-        // workspace-wide preparation and just prepare the currently needed target.
+      if other.purpose == .forIndexing && self.purpose == .forEditorFunctionality {
+        // If we're running a background indexing operation but need a target indexed for user interaction,
+        // we should prioritize the latter.
         return .cancelAndRescheduleDependency(other)
       }
       return .waitAndElevatePriorityOfDependency(other)
     }
   }
 }
+
+/// The reason why a target is being prepared. This is used to determine the `IndexProgressStatus`
+/// and to prioritize preparation tasks when several of them are running.
+package enum TargetPreparationPurpose: Comparable {
+  /// We are preparing the target so we can index files in it.
+  case forIndexing
+
+  /// We are preparing the target to provide semantic functionality in one of its files.
+  case forEditorFunctionality
+}

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -145,15 +145,6 @@ private struct InProgressPrepareForEditorTask {
   let task: Task<Void, Never>
 }
 
-/// The reason why a target is being prepared. This is used to determine the `IndexProgressStatus`.
-private enum TargetPreparationPurpose: Comparable {
-  /// We are preparing the target so we can index files in it.
-  case forIndexing
-
-  /// We are preparing the target to provide semantic functionality in one of its files.
-  case forEditorFunctionality
-}
-
 /// An entry in `SemanticIndexManager.inProgressPreparationTasks`.
 private struct InProgressPreparationTask {
   let task: OpaqueQueuedIndexTask
@@ -222,6 +213,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 size of the batches in which the `SemanticIndexManager` should dispatch preparation tasks.
+  private let indexTaskBatchSize: Int
+
   /// Callback that is called when `progressStatus` might have changed.
   private let indexProgressStatusDidChange: @Sendable () -> Void
 
@@ -261,6 +255,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 +268,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
@@ -655,6 +651,7 @@ package final actor SemanticIndexManager {
         targetsToPrepare: targetsToPrepare,
         buildServerManager: self.buildServerManager,
         preparationUpToDateTracker: preparationUpToDateTracker,
+        purpose: purpose,
         logMessageToIndexLog: logMessageToIndexLog,
         hooks: hooks
       )
@@ -877,10 +874,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]! })
 
@@ -904,10 +898,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: indexTaskBatchSize) {
               taskGroup.addTask {
                 await self.updateIndexStore(
                   for: fileBatch,

Sources/SourceKitLSP/Workspace.swift

@@ -162,12 +162,25 @@ package final class Workspace: Sendable, BuildServerManagerDelegate {
     if options.backgroundIndexingOrDefault, let uncheckedIndex,
       await buildServerManager.initializationData?.prepareProvider ?? false
     {
+      let shouldIndexInParallel = await buildServerManager.initializationData?.multiTargetPreparation?.supported ?? true
+      let batchSize: Int
+      if shouldIndexInParallel {
+        if let customBatchSize = await buildServerManager.initializationData?.multiTargetPreparation?.batchSize {
+          batchSize = customBatchSize
+        } else {
+          let processorCount = ProcessInfo.processInfo.activeProcessorCount
+          batchSize = max(1, processorCount / 2)
+        }
+      } else {
+        batchSize = 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)
         },