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.

[provinzkraut]

At first I though this was just a side effect of not having implemented cancelling when a client disconnects (#2496), but on second look, this just seems to be an unexpected behaviour in the with the interplay of route handlers and generators, as @winstxnhdw mentioned.

Even if Litestar were to cancel the route handler task on client disconnect, this wouldn't have the result you want, since the handler task finishes as soon as you return the ServerSentEvent. It's actually the generator inside you want to aclose.

Maybe we could provide some functionality to handle this special case, something like adding Request.listen_for_disconnect that you can await, which would block until the client disconnects. Then you could react to that disconnect independently, and close your resource. The benefit of an approach like this would be that it's generic.