Add support for TCP keepalive

Add configurable support for TCP keepalive for protocol port (8001) and
ntie/updates port (8002).

Turned on by default, via a flag "keepalive" on __init__.

src/zinolib/ritz.py

@@ -88,7 +88,7 @@
 import codecs
 import select
 
-from .utils import windows_codepage_cp1252, generate_authtoken
+from .utils import windows_codepage_cp1252, generate_authtoken, enable_socket_keepalive
 
 
 codecs.register_error("windows_codepage_cp1252", windows_codepage_cp1252)
@@ -347,7 +347,7 @@ class ritz:
     """
     DELIMITER = "\r\n"
 
-    def __init__(self, server, port=8001, timeout=10, username=None, password=None):
+    def __init__(self, server, port=8001, timeout=10, username=None, password=None, keepalive=True):
         """Initialize"""
         global logger
 
@@ -358,6 +358,7 @@ def __init__(self, server, port=8001, timeout=10, username=None, password=None):
         self.timeout = timeout
         self.username = username
         self.password = password
+        self.keepalive = keepalive
         self._buff = ""
 
     def __enter__(self):
@@ -463,6 +464,10 @@ def connect(self):
         else:
             raise NotConnectedError("Did not get a status code 200")
 
+        if self.keepalive:
+            enable_socket_keepalive(self._sock)
+            logger.info("Set keepalive on protocol socket")
+
         # Automaticly authenticate if username and password is supplied
         if self.username and self.password:
             self.authenticate(self.username, self.password)
@@ -1105,13 +1110,14 @@ class notifier:
     """
     DELIMITER = "\r\n"
 
-    def __init__(self, zino_session, port=8002, timeout=30):
+    def __init__(self, zino_session, port=8002, timeout=30, keepalive=True):
         self._sock = None
         self.connStatus = False
         self._buff = ""
         self.zino_session = zino_session
         self.port = port
         self.timeout = timeout
+        self.keepalive = keepalive
 
     def __enter__(self):
         self.connect()
@@ -1146,6 +1152,10 @@ def connect(self):
             else:
                 raise NotConnectedError("Key not found")
 
+            if self.keepalive:
+                enable_socket_keepalive(self._sock)
+                logger.info("Set keepalive on notifier socket")
+
     def poll(self, timeout=0):
         """Poll the notifier socket for new data
 

src/zinolib/utils.py

@@ -1,9 +1,12 @@
 import hashlib
+import platform
+import socket
 
 
 __all__ = [
     "windows_codepage_cp1252",
     "generate_authtoken",
+    "enable_socket_keepalive",
 ]
 
 
@@ -56,3 +59,43 @@ def generate_authtoken(challenge, password):
     raw_token = "%s %s" % (challenge, password)
     token = hashlib.sha1(raw_token.encode("UTF-8")).hexdigest()
     return token
+
+
+def _enable_keepalive_linux(sock, after_idle_sec, interval_sec, max_fails):
+    """Set TCP keepalive on an open socket.
+
+    It activates after 1 second (after_idle_sec) of idleness,
+    then sends a keepalive ping once every 3 seconds (interval_sec),
+    and closes the connection after 5 failed ping (max_fails), or 15 seconds
+    """
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, after_idle_sec)
+    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, interval_sec)
+    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, max_fails)
+
+
+def _enable_keepalive_osx(sock, after_idle_sec, interval_sec, max_fails):
+    """Set TCP keepalive on an open socket.
+
+    sends a keepalive ping once every 3 seconds (interval_sec)
+    """
+    # scraped from /usr/include, not exported by python's socket module
+    TCP_KEEPALIVE = 0x10
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+    sock.setsockopt(socket.IPPROTO_TCP, TCP_KEEPALIVE, interval_sec)
+
+
+def _enable_keepalive_win(sock, after_idle_sec, interval_sec, max_fails):
+    sock.ioctl(socket.SIO_KEEPALIVE_VALS, (1, after_idle_sec * 1000, interval_sec * 1000))
+
+
+def enable_socket_keepalive(sock, after_idle_sec=60, interval_sec=60, max_fails=5):
+    platforms = {
+        "Linux": _enable_keepalive_linux,
+        "Darwin": _enable_keepalive_osx,
+        "Windows": _enable_keepalive_win,
+    }
+    plat = platform.system()
+    if plat in platforms:
+        return platforms[plat](sock, after_idle_sec=60, interval_sec=60, max_fails=5)
+    raise RuntimeError('Unsupported platform: {}'.format(plat))