fix socket communications (#81)

For the `UdpWriter`, only open the socket once, during initialization. While
this has not been an issue in practice, and performs comparably, it leads to
port churn on the host that uses this. A single open socket will lead to a
single port being used for communication.

The `FileWriter` is not sufficient for communicating with Unix Domain Sockets,
and we need to add a `UnixWriter` to handle this correctly. Verified in tests.

Author: copperlight

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = netflix-spectator-py
-version = 1.0.3
+version = 1.0.4
 description = Library for reporting metrics from Python applications to SpectatorD and the Netflix Atlas Timeseries Database.
 long_description = file: README.md
 long_description_content_type = text/markdown

spectator/writer/new_writer.py

@@ -6,6 +6,7 @@
 from spectator.writer.memory_writer import MemoryWriter
 from spectator.writer.noop_writer import NoopWriter
 from spectator.writer.udp_writer import UdpWriter
+from spectator.writer.unix_writer import UnixWriter
 
 WriterUnion = Union[FileWriter, MemoryWriter, NoopWriter, UdpWriter]
 
@@ -15,7 +16,8 @@ def is_valid_output_location(location: str) -> bool:
         return False
     return location in ["none", "memory", "stdout", "stderr", "udp", "unix"] or \
         location.startswith("file://") or \
-        location.startswith("udp://")
+        location.startswith("udp://") or \
+        location.startswith("unix://")
 
 
 def new_writer(location: str) -> WriterUnion:
@@ -37,16 +39,17 @@ def new_writer(location: str) -> WriterUnion:
         writer = UdpWriter(location, address)
     elif location == "unix":
         # default unix domain socket for spectatord
-        location = "file:///run/spectatord/spectatord.unix"
-        file = open(urlparse(location).path, "a", encoding="utf-8")
-        writer = FileWriter(location, file)
+        location = "unix:///run/spectatord/spectatord.unix"
+        writer = UnixWriter(urlparse(location).path)
     elif location.startswith("file://"):
         file = open(urlparse(location).path, "a", encoding="utf-8")
         writer = FileWriter(location, file)
     elif location.startswith("udp://"):
         parsed = urlparse(location)
         address = (parsed.hostname, parsed.port)
         writer = UdpWriter(location, address)
+    elif location.startswith("unix://"):
+        writer = UnixWriter(urlparse(location).path)
     else:
         raise ValueError(f"unsupported Writer location: {location}")
 

spectator/writer/udp_writer.py

@@ -22,14 +22,15 @@ def __init__(self, location: str, address: Tuple[str, int]) -> None:
             # anything that does not appear to be an IPv4 or IPv6 address (i.e. hostnames)
             self._family = socket.AF_INET
 
+        self._sock = socket.socket(family=self._family, type=socket.SOCK_DGRAM)
+
     def write(self, line: str) -> None:
         self._logger.debug("write line=%s", line)
 
         try:
-            with socket.socket(self._family, socket.SOCK_DGRAM) as s:
-                s.sendto(bytes(line, encoding="utf-8"), self._address)
+            self._sock.sendto(bytes(line, encoding="utf-8"), self._address)
         except IOError:
             self._logger.error("failed to write line=%s", line)
 
     def close(self) -> None:
-        pass
+        self._sock.close()

spectator/writer/unix_writer.py

@@ -0,0 +1,24 @@
+import socket
+
+from spectator.writer import Writer
+
+
+class UnixWriter(Writer):
+    """Writer that outputs data to a Unix Domain Socket."""
+
+    def __init__(self, location: str) -> None:
+        super().__init__()
+        self._logger.debug("initialize UnixWriter to %s", location)
+        self._sock = socket.socket(family=socket.AF_UNIX, type=socket.SOCK_DGRAM)
+        self._location = location
+
+    def write(self, line: str) -> None:
+        self._logger.debug("write line=%s", line)
+
+        try:
+            self._sock.sendto(bytes(line, encoding="utf-8"), self._location)
+        except IOError:
+            self._logger.error("failed to write line=%s", line)
+
+    def close(self) -> None:
+        self._sock.close()

tests/udp_server.py

@@ -4,8 +4,8 @@
 
 class UdpServer:
 
-    def __init__(self, address: Tuple[str, int] = ("127.0.0.1", 0)):
-        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    def __init__(self, address: Tuple[str, int] = ("127.0.0.1", 0)) -> None:
+        self._sock = socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM)
         self._sock.bind(address)
 
     def address(self) -> str:

tests/unix_server.py

@@ -0,0 +1,24 @@
+import os
+import random
+import socket
+
+
+class UnixServer:
+
+    def __init__(self, path: str = f"/tmp/spectatord-test-{random.randint(1, 10000)}") -> None:
+        if os.path.exists(path):
+            os.remove(path)
+        self._path = path
+        self._sock = socket.socket(family=socket.AF_UNIX, type=socket.SOCK_DGRAM)
+        self._sock.bind(path)
+
+    def address(self) -> str:
+        return "unix://{}".format(self._path)
+
+    def read(self) -> str:
+        data, _ = self._sock.recvfrom(1024)
+        return data.decode(encoding="utf-8")
+
+    def close(self) -> None:
+        self._sock.close()
+        os.remove(self._path)

tests/writer/test_new_writer.py

@@ -0,0 +1,18 @@
+import unittest
+
+from spectator.writer.new_writer import is_valid_output_location
+
+
+class IsValidOutputLocationTest(unittest.TestCase):
+
+    def test_is_valid_output_location(self):
+        self.assertTrue(is_valid_output_location("none"))
+        self.assertTrue(is_valid_output_location("memory"))
+        self.assertTrue(is_valid_output_location("stdout"))
+        self.assertTrue(is_valid_output_location("stderr"))
+        self.assertTrue(is_valid_output_location("udp"))
+        self.assertTrue(is_valid_output_location("unix"))
+        self.assertTrue(is_valid_output_location("file://testfile.txt"))
+        self.assertTrue(is_valid_output_location("udp://localhost:1234"))
+        self.assertTrue(is_valid_output_location("unix:///tmp/socket.sock"))
+        self.assertFalse(is_valid_output_location("invalid"))

tests/writer/test_unix_writer.py

@@ -0,0 +1,16 @@
+import unittest
+from contextlib import closing
+
+from spectator import new_writer
+from ..unix_server import UnixServer
+
+
+class UnixWriterTest(unittest.TestCase):
+
+    def test_unix(self) -> None:
+        with closing(UnixServer()) as server:
+            with closing(new_writer(server.address())) as w:
+                w.write("foo")
+                self.assertEqual("foo", server.read())
+                w.write("bar")
+                self.assertEqual("bar", server.read())