diff --git a/CHANGELOG.md b/CHANGELOG.md index d482ae2..8f834d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,8 @@ +## unreleased + +* Limit lock acquisition attempt time + + ## v0.0.3 (2026-10-06) * Use PyPI trusted publishing diff --git a/asyncpg_lock/guard.py b/asyncpg_lock/guard.py index 4e4ad34..0227ff4 100644 --- a/asyncpg_lock/guard.py +++ b/asyncpg_lock/guard.py @@ -20,6 +20,7 @@ async def _connect() -> asyncpg.Connection: class AdvisoryLockGuard: __slots__ = ( + "__acquire_timeout", "__after_acquire_delay", "__connect", "__reacquire_delay", @@ -34,6 +35,7 @@ def __init__( reconnect_delay: float = 5, reacquire_delay: float = 5, after_acquire_delay: float = 5, + acquire_timeout: float = 5, ) -> None: if reconnect_delay < 0: raise ValueError("reconnect_delay must be non-negative") @@ -41,10 +43,13 @@ def __init__( raise ValueError("reacquire_delay must be positive") if after_acquire_delay <= 0: raise ValueError("after_acquire_delay 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.__acquire_timeout = acquire_timeout self.__connect = connect self.__tasks: set[asyncio.Task[None]] = set() @@ -105,10 +110,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 b933778..0e7bd2a 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -18,6 +18,7 @@ LOCK_ACQUIRE_GRACE_PERIOD = 0.75 LOCK_ACQUIRE_RETRY_INTERVAL = 0.5 PER_ATTEMPT_DELAY = 0.1 +LOCK_ACQUIRE_TIMEOUT = 0.25 LOCK_KEY = random.randint(0, 2**63 - 1) NON_CONFLICTING_LOCK_KEY = random.randint(-(2**63), -1) @@ -153,6 +154,7 @@ def guard(connector: PgConnector) -> asyncpg_lock.AdvisoryLockGuard: reconnect_delay=RECONNECT_DELAY, after_acquire_delay=LOCK_ACQUIRE_GRACE_PERIOD, reacquire_delay=LOCK_ACQUIRE_RETRY_INTERVAL, + acquire_timeout=LOCK_ACQUIRE_TIMEOUT, ) @@ -234,6 +236,56 @@ async def test_reacquire_lock_after_silent_disruption( assert not tracker.overlaps +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: