rabbitkit 0.9.0__py3-none-any.whl

rabbitkit/__init__.py

@@ -0,0 +1,201 @@
+"""rabbitkit — Production-grade RabbitMQ toolkit."""
+
+from rabbitkit import experimental
+from rabbitkit._version import __version__
+from rabbitkit.async_.batch import AsyncBatchPublisher
+from rabbitkit.async_.broker import AsyncBroker
+from rabbitkit.asyncapi import AsyncAPIGeneratorConfig, generate_asyncapi_doc, generate_asyncapi_json
+from rabbitkit.concurrency import AsyncWorkerPool, SyncWorkerPool
+from rabbitkit.core.app import AppState, RabbitApp
+from rabbitkit.core.config import (
+    RETRY_DISABLED,
+    BackpressureConfig,
+    BatchAckConfig,
+    BatchPublishConfig,
+    CompressionConfig,
+    ConnectionConfig,
+    ConsumerConfig,
+    DeduplicationConfig,
+    HealthCheckConfig,
+    MetricsConfig,
+    PoolConfig,
+    PublisherConfig,
+    RabbitConfig,
+    RetryConfig,
+    RetryDisabled,
+    SafetyConfig,
+    SecurityConfig,
+    SocketConfig,
+    SSLConfig,
+    WorkerConfig,
+)
+from rabbitkit.core.errors import (
+    BackpressureError,
+    ConfigurationError,
+    MissingDependencyError,
+    UnsafeTopologyError,
+)
+from rabbitkit.core.logging import DEFAULT_REDACT_KEYS, LoggingConfig, configure_structlog
+from rabbitkit.core.message import AckMessage, NackMessage, RabbitMessage, RejectMessage
+from rabbitkit.core.router import RabbitRouter
+from rabbitkit.core.topology import RabbitExchange, RabbitQueue
+from rabbitkit.core.types import (
+    AckPolicy,
+    ClassifiedError,
+    DeduplicationMarkPolicy,
+    ErrorSeverity,
+    ExchangeType,
+    MessageEnvelope,
+    PublishOutcome,
+    PublishStatus,
+    QueueType,
+    RejectWithoutDLXPolicy,
+    TopologyMode,
+)
+from rabbitkit.di import Context, ContextRepo, Depends, DIResolver, Header, Path
+from rabbitkit.di.resolver import DependencyScope
+from rabbitkit.dlq import DLQInspector, ReplayResult
+from rabbitkit.fastapi import rabbitkit_lifespan
+from rabbitkit.health import (
+    AsyncHealthWatcher,
+    BrokerHealthResult,
+    HealthStatus,
+    HealthWatcher,
+    broker_health_check,
+    broker_health_check_async,
+    broker_liveness,
+    broker_liveness_async,
+    broker_readiness,
+    broker_readiness_async,
+)
+from rabbitkit.highload.backpressure import FlowController
+from rabbitkit.highload.batch import BatchAcker, BatchPublisher
+from rabbitkit.management import ManagementConfig, RabbitManagementClient
+from rabbitkit.middleware.circuit_breaker import CircuitBreakerMiddleware, CircuitBreakerOpenError
+from rabbitkit.middleware.deduplication import DeduplicationMiddleware
+from rabbitkit.middleware.metrics import (
+    MetricsCollector,
+    MetricsMiddleware,
+    PrometheusCollector,
+    metrics_app,
+    start_metrics_server,
+)
+from rabbitkit.middleware.otel import OTelTracingMiddleware
+from rabbitkit.middleware.rate_limit import RateLimitConfig, RateLimitMiddleware
+from rabbitkit.queue_metrics import QueueMetricsPoller
+from rabbitkit.serialization.pipeline import (
+    DataclassDecoder,
+    JsonParser,
+    MessageDecoder,
+    MessageParser,
+    PydanticDecoder,
+    RawDecoder,
+    SerializationPipeline,
+)
+from rabbitkit.sync.batch import SyncBatchPublisher
+from rabbitkit.sync.broker import SyncBroker
+
+__all__ = [
+    "DEFAULT_REDACT_KEYS",
+    "RETRY_DISABLED",
+    "AckMessage",
+    "AckPolicy",
+    "AppState",
+    "AsyncAPIGeneratorConfig",
+    "AsyncBatchPublisher",
+    "AsyncBroker",
+    "AsyncHealthWatcher",
+    "AsyncWorkerPool",
+    "BackpressureConfig",
+    "BackpressureError",
+    "BatchAckConfig",
+    "BatchAcker",
+    "BatchPublishConfig",
+    "BatchPublisher",
+    "BrokerHealthResult",
+    "CircuitBreakerMiddleware",
+    "CircuitBreakerOpenError",
+    "ClassifiedError",
+    "CompressionConfig",
+    "ConfigurationError",
+    "ConnectionConfig",
+    "ConsumerConfig",
+    "Context",
+    "ContextRepo",
+    "DIResolver",
+    "DLQInspector",
+    "DataclassDecoder",
+    "DeduplicationConfig",
+    "DeduplicationMarkPolicy",
+    "DeduplicationMiddleware",
+    "DependencyScope",
+    "Depends",
+    "ErrorSeverity",
+    "ExchangeType",
+    "FlowController",
+    "Header",
+    "HealthCheckConfig",
+    "HealthStatus",
+    "HealthWatcher",
+    "JsonParser",
+    "LoggingConfig",
+    "ManagementConfig",
+    "MessageDecoder",
+    "MessageEnvelope",
+    "MessageParser",
+    "MetricsCollector",
+    "MetricsConfig",
+    "MetricsMiddleware",
+    "MissingDependencyError",
+    "NackMessage",
+    "OTelTracingMiddleware",
+    "Path",
+    "PoolConfig",
+    "PrometheusCollector",
+    "PublishOutcome",
+    "PublishStatus",
+    "PublisherConfig",
+    "PydanticDecoder",
+    "QueueMetricsPoller",
+    "QueueType",
+    "RabbitApp",
+    "RabbitConfig",
+    "RabbitExchange",
+    "RabbitManagementClient",
+    "RabbitMessage",
+    "RabbitQueue",
+    "RabbitRouter",
+    "RateLimitConfig",
+    "RateLimitMiddleware",
+    "RawDecoder",
+    "RejectMessage",
+    "RejectWithoutDLXPolicy",
+    "ReplayResult",
+    "RetryConfig",
+    "RetryDisabled",
+    "SSLConfig",
+    "SafetyConfig",
+    "SecurityConfig",
+    "SerializationPipeline",
+    "SocketConfig",
+    "SyncBatchPublisher",
+    "SyncBroker",
+    "SyncWorkerPool",
+    "TopologyMode",
+    "UnsafeTopologyError",
+    "WorkerConfig",
+    "__version__",
+    "broker_health_check",
+    "broker_health_check_async",
+    "broker_liveness",
+    "broker_liveness_async",
+    "broker_readiness",
+    "broker_readiness_async",
+    "configure_structlog",
+    "experimental",
+    "generate_asyncapi_doc",
+    "generate_asyncapi_json",
+    "metrics_app",
+    "rabbitkit_lifespan",
+    "start_metrics_server",
+]

