Make the bazel command and threads count configurable

print_args.cpp

@@ -1,13 +1,13 @@
 // Prints the arguments passed to the script
 
-#include <iostream>
+#include <stdio.h>
 
 int main(int argc, char *argv[]) {
-  std::cout << "===HEDRON_COMPILE_COMMANDS_BEGIN_ARGS===\n";
-  for (int i = 1; i < argc; ++i) {
-    std::cout << argv[i] << "\n";
-  }
-  std::cout << "===HEDRON_COMPILE_COMMANDS_END_ARGS===\n";
-  // We purposely return a non-zero exit code to have the emcc process exit after running this fake clang wrapper.
-  return 1;
+    printf("===HEDRON_COMPILE_COMMANDS_BEGIN_ARGS===\n");
+    for (int i = 1; i < argc; ++i) {
+        printf("%s\n",argv[i]);
+    }
+    printf("===HEDRON_COMPILE_COMMANDS_END_ARGS===\n");
+    // We purposely return a non-zero exit code to have the emcc process exit after running this fake clang wrapper.
+    return 1;
 }

refresh.template.py

@@ -47,6 +47,14 @@ class SGR(enum.Enum):
     FG_YELLOW = '\033[0;33m'
     FG_BLUE = '\033[0;34m'
 
+def _bazel():
+    bazelcmd = {bazel_command}
+    return bazelcmd
+
+def _threads():
+    user_max_threads = {max_threads}
+    threads = user_max_threads if user_max_threads else min(32, (os.cpu_count() or 1) + 4)
+    return threads
 
 def _log_with_sgr(sgr, colored_message, uncolored_message=''):
     """Log a message to stderr wrapped in an SGR context."""
@@ -104,7 +112,7 @@ def _get_bazel_version():
     If the version can't be determined, returns (0, 0, 0).
     """
     bazel_version_process = subprocess.run(
-        ['bazel', 'version'],
+        [_bazel(), 'version'],
         # MIN_PY=3.7: Replace PIPEs with capture_output.
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
@@ -130,7 +138,7 @@ def _get_bazel_version():
 def _get_bazel_cached_action_keys():
     """Gets the set of actionKeys cached in bazel-out."""
     action_cache_process = subprocess.run(
-        ['bazel', 'dump', '--action_cache'],
+        [_bazel(), 'dump', '--action_cache'],
         # MIN_PY=3.7: Replace PIPEs with capture_output.
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
@@ -1143,7 +1151,7 @@ def _convert_compile_commands(aquery_output):
     # Process each action from Bazelisms -> file paths and their clang commands
     # Threads instead of processes because most of the execution time is farmed out to subprocesses. No need to sidestep the GIL. Might change after https://github.com/clangd/clangd/issues/123 resolved
     with concurrent.futures.ThreadPoolExecutor(
-        max_workers=min(32, (os.cpu_count() or 1) + 4) # Backport. Default in MIN_PY=3.8. See "using very large resources implicitly on many-core machines" in https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor
+            max_workers=_threads()
     ) as threadpool:
         outputs = threadpool.map(_get_cpp_command_for_files, aquery_output.actions)
 
@@ -1200,7 +1208,7 @@ def _get_commands(target: str, flags: str):
         # For efficiency, have bazel filter out external targets (and therefore actions) before they even get turned into actions or serialized and sent to us. Note: this is a different mechanism than is used for excluding just external headers.
         target_statment = f"filter('^(//|@//)',{target_statment})"
     aquery_args = [
-        'bazel',
+        _bazel(),
         'aquery',
         # Aquery docs if you need em: https://docs.bazel.build/versions/master/aquery.html
         # Aquery output proto reference: https://github.com/bazelbuild/bazel/blob/master/src/main/protobuf/analysis_v2.proto

refresh_compile_commands.bzl

@@ -64,6 +64,7 @@ def refresh_compile_commands(
         targets = None,
         exclude_headers = None,
         exclude_external_sources = False,
+        max_threads = None,
         **kwargs):  # For the other common attributes. Tags, compatible_with, etc. https://docs.bazel.build/versions/main/be/common-definitions.html#common-attributes.
     # Convert the various, acceptable target shorthands into the dictionary format
     # In Python, `type(x) == y` is an antipattern, but [Starlark doesn't support inheritance](https://bazel.build/rules/language), so `isinstance` doesn't exist, and this is the correct way to switch on type.
@@ -89,7 +90,7 @@ def refresh_compile_commands(
 
     # Generate the core, runnable python script from refresh.template.py
     script_name = name + ".py"
-    _expand_template(name = script_name, labels_to_flags = targets, exclude_headers = exclude_headers, exclude_external_sources = exclude_external_sources, **kwargs)
+    _expand_template(name = script_name, labels_to_flags = targets, exclude_headers = exclude_headers, exclude_external_sources = exclude_external_sources, max_threads = max_threads, **kwargs)
 
     # Combine them so the wrapper calls the main script
     native.py_binary(
@@ -115,8 +116,11 @@ def _expand_template_impl(ctx):
             "{exclude_headers}": repr(ctx.attr.exclude_headers),
             "{exclude_external_sources}": repr(ctx.attr.exclude_external_sources),
             "{print_args_executable}": repr(ctx.executable._print_args_executable.path),
+            "{bazel_command}": repr(ctx.var.get("BAZEL_COMMAND", "bazel")),
+            "{max_threads}": repr(ctx.attr.max_threads),
         },
     )
+
     return DefaultInfo(files = depset([script]))
 
 _expand_template = rule(
@@ -129,6 +133,8 @@ _expand_template = rule(
         # For Windows INCLUDE. If this were eliminated, for example by the resolution of https://github.com/clangd/clangd/issues/123, we'd be able to just use a macro and skylib's expand_template rule: https://github.com/bazelbuild/bazel-skylib/pull/330
         # Once https://github.com/bazelbuild/bazel/pull/17108 is widely released, we should be able to eliminate this and get INCLUDE directly. Perhaps for 7.0? Should be released in the sucessor to 6.0
         "_cc_toolchain": attr.label(default = "@bazel_tools//tools/cpp:current_cc_toolchain"),
+        "bazel_command": attr.string(default = "bazel"),
+        "max_threads": attr.int(),
     },
     toolchains = ["@bazel_tools//tools/cpp:toolchain_type"],  # Needed for find_cpp_toolchain with --incompatible_enable_cc_toolchain_resolution
     implementation = _expand_template_impl,