queue-func 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

queue_func/broker/base.py

@@ -3,7 +3,8 @@
 import abc
 import typing as t
 
-from queue_func.serializer.base import BaseSerializer, Message
+from queue_func.serializer.base import Message
+from queue_func.shared_resource.base import Resource
 
 if t.TYPE_CHECKING:
     from queue_func.config.general import Config
@@ -22,33 +23,12 @@ class BaseBroker(t.Generic[ConfigType, Message], metaclass=abc.ABCMeta):
     """
 
     __slots__ = (
-        "_config",
-        "_serializer",
         "_producer",
         "_consumer",
     )
 
-    def __init__(self, config: ConfigType):
-
-        self._config = config
-        self._serializer: type[BaseSerializer[Message]] = self._config.serializer_cls
-        self._producer = self._create_producer()
-        self._consumer = self._create_consumer()
-
-    @abc.abstractmethod
-    def _create_producer(self) -> "BaseProducer[Message]":
-
-        pass
-
-    @abc.abstractmethod
-    def _create_consumer(self) -> "BaseConsumer[Message]":
-
-        pass
-
-    @property
-    def serializer(self) -> type[BaseSerializer[Message]]:
-
-        return self._serializer
+    _producer: "BaseProducer"
+    _consumer: "BaseConsumer"
 
     async def publish(
             self,

queue_func/broker/queue.py

@@ -4,6 +4,7 @@ import asyncio
 import typing as t
 
 from queue_func.serializer.base import Message
+from queue_func.shared_resource.queue import Resource
 from queue_func.broker.base import BaseBroker as _BaseBroker
 from queue_func.producer.queue import Producer
 from queue_func.consumer.queue import Consumer
@@ -15,30 +16,19 @@ if t.TYPE_CHECKING:
 class Broker(_BaseBroker["Config", Message]):
     """内存队列通道。队列由这里创建，再交给 Producer 和 Consumer。"""
 
-    __slots__ = (
-        "_queue",
-    )
-
     def __init__(self, config: "Config"):
-        self._queue: asyncio.Queue[Message] = asyncio.Queue(maxsize=config.maxsize)
-        super().__init__(config)
-
-    def _create_producer(self) -> Producer[Message]:
-
-        return Producer(
-            queue=self._queue,
-            publish_timeout=self._config.publish_timeout,
-            serializer=self._serializer,
+        resource = Resource(
+            serializer_cls=config.serializer_cls,
+            maxsize=config.maxsize)
+        self._producer = Producer(
+            resource=resource,
+            publish_timeout=config.publish_timeout,
         )
-
-    def _create_consumer(self) -> Consumer[Message]:
-
-        return Consumer(
-            queue=self._queue,
-            consume_timeout=self._config.consume_timeout,
-            serializer=self._serializer,
-            coro=self._config.coro,
-            enable_consumer=self._config.enable_consumer,
-            pool_executor_cls=self._config.consumer_pool_executor_cls,
-            pool_maxsize=self._config.consumer_pool_maxsize,
+        self._consumer = Consumer(
+            resource=resource,
+            coro=config.coro,
+            enable_consumer=config.enable_consumer,
+            consume_timeout=config.consume_timeout,
+            pool_executor_cls=config.consumer_pool_executor_cls,
+            pool_maxsize=config.consumer_pool_maxsize,
         )

queue_func/broker/rabbit_aio_pika.py

@@ -0,0 +1,47 @@
+# coding: utf-8
+
+import typing as t
+
+from queue_func.broker.base import BaseBroker as _BaseBroker
+from queue_func.consumer.rabbit_aio_pika import Consumer
+from queue_func.producer.rabbit_aio_pika import Producer
+from queue_func.shared_resource.rabbit_aio_pika import Resource
+
+if t.TYPE_CHECKING:
+    from queue_func.config.rabbit_aio_pika 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,
+            enable_consumer=config.enable_consumer,
+            consume_timeout=config.consume_timeout,
+            pool_executor_cls=config.consumer_pool_executor_cls,
+            pool_maxsize=config.consumer_pool_maxsize,
+        )

queue_func/config/general.py

@@ -11,18 +11,18 @@ from queue_func.pool_executor.base import BasePoolExecutor
 from queue_func.pool_executor.general import PoolExecutor
 
 
-@pyd_dataclass(frozen=True, config=ExtraIgnoreConfig)
+@pyd_dataclass(frozen=True, kw_only=True, config=ExtraIgnoreConfig)
 class Config:
 
-    # consumer
+    ## shared resource
+    serializer_cls: type[BaseSerializer] = Serializer
+
+    ## consumer
     coro: t.Callable[[...], t.Coroutine]
     enable_consumer: bool = True
     consume_timeout: float = 2.0
     consumer_pool_executor_cls: type[BasePoolExecutor] = PoolExecutor
     consumer_pool_maxsize: int = 200
 
-    # broker
-    serializer_cls: type[BaseSerializer] = Serializer
-
-    # producer
+    ## producer
     publish_timeout: float | None = None

queue_func/config/queue.py

@@ -12,6 +12,6 @@ from queue_func.serializer.raw import Serializer
 @pyd_dataclass(frozen=True, config=ExtraIgnoreConfig)
 class Config(_Config):
 
-    # broker
+    ## shared resource
     maxsize: int = 200
     serializer_cls: type[BaseSerializer] = Serializer

queue_func/config/rabbit_aio_pika.py

@@ -0,0 +1,55 @@
+# coding: utf-8
+
+from enum import unique, StrEnum, IntEnum
+
+from pydantic.dataclasses import dataclass as pyd_dataclass
+
+from queue_func.config.utils import ExtraIgnoreConfig
+from queue_func.config.general import Config as _Config
+
+
+@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

queue_func/consumer/base.py

@@ -7,36 +7,43 @@ import contextlib
 import typing as t
 from functools import wraps
 
-from queue_func.serializer.base import BaseSerializer, Message
+from queue_func.serializer.base import Message
+from queue_func.shared_resource.base import Resource
 from queue_func.pool_executor.base import BasePoolExecutor
+from queue_func.utils.hook import OnceHooks
 
 
-class BaseConsumer(t.Generic[Message], metaclass=abc.ABCMeta):
+class BaseConsumer(t.Generic[Message, Resource], metaclass=abc.ABCMeta):
 
     __slots__ = (
-        "_serializer",
+        "_resource",
         "_coro",
         "_enable_consumer",
         "_pool_executor",
-        "_task",
-        "_stop_flag",
+        "_consume_timeout",
+        "_startup",
+        "_shutdown",
+        "_stopped",
         "_logger",
     )
 
     def __init__(
             self,
-            serializer: type[BaseSerializer[Message]],
+            resource: Resource,
             coro: t.Callable[..., t.Coroutine],
             enable_consumer: bool,
+            consume_timeout: float,
             pool_executor_cls: type[BasePoolExecutor],
             pool_maxsize: int):
 
-        self._serializer = serializer
+        self._resource = resource
         self._coro = coro
         self._enable_consumer = enable_consumer
         self._pool_executor = self._init_pool_executor(pool_executor_cls, pool_maxsize)
-        self._task: asyncio.Task | None = None
-        self._stop_flag = False
+        self._consume_timeout = consume_timeout
+        self._stopped = False
+        self._startup = OnceHooks()
+        self._shutdown = OnceHooks()
         self._logger = logging.getLogger(__name__)
 
     def _init_pool_executor(
@@ -50,15 +57,15 @@ class BaseConsumer(t.Generic[Message], metaclass=abc.ABCMeta):
 
     def _run_until_stopped(
             self,
-            coro: t.Callable[..., t.Awaitable],
-            delay: float = 0.02) -> t.Callable[..., t.Coroutine]:
+            fn: t.Callable[..., t.Awaitable],
+            delay: float = 0.05) -> t.Callable[..., t.Coroutine]:
 
-        @wraps(coro)
+        @wraps(fn)
         async def wrapped(*args, **kwargs) -> None:
 
-            while not self._stop_flag:
+            while not self._stopped:
                 try:
-                    await coro(*args, **kwargs)
+                    await fn(*args, **kwargs)
                 except Exception as e:
                     self._logger.exception(e)
                 await asyncio.sleep(delay)
@@ -73,10 +80,10 @@ class BaseConsumer(t.Generic[Message], metaclass=abc.ABCMeta):
     async def consume(self) -> None:
 
         if self._enable_consumer:
-            self._stop_flag = False
-            consume_loop = self._run_until_stopped(self._consume)
-            self._task = asyncio.create_task(consume_loop())
-        return None
+            self._stopped = False
+            self._startup.register(self._run_until_stopped(self._consume))
+            self._register_shutdown()
+            await self._startup.start()
 
     @abc.abstractmethod
     async def _consume(self) -> None:  # pragma: no cover
@@ -85,13 +92,22 @@ class BaseConsumer(t.Generic[Message], metaclass=abc.ABCMeta):
 
     async def stop(self) -> None:
 
-        self._stop_flag = True
+        await self._shutdown.start(prepare=self._mark_stopped)
+
+    def _mark_stopped(self) -> None:
+
+        self._stopped = True
 
     async def join(self) -> None:
 
-        if self._task is not None:
-            with contextlib.suppress(asyncio.CancelledError):
-                await self._task
+        await self._shutdown.wait()
+
+    async def _join_pool_executor(self):
 
         if self._pool_executor is not None:
             await self._pool_executor.join()
+
+    def _register_shutdown(self) -> None:
+
+        self._shutdown.register(self._startup.wait)
+        self._shutdown.register(self._join_pool_executor)

queue_func/consumer/queue.py

@@ -3,53 +3,30 @@
 import asyncio
 import typing as t
 
-from queue_func.serializer.base import BaseSerializer, Message
+from queue_func.serializer.base import Message
+from queue_func.shared_resource.queue import Resource
 from queue_func.pool_executor.base import BasePoolExecutor
 from queue_func.consumer.base import BaseConsumer as _BaseConsumer
 
 
-class Consumer(_BaseConsumer[Message]):
-
-    __slots__ = (
-        "_queue",
-        "_consume_timeout",
-    )
-
-    def __init__(
-            self,
-            queue: asyncio.Queue[Message],
-            consume_timeout: float,
-            serializer: type[BaseSerializer[Message]],
-            coro: t.Callable[..., t.Coroutine],
-            enable_consumer: bool,
-            pool_executor_cls: type[BasePoolExecutor],
-            pool_maxsize: int):
-
-        super().__init__(
-            serializer,
-            coro,
-            enable_consumer,
-            pool_executor_cls,
-            pool_maxsize,
-        )
-        self._queue = queue
-        self._consume_timeout = consume_timeout
+class Consumer(_BaseConsumer[Message, Resource]):
 
     async def _consume(self) -> None:
 
-        while (not self._stop_flag) or (not self._queue.empty()):
+        while (not self._stopped) or (not self._resource.queue.empty()):
             try:
-                packed: Message = await asyncio.wait_for(
-                    self._queue.get(),
+                message: Message = await asyncio.wait_for(
+                    self._resource.queue.get(),
                     timeout=self._consume_timeout,
                 )
             except asyncio.TimeoutError:
                 continue
-            coro_args, coro_kwargs = self._serializer.unpack(packed)
+            coro_args, coro_kwargs = self._resource.serializer.unpack(message)
             await self._dispatch(coro_args, coro_kwargs)
-            self._queue.task_done()
+            self._resource.queue.task_done()
 
-    async def join(self) -> None:
+    def _register_shutdown(self) -> None:
 
-        await super().join()
-        await self._queue.join()
+        self._shutdown.register(self._startup.wait)
+        self._shutdown.register(self._resource.queue.join)
+        self._shutdown.register(self._join_pool_executor)

queue_func/consumer/rabbit_aio_pika.py

@@ -0,0 +1,99 @@
+# coding: utf-8
+
+import asyncio
+import typing as t
+
+from queue_func.consumer.base import BaseConsumer as _BaseConsumer
+from queue_func.pool_executor.base import BasePoolExecutor
+from queue_func.shared_resource.rabbit_aio_pika 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",
+        "_iterator",
+        "_setup_lock",
+    )
+
+    def __init__(
+            self,
+            resource: Resource,
+            prefetch_count: int,
+            prefetch_size: int,
+            global_qos: bool,
+            coro: t.Callable[..., t.Coroutine],
+            enable_consumer: bool,
+            consume_timeout: float,
+            pool_executor_cls: type[BasePoolExecutor],
+            pool_maxsize: int):
+
+        super().__init__(
+            resource=resource,
+            coro=coro,
+            enable_consumer=enable_consumer,
+            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
+        self._setup_lock = asyncio.Lock()
+
+    async def _consume(self) -> None:
+
+        await self._ensure_queue()
+        while not self._stopped:
+            async with self._queue.iterator(timeout=self._consume_timeout) as queue_iter:
+                try:
+                    message: "AbstractIncomingMessage"
+                    async for message in queue_iter:
+                        body: bytes = message.body
+                        coro_args, coro_kwargs = self._resource.serializer.unpack(body)
+                        await self._dispatch(coro_args, coro_kwargs)
+                        await message.ack()
+                except asyncio.TimeoutError:
+                    continue
+
+    async def _ensure_queue(self) -> None:
+
+        if self._queue is not None:
+            return None
+
+        async with self._setup_lock:
+            if self._queue is not None:
+                return None
+            channel = await self._resource.get_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._channel, self._queue = channel, queue
+
+    async def _release(self) -> None:
+
+        channel = self._channel
+        if channel is None:
+            return None
+        self._channel = None
+        self._queue = None
+        await self._resource.close_channel(channel)
+        await self._resource.try_close_connection()
+
+    def _register_shutdown(self) -> None:
+
+        self._shutdown.register(self._startup.wait)
+        self._shutdown.register(self._release)
+        self._shutdown.register(self._join_pool_executor)

queue_func/factory/backend.py

@@ -11,3 +11,4 @@ __all__ = (
 class Backend(StrEnum):
 
     QUEUE = "queue"
+    RABBIT_AIO_PIKA = "rabbit_aio_pika"

queue_func/factory/builders/__init__.py

@@ -3,7 +3,9 @@
 """Import builder modules so they register with :func:`register_backend`."""
 
 from queue_func.factory.builders import queue as queue
+from queue_func.factory.builders import rabbit_aio_pika as rabbit_aio_pika
 
 __all__ = (
     "queue",
+    "rabbit_aio_pika",
 )

queue_func/factory/builders/rabbit_aio_pika.py

@@ -0,0 +1,26 @@
+# coding: utf-8
+
+import typing as t
+
+from queue_func.factory.backend import Backend
+from queue_func.factory.registry import register_backend
+
+if t.TYPE_CHECKING:
+    from queue_func.broker.base 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 queue_func.broker.rabbit_aio_pika import Broker
+        from queue_func.config.rabbit_aio_pika 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))

queue_func/factory/deco.py

@@ -2,6 +2,7 @@
 
 import typing as t
 
+from queue_func.config.rabbit_aio_pika import DeliveryMode, ExchangeName
 from queue_func.factory import builders as _builders  # noqa: F401 — register backends
 from queue_func.factory.backend import Backend
 from queue_func.factory.registry import build_broker
@@ -84,6 +85,35 @@ def queue(
     ...
 
 
+@t.overload
+def queue(
+        backend: t.Literal[Backend.RABBIT_AIO_PIKA],
+        *,
+        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 = ...,
+        serializer_cls: type["BaseSerializer"] = ...,
+        enable_consumer: bool = ...,
+        consume_timeout: float = ...,
+        consumer_pool_executor_cls: type["BasePoolExecutor"] = ...,
+        consumer_pool_maxsize: int = ...,
+        prefetch_count: int = ...,
+        prefetch_size: int = ...,
+        global_qos: bool = ...,
+        exchange_name: ExchangeName | str = ...,
+        routing_key: str = ...,
+        delivery_mode: DeliveryMode = ...,
+        publish_timeout: float | None = ...,
+) -> t.Callable[[t.Callable[_P, t.Coroutine[t.Any, t.Any, _R]]], QueueFunc[_P, _R]]:
+
+    ...
+
+
 def queue(
         backend: Backend,
         **config_kwargs: t.Any) -> t.Callable[[t.Callable[_P, t.Coroutine[t.Any, t.Any, _R]]], QueueFunc[_P, _R]]:

queue_func/producer/base.py

@@ -3,23 +3,33 @@
 import abc
 import typing as t
 
-from queue_func.serializer.base import BaseSerializer, Message
+from queue_func.shared_resource.base import Resource
+from queue_func.serializer.base import Message
+from queue_func.utils.hook import OnceHooks
+from queue_func.utils.waiter import IdleWaiter
 
 
-class BaseProducer(t.Generic[Message], metaclass=abc.ABCMeta):
+class BaseProducer(t.Generic[Message, Resource], metaclass=abc.ABCMeta):
 
     __slots__ = (
-        "_serializer",
+        "_resource",
         "_publish_timeout",
+        "_stopped",
+        "_shutdown",
+        "_pub_waiter",
     )
 
     def __init__(
             self,
-            serializer: type[BaseSerializer[Message]],
+            resource: Resource,
             publish_timeout: float | None):
 
-        self._serializer = serializer
+        self._resource = resource
         self._publish_timeout = publish_timeout
+        self._stopped = False
+        self._shutdown = OnceHooks()
+        self._pub_waiter = IdleWaiter()
+        self._register_shutdown()
 
     async def publish(
             self,
@@ -28,17 +38,32 @@ class BaseProducer(t.Generic[Message], metaclass=abc.ABCMeta):
             *,
             timeout: float | None = None) -> None:
 
+
+        if self._stopped:
+            raise RuntimeError("publish called on already stopped")
+
         if timeout is None:
             timeout = self._publish_timeout
-        packed = self._serializer.pack(coro_args, coro_kwargs)
-        await self._publish(packed, timeout=timeout)
+        packed = self._resource.serializer.pack(coro_args, coro_kwargs)
+        async with self._pub_waiter.hold():
+            await self._publish(packed, timeout=timeout)
 
     @abc.abstractmethod
-    async def _publish(self, data: Message, *, timeout: float | None) -> None:
+    async def _publish(self, msg: Message, *, timeout: float | None) -> None:
         pass
 
     async def stop(self) -> None:
-        pass
+
+        await self._shutdown.start(prepare=self._mark_stopped)
+
+    def _mark_stopped(self) -> None:
+
+        self._stopped = True
 
     async def join(self) -> None:
-        pass
+
+        await self._shutdown.wait()
+
+    def _register_shutdown(self):
+
+        self._shutdown.register(self._pub_waiter.wait)