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.
354 lines
12 KiB
Python
354 lines
12 KiB
Python
"""A SQLite connection that is never closed under the thread using it.
|
|
|
|
``PromptModel`` gives each thread its own connection, but Move library and
|
|
``close_all`` must retire every thread's connection from one thread. Closing
|
|
a ``sqlite3.Connection`` while its owner is inside a statement (Python runs
|
|
``sqlite3_step`` without the GIL) frees memory the owner then reads: on
|
|
Windows the process dies with "access violation", elsewhere the owner sees
|
|
garbage such as a ``MemoryError``.
|
|
|
|
``ManagedConnection`` counts the calls its owner is inside. ``retire()`` is
|
|
safe from any thread and never raises: an idle connection is closed at once;
|
|
one in use is closed by its owner as the call ends. A retired connection
|
|
refuses new work (cursors, statements, commits) with the error a closed
|
|
connection raises.
|
|
|
|
What retiring a connection in use does and does not do:
|
|
|
|
- It calls ``sqlite3_interrupt`` (thread-safe). That ends a running
|
|
statement at its next step, but it does not wake a statement that is
|
|
waiting for another connection's lock: that one waits until the lock is
|
|
released (or the busy timeout ends it), and only then is it interrupted.
|
|
An interrupt that arrives between two statements, for example between the
|
|
implicit ``BEGIN`` and the DML or between ``executemany`` rows, is lost.
|
|
- Writes are kept out of a file the process has switched away from because
|
|
a retired connection never commits: legacy-mode DML runs inside a
|
|
transaction that only ``commit()`` (refused) or ``close()`` (rolls back)
|
|
ends. A statement that commits by itself is not covered: DDL or a PRAGMA
|
|
write outside a transaction, or an explicit ``COMMIT`` passed to
|
|
``execute``/``executescript`` (``executescript`` also commits first, in C,
|
|
on Python before 3.12). Managed connections must never be used that way
|
|
for library writes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import logging
|
|
import threading
|
|
import time
|
|
import weakref
|
|
from typing import Iterable, Iterator, List, Optional
|
|
|
|
import sqlite3
|
|
|
|
CLOSED_MESSAGE = "Cannot operate on a closed database."
|
|
# Rows fetched per step when a cursor is iterated: the connection is marked
|
|
# in use once per chunk instead of once per row (iteration costs about what
|
|
# a plain cursor costs). A loop that stops early leaves up to this many rows
|
|
# fetched; ``fetchone``/``fetchmany`` after a partial loop skip them.
|
|
ITER_CHUNK = 128
|
|
|
|
|
|
class ManagedCursor(sqlite3.Cursor):
|
|
"""A cursor whose calls into SQLite count as its connection being in use."""
|
|
|
|
def execute(self, sql, parameters=(), /):
|
|
conn = self.connection
|
|
conn._enter()
|
|
try:
|
|
return super().execute(sql, parameters)
|
|
finally:
|
|
conn._leave()
|
|
|
|
def executemany(self, sql, seq_of_parameters, /):
|
|
conn = self.connection
|
|
conn._enter()
|
|
try:
|
|
return super().executemany(sql, seq_of_parameters)
|
|
finally:
|
|
conn._leave()
|
|
|
|
def executescript(self, sql_script, /):
|
|
conn = self.connection
|
|
conn._enter()
|
|
try:
|
|
return super().executescript(sql_script)
|
|
finally:
|
|
conn._leave()
|
|
|
|
def fetchone(self):
|
|
conn = self.connection
|
|
conn._enter()
|
|
try:
|
|
return super().fetchone()
|
|
finally:
|
|
conn._leave()
|
|
|
|
def fetchmany(self, *args, **kwargs):
|
|
conn = self.connection
|
|
conn._enter()
|
|
try:
|
|
return super().fetchmany(*args, **kwargs)
|
|
finally:
|
|
conn._leave()
|
|
|
|
def fetchall(self):
|
|
conn = self.connection
|
|
conn._enter()
|
|
try:
|
|
return super().fetchall()
|
|
finally:
|
|
conn._leave()
|
|
|
|
def __iter__(self) -> Iterator:
|
|
return self._chunks()
|
|
|
|
def _chunks(self) -> Iterator:
|
|
while True:
|
|
rows = self.fetchmany(ITER_CHUNK)
|
|
if not rows:
|
|
return
|
|
yield from rows
|
|
|
|
def __next__(self):
|
|
conn = self.connection
|
|
conn._enter()
|
|
try:
|
|
return super().__next__()
|
|
finally:
|
|
conn._leave()
|
|
|
|
def close(self):
|
|
conn = self.connection
|
|
try:
|
|
conn._enter()
|
|
except sqlite3.ProgrammingError:
|
|
return super().close() # the connection is closed or closing
|
|
try:
|
|
return super().close()
|
|
finally:
|
|
conn._leave()
|
|
|
|
|
|
class ManagedConnection(sqlite3.Connection):
|
|
"""A connection only its owner closes while a call into it is running."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
# Guards _depth, _retired, _closing and _cursors. Re-entrant so that
|
|
# code running in the closing thread while it holds the lock (a
|
|
# cursor() during the cursor-set copy) is refused instead of hanging.
|
|
self._use_lock = threading.RLock()
|
|
self._depth = 0
|
|
self._retired = False
|
|
self._closing = False
|
|
self._released = threading.Event()
|
|
# Python before 3.12 keeps the file open after close() while a cursor
|
|
# still holds a statement; these are closed first (see _close_now).
|
|
self._cursors: "weakref.WeakSet[sqlite3.Cursor]" = weakref.WeakSet()
|
|
self.close_logger: Optional[logging.Logger] = None
|
|
|
|
# ── in use ───────────────────────────────────────────────────────
|
|
|
|
def _enter(self) -> None:
|
|
with self._use_lock:
|
|
if self._retired:
|
|
raise sqlite3.ProgrammingError(CLOSED_MESSAGE)
|
|
self._depth += 1
|
|
|
|
def _leave(self) -> None:
|
|
with self._use_lock:
|
|
self._depth -= 1
|
|
close_now = self._retired and self._depth == 0
|
|
if close_now:
|
|
self._close_now()
|
|
|
|
@contextlib.contextmanager
|
|
def using(self):
|
|
"""Mark the connection in use for a call this class cannot wrap.
|
|
|
|
For example as the *target* of another connection's ``backup()``.
|
|
"""
|
|
self._enter()
|
|
try:
|
|
yield self
|
|
finally:
|
|
self._leave()
|
|
|
|
def cursor(self, factory=None):
|
|
with self._use_lock:
|
|
if self._retired:
|
|
raise sqlite3.ProgrammingError(CLOSED_MESSAGE)
|
|
cursor = super().cursor(factory or ManagedCursor)
|
|
self._cursors.add(cursor)
|
|
return cursor
|
|
|
|
def execute(self, sql, parameters=(), /):
|
|
return self.cursor().execute(sql, parameters)
|
|
|
|
def executemany(self, sql, seq_of_parameters, /):
|
|
return self.cursor().executemany(sql, seq_of_parameters)
|
|
|
|
def executescript(self, sql_script, /):
|
|
return self.cursor().executescript(sql_script)
|
|
|
|
def commit(self):
|
|
self._enter()
|
|
try:
|
|
return super().commit()
|
|
finally:
|
|
self._leave()
|
|
|
|
def rollback(self):
|
|
self._enter() # a closed connection raises, as sqlite3's does
|
|
try:
|
|
return super().rollback()
|
|
finally:
|
|
self._leave()
|
|
|
|
def backup(self, target, **kwargs):
|
|
self._enter()
|
|
try:
|
|
return super().backup(target, **kwargs)
|
|
finally:
|
|
self._leave()
|
|
|
|
def create_function(self, *args, **kwargs):
|
|
self._enter()
|
|
try:
|
|
return super().create_function(*args, **kwargs)
|
|
finally:
|
|
self._leave()
|
|
|
|
@property
|
|
def in_transaction(self) -> bool:
|
|
self._enter()
|
|
try:
|
|
return super().in_transaction
|
|
finally:
|
|
self._leave()
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
# The C implementation calls its own commit/rollback, not these.
|
|
if exc_type is None:
|
|
try:
|
|
self.commit()
|
|
except BaseException:
|
|
with contextlib.suppress(sqlite3.Error):
|
|
self.rollback()
|
|
raise
|
|
else:
|
|
self.rollback()
|
|
return False
|
|
|
|
# ── closing ──────────────────────────────────────────────────────
|
|
|
|
def close(self) -> None:
|
|
"""Close now if idle, else as soon as the running call ends.
|
|
|
|
Unlike ``sqlite3.Connection.close``, a call from another thread while
|
|
the owner is inside a call returns before the file is closed. A caller
|
|
that must know the file is closed (before renaming or deleting it)
|
|
waits for ``wait_closed`` (``PromptModel.close_detached`` does).
|
|
"""
|
|
self.retire()
|
|
|
|
def retire(self) -> bool:
|
|
"""Refuse new work and close; True when the connection is closed already.
|
|
|
|
Safe from any thread, and never raises: failures are logged.
|
|
"""
|
|
try:
|
|
with self._use_lock:
|
|
if not self._retired:
|
|
self._retired = True
|
|
if self._depth:
|
|
# Ends the running statement (see the module docstring
|
|
# for what it does not end). The owner cannot start
|
|
# closing the connection while this lock is held.
|
|
self.interrupt()
|
|
return False
|
|
elif self._depth:
|
|
return False
|
|
else:
|
|
return self._released.wait(0)
|
|
except Exception:
|
|
self._log_exception("Error retiring database connection")
|
|
return self._released.wait(0)
|
|
self._close_now()
|
|
return True
|
|
|
|
def wait_closed(self, timeout: float) -> bool:
|
|
return self._released.wait(timeout)
|
|
|
|
def _close_now(self) -> None:
|
|
# Only ever called with no call into the connection running, and a
|
|
# retired connection refuses new calls and cursors: its cursors are
|
|
# idle, and the set is copied under the lock that guards it.
|
|
with self._use_lock:
|
|
if self._closing:
|
|
return # another thread closes it and sets _released
|
|
self._closing = True
|
|
cursors = list(self._cursors)
|
|
try:
|
|
for cursor in cursors:
|
|
with contextlib.suppress(sqlite3.Error):
|
|
sqlite3.Cursor.close(cursor)
|
|
super().close()
|
|
except Exception:
|
|
self._log_exception("Error closing database connection")
|
|
finally:
|
|
self._released.set()
|
|
|
|
def _log_exception(self, message: str) -> None:
|
|
if self.close_logger is not None:
|
|
self.close_logger.exception(message)
|
|
|
|
|
|
def retire_connections(
|
|
connections: Iterable[sqlite3.Connection], logger=None
|
|
) -> List[ManagedConnection]:
|
|
"""Retire every connection; return those whose owner has yet to close them.
|
|
|
|
Safe while holding a lock the owners wait on: nothing here waits for them.
|
|
One connection failing to retire is logged and never skips the rest.
|
|
"""
|
|
pending = []
|
|
for conn in connections:
|
|
try:
|
|
if isinstance(conn, ManagedConnection):
|
|
if logger is not None:
|
|
conn.close_logger = logger
|
|
if not conn.retire():
|
|
pending.append(conn)
|
|
else:
|
|
conn.close()
|
|
except Exception:
|
|
if logger is not None:
|
|
logger.exception("Error closing database connection")
|
|
return pending
|
|
|
|
|
|
def await_closed(
|
|
connections: Iterable[ManagedConnection], timeout: float, logger=None
|
|
) -> bool:
|
|
"""Wait until the owners closed ``connections`` (retired); False on timeout.
|
|
|
|
A connection still open after ``timeout`` (its owner is blocked, for
|
|
example on another process's lock after a lost interrupt) is logged and
|
|
left to close when its call ends; a Move library then reports the old
|
|
files as still in use instead of renaming them.
|
|
"""
|
|
deadline = time.monotonic() + timeout
|
|
still_open = 0
|
|
for conn in connections:
|
|
if not conn.wait_closed(max(0.0, deadline - time.monotonic())):
|
|
still_open += 1
|
|
if still_open and logger is not None:
|
|
logger.warning(
|
|
f"{still_open} database connection(s) still in use after "
|
|
f"{timeout:g}s; they close when their statement ends"
|
|
)
|
|
return not still_open
|