delphai-rpc 4.1.6__tar.gz

delphai_rpc-4.1.6/PKG-INFO

@@ -0,0 +1,22 @@
+Metadata-Version: 2.1
+Name: delphai-rpc
+Version: 4.1.6
+Summary: Queue-based RPC client and server
+Author: Anton Ryzhov
+Author-email: anton@delphai.com
+Requires-Python: >=3.8,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: aio-pika (>=9.1.4,<10.0.0)
+Requires-Dist: msgpack (>=1.0.5,<2.0.0)
+Requires-Dist: prometheus-client (>=0.20.0,<0.21.0)
+Requires-Dist: pydantic (>=2.2,<3.0)
+Description-Content-Type: text/markdown
+
+# delphai-rpc
+Queue-based RPC client and server
+

delphai_rpc-4.1.6/README.md

@@ -0,0 +1,2 @@
+# delphai-rpc
+Queue-based RPC client and server

delphai_rpc-4.1.6/delphai_rpc/__init__.py

@@ -0,0 +1,5 @@
+from .client import Options, RpcClient
+from .server import RpcServer
+
+
+__all__ = ["Options", "RpcClient", "RpcServer"]

delphai_rpc-4.1.6/delphai_rpc/client.py

@@ -0,0 +1,331 @@
+import aio_pika
+import aio_pika.exceptions
+import asyncio
+import functools
+import logging
+import socket
+import time
+import uuid
+import weakref
+
+from aio_pika.abc import (
+    AbstractChannel,
+    AbstractExchange,
+    AbstractRobustConnection,
+    AbstractQueue,
+)
+from typing import Any, Dict, Optional
+
+from . import errors
+from . import metrics
+from .connection_manager import get_connection
+from .models import Request, Response
+from .server import request_context
+from .types import AbstractOptions, IncomingMessage, Message, Priority
+from .utils import clean_service_name, fix_message_timestamp
+
+
+logger = logging.getLogger(__name__)
+
+
+class Options(AbstractOptions):
+    timeout: Optional[float] = 60
+    priority: Optional[Priority] = None
+    no_wait: bool = False
+
+
+class RpcClient:
+    def __init__(
+        self, client_name: str, connection_string: str, *args: Options, **kwargs: Any
+    ) -> None:
+        self._client_name = clean_service_name(client_name)
+        self._connection_string = connection_string
+        self._app_id = f"{self._client_name}@{socket.gethostname()}"
+
+        self._options = Options(*args, **kwargs)
+        self._client_lock = asyncio.Lock()
+        self._declare_exchange_lock = asyncio.Lock()
+
+        self._reset()
+
+    def _reset(self) -> None:
+        self._connection: Optional[AbstractRobustConnection] = None
+        self._channel: Optional[AbstractChannel] = None
+        self._reply_queue: Optional[AbstractQueue] = None
+
+        self._exchanges: Dict[str, AbstractExchange] = {}
+        self._futures: weakref.WeakValueDictionary[str, asyncio.Future] = (
+            weakref.WeakValueDictionary()
+        )
+
+    def get_service(
+        self, service_name: str, *args: Options, **kwargs: Dict[str, Any]
+    ) -> "RpcService":
+        options = self._options.update(*args, **kwargs)
+        return RpcService(self, clean_service_name(service_name), options)
+
+    __getitem__ = __getattr__ = get_service
+
+    async def _ensure_connection(self) -> None:
+        async with self._client_lock:
+            if self._connection and self._channel:
+                return
+
+            self._connection = await get_connection(
+                self._connection_string, self._client_name
+            )
+            self._channel = await self._connection.channel(on_return_raises=True)
+
+    async def _ensure_reply_queue(self) -> None:
+        await self._ensure_connection()
+
+        assert self._channel
+
+        async with self._client_lock:
+            if self._reply_queue:
+                return
+
+            queue_name = f"client.{self._client_name}.{uuid.uuid4().hex}"
+            queue = await self._channel.declare_queue(
+                name=queue_name,
+                exclusive=True,
+                auto_delete=True,
+            )
+            await queue.consume(fix_message_timestamp(self._on_message))
+
+            self._reply_queue = queue
+
+    async def stop(self) -> None:
+        channel = self._channel
+        if channel:
+            self._reset()
+            await channel.close()
+
+    async def call(
+        self,
+        service_name: str,
+        method_name: str,
+        arguments: Optional[Dict[str, Any]] = None,
+        options: Optional[Options] = None,
+    ) -> Any:
+        options = options or Options()
+
+        current_request_context = request_context.get()
+        if current_request_context:
+            if options.priority is None:
+                options = options.update(priority=current_request_context.priority)
+
+            if current_request_context.deadline is not None:
+                timeout = current_request_context.deadline - time.time()
+                if options.timeout is not None:
+                    timeout = min(options.timeout, timeout)
+
+                options = options.update(timeout=timeout)
+
+        if options.priority is None:
+            options = options.update(priority=Priority.DEFAULT)
+
+        if options.timeout:
+            deadline = time.monotonic() + options.timeout
+            waiter = lambda coro: asyncio.wait_for(  # noqa: E731
+                coro, timeout=(deadline - time.monotonic())
+            )
+        else:
+            waiter = lambda coro: coro  # noqa: E731
+
+        request = Request(method_name=method_name, arguments=arguments or {})
+
+        if options.no_wait:
+            correlation_id = None
+            future = None
+        else:
+            correlation_id = str(uuid.uuid1())
+            future = asyncio.get_running_loop().create_future()
+            self._futures[correlation_id] = future
+
+        with metrics.client_requests_in_progress.labels(
+            service=service_name,
+            method=method_name,
+        ).track_inprogress():
+            elapsed = -time.perf_counter()
+            await waiter(
+                self._send_request(service_name, request, options, correlation_id)
+            )
+
+            try:
+                if future:
+                    return await waiter(future)
+            except asyncio.CancelledError:
+                logger.warning(
+                    "Wait was cancelled but not the request itself. "
+                    "Pass `timeout` option instead of using `asyncio.wait_for` or similar"
+                )
+                raise
+            finally:
+                elapsed += time.perf_counter()
+
+                metrics.client_request_processed(
+                    priority=options.priority or 0,
+                    service=service_name,
+                    method=method_name,
+                    elapsed=elapsed,
+                )
+
+    async def _send_request(
+        self,
+        service_name: str,
+        request: Request,
+        options: Options,
+        correlation_id: Optional[str] = None,
+    ):
+        await self._ensure_connection()
+
+        assert self._channel
+
+        request_message = Message(
+            body=request,
+            app_id=self._app_id,
+            priority=options.priority,
+            expiration=options.timeout or None,
+            type="rpc.request",
+        )
+
+        if correlation_id:
+            await self._ensure_reply_queue()
+            assert self._reply_queue
+            request_message.correlation_id = correlation_id
+            request_message.reply_to = self._reply_queue.name
+
+        async with self._declare_exchange_lock:
+            if service_name not in self._exchanges:
+                try:
+                    self._exchanges[service_name] = await self._channel.get_exchange(
+                        f"service.{service_name}"
+                    )
+                except aio_pika.exceptions.ChannelNotFoundEntity:
+                    raise errors.UnknownServiceError("Exchange was not found")
+
+        try:
+            routing_key = f"method.{request.method_name}"
+            await self._exchanges[service_name].publish(
+                message=request_message,
+                routing_key=routing_key,
+            )
+            metrics.message_published(
+                exchange=self._exchanges[service_name].name,
+                routing_key=routing_key,
+                type=request_message.type or "",
+                priority=request_message.priority or 0,
+                payload_size=request_message.body_size,
+            )
+
+        except aio_pika.exceptions.PublishError:
+            raise errors.UnknownServiceError("Request was not delivered to a queue")
+
+    async def _on_message(self, message: IncomingMessage) -> None:
+        message_consumed = functools.partial(
+            metrics.message_consumed,
+            exchange=message.exchange,
+            routing_key="",  # random, causes high cardinality metric
+            type=message.type,
+            priority=message.priority,
+            redelivered=message.redelivered,
+            payload_size=message.body_size,
+        )
+
+        if message.type != "rpc.response":
+            logger.warning(
+                "[MID:%s] Unexpected message type: `%s`",
+                message.message_id,
+                message.type,
+            )
+            await message.reject()
+            message_consumed(error="WRONG_MESSAGE_TYPE")
+            return
+
+        if not message.correlation_id:
+            logger.warning("[MID:%s] `correlation_id` is not set", message.message_id)
+            await message.reject()
+            message_consumed(error="NO_CORRELATION_ID")
+            return
+
+        future = self._futures.pop(message.correlation_id, None)
+        if not future or future.done():
+            logger.warning(
+                "[MID:%s] [CID:%s] Response is not awaited (too late or duplicate)",
+                message.message_id,
+                message.correlation_id,
+            )
+            await message.reject()
+            message_consumed(error="UNKNOWN_CORRELATION_ID")
+            return
+
+        try:
+            future.set_result(await self._process_message(message))
+        except Exception as error:
+            future.set_exception(error)
+
+        await message.ack()
+        message_consumed()
+
+        logger.debug(
+            "[MID:%s] [CID:%s] Got `%s` from `%s` service",
+            message.message_id,
+            message.correlation_id,
+            message.type,
+            message.app_id or "unknown",
+        )
+
+    async def _process_message(self, message: IncomingMessage) -> Any:
+        response = Response.model_validate_message(message)
+
+        if response.error:
+            error_class = getattr(errors, response.error.type, errors.UnknownError)
+            raise error_class(response.error.message)
+
+        return response.result
+
+
+class RpcService:
+    def __init__(self, client: RpcClient, service_name: str, options: Options):
+        self._client = client
+        self._service_name = service_name
+        self._options = options
+
+    def get_method(
+        self, method_name: str, *args: Options, **kwargs: Dict[str, Any]
+    ) -> "RpcMethod":
+        options = self._options.update(*args, **kwargs)
+        return RpcMethod(self._client, self._service_name, method_name, options)
+
+    __getitem__ = __getattr__ = get_method
+
+    def __repr__(self) -> str:
+        class_ = self.__class__
+        return f"<{class_.__qualname__} `{self._service_name}`>"
+
+
+class RpcMethod:
+    def __init__(
+        self, client: RpcClient, service_name: str, method_name: str, options: Options
+    ) -> None:
+        self._client = client
+        self._service_name = service_name
+        self._method_name = method_name
+        self._options = options
+
+    def __repr__(self) -> str:
+        class_ = self.__class__
+        return (
+            f"<{class_.__qualname__} `{self._method_name}` "
+            "of service `{self._service_name}`>"
+        )
+
+    def __call__(self, *args: Options, **kwargs: Dict[str, Any]) -> Any:
+        options = self._options.update(*args)
+        return self._client.call(
+            service_name=self._service_name,
+            method_name=self._method_name,
+            arguments=kwargs,
+            options=options,
+        )