rabbitkit/_version.py

@@ -0,0 +1,3 @@
+"""rabbitkit version — single source of truth."""
+
+__version__ = "0.9.0"

rabbitkit/aio/__init__.py

@@ -0,0 +1,31 @@
+"""rabbitkit.aio — deprecated alias for ``rabbitkit.async_`` (L8).
+
+``rabbitkit.async_`` is the canonical import path (it matches the actual
+module structure — ``async_/broker.py``, ``async_/transport.py``, etc. — and
+is what every example, doc, and test in this codebase uses). This module is
+kept only so any existing ``from rabbitkit.aio import ...`` import keeps
+working; importing it emits a ``DeprecationWarning``.
+
+Prefer::
+
+    from rabbitkit.async_.broker import AsyncBroker
+    # or, once exported at top level:
+    from rabbitkit import AsyncBroker
+"""
+
+import warnings
+
+from rabbitkit.async_.broker import AsyncBroker
+from rabbitkit.async_.transport import AsyncTransportImpl
+
+warnings.warn(
+    "rabbitkit.aio is deprecated -- import from rabbitkit.async_ (or rabbitkit) instead. "
+    "rabbitkit.async_ is the canonical async import path.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+__all__ = [
+    "AsyncBroker",
+    "AsyncTransportImpl",
+]

rabbitkit/async_/__init__.py

@@ -0,0 +1,9 @@
+"""Async transport module — aio-pika-based I/O adapter."""
+
+from rabbitkit.async_.broker import AsyncBroker
+from rabbitkit.async_.transport import AsyncTransportImpl
+
+__all__ = [
+    "AsyncBroker",
+    "AsyncTransportImpl",
+]

