rabbitkit 0.9.0__py3-none-any.whl

rabbitkit/middleware/timeout.py

@@ -0,0 +1,291 @@
+"""Handler timeout middleware.
+
+Enforces a maximum processing time per message.  If the handler takes longer
+than ``TimeoutConfig.timeout_seconds``, a ``HandlerTimeoutError`` is raised.
+
+Implementation strategy
+-----------------------
+* **Async** — ``asyncio.wait_for()`` cancels the coroutine cleanly.
+* **Sync**  — handler is run in a ``daemon=True`` ``threading.Thread``; the
+  calling thread waits ``timeout_seconds`` and raises if the thread is still
+  alive.  The background thread continues to its natural end but its result is
+  discarded.  (CPython has no safe way to kill a thread, so long-running IO
+  handlers will not be interrupted — they just become detached.)
+
+``HandlerTimeoutError`` is a subclass of ``TimeoutError``.  The default error
+classifier treats it as **TRANSIENT**, so retry middleware (if configured) will
+re-queue the message.  Override the classifier if you want timeouts to go
+straight to the DLQ.
+
+Quick start
+-----------
+    from rabbitkit.middleware.timeout import TimeoutMiddleware, TimeoutConfig
+
+    # 10-second hard limit on all handlers for this route
+    timeout_mw = TimeoutMiddleware(TimeoutConfig(timeout_seconds=10.0))
+
+    @broker.subscriber(queue="slow-tasks", middlewares=[timeout_mw])
+    async def handle_task(body: bytes) -> None:
+        await some_slow_operation()
+
+Default timeout (30 s)::
+
+    timeout_mw = TimeoutMiddleware()   # uses TimeoutConfig(timeout_seconds=30.0)
+
+Combining with retry::
+
+    from rabbitkit import RetryConfig
+    from rabbitkit.middleware.retry import RetryMiddleware
+
+    retry_mw = RetryMiddleware(RetryConfig(max_retries=3, delays=(5, 30, 120)))
+    timeout_mw = TimeoutMiddleware(TimeoutConfig(timeout_seconds=15.0))
+
+    # Order: timeout wraps handler, retry wraps timeout
+    @broker.subscriber(
+        queue="jobs",
+        middlewares=[retry_mw, timeout_mw],   # retry outermost
+    )
+    async def run_job(body: bytes) -> None: ...
+
+Exception hierarchy
+-------------------
+``HandlerTimeoutError(TimeoutError)``  — carries ``.timeout_seconds`` attribute
+for logging / classification.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+from rabbitkit.core.message import RabbitMessage
+from rabbitkit.middleware.base import BaseMiddleware
+
+logger = logging.getLogger(__name__)
+
+
+class _DiscardedSettlement(BaseException):
+    """H9 internal sentinel — raised (never caught by user code, since it's
+    not an ``Exception``) from a guarded settlement stand-in to abort a
+    discarded ack()/nack()/reject() call from an abandoned timed-out handler
+    thread BEFORE ``RabbitMessage.ack()``/``nack()``/``reject()`` sets
+    ``_disposition``. Without this, the discarded call would still flip
+    ``_disposition`` away from "pending" (their guard for the real fn calls
+    it AFTER the settlement fn returns, unconditionally) — silently
+    preventing the consumer thread's own later, legitimate settlement from
+    doing anything at all, since ``RabbitMessage`` would think the message
+    was already settled. Landing back in ``_run()``'s ``except
+    BaseException`` — harmless, since the background thread's outcome is
+    already discarded by that point (see ``TimeoutMiddleware.consume_scope``)."""
+
+
+class HandlerTimeoutError(TimeoutError):
+    """Raised when a handler exceeds the configured timeout."""
+
+    def __init__(self, timeout_seconds: float) -> None:
+        super().__init__(f"Handler exceeded timeout of {timeout_seconds}s")
+        self.timeout_seconds = timeout_seconds
+
+
+@dataclass(frozen=True, slots=True)
+class TimeoutConfig:
+    """Configuration for handler timeout.
+
+    Attributes:
+        timeout_seconds: Maximum handler execution time in seconds.
+    """
+
+    timeout_seconds: float = 30.0
+
+    def __post_init__(self) -> None:
+        if self.timeout_seconds <= 0:
+            raise ValueError("timeout_seconds must be positive")
+
+
+class TimeoutMiddleware(BaseMiddleware):
+    """Enforces a maximum processing time per message.
+
+    Async: uses ``asyncio.wait_for()`` — clean cancellation.
+    Sync: runs ``call_next`` in a daemon ``threading.Thread`` and raises
+    ``HandlerTimeoutError`` if the thread is still alive at the deadline.
+
+    .. warning::
+        CPython cannot safely kill a thread, so a sync handler that exceeds the
+        timeout is **abandoned** — it keeps running detached until it finishes
+        naturally. This can leak threads and resources for long-running / blocked
+        IO handlers. Use an **async** handler for true cooperative cancellation.
+        When a sync timeout fires, a CRITICAL log record is emitted and the
+        optional ``on_timeout`` callback (if configured) is invoked so the
+        abandonment is observable (e.g. increment a metric counter).
+    """
+
+    def __init__(
+        self,
+        config: TimeoutConfig | None = None,
+        *,
+        on_timeout: Callable[[RabbitMessage, float], None] | None = None,
+    ) -> None:
+        self._config = config or TimeoutConfig()
+        self._on_timeout = on_timeout
+        # Observable counter of sync threads abandoned due to timeout.
+        self.abandoned_threads: int = 0
+
+    def consume_scope(
+        self,
+        call_next: Any,
+        message: RabbitMessage,
+    ) -> Any:
+        """Sync timeout — runs call_next in a thread with timeout.
+
+        If the deadline elapses with the thread still alive, the thread is
+        abandoned (CPython cannot kill it), a CRITICAL log is emitted, the
+        ``abandoned_threads`` counter is incremented, ``on_timeout`` (if any)
+        is invoked, and ``HandlerTimeoutError`` is raised.
+
+        H9 — settlement is exclusively from the consumer (this) thread, never
+        the background thread, even under ``AckPolicy.MANUAL`` (where the
+        handler itself calls ``message.ack()``/``nack()``/``reject()``):
+        while the background thread is running, ``message``'s settlement
+        functions are swapped for stand-ins that CAPTURE (rather than
+        execute) any settlement attempt made from that specific thread —
+        calling the real pika-backed function from there, while THIS thread
+        is blocked in ``thread.join()`` (i.e. not pumping the connection's
+        I/O loop), can deadlock: the settlement call would marshal onto this
+        thread and wait for it to drain the callback, but this thread is
+        itself waiting on the background thread to finish. If the background
+        thread finishes within the deadline, any settlement it captured is
+        replayed for real on this (consumer/owner) thread, safely, after
+        ``join()`` returns. If it does not, the guards stay installed
+        (deliberately NOT restored — the background thread may still call
+        ack()/nack()/reject() at any point later) but their thread-identity
+        check means any FUTURE call specifically from that thread is
+        discarded, while THIS thread's own subsequent settlement (e.g. via
+        AckPolicy/RetryMiddleware after ``HandlerTimeoutError`` propagates
+        below) is routed straight to the real fn — safe, since it's a
+        same-thread call.
+        """
+        result_holder: list[Any] = []
+        exception_holder: list[BaseException] = []
+        captured_settlement: list[Callable[[], None]] = []
+        timed_out = threading.Event()
+
+        real_ack, real_nack, real_reject = message._ack_fn, message._nack_fn, message._reject_fn
+
+        def _guarded_ack() -> None:
+            if threading.get_ident() != thread.ident:
+                if real_ack is not None:
+                    real_ack()
+                return
+            if timed_out.is_set():
+                logger.warning("Discarding ack() from an abandoned timed-out handler thread")
+                # Raise (rather than return) so RabbitMessage.ack() does not
+                # proceed to set _disposition="acked" for a call that never
+                # actually touched the channel — that would silently block
+                # the consumer thread's own later, real settlement.
+                raise _DiscardedSettlement
+            if real_ack is not None:
+                captured_settlement.append(real_ack)
+
+        def _guarded_nack(requeue: bool = True) -> None:
+            if threading.get_ident() != thread.ident:
+                if real_nack is not None:
+                    real_nack(requeue)
+                return
+            if timed_out.is_set():
+                logger.warning("Discarding nack() from an abandoned timed-out handler thread")
+                raise _DiscardedSettlement
+            if real_nack is not None:
+                captured_settlement.append(lambda: real_nack(requeue))
+
+        def _guarded_reject(requeue: bool = False) -> None:
+            if threading.get_ident() != thread.ident:
+                if real_reject is not None:
+                    real_reject(requeue)
+                return
+            if timed_out.is_set():
+                logger.warning("Discarding reject() from an abandoned timed-out handler thread")
+                raise _DiscardedSettlement
+            if real_reject is not None:
+                captured_settlement.append(lambda: real_reject(requeue))
+
+        # Only install a guard where a real fn exists — leaving an already-None
+        # fn as None preserves message.ack()/nack()/reject()'s own "no
+        # settlement fn set" RuntimeError for e.g. no-ack deliveries, instead
+        # of silently swallowing it inside the guard.
+        if real_ack is not None:
+            message._ack_fn = _guarded_ack
+        if real_nack is not None:
+            message._nack_fn = _guarded_nack
+        if real_reject is not None:
+            message._reject_fn = _guarded_reject
+
+        def _run() -> None:
+            try:
+                result_holder.append(call_next(message))
+            except BaseException as exc:
+                exception_holder.append(exc)
+
+        thread = threading.Thread(target=_run, daemon=True)
+        thread.start()
+        thread.join(timeout=self._config.timeout_seconds)
+
+        if thread.is_alive():
+            # CPython cannot safely kill a thread — the handler keeps running
+            # detached. Make the abandonment explicit and observable.
+            #
+            # Deliberately do NOT restore the real settlement fns here: the
+            # background thread is still running and may call
+            # ack()/nack()/reject() at any point after this method returns —
+            # if the real fn were restored, that eventual call would hit the
+            # pika channel directly from a non-owner thread. The guards stay
+            # installed; their threading.get_ident() check already routes
+            # THIS thread's own subsequent settlement (AckPolicy/
+            # RetryMiddleware, after the exception below propagates) straight
+            # to the real fn (safe, same-thread), while any FUTURE call
+            # specifically from the background thread is discarded now that
+            # timed_out is set.
+            timed_out.set()
+            self.abandoned_threads += 1
+            logger.critical(
+                "Sync handler exceeded %.1fs timeout; thread abandoned (still running). "
+                "Use an async handler for real cancellation. "
+                "abandoned_threads=%d",
+                self._config.timeout_seconds,
+                self.abandoned_threads,
+            )
+            if self._on_timeout is not None:
+                try:
+                    self._on_timeout(message, self._config.timeout_seconds)
+                except Exception:  # pragma: no cover - callback must not break flow
+                    logger.exception("on_timeout callback raised")
+            raise HandlerTimeoutError(self._config.timeout_seconds)
+
+        # Finished within the deadline — restore the real fns, then replay
+        # any settlement the background thread captured, for real, on this
+        # (consumer/owner) thread.
+        message._ack_fn, message._nack_fn, message._reject_fn = real_ack, real_nack, real_reject
+        for replay in captured_settlement:
+            replay()
+
+        if exception_holder:
+            raise exception_holder[0]
+
+        return result_holder[0] if result_holder else None
+
+    async def consume_scope_async(
+        self,
+        call_next: Any,
+        message: RabbitMessage,
+    ) -> Any:
+        """Async timeout — uses asyncio.wait_for()."""
+        try:
+            return await asyncio.wait_for(
+                call_next(message),
+                timeout=self._config.timeout_seconds,
+            )
+        except TimeoutError:
+            raise HandlerTimeoutError(self._config.timeout_seconds) from None

rabbitkit/queue_metrics.py

@@ -0,0 +1,174 @@
+"""QueueMetricsPoller — bridge the RabbitMQ management API into metrics (H5).
+
+The consume/publish counters (``MetricsMiddleware``) can only see messages
+*this process* handles — they go on reading healthy while a queue silently
+accumulates millions of messages because a consumer fell behind or died.
+Queue depth / consumer lag / DLQ growth is the #1 RabbitMQ incident signal,
+and it lives on the broker, not in-process.
+
+``QueueMetricsPoller`` periodically calls a management client's
+``list_queues()`` and emits gauges (labeled by queue) through the same
+``MetricsCollector`` the rest of rabbitkit uses:
+
+- ``{ns}_queue_messages_ready``   — backlog depth
+- ``{ns}_queue_messages_unacked`` — delivered-but-unacked
+- ``{ns}_queue_messages_total``   — ready + unacked
+- ``{ns}_queue_consumers``        — 0 means nothing is draining
+
+Usage (sync)::
+
+    from rabbitkit import RabbitManagementClient, PrometheusCollector, QueueMetricsPoller
+
+    poller = QueueMetricsPoller(
+        management_client=RabbitManagementClient(...),
+        collector=PrometheusCollector(),
+        interval=15.0,
+    )
+    poller.start()   # background daemon thread
+    ...
+    poller.stop()
+
+Async brokers use ``QueueMetricsPoller.start_async()`` with a management
+client exposing ``list_queues_async()``.
+
+Alert on ``queue_messages_ready`` growth and ``queue_consumers == 0`` — those
+are the signals rabbitkit's own metrics cannot provide.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+from collections.abc import Callable
+from typing import Any
+
+from rabbitkit.core.config import MetricsConfig
+
+logger = logging.getLogger(__name__)
+
+
+class QueueMetricsPoller:
+    """Polls the management API and emits queue-depth gauges.
+
+    Args:
+        management_client: object with ``list_queues(vhost)`` (sync) and/or
+            ``list_queues_async(vhost)`` (async) returning a list of dicts
+            with ``name``/``messages``/``messages_ready``/
+            ``messages_unacknowledged``/``consumers`` keys (rabbitkit's
+            ``RabbitManagementClient`` satisfies this).
+        collector: any ``MetricsCollector`` (needs ``set_gauge``).
+        config: metric-naming config (defaults to ``MetricsConfig()``).
+        vhost: vhost to poll (default ``"/"``).
+        interval: seconds between polls in the background loop.
+        queue_filter: optional predicate ``(queue_name) -> bool`` — only
+            matching queues emit gauges (e.g. skip rabbitkit's own delay
+            queues, or restrict to a service's queues to bound cardinality).
+    """
+
+    def __init__(
+        self,
+        management_client: Any,
+        collector: Any,
+        config: MetricsConfig | None = None,
+        *,
+        vhost: str = "/",
+        interval: float = 15.0,
+        queue_filter: Callable[[str], bool] | None = None,
+    ) -> None:
+        self._mgmt = management_client
+        self._collector = collector
+        self._cfg = config or MetricsConfig()
+        self._vhost = vhost
+        self._interval = interval
+        self._queue_filter = queue_filter
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    # ── Core: one poll → gauges ───────────────────────────────────────────
+
+    def _emit(self, queues: list[dict[str, Any]]) -> int:
+        """Set gauges for each (filtered) queue. Returns count emitted."""
+        emitted = 0
+        for q in queues:
+            name = q.get("name")
+            if not name:
+                continue
+            if self._queue_filter is not None and not self._queue_filter(name):
+                continue
+            labels = {"queue": name}
+            # The management API omits counts for a queue mid-declaration or
+            # in some states; default missing values to 0 rather than skip so
+            # a gauge reset (queue drained) is still visible.
+            ready = int(q.get("messages_ready", 0) or 0)
+            unacked = int(q.get("messages_unacknowledged", 0) or 0)
+            total = int(q.get("messages", ready + unacked) or 0)
+            consumers = int(q.get("consumers", 0) or 0)
+            self._collector.set_gauge(self._cfg.queue_messages_ready, labels, ready)
+            self._collector.set_gauge(self._cfg.queue_messages_unacked, labels, unacked)
+            self._collector.set_gauge(self._cfg.queue_messages_total, labels, total)
+            self._collector.set_gauge(self._cfg.queue_consumers, labels, consumers)
+            emitted += 1
+        return emitted
+
+    def poll_once(self) -> int:
+        """Fetch queues once (sync) and emit gauges. Returns count emitted.
+
+        Never raises — a management-API error is logged and the poll is
+        skipped, so a transient management-plane outage does not crash the
+        poller thread (the next tick retries).
+        """
+        try:
+            queues = self._mgmt.list_queues(self._vhost)
+        except Exception:
+            logger.warning("QueueMetricsPoller: list_queues failed; skipping this poll", exc_info=True)
+            return 0
+        return self._emit(queues)
+
+    async def poll_once_async(self) -> int:
+        """Async variant of :meth:`poll_once`."""
+        try:
+            queues = await self._mgmt.list_queues_async(self._vhost)
+        except Exception:
+            logger.warning("QueueMetricsPoller: list_queues_async failed; skipping this poll", exc_info=True)
+            return 0
+        return self._emit(queues)
+
+    # ── Background loops ──────────────────────────────────────────────────
+
+    def start(self) -> None:
+        """Start a background daemon thread polling every ``interval`` seconds."""
+        if self._thread is not None:
+            return
+        self._stop.clear()
+        self._thread = threading.Thread(
+            target=self._run_loop, name="rabbitkit-queue-metrics", daemon=True
+        )
+        self._thread.start()
+
+    def _run_loop(self) -> None:
+        # Poll immediately, then every interval. Event.wait doubles as the
+        # sleep and the stop signal (interruptible shutdown).
+        while not self._stop.is_set():
+            self.poll_once()
+            self._stop.wait(self._interval)
+
+    def stop(self, timeout: float = 5.0) -> None:
+        """Signal the background thread to stop and join it."""
+        self._stop.set()
+        if self._thread is not None:
+            self._thread.join(timeout=timeout)
+            self._thread = None
+
+    async def run_async(self) -> None:
+        """Async polling loop — run as a task; cancel to stop.
+
+        Usage::
+
+            task = asyncio.create_task(poller.run_async())
+            ...
+            task.cancel()
+        """
+        while True:
+            await self.poll_once_async()
+            await asyncio.sleep(self._interval)

