fix(system): make retiring a library connection race-free and cheap

Review round for the cross-thread close fix:

- cursor() now records its cursor under the lock _close_now copies the
  cursor set under, and refuses once the connection is retired; a cursor
  opened while the connection closed raised "Set changed size during
  iteration" out of retire() on Python 3.10-3.12. retire() never raises,
  retire_connections keeps going when one connection fails, and the mover
  never reports a completed switch as failed because retiring or awaiting
  old connections did.
- The restore's backup target is marked in use (ManagedConnection.using),
  and in_transaction and create_function are guarded.
- Iterating a cursor marks the connection in use once per fetchmany(128)
  chunk instead of once per row: a 1000-row scan costs 1.10x a plain
  cursor instead of 2.7x (3.14) / 3.8x (3.10).
- Document what an interrupt does not end, the statements a retired
  connection cannot hold back, close() from another thread returning
  before the file is closed, and the await timeout ending as old-in-use.

database/library_move.py

@@ -725,16 +725,14 @@ def _switch_guarded(move: _Move, location) -> list:
             # Retire, never wait, here: a connection still in use may be
             # waiting on the very locks this thread holds.
             try:
-                move.retiring.extend(
-                    PromptModel.retire_detached(
-                        move.unswitch(plan.db_path, location.db_path), logger
-                    )
-                )
+                undone = move.unswitch(plan.db_path, location.db_path)
             except Exception:
                 logger.exception("Library move: undoing the switch failed")
+            else:
+                _retire(move, undone)
             # The old file's connections a partial switch detached: the
             # models reopen it on next use, so close them now, not at GC.
-            move.retiring.extend(PromptModel.retire_detached(detached, logger))
+            _retire(move, detached)
             raise MoveFailed("switch-failed") from error
         move.checkpoint("switched")
         return detached
@@ -755,7 +753,23 @@ def _copy_and_switch(move: _Move, location, attempts: int):
     try:
         return _copy_and_switch_locked(move, location, attempts)
     finally:
-        PromptModel.await_detached(move.retiring, logger)
+        try:
+            PromptModel.await_detached(move.retiring, logger)
+        except Exception:
+            logger.exception("Library move: waiting for old connections failed")
+
+
+def _retire(move: _Move, connections: list) -> None:
+    """Retire ``connections``, adding those still in use to ``move.retiring``.
+
+    Never raises: a switch that completed must not be reported as failed
+    because one of the old connections could not be retired (logged; the
+    rename then reports the old files as still in use).
+    """
+    try:
+        move.retiring.extend(PromptModel.retire_detached(connections, logger))
+    except Exception:
+        logger.exception("Library move: retiring the old connections failed")
 
 
 def _copy_and_switch_locked(move: _Move, location, attempts: int):
@@ -781,7 +795,7 @@ def _copy_and_switch_locked(move: _Move, location, attempts: int):
                 # on it is interrupted and cannot commit, so it fails (and
                 # runs again on the new library) instead of landing in the
                 # old file. Its own thread closes the connection.
-                move.retiring.extend(PromptModel.retire_detached(detached, logger))
+                _retire(move, detached)
             finally:
                 if source.in_transaction:
                     source.execute("ROLLBACK")

database/managed_connection.py

@@ -8,12 +8,27 @@ Windows the process dies with "access violation", elsewhere the owner sees
 garbage such as a ``MemoryError``.
 
 ``ManagedConnection`` counts the calls its owner is inside. ``retire()`` is
