From 9a7f1abcdbdda9064a49b9bd98b7ca2707018ef2 Mon Sep 17 00:00:00 2001
From: John Sirois <john.sirois@gmail.com>
Date: Tue, 17 Sep 2024 11:52:33 -0700
Subject: [PATCH] Re-work testing pyenv management. (#2535)

The new setup allows for safe upgrades by just bumping pyenv Python
versions. Also re-work atomic_directory to support this cleanly, allow
obtaining an independent lock on an atomic directory after it has been
established. This can be useful when reads of the directory contents all
happen in the same code path as the writes.
---
 pex/atomic_directory.py | 38 ++++++++++++++++++++++++++++++--------
 testing/__init__.py     | 20 ++++++--------------
 2 files changed, 36 insertions(+), 22 deletions(-)

diff --git a/pex/atomic_directory.py b/pex/atomic_directory.py
index 4e1108b24..f38187eda 100644
--- a/pex/atomic_directory.py
+++ b/pex/atomic_directory.py
@@ -67,8 +67,16 @@ def __init__(
         locked=False,  # type: bool
     ):
         # type: (...) -> None
+
+        head, tail = os.path.split(os.path.normpath(target_dir))
+        self._lockfile = os.path.join(
+            head, ".{target_dir_name}.atomic_directory.lck".format(target_dir_name=tail)
+        )
+        self._work_dir = "{target_dir}.{type}.work".format(
+            target_dir=target_dir, type="lck" if locked else uuid4().hex
+        )
         self._target_dir = target_dir
-        self._work_dir = "{}.{}.work".format(target_dir, "lck" if locked else uuid4().hex)
+
         target_basename = os.path.basename(self._work_dir)
         if len(target_basename) > 143:
             # Guard against eCryptFS home dir encryption which restricts file names to 143
@@ -110,6 +118,24 @@ def is_finalized(self):
         # type: () -> bool
         return os.path.exists(self._target_dir)
 
+    @property
+    def lockfile(self):
+        # type: () -> str
+        return self._lockfile
+
+    def lock(self, lock_style=None):
+        # type: (Optional[FileLockStyle.Value]) -> Callable[[], None]
+        return _LOCK_MANAGER.lock(self._lockfile, lock_style=lock_style)
+
+    @contextmanager
+    def locked(self, lock_style=None):
+        # type: (Optional[FileLockStyle.Value]) -> Iterator[None]
+        unlock = self.lock(lock_style=lock_style)
+        try:
+            yield
+        finally:
+            unlock()
+
     def finalize(self, source=None):
         # type: (Optional[str]) -> None
         """Rename `work_dir` to `target_dir` using `os.rename()`.
@@ -183,6 +209,7 @@ def acquire(self):
 
         # N.B.: We don't actually write anything to the lock file but the fcntl file locking
         # operations only work on files opened for at least write.
+        safe_mkdir(os.path.dirname(self._path))
         lock_fd = os.open(self._path, os.O_CREAT | os.O_WRONLY)
 
         lock_api = cast(
@@ -265,12 +292,7 @@ def atomic_directory(
         yield atomic_dir
         return
 
-    head, tail = os.path.split(atomic_dir.target_dir)
-    if head:
-        safe_mkdir(head)
-    lockfile = os.path.join(head, ".{}.atomic_directory.lck".format(tail or "here"))
-
-    unlock = _LOCK_MANAGER.lock(file_path=lockfile, lock_style=lock_style)
+    unlock = atomic_dir.lock(lock_style=lock_style)
     if atomic_dir.is_finalized():
         # We lost the double-checked locking race and our work was done for us by the race
         # winner so exit early.
@@ -293,7 +315,7 @@ def atomic_directory(
             "{ident}: After obtaining an exclusive lock on {lockfile}, failed to establish a work "
             "directory at {workdir} due to: {err}".format(
                 ident=ident,
-                lockfile=lockfile,
+                lockfile=atomic_dir.lockfile,
                 workdir=atomic_dir.work_dir,
                 err=e,
             ),
diff --git a/testing/__init__.py b/testing/__init__.py
index 7ab8580da..319ec3779 100644
--- a/testing/__init__.py
+++ b/testing/__init__.py
@@ -539,22 +539,14 @@ def ensure_python_distribution(version):
 
     pip = os.path.join(interpreter_location, "bin", "pip")
 
-    with atomic_directory(target_dir=pyenv_root) as target_dir:
-        if not target_dir.is_finalized():
-            bootstrap_python_installer(target_dir.work_dir)
-
     with atomic_directory(target_dir=interpreter_location) as interpreter_target_dir:
         if not interpreter_target_dir.is_finalized():
-            subprocess.check_call(
-                [
-                    "git",
-                    "--git-dir={}".format(os.path.join(pyenv_root, ".git")),
-                    "--work-tree={}".format(pyenv_root),
-                    "reset",
-                    "--hard",
-                    "v2.3.8",
-                ]
-            )
+            with atomic_directory(target_dir=pyenv_root) as pyenv_root_target_dir:
+                if pyenv_root_target_dir.is_finalized():
+                    with pyenv_root_target_dir.locked():
+                        subprocess.check_call(args=["git", "pull", "--ff-only"], cwd=pyenv_root)
+                else:
+                    bootstrap_python_installer(pyenv_root_target_dir.work_dir)
             env = pyenv_env.copy()
             if sys.platform.lower().startswith("linux"):
                 env["CONFIGURE_OPTS"] = "--enable-shared"