fix(storage): finish SQLite transactions before propagating cancellation
This commit is contained in:
Eric Wang
@@ -0,0 +1,189 @@
|
||||
"""SQLite 请求取消后，连接必须先退出事务再交给下一次请求。"""
|
||||
|
||||
import asyncio
|
||||
from contextlib import aclosing, suppress
|
||||
|
||||
import aiosqlite
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from src.agent.graph import build_graph
|
||||
from src.backends.chat import ChatInput
|
||||
from src.backends.langgraph import LangGraphBackend
|
||||
from src.services.app_service import AccidentAppService
|
||||
from src.storage.database import SqliteDatabase
|
||||
from src.storage.runtime import StorageSettings, open_storage
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_begin_does_not_leave_transaction(tmp_path, monkeypatch):
|
||||
async with aiosqlite.connect(tmp_path/'begin.sqlite3', isolation_level=None) as raw:
|
||||
database = SqliteDatabase(raw)
|
||||
started = asyncio.Event()
|
||||
execute = raw.execute
|
||||
async def paused_begin(sql, params=()):
|
||||
cursor = await execute(sql, params)
|
||||
if sql == 'BEGIN IMMEDIATE':
|
||||
started.set()
|
||||
await asyncio.Future()
|
||||
return cursor
|
||||
monkeypatch.setattr(raw, 'execute', paused_begin)
|
||||
async def request():
|
||||
async with database.transaction():
|
||||
pass
|
||||
task = asyncio.create_task(request())
|
||||
await asyncio.wait_for(started.wait(), 3)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
monkeypatch.setattr(raw, 'execute', execute)
|
||||
assert not raw.in_transaction
|
||||
async with database.transaction() as connection:
|
||||
await connection.execute('CREATE TABLE recovered (value TEXT)')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_asgi_cancel_scope_does_not_interrupt_rollback(tmp_path):
|
||||
async with aiosqlite.connect(tmp_path/'scope.sqlite3', isolation_level=None) as raw:
|
||||
database = SqliteDatabase(raw)
|
||||
with anyio.CancelScope() as scope:
|
||||
async with database.transaction() as connection:
|
||||
await connection.execute('CREATE TABLE uncommitted (value TEXT)')
|
||||
scope.cancel()
|
||||
await anyio.sleep(0)
|
||||
# 用同一 worker 队列作为屏障，确认已排队的 SQL 已完成。
|
||||
await raw.execute('SELECT 1')
|
||||
assert not raw.in_transaction
|
||||
async with database.transaction() as connection:
|
||||
assert await connection.one("SELECT name FROM sqlite_master WHERE name='uncommitted'") is None
|
||||
async with aiosqlite.connect(tmp_path/'scope.sqlite3', isolation_level=None, timeout=0.1) as other:
|
||||
await other.execute('CREATE TABLE unlocked (value TEXT)')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_repeated_task_cancel_waits_for_rollback_before_unlock(tmp_path, monkeypatch):
|
||||
async with aiosqlite.connect(tmp_path/'repeat.sqlite3', isolation_level=None) as raw:
|
||||
database = SqliteDatabase(raw)
|
||||
entered = asyncio.Event()
|
||||
rolling_back = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
rollback = raw.rollback
|
||||
async def paused_rollback():
|
||||
rolling_back.set()
|
||||
await release.wait()
|
||||
await rollback()
|
||||
monkeypatch.setattr(raw, 'rollback', paused_rollback)
|
||||
async def request():
|
||||
async with database.transaction():
|
||||
entered.set()
|
||||
await asyncio.Future()
|
||||
task = asyncio.create_task(request())
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 3)
|
||||
task.cancel()
|
||||
await asyncio.wait_for(rolling_back.wait(), 3)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert database.lock.locked()
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert not raw.in_transaction
|
||||
assert not database.lock.locked()
|
||||
finally:
|
||||
release.set()
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('operation', ['put', 'writes', 'delete'])
|
||||
async def test_checkpoint_cancel_finishes_transaction_before_releasing_lock(tmp_path, monkeypatch, operation):
|
||||
from langgraph.checkpoint.base import empty_checkpoint
|
||||
|
||||
settings = StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'checkpoint.sqlite3'), auto_migrate=True)
|
||||
async with open_storage(settings) as storage:
|
||||
saver = storage.checkpointer
|
||||
config = {'configurable': {'thread_id': 'cancel-check', 'checkpoint_ns': ''}}
|
||||
config = await saver.aput(config, empty_checkpoint(), {}, {})
|
||||
committing = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
commit = saver.conn.commit
|
||||
async def paused_commit():
|
||||
committing.set()
|
||||
await release.wait()
|
||||
await commit()
|
||||
monkeypatch.setattr(saver.conn, 'commit', paused_commit)
|
||||
if operation == 'put':
|
||||
call = saver.aput(config, empty_checkpoint(), {}, {})
|
||||
elif operation == 'writes':
|
||||
call = saver.aput_writes(config, [('channel', 'value')], 'task')
|
||||
else:
|
||||
call = saver.adelete_thread('cancel-check')
|
||||
task = asyncio.create_task(call)
|
||||
try:
|
||||
await asyncio.wait_for(committing.wait(), 3)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert not saver.conn.in_transaction
|
||||
assert not saver.lock.locked()
|
||||
# 业务连接可以立即写入同一个数据库。
|
||||
await storage.repository.apply_update('next-request', {'hphm1': 'test'})
|
||||
finally:
|
||||
release.set()
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_write_failure_rolls_back_before_business_write(tmp_path, monkeypatch):
|
||||
from langgraph.checkpoint.base import empty_checkpoint
|
||||
|
||||
settings = StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'failure.sqlite3'), auto_migrate=True)
|
||||
async with open_storage(settings) as storage:
|
||||
saver = storage.checkpointer
|
||||
config = {'configurable': {'thread_id': 'failed-write', 'checkpoint_ns': ''}}
|
||||
commit = saver.conn.commit
|
||||
async def fail_commit():
|
||||
raise OSError('commit unavailable')
|
||||
monkeypatch.setattr(saver.conn, 'commit', fail_commit)
|
||||
with pytest.raises(OSError):
|
||||
await saver.aput(config, empty_checkpoint(), {}, {})
|
||||
assert not saver.conn.in_transaction
|
||||
monkeypatch.setattr(saver.conn, 'commit', commit)
|
||||
await storage.repository.apply_update('healthy', {'hphm1': 'test'})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_asgi_stream_cancel_keeps_snapshot_and_next_requests_work(tmp_path):
|
||||
class LLM:
|
||||
async def complete(self, *args, **kwargs):
|
||||
return '{}'
|
||||
async def stream_tokens(self, *args, **kwargs):
|
||||
yield '<state>1002</state>partial'
|
||||
await asyncio.Future()
|
||||
|
||||
settings = StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'stream.sqlite3'), auto_migrate=True)
|
||||
async with open_storage(settings) as storage:
|
||||
backend = LangGraphBackend(build_graph(storage.checkpointer), llm=LLM())
|
||||
service = AccidentAppService(backend, storage.repository, locks=storage.locks)
|
||||
await service.chat(ChatInput('session', 'start'))
|
||||
with anyio.CancelScope() as scope:
|
||||
async with aclosing(service.stream(ChatInput('session', 'cancel-me'))) as events:
|
||||
await anext(events)
|
||||
scope.cancel()
|
||||
await anyio.sleep(0)
|
||||
assert not storage.repository.database.connection.in_transaction
|
||||
assert not storage.checkpointer.conn.in_transaction
|
||||
async with storage.repository.database.transaction() as conn:
|
||||
row = await conn.one('SELECT version FROM agent_sessions WHERE session_id=?', ('session',))
|
||||
assert row['version'] == 1
|
||||
assert not service._locks._entries
|
||||
assert await service.set_info('other', 'hphm1', 'working') == ('200', '')
|
||||
await service.delete_session('session', backend.reset_checkpoint)
|
||||
await service.delete_session('other', backend.reset_checkpoint)
|
||||
Reference in New Issue
Block a user