From 99e024f84bbdb4b4252ea16e43d5c894a56c1795 Mon Sep 17 00:00:00 2001 From: Luis Medel Date: Wed, 4 Dec 2024 16:21:26 +0100 Subject: [PATCH 1/3] Add is_ready and reset() to clients --- statsd/client/base.py | 9 +++++++++ statsd/client/stream.py | 13 +++++++++++++ statsd/client/udp.py | 40 ++++++++++++++++++++++++++++++++++------ 3 files changed, 56 insertions(+), 6 deletions(-) diff --git a/statsd/client/base.py b/statsd/client/base.py index 4fbbe5cb4..d64c653b8 100644 --- a/statsd/client/base.py +++ b/statsd/client/base.py @@ -8,6 +8,15 @@ class StatsClientBase: """A Base class for various statsd clients.""" + @property + def is_ready(self): + """Returns True if the client is ready to send data.""" + raise NotImplementedError() + + def reset(self): + """Reset the client.""" + raise NotImplementedError() + def close(self): """Used to close and clean up any underlying resources.""" raise NotImplementedError() diff --git a/statsd/client/stream.py b/statsd/client/stream.py index 79951d255..d385086f8 100644 --- a/statsd/client/stream.py +++ b/statsd/client/stream.py @@ -18,6 +18,9 @@ def close(self): self._sock.close() self._sock = None + def reset(self): + self.reconnect() + def reconnect(self): self.close() self.connect() @@ -48,6 +51,11 @@ def __init__(self, host='localhost', port=8125, prefix=None, self._prefix = prefix self._sock = None + @property + def is_ready(self): + """Return True if the client is ready to send data.""" + return self._sock is not None + def connect(self): fam = socket.AF_INET6 if self._ipv6 else socket.AF_INET family, _, _, _, addr = socket.getaddrinfo( @@ -67,6 +75,11 @@ def __init__(self, socket_path, prefix=None, timeout=None): self._prefix = prefix self._sock = None + @property + def is_ready(self): + """Return True if the client is ready to send data.""" + return self._sock is not None + def connect(self): self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) self._sock.settimeout(self._timeout) diff --git a/statsd/client/udp.py b/statsd/client/udp.py index ec10fc737..eda70a6ef 100644 --- a/statsd/client/udp.py +++ b/statsd/client/udp.py @@ -26,15 +26,39 @@ class StatsClient(StatsClientBase): """A client for statsd.""" def __init__(self, host='localhost', port=8125, prefix=None, - maxudpsize=512, ipv6=False): + maxudpsize=512, ipv6=False, ignore_socket_errors=False): """Create a new client.""" - fam = socket.AF_INET6 if ipv6 else socket.AF_INET - family, _, _, _, addr = socket.getaddrinfo( - host, port, fam, socket.SOCK_DGRAM)[0] - self._addr = addr - self._sock = socket.socket(family, socket.SOCK_DGRAM) + self._host = host + self._port = port + self._ipv6 = ipv6 self._prefix = prefix self._maxudpsize = maxudpsize + self._addr = None + self._sock = None + self._ignore_socket_errors = ignore_socket_errors + self._ready = False + self._open() + + @property + def is_ready(self): + """Return True if the client is ready to send data.""" + return self._ready + + def _open(self): + fam = socket.AF_INET6 if self._ipv6 else socket.AF_INET + + try: + family, _, _, _, addr = socket.getaddrinfo( + self._host, self._port, fam, socket.SOCK_DGRAM)[0] + self._addr = addr + self._sock = socket.socket(family, socket.SOCK_DGRAM) + self._prefix = self._prefix + self._maxudpsize = self._maxudpsize + self._ready = True + except (socket.gaierror, OSError): + if not self._ignore_socket_errors: + raise + self._ready = False def _send(self, data): """Send data to statsd.""" @@ -44,6 +68,10 @@ def _send(self, data): # No time for love, Dr. Jones! pass + def reset(self): + self.close() + self._open() + def close(self): if self._sock and hasattr(self._sock, 'close'): self._sock.close() From d12dbf36b3d582dc922f621529a47845067b47f4 Mon Sep 17 00:00:00 2001 From: Luis Medel Date: Wed, 4 Dec 2024 16:21:34 +0100 Subject: [PATCH 2/3] Add tests --- statsd/tests.py | 46 ++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 40 insertions(+), 6 deletions(-) diff --git a/statsd/tests.py b/statsd/tests.py index 568897723..0c4bdca13 100644 --- a/statsd/tests.py +++ b/statsd/tests.py @@ -1,5 +1,6 @@ import asyncio import functools +import os import random import re import socket @@ -35,12 +36,18 @@ def eq_(a, b, msg=None): assert a == b, msg -def _udp_client(prefix=None, addr=None, port=None, ipv6=False): +def _is_travis(): + return os.environ.get('TRAVIS') in ('true', '1') + + +def _udp_client(prefix=None, addr=None, port=None, ipv6=False, + ignore_socket_errors=False): if not addr: addr = ADDR[0] if not port: port = ADDR[1] - sc = StatsClient(host=addr, port=port, prefix=prefix, ipv6=ipv6) + sc = StatsClient(host=addr, port=port, prefix=prefix, ipv6=ipv6, + ignore_socket_errors=ignore_socket_errors) sc._sock = mock.Mock() return sc @@ -268,20 +275,47 @@ def _test_resolution(cl, proto, addr): _sock_check(cl._sock, 1, proto, 'foo:1|c', addr=addr) +def test_ipv6_resolution_tcp(): + cl = _tcp_client(addr='localhost', ipv6=True) + _test_resolution(cl, 'tcp', ('::1', 8125, 0, 0)) + + def test_ipv6_resolution_udp(): - raise SkipTest('IPv6 resolution is broken on Travis') + if _is_travis(): + raise SkipTest('IPv6 resolution is broken on Travis') cl = _udp_client(addr='localhost', ipv6=True) _test_resolution(cl, 'udp', ('::1', 8125, 0, 0)) + assert cl.is_ready is True -def test_ipv6_resolution_tcp(): - cl = _tcp_client(addr='localhost', ipv6=True) - _test_resolution(cl, 'tcp', ('::1', 8125, 0, 0)) +def test_ipv6_error_in_resolution_udp(): + if _is_travis(): + raise SkipTest('IPv6 resolution is broken on Travis') + with assert_raises(socket.gaierror): + _ = _udp_client(addr='fakehost', ipv6=True) + + +def test_ipv6_error_in_resolution_ignored_udp(): + if _is_travis(): + raise SkipTest('IPv6 resolution is broken on Travis') + cl = _udp_client(addr='fakehost', ipv6=True, ignore_socket_errors=True) + assert cl.is_ready is False def test_ipv4_resolution_udp(): cl = _udp_client(addr='localhost') _test_resolution(cl, 'udp', ('127.0.0.1', 8125)) + assert cl.is_ready is True + + +def test_ipv4_error_in_resolution_udp(): + with assert_raises(socket.gaierror): + _ = _udp_client(addr='fakehost') + + +def test_ipv4_error_in_resolution_ignored_udp(): + cl = _udp_client(addr='fakehost', ignore_socket_errors=True) + assert cl.is_ready is False def test_ipv4_resolution_tcp(): From 479730b3217d19ba22d0a5c6dfc54a046124ee16 Mon Sep 17 00:00:00 2001 From: Luis Medel Date: Wed, 4 Dec 2024 16:21:40 +0100 Subject: [PATCH 3/3] Update docs --- docs/reference.rst | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/docs/reference.rst b/docs/reference.rst index d85911b36..677589422 100644 --- a/docs/reference.rst +++ b/docs/reference.rst @@ -14,7 +14,7 @@ server supports. information. -.. py:class:: StatsClient(host='localhost', port=8125, prefix=None, maxudpsize=512) +.. py:class:: StatsClient(host='localhost', port=8125, prefix=None, maxudpsize=512, ignore_socket_errors=False) Create a new ``StatsClient`` instance with the appropriate connection and prefix information. @@ -27,11 +27,22 @@ server supports. :param int maxudpsize: the largest safe UDP packet to send. 512 is generally considered safe for the public internet, but private networks may support larger packet sizes. + :param bool ignore_socket_errors: whether or not to ignore socket errors + when initializing the client. If set to ``True``, the client will ignore + the error and continue. If set to ``False``, the client will raise the error. + You can check the ``is_ready`` attribute to see if the client is ready to send. .. py:method:: StatsClient.close() Close the underlying UDP socket. +.. py:method:: StatsClient.reset() + + Resets the underlying UDP socket. This is useful if the connection is + not ready due to network issues or if the connection was closed. This + method will honor the ``ignore_socket_errors`` parameter in the constructor. + You can check the ``is_ready`` attribute to see if the client is ready to send. + .. py:method:: StatsClient.incr(stat, count=1, rate=1) Increment a :ref:`counter `.