Implement an AsyncStatsClient to send statsd communications asynchronously

2 files changed, 49 insertions, 3 deletions

diff --git a/bot/async_stats.py b/bot/async_stats.py
new file mode 100644
index 000000000..58a80f528
--- /dev/null
+++ b/bot/async_stats.py
@@ -0,0 +1,39 @@
+import asyncio
+import socket
+
+from statsd.client.base import StatsClientBase
+
+
+class AsyncStatsClient(StatsClientBase):
+    """An async transport method for statsd communication."""
+
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop,
+        host: str = 'localhost',
+        port: int = 8125,
+        prefix: str = None
+    ):
+        """Create a new client."""
+        family, _, _, _, addr = socket.getaddrinfo(
+            host, port, socket.AF_INET, socket.SOCK_DGRAM)[0]
+        self._addr = addr
+        self._prefix = prefix
+        self._loop = loop
+        self._transport = None
+
+    async def create_socket(self) -> None:
+        """Use the loop.create_datagram_endpoint method 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."""
+        self._loop.create_task(self._async_send(data))
+
+    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)
diff --git a/bot/bot.py b/bot/bot.py
index 65081e438..c5d490409 100644
--- a/bot/bot.py
+++ b/bot/bot.py
@@ -6,10 +6,10 @@ from typing import Optional
 
 import aiohttp
 import discord
-import statsd
 from discord.ext import commands
 
 from bot import DEBUG_MODE, api, constants
+from bot.async_stats import AsyncStatsClient
 
 log = logging.getLogger('bot')
 
@@ -41,7 +41,7 @@ class Bot(commands.Bot):
             # will effectively disable stats.
             statsd_url = "127.0.0.1"
 
-        self.stats = statsd.StatsClient(statsd_url, 8125, prefix="bot")
+        self.stats = AsyncStatsClient(self.loop, statsd_url, 8125, prefix="bot")
 
     def add_cog(self, cog: commands.Cog) -> None:
         """Adds a "cog" to the bot and logs the operation."""
@@ -60,7 +60,7 @@ class Bot(commands.Bot):
         super().clear()
 
     async def close(self) -> None:
-        """Close the Discord connection and the aiohttp session, connector, and resolver."""
+        """Close the Discord connection and the aiohttp session, connector, statsd client, and resolver."""
         await super().close()
 
         await self.api_client.close()
@@ -74,6 +74,9 @@ class Bot(commands.Bot):
         if self._resolver:
             await self._resolver.close()
 
+        if self.stats._transport:
+            await self.stats._transport.close()
+
     async def login(self, *args, **kwargs) -> None:
         """Re-create the connector and set up sessions before logging into Discord."""
         self._recreate()
@@ -111,6 +114,10 @@ class Bot(commands.Bot):
         self.http_session = aiohttp.ClientSession(connector=self._connector)
         self.api_client.recreate(force=True, connector=self._connector)
 
+    async def on_ready(self) -> None:
+        """Construct an asynchronous transport for the statsd client."""
+        await self.stats.create_socket()
+
     async def on_guild_available(self, guild: discord.Guild) -> None:
         """
         Set the internal guild available event when constants.Guild.id becomes available.