fix(system): never close another thread's library connection during a switch
Move library closed every thread's old connection from the mover thread, including one whose owner was inside a statement waiting for the old file's write lock. Closing a sqlite3 connection under a running sqlite3_step is a use-after-free: Windows CI died with "access violation" in record_prompt_use, and on Linux the writer got a MemoryError once its busy timeout ran out. Managed connections now count the calls their owner is inside. Retiring one from another thread closes it at once only when idle; otherwise it is interrupted (sqlite3_interrupt is thread-safe) and its owner closes it as the call ends. A retired connection refuses new work and never commits, so a write that waited for the lock is rolled back and survive_library_switch runs it again on the new library. The mover retires connections while it holds the locks and waits for the owners to close them only after releasing them, before the old files are renamed or a failed copy is removed. Cursors are closed before the connection so Python 3.10/3.11 release the file handle too.
This commit is contained in:
@@ -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,19 @@ 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
|
||||
move.retiring.extend(
|
||||
PromptModel.retire_detached(
|
||||
move.unswitch(plan.db_path, location.db_path), logger
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Library move: undoing the switch failed")
|
||||
# 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)
|
||||
move.retiring.extend(PromptModel.retire_detached(detached, logger))
|
||||
raise MoveFailed("switch-failed") from error
|
||||
move.checkpoint("switched")
|
||||
return detached
|
||||
@@ -737,6 +745,20 @@ 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:
|
||||
PromptModel.await_detached(move.retiring, logger)
|
||||
|
||||
|
||||
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 +778,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.
|
||||
move.retiring.extend(PromptModel.retire_detached(detached, logger))
|
||||
finally:
|
||||
if source.in_transaction:
|
||||
source.execute("ROLLBACK")
|
||||
|
||||
@@ -0,0 +1,260 @@
|
||||
"""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: 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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from typing import Iterable, List, Optional
|
||||
|
||||
import sqlite3
|
||||
|
||||
CLOSED_MESSAGE = "Cannot operate on a closed database."
|
||||
|
||||
|
||||
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 __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)
|
||||
self._use_lock = threading.Lock()
|
||||
self._depth = 0
|
||||
self._retired = 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()
|
||||
|
||||
def cursor(self, factory=None):
|
||||
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 __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 (any thread)."""
|
||||
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()
|
||||
return False
|
||||
elif self._depth:
|
||||
return False
|
||||
else:
|
||||
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 ones: its cursors are idle.
|
||||
for cursor in list(self._cursors):
|
||||
with contextlib.suppress(sqlite3.Error):
|
||||
sqlite3.Cursor.close(cursor)
|
||||
try:
|
||||
super().close()
|
||||
except sqlite3.Error as error:
|
||||
if self.close_logger is not None:
|
||||
self.close_logger.warning(f"Error closing database connection: {error}")
|
||||
finally:
|
||||
self._released.set()
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
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 logger is not None:
|
||||
logger.warning(f"Error closing database connection: {error}")
|
||||
return pending
|
||||
|
||||
|
||||
def await_closed(
|
||||
connections: Iterable[ManagedConnection], timeout: float, logger=None
|
||||
) -> bool:
|
||||
"""Wait until the owners closed ``connections`` (retired); False on timeout."""
|
||||
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
|
||||
+42
-18
@@ -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:
|
||||
"""
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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,87 @@ 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)
|
||||
|
||||
@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):
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
"""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
|
||||
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 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()
|
||||
Reference in New Issue
Block a user