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.
333 lines
12 KiB
Python
333 lines
12 KiB
Python
"""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()
|