Add setters to configure the global instance (#790)

* Add method to enable background sender mode on the global instance

* Add method to set socket timeout on the global instance

Author: vickenty

datadog/dogstatsd/base.py

@@ -415,6 +415,37 @@ def __init__(
 
         self._queue = None
         if not disable_background_sender:
+            self.enable_background_sender(sender_queue_size, sender_queue_timeout)
+
+    def enable_background_sender(self, sender_queue_size=0, sender_queue_timeout=0):
+        """
+        Use a background thread to communicate with the dogstatsd server.
+        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.
+
+        This method is not thread safe and should not be called concurrently with other methods on the current object.
+        Normally, this should be called shortly after process initialization (for example from a post-fork hook in a
+        forking server).
+
+        :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
+        """
+
+        # Avoid a race on _queue with the background buffer flush thread that reads _queue.
+        with self._buffer_lock:
+            if self._queue is not None:
+                return
+
             self._queue = queue.Queue(sender_queue_size)
             self._start_sender_thread()
             if sender_queue_timeout is None:
@@ -559,6 +590,18 @@ def get_socket(self, telemetry=False):
 
             return self.socket
 
+    def set_socket_timeout(self, timeout):
+        """
+        Set timeout for socket operations, in seconds.
+
+        If set 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.
+        """
+        with self._socket_lock:
+            self.socket_timeout = timeout
+            if self.socket:
+                self.socket.settimeout(timeout)
+
     @classmethod
     def _ensure_min_send_buffer_size(cls, sock, min_size=MIN_SEND_BUFFER_SIZE):
         # Increase the receiving buffer size where needed (e.g. MacOS has 4k RX

tests/integration/dogstatsd/test_statsd_sender.py

@@ -37,3 +37,12 @@ def reader_thread():
 
     t.join(timeout=10)
     assert not t.is_alive()
+
+def test_set_socket_timeout():
+    statsd = DogStatsd(socket_timeout=0)
+    assert statsd.get_socket().gettimeout() == 0
+    statsd.set_socket_timeout(1)
+    assert statsd.get_socket().gettimeout() == 1
+    statsd.close_socket()
+    assert statsd.get_socket().gettimeout() == 1
+

tests/unit/dogstatsd/test_statsd.py

@@ -46,6 +46,7 @@ def __init__(self, flush_interval=DEFAULT_FLUSH_INTERVAL):
 
         self._flush_interval = flush_interval
         self._flush_wait = False
+        self.timeout = () # unit tuple = settimeout was not called
 
     def send(self, payload):
         if is_p3k():
@@ -77,6 +78,8 @@ def close(self):
     def __repr__(self):
         return str(self.payloads)
 
+    def settimeout(self, timeout):
+        self.timeout = timeout
 
 class BrokenSocket(FakeSocket):
     def __init__(self, error_number=None):
@@ -1875,6 +1878,9 @@ def test_sender_mode(self):
         statsd = DogStatsd(disable_background_sender=True)
         self.assertIsNone(statsd._queue)
 
+        statsd.enable_background_sender()
+        self.assertIsNotNone(statsd._queue)
+
         statsd = DogStatsd(disable_background_sender=False)
         self.assertIsNotNone(statsd._queue)
 
@@ -1886,3 +1892,10 @@ def test_sender_calls_task_done(self):
 
     def test_sender_queue_no_timeout(self):
         statsd = DogStatsd(disable_background_sender=False, sender_queue_timeout=None)
+
+    def test_set_socket_timeout(self):
+        statsd = DogStatsd(disable_background_sender=False)
+        statsd.socket = FakeSocket()
+        statsd.set_socket_timeout(1)
+        self.assertEqual(statsd.socket.timeout, 1)
+        self.assertEqual(statsd.socket_timeout, 1)