Merge remote-tracking branch 'origin/main' into feat/249-system-s7-maintenance-ui

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);

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):

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):

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()