Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add CRT process lock utility #281

Merged
merged 1 commit into from
Nov 14, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions s3transfer/crt.py
Original file line number Diff line number Diff line change
@@ -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):
41 changes: 41 additions & 0 deletions tests/unit/test_crt.py
Original file line number Diff line number Diff line change
@@ -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):