danielperna84 · April 3, 2019 16:10 · Apr 3, 2019
diff --git a/uds_socketserver.py b/uds_socketserver.py
@@ -0,0 +1,80 @@
+#!/usr/bin/python
+import os
+import sys
+import time
+import signal
+import logging
+import multiprocessing
+from multiprocessing import Process
+import SocketServer
+
+SERVER = None
+LOGLEVEL_MAPPING = {
+    "critical": logging.CRITICAL,
+    "error": logging.ERROR,
+    "warning": logging.WARNING,
+    "info": logging.INFO,
+    "debug": logging.DEBUG
+}
+DEFAULT_LOGLEVEL = "debug"
+LOGLEVEL = LOGLEVEL_MAPPING.get(os.environ.get("LOGLEVEL", DEFAULT_LOGLEVEL))
+LOG = logging.getLogger(__name__)
+LOG.setLevel(LOGLEVEL)
+SO = logging.StreamHandler(sys.stdout)
+SO.setLevel(LOGLEVEL)
+SO.setFormatter(
+    logging.Formatter('%(levelname)s:%(asctime)s:%(name)s:%(message)s'))
+LOG.addHandler(SO)
+SERVER_ADDRESS = './uds.sock'
+
+try:
+    os.unlink(SERVER_ADDRESS)
+except OSError:
+    if os.path.exists(SERVER_ADDRESS):
+        raise
+
+def signal_handler(sig, frame):
+    """Handle signal to shut down server."""
+    global SERVER
+    LOG.warning("Got signal: %s. Shutting down server." % str(sig))
+    SERVER.server_close()
+    for child in multiprocessing.active_children():
+        LOG.warning("Terminating: %s" % child)
+        child.terminate()
+    sys.exit(0)
+
+def sleeper(duration):
+    """Dummy process that would block a single-threaded server."""
+    LOG.debug("sleeping for %s seconds" % duration)
+    time.sleep(int(duration))
+    LOG.debug("done sleeping")
+
+class ThreadingUDSServer(SocketServer.ThreadingMixIn, SocketServer.UnixStreamServer):
+    pass
+
+class SockHandler(SocketServer.BaseRequestHandler):
+    """
+    The request handler class for our server.
+    """
+    def handle(self):
+        self.data = self.request.recv(1024).strip()
+        LOG.debug(self.data)
+        if "status" in self.data.split(':'):
+            self.request.sendall("processes:%s" % len(multiprocessing.active_children()))
+        elif "sleep" in self.data.split(':'):
+            _, duration = self.data.split(':')
+            p = Process(target=sleeper, args=(duration,))
+            p.start()
+            self.request.sendall("started:%s" % p.name)
+        else:
+            self.request.sendall("unknown-command")
+
+def main():
+    global SERVER
+    signal.signal(signal.SIGINT, signal_handler)
+    SERVER = ThreadingUDSServer(SERVER_ADDRESS, SockHandler)
+    LOG.info("Created socket at %s" % SERVER_ADDRESS)
+    SERVER.serve_forever()
+
+if __name__ == "__main__":
+    main()