Fixed ModuleStorage commit race by serializing concurrent operations with an asyncio lock.
CI / Formatting (push) Failing after 6s
CI / Linting (push) Successful in 5s
CI / Tests (Python 3.12) (push) Successful in 14s
CI / Tests (Python 3.13) (push) Successful in 14s
CI / Tests (Python 3.14) (push) Successful in 12s
CI / Type Checking (push) Successful in 9s
CI / Spelling (push) Successful in 5s

This commit is contained in:
2026-04-17 14:04:15 -04:00
parent f0e86cb4ad
commit 9c86b8bd75
2 changed files with 85 additions and 20 deletions
+23 -20
View File
@@ -66,6 +66,7 @@ class ModuleStorage:
self._module_name = module_name
self._conn: aiosqlite.Connection | None = None
self._closed = False
self._lock = asyncio.Lock()
self._txn_conn: contextvars.ContextVar[aiosqlite.Connection | None] = (
contextvars.ContextVar(f"_txn_conn_{module_name}", default=None)
)
@@ -114,19 +115,20 @@ class ModuleStorage:
raise RuntimeError(
"transaction() cannot be nested. Already inside an active transaction."
)
conn = await self._ensure_connection()
token = self._txn_conn.set(conn)
self._logger.debug("Explicit transaction started.")
try:
yield self
await conn.commit()
self._logger.debug("Transaction committed.")
except BaseException:
await conn.rollback()
self._logger.debug("Transaction rolled back.")
raise
finally:
self._txn_conn.reset(token)
async with self._lock:
conn = await self._ensure_connection()
token = self._txn_conn.set(conn)
self._logger.debug("Explicit transaction started.")
try:
yield self
await conn.commit()
self._logger.debug("Transaction committed.")
except BaseException:
await conn.rollback()
self._logger.debug("Transaction rolled back.")
raise
finally:
self._txn_conn.reset(token)
async def execute(
self,
@@ -273,13 +275,14 @@ class ModuleStorage:
yield existing
return
conn = await self._ensure_connection()
try:
yield conn
await conn.commit()
except BaseException:
await conn.rollback()
raise
async with self._lock:
conn = await self._ensure_connection()
try:
yield conn
await conn.commit()
except BaseException:
await conn.rollback()
raise
async def _close(self) -> None:
"""Close the connection (internal use by bot)."""
+62
View File
@@ -279,6 +279,68 @@ class TestTransaction:
assert row is None
class TestConcurrency:
"""Storage serializes concurrent operations on the shared connection."""
async def test_concurrent_ops_do_not_race(
self, storage_with_table: ModuleStorage
) -> None:
"""Concurrent writers and readers do not trip sqlite's 'in progress' guard.
Without serialization, one task's commit can run while another task's
cursor is still open on the shared connection, raising OperationalError.
"""
await storage_with_table.execute(
"INSERT INTO items (id, name, value) VALUES (?, ?, ?)", (1, "a", 0)
)
async def updater() -> None:
await storage_with_table.execute(
"UPDATE items SET value = 1 WHERE id = 1"
)
async def returner() -> None:
row = await storage_with_table.fetch_one(
"UPDATE items SET value = 2 WHERE id = 1 RETURNING value"
)
assert row is not None
await asyncio.gather(updater(), returner())
async def test_transaction_blocks_concurrent_ops(
self, storage_with_table: ModuleStorage
) -> None:
"""A held transaction serializes non-transactional ops from other tasks."""
gate = asyncio.Event()
released = asyncio.Event()
async def holder() -> None:
async with storage_with_table.transaction():
await storage_with_table.execute(
"INSERT INTO items (id, name, value) VALUES (?, ?, ?)",
(1, "a", 1),
)
gate.set()
await released.wait()
async def waiter() -> None:
await gate.wait()
await storage_with_table.execute(
"INSERT INTO items (id, name, value) VALUES (?, ?, ?)", (2, "b", 2)
)
holder_task = asyncio.create_task(holder())
waiter_task = asyncio.create_task(waiter())
await gate.wait()
await asyncio.sleep(0.01)
assert not waiter_task.done()
released.set()
await asyncio.gather(holder_task, waiter_task)
rows = await storage_with_table.fetch_all("SELECT id FROM items ORDER BY id")
assert [r["id"] for r in rows] == [1, 2]
class TestStorageError:
"""Invalid SQL raises StorageError wrapping the underlying error."""