py-messor 0.2.4__py3-none-any.whl

messor/__init__.py

@@ -0,0 +1,7 @@
+# coding: utf-8
+
+from messor.errors import PublishTimeoutError
+
+__all__ = (
+    "PublishTimeoutError",
+)

messor/backends/__init__.py

@@ -0,0 +1 @@
+# coding: utf-8

messor/backends/queue/__init__.py

@@ -0,0 +1 @@
+# coding: utf-8

messor/backends/queue/broker.py

@@ -0,0 +1,35 @@
+# coding: utf-8
+
+import asyncio
+import typing as t
+
+from messor.serializer.base import Message
+from messor.backends.queue.resource import Resource
+from messor.core.broker import BaseBroker as _BaseBroker
+from messor.backends.queue.producer import Producer
+from messor.backends.queue.consumer import Consumer
+
+if t.TYPE_CHECKING:
+    from messor.backends.queue.config import Config
+
+
+class Broker(_BaseBroker["Config", Message]):
+    """内存队列通道。队列由这里创建，再交给 Producer 和 Consumer。"""
+
+    IS_MEMORY = True
+
+    def __init__(self, config: "Config"):
+        resource = Resource(
+            serializer_cls=config.serializer_cls,
+            maxsize=config.maxsize)
+        self._producer = Producer(
+            resource=resource,
+            publish_timeout=config.publish_timeout,
+        )
+        self._consumer = Consumer(
+            resource=resource,
+            coro=config.coro,
+            consume_timeout=config.consume_timeout,
+            pool_executor_cls=config.consumer_pool_executor_cls,
+            pool_maxsize=config.consumer_pool_maxsize,
+        )

messor/backends/queue/build.py

@@ -0,0 +1,20 @@
+# coding: utf-8
+
+import typing as t
+
+from messor.factory.backend import Backend
+from messor.factory.registry import register_backend
+
+if t.TYPE_CHECKING:
+    from messor.core.broker import BaseBroker
+
+
+@register_backend(Backend.QUEUE)
+def build_queue_broker(
+        coro: t.Callable[..., t.Coroutine],
+        **kwargs: t.Any) -> "BaseBroker[t.Any, t.Any]":
+
+    from messor.backends.queue.broker import Broker
+    from messor.backends.queue.config import Config
+
+    return Broker(Config(coro=coro, **kwargs))

messor/backends/queue/config.py

@@ -0,0 +1,17 @@
+# coding: utf-8
+
+import typing as t
+from pydantic.dataclasses import dataclass as pyd_dataclass
+
+from messor.core.config import Config as _Config
+from messor.core.config import ExtraIgnoreConfig
+from messor.serializer.base import BaseSerializer
+from messor.serializer.raw import Serializer
+
+
+@pyd_dataclass(frozen=True, config=ExtraIgnoreConfig)
+class Config(_Config):
+
+    ## shared resource
+    maxsize: int = 200
+    serializer_cls: type[BaseSerializer] = Serializer

messor/backends/queue/consumer.py

@@ -0,0 +1,47 @@
+# coding: utf-8
+
+import asyncio
+
+from messor.serializer.base import Message
+from messor.backends.queue.resource import Resource
+from messor.core._inflight import dec
+from messor.core.consumer import BaseConsumer as _BaseConsumer
+
+
+class Consumer(_BaseConsumer[Message, Resource]):
+
+    __slots__ = ()
+
+    def __init__(self, *args, **kwargs):
+
+        super().__init__(*args, **kwargs)
+        # coro 执行完成后 -1
+        self._on_coro_done.register(dec)
+    
+    async def _consume_once(self) -> Message | None:
+
+        queue = self._resource.queue
+        try:
+            message: Message = await asyncio.wait_for(
+                queue.get(),
+                timeout=self._consume_timeout,
+            )
+        except asyncio.TimeoutError:
+            return None
+        else:
+            await self._dispatch(message)
+            queue.task_done()
+            return message
+
+    async def _drain(self) -> None:
+
+        queue = self._resource.queue
+        while not queue.empty():
+            message: Message = queue.get_nowait()
+            await self._dispatch(message)
+            queue.task_done()
+
+    async def _on_shutdown(self) -> None:
+
+        await self._resource.queue.join()
+        await self._pool_executor.obj.join()

