Add support for keepalive

src/zinolib/ritz.py

@@ -87,7 +87,7 @@
 import select
 
 from .config.tcl import parse_tcl_config  # noqa: F401 (used to be in this file)
-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)
@@ -328,7 +328,7 @@
     """
     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
 
@@ -340,6 +340,7 @@
         self.timeout = timeout
         self.username = username
         self.password = password
+        self.keepalive = keepalive
         self._buff = ""
 
     def __enter__(self):
@@ -448,6 +449,10 @@
         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)
@@ -1090,13 +1095,14 @@
     """
     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()
@@ -1133,6 +1139,10 @@
             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,10 +1,13 @@
 import functools
 import hashlib
+import platform
+import socket
 
 
 __all__ = [
     "windows_codepage_cp1252",
     "generate_authtoken",
+    "enable_socket_keepalive",
 ]
 
 
@@ -80,3 +83,42 @@
                 return return_value
         return wrapper
     return inner
+
+
+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,
+    }
+    if (plat := platform.system()) in platforms:
+        return platforms[plat](sock, after_idle_sec, interval_sec, max_fails)
+    raise RuntimeError('Unsupported platform: {}'.format(plat))