Vito Sansevero
1dcebd640c
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.
221 lines
7.8 KiB
Python
221 lines
7.8 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
|
|
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 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()
|