Merge pull request #112 from jsocol/unix-sockets

Introduce UnixSocketStatsClient class.

Fixes #76.

6 files changed, 119 insertions, 1 deletions

diff --git a/docs/configure.rst b/docs/configure.rst
index 27294f1..903b407 100644
--- a/docs/configure.rst
+++ b/docs/configure.rst
@@ -82,6 +82,13 @@ which defaults to ``None``, and is passed to `settimeout
 <https://docs.python.org/2/library/socket.html#socket.socket.settimeout>`.
 
 
+UnixSocket Clients
+------------------
+
+:ref:`UnixSocket-based clients <unix-socket-chapter>` have a single required
+``socket_path`` argument instead of ``host`` and ``port``.
+
+
 In Django
 =========
 
diff --git a/docs/index.rst b/docs/index.rst
index 6e4b5ec..2b58d11 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -71,6 +71,7 @@ Contents
    timing.rst
    pipeline.rst
    tcp.rst
+   unix_socket.rst
    reference.rst
    contributing.rst
 
diff --git a/docs/unix_socket.rst b/docs/unix_socket.rst
new file mode 100644
index 0000000..ac8e4be
--- /dev/null
+++ b/docs/unix_socket.rst
@@ -0,0 +1,16 @@
+.. _unix-socket-chapter:
+
+=====================
+UnixSocketStatsClient
+=====================
+
+The ``UnixSocketStatsClient`` class has a very similar interface to
+``TCPStatsClient``, but internally it uses Unix Domain sockets instead of TCP.
+These are the main differencies when using ``TCPStatsClient`` compared
+to ``UnixSocketStatsClient``:
+
+* Instead of host and port params UnixStatsSocket constructor accepts only socket_path and there is no default value for it.
+
+* There is not ``ipv6`` parameter in constructor.
+
+* You need to make sure that you have correct permissions to write to provided socket.
diff --git a/statsd/__init__.py b/statsd/__init__.py
index e3de0a9..7c56a47 100644
--- a/statsd/__init__.py
+++ b/statsd/__init__.py
@@ -2,8 +2,9 @@ from __future__ import absolute_import
 
 from .client import StatsClient
 from .client import TCPStatsClient
+from .client import UnixSocketStatsClient
 
 
 VERSION = (3, 2, 1)
 __version__ = '.'.join(map(str, VERSION))
-__all__ = ['StatsClient', 'TCPStatsClient']
+__all__ = ['StatsClient', 'TCPStatsClient', 'UnixSocketStatsClient']
diff --git a/statsd/client.py b/statsd/client.py
index 9e46dbd..cb54a00 100644
--- a/statsd/client.py
+++ b/statsd/client.py
@@ -213,6 +213,42 @@ class TCPStatsClient(StatsClientBase):
         self.connect()
 
 
+class UnixSocketStatsClient(StatsClientBase):
+    """Unix domain socket version of StatsClient."""
+
+    def __init__(self, socket_path, prefix=None, timeout=None):
+        """Create a new client."""
+        self._socket_path = socket_path
+        self._timeout = timeout
+        self._prefix = prefix
+        self._sock = None
+
+    def connect(self):
+        self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        self._sock.settimeout(self._timeout)
+        self._sock.connect(self._socket_path)
+
+    def _send(self, data):
+        """Send data to statsd."""
+        if not self._sock:
+            self.connect()
+        self._do_send(data)
+
+    def _do_send(self, data):
+        self._sock.sendall(data.encode('ascii') + b'\n')
+
+    def close(self):
+        self._sock.close()
+        self._sock = None
+
+    def reconnect(self, data):
+        self.close()
+        self.connect()
+
+    def pipeline(self):
+        return TCPPipeline(self)
+
+
 class PipelineBase(StatsClientBase):
 
     def __init__(self, client):
diff --git a/statsd/tests.py b/statsd/tests.py
index 69d5c4f..88dde22 100644
--- a/statsd/tests.py
+++ b/statsd/tests.py
@@ -11,15 +11,18 @@ from nose.tools import eq_
 
 from statsd import StatsClient
 from statsd import TCPStatsClient
+from statsd import UnixSocketStatsClient
 
 
 ADDR = (socket.gethostbyname('localhost'), 8125)
