qena-shared-lib 0.1.0__py3-none-any.whl

qena_shared_lib/rabbitmq/_pool.py

@@ -0,0 +1,74 @@
+from asyncio import Lock
+from uuid import UUID
+
+from pika.adapters.asyncio_connection import AsyncioConnection
+
+from ..utils import AsyncEventLoopMixin
+from ._channel import BaseChannel
+
+__all__ = ["ChannelPool"]
+
+
+class ChannelPool(AsyncEventLoopMixin):
+    def __init__(
+        self,
+        initial_pool_size: int = 50,
+    ):
+        self._initial_pool_size = initial_pool_size
+        self._pool: dict[UUID, BaseChannel] = {}
+        self._pool_lock = Lock()
+
+    async def fill(
+        self,
+        connection: AsyncioConnection,
+    ):
+        self._connection = connection
+
+        async with self._pool_lock:
+            for _ in range(self._initial_pool_size):
+                await self._add_new_channel()
+
+    async def dain(self):
+        async with self._pool_lock:
+            for channel_base in self._pool.values():
+                with channel_base as channel:
+                    if channel.is_closing or channel.is_closed:
+                        continue
+
+                    channel.close()
+
+            self._pool.clear()
+
+    async def _add_new_channel(self) -> BaseChannel:
+        channel = BaseChannel(
+            connection=self._connection,
+        )
+        channel_id = await channel.open()
+        self._pool[channel_id] = channel
+
+        return channel
+
+    async def get(self) -> BaseChannel:
+        async with self._pool_lock:
+            _channel_base = None
+
+            for channel_base in self._pool.values():
+                if not channel_base.healthy:
+                    _ = self._pool.pop(channel_base.channel_id)
+
+                    continue
+
+                if not channel_base.reserved:
+                    _channel_base = channel_base
+
+                    break
+
+            if _channel_base is None:
+                _channel_base = await self._add_new_channel()
+
+            _channel_base.reserve()
+
+            return _channel_base
+
+    def __len__(self) -> int:
+        return len(self._pool)

qena_shared_lib/rabbitmq/_publisher.py

@@ -0,0 +1,73 @@
+from typing import Any, Callable
+
+from pika import BasicProperties
+from prometheus_client import Counter
+from pydantic_core import to_json
+
+from ..logging import LoggerProvider
+from ._pool import ChannelPool
+
+__all__ = ["Publisher"]
+
+
+class Publisher:
+    PUBLISHED_MESSAGES = Counter(
+        name="published_messages",
+        documentation="Published messages",
+        labelnames=["routing_key", "target"],
+    )
+
+    def __init__(
+        self,
+        routing_key: str,
+        channel_pool: ChannelPool,
+        blocked_connection_check_callback: Callable[[], bool],
+        exchange: str | None = None,
+        target: str | None = None,
+        headers: dict[str, str] | None = None,
+    ):
+        self._routing_key = routing_key
+        self._exchange = exchange or ""
+        self._headers = headers or {}
+        self._target = target if target is not None else "__default__"
+
+        self._headers.update({"target": self._target})
+
+        self._channel_pool = channel_pool
+        self._blocked_connection_check_callback = (
+            blocked_connection_check_callback
+        )
+        self._logger = LoggerProvider.default().get_logger("rabbitmq.publisher")
+
+    async def publish_as_arguments(self, *args, **kwargs):
+        await self._get_channel_and_publish({"args": args, "kwargs": kwargs})
+
+    async def publish(self, message: Any | None = None):
+        await self._get_channel_and_publish(message)
+
+    async def _get_channel_and_publish(self, message: Any):
+        if self._blocked_connection_check_callback():
+            raise RuntimeError(
+                "rabbitmq broker is not able to accept message right now"
+            )
+
+        with await self._channel_pool.get() as channel:
+            channel.basic_publish(
+                exchange=self._exchange,
+                routing_key=self._routing_key,
+                body=to_json(message),
+                properties=BasicProperties(
+                    content_type="application/json",
+                    headers=self._headers,
+                ),
+            )
+
+        self._logger.debug(
+            "message published to exchange `%s`, routing key `%s` and target `%s`",
+            self._exchange,
+            self._routing_key,
+            self._target,
+        )
+        self.PUBLISHED_MESSAGES.labels(
+            routing_key=self._routing_key, target=self._target
+        ).inc()

qena_shared_lib/rabbitmq/_rpc_client.py

