From f2b103712c4fde78fa6f49c9e9f552ca9f30aa18 Mon Sep 17 00:00:00 2001 From: vvanglro Date: Fri, 25 Oct 2024 16:01:52 +0800 Subject: [PATCH] feat(config): add socket load balance option Co-authored-by: Thomas Grainger --- docs/deployment.md | 1 + docs/index.md | 1 + uvicorn/_subprocess.py | 8 ++++++-- uvicorn/config.py | 9 +++++++++ uvicorn/main.py | 17 ++++++++++++++++- uvicorn/supervisors/multiprocess.py | 4 ++-- 6 files changed, 35 insertions(+), 5 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index d69fcf88e..e0da7a69e 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -135,6 +135,7 @@ Options: buffer of an incomplete event. --factory Treat APP as an application factory, i.e. a () -> callable. + --socket-load-balance Use kernel support for socket load balancing --help Show this message and exit. ``` diff --git a/docs/index.md b/docs/index.md index bb6fc321a..e9bf5b3b5 100644 --- a/docs/index.md +++ b/docs/index.md @@ -205,6 +205,7 @@ Options: buffer of an incomplete event. --factory Treat APP as an application factory, i.e. a () -> callable. + --socket-load-balance Use kernel support for socket load balancing --help Show this message and exit. ``` diff --git a/uvicorn/_subprocess.py b/uvicorn/_subprocess.py index 1c06844de..1cf0ac3ff 100644 --- a/uvicorn/_subprocess.py +++ b/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) diff --git a/uvicorn/config.py b/uvicorn/config.py index 65dfe651e..d2d587222 100644 --- a/uvicorn/config.py +++ b/uvicorn/config.py @@ -223,6 +223,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 @@ -268,6 +269,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() @@ -510,6 +512,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 @@ -527,3 +531,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") diff --git a/uvicorn/main.py b/uvicorn/main.py index 96a10d538..637fbf95a 100644 --- a/uvicorn/main.py +++ b/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 () -> 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() diff --git a/uvicorn/supervisors/multiprocess.py b/uvicorn/supervisors/multiprocess.py index e198fe780..27268518b 100644 --- a/uvicorn/supervisors/multiprocess.py +++ b/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