diff --git a/database/library_move.py b/database/library_move.py index 896a4bb2..1cec52e6 100644 --- a/database/library_move.py +++ b/database/library_move.py @@ -438,8 +438,8 @@ def switch_process(old_db: str, location, detached: list) -> None: """Point this process at ``location``: resolver, every model, observers. Adds the old file's connections, detached and still open, to - ``detached`` as soon as they are detached, so the mover closes them - (``PromptModel.close_detached``) whether or not the rest succeeds. + ``detached`` as soon as they are detached, so the mover retires them + (``PromptModel.retire_detached``) whether or not the rest succeeds. """ from .execution_history import rekey_history_restore from .review_changes import rekey_review_changes @@ -474,6 +474,8 @@ class _Move: self.published: List[str] = [] self.marker: Optional[str] = None self.moved = False + # Retired connections still in use; awaited once the locks are free. + self.retiring: list = [] def temp(self, suffix: str) -> str: path = os.path.join( @@ -704,7 +706,9 @@ def _switch_guarded(move: _Move, location) -> list: The copy's write lock keeps every repointed thread from writing to it until the switch is whole, so undoing a partial switch loses nothing. - Returns the old file's detached connections (still open). + Returns the old file's detached connections (still open). On failure + every detached connection is retired; those still in use are added to + ``move.retiring`` for the caller to await once the locks are released. """ plan = move.plan guard = sqlite3.connect( @@ -718,15 +722,17 @@ def _switch_guarded(move: _Move, location) -> list: except Exception as error: logger.exception("Library move: switching failed; undoing it") move.moved = False # the copy and its marker are removed + # Retire, never wait, here: a connection still in use may be + # waiting on the very locks this thread holds. try: - PromptModel.close_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. - PromptModel.close_detached(detached, logger) + _retire(move, detached) raise MoveFailed("switch-failed") from error move.checkpoint("switched") return detached @@ -737,6 +743,36 @@ def _switch_guarded(move: _Move, location) -> list: def _copy_and_switch(move: _Move, location, attempts: int): + """Copy, then switch under both locks; retired connections are awaited. + + A connection that was inside a statement when it was retired is closed + by its own thread once the statement ends, which may need the locks + released first; the old files are renamed (and a failed copy removed) + only after that, since Windows refuses to rename or delete an open file. + """ + try: + return _copy_and_switch_locked(move, location, attempts) + finally: + 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): plan = move.plan for attempt in range(1, attempts + 1): source = sqlite3.connect( @@ -756,9 +792,10 @@ def _copy_and_switch(move: _Move, location, attempts: int): files = move.publish(temp_db, copies) detached = _switch_guarded(move, location) # Still holding the old file's lock: a write that was waiting - # on it fails when its connection closes (and runs again on - # the new library) instead of landing in the old file. - PromptModel.close_detached(detached, logger) + # 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. + _retire(move, detached) finally: if source.in_transaction: source.execute("ROLLBACK") diff --git a/database/managed_connection.py b/database/managed_connection.py new file mode 100644 index 00000000..fb71fc6d --- /dev/null +++ b/database/managed_connection.py @@ -0,0 +1,353 @@ +"""A SQLite connection that is never closed under the thread using it. + +``PromptModel`` gives each thread its own connection, but Move library and +``close_all`` must retire every thread's connection from one thread. Closing +a ``sqlite3.Connection`` while its owner is inside a statement (Python runs +``sqlite3_step`` without the GIL) frees memory the owner then reads: on +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 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 + +import contextlib +import logging +import threading +import time +import weakref +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): + """A cursor whose calls into SQLite count as its connection being in use.""" + + def execute(self, sql, parameters=(), /): + conn = self.connection + conn._enter() + try: + return super().execute(sql, parameters) + finally: + conn._leave() + + def executemany(self, sql, seq_of_parameters, /): + conn = self.connection + conn._enter() + try: + return super().executemany(sql, seq_of_parameters) + finally: + conn._leave() + + def executescript(self, sql_script, /): + conn = self.connection + conn._enter() + try: + return super().executescript(sql_script) + finally: + conn._leave() + + def fetchone(self): + conn = self.connection + conn._enter() + try: + return super().fetchone() + finally: + conn._leave() + + def fetchmany(self, *args, **kwargs): + conn = self.connection + conn._enter() + try: + return super().fetchmany(*args, **kwargs) + finally: + conn._leave() + + def fetchall(self): + conn = self.connection + conn._enter() + try: + return super().fetchall() + 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() + try: + return super().__next__() + finally: + conn._leave() + + def close(self): + conn = self.connection + try: + conn._enter() + except sqlite3.ProgrammingError: + return super().close() # the connection is closed or closing + try: + return super().close() + finally: + conn._leave() + + +class ManagedConnection(sqlite3.Connection): + """A connection only its owner closes while a call into it is running.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # 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). + self._cursors: "weakref.WeakSet[sqlite3.Cursor]" = weakref.WeakSet() + self.close_logger: Optional[logging.Logger] = None + + # ── in use ─────────────────────────────────────────────────────── + + def _enter(self) -> None: + with self._use_lock: + if self._retired: + raise sqlite3.ProgrammingError(CLOSED_MESSAGE) + self._depth += 1 + + def _leave(self) -> None: + with self._use_lock: + self._depth -= 1 + close_now = self._retired and self._depth == 0 + 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): + 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=(), /): + return self.cursor().execute(sql, parameters) + + def executemany(self, sql, seq_of_parameters, /): + return self.cursor().executemany(sql, seq_of_parameters) + + def executescript(self, sql_script, /): + return self.cursor().executescript(sql_script) + + def commit(self): + self._enter() + try: + return super().commit() + finally: + self._leave() + + def rollback(self): + self._enter() # a closed connection raises, as sqlite3's does + try: + return super().rollback() + finally: + self._leave() + + def backup(self, target, **kwargs): + self._enter() + try: + return super().backup(target, **kwargs) + 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: + try: + self.commit() + except BaseException: + with contextlib.suppress(sqlite3.Error): + self.rollback() + raise + else: + self.rollback() + return False + + # ── closing ────────────────────────────────────────────────────── + + def close(self) -> None: + """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. + + 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 + 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 + + def wait_closed(self, timeout: float) -> bool: + return self._released.wait(timeout) + + def _close_now(self) -> None: + # Only ever called with no call into the connection running, and a + # 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 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 +) -> List[ManagedConnection]: + """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: + try: + 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.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. + + 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: + if not conn.wait_closed(max(0.0, deadline - time.monotonic())): + still_open += 1 + if still_open and logger is not None: + logger.warning( + f"{still_open} database connection(s) still in use after " + f"{timeout:g}s; they close when their statement ends" + ) + return not still_open diff --git a/database/models.py b/database/models.py index 0c0210e3..78c0fc99 100644 --- a/database/models.py +++ b/database/models.py @@ -15,6 +15,11 @@ from contextlib import closing from pathlib import Path from typing import Optional, Set, Tuple +from .managed_connection import ( + ManagedConnection, + await_closed, + retire_connections, +) from .library_schema import library_schema_error, migrate_library_state from .image_favorites import image_favorites_schema_error, migrate_image_favorites from .annotation_schema import annotation_schema_error, migrate_annotation_state @@ -615,10 +620,15 @@ class PromptModel: # meanwhile: otherwise it would write to the old file. with self._conn_lock: path, generation = self.db_path, self._generation - # check_same_thread=False only so close_all can clean up other - # threads' connections; each connection is still used by one thread. + # check_same_thread=False only so close_all and Move library can + # retire other threads' connections; each connection is still + # used by one thread, and ManagedConnection never closes it under + # a running statement (see managed_connection). conn = sqlite3.connect( - path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS + path, + check_same_thread=False, + timeout=BUSY_TIMEOUT_SECONDS, + factory=ManagedConnection, ) try: conn.row_factory = sqlite3.Row @@ -654,9 +664,9 @@ class PromptModel: Every thread opens the new file on its next ``get_connection``. The old connections are detached, not closed, and returned: the caller - closes them once nothing can block on the new file any more. Closing - a connection another thread is using waits for its statement; a write - waiting on the mover's lock then fails (and, through + retires them (``retire_detached``) while it still holds the old + file's write lock. A write waiting on that lock is then interrupted + and can never commit, so it fails (and, through ``PromptDatabase.survive_library_switch``, runs again on the new file). """ old_key = os.path.normcase(os.path.realpath(old_path)) @@ -671,14 +681,31 @@ class PromptModel: return detached @staticmethod - def close_detached(connections, logger=None) -> None: - """Close connections ``repoint_instances`` detached.""" - for conn in connections: - try: - conn.close() - except sqlite3.Error as e: - if logger is not None: - logger.warning(f"Error closing database connection: {e}") + def retire_detached(connections, logger=None) -> list: + """Retire detached connections without waiting for their owners. + + Idle ones close now. One whose owner thread is inside a statement is + never closed from here (that frees memory the owner is using: an + access violation on Windows); it is interrupted, refuses to commit and + is closed by its owner as the statement ends. Returns those still + open, for ``await_detached`` once no lock the owners wait on is held. + """ + return retire_connections(connections, logger) + + @staticmethod + def await_detached(connections, logger=None) -> bool: + """Wait until retired connections are closed; False if some are not.""" + return await_closed(connections, BUSY_TIMEOUT_SECONDS, logger) + + @staticmethod + def close_detached(connections, logger=None) -> bool: + """Retire connections ``repoint_instances`` detached and await them. + + Do not call it while holding a lock their owners may wait on; retire + them there (``retire_detached``) and await them after releasing it. + """ + pending = retire_connections(connections, logger) + return await_closed(pending, BUSY_TIMEOUT_SECONDS, logger) def close_all(self) -> None: """Close every thread's connection; they reopen lazily on next use. @@ -714,10 +741,7 @@ class PromptModel: self._local.conn = None with self._conn_lock: self._connections.discard(conn) - try: - conn.close() - except sqlite3.Error as e: - self.logger.warning(f"Error closing database connection: {e}") + self.close_detached([conn], self.logger) def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None: """ @@ -1223,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/browser/workspace-tagging.spec.mjs b/tests/browser/workspace-tagging.spec.mjs index 15de079f..525e2ef4 100644 --- a/tests/browser/workspace-tagging.spec.mjs +++ b/tests/browser/workspace-tagging.spec.mjs @@ -602,6 +602,10 @@ test('a save under a Tags filter re-reads the listing: the selection stays while test('a save under a search keeps the scroll position and the loaded range', async ({ page, request }) => { await expect(await request.post('/__test__/tagging/reset', { data: { extra_images: 140 } })).toBeOK(); + // Listing pages still in flight: one that lands after the scroll below would grow the loaded range. + const inFlight = new Set(); + page.on('request', item => { if (item.url().includes('/images/output?')) inFlight.add(item); }); + for (const done of ['requestfinished', 'requestfailed']) page.on(done, item => inFlight.delete(item)); await page.goto(images); const searched = listingOf(page, query => query.get('q') === 'bulk'); await page.getByRole('searchbox', { name: 'Search workspace', exact: true }).fill('bulk'); @@ -615,6 +619,9 @@ test('a save under a search keeps the scroll position and the loaded range', asy await expect.poll(async () => { await continuation.scrollIntoViewIfNeeded(); return loadedCount(); }).toBeGreaterThanOrEqual(80); await page.evaluate(() => window.scrollTo(0, Math.round(document.documentElement.scrollHeight / 3))); await expect(continuation).not.toBeInViewport(); + // A load the continuation scheduled for the next frame has started by now; wait for every page to land. + await page.evaluate(() => new Promise(done => requestAnimationFrame(() => requestAnimationFrame(done)))); + await expect.poll(() => inFlight.size).toBe(0); const loaded = await continuation.textContent(); await expect.poll(() => continuation.textContent()).toBe(loaded); const y = await page.evaluate(() => window.scrollY); diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 7eddca75..e4aa9a88 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -293,6 +293,8 @@ class TestRestoreFromFile(BackupTestCase): super().backup(target, *args, **kwargs) def connect_with_failure(*args, **kwargs) -> sqlite3.Connection: + if "factory" in kwargs: # the model's own managed connections + return connect(*args, **kwargs) return connect(*args, **kwargs, factory=FailingLiveBackup) with patch("database.models.sqlite3.connect", side_effect=connect_with_failure): diff --git a/tests/test_library_move_review.py b/tests/test_library_move_review.py index 1e4fa823..14e5d521 100644 --- a/tests/test_library_move_review.py +++ b/tests/test_library_move_review.py @@ -7,6 +7,9 @@ because Move library switched runs once more, on the new library; a write racing the switch on a real connection (a run count, an image link) lands in the new library only. +- The switching thread never closes a connection another thread is inside + a statement on (a use-after-free in SQLite: "Windows fatal exception: + access violation"); it interrupts it and the owner closes it. - M2: the copy is verified through a file URI that survives ``#`` and ``%``. - M6: a restart after a crash between the switch and the rename renames the old files, never deletes them. @@ -46,6 +49,39 @@ def _run_count(path, text): return row.fetchone()[0] +class _CrossThreadCloseProbe: + """Record every ``close()`` this thread calls on a connection in use elsewhere. + + The writer marks its connection with ``using`` when its statement starts. + Closing that connection from the switching thread while the writer is + inside the statement frees memory SQLite is still using; on Windows the + process dies with an access violation, on Linux the writer sees garbage. + """ + + def __init__(self): + self.in_use = set() + self.violations = [] + + def using(self, conn): + self.in_use.add(conn) + + def _profile(self, _frame, event, arg): + if event != "c_call" or getattr(arg, "__name__", "") != "close": + return + target = getattr(arg, "__self__", None) + if isinstance(target, sqlite3.Connection) and target in self.in_use: + self.violations.append(threading.current_thread().name) + + @contextlib.contextmanager + def watching(self): + previous = sys.getprofile() + sys.setprofile(self._profile) + try: + yield self + finally: + sys.setprofile(previous) + + class TestFailedSwitchIsUndone(_MoveCase): """I1: whatever part of the switch fails, nothing was moved.""" @@ -166,44 +202,46 @@ class TestBackgroundWritesSurviveTheSwitch(_MoveCase): self.db.model.close_all() image = os.path.join(self.comfy, "output", "late.png") - # A short busy timeout so the blocked writer gives up in seconds. - with mock.patch.object(models, "BUSY_TIMEOUT_SECONDS", 2): - writer_db = PromptDatabase(self.old_db) - handler = ImageGenerationHandler(writer_db, mock.MagicMock()) - started = threading.Event() - outcome = {} + writer_db = PromptDatabase(self.old_db) + handler = ImageGenerationHandler(writer_db, mock.MagicMock()) + started = threading.Event() + probe = _CrossThreadCloseProbe() + outcome = {} - def write(): - conn = writer_db.model.get_connection() - conn.set_trace_callback( - lambda sql: ( - started.set() - if sql.lstrip().upper().startswith(("BEGIN", "INSERT")) - else None - ) - ) - try: - # No metadata: the fallback link, not the resolved pair. - handler.link_image_to_prompt(image, {"id": 1}, {}) - outcome["ok"] = True - except Exception as error: # pragma: no cover - reported below - outcome["error"] = error - finally: - writer_db.model.close() + def write(): + conn = writer_db.model.get_connection() - thread = threading.Thread(target=write) + def trace(sql): + if sql.lstrip().upper().startswith(("BEGIN", "INSERT")): + probe.using(conn) + started.set() - def checkpoint(step): - if step == "locked": - # The old file is write-locked now: start the writer and - # wait until its first write is running (it then waits on us). - thread.start() - self.assertTrue(started.wait(WAIT)) + conn.set_trace_callback(trace) + try: + # No metadata: the fallback link, not the resolved pair. + handler.link_image_to_prompt(image, {"id": 1}, {}) + outcome["ok"] = True + except Exception as error: # pragma: no cover - reported below + outcome["error"] = error + finally: + writer_db.model.close() - self.move(checkpoint=checkpoint) - thread.join(WAIT) + thread = threading.Thread(target=write) + + def checkpoint(step): + if step == "locked": + # The old file is write-locked now: start the writer and + # wait until its first write is running (it then waits on us). + thread.start() + self.assertTrue(started.wait(WAIT)) + + with probe.watching(): + result = self.move(checkpoint=checkpoint) + thread.join(WAIT) self.assertFalse(thread.is_alive()) + self.assertEqual(probe.violations, []) self.assertEqual(outcome, {"ok": True}) + self.assertIsNone(result.retire_problem) moved = [ row for row in _rows(self.new_db)["generated_images"] @@ -229,45 +267,113 @@ class TestBackgroundWritesSurviveTheSwitch(_MoveCase): text = "a racing prompt" self.db.save_prompt(text, prompt_hash=generate_prompt_hash(text)) self.db.model.close_all() - # A short busy timeout so the blocked writer gives up in seconds. - with mock.patch.object(models, "BUSY_TIMEOUT_SECONDS", 2): - writer_db = PromptDatabase(self.old_db) - started = threading.Event() - outcome = {} + writer_db = PromptDatabase(self.old_db) + thread, started, probe, outcome = self._racing_run_count(writer_db, text) - def write(): - conn = writer_db.model.get_connection() - conn.set_trace_callback( - lambda sql: started.set() if sql.startswith("UPDATE") else None - ) - try: - usage_tracking.count_runs( - [(generate_prompt_hash(text), text)], - writer_db, - usage_tracking.PendingFirstUse(), - ) - outcome["ok"] = True - except Exception as error: # pragma: no cover - reported below - outcome["error"] = error - finally: - writer_db.model.close() + def checkpoint(step): + if step == "locked": + # The old file is write-locked now: start the writer and + # wait until its UPDATE is running (it then waits on us). + thread.start() + self.assertTrue(started.wait(WAIT)) - thread = threading.Thread(target=write) - - def checkpoint(step): - if step == "locked": - # The old file is write-locked now: start the writer and - # wait until its UPDATE is running (it then waits on us). - thread.start() - self.assertTrue(started.wait(WAIT)) - - self.move(checkpoint=checkpoint) - thread.join(WAIT) + with probe.watching(): + result = self.move(checkpoint=checkpoint) + # The writer's old connection was closed before the old file was + # renamed: Windows refuses to rename an open file. + self.assertIsNone(result.retire_problem) + thread.join(WAIT) self.assertFalse(thread.is_alive()) + self.assertEqual(probe.violations, []) self.assertEqual(outcome, {"ok": True}) self.assertEqual(_run_count(self.new_db, text), 1) renamed = os.path.join(self.ext, "prompts.db.moved-to-__promptmanager") self.assertEqual(_run_count(renamed, text), 0) + assert_closed(self, renamed) + + def test_a_write_racing_a_failed_switch_lands_in_the_old_library(self): + text = "a prompt racing a failed switch" + self.db.save_prompt(text, prompt_hash=generate_prompt_hash(text)) + self.db.model.close_all() + writer_db = PromptDatabase(self.old_db) + thread, started, probe, outcome = self._racing_run_count(writer_db, text) + + def checkpoint(step): + if step == "locked": + thread.start() + self.assertTrue(started.wait(WAIT)) + + def switch(old_db, location, detached): + lm.switch_process(old_db, location, detached) + raise RuntimeError("switch failed at its very end") + + with probe.watching(): + with self.assertRaises(lm.MoveFailed): + self.move(checkpoint=checkpoint, switch=switch) + # The copy was removed: no connection was left open on it. + self.assert_target_clean() + assert_closed(self, self.target) + thread.join(WAIT) + self.assertFalse(thread.is_alive()) + self.assertEqual(probe.violations, []) + 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.""" + started = threading.Event() + probe = _CrossThreadCloseProbe() + outcome = {} + + def write(): + conn = writer_db.model.get_connection() + + def trace(sql): + if sql.startswith("UPDATE"): + probe.using(conn) + started.set() + + conn.set_trace_callback(trace) + try: + usage_tracking.count_runs( + [(generate_prompt_hash(text), text)], + writer_db, + usage_tracking.PendingFirstUse(), + ) + outcome["ok"] = True + except Exception as error: # pragma: no cover - reported below + outcome["error"] = error + finally: + writer_db.model.close() + + return threading.Thread(target=write), started, probe, outcome class TestVerifyUri(unittest.TestCase): diff --git a/tests/test_managed_connection.py b/tests/test_managed_connection.py new file mode 100644 index 00000000..b33d437a --- /dev/null +++ b/tests/test_managed_connection.py @@ -0,0 +1,332 @@ +"""ManagedConnection: never closed under the thread using it. + +Closing a sqlite3 connection from one thread while another is inside a +statement on it is a use-after-free in SQLite ("Windows fatal exception: +access violation" in CI). A retired connection is closed at once when idle, +otherwise interrupted and closed by its owner as the call ends; it never +commits again. +""" + +import contextlib +import os +import sqlite3 +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__)))) + +from database.managed_connection import ( # noqa: E402 + ManagedConnection, + ManagedCursor, + await_closed, + retire_connections, +) +from database.models import PromptModel # noqa: E402 +from tests.open_handles import assert_closed # noqa: E402 + +WAIT = 30 + + +class _Case(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.path = os.path.join(self._tmp.name, "managed.db") + with contextlib.closing(sqlite3.connect(self.path)) as conn: + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("CREATE TABLE t(x INTEGER)") + conn.execute("INSERT INTO t VALUES (0)") + conn.commit() + self.opened = [] + + def tearDown(self): + for conn in self.opened: + conn.close() + try: + assert_closed(self, self._tmp.name) + finally: + self._tmp.cleanup() + + def connect(self, **kwargs): + conn = sqlite3.connect( + self.path, + check_same_thread=False, + timeout=WAIT, + factory=ManagedConnection, + **kwargs, + ) + self.opened.append(conn) + return conn + + def value(self): + with contextlib.closing(sqlite3.connect(self.path)) as conn: + return conn.execute("SELECT x FROM t").fetchone()[0] + + +class TestRetire(_Case): + def test_an_idle_connection_retired_from_another_thread_closes_at_once(self): + conn = self.connect() + conn.execute("SELECT 1").fetchone() + closed = [] + thread = threading.Thread(target=lambda: closed.append(conn.retire())) + thread.start() + thread.join(WAIT) + self.assertEqual(closed, [True]) + self.assertTrue(conn.wait_closed(0)) + with self.assertRaises(sqlite3.ProgrammingError): + conn.execute("SELECT 1") + self.assertTrue(conn.retire()) # retiring again is harmless + + def test_a_connection_in_use_is_closed_by_its_owner_when_the_call_ends(self): + conn = self.connect() + holder = sqlite3.connect(self.path, isolation_level=None) + self.opened.append(holder) + holder.execute("BEGIN IMMEDIATE") + started = threading.Event() + outcome = {} + conn.set_trace_callback( + lambda sql: started.set() if sql.startswith("UPDATE") else None + ) + + def write(): + try: + with conn: + conn.execute("UPDATE t SET x = 1") + outcome["ok"] = True + except sqlite3.Error as error: + outcome["error"] = error + + thread = threading.Thread(target=write) + thread.start() + self.assertTrue(started.wait(WAIT)) + # The owner is waiting on our lock inside its UPDATE: not closed here. + self.assertFalse(conn.retire()) + self.assertFalse(conn.retire()) + self.assertFalse(conn.wait_closed(0)) + self.assertEqual(retire_connections([conn]), [conn]) + holder.execute("ROLLBACK") + thread.join(WAIT) + self.assertFalse(thread.is_alive()) + self.assertTrue(conn.wait_closed(WAIT)) + self.assertIsInstance(outcome.get("error"), sqlite3.Error) + self.assertEqual(self.value(), 0) + + def test_a_write_retired_before_its_commit_is_rolled_back(self): + conn = self.connect() + conn.execute("UPDATE t SET x = 2") # an implicit transaction is open + thread = threading.Thread(target=conn.retire) + thread.start() + thread.join(WAIT) + with self.assertRaises(sqlite3.ProgrammingError): + conn.commit() + with self.assertRaises(sqlite3.ProgrammingError): + with conn: + pass + with self.assertRaises(sqlite3.ProgrammingError): + conn.rollback() + self.assertEqual(self.value(), 0) + + def test_close_is_retire(self): + conn = self.connect() + conn.close() + self.assertTrue(conn.wait_closed(0)) + with self.assertRaises(sqlite3.ProgrammingError): + conn.commit() + + def test_plain_connections_are_closed_directly(self): + plain = sqlite3.connect(self.path) + self.assertEqual(retire_connections([plain]), []) + with self.assertRaises(sqlite3.ProgrammingError): + plain.execute("SELECT 1") + + def test_awaiting_a_connection_still_in_use_times_out_with_a_warning(self): + conn = self.connect() + conn._enter() # its owner is inside a call + self.assertEqual(retire_connections([conn]), [conn]) + logger = mock.Mock() + self.assertFalse(await_closed([conn], 0, logger)) + logger.warning.assert_called_once() + conn._leave() + self.assertTrue(await_closed([conn], WAIT, logger)) + + +class TestCalls(_Case): + def test_every_cursor_call_works_and_is_refused_once_retired(self): + conn = self.connect() + conn.row_factory = sqlite3.Row + cursor = conn.cursor() + self.assertIsInstance(cursor, ManagedCursor) + conn.executemany("INSERT INTO t VALUES (?)", [(1,), (2,)]) + conn.executescript("INSERT INTO t VALUES (3);") + self.assertEqual( + [row["x"] for row in conn.execute("SELECT x FROM t ORDER BY x")], + [0, 1, 2, 3], + ) + cursor.execute("SELECT x FROM t ORDER BY x") + self.assertEqual(cursor.fetchone()["x"], 0) + self.assertEqual([r["x"] for r in cursor.fetchmany(2)], [1, 2]) + self.assertEqual([r["x"] for r in cursor.fetchall()], [3]) + cursor.executemany("UPDATE t SET x = x WHERE x = ?", [(1,)]) + cursor.executescript("SELECT 1;") + cursor.close() + conn.commit() + conn.rollback() + with conn: + conn.execute("DELETE FROM t WHERE x = 3") + with self.assertRaises(ZeroDivisionError): + with conn: + conn.execute("DELETE FROM t") + raise ZeroDivisionError + self.assertEqual(self.value(), 0) + copy = sqlite3.connect(":memory:") + self.opened.append(copy) + conn.backup(copy) + self.assertEqual(copy.execute("SELECT COUNT(*) FROM t").fetchone()[0], 3) + + reader = conn.execute("SELECT x FROM t") + conn.retire() + for call in ( + lambda: reader.fetchone(), + lambda: reader.fetchmany(1), + lambda: reader.fetchall(), + lambda: next(reader), + lambda: reader.execute("SELECT 1"), + lambda: reader.executemany("SELECT ?", [(1,)]), + lambda: reader.executescript("SELECT 1;"), + lambda: conn.backup(copy), + ): + with self.assertRaises(sqlite3.ProgrammingError): + call() + # As with sqlite3's own cursors (the result varies by Python version). + with contextlib.suppress(sqlite3.ProgrammingError): + 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")) + conn = model.get_connection() + self.assertIsInstance(conn, ManagedConnection) + model.close() + self.assertTrue(conn.wait_closed(0)) + conn = model.get_connection() + model.close_all() + self.assertTrue(conn.wait_closed(0)) + + +if __name__ == "__main__": + unittest.main()