binomic 0.1.0__py3-none-any.whl

binomic/base/__init__.py

@@ -0,0 +1,13 @@
+from .exceptions import (
+    BinomicError,
+    DeserializationError,
+    SerializationError,
+)
+from .schemas import BaseStruct
+
+__all__ = (
+    "BaseStruct",
+    "BinomicError",
+    "DeserializationError",
+    "SerializationError",
+)

binomic/base/constants.py

@@ -0,0 +1,3 @@
+from typing import Final
+
+APP_NAME: Final[str] = "binomic"

binomic/base/exceptions.py

@@ -0,0 +1,17 @@
+__all__ = (
+    "BinomicError",
+    "DeserializationError",
+    "SerializationError",
+)
+
+
+class BinomicError(Exception):
+    """Base class for binomic exceptions."""
+
+
+class SerializationError(BinomicError):
+    """Exception raised for errors during serialization."""
+
+
+class DeserializationError(BinomicError):
+    """Exception raised for errors during deserialization."""

binomic/base/schemas.py

@@ -0,0 +1,39 @@
+from typing import Self
+
+from msgspec import Struct, json
+
+from .exceptions import DeserializationError, SerializationError
+
+__all__ = ("BaseStruct",)
+
+
+class BaseStruct(Struct):
+    """Base data structure."""
+
+    def to_json(self) -> str:
+        """Convert the data structure to a JSON string."""
+
+        try:
+            return json.encode(self).decode("utf-8")
+
+        except Exception as exc:
+            raise SerializationError(
+                f"Failed to serialize data structure: {exc}",
+            ) from exc
+
+    @classmethod
+    def from_json(cls, json_str: str) -> Self:
+        """Create a data structure from a JSON string."""
+
+        try:
+            return json.decode(json_str, type=cls)
+
+        except Exception as exc:
+            raise DeserializationError(
+                f"Failed to deserialize data structure: {exc}",
+            ) from exc
+
+    def __repr__(self) -> str:
+        """Return a string representation of the data structure."""
+
+        return self.to_json()

binomic/broker/__init__.py

@@ -0,0 +1,11 @@
+from .handlers import BrokerFactory
+from .protocols import Broker
+from .redis import AsyncredisBroker
+from .types import Entry
+
+__all__ = (
+    "AsyncredisBroker",
+    "Broker",
+    "BrokerFactory",
+    "Entry",
+)

binomic/broker/handlers.py

@@ -0,0 +1,33 @@
+from typing import TYPE_CHECKING
+
+from .redis import AsyncredisBroker
+
+if TYPE_CHECKING:
+    from .protocols import Broker
+
+__all__ = ("BrokerFactory",)
+
+
+class BrokerFactory:
+    """A factory class for creating broker instances."""
+
+    def __init__(self, dsn: str, queues: list[str]) -> None:
+
+        self._dsn = dsn
+        self._queues = queues
+
+    def __call__(self) -> "Broker":
+
+        match self._dsn.split(":")[0]:
+            case "redis":
+                return AsyncredisBroker(
+                    dsn=self._dsn,
+                    queues=self._queues,
+                    decode_responses=True,
+                )
+
+            case "amqp":
+                raise NotImplementedError("AMQP is not supported yet.")
+
+            case _:
+                raise ValueError("Unsupported broker type.")

binomic/broker/protocols.py

@@ -0,0 +1,27 @@
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from uuid import UUID
+
+    from binomic.message import Message
+
+    from .types import Entry
+
+
+__all__ = ("Broker",)
+
+
+class Broker(Protocol):
+    """Broker protocol class."""
+
+    async def initialize(self) -> None: ...
+
+    async def enqueue(self, msg: "Message") -> "UUID": ...
+
+    async def acquire(self, consumer: str, *, count: int) -> list["Entry"]: ...
+
+    async def ack(self, entry: "Entry") -> int: ...
+
+    async def reclaim(self, consumer: str, *, min_idle_ms: int, count: int) -> int: ...
+
+    async def aclose(self) -> None: ...

binomic/broker/redis.py

