mirror of
https://github.com/d3vyce/taskiq-deduplication.git
synced 2026-08-04 19:14:07 +00:00
fix: close leaked Redis clients on failed startup retries (#43)
This commit is contained in:
@@ -60,12 +60,14 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware):
|
||||
async def startup(self) -> None:
|
||||
last_error: BaseException | None = None
|
||||
for attempt in range(self.startup_retries):
|
||||
client = Redis.from_url(self.redis_url)
|
||||
try:
|
||||
self._redis = Redis.from_url(self.redis_url)
|
||||
await cast(Awaitable[bool], self._redis.ping())
|
||||
await cast(Awaitable[bool], client.ping())
|
||||
self._redis = client
|
||||
self._release_script = self._redis.register_script(RELEASE_LUA_SCRIPT)
|
||||
return
|
||||
except Exception as exc:
|
||||
await client.aclose()
|
||||
last_error = exc
|
||||
if attempt < self.startup_retries - 1:
|
||||
delay = self.startup_retry_delay * (2**attempt)
|
||||
|
||||
@@ -455,6 +455,46 @@ class TestStartupRetry:
|
||||
assert mock_sleep.call_args_list[0].args[0] == 0.01
|
||||
assert mock_sleep.call_args_list[1].args[0] == 0.02
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_startup_closes_failed_clients(self):
|
||||
mw = RedisDeduplicationMiddleware(
|
||||
redis_url="redis://localhost",
|
||||
startup_retries=3,
|
||||
startup_retry_delay=0.01,
|
||||
)
|
||||
failed1, failed2, good = AsyncMock(), AsyncMock(), AsyncMock()
|
||||
failed1.ping.side_effect = ConnectionError("fail")
|
||||
failed2.ping.side_effect = ConnectionError("fail")
|
||||
good.register_script = MagicMock()
|
||||
|
||||
with patch(
|
||||
"redis.asyncio.Redis.from_url", side_effect=[failed1, failed2, good]
|
||||
):
|
||||
await mw.startup()
|
||||
|
||||
failed1.aclose.assert_called_once()
|
||||
failed2.aclose.assert_called_once()
|
||||
good.aclose.assert_not_called()
|
||||
assert mw._redis is good
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_startup_closes_client_when_all_retries_exhausted(self):
|
||||
mw = RedisDeduplicationMiddleware(
|
||||
redis_url="redis://localhost",
|
||||
startup_retries=2,
|
||||
startup_retry_delay=0.01,
|
||||
)
|
||||
failed1, failed2 = AsyncMock(), AsyncMock()
|
||||
failed1.ping.side_effect = ConnectionError("fail")
|
||||
failed2.ping.side_effect = ConnectionError("fail")
|
||||
|
||||
with patch("redis.asyncio.Redis.from_url", side_effect=[failed1, failed2]):
|
||||
with pytest.raises(ConnectionError):
|
||||
await mw.startup()
|
||||
|
||||
failed1.aclose.assert_called_once()
|
||||
failed2.aclose.assert_called_once()
|
||||
|
||||
|
||||
class TestLabelTypeParsing:
|
||||
@pytest.mark.anyio
|
||||
|
||||
Reference in New Issue
Block a user