+UNIX_SOCKET = 'tmp.socket'
 
 
 # proto specific methods to get the socket method to send data
 send_method = {
     'udp': lambda x: x.sendto,
     'tcp': lambda x: x.sendall,
+    'unix': lambda x: x.sendall,
 }
 
 
@@ -27,6 +30,7 @@ send_method = {
 make_val = {
     'udp': lambda x, addr: mock.call(str.encode(x), addr),
     'tcp': lambda x, addr: mock.call(str.encode(x + '\n')),
+    'unix': lambda x, addr: mock.call(str.encode(x + '\n')),
 }
 
 
@@ -51,6 +55,15 @@ def _tcp_client(prefix=None, addr=None, port=None, timeout=None, ipv6=False):
     return sc
 
 
+def _unix_socket_client(prefix=None, socket_path=None):
+    if not socket_path:
+        socket_path = UNIX_SOCKET
+
+    sc = UnixSocketStatsClient(socket_path=socket_path, prefix=prefix)
+    sc._sock = mock.Mock()
+    return sc
+
+
 def _timer_check(sock, count, proto, start, end):
     send = send_method[proto](sock)
     eq_(send.call_count, count)
@@ -156,6 +169,13 @@ def test_incr_tcp():
     _test_incr(cl, 'tcp')
 
 
+@mock.patch.object(random, 'random', lambda: -1)
+def test_incr_unix_socket():
+    """TCPStatsClient.incr works."""
+    cl = _unix_socket_client()
+    _test_incr(cl, 'unix')
+
+
 def _test_decr(cl, proto):
     cl.decr('foo')
     _sock_check(cl._sock, 1, proto, 'foo:-1|c')
@@ -184,6 +204,13 @@ def test_decr_tcp():
     _test_decr(cl, 'tcp')
 
 
+@mock.patch.object(random, 'random', lambda: -1)
+def test_decr_unix_socket():
+    """TCPStatsClient.decr works."""
+    cl = _unix_socket_client()
+    _test_decr(cl, 'unix')
+
+
 def _test_gauge(cl, proto):
     cl.gauge('foo', 30)
     _sock_check(cl._sock, 1, proto, 'foo:30|g')
@@ -209,6 +236,13 @@ def test_gauge_tcp():
     _test_gauge(cl, 'tcp')
 
 
+@mock.patch.object(random, 'random', lambda: -1)
+def test_gauge_unix_socket():
+    """TCPStatsClient.decr works."""
+    cl = _unix_socket_client()
+    _test_gauge(cl, 'unix')
+
+
 def _test_ipv6(cl, proto, addr):
     cl.gauge('foo', 30)
     _sock_check(cl._sock, 1, proto, 'foo:30|g', addr=addr)
@@ -393,6 +427,13 @@ def test_timing_supports_timedelta():
     _sock_check(cl._sock, 2, proto, 'foo:129600000.000000|ms')
 
 
+@mock.patch.object(random, 'random', lambda: -1)
+def test_timing_unix_socket():
+    """UnixSocketStatsClient.timing works."""
+    cl = _unix_socket_client()
+    _test_timing(cl, 'unix')
+
+
 def _test_prepare(cl, proto):
     tests = (
         ('foo:1|c', ('foo', '1|c', 1)),
@@ -441,6 +482,13 @@ def test_prefix_tcp():
     _test_prefix(cl, 'tcp')
 
 
+@mock.patch.object(random, 'random', lambda: -1)
+def test_prefix_unix_socket():
+    """UnixSocketStatsClient.incr works."""
+    cl = _unix_socket_client(prefix='foo')
+    _test_prefix(cl, 'unix')
+
+
 def _test_timer_manager(cl, proto):
     with cl.timer('foo'):
         pass
@@ -982,3 +1030,12 @@ def test_tcp_timeout(mock_socket):
     cl = TCPStatsClient(timeout=test_timeout)
     cl.incr('foo')
     cl._sock.settimeout.assert_called_once_with(test_timeout)
+
+
+@mock.patch.object(socket, 'socket')
+def test_unix_socket_timeout(mock_socket):
+    """Timeout on UnixSocketStatsClient should be set on socket."""
+    test_timeout = 321
+    cl = UnixSocketStatsClient(UNIX_SOCKET, timeout=test_timeout)
+    cl.incr('foo')
+    cl._sock.settimeout.assert_called_once_with(test_timeout)