Drain accepted writes consistently during database shutdown

tests/test_write_budget.py

@@ -6,7 +6,6 @@ import time
 import pytest
 
 from datasette.app import Datasette
-from datasette.database import DatasetteClosedError
 from datasette.write_budget import (
     DatabaseAdmissionTimeout,
     DatabaseQueueFull,
@@ -126,8 +125,10 @@ async def test_queue_full_expiration_and_internal_progress(tmp_path):
 
 
 @pytest.mark.asyncio
-async def test_close_waiting_database_releases_pending(tmp_path):
-    ds, dbs = make_databases(tmp_path, max_write_connections=1)
+async def test_close_waiting_database_preserves_admission_deadline(tmp_path):
+    ds, dbs = make_databases(
+        tmp_path, max_write_connections=1, write_queue_timeout_ms=50
+    )
     entered, release = threading.Event(), threading.Event()
 
     def slow(conn):
@@ -142,7 +143,7 @@ async def test_close_waiting_database_releases_pending(tmp_path):
         )
         await wait_for(lambda: ds._write_budget.snapshot()["pending"] == 1)
         dbs[1].close()
-        with pytest.raises(DatasetteClosedError):
+        with pytest.raises(DatabaseAdmissionTimeout):
             await pending
         assert ds._write_budget.snapshot()["pending"] == 0
     finally:

tests/test_write_budget_shutdown.py