@@ -0,0 +1,165 @@
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Final, cast
+
+from redis.asyncio import BlockingConnectionPool
+from redis.asyncio import Redis as AsyncRedis
+from redis.exceptions import ResponseError
+
+from binomic.base.constants import APP_NAME
+from binomic.message import Message
+
+from .protocols import Broker
+from .types import Entry
+
+if TYPE_CHECKING:
+    from uuid import UUID
+
+    from redis.asyncio import ConnectionPool
+
+    from .types import Fields
+
+
+__all__ = ("AsyncredisBroker",)
+
+
+GROUP_NAMESPACE: Final[str] = APP_NAME
+
+
+class AsyncredisBroker(Broker):
+    """Broker based on async redis."""
+
+    def __init__(
+        self,
+        dsn: str,
+        queues: Sequence[str],
+        **config: Any,
+    ) -> None:
+
+        self._dsn = dsn
+        self._queues = queues
+        self._config = config
+        self._group = GROUP_NAMESPACE
+
+        self._connpool: ConnectionPool | None = None
+
+    def get_stream_key(self, queue: str) -> str:
+        """Namespace a queue name into its Redis stream key."""
+
+        return f"{self._group}:{queue}"
+
+    @property
+    def client(self) -> "AsyncRedis":
+
+        if self._connpool is None:
+            self._connpool = BlockingConnectionPool.from_url(
+                url=self._dsn,
+                **self._config,
+            )
+
+        return AsyncRedis.from_pool(self._connpool)
+
+    async def initialize(self) -> None:
+
+        for queue in self._queues:
+            try:
+                await self.client.xgroup_create(
+                    self.get_stream_key(queue),
+                    groupname=self._group,
+                    id="0",
+                    mkstream=True,
+                )
+
+            except ResponseError as err:
+                if "BUSYGROUP" not in str(err):
+                    raise
+
+    async def enqueue(self, msg: "Message") -> "UUID":
+
+        await self.client.xadd(
+            name=self.get_stream_key(msg.queue),
+            fields={
+                "id": str(msg.id),
+                "message": msg.to_json(),
+            },
+        )
+
+        return msg.id
+
+    async def acquire(self, consumer: str, *, count: int) -> list["Entry"]:
+
+        result = await self.client.xreadgroup(
+            groupname=self._group,
+            consumername=consumer,
+            streams={self.get_stream_key(q): ">" for q in self._queues},
+            count=count,
+        )
+
+        entries: list[Entry] = []
+        if not result:
+            return entries
+
+        queue_by_key = {self.get_stream_key(q): q for q in self._queues}
+
+        for key, msgs in cast(
+            "list[tuple[str, list[tuple[str, Fields]]]]",
+            result.items() if isinstance(result, dict) else result,
+        ):
+            queue = queue_by_key[key]
+            for msg_id, fields in msgs:
+                entries.append(Entry(queue, msg_id, fields))
+
+        return entries
+
+    async def ack(self, entry: "Entry") -> int:
+
+        queue, msg_id, _ = entry
+        return await self.client.xack(
+            self.get_stream_key(queue),
+            self._group,
+            msg_id,
+        )
+
+    async def reclaim(
+        self,
+        consumer: str,
+        *,
+        min_idle_ms: int,
+        count: int,
+    ) -> int:
+
+        reclaimed = 0
+        for queue in self._queues:
+            try:
+                result = await self.client.xautoclaim(
+                    name=self.get_stream_key(queue),
+                    groupname=self._group,
+                    consumername=consumer,
+                    min_idle_time=min_idle_ms,
+                    start_id="0-0",
+                    count=count,
+                )
+
+            except ResponseError as err:
+                if "NOGROUP" in str(err):
+                    continue
+
+                raise
+
+            msgs = result[1] if result else []
+            for msg_id, fields in cast(
+                "list[tuple[str, Fields]]",
+                msgs,
+            ):
+                msg = Message.from_json(fields.get("message"))
+                await self.enqueue(msg)
+                await self.ack(Entry(queue, msg_id, fields))
+
+                reclaimed += 1
+
+        return reclaimed
+
+    async def aclose(self) -> None:
+
+        if self._connpool is not None:
+            await self._connpool.disconnect()
+            self._connpool = None

binomic/broker/types.py

@@ -0,0 +1,21 @@
+from typing import NamedTuple, TypedDict
+
+__all__ = (
+    "Entry",
+    "Fields",
+)
+
+
+class Fields(TypedDict):
+    """Fields in the Entry."""
+
+    id: str
+    message: str
+
+
+class Entry(NamedTuple):
+    """Entry in the broker."""
+
+    queue: str
+    msg_id: str
+    fields: Fields

