Skip to content

Commit

Permalink
add additional testing for the stream socket addition
Browse files Browse the repository at this point in the history
  • Loading branch information
ddrthall committed Dec 4, 2024
1 parent c9dd150 commit e3b82a9
Show file tree
Hide file tree
Showing 3 changed files with 55 additions and 6 deletions.
18 changes: 12 additions & 6 deletions datadog/dogstatsd/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -503,9 +503,12 @@ def socket(self):
def socket(self, new_socket):
self._socket = new_socket
if new_socket:
self._socket_kind = new_socket.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE)
else:
self._socket_kind = None
try:
self._socket_kind = new_socket.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE)
return
except AttributeError: # _socket can't have a type if it doesn't have sockopts
pass
self._socket_kind = None

@property
def telemetry_socket(self):
Expand All @@ -515,9 +518,12 @@ def telemetry_socket(self):
def telemetry_socket(self, t_socket):
self._telemetry_socket = t_socket
if t_socket:
self._telemetry_socket_kind = t_socket.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE)
else:
self._telemetry_socket_kind = None
try:
self._telemetry_socket_kind = t_socket.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE)
return
except AttributeError:# _telemetry_socket can't have a kind if it doesn't have sockopts
pass
self._telemetry_socket_kind = None

def enable_background_sender(self, sender_queue_size=0, sender_queue_timeout=0):
"""
Expand Down
21 changes: 21 additions & 0 deletions tests/integration/dogstatsd/test_statsd_sender.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,27 @@ def test_set_socket_timeout():
statsd.close_socket()
assert statsd.get_socket().gettimeout() == 1

def test_stream_cleanup():
foo, _ = socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM, 0)

foo.settimeout(0)
statsd = DogStatsd(disable_buffering=True)
statsd.socket = foo
statsd.increment("test", 1)
statsd.increment("test", 1)
statsd.increment("test", 1)
assert statsd.socket is not None

foo.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1) # different os's have different mins, e.g. this sets the buffer size to 2304 on certain linux variants

try:
foo.sendall(os.urandom(5000)) # pre-emptively clog the buffer
except BlockingIOError:
pass

statsd.increment("test", 1)

assert statsd.socket is None

@pytest.mark.parametrize(
"disable_background_sender, disable_buffering",
Expand Down
22 changes: 22 additions & 0 deletions tests/unit/dogstatsd/test_statsd.py
Original file line number Diff line number Diff line change
Expand Up @@ -2116,3 +2116,25 @@ def inner():
t.start()
t.join(timeout=5)
self.assertFalse(t.is_alive())

def test_fake_sockets(self):
"""
To support legacy behavior wherein customers were able to set sockets directly as long as they supported a .send interface,
ensure that arbitrary values passed to these properties are allowed and are handled correctly
"""
statsd = DogStatsd(disable_buffering=True)

class fake_sock:
def __init__(self, id):
self.id = id
def send(self, _):
pass
statsd.socket = fake_sock(5)
statsd.telemetry_socket = fake_sock(10)

assert statsd.socket.id == 5
assert statsd.telemetry_socket.id == 10

statsd.increment("test", 1)

assert statsd.socket is not None

0 comments on commit e3b82a9

Please sign in to comment.