diff options
author | Marko Kreen | 2011-06-28 15:28:50 +0000 |
---|---|---|
committer | Marko Kreen | 2011-06-29 07:11:05 +0000 |
commit | e3abfa3f014019fcca3c667209938036f6dc496b (patch) | |
tree | fe5e2e142471a09cbeec244021fb115fdb8ecb70 | |
parent | e44881488d19e2c928fbdef4949f40530a268797 (diff) |
skytools.sockutil: separate module for low-level socket stuff
Move set_tcp_keepalive() here, plus set_nonblocking() and
set_cloexec().
-rw-r--r-- | Makefile | 2 | ||||
-rw-r--r-- | python/skytools/__init__.py | 5 | ||||
-rw-r--r-- | python/skytools/psycopgwrapper.py | 51 | ||||
-rw-r--r-- | python/skytools/sockutil.py | 131 |
4 files changed, 139 insertions, 50 deletions
@@ -24,7 +24,7 @@ SQLDIR = $(prefix)/share/skytools$(SUFFIX) # modules that use doctest for regtests DOCTESTMODS = skytools.quoting skytools.parsing skytools.timeutil \ skytools.sqltools skytools.querybuilder skytools.natsort \ - skytools.utf8 + skytools.utf8 skytools.sockutil all: python-all sub-all config.mak diff --git a/python/skytools/__init__.py b/python/skytools/__init__.py index 837a57e2..22d68300 100644 --- a/python/skytools/__init__.py +++ b/python/skytools/__init__.py @@ -60,7 +60,10 @@ _symbols = { # skytools.psycopgwrapper 'connect_database': 'skytools.psycopgwrapper:connect_database', 'DBError': 'skytools.psycopgwrapper:DBError', - 'set_tcp_keepalive': 'skytools.psycopgwrapper:set_tcp_keepalive', + # skytools.sockutil + 'set_tcp_keepalive': 'skytools.sockutil:set_tcp_keepalive', + 'set_cloexec': 'skytools.sockutil:set_cloexec', + 'set_nonblocking': 'skytools.sockutil:set_nonblocking', # skytools.scripting 'BaseScript': 'skytools.scripting:BaseScript', 'daemonize': 'skytools.scripting:daemonize', diff --git a/python/skytools/psycopgwrapper.py b/python/skytools/psycopgwrapper.py index 753700cd..923b6f39 100644 --- a/python/skytools/psycopgwrapper.py +++ b/python/skytools/psycopgwrapper.py @@ -55,7 +55,7 @@ Plain .fetchall() / .fetchone() give exact same result. """ # no exports -__all__ = ['connect_database', 'set_tcp_keepalive', 'DBError'] +__all__ = ['connect_database', 'DBError'] ##from psycopg2.psycopg1 import connect as _pgconnect # psycopg2.psycopg1.cursor is too backwards compatible, @@ -67,6 +67,8 @@ import psycopg2.extensions, psycopg2.extras from psycopg2 import Error as DBError import skytools +from skytools.sockutil import set_tcp_keepalive + class _CompatRow(psycopg2.extras.DictRow): """Make DictRow more dict-like.""" __slots__ = ('_index',) @@ -104,53 +106,6 @@ class _CompatConnection(psycopg2.extensions.connection): def cursor(self): return psycopg2.extensions.connection.cursor(self, cursor_factory = _CompatCursor) -def set_tcp_keepalive(fd, keepalive = True, - tcp_keepidle = 4 * 60, - tcp_keepcnt = 4, - tcp_keepintvl = 15): - """Turn on TCP keepalive. The fd can be either numeric or socket - object with 'fileno' method. - - OS defaults for SO_KEEPALIVE=1: - - Linux: (7200, 9, 75) - can configure all. - - MacOS: (7200, 8, 75) - can configure only tcp_keepidle. - - Win32: (7200, 5|10, 1) - can configure tcp_keepidle and tcp_keepintvl. - Python needs SIO_KEEPALIVE_VALS support in socket.ioctl to enable it. - - Our defaults: (240, 4, 15). - """ - - # usable on this OS? - if not hasattr(socket, 'SO_KEEPALIVE') or not hasattr(socket, 'fromfd'): - return - - # get numeric fd and cast to socket - if hasattr(fd, 'fileno'): - fd = fd.fileno() - s = socket.fromfd(fd, socket.AF_INET, socket.SOCK_STREAM) - - # skip if unix socket - if type(s.getsockname()) != type(()): - return - - # turn on keepalive on the connection - if keepalive: - s.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) - if hasattr(socket, 'TCP_KEEPCNT'): - s.setsockopt(socket.IPPROTO_TCP, getattr(socket, 'TCP_KEEPIDLE'), tcp_keepidle) - s.setsockopt(socket.IPPROTO_TCP, getattr(socket, 'TCP_KEEPCNT'), tcp_keepcnt) - s.setsockopt(socket.IPPROTO_TCP, getattr(socket, 'TCP_KEEPINTVL'), tcp_keepintvl) - elif hasattr(socket, 'TCP_KEEPALIVE'): - s.setsockopt(socket.IPPROTO_TCP, getattr(socket, 'TCP_KEEPALIVE'), tcp_keepidle) - elif sys.platform == 'darwin': - TCP_KEEPALIVE = 0x10 - s.setsockopt(socket.IPPROTO_TCP, TCP_KEEPALIVE, tcp_keepidle) - elif sys.platform == 'win32': - #s.ioctl(SIO_KEEPALIVE_VALS, (1, tcp_keepidle*1000, tcp_keepintvl*1000)) - pass - else: - s.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 0) - def connect_database(connstr, keepalive = True, tcp_keepidle = 4 * 60, # 7200 tcp_keepcnt = 4, # 9 diff --git a/python/skytools/sockutil.py b/python/skytools/sockutil.py new file mode 100644 index 00000000..f950f5a0 --- /dev/null +++ b/python/skytools/sockutil.py @@ -0,0 +1,131 @@ +"""Various low-level utility functions for sockets.""" + +__all__ = ['set_tcp_keepalive', 'set_nonblocking', 'set_cloexec'] + +import sys +import os +import socket + +try: + import fcntl +except ImportError: + pass + +__all__ = ['set_tcp_keepalive', 'set_nonblocking', 'set_cloexec'] + +def set_tcp_keepalive(fd, keepalive = True, + tcp_keepidle = 4 * 60, + tcp_keepcnt = 4, + tcp_keepintvl = 15): + """Turn on TCP keepalive. The fd can be either numeric or socket + object with 'fileno' method. + + OS defaults for SO_KEEPALIVE=1: + - Linux: (7200, 9, 75) - can configure all. + - MacOS: (7200, 8, 75) - can configure only tcp_keepidle. + - Win32: (7200, 5|10, 1) - can configure tcp_keepidle and tcp_keepintvl. + Python needs SIO_KEEPALIVE_VALS support in socket.ioctl to enable it. + + Our defaults: (240, 4, 15). + + >>> import socket + >>> s = socket.socket() + >>> set_tcp_keepalive(s) + """ + + # usable on this OS? + if not hasattr(socket, 'SO_KEEPALIVE') or not hasattr(socket, 'fromfd'): + return + + # get numeric fd and cast to socket + if hasattr(fd, 'fileno'): + fd = fd.fileno() + s = socket.fromfd(fd, socket.AF_INET, socket.SOCK_STREAM) + + # skip if unix socket + if type(s.getsockname()) != type(()): + return + + # turn on keepalive on the connection + if keepalive: + s.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + if hasattr(socket, 'TCP_KEEPCNT'): + s.setsockopt(socket.IPPROTO_TCP, getattr(socket, 'TCP_KEEPIDLE'), tcp_keepidle) + s.setsockopt(socket.IPPROTO_TCP, getattr(socket, 'TCP_KEEPCNT'), tcp_keepcnt) + s.setsockopt(socket.IPPROTO_TCP, getattr(socket, 'TCP_KEEPINTVL'), tcp_keepintvl) + elif hasattr(socket, 'TCP_KEEPALIVE'): + s.setsockopt(socket.IPPROTO_TCP, getattr(socket, 'TCP_KEEPALIVE'), tcp_keepidle) + elif sys.platform == 'darwin': + TCP_KEEPALIVE = 0x10 + s.setsockopt(socket.IPPROTO_TCP, TCP_KEEPALIVE, tcp_keepidle) + elif sys.platform == 'win32': + #s.ioctl(SIO_KEEPALIVE_VALS, (1, tcp_keepidle*1000, tcp_keepintvl*1000)) + pass + else: + s.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 0) + + +def set_nonblocking(fd, onoff=True): + """Toggle the O_NONBLOCK flag. + + If onoff==None then return current setting. + + Actual sockets from 'socket' module should use .setblocking() method, + this is for situations where it is not available. Eg. pipes + from 'subprocess' module. + + >>> import socket + >>> s = socket.socket() + >>> set_nonblocking(s, None) + False + >>> set_nonblocking(s, 1) + >>> set_nonblocking(s, None) + True + """ + + flags = fcntl.fcntl(fd, fcntl.F_GETFL) + if onoff is None: + return (flags & os.O_NONBLOCK) > 0 + if onoff: + flags |= os.O_NONBLOCK + else: + flags &= ~os.O_NONBLOCK + fcntl.fcntl(fd, fcntl.F_SETFL, flags) + +def set_cloexec(fd, onoff=True): + """Toggle the FD_CLOEXEC flag. + + If onoff==None then return current setting. + + Some libraries do it automatically (eg. libpq). + Others do not (Python stdlib). + + >>> import os + >>> f = open(os.devnull, 'rb') + >>> set_cloexec(f, None) + False + >>> set_cloexec(f, True) + >>> set_cloexec(f, None) + True + >>> import socket + >>> s = socket.socket() + >>> set_cloexec(s, None) + False + >>> set_cloexec(s) + >>> set_cloexec(s, None) + True + """ + + flags = fcntl.fcntl(fd, fcntl.F_GETFD) + if onoff is None: + return (flags & fcntl.FD_CLOEXEC) > 0 + if onoff: + flags |= fcntl.FD_CLOEXEC + else: + flags &= ~fcntl.FD_CLOEXEC + fcntl.fcntl(fd, fcntl.F_SETFD, flags) + +if __name__ == '__main__': + import doctest + doctest.testmod() + |