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

redis_stream_queue/group.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import logging
+import random
+import socket
+from dataclasses import dataclass
+from typing import ClassVar
+
+from .client import StreamClient
+from .message import ConsumerInfo, PendingEntry, StreamStats
+
+logger = logging.getLogger(__name__)
+
+
+def _default_worker_name(group: str) -> str:
+    hostname = socket.gethostname()
+    suffix = f"{random.randint(0, 9999):04d}"
+    # Truncate group name, not hostname/suffix, to keep uniqueness guarantees
+    tail = f"_{hostname}_{suffix}"
+    max_group = 64 - len(tail)
+    return f"{group[:max_group]}{tail}"
+
+
+@dataclass
+class ConsumerConfig:
+    """Configuration for a stream consumer."""
+    group: str
+    worker_name: str | None = None
+    dlq_stream: str | None = None
+    dlq_group: str | None = None
+    batch_size: int = 100
+    block_ms: int = 5_000
+    min_idle_claim_ms: int = 10_000
+    max_deliveries: int = 3
+    max_stream_size: int = 100_000
+    max_claim_passes: int | None = None  # None = sweep full PEL per iteration
+
+    def resolved_worker_name(self) -> str:
+        return self.worker_name or _default_worker_name(self.group)
+
+
+class ConsumerGroup:
+    """
+    Manages and inspects a Redis Stream consumer group.
+
+    Provides group creation (idempotent), stats, health checks,
+    and pending entry inspection.
+
+    Group existence is tracked in a class-level registry keyed by
+    (pool_key, stream, group), so multiple instances targeting the same
+    group share the cache. Call reset() to force re-creation (e.g. after
+    external FLUSHALL or XGROUP DESTROY).
+    """
+
+    _known: ClassVar[set[tuple]] = set()
+
+    def __init__(self, client: StreamClient, stream: str, group: str) -> None:
+        self._client = client
+        self._stream = stream
+        self._group = group
+
+    def _key(self) -> tuple:
+        return (self._client._pool_key(), self._stream, self._group)
+
+    async def ensure(self, dlq_stream: str | None = None, dlq_group: str | None = None) -> None:
+        """Create consumer group (and DLQ group) if not already known. No-op otherwise."""
+        key = self._key()
+        if key in ConsumerGroup._known:
+            return
+        await self._client.create_group(self._stream, self._group)
+        if dlq_stream and dlq_group:
+            await self._client.create_group(dlq_stream, dlq_group)
+        ConsumerGroup._known.add(key)
+
+    def reset(self) -> None:
+        """Remove this group from the known registry — next ensure() will re-create it."""
+        ConsumerGroup._known.discard(self._key())
+
+    async def stats(self, dlq_stream: str | None = None) -> StreamStats:
+        """Return stream length, lag, PEL size, and per-consumer info."""
+        stream_length = await self._client.stream_len(self._stream)
+        groups = await self._client.group_info(self._stream)
+        consumers_raw = await self._client.consumer_info(self._stream, self._group)
+
+        group_data = next(
+            (g for g in groups if _decode(g.get("name")) == self._group), None
+        )
+        pel_size = int(group_data.get("pending", 0)) if group_data else 0
+        lag = int(group_data.get("lag") or 0) if group_data else 0
+        last_delivered_id = _decode(group_data.get("last-delivered-id", b"0-0")) if group_data else "0-0"
+
+        consumers = [
+            ConsumerInfo(
+                name=_decode(c.get("name", b"")),
+                pending=int(c.get("pending", 0)),
+                idle_ms=int(c.get("idle", 0)),
+            )
+            for c in consumers_raw
+        ]
+
+        dlq_length: int | None = None
+        if dlq_stream:
+            try:
+                dlq_length = await self._client.stream_len(dlq_stream)
+            except Exception:
+                pass
+
+        return StreamStats(
+            stream=self._stream,
+            group=self._group,
+            stream_length=stream_length,
+            group_pel_size=pel_size,
+            lag=lag,
+            last_delivered_id=last_delivered_id,
+            consumers=consumers,
+            dlq_length=dlq_length,
+        )
+
+    async def health_check(
+        self,
+        max_lag: int = 10_000,
+        max_idle_ms: int = 60_000,
+        dlq_stream: str | None = None,
+    ) -> dict:
+        """Return health status dict with issue descriptions."""
+        s = await self.stats(dlq_stream=dlq_stream)
+        issues: list[str] = []
+        if s.lag > max_lag:
+            issues.append(f"lag={s.lag} exceeds max_lag={max_lag}")
+        for c in s.consumers:
+            if c.idle_ms > max_idle_ms:
+                issues.append(f"consumer '{c.name}' idle for {c.idle_ms}ms (max {max_idle_ms}ms)")
+        return {"healthy": not issues, "issues": issues, "stats": s}
+
+    async def pending_details(self, count: int = 100) -> list[PendingEntry]:
+        """Return details for all pending (unacknowledged) entries."""
+        return await self._client.pending_range(self._stream, self._group, count)
+
+
+def _decode(value) -> str:
+    if isinstance(value, bytes):
+        return value.decode()
+    return str(value) if value is not None else ""

