Vito
9601837aa9
* 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. * test(tagging): wait for in-flight pages before reading the loaded range The scroll-position test read the loaded count right after scrolling the continuation out of view, while a listing page the continuation had already requested could still be in flight. When it landed the count grew (80 -> 120 in CI) and the stability poll failed. The test now waits a frame for any scheduled load to start and then for every listing request to finish before it reads the count. * 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.
459 lines
18 KiB
Python
459 lines
18 KiB
Python
"""Move library, final review round (#249 S6b PR #318).
|
|
|
|
- I1: a switch that fails at any point is undone while both files are still
|
|
locked: the copy is removed, the old library stays in use and a restart
|
|
keeps it. A follow-up after a complete switch only logs.
|
|
- I2: a background write (run counting, the image monitor) that fails
|
|
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.
|
|
"""
|
|
|
|
import contextlib
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import threading
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from database import library_move as lm # noqa: E402
|
|
from database import models # noqa: E402
|
|
from database.operations import PromptDatabase # noqa: E402
|
|
from tests.open_handles import assert_closed # noqa: E402
|
|
from utils import library_location as ll # noqa: E402
|
|
from utils import usage_tracking # noqa: E402
|
|
from utils.hashing import generate_prompt_hash # noqa: E402
|
|
|
|
try:
|
|
from tests.test_library_move import KEY, SimulatedCrash, _MoveCase, _rows
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from test_library_move import KEY, SimulatedCrash, _MoveCase, _rows
|
|
|
|
WAIT = 30
|
|
|
|
|
|
def _run_count(path, text):
|
|
uri = "file:" + path.replace(os.sep, "/") + "?mode=ro"
|
|
with contextlib.closing(sqlite3.connect(uri, uri=True)) as conn:
|
|
row = conn.execute("SELECT run_count FROM prompts WHERE text=?", (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."""
|
|
|
|
def _assert_not_moved(self, caught):
|
|
self.assertEqual(caught.exception.code, "switch-failed")
|
|
self.assertIn("Nothing was moved", caught.exception.message)
|
|
self.assertIsNone(ll.current_override())
|
|
self.assertEqual(self.db.model.db_path, self.old_db)
|
|
# Writes keep landing in the old library, and a restart keeps it.
|
|
self.db.save_prompt("written after the failed switch")
|
|
self.assertIn("written after the failed switch", repr(_rows(self.old_db)))
|
|
self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER)
|
|
self.assertNotEqual(ll.read_marker(self.target).get("state"), "moved")
|
|
self.assert_target_clean()
|
|
self.assertTrue(os.path.isfile(self.config))
|
|
assert_closed(self, self.target)
|
|
|
|
def _failing(self, point):
|
|
real_switch_to = ll.switch_to
|
|
real_repoint = models.PromptModel.repoint_instances
|
|
|
|
def switch(old_db, location, detached):
|
|
if point == "before":
|
|
raise RuntimeError("switch failed before anything")
|
|
if point == "resolver":
|
|
real_switch_to(location)
|
|
raise RuntimeError("switch failed after the resolver")
|
|
if point == "repoint":
|
|
real_switch_to(location)
|
|
detached.extend(real_repoint(old_db, location.db_path))
|
|
raise RuntimeError("switch failed after the models")
|
|
lm.switch_process(old_db, location, detached)
|
|
raise RuntimeError("switch failed at its very end")
|
|
|
|
return switch
|
|
|
|
def test_a_switch_failing_at_any_point_is_undone(self):
|
|
node = PromptDatabase(self.old_db)
|
|
node.get_prompt_by_id(1)
|
|
for point in ("before", "resolver", "repoint", "complete"):
|
|
with self.subTest(point=point):
|
|
with self.assertRaises(lm.MoveFailed) as caught:
|
|
self.move(switch=self._failing(point))
|
|
self._assert_not_moved(caught)
|
|
self.assertEqual(node.model.db_path, self.old_db)
|
|
self.assertEqual(node.get_prompt_by_id(1)["text"], "a cat number 0")
|
|
|
|
def test_a_failed_switch_closes_the_old_connections_it_detached(self):
|
|
node = PromptDatabase(self.old_db)
|
|
for point in ("repoint", "complete"):
|
|
with self.subTest(point=point):
|
|
before = node.model.get_connection()
|
|
with self.assertRaises(lm.MoveFailed):
|
|
self.move(switch=self._failing(point))
|
|
with self.assertRaises(sqlite3.ProgrammingError):
|
|
before.execute("SELECT 1")
|
|
# The model reopens the old library on next use.
|
|
self.assertEqual(node.get_prompt_by_id(1)["text"], "a cat number 0")
|
|
|
|
def test_a_failing_follow_up_after_the_switch_only_logs(self):
|
|
def follow_up():
|
|
raise RuntimeError("the health report could not be dropped")
|
|
|
|
with self.assertLogs("prompt_manager.database.library_move", "ERROR"):
|
|
result = self.move(after_switch=follow_up)
|
|
self.assertIsNone(result.retire_problem)
|
|
self.assertEqual(self.db.model.db_path, self.new_db)
|
|
self.assertTrue(
|
|
os.path.isfile(
|
|
os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
|
)
|
|
)
|
|
|
|
|
|
class _Flaky:
|
|
"""A write that fails once as if Move library switched while it waited."""
|
|
|
|
def __init__(self, db, name):
|
|
self.db, self.name, self.calls = db, name, 0
|
|
self.real = getattr(db, name)
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
self.db.model.close_all() # what the switch does: a new generation
|
|
raise sqlite3.ProgrammingError("Cannot operate on a closed database.")
|
|
return self.real(*args, **kwargs)
|
|
|
|
|
|
class TestBackgroundWritesSurviveTheSwitch(_MoveCase):
|
|
"""I2: run counting and the image monitor write again on a fresh connection."""
|
|
|
|
def test_run_counting_retries_once_after_a_switch(self):
|
|
db = PromptDatabase(self.old_db)
|
|
text = "a counted prompt"
|
|
db.save_prompt(text, prompt_hash=generate_prompt_hash(text))
|
|
flaky = _Flaky(db, "record_prompt_use")
|
|
with mock.patch.object(db, "record_prompt_use", flaky):
|
|
pending = usage_tracking.PendingFirstUse()
|
|
usage_tracking.count_runs([(generate_prompt_hash(text), text)], db, pending)
|
|
self.assertEqual(flaky.calls, 2)
|
|
self.assertEqual(_run_count(self.old_db, text), 1)
|
|
|
|
def test_a_failure_without_a_switch_is_not_retried(self):
|
|
db = PromptDatabase(self.old_db)
|
|
calls = []
|
|
|
|
def broken():
|
|
calls.append(1)
|
|
raise sqlite3.OperationalError("disk I/O error")
|
|
|
|
with self.assertRaises(sqlite3.OperationalError):
|
|
db.survive_library_switch(broken)
|
|
self.assertEqual(calls, [1])
|
|
|
|
def test_an_image_link_racing_the_switch_lands_in_the_new_library(self):
|
|
from utils.image_monitor import ImageGenerationHandler
|
|
|
|
self.db.model.close_all()
|
|
image = os.path.join(self.comfy, "output", "late.png")
|
|
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()
|
|
|
|
def trace(sql):
|
|
if sql.lstrip().upper().startswith(("BEGIN", "INSERT")):
|
|
probe.using(conn)
|
|
started.set()
|
|
|
|
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()
|
|
|
|
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"]
|
|
if "late.png" in repr(row)
|
|
]
|
|
self.assertEqual(len(moved), 1)
|
|
renamed = os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
|
self.assertNotIn("late.png", repr(_rows(renamed).get("generated_images")))
|
|
|
|
def test_an_image_that_was_not_linked_is_not_reported_as_linked(self):
|
|
from utils.image_monitor import ImageGenerationHandler
|
|
|
|
db = PromptDatabase(self.old_db)
|
|
handler = ImageGenerationHandler(db, mock.MagicMock())
|
|
image = os.path.join(self.comfy, "output", "orphan.png")
|
|
with mock.patch.object(handler, "logger") as log:
|
|
handler.link_image_to_prompt(image, {"id": 99999}, {})
|
|
messages = [str(call) for call in log.debug.call_args_list]
|
|
self.assertFalse([m for m in messages if "Successfully linked" in m])
|
|
self.assertNotIn("orphan.png", repr(_rows(self.old_db)["generated_images"]))
|
|
|
|
def test_a_write_racing_the_switch_lands_in_the_new_library(self):
|
|
text = "a racing prompt"
|
|
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":
|
|
# 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))
|
|
|
|
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):
|
|
"""M2: a copy whose path holds ``#`` and ``%`` verifies."""
|
|
|
|
def test_copy_path_with_uri_characters_verifies(self):
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
folder = os.path.join(tmp, "lib #1 %20")
|
|
os.makedirs(folder)
|
|
path = os.path.join(folder, "copy.db")
|
|
db = PromptDatabase(path)
|
|
db.save_prompt("a prompt")
|
|
db.model.close_all()
|
|
with contextlib.closing(sqlite3.connect(path)) as conn:
|
|
conn.execute("PRAGMA journal_mode=DELETE")
|
|
counts, documents = lm._describe(conn)
|
|
lm._verify_copy(path, counts, documents)
|
|
models.PromptModel.close_all_instances()
|
|
|
|
|
|
class TestRestartFinishesTheRename(_MoveCase):
|
|
"""M6: the marker names the old files; a restart renames them."""
|
|
|
|
def _crash_after_switch(self):
|
|
def checkpoint(step):
|
|
if step == "switched":
|
|
raise SimulatedCrash(step)
|
|
|
|
with self.assertRaises(SimulatedCrash):
|
|
self.move(checkpoint=checkpoint)
|
|
models.PromptModel.close_all_instances()
|
|
ll.clear_override()
|
|
|
|
def test_a_restart_renames_the_old_files_never_deletes(self):
|
|
self._crash_after_switch()
|
|
self.assertTrue(os.path.isfile(self.old_db))
|
|
restarted = self.location()
|
|
self.assertEqual(restarted.db_path, self.new_db)
|
|
marker = ll.read_marker(self.target)
|
|
self.assertEqual(
|
|
(marker["from"], marker["config_from"]), (self.old_db, self.config)
|
|
)
|
|
renamed = lm.finish_interrupted_move(restarted)
|
|
self.assertEqual(
|
|
sorted(renamed),
|
|
[
|
|
"config.json.moved-to-__promptmanager",
|
|
"prompts.db.moved-to-__promptmanager",
|
|
],
|
|
)
|
|
self.assertFalse(os.path.exists(self.old_db))
|
|
moved_config = os.path.join(self.ext, "config.json.moved-to-__promptmanager")
|
|
with open(moved_config, encoding="utf-8") as handle:
|
|
self.assertEqual(
|
|
json.load(handle)["integrations"]["lora_manager"]["civitai_api_key"],
|
|
KEY,
|
|
)
|
|
self.assertEqual(lm.finish_interrupted_move(self.location()), [])
|
|
|
|
def test_nothing_is_renamed_while_the_old_file_is_in_use(self):
|
|
self._crash_after_switch()
|
|
reader = sqlite3.connect(self.old_db)
|
|
try:
|
|
reader.execute("BEGIN")
|
|
reader.execute("SELECT COUNT(*) FROM prompts").fetchone()
|
|
PromptDatabase(self.old_db).save_prompt("in the wal")
|
|
self.assertEqual(lm.finish_interrupted_move(self.location()), [])
|
|
self.assertTrue(os.path.isfile(self.old_db))
|
|
finally:
|
|
reader.rollback()
|
|
reader.close()
|
|
models.PromptModel.close_all_instances()
|
|
|
|
def test_a_node_folder_library_has_nothing_to_finish(self):
|
|
self.assertEqual(lm.finish_interrupted_move(self.location()), [])
|
|
self.assertTrue(os.path.isfile(self.old_db))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|