pw_build: Introduce a pip lock

The pip lock prevents pip from being run concurrently with any other
Python action. This should reduce flakiness due to modules being
(un)installed while they're imported by other parts of the build.

This version only works on Linux and Mac: it's a no-op on Windows. I'll
develop a Windows version as a followup, but submitting this now to
reduce user pain.

Bug: b/227670947
Change-Id: Ife2cb66b2fbadbee8d61c0641165909dd50accba
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/88861
Reviewed-by: Keir Mierle <keir@google.com>
Commit-Queue: Ted Pudlik <tpudlik@google.com>

2 files changed, 76 insertions, 1 deletions

diff --git a/pw_build/py/pw_build/python_runner.py b/pw_build/py/pw_build/python_runner.py
index 67a105ade..8591c4ffc 100755
--- a/pw_build/py/pw_build/python_runner.py
+++ b/pw_build/py/pw_build/python_runner.py
@@ -18,6 +18,7 @@ the command.
 """
 
 import argparse
+import atexit
 from dataclasses import dataclass
 import enum
 import logging
@@ -27,10 +28,16 @@ import re
 import shlex
 import subprocess
 import sys
+import time
 from typing import Callable, Dict, Iterable, Iterator, List, NamedTuple
 from typing import Optional, Tuple
 
+if sys.platform != 'win32':
+    import fcntl  # pylint: disable=import-error
+    # TODO(b/227670947): Support Windows.
+
 _LOG = logging.getLogger(__name__)
+_LOCK_ACQUISITION_TIMEOUT = 30 * 60  # 30 minutes in seconds
 
 
 def _parse_args() -> argparse.Namespace:
@@ -76,6 +83,12 @@ def _parse_args() -> argparse.Namespace:
         nargs=argparse.REMAINDER,
         help='Python script with arguments to run',
     )
+    parser.add_argument(
+        '--lockfile',
+        type=Path,
+        required=True,
+        help=('Path to a pip lockfile. Any pip execution will aquire an '
+              'exclusive lock on it, any other module a shared lock.'))
     return parser.parse_args()
 
 
@@ -442,7 +455,57 @@ def expand_expressions(paths: GnPaths, arg: str) -> Iterable[str]:
     return (''.join(arg) for arg in expanded_args if arg)
 
 
-def main(
+class LockAcquisitionTimeoutError(Exception):
+    """Raised on a timeout."""
+
+
+def acquire_lock(lockfile: Path, exclusive: bool):
+    """Attempts to acquire the lock.
+
+    Args:
+      lockfile: pathlib.Path to the lock.
+      exclusive: whether this needs to be an exclusive lock.
+
+    Raises:
+      LockAcquisitionTimeoutError: If the lock is not acquired after a
+        reasonable time.
+    """
+    if sys.platform == 'win32':
+        # No-op on Windows, which doesn't have POSIX file locking.
+        # TODO(b/227670947): Get this working on Windows, too.
+        return
+
+    start_time = time.monotonic()
+    if exclusive:
+        lock_type = fcntl.LOCK_EX  # type: ignore[name-defined]
+    else:
+        lock_type = fcntl.LOCK_SH  # type: ignore[name-defined]
+    fd = os.open(lockfile, os.O_RDWR | os.O_CREAT)
+
+    # Make sure we close the file when the process exits. If we manage to
+    # acquire the lock below, closing the file will release it.
+    def cleanup():
+        os.close(fd)
+
+    atexit.register(cleanup)
+
+    backoff = 1
+    while time.monotonic() - start_time < _LOCK_ACQUISITION_TIMEOUT:
+        try:
+            fcntl.flock(  # type: ignore[name-defined]
+                fd, lock_type | fcntl.LOCK_NB)  # type: ignore[name-defined]
+            return  # Lock acquired!
+        except BlockingIOError:
+            pass  # Keep waiting.
+
+        time.sleep(backoff * 0.05)
+        backoff += 1
+
+    raise LockAcquisitionTimeoutError(
+        f"Failed to acquire lock {lockfile} in {_LOCK_ACQUISITION_TIMEOUT}")
+
+
+def main(  # pylint: disable=too-many-arguments
     gn_root: Path,
     current_path: Path,
     original_cmd: List[str],
@@ -453,6 +516,7 @@ def main(
     capture_output: bool,
     touch: Optional[Path],
     working_directory: Optional[Path],
+    lockfile: Path,
 ) -> int:
     """Script entry point."""
 
@@ -497,6 +561,12 @@ def main(
     if working_directory:
         run_args['cwd'] = working_directory
 
+    try:
+        acquire_lock(lockfile, module == 'pip')
+    except LockAcquisitionTimeoutError as exception:
+        _LOG.error('%s', exception)
+        return 1
+
     _LOG.debug('RUN %s', ' '.join(shlex.quote(arg) for arg in command))
 
     completed_process = subprocess.run(command, **run_args)
diff --git a/pw_build/python_action.gni b/pw_build/python_action.gni
index 75973cd6b..1110535cf 100644
--- a/pw_build/python_action.gni
+++ b/pw_build/python_action.gni
@@ -65,6 +65,11 @@ template("pw_python_action") {
     "--current-path",
     rebase_path(".", root_build_dir),
 
+    # pip lockfile, prevents pip from running in parallel with other Python
+    # actions.
+    "--lockfile",
+    rebase_path("$root_out_dir/pip.lock", root_build_dir),
+
     "--default-toolchain=$default_toolchain",
     "--current-toolchain=$current_toolchain",
   ]