Add blocking socket mode and background sender (#787)

* Add option to use socket timeouts.

settimeout(0) is equivalent to setblocking(0), so need only one call.

Neither UDP or nor unix datagram sockets block on connect, so handling
timeout in get_socket is not necessary.

* Add background sender thread and a queue

* Separate telemetry for dropped packets on write

Align reported metrics with the go client.

* Ignore missing queue type annotations

* Add simple integration test for statsd

* Fix import error on pypy

* Fix formatting

* Use py2.7 compatible syntax

Python2.7 does not allow other positional parameters after *expression.

* Update datadog/dogstatsd/base.py

Co-authored-by: Rémy Mathieu <remy.mathieu@datadoghq.com>

* Update datadog/dogstatsd/base.py

Co-authored-by: Rémy Mathieu <remy.mathieu@datadoghq.com>

* Make parameter name more descriptive

* Mention wait_for_pending in disable_background_sender documentation

* Update tests

---------

Co-authored-by: Rémy Mathieu <remy.mathieu@datadoghq.com>

Author: vickenty

datadog/dogstatsd/base.py

@@ -17,6 +17,12 @@
 import time
 from threading import Lock, RLock
 
+try:
+    import queue
+except ImportError:
+    # pypy has the same module, but capitalized.
+    import Queue as queue  # type: ignore[no-redef]
+
 from typing import Optional, List, Text, Union
 
 # Datadog libraries
@@ -78,8 +84,12 @@
         "datadog.dogstatsd.client.service_checks:%s|c|#%s",
         "datadog.dogstatsd.client.bytes_sent:%s|c|#%s",
         "datadog.dogstatsd.client.bytes_dropped:%s|c|#%s",
+        "datadog.dogstatsd.client.bytes_dropped_queue:%s|c|#%s",
+        "datadog.dogstatsd.client.bytes_dropped_writer:%s|c|#%s",
         "datadog.dogstatsd.client.packets_sent:%s|c|#%s",
         "datadog.dogstatsd.client.packets_dropped:%s|c|#%s",
+        "datadog.dogstatsd.client.packets_dropped_queue:%s|c|#%s",
+        "datadog.dogstatsd.client.packets_dropped_writer:%s|c|#%s",
     ]
 ) + "\n"
 
@@ -110,6 +120,11 @@ def __init__(
         max_buffer_len=0,                       # type: int
         container_id=None,                      # type: Optional[Text]
         origin_detection_enabled=True,          # type: bool
+        socket_timeout=0,                       # type: Optional[float]
+        telemetry_socket_timeout=0,             # type: Optional[float]
+        disable_background_sender=True,         # type: bool
+        sender_queue_size=0,                    # type: int
+        sender_queue_timeout=0,                 # type: Optional[float]
     ):  # type: (...) -> None
         """
         Initialize a DogStatsd object.
@@ -239,6 +254,35 @@ def __init__(
         Default: True.
         More on this: https://docs.datadoghq.com/developers/dogstatsd/?tab=kubernetes#origin-detection-over-udp
         :type origin_detection_enabled: boolean
+
+        :param socket_timeout: Set timeout for socket operations, in seconds. Optional.
+        If sets to zero, never wait if operation can not be completed immediately. If set to None, wait forever.
+        This option does not affect hostname resolution when using UDP.
+        :type socket_timeout: float
+
+        :param telemetry_socket_timeout: Set timeout for the telemetry socket operations. Optional.
+        Effective only if either telemetry_host or telemetry_socket_path are set.
+        If sets to zero, never wait if operation can not be completed immediately. If set to None, wait forever.
+        This option does not affect hostname resolution when using UDP.
+        :type telemetry_socket_timeout: float
+
+        :param disable_background_sender: Use a background thread to communicate with the dogstatsd server. Optional.
+        When enabled, a background thread will be used to send metric payloads to the Agent.
+        Applications should call wait_for_pending() before exiting to make sure all pending payloads are sent.
+        Default: True.
+        :type disable_background_sender: boolean
+
+        :param sender_queue_size: Set the maximum number of packets to queue for the sender. Optional
+        How may packets to queue before blocking or dropping the packet if the packet queue is already full.
+        Default: 0 (unlimited).
+        :type sender_queue_size: integer
+
+        :param sender_queue_timeout: Set timeout for packet queue operations, in seconds. Optional.
+        How long the application thread is willing to wait for the queue clear up before dropping the metric packet.
+        If set to None, wait forever.
+        If set to zero drop the packet immediately if the queue is full.
+        Default: 0 (no wait)
+        :type sender_queue_timeout: float
         """
 
         self._socket_lock = Lock()
