mirror of
https://github.com/d3vyce/taskiq-deduplication.git
synced 2026-08-04 19:14:07 +00:00
feat: add startup retry with exponential backoff for Redis connection (#20)
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user