delphai_rpc-4.1.6/delphai_rpc/connection_manager.py

@@ -0,0 +1,38 @@
+import aio_pika.exceptions
+import asyncio
+import importlib.metadata
+import logging
+import weakref
+
+from aio_pika.connection import URL
+from aio_pika.robust_connection import AbstractRobustConnection
+from typing import Union
+
+
+logger = logging.getLogger(__name__)
+
+_connect_lock = asyncio.Lock()
+_connections: weakref.WeakValueDictionary[URL, AbstractRobustConnection] = (
+    weakref.WeakValueDictionary()
+)
+
+
+async def get_connection(
+    connection_string: Union[str, URL, None], service_name: str = "?"
+) -> AbstractRobustConnection:
+    connection_url = aio_pika.connection.make_url(connection_string)
+    if not connection_url.query.get("name"):
+        package_name = __package__.split(".")[0]
+        package_name_version = importlib.metadata.version(package_name)
+
+        connection_url %= {
+            "name": f"{service_name} ({package_name} v{package_name_version})"
+        }
+
+    async with _connect_lock:
+        connection = _connections.get(connection_url)
+        if connection is None:
+            connection = await aio_pika.connect_robust(connection_url)
+            _connections[connection_url] = connection
+
+    return connection

delphai_rpc-4.1.6/delphai_rpc/errors.py