@@ -0,0 +1,286 @@
+from asyncio import Future, Lock
+from functools import partial
+from importlib import import_module
+from time import time
+from typing import Any, Callable
+from uuid import uuid4
+
+from pika import BasicProperties
+from pika.channel import Channel
+from pika.frame import Method
+from pika.spec import Basic
+from prometheus_client import Counter, Summary
+from pydantic_core import from_json, to_json
+
+from ..logging import LoggerProvider
+from ..utils import AsyncEventLoopMixin
+from ._exceptions import RabbitMQException
+from ._pool import ChannelPool
+from ._utils import TypeAdapterCache
+
+__all__ = ["RpcClient"]
+
+
+class ExitHandler:
+    _exiting = False
+    _rpc_futures = []
+    _original_exit_handler: Callable
+
+    @classmethod
+    def is_exising(cls):
+        return cls._exiting
+
+    @classmethod
+    def add_rpc_future(cls, rpc_future: Future):
+        cls._rpc_futures.append(rpc_future)
+
+    @classmethod
+    def remove_rpc_future(cls, rpc_future: Future):
+        try:
+            cls._rpc_futures.remove(rpc_future)
+        except:
+            pass
+
+    @classmethod
+    def cancel_futures(cls):
+        cls._exiting = True
+
+        for rpc_future in cls._rpc_futures:
+            if not rpc_future.done():
+                rpc_future.cancel()
+
+    @staticmethod
+    def patch_exit_handler():
+        try:
+            Server = import_module("uvicorn.server").Server
+        except ModuleNotFoundError:
+            return
+
+        ExitHandler._original_exit_handler = Server.handle_exit
+        Server.handle_exit = ExitHandler.handle_exit
+
+    @staticmethod
+    def notify_clients():
+        ExitHandler.cancel_futures()
+
+    @staticmethod
+    def handle_exit(*args, **kwargs):
+        ExitHandler.notify_clients()
+        ExitHandler._original_exit_handler(*args, **kwargs)
+
+
+ExitHandler.patch_exit_handler()
+
+
+class RpcClient(AsyncEventLoopMixin):
+    SUCCEEDED_RPC_CALLS = Counter(
+        name="succeeded_rpc_calls",
+        documentation="RPC calls made",
+        labelnames=["routing_key", "procedure"],
+    )
+    FAILED_RPC_CALL = Counter(
+        name="failed_rpc_call",
+        documentation="Failed RPC calls",
+        labelnames=["routing_key", "procedure", "exception"],
+    )
+    RPC_CALL_LATENCY = Summary(
+        name="rpc_call_latency",
+        documentation="Time it took for RPC calls",
+        labelnames=["routing_key", "procedure"],
+    )
+
+    def __init__(
+        self,
+        routing_key: str,
+        channel_pool: ChannelPool,
+        blocked_connection_check_callback: Callable[[], bool],
+        exchange: str | None = None,
+        procedure: str | None = None,
+        headers: dict[str, str] | None = None,
+        return_type: type | None = None,
+        timeout: float = 0,
+    ):
+        self._routing_key = routing_key
+        self._exchange = exchange or ""
+        self._headers = headers or {}
+        self._procedure = procedure if procedure is not None else "__default__"
+
+        self._headers.update({"procedure": self._procedure})
+
+        self._return_type = return_type
+        self._timeout = timeout
+        self._channel_pool = channel_pool
+        self._blocked_connection_check_callback = (
+            blocked_connection_check_callback
+        )
+        self._rpc_future = None
+        self._rpc_call_start_time = None
+        self._rpc_call_lock = Lock()
+        self._rpc_call_pending = False
+        self._logger = LoggerProvider.default().get_logger(
+            "rabbitmq.rpc_client"
+        )
+
+    async def call_with_arguments(self, *args, **kwargs) -> Any:
+        return await self._get_channel_and_call(
+            ({"args": args, "kwargs": kwargs})
+        )
+
+    async def call(self, message: Any | None = None) -> Any:
+        return await self._get_channel_and_call(message)
+
+    async def _get_channel_and_call(self, message: Any) -> Any:
+        if self._blocked_connection_check_callback():
+            raise RuntimeError(
+                "rabbitmq broker is not able to accept message right now"
+            )
+
+        async with self._rpc_call_lock:
+            if self._rpc_call_pending:
+                raise RuntimeError("previous rpc request not done yet")
+
+            self._rpc_call_pending = True
+
+        self._rpc_call_start_time = time()
+        self._channel_base = await self._channel_pool.get()
+        self._channel = self._channel_base.channel
+        self._rpc_future = self.loop.create_future()
+
+        ExitHandler.add_rpc_future(self._rpc_future)
+        self._channel.queue_declare(
+            queue="",
+            exclusive=True,
+            auto_delete=True,
+            callback=partial(self._on_queue_declared, message),
+        )
+
+        return await self._rpc_future
+
+    def _on_queue_declared(self, message: Any, method: Method):
+        try:
+            self._rpc_reply_consumer_tag = self._channel.basic_consume(
+                queue=method.method.queue,
+                on_message_callback=self._on_reply_message,
+                auto_ack=True,
+            )
+        except Exception as e:
+            self._finalize_call(exception=e)
+
+            return
+
+        self._correlation_id = str(uuid4())
+
+        try:
+            self._channel.basic_publish(
+                exchange=self._exchange,
+                routing_key=self._routing_key,
+                properties=BasicProperties(
+                    content_type="application/json",
+                    reply_to=method.method.queue,
+                    correlation_id=self._correlation_id,
+                    headers=self._headers,
+                ),
+                body=to_json(message),
+            )
+        except Exception as e:
+            self._finalize_call(exception=e)
+
+            return
+
+        if self._timeout > 0:
+            _ = self.loop.call_later(
+                delay=self._timeout, callback=self._on_timeout
+            )
+
+        self._logger.debug(
+            "rpc request sent to exchange `%s`, routing key `%s` and procedure `%s`",
+            self._exchange,
+            self._routing_key,
+            self._procedure,
+        )
+
+    def _on_timeout(self):
+        self._finalize_call(
+            exception=TimeoutError(
+                f"rpc worker didn't responed in a timely manner within `{self._timeout}` seconds"
+            )
+        )
+
+    def _on_reply_message(
+        self,
+        channel: Channel,
+        method: Basic.Deliver,
+        properties: BasicProperties,
+        body: bytes,
+    ):
+        del channel, method
+
+        if properties.correlation_id != self._correlation_id:
+            self._finalize_call(
+                exception=ValueError(
+                    f"correlation id {properties.correlation_id} from rpc worker doesn't match {self._correlation_id}"
+                )
+            )
+
+            return
+
+        try:
+            response = from_json(body)
+        except Exception as e:
+            self._finalize_call(exception=e)
+
+            return
+
+        if isinstance(response, dict) and "exception" in response:
+            self._finalize_call(
+                exception=RabbitMQException(
+                    code=response.get("code") or 0,
+                    message=response.get("message")
+                    or "unknown error occured from the rpc worker side",
+                )
+            )
+
+        if self._return_type is not None:
+            type_adapter = TypeAdapterCache.get_type_adapter(self._return_type)
+
+            try:
+                response = type_adapter.validate_python(response)
+            except Exception as e:
+                self._finalize_call(exception=e)
+
+                return
+
+        self._finalize_call(response=response)
+
+    def _finalize_call(
+        self,
+        response: Any | None = None,
+        exception: BaseException | None = None,
+    ):
+        self._rpc_call_pending = False
+        self._channel.basic_cancel(self._rpc_reply_consumer_tag)
+        self._channel_base.release()
+
+        if self._rpc_future is None:
+            return
+
+        ExitHandler.remove_rpc_future(self._rpc_future)
+
+        if self._rpc_future.done():
+            return
+        elif exception is not None:
+            self._rpc_future.set_exception(exception)
+            self.FAILED_RPC_CALL.labels(
+                routing_key=self._routing_key,
+                procedure=self._procedure,
+                exception=exception.__class__.__name__,
+            ).inc()
+        else:
+            self._rpc_future.set_result(response)
+            self.SUCCEEDED_RPC_CALLS.labels(
+                routing_key=self._routing_key, procedure=self._procedure
+            ).inc()
+
+        self.RPC_CALL_LATENCY.labels(
+            routing_key=self._routing_key, procedure=self._procedure
+        ).observe((self._rpc_call_start_time or time()) - time())

qena_shared_lib/rabbitmq/_utils.py

@@ -0,0 +1,18 @@
+from pydantic import TypeAdapter
+
+__all__ = ["TypeAdapterCache"]
+
+
+class TypeAdapterCache:
+    _cache = {}
+
+    @classmethod
+    def cache_annotation(cls, annotation: type):
+        if annotation not in cls._cache:
+            cls._cache[annotation] = TypeAdapter(annotation)
+
+    @classmethod
+    def get_type_adapter(cls, annotation: type) -> TypeAdapter:
+        cls.cache_annotation(annotation)
+
+        return cls._cache[annotation]