Index batching: Use upstream version

Example/.gitignore

@@ -17,4 +17,6 @@
 /.sass-cache
 user.bazelrc
 /.bsp/skbsp.json
-/.bsp/sourcekit-bazel-bsp
+/.bsp/sourcekit-bazel-bsp
+/.sourcekit-lsp/config.json
+/.vscode/settings.json

Example/HelloWorld/BUILD

@@ -221,6 +221,7 @@ setup_sourcekit_bsp(
         "config=index_build",
     ],
     compile_top_level = True,
+    index_build_batch_size = 10,
     targets = [
         "//HelloWorld/...",
     ],

Sources/SourceKitBazelBSP/RequestHandlers/InitializeHandler.swift

@@ -172,18 +172,6 @@ final class InitializeHandler {
         } else {
             watchers = nil
         }
-        let batchSize: Int? = {
-            let compileTopLevel = initializedConfig.baseConfig.compileTopLevel
-            let indexBuildBatchSize = initializedConfig.baseConfig.indexBuildBatchSize
-            if compileTopLevel == false && indexBuildBatchSize != nil {
-                // FIXME: It's possible to support this, just doing one thing at a time for simplicity.
-                logger.warning("Ignoring indexBuildBatchSize: Currently only supported for --compile-top-level.")
-            }
-            guard compileTopLevel == false else {
-                return indexBuildBatchSize
-            }
-            return 1
-        }()
         return InitializeBuildResponse(
             displayName: "sourcekit-bazel-bsp",
             version: sourcekitBazelBSPVersion,
@@ -204,11 +192,6 @@ final class InitializeHandler {
             data: SourceKitInitializeBuildResponseData(
                 indexDatabasePath: initializedConfig.indexDatabasePath,
                 indexStorePath: initializedConfig.indexStorePath,
-                // Commented out on purpose while we get batching working upstream
-                //                multiTargetPreparation: MultiTargetPreparationSupport(
-                //                    supported: true,
-                //                    batchSize: batchSize
-                //                ),
                 outputPathsProvider: nil,
                 prepareProvider: true,
                 sourceKitOptionsProvider: true,

Sources/SourceKitBazelBSP/RequestHandlers/PrepareHandler.swift

@@ -26,11 +26,13 @@ import struct os.OSAllocatedUnfairLock
 private let logger = makeFileLevelBSPLogger()
 
 enum PrepareHandlerError: Error, LocalizedError {
-    case unexpectedBatching([String])
+    case unexpectedBatching
 
     var errorDescription: String? {
         switch self {
-        case .unexpectedBatching(let labels): return "Unexpected batching of targets: \(labels.joined(separator: ", "))"
+        case .unexpectedBatching:
+            return
+                "Your sourcekit-lsp instance appears to be configured for batching, but the BSP currently only supports this when passing --compile-top-level. To fix this, either disable batching or enable --compile-top-level for the BSP."
         }
     }
 }
@@ -108,9 +110,7 @@ final class PrepareHandler {
                 extraArgs = []  // Not applicable in this case
             } else {
                 guard platformInfo.count == 1 else {
-                    // Should not happen as we force the batch size to 1 in this case,
-                    // but catching it just in case.
-                    throw PrepareHandlerError.unexpectedBatching(platformInfo.map { $0.label })
+                    throw PrepareHandlerError.unexpectedBatching
                 }
                 let infoToBuild = platformInfo[0]
                 labelsToBuild = [infoToBuild.label]

Sources/SourceKitBazelBSP/Server/BaseServerConfig.swift

@@ -29,7 +29,6 @@ package struct BaseServerConfig: Equatable {
     let targets: [String]
     let indexFlags: [String]
     let filesToWatch: String?
-    let indexBuildBatchSize: Int?
     let compileTopLevel: Bool
 
     package init(
@@ -38,13 +37,11 @@ package struct BaseServerConfig: Equatable {
         indexFlags: [String],
         filesToWatch: String?,
         compileTopLevel: Bool,
-        indexBuildBatchSize: Int? = nil
     ) {
         self.bazelWrapper = bazelWrapper
         self.indexFlags = indexFlags
         self.filesToWatch = filesToWatch
         self.compileTopLevel = compileTopLevel
-        self.indexBuildBatchSize = indexBuildBatchSize
 
         // We need to post-process the target list provided by the user
         // because the queries will always return the "full" label.

Sources/sourcekit-bazel-bsp/Commands/Serve.swift

@@ -52,12 +52,6 @@ struct Serve: ParsableCommand {
     )
     var topLevelRuleToDiscover: [TopLevelRuleType] = []
 
-    @Option(
-        help:
-            "The number of targets to prepare in parallel. If not specified, SourceKit-LSP will calculate an appropriate value based on the environment. Requires --compile-top-level and using the pre-built SourceKit-LSP binary from the release archive.",
-    )
-    var indexBuildBatchSize: Int? = nil
-
     @Flag(
         help:
             "Instead of attempting to build targets individually, build the top-level parent. If your project contains build_test targets for your individual libraries and you're passing them as the top-level targets for the BSP, you can use this flag to build those targets directly for better predictability and caching."
@@ -106,8 +100,7 @@ struct Serve: ParsableCommand {
             targets: targets,
             indexFlags: indexFlags,
             filesToWatch: filesToWatch,
-            compileTopLevel: compileTopLevel,
-            indexBuildBatchSize: indexBuildBatchSize
+            compileTopLevel: compileTopLevel
         )
 
         logger.debug("Initializing BSP with targets: \(targets)")

rules/setup_sourcekit_bsp.bzl

@@ -1,5 +1,10 @@
 def _setup_sourcekit_bsp_impl(ctx):
+    # Config of the BSP itself, a.k.a .bsp/skbsp.json
     rendered_bsp_config = ctx.actions.declare_file("skbsp.json")
+    # Configs for sourcekit-lsp, a.k.a .sourcekit-lsp/config.json
+    rendered_lsp_config = ctx.actions.declare_file("config.json")
+
+    # BSP setup bits
     bsp_config_argv = [
         ".bsp/sourcekit-bazel-bsp",
         "serve",
@@ -9,9 +14,6 @@ def _setup_sourcekit_bsp_impl(ctx):
         bsp_config_argv.append(target)
     bsp_config_argv.append("--bazel-wrapper")
     bsp_config_argv.append(ctx.attr.bazel_wrapper)
-    if ctx.attr.index_build_batch_size:
-        bsp_config_argv.append("--index-build-batch-size")
-        bsp_config_argv.append(ctx.attr.index_build_batch_size)
     if ctx.attr.compile_top_level:
         bsp_config_argv.append("--compile-top-level")
     for index_flag in ctx.attr.index_flags:
@@ -31,20 +33,41 @@ def _setup_sourcekit_bsp_impl(ctx):
             "%argv%": ",\n        ".join(["\"%s\"" % arg for arg in bsp_config_argv]),
         },
     )
+
+    # sourcekit-lsp setup bits
+    lsp_config_json = {
+        "backgroundIndexing": True,
+        "backgroundPreparationMode": "build",
+        "defaultWorkspaceType": "buildServer",
+        "buildServerWorkspaceRequestsTimeout": 999, # Temporary while we don't follow the recommendation of returning from buildTargets as fast as possible + use notifications
+        "buildSettingsTimeout": 999, # Temporary while we don't follow the recommendation of returning from buildTargets as fast as possible + use notifications
+    }
+    if ctx.attr.index_build_batch_size:
+        if not ctx.attr.compile_top_level:
+            fail("index_build_batch_size is currently only supported when compile_top_level is true.")
+        lsp_config_json["preparationBatchingStrategy"] = {
+            "strategy": "fixedTargetBatchSize",
+            "batchSize": ctx.attr.index_build_batch_size
+        }
+    ctx.actions.write(rendered_lsp_config, json.encode(lsp_config_json))
+
+    # Generating the script that ties everything together
     executable = ctx.actions.declare_file("setup_sourcekit_bsp.sh")
     ctx.actions.expand_template(
         template = ctx.file._setup_sourcekit_bsp_script,
         is_executable = True,
         output = executable,
         substitutions = {
             "%bsp_config_path%": rendered_bsp_config.short_path,
+            "%lsp_config_path%": rendered_lsp_config.short_path,
             "%sourcekit_bazel_bsp_path%": ctx.executable.sourcekit_bazel_bsp.short_path,
         },
     )
     tools_runfiles = ctx.runfiles(
         files = [
             ctx.executable.sourcekit_bazel_bsp,
             rendered_bsp_config,
+            rendered_lsp_config,
         ],
     )
     return DefaultInfo(
@@ -97,7 +120,7 @@ setup_sourcekit_bsp = rule(
             default = [],
         ),
         "index_build_batch_size": attr.int(
-            doc = "The number of targets to prepare in parallel. If not specified, SourceKit-LSP will calculate an appropriate value based on the environment. Requires compile_top_level and using the pre-built SourceKit-LSP binary from the release archive.",
+            doc = "The number of targets to prepare in parallel. Requires compile_top_level and using the pre-built SourceKit-LSP binary from the release archive.",
             mandatory = False,
         ),
         "compile_top_level": attr.bool(

rules/setup_sourcekit_bsp.sh.tpl

@@ -4,19 +4,25 @@ set -euo pipefail
 
 sourcekit_bazel_bsp_path="%sourcekit_bazel_bsp_path%"
 bsp_config_path="%bsp_config_path%"
+lsp_config_path="%lsp_config_path%"
 
 bsp_folder_path="$BUILD_WORKSPACE_DIRECTORY/.bsp"
+lsp_folder_path="$BUILD_WORKSPACE_DIRECTORY/.sourcekit-lsp"
 
 mkdir -p "$bsp_folder_path"
+mkdir -p "$lsp_folder_path"
 
 target_bsp_config_path="$bsp_folder_path/skbsp.json"
 target_sourcekit_bazel_bsp_path="$bsp_folder_path/sourcekit-bazel-bsp"
+target_lsp_config_path="$lsp_folder_path/config.json"
 
 cp "$bsp_config_path" "$target_bsp_config_path"
+cp "$lsp_config_path" "$target_lsp_config_path"
 
 # Delete the existing binary first to avoid issues when running this while a server is running.
 rm -f "$target_sourcekit_bazel_bsp_path" || true
 cp "$sourcekit_bazel_bsp_path" "$target_sourcekit_bazel_bsp_path"
 
 chmod +w "$target_bsp_config_path"
 chmod +w "$target_sourcekit_bazel_bsp_path"
+chmod +w "$target_lsp_config_path"