Bug: No cleanup for SSE-endpoints if client tears connection down

[vrslev]

Description

This leads to resource leakage. Possibly relevant to #2958.

URL to code causing the issue

https://github.com/vrslev/litestar-sse-cleanup-issue

MCVE

No response

Steps to reproduce

1. Follow the instructions in README in the repository
2. See that logs do not contain `closed redis client`—which is expected.

Screenshots

No response

Logs

INFO:     Started server process [54771]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
INFO:     127.0.0.1:54353 - "GET /sse HTTP/1.1" 200 OK
initialized redis client
caught an exception from blpop: Cancelled by cancel scope 1041246b0
^CINFO:     Shutting down
INFO:     Waiting for application shutdown.
INFO:     Application shutdown complete.
INFO:     Finished server process [54771]

Litestar Version

2.12.1

Platform

• Linux
• Mac
• Windows
• Other (Please specify in the description above)

---

Note

While we are open for sponsoring on GitHub Sponsors and
OpenCollective, we also utilize Polar.sh to engage in pledge-based sponsorship.

Check out all issues funded or available for funding on our Polar.sh dashboard

• If you would like to see an issue prioritized, make a pledge towards it!
• We receive the pledge once the issue is completed & verified
• This, along with engagement in the community, helps us know which features are a priority to our users.

[vrslev]

Temporary solution is to shield __aexit__ in entered async context manager:

import contextlib
import sys
import typing

import anyio
import litestar
import redis.asyncio
from litestar.response import ServerSentEvent

AsyncContextManagerReturnType = typing.TypeVar("AsyncContextManagerReturnType")


def shield_async_context_manager_aexit(
    async_context_manager: typing.AsyncContextManager[AsyncContextManagerReturnType],
) -> typing.AsyncContextManager[AsyncContextManagerReturnType]:
    @contextlib.asynccontextmanager
    async def _inner() -> typing.AsyncGenerator[AsyncContextManagerReturnType, None]:
        async_context_manager_result: typing.Final = (
            await async_context_manager.__aenter__()
        )
        try:
            yield async_context_manager_result
        finally:
            with anyio.CancelScope(shield=True):
                await async_context_manager.__aexit__(*sys.exc_info())

    return _inner()


@contextlib.asynccontextmanager
async def create_redis_client() -> typing.AsyncGenerator[redis.asyncio.Redis, None]:
    redis_client = redis.asyncio.Redis(host="localhost", port="6379")
    try:
        await redis_client.initialize()
        print("initialized redis client")
        yield redis_client
    finally:
        await redis_client.aclose()
        print("closed redis client")


redis_list_key = "whatever"


async def _iter_sse_session_events_as_str() -> typing.AsyncIterable[str]:
    async with shield_async_context_manager_aexit(
        create_redis_client()
    ) as redis_client:
        while True:
            try:
                # BLPOP blocks redis client until an item in list is available,
                # i. e. you can't do anything with the client while waiting here.
                _list_key, event_content = await redis_client.blpop(redis_list_key)
            except BaseException as exception:
                print("caught an exception from blpop:", exception)
                raise exception

            yield event_content


@litestar.get("/sse")
async def listen_to_sse_session_events() -> ServerSentEvent:
    return ServerSentEvent(_iter_sse_session_events_as_str())


app = litestar.Litestar([listen_to_sse_session_events])

[euri10]

any reason you create the redis connection inside the generator ? that seems highly inneficient, usually that's what lifespan is for

[vrslev]

BLPOP is a blocking operation. To serve to more than one request we would have to create connection-per-request.

[winstxnhdw]

Running into this issue as well

[winstxnhdw]

Anyways, I found out that this is intended and simply a quirk with async generators and async context managers in Python. Nothing you can do about it except using contextlib.aclosing.