rabbitkit/results/__init__.py

@@ -0,0 +1,6 @@
+"""Result storage backends for rabbitkit."""
+
+from rabbitkit.results.backend import RedisResultBackend, ResultBackend
+from rabbitkit.results.middleware import ResultMiddleware
+
+__all__ = ["RedisResultBackend", "ResultBackend", "ResultMiddleware"]

rabbitkit/results/backend.py

@@ -0,0 +1,102 @@
+"""Result backend protocol and implementations.
+
+A **result backend** stores handler return values so callers can retrieve
+them later using a ``correlation_id``.  This is the server side of the
+request/result pattern — the broker stores results; clients fetch them.
+
+Protocol
+--------
+``ResultBackend`` is a ``@runtime_checkable`` generic ``Protocol[T]``.  ``T``
+is the type of the stored result, so the type flows from
+``store(correlation_id, result: T)`` into ``fetch(correlation_id) -> T | None``.
+Any object with ``store`` / ``fetch`` (sync) and ``store_async`` / ``fetch_async``
+(async) methods qualifies.
+
+Built-in implementation
+-----------------------
+``RedisResultBackend`` stores results as Redis strings with configurable TTL,
+so it satisfies ``ResultBackend[bytes]``:
+
+    import redis
+    from rabbitkit.results.backend import RedisResultBackend
+
+    r = redis.Redis(host="redis")
+    backend = RedisResultBackend(r, key_prefix="myapp:result:")
+
+Keys are stored as ``{key_prefix}{correlation_id}``.  Default TTL is 3600 s (1 h).
+
+Async variant (redis-py >= 4.2 async client)::
+
+    import redis.asyncio as aioredis
+    r = aioredis.Redis(host="redis")
+    backend = RedisResultBackend(r)
+    await backend.store_async("corr-123", b'{"status": "done"}', ttl=600)
+    result = await backend.fetch_async("corr-123")
+
+Custom backend
+--------------
+Implement the protocol to use any other storage:
+
+    class PostgresResultBackend:
+        def store(self, correlation_id: str, result: bytes, ttl: int = 3600) -> None:
+            db.execute("INSERT INTO results ...", (correlation_id, result))
+
+        def fetch(self, correlation_id: str, timeout: float = 5.0) -> bytes | None:
+            row = db.fetchone("SELECT result FROM results WHERE id=%s", (correlation_id,))
+            return row[0] if row else None
+
+        async def store_async(...): ...
+        async def fetch_async(...): ...
+
+See also
+--------
+``ResultMiddleware`` — middleware that wires this backend to the pipeline so
+handler return values are stored automatically.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Protocol, TypeVar, runtime_checkable
+
+T = TypeVar("T")
+
+
+@runtime_checkable
+class ResultBackend(Protocol[T]):
+    """Protocol for result storage backends.
+
+    Generic in ``T`` (the stored result type) so the type flows from
+    ``store`` into ``fetch``.  ``RedisResultBackend`` stores ``bytes``, so it
+    satisfies ``ResultBackend[bytes]``.
+    """
+
+    def store(self, correlation_id: str, result: T, ttl: int = 3600) -> None: ...
+    def fetch(self, correlation_id: str, timeout: float = 5.0) -> T | None: ...
+    async def store_async(self, correlation_id: str, result: T, ttl: int = 3600) -> None: ...
+    async def fetch_async(self, correlation_id: str, timeout: float = 5.0) -> T | None: ...
+
+
+class RedisResultBackend:
+    """Redis-based result backend using GET/SET with TTL.
+
+    Stores results as raw bytes, so it satisfies ``ResultBackend[bytes]``.
+    """
+
+    def __init__(self, redis_client: Any, key_prefix: str = "rabbitkit:result:") -> None:
+        self._redis = redis_client
+        self._prefix = key_prefix
+
+    def _key(self, correlation_id: str) -> str:
+        return f"{self._prefix}{correlation_id}"
+
+    def store(self, correlation_id: str, result: bytes, ttl: int = 3600) -> None:
+        self._redis.set(self._key(correlation_id), result, ex=ttl)
+
+    def fetch(self, correlation_id: str, timeout: float = 5.0) -> bytes | None:
+        return self._redis.get(self._key(correlation_id))  # type: ignore[no-any-return]
+
+    async def store_async(self, correlation_id: str, result: bytes, ttl: int = 3600) -> None:
+        await self._redis.set(self._key(correlation_id), result, ex=ttl)
+
+    async def fetch_async(self, correlation_id: str, timeout: float = 5.0) -> bytes | None:
+        return await self._redis.get(self._key(correlation_id))  # type: ignore[no-any-return]