@@ -0,0 +1,34 @@
+class RpcError(Exception):
+    pass
+
+
+class TemporaryError(RpcError):
+    pass
+
+
+class FinalError(RpcError):
+    pass
+
+
+class UnknownError(FinalError):
+    pass
+
+
+class UnhandledError(FinalError):
+    pass
+
+
+class ParsingError(FinalError):
+    pass
+
+
+class UnknownServiceError(FinalError):
+    pass
+
+
+class UnknownMethodError(FinalError):
+    pass
+
+
+class ExecutionError(FinalError):
+    pass

delphai_rpc-4.1.6/delphai_rpc/metrics.py

@@ -0,0 +1,223 @@
+import importlib.metadata
+
+from .models import ResponseError
+
+from prometheus_client import Counter, Gauge, Histogram, Info, Summary
+from typing import Optional
+
+
+application_info = Info("application", "Application info")
+
+
+def set_application_info(**kwargs) -> None:
+    mertics_package = __package__.split(".")[0]
+    application_info.info(
+        {
+            "mertics_package": mertics_package,
+            "mertics_package_version": importlib.metadata.version(mertics_package),
+            **kwargs,
+        }
+    )
+
+
+set_application_info()
+
+
+messages_published_count = Counter(
+    name="queue_rpc_messages_published_total",
+    documentation="Total number of published messages",
+    labelnames=["exchange", "routing_key", "type", "priority"],
+)
+
+messages_published_payload_size = Summary(
+    name="queue_rpc_messages_published_payload_size_bytes",
+    documentation="Payload size of published messages",
+    labelnames=["exchange", "routing_key", "type"],
+)
+
+
+def message_published(
+    *, exchange: str, routing_key: str, type: str, priority: int, payload_size: int
+) -> None:
+    labels = dict(
+        exchange=exchange,
+        routing_key=routing_key,
+        type=type,
+    )
+    messages_published_count.labels(priority=priority, **labels).inc()
+    messages_published_payload_size.labels(**labels).observe(payload_size)
+
+
+messages_consumed_count = Counter(
+    name="queue_rpc_messages_consumed_total",
+    documentation="Total number of consumed messages",
+    labelnames=["exchange", "routing_key", "redelivered", "type", "priority", "error"],
+)
+
+messages_consumed_payload_size = Summary(
+    name="queue_rpc_messages_consumed_payload_size_bytes",
+    documentation="Payload size of consumed messages",
+    labelnames=["exchange", "routing_key", "redelivered", "type"],
+)
+
+
+def message_consumed(
+    *,
+    exchange: str,
+    routing_key: str,
+    redelivered: bool,
+    type: str,
+    priority: int,
+    payload_size: int,
+    error: Optional[str] = None,
+) -> None:
+    labels = dict(
+        exchange=exchange,
+        routing_key=routing_key,
+        redelivered=redelivered,
+        type=type,
+    )
+    messages_consumed_count.labels(priority=priority, error=error or "", **labels).inc()
+    messages_consumed_payload_size.labels(**labels).observe(payload_size)
+
+
+server_requests_count = Counter(
+    name="queue_rpc_server_requests_total",
+    documentation="Total number of requests",
+    labelnames=["priority", "method", "error"],
+)
+
+server_requests_in_progress = Gauge(
+    name="queue_rpc_server_requests_in_progress",
+    documentation="Number of requests in progress",
+    labelnames=["method"],
+)
+
+server_request_waiting_time = Histogram(
+    name="queue_rpc_server_request_waiting_seconds",
+    documentation="Time request spent in queue",
+    labelnames=["priority", "method"],
+    buckets=(
+        0.1,
+        0.3,
+        0.5,
+        1,
+        3,
+        5,
+        10,
+        30,
+        1 * 60,
+        3 * 60,
+        5 * 60,
+        10 * 60,
+        30 * 60,
+        1 * 3600,
+        3 * 3600,
+        5 * 3600,
+        10 * 3600,
+        30 * 3600,
+    ),
+)
+
+server_request_processing_time = Histogram(
+    name="queue_rpc_server_request_processing_seconds",
+    documentation="Time spent processing request",
+    labelnames=["method"],
+    buckets=(
+        0.1,
+        0.25,
+        0.5,
+        0.75,
+        1.0,
+        2.5,
+        5.0,
+        7.5,
+        10,
+        25,
+        50,
+        75,
+        100,
+        250,
+        500,
+        750,
+        1000,
+        2500,
+    ),
+)
+
+
+def server_request_processed(
+    priority: int,
+    method: str,
+    error: Optional[ResponseError],
+    queued_for: float,
+    elapsed: float,
+) -> None:
+    server_requests_count.labels(
+        priority=priority,
+        method=method,
+        error=error and error.type or "",
+    ).inc()
+
+    if queued_for is not None:
+        server_request_waiting_time.labels(
+            priority=priority,
+            method=method,
+        ).observe(queued_for)
+
+    server_request_processing_time.labels(
+        method=method,
+    ).observe(elapsed)
+
+
+client_requests_count = Counter(
+    name="queue_rpc_client_requests_total",
+    documentation="Total number of requests",
+    labelnames=["priority", "service", "method"],
+)
+
+client_requests_in_progress = Gauge(
+    name="queue_rpc_client_requests_in_progress",
+    documentation="Number of requests in progress",
+    labelnames=["service", "method"],
+)
+
+client_request_time = Histogram(
+    name="queue_rpc_client_request_seconds",
+    documentation="Time request took",
+    labelnames=["priority", "service", "method"],
+    buckets=(
+        0.1,
+        0.3,
+        0.5,
+        1,
+        3,
+        5,
+        10,
+        30,
+        1 * 60,
+        3 * 60,
+        5 * 60,
+        10 * 60,
+        30 * 60,
+        1 * 3600,
+        3 * 3600,
+        5 * 3600,
+        10 * 3600,
+        30 * 3600,
+    ),
+)
+
+
+def client_request_processed(
+    priority: int, service: str, method: str, elapsed: float
+) -> None:
+    labels = dict(
+        priority=int(priority or 0),
+        service=service,
+        method=method,
+    )
+
+    client_requests_count.labels(**labels).inc()
+
+    client_request_time.labels(**labels).observe(elapsed)

