diff --git a/CHANGELOG.md b/CHANGELOG.md index 6ef0ebc..45bf9fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## unreleased * Wait for a connection closure with `close_timeout`, require asyncpg 0.32.0 +* Limit lock acquisition attempt time ## v0.0.3 (2026-10-06) diff --git a/asyncpg_lock/guard.py b/asyncpg_lock/guard.py index c9f1276..7ef960e 100644 --- a/asyncpg_lock/guard.py +++ b/asyncpg_lock/guard.py @@ -21,6 +21,7 @@ async def _connect() -> asyncpg.Connection: class AdvisoryLockGuard: __slots__ = ( + "__acquire_timeout", "__after_acquire_delay", "__close_timeout", "__connect", @@ -36,6 +37,7 @@ def __init__( reacquire_delay: float = 5, after_acquire_delay: float = 5, close_timeout: float = 5, + acquire_timeout: float = 5, ) -> None: if reconnect_delay < 0: raise ValueError("reconnect_delay must be non-negative") @@ -45,11 +47,14 @@ def __init__( raise ValueError("after_acquire_delay must be positive") if close_timeout <= 0: raise ValueError("close_timeout must be positive") + if acquire_timeout <= 0: + raise ValueError("acquire_timeout must be positive") self.__reconnect_delay = reconnect_delay self.__after_acquire_delay = after_acquire_delay self.__reacquire_delay = reacquire_delay self.__close_timeout = close_timeout + self.__acquire_timeout = acquire_timeout self.__connect = connect async def run( @@ -108,10 +113,15 @@ async def __acquire_lock(self, connection: asyncpg.Connection, key: int | tuple[ while True: try: if isinstance(key, int): - acquired = await connection.fetchval("SELECT pg_try_advisory_lock($1)", key) + acquired = await connection.fetchval( + "SELECT pg_try_advisory_lock($1)", key, timeout=self.__acquire_timeout + ) else: - acquired = await connection.fetchval("SELECT pg_try_advisory_lock($1, $2)", key[0], key[1]) + acquired = await connection.fetchval( + "SELECT pg_try_advisory_lock($1, $2)", key[0], key[1], timeout=self.__acquire_timeout + ) except Exception: + logger.warning("Failed to acquire lock %s", key, exc_info=True) raise Exception(f"Lock {key} not acquired") if acquired: return diff --git a/tests/test_guard.py b/tests/test_guard.py index cbbf43c..5d266e5 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -19,6 +19,7 @@ LOCK_ACQUIRE_RETRY_INTERVAL = 0.5 PER_ATTEMPT_DELAY = 0.1 CLOSE_TIMEOUT = 0.25 +LOCK_ACQUIRE_TIMEOUT = 0.25 LOCK_KEY = random.randint(0, 2**63 - 1) NON_CONFLICTING_LOCK_KEY = random.randint(-(2**63), -1) @@ -155,6 +156,7 @@ def guard(connector: PgConnector) -> asyncpg_lock.AdvisoryLockGuard: after_acquire_delay=LOCK_ACQUIRE_GRACE_PERIOD, reacquire_delay=LOCK_ACQUIRE_RETRY_INTERVAL, close_timeout=CLOSE_TIMEOUT, + acquire_timeout=LOCK_ACQUIRE_TIMEOUT, ) @@ -268,6 +270,56 @@ async def test_stop_on_silently_dead_connection( await holder.close() +async def test_acquire_lock_after_silent_disruption_while_waiting( + guard: asyncpg_lock.AdvisoryLockGuard, connector: PgConnector, proxy: TcpProxy, pg_14: pytest_pg.PG +) -> None: + holder = await asyncpg.connect(host=pg_14.host, port=pg_14.port, user=pg_14.user, database=pg_14.database) + await holder.execute("SELECT pg_advisory_lock($1)", LOCK_KEY) + + tracker = ExecutionTracker() + task = asyncio.create_task(guard.run(LOCK_KEY, tracker)) + try: + await asyncio.sleep(LOCK_ACQUIRE_RETRY_INTERVAL) + await proxy.freeze_connections() + await asyncio.sleep(LOCK_ACQUIRE_RETRY_INTERVAL * 2) + await holder.close() + async with asyncio.timeout( + (LOCK_ACQUIRE_RETRY_INTERVAL + LOCK_ACQUIRE_TIMEOUT + LOCK_ACQUIRE_GRACE_PERIOD + PER_ATTEMPT_DELAY) * 2 + ): + await tracker.min_completed_executions_event.wait() + assert not task.done() + finally: + await cancel_and_wait(task) + await holder.close() + + assert connector.total_open_connections == 2 + assert not tracker.overlaps + + +async def test_log_failed_lock_acquisition_attempt( + guard: asyncpg_lock.AdvisoryLockGuard, proxy: TcpProxy, pg_14: pytest_pg.PG, caplog: pytest.LogCaptureFixture +) -> None: + holder = await asyncpg.connect(host=pg_14.host, port=pg_14.port, user=pg_14.user, database=pg_14.database) + await holder.execute("SELECT pg_advisory_lock($1)", LOCK_KEY) + + task = asyncio.create_task(guard.run(LOCK_KEY, ExecutionTracker())) + try: + await asyncio.sleep(LOCK_ACQUIRE_RETRY_INTERVAL) + await proxy.freeze_connections() + await asyncio.sleep((LOCK_ACQUIRE_RETRY_INTERVAL + LOCK_ACQUIRE_TIMEOUT) * 2) + finally: + await cancel_and_wait(task) + await holder.close() + + assert any( + x.name == "asyncpg_lock" + and x.levelname == "WARNING" + and x.exc_info is not None + and x.exc_info[0] is TimeoutError + for x in caplog.records + ) + + async def test_no_overlapping_execution_for_same_keys( guard: asyncpg_lock.AdvisoryLockGuard, connector: PgConnector ) -> None: