diff --git a/src/zinolib/ritz.py b/src/zinolib/ritz.py index c401f01..6067e70 100644 --- a/src/zinolib/ritz.py +++ b/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 @@ 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 @@ -340,6 +340,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): @@ -448,6 +449,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) @@ -1090,13 +1095,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() @@ -1133,6 +1139,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 diff --git a/src/zinolib/utils.py b/src/zinolib/utils.py index 7c07f21..2b95b26 100644 --- a/src/zinolib/utils.py +++ b/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 @@ def wrapper(*args, **kwargs): 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))