rabbitkit 0.9.0__py3-none-any.whl

rabbitkit/async_/connection.py

@@ -0,0 +1,274 @@
+"""aio-pika-specific connection parameter builders and error tuples.
+
+This is where all aio-pika imports live — core/ stays clean.
+Provides helpers to build aio_pika.connect_robust() kwargs from rabbitkit config objects.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import random
+import ssl
+from typing import Any
+from urllib.parse import quote
+
+from rabbitkit.core.config import ConnectionConfig, SecurityConfig, SSLConfig
+
+logger = logging.getLogger(__name__)
+
+
+# ── Transport-specific connection errors ──────────────────────────────────
+
+
+def get_connection_errors() -> tuple[type[BaseException], ...]:
+    """Get aio-pika-specific connection error tuple.
+
+    Returns generic stdlib errors if aio-pika is not installed.
+    """
+    base_errors: tuple[type[BaseException], ...] = (
+        ConnectionResetError,
+        BrokenPipeError,
+        ConnectionAbortedError,
+        ConnectionRefusedError,
+        TimeoutError,
+        EOFError,
+        OSError,
+    )
+
+    try:
+        import aio_pika.exceptions
+
+        aio_pika_errors: tuple[type[BaseException], ...] = (
+            aio_pika.exceptions.AMQPConnectionError,
+            aio_pika.exceptions.ChannelClosed,
+            aio_pika.exceptions.ConnectionClosed,
+        )
+        return aio_pika_errors + base_errors
+    except (ImportError, AttributeError):
+        return base_errors
+
+
+def build_ssl_context(ssl_config: SSLConfig) -> ssl.SSLContext | None:
+    """Build stdlib ssl.SSLContext from SSLConfig.
+
+    Returns None if SSL is not enabled.
+    Shared with sync/connection.py — same logic.
+    """
+    if not ssl_config.enabled:
+        return None
+
+    cert_reqs_map = {
+        "CERT_REQUIRED": ssl.CERT_REQUIRED,
+        "CERT_OPTIONAL": ssl.CERT_OPTIONAL,
+        "CERT_NONE": ssl.CERT_NONE,
+    }
+    cert_reqs = cert_reqs_map.get(ssl_config.cert_reqs, ssl.CERT_REQUIRED)
+
+    # M13: disabling certificate verification makes the connection MITM-able —
+    # warn loudly (see sync/connection.py for rationale).
+    if cert_reqs == ssl.CERT_NONE:
+        import warnings
+
+        warnings.warn(
+            "SSLConfig(cert_reqs='CERT_NONE') disables TLS certificate and hostname "
+            "verification — the connection is encrypted but MITM-able. Use "
+            "'CERT_REQUIRED' (the default) with a proper ca_certs bundle in production.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+
+    # Defense in depth: never negotiate below TLS 1.2.
+    ctx.minimum_version = ssl.TLSVersion.TLSv1_2
+
+    # check_hostname must be disabled BEFORE setting verify_mode=CERT_NONE
+    if cert_reqs == ssl.CERT_NONE:
+        ctx.check_hostname = False
+
+    ctx.verify_mode = cert_reqs
+
+    if ssl_config.ca_certs:
+        ctx.load_verify_locations(ssl_config.ca_certs)
+    elif cert_reqs == ssl.CERT_REQUIRED:
+        # No explicit CA bundle configured — fall back to the system trust
+        # store so verification actually succeeds against broker certs
+        # signed by a well-known CA. Guarded so the explicit-ca_certs path
+        # above is unchanged.
+        try:
+            ctx.load_default_certs()
+        except Exception:  # pragma: no cover — best effort, platform-dependent
+            try:
+                ctx.set_default_verify_paths()
+            except Exception:  # pragma: no cover
+                pass
+
+    if ssl_config.certfile:
+        ctx.load_cert_chain(
+            certfile=ssl_config.certfile,
+            keyfile=ssl_config.keyfile,
+        )
+
+    return ctx
+
+
+def make_aio_pika_connect_kwargs(
+    connection: ConnectionConfig,
+    security: SecurityConfig,
+    *,
+    host_override: str | None = None,
+    port_override: int | None = None,
+) -> dict[str, Any]:
+    """Build kwargs for aio_pika.connect_robust().
+
+    Returns a dict of keyword arguments.
+    Raises ImportError if aio-pika is not installed.
+
+    Note: aio-pika has no native ``blocked_connection_timeout`` knob (unlike
+    pika's ``ConnectionParameters.blocked_connection_timeout``). To honour
+    ``ConnectionConfig.blocked_connection_timeout`` on the async side, call
+    :func:`install_blocked_connection_watchdog` on the returned connection
+    after ``connect_robust`` succeeds. That helper drives a timer task that
+    closes the connection when a ``connection.blocked`` alarm is not cleared
+    by ``connection.unblocked`` within the configured timeout, forcing a
+    reconnect instead of stalling publishes indefinitely.
+    """
+    try:
+        import aio_pika  # noqa: F401
+    except ImportError:
+        raise ImportError(
+            "aio-pika is required for async transport. Install it with: pip install rabbitkit[async]"
+        ) from None
+
+    # Build URL — carry heartbeat as a query param. aio-pika/aiormq read heartbeat
+    # from the URL; passing it as a kwarg is not portable across versions. Without
+    # this the configured ConnectionConfig.heartbeat was silently dropped on async
+    # (the sync transport already honors it).
+    #
+    # URL-encode username/password so credentials containing special characters
+    # (e.g. ":", "@", "/", "+") don't corrupt the AMQP URL or leak as plaintext
+    # delimiters. ConnectionConfig.url (core/config.py) is NOT in this module's
+    # scope, so we rebuild a safe URL here.
+    # M13: resolve credentials via credentials_provider if set (rotation).
+    raw_user, raw_pwd = connection.resolve_credentials()
+    user = quote(raw_user, safe="")
+    pwd = quote(raw_pwd, safe="")
+    vhost = connection.vhost
+    if vhost == "/":
+        vhost = "%2F"
+    # M9: allow the caller (pool) to target a specific cluster node.
+    host = host_override if host_override is not None else connection.host
+    port = port_override if port_override is not None else connection.port
+    base_url = f"amqp://{user}:{pwd}@{host}:{port}/{vhost}"
+    sep = "&" if "?" in base_url else "?"
+    url = f"{base_url}{sep}heartbeat={connection.heartbeat}"
+
+    # H4: aio-pika's connect_robust uses a FIXED reconnect_interval — no
+    # exponential backoff. A fleet of consumers restarted together (e.g. a
+    # broker bounce under 200 pods) would otherwise retry in lockstep every
+    # `reconnect_backoff_base` seconds, hammering the recovering node in a
+    # thundering herd. We can't inject exponential backoff into connect_robust,
+    # but randomizing the interval PER PROCESS (full jitter over
+    # [base, min(base*2, backoff_max)]) de-synchronizes the herd so retries
+    # spread across the window instead of arriving as a spike. Each process
+    # picks its interval once at connect time; `random` needs no seeding for
+    # inter-process spread since PYTHONHASHSEED/os entropy differ per process.
+    base = connection.reconnect_backoff_base
+    jitter_span = max(0.0, min(base, connection.reconnect_backoff_max - base))
+    jittered_interval = base + random.uniform(0.0, jitter_span)  # noqa: S311 — jitter, not crypto
+    kwargs: dict[str, Any] = {
+        "url": url,
+        "timeout": connection.socket_timeout,
+        "reconnect_interval": jittered_interval,
+    }
+
+    # SSL
+    ssl_context = build_ssl_context(security.ssl)
+    if ssl_context is not None:
+        kwargs["ssl_context"] = ssl_context
+
+    # Client properties
+    if connection.connection_name:
+        kwargs["client_properties"] = {
+            "connection_name": connection.connection_name,
+        }
+
+    return kwargs
+
+
+async def install_blocked_connection_watchdog(connection: Any, blocked_timeout: float) -> None:
+    """Install a watchdog that closes *connection* when a blocked alarm lingers.
+
+    aio-pika has no native ``blocked_connection_timeout`` (I-11), so without
+    this a ``connection.blocked`` alarm from RabbitMQ (memory/disk pressure)
+    can stall publishes indefinitely while the connection stays "open". This
+    helper registers ``connection.connection_blocked`` /
+    ``connection.connection_unblocked`` callbacks that drive a timer task:
+
+    - on ``blocked``: start (or replace) a task that sleeps *blocked_timeout*
+      then closes the connection (forcing ``connect_robust`` to reconnect).
+    - on ``unblocked``: cancel the pending timer so a transient alarm does not
+      tear down a recovered connection.
+
+    Safe to call with ``blocked_timeout <= 0`` (no-op) or on a connection
+    that does not expose the callback collections (logged at debug, no raise).
+    Must be called from the event loop that owns *connection*.
+    """
+    if blocked_timeout <= 0:
+        return
+
+    blocked_cb_collection = getattr(connection, "connection_blocked", None)
+    unblocked_cb_collection = getattr(connection, "connection_unblocked", None)
+    if blocked_cb_collection is None or unblocked_cb_collection is None:
+        logger.debug(
+            "connection does not expose connection_blocked/connection_unblocked "
+            "callback collections; blocked-connection watchdog not installed"
+        )
+        return
+
+    loop = asyncio.get_event_loop()
+    state: dict[str, asyncio.Task[None] | None] = {"timer": None}
+
+    async def _on_blocked(*_args: Any) -> None:
+        # Replace any pending timer with a fresh one.
+        existing = state.get("timer")
+        if existing is not None and not existing.done():
+            existing.cancel()
+        logger.warning(
+            "Connection blocked by RabbitMQ; will close in %.1fs if not unblocked",
+            blocked_timeout,
+        )
+        state["timer"] = asyncio.ensure_future(_close_after(blocked_timeout))
+
+    async def _on_unblocked(*_args: Any) -> None:
+        existing = state.get("timer")
+        if existing is not None and not existing.done():
+            existing.cancel()
+        state["timer"] = None
+        logger.info("Connection unblocked; watchdog timer cancelled")
+
+    async def _close_after(delay: float) -> None:
+        try:
+            await asyncio.sleep(delay)
+        except asyncio.CancelledError:
+            return
+        logger.warning("Connection blocked for > %.1fs; closing to force reconnect", delay)
+        try:
+            close = connection.close
+            result = close()
+            if hasattr(result, "__await__"):
+                await result
+        except Exception:  # pragma: no cover — best effort; connect_robust will retry
+            logger.debug("watchdog close raised", exc_info=True)
+
+    # aio-pika CallbackCollection.add_callback accepts a coroutine fn.
+    try:
+        blocked_cb_collection.add_callback(_on_blocked)
+        unblocked_cb_collection.add_callback(_on_unblocked)
+    except Exception:  # pragma: no cover — defensive across aio-pika versions
+        logger.debug("Could not register blocked/unblocked watchdog callbacks", exc_info=True)
+        return
+    # Keep a reference so the timer is not GC'd and the callbacks are traceable.
+    connection._rabbitkit_blocked_watchdog = state
+    connection._rabbitkit_blocked_watchdog_loop = loop

rabbitkit/async_/pool.py

@@ -0,0 +1,363 @@
+"""Async connection and channel pools.
+
+Minimal in 0.1.0 — internal performance utilities.
+Do not oversell as a promised optimization layer.
+
+Uses asyncio.Queue for channel pooling and dedicated
+connections for publisher vs consumer separation.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import random
+from contextlib import asynccontextmanager
+from typing import Any
+
+from rabbitkit.async_.connection import get_connection_errors, make_aio_pika_connect_kwargs
+from rabbitkit.core.config import ConnectionConfig, PoolConfig, SecurityConfig
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncChannelPool:
+    """Async channel pool.
+
+    Manages a pool of aio-pika channels on a single connection.
+    Channels are acquired and released by callers.
+
+    ``acquire_timeout`` controls how long to wait when all channels are
+    checked out.  Raises ``asyncio.TimeoutError`` on exhaustion rather than
+    blocking forever (which would deadlock if a handler tries to publish).
+    """
+
+    def __init__(
+        self,
+        connection: Any,  # aio_pika.RobustConnection
+        pool_size: int = 10,
+        acquire_timeout: float = 10.0,
+        publisher_confirms: bool = True,
+    ) -> None:
+        self._connection = connection
+        self._pool_size = pool_size
+        self._acquire_timeout = acquire_timeout
+        self._publisher_confirms = publisher_confirms
+        self._pool: asyncio.Queue[Any] = asyncio.Queue(maxsize=pool_size)
+        self._lock = asyncio.Lock()
+        self._created = 0
+        # Channels currently checked out by callers; closed in close_all() so they
+        # are not orphaned if a caller forgets to release (leak detection).
+        self._in_use: set[Any] = set()
+
+    async def acquire(self) -> Any:
+        """Acquire a channel from the pool.
+
+        Creates a new channel if the pool is empty and under the size limit.
+        Waits up to ``acquire_timeout`` seconds when all channels are in use;
+        raises ``asyncio.TimeoutError`` if the wait expires, preventing
+        deadlocks when handlers publish while processing messages.
+        """
+        try:
+            channel = self._pool.get_nowait()
+            if not channel.is_closed:
+                async with self._lock:
+                    self._in_use.add(channel)
+                return channel
+            # I-6: a pooled channel was found closed — it still counted
+            # against _created when it was released, so decrement before
+            # discarding it (otherwise acquire() leaks a slot each time a
+            # closed-idle channel is pulled from the pool).
+            async with self._lock:
+                self._created = max(0, self._created - 1)
+        except asyncio.QueueEmpty:
+            pass
+
+        # perf-M-2: create the channel OUTSIDE the lock (network round-trip) so
+        # concurrent acquire() calls don't serialize on channel creation during
+        # warmup/refill. The slot is reserved atomically under the lock (so we
+        # never over-create); the channel-open I/O happens outside the lock, and
+        # we re-acquire only to publish _in_use. If creation fails, the reserved
+        # slot is returned. The I-6 closed-idle decrement above stays under the
+        # lock.
+        async with self._lock:
+            if self._created < self._pool_size:
+                need_create = True
+                self._created += 1
+            else:
+                need_create = False
+
+        if need_create:
+            try:
+                channel = await self._connection.channel(publisher_confirms=self._publisher_confirms)
+            except BaseException:
+                # creation failed — give the reserved slot back.
+                async with self._lock:
+                    self._created = max(0, self._created - 1)
+                raise
+            async with self._lock:
+                self._in_use.add(channel)
+                return channel
+
+        # Pool exhausted — wait with timeout to avoid deadlocks. R-timeout:
+        # ``asyncio.timeout`` (3.11+) replaces ``asyncio.wait_for`` to avoid
+        # the wrapper-task overhead.
+        logger.warning(
+            "Channel pool exhausted (pool_size=%d, created=%d). "
+            "Waiting up to %.1fs for a channel to be released. "
+            "Consider increasing PoolConfig.channel_pool_size.",
+            self._pool_size,
+            self._created,
+            self._acquire_timeout,
+        )
+        try:
+            async with asyncio.timeout(self._acquire_timeout):
+                channel = await self._pool.get()
+        except TimeoutError:
+            raise TimeoutError(
+                f"Timed out after {self._acquire_timeout}s waiting for a channel "
+                f"from the pool (pool_size={self._pool_size})."
+            ) from None
+        if channel.is_closed:
+            async with self._lock:
+                self._created = max(0, self._created - 1)
+            return await self.acquire()
+        async with self._lock:
+            self._in_use.add(channel)
+        return channel
+
+    async def release(self, channel: Any) -> None:
+        """Release a channel back to the pool."""
+        async with self._lock:
+            self._in_use.discard(channel)
+        if not channel.is_closed:
+            try:
+                self._pool.put_nowait(channel)
+                return
+            except asyncio.QueueFull:
+                pass
+        # Channel is closed or pool is full — discard
+        try:
+            if not channel.is_closed:
+                await channel.close()
+        except Exception:
+            pass
+        async with self._lock:
+            self._created = max(0, self._created - 1)
+
+    @asynccontextmanager
+    async def acquire_ctx(self) -> Any:
+        """Async context manager for acquire/release — prevents leaks.
+
+        Usage::
+
+            async with pool.acquire_ctx() as ch:
+                await ch.publish(...)
+        """
+        channel = await self.acquire()
+        try:
+            yield channel
+        finally:
+            await self.release(channel)
+
+    async def close_all(self) -> None:
+        """Close all channels in the pool, including checked-out ones."""
+        async with self._lock:
+            in_use = list(self._in_use)
+            self._in_use.clear()
+        for channel in in_use:
+            try:
+                if not channel.is_closed:
+                    await channel.close()
+            except Exception:  # pragma: no cover — best effort
+                pass
+            async with self._lock:
+                self._created = max(0, self._created - 1)
+        while not self._pool.empty():
+            try:
+                channel = self._pool.get_nowait()
+                if not channel.is_closed:
+                    await channel.close()
+            except (asyncio.QueueEmpty, Exception):
+                pass
+        async with self._lock:
+            self._created = 0
+
+    @property
+    def size(self) -> int:
+        """Number of channels currently in the pool (available)."""
+        return self._pool.qsize()
+
+    @property
+    def created_count(self) -> int:
+        """Total number of channels created."""
+        return self._created
+
+
+class AsyncConnectionPool:
+    """Separate publisher/consumer connections with channel pools.
+
+    Provides dedicated connections for publishing and consuming to avoid
+    head-of-line blocking, and exposes ``AsyncChannelPool`` instances so
+    callers never share a single channel across concurrent operations.
+
+    Usage::
+
+        pool = AsyncConnectionPool(connection_config, security_config, pool_config)
+        await pool.connect()
+
+        async with pool.acquire_publisher_channel() as ch:
+            await ch.publish(...)
+
+        await pool.close_all()
+    """
+
+    def __init__(
+        self,
+        connection_config: ConnectionConfig,
+        security_config: SecurityConfig,
+        pool_config: PoolConfig | None = None,
+        publisher_confirms: bool = True,
+    ) -> None:
+        self._connection_config = connection_config
+        self._security_config = security_config
+        self._pool_config = pool_config or PoolConfig()
+        self._publisher_confirms = publisher_confirms
+
+        self._publisher_connection: Any | None = None
+        self._consumer_connection: Any | None = None
+        self._publisher_channel_pool: AsyncChannelPool | None = None
+        self._lock = asyncio.Lock()
+        self._prewarmed = False
+
+    async def connect(self) -> None:
+        """Establish publisher and consumer connections eagerly."""
+        do_prewarm = False
+        async with self._lock:
+            if self._publisher_connection is None:
+                self._publisher_connection = await self._create_connection()
+                self._publisher_channel_pool = AsyncChannelPool(
+                    self._publisher_connection,
+                    pool_size=self._pool_config.channel_pool_size,
+                    acquire_timeout=self._pool_config.channel_acquire_timeout,
+                    publisher_confirms=self._publisher_confirms,
+                )
+            if self._consumer_connection is None:
+                self._consumer_connection = await self._create_connection()
+            if self._pool_config.prewarm_channels and not self._prewarmed:
+                self._prewarmed = True
+                do_prewarm = True
+
+        if do_prewarm and self._publisher_channel_pool is not None:
+            pool = self._publisher_channel_pool
+            channels = await asyncio.gather(
+                *(pool.acquire() for _ in range(self._pool_config.channel_pool_size)),
+                return_exceptions=True,
+            )
+            for ch in channels:
+                if not isinstance(ch, BaseException):
+                    await pool.release(ch)
+
+    async def get_publisher_connection(self) -> Any:
+        """Get a connection dedicated for publishing.
+
+        Creates the connection lazily on first call.
+        """
+        async with self._lock:
+            if self._publisher_connection is None:
+                self._publisher_connection = await self._create_connection()
+                self._publisher_channel_pool = AsyncChannelPool(
+                    self._publisher_connection,
+                    pool_size=self._pool_config.channel_pool_size,
+                    acquire_timeout=self._pool_config.channel_acquire_timeout,
+                    publisher_confirms=self._publisher_confirms,
+                )
+            return self._publisher_connection
+
+    async def get_consumer_connection(self) -> Any:
+        """Get a connection dedicated for consuming.
+
+        Creates the connection lazily on first call.
+        """
+        async with self._lock:
+            if self._consumer_connection is None:
+                self._consumer_connection = await self._create_connection()
+            return self._consumer_connection
+
+    async def acquire_publisher_channel(self) -> Any:
+        """Acquire a channel from the publisher channel pool."""
+        if self._publisher_channel_pool is None:
+            await self.get_publisher_connection()
+        assert self._publisher_channel_pool is not None
+        return await self._publisher_channel_pool.acquire()
+
+    async def release_publisher_channel(self, channel: Any) -> None:
+        """Return a publisher channel to the pool."""
+        if self._publisher_channel_pool is not None:
+            await self._publisher_channel_pool.release(channel)
+
+    async def close_all(self) -> None:
+        """Close all channel pools and connections."""
+        async with self._lock:
+            if self._publisher_channel_pool is not None:
+                await self._publisher_channel_pool.close_all()
+                self._publisher_channel_pool = None
+
+            for conn in [self._publisher_connection, self._consumer_connection]:
+                if conn is not None:
+                    try:
+                        if not conn.is_closed:
+                            await conn.close()
+                    except Exception as e:
+                        logger.warning("Error closing connection: %s", e)
+
+            self._publisher_connection = None
+            self._consumer_connection = None
+
+    async def _create_connection(self) -> Any:
+        """Create a new aio-pika connection."""
+        try:
+            import aio_pika
+        except ImportError:
+            raise ImportError(
+                "aio-pika is required for async transport. Install it with: pip install rabbitkit[async]"
+            ) from None
+
+        # M9: cycle through cluster endpoints on the initial connect so a dead
+        # configured primary doesn't take the client down at startup. Once
+        # connect_robust succeeds it pins to that node for reconnects (aio-pika
+        # has no multi-host reconnect) — put a load balancer / DNS in front for
+        # per-reconnect failover across nodes.
+        endpoints = self._connection_config.cluster_endpoints()
+
+        # H-SRE3: connect_robust handles reconnects AFTER the first connection
+        # with a FIXED interval, so a fleet of clients starting at once thunder
+        # the broker. Apply an outer retry with full jitter for the INITIAL
+        # connect only; bounded so we never spin forever.
+        backoff = self._connection_config.reconnect_backoff_base
+        max_backoff = self._connection_config.reconnect_backoff_max
+        connection_errors = get_connection_errors()
+        max_attempts = 30
+        for attempt in range(1, max_attempts + 1):
+            host, port = endpoints[(attempt - 1) % len(endpoints)]
+            kwargs = make_aio_pika_connect_kwargs(
+                self._connection_config,
+                self._security_config,
+                host_override=host,
+                port_override=port,
+            )
+            try:
+                return await aio_pika.connect_robust(**kwargs)
+            except connection_errors as e:
+                if attempt == max_attempts:
+                    raise
+                sleep_for = random.uniform(0.0, backoff)  # noqa: S311
+                logger.warning(
+                    "aio-pika initial connect failed (attempt %d), retrying in %.2fs: %s",
+                    attempt,
+                    sleep_for,
+                    e,
+                )
+                await asyncio.sleep(sleep_for)
+                backoff = min(backoff * 2, max_backoff)
+        raise RuntimeError("unreachable")  # pragma: no cover