delphai_rpc-4.1.6/delphai_rpc/models.py

@@ -0,0 +1,89 @@
+import aio_pika.message
+import functools
+import logging
+import msgpack
+import pydantic
+import zlib
+
+from typing import Callable, Dict, Any, List, Optional, Tuple, Type, TypeVar
+
+from . import errors
+
+
+logger = logging.getLogger(__name__)
+
+
+TBaseModel = TypeVar("TBaseModel", bound="BaseModel")
+
+
+class BaseModel(pydantic.BaseModel):
+    @classmethod
+    def model_validate_message(
+        cls: Type[TBaseModel], message: aio_pika.message.Message
+    ) -> "TBaseModel":
+        body = message.body
+        if message.content_encoding == "deflate":
+            try:
+                body = zlib.decompress(body)
+            except Exception as error:
+                raise errors.ParsingError(f"Message decompression failed: `{error!r}`")
+
+        elif message.content_encoding:
+            raise errors.ParsingError(
+                f"Unknown content_encoding: `{message.content_encoding}`"
+            )
+
+        if message.content_type != "application/msgpack":
+            raise errors.ParsingError(
+                f"Got a message with unknown content type: {message.content_type}"
+            )
+
+        try:
+            return cls.model_validate(msgpack.loads(body))
+        except ValueError as error:
+            raise errors.ParsingError(f"Message deserialization failed: `{error!r}`")
+
+    def model_dump_msgpack(self, **kwargs) -> bytes:
+        kwargs.setdefault("exclude_defaults", True)
+        return msgpack.dumps(self.model_dump(**kwargs))
+
+
+class Request(BaseModel):
+    method_name: str
+    arguments: Dict[str, Any] = {}
+    context: Optional[Any] = None
+    timings: List[Tuple[str, float]] = []
+
+
+class ResponseError(BaseModel):
+    type: str
+    message: Optional[str] = None
+
+
+class Response(BaseModel):
+    result: Optional[Any] = None
+    error: Optional[ResponseError] = None
+    context: Optional[Any] = None
+    timings: List[Tuple[str, float]] = []
+
+    @classmethod
+    def wrap_errors(cls, func: Callable) -> Callable:
+        @functools.wraps(func)
+        async def inner(*args, **kwargs):
+            try:
+                response = await func(*args, **kwargs)
+                if not isinstance(response, cls):
+                    raise TypeError(f"Incorrect response type, got: {type(response)}")
+
+            except errors.RpcError as error:
+                response = cls(
+                    error={"type": type(error).__name__, "message": error.args[0]}
+                )
+
+            except Exception as error:
+                logger.exception("Unhandled error")
+                response = cls(error={"type": "UnhandledError", "message": repr(error)})
+
+            return response
+
+        return inner

