[cherry-pick] Implement global index store cache

bazelbuild#567

swift and clang write index data based on the status of what resides in
`index-store-path`. When using a transient per-swift-library index, it
writes O( M imports * N libs) indexer data and slows down compilation
significantly: to the tune of 300% slow down and 6GB+ index data in my
testing. Bazel likes to use per `swift_library` data in order to remote
cache them, which needs extra consideration to work with the design of
swift and clang and preserve the performance

This PR does 2 things to make index while building use the original
model of the index-store so we don't have to patch clang and swift for
this yet. When the flag `-Xwrapped-swift-enable-global-index-store` is
set:

1. it writes to a global index cache, `bazel-out/global_index_store`
2. it copies indexstore data (records and units) to where Bazel
   expected them for caching: e.g. the original location with
   `swift.index_while_building`

Note that while this uses a `index-import` binary, it is very different
than how index-import currently works. In the variant this expects, it
is program that can _copy_ index data for specific compiler outputs.
This has the effect fo remote caching the subset that was used in
compilation for this library.

I added a detailed description of this feature in the RFC to add it
there MobileNativeFoundation/index-import#53.  In practice,
transitive imports will be indexed by deps and if they aren't cached
then the fallback is Xcode wil index transitive files if we set those as
deps on the unit files

Author: jerrymarino

swift/internal/actions.bzl

@@ -16,7 +16,8 @@
 
 load("@bazel_skylib//lib:partial.bzl", "partial")
 load("@bazel_skylib//lib:types.bzl", "types")
-load(":features.bzl", "are_all_features_enabled")
+load(":feature_names.bzl", "SWIFT_FEATURE_INDEX_WHILE_BUILDING_V2")
+load(":features.bzl", "are_all_features_enabled", "is_feature_enabled")
 load(":toolchain_config.bzl", "swift_toolchain_config")
 
 # The names of actions currently supported by the Swift build rules.
@@ -218,6 +219,19 @@ def run_toolchain_action(
         ):
             execution_requirements["supports-workers"] = "1"
 
+        # For index_while_building_v2 we import a subset of the indexed outputs
+        # from the global index store cache into bazel-out. The program that
+        # handles this is an external program
+        if is_feature_enabled(
+            feature_configuration = feature_configuration,
+            feature_name = SWIFT_FEATURE_INDEX_WHILE_BUILDING_V2,
+        ) and len(swift_toolchain.swift_worker_output_processor[DefaultInfo].files.to_list()):
+            processor = swift_toolchain.swift_worker_output_processor[DefaultInfo].files.to_list()[0]
+            tools.append(processor)
+            args.add_all([
+               "-Xwrapped-swift-output-processor",
+               processor.path])
+
         executable = swift_toolchain.swift_worker
         tool_executable_args.add(tool_config.executable)
         if not types.is_string(tool_config.executable):

swift/internal/compiling.bzl

