From ee1bb920d74cb23230322d338a04501d883fecaf Mon Sep 17 00:00:00 2001 From: Vlad Yermakov Date: Mon, 27 May 2024 13:17:19 +0300 Subject: [PATCH] Reconnect to socket on exception --- statsd/client/stream.py | 6 +++++- statsd/tests.py | 7 ++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/statsd/client/stream.py b/statsd/client/stream.py index 79951d255..6d71af647 100644 --- a/statsd/client/stream.py +++ b/statsd/client/stream.py @@ -29,7 +29,11 @@ def _send(self, data): """Send data to statsd.""" if not self._sock: self.connect() - self._do_send(data) + try: + self._do_send(data) + except (OSError, RuntimeError): + self.reconnect() + self._do_send(data) def _do_send(self, data): self._sock.sendall(data.encode('ascii') + b'\n') diff --git a/statsd/tests.py b/statsd/tests.py index 568897723..23cc855ac 100644 --- a/statsd/tests.py +++ b/statsd/tests.py @@ -1037,15 +1037,16 @@ def test_pipeline_packet_size(): @mock.patch.object(socket, 'socket') -def test_tcp_raises_exception_to_user(mock_socket): +def test_tcp_reconnect_on_exception(mock_socket): """Socket errors in TCPStatsClient should be raised to user.""" addr = ('127.0.0.1', 1234) cl = _tcp_client(addr=addr[0], port=addr[1]) + cl.reconnect = mock.Mock(wraps=cl.reconnect) cl.incr('foo') eq_(1, cl._sock.sendall.call_count) cl._sock.sendall.side_effect = socket.error - with assert_raises(socket.error): - cl.incr('foo') + cl.incr('foo') + assert 1 == cl.reconnect.call_count @mock.patch.object(socket, 'socket')