skytools.sockutil: separate module for low-level socket stuff

Move set_tcp_keepalive() here, plus set_nonblocking() and
set_cloexec().

4 files changed, 139 insertions, 50 deletions

diff --git a/Makefile b/Makefile
index fc604098..b8362d1d 100644
--- a/Makefile
+++ b/Makefile
@@ -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()
+