redis_stream_queue/message.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import time
+from collections import deque
+from dataclasses import dataclass, field
+
+_TPS_WINDOW = 60.0  # seconds
+
+
+class _TpsTracker:
+    """Sliding-window TPS counter. Stores (timestamp, batch_count) pairs."""
+
+    def __init__(self, window_secs: float = _TPS_WINDOW) -> None:
+        self._window = window_secs
+        self._batches: deque[tuple[float, int]] = deque()
+
+    def record(self, count: int) -> None:
+        if count > 0:
+            self._batches.append((time.monotonic(), count))
+
+    def tps(self) -> float:
+        now = time.monotonic()
+        cutoff = now - self._window
+        while self._batches and self._batches[0][0] < cutoff:
+            self._batches.popleft()
+        if not self._batches:
+            return 0.0
+        total = sum(c for _, c in self._batches)
+        span = now - self._batches[0][0]
+        # span < 1s → assume 1s to avoid artificially inflated rate on startup
+        return total / max(span, 1.0)
+
+
+@dataclass
+class StreamMessage:
+    id: str
+    data: dict
+    delivery_count: int = 0
+
+
+@dataclass
+class PendingEntry:
+    id: str
+    consumer: str
+    idle_ms: int
+    delivery_count: int
+
+
+@dataclass
+class ConsumerInfo:
+    name: str
+    pending: int
+    idle_ms: int
+
+
+@dataclass
+class StreamStats:
+    stream: str
+    group: str
+    stream_length: int
+    group_pel_size: int
+    lag: int
+    last_delivered_id: str
+    consumers: list[ConsumerInfo] = field(default_factory=list)
+    dlq_length: int | None = None
+
+
+@dataclass
+class ConsumerMetrics:
+    """Runtime throughput and error counters for a StreamConsumer."""
+    total_read: int        # messages pulled from stream (incl. reclaimed)
+    total_acked: int       # successfully processed + ACKed by handler
+    total_dlq: int         # routed to DLQ (decode_error + max_deliveries)
+    total_errors: int      # handler exceptions (message stays in PEL)
+    tps_in: float          # reads/sec — XREADGROUP + autoclaim, sliding 60s window
+    tps_out: float         # acked/sec — sliding 60s window
+    tps_total: float       # tps_in + tps_out
+    avg_tps: float         # total_acked / uptime_secs since first message
+    uptime_secs: float     # seconds since first message was processed
+
+
+@dataclass
+class ProducerMetrics:
+    """Runtime throughput counters for a StreamProducer."""
+    total_pushed: int      # messages pushed since instance creation
+    tps: float             # pushed/sec — sliding 60s window
+    avg_tps: float         # total_pushed / uptime_secs since first push
+    uptime_secs: float     # seconds since first push

