fix: incorrect concurrent usage of connection and transaction

[zevisert]

Note

EDIT: This PR contains some back and forth on how databases should handle concurrent connections across tasks. If you're looking into databases concurrency model, it's worth taking a read through all of our comments. The specific issue with the @Database.transaction decorator was fixed, but it required digging into and changing how connection state is tracked across asyncio.Tasks using asyncio.ContextVars

---

This commit should fix an issue with shared state with transactions.

The main issue is very well described by @circlingthesun here: #123 (comment) - we also ran into this most reliably using locust to load test as well, but did find this issue a few times in normal usage.

The big problem is that using the @db.transaction decorator creates a new Databases.Transaction instance that, when called creates a context manager for a transaction. The problem is that in concurrent frameworks, such as ASGI servers like starlette / fastapi, it's possible to have two concurrent tasks running. Since the Transaction instance is shared between tasks, it's self variable is comparable to shared global state, concurrent code can't trust that no race condition has occurred without some helper mechanism like a ContextVar.

Here's a mermaid diagram trying to show this:

sequenceDiagram
    actor User1
    actor User2
    participant API
    participant TXN
    API->>TXN: (startup) @db.transaction creates new shared Transaction(...) for route
    User1->>API: (1): Request to API route creates new asyncio.Task(...)
    API->>TXN: (1): Transaction.__call__ assignes self._connection
    User2->>API: (2): Request to same API route creates new asyncio.Task(...)
    API->>TXN: (2): PROBLEM -- Transaction.__call__ obtains new Databases.Connection for this asyncio.Task,<br/> like above, but overwrites self._connection
    TXN->>API: (1): PROBLEM -- Transaction.__aexit__ uses self._connection of task (2)<br/>since it was not retrieved from a ContextVar
    API->>User1: (1): Responding here closes the transaction's _connection
    TXN->>API: (2): FAULT -- Transaction.__aexit__ expects to also close self._connection, <br/>but since it also tries to retrieve the connection from self._connection instead of the ContextVar,<br/>it retireves (1)'s already closed connection
    API->>User2: (2): FAULT -- Propigated exception results in 500 response code

Loading

All I've done with this PR is move the shared global state of self._connection and self._transaction into local or contextvars instead of instance variables.

Related issues are: #123 #134
Possibly related: #340
Related meta issue: #456

[zanieb]

Welcome and thanks for contributing! For some context, I'm a member of encode but not a maintainer of this project. Reading this out of curiosity and just have a couple notes about proper usage of the context var. Someone else with more context for the project will have to do a review of the changes too.

[zevisert]

Thanks for the review. I'm eager to move this forward. We use databases in one of our products, and when load testing this certainly fixes an issue with concurrency. Even your two comments were pretty insightful, so if you have another encode member who is interested in taking a look that would be great.

I'll also try the module level context var solution you helped me discover in our codebase to see if that can still fix the issue. As is stands Databases was (and continues) leaking memory from it's usage of non-module level contextvars, but if I can fix that here too then that's great!

[zevisert]

Also related #155

[zanieb]

Do you have a minimal executable example I could use to explore this?

[zevisert]

I think #155 should be close enough, but I'll try and get a dedicated MCVE up here today

[zevisert]

@madkinsz Here's a minimal example as a test - turns out I could get it pretty small. This fails on master, but passes perfectly fine here:

# file: tests/test_concurrent_tasks.py
import pytest
import asyncio
from tests.test_databases import async_adapter, DATABASE_URLS
from databases import Database


@pytest.mark.parametrize("database_url", DATABASE_URLS)
@async_adapter
async def test_concurrent_tasks(database_url: str):
    """
    Test concurrent tasks.
    """

    async with Database(database_url) as database:
        # This could be an eg, fastapi / starlette endpoint
        @database.transaction()
        async def read_items() -> dict:
            return dict(await database.fetch_one(query="SELECT 1 AS value"))

        # Run 10 concurrent **tasks**
        tasks = (asyncio.create_task(read_items()) for _ in range(10))
        responses = await asyncio.gather(*tasks, return_exceptions=True)

        # Check the responses
        assert all(isinstance(response, dict) for response in responses)
        assert all(response["value"] == 1 for response in responses)

I'm using docker compose and pytest to run this

