"""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. 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 # need socket object if isinstance(fd, socket.SocketType): s = fd else: 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 # no keepalive? if not keepalive: s.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 0) return # basic keepalive s.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) # detect available options TCP_KEEPCNT = getattr(socket, 'TCP_KEEPCNT', None) TCP_KEEPINTVL = getattr(socket, 'TCP_KEEPINTVL', None) TCP_KEEPIDLE = getattr(socket, 'TCP_KEEPIDLE', None) TCP_KEEPALIVE = getattr(socket, 'TCP_KEEPALIVE', None) SIO_KEEPALIVE_VALS = getattr(socket, 'SIO_KEEPALIVE_VALS', None) if TCP_KEEPIDLE is None and TCP_KEEPALIVE is None and sys.platform == 'darwin': TCP_KEEPALIVE = 0x10 # configure if TCP_KEEPCNT is not None: s.setsockopt(socket.IPPROTO_TCP, TCP_KEEPCNT, tcp_keepcnt) if TCP_KEEPINTVL is not None: s.setsockopt(socket.IPPROTO_TCP, TCP_KEEPINTVL, tcp_keepintvl) if TCP_KEEPIDLE is not None: s.setsockopt(socket.IPPROTO_TCP, TCP_KEEPIDLE, tcp_keepidle) elif TCP_KEEPALIVE is not None: s.setsockopt(socket.IPPROTO_TCP, TCP_KEEPALIVE, tcp_keepidle) elif SIO_KEEPALIVE_VALS is not None: s.ioctl(SIO_KEEPALIVE_VALS, (1, tcp_keepidle*1000, tcp_keepintvl*1000)) 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()