Vito
b1e7abd380
* fix(threads): stop and join the prompt cleanup thread, the image worker and the observer (#302, #370) - PromptTracker: the cleanup thread waits on a stop event instead of time.sleep; new shutdown(timeout=2) joins it and returns whether it ended. - ImageGenerationHandler.shutdown(timeout=5) now joins the worker and returns a bool; new close() is final: schedule() refuses and a sealed run's history re-check is dropped, so no worker restarts afterwards. - ImageMonitor.stop_monitoring(timeout=5) closes the handler first, then stops and joins the watchdog observer within one deadline. - Every test that builds a real tracker, handler or monitor ends with its final stop and an assertion, before its folder and databases go. - #370: the wrapper/session history test drains before dispose(); the held-write bound test uses 30 s ceilings instead of tight wall clocks. * test(history): drain the writer before asserting a bounded close it does not test (#370) * fix(threads): a later stop waits again for a worker or observer an earlier stop left running (#302) Review findings on #372: - ImageGenerationHandler keeps the workers a stop signalled that had not ended; every later shutdown/close joins them again, so a second call no longer returns True while the first worker still runs. - ImageMonitor.stop_monitoring keeps an observer or handler it left running and stops and joins it again on the next call. - A closed handler takes no awaiting history path, and a re-check that raced close() is not kept due. - Tests: second stops of a busy worker and a live observer, the closed handler's re-check paths, the handle test asserts the worker ended and stops before it asserts, and a no-new-thread check instead of set equality. * test(browser): the System reset restores every settings value, safety_backup included The frontend coverage shard of #372 failed in workspace-system-tagging: the review read "No safety backup (the policy is off)". The backups spec saves safety_backup=false, and /__test__/system/reset restored config.json but not PromptManagerConfig.SAFETY_BACKUP (nor INFINITE_SCROLL), so a later spec on the same worker server saw the policy off. The reset now captures and restores every non-tagging value POST /settings can change.
462 lines
18 KiB
Python
462 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.background_stops import close_handler # 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())
|
|
self.addCleanup(close_handler, self, handler)
|
|
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())
|
|
self.addCleanup(close_handler, self, handler)
|
|
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()
|