rabbitkit 0.9.0__py3-none-any.whl

rabbitkit/middleware/otel.py

@@ -0,0 +1,203 @@
+"""OTelTracingMiddleware — native OpenTelemetry tracing.
+
+Speaks ``opentelemetry-api`` directly so publish→consume trace propagation
+works for anyone, using W3C ``traceparent``/``tracestate`` headers over AMQP.
+Install with the optional extra::
+
+    pip install rabbitkit[otel]
+
+Lazy/no-op: if ``opentelemetry`` is not importable, every span is a
+passthrough — and the middleware warns ONCE at construction (a caller who
+adds it is opting into tracing; silently no-oping forever reads as
+"nothing to trace yet" instead of "tracing was never active").
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Awaitable, Callable
+from dataclasses import replace
+from typing import Any
+
+from rabbitkit.core.message import RabbitMessage
+from rabbitkit.core.types import MessageEnvelope
+from rabbitkit.middleware.base import BaseMiddleware
+
+logger = logging.getLogger(__name__)
+
+
+def _get_otel() -> Any:
+    """Lazy import of the ``opentelemetry`` API — returns module pair or ``None``."""
+    try:
+        from opentelemetry import propagate, trace
+
+        return trace, propagate
+    except ImportError:
+        return None
+
+
+class OTelTracingMiddleware(BaseMiddleware):
+    """Wrap handler execution and publishes in standard OpenTelemetry spans.
+
+    - Consume: extracts W3C trace context from message headers and starts a
+      ``CONSUMER``-kind span as its child. Handler exceptions are recorded on
+      the span (status ``ERROR``) and re-raised.
+    - Publish: starts a ``PRODUCER``-kind span and injects the current trace
+      context into a COPY of the envelope's headers (envelopes are frozen).
+
+    Span names follow the OTel messaging convention: ``{destination} receive``
+    / ``{destination} send``.
+    """
+
+    def __init__(self, service_name: str = "rabbitkit") -> None:
+        self._service_name = service_name
+        otel = _get_otel()
+        if otel is None:
+            self._trace: Any = None
+            self._propagate: Any = None
+            self._tracer: Any = None
+            logger.warning(
+                "OTelTracingMiddleware(service_name=%r) added but opentelemetry is not "
+                "installed -- every consume/publish span will be a silent no-op. "
+                "Install with `pip install rabbitkit[otel]`, or remove this middleware.",
+                service_name,
+            )
+        else:
+            self._trace, self._propagate = otel
+            self._tracer = self._trace.get_tracer("rabbitkit", instrumenting_library_version="1")
+
+    @property
+    def is_available(self) -> bool:
+        """True if the opentelemetry API is importable."""
+        return self._tracer is not None
+
+    # ── Internal helpers ──────────────────────────────────────────────────
+
+    @staticmethod
+    def _str_carrier(headers: dict[str, Any]) -> dict[str, str]:
+        """AMQP headers may hold non-str values; propagators want str→str."""
+        return {k: v for k, v in headers.items() if isinstance(v, str)}
+
+    def _consume_attributes(self, message: RabbitMessage) -> dict[str, str]:
+        attrs: dict[str, str] = {
+            "messaging.system": "rabbitmq",
+            "messaging.operation": "receive",
+        }
+        if message.routing_key:
+            attrs["messaging.rabbitmq.destination.routing_key"] = message.routing_key
+        queue = message.headers.get("x-rabbitkit-original-queue", "")
+        if queue:
+            attrs["messaging.destination.name"] = str(queue)
+        if message.message_id:
+            attrs["messaging.message.id"] = message.message_id
+        if message.correlation_id:
+            attrs["messaging.message.conversation_id"] = message.correlation_id
+        retry_count = message.headers.get("x-rabbitkit-retry-count")
+        if retry_count is not None:
+            attrs["messaging.rabbitmq.retry_count"] = str(retry_count)
+        return attrs
+
+    def _publish_attributes(self, envelope: MessageEnvelope) -> dict[str, str]:
+        attrs: dict[str, str] = {
+            "messaging.system": "rabbitmq",
+            "messaging.operation": "send",
+        }
+        if envelope.routing_key:
+            attrs["messaging.rabbitmq.destination.routing_key"] = envelope.routing_key
+        if envelope.exchange:
+            attrs["messaging.destination.name"] = envelope.exchange
+        if envelope.message_id:
+            attrs["messaging.message.id"] = envelope.message_id
+        if envelope.correlation_id:
+            attrs["messaging.message.conversation_id"] = envelope.correlation_id
+        return attrs
+
+    def _consume_span(self, message: RabbitMessage) -> Any:
+        ctx = self._propagate.extract(self._str_carrier(message.headers))
+        name = f"{message.headers.get('x-rabbitkit-original-queue', message.routing_key) or 'queue'} receive"
+        return self._tracer.start_as_current_span(
+            name,
+            context=ctx,
+            kind=self._trace.SpanKind.CONSUMER,
+            attributes=self._consume_attributes(message),
+        )
+
+    def _publish_span(self, envelope: MessageEnvelope) -> Any:
+        name = f"{envelope.exchange or envelope.routing_key or 'exchange'} send"
+        return self._tracer.start_as_current_span(
+            name,
+            kind=self._trace.SpanKind.PRODUCER,
+            attributes=self._publish_attributes(envelope),
+        )
+
+    def _record_failure(self, span: Any, exc: BaseException) -> None:
+        span.record_exception(exc)
+        span.set_status(self._trace.Status(self._trace.StatusCode.ERROR, str(exc)))
+
+    def _envelope_with_context(self, envelope: MessageEnvelope) -> MessageEnvelope:
+        """Copy of *envelope* with the CURRENT trace context injected."""
+        carrier: dict[str, str] = {}
+        self._propagate.inject(carrier)
+        if not carrier:
+            return envelope
+        return replace(envelope, headers={**envelope.headers, **carrier})
+
+    # ── Consume-side hooks ────────────────────────────────────────────────
+
+    def consume_scope(
+        self,
+        call_next: Callable[[RabbitMessage], Any],
+        message: RabbitMessage,
+    ) -> Any:
+        if self._tracer is None:
+            return call_next(message)
+        with self._consume_span(message) as span:
+            try:
+                return call_next(message)
+            except BaseException as exc:
+                self._record_failure(span, exc)
+                raise
+
+    async def consume_scope_async(
+        self,
+        call_next: Callable[[RabbitMessage], Awaitable[Any]],
+        message: RabbitMessage,
+    ) -> Any:
+        if self._tracer is None:
+            return await call_next(message)
+        with self._consume_span(message) as span:
+            try:
+                return await call_next(message)
+            except BaseException as exc:
+                self._record_failure(span, exc)
+                raise
+
+    # ── Publish-side hooks ────────────────────────────────────────────────
+
+    def publish_scope(
+        self,
+        call_next: Callable[[MessageEnvelope], Any],
+        envelope: MessageEnvelope,
+    ) -> Any:
+        if self._tracer is None:
+            return call_next(envelope)
+        with self._publish_span(envelope) as span:
+            try:
+                return call_next(self._envelope_with_context(envelope))
+            except BaseException as exc:
+                self._record_failure(span, exc)
+                raise
+
+    async def publish_scope_async(
+        self,
+        call_next: Callable[[MessageEnvelope], Awaitable[Any]],
+        envelope: MessageEnvelope,
+    ) -> Any:
+        if self._tracer is None:
+            return await call_next(envelope)
+        with self._publish_span(envelope) as span:
+            try:
+                return await call_next(self._envelope_with_context(envelope))
+            except BaseException as exc:
+                self._record_failure(span, exc)
+                raise

rabbitkit/middleware/rate_limit.py

@@ -0,0 +1,247 @@
+"""Rate-limiting middleware for message consumption.
+
+Limits the rate of message processing using a **token bucket** algorithm.
+Thread-safe for sync consumers; uses ``asyncio.sleep`` for async consumers.
+
+Three behaviours when the rate is exceeded (``on_limited``):
+
+* ``"wait"``  — sleep until a token is available (default, back-pressures the consumer)
+* ``"nack"``  — nack with ``requeue=True`` so another consumer / retry can handle it
+* ``"drop"``  — nack with ``requeue=False`` (message discarded / sent to DLQ)
+
+Quick start
+-----------
+    from rabbitkit.middleware.rate_limit import RateLimitMiddleware, RateLimitConfig
+
+    rate_mw = RateLimitMiddleware(
+        RateLimitConfig(max_rate=100.0, burst=10, on_limited="wait")
+    )
+
+    @broker.subscriber(queue="events", middlewares=[rate_mw])
+    async def handle_event(body: bytes) -> None:
+        ...
+
+Per-consumer vs broker-wide
+----------------------------
+Attach as a per-route middleware to scope the limit to one queue:
+
+    @broker.subscriber(queue="high-volume", middlewares=[rate_mw])
+    def handle(body: bytes) -> None: ...
+
+Or attach broker-wide by passing it in the broker constructor's middleware
+list (if supported by your broker version).
+
+Combining with FlowController
+------------------------------
+``RateLimitMiddleware`` limits the *consumer* side (processing rate).
+``FlowController`` / ``BackpressureConfig`` limits the *publisher* side.
+Use both together for full end-to-end flow control:
+
+    from rabbitkit import FlowController, BackpressureConfig
+    from rabbitkit.middleware.rate_limit import RateLimitMiddleware, RateLimitConfig
+
+    # Publisher-side: max 5 000 msgs/s
+    fc = FlowController(BackpressureConfig(rate_limit=5000))
+
+    # Consumer-side: process max 200 msgs/s, nack the rest
+    rate_mw = RateLimitMiddleware(RateLimitConfig(max_rate=200, on_limited="nack"))
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+import time
+from dataclasses import dataclass
+from typing import Any
+
+from rabbitkit.core.message import RabbitMessage
+from rabbitkit.middleware.base import BaseMiddleware
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class RateLimitConfig:
+    """Configuration for rate limiting.
+
+    Attributes:
+        max_rate: Maximum messages per second.
+        burst: Maximum burst size above steady rate.
+        on_limited: Behavior when rate exceeded: "wait", "nack", or "drop".
+    """
+
+    max_rate: float
+    burst: int = 1
+    on_limited: str = "wait"  # "wait" | "nack" | "drop"
+
+    def __post_init__(self) -> None:
+        if self.max_rate <= 0:
+            raise ValueError("max_rate must be positive")
+        if self.burst < 1:
+            raise ValueError("burst must be >= 1")
+        if self.on_limited not in ("wait", "nack", "drop"):
+            raise ValueError(f"on_limited must be 'wait', 'nack', or 'drop', got '{self.on_limited}'")
+
+
+class _TokenBucket:
+    """Thread-safe token bucket for rate limiting."""
+
+    __slots__ = ("_capacity", "_last_refill", "_lock", "_rate", "_tokens")
+
+    def __init__(self, rate: float, capacity: int) -> None:
+        self._rate = rate
+        self._capacity = capacity
+        self._tokens = float(capacity)
+        self._last_refill = time.monotonic()
+        self._lock = threading.Lock()
+
+    def _refill(self) -> None:
+        now = time.monotonic()
+        elapsed = now - self._last_refill
+        self._tokens = min(self._capacity, self._tokens + elapsed * self._rate)
+        self._last_refill = now
+
+    def try_acquire(self) -> bool:
+        """Try to acquire a token without blocking. Returns True if acquired."""
+        with self._lock:
+            self._refill()
+            if self._tokens >= 1.0:
+                self._tokens -= 1.0
+                return True
+            return False
+
+    def wait_time(self) -> float:
+        """Return seconds to wait until a token is available."""
+        with self._lock:
+            self._refill()
+            if self._tokens >= 1.0:
+                return 0.0
+            deficit = 1.0 - self._tokens
+            return deficit / self._rate
+
+
+class RateLimitMiddleware(BaseMiddleware):
+    """Limits message processing rate using a token bucket.
+
+    Behavior when rate is exceeded (configurable via on_limited):
+    - "wait": Sleep until a token is available (default)
+    - "nack": Reject message with requeue=True (another consumer can try)
+    - "drop": Reject message with requeue=False (message is lost/goes to DLQ)
+
+    L5 — per-process scoping: the token bucket lives in this process's
+    memory. With ``worker_count`` > 1 consumer processes/pods sharing a
+    queue, each gets its own independent bucket — the effective CLUSTER
+    rate is ``workers x max_rate``, not ``max_rate``. Size ``max_rate``
+    accordingly, or put a cluster-wide limiter in front (e.g. a
+    Redis-backed token bucket) if you need a hard global cap. Under
+    sustained overload, "wait" blocks for up to ``_wait_deadline`` (30s by
+    default) before falling back to a drop (nack, no requeue) — every
+    nack/drop/wait-timeout is logged at WARNING and, if you pass
+    ``metrics_collector``, increments ``rate_limit_dropped_total`` labeled
+    by ``reason`` so sustained overload is observable instead of silent.
+    """
+
+    def __init__(
+        self,
+        config: RateLimitConfig,
+        *,
+        metrics_collector: Any | None = None,
+        metrics_config: Any | None = None,
+    ) -> None:
+        self._config = config
+        self._bucket = _TokenBucket(config.max_rate, config.burst)
+        # Maximum time (s) a "wait" policy will block for a token before falling
+        # back to drop semantics. Override per-instance if needed.
+        self._wait_deadline: float = 30.0
+        self._metrics_collector = metrics_collector
+        self._metrics_config = metrics_config
+
+    def _record_drop(self, reason: str) -> None:
+        """L5: log + optional metric every time a message is settled without
+        the handler running, so sustained rate-limit pressure is observable."""
+        logger.warning("RateLimitMiddleware dropped a message (reason=%s)", reason)
+        if self._metrics_collector is not None and self._metrics_config is not None:
+            self._metrics_collector.inc_counter(
+                self._metrics_config.rate_limit_dropped_total, {"reason": reason}
+            )
+
+    def consume_scope(
+        self,
+        call_next: Any,
+        message: RabbitMessage,
+    ) -> Any:
+        """Rate-limit sync message processing.
+
+        For ``on_limited="wait"`` the loop polls the token bucket until a token
+        is acquired **or** ``wait_deadline`` (seconds, monotonic) expires. If the
+        deadline elapses with no token, the message falls back to the configured
+        drop/nack semantics so the handler is **never** invoked without a token.
+        """
+        if self._bucket.try_acquire():
+            return call_next(message)
+
+        if self._config.on_limited == "nack":
+            self._record_drop("nack")
+            if not message.is_settled:
+                message.nack(requeue=True)
+            return None
+        if self._config.on_limited == "drop":
+            self._record_drop("drop")
+            if not message.is_settled:
+                message.nack(requeue=False)
+            return None
+
+        # "wait" — bounded loop; only proceed once a token is actually acquired.
+        deadline = time.monotonic() + self._wait_deadline
+        while not self._bucket.try_acquire():
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                # No token within the deadline — fall back to drop semantics so
+                # the handler is NOT called without a token.
+                self._record_drop("wait_deadline_exceeded")
+                if not message.is_settled:
+                    message.nack(requeue=False)
+                return None
+            time.sleep(min(self._bucket.wait_time(), remaining))
+
+        return call_next(message)
+
+    async def consume_scope_async(
+        self,
+        call_next: Any,
+        message: RabbitMessage,
+    ) -> Any:
+        """Rate-limit async message processing.
+
+        Mirrors the sync logic: the "wait" loop is bounded by
+        ``self._wait_deadline`` and falls back to drop semantics if no token is
+        acquired in time, so the handler is never called without a token.
+        """
+        if self._bucket.try_acquire():
+            return await call_next(message)
+
+        if self._config.on_limited == "nack":
+            self._record_drop("nack")
+            if not message.is_settled:
+                await message.nack_async(requeue=True)
+            return None
+        if self._config.on_limited == "drop":
+            self._record_drop("drop")
+            if not message.is_settled:
+                await message.nack_async(requeue=False)
+            return None
+
+        # "wait" — bounded loop; only proceed once a token is actually acquired.
+        deadline = time.monotonic() + self._wait_deadline
+        while not self._bucket.try_acquire():
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                self._record_drop("wait_deadline_exceeded")
+                if not message.is_settled:
+                    await message.nack_async(requeue=False)
+                return None
+            await asyncio.sleep(min(self._bucket.wait_time(), remaining))
+
+        return await call_next(message)