loop: Add support for reuse_port in create_server/create_datagram

Author: 1st1

tests/test_tcp.py

@@ -66,6 +66,7 @@ async def start_server():
             srv = await asyncio.start_server(
                 handle_client,
                 addrs, 0,
+                reuse_port=getattr(socket, 'SO_REUSEPORT', None),
                 family=socket.AF_INET,
                 loop=self.loop)
 

tests/test_udp.py

@@ -68,6 +68,7 @@ def datagram_received(self, data, addr):
                 coro = self.loop.create_datagram_endpoint(
                     lambda: MyDatagramProto(loop=self.loop),
                     family=socket.AF_INET,
+                    reuse_port=getattr(socket, 'SO_REUSEPORT', None),
                     remote_addr=None if lc is None else (host, port))
                 transport, client = self.loop.run_until_complete(coro)
 

uvloop/handles/tcp.pxd

@@ -4,7 +4,7 @@ cdef class TCPServer(UVStreamServer):
 
     @staticmethod
     cdef TCPServer new(Loop loop, object protocol_factory, Server server,
-                       object ssl)
+                       object ssl, unsigned int flags)
 
 
 cdef class TCPTransport(UVStream):

uvloop/handles/tcp.pyx

@@ -1,4 +1,4 @@
-cdef __tcp_init_uv_handle(UVStream handle, Loop loop):
+cdef __tcp_init_uv_handle(UVStream handle, Loop loop, unsigned int flags):
     cdef int err
 
     handle._handle = <uv.uv_handle_t*> \
@@ -7,7 +7,9 @@ cdef __tcp_init_uv_handle(UVStream handle, Loop loop):
         handle._abort_init()
         raise MemoryError()
 
-    err = uv.uv_tcp_init(handle._loop.uvloop, <uv.uv_tcp_t*>handle._handle)
+    err = uv.uv_tcp_init_ex(handle._loop.uvloop,
+                            <uv.uv_tcp_t*>handle._handle,
+                            flags)
     if err < 0:
         handle._abort_init()
         raise convert_error(err)
@@ -56,12 +58,12 @@ cdef class TCPServer(UVStreamServer):
 
     @staticmethod
     cdef TCPServer new(Loop loop, object protocol_factory, Server server,
-                         object ssl):
+                       object ssl, unsigned int flags):
 
         cdef TCPServer handle
         handle = TCPServer.__new__(TCPServer)
         handle._init(loop, protocol_factory, server, ssl)
-        __tcp_init_uv_handle(<UVStream>handle, loop)
+        __tcp_init_uv_handle(<UVStream>handle, loop, flags)
         return handle
 
     cdef _new_socket(self):
@@ -101,7 +103,7 @@ cdef class TCPTransport(UVStream):
         cdef TCPTransport handle
         handle = TCPTransport.__new__(TCPTransport)
         handle._init(loop, protocol, server, waiter)
-        __tcp_init_uv_handle(<UVStream>handle, loop)
+        __tcp_init_uv_handle(<UVStream>handle, loop, uv.AF_UNSPEC)
         handle.__peername_set = 0
         handle.__sockname_set = 0
         return handle

uvloop/includes/stdlib.pxi

@@ -46,6 +46,7 @@ cdef iter_chain = itertools.chain
 cdef int has_AF_INET6 = hasattr(socket, 'AF_INET6')
 cdef int has_SO_REUSEPORT = hasattr(socket, 'SO_REUSEPORT')
 cdef int has_IPPROTO_IPV6 = hasattr(socket, 'IPPROTO_IPV6')
+cdef int SO_REUSEPORT = getattr(socket, 'SO_REUSEPORT', 0)
 
 cdef socket_gaierror = socket.gaierror
 cdef socket_error = socket.error

uvloop/includes/system.pxd

@@ -43,6 +43,9 @@ cdef extern from "sys/socket.h" nogil:
 
     ssize_t send(int sockfd, const void *buf, size_t len, int flags);
 
+    int setsockopt(int socket, int level, int option_name,
+                   const void *option_value, int option_len)
+
 
 cdef extern from "errno.h" nogil:
 

uvloop/includes/uv.pxd

@@ -319,7 +319,7 @@ cdef extern from "../vendor/libuv/include/uv.h" nogil:
 
     # TCP
 
