Asyncio client

requirements/test.in

@@ -2,3 +2,4 @@ pytest
 pytest-cov
 pytest-xdist
 coverage
+pytest-asyncio

requirements/test.txt

@@ -17,8 +17,11 @@ pluggy==1.5.0
 pytest==8.2.2
     # via
     #   -r requirements/test.in
+    #   pytest-asyncio
     #   pytest-cov
     #   pytest-xdist
+pytest-asyncio==0.23.7
+    # via -r requirements/test.in
 pytest-cov==5.0.0
     # via -r requirements/test.in
 pytest-xdist==3.6.1

src/statsd/__init__.py

@@ -4,6 +4,8 @@
 Version: v\ |version|.
 """
 
+from .async_client import BaseAsyncStatsdClient, DebugAsyncStatsdClient
+from .base import Sample
 from .client import (
     BaseStatsdClient,
     DebugStatsdClient,
@@ -14,8 +16,11 @@
 
 
 __all__ = (
+    "BaseAsyncStatsdClient",
     "BaseStatsdClient",
+    "DebugAsyncStatsdClient",
     "DebugStatsdClient",
+    "Sample",
     "StatsdClient",
     "UDPStatsdClient",
     "__version__",

src/statsd/async_client.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import abc
+import contextlib
+import functools
+import logging
+import time
+from typing import Any, AsyncIterator, Awaitable, Callable, Mapping, TypeVar
+from typing_extensions import ParamSpec
+
+from statsd.base import AbstractStatsdClient
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+U = TypeVar("U")
+
+logger = logging.getLogger("statsd")
+
+
+class BaseAsyncStatsdClient(AbstractStatsdClient[Awaitable[None]]):
+    """
+    Base async client.
+
+    This class exposes the public interface and takes care of packet formatting
+    as well as sampling. It does not actually send packets anywhere, which is
+    left to concrete subclasses implementing :meth:`_emit`.
+    """
+
+    @abc.abstractmethod
+    async def _emit(self, packets: list[str]) -> None:
+        """
+        Async send implementation.
+
+        This method is responsible for actually sending the formatted packets
+        and should be implemented by all subclasses.
+
+        It may batch or buffer packets but should not modify them in any way. It
+        should be agnostic to the Statsd format.
+        """
+        raise NotImplementedError()
+
+    def timed(
+        self,
+        name: str | None = None,
+        *,
+        tags: Mapping[str, str] | None = None,
+        sample_rate: float | None = None,
+        use_distribution: bool = False,
+    ) -> Callable[[Callable[P, Awaitable[U]]], Callable[P, Awaitable[U]]]:
+        """
+        Wrap a function to record its execution time.
+
+        This just wraps the function call with a :meth:`timer` context manager.
+
+        If a name is not provided, the function name will be used.
+
+        Passing ``use_distribution=True`` will report the value as a globally
+        aggregated :meth:`distribution` metric instead of a :meth:`timing`
+        metric.
+
+        >>> client = AsyncStatsdClient()
+        >>> @client.timed()
+        ... async def do_something():
+        ...     pass
+        """
+
+        def decorator(
+            fn: Callable[P, Awaitable[U]],
+        ) -> Callable[P, Awaitable[U]]:
+            # TODO: Should the fallback include the module? Class (for methods)?
+            # or func.__name__
+            metric_name = name or fn.__name__
+
+            @functools.wraps(fn)
+            async def wrapped(*args: P.args, **kwargs: P.kwargs) -> U:
+                async with self.timer(
+                    metric_name,
+                    tags=tags,
+                    use_distribution=use_distribution,
+                    sample_rate=sample_rate,
+                ):
+                    return await fn(*args, **kwargs)
+
+            return wrapped
+
+        return decorator
+
+    @contextlib.asynccontextmanager
+    async def timer(
+        self,
+        name: str,
+        *,
+        tags: Mapping[str, str] | None = None,
+        sample_rate: float | None = None,
+        use_distribution: bool = False,
+    ) -> AsyncIterator[None]:
+        """
+        Context manager to measure the execution time of an async block.
+
+        Passing ``use_distribution=True`` will report the value as a globally
+        aggregated :meth:`distribution` metric instead of a :meth:`timing`
+        metric.
+
+        >>> client = AsyncStatsdClient()
+        >>> async def operation():
+        ...     async with client.timer("download_duration"):
+        ...         pass
+        """
+        start = time.perf_counter()
+        try:
+            yield
+        finally:
+            duration_ms = int(1000 * (time.perf_counter() - start))
+            if use_distribution:
+                await self.distribution(
+                    name,
+                    duration_ms,
+                    tags=tags,
+                    sample_rate=sample_rate,
+                )
+            else:
+                await self.timing(
+                    name,
+                    duration_ms,
+                    tags=tags,
+                    sample_rate=sample_rate,
+                )
+
+
+class DebugAsyncStatsdClient(BaseAsyncStatsdClient):
+    """
+    Verbose client for development or debugging purposes.
+
+    All Statsd packets will be logged and optionally forwarded to a wrapped
+    client.
+    """
+
+    def __init__(
+        self,
+        level: int = logging.INFO,
+        logger: logging.Logger = logger,
+        inner: BaseAsyncStatsdClient | None = None,
+        **kwargs: Any,
+    ) -> None:
+        r"""
+        Initialize DebugStatsdClient.
+
+        :param level: Log level to use, defaults to ``INFO``.
+
+        :param logger: Logger instance to use, defaults to ``statsd``.
+
+        :param inner: Wrapped client.
+
+        :param \**kwargs: Extra arguments forwarded to :class:`BaseAsyncStatsdClient`.
+        """
+        super().__init__(**kwargs)
+        self.level = level
+        self.logger = logger
+        self.inner = inner
+
+    async def _emit(self, packets: list[str]) -> None:
+        for packet in packets:
+            self.logger.log(self.level, "> %s", packet)
+        if self.inner:
+            await self.inner._emit(packets)
+
+
+AsyncStatsdClient = DebugAsyncStatsdClient