SERVER-99469 fallback to normal bazel call on wrapper script failure (#31245)

GitOrigin-RevId: 7da7cf364cb2b63c2895754cd93265fb4f88fc2b

bazel/wrapper_hook.py

@@ -1,5 +1,6 @@
 import hashlib
 import os
+import platform
 import shutil
 import subprocess
 import sys
@@ -38,30 +39,38 @@ class DuplicateSourceNames(Exception):
 wrapper_debug(f"wrapper hook script is using {sys.executable}")
 
 
-def search_for_modules(deps, deps_installed, deps_needed, lockfile_changed=False):
+def get_deps_dirs(deps):
     bazel_out_dir = os.path.join(REPO_ROOT, "bazel-out")
-    if os.path.exists(bazel_out_dir):
-        for dep in deps:
-            found = False
-            for entry in os.listdir(bazel_out_dir):
-                if os.path.exists(f"{bazel_out_dir}/{entry}/bin/external/poetry/{dep}"):
-                    if not lockfile_changed:
-                        deps_installed.append(dep)
-                        sys.path.append(f"{bazel_out_dir}/{entry}/bin/external/poetry/{dep}")
-                        found = True
-                        break
-                    else:
-                        os.chmod(f"{bazel_out_dir}/{entry}/bin/external/poetry/{dep}", 0o777)
-                        for root, dirs, files in os.walk(
-                            f"{bazel_out_dir}/{entry}/bin/external/poetry/{dep}"
-                        ):
-                            for somedir in dirs:
-                                os.chmod(os.path.join(root, somedir), 0o777)
-                            for file in files:
-                                os.chmod(os.path.join(root, file), 0o777)
-                        shutil.rmtree(f"{bazel_out_dir}/{entry}/bin/external/poetry/{dep}")
-            if not found:
-                deps_needed.append(dep)
+    bazel_bin = os.path.join(REPO_ROOT, "bazel-bin")
+    for dep in deps:
+        for child in os.listdir(bazel_out_dir):
+            yield f"{bazel_out_dir}/{child}/bin/external/poetry/{dep}", dep
+        yield f"{bazel_bin}/bin/external/poetry/{dep}", dep
+
+
+def search_for_modules(deps, deps_installed, lockfile_changed=False):
+    deps_not_found = deps.copy()
+    for target_dir, dep in get_deps_dirs(deps):
+        if dep in deps_installed:
+            continue
+
+        if not os.path.exists(target_dir):
+            continue
+
+        if not lockfile_changed:
+            deps_installed.append(dep)
+            deps_not_found.remove(dep)
+            sys.path.append(target_dir)
+        else:
+            os.chmod(target_dir, 0o777)
+            for root, dirs, files in os.walk(target_dir):
+                for somedir in dirs:
+                    os.chmod(os.path.join(root, somedir), 0o777)
+                for file in files:
+                    os.chmod(os.path.join(root, file), 0o777)
+            shutil.rmtree(target_dir)
+
+    return deps_not_found
 
 
 def install_modules(bazel):
@@ -81,8 +90,9 @@ def install_modules(bazel):
 
     deps = ["retry"]
     deps_installed = []
-    deps_needed = []
-    search_for_modules(deps, deps_installed, deps_needed, lockfile_changed=old_hash != current_hash)
+    deps_needed = search_for_modules(
+        deps, deps_installed, lockfile_changed=old_hash != current_hash
+    )
 
     if deps_needed:
         need_to_install = True
@@ -95,8 +105,7 @@ def install_modules(bazel):
         subprocess.run(
             [bazel, "build", "--config=local"] + ["@poetry//:install_" + dep for dep in deps_needed]
         )
-        deps_missing = []
-        search_for_modules(deps_needed, deps_installed, deps_missing)
+        deps_missing = search_for_modules(deps_needed, deps_installed)
         if deps_missing:
             raise Exception(f"Failed to install python deps {deps_missing}")
 
@@ -104,9 +113,10 @@ def install_modules(bazel):
 def get_buildozer_output(autocomplete_query):
     from buildscripts.install_bazel import install_bazel
 
-    buildozer = shutil.which("buildozer")
+    buildozer_name = "buildozer" if not platform.system() == "Windows" else "buildozer.exe"
+    buildozer = shutil.which(buildozer_name)
     if not buildozer:
-        buildozer = os.path.expanduser("~/.local/bin/buildozer")
+        buildozer = os.path.expanduser(f"~/.local/bin/{buildozer_name}")
         if not os.path.exists(buildozer):
             bazel_bin_dir = os.path.expanduser("~/.local/bin")
             if not os.path.exists(bazel_bin_dir):

site_scons/site_tools/integrate_bazel.py

@@ -412,13 +412,16 @@ def bazel_server_timeout_dumper(jvm_out, proc_pid, project_root):
 
 def bazel_build_subproc_func(**kwargs):
     project_root = os.path.abspath(".")
-    output_base = subprocess.run(
-        [Globals.bazel_executable, "info", "output_base"],
-        capture_output=True,
-        text=True,
-        check=True,
-        env=kwargs["env"],
-    ).stdout.strip()
+    try:
+        output_base = subprocess.run(
+            [Globals.bazel_executable, "info", "output_base"],
+            capture_output=True,
+            text=True,
+            check=True,
+            env=kwargs["env"],
+        ).stdout.strip()
+    except subprocess.CalledProcessError:
+        output_base = None
 
     if os.environ.get("CI"):
         if os.path.exists(".bazel_real"):
@@ -429,14 +432,15 @@ def bazel_build_subproc_func(**kwargs):
 
     bazel_proc = subprocess.Popen(**kwargs)
 
-    t = threading.Thread(
-        target=bazel_server_timeout_dumper,
-        args=(jvm_out, bazel_proc.pid, project_root),
-    )
+    if output_base:
+        t = threading.Thread(
+            target=bazel_server_timeout_dumper,
+            args=(jvm_out, bazel_proc.pid, project_root),
+        )
 
-    # the bazel calls are wrapped in retries so we can rely on them to restart the attempt.
-    t.daemon = True
-    t.start()
+        # the bazel calls are wrapped in retries so we can rely on them to restart the attempt.
+        t.daemon = True
+        t.start()
 
     return bazel_proc
 
@@ -537,6 +541,17 @@ def perform_non_tty_bazel_build(bazel_cmd: str) -> None:
 
 
 def run_bazel_command(env, bazel_cmd, tries_so_far=0):
+    try:
+        server_pid = subprocess.run(
+            [Globals.bazel_executable, "info", "server_pid"],
+            capture_output=True,
+            text=True,
+            check=True,
+            env={**os.environ.copy(), **Globals.bazel_env_variables},
+        ).stdout.strip()
+    except subprocess.CalledProcessError:
+        server_pid = None
+
     try:
         tty_import_fail = False
         try:
@@ -565,13 +580,21 @@ def run_bazel_command(env, bazel_cmd, tries_so_far=0):
                 print(
                     "Killing Bazel between retries on Windows to work around file access deadlock"
                 )
-                subprocess.run(
-                    [os.path.abspath(Globals.bazel_executable), "shutdown"],
-                    check=True,
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.STDOUT,
-                    env={**os.environ.copy(), **Globals.bazel_env_variables},
-                )
+                try:
+                    subprocess.run(
+                        [os.path.abspath(Globals.bazel_executable), "shutdown"],
+                        check=True,
+                        stdout=subprocess.PIPE,
+                        stderr=subprocess.STDOUT,
+                        env={**os.environ.copy(), **Globals.bazel_env_variables},
+                    )
+                except subprocess.CalledProcessError:
+                    proc = psutil.Process(server_pid)
+                    if proc.is_running():
+                        proc.terminate()
+                        proc.wait(timeout=10)
+                        proc.kill()
+
             linker_jobs = 4
             sanitizers = env.GetOption("sanitize")
             if sanitizers is not None and "fuzzer" in sanitizers.split(","):

tools/bazel

@@ -1,7 +1,5 @@
 #!/usr/bin/env bash
 
-set -e
-
 # Whenever Bazel is invoked, it first calls this script setting "BAZEL_REAL" to the path of the real Bazel binary.
 # Use this file as a wrapper for any logic that should run before bazel itself is executed.
 
@@ -76,14 +74,12 @@ if [[ $MONGO_BAZEL_WRAPPER_DEBUG == 1 ]]; then
     wrapper_start_time="$(date -u +%s.%N)"
 fi
  
-
 if [[ "$OSTYPE" == "linux"* ]]; then
   os="linux"
 elif [[ "$OSTYPE" == "darwin"* ]]; then
   os="macos"
 else
-  echo "Unsupported OS $OSTYPE"
-  exit 1
+  os="unknown"
 fi
 
 ARCH=$(uname -m)
@@ -119,6 +115,10 @@ python="$REPO_ROOT/bazel-$cur_dir/external/py_${os}_${ARCH}/dist/bin/python3"
 if [ ! -f $python ]; then
   >&2 echo "python prereq missing, using bazel to install python..."
   >&2 $bazel_real build --config=local @py_${os}_${ARCH}//:all
+  if [[ $? != 0 ]]; then
+    >&2 echo "wrapper script failed to install python! falling back to normal bazel call..."
+    exec "$bazel_real" $@
+  fi
 fi
 
 autocomplete_query=0
@@ -137,7 +137,10 @@ MONGO_BAZEL_WRAPPER_ARGS=$(mktemp)
 MONGO_BAZEL_WRAPPER_ARGS=$MONGO_BAZEL_WRAPPER_ARGS \
 MONGO_AUTOCOMPLETE_QUERY=$autocomplete_query \
 $python $REPO_ROOT/bazel/wrapper_hook.py $bazel_real "$@"
-
+if [[ $? != 0 ]]; then
+  >&2 echo "wrapper script failed! falling back to normal bazel call..."
+  exec "$bazel_real" $@
+fi
 new_args=$(<$MONGO_BAZEL_WRAPPER_ARGS)
 
 if [[ $MONGO_BAZEL_WRAPPER_DEBUG == 1 ]] && [[ $autocomplete_query == 0 ]]; then

tools/bazel.bat

@@ -3,7 +3,7 @@ setlocal EnableDelayedExpansion
 
 echo common --//bazel/config:running_through_bazelisk > .bazelrc.bazelisk
 
-set REPO_ROOT=%~dp0\..
+set REPO_ROOT=%~dp0..
 
 for %%I in (%REPO_ROOT%) do set cur_dir=%%~nxI
 
@@ -12,12 +12,22 @@ set python="%REPO_ROOT%\bazel-%cur_dir%\external\py_windows_x86_64\dist\python.e
 if not exist "%python%" (
     echo python prereq missing, using bazel to install python... 1>&2
     "%BAZEL_REAL%" build --config=local @py_windows_x86_64//:all 1>&2
+    if %ERRORLEVEL% NEQ 0 (
+        echo wrapper script failed to install python! falling back to normal bazel call...
+        "%BAZEL_REAL%" %*
+        exit %ERRORLEVEL%
+    )
 )
 SET STARTTIME=%TIME%
 
 set "MONGO_BAZEL_WRAPPER_ARGS=%tmp%\bat~%RANDOM%.tmp"
 echo "" > %MONGO_BAZEL_WRAPPER_ARGS%
 %python% %REPO_ROOT%/bazel/wrapper_hook.py "%BAZEL_REAL%" %*
+if %ERRORLEVEL% NEQ 0 (
+    echo wrapper script failed! falling back to normal bazel call...
+    "%BAZEL_REAL%" %*
+    exit %ERRORLEVEL%
+)
 for /f "Tokens=* Delims=" %%x in ( %MONGO_BAZEL_WRAPPER_ARGS% ) do set "new_args=!new_args!%%x"
 del %MONGO_BAZEL_WRAPPER_ARGS%