feat: socket lb

uvicorn/_subprocess.py

@@ -21,7 +21,7 @@
 def get_subprocess(
     config: Config,
     target: Callable[..., None],
-    sockets: list[socket],
+    sockets: list[socket] | list[Callable[[], socket]],
 ) -> SpawnProcess:
     """
     Called in the parent process, to instantiate a new child process instance.
@@ -54,7 +54,7 @@ def get_subprocess(
 def subprocess_started(
     config: Config,
     target: Callable[..., None],
-    sockets: list[socket],
+    sockets: list[socket] | list[Callable[[], socket]],
     stdin_fileno: int | None,
 ) -> None:
     """
@@ -75,6 +75,10 @@ def subprocess_started(
     # Logging needs to be setup again for each child.
     config.configure_logging()
 
+    for idx, sock in enumerate(sockets):
+        if callable(sock):
+            sockets[idx] = sock()
+
     try:
         # Now we can call into `Server.run(sockets=sockets)`
         target(sockets=sockets)

uvicorn/config.py

@@ -224,6 +224,7 @@ def __init__(
         headers: list[tuple[str, str]] | None = None,
         factory: bool = False,
         h11_max_incomplete_event_size: int | None = None,
+        socket_load_balance: bool = False,
     ):
         self.app = app
         self.host = host
@@ -269,6 +270,7 @@ def __init__(
         self.encoded_headers: list[tuple[bytes, bytes]] = []
         self.factory = factory
         self.h11_max_incomplete_event_size = h11_max_incomplete_event_size
+        self.socket_load_balance = socket_load_balance
 
         self.loaded = False
         self.configure_logging()
@@ -511,6 +513,8 @@ def bind_socket(self) -> socket.socket:
 
             sock = socket.socket(family=family)
             sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            if self.socket_load_balance:
+                sock.setsockopt(socket.SOL_SOCKET, getattr(socket, "SO_REUSEPORT_LB", socket.SO_REUSEPORT), 1)
             try:
                 sock.bind((self.host, self.port))
             except OSError as exc:  # pragma: full coverage
@@ -528,3 +532,8 @@ def bind_socket(self) -> socket.socket:
     @property
     def should_reload(self) -> bool:
         return isinstance(self.app, str) and self.reload
+
+    @staticmethod
+    def check_socket_lb() -> None:
+        if not (sys.platform == "linux" and hasattr(socket, "SO_REUSEPORT")) or hasattr(socket, "SO_REUSEPORT_LB"):
+            raise RuntimeError("socket_load_balance not supported")

uvicorn/main.py

@@ -360,6 +360,13 @@ def print_version(ctx: click.Context, param: click.Parameter, value: bool) -> No
     help="Treat APP as an application factory, i.e. a () -> <ASGI app> callable.",
     show_default=True,
 )
+@click.option(
+    "--socket-load-balance",
+    is_flag=True,
+    default=False,
+    help="Use kernel support for socket load balancing",
+    show_default=True,
+)
 def main(
     app: str,
     host: str,
@@ -408,6 +415,7 @@ def main(
     app_dir: str,
     h11_max_incomplete_event_size: int | None,
     factory: bool,
+    socket_load_balance: bool = False,
 ) -> None:
     run(
         app,
@@ -457,6 +465,7 @@ def main(
         factory=factory,
         app_dir=app_dir,
         h11_max_incomplete_event_size=h11_max_incomplete_event_size,
+        socket_load_balance=socket_load_balance,
     )
 
 
@@ -509,6 +518,7 @@ def run(
     app_dir: str | None = None,
     factory: bool = False,
     h11_max_incomplete_event_size: int | None = None,
+    socket_load_balance: bool = False,
 ) -> None:
     if app_dir is not None:
         sys.path.insert(0, app_dir)
@@ -560,6 +570,7 @@ def run(
         use_colors=use_colors,
         factory=factory,
         h11_max_incomplete_event_size=h11_max_incomplete_event_size,
+        socket_load_balance=socket_load_balance,
     )
     server = Server(config=config)
 
@@ -573,7 +584,11 @@ def run(
             sock = config.bind_socket()
             ChangeReload(config, target=server.run, sockets=[sock]).run()
         elif config.workers > 1:
-            sock = config.bind_socket()
+            if config.socket_load_balance:
+                config.check_socket_lb()
+                sock = config.bind_socket  # type: ignore[assignment]
+            else:
+                sock = config.bind_socket()
             Multiprocess(config, target=server.run, sockets=[sock]).run()
         else:
             server.run()

uvicorn/supervisors/multiprocess.py

@@ -27,7 +27,7 @@ def __init__(
         self,
         config: Config,
         target: Callable[[list[socket] | None], None],
-        sockets: list[socket],
+        sockets: list[socket] | list[Callable[[], socket]],
     ) -> None:
         self.real_target = target
 
@@ -104,7 +104,7 @@ def __init__(
         self,
         config: Config,
         target: Callable[[list[socket] | None], None],
-        sockets: list[socket],
+        sockets: list[socket] | list[Callable[[], socket]],
     ) -> None:
         self.config = config
         self.target = target