From 79993de83dbdfaa772642766be6424cc78ecd1b0 Mon Sep 17 00:00:00 2001
From: John Sirois <john.sirois@gmail.com>
Date: Fri, 21 Oct 2022 14:37:02 -0700
Subject: [PATCH] Add a private `_PEX_FILE_LOCK_STYLE` knob.

This allows experimenting with making all Pex file locks use BSD style
locks to help debug file locking issues. This environment variable knob
can be removed at any time Pex sees fit.
---
 pex/atomic_directory.py        | 20 ++++++++++++++++++--
 tests/test_atomic_directory.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_atomic_directory.py

diff --git a/pex/atomic_directory.py b/pex/atomic_directory.py
index bfd3dc945..8354b76e6 100644
--- a/pex/atomic_directory.py
+++ b/pex/atomic_directory.py
@@ -83,10 +83,26 @@ class Value(Enum.Value):
     POSIX = Value("posix")
 
 
+def _is_bsd_lock(lock_style=None):
+    # type: (Optional[FileLockStyle.Value]) -> bool
+
+    # The atomic_directory file locking has used POSIX locks since inception. These have maximum
+    # compatibility across OSes and stand a decent chance of working over modern NFS. With the
+    # introduction of `pex3 lock ...` a limited set of atomic_directory uses started asking for BSD
+    # locks since they operate in a thread pool. Only those uses actually pass an explicit value for
+    # `lock_style` to atomic_directory. In order to allow experimenting with / debugging possible
+    # file locking bugs, we allow a `_PEX_FILE_LOCK_STYLE` back door private ~API to upgrade all
+    # locks to BSD style locks. This back door can be removed at any time.
+    file_lock_style = lock_style or FileLockStyle.for_value(
+        os.environ.get("_PEX_FILE_LOCK_STYLE", FileLockStyle.POSIX.value)
+    )
+    return file_lock_style is FileLockStyle.BSD
+
+
 @contextmanager
 def atomic_directory(
     target_dir,  # type: str
-    lock_style=FileLockStyle.POSIX,  # type: FileLockStyle.Value
+    lock_style=None,  # type: Optional[FileLockStyle.Value]
     source=None,  # type: Optional[str]
 ):
     # type: (...) -> Iterator[AtomicDirectory]
@@ -125,7 +141,7 @@ def atomic_directory(
 
     lock_api = cast(
         "Callable[[int, int], None]",
-        fcntl.flock if lock_style is FileLockStyle.BSD else fcntl.lockf,
+        fcntl.flock if _is_bsd_lock(lock_style) else fcntl.lockf,
     )
 
     def unlock():
diff --git a/tests/test_atomic_directory.py b/tests/test_atomic_directory.py
new file mode 100644
index 000000000..811400ac8
--- /dev/null
+++ b/tests/test_atomic_directory.py
@@ -0,0 +1,28 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from pex.atomic_directory import FileLockStyle, _is_bsd_lock
+from pex.testing import environment_as
+
+
+def test_is_bsd_lock():
+    # type: () -> None
+
+    assert not _is_bsd_lock(
+        lock_style=None
+    ), "Expected the default lock style to be POSIX for maximum compatibility."
+    assert not _is_bsd_lock(lock_style=FileLockStyle.POSIX)
+    assert _is_bsd_lock(lock_style=FileLockStyle.BSD)
+
+    # The hard-coded default is already POSIX, so setting the env var default changes nothing.
+    with environment_as(_PEX_FILE_LOCK_STYLE="posix"):
+        assert not _is_bsd_lock(lock_style=None)
+        assert not _is_bsd_lock(lock_style=FileLockStyle.POSIX)
+        assert _is_bsd_lock(lock_style=FileLockStyle.BSD)
+
+    with environment_as(_PEX_FILE_LOCK_STYLE="bsd"):
+        assert _is_bsd_lock(
+            lock_style=None
+        ), "Expected the default lock style to be taken from the environment when defined."
+        assert not _is_bsd_lock(lock_style=FileLockStyle.POSIX)
+        assert _is_bsd_lock(lock_style=FileLockStyle.BSD)