Drain accepted writes consistently during database shutdown

datasette/app.py

@@ -1081,7 +1081,13 @@ class Datasette:
         # Stop reporting metrics before closing databases
         unregister_datasette(self)
         first_exception = None
-        for db in list(self.databases.values()):
+        databases = list(self.databases.values())
+        # Release every user database's idle or pinned writer before
+        # waiting for accepted work to obtain the global writer budget.
+        # The internal catalog remains available until user drains finish.
+        for db in databases:
+            db._begin_close()
+        for db in databases:
             try:
                 db.close()
             except Exception as e:  # noqa: BLE001

datasette/database.py

@@ -86,6 +86,7 @@ from .write_budget import (
 logger = logging.getLogger(__name__)
 
 EXECUTE_WRITE_RETURNING_LIMIT = 10
+WRITE_SHUTDOWN_TIMEOUT = 10
 
 AttachedDatabase = namedtuple("AttachedDatabase", ("seq", "name", "file"))
 
@@ -95,6 +96,7 @@ class DatasetteClosedError(RuntimeError):
 
 
 _SHUTDOWN = object()
+_write_connection_owner = contextvars.ContextVar("write_connection_owner", default=None)
 
 
 class Database:
@@ -163,6 +165,10 @@ class Database:
         self._write_threads_started = 0
         self._write_thread_exits_averted = 0
         self._closed = False
+        self._close_lock = threading.RLock()
+        self._close_complete = False
+        self._closing_execute_futures = ()
+        self._cleanup_temp_file_on_writer_exit = False
         self._pending_execute_futures = set()
         self._pending_execute_futures_lock = threading.Lock()
         # The write connection in non-threaded mode (reads go through
@@ -229,6 +235,21 @@ class Database:
         except ValueError:
             pass
 
+    def _connect_write_thread(self):
+        # Publish the connection only after marking its owning writer. A
+        # timed-out close() must never see a half-registered write connection
+        # and close it before the callback can use it.
+        token = _write_connection_owner.set(self)
+        try:
+            return self.connect(write=True)
+        finally:
+            _write_connection_owner.reset(token)
+
+    def _track_connection(self, conn):
+        if _write_connection_owner.get() is self:
+            self._write_thread_connections.add(conn)
+        self._all_connections.append(conn)
+
     def _untracked_connection_opened(self):
         """Call before opening a connection with ``connect(track=False)``.
 
@@ -361,12 +382,12 @@ class Database:
             if not write:
                 conn.execute("PRAGMA query_only=1")
             if track:
-                self._all_connections.append(conn)
+                self._track_connection(conn)
             return conn
         if self.is_memory:
             conn = sqlite3.connect(":memory:", uri=True, check_same_thread=False)
             if track:
-                self._all_connections.append(conn)
+                self._track_connection(conn)
             return conn
 
         # mode=ro or immutable=1?
@@ -397,7 +418,7 @@ class Database:
         )
         self._file_seen = True
         if track:
-            self._all_connections.append(conn)
+            self._track_connection(conn)
         if self.is_temp_disk and not self._wal_enabled:
             conn.execute("PRAGMA journal_mode=WAL")
             self._wal_enabled = True
@@ -407,45 +428,63 @@ class Database:
         """Release all resources held by this database.
 
         Idempotent. After close() further calls to execute()/execute_fn()/
-        execute_write()/execute_write_fn() raise DatasetteClosedError.
+        execute_write()/execute_write_fn() raise DatasetteClosedError. Accepted
+        writes drain through normal admission, with their existing deadlines.
+        Budgeted writes still awaiting admission after ten seconds are
+        cancelled; running writers retain their connections until finished.
         """
-        if self._closed:
-            return
-        with self._pending_execute_futures_lock, self._untracked_cond:
+        with self._close_lock:
+            if self._close_complete:
+                return
+            self._begin_close()
+            self._finish_close()
+            self._close_complete = True
+
+    def _begin_close(self):
+        # Datasette begins closing every user database before waiting for
+        # any one of them, releasing pinned writers needed by queued work.
+        with self._close_lock:
             if self._closed:
                 return
-            self._closed = True
-            pending_execute_futures = tuple(self._pending_execute_futures)
-        # Shut down the write thread, if any, via a sentinel. The thread
-        # drains any writes already queued before the sentinel and then
-        # closes its own write connection and returns. _closed is already
-        # set, so once we hold the lock no further task can be enqueued.
+            with self._pending_execute_futures_lock, self._untracked_cond:
+                self._closed = True
+                self._closing_execute_futures = tuple(self._pending_execute_futures)
+            with self._write_thread_lock:
+                # Include writers awaiting admission: the scheduler must
+                # drain accepted tasks before consuming this sentinel.
+                if self._write_queue is not None:
+                    self._write_queue.put(_SHUTDOWN)
+
+    def _finish_close(self):
+        budget = self.ds.__dict__.get("_write_budget_instance")
+        still_running = False
+        if self._budgeted_writer and budget is not None:
+            still_running = not budget.drain_database(self, WRITE_SHUTDOWN_TIMEOUT)
         with self._write_thread_lock:
             write_thread = self._write_thread
             retiring_threads = tuple(self._retiring_write_threads)
-            if write_thread is not None and write_thread.is_alive():
-                self._write_queue.put(_SHUTDOWN)
-            elif self._budgeted_writer:
-                self.ds._write_budget.cancel_database(self)
-        still_running = False
-        for thread in ((write_thread,) if write_thread else ()) + retiring_threads:
-            thread.join(timeout=10)
-            if thread.is_alive():
-                still_running = True
-                sys.stderr.write(
-                    f"Datasette: write thread for {self.name!r} did not exit within 10s\n"
-                )
-                sys.stderr.flush()
+        if not self._budgeted_writer:
+            deadline = time.monotonic() + WRITE_SHUTDOWN_TIMEOUT
+            for thread in ((write_thread,) if write_thread else ()) + retiring_threads:
+                thread.join(timeout=max(0, deadline - time.monotonic()))
+                still_running = still_running or thread.is_alive()
+        if still_running:
+            sys.stderr.write(
+                f"Datasette: write thread for {self.name!r} did not exit within "
+                f"{WRITE_SHUTDOWN_TIMEOUT:g}s\n"
+            )
+            sys.stderr.flush()
         # A write thread that is still running owns its connection and closes
         # it when it finishes; closing it here could free it mid-query
         # (membership tests only: copying the set could race the thread)
         in_use = self._write_thread_connections if still_running else ()
-        for future in pending_execute_futures:
+        for future in self._closing_execute_futures:
             try:
                 future.result()
             except Exception:  # noqa: BLE001, S110
                 # Shutdown teardown - a failed pending write must not block close()
                 pass
+        self._closing_execute_futures = ()
         # Close idle pooled read connections; any still leased are closed
         # by the pool when they are released
         read_pool = getattr(self.ds, "_read_pool_or_none", None)
@@ -468,7 +507,13 @@ class Database:
                 pass
             self._write_connection = None
         if self.is_temp_disk:
-            self._cleanup_temp_file()
+            with self._write_thread_lock:
+                if (
+                    self._write_thread is not None and self._write_thread.is_alive()
+                ) or self._retiring_write_threads:
+                    self._cleanup_temp_file_on_writer_exit = True
+                else:
+                    self._cleanup_temp_file()
 
     def _cleanup_temp_file(self):
         if self.is_temp_disk and self.path:
@@ -884,6 +929,14 @@ class Database:
                     self._write_thread = None
                     if not self._closed and not self._write_queue.empty():
                         self._start_write_thread()
+                if self._closed and self._write_thread is current:
+                    self._write_thread = None
+                if (
+                    self._cleanup_temp_file_on_writer_exit
+                    and self._write_thread is None
+                    and not self._retiring_write_threads
+                ):
+                    self._cleanup_temp_file()
 
     def _process_write_queue(self, writer):
         configured_idle_timeout = self._write_thread_idle_timeout()
@@ -913,7 +966,11 @@ class Database:
                         self._write_thread = None
                     self._retiring_write_threads.add(current)
                 return
-            if task is _SHUTDOWN or task is RETIRE_WRITER:
+            if task is _SHUTDOWN:
+                return
+            if task is RETIRE_WRITER:
+                if self._closed:
+                    continue
                 return
             if self._budgeted_writer and not self.ds._write_budget.take(self, task):
                 continue
@@ -974,10 +1031,9 @@ class Database:
                                 task.isolated_connection,
                             )
                             span.set_attribute(TRANSACTION, task.transaction)
-                            isolated_connection = self.connect(write=True)
+                            isolated_connection = self._connect_write_thread()
                             # Owned by this thread: close() must not close it
                             # if this thread outlives close()'s join timeout
-                            self._write_thread_connections.add(isolated_connection)
                             try:
                                 result = _call_with_lease(
                                     task.fn, isolated_connection, self.name, "isolated"
@@ -1627,8 +1683,7 @@ class _WriteConnection:
         self.generation = db._conn_generation
         self.exception = None
         try:
-            self.conn = db.connect(write=True)
-            db._write_thread_connections.add(self.conn)
+            self.conn = db._connect_write_thread()
             _pin_write_connection_pragmas(self.conn)
             # Threads do not inherit the caller's context, so any spans
             # created by prepare_connection hooks here are root spans

datasette/write_budget.py

@@ -195,6 +195,27 @@ class WriteBudget:
         with self._condition:
             self._condition.notify_all()
 
+    def drain_database(self, db, timeout):
+        """Wait for accepted work and its owning threads, including admission.
+
+        The caller has stopped submissions and queued the shutdown sentinel.
+        Normal admission deadlines still apply. If shutdown itself times out,
+        only callbacks that have not started are cancelled; running callbacks
+        keep ownership of their connections until they return.
+        """
+        with self._condition:
+            drained = self._condition.wait_for(
+                lambda: db not in self._active
+                and not any(owner is db for owner, _ in self._pending.values()),
+                timeout=timeout,
+            )
+        if not drained:
+            self.cancel_database(db)
+        else:
+            with self._inline_lock:
+                self._inline.pop(db, None)
+        return drained
+
     def cancel_database(self, db):
         from .database import DatasetteClosedError
 
@@ -262,6 +283,9 @@ class WriteBudget:
                     stop = True
                 else:
                     stop = False
+                # close() may be waiting for a starting or retiring worker,
+                # not just for callbacks to complete.
+                self._condition.notify_all()
             for task in expired:
                 self._fail(
                     task,
@@ -272,12 +296,11 @@ class WriteBudget:
             for db in start:
                 try:
                     with db._write_thread_lock:
-                        if db._closed:
-                            self.cancel_database(db)
-                            thread = None
-                        else:
-                            db._start_write_thread()
-                            thread = db._write_thread
+                        # close() rejects new submissions, but accepted work
+                        # remains eligible for this normally budgeted slot.
+                        # Its shutdown sentinel follows the accepted tasks.
+                        db._start_write_thread()
+                        thread = db._write_thread
                     with self._condition:
                         if thread is None:
                             self._active.pop(db, None)

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()