Run datasette serve startup and uvicorn on a single event loop (#2886)

* Run datasette serve startup and uvicorn on a single event loop
* Move the serve-subprocess test plumbing into a conftest fixture
* Fix datasette-litestream URL and trim marker-task test comments
* Explain why serve_with_plugins needs a subprocess and plugin files
* Apply ruff 0.16 and black fixes
* Tweaked some comments

tests/conftest.py

@@ -2,6 +2,7 @@ import importlib.metadata
 import os
 import pathlib
 import re
+import socket
 import subprocess
 import sys
 import tempfile
@@ -32,17 +33,31 @@ UNDOCUMENTED_PERMISSIONS = {
 }
 
 
-def wait_until_responds(url, timeout=5.0, client=httpx, **kwargs):
+def wait_until_responds(url, timeout=5.0, client=httpx, process=None, **kwargs):
     start = time.time()
     while time.time() - start < timeout:
+        # If the server died there is no point waiting out the timeout - fail
+        # now, with its output, instead of after `timeout` seconds of silence
+        if process is not None and process.poll() is not None:
+            raise AssertionError(
+                "Server exited early with returncode {}\n{}".format(
+                    process.returncode, process.stdout.read().decode("utf-8")
+                )
+            )
         try:
             client.get(url, **kwargs)
             return
-        except httpx.ConnectError:
+        except httpx.TransportError:
             time.sleep(0.1)
     raise AssertionError(f"Timed out waiting for {url} to respond")
 
 
+def find_free_port():
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
 @pytest.fixture
 def bare_ds():
     """
@@ -301,6 +316,71 @@ def ds_unix_domain_socket_server(tmp_path_factory):
             pass
 
 
+@pytest.fixture
+def serve_with_plugins(tmp_path):
+    """Factory fixture for starting ``datasette serve`` in a subprocess with
+    plugins written to a temporary ``--plugins-dir``.
+
+    For tests that need the real serve path: event-loop wiring, exit codes,
+    signals. The usual in-process ``pm.register`` plugin pattern can't reach
+    a subprocess, so plugin source is written out as importable files instead.
+
+    Unlike ``ds_localhost_http_server`` this is function-scoped and takes a
+    fresh port each time, because each test needs its own plugins. Call it as::
+
+        proc, port = serve_with_plugins({"my_plugin": PLUGIN_SOURCE})
+
+    ``plugins`` maps module name to Python source. Pass
+    ``wait_for_startup=False`` when the server is expected to fail during
+    startup rather than begin serving. Extra CLI arguments are passed through.
+    Every process started is terminated when the test ends.
+    """
+    processes = []
+
+    def start(plugins, *extra_args, wait_for_startup=True):
+        plugins_dir = tmp_path / "plugins"
+        plugins_dir.mkdir(exist_ok=True)
+        for module_name, source in plugins.items():
+            (plugins_dir / f"{module_name}.py").write_text(source, "utf-8")
+        port = find_free_port()
+        proc = subprocess.Popen(
+            [
+                sys.executable,
+                "-m",
+                "datasette",
+                "--memory",
+                "--plugins-dir",
+                str(plugins_dir),
+                "-h",
+                "127.0.0.1",
+                "-p",
+                str(port),
+                *extra_args,
+            ],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            # Avoid FileNotFoundError: [Errno 2] No such file or directory:
+            cwd=tempfile.gettempdir(),
+        )
+        processes.append(proc)
+        if wait_for_startup:
+            wait_until_responds(
+                f"http://127.0.0.1:{port}/-/versions.json", process=proc
+            )
+        return proc, port
+
+    yield start
+
+    for proc in processes:
+        if proc.poll() is None:
+            proc.terminate()
+            try:
+                proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+                proc.wait()
+
+
 # Import fixtures from fixtures.py to make them available
 from .fixtures import (  # noqa: F401
     TEMP_PLUGIN_SECRET_FILE,

tests/test_cli_serve_server.py

@@ -1,4 +1,5 @@
 import socket
+import time
 
 import httpx
 import pytest
@@ -28,3 +29,119 @@ def test_serve_unix_domain_socket(ds_unix_domain_socket_server):
         "path": "/_memory",
         "tables": [],
     }.items() <= response.json().items()
+
+
+# Shaped after datasette-litestream's startup hook, which schedules a
+# background task with asyncio.get_running_loop().create_task(...):
+# https://github.com/datasette/datasette-litestream
+MARKER_TASK_PLUGIN = """
+import asyncio
+from datasette import hookimpl
+from datasette.utils.asgi import Response
+
+
+@hookimpl
+def startup(datasette):
+    datasette._startup_calls = getattr(datasette, "_startup_calls", 0) + 1
+
+    async def _mark():
+        # Must await before setting the flag: a task with no internal
+        # await point could finish on the throwaway loop before it
+        # closed, masking the regression this test guards against.
+        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),
+                "startup_calls": getattr(datasette, "_startup_calls", 0),
+            }
+        )
+
+    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(serve_with_plugins):
+    """
+    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.
+    """
+    _, port = serve_with_plugins({"marker_task_plugin": MARKER_TASK_PLUGIN})
+    # The fixture has already waited for the server to answer requests. The
+    # marker task deliberately awaits before setting its flag, so poll for a
+    # moment rather than assuming it landed before the first request arrived.
+    deadline = time.time() + 3.0
+    payload = {}
+    while time.time() < deadline:
+        payload = httpx.get(
+            f"http://127.0.0.1:{port}/-/marker-task-ran", timeout=1.0
+        ).json()
+        if payload["marker_task_ran"]:
+            break
+        time.sleep(0.05)
+    assert payload.get("marker_task_ran"), (
+        "The startup hook's asyncio.create_task(...) never ran - "
+        "invoke_startup() and the server are not sharing an event loop"
+    )
+    # Polling above means this test would also pass if the startup hook were
+    # re-run on the serving loop by the first-request fallback - which would
+    # hide exactly the bug being tested. invoke_startup() is idempotent today
+    # so that cannot happen; assert it explicitly so that if the idempotency
+    # guard is ever removed this test fails loudly instead of silently
+    # becoming a no-op.
+    assert payload["startup_calls"] == 1, (
+        "startup hook ran {} times - the marker may have been set by a "
+        "re-run on the serving loop rather than by the original task".format(
+            payload["startup_calls"]
+        )
+    )
+
+
+@pytest.mark.serial
+def test_startup_error_fails_fast_before_port_binds(serve_with_plugins):
+    """
+    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.
+    """
+    proc, port = serve_with_plugins(
+        {"startup_error_plugin": STARTUP_ERROR_PLUGIN}, wait_for_startup=False
+    )
+    stdout, _ = proc.communicate(timeout=15)
+    output = stdout.decode("utf-8")
+    assert proc.returncode not in (0, None), output
+    assert "boom from plugin" in output, output
+
+    # Nothing is listening on the port now the process has exited. This
+    # confirms the socket was not left bound; on its own it cannot prove the
+    # failure preceded the bind, since a port nothing ever touched also
+    # refuses connections.
+    with (
+        pytest.raises(OSError),
+        socket.create_connection(("127.0.0.1", port), timeout=0.2),
+    ):
+        pass