"""Concurrent access to PromptDatabase from several threads. ComfyUI touches the database from the server loop, executor threads, the watchdog image monitor and the queue hook at the same time. Every thread must get its own connection so one thread's commit or rollback never lands on another thread's half-done write. """ import os import sqlite3 import sys import tempfile import threading import unittest from unittest.mock import patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from database import models from database.operations import PromptDatabase from utils.hashing import generate_prompt_hash THREADS = 8 CALLS_PER_THREAD = 50 class ConcurrencyTestCase(unittest.TestCase): def setUp(self): fd, self.path = tempfile.mkstemp(suffix=".db") os.close(fd) self.db = PromptDatabase(self.path) self.addCleanup(self._cleanup) def _cleanup(self): self.db.close_all() for suffix in ("", "-wal", "-shm"): if os.path.exists(self.path + suffix): os.unlink(self.path + suffix) def _run_threads(self, target): errors = [] def wrapped(index): try: target(index) except Exception as exc: # noqa: BLE001 - collected for the assertion errors.append(exc) finally: self.db.close() threads = [threading.Thread(target=wrapped, args=(i,)) for i in range(THREADS)] for thread in threads: thread.start() for thread in threads: thread.join() return errors class TestPerThreadConnections(ConcurrencyTestCase): def test_each_thread_gets_its_own_connection(self): main_conn = self.db.model.get_connection() self.assertIs(self.db.model.get_connection(), main_conn) seen = [] def grab(_index): first = self.db.model.get_connection() seen.append((first, self.db.model.get_connection())) self.assertEqual(self._run_threads(grab), []) self.assertEqual(len(seen), THREADS) for first, second in seen: self.assertIs(first, second) self.assertIsNot(first, main_conn) self.assertEqual(len({id(first) for first, _ in seen}), THREADS) def test_close_releases_only_the_calling_threads_connection(self): main_conn = self.db.model.get_connection() other = [] def grab(_index): other.append(self.db.model.get_connection()) self.db.close() self.assertEqual(self._run_threads(grab), []) # The main thread's connection survives closes made by other threads self.assertIs(self.db.model.get_connection(), main_conn) main_conn.execute("SELECT 1").fetchone() self.db.close() self.assertIsNot(self.db.model.get_connection(), main_conn) class TestConcurrentWrites(ConcurrencyTestCase): def test_distinct_prompts_saved_and_counted_from_many_threads(self): def work(index): for n in range(CALLS_PER_THREAD): text = f"thread {index} prompt {n}" pid = self.db.save_prompt( text=text, prompt_hash=generate_prompt_hash(text) ) self.assertTrue(self.db.record_prompt_use(pid)) self.assertEqual(self._run_threads(work), []) info = self.db.model.get_database_info() self.assertEqual(info["total_prompts"], THREADS * CALLS_PER_THREAD) prompts = self.db.search_prompts(limit=THREADS * CALLS_PER_THREAD + 1) self.assertEqual(len(prompts), THREADS * CALLS_PER_THREAD) self.assertTrue(all(p["run_count"] == 1 for p in prompts)) def test_saving_the_same_prompt_from_many_threads_keeps_one_row(self): text = "the same prompt from every thread" prompt_hash = generate_prompt_hash(text) ids = [] def work(_index): for _ in range(CALLS_PER_THREAD): pid = self.db.save_prompt(text=text, prompt_hash=prompt_hash) ids.append(pid) self.assertTrue(self.db.record_prompt_use(pid)) self.assertEqual(self._run_threads(work), []) self.assertEqual(len(ids), THREADS * CALLS_PER_THREAD) self.assertEqual(len(set(ids)), 1) self.assertEqual(self.db.model.get_database_info()["total_prompts"], 1) prompt = self.db.get_prompt_by_hash(prompt_hash) self.assertEqual(prompt["id"], ids[0]) self.assertEqual(prompt["run_count"], THREADS * CALLS_PER_THREAD) def connection_settings(conn): return tuple( conn.execute(f"PRAGMA {name}").fetchone()[0] for name in ("journal_mode", "synchronous", "busy_timeout", "foreign_keys") ) # WAL journal, synchronous=NORMAL (1), the shared busy timeout, foreign keys on. MANAGED_SETTINGS = ("wal", 1, models.BUSY_TIMEOUT_SECONDS * 1000, 1) class TestConnectionConfiguration(ConcurrencyTestCase): """Every managed connection shares one locking and durability setup. With synchronous=FULL each WAL commit waits for a disk flush while it holds SQLite's single write lock. On slow disks (Windows runners) the other writers' polling busy handlers then starve past the busy timeout and fail with "database is locked" (#242). In WAL mode synchronous=NORMAL keeps the database consistent after an application crash and only flushes at checkpoints, so commits stop holding the lock across a flush. """ def test_every_thread_connection_uses_the_managed_settings(self): seen = [] def grab(_index): seen.append(connection_settings(self.db.model.get_connection())) self.assertEqual(self._run_threads(grab), []) seen.append(connection_settings(self.db.model.get_connection())) self.db.close_all() seen.append(connection_settings(self.db.model.get_connection())) self.assertEqual(seen, [MANAGED_SETTINGS] * (THREADS + 2)) def test_schema_connection_uses_the_managed_settings(self): fd, path = tempfile.mkstemp(suffix=".db") os.close(fd) seen = [] create_tables = models.PromptModel._create_tables def record(model, conn): seen.append(connection_settings(conn)) return create_tables(model, conn) fresh = None try: with patch.object( models.PromptModel, "_create_tables", autospec=True, side_effect=record ): fresh = PromptDatabase(path) finally: if fresh is not None: fresh.close_all() for suffix in ("", "-wal", "-shm"): if os.path.exists(path + suffix): os.unlink(path + suffix) self.assertEqual(seen, [MANAGED_SETTINGS]) def test_normal_sync_is_only_used_with_a_wal_journal(self): # An in-memory database cannot use WAL; keep SQLite's FULL default there. conn = sqlite3.connect(":memory:") try: models.configure_connection(conn) self.assertEqual( connection_settings(conn), ("memory", 2, *MANAGED_SETTINGS[2:]) ) finally: conn.close() if __name__ == "__main__": unittest.main()