redis-stream-queue 0.1.0__py3-none-any.whl

redis_stream_queue/__init__.py

@@ -0,0 +1,27 @@
+from .client import StreamClient
+from .consumer import StreamConsumer
+from .group import ConsumerConfig, ConsumerGroup
+from .message import ConsumerInfo, ConsumerMetrics, PendingEntry, ProducerMetrics, StreamMessage, StreamStats
+from .producer import StreamProducer
+from .retry import DLQHandler, RetryHandler
+from .serializers import JsonSerializer, MsgpackSerializer, PickleSerializer, Serializer
+
+__all__ = [
+    "StreamClient",
+    "StreamConsumer",
+    "StreamProducer",
+    "ConsumerConfig",
+    "ConsumerGroup",
+    "RetryHandler",
+    "StreamMessage",
+    "PendingEntry",
+    "StreamStats",
+    "ConsumerInfo",
+    "ConsumerMetrics",
+    "ProducerMetrics",
+    "DLQHandler",
+    "Serializer",
+    "JsonSerializer",
+    "MsgpackSerializer",
+    "PickleSerializer",
+]

redis_stream_queue/client.py

@@ -0,0 +1,306 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import ClassVar
+
+import redis.asyncio as aioredis
+
+from .message import PendingEntry, StreamMessage
+
+logger = logging.getLogger(__name__)
+
+_BUSYGROUP = "BUSYGROUP"
+
+
+class StreamClient:
+    """
+    Async Redis Stream client with singleton connection pool per (host, port, db).
+    Supports single-node, cluster, and URL connection modes.
+    """
+
+    _pools: ClassVar[dict[tuple, aioredis.Redis]] = {}
+    _lock: ClassVar[asyncio.Lock | None] = None
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 6379,
+        db: int = 0,
+        username: str | None = None,
+        password: str | None = None,
+        prefix: str = "",
+        max_connections: int = 1000,
+        pool_timeout: int = 5,
+        client_name: str | None = None,
+        ssl: bool = False,
+    ) -> None:
+        self._host = host
+        self._port = port
+        self._db = db
+        self._username = username
+        self._password = password
+        self._prefix = prefix
+        self._max_connections = max_connections
+        self._pool_timeout = pool_timeout
+        self._client_name = client_name
+        self._ssl = ssl
+        self._cluster = False
+        self._startup_nodes: list[dict] | None = None
+        self._url: str | None = None
+
+    @classmethod
+    def from_url(cls, url: str, **kwargs) -> "StreamClient":
+        instance = cls.__new__(cls)
+        instance._url = url
+        instance._cluster = url.startswith(("redis+cluster://", "rediss+cluster://"))
+        instance._startup_nodes = None
+        instance._prefix = kwargs.get("prefix", "")
+        instance._max_connections = kwargs.get("max_connections", 1000)
+        instance._pool_timeout = kwargs.get("pool_timeout", 5)
+        instance._client_name = kwargs.get("client_name", None)
+        instance._host = None
+        instance._port = None
+        instance._db = kwargs.get("db", 0)
+        instance._username = None
+        instance._password = None
+        instance._ssl = False
+        return instance
+
+    @classmethod
+    def from_cluster(
+        cls,
+        startup_nodes: list[dict],
+        password: str | None = None,
+        username: str | None = None,
+        prefix: str = "",
+        **kwargs,
+    ) -> "StreamClient":
+        instance = cls.__new__(cls)
+        instance._startup_nodes = startup_nodes
+        instance._cluster = True
+        instance._password = password
+        instance._username = username
+        instance._prefix = prefix
+        instance._max_connections = kwargs.get("max_connections", 1000)
+        instance._pool_timeout = kwargs.get("pool_timeout", 5)
+        instance._client_name = kwargs.get("client_name", None)
+        instance._host = None
+        instance._port = None
+        instance._db = 0
+        instance._ssl = kwargs.get("ssl", False)
+        instance._url = None
+        return instance
+
+    @classmethod
+    def _get_lock(cls) -> asyncio.Lock:
+        if cls._lock is None:
+            cls._lock = asyncio.Lock()
+        return cls._lock
+
+    def _pool_key(self) -> tuple:
+        if self._cluster and self._startup_nodes:
+            return ("cluster", tuple(sorted(str(n) for n in self._startup_nodes)))
+        return (self._url or f"{self._host}:{self._port}:{self._db}", self._prefix)
+
+    async def _get_client(self) -> aioredis.Redis:
+        key = self._pool_key()
+        async with self._get_lock():
+            if key not in self._pools:
+                self._pools[key] = await self._create_client()
+            return self._pools[key]
+
+    async def _create_client(self) -> aioredis.Redis:
+        if self._cluster and self._startup_nodes:
+            from redis.asyncio.cluster import ClusterNode, RedisCluster
+            nodes = [ClusterNode(n["host"], n["port"]) for n in self._startup_nodes]
+            return RedisCluster(startup_nodes=nodes, password=self._password, decode_responses=False)
+
+        if self._url:
+            pool = aioredis.BlockingConnectionPool.from_url(
+                self._url,
+                max_connections=self._max_connections,
+                timeout=self._pool_timeout,
+                decode_responses=False,
+            )
+        else:
+            pool_kwargs: dict = dict(
+                host=self._host,
+                port=self._port,
+                db=self._db,
+                username=self._username,
+                password=self._password,
+                client_name=self._client_name,
+                max_connections=self._max_connections,
+                timeout=self._pool_timeout,
+                decode_responses=False,
+            )
+            if self._ssl:
+                pool_kwargs["ssl"] = True
+            pool = aioredis.BlockingConnectionPool(**pool_kwargs)
+        return aioredis.Redis(connection_pool=pool)
+
+    def _prefixed(self, key: str) -> str:
+        return f"{self._prefix}_{key}" if self._prefix else key
+
+    # ── Stream operations ────────────────────────────────────────────────────
+
+    async def push(self, stream: str, data: bytes, max_len: int = 100_000) -> str:
+        client = await self._get_client()
+        msg_id = await client.xadd(self._prefixed(stream), {"data": data}, maxlen=max_len, approximate=True)
+        return msg_id.decode() if isinstance(msg_id, bytes) else msg_id
+
+    async def push_many(self, stream: str, items: list[bytes], max_len: int = 100_000) -> list[str]:
+        if not items:
+            return []
+        client = await self._get_client()
+        async with client.pipeline(transaction=False) as pipe:
+            for data in items:
+                pipe.xadd(self._prefixed(stream), {"data": data}, maxlen=max_len, approximate=True)
+            results = await pipe.execute()
+        return [r.decode() if isinstance(r, bytes) else r for r in results]
+
+    async def create_group(self, stream: str, group: str, start_id: str = "0") -> None:
+        client = await self._get_client()
+        try:
+            await client.xgroup_create(self._prefixed(stream), group, id=start_id, mkstream=True)
+        except Exception as e:
+            if _BUSYGROUP not in str(e):
+                raise
+
+    async def read(
+        self, stream: str, group: str, consumer: str, count: int, block_ms: int
+    ) -> list[StreamMessage]:
+        client = await self._get_client()
+        resp = await client.xreadgroup(
+            groupname=group,
+            consumername=consumer,
+            streams={self._prefixed(stream): ">"},
+            count=count,
+            block=block_ms,
+        )
+        if not resp:
+            return []
+        messages = []
+        for _stream_name, entries in resp:
+            for msg_id, fields in entries:
+                raw_id = msg_id.decode() if isinstance(msg_id, bytes) else msg_id
+                raw_data = fields.get(b"data", fields.get("data", b""))
+                messages.append(StreamMessage(id=raw_id, data=raw_data, delivery_count=1))
+        return messages
+
+    async def ack(self, stream: str, group: str, *ids: str) -> int:
+        if not ids:
+            return 0
+        client = await self._get_client()
+        return await client.xack(self._prefixed(stream), group, *ids)
+
+    async def autoclaim(
+        self,
+        stream: str,
+        group: str,
+        consumer: str,
+        min_idle_ms: int,
+        count: int,
+        max_passes: int | None = None,
+    ) -> list[StreamMessage]:
+        """
+        Reclaim idle PEL entries via XAUTOCLAIM.
+
+        Follows the cursor until the full PEL is swept (cursor returns "0-0").
+        max_passes caps the number of XAUTOCLAIM calls per invocation; None = unlimited.
+        Stall detection (cursor unchanged) guards against misbehaving Redis forks.
+        """
+        r = await self._get_client()
+        messages: list[StreamMessage] = []
+        cursor = "0-0"
+        passes = 0
+        while True:
+            if max_passes is not None and passes >= max_passes:
+                break
+            next_cursor, entries, _deleted = await r.xautoclaim(
+                name=self._prefixed(stream),
+                groupname=group,
+                consumername=consumer,
+                min_idle_time=min_idle_ms,
+                start_id=cursor,
+                count=count,
+            )
+            passes += 1
+            for msg_id, fields in entries:
+                raw_id = msg_id.decode() if isinstance(msg_id, bytes) else msg_id
+                raw_data = fields.get(b"data", fields.get("data", b""))
+                # XAUTOCLAIM doesn't return delivery_count in entries — 1 is minimum
+                # since reclaimed msgs were already delivered (real count lives in PEL)
+                messages.append(StreamMessage(id=raw_id, data=raw_data, delivery_count=1))
+            next_str = next_cursor.decode() if isinstance(next_cursor, bytes) else str(next_cursor)
+            if next_str == "0-0" or next_str == cursor:
+                # "0-0"           → full PEL sweep complete (real Redis)
+                # next_str==cursor → no progress / stall (fakeredis bug guard)
+                break
+            cursor = next_str
+        return messages
+
+    async def pending_range(
+        self, stream: str, group: str, count: int, consumer: str | None = None
+    ) -> list[PendingEntry]:
+        client = await self._get_client()
+        resp = await client.xpending_range(
+            name=self._prefixed(stream), groupname=group, min="-", max="+",
+            count=count, consumername=consumer,
+        )
+        result = []
+        for entry in resp:
+            msg_id = entry["message_id"]
+            raw_id = msg_id.decode() if isinstance(msg_id, bytes) else msg_id
+            consumer_name = entry["consumer"]
+            if isinstance(consumer_name, bytes):
+                consumer_name = consumer_name.decode()
+            result.append(PendingEntry(
+                id=raw_id, consumer=consumer_name,
+                idle_ms=entry.get("time_since_delivered", 0),
+                delivery_count=entry.get("times_delivered", 0),
+            ))
+        return result
+
+    async def fetch_by_ids(self, stream: str, ids: list[str]) -> list[StreamMessage]:
+        if not ids:
+            return []
+        client = await self._get_client()
+        results = []
+        for msg_id in ids:
+            resp = await client.xrange(self._prefixed(stream), min=msg_id, max=msg_id, count=1)
+            for entry_id, fields in resp:
+                raw_id = entry_id.decode() if isinstance(entry_id, bytes) else entry_id
+                raw_data = fields.get(b"data", fields.get("data", b""))
+                results.append(StreamMessage(id=raw_id, data=raw_data))
+        return results
+
+    async def stream_len(self, stream: str) -> int:
+        client = await self._get_client()
+        return await client.xlen(self._prefixed(stream))
+
+    async def group_info(self, stream: str) -> list[dict]:
+        client = await self._get_client()
+        return await client.xinfo_groups(self._prefixed(stream))
+
+    async def consumer_info(self, stream: str, group: str) -> list[dict]:
+        client = await self._get_client()
+        return await client.xinfo_consumers(self._prefixed(stream), group)
+
+    async def close(self) -> None:
+        key = self._pool_key()
+        async with self._get_lock():
+            client = self._pools.pop(key, None)
+            if client:
+                await client.aclose()
+
+    @classmethod
+    async def close_all(cls) -> None:
+        async with cls._get_lock():
+            for client in cls._pools.values():
+                try:
+                    await client.aclose()
+                except Exception:
+                    pass
+            cls._pools.clear()

