diff --git a/README.md b/README.md index a1cf079..f238199 100644 --- a/README.md +++ b/README.md @@ -82,7 +82,8 @@ pip install pydocket ``` Docket requires a [Redis](http://redis.io/) server with Streams support (which was -introduced in Redis 5.0.0). Docket is tested with: +introduced in Redis 5.0.0). Reliable message queues require Redis 6.2 or newer. +Docket is tested with: - Redis 6.2, 7.4, and 8.6 (standalone and cluster modes) - [Valkey](https://valkey.io/) 8.1 diff --git a/docs/api-reference.md b/docs/api-reference.md index f014f37..61727af 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -1,3 +1,5 @@ # API Reference ::: docket + +::: docket.queue diff --git a/docs/queues.md b/docs/queues.md new file mode 100644 index 0000000..89d7c8a --- /dev/null +++ b/docs/queues.md @@ -0,0 +1,101 @@ +# Reliable Message Queues + +Docket queues provide durable, at-least-once message delivery for systems that +own their own execution and result model. Unlike Docket tasks, queue messages +are opaque bytes: Docket does not import a function, execute the payload, or +store a result. + +This is useful when Docket is the delivery layer beneath another runtime: + +```python +from datetime import timedelta + +from docket import Docket + +async with Docket(name="orders") as docket: + queue = docket.queue("commands") + await queue.put( + "scheduled", + b'{"order_id": "123"}', + key="order:123:charge", + ) +``` + +A subscription competes with other subscriptions in the same consumer group. +Only one receives a given delivery: + +```python +async with Docket(name="orders") as docket: + queue = docket.queue("commands") + async with queue.subscribe( + {"retry": 0, "scheduled": 1}, + visibility_timeout=timedelta(minutes=5), + ) as subscription: + while True: + message = await subscription.receive() + try: + await execute_in_my_runtime(message.data) + except RetryableError: + await message.release("retry") + else: + await message.acknowledge() +``` + +Lower numeric topic priorities are returned first when multiple claimed +messages are ready. Each topic is FIFO. `release()` atomically moves a message +to another topic, which supports an immediate retry lane without losing the +original delivery. + +## Delivery guarantees + +Queue delivery is at least once: + +- A claimed message remains in Redis until it is acknowledged. +- The subscription renews visibility while the message is outstanding. +- If the subscriber exits or loses its Redis connection, another subscriber + can reclaim the message after `visibility_timeout`. +- `acknowledge()` removes the message only after the downstream runtime has + accepted it. + +Choose a visibility timeout longer than normal processing stalls and Redis +failovers. Consumers must still be idempotent because a process can finish its +side effect and fail before acknowledging the message. + +Message keys are deduplicated across every topic in a queue while the message +is queued or in flight. For repair loops that may rediscover accepted work, +retain a short acknowledgement tombstone: + +```python +queue = docket.queue( + "commands", + acknowledgement_ttl=timedelta(hours=1), +) +``` + +Publishing the same key during that period returns `False`; a new publication +returns `True`. + +## Backpressure + +Set `max_size` on `put()` to bound the number of queued and in-flight messages +in a topic: + +```python +await queue.put("scheduled", payload, max_size=1_000) +``` + +The publisher waits until capacity is available. The same option on +`release()` prevents an immediate-retry lane from exceeding its bound while +keeping the source delivery claimable until the atomic move succeeds. + +## Operations + +Queues use Redis Streams consumer groups and require Redis 6.2 or newer. +Subscriptions retry transient Redis errors, recreate expired consumer groups, +renew claims, and reclaim abandoned deliveries. They use the Docket's existing +standalone, cluster, Sentinel, authentication, and connection-pool +configuration. + +Use the same Docket name, queue name, and consumer group for replicas that +should share work. A queue supports one logical consumer group: acknowledged +messages are deleted rather than broadcast to independent groups. diff --git a/mkdocs.yml b/mkdocs.yml index 4f8e43a..e9607ab 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -76,6 +76,7 @@ nav: - Dependency Injection: dependency-injection.md - Task Design Patterns: task-patterns.md - Task Observability: observability.md + - Reliable Message Queues: queues.md - Testing with Docket: testing.md - Docket in Production: production.md - API Reference: api-reference.md diff --git a/pyproject.toml b/pyproject.toml index 75c79a0..417e248 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,6 +95,9 @@ docket = "docket.__main__:app" [tool.hatch.version] source = "vcs" +[tool.hatch.version.raw-options] +fallback_version = "0.0.0" + [tool.hatch.metadata] allow-direct-references = true diff --git a/src/docket/__init__.py b/src/docket/__init__.py index 5cc7c31..bf90d05 100644 --- a/src/docket/__init__.py +++ b/src/docket/__init__.py @@ -39,6 +39,7 @@ TaskCall, ) from .strikelist import StrikeList +from .queue import Queue, QueueMessage, QueueSubscription from .worker import Worker from . import testing @@ -63,6 +64,9 @@ "Logged", "Perpetual", "Progress", + "Queue", + "QueueMessage", + "QueueSubscription", "Retry", "Shared", "StrikeList", diff --git a/src/docket/_queue_scripts.py b/src/docket/_queue_scripts.py new file mode 100644 index 0000000..eeb69ef --- /dev/null +++ b/src/docket/_queue_scripts.py @@ -0,0 +1,99 @@ +"""Atomic Redis operations used by reliable message queues.""" + +from ._lua import Arg, Key, redis_script +from ._redis import RedisClient + + +@redis_script +async def put_message( + redis: RedisClient, + *, + stream_key: Key[str], + deduplication_key: Key[str], + message_key: Arg[str], + data: Arg[bytes], + max_size: Arg[int], + now_timestamp: Arg[float], +) -> bytes: + """ + redis.call('ZREMRANGEBYSCORE', deduplication_key, 1, now_timestamp) + if redis.call('ZSCORE', deduplication_key, message_key) then + return 'DUPLICATE' + end + if max_size > 0 and redis.call('XLEN', stream_key) >= max_size then + return 'FULL' + end + + local message_id = redis.call( + 'XADD', stream_key, '*', 'key', message_key, 'data', data + ) + redis.call('EXPIRE', stream_key, 2147483647) + redis.call('ZADD', deduplication_key, 0, message_key) + redis.call('EXPIRE', deduplication_key, 2147483647) + return message_id + """ + ... + + +@redis_script +async def acknowledge_message( + redis: RedisClient, + *, + stream_key: Key[str], + deduplication_key: Key[str], + group_name: Arg[str], + message_id: Arg[bytes], + message_key: Arg[str], + idle_ttl_seconds: Arg[int], + acknowledged_until: Arg[float], +) -> int: + """ + redis.call('XACK', stream_key, group_name, message_id) + redis.call('XDEL', stream_key, message_id) + if acknowledged_until > 0 then + redis.call('ZADD', deduplication_key, acknowledged_until, message_key) + redis.call('EXPIRE', deduplication_key, 2147483647) + else + redis.call('ZREM', deduplication_key, message_key) + end + if redis.call('XLEN', stream_key) == 0 then + redis.call('EXPIRE', stream_key, idle_ttl_seconds) + end + return 1 + """ + ... + + +@redis_script +async def release_message( + redis: RedisClient, + *, + source_stream_key: Key[str], + destination_stream_key: Key[str], + group_name: Arg[str], + message_id: Arg[bytes], + message_key: Arg[str], + data: Arg[bytes], + max_size: Arg[int], + idle_ttl_seconds: Arg[int], +) -> bytes: + """ + if source_stream_key ~= destination_stream_key + and max_size > 0 + and redis.call('XLEN', destination_stream_key) >= max_size + then + return 'FULL' + end + + redis.call('XACK', source_stream_key, group_name, message_id) + redis.call('XDEL', source_stream_key, message_id) + local new_message_id = redis.call( + 'XADD', destination_stream_key, '*', 'key', message_key, 'data', data + ) + redis.call('EXPIRE', destination_stream_key, 2147483647) + if redis.call('XLEN', source_stream_key) == 0 then + redis.call('EXPIRE', source_stream_key, idle_ttl_seconds) + end + return new_message_id + """ + ... diff --git a/src/docket/docket.py b/src/docket/docket.py index 165eddb..9aa9085 100644 --- a/src/docket/docket.py +++ b/src/docket/docket.py @@ -60,6 +60,7 @@ Strike, StrikeList, ) +from .queue import DocketQueueMixin logger: logging.Logger = logging.getLogger(__name__) tracer: trace.Tracer = trace.get_tracer(__name__) @@ -126,7 +127,7 @@ async def _cancel_task( TaskCollection = Iterable[TaskFunction] -class Docket(DocketSnapshotMixin): +class Docket(DocketQueueMixin, DocketSnapshotMixin): """A Docket represents a collection of tasks that may be scheduled for later execution. With a Docket, you can add, replace, and cancel tasks. Example: @@ -199,7 +200,6 @@ def worker_group_name(self) -> str: @property def prefix(self) -> str: """Return the key prefix for this docket. - All Redis keys for this docket are prefixed with this value. For Redis Cluster mode, returns a hash-tagged prefix like "{myapp}" diff --git a/src/docket/queue.py b/src/docket/queue.py new file mode 100644 index 0000000..7e7ace0 --- /dev/null +++ b/src/docket/queue.py @@ -0,0 +1,463 @@ +"""Reliable, keyed message delivery backed by Redis Streams.""" + +# pyright: reportPrivateUsage=false + +from __future__ import annotations + +import asyncio +import itertools +import logging +import time +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from datetime import timedelta +from typing import TYPE_CHECKING, cast +from uuid import uuid4 + +from redis.exceptions import ConnectionError, ResponseError, TimeoutError + +from ._queue_scripts import ( + acknowledge_message, + put_message, + release_message, +) + +if TYPE_CHECKING: + from .docket import Docket + +__all__ = ["Queue", "QueueMessage", "QueueSubscription"] + +logger: logging.Logger = logging.getLogger(__name__) + + +@dataclass(eq=False) +class QueueMessage: + """A message claimed by a :class:`QueueSubscription`. + + Call :meth:`acknowledge` after the downstream consumer accepts the + message. Call :meth:`release` to atomically move it to another topic for + immediate redelivery. If neither method is called, another subscription + can reclaim it after the visibility timeout. + """ + + data: bytes + key: str + topic: str + _subscription: QueueSubscription = field(repr=False) + _message_id: bytes = field(repr=False) + _settled: asyncio.Event = field(default_factory=asyncio.Event, repr=False) + _acknowledged: bool = field(default=False, repr=False) + + async def acknowledge(self) -> None: + """Permanently remove this delivery from the queue.""" + await self._subscription._acknowledge(self) + + async def release(self, topic: str, *, max_size: int = 0) -> None: + """Move this delivery to ``topic`` for immediate redelivery. + + If ``max_size`` is positive, wait until the destination topic has + capacity. The move is atomic: the original remains claimable until + the destination accepts it. + """ + await self._subscription._release(self, topic, max_size=max_size) + + +class Queue: + """A durable keyed message queue within a :class:`Docket`. + + Message keys are unique across the queue until acknowledgement. Publishing + the same key again while it is queued or in flight is an idempotent no-op. + Topics provide independent FIFO lanes and capacity limits. + """ + + def __init__( + self, + docket: Docket, + name: str, + *, + idle_ttl: timedelta = timedelta(hours=1), + acknowledgement_ttl: timedelta = timedelta(0), + ) -> None: + self.docket: Docket = docket + self.name: str = name + self.idle_ttl: timedelta = idle_ttl + self.acknowledgement_ttl: timedelta = acknowledgement_ttl + + async def put( + self, + topic: str, + data: bytes, + *, + key: str | None = None, + max_size: int = 0, + ) -> bool: + """Publish a message, waiting for topic capacity when bounded. + + Args: + topic: FIFO lane that subscriptions consume from. + data: Opaque message payload. + key: Idempotency key. Defaults to a fresh UUID. + max_size: Maximum messages queued or in flight on this topic. + Zero means unbounded. + + Returns: + ``True`` when published, or ``False`` when ``key`` was already + present in this queue. + """ + if max_size < 0: + raise ValueError("max_size must be non-negative") + if self.acknowledgement_ttl < timedelta(0): + raise ValueError("acknowledgement_ttl must be non-negative") + message_key = key or str(uuid4()) + while True: + async with self.docket.redis() as redis: + result = await put_message( + redis, + stream_key=self._stream_key(topic), + deduplication_key=self._deduplication_key, + message_key=message_key, + data=data, + max_size=max_size, + now_timestamp=time.time(), + ) + if result not in (b"FULL", "FULL"): + return result not in (b"DUPLICATE", "DUPLICATE") + await asyncio.sleep(0.01) + + def subscribe( + self, + topics: Sequence[str] | Mapping[str, int], + *, + visibility_timeout: timedelta, + group: str = "consumers", + ) -> QueueSubscription: + """Create a subscription to one or more topics. + + A mapping assigns lower numeric priorities to topics that should be + delivered first. A sequence gives every topic equal priority. + """ + priorities = ( + dict(topics) + if isinstance(topics, Mapping) + else {topic: 0 for topic in topics} + ) + if not priorities: + raise ValueError("at least one topic is required") + if visibility_timeout <= timedelta(0): + raise ValueError("visibility_timeout must be positive") + return QueueSubscription(self, priorities, visibility_timeout, group) + + @property + def _deduplication_key(self) -> str: + return self.docket.key(f"queues:{self.name}:messages") + + def _stream_key(self, topic: str) -> str: + return self.docket.key(f"queues:{self.name}:topics:{topic}") + + @property + def _idle_ttl_seconds(self) -> int: + return max(1, int(self.idle_ttl.total_seconds())) + + @property + def _acknowledged_until(self) -> float: + if not self.acknowledgement_ttl: + return 0 + return time.time() + self.acknowledgement_ttl.total_seconds() + + +class QueueSubscription: + """A competing-consumer subscription with visibility-based redelivery.""" + + def __init__( + self, + queue: Queue, + priorities: Mapping[str, int], + visibility_timeout: timedelta, + group: str, + ) -> None: + self.queue: Queue = queue + self.priorities: dict[str, int] = dict(priorities) + self.visibility_timeout: timedelta = visibility_timeout + self.group: str = group + self.consumer: str = str(uuid4()) + self._available: asyncio.PriorityQueue[tuple[int, int, QueueMessage]] = ( + asyncio.PriorityQueue(maxsize=1) + ) + self._sequence = itertools.count() + self._outstanding: set[QueueMessage] = set() + self._claimed: list[QueueMessage] = [] + self._tasks: list[asyncio.Task[None]] = [] + self._initialized_streams: set[str] = set() + self._next_recovery: dict[str, float] = {} + self._entered = False + + async def __aenter__(self) -> QueueSubscription: + if self._entered: + raise RuntimeError("queue subscription is already active") + self._entered = True + self._tasks.append(asyncio.create_task(self._consume())) + self._tasks.append(asyncio.create_task(self._renew_visibility())) + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> None: + for task in self._tasks: + task.cancel() + await asyncio.gather(*self._tasks, return_exceptions=True) + for message in list(self._outstanding): + try: + await self._release(message, message.topic, max_size=0) + except (ConnectionError, ResponseError, TimeoutError): + logger.warning( + "Queue consumer %s could not release a claimed message", + self.consumer, + exc_info=True, + ) + self._claimed.clear() + while not self._available.empty(): + self._available.get_nowait() + self._tasks.clear() + self._entered = False + + async def receive(self, *, timeout: float | None = None) -> QueueMessage: + """Wait for the next claimed message.""" + if not self._entered: + raise RuntimeError("queue subscription is not active") + if timeout is None: + _, _, message = await self._available.get() + else: + _, _, message = await asyncio.wait_for( + self._available.get(), timeout=timeout + ) + return message + + async def _consume(self) -> None: + while True: + try: + message = await self._claim() + self._outstanding.add(message) + await self._available.put( + ( + self.priorities[message.topic], + next(self._sequence), + message, + ) + ) + await message._settled.wait() + except asyncio.CancelledError: + raise + except (ConnectionError, ResponseError, TimeoutError): + logger.warning( + "Queue consumer %s lost its Redis connection; retrying", + self.consumer, + exc_info=True, + ) + await asyncio.sleep(0.5) + + async def _claim(self) -> QueueMessage: + if self._claimed: + return self._claimed.pop(0) + + streams = {self.queue._stream_key(topic): topic for topic in self.priorities} + while True: + async with self.queue.docket.redis() as redis: + for stream_key in streams: + if stream_key not in self._initialized_streams: + try: + await redis.xgroup_create( + stream_key, + self.group, + id="0", + mkstream=True, + ) + except ResponseError as exc: + if "BUSYGROUP" not in str(exc): + raise + if await redis.xlen(stream_key) == 0: + await redis.expire( + stream_key, + self.queue._idle_ttl_seconds, + ) + self._initialized_streams.add(stream_key) + + loop_time = asyncio.get_running_loop().time() + for topic in sorted(self.priorities, key=self.priorities.get): + if loop_time >= self._next_recovery.get(topic, 0): + recovered = await redis.xautoclaim( + self.queue._stream_key(topic), + self.group, + self.consumer, + min_idle_time=int( + self.visibility_timeout.total_seconds() * 1000 + ), + start_id="0-0", + count=1, + ) + if recovered[1]: + message_id, fields = recovered[1][0] + return self._message(topic, message_id, fields) + self._next_recovery[topic] = loop_time + min( + 1, self.visibility_timeout.total_seconds() / 2 + ) + + try: + result = await redis.xreadgroup( + self.group, + self.consumer, + streams={stream_key: ">" for stream_key in streams}, + count=1, + block=1000, + ) + except ResponseError as exc: + if "NOGROUP" not in str(exc): + raise + self._initialized_streams.clear() + continue + if result: + messages = [ + self._message( + streams[ + ( + raw_stream_key.decode() + if isinstance(raw_stream_key, bytes) + else raw_stream_key + ) + ], + message_id, + fields, + ) + for raw_stream_key, stream_messages in result + for message_id, fields in stream_messages + ] + messages.sort(key=lambda message: self.priorities[message.topic]) + self._outstanding.update(messages) + self._claimed.extend(messages[1:]) + return messages[0] + + def _message( + self, topic: str, message_id: bytes, fields: Mapping[bytes, bytes] + ) -> QueueMessage: + try: + key = fields[b"key"].decode() + data = fields[b"data"] + except KeyError as exc: + raise ValueError( + f"queue message {message_id!r} is missing {exc.args[0]!r}" + ) from exc + return QueueMessage( + data=data, + key=key, + topic=topic, + _subscription=self, + _message_id=message_id, + ) + + async def _acknowledge(self, message: QueueMessage) -> None: + """Acknowledge a message claimed by this subscription.""" + if message._acknowledged: + return + if not message._settled.is_set(): + self._settle(message) + async with self.queue.docket.redis() as redis: + await acknowledge_message( + redis, + stream_key=self.queue._stream_key(message.topic), + deduplication_key=self.queue._deduplication_key, + group_name=self.group, + message_id=message._message_id, + message_key=message.key, + idle_ttl_seconds=self.queue._idle_ttl_seconds, + acknowledged_until=self.queue._acknowledged_until, + ) + message._acknowledged = True + + async def _release( + self, message: QueueMessage, topic: str, *, max_size: int + ) -> None: + """Release a message to another topic for immediate redelivery.""" + if max_size < 0: + raise ValueError("max_size must be non-negative") + if message._settled.is_set(): + return + while True: + async with self.queue.docket.redis() as redis: + result = await release_message( + redis, + source_stream_key=self.queue._stream_key(message.topic), + destination_stream_key=self.queue._stream_key(topic), + group_name=self.group, + message_id=message._message_id, + message_key=message.key, + data=message.data, + max_size=max_size, + idle_ttl_seconds=self.queue._idle_ttl_seconds, + ) + if result not in (b"FULL", "FULL"): + self._settle(message) + return + await asyncio.sleep(0.01) + + def _settle(self, message: QueueMessage) -> None: + self._outstanding.discard(message) + message._settled.set() + + async def _renew_visibility(self) -> None: + interval = max(0.01, self.visibility_timeout.total_seconds() / 4) + while True: + await asyncio.sleep(interval) + by_topic: dict[str, list[bytes]] = {} + for message in self._outstanding: + by_topic.setdefault(message.topic, []).append(message._message_id) + if not by_topic: + continue + try: + async with self.queue.docket.redis() as redis: + for topic, message_ids in by_topic.items(): + await redis.xclaim( + self.queue._stream_key(topic), + self.group, + self.consumer, + min_idle_time=0, + message_ids=message_ids, + idle=0, + justid=True, + ) + except (ConnectionError, ResponseError, TimeoutError): + logger.warning( + "Queue consumer %s could not renew message visibility", + self.consumer, + exc_info=True, + ) + + +class DocketQueueMixin: + """Construct reliable message queues scoped to a Docket.""" + + def queue( + self, + name: str, + *, + acknowledgement_ttl: timedelta = timedelta(0), + ) -> Queue: + """Return a reliable message queue scoped to this Docket. + + Queues provide keyed, at-least-once delivery independently of task + execution. Use them when another runtime owns the work but needs + Docket's Redis-backed publication, acknowledgement, and crash + redelivery semantics. + + Args: + name: Stable queue name shared by publishers and subscribers. + acknowledgement_ttl: How long acknowledged message keys remain + deduplicated. This supports repair loops that may briefly + rediscover already-delivered work. + """ + return Queue( + cast("Docket", self), + name, + acknowledgement_ttl=acknowledgement_ttl, + ) diff --git a/tests/test_queue.py b/tests/test_queue.py new file mode 100644 index 0000000..7c72200 --- /dev/null +++ b/tests/test_queue.py @@ -0,0 +1,492 @@ +# pyright: reportPrivateUsage=false + +import asyncio +import logging +from contextlib import asynccontextmanager +from datetime import timedelta +from typing import AsyncGenerator +from unittest.mock import AsyncMock, patch + +import pytest +from redis.exceptions import ConnectionError, ResponseError + +from docket import Docket +from docket.queue import Queue, QueueSubscription + + +async def test_queue_round_trip_is_fifo_and_idempotent(docket: Docket) -> None: + queue = docket.queue("jobs") + + assert await queue.put("alpha", b"one", key="one") + assert not await queue.put("alpha", b"duplicate", key="one") + assert await queue.put("alpha", b"two", key="two") + + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) as subscription: + first = await subscription.receive(timeout=1) + assert (first.key, first.data, first.topic) == ("one", b"one", "alpha") + await first.acknowledge() + + second = await subscription.receive(timeout=1) + assert (second.key, second.data) == ("two", b"two") + await second.acknowledge() + + assert await queue.put("alpha", b"one-again", key="one") + + +async def test_receive_without_timeout(docket: Docket) -> None: + queue = docket.queue("jobs") + assert await queue.put("alpha", b"one") + + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) as subscription: + message = await subscription.receive() + await message.acknowledge() + + +async def test_acknowledgement_tombstone_deduplicates_repairs( + docket: Docket, +) -> None: + queue = docket.queue( + "jobs", + acknowledgement_ttl=timedelta(milliseconds=50), + ) + assert await queue.put("alpha", b"one", key="one") + + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) as subscription: + message = await subscription.receive(timeout=1) + await message.acknowledge() + await message.acknowledge() + + assert not await queue.put("alpha", b"too-soon", key="one") + await asyncio.sleep(0.06) + assert await queue.put("alpha", b"after-expiry", key="one") + + +async def test_failed_acknowledgement_is_redelivered(docket: Docket) -> None: + queue = docket.queue("jobs") + assert await queue.put("alpha", b"one", key="one") + + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(milliseconds=50) + ) as first: + message = await first.receive(timeout=1) + with ( + patch( + "docket.queue.acknowledge_message", + new=AsyncMock(side_effect=ConnectionError("offline")), + ), + pytest.raises(ConnectionError, match="offline"), + ): + await message.acknowledge() + + # The failed durable acknowledgement is locally settled, so it is no + # longer renewed by the first subscriber. Allow margin for coarse + # event-loop clocks before another subscriber attempts to reclaim it. + await asyncio.sleep(0.1) + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(milliseconds=50) + ) as second: + redelivered = await second.receive(timeout=2) + assert redelivered.key == "one" + await redelivered.acknowledge() + + +async def test_acknowledgement_can_retry_after_local_settlement( + docket: Docket, +) -> None: + queue = docket.queue("jobs") + assert await queue.put("alpha", b"one") + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) as subscription: + message = await subscription.receive(timeout=1) + subscription._settle(message) + await message.acknowledge() + + +async def test_bounded_put_waits_for_acknowledgement(docket: Docket) -> None: + queue = docket.queue("jobs") + assert await queue.put("alpha", b"one", key="one", max_size=1) + blocked = asyncio.create_task(queue.put("alpha", b"two", key="two", max_size=1)) + await asyncio.sleep(0.02) + assert not blocked.done() + + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) as subscription: + message = await subscription.receive(timeout=1) + await message.acknowledge() + assert await asyncio.wait_for(blocked, timeout=1) + message = await subscription.receive(timeout=1) + await message.acknowledge() + + +async def test_release_moves_message_to_priority_topic(docket: Docket) -> None: + queue = docket.queue("jobs") + assert await queue.put("scheduled", b"one", key="one") + + async with queue.subscribe( + {"retry": 0, "scheduled": 1}, + visibility_timeout=timedelta(seconds=1), + ) as subscription: + message = await subscription.receive(timeout=1) + await message.release("retry", max_size=1) + + redelivered = await subscription.receive(timeout=1) + assert redelivered.key == "one" + assert redelivered.topic == "retry" + await redelivered.acknowledge() + + +async def test_bounded_release_waits_for_destination_capacity( + docket: Docket, +) -> None: + queue = docket.queue("jobs") + assert await queue.put("scheduled", b"one", key="one") + assert await queue.put("retry", b"blocker", key="blocker") + + async with queue.subscribe( + ["scheduled"], visibility_timeout=timedelta(seconds=1) + ) as scheduled: + message = await scheduled.receive(timeout=1) + release = asyncio.create_task(message.release("retry", max_size=1)) + await asyncio.sleep(0.02) + assert not release.done() + + async with queue.subscribe( + ["retry"], visibility_timeout=timedelta(seconds=1) + ) as retry: + blocker = await retry.receive(timeout=1) + await blocker.acknowledge() + await asyncio.wait_for(release, timeout=1) + + redelivered = await retry.receive(timeout=1) + assert redelivered.key == "one" + await redelivered.acknowledge() + + +async def test_bounded_release_to_same_topic_does_not_deadlock( + docket: Docket, +) -> None: + queue = docket.queue("jobs") + assert await queue.put("retry", b"one", max_size=1) + async with queue.subscribe( + ["retry"], visibility_timeout=timedelta(seconds=1) + ) as subscription: + message = await subscription.receive(timeout=1) + await asyncio.wait_for(message.release("retry", max_size=1), timeout=1) + redelivered = await subscription.receive(timeout=1) + await redelivered.acknowledge() + + +async def test_unacknowledged_message_is_reclaimed_after_visibility_timeout( + docket: Docket, +) -> None: + queue = docket.queue("jobs") + assert await queue.put("alpha", b"one", key="one") + + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(milliseconds=50) + ) as first: + claimed = await first.receive(timeout=1) + assert claimed.key == "one" + + await asyncio.sleep(0.06) + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(milliseconds=50) + ) as second: + reclaimed = await second.receive(timeout=1) + assert reclaimed.key == "one" + await reclaimed.acknowledge() + + +async def test_visibility_is_renewed_while_subscription_is_active( + docket: Docket, +) -> None: + queue = docket.queue("jobs") + assert await queue.put("alpha", b"one", key="one") + + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(milliseconds=50) + ) as first: + claimed = await first.receive(timeout=1) + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(milliseconds=50) + ) as second: + await asyncio.sleep(0.12) + with pytest.raises(asyncio.TimeoutError): + await second.receive(timeout=0.05) + await claimed.acknowledge() + + +async def test_topics_route_messages_to_matching_subscriptions( + docket: Docket, +) -> None: + queue = docket.queue("jobs") + assert await queue.put("alpha", b"one", key="one") + + async with ( + queue.subscribe(["beta"], visibility_timeout=timedelta(seconds=1)) as wrong, + queue.subscribe(["alpha"], visibility_timeout=timedelta(seconds=1)) as right, + ): + message = await right.receive(timeout=1) + with pytest.raises(asyncio.TimeoutError): + await wrong.receive(timeout=0.05) + await message.acknowledge() + + +async def test_subscription_reads_many_topics_with_constant_background_tasks( + docket: Docket, +) -> None: + queue = docket.queue("jobs") + topics = [f"topic-{index}" for index in range(25)] + for topic in topics: + assert await queue.put(topic, topic.encode(), key=topic) + + async with queue.subscribe( + topics, visibility_timeout=timedelta(seconds=1) + ) as subscription: + assert len(subscription._tasks) == 2 + received: list[str] = [] + for _ in topics: + message = await subscription.receive(timeout=1) + received.append(message.topic) + await message.acknowledge() + + assert set(received) == set(topics) + + +async def test_lower_priority_number_is_delivered_first(docket: Docket) -> None: + queue = docket.queue("jobs") + assert await queue.put("scheduled", b"scheduled") + assert await queue.put("retry", b"retry") + + async with queue.subscribe( + {"retry": 0, "scheduled": 1}, + visibility_timeout=timedelta(seconds=1), + ) as subscription: + first = await subscription.receive(timeout=1) + await first.acknowledge() + second = await subscription.receive(timeout=1) + await second.acknowledge() + + assert (first.topic, second.topic) == ("retry", "scheduled") + + +async def test_competing_consumers_do_not_reclaim_prefetched_messages( + docket: Docket, +) -> None: + queue = docket.queue("jobs") + for index in range(10): + assert await queue.put("alpha", str(index).encode(), key=str(index)) + + first_data: list[bytes] = [] + second_data: list[bytes] = [] + async with ( + queue.subscribe( + ["alpha"], visibility_timeout=timedelta(milliseconds=50) + ) as first, + queue.subscribe( + ["alpha"], visibility_timeout=timedelta(milliseconds=50) + ) as second, + ): + for _ in range(5): + first_message = await first.receive(timeout=1) + first_data.append(first_message.data) + await first_message.acknowledge() + second_message = await second.receive(timeout=1) + second_data.append(second_message.data) + await second_message.acknowledge() + + assert set(first_data).isdisjoint(second_data) + assert set(first_data + second_data) == {str(index).encode() for index in range(10)} + + +async def test_queue_validates_configuration(docket: Docket) -> None: + queue = docket.queue("jobs") + + with pytest.raises(ValueError, match="max_size must be non-negative"): + await queue.put("alpha", b"one", max_size=-1) + with pytest.raises(ValueError, match="at least one topic"): + queue.subscribe([], visibility_timeout=timedelta(seconds=1)) + with pytest.raises(ValueError, match="visibility_timeout must be positive"): + queue.subscribe(["alpha"], visibility_timeout=timedelta(0)) + + invalid_tombstone = Queue( + docket, + "invalid", + acknowledgement_ttl=timedelta(seconds=-1), + ) + with pytest.raises(ValueError, match="acknowledgement_ttl"): + await invalid_tombstone.put("alpha", b"one") + + +async def test_message_validates_release_and_settlement(docket: Docket) -> None: + queue = docket.queue("jobs") + assert await queue.put("alpha", b"one") + async with queue.subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) as subscription: + message = await subscription.receive(timeout=1) + with pytest.raises(ValueError, match="max_size must be non-negative"): + await message.release("alpha", max_size=-1) + await message.acknowledge() + await message.release("alpha") + + +async def test_corrupt_message_is_rejected(docket: Docket) -> None: + subscription = docket.queue("jobs").subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) + with pytest.raises(ValueError, match="is missing"): + subscription._message("alpha", b"1-0", {b"key": b"one"}) + + +async def test_consumer_retries_connection_errors( + docket: Docket, caplog: pytest.LogCaptureFixture +) -> None: + subscription = docket.queue("jobs").subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) + subscription._claim = AsyncMock( + side_effect=[ConnectionError("offline"), asyncio.CancelledError()] + ) + with ( + caplog.at_level(logging.WARNING), + patch("docket.queue.asyncio.sleep", new=AsyncMock()), + pytest.raises(asyncio.CancelledError), + ): + await subscription._consume() + assert "lost its Redis connection" in caplog.text + + +@asynccontextmanager +async def _redis_connection(redis: AsyncMock) -> AsyncGenerator[AsyncMock, None]: + yield redis + + +async def test_claim_recovers_a_missing_group(docket: Docket) -> None: + subscription = docket.queue("jobs").subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) + stream_key = subscription.queue._stream_key("alpha") + subscription._initialized_streams.add(stream_key) + subscription._next_recovery["alpha"] = float("inf") + redis = AsyncMock() + redis.xreadgroup.side_effect = [ + ResponseError("NOGROUP no such key"), + [], + [(stream_key, [(b"1-0", {b"key": b"one", b"data": b"payload"})])], + ] + + with patch.object(docket, "redis", side_effect=lambda: _redis_connection(redis)): + message = await subscription._claim() + + assert message.key == "one" + redis.xgroup_create.assert_awaited_once() + + +async def test_claim_propagates_other_redis_errors(docket: Docket) -> None: + subscription = docket.queue("jobs").subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) + stream_key = subscription.queue._stream_key("alpha") + subscription._initialized_streams.add(stream_key) + subscription._next_recovery["alpha"] = float("inf") + redis = AsyncMock() + redis.xreadgroup.side_effect = ResponseError("WRONGTYPE") + + with ( + patch.object(docket, "redis", side_effect=lambda: _redis_connection(redis)), + pytest.raises(ResponseError, match="WRONGTYPE"), + ): + await subscription._claim() + + +async def test_group_creation_propagates_other_redis_errors( + docket: Docket, +) -> None: + subscription = docket.queue("jobs").subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) + redis = AsyncMock() + redis.xgroup_create.side_effect = ResponseError("WRONGTYPE") + + with ( + patch.object(docket, "redis", side_effect=lambda: _redis_connection(redis)), + pytest.raises(ResponseError, match="WRONGTYPE"), + ): + await subscription._claim() + + +async def test_visibility_renewal_retries_connection_errors( + docket: Docket, caplog: pytest.LogCaptureFixture +) -> None: + subscription: QueueSubscription = docket.queue("jobs").subscribe( + ["alpha"], visibility_timeout=timedelta(milliseconds=40) + ) + message = subscription._message( + "alpha", b"1-0", {b"key": b"one", b"data": b"payload"} + ) + subscription._outstanding.add(message) + redis = AsyncMock() + redis.xclaim.side_effect = ConnectionError("offline") + + with ( + caplog.at_level(logging.WARNING), + patch.object(docket, "redis", side_effect=lambda: _redis_connection(redis)), + patch( + "docket.queue.asyncio.sleep", + new=AsyncMock(side_effect=[None, asyncio.CancelledError()]), + ), + pytest.raises(asyncio.CancelledError), + ): + await subscription._renew_visibility() + assert "could not renew message visibility" in caplog.text + + +async def test_subscription_exit_tolerates_release_errors_and_clears_buffer( + docket: Docket, caplog: pytest.LogCaptureFixture +) -> None: + subscription = docket.queue("jobs").subscribe( + ["alpha"], visibility_timeout=timedelta(seconds=1) + ) + message = subscription._message( + "alpha", b"1-0", {b"key": b"one", b"data": b"payload"} + ) + subscription._outstanding.add(message) + subscription._available.put_nowait((0, 0, message)) + subscription._release = AsyncMock(side_effect=ConnectionError("offline")) + + with caplog.at_level(logging.WARNING): + await subscription.__aexit__(None, None, None) + + assert "could not release a claimed message" in caplog.text + assert subscription._available.empty() + + +async def test_subscription_must_be_active_and_cannot_be_reentered( + docket: Docket, +) -> None: + queue = docket.queue("jobs") + subscription = queue.subscribe(["alpha"], visibility_timeout=timedelta(seconds=1)) + + with pytest.raises(RuntimeError, match="not active"): + await subscription.receive(timeout=0.01) + + async with subscription: + with pytest.raises(RuntimeError, match="already active"): + await subscription.__aenter__() + + assert await queue.put("alpha", b"one") + async with subscription: + first = await subscription.receive(timeout=1) + async with subscription: + redelivered = await subscription.receive(timeout=1) + await redelivered.acknowledge() + assert redelivered.key == first.key