From 4cb6d810932cf637c6e35fcec9c1a0b2488249d5 Mon Sep 17 00:00:00 2001 From: euri10 Date: Mon, 21 Jun 2021 19:53:40 +0200 Subject: [PATCH] Adapt bind_socket to make it usable with multiple processes (#1009) * Adapt bind_socket to make it usable with multple processes * Lint * we are in single proc here * Remove socket file also in reloader * Tests * Tests * Remove socket on close * Changed if / else to check uds / fds * Refactored logging inside bind_socket * Added test for fd * Refactored removal of socket * Minimized diff post merge --- tests/test_config.py | 44 ++++++++++++++++++++++++ uvicorn/config.py | 82 +++++++++++++++++++++++++++++--------------- uvicorn/main.py | 3 ++ 3 files changed, 101 insertions(+), 28 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 761eec4048..1e26c4cc51 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -355,3 +355,47 @@ def test_ws_max_size() -> None: config = Config(app=asgi_app, ws_max_size=1000) config.load() assert config.ws_max_size == 1000 + + +@pytest.mark.parametrize( + "reload, workers", + [ + (True, 1), + (False, 2), + ], + ids=["--reload=True --workers=1", "--reload=False --workers=2"], +) +@pytest.mark.skipif(sys.platform == "win32", reason="require unix-like system") +def test_bind_unix_socket_works_with_reload_or_workers(tmp_path, reload, workers): + uds_file = tmp_path / "uvicorn.sock" + config = Config( + app=asgi_app, uds=uds_file.as_posix(), reload=reload, workers=workers + ) + config.load() + sock = config.bind_socket() + assert isinstance(sock, socket.socket) + assert sock.family == socket.AF_UNIX + assert sock.getsockname() == uds_file.as_posix() + sock.close() + + +@pytest.mark.parametrize( + "reload, workers", + [ + (True, 1), + (False, 2), + ], + ids=["--reload=True --workers=1", "--reload=False --workers=2"], +) +@pytest.mark.skipif(sys.platform == "win32", reason="require unix-like system") +def test_bind_fd_works_with_reload_or_workers(reload, workers): + fdsock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + fd = fdsock.fileno() + config = Config(app=asgi_app, fd=fd, reload=reload, workers=workers) + config.load() + sock = config.bind_socket() + assert isinstance(sock, socket.socket) + assert sock.family == socket.AF_UNIX + assert sock.getsockname() == "" + sock.close() + fdsock.close() diff --git a/uvicorn/config.py b/uvicorn/config.py index 173492e5ae..2846d104e8 100644 --- a/uvicorn/config.py +++ b/uvicorn/config.py @@ -395,37 +395,63 @@ def setup_event_loop(self) -> None: loop_setup() def bind_socket(self) -> socket.socket: - family = socket.AF_INET - addr_format = "%s://%s:%d" + logger_args: List[Union[str, int]] + if self.uds: + path = self.uds + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + sock.bind(path) + uds_perms = 0o666 + os.chmod(self.uds, uds_perms) + except OSError as exc: + logger.error(exc) + sys.exit(1) - if self.host and ":" in self.host: - # It's an IPv6 address. - family = socket.AF_INET6 - addr_format = "%s://[%s]:%d" + message = "Uvicorn running on unix socket %s (Press CTRL+C to quit)" + sock_name_format = "%s" + color_message = ( + "Uvicorn running on " + + click.style(sock_name_format, bold=True) + + " (Press CTRL+C to quit)" + ) + logger_args = [self.uds] + elif self.fd: + sock = socket.fromfd(self.fd, socket.AF_UNIX, socket.SOCK_STREAM) + message = "Uvicorn running on socket %s (Press CTRL+C to quit)" + fd_name_format = "%s" + color_message = ( + "Uvicorn running on " + + click.style(fd_name_format, bold=True) + + " (Press CTRL+C to quit)" + ) + logger_args = [sock.getsockname()] + else: + family = socket.AF_INET + addr_format = "%s://%s:%d" + + if self.host and ":" in self.host: + # It's an IPv6 address. + family = socket.AF_INET6 + addr_format = "%s://[%s]:%d" + + sock = socket.socket(family=family) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind((self.host, self.port)) + except OSError as exc: + logger.error(exc) + sys.exit(1) - sock = socket.socket(family=family) - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - try: - sock.bind((self.host, self.port)) - except OSError as exc: - logger.error(exc) - sys.exit(1) + message = f"Uvicorn running on {addr_format} (Press CTRL+C to quit)" + color_message = ( + "Uvicorn running on " + + click.style(addr_format, bold=True) + + " (Press CTRL+C to quit)" + ) + protocol_name = "https" if self.is_ssl else "http" + logger_args = [protocol_name, self.host, self.port] + logger.info(message, *logger_args, extra={"color_message": color_message}) sock.set_inheritable(True) - - message = f"Uvicorn running on {addr_format} (Press CTRL+C to quit)" - color_message = ( - "Uvicorn running on " - + click.style(addr_format, bold=True) - + " (Press CTRL+C to quit)" - ) - protocol_name = "https" if self.is_ssl else "http" - logger.info( - message, - protocol_name, - self.host, - self.port, - extra={"color_message": color_message}, - ) return sock @property diff --git a/uvicorn/main.py b/uvicorn/main.py index 1d3260ade8..42d3415f7d 100644 --- a/uvicorn/main.py +++ b/uvicorn/main.py @@ -1,4 +1,5 @@ import logging +import os import platform import ssl import sys @@ -426,6 +427,8 @@ def run(app: typing.Union[ASGIApplication, str], **kwargs: typing.Any) -> None: Multiprocess(config, target=server.run, sockets=[sock]).run() else: server.run() + if config.uds: + os.remove(config.uds) if __name__ == "__main__":