# file: test.docker-compose.yaml
version: "3"
services:
  postgres:
    image: postgres:14
    environment:
      POSTGRES_USER: username
      POSTGRES_PASSWORD: password
      POSTGRES_DB: testsuite
    ports:
      - 5432:5432

  mysql:
    image: mysql:5.7
    environment:
      MYSQL_USER: username
      MYSQL_PASSWORD: password
      MYSQL_ROOT_PASSWORD: password
      MYSQL_DATABASE: testsuite
    ports:
      - 3306:3306

$  docker compose --file test.docker-compose.yaml up --detach 
[+] Running 2/2
 ⠿ Container databases-mysql-1     Started
 ⠿ Container databases-postgres-1  Started   

$ git log -1 --oneline
b6eba5f (HEAD -> master, upstream/master, upstream/HEAD, github/master) Bump wheel from 0.37.1 to 0.38.1 (#524)

$ DATABASE_URLS=(
>     'sqlite:///testsuite'
>     'sqlite+aiosqlite:///testsuite'
>     'mysql://username:password@localhost:3306/testsuite'
>     'mysql+aiomysql://username:password@localhost:3306/testsuite'
>     'mysql+asyncmy://username:password@localhost:3306/testsuite'
>     'postgresql://username:password@localhost:5432/testsuite'
>     'postgresql+aiopg://username:[email protected]:5432/testsuite'
>     'postgresql+asyncpg://username:password@localhost:5432/testsuite'
> )
$ export TEST_DATABASE_URLS=$(IFS=, ; echo "${DATABASE_URLS[*]}")

$ pytest -k test_concurrent_tasks --verbose --exitfirst  # --exitfirst for brevity, they all fail or hang
==================================== test session starts =====================================
platform linux -- Python 3.10.8, pytest-7.1.2, pluggy-1.0.0 -- /opt/tooling/pyenv/versions/encode-databases/bin/python3.10
cachedir: .pytest_cache
rootdir: /code/oss/github/encode/databases
plugins: cov-3.0.0, anyio-3.6.2
collected 360 items / 352 deselected / 8 selected                                            

tests/test_concurrent_fastapi.py::test_concurrent_tasks[sqlite:///testsuite] FAILED    [ 12%]

========================================== FAILURES ==========================================
_________________________ test_concurrent_tasks[sqlite:///testsuite] _________________________

database_url = 'sqlite:///testsuite'

    @pytest.mark.parametrize("database_url", DATABASE_URLS)
    @async_adapter
    async def test_concurrent_tasks(database_url: str):
        """
        Test concurrent tasks.
        """
    
        async with Database(database_url) as database:
            # This could be an eg, fastapi / starlette endpoint
            @database.transaction()
            async def read_items() -> dict:
                return dict(await database.fetch_one(query="SELECT 1 AS value"))
    
            # Run 10 concurrent **tasks**
            tasks = (asyncio.create_task(read_items()) for _ in range(10))
            responses = await asyncio.gather(*tasks, return_exceptions=True)
    
            # Check the responses
E >         assert all(isinstance(response, dict) for response in responses)
E           assert False
E            +  where False = all(<generator object test_concurrent_tasks.<locals>.<genexpr> at 0x7f1b26e86420>)

tests/test_concurrent_fastapi.py:25: AssertionError
================================== short test summary info ===================================
FAILED tests/test_concurrent_fastapi.py::test_concurrent_tasks[sqlite:///testsuite] - asser...
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 1 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
======================== 1 failed, 352 deselected, 1 warning in 0.28s ========================
^CException ignored in: <module 'threading' from '/opt/tooling/pyenv/versions/3.10.8/lib/python3.10/threading.py'>
Traceback (most recent call last):
  File "/opt/tooling/pyenv/versions/3.10.8/lib/python3.10/threading.py", line 1567, in _shutdown
    lock.acquire()
KeyboardInterrupt: 

$ git switch -  
Switched to branch 'fix-transaction-contextvar'
Your branch is up to date with 'github/fix-transaction-contextvar'.

$ git log -1 --oneline
3c0790d (HEAD -> fix-transaction-contextvar, github/fix-transaction-contextvar) fix: contextvar.get takes no keyword arguments

$ pytest -k test_concurrent_tasks --verbose --exitfirst
==================================== test session starts =====================================
platform linux -- Python 3.10.8, pytest-7.1.2, pluggy-1.0.0 -- /opt/tooling/pyenv/versions/encode-databases/bin/python3.10
cachedir: .pytest_cache
rootdir: /code/oss/github/encode/databases
plugins: cov-3.0.0, anyio-3.6.2
collected 360 items / 352 deselected / 8 selected                                            

tests/test_concurrent_fastapi.py::test_concurrent_tasks[sqlite:///testsuite] PASSED    [ 12%]
tests/test_concurrent_fastapi.py::test_concurrent_tasks[sqlite+aiosqlite:///testsuite] PASSED [ 25%]
tests/test_concurrent_fastapi.py::test_concurrent_tasks[mysql://username:password@localhost:3306/testsuite] PASSED [ 37%]
tests/test_concurrent_fastapi.py::test_concurrent_tasks[mysql+aiomysql://username:password@localhost:3306/testsuite] PASSED [ 50%]
tests/test_concurrent_fastapi.py::test_concurrent_tasks[mysql+asyncmy://username:password@localhost:3306/testsuite] PASSED [ 62%]
tests/test_concurrent_fastapi.py::test_concurrent_tasks[postgresql://username:password@localhost:5432/testsuite] PASSED [ 75%]
tests/test_concurrent_fastapi.py::test_concurrent_tasks[postgresql+aiopg://username:[email protected]:5432/testsuite] PASSED [ 87%]
tests/test_concurrent_fastapi.py::test_concurrent_tasks[postgresql+asyncpg://username:password@localhost:5432/testsuite] PASSED [100%]

======================= 8 passed, 352 deselected, 21 warnings in 0.68s =======================

[zevisert]

Rebased to sign commits, all new and old tests passing, switched to module level contextvars to allow for the garbage collector to clean up unreferenced objects.

Edit: And now all linting checks pass too

[zanieb]

Hey! I took some time to poke around at this today. It seems a little wild to use a ContextVar to store a dictionary that includes the asyncio.Task in it (the whole point of them is that they handle task locality for you). I'm not sure what we gain by using a context variable in this case; we can just use an instance dictionary and remove a lot of complexity.

For example, I've implemented this at master...madkinsz:example/instance-safe and it passes the MRE you provided (I only checked SQLite because I didn't want to deal with docker-compose)

Curious for your thoughts

[zevisert]

we can just use an instance dictionary and remove a lot of complexity

I like that take too. Looking back I was a little too committed to getting the ContextVars to work properly. With the constraint of only using context variables in the module scope, and the possibility of multiple database instances with connections to different db's, I needed to introduce the current task back into the context variable to keep things separated correctly. I think that what you're showing with it being stored on the instance and using asyncio.current_task instead of a ContextVar altogether should work just fine!

I'll give it a go with the other databases right now, and in the project were we initially ran into this problem.

[zevisert]

Yeah, this seems great. I've adopted/merged your changes into this MR to accompany the new tests I'm contributing here. I kept a few of the assertions I had added when resolving conflicts - I saw a few of them in the existing code. Let me know what encode's stance is on assert as well as del some_dict[some_key] vs some_dict.pop(some_key) as it seems like we have different styles there too.

Thanks for your help on this, I'm excited for sentry to finally stop notifying us about the errors from this issue!

[zanieb]

Glad you think it makes sense too!

I think my last concern is the event-loop compatibility bit (e.g. trio).

@tomchristie Do you think you'd be willing to give this a look soon or point me to a better reviewer?

[zevisert]

I don't think Databases supports other event loop implementations like trio, curio, etc. I see an old PR #357 that was working on adding anyio support, but it's in draft still.

[zevisert]

Oh and another thing I found just now with the current implementation - context variables undergo a shallow clone when context is copied for new tasks. That means that we had been undoing any isolation that context variables had been providing by mutating the WeakKeyDictionary if it had been created before any descendant tasks started. Simple fix though, just need to treat those dictionaries as immutable.

[zevisert]

Anything more I can do for anyone here to get this reviewed and released?

[zanieb]

I think this has landed in a good place. However, I am not a frequent maintainer of this repository and we'll need someone else to look it over before it can be merged.

[zanieb]

Thanks for your patience!

[zevisert]

Thanks for the merge! Could we get @Kludex or @aminalaee to cut a new release so that we can actually have this bug corrected in our codebases? Is there a changelog or blog article I can help write for this? Is there a release cadence or policy that my team can plan from?

[zanieb]

@zevisert as far as I can tell, I can make a release. I do not think anyone else is available to maintain this project.

I'd definitely appreciate if you opened a pull request for a changelog entry at https://github.com/encode/databases/blob/master/CHANGELOG.md

There is no planned release cadence for this project. I'd like to release this change and some other fixes but there is little active development here.