messor/backends/queue/producer.py

@@ -0,0 +1,28 @@
+# coding: utf-8
+
+import asyncio
+
+from messor.errors import PublishTimeoutError
+from messor.serializer.base import Message
+from messor.core._inflight import inc
+from messor.core.producer import BaseProducer
+from messor.backends.queue.resource import Resource
+
+
+class Producer(BaseProducer[Message, Resource]):
+
+    def __init__(self, *args, **kwargs):
+
+        super().__init__(*args, **kwargs)
+        # push 成功后 +1
+        self._on_publish_success.register(inc)
+
+    async def _publish(self, msg: Message, *, timeout: float | None) -> None:
+
+        if timeout:
+            try:
+                await asyncio.wait_for(self._resource.queue.put(msg), timeout)
+            except asyncio.TimeoutError as exc:
+                raise PublishTimeoutError("publish timed out") from exc
+            return None
+        await self._resource.queue.put(msg)

messor/backends/queue/resource.py

@@ -0,0 +1,23 @@
+# coding: utf-8
+
+import asyncio
+
+from messor.serializer.base import Message, BaseSerializer
+from messor.core.resource import BaseResource as _Base
+
+
+class Resource(_Base[Message]):
+
+    __slots__ = (
+        "_queue",
+    )
+
+    def __init__(self, serializer_cls: type[BaseSerializer], maxsize: int) -> None:
+
+        super().__init__(serializer_cls)
+        self._queue: asyncio.Queue[Message] = asyncio.Queue(maxsize=maxsize)
+
+    @property
+    def queue(self) -> asyncio.Queue[Message]:
+
+        return self._queue

messor/backends/rabbit_aio_pika/__init__.py

@@ -0,0 +1 @@
+# coding: utf-8

messor/backends/rabbit_aio_pika/broker.py

@@ -0,0 +1,46 @@
+# coding: utf-8
+
+import typing as t
+
+from messor.core.broker import BaseBroker as _BaseBroker
+from messor.backends.rabbit_aio_pika.consumer import Consumer
+from messor.backends.rabbit_aio_pika.producer import Producer
+from messor.backends.rabbit_aio_pika.resource import Resource
+
+if t.TYPE_CHECKING:
+    from messor.backends.rabbit_aio_pika.config import Config
+
+
+class Broker(_BaseBroker["Config", bytes]):
+    """RabbitMQ 通道。一条连接上分开发布和消费，再交给 Producer 和 Consumer。"""
+
+    def __init__(self, config: "Config"):
+
+        resource = Resource(
+            serializer_cls=config.serializer_cls,
+            url=config.url,
+            queue_name=config.queue_name,
+            enable_declare_queue=config.enable_declare_queue,
+            queue_durable=config.queue_durable,
+            queue_exclusive=config.queue_exclusive,
+            queue_passive=config.queue_passive,
+            queue_auto_delete=config.queue_auto_delete,
+            queue_arguments=config.queue_arguments,
+        )
+        self._producer = Producer(
+            resource=resource,
+            publish_timeout=config.publish_timeout,
+            exchange_name=config.exchange_name,
+            routing_key=config.routing_key,
+            delivery_mode=config.delivery_mode,
+        )
+        self._consumer = Consumer(
+            resource=resource,
+            prefetch_count=config.prefetch_count,
+            prefetch_size=config.prefetch_size,
+            global_qos=config.global_qos,
+            coro=config.coro,
+            consume_timeout=config.consume_timeout,
+            pool_executor_cls=config.consumer_pool_executor_cls,
+            pool_maxsize=config.consumer_pool_maxsize,
+        )

messor/backends/rabbit_aio_pika/build.py

