rabbitkit 0.9.0__py3-none-any.whl

rabbitkit/sync/connection.py

@@ -0,0 +1,209 @@
+"""Pika-specific connection parameter builders and error tuples.
+
+This is where all pika imports live — core/ stays clean.
+Provides helpers to build pika.ConnectionParameters from rabbitkit config objects.
+"""
+
+from __future__ import annotations
+
+import logging
+import socket
+import ssl
+from typing import Any
+
+from rabbitkit.core.config import ConnectionConfig, SecurityConfig, SocketConfig, SSLConfig
+
+logger = logging.getLogger(__name__)
+
+
+# ── Transport-specific connection errors ──────────────────────────────────
+# These extend the core TRANSIENT_ERRORS for pika-specific exceptions.
+# Lazy-loaded to avoid import errors when pika is not installed.
+
+
+def get_connection_errors() -> tuple[type[BaseException], ...]:
+    """Get pika-specific connection error tuple.
+
+    Returns generic stdlib errors if pika is not installed.
+    """
+    base_errors: tuple[type[BaseException], ...] = (
+        ConnectionResetError,
+        BrokenPipeError,
+        ConnectionAbortedError,
+        ConnectionRefusedError,
+        TimeoutError,
+        EOFError,
+        OSError,
+    )
+
+    try:
+        import pika.exceptions
+
+        pika_errors: tuple[type[BaseException], ...] = (
+            pika.exceptions.StreamLostError,
+            pika.exceptions.AMQPConnectionError,
+            pika.exceptions.ConnectionClosedByBroker,
+            pika.exceptions.ChannelWrongStateError,
+            pika.exceptions.ChannelClosedByBroker,
+            pika.exceptions.AMQPChannelError,
+        )
+        return pika_errors + base_errors
+    except ImportError:
+        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.
+    """
+    if not ssl_config.enabled:
+        return None
+
+    # Determine cert_reqs
+    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, since it's a copy-paste "make TLS errors go
+    # away" footgun that otherwise ships silently to production.
+    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
+    # (Python 3.12+ raises ValueError otherwise)
+    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. Without this, CERT_REQUIRED + no ca_certs
+        # silently leaves the context with zero trusted CAs → every handshake
+        # fails. 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_pika_connection_params(
+    connection: ConnectionConfig,
+    socket_config: SocketConfig,
+    security: SecurityConfig,
+) -> Any:
+    """Build pika.ConnectionParameters with TCP tuning, SSL, heartbeat.
+
+    Returns a pika.ConnectionParameters object.
+    Raises ImportError if pika is not installed.
+    """
+    try:
+        import pika
+    except ImportError:
+        raise ImportError("pika is required for sync transport. Install it with: pip install rabbitkit[sync]") from None
+
+    # SSL context
+    ssl_context = build_ssl_context(security.ssl)
+    ssl_options = None
+    if ssl_context is not None:
+        ssl_options = pika.SSLOptions(
+            context=ssl_context,
+            server_hostname=security.ssl.server_hostname or connection.host,
+        )
+
+    # Credentials (M13: resolve via credentials_provider if set, so a rotated
+    # secret is picked up on this (re)connect).
+    username, password = connection.resolve_credentials()
+    credentials = pika.PlainCredentials(
+        username=username,
+        password=password,
+    )
+
+    # Client properties
+    client_properties: dict[str, str] = {}
+    if connection.connection_name:
+        client_properties["connection_name"] = connection.connection_name
+
+    def _params_for(host: str, port: int) -> Any:
+        return pika.ConnectionParameters(
+            host=host,
+            port=port,
+            virtual_host=connection.vhost,
+            credentials=credentials,
+            heartbeat=connection.heartbeat,
+            socket_timeout=connection.socket_timeout,
+            blocked_connection_timeout=connection.blocked_connection_timeout,
+            ssl_options=ssl_options,
+            client_properties=client_properties if client_properties else None,
+        )
+
+    endpoints = connection.cluster_endpoints()
+    if len(endpoints) == 1:
+        return _params_for(*endpoints[0])
+    # M9: pika.BlockingConnection accepts a LIST of ConnectionParameters and
+    # tries each in order until one connects — native cluster failover.
+    return [_params_for(host, port) for host, port in endpoints]
+
+
+def apply_socket_options(sock: socket.socket, config: SocketConfig) -> None:
+    """Apply TCP_NODELAY, keepalive, buffer sizes to a socket.
+
+    Best-effort — not all options are universally guaranteed
+    depending on OS and backend internals.
+    """
+    try:
+        # TCP_NODELAY — disable Nagle's algorithm
+        if config.tcp_nodelay:
+            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+
+        # TCP keepalive
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+
+        # Platform-specific keepalive options
+        if hasattr(socket, "TCP_KEEPIDLE"):
+            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, config.tcp_keepidle)
+        if hasattr(socket, "TCP_KEEPINTVL"):
+            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, config.tcp_keepintvl)
+        if hasattr(socket, "TCP_KEEPCNT"):
+            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, config.tcp_keepcnt)
+
+        # Buffer sizes
+        if config.tcp_sndbuf > 0:
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, config.tcp_sndbuf)
+        if config.tcp_rcvbuf > 0:
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, config.tcp_rcvbuf)
+
+    except OSError as e:
+        logger.warning("Failed to apply socket option: %s (best-effort, continuing)", e)

rabbitkit/sync/pool.py

@@ -0,0 +1,262 @@
+"""Sync connection and channel pools.
+
+Minimal in 0.1.0 — internal performance utilities.
+Do not oversell as a promised optimization layer.
+
+Model A: one-connection-per-thread. Pools enforce this by assigning
+dedicated connections per role (publisher vs consumer).
+"""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+from contextlib import contextmanager
+from typing import Any
+
+from rabbitkit.core.config import ConnectionConfig, PoolConfig, SecurityConfig, SocketConfig
+from rabbitkit.sync.connection import make_pika_connection_params
+
+logger = logging.getLogger(__name__)
+
+
+class SyncChannelPool:
+    """Thread-safe channel pool.
+
+    Manages a pool of pika channels on a single connection.
+    Channels are acquired and released by callers.
+
+    ``acquire_timeout`` bounds how long ``acquire()`` blocks when all channels
+    are checked out; it raises ``TimeoutError`` on expiry (mirrors the async
+    pool) rather than blocking forever — which would deadlock a worker that
+    tries to publish while processing.
+
+    Note: the default ``SyncTransport`` does not route publishes through this
+    pool (it uses a single dedicated publisher channel); the pool is kept as a
+    reusable utility for advanced/pooled usage and is covered by unit tests so
+    it is not silently dead code.
+    """
+
+    def __init__(
+        self,
+        connection: Any,  # pika.BlockingConnection
+        pool_size: int = 10,
+        acquire_timeout: float = 10.0,
+    ) -> None:
+        self._connection = connection
+        self._pool_size = pool_size
+        self._acquire_timeout = acquire_timeout
+        self._pool: queue.Queue[Any] = queue.Queue(maxsize=pool_size)
+        self._lock = threading.Lock()
+        self._created = 0
+        # Channels currently checked out by callers (leak detection / close_all).
+        self._in_use: set[Any] = set()
+
+    def acquire(self) -> Any:
+        """Acquire a channel from the pool.
+
+        Creates a new channel if the pool is empty and under the size limit.
+        Blocks up to ``acquire_timeout`` if the pool is exhausted; raises
+        ``TimeoutError`` on expiry so callers don't deadlock.
+        """
+        try:
+            channel = self._pool.get_nowait()
+            if channel.is_open:
+                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).
+            with self._lock:
+                self._created = max(0, self._created - 1)
+        except queue.Empty:
+            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.
+        with self._lock:
+            if self._created < self._pool_size:
+                need_create = True
+                self._created += 1
+            else:
+                need_create = False
+
+        if need_create:
+            try:
+                channel = self._connection.channel()
+            except BaseException:
+                with self._lock:
+                    self._created = max(0, self._created - 1)
+                raise
+            with self._lock:
+                self._in_use.add(channel)
+                return channel
+
+        # Pool exhausted — block until one is released, bounded by acquire_timeout
+        logger.warning(
+            "SyncChannelPool 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:
+            channel = self._pool.get(timeout=self._acquire_timeout)
+        except queue.Empty as e:
+            raise TimeoutError(
+                f"Timed out after {self._acquire_timeout}s waiting for a pooled "
+                "channel. Increase PoolConfig.channel_pool_size."
+            ) from e
+        if not channel.is_open:
+            # Discard the stale channel and recurse to try again.
+            with self._lock:
+                self._created = max(0, self._created - 1)
+            return self.acquire()
+        with self._lock:
+            self._in_use.add(channel)
+        return channel
+
+    def release(self, channel: Any) -> None:
+        """Release a channel back to the pool."""
+        with self._lock:
+            self._in_use.discard(channel)
+        if channel.is_open:
+            try:
+                self._pool.put_nowait(channel)
+                return
+            except queue.Full:
+                pass
+        # Channel is closed or pool is full — discard
+        try:
+            if channel.is_open:
+                channel.close()
+        except Exception:
+            pass
+        with self._lock:
+            self._created = max(0, self._created - 1)
+
+    @contextmanager
+    def acquire_ctx(self) -> Any:
+        """Context manager for acquire/release — prevents leaks."""
+        channel = self.acquire()
+        try:
+            yield channel
+        finally:
+            self.release(channel)
+
+    def close_all(self) -> None:
+        """Close all channels in the pool (including checked-out ones)."""
+        with self._lock:
+            in_use = list(self._in_use)
+            self._in_use.clear()
+        for channel in in_use:
+            try:
+                if channel.is_open:
+                    channel.close()
+            except Exception:  # pragma: no cover — best effort
+                pass
+            with self._lock:
+                self._created = max(0, self._created - 1)
+        while not self._pool.empty():
+            try:
+                channel = self._pool.get_nowait()
+                if channel.is_open:
+                    channel.close()
+            except (queue.Empty, Exception):
+                pass
+        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 SyncConnectionPool:
+    """Separate publisher/consumer connections.
+
+    Provides dedicated connections for publishing and consuming
+    to avoid head-of-line blocking.
+    """
+
+    def __init__(
+        self,
+        connection_config: ConnectionConfig,
+        socket_config: SocketConfig,
+        security_config: SecurityConfig,
+        pool_config: PoolConfig | None = None,
+    ) -> None:
+        self._connection_config = connection_config
+        self._socket_config = socket_config
+        self._security_config = security_config
+        self._pool_config = pool_config or PoolConfig()
+
+        self._publisher_connections: list[Any] = []
+        self._consumer_connections: list[Any] = []
+        self._lock = threading.Lock()
+
+    def get_publisher_connection(self) -> Any:
+        """Get a connection dedicated for publishing.
+
+        Creates the connection lazily on first call.
+        """
+        with self._lock:
+            if not self._publisher_connections:
+                conn = self._create_connection()
+                self._publisher_connections.append(conn)
+            return self._publisher_connections[0]
+
+    def get_consumer_connection(self) -> Any:
+        """Get a connection dedicated for consuming.
+
+        Creates the connection lazily on first call.
+        """
+        with self._lock:
+            if not self._consumer_connections:
+                conn = self._create_connection()
+                self._consumer_connections.append(conn)
+            return self._consumer_connections[0]
+
+    def close_all(self) -> None:
+        """Close all connections."""
+        with self._lock:
+            for conn in self._publisher_connections + self._consumer_connections:
+                try:
+                    if conn.is_open:
+                        conn.close()
+                except Exception as e:
+                    logger.warning("Error closing connection: %s", e)
+
+            self._publisher_connections.clear()
+            self._consumer_connections.clear()
+
+    def _create_connection(self) -> Any:
+        """Create a new pika connection."""
+        try:
+            import pika
+        except ImportError:
+            raise ImportError(
+                "pika is required for sync transport. Install it with: pip install rabbitkit[sync]"
+            ) from None
+
+        params = make_pika_connection_params(
+            self._connection_config,
+            self._socket_config,
+            self._security_config,
+        )
+        return pika.BlockingConnection(params)