From 3249762f00c796e38afb21f6cf4fbd7941beda5c Mon Sep 17 00:00:00 2001 From: Nir Geller Date: Tue, 16 Jul 2024 23:54:39 +0300 Subject: [PATCH 1/3] Allow passing a custom loop setup function instead of None --- uvicorn/config.py | 20 ++++++++++++++++++-- uvicorn/main.py | 11 +++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/uvicorn/config.py b/uvicorn/config.py index 9aff8c968..cad7112bd 100644 --- a/uvicorn/config.py +++ b/uvicorn/config.py @@ -26,7 +26,7 @@ HTTPProtocolType = Literal["auto", "h11", "httptools"] WSProtocolType = Literal["auto", "none", "websockets", "wsproto"] LifespanType = Literal["auto", "on", "off"] -LoopSetupType = Literal["none", "auto", "asyncio", "uvloop"] +LoopSetupType = Literal["custom", "none", "auto", "asyncio", "uvloop"] InterfaceType = Literal["auto", "asgi3", "asgi2", "wsgi"] LOG_LEVELS: dict[str, int] = { @@ -55,6 +55,7 @@ } LOOP_SETUPS: dict[LoopSetupType, str | None] = { "none": None, + "custom": None, "auto": "uvicorn.loops.auto:auto_loop_setup", "asyncio": "uvicorn.loops.asyncio:asyncio_setup", "uvloop": "uvicorn.loops.uvloop:uvloop_setup", @@ -181,6 +182,7 @@ def __init__( uds: str | None = None, fd: int | None = None, loop: LoopSetupType = "auto", + loop_setup: str | None = None, http: type[asyncio.Protocol] | HTTPProtocolType = "auto", ws: type[asyncio.Protocol] | WSProtocolType = "auto", ws_max_size: int = 16 * 1024 * 1024, @@ -230,6 +232,7 @@ def __init__( self.uds = uds self.fd = fd self.loop = loop + self.loop_setup = loop_setup self.http = http self.ws = ws self.ws_max_size = ws_max_size @@ -472,7 +475,20 @@ def load(self) -> None: self.loaded = True def setup_event_loop(self) -> None: - loop_setup: Callable | None = import_from_string(LOOP_SETUPS[self.loop]) + loop_setup: Callable | None = None + + if self.loop == "custom": + if self.loop_setup is None: + logger.error("Custom loop setup is selected but no loop setup callable was provided.") + sys.exit(1) + try: + loop_setup = import_from_string(self.loop_setup) + except ImportFromStringError as exc: + logger.error("Error loading custom loop setup function. %s" % exc) + sys.exit(1) + else: + loop_setup = import_from_string(LOOP_SETUPS[self.loop]) + if loop_setup is not None: loop_setup(use_subprocess=self.use_subprocess) diff --git a/uvicorn/main.py b/uvicorn/main.py index 4352efbca..dcc090e58 100644 --- a/uvicorn/main.py +++ b/uvicorn/main.py @@ -122,6 +122,13 @@ def print_version(ctx: click.Context, param: click.Parameter, value: bool) -> No help="Event loop implementation.", show_default=True, ) +@click.option( + "--loop-setup", + type=str, + default=None, + help="Import path to event loop setup function., i.e. a (use_subprocess: bool) -> None callable.", + show_default=True, +) @click.option( "--http", type=HTTP_CHOICES, @@ -365,6 +372,7 @@ def main( uds: str, fd: int, loop: LoopSetupType, + loop_setup: str | None, http: HTTPProtocolType, ws: WSProtocolType, ws_max_size: int, @@ -414,6 +422,7 @@ def main( uds=uds, fd=fd, loop=loop, + loop_setup=loop_setup, http=http, ws=ws, ws_max_size=ws_max_size, @@ -466,6 +475,7 @@ def run( uds: str | None = None, fd: int | None = None, loop: LoopSetupType = "auto", + loop_setup: str | None = None, http: type[asyncio.Protocol] | HTTPProtocolType = "auto", ws: type[asyncio.Protocol] | WSProtocolType = "auto", ws_max_size: int = 16777216, @@ -518,6 +528,7 @@ def run( uds=uds, fd=fd, loop=loop, + loop_setup=loop_setup, http=http, ws=ws, ws_max_size=ws_max_size, From 79d556b2f7c87c033fde9c24b2d465f622f89002 Mon Sep 17 00:00:00 2001 From: Nir Geller Date: Tue, 16 Jul 2024 23:54:39 +0300 Subject: [PATCH 2/3] Update documentation --- docs/deployment.md | 6 +++++- docs/index.md | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index 7a2c7972c..f5b2d07ba 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -57,7 +57,11 @@ Options: --workers INTEGER Number of worker processes. Defaults to the $WEB_CONCURRENCY environment variable if available, or 1. Not valid with --reload. - --loop [auto|asyncio|uvloop] Event loop implementation. [default: auto] + --loop [custom|auto|asyncio|uvloop] + Event loop implementation. [default: auto] + --loop-setup TEXT Import path to event loop setup function., + i.e. a (use_subprocess: bool) -> None + callable. --http [auto|h11|httptools] HTTP protocol implementation. [default: auto] --ws [auto|none|websockets|wsproto] diff --git a/docs/index.md b/docs/index.md index 5d805316b..d3e40c33f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -127,7 +127,11 @@ Options: --workers INTEGER Number of worker processes. Defaults to the $WEB_CONCURRENCY environment variable if available, or 1. Not valid with --reload. - --loop [auto|asyncio|uvloop] Event loop implementation. [default: auto] + --loop [custom|auto|asyncio|uvloop] + Event loop implementation. [default: auto] + --loop-setup TEXT Import path to event loop setup function., + i.e. a (use_subprocess: bool) -> None + callable. --http [auto|h11|httptools] HTTP protocol implementation. [default: auto] --ws [auto|none|websockets|wsproto] From 0b14e3c999f2668f2320890601c61d5c89ccd7ef Mon Sep 17 00:00:00 2001 From: Nir Geller Date: Wed, 17 Jul 2024 01:01:51 +0300 Subject: [PATCH 3/3] Add unit tests for custom loop parameter --- tests/test_config.py | 78 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/tests/test_config.py b/tests/test_config.py index e16cc5d56..eb95584a3 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,6 +1,8 @@ from __future__ import annotations +import asyncio import configparser +import contextlib import io import json import logging @@ -8,6 +10,8 @@ import socket import sys import typing +from asyncio.events import BaseDefaultEventLoopPolicy +from contextlib import closing from pathlib import Path from typing import Any, Literal from unittest.mock import MagicMock @@ -545,3 +549,77 @@ def test_warn_when_using_reload_and_workers(caplog: pytest.LogCaptureFixture) -> Config(app=asgi_app, reload=True, workers=2) assert len(caplog.records) == 1 assert '"workers" flag is ignored when reloading is enabled.' in caplog.records[0].message + + +def custom_loop(use_subprocess: bool): + asyncio.set_event_loop_policy(CustomEventLoopPolicy()) + + +@contextlib.contextmanager +def with_event_loop_cleanup() -> typing.Generator[None, None, None]: + """ + Cleanup the event loop policy after the test. + """ + yield + asyncio.set_event_loop_policy(None) + + +class CustomLoop(asyncio.SelectorEventLoop): + pass + + +class CustomEventLoopPolicy(BaseDefaultEventLoopPolicy): + def set_child_watcher(self, watcher): + raise NotImplementedError + + def get_child_watcher(self): + raise NotImplementedError + + def _loop_factory(self) -> CustomLoop: + return CustomLoop() + + +def test_custom_loop__importable_custom_loop_setup_function() -> None: + with with_event_loop_cleanup(): + config = Config(app=asgi_app, loop="custom", loop_setup="tests.test_config:custom_loop") + config.load() + config.setup_event_loop() + event_loop = asyncio.new_event_loop() + with closing(event_loop): + assert event_loop is not None + assert isinstance(event_loop, CustomLoop) + + +@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") +def test_custom_loop__not_importable_custom_loop_setup_function(caplog: pytest.LogCaptureFixture) -> None: + with with_event_loop_cleanup(): + config = Config(app=asgi_app, loop="custom", loop_setup="tests.test_config:non_existing_setup_function") + config.load() + with pytest.raises(SystemExit): + config.setup_event_loop() + error_messages = [ + record.message + for record in caplog.records + if record.name == "uvicorn.error" and record.levelname == "ERROR" + ] + assert ( + 'Error loading custom loop setup function. Attribute "non_existing_setup_function" not found in module "tests.test_config".' # noqa: E501 + == error_messages.pop(0) + ) + + +def test_custom_loop__no_loop_setup_passed(caplog: pytest.LogCaptureFixture) -> None: + with with_event_loop_cleanup(): + config = Config(app=asgi_app, loop="custom") + config.load() + with pytest.raises(SystemExit): + config.setup_event_loop() + error_messages = [ + record.message + for record in caplog.records + if record.name == "uvicorn.error" and record.levelname == "ERROR" + ] + assert ( + "Custom loop setup is selected but no loop setup callable was provided." # noqa: E501 + == error_messages.pop(0) + )