@@ -276,6 +320,7 @@ def __init__(
 
         # Connection
         self._max_payload_size = max_buffer_len
+        self.socket_timeout = socket_timeout
         if socket_path is not None:
             self.socket_path = socket_path  # type: Optional[text]
             self.host = None
@@ -294,6 +339,7 @@ def __init__(
         self.telemetry_socket_path = telemetry_socket_path
         self.telemetry_host = None
         self.telemetry_port = None
+        self.telemetry_socket_timeout = telemetry_socket_timeout
         if not telemetry_socket_path and telemetry_host:
             self.telemetry_socket_path = None
             self.telemetry_host = self.resolve_host(telemetry_host, use_default_route)
@@ -367,6 +413,17 @@ def __init__(
         self._flush_thread_stop = threading.Event()
         self._start_flush_thread(self._flush_interval)
 
+        self._queue = None
+        if not disable_background_sender:
+            self._queue = queue.Queue(sender_queue_size)
+            self._start_sender_thread()
+            if sender_queue_timeout is None:
+                self._queue_blocking = True
+                self._queue_timeout = None
+            else:
+                self._queue_blocking = sender_queue_timeout > 0
+                self._queue_timeout = max(0, sender_queue_timeout)
+
     def disable_telemetry(self):
         self._telemetry = False
 
@@ -479,22 +536,25 @@ def get_socket(self, telemetry=False):
                     if self.telemetry_socket_path is not None:
                         self.telemetry_socket = self._get_uds_socket(
                             self.telemetry_socket_path,
+                            self.telemetry_socket_timeout,
                         )
                     else:
                         self.telemetry_socket = self._get_udp_socket(
                             self.telemetry_host,
                             self.telemetry_port,
+                            self.telemetry_socket_timeout,
                         )
 
                 return self.telemetry_socket
 
             if not self.socket:
                 if self.socket_path is not None:
-                    self.socket = self._get_uds_socket(self.socket_path)
+                    self.socket = self._get_uds_socket(self.socket_path, self.socket_timeout)
                 else:
                     self.socket = self._get_udp_socket(
                         self.host,
                         self.port,
+                        self.socket_timeout,
                     )
 
             return self.socket
@@ -513,17 +573,17 @@ def _ensure_min_send_buffer_size(cls, sock, min_size=MIN_SEND_BUFFER_SIZE):
                 pass
 
     @classmethod
-    def _get_uds_socket(cls, socket_path):
+    def _get_uds_socket(cls, socket_path, timeout):
         sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
-        sock.setblocking(0)
+        sock.settimeout(timeout)
         cls._ensure_min_send_buffer_size(sock)
         sock.connect(socket_path)
         return sock
 
     @classmethod
-    def _get_udp_socket(cls, host, port):
+    def _get_udp_socket(cls, host, port, timeout):
         sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        sock.setblocking(0)
+        sock.settimeout(timeout)
         cls._ensure_min_send_buffer_size(sock)
         sock.connect((host, port))
 
@@ -804,9 +864,12 @@ def _reset_telemetry(self):
         self.events_count = 0
         self.service_checks_count = 0
         self.bytes_sent = 0
-        self.bytes_dropped = 0
+        self.bytes_dropped_queue = 0
+        self.bytes_dropped_writer = 0
         self.packets_sent = 0
         self.packets_dropped = 0
+        self.packets_dropped_queue = 0
+        self.packets_dropped_writer = 0
         self._last_flush_time = time.time()
 
     def _flush_telemetry(self):
@@ -821,11 +884,19 @@ def _flush_telemetry(self):
             telemetry_tags,
             self.bytes_sent,
             telemetry_tags,
-            self.bytes_dropped,
+            self.bytes_dropped_queue + self.bytes_dropped_writer,
+            telemetry_tags,
+            self.bytes_dropped_queue,
+            telemetry_tags,
+            self.bytes_dropped_writer,
             telemetry_tags,
             self.packets_sent,
             telemetry_tags,
-            self.packets_dropped,
+            self.packets_dropped_queue + self.packets_dropped_writer,
+            telemetry_tags,
+            self.packets_dropped_queue,
+            telemetry_tags,
+            self.packets_dropped_writer,
             telemetry_tags,
         )
 
@@ -834,7 +905,19 @@ def _is_telemetry_flush_time(self):
             self._last_flush_time + self._telemetry_flush_interval < time.time()
 
     def _send_to_server(self, packet):
-        self._xmit_packet(packet + '\n', False)
+        if self._queue is not None:
+            try:
+                self._queue.put(packet + '\n', self._queue_blocking, self._queue_timeout)
+            except queue.Full:
+                self.packets_dropped_queue += 1
+                self.bytes_dropped_queue += 1
+            return
+
+        self._xmit_packet_with_telemetry(packet + '\n')
+
+    def _xmit_packet_with_telemetry(self, packet):
+        self._xmit_packet(packet, False)
+
         if self._is_telemetry_flush_time():
             telemetry = self._flush_telemetry()
             if self._xmit_packet(telemetry, True):
@@ -844,8 +927,8 @@ def _send_to_server(self, packet):
             else:
                 # Telemetry packet has been dropped, keep telemetry data for the next flush
                 self._last_flush_time = time.time()
-                self.bytes_dropped += len(telemetry)
-                self.packets_dropped += 1
+                self.bytes_dropped_writer += len(telemetry)
+                self.packets_dropped_writer += 1
 
     def _xmit_packet(self, packet, is_telemetry):
         try:
@@ -892,8 +975,8 @@ def _xmit_packet(self, packet, is_telemetry):
             log.error("Unexpected error: %s", str(exc))
 
         if not is_telemetry and self._telemetry:
-            self.bytes_dropped += len(packet)
-            self.packets_dropped += 1
+            self.bytes_dropped_writer += len(packet)
+            self.packets_dropped_writer += 1
 
         return False
 
@@ -1066,5 +1149,28 @@ def _set_container_id(self, container_id, origin_detection_enabled):
                 log.debug("Couldn't get container ID: %s", str(e))
                 self._container_id = None
 
+    def _start_sender_thread(self):
+        log.debug("Starting background sender thread")
+        self._sender_thread = threading.Thread(
+            name="{}_sender_thread".format(self.__class__.__name__),
+            target=self._sender_main_loop,
+        )
+        self._sender_thread.daemon = True
+        self._sender_thread.start()
+
+    def _sender_main_loop(self):
+        while True:
+            self._xmit_packet_with_telemetry(self._queue.get())
+            self._queue.task_done()
+
+    def wait_for_pending(self):
+        """
+        Flush the buffer and wait for all queued payloads to be written to the server.
+        """
+
+        self.flush()
+        if self._queue is not None:
+            self._queue.join()
+
 
 statsd = DogStatsd()

mypy.ini

@@ -2,9 +2,9 @@
 
 # Ignore Mypy errors about packages with missing type hints.
 #
-# Note: If Python 2 support is dropped in the future, `pkg_resources` and
-# `configparser` will no longer need to be ignored because the Python 3
-# versions have type hints.
+# Note: If Python 2 support is dropped in the future, `pkg_resources`,
+# `configparser` and `queue` will no longer need to be ignored because
+# the Python 3 versions have type hints.
 
 [mypy-boto.*]
 ignore_missing_imports = True
@@ -20,3 +20,9 @@ ignore_missing_imports = True
 
 [mypy-pkg_resources.*]
 ignore_missing_imports = True
+
+[mypy-queue.*]
+ignore_missing_imports = True
+
+[mypy-Queue.*]
+ignore_missing_imports = True

tests/integration/dogstatsd/test_statsd_sender.py

@@ -0,0 +1,39 @@
+import itertools
+import socket
+from threading import Thread
+
+import pytest
+
+from datadog.dogstatsd.base import DogStatsd
+
+@pytest.mark.parametrize(
+    "disable_background_sender, disable_buffering, wait_for_pending, socket_timeout",
+    list(itertools.product([True, False], [True, False], [True, False], [0, 1])),
+)
+def test_sender_mode(disable_background_sender, disable_buffering, wait_for_pending, socket_timeout):
+    # Test basic sender operation with an assortment of options
+    foo, bar = socket.socketpair(socket.AF_UNIX, socket.SOCK_DGRAM, 0)
+    statsd = DogStatsd(
+        telemetry_min_flush_interval=0,
+        disable_background_sender=disable_background_sender,
+        disable_buffering=disable_buffering,
+        socket_timeout=socket_timeout,
+    )
+
+    statsd.socket = foo
+    statsd._reset_telemetry()
+
+    def reader_thread():
+        msg = bar.recv(8192)
+        assert msg == b"test.metric:1|c\n"
+
+    t = Thread(target=reader_thread, name="test_sender_mode/reader_thread")
+    t.daemon = True
+    t.start()
+
+    statsd.increment("test.metric")
+    if wait_for_pending:
+        statsd.wait_for_pending()
+
+    t.join(timeout=10)
+    assert not t.is_alive()