From e1e3bbc6913132b8bf9c07517e4cb7bde3442719 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 24 Oct 2026 18:42:14 -0700 Subject: [PATCH] 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 | 30 +++-- database/managed_connection.py | 175 +++++++++++++++++++++++------- database/models.py | 6 +- tests/test_library_move_review.py | 26 +++++ tests/test_managed_connection.py | 112 +++++++++++++++++++ 5 files changed, 299 insertions(+), 50 deletions(-) diff --git a/database/library_move.py b/database/library_move.py index 269a869..1cec52e 100644 --- a/database/library_move.py +++ b/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") diff --git a/database/managed_connection.py b/database/managed_connection.py index 4eda698..fb71fc6 100644 --- a/database/managed_connection.py +++ b/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: diff --git a/database/models.py b/database/models.py index 443d4ac..78c0fc9 100644 --- a/database/models.py +++ b/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() diff --git a/tests/test_library_move_review.py b/tests/test_library_move_review.py index 4cfaa0f..14e5d52 100644 --- a/tests/test_library_move_review.py +++ b/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.""" diff --git a/tests/test_managed_connection.py b/tests/test_managed_connection.py index 8303623..b33d437 100644 --- a/tests/test_managed_connection.py +++ b/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"))