From bc02db0f48c088581b9f7c08e130b98eab39e402 Mon Sep 17 00:00:00 2001
From: Dave Shawley <daveshawley@gmail.com>
Date: Mon, 2 Sep 2019 09:14:02 -0400
Subject: [PATCH] Add StatsDCollector.close method.

---
 docs/history.rst                   |  1 +
 sprockets/mixins/metrics/statsd.py | 18 ++++++++++++++----
 tests.py                           | 23 +++++++++++++++++++++++
 3 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/docs/history.rst b/docs/history.rst
index 515cc75..12321fa 100644
--- a/docs/history.rst
+++ b/docs/history.rst
@@ -8,6 +8,7 @@ Release History
 - Add configuration documentation
 - Exclude Tornado >6 (as-yet-unreleased version)
 - Add :func:`sprockets.mixins.metrics.statsd.get_client` function
+- Add :meth:`sprockets.mixins.metrics.statsd.StatsDCollector.close` method
 
 `4.0.0`_ (06-Feb-2019)
 ----------------------
diff --git a/sprockets/mixins/metrics/statsd.py b/sprockets/mixins/metrics/statsd.py
index 14952cc..80ae5bf 100644
--- a/sprockets/mixins/metrics/statsd.py
+++ b/sprockets/mixins/metrics/statsd.py
@@ -113,6 +113,7 @@ class StatsDCollector:
         self._namespace = namespace
         self._prepend_metric_type = prepend_metric_type
         self._tcp_reconnect_sleep = 5
+        self._closing = False
 
         if protocol == 'tcp':
             self._tcp = True
@@ -137,10 +138,19 @@ class StatsDCollector:
 
     async def _tcp_on_closed(self):
         """Invoked when the socket is closed."""
-        LOGGER.warning('Not connected to statsd, connecting in %s seconds',
-                       self._tcp_reconnect_sleep)
-        await asyncio.sleep(self._tcp_reconnect_sleep)
-        self._sock = self._tcp_socket()
+        if self._closing:
+            LOGGER.info('Statsd socket closed')
+        else:
+            LOGGER.warning('Not connected to statsd, connecting in %s seconds',
+                           self._tcp_reconnect_sleep)
+            await asyncio.sleep(self._tcp_reconnect_sleep)
+            self._sock = self._tcp_socket()
+
+    def close(self):
+        """Gracefully close the socket."""
+        if not self._closing:
+            self._closing = True
+            self._sock.close()
 
     def send(self, path, value, metric_type):
         """Send a metric to Statsd.
diff --git a/tests.py b/tests.py
index c4b86d9..a16d979 100644
--- a/tests.py
+++ b/tests.py
@@ -169,6 +169,29 @@ class TCPStatsdMetricCollectionTests(testing.AsyncHTTPTestCase):
         response = self.fetch('/', method='POST', body='')
         self.assertEqual(response.code, 204)
 
+    def test_that_client_closes_socket(self):
+        response = self.fetch('/status_code')
+        self.assertEqual(response.code, 200)
+
+        self.application.statsd.close()
+        response = self.fetch('/status_code')
+        self.assertEqual(response.code, 200)
+        self.assertTrue(self.application.statsd._sock.closed())
+
+    def test_that_client_can_be_closed_multiple_times(self):
+        response = self.fetch('/status_code')
+        self.assertEqual(response.code, 200)
+
+        self.application.statsd.close()
+        response = self.fetch('/status_code')
+        self.assertEqual(response.code, 200)
+        self.assertTrue(self.application.statsd._sock.closed())
+
+        self.application.statsd.close()
+        response = self.fetch('/status_code')
+        self.assertEqual(response.code, 200)
+        self.assertTrue(self.application.statsd._sock.closed())
+
 
 class TCPStatsdConfigurationTests(testing.AsyncHTTPTestCase):