diff --git a/datasette/cli.py b/datasette/cli.py index 57db83b6..06fa6199 100644 --- a/datasette/cli.py +++ b/datasette/cli.py @@ -663,16 +663,6 @@ def serve( # Private utility mechanism for writing unit tests return ds - # Run async soundness checks before startup hooks, since invoke_startup - # now populates internal tables which requires querying each database - run_sync(lambda: check_databases(ds)) - - # Run the "startup" plugin hooks - try: - run_sync(ds.invoke_startup) - except StartupError as e: - raise click.ClickException(e.args[0]) - if headers and not get: raise click.ClickException("--headers can only be used with --get") @@ -680,6 +670,16 @@ def serve( raise click.ClickException("--token can only be used with --get") if get: + # Run async soundness checks before startup hooks, since invoke_startup + # now populates internal tables which requires querying each database + run_sync(lambda: check_databases(ds)) + + # Run the "startup" plugin hooks + try: + run_sync(ds.invoke_startup) + except StartupError as e: + raise click.ClickException(e.args[0]) + client = TestClient(ds) request_headers = {} if token: @@ -704,34 +704,52 @@ def serve( sys.exit(exit_code) return - # Start the server - url = None - if root: - ds.root_enabled = True - url = "http://{}:{}{}?token={}".format( - host, port, ds.urls.path("-/auth-token"), ds._root_token - ) - click.echo(url) - if open_browser: - if url is None: - # Figure out most convenient URL - to table, database or homepage - path = run_sync(lambda: initial_path_for_datasette(ds)) - url = f"http://{host}:{port}{path}" - webbrowser.open(url) - uvicorn_kwargs = { - "host": host, - "port": port, - "log_level": "info", - "lifespan": "on", - "workers": 1, - } - if uds: - uvicorn_kwargs["uds"] = uds - if ssl_keyfile: - uvicorn_kwargs["ssl_keyfile"] = ssl_keyfile - if ssl_certfile: - uvicorn_kwargs["ssl_certfile"] = ssl_certfile - uvicorn.run(ds.app(), **uvicorn_kwargs) + # check_databases, invoke_startup() and the uvicorn server all run on a + # single event loop, so that anything a plugin's "startup" hook schedules + # on the loop (asyncio.create_task, Lock/Queue/Event objects, ...) is + # still alive when the server starts handling requests. + async def _serve_async(): + # Run async soundness checks before startup hooks, since invoke_startup + # now populates internal tables which requires querying each database + await check_databases(ds) + + # Run the "startup" plugin hooks + try: + await ds.invoke_startup() + except StartupError as e: + raise click.ClickException(e.args[0]) + + # Start the server + url = None + if root: + ds.root_enabled = True + url = "http://{}:{}{}?token={}".format( + host, port, ds.urls.path("-/auth-token"), ds._root_token + ) + click.echo(url) + if open_browser: + if url is None: + # Figure out most convenient URL - to table, database or homepage + path = await initial_path_for_datasette(ds) + url = f"http://{host}:{port}{path}" + webbrowser.open(url) + uvicorn_kwargs = { + "host": host, + "port": port, + "log_level": "info", + "lifespan": "on", + "workers": 1, + } + if uds: + uvicorn_kwargs["uds"] = uds + if ssl_keyfile: + uvicorn_kwargs["ssl_keyfile"] = ssl_keyfile + if ssl_certfile: + uvicorn_kwargs["ssl_certfile"] = ssl_certfile + server = uvicorn.Server(uvicorn.Config(ds.app(), **uvicorn_kwargs)) + await server.serve() + + asyncio.run(_serve_async()) @cli.command() diff --git a/pyproject.toml b/pyproject.toml index cf5db905..e658955f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ dependencies = [ "hupper>=1.9", "httpx>=0.20,<1.0", "pluggy>=1.0", - "uvicorn>=0.11", + "uvicorn>=0.29", "aiofiles>=0.4", "PyYAML>=5.3", "mergedeep>=1.1.1", diff --git a/tests/test_cli_serve_server.py b/tests/test_cli_serve_server.py index b7604bb8..19fc9198 100644 --- a/tests/test_cli_serve_server.py +++ b/tests/test_cli_serve_server.py @@ -1,4 +1,8 @@ import socket +import subprocess +import sys +import tempfile +import time import httpx import pytest @@ -28,3 +32,180 @@ def test_serve_unix_domain_socket(ds_unix_domain_socket_server): "path": "/_memory", "tables": [], }.items() <= response.json().items() + + +def _find_free_port(): + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +# Shaped after datasette-litestream's (sync) startup hook, which schedules a +# background task with asyncio.get_running_loop().create_task(...): +# https://github.com/simonw/datasette-litestream/blob/main/datasette_litestream/__init__.py +# That only has a chance to actually run if invoke_startup() executes on the +# same event loop that goes on to serve requests - if it runs on a throwaway +# loop that gets closed straight after (as on unmodified main), the task is +# scheduled but never gets a turn before the loop is torn down. +MARKER_TASK_PLUGIN = ''' +import asyncio +from datasette import hookimpl +from datasette.utils.asgi import Response + + +@hookimpl +def startup(datasette): + async def _mark(): + # The await is essential to the regression: a task with no + # internal await point can complete during the brief window + # between run_until_complete()'s coroutine finishing and the + # temporary loop actually stopping, masking the bug this test + # guards against. Real background tasks (like + # datasette-litestream's credential_refresh_loop) always have an + # internal await, and never get to resume once their throwaway + # loop is closed. + await asyncio.sleep(0.2) + datasette._marker_task_ran = True + + asyncio.get_running_loop().create_task(_mark()) + + +@hookimpl +def register_routes(): + async def marker_status(datasette): + return Response.json( + {"marker_task_ran": getattr(datasette, "_marker_task_ran", False)} + ) + + return [(r"^/-/marker-task-ran$", marker_status)] +''' + + +STARTUP_ERROR_PLUGIN = ''' +from datasette import hookimpl +from datasette.utils import StartupError + + +@hookimpl +def startup(datasette): + raise StartupError("boom from plugin") +''' + + +@pytest.mark.serial +def test_startup_hook_background_task_runs_on_serving_loop(tmp_path): + """ + Litestream-shaped regression test: a startup hook that does + asyncio.get_running_loop().create_task(...) must have that task + actually execute before/while the server is handling requests. This + only holds if invoke_startup() and uvicorn.Server.serve() share one + event loop. This test fails against unmodified main, where + invoke_startup() runs on a throwaway loop that is closed before + uvicorn opens its own loop to serve. + """ + plugins_dir = tmp_path / "plugins" + plugins_dir.mkdir() + (plugins_dir / "marker_task_plugin.py").write_text(MARKER_TASK_PLUGIN, "utf-8") + + port = _find_free_port() + ds_proc = subprocess.Popen( + [ + sys.executable, + "-m", + "datasette", + "--memory", + "--plugins-dir", + str(plugins_dir), + "-p", + str(port), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + cwd=tempfile.gettempdir(), + ) + try: + url = f"http://localhost:{port}/-/marker-task-ran" + deadline = time.time() + 15.0 + marker_task_ran = False + while time.time() < deadline: + if ds_proc.poll() is not None: + raise AssertionError( + "datasette serve exited early\n" + + ds_proc.stdout.read().decode("utf-8") + ) + try: + response = httpx.get(url, timeout=1.0) + except httpx.TransportError: + time.sleep(0.1) + continue + if response.status_code == 200 and response.json().get( + "marker_task_ran" + ): + marker_task_ran = True + break + time.sleep(0.1) + assert marker_task_ran, ( + "The startup hook's asyncio.create_task(...) never ran - " + "invoke_startup() and the server are not sharing an event loop" + ) + finally: + ds_proc.terminate() + try: + ds_proc.wait(timeout=5) + except subprocess.TimeoutExpired: + ds_proc.kill() + ds_proc.wait() + + +@pytest.mark.serial +def test_startup_error_fails_fast_before_port_binds(tmp_path): + """ + A "startup" plugin hook that raises StartupError must fail fast: print + the message, exit non-zero, and never accept a connection on the port - + the failure must happen before uvicorn.Server binds the socket. + """ + plugins_dir = tmp_path / "plugins" + plugins_dir.mkdir() + (plugins_dir / "startup_error_plugin.py").write_text( + STARTUP_ERROR_PLUGIN, "utf-8" + ) + + port = _find_free_port() + ds_proc = subprocess.Popen( + [ + sys.executable, + "-m", + "datasette", + "--memory", + "--plugins-dir", + str(plugins_dir), + "-p", + str(port), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + cwd=tempfile.gettempdir(), + ) + try: + deadline = time.time() + 15.0 + # While the process is still alive (it should crash almost + # immediately) repeatedly confirm nothing is listening on the port + while ds_proc.poll() is None and time.time() < deadline: + with pytest.raises(OSError): + with socket.create_connection(("127.0.0.1", port), timeout=0.2): + pass + time.sleep(0.05) + + stdout, _ = ds_proc.communicate(timeout=5) + output = stdout.decode("utf-8") + assert ds_proc.returncode not in (0, None), output + assert "boom from plugin" in output, output + + # And confirm it never accepted a connection even now it has exited + with pytest.raises(OSError): + with socket.create_connection(("127.0.0.1", port), timeout=0.2): + pass + finally: + if ds_proc.poll() is None: + ds_proc.kill() + ds_proc.wait()