from __future__ import annotations import json import re import time from asyncio import Event from dataclasses import dataclass, field, replace from datetime import datetime, timedelta from enum import Enum from typing import ( Any, AsyncContextManager, Callable, Dict, List, Optional, Protocol, overload, ) from nats.aio.client import Client from nats.aio.msg import Msg from nats.aio.subscription import Subscription from .request import Handler, Request, ServiceError DEFAULT_QUEUE_GROUP = "q" """Queue Group name used across all services.""" DEFAULT_PREFIX = "$SRV" """The root of all control subjects.""" INFO_RESPONSE_TYPE = "io.nats.micro.v1.info_response" PING_RESPONSE_TYPE = "io.nats.micro.v1.ping_response" STATS_RESPONSE_TYPE = "io.nats.micro.v1.stats_response" SEMVER_REGEX = re.compile( r"^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)(?:-((?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\+([0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?$" ) NAME_REGEX = re.compile(r"^[A-Za-z0-9\-_]+$") SUBJECT_REGEX = re.compile(r"^[^ >]*[>]?$") class ServiceVerb(str, Enum): PING = "PING" STATS = "STATS" INFO = "INFO" @dataclass class EndpointConfig: name: str """An alphanumeric human-readable string used to describe the endpoint. Multiple endpoints can have the same names. """ handler: Handler """The handler of the endpoint.""" queue_group: Optional[str] = None """The queue group of the endpoint. When queue group is not set, it defaults to the queue group of the parent group or service.""" subject: Optional[str] = None """The subject of the endpoint. When subject is not set, it defaults to the name of the endpoint.""" metadata: Optional[Dict[str, str]] = None """The metadata of the endpoint.""" def __post_init__(self) -> None: if not self.name: raise ValueError("Name cannot be empty.") if not NAME_REGEX.match(self.name): raise ValueError("Invalid name. Name must contain only alphanumeric characters, underscores, and hyphens.") if self.subject: if not SUBJECT_REGEX.match(self.subject): raise ValueError("Invalid subject. Subject must not contain spaces, and can only have '>' at the end.") if self.queue_group: if not SUBJECT_REGEX.match(self.queue_group): raise ValueError("Invalid queue group. Queue group must not contain spaces.") @dataclass class EndpointStats: """ Statistics about a specific service endpoint """ name: str """ The endpoint name """ subject: str """ The subject the endpoint listens on """ queue_group: str """ The queue group this endpoint listens on for requests """ num_requests: int """ The number of requests this endpoint has received """ num_errors: int """ The number of errors this endpoint has encountered """ processing_time: int """ The total processing time of requests in nanoseconds """ average_processing_time: int """ The average processing time of requests in nanoseconds """ last_error: Optional[str] = None """ The last error the service encountered, if any. """ data: Optional[Any] = None """ Additional statistics the endpoint makes available """ @classmethod def from_dict(cls, data: Dict[str, Any]) -> EndpointStats: return cls( name=data["name"], subject=data["subject"], queue_group=data["queue_group"], num_requests=data["num_requests"], num_errors=data["num_errors"], last_error=data["last_error"], processing_time=data["processing_time"], average_processing_time=data["average_processing_time"], data=data["data"], ) def to_dict(self) -> Dict[str, Any]: return { "name": self.name, "subject": self.subject, "queue_group": self.queue_group, "num_requests": self.num_requests, "num_errors": self.num_errors, "last_error": self.last_error or "", "processing_time": self.processing_time, "average_processing_time": self.average_processing_time, "data": self.data, } @dataclass class EndpointInfo: """The information of an endpoint.""" name: str """ The endopoint name """ subject: str """ The subject the endpoint listens on """ queue_group: str """ The queue group this endpoint listens on for requests """ metadata: Dict[str, Any] = field(default_factory=dict) """ The endpoint metadata. """ @classmethod def from_dict(cls, data: Dict[str, Any]) -> EndpointInfo: return cls( name=data["name"], subject=data["subject"], queue_group=data["queue_group"], metadata=data["metadata"], ) def to_dict(self) -> Dict[str, Any]: return { "name": self.name, "subject": self.subject, "queue_group": self.queue_group, "metadata": self.metadata, } class Endpoint: """Endpoint manages a service endpoint.""" def __init__(self, service: Service, config: EndpointConfig) -> None: self._service = service self._name = config.name self._subject = config.subject or config.name self._queue_group = config.queue_group or DEFAULT_QUEUE_GROUP self._handler = config.handler self._metadata = config.metadata self._num_requests = 0 self._num_errors = 0 self._processing_time = 0 self._average_processing_time = 0 self._last_error = None self._subscription: Optional[Subscription] = None async def _start(self) -> None: assert not self._subscription self._subscription = await self._service._client.subscribe( subject=self._subject, queue=self._queue_group, cb=self._handle_request, ) async def _stop(self) -> None: assert self._subscription await self._subscription.unsubscribe() self._subscription = None def _reset(self) -> None: """Reset the endpoint statistics.""" self._num_requests = 0 self._num_errors = 0 self._processing_time = 0 self._average_processing_time = timedelta() async def _handle_request(self, msg: Msg) -> None: """Handle an endpoint message.""" start_time = time.perf_counter_ns() self._num_requests += 1 request = Request(msg) try: await self._handler(request) except Exception as error: self._num_errors += 1 self._last_error = repr(error) if isinstance(error, ServiceError): await request.respond_error(error.code, error.description) else: await request.respond_error("500", repr(error)) current_time = time.perf_counter_ns() elapsed_time = current_time - start_time self._processing_time += elapsed_time self._average_processing_time = int(self._processing_time / self._num_requests) @dataclass class GroupConfig: """The configuration of a group.""" name: str """The name of the group.""" queue_group: Optional[str] = None """The default queue group of the group.""" class EndpointManager(Protocol): """ Manages the endpoints of a service. """ @overload async def add_endpoint(self, config: EndpointConfig) -> None: ... @overload async def add_endpoint( self, *, name: str, handler: Handler, queue_group: Optional[str] = None, subject: Optional[str] = None, metadata: Optional[Dict[str, str]] = None, ) -> None: ... async def add_endpoint(self, config: Optional[EndpointConfig] = None, **kwargs) -> None: ... class GroupManager(Protocol): """ Manages the groups of a service. """ @overload def add_group(self, *, name: str, queue_group: Optional[str] = None) -> Group: ... @overload def add_group(self, config: GroupConfig) -> Group: ... def add_group(self, config: Optional[GroupConfig] = None, **kwargs) -> Group: ... class Group(GroupManager, EndpointManager): def __init__(self, service: "Service", config: GroupConfig) -> None: self._service = service self._prefix = config.name self._queue_group = config.queue_group @overload async def add_endpoint(self, config: EndpointConfig) -> None: ... @overload async def add_endpoint( self, *, name: str, handler: Handler, queue_group: Optional[str] = None, subject: Optional[str] = None, metadata: Optional[Dict[str, str]] = None, ) -> None: ... async def add_endpoint(self, config: Optional[EndpointConfig] = None, **kwargs) -> None: if config is None: config = EndpointConfig(**kwargs) else: config = replace(config, **kwargs) config = replace( config, subject=f"{self._prefix.strip('.')}.{config.subject or config.name}".strip("."), queue_group=config.queue_group or self._queue_group, ) await self._service.add_endpoint(config) @overload def add_group(self, *, name: str, queue_group: Optional[str] = None) -> Group: ... @overload def add_group(self, config: GroupConfig) -> Group: ... def add_group(self, config: Optional[GroupConfig] = None, **kwargs) -> Group: if config: config = replace(config, **kwargs) else: config = GroupConfig(**kwargs) config = replace( config, name=f"{self._prefix}.{config.name}", queue_group=config.queue_group or self._queue_group, ) return Group(self._service, config) StatsHandler = Callable[[EndpointStats], Any] """ A handler function used to configure a custom *STATS* endpoint. """ @dataclass class ServiceConfig: """The configuration of a service.""" name: str """The kind of the service. Shared by all services that have the same name. This name can only have A-Z, a-z, 0-9, dash, underscore.""" version: str """The version of the service. This version must be a valid semantic version.""" description: Optional[str] = None """The description of the service.""" metadata: Optional[Dict[str, str]] = None """The metadata of the service.""" queue_group: Optional[str] = None """The default queue group of the service.""" stats_handler: Optional[StatsHandler] = None """ A handler function used to configure a custom *STATS* endpoint. """ def __post_init__(self) -> None: if not self.name: raise ValueError("Name cannot be empty.") if not NAME_REGEX.match(self.name): raise ValueError("Invalid name. It must contain only alphanumeric characters, dashes, and underscores.") if not self.version: raise ValueError("Version cannot be empty.") if not SEMVER_REGEX.match(self.version): raise ValueError("Invalid version. It must follow semantic versioning (e.g., 1.0.0, 2.1.3-alpha.1).") if self.queue_group: if not SUBJECT_REGEX.match(self.queue_group): raise ValueError( "Invalid queue group. It must contain only alphanumeric characters, dashes, and underscores." ) class ServiceIdentity(Protocol): """ Defines fields helping to identity a service instance. """ id: str name: str version: str metadata: Dict[str, str] @dataclass class ServicePing(ServiceIdentity): """The response to a ping message.""" id: str name: str version: str metadata: Dict[str, str] = field(default_factory=dict) type: str = PING_RESPONSE_TYPE @classmethod def from_dict(cls, data: Dict[str, Any]) -> ServicePing: return cls( id=data["id"], name=data["name"], version=data["version"], metadata=data.get("metadata", {}), ) def to_dict(self) -> Dict[str, Any]: return { "type": self.type, "id": self.id, "name": self.name, "version": self.version, "metadata": self.metadata, } @dataclass class ServiceStats(ServiceIdentity): """The statistics of a service.""" id: str """ A unique ID for this instance of a service """ name: str """ The kind of the service. Shared by all the services that have the same name """ version: str """ The version of the service """ started: datetime """ The time the service was started """ endpoints: List[EndpointStats] = field(default_factory=list) """ Statistics for each known endpoint """ metadata: Dict[str, str] = field(default_factory=dict) """Service metadata.""" type: str = STATS_RESPONSE_TYPE """ The schema type of the message """ @classmethod def from_dict(cls, data: Dict[str, Any]) -> ServiceStats: result = cls( id=data["id"], name=data["name"], version=data["version"], started=datetime.strptime(data["started"], "%Y-%m-%dT%H:%M:%S.%fZ"), endpoints=[EndpointStats.from_dict(endpoint) for endpoint in data["endpoints"]], metadata=data["metadata"], ) return result def to_dict(self) -> Dict[str, Any]: return { "type": self.type, "name": self.name, "id": self.id, "version": self.version, "started": self.started.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z", "endpoints": [endpoint.to_dict() for endpoint in self.endpoints], "metadata": self.metadata, } @dataclass class ServiceInfo: """The information of a service.""" name: str """ The name of the service. Shared by all the services that have the same name """ id: str """ A unique ID for this instance of a service """ version: str """ The version of the service """ description: Optional[str] = None """ The description of the service supplied as configuration while creating the service """ endpoints: List[EndpointInfo] = field(default_factory=list) """ Information for all service endpoints """ metadata: Dict[str, str] = field(default_factory=dict) """ The service metadata """ type: str = INFO_RESPONSE_TYPE """ The type of the message """ @classmethod def from_dict(cls, data: Dict[str, Any]) -> ServiceInfo: """ Create a `ServiceInfo` from a dictionary. Notes: - Unknown fields are ignored. """ return cls( id=data["id"], name=data["name"], version=data["version"], description=data.get("description"), endpoints=[EndpointInfo.from_dict(endpoint) for endpoint in data["endpoints"]], metadata=data["metadata"], type=data.get("type", "io.nats.micro.v1.info_response"), ) def to_dict(self) -> Dict[str, Any]: """ Convert the `ServiceInfo` to a dictionary. """ return { "type": self.type, "id": self.id, "name": self.name, "version": self.version, "description": self.description, "endpoints": [endpoint.to_dict() for endpoint in self.endpoints], "metadata": self.metadata, } class Service(AsyncContextManager): def __init__(self, client: Client, config: ServiceConfig) -> None: self._id = client._nuid.next().decode() self._name = config.name self._version = config.version self._description = config.description self._metadata = config.metadata or {} self._queue_group = config.queue_group self._stats_handler = config.stats_handler self._client = client self._subscriptions = {} self._endpoints = [] self._started = datetime.utcnow() self._stopped = Event() self._prefix = DEFAULT_PREFIX @property def id(self) -> str: """ Returns the unique identifier of the service. """ return self._id async def start(self) -> None: if self._subscriptions: return verb_request_handlers = { ServiceVerb.PING: self._handle_ping_request, ServiceVerb.INFO: self._handle_info_request, ServiceVerb.STATS: self._handle_stats_request, } for verb, verb_handler in verb_request_handlers.items(): verb_subjects = [ ( f"{verb}-all", control_subject(verb, name=None, id=None, prefix=self._prefix), ), ( f"{verb}-kind", control_subject(verb, name=self._name, id=None, prefix=self._prefix), ), ( verb, control_subject(verb, name=self._name, id=self._id, prefix=self._prefix), ), ] for key, subject in verb_subjects: self._subscriptions[key] = await self._client.subscribe(subject, cb=verb_handler) self._started = datetime.utcnow() await self._client.flush() @overload async def add_endpoint(self, config: EndpointConfig) -> None: ... @overload async def add_endpoint( self, *, name: str, handler: Handler, queue_group: Optional[str] = None, subject: Optional[str] = None, metadata: Optional[Dict[str, str]] = None, ) -> None: ... async def add_endpoint(self, config: Optional[EndpointConfig] = None, **kwargs) -> None: if config is None: config = EndpointConfig(**kwargs) else: config = replace(config, **kwargs) config = replace(config, queue_group=config.queue_group or self._queue_group) endpoint = Endpoint(self, config) await endpoint._start() self._endpoints.append(endpoint) @overload def add_group(self, *, name: str, queue_group: Optional[str] = None) -> Group: ... @overload def add_group(self, config: GroupConfig) -> Group: ... def add_group(self, config: Optional[GroupConfig] = None, **kwargs) -> Group: if config: config = replace(config, **kwargs) else: config = GroupConfig(**kwargs) config = replace(config, queue_group=config.queue_group or self._queue_group) return Group(self, config) def stats(self) -> ServiceStats: """ Returns the statistics for the service. """ stats = ServiceStats( id=self._id, name=self._name, version=self._version, metadata=self._metadata, endpoints=[ EndpointStats( name=endpoint._name, subject=endpoint._subject, queue_group=endpoint._queue_group, num_requests=endpoint._num_requests, num_errors=endpoint._num_errors, last_error=endpoint._last_error, processing_time=endpoint._processing_time, average_processing_time=endpoint._average_processing_time, ) for endpoint in (self._endpoints or []) ], started=self._started, ) if self._stats_handler: for endpoint_stats in stats.endpoints: endpoint_stats.data = self._stats_handler(endpoint_stats) return stats def info(self) -> ServiceInfo: return ServiceInfo( id=self._id, name=self._name, version=self._version, metadata=self._metadata, description=self._description, endpoints=[ EndpointInfo( name=endpoint._name, subject=endpoint._subject, queue_group=endpoint._queue_group, metadata=endpoint._metadata, ) for endpoint in self._endpoints ], ) async def reset(self) -> None: """ Resets all statistics for all endpoints on a service instance. """ self._started = datetime.utcnow() async def stop(self) -> None: if self._stopped.is_set(): return for endpoint in self._endpoints: await endpoint._stop() for subscription in self._subscriptions.values(): await subscription.drain() self._endpoints = [] self._subscriptions = {} self._stopped.set() @property def stopped(self) -> Event: """ Indicates whether `stop` was executed on the service. """ return self._stopped async def __aenter__(self) -> "Service": await self.start() return self async def __aexit__(self, exc_type, exc, tb) -> None: await self.stop() async def _handle_ping_request(self, msg: Msg) -> None: """Handle a ping message.""" ping = ServicePing( id=self._id, name=self._name, version=self._version, metadata=self._metadata, ).to_dict() await msg.respond(data=json.dumps(ping).encode()) async def _handle_info_request(self, msg: Msg) -> None: """Handle an info message.""" info = self.info().to_dict() await msg.respond(data=json.dumps(info).encode()) async def _handle_stats_request(self, msg: Msg) -> None: """Handle a stats message.""" stats = self.stats().to_dict() await msg.respond(data=json.dumps(stats).encode()) def control_subject( verb: ServiceVerb, name: Optional[str] = None, id: Optional[str] = None, prefix=DEFAULT_PREFIX, ) -> str: if name is None and id is None: return f"{prefix}.{verb.value}" elif id is None: return f"{prefix}.{verb.value}.{name}" else: return f"{prefix}.{verb.value}.{name}.{id}"