redis_stream_queue/consumer.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+import weakref
+from typing import Awaitable, Callable, ClassVar
+
+from .client import StreamClient
+from .group import ConsumerConfig, ConsumerGroup
+from .message import ConsumerMetrics, StreamMessage, _TpsTracker
+from .retry import DLQHandler, RetryHandler
+from .serializers import JsonSerializer, Serializer
+
+logger = logging.getLogger(__name__)
+
+MessageHandler = Callable[[list[StreamMessage]], Awaitable[list[str]]]
+ErrorHook = Callable[[Exception, list[StreamMessage]], Awaitable[None]]
+
+_NOGROUP = "NOGROUP"
+
+
+class StreamConsumer:
+    """
+    Async Redis Stream consumer using consumer groups.
+
+    Lifecycle per iteration (run_once):
+      1. Ensure consumer group exists (idempotent — no-op after first call)
+      2. XREADGROUP ">" — fetch new messages, pass to handler, ACK returned IDs
+      3. XAUTOCLAIM — reclaim orphaned messages from crashed consumers, re-process
+      4. Poison-pill sweep — messages exceeding max_deliveries go to DLQ via RetryHandler
+
+    handler(messages) must return the list of IDs it successfully processed.
+    Messages whose IDs are NOT returned stay in PEL for reclaim recovery.
+
+    Safe for multiple pods: each pod gets a unique worker_name; consumer groups
+    distribute messages and XAUTOCLAIM provides crash recovery across pods.
+
+    Call metrics() for this instance or all_metrics() for all live instances in this process.
+    """
+
+    _instances: ClassVar[weakref.WeakSet[StreamConsumer]] = weakref.WeakSet()
+
+    def __init__(
+        self,
+        client: StreamClient,
+        stream: str,
+        config: ConsumerConfig,
+        handler: MessageHandler,
+        dlq_handler: DLQHandler | None = None,
+        serializer: Serializer | None = None,
+        error_hook: ErrorHook | None = None,
+    ) -> None:
+        self._client = client
+        self._stream = stream
+        self._config = config
+        self._handler = handler
+        self._serializer = serializer or JsonSerializer()
+        self._error_hook = error_hook
+        self._worker_name = config.resolved_worker_name()
+
+        self._cg = ConsumerGroup(client, stream, config.group)
+        self._retry = RetryHandler(
+            client=client,
+            stream=stream,
+            group=config.group,
+            dlq_handler=dlq_handler,
+            max_deliveries=config.max_deliveries,
+            batch_size=config.batch_size,
+            serializer=self._serializer,
+        )
+        self._stop = False
+        StreamConsumer._instances.add(self)
+
+        # Throughput counters
+        self._total_read: int = 0
+        self._total_acked: int = 0
+        self._total_dlq: int = 0
+        self._total_errors: int = 0
+        self._started_at: float | None = None
+        self._tps_in = _TpsTracker()
+        self._tps_out = _TpsTracker()
+
+    @classmethod
+    def all_metrics(cls) -> list[ConsumerMetrics]:
+        """Return metrics for all live StreamConsumer instances in this process."""
+        return [inst.metrics() for inst in cls._instances]
+
+    def stop(self) -> None:
+        """Signal the consumer to exit after the current iteration."""
+        self._stop = True
+
+    def metrics(self) -> ConsumerMetrics:
+        """Return current throughput counters and TPS (non-blocking, no Redis calls)."""
+        uptime = (time.monotonic() - self._started_at) if self._started_at else 0.0
+        avg_tps = self._total_acked / uptime if uptime >= 1.0 else 0.0
+        tps_in = round(self._tps_in.tps(), 2)
+        tps_out = round(self._tps_out.tps(), 2)
+        return ConsumerMetrics(
+            total_read=self._total_read,
+            total_acked=self._total_acked,
+            total_dlq=self._total_dlq,
+            total_errors=self._total_errors,
+            tps_in=tps_in,
+            tps_out=tps_out,
+            tps_total=round(tps_in + tps_out, 2),
+            avg_tps=round(avg_tps, 2),
+            uptime_secs=round(uptime, 1),
+        )
+
+    async def run(self) -> None:
+        """Infinite loop. Catches all exceptions and restarts. Exits on stop() or CancelledError."""
+        while not self._stop:
+            try:
+                await self.run_once()
+            except asyncio.CancelledError:
+                raise
+            except Exception as e:
+                if _NOGROUP in str(e):
+                    # Stream or group was externally deleted — reset registry so ensure() re-creates
+                    self._cg.reset()
+                    logger.warning("Consumer group lost (NOGROUP) — will re-create on next iteration")
+                else:
+                    logger.exception("Consumer loop error: %s", e)
+                await asyncio.sleep(1)
+
+    async def run_once(self) -> None:
+        """Execute one full consumer iteration."""
+        cfg = self._config
+
+        await self._cg.ensure(dlq_stream=cfg.dlq_stream, dlq_group=cfg.dlq_group)
+
+        # 1. New messages
+        raw_messages = await self._client.read(
+            stream=self._stream,
+            group=cfg.group,
+            consumer=self._worker_name,
+            count=cfg.batch_size,
+            block_ms=cfg.block_ms,
+        )
+        if raw_messages:
+            if self._started_at is None:
+                self._started_at = time.monotonic()
+            self._total_read += len(raw_messages)
+            self._tps_in.record(len(raw_messages))
+            await self._process_batch(raw_messages)
+
+        # 2. Orphan recovery — reclaim from crashed consumers
+        claimed = await self._client.autoclaim(
+            stream=self._stream,
+            group=cfg.group,
+            consumer=self._worker_name,
+            min_idle_ms=cfg.min_idle_claim_ms,
+            count=cfg.batch_size,
+            max_passes=cfg.max_claim_passes,
+        )
+        if claimed:
+            if self._started_at is None:
+                self._started_at = time.monotonic()
+            self._total_read += len(claimed)
+            self._tps_in.record(len(claimed))
+            await self._process_batch(claimed)
+
+        # 3. Poison-pill sweep
+        dlq_count = await self._retry.handle_poison_pills()
+        self._total_dlq += dlq_count
+
+    # ── Internal ──────────────────────────────────────────────────────────────
+
+    async def _process_batch(self, messages: list[StreamMessage]) -> None:
+        cfg = self._config
+        decoded: list[StreamMessage] = []
+        bad_ids: list[str] = []
+
+        for msg in messages:
+            try:
+                decoded.append(self._decode(msg))
+            except Exception as e:
+                logger.warning("Decode error for msg %s: %s", msg.id, e)
+                bad_ids.append(msg.id)
+                await self._retry.send_to_dlq(msg, "decode_error")
+
+        if bad_ids:
+            self._total_dlq += len(bad_ids)
+            await self._client.ack(self._stream, cfg.group, *bad_ids)
+
+        if not decoded:
+            return
+
+        try:
+            acked_ids = await self._handler(decoded)
+        except asyncio.CancelledError:
+            raise
+        except Exception as e:
+            logger.exception("Handler error: %s", e)
+            self._total_errors += 1
+            if self._error_hook:
+                await self._error_hook(e, decoded)
+            return
+
+        if acked_ids is None:
+            logger.warning(
+                "Handler returned None instead of a list — no messages ACKed. "
+                "Return [] to explicitly ACK nothing, or [msg.id, ...] to ACK."
+            )
+            return
+
+        if acked_ids:
+            await self._client.ack(self._stream, cfg.group, *acked_ids)
+            self._total_acked += len(acked_ids)
+            self._tps_out.record(len(acked_ids))
+
+    def _decode(self, msg: StreamMessage) -> StreamMessage:
+        raw = msg.data
+        if isinstance(raw, str):
+            raw = raw.encode()
+        decoded_data = self._serializer.decode(raw)
+        return StreamMessage(id=msg.id, data=decoded_data, delivery_count=msg.delivery_count)

redis_stream_queue/exceptions.py

@@ -0,0 +1,14 @@
+class RedisStreamError(Exception):
+    pass
+
+
+class DLQError(RedisStreamError):
+    pass
+
+
+class SerializationError(RedisStreamError):
+    pass
+
+
+class ClusterSlotError(RedisStreamError):
+    pass