-safe from any thread: an idle connection is closed at once; one in use is
-interrupted (``sqlite3_interrupt`` is thread-safe) and closed by its owner as
-the call ends. A retired connection refuses new work with the error a closed
-connection raises, and never commits: a write that started before the
-retirement and waited for the database lock is rolled back, so it cannot
-land in a file the process has switched away from.
+safe from any thread and never raises: an idle connection is closed at once;
+one in use is closed by its owner as the call ends. A retired connection
+refuses new work (cursors, statements, commits) with the error a closed
+connection raises.
+
+What retiring a connection in use does and does not do:
+
+- It calls ``sqlite3_interrupt`` (thread-safe). That ends a running
+  statement at its next step, but it does not wake a statement that is
+  waiting for another connection's lock: that one waits until the lock is
+  released (or the busy timeout ends it), and only then is it interrupted.
+  An interrupt that arrives between two statements, for example between the
+  implicit ``BEGIN`` and the DML or between ``executemany`` rows, is lost.
+- Writes are kept out of a file the process has switched away from because
+  a retired connection never commits: legacy-mode DML runs inside a
+  transaction that only ``commit()`` (refused) or ``close()`` (rolls back)
+  ends. A statement that commits by itself is not covered: DDL or a PRAGMA
+  write outside a transaction, or an explicit ``COMMIT`` passed to
+  ``execute``/``executescript`` (``executescript`` also commits first, in C,
+  on Python before 3.12). Managed connections must never be used that way
+  for library writes.
 """
 
 from __future__ import annotations
@@ -23,11 +38,16 @@ import logging
 import threading
 import time
 import weakref
-from typing import Iterable, List, Optional
+from typing import Iterable, Iterator, List, Optional
 
 import sqlite3
 
 CLOSED_MESSAGE = "Cannot operate on a closed database."
+# Rows fetched per step when a cursor is iterated: the connection is marked
+# in use once per chunk instead of once per row (iteration costs about what
+# a plain cursor costs). A loop that stops early leaves up to this many rows
+# fetched; ``fetchone``/``fetchmany`` after a partial loop skip them.
+ITER_CHUNK = 128
 
 
 class ManagedCursor(sqlite3.Cursor):
@@ -81,6 +101,16 @@ class ManagedCursor(sqlite3.Cursor):
         finally:
             conn._leave()
 
+    def __iter__(self) -> Iterator:
+        return self._chunks()
+
+    def _chunks(self) -> Iterator:
+        while True:
+            rows = self.fetchmany(ITER_CHUNK)
+            if not rows:
+                return
+            yield from rows
+
     def __next__(self):
         conn = self.connection
         conn._enter()
@@ -106,9 +136,13 @@ class ManagedConnection(sqlite3.Connection):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._use_lock = threading.Lock()
+        # Guards _depth, _retired, _closing and _cursors. Re-entrant so that
+        # code running in the closing thread while it holds the lock (a
+        # cursor() during the cursor-set copy) is refused instead of hanging.
+        self._use_lock = threading.RLock()
         self._depth = 0
         self._retired = False
+        self._closing = False
         self._released = threading.Event()
         # Python before 3.12 keeps the file open after close() while a cursor
         # still holds a statement; these are closed first (see _close_now).
@@ -130,9 +164,24 @@ class ManagedConnection(sqlite3.Connection):
         if close_now:
             self._close_now()
 
+    @contextlib.contextmanager
+    def using(self):
+        """Mark the connection in use for a call this class cannot wrap.
+
+        For example as the *target* of another connection's ``backup()``.
+        """
+        self._enter()
+        try:
+            yield self
+        finally:
+            self._leave()
+
     def cursor(self, factory=None):
-        cursor = super().cursor(factory or ManagedCursor)
-        self._cursors.add(cursor)
+        with self._use_lock:
+            if self._retired:
+                raise sqlite3.ProgrammingError(CLOSED_MESSAGE)
+            cursor = super().cursor(factory or ManagedCursor)
+            self._cursors.add(cursor)
         return cursor
 
     def execute(self, sql, parameters=(), /):
@@ -165,6 +214,21 @@ class ManagedConnection(sqlite3.Connection):
         finally:
             self._leave()
 
+    def create_function(self, *args, **kwargs):
+        self._enter()
+        try:
+            return super().create_function(*args, **kwargs)
+        finally:
+            self._leave()
+
+    @property
+    def in_transaction(self) -> bool:
+        self._enter()
+        try:
+            return super().in_transaction
+        finally:
+            self._leave()
+
     def __exit__(self, exc_type, exc_value, traceback):
         # The C implementation calls its own commit/rollback, not these.
         if exc_type is None:
@@ -181,24 +245,37 @@ class ManagedConnection(sqlite3.Connection):
     # ── closing ──────────────────────────────────────────────────────
 
     def close(self) -> None:
-        """Close now if idle, else as soon as the running call ends (any thread)."""
+        """Close now if idle, else as soon as the running call ends.
+
+        Unlike ``sqlite3.Connection.close``, a call from another thread while
+        the owner is inside a call returns before the file is closed. A caller
+        that must know the file is closed (before renaming or deleting it)
+        waits for ``wait_closed`` (``PromptModel.close_detached`` does).
+        """
         self.retire()
 
     def retire(self) -> bool:
-        """Refuse new work and close; True when the connection is closed already."""
-        with self._use_lock:
-            if not self._retired:
-                self._retired = True
-                if self._depth:
-                    # Aborts the running statement, or the one waiting for
-                    # the lock as soon as it gets it. Thread-safe; the owner
-                    # cannot close the connection while this lock is held.
-                    self.interrupt()
+        """Refuse new work and close; True when the connection is closed already.
+
+        Safe from any thread, and never raises: failures are logged.
+        """
+        try:
+            with self._use_lock:
+                if not self._retired:
+                    self._retired = True
+                    if self._depth:
+                        # Ends the running statement (see the module docstring
+                        # for what it does not end). The owner cannot start
+                        # closing the connection while this lock is held.
+                        self.interrupt()
+                        return False
+                elif self._depth:
                     return False
-            elif self._depth:
-                return False
-            else:
-                return self._released.wait(0)
+                else:
+                    return self._released.wait(0)
+        except Exception:
+            self._log_exception("Error retiring database connection")
+            return self._released.wait(0)
         self._close_now()
         return True
 
@@ -207,18 +284,27 @@ class ManagedConnection(sqlite3.Connection):
 
     def _close_now(self) -> None:
         # Only ever called with no call into the connection running, and a
-        # retired connection refuses new ones: its cursors are idle.
-        for cursor in list(self._cursors):
-            with contextlib.suppress(sqlite3.Error):
-                sqlite3.Cursor.close(cursor)
+        # retired connection refuses new calls and cursors: its cursors are
+        # idle, and the set is copied under the lock that guards it.
+        with self._use_lock:
+            if self._closing:
+                return  # another thread closes it and sets _released
+            self._closing = True
+            cursors = list(self._cursors)
         try:
+            for cursor in cursors:
+                with contextlib.suppress(sqlite3.Error):
+                    sqlite3.Cursor.close(cursor)
             super().close()
-        except sqlite3.Error as error:
-            if self.close_logger is not None:
-                self.close_logger.warning(f"Error closing database connection: {error}")
+        except Exception:
+            self._log_exception("Error closing database connection")
         finally:
             self._released.set()
 
+    def _log_exception(self, message: str) -> None:
+        if self.close_logger is not None:
+            self.close_logger.exception(message)
+
 
 def retire_connections(
     connections: Iterable[sqlite3.Connection], logger=None
@@ -226,27 +312,34 @@ def retire_connections(
     """Retire every connection; return those whose owner has yet to close them.
 
     Safe while holding a lock the owners wait on: nothing here waits for them.
