[AMLII-2170] fix removed/renamed function flush (#868)

* fix missing function that was renamed

* fix lint

* change variable name back to original

Author: andrewqian2001datadog

datadog/dogstatsd/base.py

@@ -13,6 +13,7 @@
 import os
 import socket
 import errno
+import struct
 import threading
 import time
 from threading import Lock, RLock
@@ -49,8 +50,13 @@
 DEFAULT_HOST = "localhost"
 DEFAULT_PORT = 8125
 
+# Socket prefixes
+UNIX_ADDRESS_SCHEME = "unix://"
+UNIX_ADDRESS_DATAGRAM_SCHEME = "unixgram://"
+UNIX_ADDRESS_STREAM_SCHEME = "unixstream://"
+
 # Buffering-related values (in seconds)
-DEFAULT_FLUSH_INTERVAL = 0.3
+DEFAULT_BUFFERING_FLUSH_INTERVAL = 0.3
 MIN_FLUSH_INTERVAL = 0.0001
 
 # Env var to enable/disable sending the container ID field
@@ -145,7 +151,7 @@ def __init__(
         host=DEFAULT_HOST,                      # type: Text
         port=DEFAULT_PORT,                      # type: int
         max_buffer_size=None,                   # type: None
-        flush_interval=DEFAULT_FLUSH_INTERVAL,  # type: float
+        flush_interval=DEFAULT_BUFFERING_FLUSH_INTERVAL,  # type: float
         disable_aggregation=True,               # type: bool
         disable_buffering=True,                 # type: bool
         namespace=None,                         # type: Optional[Text]
@@ -489,6 +495,30 @@ def socket_path(self, path):
                 self._transport = "uds"
                 self._max_payload_size = self._max_buffer_len or UDS_OPTIMAL_PAYLOAD_LENGTH
 
+    @property
+    def socket(self):
+        return self._socket
+
+    @socket.setter
+    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
+
+    @property
+    def telemetry_socket(self):
+        return self._telemetry_socket
+
+    @telemetry_socket.setter
+    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
+
     def enable_background_sender(self, sender_queue_size=0, sender_queue_timeout=0):
         """
         Use a background thread to communicate with the dogstatsd server.
@@ -643,7 +673,7 @@ def disable_aggregation(self):
                 self._stop_flush_thread()
             log.debug("Statsd aggregation is disabled")
 
-    def enable_aggregation(self, flush_interval=DEFAULT_FLUSH_INTERVAL):
+    def enable_aggregation(self, flush_interval=DEFAULT_BUFFERING_FLUSH_INTERVAL):
         with self._config_lock:
             if not self._disable_aggregation:
                 return
@@ -731,11 +761,37 @@ def _ensure_min_send_buffer_size(cls, sock, min_size=MIN_SEND_BUFFER_SIZE):
 
     @classmethod
     def _get_uds_socket(cls, socket_path, timeout):
-        sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
-        sock.settimeout(timeout)
-        cls._ensure_min_send_buffer_size(sock)
-        sock.connect(socket_path)
-        return sock
+        valid_socket_kinds = [socket.SOCK_DGRAM, socket.SOCK_STREAM]
+        if socket_path.startswith(UNIX_ADDRESS_DATAGRAM_SCHEME):
+            valid_socket_kinds = [socket.SOCK_DGRAM]
+            socket_path = socket_path[len(UNIX_ADDRESS_DATAGRAM_SCHEME):]
+        elif socket_path.startswith(UNIX_ADDRESS_STREAM_SCHEME):
+            valid_socket_kinds = [socket.SOCK_STREAM]
+            socket_path = socket_path[len(UNIX_ADDRESS_STREAM_SCHEME):]
+        elif socket_path.startswith(UNIX_ADDRESS_SCHEME):
+            socket_path = socket_path[len(UNIX_ADDRESS_SCHEME):]
+
+        last_error = ValueError("Invalid socket path")
+        for socket_kind in valid_socket_kinds:
+            # py2 stores socket kinds differently than py3, determine the name independently from version
+            sk_name = { socket.SOCK_STREAM: "stream", socket.SOCK_DGRAM: "datagram" }[socket_kind]
+
+            try:
+                sock = socket.socket(socket.AF_UNIX, socket_kind)
+                sock.settimeout(timeout)
+                cls._ensure_min_send_buffer_size(sock)
+                sock.connect(socket_path)
+                log.debug("Connected to socket %s with kind %s", socket_path, sk_name)
+                return sock
+            except Exception as e:
+                if sock is not None:
+                    sock.close()
+                log.debug("Failed to connect to %s with kind %s: %s", socket_path, sk_name, e)
+                if e.errno == errno.EPROTOTYPE:
+                    last_error = e
+                    continue
+                raise e
+        raise last_error
 
     @classmethod
     def _get_udp_socket(cls, host, port, timeout):
@@ -805,6 +861,9 @@ def _reset_buffer(self):
             self._current_buffer_total_size = 0
             self._buffer = []
 
+    def flush(self):
+        self.flush_buffered_metrics()
+
     def flush_buffered_metrics(self):
         """
         Flush the metrics buffer by sending the data to the server.
@@ -1216,10 +1275,14 @@ def _xmit_packet(self, packet, is_telemetry):
         try:
             if is_telemetry and self._dedicated_telemetry_destination():
                 mysocket = self.telemetry_socket or self.get_socket(telemetry=True)
+                socket_kind = self._telemetry_socket_kind
             else:
                 # If set, use socket directly
                 mysocket = self.socket or self.get_socket()
+                socket_kind = self._socket_kind
 
+            if socket_kind == socket.SOCK_STREAM:
+                mysocket.send(struct.pack('<I', len(packet)))
             mysocket.send(packet.encode(self.encoding))
 
             if not is_telemetry and self._telemetry:
@@ -1253,7 +1316,7 @@ def _xmit_packet(self, packet, is_telemetry):
                 )
                 self.close_socket()
         except Exception as exc:
-            print("Unexpected error: %s", exc)
+            print("Unexpected error: ", exc)
             log.error("Unexpected error: %s", str(exc))
 
         if not is_telemetry and self._telemetry:

tests/integration/dogstatsd/test_statsd_sender.py

@@ -1,18 +1,23 @@
+from contextlib import closing
 import itertools
+import os
+import shutil
 import socket
+import tempfile
 from threading import Thread
+import uuid
 
 import pytest
 
 from datadog.dogstatsd.base import DogStatsd
 
 @pytest.mark.parametrize(
-    "disable_background_sender, disable_buffering, wait_for_pending, socket_timeout, stop",
-    list(itertools.product([True, False], [True, False], [True, False], [0, 1], [True, False])),
+    "disable_background_sender, disable_buffering, wait_for_pending, socket_timeout, stop, socket_kind",
+    list(itertools.product([True, False], [True, False], [True, False], [0, 1], [True, False], [socket.SOCK_DGRAM, socket.SOCK_STREAM])),
 )
-def test_sender_mode(disable_background_sender, disable_buffering, wait_for_pending, socket_timeout, stop):
+def test_sender_mode(disable_background_sender, disable_buffering, wait_for_pending, socket_timeout, stop, socket_kind):
     # Test basic sender operation with an assortment of options
-    foo, bar = socket.socketpair(socket.AF_UNIX, socket.SOCK_DGRAM, 0)
+    foo, bar = socket.socketpair(socket.AF_UNIX, socket_kind, 0)
     statsd = DogStatsd(
         telemetry_min_flush_interval=0,
         disable_background_sender=disable_background_sender,
@@ -101,3 +106,41 @@ def test_buffering_with_context():
     bar.settimeout(5)
     msg = bar.recv(8192)
     assert msg == b"first:1|c\n"
+
+@pytest.fixture()
+def socket_dir():
+    tempdir = tempfile.mkdtemp()
+    yield tempdir
+    shutil.rmtree(tempdir)
+
+@pytest.mark.parametrize(
+        "socket_prefix, socket_kind, success",
+        [
+            ("", socket.SOCK_DGRAM, True),
+            ("", socket.SOCK_STREAM, True),
+            ("unix://", socket.SOCK_DGRAM, True),
+            ("unix://", socket.SOCK_STREAM, True),
+            ("unixstream://", socket.SOCK_DGRAM, False),
+            ("unixstream://", socket.SOCK_STREAM, True),
+            ("unixgram://", socket.SOCK_DGRAM, True),
+            ("unixgram://", socket.SOCK_STREAM, False)
+        ]
+)
+def test_socket_connection(socket_dir, socket_prefix, socket_kind, success):
+    socket_path = os.path.join(socket_dir, str(uuid.uuid1()) + ".sock")
+    listener_socket = socket.socket(socket.AF_UNIX, socket_kind)
+    listener_socket.bind(socket_path)
+
+    if socket_kind == socket.SOCK_STREAM:
+        listener_socket.listen(1)
+
+    with closing(listener_socket):
+        statsd = DogStatsd(
+            socket_path = socket_prefix + socket_path
+        )
+
+        if success:
+            assert statsd.get_socket() is not None
+        else:
+            with pytest.raises(socket.error):
+                statsd.get_socket()