@@ -45,6 +45,7 @@ load(
     "SWIFT_FEATURE_FULL_DEBUG_INFO",
     "SWIFT_FEATURE_IMPLICIT_MODULES",
     "SWIFT_FEATURE_INDEX_WHILE_BUILDING",
+    "SWIFT_FEATURE_INDEX_WHILE_BUILDING_V2",
     "SWIFT_FEATURE_MINIMAL_DEPS",
     "SWIFT_FEATURE_MODULE_MAP_HOME_IS_CWD",
     "SWIFT_FEATURE_NO_EMBED_DEBUG_MODULE",
@@ -1405,6 +1406,15 @@ def compile(
     else:
         vfsoverlay_file = None
 
+    # For index while building v2 it uses a global index store which is
+    # configured by this flag
+    additional_copts = []
+    if is_feature_enabled(
+        feature_configuration = feature_configuration,
+        feature_name = SWIFT_FEATURE_INDEX_WHILE_BUILDING_V2,
+    ):
+        additional_copts.append("-Xwrapped-swift-enable-global-index-store")
+
     prerequisites = struct(
         additional_inputs = additional_inputs,
         bin_dir = bin_dir,
@@ -1421,7 +1431,7 @@ def compile(
         transitive_defines = merged_providers.swift_info.transitive_defines,
         transitive_modules = transitive_modules,
         transitive_swiftmodules = transitive_swiftmodules,
-        user_compile_flags = copts + swift_toolchain.command_line_copts,
+        user_compile_flags = copts + swift_toolchain.command_line_copts + additional_copts,
         vfsoverlay_file = vfsoverlay_file,
         vfsoverlay_search_path = _SWIFTMODULES_VFS_ROOT,
         # Merge the compile outputs into the prerequisites.
@@ -1756,6 +1766,9 @@ def _declare_compile_outputs(
     index_while_building = is_feature_enabled(
         feature_configuration = feature_configuration,
         feature_name = SWIFT_FEATURE_INDEX_WHILE_BUILDING,
+    ) or is_feature_enabled(
+        feature_configuration = feature_configuration,
+        feature_name = SWIFT_FEATURE_INDEX_WHILE_BUILDING_V2,
     )
     if (
         index_while_building and

swift/internal/feature_names.bzl

@@ -105,6 +105,10 @@ SWIFT_FEATURE_FULL_DEBUG_INFO = "swift.full_debug_info"
 # explicitly.
 SWIFT_FEATURE_IMPLICIT_MODULES = "swift.implicit_modules"
 
+# Index while building - using a global index store cache
+# https://docs.google.com/document/d/1cH2sTpgSnJZCkZtJl1aY-rzy4uGPcrI-6RrUpdATO2Q/
+SWIFT_FEATURE_INDEX_WHILE_BUILDING_V2 = "swift.index_while_building_v2"
+
 # If enabled, the compilation action for a target will produce an index store.
 SWIFT_FEATURE_INDEX_WHILE_BUILDING = "swift.index_while_building"
 

swift/internal/providers.bzl

@@ -200,6 +200,9 @@ target.
 `File`. The executable representing the worker executable used to invoke the
 compiler and other Swift tools (for both incremental and non-incremental
 compiles).
+""",
+        "swift_worker_output_processor": """\
+`File`. An executable that runs to manage llvm indexes for index while building v2.
 """,
         "system_name": """\
 `String`. The name of the operating system that the toolchain is targeting.

swift/internal/swift_toolchain.bzl

@@ -215,6 +215,7 @@ def _swift_toolchain_impl(ctx):
             root_dir = toolchain_root,
             supports_objc_interop = False,
             swift_worker = ctx.executable._worker,
+            swift_worker_output_processor = ctx.attr._worker_output_processor,
             system_name = ctx.attr.os,
             test_configuration = struct(
                 env = {},
@@ -273,6 +274,14 @@ for incremental compilation using a persistent mode.
 """,
                 executable = True,
             ),
+            "_worker_output_processor": attr.label(
+                cfg = "host",
+                allow_files = True,
+                default = "@build_bazel_rules_swift//tools/worker:worker_output_processor",
+                doc = """\
+An executable used to operate on llvm indexs for index while building
+""",
+            ),
         },
     ),
     doc = "Represents a Swift compiler toolchain.",

swift/internal/xcode_swift_toolchain.bzl

@@ -691,6 +691,7 @@ def _xcode_swift_toolchain_impl(ctx):
             required_implicit_deps = [],
             supports_objc_interop = True,
             swift_worker = ctx.executable._worker,
+            swift_worker_output_processor = ctx.attr._worker_output_processor,
             system_name = "darwin",
             test_configuration = struct(
                 env = env,
@@ -726,6 +727,14 @@ for incremental compilation using a persistent mode.
 """,
                 executable = True,
             ),
+            "_worker_output_processor": attr.label(
+                cfg = "host",
+                allow_files = True,
+                default = "@build_bazel_rules_swift//tools/worker:worker_output_processor",
+                doc = """\
+An executable used to operate on llvm indexs for index while building.
+""",
+            ),
             "_xcode_config": attr.label(
                 default = configuration_field(
                     name = "xcode_config_label",

tools/worker/BUILD

@@ -77,6 +77,29 @@ cc_binary(
     }),
 )
 
+# Internal hinge for index while building V2 feature
+config_setting(
+    name = "index_while_building_v2",
+    values = {
+        "features": "swift.index_while_building_v2"
+    },
+)
+
+
+filegroup(
+   name = "worker_output_processor",
+   # with index_while_building_v2 we want to add a dependency on the tool that
+   # manages the per swift_library index. Otherwise, don't load it
+   # This is mainly an internal capability right now - check
+   # tools/worker/work_processor.cc for usage:
+   # build_bazel_rules_swift_index_import is be provided by the user for now
+   srcs = select({
+       ":index_while_building_v2": ["@build_bazel_rules_swift_index_import//:index_import"],
+       "//conditions:default": []
+   }),
+   visibility = ["//visibility:public"],
+)
+
 # Consumed by Bazel integration tests.
 filegroup(
     name = "for_bazel_tests",

tools/worker/work_processor.cc

@@ -22,9 +22,11 @@
 #include <nlohmann/json.hpp>
 #include <sstream>
 #include <string>
+#include <vector>
 
 #include "tools/common/file_system.h"
 #include "tools/common/path_utils.h"
+#include "tools/common/process.h"
 #include "tools/common/string_utils.h"
 #include "tools/common/temp_file.h"
 #include "tools/worker/output_file_map.h"
@@ -64,6 +66,13 @@ void WorkProcessor::ProcessWorkRequest(
   bool is_wmo = false;
 
   std::string prev_arg;
+
+  std::string bazel_index_store_path;
+  bool enable_global_index_store = false;
+
+
+  std::string output_processor;
+
   for (auto arg : request.arguments()) {
     auto original_arg = arg;
     // Peel off the `-output-file-map` argument, so we can rewrite it if
@@ -77,16 +86,49 @@ void WorkProcessor::ProcessWorkRequest(
       is_wmo = true;
     }
 
+    // The output processor is passed as an argument
+    if (arg == "-Xwrapped-swift-output-processor") {
+      arg.clear();
+    } else if (prev_arg == "-Xwrapped-swift-output-processor") {
+      output_processor = arg;
+      arg.clear();
+    }
+
+    // Peel off index-store-path, we will conditionally pop on a new one if
+    // it passes -Xwrapped-swift-enable-global-index-store
+    if (arg == "-index-store-path") {
+      arg.clear();
+    } else if (prev_arg == "-index-store-path") {
+      bazel_index_store_path = arg;
+      arg.clear();
+    }
+    if (arg == "-Xwrapped-swift-enable-global-index-store") {
+      arg.clear();
+      enable_global_index_store = true;
+    }
+
     if (!arg.empty()) {
       params_file_stream << arg << '\n';
     }
 
     prev_arg = original_arg;
   }
 
+  // Material within this index store is namespaced to arcitecture, os, etc
+  auto exec_root = GetCurrentDirectory();
+  auto global_index_store = exec_root + "/bazel-out/global_index_store";
+  if (enable_global_index_store) {
+    processed_args.push_back("-index-store-path");
+    processed_args.push_back(global_index_store);
+  } else if (!bazel_index_store_path.empty()) {
+    processed_args.push_back("-index-store-path");
+    processed_args.push_back(bazel_index_store_path);
+  }
+
   if (!output_file_map_path.empty()) {
+    output_file_map.ReadFromPath(output_file_map_path);
+
     if (!is_wmo) {
-      output_file_map.ReadFromPath(output_file_map_path);
 
       // Rewrite the output file map to use the incremental storage area and
       // pass the compiler the path to the rewritten file.
@@ -130,7 +172,6 @@ void WorkProcessor::ProcessWorkRequest(
   SwiftRunner swift_runner(processed_args, /*force_response_file=*/true);
 
   int exit_code = swift_runner.Run(&stderr_stream, /*stdout_to_stderr=*/true);
-
   if (!is_wmo) {
     // Copy the output files from the incremental storage area back to the
     // locations where Bazel declared the files.
@@ -144,6 +185,35 @@ void WorkProcessor::ProcessWorkRequest(
     }
   }
 
+  // Import global indexes if compilation succeeds
+  if (exit_code != EXIT_FAILURE && enable_global_index_store) {
+    std::vector<std::string> ii_args;
+
+    auto index_import_path = exec_root + "/" + output_processor;
+    ii_args.push_back(index_import_path);
+
+    // Use the output path map to deterine what compilation ouputs to import
+    auto outputs  = output_file_map.incremental_outputs();
+    std::map<std::string, std::string>::iterator it;
+    for (it = outputs.begin(); it != outputs.end(); it++) {
+      auto output_path = it->first;
+      auto file_type = output_path.substr(output_path.find_last_of(".") + 1);
+      if (file_type == "o") {
+        ii_args.push_back("-import-output-file");
+        ii_args.push_back(output_path);
+      }
+    }
+
+    // Copy back from the global index store to bazel's index store
+    ii_args.push_back(global_index_store);
+
+    ii_args.push_back(exec_root + "/" + bazel_index_store_path);
+
+    const std::vector<std::string>& ii_args_ = ii_args;
+    // Dupes the output into the action for failure cases
+    exit_code = RunSubProcess(ii_args_, &stderr_stream, /*stdout_to_stderr=*/true);
+  }
+
   response->set_exit_code(exit_code);
   response->set_output(stderr_stream.str());
 }