binomic/client.py

@@ -0,0 +1,114 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Self
+
+import anyio
+from anyio import AsyncContextManagerMixin
+
+from binomic.broker import Broker, BrokerFactory
+from binomic.worker import Master, MasterPolicy, WorkerPolicy
+
+if TYPE_CHECKING:
+    from binomic.config import BinomicConfig
+    from binomic.message import Message
+
+
+__all__ = (
+    "Binomic",
+    "BinomicFactory",
+)
+
+
+class Binomic(AsyncContextManagerMixin):
+    """Binomic working client.
+
+    Entering the async context starts the master process group and keeps
+    it running until the context exits.
+    """
+
+    def __init__(
+        self,
+        broker_dsn: str,
+        redis_dsn: str,
+        module_name: str,
+        config: "BinomicConfig",
+    ) -> None:
+
+        self._broker_dsn = broker_dsn
+        self._redis_dsn = redis_dsn
+        self._module_name = module_name
+        self._config = config
+
+        self._broker: Broker | None = None
+
+    async def submit(self, msg: "Message") -> None:
+        """Enqueue a message to the broker stream.
+
+        The broker is created lazily on the first submit, so constructing
+        the client performs no IO.
+        """
+
+        if self._broker is None:
+            self._broker = BrokerFactory(
+                self._broker_dsn,
+                self._config.queues,
+            )()
+
+        await self._broker.enqueue(msg)
+
+    @asynccontextmanager
+    async def __asynccontextmanager__(self) -> AsyncGenerator[Self, None]:
+
+        master = Master(
+            broker_dsn=self._broker_dsn,
+            redis_dsn=self._redis_dsn,
+            module_name=self._module_name,
+            policy=MasterPolicy(
+                workers=self._config.workers,
+                worker=WorkerPolicy(
+                    queues=self._config.queues,
+                    concurrency=self._config.concurrency,
+                ),
+            ),
+        )
+
+        try:
+            async with anyio.create_task_group() as tg:
+                tg.start_soon(master.arun)
+                yield self
+                tg.cancel_scope.cancel()
+
+        finally:
+            await master.aclose()
+
+
+class BinomicFactory:
+    """Binomic working client factory."""
+
+    def __init__(
+        self,
+        broker_dsn: str,
+        redis_dsn: str,
+        module_name: str,
+        config: "BinomicConfig",
+    ) -> None:
+
+        self._broker_dsn = broker_dsn
+        self._redis_dsn = redis_dsn
+        self._module_name = module_name
+        self._config = config
+
+        self._binomic: Binomic | None = None
+
+    def create(self) -> "Binomic":
+        """Return the client, creating it on the first call."""
+
+        if self._binomic is None:
+            self._binomic = Binomic(
+                broker_dsn=self._broker_dsn,
+                redis_dsn=self._redis_dsn,
+                module_name=self._module_name,
+                config=self._config,
+            )
+
+        return self._binomic

binomic/config.py

@@ -0,0 +1,16 @@
+from binomic.base import BaseStruct
+
+__all__ = ("BinomicConfig",)
+
+
+class BinomicConfig(BaseStruct):
+    """Binomic configuration.
+
+    ``workers`` is the number of worker subprocesses the master spawns and
+    ``concurrency`` the number of tasks each worker executes concurrently;
+    every stream in ``queues`` is consumed by all workers.
+    """
+
+    queues: list[str]
+    workers: int = 1
+    concurrency: int = 5

binomic/message/__init__.py

@@ -0,0 +1,3 @@
+from .schemas import Message
+
+__all__ = ("Message",)

binomic/message/schemas.py

@@ -0,0 +1,54 @@
+from typing import Any
+from uuid import UUID
+
+from msgspec import field, json
+from uuid_utils.compat import uuid7
+
+from binomic.base import BaseStruct
+
+__all__ = ("Message",)
+
+
+class Message(BaseStruct):
+    """Message for binomic.
+
+    ``name`` is the registered task name, ``queue`` the target stream and
+    ``enqueued_at`` the submission timestamp; ``args`` and ``kwargs`` carry
+    the task call arguments.
+    """
+
+    name: str
+    queue: str
+    enqueued_at: float
+
+    id: "UUID" = field(default_factory=uuid7)
+    args: list[Any] = field(default_factory=list)
+    kwargs: dict[str, Any] = field(default_factory=dict)
+
+    def to_json(self) -> str:
+
+        return json.encode(
+            {
+                "id": str(self.id),
+                "name": self.name,
+                "queue": self.queue,
+                "enqueued_at": self.enqueued_at,
+                "args": self.args,
+                "kwargs": self.kwargs,
+            },
+            order="sorted",
+        ).decode("utf-8")
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "Message":
+
+        data = json.decode(json_str)
+
+        return cls(
+            id=UUID(data["id"]),
+            name=data["name"],
+            queue=data["queue"],
+            enqueued_at=data["enqueued_at"],
+            args=data["args"],
+            kwargs=data["kwargs"],
+        )

