Merge pull request #21643 from lidizheng/aio-reuseport

[Aio] Support SO_REUSEPORT channel option

src/core/lib/iomgr/tcp_custom.h

@@ -24,6 +24,9 @@
 #include "src/core/lib/iomgr/endpoint.h"
 #include "src/core/lib/iomgr/sockaddr.h"
 
+// Same number as the micro of SO_REUSEPORT in kernel
+#define GRPC_CUSTOM_SOCKET_OPT_SO_REUSEPORT (0x00000200u)
+
 typedef struct grpc_tcp_listener grpc_tcp_listener;
 typedef struct grpc_custom_tcp_connect grpc_custom_tcp_connect;
 

src/core/lib/iomgr/tcp_server_custom.cc

@@ -26,6 +26,7 @@
 #include <grpc/support/alloc.h>
 #include <grpc/support/log.h>
 
+#include "src/core/lib/channel/channel_args.h"
 #include "src/core/lib/iomgr/error.h"
 #include "src/core/lib/iomgr/exec_ctx.h"
 #include "src/core/lib/iomgr/iomgr_custom.h"
@@ -72,6 +73,7 @@ struct grpc_tcp_server {
   grpc_closure* shutdown_complete;
 
   bool shutdown;
+  bool so_reuseport;
 
   grpc_resource_quota* resource_quota;
 };
@@ -80,8 +82,13 @@ static grpc_error* tcp_server_create(grpc_closure* shutdown_complete,
                                      const grpc_channel_args* args,
                                      grpc_tcp_server** server) {
   grpc_tcp_server* s = (grpc_tcp_server*)gpr_malloc(sizeof(grpc_tcp_server));
+  // Let the implementation decide if so_reuseport can be enabled or not.
+  s->so_reuseport = true;
   s->resource_quota = grpc_resource_quota_create(nullptr);
   for (size_t i = 0; i < (args == nullptr ? 0 : args->num_args); i++) {
+    if (!grpc_channel_args_find_bool(args, GRPC_ARG_ALLOW_REUSEPORT, true)) {
+      s->so_reuseport = false;
+    }
     if (0 == strcmp(GRPC_ARG_RESOURCE_QUOTA, args->args[i].key)) {
       if (args->args[i].type == GRPC_ARG_POINTER) {
         grpc_resource_quota_unref_internal(s->resource_quota);
@@ -280,9 +287,15 @@ static grpc_error* add_socket_to_server(grpc_tcp_server* s,
   grpc_error* error;
   grpc_resolved_address sockname_temp;
 
-  // The last argument to uv_tcp_bind is flags
+  // NOTE(lidiz) The last argument is "flags" which is unused by other
+  // implementations. Python IO managers uses it to specify SO_REUSEPORT.
+  int flags = 0;
+  if (s->so_reuseport) {
+    flags |= GRPC_CUSTOM_SOCKET_OPT_SO_REUSEPORT;
+  }
+
   error = grpc_custom_socket_vtable->bind(socket, (grpc_sockaddr*)addr->addr,
-                                          addr->len, 0);
+                                          addr->len, flags);
   if (error != GRPC_ERROR_NONE) {
     return error;
   }

src/python/grpcio/grpc/_cython/_cygrpc/aio/iomgr/iomgr.pyx.pxi

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 
+import platform
+
 from cpython cimport Py_INCREF, Py_DECREF
 from libc cimport string
 
@@ -122,11 +124,15 @@ cdef grpc_error* asyncio_socket_listen(grpc_custom_socket* grpc_socket) with gil
     return grpc_error_none()
 
 
-def _asyncio_apply_socket_options(object socket):
-    # TODO(https://github.com/grpc/grpc/issues/20667)
-    # Connects the so_reuse_port option to channel arguments
-    socket.setsockopt(native_socket.SOL_SOCKET, native_socket.SO_REUSEADDR, 1)
-    socket.setsockopt(native_socket.IPPROTO_TCP, native_socket.TCP_NODELAY, True)
+def _asyncio_apply_socket_options(object s, int flags):
+    # Turn SO_REUSEADDR on for TCP sockets; if we want to support UDS, we will
+    # need to update this function.
+    s.setsockopt(native_socket.SOL_SOCKET, native_socket.SO_REUSEADDR, 1)
+    # SO_REUSEPORT only available in POSIX systems.
+    if platform.system() != 'Windows':
+        if GRPC_CUSTOM_SOCKET_OPT_SO_REUSEPORT & flags:
+            s.setsockopt(native_socket.SOL_SOCKET, native_socket.SO_REUSEPORT, 1)
+    s.setsockopt(native_socket.IPPROTO_TCP, native_socket.TCP_NODELAY, True)
 
 
 cdef grpc_error* asyncio_socket_bind(
@@ -142,7 +148,7 @@ cdef grpc_error* asyncio_socket_bind(
             family = native_socket.AF_INET
 
         socket = native_socket.socket(family=family)
-        _asyncio_apply_socket_options(socket)
+        _asyncio_apply_socket_options(socket, flags)
         socket.bind((host, port))
     except IOError as io_error:
         return socket_error("bind", str(io_error))

src/python/grpcio/grpc/_cython/_cygrpc/iomgr.pxd.pxi

@@ -54,6 +54,8 @@ cdef extern from "src/core/lib/iomgr/resolve_address_custom.h":
                                     grpc_error* error);
 
 cdef extern from "src/core/lib/iomgr/tcp_custom.h":
+  cdef int GRPC_CUSTOM_SOCKET_OPT_SO_REUSEPORT
+
   struct grpc_custom_socket:
     void* impl
     # We don't care about the rest of the fields

src/python/grpcio_tests/tests_aio/unit/channel_argument_test.py

@@ -89,7 +89,6 @@ class TestChannelArgument(AioTestBase):
     async def setUp(self):
         random.seed(_RANDOM_SEED)
 
-    @unittest.skip('https://github.com/grpc/grpc/issues/20667')
     @unittest.skipIf(platform.system() == 'Windows',
                      'SO_REUSEPORT only available in Linux-like OS.')
     async def test_server_so_reuse_port_is_set_properly(self):