-    int uv_tcp_init(uv_loop_t*, uv_tcp_t* handle)
+    int uv_tcp_init_ex(uv_loop_t*, uv_tcp_t* handle, unsigned int flags)
     int uv_tcp_nodelay(uv_tcp_t* handle, int enable)
     int uv_tcp_keepalive(uv_tcp_t* handle, int enable, unsigned int delay)
     int uv_tcp_open(uv_tcp_t* handle, uv_os_sock_t sock)

uvloop/loop.pxd

@@ -159,6 +159,8 @@ cdef class Loop:
     cdef _sock_connect(self, fut, sock, address)
     cdef _sock_connect_cb(self, fut, sock, address)
 
+    cdef _sock_set_reuseport(self, int fd)
+
     cdef _set_coroutine_wrapper(self, bint enabled)
 
 

uvloop/loop.pyx

@@ -689,6 +689,21 @@ cdef class Loop:
         else:
             fut.set_result(None)
 
+    cdef _sock_set_reuseport(self, int fd):
+        cdef:
+            int err
+            int reuseport_flag = 1
+
+        err = system.setsockopt(
+            fd,
+            uv.SOL_SOCKET,
+            SO_REUSEPORT,
+            <char*>&reuseport_flag,
+            sizeof(reuseport_flag))
+
+        if err < 0:
+            raise convert_error(-errno.errno)
+
     cdef _set_coroutine_wrapper(self, bint enabled):
         enabled = bool(enabled)
         if self._coroutine_wrapper_set == enabled:
@@ -1003,7 +1018,7 @@ cdef class Loop:
                             int backlog=100,
                             ssl=None,
                             reuse_address=None,  # ignored, libuv sets it
-                            reuse_port=None):    # ignored
+                            reuse_port=None):
 
         cdef:
             TCPServer tcp
@@ -1018,6 +1033,11 @@ cdef class Loop:
                 raise ValueError(
                     'host/port and sock can not be specified at the same time')
 
+            reuse_port = bool(reuse_port)
+            if reuse_port and not has_SO_REUSEPORT:
+                raise ValueError(
+                    'reuse_port not supported by socket module')
+
             if host == '':
                 hosts = [None]
             elif (isinstance(host, str) or not isinstance(host, col_Iterable)):
@@ -1036,8 +1056,15 @@ cdef class Loop:
                 for info in infos:
                     addrinfo = (<AddrInfo>info).data
                     while addrinfo != NULL:
+                        if addrinfo.ai_family == uv.AF_UNSPEC:
+                            raise RuntimeError('AF_UNSPEC in DNS results')
+
                         tcp = TCPServer.new(
-                            self, protocol_factory, server, ssl)
+                            self, protocol_factory, server, ssl,
+                            addrinfo.ai_family)
+
+                        if reuse_port:
+                            self._sock_set_reuseport(tcp._fileno())
 
                         try:
                             tcp.bind(addrinfo.ai_addr)
@@ -1057,7 +1084,8 @@ cdef class Loop:
         else:
             if sock is None:
                 raise ValueError('Neither host/port nor sock were specified')
-            tcp = TCPServer.new(self, protocol_factory, server, ssl)
+            tcp = TCPServer.new(self, protocol_factory, server, ssl,
+                                uv.AF_UNSPEC)
             fileno = os_dup(sock.fileno())
             try:
                 tcp.open(fileno)
@@ -1684,6 +1712,10 @@ cdef class Loop:
             udp._attach_fileobj(sock)
         else:
             reuse_address = bool(reuse_address)
+            reuse_port = bool(reuse_port)
+            if reuse_port and not has_SO_REUSEPORT:
+                raise ValueError(
+                    'reuse_port not supported by socket module')
 
             lads = None
             if local_addr is not None:
@@ -1720,6 +1752,9 @@ cdef class Loop:
                     udp = UDPTransport.__new__(UDPTransport)
                     udp._init(self, family)
 
+                if reuse_port:
+                    self._sock_set_reuseport(udp._fileno())
+
                 socket = udp._get_socket()
                 socket.bind(('0.0.0.0', 0))
             else:
@@ -1728,6 +1763,8 @@ cdef class Loop:
                     try:
                         udp = UDPTransport.__new__(UDPTransport)
                         udp._init(self, lai.ai_family)
+                        if reuse_port:
+                            self._sock_set_reuseport(udp._fileno())
                         udp._bind(lai.ai_addr, reuse_address)
                     except Exception as ex:
                         lai = lai.ai_next