From 8528c4170236ba9ee6ee445d359069be82df171e Mon Sep 17 00:00:00 2001
From: kyleknap <kyknapp1@gmail.com>
Date: Tue, 14 Nov 2023 12:49:50 -0800
Subject: [PATCH] Add CRT process lock utility

The CRT S3 client performs best when there is only one
instance of it running on a host. This lock allows an application to
signal across processes whether there is another process of the same
application using the CRT S3 client and prevent spawning more than one
CRT S3 clients running on the system for that application.
---
 s3transfer/crt.py      | 27 +++++++++++++++++++++++++++
 tests/unit/test_crt.py | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/s3transfer/crt.py b/s3transfer/crt.py
index ed5a0864..7817e733 100644
--- a/s3transfer/crt.py
+++ b/s3transfer/crt.py
@@ -38,6 +38,33 @@
 
 logger = logging.getLogger(__name__)
 
+CRT_S3_PROCESS_LOCK = None
+
+
+def acquire_crt_s3_process_lock(name):
+    # Currently, the CRT S3 client performs best when there is only one
+    # instance of it running on a host. This lock allows an application to
+    # signal across processes whether there is another process of the same
+    # application using the CRT S3 client and prevent spawning more than one
+    # CRT S3 clients running on the system for that application.
+    #
+    # NOTE: When acquiring the CRT process lock, the lock automatically is
+    # released when the lock object is garbage collected. So, the CRT process
+    # lock is set as a global so that it is not unintentionally garbage
+    # collected/released if reference of the lock is lost.
+    global CRT_S3_PROCESS_LOCK
+    if CRT_S3_PROCESS_LOCK is None:
+        crt_lock = awscrt.s3.CrossProcessLock(name)
+        try:
+            crt_lock.acquire()
+        except RuntimeError:
+            # If there is another process that is holding the lock, the CRT
+            # returns a RuntimeError. We return None here to signal that our
+            # current process was not able to acquire the lock.
+            return None
+        CRT_S3_PROCESS_LOCK = crt_lock
+    return CRT_S3_PROCESS_LOCK
+
 
 class CRTCredentialProviderAdapter:
     def __init__(self, botocore_credential_provider):
diff --git a/tests/unit/test_crt.py b/tests/unit/test_crt.py
index aadd3827..8bd5288e 100644
--- a/tests/unit/test_crt.py
+++ b/tests/unit/test_crt.py
@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 import io
 
+import pytest
 from botocore.credentials import CredentialResolver, ReadOnlyCredentials
 from botocore.session import Session
 
@@ -25,10 +26,50 @@
     import s3transfer.crt
 
 
+@pytest.fixture
+def mock_crt_process_lock(monkeypatch):
+    # The process lock is cached at the module layer whenever the
+    # cross process lock is successfully acquired. This patch ensures that
+    # test cases will start off with no previously cached process lock and
+    # if a cross process is instantiated/acquired it will be the mock that
+    # can be used for controlling lock behavior.
+    monkeypatch.setattr('s3transfer.crt.CRT_S3_PROCESS_LOCK', None)
+    with mock.patch('awscrt.s3.CrossProcessLock', spec=True) as mock_lock:
+        yield mock_lock
+
+
 class CustomFutureException(Exception):
     pass
 
 
+@pytest.mark.skipif(
+    not HAS_CRT, reason="Test requires awscrt to be installed."
+)
+class TestCRTProcessLock:
+    def test_acquire_crt_s3_process_lock(self, mock_crt_process_lock):
+        lock = s3transfer.crt.acquire_crt_s3_process_lock('app-name')
+        assert lock is s3transfer.crt.CRT_S3_PROCESS_LOCK
+        assert lock is mock_crt_process_lock.return_value
+        mock_crt_process_lock.assert_called_once_with('app-name')
+        mock_crt_process_lock.return_value.acquire.assert_called_once_with()
+
+    def test_unable_to_acquire_lock_returns_none(self, mock_crt_process_lock):
+        mock_crt_process_lock.return_value.acquire.side_effect = RuntimeError
+        assert s3transfer.crt.acquire_crt_s3_process_lock('app-name') is None
+        assert s3transfer.crt.CRT_S3_PROCESS_LOCK is None
+        mock_crt_process_lock.assert_called_once_with('app-name')
+        mock_crt_process_lock.return_value.acquire.assert_called_once_with()
+
+    def test_multiple_acquires_return_same_lock(self, mock_crt_process_lock):
+        lock = s3transfer.crt.acquire_crt_s3_process_lock('app-name')
+        assert s3transfer.crt.acquire_crt_s3_process_lock('app-name') is lock
+        assert lock is s3transfer.crt.CRT_S3_PROCESS_LOCK
+
+        # The process lock should have only been instantiated and acquired once
+        mock_crt_process_lock.assert_called_once_with('app-name')
+        mock_crt_process_lock.return_value.acquire.assert_called_once_with()
+
+
 @requires_crt
 class TestBotocoreCRTRequestSerializer(unittest.TestCase):
     def setUp(self):