diff --git a/README.md b/README.md index be73c48..7212e99 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,7 @@ except DuplicateTaskError: - **Explicit lock key** — pin any task to a fixed Redis key with `deduplication_key`, bypassing fingerprint computation entirely. - **Partial fingerprint** — deduplicate on a subset of kwargs with `deduplication_key_fields`, ignoring irrelevant arguments. - **Per-task opt-out** — disable deduplication for individual tasks with the `deduplication` label. +- **Startup resilience** — automatic reconnection with exponential backoff if Redis is unavailable at broker startup. ## License diff --git a/docs/index.md b/docs/index.md index 2527d55..406b8fb 100644 --- a/docs/index.md +++ b/docs/index.md @@ -56,6 +56,7 @@ except DuplicateTaskError: - **Explicit lock key** — pin any task to a fixed Redis key with `deduplication_key`, bypassing fingerprint computation entirely. - **Partial fingerprint** — deduplicate on a subset of kwargs with `deduplication_key_fields`, ignoring irrelevant arguments. - **Per-task opt-out** — disable deduplication for individual tasks with the `deduplication` label. +- **Startup resilience** — automatic reconnection with exponential backoff if Redis is unavailable at broker startup. ## License diff --git a/docs/usage.md b/docs/usage.md index 326dad1..e9262d7 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -21,6 +21,8 @@ broker = ListQueueBroker("redis://localhost:6379").with_middlewares( | `default_deduplication` | `bool` | `True` | Whether deduplication is enabled for all tasks by default. Set `False` to opt-in per task instead of opting out. | | `default_ttl` | `int` | `300` | Default lock TTL in seconds. Overridden per task with the `deduplication_ttl` label. | | `key_prefix` | `str` | `"taskiq:deduplication"` | Prefix for all Redis lock keys. | +| `startup_retries` | `int` | `3` | Number of connection attempts during broker startup. | +| `startup_retry_delay` | `float` | `1.0` | Base delay in seconds between retries (exponential backoff: delay × 2^n). | ```python broker = ListQueueBroker("redis://localhost:6379").with_middlewares( @@ -29,10 +31,29 @@ broker = ListQueueBroker("redis://localhost:6379").with_middlewares( default_deduplication=True, default_ttl=60, key_prefix="myapp:dedup", + startup_retries=5, + startup_retry_delay=0.5, ), ) ``` +## Startup resilience + +On startup the middleware verifies the Redis connection with a `PING`. If Redis is +temporarily unavailable, it retries with exponential backoff. +After all attempts are exhausted a `ConnectionError` is raised and the broker +fails to start. + +Adjust `startup_retries` and `startup_retry_delay` to suit your deployment: + +```python +RedisDeduplicationMiddleware( + redis_url="redis://localhost:6379", + startup_retries=5, + startup_retry_delay=2.0, +) +``` + ## How it works When a task is dispatched, the middleware acquires a Redis lock keyed on the task's diff --git a/src/taskiq_deduplication/middleware.py b/src/taskiq_deduplication/middleware.py index 8f79d12..ff5b663 100644 --- a/src/taskiq_deduplication/middleware.py +++ b/src/taskiq_deduplication/middleware.py @@ -1,7 +1,8 @@ +import asyncio import hashlib import json import logging -from typing import Any +from typing import Any, Awaitable, cast from redis.asyncio import Redis from taskiq import TaskiqMessage, TaskiqResult @@ -42,15 +43,43 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware): default_deduplication: bool = True, default_ttl: int = 300, key_prefix: str = "taskiq:deduplication", + startup_retries: int = 3, + startup_retry_delay: float = 1.0, ) -> None: self.redis_url = redis_url self.default_deduplication = default_deduplication self.default_ttl = default_ttl self.key_prefix = key_prefix + self.startup_retries = startup_retries + self.startup_retry_delay = startup_retry_delay self._redis: Redis | None = None async def startup(self) -> None: - self._redis = Redis.from_url(self.redis_url) + last_error: BaseException | None = None + for attempt in range(self.startup_retries): + try: + self._redis = Redis.from_url(self.redis_url) + await cast(Awaitable[bool], self._redis.ping()) + return + except Exception as exc: + last_error = exc + if attempt < self.startup_retries - 1: + delay = self.startup_retry_delay * (2**attempt) + logger.warning( + "Failed to connect to Redis (attempt %d/%d): %s. " + "Retrying in %.1fs...", + attempt + 1, + self.startup_retries, + exc, + delay, + ) + await asyncio.sleep(delay) + logger.error( + "Failed to connect to Redis after %d attempts.", self.startup_retries + ) + raise ConnectionError( + f"Could not connect to Redis after {self.startup_retries} attempts" + ) from last_error async def shutdown(self) -> None: if self._redis is not None: diff --git a/tests/test_middleware.py b/tests/test_middleware.py index 7387ba3..2dba6e3 100644 --- a/tests/test_middleware.py +++ b/tests/test_middleware.py @@ -332,3 +332,75 @@ class TestLifecycle: mw = RedisDeduplicationMiddleware(redis_url="redis://localhost") with pytest.raises(RuntimeError, match="startup"): await mw.pre_send(make_message()) + + +class TestStartupRetry: + @pytest.mark.anyio + async def test_startup_succeeds_after_retries(self): + mw = RedisDeduplicationMiddleware( + redis_url="redis://localhost", + startup_retries=3, + startup_retry_delay=0.01, + ) + with patch("redis.asyncio.Redis.from_url") as mock_from_url: + mock_client = AsyncMock() + mock_client.ping.side_effect = [ + ConnectionError("fail"), + ConnectionError("fail"), + None, + ] + mock_from_url.return_value = mock_client + await mw.startup() + assert mw._redis is mock_client + assert mock_client.ping.call_count == 3 + + @pytest.mark.anyio + async def test_startup_raises_after_all_retries_exhausted(self): + mw = RedisDeduplicationMiddleware( + redis_url="redis://localhost", + startup_retries=2, + startup_retry_delay=0.01, + ) + with patch("redis.asyncio.Redis.from_url") as mock_from_url: + mock_client = AsyncMock() + mock_client.ping.side_effect = ConnectionError("refused") + mock_from_url.return_value = mock_client + with pytest.raises(ConnectionError, match="2 attempts"): + await mw.startup() + assert mock_client.ping.call_count == 2 + + @pytest.mark.anyio + async def test_startup_no_retry_on_first_success(self): + mw = RedisDeduplicationMiddleware( + redis_url="redis://localhost", + startup_retries=3, + startup_retry_delay=0.01, + ) + with patch("redis.asyncio.Redis.from_url") as mock_from_url: + mock_client = AsyncMock() + mock_from_url.return_value = mock_client + await mw.startup() + mock_client.ping.assert_called_once() + + @pytest.mark.anyio + async def test_startup_retry_delay_exponential(self): + mw = RedisDeduplicationMiddleware( + redis_url="redis://localhost", + startup_retries=3, + startup_retry_delay=0.01, + ) + with ( + patch("redis.asyncio.Redis.from_url") as mock_from_url, + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + ): + mock_client = AsyncMock() + mock_client.ping.side_effect = [ + ConnectionError("fail"), + ConnectionError("fail"), + None, + ] + mock_from_url.return_value = mock_client + await mw.startup() + assert mock_sleep.call_count == 2 + assert mock_sleep.call_args_list[0].args[0] == 0.01 + assert mock_sleep.call_args_list[1].args[0] == 0.02