@@ -0,0 +1,26 @@
+# coding: utf-8
+
+import typing as t
+
+from messor.factory.backend import Backend
+from messor.factory.registry import register_backend
+
+if t.TYPE_CHECKING:
+    from messor.core.broker import BaseBroker
+
+
+@register_backend(Backend.RABBIT_AIO_PIKA)
+def build_rabbit_aio_pika_broker(
+        coro: t.Callable[..., t.Coroutine],
+        **kwargs: t.Any) -> "BaseBroker[t.Any, t.Any]":
+
+    try:
+        from messor.backends.rabbit_aio_pika.broker import Broker
+        from messor.backends.rabbit_aio_pika.config import Config
+    except ImportError as exc:
+        raise ImportError(
+            "Backend.RABBIT_AIO_PIKA requires aio-pika. "
+            "Install it with the rabbit extra.",
+        ) from exc
+
+    return Broker(Config(coro=coro, **kwargs))

messor/backends/rabbit_aio_pika/config.py

@@ -0,0 +1,55 @@
+# coding: utf-8
+
+from enum import unique, StrEnum, IntEnum
+
+from pydantic.dataclasses import dataclass as pyd_dataclass
+
+from messor.core.config import Config as _Config
+from messor.core.config import ExtraIgnoreConfig
+
+
+@unique
+class ExchangeName(StrEnum):
+
+    DEFAULT = ""
+    DIRECT = "amq.direct"
+    TOPIC = "amq.topic"
+    FANOUT = "amq.fanout"
+    HEADERS = "amq.headers"
+
+
+@unique
+class DeliveryMode(IntEnum):
+
+    NOT_PERSISTENT = 1
+    PERSISTENT = 2
+
+
+@pyd_dataclass(frozen=True, config=ExtraIgnoreConfig, kw_only=True)
+class Config(_Config):
+
+    ## shared resource
+    # amqp://user:password@host/vhost?heartbeat=60&x=y
+    url: str
+    queue_name: str
+    enable_declare_queue: bool = False
+    # broker 重启后队列仍在；消息是否保留看 delivery_mode
+    queue_durable: bool = False
+    # 仅当前连接可使用，连接断开后队列删除
+    queue_exclusive: bool = False
+    # 不创建队列，只确认同名队列已存在，不存在则失败
+    queue_passive: bool = False
+    # 最后一个消费者取消或断开后删除队列
+    queue_auto_delete: bool = False
+    # 声明时的附加参数，例如 x-message-ttl、x-dead-letter-exchange
+    queue_arguments: dict | None = None
+
+    ## producer
+    exchange_name: ExchangeName | str = ExchangeName.DEFAULT
+    routing_key: str = ""
+    delivery_mode: DeliveryMode = DeliveryMode.PERSISTENT
+
+    ## consumer
+    prefetch_count: int = 0
+    prefetch_size: int = 0
+    global_qos: bool = False

messor/backends/rabbit_aio_pika/consumer.py

@@ -0,0 +1,77 @@
+# coding: utf-8
+
+import asyncio
+import typing as t
+from functools import partial
+
+from messor.core.consumer import BaseConsumer as _BaseConsumer
+from messor.pool_executor.base import BasePoolExecutor
+from messor.backends.rabbit_aio_pika.resource import Resource
+
+if t.TYPE_CHECKING:
+    from aio_pika.abc import AbstractIncomingMessage, AbstractQueue
+
+
+class Consumer(_BaseConsumer[bytes, Resource]):
+
+    __slots__ = (
+        "_prefetch_count",
+        "_prefetch_size",
+        "_global_qos",
+        # "_channel",
+        "_queue",
+    )
+
+    def __init__(
+            self,
+            resource: Resource,
+            prefetch_count: int,
+            prefetch_size: int,
+            global_qos: bool,
+            coro: t.Callable[..., t.Coroutine],
+            consume_timeout: float,
+            pool_executor_cls: type[BasePoolExecutor],
+            pool_maxsize: int):
+
+        super().__init__(
+            resource=resource,
+            coro=coro,
+            consume_timeout=consume_timeout,
+            pool_executor_cls=pool_executor_cls,
+            pool_maxsize=pool_maxsize,
+        )
+        self._prefetch_count = prefetch_count
+        self._prefetch_size = prefetch_size
+        self._global_qos = global_qos
+        # self._channel: "AbstractChannel | None" = None
+        self._queue: "AbstractQueue | None" = None
+
+    async def _consume_once(self) -> None:
+
+        queue = await self._ensure_queue()
+        async with queue.iterator(timeout=self._consume_timeout) as queue_iter:
+            message: "AbstractIncomingMessage"
+            try:
+                async for message in queue_iter:
+                    body: bytes = message.body
+                    await self._dispatch(body)
+                    await message.ack()
+            except asyncio.TimeoutError:
+                pass
+
+    async def _ensure_queue(self) -> "AbstractQueue":
+
+        if self._queue is not None:
+            return self._queue
+
+        self._on_cleanup.register(self._resource.try_close_connection)
+        channel = await self._resource.get_channel()
+        self._on_cleanup.register(partial(self._resource.close_channel, channel))
+        await channel.set_qos(
+            prefetch_count=self._prefetch_count,
+            prefetch_size=self._prefetch_size,
+            global_=self._global_qos,
+        )
+        queue = await self._resource.declare_queue(channel)
+        self._queue = queue
+        return queue