@@ -0,0 +1,250 @@
+import asyncio
+import sqlite3
+import threading
+import time
+from pathlib import Path
+
+import pytest
+
+from datasette.app import Datasette
+from datasette.database import Database, DatasetteClosedError
+from datasette.write_budget import RETIRE_WRITER, DatabaseAdmissionTimeout
+
+
+def make_databases(tmp_path, count=2, **settings):
+    paths = []
+    for i in range(count):
+        path = tmp_path / f"shutdown{i}.db"
+        with sqlite3.connect(path) as conn:
+            conn.execute("create table t(id)")
+        conn.close()
+        paths.append(str(path))
+    ds = Datasette(paths, settings={"schema_watch_interval_ms": 0, **settings})
+    return ds, [ds.get_database(f"shutdown{i}") for i in range(count)]
+
+
+async def wait_for(predicate):
+    deadline = time.monotonic() + 5
+    while not predicate():
+        assert time.monotonic() < deadline
+        await asyncio.sleep(0.005)
+
+
+def rows(db):
+    with sqlite3.connect(db.path) as conn:
+        result = conn.execute("select id from t order by rowid").fetchall()
+    conn.close()
+    return [row[0] for row in result]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("retirement_requested", [False, True])
+async def test_close_drains_writes_before_worker_admission(
+    tmp_path, monkeypatch, retirement_requested
+):
+    ds, (db, _) = make_databases(tmp_path, max_write_connections=1)
+    admitted = threading.Event()
+    release = threading.Event()
+    run = ds._write_budget._run
+
+    def paused_scheduler():
+        admitted.set()
+        assert release.wait(5)
+        run()
+
+    monkeypatch.setattr(ds._write_budget, "_run", paused_scheduler)
+    replies = []
+    closing = None
+    try:
+        for i in range(20):
+            _, reply = await db._send_to_write_thread(
+                lambda conn, i=i: conn.execute("insert into t values (?)", [i]),
+                block=False,
+            )
+            replies.append(reply)
+        assert admitted.wait(1)
+        assert db._write_thread is None
+        if retirement_requested:
+            # A previous retirement request can precede tasks accepted later.
+            with db._write_queue.mutex:
+                db._write_queue.queue.appendleft(RETIRE_WRITER)
+        closing = asyncio.create_task(asyncio.to_thread(db.close))
+        await wait_for(lambda: db._closed)
+        release.set()
+        await asyncio.wait_for(closing, 5)
+        await asyncio.gather(*replies)
+        assert rows(db) == list(range(20))
+        assert ds._write_budget.snapshot()["writers"] == 0
+        assert ds._write_budget.snapshot()["stats"]["peak_writers"] == 1
+        with pytest.raises(DatasetteClosedError):
+            await db.execute_write("insert into t values (21)")
+    finally:
+        release.set()
+        if closing is not None:
+            await closing
+        await asyncio.gather(*replies, return_exceptions=True)
+        ds.close()
+
+
+@pytest.mark.asyncio
+async def test_close_waits_for_budget_capacity_then_drains(tmp_path):
+    ds, (first, waiting) = make_databases(tmp_path, max_write_connections=1)
+    entered, release = threading.Event(), threading.Event()
+
+    def slow(conn):
+        entered.set()
+        assert release.wait(5)
+
+    _, first_reply = await first._send_to_write_thread(slow, block=False)
+    closing = None
+    waiting_reply = None
+    try:
+        assert await asyncio.to_thread(entered.wait, 5)
+        _, waiting_reply = await waiting._send_to_write_thread(
+            lambda conn: conn.execute("insert into t values (1)"), block=False
+        )
+        closing = asyncio.create_task(asyncio.to_thread(waiting.close))
+        await wait_for(lambda: waiting._closed)
+        assert ds._write_budget.snapshot()["pending"] == 1
+        release.set()
+        await asyncio.wait_for(closing, 5)
+        await asyncio.gather(first_reply, waiting_reply)
+        assert rows(waiting) == [1]
+        assert ds._write_budget.snapshot()["stats"]["peak_writers"] == 1
+    finally:
+        release.set()
+        if closing is not None:
+            await closing
+        await asyncio.gather(
+            first_reply,
+            *([waiting_reply] if waiting_reply else []),
+            return_exceptions=True,
+        )
+        ds.close()
+
+
+@pytest.mark.asyncio
+async def test_close_preserves_admission_deadline_while_capacity_pinned(tmp_path):
+    ds, (waiting, _) = make_databases(
+        tmp_path, max_write_connections=1, write_queue_timeout_ms=50
+    )
+    memory = ds.add_memory_database("pinned")
+    ran = []
+    try:
+        await memory.execute_write("create table t(id)")
+        _, reply = await waiting._send_to_write_thread(
+            lambda conn: ran.append(True), block=False
+        )
+        await asyncio.wait_for(asyncio.to_thread(waiting.close), 2)
+        with pytest.raises(DatabaseAdmissionTimeout):
+            await reply
+        assert not ran
+        assert ds._write_budget.snapshot()["pending"] == 0
+        assert ds._write_budget.snapshot()["stats"]["peak_writers"] == 1
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+async def test_datasette_shutdown_releases_later_memory_writer_before_drain(tmp_path):
+    ds, (waiting, _) = make_databases(
+        tmp_path, max_write_connections=1, write_queue_timeout_ms=1000
+    )
+    memory = ds.add_memory_database("pinned")
+    try:
+        await memory.execute_write("create table t(id)")
+        _, reply = await waiting._send_to_write_thread(
+            lambda conn: conn.execute("insert into t values (1)"), block=False
+        )
+        await asyncio.wait_for(asyncio.to_thread(ds.close), 5)
+        await reply
+        assert rows(waiting) == [1]
+        assert ds._write_budget.snapshot()["writers"] == 0
+        assert ds._write_budget.snapshot()["stats"]["peak_writers"] == 1
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("database_kind", ["file", "temporary", "internal"])
+async def test_close_timeout_preserves_running_connection(
+    tmp_path, monkeypatch, database_kind
+):
+    import datasette.database
+
+    monkeypatch.setattr(datasette.database, "WRITE_SHUTDOWN_TIMEOUT", 0.05)
+    ds, (db, _) = make_databases(tmp_path, max_write_connections=1)
+    if database_kind == "temporary":
+        db = Database(ds, is_temp_disk=True)
+        ds.add_database(db, name="temporary")
+        await db.execute_write("create table t(id)")
+    elif database_kind == "internal":
+        db = ds._internal_database
+        await db.execute_write("create table t(id)")
+    entered, release = threading.Event(), threading.Event()
+
+    def slow(conn):
+        entered.set()
+        assert release.wait(5)
+        conn.execute("insert into t values (1)")
+        return conn.execute("select id from t").fetchall()[0][0]
+
+    _, running = await db._send_to_write_thread(slow, block=False)
+    queued = None
+    try:
+        assert await asyncio.to_thread(entered.wait, 5)
+        if database_kind != "internal":
+            _, queued = await db._send_to_write_thread(
+                lambda conn: pytest.fail("cancelled callback ran"), block=False
+            )
+        await asyncio.wait_for(asyncio.to_thread(db.close), 2)
+        if queued is not None:
+            with pytest.raises(DatasetteClosedError):
+                await queued
+        assert Path(db.path).exists()
+        assert not running.done()
+        release.set()
+        assert await running == 1
+        if database_kind == "file":
+            assert rows(db) == [1]
+        else:
+            await wait_for(lambda: not Path(db.path).exists())
+        await wait_for(lambda: db._write_thread is None)
+    finally:
+        release.set()
+        await asyncio.gather(
+            running, *([queued] if queued else []), return_exceptions=True
+        )
+        ds.close()
+
+
+@pytest.mark.asyncio
+async def test_close_timeout_during_write_connection_open(tmp_path, monkeypatch):
+    import datasette.database
+
+    monkeypatch.setattr(datasette.database, "WRITE_SHUTDOWN_TIMEOUT", 0.05)
+    ds, (db, _) = make_databases(tmp_path)
+    opened, release = threading.Event(), threading.Event()
+    connect = db.connect
+
+    def paused_connect(**kwargs):
+        conn = connect(**kwargs)
+        opened.set()
+        assert release.wait(5)
+        return conn
+
+    monkeypatch.setattr(db, "connect", paused_connect)
+    _, reply = await db._send_to_write_thread(
+        lambda conn: conn.execute("insert into t values (1)"), block=False
+    )
+    try:
+        assert await asyncio.to_thread(opened.wait, 5)
+        await asyncio.wait_for(asyncio.to_thread(db.close), 2)
+        release.set()
+        await reply
+        assert rows(db) == [1]
+        await wait_for(lambda: db._write_thread is None)
+    finally:
+        release.set()
+        await asyncio.gather(reply, return_exceptions=True)
+        ds.close()