Keep tools up-to-date when run

bazel_env.bzl

@@ -160,6 +160,7 @@ def _tool_impl(ctx):
         substitutions = {
             "{{bazel_env_label}}": str(ctx.label).removeprefix("@@").removesuffix("/bin/" + name),
             "{{rlocation_path}}": rlocation_path,
+            "{{update_when_run}}": ctx.attr.update_when_run,
         },
     )
 
@@ -188,6 +189,9 @@ _tool = rule(
             cfg = _flip_output_dir,
             allow_files = True,
         ),
+        "update_when_run": attr.string(
+            values = ["auto", "yes", "no"],
+        ),
         "_launcher": attr.label(
             allow_single_file = True,
             cfg = _flip_output_dir,
@@ -339,7 +343,7 @@ _bazel_env_rule = rule(
 
 _FORBIDDEN_TOOL_NAMES = ["direnv", "bazel", "bazelisk"]
 
-def bazel_env(*, name, tools = {}, toolchains = {}, **kwargs):
+def bazel_env(*, name, tools = {}, toolchains = {}, update_when_run = "auto", **kwargs):
     # type: (string, dict[string, string | Label], dict[string, string | Label]) -> None
     """Makes Bazel-managed tools and toolchains available under stable paths.
 
@@ -369,6 +373,20 @@ def bazel_env(*, name, tools = {}, toolchains = {}, **kwargs):
             basename of the toolchain directory in the `toolchains` directory. The directory is
             a symlink to the repository root of the (single) repository containing the toolchain.
 
+        update_when_run: Whether to always update the tools via `bazel build` before running them.
+            This is a trade-off between performance (startup latency of a tool) and correctness.
+            The supported values are:
+            <ul>
+            <li>"auto" (default): Always try to update the tool, but skip the update with a warning if
+              it would result in excessive latency (e.g. because there is a concurrent Bazel
+              invocation).
+            <li>"yes": Always try to update the tool and fail with an error if it would result in
+              excessive latency.
+            <li>"no": Use the tool in the state of the last build of the `bazel_env` target. This is
+              the fastest option as it never invokes Bazel, but it may run the an old version of a
+              tool if the `bazel_env` target or the tool itself has changed since the last build.
+            </ul>
+
         **kwargs: Additional arguments to pass to the main `bazel_env` target. It is usually not
             necessary to provide any and the target should have private visibility.
     """
@@ -429,6 +447,7 @@ def bazel_env(*, name, tools = {}, toolchains = {}, **kwargs):
             name = tool_target_name,
             raw_tool = str(tool),
             toolchain_targets = reversed_toolchains,
+            update_when_run = update_when_run,
             visibility = ["//visibility:private"],
             tags = ["manual"],
             **tool_kwargs

docs-gen/bazel_env.md

@@ -7,7 +7,7 @@
 ## bazel_env
 
 <pre>
-bazel_env(<a href="#bazel_env-name">name</a>, <a href="#bazel_env-tools">tools</a>, <a href="#bazel_env-toolchains">toolchains</a>, <a href="#bazel_env-kwargs">kwargs</a>)
+bazel_env(<a href="#bazel_env-name">name</a>, <a href="#bazel_env-tools">tools</a>, <a href="#bazel_env-toolchains">toolchains</a>, <a href="#bazel_env-update_when_run">update_when_run</a>, <a href="#bazel_env-kwargs">kwargs</a>)
 </pre>
 
 Makes Bazel-managed tools and toolchains available under stable paths.
@@ -31,6 +31,7 @@ well as cleans up stale tools.
 | <a id="bazel_env-name"></a>name |  The name of the rule.   |  none |
 | <a id="bazel_env-tools"></a>tools |  A dictionary mapping tool names to their targets or paths. The name is used as the basename of the tool in the `bin` directory and will be available on `PATH`.<br><br>If a target is provided, the corresponding executable is staged in the `bin` directory together with its runfiles.<br><br>If a path is provided, Make variables provided by `toolchains` are expanded in it and all the files of referenced toolchains are staged as runfiles.   |  `{}` |
 | <a id="bazel_env-toolchains"></a>toolchains |  A dictionary mapping toolchain names to their targets. The name is used as the basename of the toolchain directory in the `toolchains` directory. The directory is a symlink to the repository root of the (single) repository containing the toolchain.   |  `{}` |
+| <a id="bazel_env-update_when_run"></a>update_when_run |  Whether to always update the tools via `bazel build` before running them. This is a trade-off between performance (startup latency of a tool) and correctness. The supported values are: <ul> <li>"auto" (default): Always try to update the tool, but skip the update with a warning if   it would result in excessive latency (e.g. because there is a concurrent Bazel   invocation). <li>"yes": Always try to update the tool and fail with an error if it would result in   excessive latency. <li>"no": Use the tool in the state of the last build of the `bazel_env` target. This is   the fastest option as it never invokes Bazel, but it may run the an old version of a   tool if the `bazel_env` target or the tool itself has changed since the last build. </ul>   |  `"auto"` |
 | <a id="bazel_env-kwargs"></a>kwargs |  Additional arguments to pass to the main `bazel_env` target. It is usually not necessary to provide any and the target should have private visibility.   |  none |
 
 

examples/BUILD.bazel

@@ -55,6 +55,8 @@ sh_test(
         "MODULE.bazel",
         # Enforces that the bazel_env has been built.
         ":bazel_env",
+        # Replaces "bazel" in the test script.
+        "fake_bazel.sh",
     ],
     tags = [
         # The test only depends on the bazel_env fake outputs, not the individual tools.

examples/bazel_env_test.sh

@@ -8,17 +8,38 @@ build_workspace_directory="$(dirname "$(readlink -f MODULE.bazel)")"
 # output, possibly with wildcards.
 function assert_cmd_output() {
   local -r cmd="$1"
-  local -r expected_output="$2"
+  local -r expected_first_line="$2"
   local -r extra_path="${3:-}"
+  local -r no_bazel_check="${4:-}"
 
   local -r bazel_env="$build_workspace_directory/bazel-out/bazel_env-opt/bin/bazel_env/bin"
-  local -r actual_output="$(PATH="$bazel_env:/bin:/usr/bin$extra_path" $cmd 2>&1 | head -n 1 || true)"
+  local -r fake_bazel_marker_file=$(mktemp)
+  # The env var is no longer defined when the trap runs, so expand it early.
+  # shellcheck disable=SC2064
+  trap "rm '$fake_bazel_marker_file'" EXIT
+  if ! full_output="$( \
+      FAKE_BAZEL_MARKER_FILE="$fake_bazel_marker_file" \
+      BAZEL=./fake_bazel.sh \
+      PATH="$bazel_env:/bin:/usr/bin$extra_path" \
+      $cmd 2>&1)"; then
+    echo "Command $cmd failed:"
+    echo "$full_output"
+    exit 1
+  fi
+  if [[ $# -ne 4 || "$no_bazel_check" != "no_bazel_check" ]]; then
+    [[ "$(cat "$fake_bazel_marker_file")" == "1" ]] || {
+      echo "Expected to run Bazel with $1, but didn't."
+      exit 1
+    }
+  fi
 
-  # Allow for wildcard matching first.
-  if [[ $actual_output == $expected_output ]]; then
+  local -r actual_first_line="$(echo "$full_output" | head -n 1)"
+  # Allow for wildcard matching and print a diff if the output doesn't match.
+  # shellcheck disable=SC2053
+  if [[ $actual_first_line == $expected_first_line ]]; then
     return
   fi
-  diff <(echo "$expected_output") <(echo "$actual_output") || exit 1
+  diff <(echo "$expected_first_line") <(echo "$actual_first_line") || exit 1
 }
 
 function assert_contains() {
@@ -75,7 +96,7 @@ diff <(echo "$expected_output") <(echo "$status_out") || exit 1
 
 #### Tools ####
 
-assert_cmd_output "bazel-cc" "* error: no input files"
+assert_cmd_output "bazel-cc --version" "@(*gcc*|*clang*)"
 assert_cmd_output "buildifier --version" "buildifier version: 6.4.0 "
 assert_cmd_output "buildozer --version" "buildozer version: 7.1.2 "
 case "$(arch)" in
@@ -90,8 +111,17 @@ assert_cmd_output "python --version" "Python 3.11.8"
 # Bazel's Python launcher requires a system installation of python3.
 assert_cmd_output "python_tool" "python_tool version 0.0.1" ":$(dirname "$(which python3)")"
 
+# Verify that a failing Bazel invocation doesn't prevent the tool from running.
+python_ran=$(mktemp)
+rm "$python_ran"
+FAKE_BAZEL_EXIT_CODE=1 assert_cmd_output "python -c open(\"${python_ran}\",\"w+\")" "WARNING[bazel_env.bzl]: Failed to keep 'python' up-to-date with 'bazel build //:bazel_env':" ":$(dirname "$(which python3)")"
+[[ -f "$python_ran" ]] || {
+  echo "Expected python to run despite Bazel failing"
+  exit 1
+}
+
 #### Toolchains ####
 
 [[ -d "$build_workspace_directory/bazel-out/bazel_env-opt/bin/bazel_env/toolchains/cc_toolchain" ]]
-assert_cmd_output "$build_workspace_directory/bazel-out/bazel_env-opt/bin/bazel_env/toolchains/jdk/bin/java --version" "openjdk 17.0.8.1 2023-08-24 LTS"
-assert_cmd_output "$build_workspace_directory/bazel-out/bazel_env-opt/bin/bazel_env/toolchains/python/bin/python3 --version" "Python 3.11.8"
+assert_cmd_output "$build_workspace_directory/bazel-out/bazel_env-opt/bin/bazel_env/toolchains/jdk/bin/java --version" "openjdk 17.0.8.1 2023-08-24 LTS" "" "no_bazel_check"
+assert_cmd_output "$build_workspace_directory/bazel-out/bazel_env-opt/bin/bazel_env/toolchains/python/bin/python3 --version" "Python 3.11.8" "" "no_bazel_check"

examples/fake_bazel.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+if [[ "${*: -1}" != "//:bazel_env" ]]; then
+  echo "Expected last argument to be //:bazel_env, got ${*: -1}" >&2
+  exit 1
+fi
+
+echo "Fake Bazel stdout"
+echo "Fake Bazel stderr" >&2
+echo "1" >> "${FAKE_BAZEL_MARKER_FILE:-/dev/null}"
+exit "${FAKE_BAZEL_EXIT_CODE:-0}"

launcher.sh.tpl

@@ -21,17 +21,74 @@ _bazel__get_workspace_path() {
   echo "$workspace"
 }
 
+log_colored() {
+  if [[ -t 2 ]]; then
+    local -r colored="$1"
+    local -r normal="\e[m"
+  else
+    local -r colored=""
+    local -r normal=""
+  fi
+  printf "${colored}%s${normal}\n" "$2" >&2
+}
+
+error_color="\e[38;5;1m"
+warning_color="\e[38;5;3m"
+
+log_error() {
+  log_colored "$error_color" "ERROR[bazel_env.bzl]: $1"
+}
+
+log_warning() {
+  log_colored "$warning_color" "WARNING[bazel_env.bzl]: $1"
+}
+
 case "${BASH_SOURCE[0]}" in
   /*) own_path="${BASH_SOURCE[0]}" ;;
   *) own_path="$PWD/${BASH_SOURCE[0]}" ;;
 esac
 own_dir="$(dirname "$own_path")"
 own_name="$(basename "$own_path")"
 if ! grep -q -F "$own_name" "$own_dir/_all_tools.txt"; then
-  echo "ERROR: $own_name has been removed from bazel_env, run 'bazel run {{bazel_env_label}}' to remove it from PATH." >&2
+  log_error "'$own_name' has been removed from '{{bazel_env_label}}', run 'bazel run {{bazel_env_label}}' to remove it from PATH."
   exit 1
 fi
 
+if [[ {{update_when_run}} != no && "${BAZEL_ENV_INTERNAL_EXEC:-False}" != True ]]; then
+  if [[ -t 2 ]]; then
+    color=yes
+    warning_prefix="\e[38;5;3m>\e[m "
+  else
+    color=no
+    warning_prefix="> "
+  fi
+  # Minimize latency (this tool may be run by another tool or even an IDE) by
+  # not waiting for concurrent Bazel commands and also avoid thrashing the
+  # analysis cache, which could silently slow down subsequent Bazel commands.
+  if ! bazel_output=$(\
+         "${BAZEL:-bazel}" \
+           --noblock_for_lock \
+         build \
+           --color=$color \
+           --noallow_analysis_cache_discard \
+           {{bazel_env_label}} 2>&1); then
+    msg="Failed to keep '$own_name' up-to-date with 'bazel build {{bazel_env_label}}':"
+    if [[ {{update_when_run}} == auto ]]; then
+      log_warning "$msg"
+    else
+      log_error "$msg"
+    fi
+    echo "$bazel_output" | while IFS= read -r line; do printf "$warning_prefix%s\n" "$line" >&2; done
+    if [[ {{update_when_run}} == yes ]]; then
+      # Use an abnormal exit code (SIGABRT) that can't be confused with a
+      # legitimate exit code of the wrapped tool.
+      exit 134
+    fi
+  fi
+  # Re-exec the script, which might have been replaced by Bazel.
+  BAZEL_ENV_INTERNAL_EXEC=True exec "$own_path" "$@"
+fi
+
 # Set up an environment similar to 'bazel run' to support tools designed to be
 # run with it.
 # Since tools may cd into BUILD_WORKSPACE_DIRECTORY, ensure that RUNFILES_DIR