delphai_rpc-4.1.6/delphai_rpc/server.py

@@ -0,0 +1,364 @@
+import aio_pika
+import aio_pika.connection
+import asyncio
+import contextvars
+import functools
+import inspect
+import logging
+import pydantic
+import socket
+import time
+
+from aio_pika.abc import (
+    AbstractChannel,
+    AbstractExchange,
+    AbstractRobustConnection,
+    AbstractQueue,
+)
+from typing import Any, Callable, Dict, Optional, cast
+
+from . import errors
+from . import metrics
+from .connection_manager import get_connection
+from .models import Request, Response
+from .types import IncomingMessage, Message, Priority, RequestContext
+from .utils import clean_service_name, fix_message_timestamp
+
+
+logger = logging.getLogger(__name__)
+
+request_context: contextvars.ContextVar[Optional[RequestContext]] = (
+    contextvars.ContextVar("request_context", default=None)
+)
+
+
+class RpcServer:
+    def __init__(self, service_name) -> None:
+        self._service_name = clean_service_name(service_name)
+        self._app_id = f"{self._service_name}@{socket.gethostname()}"
+
+        self._handlers: Dict[str, Callable] = {}
+
+        self._reset()
+
+        self.bind(self._ping)
+        self.bind(self._help)
+
+    def _reset(self) -> None:
+        self._connection: Optional[AbstractRobustConnection] = None
+        self._channel: Optional[AbstractChannel] = None
+        self._exchange: Optional[AbstractExchange] = None
+        self._queue: Optional[AbstractQueue] = None
+
+    def bind(
+        self, handler: Optional[Callable] = None, *, name: Optional[str] = None
+    ) -> Callable:
+        """
+        Binds to be exposed handlers (functions) to RPC server instance:
+
+        @server.bind
+        def add(*, a: float, b: float) -> float:
+            ...
+
+        # or
+
+        def sub(*, a: float, b: float) -> float:
+            ...
+
+        server.bind(sub)
+
+        # or
+
+        @server.bind(name="mul")
+        def multiply(*, a: float, b: float) -> float:
+            ...
+
+        """
+
+        def decorator(handler):
+            self._bind_handler(handler=handler, name=name)
+            return handler
+
+        return decorator(handler) if handler else decorator
+
+    def _bind_handler(self, *, handler: Callable, name: Optional[str] = None) -> None:
+        handler_name = name or handler.__name__
+        if handler_name in self._handlers:
+            raise KeyError(f"Handler {handler_name} already defined")
+
+        if hasattr(handler, "raw_function"):
+            # Unwrap `pydantic.validate_call` decorator
+            handler = handler.raw_function
+
+        self._validate_handler(handler)
+
+        self._handlers[handler_name] = pydantic.validate_call(validate_return=True)(
+            handler
+        )
+
+    def _validate_handler(self, handler: Callable) -> None:
+        if not inspect.iscoroutinefunction(handler):
+            raise TypeError("Handlers must be coroutine functions")
+
+        positional_only = []
+        positional_or_keyword = []
+
+        for parameter_name, parameter in inspect.signature(handler).parameters.items():
+            if parameter.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD:
+                positional_or_keyword.append(parameter_name)
+
+            elif parameter.kind in [
+                inspect.Parameter.POSITIONAL_ONLY,
+                inspect.Parameter.VAR_POSITIONAL,
+            ]:
+                positional_only.append(parameter_name)
+
+        if positional_only:
+            raise TypeError(
+                "{} has positional-only parameters {} that are not supported in {}".format(
+                    handler,
+                    positional_only,
+                    self.__class__,
+                )
+            )
+
+        if positional_or_keyword:
+            logger.warning(
+                "%s has positional parameters %s, only keyword parameters are supported in %s",
+                handler,
+                positional_or_keyword,
+                self.__class__,
+            )
+
+    async def start(self, connection_string: str, prefetch_count: int = 1) -> None:
+        if self._connection:
+            raise RuntimeError("Already started")
+
+        connection = self._connection = await get_connection(
+            connection_string, self._service_name
+        )
+        channel = self._channel = await connection.channel()
+        await channel.set_qos(prefetch_count=prefetch_count)
+
+        exchange = self._exchange = await channel.declare_exchange(
+            name=f"service.{self._service_name}",
+            type=aio_pika.ExchangeType.TOPIC,
+            durable=True,
+        )
+
+        queue = self._queue = await channel.declare_queue(
+            name=f"service.{self._service_name}",
+            durable=True,
+            arguments={
+                "x-max-priority": max(Priority),
+                "x-dead-letter-exchange": f"service.{self._service_name}.dlx",
+            },
+        )
+
+        await queue.bind(exchange, "#")
+        await queue.consume(fix_message_timestamp(self._on_message))
+
+        logger.info("RPC server is consuming messages from `%s`", queue)
+
+    async def stop(self) -> None:
+        channel = self._channel
+        if channel:
+            self._reset()
+            await channel.close()
+
+    async def serve_forever(
+        self, connection_string: str, prefetch_count: int = 1
+    ) -> None:
+        await self.start(connection_string, prefetch_count=prefetch_count)
+        try:
+            await asyncio.Future()
+        finally:
+            await self.stop()
+
+    async def _on_message(self, message: IncomingMessage) -> None:
+        consumed_timestamp = time.time()
+
+        logger.debug(
+            "[MID:%s] Got `%s` message from `%s` service",
+            message.message_id,
+            message.type or "[untyped]",
+            message.app_id or "unknown",
+        )
+
+        message_consumed = functools.partial(
+            metrics.message_consumed,
+            exchange=message.exchange,
+            routing_key=message.routing_key,
+            type=message.type,
+            priority=message.priority,
+            redelivered=message.redelivered,
+            payload_size=message.body_size,
+        )
+
+        if message.type != "rpc.request":
+            logger.warning(
+                "[MID:%s] Unexpected message type: `%s`",
+                message.message_id,
+                message.type,
+            )
+            await message.reject()
+            message_consumed(error="WRONG_MESSAGE_TYPE")
+            return
+
+        if message.reply_to and not message.correlation_id:
+            logger.warning("[MID:%s] `correlation_id` is not set", message.message_id)
+            message_consumed(error="NO_CORRELATION_ID")
+            await message.reject()
+            return
+
+        published_timestamp = None
+        if message.timestamp:
+            published_timestamp = message.timestamp.timestamp()
+
+        deadline = None
+        if published_timestamp and message.expiration:
+            deadline = published_timestamp + cast(float, message.expiration)
+
+        request_context.set(
+            RequestContext(
+                deadline=deadline,
+                priority=Priority(message.priority or 0),
+            )
+        )
+
+        try:
+            response = await asyncio.wait_for(
+                self._process_message(message, consumed_timestamp, published_timestamp),
+                timeout=(deadline - time.time()) if deadline else None,
+            )
+        except asyncio.TimeoutError:
+            logger.info(
+                "[MID:%s] [CID:%s] Execution of `%s` from `%s` service timed out",  # noqa: E501
+                message.message_id,
+                message.correlation_id,
+                message.type,
+                message.app_id or "unknown",
+            )
+
+            await message.ack()
+            message_consumed(error="TIMEOUT")
+            return
+
+        if message.reply_to:
+            response_message = Message(
+                body=response,
+                app_id=self._app_id,
+                priority=message.priority,
+                correlation_id=message.correlation_id,
+                expiration=(deadline - time.time()) if deadline else None,
+                type="rpc.response",
+            )
+
+            await message.channel.basic_publish(
+                body=response_message.body,
+                routing_key=message.reply_to,
+                properties=response_message.properties,
+            )
+            metrics.message_published(
+                exchange="",
+                routing_key="",  # random, causes high cardinality metric
+                type=response_message.type or "",
+                priority=response_message.priority or 0,
+                payload_size=response_message.body_size,
+            )
+
+        await message.ack()
+        message_consumed()
+
+        logger.debug(
+            "[MID:%s] [CID:%s] Handled `%s` from `%s` service, success: %s",  # noqa: E501
+            message.message_id,
+            message.correlation_id,
+            message.type,
+            message.app_id or "unknown",
+            response.error is None,
+        )
+
+    @Response.wrap_errors
+    async def _process_message(
+        self,
+        message: IncomingMessage,
+        consumed_timestamp: float,
+        published_timestamp: Optional[float],
+    ) -> Response:
+        request = Request.model_validate_message(message)
+
+        timings = request.timings
+        queued_for = None
+        if published_timestamp is not None:
+            timings.append(("queue.published", published_timestamp))
+            queued_for = consumed_timestamp - published_timestamp
+
+        timings.append((f"queue.consumed by {self._app_id}", consumed_timestamp))
+
+        with metrics.server_requests_in_progress.labels(
+            method=request.method_name
+        ).track_inprogress():
+            elapsed = -time.perf_counter()
+            response = await self._process_request(request)
+            elapsed += time.perf_counter()
+        timings.append(("execution.completed", consumed_timestamp + elapsed))
+
+        response.context = request.context
+        response.timings = timings
+
+        metrics.server_request_processed(
+            priority=message.priority or 0,
+            method=request.method_name,
+            error=response.error,
+            queued_for=queued_for or 0,
+            elapsed=elapsed,
+        )
+
+        logger.info(
+            "[MID:%s] Processed `%s` from `%s` service to method `%s`. In queue: %ims, execution: %ims, success: %s%s",
+            message.message_id,
+            message.type,
+            message.app_id or "unknown",
+            request.method_name,
+            (None if queued_for is None else max(queued_for * 1000, 0)),
+            elapsed * 1000,
+            response.error is None,
+            (f", error: {response.error.message}" if response.error else ""),
+        )
+
+        return response
+
+    @Response.wrap_errors
+    async def _process_request(self, request: Request) -> Response:
+        handler = self._handlers.get(request.method_name)
+        if handler is None:
+            raise errors.UnknownMethodError(request.method_name)
+
+        try:
+            result = await handler(**request.arguments)
+        except Exception as error:
+            raise errors.ExecutionError(repr(error))
+
+        if isinstance(result, pydantic.BaseModel):
+            result = result.model_dump()
+
+        return Response(result=result)
+
+    async def _ping(self) -> None:
+        return None
+
+    async def _help(self) -> Dict[str, Any]:
+        """
+        Returns methods list
+        """
+        return {
+            "methods": [
+                {
+                    "method_name": method_name,
+                    "signature": f"{method_name}{inspect.signature(handler)}",
+                    "description": handler.__doc__,
+                }
+                for method_name, handler in self._handlers.items()
+            ]
+        }

