Merge spike/write-idle-thread into spike/integrated

Settings-list conflicts resolved as a union; the settings are
consolidated in a later commit.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EwkLo2rJWidpkDVT3HkvNv

datasette/app.py

@@ -250,6 +250,11 @@ SETTINGS = (
         True,
         "Reuse read connections between queries - turn off to open a new connection for every query",
     ),
+    Setting(
+        "write_thread_idle_timeout_ms",
+        30000,
+        "Milliseconds a database's write thread waits for a write before closing its connection and exiting - set 0 to keep it running",
+    ),
     Setting("sql_time_limit_ms", 1000, "Time limit for a SQL query in milliseconds"),
     Setting(
         "default_facet_size", 30, "Number of values to return for requested facets"

datasette/database.py

@@ -125,6 +125,18 @@ class Database:
         self._cached_derived_table_dependencies = None
         self._write_thread = None
         self._write_queue = None
+        # Guards the write thread lifecycle. A task is only ever enqueued,
+        # and a write thread only ever decides to exit, while holding this
+        # lock - see _send_to_write_thread() and _execute_writes().
+        self._write_thread_lock = threading.Lock()
+        # Write threads that have deregistered (idle exit) but may still be
+        # closing their connection. close() joins them.
+        self._retiring_write_threads = set()
+        # Number of write threads started over the lifetime of this Database,
+        # and how often a thread's idle timeout fired but it found a newly
+        # queued task once it held the lock (so it kept running)
+        self._write_threads_started = 0
+        self._write_thread_exits_averted = 0
         self._closed = False
         self._pending_execute_futures = set()
         self._pending_execute_futures_lock = threading.Lock()
@@ -229,12 +241,16 @@ class Database:
             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.
-        write_thread = self._write_thread
-        if write_thread is not None and self._write_queue is not None:
-            self._write_queue.put(_SHUTDOWN)
-            write_thread.join(timeout=10)
-            if write_thread.is_alive():
+        # 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._write_thread_lock:
+            write_thread = self._write_thread
+            retiring_threads = tuple(self._retiring_write_threads)
+            if write_thread is not None:
+                self._write_queue.put(_SHUTDOWN)
+        for thread in ((write_thread,) if write_thread else ()) + retiring_threads:
+            thread.join(timeout=10)
+            if thread.is_alive():
                 sys.stderr.write(
                     f"Datasette: write thread for {self.name!r} did not exit within 10s\n"
                 )
@@ -545,58 +561,134 @@ class Database:
     async def _send_to_write_thread(
         self, fn, block=True, isolated_connection=False, transaction=True
     ):
-        if self._write_queue is None:
-            self._write_queue = queue.Queue()
-        if self._write_thread is None:
-            self._write_thread = threading.Thread(
-                target=self._execute_writes, daemon=True
-            )
-            self._write_thread.name = f"_execute_writes for database {self.name}"
-            self._write_thread.start()
         task_id = uuid.uuid4()
         loop = asyncio.get_running_loop()
         reply_future = loop.create_future()
-        # Capture the OpenTelemetry context and enqueue time for the write thread
-        self._write_queue.put(
-            WriteTask(
-                fn,
-                task_id,
-                loop,
-                reply_future,
-                isolated_connection,
-                transaction,
-                otel_context_api.get_current(),
-                time.time_ns(),
-                block,
+        with self._write_thread_lock:
+            # Enqueueing and checking for a live write thread happen under
+            # one lock. A write thread only decides to exit while holding
+            # this lock and after seeing an empty queue, so a task can never
+            # be left in the queue with no thread to run it: either the
+            # thread sees this task and keeps going, or it has already
+            # deregistered and we start a new one.
+            if self._closed:
+                raise DatasetteClosedError(f"Database {self.name!r} has been closed")
+            if self._write_queue is None:
+                self._write_queue = queue.Queue()
+            if self._write_thread is None:
+                # Start the thread before enqueueing: if start() fails (for
+                # example "can't start new thread") nothing is stranded
+                self._start_write_thread()
+            # Capture the OpenTelemetry context and enqueue time for the write thread
+            self._write_queue.put(
+                WriteTask(
+                    fn,
+                    task_id,
+                    loop,
+                    reply_future,
+                    isolated_connection,
+                    transaction,
+                    otel_context_api.get_current(),
+                    time.time_ns(),
+                    block,
+                )
             )
-        )
         if block:
             return await reply_future
         else:
             return task_id, reply_future
 
+    def _start_write_thread(self):
+        # Caller must hold self._write_thread_lock
+        thread = threading.Thread(
+            target=self._execute_writes,
+            daemon=True,
+            name=f"_execute_writes for database {self.name}",
+        )
+        thread.start()
+        self._write_thread = thread
+        self._write_threads_started += 1
+
+    def _write_thread_idle_timeout(self):
+        """Seconds the write thread may wait for a task before exiting.
+
+        None means wait forever (the thread lives until close()).
+        """
+        if self.is_memory:
+            # Closing the last connection to a memory database discards
+            # its contents, so its write connection must stay open
+            return None
+        setting = getattr(self.ds, "setting", None)
+        timeout_ms = setting("write_thread_idle_timeout_ms") if setting else None
+        if not timeout_ms or timeout_ms <= 0:
+            return None
+        return timeout_ms / 1000
+
     def _execute_writes(self):
-        # Infinite looping thread that protects the single write connection
-        # to this database
+        # Thread that protects the single write connection to this database.
+        # It runs until close(), or until it has waited idle_timeout seconds
+        # with no task, in which case it closes its connection and exits.
+        # The next write starts a new thread.
+        idle_timeout = self._write_thread_idle_timeout()
         conn_exception = None
         conn = None
+        deregistered = False
         try:
             conn = self.connect(write=True)
             # Threads do not inherit the caller's context, so any spans
             # created by prepare_connection hooks here are root spans
             self.ds._prepare_connection(conn, self.name)
         except Exception as e:  # noqa: BLE001
-            # Stored and re-raised to whoever queues the next write
+            # Stored and re-raised to every write already queued. With an
+            # idle timeout the thread then exits as soon as the queue is
+            # empty, so the next write tries to connect again.
             conn_exception = e
+            if idle_timeout is not None:
+                idle_timeout = 0
+        try:
+            self._process_write_queue(conn, conn_exception, idle_timeout)
+            deregistered = True
+        finally:
+            if conn is not None:
+                try:
+                    conn.close()
+                except Exception:  # noqa: BLE001, S110
+                    # Best-effort close as the write thread exits
+                    pass
+                try:
+                    self._all_connections.remove(conn)
+                except ValueError:
+                    # May already have been cleared by close().
+                    pass
+            current = threading.current_thread()
+            with self._write_thread_lock:
+                self._retiring_write_threads.discard(current)
+                if not deregistered and self._write_thread is current:
+                    # Died from an unexpected BaseException. Deregister so
+                    # the next write starts a fresh thread, and start one
+                    # now if tasks are already waiting.
+                    self._write_thread = None
+                    if not self._closed and not self._write_queue.empty():
+                        self._start_write_thread()
+
+    def _process_write_queue(self, conn, conn_exception, idle_timeout):
         while True:
-            task = self._write_queue.get()
+            try:
+                task = self._write_queue.get(timeout=idle_timeout)
+            except queue.Empty:
+                with self._write_thread_lock:
+                    if not self._write_queue.empty():
+                        # A task was enqueued after the timeout fired
+                        self._write_thread_exits_averted += 1
+                        continue
+                    # Deregister while holding the lock: any write enqueued
+                    # from now on will start a new thread
+                    current = threading.current_thread()
+                    if self._write_thread is current:
+                        self._write_thread = None
+                    self._retiring_write_threads.add(current)
+                return
             if task is _SHUTDOWN:
-                if conn is not None:
-                    try:
-                        conn.close()
-                    except Exception:  # noqa: BLE001, S110
-                        # Best-effort close as the write thread exits
-                        pass
                 return
             # block=True: the caller awaits the result, so the write spans
             # are children of the caller's span. The token must be detached

docs/settings.rst

@@ -280,6 +280,29 @@ Default: ``on``
 
 Reuse read connections between queries. Turn this off to open, prepare and close a new connection for every read query.
 
+.. _setting_write_thread_idle_timeout_ms:
+
+write_thread_idle_timeout_ms
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. [[[cog
+    setting_default(cog, "write_thread_idle_timeout_ms")
+.. ]]]
+
+Default: ``30000``
+
+.. [[[end]]]
+
+Each database that receives writes gets its own write thread, holding a single write connection. Writes to that database are queued and executed one at a time by that thread.
+
+If a write thread has had no work for this many milliseconds it closes its connection and exits. The next write to that database starts a new thread and opens a new connection. This keeps an instance with many databases from holding a thread and a file descriptor open for every database it has ever written to.
+
+Anything a write function attaches to the write connection itself - temporary tables, ``ATTACH`` statements, ``PRAGMA`` settings - is lost when the thread exits. In-memory databases are exempt, since closing their only write connection would discard their contents.
+
+Set this to 0 to keep write threads running until Datasette shuts down::
+
+    datasette mydatabase.db --setting write_thread_idle_timeout_ms 0
+
 .. _setting_allow_facet:
 
 allow_facet

tests/test_api.py

@@ -763,6 +763,7 @@ async def test_settings_json(ds_client):
         "connection_idle_timeout": 60,
         "connection_pool_wait_ms": 0,
         "pool_read_connections": True,
+        "write_thread_idle_timeout_ms": 30000,
         "cache_size_kb": 0,
         "allow_csv_stream": True,
         "max_csv_mb": 100,

tests/test_write_thread_idle.py

@@ -0,0 +1,355 @@
+"""
+Tests for the write_thread_idle_timeout_ms setting: a database's write thread
+closes its connection and exits after sitting idle, and the next write
+starts a new one.
+"""
+
+import asyncio
+import queue
+import random
+import sqlite3
+import threading
+import time
+
+import pytest
+
+from datasette.app import Datasette
+from datasette.database import Database, DatasetteClosedError
+
+
+def _make_db(tmp_path, name="idle", **settings):
+    path = tmp_path / f"{name}.db"
+    conn = sqlite3.connect(path)
+    conn.execute("create table t (id integer primary key, writer integer, seq integer)")
+    conn.close()
+    ds = Datasette([str(path)], settings=settings)
+    return ds, ds.get_database(name)
+
+
+async def _wait_for(predicate, timeout=5):
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if predicate():
+            return True
+        await asyncio.sleep(0.005)
+    return predicate()
+
+
+def _write_threads(name):
+    thread_name = f"_execute_writes for database {name}"
+    return [t for t in threading.enumerate() if t.name == thread_name]
+
+
+@pytest.mark.asyncio
+async def test_write_thread_exits_when_idle_and_restarts(tmp_path):
+    ds, db = _make_db(tmp_path, write_thread_idle_timeout_ms=50)
+    try:
+        await db.execute_write("insert into t (writer, seq) values (1, 1)")
+        thread = db._write_thread
+        assert thread is not None and thread.is_alive()
+        connections_with_writer = len(db._all_connections)
+        assert await _wait_for(lambda: db._write_thread is None)
+        thread.join(timeout=5)
+        assert not thread.is_alive()
+        # Its write connection was closed and is no longer tracked
+        assert len(db._all_connections) == connections_with_writer - 1
+        assert not _write_threads("idle")
+        # The next write starts a new thread
+        await db.execute_write("insert into t (writer, seq) values (1, 2)")
+        assert db._write_thread is not None and db._write_thread is not thread
+        assert db._write_threads_started == 2
+        count = (await db.execute("select count(*) from t")).single_value()
+        assert count == 2
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+async def test_write_thread_idle_timeout_zero_keeps_thread(tmp_path):
+    ds, db = _make_db(tmp_path, write_thread_idle_timeout_ms=0)
+    try:
+        await db.execute_write("insert into t (writer, seq) values (1, 1)")
+        thread = db._write_thread
+        await asyncio.sleep(0.2)
+        assert db._write_thread is thread
+        assert thread.is_alive()
+    finally:
+        ds.close()
+    assert not thread.is_alive()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("named", (True, False))
+async def test_memory_database_write_thread_never_idles(named):
+    # Closing the write connection of a memory database would lose its data
+    ds = Datasette(settings={"write_thread_idle_timeout_ms": 10})
+    if named:
+        db = ds.add_memory_database("idle_mem")
+    else:
+        db = ds.add_database(Database(ds, is_memory=True), name="idle_unnamed")
+    try:
+        await db.execute_write("create table t (id integer primary key)")
+        await db.execute_write("insert into t (id) values (1)")
+        thread = db._write_thread
+        await asyncio.sleep(0.1)
+        assert db._write_thread is thread and thread.is_alive()
+        count = await db.execute_write_fn(
+            lambda conn: conn.execute("select count(*) from t").fetchone()[0]
+        )
+        assert count == 1
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+async def test_close_after_idle_exit(tmp_path):
+    ds, db = _make_db(tmp_path, write_thread_idle_timeout_ms=10)
+    await db.execute_write("insert into t (writer, seq) values (1, 1)")
+    thread = db._write_thread
+    assert await _wait_for(lambda: db._write_thread is None)
+    ds.close()
+    thread.join(timeout=5)
+    assert not thread.is_alive()
+    with pytest.raises(DatasetteClosedError):
+        await db.execute_write("insert into t (writer, seq) values (1, 2)")
+
+
+@pytest.mark.asyncio
+async def test_close_drains_queued_writes_with_idle_timeout(tmp_path):
+    ds, db = _make_db(tmp_path, write_thread_idle_timeout_ms=10)
+    release = threading.Event()
+
+    def slow(conn):
+        release.wait(5)
+        conn.execute("insert into t (writer, seq) values (0, 0)")
+
+    await db.execute_write_fn(slow, block=False)
+    futures = []
+    for i in range(20):
+        _, future = await db._send_to_write_thread(
+            lambda conn, i=i: conn.execute(
+                "insert into t (writer, seq) values (1, ?)", [i]
+            ),
+            block=False,
+        )
+        futures.append(future)
+    release.set()
+    # close() is synchronous, so the queued writes complete before it returns
+    db.close()
+    await asyncio.wait_for(asyncio.gather(*futures), timeout=5)
+    conn = sqlite3.connect(db.path)
+    try:
+        assert conn.execute("select count(*) from t").fetchone()[0] == 21
+    finally:
+        conn.close()
+        ds.close()
+
+
+@pytest.mark.asyncio
+async def test_connect_failure_retried_after_queue_drains(tmp_path):
+    # With an idle timeout, a failed connect is reported to the writes that
+    # were queued, then the thread exits and the next write tries again
+    ds, db = _make_db(tmp_path, write_thread_idle_timeout_ms=30000)
+
+    class ConnectError(Exception):
+        pass
+
+    original_connect = db.connect
+
+    def broken_connect(write=False):
+        if write:
+            raise ConnectError("no")
+        return original_connect(write=write)
+
+    db.connect = broken_connect
+    try:
+        with pytest.raises(ConnectError):
+            await db.execute_write("insert into t (writer, seq) values (1, 1)")
+    finally:
+        db.connect = original_connect
+    assert await _wait_for(lambda: db._write_thread is None)
+    await db.execute_write("insert into t (writer, seq) values (1, 2)")
+    assert db._write_threads_started == 2
+    ds.close()
+
+
+class _RacyQueue(queue.Queue):
+    """
+    The first time the write thread asks for a task while the queue is
+    empty, behave as if the idle timeout fired - but first pause, so the
+    test can enqueue a task in the window between the timeout and the
+    thread's exit decision.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.armed = False
+        self.window_open = threading.Event()
+        self.proceed = threading.Event()
+
+    def get(self, block=True, timeout=None):
+        if self.armed and self.empty():
+            self.armed = False
+            self.window_open.set()
+            self.proceed.wait(5)
+            raise queue.Empty
+        return super().get(block, timeout)
+
+
+@pytest.mark.asyncio
+async def test_task_enqueued_between_timeout_and_exit_is_not_stranded(tmp_path):
+    ds, db = _make_db(tmp_path, write_thread_idle_timeout_ms=60000)
+    racy = _RacyQueue()
+    db._write_queue = racy
+    try:
+        racy.armed = True
+        await db.execute_write("insert into t (writer, seq) values (1, 1)")
+        # The thread has finished that write and is now inside the window
+        assert await asyncio.to_thread(racy.window_open.wait, 5)
+        thread = db._write_thread
+        assert thread is not None
+        _, future = await db._send_to_write_thread(
+            lambda conn: conn.execute("insert into t (writer, seq) values (1, 2)"),
+            block=False,
+        )
+        # The thread is still registered, so no new thread was started
+        assert db._write_threads_started == 1
+        racy.proceed.set()
+        await asyncio.wait_for(future, timeout=5)
+        assert db._write_thread is thread
+        count = (await db.execute("select count(*) from t")).single_value()
+        assert count == 2
+    finally:
+        racy.proceed.set()
+        ds.close()
+
+
+def _recorder(log, writer, seq):
+    # A cheap task (no disk write), so tasks finish in microseconds and the
+    # idle timeout keeps firing between bursts
+    def record(conn):
+        conn.execute("select 1").fetchone()
+        log.append((writer, seq, threading.get_ident()))
+
+    return record
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("idle_timeout_ms", (1, 5))
+async def test_idle_exit_race_stress(tmp_path, idle_timeout_ms):
+    """
+    Many concurrent writers working in rounds. Each round starts once the
+    previous round's tasks are done, so the write thread's idle clock starts
+    at about the same moment, and the writers then wait about 0.8x-1.3x of
+    the idle timeout before queueing a burst. Tasks therefore keep landing
+    on both sides of, and right on top of, the thread's exit decision.
+    Every task must complete, and each writer's tasks must run in the order
+    they were queued.
+    """
+    ds, db = _make_db(tmp_path, write_thread_idle_timeout_ms=idle_timeout_ms)
+    idle_timeout = idle_timeout_ms / 1000
+    rng = random.Random(42)
+    writers = 12
+    rounds = 150 if idle_timeout_ms == 1 else 60
+    seqs = [0] * writers
+    log = []
+
+    async def writer_round(w, delay, burst):
+        await asyncio.sleep(delay)
+        pending = []
+        for _ in range(burst):
+            seq = seqs[w]
+            seqs[w] += 1
+            fn = _recorder(log, w, seq)
+            if seq % 3 == 0:
+                # Fire-and-forget write
+                _, future = await db._send_to_write_thread(fn, block=False)
+                pending.append(future)
+            else:
+                await db.execute_write_fn(fn, transaction=bool(seq % 2))
+        await asyncio.gather(*pending)
+
+    async def all_rounds():
+        for _ in range(rounds):
+            # One shared delay per round close to the timeout, plus a little
+            # jitter per writer, so arrivals cluster around the exit decision
+            base = idle_timeout * rng.uniform(0.8, 1.2)
+            await asyncio.gather(
+                *(
+                    writer_round(
+                        w, base + idle_timeout * rng.uniform(0, 0.1), rng.randint(1, 3)
+                    )
+                    for w in range(writers)
+                )
+            )
+
+    try:
+        await asyncio.wait_for(all_rounds(), timeout=120)
+        assert len(log) == sum(seqs)
+        for w in range(writers):
+            assert [seq for writer_id, seq, _ in log if writer_id == w] == list(
+                range(seqs[w])
+            )
+        print(
+            f"idle_timeout_ms={idle_timeout_ms} tasks={len(log)} "
+            f"threads_started={db._write_threads_started} "
+            f"distinct_thread_idents={len({ident for _, _, ident in log})} "
+            f"exits_averted={db._write_thread_exits_averted}"
+        )
+        # Only meaningful if the thread really did churn
+        assert db._write_threads_started > rounds // 5
+    finally:
+        ds.close()
+    assert not _write_threads("idle")
+
+
+def test_idle_exit_race_stress_many_event_loops(tmp_path):
+    """
+    Producers on several OS threads, each with its own event loop, race
+    with the write thread's exit decision for real (no single-loop
+    ordering between producers). Rounds are synchronised with a barrier,
+    then each producer waits a random 0.5x-1.5x of the idle timeout.
+    """
+    ds, db = _make_db(tmp_path, write_thread_idle_timeout_ms=1)
+    loops = 6
+    rounds = 200
+    barrier = threading.Barrier(loops)
+    errors = []
+    log = []
+
+    def run_loop(w):
+        rng = random.Random(w)
+
+        async def go():
+            for seq in range(rounds):
+                barrier.wait(timeout=30)
+                await asyncio.sleep(0.001 * rng.uniform(0.5, 1.5))
+                await asyncio.wait_for(
+                    db.execute_write_fn(_recorder(log, w, seq)), timeout=10
+                )
+
+        try:
+            asyncio.run(go())
+        except BaseException as e:  # noqa: BLE001
+            errors.append(e)
+            barrier.abort()
+
+    threads = [threading.Thread(target=run_loop, args=(w,)) for w in range(loops)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(timeout=120)
+    try:
+        assert not errors, errors
+        assert len(log) == loops * rounds
+        for w in range(loops):
+            assert [seq for writer_id, seq, _ in log if writer_id == w] == list(
+                range(rounds)
+            )
+        print(
+            f"tasks={len(log)} threads_started={db._write_threads_started} "
+            f"exits_averted={db._write_thread_exits_averted}"
+        )
+        assert db._write_threads_started > rounds // 5
+    finally:
+        ds.close()