diff --git a/datadog/dogstatsd/base.py b/datadog/dogstatsd/base.py index c32f32058..4c64e30db 100644 --- a/datadog/dogstatsd/base.py +++ b/datadog/dogstatsd/base.py @@ -505,11 +505,11 @@ def socket(self, new_socket): if new_socket: try: self._socket_kind = new_socket.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) - if self._socket_kind == socket.SOCK_STREAM: - self._transport = "uds-stream" - self._max_payload_size = self._max_buffer_len or UDS_OPTIMAL_PAYLOAD_LENGTH - elif self._socket_kind == socket.SOCK_DGRAM: - self._transport = "uds" + if new_socket.family == socket.AF_UNIX: + if self._socket_kind == socket.SOCK_STREAM: + self._transport = "uds-stream" + else: + self._transport = "uds" self._max_payload_size = self._max_buffer_len or UDS_OPTIMAL_PAYLOAD_LENGTH else: self._transport = "udp" @@ -518,6 +518,8 @@ def socket(self, new_socket): except AttributeError: # _socket can't have a type if it doesn't have sockopts log.info("Unexpected socket provided with no support for getsockopt") self._socket_kind = None + # When the socket is None, we use the UDP optimal payload length + self._max_payload_size = UDP_OPTIMAL_PAYLOAD_LENGTH @property def telemetry_socket(self): diff --git a/tests/unit/dogstatsd/test_statsd.py b/tests/unit/dogstatsd/test_statsd.py index c817ac087..a940a5058 100644 --- a/tests/unit/dogstatsd/test_statsd.py +++ b/tests/unit/dogstatsd/test_statsd.py @@ -42,7 +42,7 @@ class FakeSocket(object): FLUSH_GRACE_PERIOD = 0.2 - def __init__(self, flush_interval=DEFAULT_BUFFERING_FLUSH_INTERVAL, socket_kind=socket.SOCK_DGRAM): + def __init__(self, flush_interval=DEFAULT_BUFFERING_FLUSH_INTERVAL, socket_kind=socket.SOCK_DGRAM, socket_path=None): self.payloads = deque() self._flush_interval = flush_interval @@ -50,6 +50,11 @@ def __init__(self, flush_interval=DEFAULT_BUFFERING_FLUSH_INTERVAL, socket_kind= self._socket_kind = socket_kind self.timeout = () # unit tuple = settimeout was not called + if socket_path: + self.family = socket.AF_UNIX + else: + self.family = socket.AF_INET + def sendall(self, payload): self.send(payload) @@ -768,19 +773,19 @@ def test_udp_socket_ensures_min_receive_buffer(self, mock_socket_create): MIN_SEND_BUFFER_SIZE, ) - def test_socket_path_updates_telemetry(self): + def test_socket_updates_telemetry(self): # Test UDP self.statsd.gauge("foo", 1) self.assert_equal_telemetry("foo:1|g\n", self.recv(2), transport="udp") # Test UDS - self.statsd.socket_path = "/fake/path" + self.statsd.socket = FakeSocket(socket_path="/fake/path") self.statsd._reset_telemetry() self.statsd.gauge("foo", 2) self.assert_equal_telemetry("foo:2|g\n", self.recv(2), transport="uds") # Test UDS stream - self.statsd.socket_path = "unixstream://fake/path" + self.statsd.socket = FakeSocket(socket_path="unixstream://fake/path", socket_kind=socket.SOCK_STREAM) self.statsd._reset_telemetry() self.statsd.gauge("foo", 2) self.assert_equal_telemetry("foo:2|g\n", self.recv(2), transport="uds-stream") @@ -1904,7 +1909,7 @@ def test_default_max_uds_packet_size(self): flush_interval=10000, disable_telemetry=True, ) - dogstatsd.socket = FakeSocket() + dogstatsd.socket = FakeSocket(socket_path=dogstatsd.socket_path) for _ in range(10000): dogstatsd.increment('val') @@ -2110,7 +2115,13 @@ def test_telemetry_api(self): def test_max_payload_size(self): statsd = DogStatsd(socket_path=None, port=8125) self.assertEqual(statsd._max_payload_size, UDP_OPTIMAL_PAYLOAD_LENGTH) - statsd.socket_path = "/foo" + + test_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + statsd.socket = test_socket + self.assertEqual(statsd._max_payload_size, UDP_OPTIMAL_PAYLOAD_LENGTH) + + test_socket = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) + statsd.socket = test_socket self.assertEqual(statsd._max_payload_size, UDS_OPTIMAL_PAYLOAD_LENGTH) def test_post_fork_locks(self):