summaryrefslogtreecommitdiff
path: root/python/skytools/sockutil.py
blob: dbcd021b28638c2d0530b6978ba1471bb5b8c72f (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
"""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

    # 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

    # 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()