messor/backends/rabbit_aio_pika/producer.py

@@ -0,0 +1,68 @@
+# coding: utf-8
+
+import asyncio
+import typing as t
+from functools import partial
+
+import aio_pika
+
+from messor.backends.rabbit_aio_pika.config import DeliveryMode, ExchangeName
+from messor.errors import PublishTimeoutError
+from messor.core.producer import BaseProducer
+from messor.backends.rabbit_aio_pika.resource import Resource
+
+if t.TYPE_CHECKING:
+    from aio_pika.abc import AbstractExchange
+
+
+class Producer(BaseProducer[bytes, Resource]):
+
+    __slots__ = (
+        "_exchange_name",
+        "_exchange_type",
+        "_routing_key",
+        "_delivery_mode",
+        "_exchange",
+    )
+
+    def __init__(
+            self,
+            resource: Resource,
+            publish_timeout: float | None,
+            exchange_name: ExchangeName | str,
+            routing_key: str,
+            delivery_mode: DeliveryMode):
+        super().__init__(resource, publish_timeout)
+        self._exchange_name = exchange_name
+        self._exchange: "AbstractExchange | None" = None
+        self._routing_key = routing_key
+        self._delivery_mode = delivery_mode
+
+    async def _publish(self, msg: bytes, *, timeout: float | None) -> None:
+
+        routing_key = self._routing_key or self._resource.queue_name
+        message = aio_pika.Message(msg, delivery_mode=self._delivery_mode)
+
+        await self._ensure_exchange()
+        try:
+            await self._exchange.publish(message, routing_key=routing_key, timeout=timeout)
+        except TimeoutError as exc:
+            raise PublishTimeoutError("publish timed out") from exc
+
+    async def _ensure_exchange(self) -> "AbstractExchange":
+
+        if self._exchange is not None:
+            return self._exchange
+
+        async with self._lock:
+            if self._exchange is None:
+                self._on_cleanup.register(self._resource.try_close_connection)
+                channel = await self._resource.get_channel()
+                self._on_cleanup.register(partial(self._resource.close_channel, channel))
+                await self._resource.declare_queue(channel)
+                if self._exchange_name == ExchangeName.DEFAULT:
+                    exchange = channel.default_exchange
+                else:
+                    exchange = await channel.get_exchange(self._exchange_name)
+                self._exchange = exchange
+                return exchange

messor/backends/rabbit_aio_pika/resource.py

