Merge pull request #281 from kyleknap/crt-process-lock

Add CRT process lock utility

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):

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):