redis_stream_queue/producer.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import time
+import weakref
+from typing import ClassVar
+
+from .client import StreamClient
+from .message import ProducerMetrics, _TpsTracker
+from .serializers import JsonSerializer, Serializer
+
+
+class StreamProducer:
+    """
+    Pushes messages onto a Redis Stream.
+
+    Safe to use from multiple pods simultaneously — Redis XADD is atomic.
+    Call metrics() for this instance or all_metrics() for all live instances in this process.
+    """
+
+    _instances: ClassVar[weakref.WeakSet[StreamProducer]] = weakref.WeakSet()
+
+    def __init__(
+        self,
+        client: StreamClient,
+        stream: str,
+        group: str | None = None,
+        max_len: int = 100_000,
+        serializer: Serializer | None = None,
+    ) -> None:
+        self._client = client
+        self._stream = stream
+        self._group = group
+        self._max_len = max_len
+        self._serializer = serializer or JsonSerializer()
+
+        self._total_pushed: int = 0
+        self._started_at: float | None = None
+        self._tps = _TpsTracker()
+        StreamProducer._instances.add(self)
+
+    async def push(self, data: dict) -> str:
+        """Serialize and publish one message. Returns the Redis stream entry ID."""
+        encoded = self._serializer.encode(data)
+        msg_id = await self._client.push(self._stream, encoded, max_len=self._max_len)
+        if self._started_at is None:
+            self._started_at = time.monotonic()
+        self._total_pushed += 1
+        self._tps.record(1)
+        return msg_id
+
+    async def push_many(self, data: list[dict]) -> list[str]:
+        """Serialize and publish multiple messages via pipeline (one round-trip). Returns list of entry IDs."""
+        if not data:
+            return []
+        encoded = [self._serializer.encode(d) for d in data]
+        ids = await self._client.push_many(self._stream, encoded, max_len=self._max_len)
+        if self._started_at is None:
+            self._started_at = time.monotonic()
+        self._total_pushed += len(data)
+        self._tps.record(len(data))
+        return ids
+
+    async def ensure_group(self, group: str | None = None) -> None:
+        """Idempotently create the consumer group for this stream."""
+        g = group or self._group
+        if not g:
+            raise ValueError("group name is required")
+        await self._client.create_group(self._stream, g)
+
+    @classmethod
+    def all_metrics(cls) -> list[ProducerMetrics]:
+        """Return metrics for all live StreamProducer instances in this process."""
+        return [inst.metrics() for inst in cls._instances]
+
+    def metrics(self) -> ProducerMetrics:
+        """Return current push 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_pushed / uptime if uptime >= 1.0 else 0.0
+        return ProducerMetrics(
+            total_pushed=self._total_pushed,
+            tps=round(self._tps.tps(), 2),
+            avg_tps=round(avg_tps, 2),
+            uptime_secs=round(uptime, 1),
+        )

redis_stream_queue/retry.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import logging
+from typing import Awaitable, Callable
+
+from .client import StreamClient
+from .message import StreamMessage
+from .serializers import JsonSerializer, Serializer
+
+logger = logging.getLogger(__name__)
+
+DLQHandler = Callable[[StreamMessage, str], Awaitable[None]]
+
+
+class RetryHandler:
+    """
+    Detects poison pills (messages exceeding max_deliveries) and routes them to DLQ.
+
+    Integrates with StreamConsumer to offload all retry/DLQ concerns from the main loop.
+    """
+
+    def __init__(
+        self,
+        client: StreamClient,
+        stream: str,
+        group: str,
+        dlq_handler: DLQHandler | None = None,
+        max_deliveries: int = 3,
+        batch_size: int = 100,
+        serializer: Serializer | None = None,
+    ) -> None:
+        self._client = client
+        self._stream = stream
+        self._group = group
+        self._dlq_handler = dlq_handler
+        self._max_deliveries = max_deliveries
+        self._batch_size = batch_size
+        self._serializer = serializer or JsonSerializer()
+
+    async def handle_poison_pills(self) -> int:
+        """
+        Find messages in PEL that exceeded max_deliveries, route to DLQ, then ACK.
+        Returns count of poison pills processed. Called once per consumer loop iteration.
+        """
+        pending = await self._client.pending_range(
+            self._stream, self._group, self._batch_size
+        )
+        poison_ids = [e.id for e in pending if e.delivery_count >= self._max_deliveries]
+        if not poison_ids:
+            return 0
+
+        if not self._dlq_handler:
+            logger.warning(
+                "Poison pills detected (%d msgs) but no dlq_handler set — "
+                "messages will be ACKed and lost: %s",
+                len(poison_ids), poison_ids[:5],
+            )
+
+        fetched = await self._client.fetch_by_ids(self._stream, poison_ids)
+        fetched_by_id = {m.id: m for m in fetched}
+
+        for msg_id in poison_ids:
+            raw_msg = fetched_by_id.get(msg_id)
+            if raw_msg is None:
+                logger.warning("Poison pill msg %s missing from stream (deleted?) — ACKing anyway", msg_id)
+            else:
+                decoded = self._try_decode(raw_msg)
+                await self.send_to_dlq(decoded, "max_deliveries")
+
+        # Batch ACK after all DLQ attempts — send_to_dlq swallows errors
+        await self._client.ack(self._stream, self._group, *poison_ids)
+        return len(poison_ids)
+
+    async def send_to_dlq(self, msg: StreamMessage, reason: str) -> None:
+        """Route a single message to the DLQ handler. Swallows handler errors."""
+        if self._dlq_handler:
+            try:
+                await self._dlq_handler(msg, reason)
+            except Exception as e:
+                logger.error("DLQ handler error for msg %s: %s", msg.id, e)
+
+    def _try_decode(self, msg: StreamMessage) -> StreamMessage:
+        raw = msg.data
+        if isinstance(raw, (bytes, str)):
+            try:
+                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)
+            except Exception:
+                pass
+        return msg

redis_stream_queue/serializers.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import json
+import pickle
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class Serializer(Protocol):
+    def encode(self, data: dict) -> bytes: ...
+    def decode(self, raw: bytes) -> dict: ...
+
+
+class JsonSerializer:
+    def encode(self, data: dict) -> bytes:
+        return json.dumps(data, ensure_ascii=False).encode("utf-8")
+
+    def decode(self, raw: bytes) -> dict:
+        return json.loads(raw)
+
+
+class MsgpackSerializer:
+    def __init__(self) -> None:
+        try:
+            import msgpack  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                "msgpack is required: pip install redis-stream-queue[msgpack]"
+            )
+
+    def encode(self, data: dict) -> bytes:
+        import msgpack
+        return msgpack.packb(data, use_bin_type=True)
+
+    def decode(self, raw: bytes) -> dict:
+        import msgpack
+        return msgpack.unpackb(raw, raw=False)
+
+
+class PickleSerializer:
+    def encode(self, data: dict) -> bytes:
+        return pickle.dumps(data)
+
+    def decode(self, raw: bytes) -> dict:
+        return pickle.loads(raw)  # noqa: S301