@@ -0,0 +1,111 @@
+# coding: utf-8
+
+import asyncio
+import typing as t
+
+import aio_pika
+
+from messor.utils.deco import before
+from messor.core.resource import BaseResource as _Base
+from messor.serializer.base import Message, BaseSerializer
+
+if t.TYPE_CHECKING:
+    from aio_pika.abc import AbstractConnection, AbstractChannel, AbstractQueue
+
+
+class Resource(_Base[Message]):
+
+    __slots__ = (
+        "_queue_name",
+        "_url",
+        "_conn",
+        "_lock",
+        "_active_channel_num",
+        "_enable_declare_queue",
+        "_queue_durable",
+        "_queue_exclusive",
+        "_queue_passive",
+        "_queue_auto_delete",
+        "_queue_arguments",
+    )
+
+    def __init__(
+            self,
+            serializer_cls: type[BaseSerializer],
+            url: str,
+            queue_name: str,
+            enable_declare_queue: bool,
+            queue_durable: bool,
+            queue_exclusive: bool,
+            queue_passive: bool,
+            queue_auto_delete: bool,
+            queue_arguments: dict[str, t.Any] | None):
+
+        super().__init__(serializer_cls)
+        self._queue_name = queue_name
+        self._url = url
+        self._conn: "AbstractConnection | None" = None
+        self._lock = asyncio.Lock()
+        self._active_channel_num = 0
+        self._enable_declare_queue = enable_declare_queue
+        self._queue_durable = queue_durable
+        self._queue_exclusive = queue_exclusive
+        self._queue_passive = queue_passive
+        self._queue_auto_delete = queue_auto_delete
+        self._queue_arguments = queue_arguments
+
+    @property
+    def queue_name(self) -> str:
+
+        return self._queue_name
+
+    async def declare_queue(self, channel: "AbstractChannel") -> "AbstractQueue":
+
+        if self._enable_declare_queue:
+            return await channel.declare_queue(
+                self._queue_name,
+                durable=self._queue_durable,
+                exclusive=self._queue_exclusive,
+                passive=self._queue_passive,
+                auto_delete=self._queue_auto_delete,
+                arguments=self._queue_arguments,
+            )
+        return await channel.declare_queue(self._queue_name, passive=True)
+
+    async def _ensure_conn(self) -> None:
+
+        if self._conn is not None:
+            return None
+
+        async with self._lock:
+            if self._conn is None:
+                self._conn = await aio_pika.connect_robust(self._url)
+
+    @property
+    def _required_conn(self) -> "AbstractConnection":
+
+        if self._conn is None:
+            raise RuntimeError("RabbitMQ connection not established")
+        return self._conn
+
+    @before(_ensure_conn)
+    async def get_channel(self) -> "AbstractChannel":
+
+        async with self._lock:
+            channel = await self._required_conn.channel()
+            self._active_channel_num += 1
+            return channel
+
+    async def close_channel(self, channel: "AbstractChannel") -> None:
+
+        async with self._lock:
+            await channel.close()
+            self._active_channel_num -= 1
+
+    async def try_close_connection(self) -> None:
+
+        async with self._lock:
+            if self._active_channel_num > 0 or self._conn is None:
+                return None
+            await self._required_conn.close()
+            self._conn = None

messor/core/__init__.py

@@ -0,0 +1 @@
+# coding: utf-8

messor/core/_inflight.py

@@ -0,0 +1,74 @@
+# coding: utf-8
+
+import asyncio
+
+__all__ = (
+    "inc",
+    "dec",
+    "wait_drained",
+)
+
+
+class _Inflight:
+    """进程级全局在途消息计数（内存队列专用）。
+
+    push 成功后 +1，coro 执行完成后 -1；
+    计数归 0 表示无在途消息，可通过 wait_drained 等待该状态。
+    """
+
+    __slots__ = (
+        "_count",
+        "_event",
+    )
+
+    def __init__(self) -> None:
+
+        self._count = 0
+        self._event = asyncio.Event()
+        self._event.set()
+
+    def inc(self, n: int = 1) -> None:
+
+        self._count += n
+        # 有在途消息，阻塞 wait_drained
+        self._event.clear()
+
+    def dec(self, n: int = 1) -> None:
+
+        self._count = max(0, self._count - n)
+        if self._count == 0:
+            # 无在途消息，放行 wait_drained
+            self._event.set()
+
+    async def wait_drained(self) -> None:
+
+        while self._count > 0:
+            await self._event.wait()
+
+    def reset(self) -> None:
+        """清空计数（测试隔离用，不对外暴露）。"""
+
+        self._count = 0
+        self._event.set()
+
+
+# 全局单例：内存队列 producer/consumer 共享，不对外暴露
+_INFLIGHT = _Inflight()
+
+
+def inc(n: int = 1) -> None:
+    """push 成功后 +n。"""
+
+    _INFLIGHT.inc(n)
+
+
+def dec(n: int = 1) -> None:
+    """coro 执行完成后 -n（默认 -1，可传批量值）。"""
+
+    _INFLIGHT.dec(n)
+
+
+async def wait_drained() -> None:
+    """等待在途消息全部处理完成。"""
+
+    await _INFLIGHT.wait_drained()