mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-30 20:35:22 +08:00
fix(persistence): serialize SQLite runtime write transactions
This commit is contained in:
+2
-2
@@ -272,13 +272,13 @@ Maintenance rule: when adding, removing, or changing an HTTP endpoint that shoul
|
||||
|
||||
Persistence is centered on `pkg/persistence/mgr.py`.
|
||||
|
||||
- SQLite is the default database; PostgreSQL is supported.
|
||||
- SQLite is the default database; PostgreSQL is supported. The shared SQLite runtime engine admits one writer at a time from the first write through physical commit/rollback/close, including ORM and direct connection writes. Readers do not acquire the gate. It is process-local and does not coordinate other engines or processes; PostgreSQL is unchanged. An engine must stay on its original event loop, even after disposal. A connection-close failure makes the gate reject subsequent SQL explicitly until the process is restarted; queued callers are woken to fail rather than left blocked.
|
||||
- Every SQLite connection sets a 64 MiB `journal_size_limit`, including connections recreated after migration or pool disposal. SQLite's default automatic checkpoint policy and synchronous setting remain unchanged. The limit reduces retained WAL space when the WAL resets; it does not cap active transactions or WAL pinned by readers. Online maintenance does not periodically force checkpoints or truncate WAL files. Diagnose persistent excess growth by checking checkpoint progress and actual transaction lifetimes, then shorten the responsible transaction rather than forcing readers to finish.
|
||||
- After schema initialization and migration, every SQLite startup completes `wal_checkpoint(TRUNCATE)` before runtime setup and producers start. Lock waits allow up to 60 seconds; checkpoint I/O may take longer. An incomplete checkpoint fails startup without undoing successful migrations, and the next startup retries even at schema head. The connection busy timeout is restored afterward. PostgreSQL is unaffected.
|
||||
- SQLite files use incremental auto-vacuum. Migration `0042_sqlite_incremental_vacuum` converts existing files before producers start, with a verified backup and restoration on conversion failure when a full rebuild is required. Fresh files enable the format before creating tables. Monitoring cleanup reclaims free pages in independent adaptive batches (1–512 pages, starting at 8, targeting 25 ms), pausing 250 ms between batches and starting no new batches after a two-second round budget. Busy writers take priority. No full `VACUUM` or forced WAL checkpoint runs during online cleanup.
|
||||
- Models live under `pkg/entity/persistence/`.
|
||||
- `pkg/persistence/media.py::MediaReferenceText` validates bound conversation payloads on ORM/Core writes while preserving the existing TEXT schema and string API. It rejects recognized structured inline images: retained images go through `MediaCache.externalize_chain_dump()` before serialization, while metadata-only snapshots and transcripts use `redact_inline_images()`. Reads and comparisons preserve historical values for incremental cleanup. Runtime/adapter messages, arbitrary diagnostic text, and opaque strings are outside this contract.
|
||||
- `pkg/persistence/resource_urls.py` interns structured resource URL fields in retained messages, history and monitoring at the typed SQL boundary, in the source transaction. `MediaReferenceJSON` preserves native JSON monitoring details. Payloads store opaque references to the global `resource_urls` catalog; typed SQL reads and comparisons expand them to current URLs, preserving API shapes and source-record authorization. Configuration URLs and URLs in free text are excluded. Catalog lookups use only the primary-key and URL-hash unique indexes. Successful refreshes overwrite the current URL and persist the check time even when standard `get_msg` returns the same URL. Converging URLs retain only ID redirects, with no old URL. PostgreSQL row locks and SQLite's immediate write transaction serialize concurrent refreshes; SQLite holds its write lock during the upstream call. Bounded in-memory aliases let reused Runner payloads resolve old URLs without retaining those URLs in the database. Migration `0041_resource_urls` normalizes existing payloads using their source creation timestamps before producers start; downgrade expands references before dropping the catalog.
|
||||
- `pkg/persistence/resource_urls.py` interns structured resource URL fields in retained messages, history and monitoring at the typed SQL boundary, in the source transaction. `MediaReferenceJSON` preserves native JSON monitoring details. Payloads store opaque references to the global `resource_urls` catalog; typed SQL reads and comparisons expand them to current URLs, preserving API shapes and source-record authorization. Configuration URLs and URLs in free text are excluded. Catalog lookups use only the primary-key and URL-hash unique indexes. Successful refreshes overwrite the current URL and persist the check time even when standard `get_msg` returns the same URL. Converging URLs retain only ID redirects, with no old URL. PostgreSQL row locks serialize concurrent refreshes. SQLite serializes refreshes within the engine, performs upstream calls outside its write transaction, and rechecks catalog rows before applying results in a short immediate transaction. Bounded in-memory aliases let reused Runner payloads resolve old URLs without retaining those URLs in the database. Migration `0041_resource_urls` normalizes existing payloads using their source creation timestamps before producers start; downgrade expands references before dropping the catalog.
|
||||
- Timezone-less SQL `DateTime` columns store UTC-naive values on both backends. Normalize aware values with `pkg/persistence/datetime_utils.py::as_naive_utc` at the bind boundary, including deadlines, leases, event times, and query/retention cutoffs; do not merely strip an offset. Existing naive rows represent UTC. Restore UTC awareness for application comparisons and epoch serialization. Runner ledger/event/transcript stores follow this existing schema contract without a data migration or changes to journal/tenant authorization.
|
||||
- Fresh schemas are created from current metadata, then Alembic migrations run to head. LangBot 4.x does not upgrade 3.x databases.
|
||||
- New schema changes use Alembic under `pkg/persistence/alembic/versions/`; there is no legacy migration chain in 4.x.
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import sqlalchemy
|
||||
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
|
||||
|
||||
from .. import database
|
||||
from .. import database, sqlite_write_gate
|
||||
|
||||
|
||||
WAL_RETAIN_BYTES = 64 * 1024 * 1024
|
||||
@@ -27,6 +27,7 @@ class SQLiteDatabaseManager(database.BaseDatabaseManager):
|
||||
db_file_path = self.ap.instance_config.data.get('database', {}).get('sqlite', {}).get('path', 'data/langbot.db')
|
||||
engine_url = f'sqlite+aiosqlite:///{db_file_path}'
|
||||
self.engine = sqlalchemy_asyncio.create_async_engine(engine_url)
|
||||
self.write_gate = sqlite_write_gate.install(self.engine)
|
||||
sqlalchemy.event.listen(self.engine.sync_engine, 'connect', _configure_connection)
|
||||
# Reader transactions must not block independent plugin callback commits.
|
||||
# Configure the persistent journal before serving requests, not per write.
|
||||
@@ -35,3 +36,6 @@ class SQLiteDatabaseManager(database.BaseDatabaseManager):
|
||||
await connection.execute(sqlalchemy.text('PRAGMA auto_vacuum=INCREMENTAL'))
|
||||
await connection.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
|
||||
await connection.commit()
|
||||
|
||||
def resource_stats(self) -> dict[str, int]:
|
||||
return self.write_gate.resource_stats()
|
||||
|
||||
@@ -5,6 +5,7 @@ ORM and Core callers retain their existing string/JSON contracts. The catalog
|
||||
is shared; authorization continues to belong to the enclosing business record.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import hashlib
|
||||
import json
|
||||
@@ -255,12 +256,23 @@ def create_postgres_expansion(connection):
|
||||
|
||||
|
||||
async def refresh_images(engine, peer, payload, *, is_current=None):
|
||||
"""Serialize refreshes on catalog identities and overwrite current URLs atomically."""
|
||||
"""Serialize refreshes without holding SQLite's writer during upstream calls."""
|
||||
from ..platform.adapters.aiocqhttp.image_urls import image_nodes
|
||||
|
||||
nodes = list(image_nodes(payload))
|
||||
if not nodes:
|
||||
return
|
||||
if engine.dialect.name == 'sqlite':
|
||||
sync_engine = engine.sync_engine
|
||||
if not hasattr(sync_engine, '_langbot_image_refresh_lock'):
|
||||
sync_engine._langbot_image_refresh_lock = asyncio.Lock()
|
||||
async with sync_engine._langbot_image_refresh_lock:
|
||||
return await _refresh_images(engine, peer, payload, nodes, is_current=is_current)
|
||||
return await _refresh_images(engine, peer, payload, nodes, is_current=is_current)
|
||||
|
||||
|
||||
async def _refresh_images(engine, peer, payload, nodes, *, is_current=None):
|
||||
from ..platform.adapters.aiocqhttp.image_urls import nt_image_type
|
||||
|
||||
def connected():
|
||||
return not peer.closed and (is_current is None or is_current())
|
||||
@@ -273,10 +285,7 @@ async def refresh_images(engine, peer, payload, *, is_current=None):
|
||||
await db.run_sync(lambda sync: intern(sync, json.dumps(payload), identities_out=references))
|
||||
identities = set(references.values())
|
||||
async with engine.connect() as db:
|
||||
if engine.dialect.name == 'sqlite':
|
||||
await db.exec_driver_sql('BEGIN IMMEDIATE')
|
||||
else:
|
||||
await db.begin()
|
||||
await db.begin()
|
||||
try:
|
||||
rows = (
|
||||
(
|
||||
@@ -301,6 +310,8 @@ async def refresh_images(engine, peer, payload, *, is_current=None):
|
||||
# Capture origins even when renewal is not due or identification fails.
|
||||
for node, message_id, image_id in nodes:
|
||||
peer.image_urls.remember(by_url[node['url']]['url'], message_id, image_id)
|
||||
if engine.dialect.name == 'sqlite':
|
||||
await db.rollback()
|
||||
# Identification and upstream credential handling remain connection-local.
|
||||
if not connected() or not await peer.image_urls.identify():
|
||||
return
|
||||
@@ -327,11 +338,39 @@ async def refresh_images(engine, peer, payload, *, is_current=None):
|
||||
return
|
||||
if not connected():
|
||||
return
|
||||
if engine.dialect.name == 'sqlite':
|
||||
cached = {}
|
||||
for previous, entry in zip(old, batch):
|
||||
origin = peer.image_urls.references.get(previous)
|
||||
entries = [
|
||||
(peer.image_urls.urls, previous),
|
||||
(peer.image_urls.urls, entry['image']['url']),
|
||||
(peer.image_urls.messages, str(origin[0]) if origin else None),
|
||||
(peer.image_urls.keys, nt_image_type(previous)),
|
||||
]
|
||||
cached[previous] = [(cache, key, cache[key]) for cache, key in entries if key in cache]
|
||||
await db.exec_driver_sql('BEGIN IMMEDIATE')
|
||||
stamp = now()
|
||||
updated = set()
|
||||
for (node, _, _, row), entry, previous in zip(due, batch, old):
|
||||
if previous not in completed:
|
||||
continue
|
||||
if engine.dialect.name == 'sqlite' and row['id'] not in updated:
|
||||
current = (await db.execute(sa.select(table).where(table.c.id == row['id']))).mappings().one()
|
||||
if current != row:
|
||||
# Another engine may have renewed or redirected this identity.
|
||||
while current['canonical_id'] is not None:
|
||||
current = (
|
||||
(await db.execute(sa.select(table).where(table.c.id == current['canonical_id'])))
|
||||
.mappings()
|
||||
.one()
|
||||
)
|
||||
node['url'] = current['url']
|
||||
# A rejected result must not return from the peer cache on the next refresh.
|
||||
for cache, key, value in cached[previous]:
|
||||
if cache.get(key) is value:
|
||||
del cache[key]
|
||||
continue
|
||||
fresh = entry['image']['url']
|
||||
node['url'] = fresh
|
||||
if napcat and fresh == previous:
|
||||
|
||||
@@ -0,0 +1,252 @@
|
||||
"""Serialize SQLite writers on one application engine, without locking readers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import time
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.sqlite.aiosqlite import AsyncAdapt_aiosqlite_connection
|
||||
from sqlalchemy.util.concurrency import await_only
|
||||
|
||||
|
||||
_TOKENS = re.compile(
|
||||
r"\ufeff|--[^\n]*(?:\n|$)|/\*.*?\*/|'(?:''|[^'])*'|\"(?:\"\"|[^\"])*\"|`[^`]*`|\[[^\]]*\]|[A-Za-z_\x80-\U0010ffff][A-Za-z_0-9$\x80-\U0010ffff]*|[(),.=]",
|
||||
re.S,
|
||||
)
|
||||
_WRITES = {'INSERT', 'UPDATE', 'DELETE', 'REPLACE', 'CREATE', 'ALTER', 'DROP', 'VACUUM', 'REINDEX', 'ATTACH', 'DETACH'}
|
||||
|
||||
|
||||
def _writes(statement: str) -> bool:
|
||||
words = [
|
||||
m.group()
|
||||
for m in _TOKENS.finditer(statement)
|
||||
if m.group() != '\ufeff' and not m.group().startswith(('--', '/*'))
|
||||
]
|
||||
if not words:
|
||||
return False
|
||||
first = words[0].upper()
|
||||
if first in _WRITES:
|
||||
return True
|
||||
if first == 'BEGIN':
|
||||
return any(word.upper() in {'IMMEDIATE', 'EXCLUSIVE'} for word in words[1:])
|
||||
if first == 'PRAGMA':
|
||||
# SQLite accepts both name=value and name(value) for writable PRAGMAs.
|
||||
read_arguments = {
|
||||
'table_info',
|
||||
'table_xinfo',
|
||||
'table_list',
|
||||
'index_info',
|
||||
'index_xinfo',
|
||||
'index_list',
|
||||
'foreign_key_list',
|
||||
'foreign_key_check',
|
||||
'integrity_check',
|
||||
'quick_check',
|
||||
}
|
||||
# Only the pragma name is an identifier; arguments/comments may contain
|
||||
# '=', keywords or parentheses without changing a read into a write.
|
||||
name_index = 3 if len(words) > 3 and words[2] == '.' else 1
|
||||
if len(words) <= name_index:
|
||||
return False
|
||||
name = words[name_index].strip('"\'`[]').lower()
|
||||
argument = words[name_index + 1 : name_index + 2]
|
||||
return name in {'incremental_vacuum', 'wal_checkpoint', 'optimize'} or (
|
||||
argument in (['='], ['(']) and name not in read_arguments
|
||||
)
|
||||
if first == 'WITH':
|
||||
index = 2 if len(words) > 1 and words[1].upper() == 'RECURSIVE' else 1
|
||||
|
||||
def group_end(start):
|
||||
depth = 0
|
||||
for position in range(start, len(words)):
|
||||
if words[position] == '(':
|
||||
depth += 1
|
||||
elif words[position] == ')':
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return position + 1
|
||||
return len(words)
|
||||
|
||||
while index < len(words):
|
||||
index += 1 # CTE name, including identifiers that resemble keywords.
|
||||
if index < len(words) and words[index] == '(':
|
||||
index = group_end(index) # Optional column names.
|
||||
if index >= len(words) or words[index].upper() != 'AS':
|
||||
return False
|
||||
index += 1
|
||||
while index < len(words) and words[index].upper() in {'NOT', 'MATERIALIZED'}:
|
||||
index += 1
|
||||
if index >= len(words) or words[index] != '(':
|
||||
return False
|
||||
index = group_end(index)
|
||||
if index < len(words) and words[index] == ',':
|
||||
index += 1
|
||||
continue
|
||||
return index < len(words) and words[index].upper() in _WRITES
|
||||
return False
|
||||
|
||||
|
||||
class SQLiteWriteGate:
|
||||
def __init__(self):
|
||||
self.lock = asyncio.Lock()
|
||||
self.loop = None
|
||||
self.failed = False
|
||||
self.waiting = 0
|
||||
self.acquisitions = 0
|
||||
self.wait_seconds = 0.0
|
||||
self.max_wait_seconds = 0.0
|
||||
|
||||
def check(self):
|
||||
loop = asyncio.get_running_loop()
|
||||
if self.loop is None:
|
||||
self.loop = loop
|
||||
elif self.loop is not loop:
|
||||
raise RuntimeError('SQLite writer gate belongs to another event loop; create a new engine for this loop')
|
||||
if self.failed:
|
||||
raise RuntimeError(
|
||||
'SQLite writer gate unavailable after connection close failure; restart the process before writing'
|
||||
)
|
||||
|
||||
async def acquire(self):
|
||||
self.check()
|
||||
start = time.monotonic()
|
||||
self.waiting += 1
|
||||
try:
|
||||
await self.lock.acquire()
|
||||
finally:
|
||||
self.waiting -= 1
|
||||
try:
|
||||
self.check()
|
||||
except BaseException:
|
||||
self.lock.release()
|
||||
raise
|
||||
elapsed = time.monotonic() - start
|
||||
self.acquisitions += 1
|
||||
self.wait_seconds += elapsed
|
||||
self.max_wait_seconds = max(self.max_wait_seconds, elapsed)
|
||||
|
||||
def resource_stats(self):
|
||||
return {
|
||||
'sqlite_writer_active': int(self.lock.locked()),
|
||||
'sqlite_writer_waiting': self.waiting,
|
||||
'sqlite_writer_acquisitions': self.acquisitions,
|
||||
'sqlite_writer_wait_ms_total': int(self.wait_seconds * 1000),
|
||||
'sqlite_writer_wait_ms_max': int(self.max_wait_seconds * 1000),
|
||||
}
|
||||
|
||||
|
||||
class _GatedConnection(AsyncAdapt_aiosqlite_connection):
|
||||
"""Release after DBAPI completion, not SQLAlchemy's pre-commit event."""
|
||||
|
||||
__slots__ = ('gate', 'writer')
|
||||
|
||||
def __init__(self, dbapi, connection, gate):
|
||||
super().__init__(dbapi, connection)
|
||||
self.gate = gate
|
||||
self.writer = False
|
||||
|
||||
def release(self):
|
||||
if self.writer:
|
||||
self.writer = False
|
||||
self.gate.lock.release()
|
||||
|
||||
def acquire(self):
|
||||
self.gate.check()
|
||||
if not self.writer:
|
||||
await_only(self.gate.acquire())
|
||||
self.writer = True
|
||||
|
||||
def fail_closed(self):
|
||||
# Wake queued callers to fail explicitly, never admit a new writer when
|
||||
# physical close completion is uncertain. Recovery requires process restart.
|
||||
self.gate.failed = True
|
||||
self.release()
|
||||
|
||||
def _terminate_force_close(self):
|
||||
try:
|
||||
stop = getattr(self.driver_connection, 'stop', None)
|
||||
closed = (
|
||||
stop() if stop is not None else asyncio.get_running_loop().create_task(self._terminate_graceful_close())
|
||||
)
|
||||
if closed is None:
|
||||
self.fail_closed()
|
||||
else:
|
||||
closed.add_done_callback(self._closed)
|
||||
except BaseException:
|
||||
self.fail_closed()
|
||||
raise
|
||||
|
||||
def _closed(self, future):
|
||||
if future.cancelled() or future.exception() is not None:
|
||||
self.fail_closed()
|
||||
else:
|
||||
self.release()
|
||||
|
||||
def commit(self):
|
||||
super().commit()
|
||||
self.release()
|
||||
|
||||
def rollback(self):
|
||||
super().rollback()
|
||||
self.release()
|
||||
|
||||
def close(self):
|
||||
try:
|
||||
super().close()
|
||||
except BaseException:
|
||||
self.fail_closed()
|
||||
raise
|
||||
self.release()
|
||||
|
||||
async def _terminate_graceful_close(self):
|
||||
try:
|
||||
await self.driver_connection.close()
|
||||
except BaseException:
|
||||
self.fail_closed()
|
||||
raise
|
||||
self.release()
|
||||
|
||||
|
||||
def install(engine):
|
||||
"""Cover ORM flushes, Core statements and direct SQL on this async engine."""
|
||||
if engine.dialect.name != 'sqlite':
|
||||
return None
|
||||
# shortcut: one runtime engine/process; independent writers need PostgreSQL.
|
||||
gate = SQLiteWriteGate()
|
||||
|
||||
def connect(dialect, connection_record, args, kwargs):
|
||||
gate.check()
|
||||
adapted = dialect.loaded_dbapi.connect(*args, **kwargs)
|
||||
return _GatedConnection(dialect.loaded_dbapi, adapted.driver_connection, gate)
|
||||
|
||||
def before(connection, cursor, statement, parameters, context, executemany):
|
||||
gate.check()
|
||||
if context.compiled is not None and context.compiled.statement.is_select:
|
||||
return
|
||||
if context.isinsert or context.isupdate or context.isdelete or _writes(statement):
|
||||
connection.connection.dbapi_connection.acquire()
|
||||
|
||||
def after(connection, cursor, statement, parameters, context, executemany):
|
||||
adapted = connection.connection.dbapi_connection
|
||||
# Autocommit DDL/DML and releasing the outermost SAVEPOINT end the
|
||||
# physical transaction even if SQLAlchemy still has an autobegin object.
|
||||
if not adapted.driver_connection.in_transaction:
|
||||
adapted.release()
|
||||
|
||||
def reset(connection, record, state):
|
||||
# A failed SQLite COMMIT can retain its transaction after SQLAlchemy
|
||||
# marks it inactive. Pool reset must not skip that physical rollback.
|
||||
if connection.writer:
|
||||
if state.asyncio_safe:
|
||||
connection.rollback()
|
||||
elif not engine.dialect.has_terminate:
|
||||
# Older SQLAlchemy/aiosqlite cannot terminate from GC themselves.
|
||||
connection._terminate_force_close()
|
||||
|
||||
sa.event.listen(engine.sync_engine, 'reset', reset)
|
||||
sa.event.listen(engine.sync_engine, 'do_connect', connect)
|
||||
sa.event.listen(engine.sync_engine, 'before_cursor_execute', before)
|
||||
sa.event.listen(engine.sync_engine, 'after_cursor_execute', after)
|
||||
return gate
|
||||
@@ -20,6 +20,7 @@ from langbot.pkg.entity.persistence.event_log import EventLog
|
||||
from langbot.pkg.entity.persistence.monitoring import MonitoringLLMCall
|
||||
from langbot.pkg.persistence.media import MediaReferenceText, MediaReferenceJSON
|
||||
from langbot.pkg.persistence.resource_urls import digest, refresh_images
|
||||
from langbot.pkg.persistence.sqlite_write_gate import install as install_write_gate
|
||||
from langbot.pkg.platform.adapters.aiocqhttp.image_urls import OneBotImageURLs
|
||||
|
||||
|
||||
@@ -46,6 +47,7 @@ async def database(tmp_path, request):
|
||||
engine = create_async_engine(url, connect_args={'server_settings': {'search_path': schema}})
|
||||
else:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "catalog.db"}')
|
||||
install_write_gate(engine)
|
||||
try:
|
||||
async with engine.begin() as db:
|
||||
await db.run_sync(ResourceURL.__table__.create)
|
||||
@@ -398,3 +400,74 @@ async def test_migration_preserves_source_creation_and_downgrade(database):
|
||||
# The same engine can write and read literal payloads after downgrade.
|
||||
await append(db, 'after', NEW)
|
||||
assert NEW in await db.scalar(sa.select(EventLog.input_json).where(EventLog.event_id == 'after'))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('napcat', [False, True])
|
||||
@pytest.mark.parametrize('blocked_action', ['get_version_info', 'get_msg'])
|
||||
@pytest.mark.parametrize('converge', [False, True])
|
||||
async def test_sqlite_refresh_releases_writer_and_rejects_stale_result(database, blocked_action, converge, napcat):
|
||||
if database.dialect.name != 'sqlite':
|
||||
pytest.skip('SQLite transaction boundary')
|
||||
async with database.begin() as db:
|
||||
await append(db, 'a')
|
||||
if converge:
|
||||
await append(db, 'b', NEW)
|
||||
await db.execute(sa.update(ResourceURL).values(created_at=dt.datetime(2000, 1, 1)))
|
||||
source_id = await db.scalar(sa.select(ResourceURL.id).where(ResourceURL.url == OLD))
|
||||
destination_id = await db.scalar(sa.select(ResourceURL.id).where(ResourceURL.url == NEW))
|
||||
entered = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
upstream_url = OLD.replace('rkey=old', 'rkey=stale')
|
||||
renewal_action = 'nc_get_rkey' if napcat else 'get_msg'
|
||||
if blocked_action == 'get_msg':
|
||||
blocked_action = renewal_action
|
||||
|
||||
async def action(name, **params):
|
||||
if name == blocked_action:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
if name == 'get_version_info':
|
||||
return {'app_name': 'NapCat.Onebot' if napcat else 'standard'}
|
||||
if napcat:
|
||||
return [{'type': 20, 'time': 1, 'ttl': 10**12, 'rkey': upstream_url.split('rkey=')[1]}]
|
||||
return {'message': [{'type': 'image', 'data': {'file': 'one.jpg', 'url': upstream_url}}]}
|
||||
|
||||
peer = SimpleNamespace(closed=False, call_action=AsyncMock(side_effect=action))
|
||||
peer.image_urls = OneBotImageURLs(peer)
|
||||
payload = chain()
|
||||
task = asyncio.create_task(refresh_images(database, peer, payload))
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 5)
|
||||
# Plain text must retain its fast path while an image renewal waits.
|
||||
await asyncio.wait_for(refresh_images(database, peer, [{'type': 'Plain', 'text': 'hello'}]), 1)
|
||||
|
||||
async def concurrent_writer():
|
||||
async with database.begin() as db:
|
||||
await append(db, 'unrelated', 'https://example.com/unrelated.png')
|
||||
values = dict(url=NEW, url_hash=digest(NEW))
|
||||
if converge:
|
||||
values = dict(url=None, url_hash=None, canonical_id=destination_id)
|
||||
await db.execute(sa.update(ResourceURL).where(ResourceURL.id == source_id).values(**values))
|
||||
|
||||
# It must commit while the external call is still blocked.
|
||||
await asyncio.wait_for(concurrent_writer(), 2)
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(task, 5)
|
||||
assert payload[1]['url'] == NEW
|
||||
async with database.connect() as db:
|
||||
assert NEW in await db.scalar(sa.select(EventLog.input_json).where(EventLog.event_id == 'a'))
|
||||
assert await db.scalar(sa.select(EventLog.event_id).where(EventLog.event_id == 'unrelated')) == 'unrelated'
|
||||
assert await db.scalar(sa.select(ResourceURL.last_checked_at).where(ResourceURL.id == source_id)) is None
|
||||
|
||||
# The rejected response also populated OneBot's connection-local caches.
|
||||
# A later due refresh must ask upstream again instead of restoring that stale URL.
|
||||
upstream_url = NEW
|
||||
second_payload = chain(NEW)
|
||||
await refresh_images(database, peer, second_payload)
|
||||
assert second_payload[1]['url'] == NEW
|
||||
assert sum(call.args[0] == renewal_action for call in peer.call_action.call_args_list) == 2
|
||||
async with database.connect() as db:
|
||||
assert NEW in await db.scalar(sa.select(EventLog.input_json).where(EventLog.event_id == 'a'))
|
||||
assert not any('rkey=stale' in (url or '') for url in (await db.execute(sa.select(ResourceURL.url))).scalars())
|
||||
|
||||
@@ -0,0 +1,359 @@
|
||||
"""Writer admission on real pooled connections, including failure/cleanup paths."""
|
||||
|
||||
import asyncio
|
||||
import gc
|
||||
import sqlite3
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||||
|
||||
from langbot.pkg.persistence.sqlite_write_gate import install
|
||||
|
||||
|
||||
class Base(DeclarativeBase):
|
||||
pass
|
||||
|
||||
|
||||
class Row(Base):
|
||||
__tablename__ = 'gate_rows'
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def database(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "gate.db"}')
|
||||
gate = install(engine)
|
||||
async with engine.begin() as connection:
|
||||
await connection.exec_driver_sql('PRAGMA journal_mode=WAL')
|
||||
await connection.run_sync(Base.metadata.create_all)
|
||||
yield engine, gate
|
||||
await engine.dispose()
|
||||
assert not gate.lock.locked()
|
||||
assert gate.waiting == 0
|
||||
|
||||
|
||||
async def queued(gate):
|
||||
async with asyncio.timeout(2):
|
||||
while gate.waiting == 0:
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_flush_and_direct_writer_share_gate_without_blocking_reads(database):
|
||||
engine, gate = database
|
||||
async with AsyncSession(engine) as session, engine.connect() as direct:
|
||||
session.add(Row(id=1))
|
||||
await session.flush()
|
||||
writer = asyncio.create_task(direct.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'))
|
||||
await queued(gate)
|
||||
async with engine.connect() as reader:
|
||||
assert await asyncio.wait_for(reader.scalar(sa.select(sa.func.count()).select_from(Row)), 1) == 0
|
||||
await session.commit()
|
||||
await asyncio.wait_for(writer, 1)
|
||||
assert gate.lock.locked()
|
||||
await direct.commit()
|
||||
async with engine.connect() as reader:
|
||||
assert await reader.scalar(sa.select(sa.func.count()).select_from(Row)) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_waiter_does_not_unlock_owner(database):
|
||||
engine, gate = database
|
||||
async with engine.connect() as owner, engine.connect() as waiting:
|
||||
await owner.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
||||
task = asyncio.create_task(waiting.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'))
|
||||
await queued(gate)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert gate.lock.locked() and gate.waiting == 0
|
||||
await owner.rollback()
|
||||
await waiting.exec_driver_sql('INSERT INTO gate_rows VALUES (3)')
|
||||
await waiting.commit()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_savepoint_rollback_keeps_gate_until_outer_transaction_ends(database):
|
||||
engine, gate = database
|
||||
async with engine.connect() as connection:
|
||||
await connection.exec_driver_sql('BEGIN')
|
||||
nested = await connection.begin_nested()
|
||||
await connection.execute(sa.insert(Row).values(id=1))
|
||||
await nested.rollback()
|
||||
assert gate.lock.locked()
|
||||
await connection.rollback()
|
||||
assert not gate.lock.locked()
|
||||
nested = await connection.begin_nested()
|
||||
await connection.execute(sa.insert(Row).values(id=2))
|
||||
await nested.commit()
|
||||
assert not gate.lock.locked() # SQLite released the outermost SAVEPOINT.
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_statement_and_commit_release_on_rollback(database):
|
||||
engine, gate = database
|
||||
async with engine.connect() as connection:
|
||||
await connection.exec_driver_sql('PRAGMA foreign_keys=ON')
|
||||
await connection.exec_driver_sql('CREATE TABLE parent (id INTEGER PRIMARY KEY)')
|
||||
await connection.exec_driver_sql(
|
||||
'CREATE TABLE child (id INTEGER REFERENCES parent(id) DEFERRABLE INITIALLY DEFERRED)'
|
||||
)
|
||||
await connection.commit()
|
||||
await connection.exec_driver_sql('INSERT INTO child VALUES (123)')
|
||||
with pytest.raises(sa.exc.IntegrityError):
|
||||
await connection.commit()
|
||||
assert gate.lock.locked()
|
||||
# A failed commit leaves SQLite's physical transaction open. Closing
|
||||
# must roll it back before another connection is admitted.
|
||||
assert not gate.lock.locked()
|
||||
async with engine.begin() as connection:
|
||||
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
||||
with pytest.raises(sa.exc.IntegrityError):
|
||||
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
||||
assert gate.lock.locked()
|
||||
assert not gate.lock.locked()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispose_reconnect_and_cte_dml(database):
|
||||
engine, gate = database
|
||||
for _ in range(2):
|
||||
async with engine.connect() as connection:
|
||||
count = gate.acquisitions
|
||||
await connection.exec_driver_sql('/* INSERT ignored */ WITH x(v) AS (SELECT 1) SELECT v FROM x')
|
||||
assert gate.acquisitions == count
|
||||
await connection.exec_driver_sql('WITH x(v) AS (SELECT 5) INSERT OR REPLACE INTO gate_rows SELECT v FROM x')
|
||||
# SQLite's legacy mode autocommits a leading WITH; no retained lock.
|
||||
await connection.commit()
|
||||
await connection.exec_driver_sql('PRAGMA user_version=3')
|
||||
assert await connection.scalar(sa.text('PRAGMA user_version')) == 3
|
||||
count = gate.acquisitions
|
||||
await connection.exec_driver_sql('PRAGMA user_version(4)')
|
||||
assert gate.acquisitions == count + 1
|
||||
assert await connection.scalar(sa.text('PRAGMA user_version')) == 4
|
||||
count = gate.acquisitions
|
||||
await connection.exec_driver_sql('PRAGMA table_info(gate_rows)')
|
||||
assert gate.acquisitions == count
|
||||
await connection.exec_driver_sql('PRAGMA incremental_vacuum(1)')
|
||||
assert not gate.lock.locked()
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gc_force_close_releases_after_driver_close(database):
|
||||
engine, gate = database
|
||||
connection = await engine.connect()
|
||||
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
||||
del connection
|
||||
with pytest.warns(sa.exc.SAWarning, match='garbage collector'):
|
||||
gc.collect()
|
||||
async with asyncio.timeout(2):
|
||||
while gate.lock.locked():
|
||||
await asyncio.sleep(0.01)
|
||||
async with engine.begin() as connection:
|
||||
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (2)')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_statement_drains_before_release(database):
|
||||
engine, gate = database
|
||||
started = asyncio.Event()
|
||||
loop = asyncio.get_running_loop()
|
||||
import time
|
||||
|
||||
def slow():
|
||||
loop.call_soon_threadsafe(started.set)
|
||||
time.sleep(0.1)
|
||||
return 1
|
||||
|
||||
async with engine.connect() as connection:
|
||||
raw = (await connection.get_raw_connection()).driver_connection
|
||||
await raw.create_function('slow', 0, slow)
|
||||
task = asyncio.create_task(connection.exec_driver_sql('INSERT INTO gate_rows VALUES (slow())'))
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
async with engine.begin() as connection:
|
||||
await asyncio.wait_for(connection.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'), 2)
|
||||
|
||||
|
||||
def test_non_sqlite_engine_is_not_touched():
|
||||
# No event API or sync engine is needed on the PostgreSQL path.
|
||||
assert install(SimpleNamespace(dialect=SimpleNamespace(name='postgresql'))) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_commit_retains_gate_until_connection_is_closed(database):
|
||||
engine, gate = database
|
||||
entered = asyncio.Event()
|
||||
async with engine.connect() as connection:
|
||||
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
||||
driver = (await connection.get_raw_connection()).driver_connection
|
||||
|
||||
async def blocked_commit():
|
||||
entered.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
driver.commit = blocked_commit
|
||||
task = asyncio.create_task(connection.commit())
|
||||
await asyncio.wait_for(entered.wait(), 1)
|
||||
assert gate.lock.locked()
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
async with engine.begin() as connection:
|
||||
await asyncio.wait_for(connection.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'), 1)
|
||||
assert await connection.scalar(sa.select(sa.func.count()).select_from(Row)) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_rollback_is_retried_before_pool_reuse(database):
|
||||
engine, gate = database
|
||||
async with engine.connect() as connection:
|
||||
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
||||
driver = (await connection.get_raw_connection()).driver_connection
|
||||
original = driver.rollback
|
||||
|
||||
async def fail_once():
|
||||
driver.rollback = original
|
||||
raise sqlite3.OperationalError('injected rollback failure')
|
||||
|
||||
driver.rollback = fail_once
|
||||
with pytest.raises(sa.exc.OperationalError, match='injected rollback failure'):
|
||||
await connection.rollback()
|
||||
assert gate.lock.locked()
|
||||
async with engine.begin() as connection:
|
||||
await asyncio.wait_for(connection.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'), 1)
|
||||
assert await connection.scalar(sa.select(sa.func.count()).select_from(Row)) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pool_reset_rollback_failure_closes_before_admitting_writer(database):
|
||||
engine, gate = database
|
||||
connection = await engine.connect()
|
||||
await connection.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
||||
driver = (await connection.get_raw_connection()).driver_connection
|
||||
|
||||
async def broken_rollback():
|
||||
raise sqlite3.OperationalError('persistent rollback failure')
|
||||
|
||||
driver.rollback = broken_rollback
|
||||
with pytest.raises(sa.exc.OperationalError, match='persistent rollback failure'):
|
||||
await connection.rollback()
|
||||
await connection.close()
|
||||
async with engine.begin() as connection:
|
||||
await asyncio.wait_for(connection.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'), 1)
|
||||
assert await connection.scalar(sa.select(sa.func.count()).select_from(Row)) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'statement,writes',
|
||||
[
|
||||
('WITH select1 AS (SELECT 7) INSERT INTO gate_rows SELECT * FROM select1', True),
|
||||
('\ufeffINSERT INTO gate_rows VALUES (1)', True),
|
||||
('WITH x AS (SELECT 2)\ufeff INSERT INTO gate_rows SELECT * FROM x', True),
|
||||
('WITH 中SELECT文 AS (SELECT 8) INSERT INTO gate_rows SELECT * FROM 中SELECT文', True),
|
||||
('WITH update1 AS (SELECT 9) SELECT * FROM update1', False),
|
||||
('WITH a😺SELECT AS (SELECT 8) INSERT INTO gate_rows SELECT * FROM a😺SELECT', True),
|
||||
('WITH a😺UPDATE AS (SELECT 9) SELECT * FROM a😺UPDATE', False),
|
||||
('WITH "insert" AS (SELECT 9) SELECT * FROM "insert"', False),
|
||||
('WITH replace AS (SELECT 9) SELECT * FROM replace', False),
|
||||
(
|
||||
'WITH RECURSIVE select1(n) AS (VALUES(1) UNION ALL SELECT n+1 FROM select1 WHERE n<3), '
|
||||
'update1 AS NOT MATERIALIZED (SELECT n FROM select1) INSERT INTO gate_rows SELECT n FROM update1',
|
||||
True,
|
||||
),
|
||||
('PRAGMA "incremental_vacuum"', True),
|
||||
('PRAGMA main."wal_checkpoint"', True),
|
||||
('PRAGMA [user_version](4)', True),
|
||||
("PRAGMA table_info('a=b')", False),
|
||||
('PRAGMA table_info(gate_rows) /* = */', False),
|
||||
],
|
||||
)
|
||||
async def test_direct_sql_classification_matches_sqlite(database, statement, writes):
|
||||
engine, gate = database
|
||||
async with engine.connect() as connection:
|
||||
acquired = gate.acquisitions
|
||||
result = await connection.exec_driver_sql(statement)
|
||||
if result.returns_rows:
|
||||
result.all()
|
||||
await connection.commit()
|
||||
assert gate.acquisitions - acquired == int(writes)
|
||||
if writes and statement.startswith('WITH'):
|
||||
assert await connection.scalar(sa.select(sa.func.count()).select_from(Row)) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('physically_closed', [False, True])
|
||||
@pytest.mark.parametrize('method', ['close', 'invalidate'])
|
||||
async def test_close_failure_fails_queued_and_future_writers(database, physically_closed, method):
|
||||
engine, gate = database
|
||||
owner = await engine.connect()
|
||||
waiting = await engine.connect()
|
||||
await owner.exec_driver_sql('INSERT INTO gate_rows VALUES (1)')
|
||||
driver = (await owner.get_raw_connection()).driver_connection
|
||||
original_close = driver.close
|
||||
|
||||
async def broken_close():
|
||||
if physically_closed:
|
||||
await original_close()
|
||||
raise sqlite3.OperationalError('injected close completion error')
|
||||
|
||||
driver.close = broken_close
|
||||
task = asyncio.create_task(waiting.exec_driver_sql('INSERT INTO gate_rows VALUES (2)'))
|
||||
await queued(gate)
|
||||
try:
|
||||
if method == 'invalidate':
|
||||
# Pool invalidation logs/swallow close errors; the gate must still fail closed.
|
||||
await owner.invalidate()
|
||||
else:
|
||||
with pytest.raises(sqlite3.OperationalError, match='injected close completion error'):
|
||||
await owner.run_sync(lambda connection: connection.connection.dbapi_connection.close())
|
||||
with pytest.raises(RuntimeError, match='gate unavailable.*restart the process'):
|
||||
await asyncio.wait_for(task, 1)
|
||||
with pytest.raises(RuntimeError, match='gate unavailable.*restart the process'):
|
||||
await waiting.exec_driver_sql('INSERT INTO gate_rows VALUES (3)')
|
||||
assert gate.failed and not gate.lock.locked()
|
||||
finally:
|
||||
driver.close = original_close
|
||||
await original_close()
|
||||
await owner.invalidate()
|
||||
await owner.close()
|
||||
await waiting.close()
|
||||
|
||||
|
||||
def test_engine_reuse_on_another_loop_fails_explicitly_even_after_dispose(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "loops.db"}')
|
||||
gate = install(engine)
|
||||
|
||||
async def first_loop():
|
||||
async with engine.begin() as connection:
|
||||
await connection.exec_driver_sql('CREATE TABLE example (id INTEGER)')
|
||||
# Force asyncio.Lock itself to bind through contention too.
|
||||
await gate.acquire()
|
||||
waiter = asyncio.create_task(gate.acquire())
|
||||
await queued(gate)
|
||||
gate.lock.release()
|
||||
await waiter
|
||||
gate.lock.release()
|
||||
await engine.dispose()
|
||||
|
||||
async def other_loop():
|
||||
with pytest.raises(RuntimeError, match='another event loop; create a new engine'):
|
||||
async with engine.connect():
|
||||
pass
|
||||
assert not gate.lock.locked()
|
||||
await engine.dispose()
|
||||
|
||||
for scenario in (first_loop, other_loop):
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
loop.run_until_complete(scenario())
|
||||
loop.run_until_complete(loop.shutdown_asyncgens())
|
||||
finally:
|
||||
loop.close()
|
||||
Reference in New Issue
Block a user