Add async_stats subclass of statsd

This is used by the bots to push stats to statsd in a non-blocking way.

Co-authored-by: Joe Banks <[email protected]>

1 files changed, 54 insertions, 0 deletions

diff --git a/botcore/async_stats.py b/botcore/async_stats.py
new file mode 100644
index 00000000..3ef28e1a
--- /dev/null
+++ b/botcore/async_stats.py
@@ -0,0 +1,54 @@
+"""An async transport method for statsd communication."""
+
+import asyncio
+import socket
+from typing import Optional
+
+from statsd.client.base import StatsClientBase
+
+from botcore.utils import scheduling
+
+
+class AsyncStatsClient(StatsClientBase):
+    """An async implementation of :obj:`statsd.client.base.StatsClientBase` that supports async stat communication."""
+
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop,
+        host: str = 'localhost',
+        port: int = 8125,
+        prefix: str = None
+    ):
+        """
+        Create a new :obj:`AsyncStatsClient`.
+
+        Args:
+            loop (asyncio.AbstractEventLoop): The event loop to use when creating the
+                :obj:`asyncio.loop.create_datagram_endpoint`.
+            host: The host to connect to.
+            port: The port to connect to.
+            prefix: The prefix to use for all stats.
+        """
+        _, _, _, _, addr = socket.getaddrinfo(
+            host, port, socket.AF_INET, socket.SOCK_DGRAM
+        )[0]
+        self._addr = addr
+        self._prefix = prefix
+        self._loop = loop
+        self._transport: Optional[asyncio.DatagramTransport] = None
+
+    async def create_socket(self) -> None:
+        """Use :obj:`asyncio.loop.create_datagram_endpoint` from the loop given on init to create a socket."""
+        self._transport, _ = await self._loop.create_datagram_endpoint(
+            asyncio.DatagramProtocol,
+            family=socket.AF_INET,
+            remote_addr=self._addr
+        )
+
+    def _send(self, data: str) -> None:
+        """Start an async task to send data to statsd."""
+        scheduling.create_task(self._async_send(data), event_loop=self._loop)
+
+    async def _async_send(self, data: str) -> None:
+        """Send data to the statsd server using the async transport."""
+        self._transport.sendto(data.encode('ascii'), self._addr)