tortoise / tortoise-orm

Familiar asyncio ORM for python, built with relations in mind
https://tortoise.github.io
Apache License 2.0
4.37k stars 355 forks source link

Race condition bug in ConnectionWrapper #1655

Closed luyiming closed 1 week ago

luyiming commented 1 week ago

There is a race condition bug in ConnectionWrapper implementation:

class ConnectionWrapper:
    __slots__ = ("connection", "lock", "client")

    def __init__(self, lock: asyncio.Lock, client: Any) -> None:
        """Wraps the connections with a lock to facilitate safe concurrent access."""
        self.lock: asyncio.Lock = lock
        self.client = client
        self.connection: Any = client._connection

    async def ensure_connection(self) -> None:
        if not self.connection:
            await self.client.create_connection(with_db=True)
            self.connection = self.client._connection

    async def __aenter__(self):
        await self.ensure_connection()
        await self.lock.acquire()
        return self.connection

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        self.lock.release()

We should acquire the lock before ensure_connection:

    async def __aenter__(self):
        await self.lock.acquire()
        await self.ensure_connection()
        return self.connection

In race condition, two concurrent invocation of async with ConnectionWrapper(...) might fail, for example the process of create_connection() in SqliteClient is as follows:

    async def create_connection(self, with_db: bool) -> None:
        if not self._connection:  # pragma: no branch
            self._connection = aiosqlite.connect(self.filename, isolation_level=None)
            self._connection.start()
            await self._connection._connect()
            self._connection._conn.row_factory = sqlite3.Row
            for pragma, val in self.pragmas.items():
                cursor = await self._connection.execute(f"PRAGMA {pragma}={val}")
                await cursor.close()
            self.log.debug(
                "Created connection %s with params: filename=%s %s",
                self._connection,
                self.filename,
                " ".join(f"{k}={v}" for k, v in self.pragmas.items()),
            )

You can see that self._connection = aiosqlite.connect(self.filename, isolation_level=None) is given a value first, then it yields control when trying to connect to the database await self._connection._connect(), in the meantime, if a second task trying to acquire connection, it finds that self._connection already has a value (but it is not connected yet!), the second task passes self.ensure_connection() and acquire the lock! However, self.connection is not connected yet, so I got the following error reports:

  File "/home/ubuntu/miniconda3/envs/crptrader-dev/lib/python3.11/site-packages/tortoise/models.py", line 957, in save
    await executor.execute_update(self, update_fields)
  File "/home/ubuntu/miniconda3/envs/crptrader-dev/lib/python3.11/site-packages/tortoise/backends/base/executor.py", line 317, in execute_update
    await self.db.execute_query(
  File "/home/ubuntu/miniconda3/envs/crptrader-dev/lib/python3.11/site-packages/tortoise/backends/sqlite/client.py", line 34, in translate_exceptions_
    return await func(self, query, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/miniconda3/envs/crptrader-dev/lib/python3.11/site-packages/tortoise/backends/sqlite/client.py", line 142, in execute_query
    start = connection.total_changes
            ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/miniconda3/envs/crptrader-dev/lib/python3.11/site-packages/aiosqlite/core.py", line 286, in total_changes
    return self._conn.total_changes
           ^^^^^^^^^^
  File "/home/ubuntu/miniconda3/envs/crptrader-dev/lib/python3.11/site-packages/aiosqlite/core.py", line 67, in _conn
    raise ValueError("no active connection")
ValueError: no active connection
abondar commented 1 week ago

Hi!

Thank you for your find

Would you like to prepare PR for that fix? I'll gladly review it