rabbitkit/async_/batch.py

@@ -0,0 +1,213 @@
+"""BatchPublisher — transparent batch publish with amortized confirm wait.
+
+Collects messages from concurrent callers into fixed-size batches, then
+publishes the whole batch on a single pooled channel.  The broker-side confirm
+wait is paid once per batch rather than once per message, dramatically reducing
+the per-message cost at high concurrency.
+
+Usage::
+
+    from rabbitkit.core.config import BatchPublishConfig
+
+    broker = AsyncBroker(
+        config,
+        batch_config=BatchPublishConfig(batch_size=64, flush_interval_ms=20),
+    )
+    await broker.start()
+    # broker.publish() is now transparently batched
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from typing import TYPE_CHECKING, Any
+
+from rabbitkit.core.config import BatchPublishConfig
+from rabbitkit.core.types import MessageEnvelope, PublishOutcome, PublishStatus
+
+if TYPE_CHECKING:
+    from rabbitkit.async_.transport import AsyncTransportImpl
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncBatchPublisher:
+    """Transparent batch-publish wrapper for AsyncTransportImpl.
+
+    N concurrent ``broker.publish()`` calls are coalesced into one batch,
+    published on a single channel, and their confirms are gathered together.
+    Reduces per-message pool acquire/release and amortizes confirm round-trips.
+
+    The caller's coroutine blocks until its message is included in a flushed
+    batch and the confirm resolves — semantics are identical to direct publish.
+
+    Unlike ``highload.batch.BatchPublisher`` (a timing/buffering helper that
+    publishes each message individually), this class pipelines confirms: all
+    messages in a batch share one channel and their ACKs are awaited together.
+
+    Blast-radius note (M17): because a batch shares one channel, an
+    AMQP-channel-level failure (e.g. a confirm timeout that closes the channel)
+    fails EVERY in-flight publish in that batch — each caller's future is
+    settled with the error (never silently lost), and the channel is replaced
+    before the next batch, but a single slow confirm can amplify into
+    batch-wide errors. This is inherent to sharing a channel for pipelined
+    confirms; there is no per-message isolation without giving up batching. To
+    bound the blast radius, lower ``BatchPublishConfig.batch_size`` (fewer
+    siblings share a channel) — or use direct ``broker.publish()`` for
+    publishes that must fail independently.
+    """
+
+    def __init__(self, transport: AsyncTransportImpl, config: BatchPublishConfig) -> None:
+        self._transport = transport
+        self._config = config
+        self._pending: asyncio.Queue[tuple[MessageEnvelope, asyncio.Future[PublishOutcome]]] = (
+            asyncio.Queue(maxsize=config.max_in_flight)
+        )
+        self._flush_tasks: list[asyncio.Task[None]] = []
+
+    def _worker_count(self) -> int:
+        if self._config.flush_workers > 0:
+            return self._config.flush_workers
+        return min(16, max(1, self._config.max_in_flight // self._config.batch_size))
+
+    async def start(self) -> None:
+        """Start N concurrent flush loops (one per channel slot)."""
+        n = self._worker_count()
+        self._flush_tasks = [
+            asyncio.create_task(self._flush_loop(), name=f"rabbitkit.batch-flush-{i}")
+            for i in range(n)
+        ]
+
+    async def stop(self) -> None:
+        """Cancel all flush loops and drain any remaining queued messages."""
+        for task in self._flush_tasks:
+            task.cancel()
+        for task in self._flush_tasks:
+            with contextlib.suppress(Exception, asyncio.CancelledError):
+                await task
+        self._flush_tasks = []
+        while not self._pending.empty():
+            try:
+                _, fut = self._pending.get_nowait()
+                if not fut.done():
+                    fut.set_exception(RuntimeError("BatchPublisher stopped before flush"))
+            except asyncio.QueueEmpty:  # pragma: no cover — TOCTOU guard, unreachable in asyncio
+                break
+
+    async def publish(self, envelope: MessageEnvelope) -> PublishOutcome:
+        """Enqueue *envelope* and wait for it to be included in a batch flush."""
+        loop = asyncio.get_running_loop()
+        fut: asyncio.Future[PublishOutcome] = loop.create_future()
+        await self._pending.put((envelope, fut))
+        return await fut
+
+    async def _flush_loop(self) -> None:
+        """Each worker holds one channel for its lifetime — no acquire/release per batch."""
+        interval = self._config.flush_interval_ms / 1000.0
+        channel: Any = None
+        batch: list[tuple[MessageEnvelope, asyncio.Future[PublishOutcome]]] = []
+        try:
+            channel = await self._transport._conn_pool.acquire_publisher_channel()
+            while True:
+                batch = []
+
+                # Block until the first item arrives
+                batch.append(await self._pending.get())
+
+                # Fast drain: grab all immediately-available items without yielding.
+                # At high concurrency the queue is almost always non-empty here, so
+                # this avoids the coroutine/timeout overhead of repeated wait_for calls.
+                while len(batch) < self._config.batch_size:
+                    try:
+                        batch.append(self._pending.get_nowait())
+                    except asyncio.QueueEmpty:
+                        break
+
+                # If still under batch_size, wait briefly for stragglers
+                if len(batch) < self._config.batch_size:
+                    loop = asyncio.get_running_loop()
+                    deadline = loop.time() + interval
+                    while len(batch) < self._config.batch_size:
+                        remaining = deadline - loop.time()
+                        if remaining <= 0:
+                            break
+                        try:
+                            async with asyncio.timeout(remaining):
+                                batch.append(await self._pending.get())
+                        except TimeoutError:
+                            break
+
+                try:
+                    await self._flush(channel, batch)
+                except BaseException as exc:
+                    # Settle any futures _flush didn't resolve (including on
+                    # CancelledError). Convert CancelledError to RuntimeError so
+                    # it can be set on the Future (set_exception rejects BaseException).
+                    err: BaseException = (
+                        exc if isinstance(exc, Exception) else RuntimeError("Batch publisher cancelled")
+                    )
+                    for _, fut in batch:
+                        if not fut.done():
+                            fut.set_exception(err)
+                    if isinstance(exc, Exception):
+                        # Channel error — replace it. Set channel=None first so
+                        # the finally block doesn't double-release if acquire fails.
+                        old, channel = channel, None
+                        with contextlib.suppress(Exception):
+                            await self._transport._conn_pool.release_publisher_channel(old)
+                        channel = await self._transport._conn_pool.acquire_publisher_channel()
+                    else:
+                        raise  # CancelledError must propagate
+                else:
+                    # _publish_on_channel closes the channel on confirm timeout;
+                    # detect that and replace it before the next batch.
+                    if channel.is_closed:
+                        old, channel = channel, None
+                        with contextlib.suppress(Exception):
+                            await self._transport._conn_pool.release_publisher_channel(old)
+                        channel = await self._transport._conn_pool.acquire_publisher_channel()
+        except BaseException:
+            # Any exception escaping the loop (e.g. CancelledError at pending.get()
+            # or straggler wait) — settle any dequeued-but-unresolved futures.
+            for _, fut in batch:
+                if not fut.done():
+                    fut.set_exception(RuntimeError("Batch publisher cancelled"))
+            raise
+        finally:
+            if channel is not None:
+                with contextlib.suppress(Exception):
+                    await self._transport._conn_pool.release_publisher_channel(channel)
+
+    async def _flush(
+        self,
+        channel: Any,
+        batch: list[tuple[MessageEnvelope, asyncio.Future[PublishOutcome]]],
+    ) -> None:
+        """Publish all envelopes on the worker's persistent channel and resolve futures."""
+        outcomes: list[Any] = await asyncio.gather(
+            *[self._transport._publish_on_channel(channel, env) for env, _ in batch],
+            return_exceptions=True,
+        )
+        for (_, fut), outcome in zip(batch, outcomes, strict=False):
+            if not fut.done():
+                if isinstance(outcome, BaseException):
+                    fut.set_exception(outcome)
+                else:
+                    fut.set_result(outcome)
+
+        # M17: _publish_on_channel no longer closes the channel itself on a
+        # confirm timeout (that used to fire from INSIDE one of the concurrent
+        # calls above, corrupting sibling in-flight confirms sharing this same
+        # channel). Instead, close it HERE — every concurrent publish in this
+        # batch has already resolved by this point, so closing can no longer
+        # amplify this batch's failure onto a message that would have
+        # confirmed fine. The flush loop's post-_flush check (channel.is_closed)
+        # still replaces it for the next batch, preserving the original
+        # "don't reuse a wedged channel" protection.
+        if not channel.is_closed and any(
+            not isinstance(o, BaseException) and o.status == PublishStatus.TIMEOUT for o in outcomes
+        ):
+            with contextlib.suppress(Exception):
+                await channel.close()