+    One connection failing to retire is logged and never skips the rest.
     """
     pending = []
     for conn in connections:
-        if isinstance(conn, ManagedConnection):
-            if logger is not None:
-                conn.close_logger = logger
-            if not conn.retire():
-                pending.append(conn)
-            continue
         try:
-            conn.close()
-        except sqlite3.Error as error:
+            if isinstance(conn, ManagedConnection):
+                if logger is not None:
+                    conn.close_logger = logger
+                if not conn.retire():
+                    pending.append(conn)
+            else:
+                conn.close()
+        except Exception:
             if logger is not None:
-                logger.warning(f"Error closing database connection: {error}")
+                logger.exception("Error closing database connection")
     return pending
 
 
 def await_closed(
     connections: Iterable[ManagedConnection], timeout: float, logger=None
 ) -> bool:
-    """Wait until the owners closed ``connections`` (retired); False on timeout."""
+    """Wait until the owners closed ``connections`` (retired); False on timeout.
+
+    A connection still open after ``timeout`` (its owner is blocked, for
+    example on another process's lock after a lost interrupt) is logged and
+    left to close when its call ends; a Move library then reports the old
+    files as still in use instead of renaming them.
+    """
     deadline = time.monotonic() + timeout
     still_open = 0
     for conn in connections:

database/models.py

@@ -1247,7 +1247,11 @@ class PromptModel:
             Path(src_path).resolve().as_uri() + "?mode=ro", uri=True
         )
         try:
-            source.backup(self.get_connection())
+            target = self.get_connection()
+            # The target is written under backup's step too: mark it in use
+            # so a concurrent close_all retires it instead of closing it.
+            with target.using():
+                source.backup(target)
         finally:
             source.close()
 

tests/test_library_move_review.py

@@ -319,6 +319,32 @@ class TestBackgroundWritesSurviveTheSwitch(_MoveCase):
         self.assertEqual(outcome, {"ok": True})
         self.assertEqual(_run_count(self.old_db, text), 1)
 
+    def test_a_retire_failure_after_a_complete_switch_does_not_fail_the_move(self):
+        node = PromptDatabase(self.old_db)
+        node.get_prompt_by_id(1)
+        real = models.PromptModel.retire_detached
+        calls = []
+
+        def failing(connections, logger=None):
+            calls.append(len(connections))
+            real(connections, logger)  # closed, but the call still fails
+            raise RuntimeError("retiring failed")
+
+        with mock.patch.object(
+            models.PromptModel, "retire_detached", staticmethod(failing)
+        ):
+            with self.assertLogs("prompt_manager.database.library_move", "ERROR"):
+                result = self.move()
+        self.assertEqual(len(calls), 1)
+        self.assertGreater(calls[0], 0)
+        self.assertEqual(node.model.db_path, self.new_db)
+        self.assertEqual(result.location.db_path, self.new_db)
+        self.assertTrue(
+            os.path.isfile(
+                os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
+            )
+        )
+
     @staticmethod
     def _racing_run_count(writer_db, text):
         """A run count whose UPDATE marks its connection in use for the probe."""

tests/test_managed_connection.py

@@ -14,6 +14,7 @@ import sys
 import tempfile
 import threading
 import unittest
+import weakref
 from unittest import mock
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -204,6 +205,117 @@ class TestCalls(_Case):
             reader.close()
 
 
+class _CursorOpenedDuringClose(weakref.WeakSet):
+    """The cursor set of a connection being closed; a cursor opens mid-iteration.
+
+    ``hook`` runs after the first cursor is yielded, where a thread switch
+    could let the owner call ``conn.cursor()`` (Python 3.10-3.12 copy a
+    WeakSet one item at a time). It runs in the closing thread itself, so
+    the result does not depend on scheduling.
+    """
+
+    def __init__(self, items, hook):
+        super().__init__(items)
+        self.hook = hook
+
+    def __iter__(self):
+        hooked = False
+        for item in super().__iter__():
+            yield item
+            if not hooked:
+                hooked = True
+                self.hook()
+
+
+class TestReviewRound(_Case):
+    """PR #321 review: I1 (cursor set race), M1 and M2 (unguarded paths)."""
+
+    def test_a_cursor_opened_while_the_connection_closes_cannot_break_retire(self):
+        conn = self.connect()
+        kept = [conn.cursor() for _ in range(3)]
+        refused = []
+
+        def open_cursor():
+            try:
+                kept.append(conn.cursor())
+            except sqlite3.ProgrammingError as error:
+                refused.append(error)
+
+        conn._cursors = _CursorOpenedDuringClose(conn._cursors, open_cursor)
+        self.assertTrue(conn.retire())  # no "Set changed size during iteration"
+        self.assertTrue(conn.wait_closed(0))
+        self.assertEqual(len(refused), 1)
+        self.assertEqual(len(kept), 3)
+
+    def test_a_retired_connection_refuses_new_cursors_and_scripts(self):
+        conn = self.connect()
+        conn._enter()  # its owner is inside a call: retired, not yet closed
+        try:
+            self.assertFalse(conn.retire())
+            for call in (
+                conn.cursor,
+                lambda: conn.executescript("CREATE TABLE late(x);"),
+                lambda: conn.in_transaction,
+                lambda: conn.create_function("f", 0, lambda: 1),
+            ):
+                with self.subTest(call=call):
+                    with self.assertRaises(sqlite3.ProgrammingError):
+                        call()
+        finally:
+            conn._leave()
+        self.assertTrue(conn.wait_closed(0))
+        with contextlib.closing(sqlite3.connect(self.path)) as check:
+            tables = check.execute(
+                "SELECT name FROM sqlite_master WHERE name='late'"
+            ).fetchall()
+        self.assertEqual(tables, [])
+
+    def test_one_connection_failing_to_retire_does_not_skip_the_rest(self):
+        class Broken(ManagedConnection):
+            def retire(self):
+                raise RuntimeError("retire failed")
+
+        broken = sqlite3.connect(self.path, check_same_thread=False, factory=Broken)
+        try:
+            after = self.connect()
+            logger = mock.Mock()
+            pending = retire_connections([broken, after], logger)
+            self.assertEqual(pending, [])
+            self.assertTrue(after.wait_closed(0))
+            logger.exception.assert_called_once()
+        finally:
+            ManagedConnection.retire(broken)
+
+    def test_a_restore_marks_its_backup_target_in_use(self):
+        model = PromptModel(os.path.join(self._tmp.name, "live.db"))
+        try:
+            self._restore_marks_its_target(model)
+        finally:
+            model.close_all()
+
+    def _restore_marks_its_target(self, model):
+        source_path = os.path.join(self._tmp.name, "source.db")
+        PromptModel(source_path).close_all()
+        target = model.get_connection()
+        depths = []
+        connect = sqlite3.connect
+
+        class Source(sqlite3.Connection):
+            def backup(self, into, *args, **kwargs):
+                depths.append(into._depth)
+                return super().backup(into, *args, **kwargs)
+
+        def connect_source(*args, **kwargs):
+            if "factory" in kwargs:
+                return connect(*args, **kwargs)
+            return connect(*args, **kwargs, factory=Source)
+
+        with mock.patch("database.models.sqlite3.connect", connect_source):
+            model._replace_contents_from(source_path)
+        self.assertEqual(depths, [1])
+        self.assertEqual(target._depth, 0)
+
+
 class TestModelConnections(_Case):
     def test_the_model_hands_out_managed_connections_and_closes_them(self):
         model = PromptModel(os.path.join(self._tmp.name, "model.db"))