delphai_rpc-4.1.6/delphai_rpc/types.py

@@ -0,0 +1,112 @@
+import aio_pika.message
+import pydantic
+import time
+import uuid
+
+from aio_pika.message import IncomingMessage
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import Dict, Any, Optional, TypeVar, Union
+
+from .models import BaseModel
+
+
+__all__ = [
+    "AbstractOptions",
+    "IncomingMessage",
+    "Message",
+    "Priority",
+    "RequestContext",
+]
+
+
+TAbstractOptions = TypeVar("TAbstractOptions", bound="AbstractOptions")
+
+
+class AbstractOptions(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(extra="forbid", frozen=True)
+
+    def __init__(self, *args: "AbstractOptions", **kwargs: Dict[str, Any]) -> None:
+        if args:
+            merged = {}
+            self_class = type(self)
+            for options in args:
+                if not isinstance(options, self_class):
+                    raise TypeError(
+                        f"Positional arguments must be {self_class} instances"
+                    )
+
+                merged.update(options.model_dump(exclude_unset=True))
+            merged.update(**kwargs)
+
+            kwargs = merged
+
+        return super().__init__(**kwargs)
+
+    def update(
+        self: TAbstractOptions, *args: "TAbstractOptions", **kwargs: Any
+    ) -> "TAbstractOptions":
+        if not args and not kwargs:
+            return self
+
+        return self.__class__(self, *args, **kwargs)
+
+
+class Priority(IntEnum):
+    LOW = 0
+    NORMAL = 1
+    DEFAULT = 1
+    HIGH = 2
+    INTERACTIVE = 3
+    SYSTEM = 4
+
+
+class Message(aio_pika.message.Message):
+    """AMQP message abstraction"""
+
+    __slots__ = ()
+
+    def __init__(
+        self,
+        body: Union[bytes, BaseModel],
+        *,
+        headers: Optional[aio_pika.message.HeadersType] = None,
+        content_type: Optional[str] = None,
+        content_encoding: Optional[str] = None,
+        delivery_mode: Union[aio_pika.message.DeliveryMode, int, None] = None,
+        priority: Optional[int] = None,
+        correlation_id: Optional[str] = None,
+        reply_to: Optional[str] = None,
+        expiration: Optional[aio_pika.message.DateType] = None,
+        message_id: Optional[str] = None,
+        timestamp: Optional[aio_pika.message.DateType] = None,
+        type: Optional[str] = None,
+        user_id: Optional[str] = None,
+        app_id: Optional[str] = None,
+    ) -> None:
+        if isinstance(body, BaseModel):
+            body = body.model_dump_msgpack()
+            content_type = None
+
+        super().__init__(
+            body=body,
+            headers=headers,
+            content_type=content_type or "application/msgpack",
+            content_encoding=content_encoding,
+            delivery_mode=delivery_mode or aio_pika.DeliveryMode.PERSISTENT,
+            priority=priority,
+            correlation_id=correlation_id,
+            reply_to=reply_to,
+            expiration=expiration,
+            message_id=message_id or str(uuid.uuid1()),
+            timestamp=timestamp or time.time(),
+            type=type,
+            user_id=user_id,
+            app_id=app_id,
+        )
+
+
+@dataclass(frozen=True)
+class RequestContext:
+    deadline: Optional[float]
+    priority: Priority

delphai_rpc-4.1.6/delphai_rpc/utils.py

@@ -0,0 +1,21 @@
+import datetime
+import functools
+import re
+
+from typing import Callable
+
+
+def fix_message_timestamp(func: Callable) -> Callable:
+    @functools.wraps(func)
+    def inner(message):
+        # Fix `pamqp` naive timestamp
+        if message.timestamp:
+            message.timestamp = message.timestamp.replace(tzinfo=datetime.timezone.utc)
+
+        return func(message)
+
+    return inner
+
+
+def clean_service_name(service_name: str) -> str:
+    return re.sub("[^a-z0-9-]+", "-", service_name.strip().lower())

delphai_rpc-4.1.6/pyproject.toml

@@ -0,0 +1,21 @@
+[tool.poetry]
+name = "delphai-rpc"
+version = "4.1.6"
+description = "Queue-based RPC client and server"
+authors = ["Anton Ryzhov <anton@delphai.com>"]
+readme = "README.md"
+packages = [{include = "delphai_rpc"}]
+
+[tool.poetry.dependencies]
+python = "^3.8"
+aio-pika = { version = "^9.1.4" }
+pydantic = { version = "^2.2" }
+msgpack = { version = "^1.0.5" }
+prometheus-client = "^0.20.0"
+
+[tool.poetry.group.dev.dependencies]
+ruff = "^0.5.1"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"