From 0cf9920012c53b70f6c5e654a77ffe7732671f4c Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Thu, 8 Oct 2026 15:41:02 -0700 Subject: [PATCH] Drain accepted writes consistently during database shutdown --- datasette/app.py | 8 +- datasette/database.py | 121 ++++++++++---- datasette/write_budget.py | 35 +++- tests/test_write_budget.py | 9 +- tests/test_write_budget_shutdown.py | 250 ++++++++++++++++++++++++++++ 5 files changed, 379 insertions(+), 44 deletions(-) create mode 100644 tests/test_write_budget_shutdown.py diff --git a/datasette/app.py b/datasette/app.py index 0ecf80b3..6a571012 100644 --- a/datasette/app.py +++ b/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 diff --git a/datasette/database.py b/datasette/database.py index 42f1f528..3096ac87 100644 --- a/datasette/database.py +++ b/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 diff --git a/datasette/write_budget.py b/datasette/write_budget.py index 4ce29a13..3ab06fc2 100644 --- a/datasette/write_budget.py +++ b/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) diff --git a/tests/test_write_budget.py b/tests/test_write_budget.py index 1eb33c4b..ba8deb99 100644 --- a/tests/test_write_budget.py +++ b/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: diff --git a/tests/test_write_budget_shutdown.py b/tests/test_write_budget_shutdown.py new file mode 100644 index 00000000..86297286 --- /dev/null +++ b/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()