binomic/plugin/litestar.py

@@ -0,0 +1,100 @@
+from typing import TYPE_CHECKING, ClassVar, cast
+
+from binomic.client import Binomic, BinomicFactory
+
+if TYPE_CHECKING:
+    from binomic.config import BinomicConfig
+
+try:
+    import litestar  # noqa: F401
+    from litestar.di import Provide
+    from litestar.plugins import InitPluginProtocol
+
+    if TYPE_CHECKING:
+        from litestar.config.app import AppConfig
+        from litestar.datastructures import State
+
+except ImportError:
+    raise ImportError(  # noqa: B904
+        "litestar is not installed. Please install it using pip install litestar."
+    )
+
+
+__all__ = ("BinomicPlugin",)
+
+
+class BinomicPlugin(InitPluginProtocol):
+    """Binomic plugin for litestar."""
+
+    _binomic_factory_state_key: ClassVar[str] = "binomic_factory"
+
+    def __init__(
+        self,
+        app_name: str,
+        broker_dsn: str,
+        redis_dsn: str,
+        config: "BinomicConfig",
+    ) -> None:
+
+        self._app_name = app_name
+        self._broker_dsn = broker_dsn
+        self._redis_dsn = redis_dsn
+        self._config = config
+
+    def on_app_init(self, app_config: "AppConfig") -> "AppConfig":
+
+        self.setup_signature_namespaces(app_config)
+        self.setup_states(app_config)
+        self.setup_dependencies(app_config)
+
+        return app_config
+
+    def provide_binomic(self, state: "State") -> "Binomic":
+
+        if (
+            factory := cast(
+                "BinomicFactory | None",
+                state.get(self._binomic_factory_state_key),
+            )
+        ) is None:
+            raise RuntimeError(
+                "BinomicFactory is not provided. "
+                "Please ensure that the BinomicPlugin is properly configured."
+            )
+
+        return factory.create()
+
+    def setup_signature_namespaces(self, app_config: "AppConfig") -> None:
+
+        app_config.signature_namespace.update(
+            {
+                "Binomic": Binomic,
+                "BinomicFactory": BinomicFactory,
+            }
+        )
+
+    def setup_states(self, app_config: "AppConfig") -> None:
+
+        need_setup = {}
+
+        state = app_config.state
+        if self._binomic_factory_state_key not in state:
+            need_setup[self._binomic_factory_state_key] = BinomicFactory(
+                broker_dsn=self._broker_dsn,
+                redis_dsn=self._redis_dsn,
+                module_name=self._app_name,
+                config=self._config,
+            )
+
+        state.update(need_setup)
+
+    def setup_dependencies(self, app_config: "AppConfig") -> None:
+
+        app_config.dependencies.update(
+            {
+                "binomic": Provide(
+                    dependency=self.provide_binomic,
+                    sync_to_thread=True,
+                )
+            }
+        )

binomic/task/__init__.py

@@ -0,0 +1,12 @@
+from .exceptions import DuplicateTaskError, TaskNotFoundError
+from .registry import TaskRegistry, TaskSpec
+from .wrappers import autodiscover, task
+
+__all__ = (
+    "DuplicateTaskError",
+    "TaskNotFoundError",
+    "TaskRegistry",
+    "TaskSpec",
+    "autodiscover",
+    "task",
+)

binomic/task/exceptions.py

@@ -0,0 +1,14 @@
+from binomic.base import BinomicError
+
+__all__ = (
+    "DuplicateTaskError",
+    "TaskNotFoundError",
+)
+
+
+class DuplicateTaskError(BinomicError):
+    """Duplicate task error."""
+
+
+class TaskNotFoundError(BinomicError):
+    """Task not found error."""