Channels and ServerSentEvents

[aranvir]

Based on my question here, I also looked into using Channels with ServerSentEvents when I just want to push messages from the server to the client and the full-duplex capability of websockets is not needed. Plus, since SSE is covered by a regular http request I don't need to add custom session/cookie handling as I had to for the websockets.

I ran into a similar problem as for the websockets that is: how to safely iterate over the events ;) Using subscriber.iter_events() worked but it would not close when the client connection closed so that was a problem (e.g., shutting down the uvicorn process would actually get stuck and it needed to be killed manually). I had the impression that the lifespan and asynccontextmanager approach would work here (https://github.com/orgs/litestar-org/discussions/3084#discussioncomment-8405508)

So I came up with my own solution with inspiration from a starlette-sse pattern. Since litestar does not have the option to pass a data_sender_callable to the ServerSentEvent response I instead start the pushing loop as a separate task. Since this task is outside the task group that handles the stream, the disconnect is not registered by this task. So instead, I pass a simple future that is set done in a BackgroundTask once the sse client disconnects and the loop should be stopped.

That seemed to work fine (e.g., log messages indicating completion of task), but I'm not sure if there are better or more efficient ways to do this.

import asyncio

import anyio
from anyio.streams.memory import MemoryObjectSendStream
from litestar import Controller, Request, get
from litestar.background_tasks import BackgroundTask
from litestar.response import ServerSentEvent
from litestar.channels import ChannelsPlugin
from litestar.logging.config import LoggingConfig
from litestar.channels.subscriber import Subscriber

from mydemo.channels import ChannelsEnum

logger = LoggingConfig().configure()()


async def event_publisher(inner_send_chan: MemoryObjectSendStream, channels: ChannelsPlugin, subscriptions: list[str], stop: asyncio.Future):
    """Graceful looping with cleanup of messages.

    Inspired by https://github.com/sysid/sse-starlette/pull/56
    """
    async with inner_send_chan:
        async with channels.start_subscription(subscriptions) as subscriber:
            subscriber: Subscriber
            try:
                async with subscriber.run_in_background(inner_send_chan.send):
                    await stop
            except anyio.get_cancelled_exc_class() as e:
                logger.info("Disconnected from client.")
                await subscriber.stop()
                with anyio.move_on_after(1, shield=True):
                    await inner_send_chan.send(dict(closing=True))
                    raise e
    logger.info("Task completed.")


class NotificationsSSE(Controller):
    @get("/sse/updates")
    async def sse_handler(self, request: Request, channels: ChannelsPlugin) -> ServerSentEvent:
        channels_to_subscribe = [ChannelsEnum.public]
        if request.session.get("clearance"):
            channels_to_subscribe.append(ChannelsEnum.secret)
        logger.info(f"Starting sse. Has clearance: {request.session.get('clearance')}")

        stop = asyncio.get_running_loop().create_future()
        send_chan, recv_chan = anyio.create_memory_object_stream(10)
        asyncio.create_task(
            event_publisher(inner_send_chan=send_chan, channels=channels, subscriptions=channels_to_subscribe, stop=stop)
        )

        async def bg():
            stop.set_result(None)
            logger.info("SSE finished!")

        return ServerSentEvent(recv_chan, background=BackgroundTask(bg))

[provinzkraut]

Using subscriber.iter_events() worked but it would not close when the client connection closed

Can you provide an MCVE for this?

[aranvir]

Tried to recreate it but cannot recreate the error reliably with a simple example. I have a SQLAlchemyInitPlugin and SessionAuth in my app, so I thought that there might be some problems. I added the sqlalchemy plugin to the simple example and had the lockup again, but only once. Afterwards, even when deleting the sqlite dummy file it would not happen again.

It also does not happen every time for my main application...

I'll keep an eye out for it if it happens again and if I can trace it down. Below is the example code that I made to reproduce the error.... but as mentioned, does not error reliably.

import anyio
import asyncio
import json
from enum import Enum
from collections.abc import AsyncGenerator

from litestar import Litestar, Request, get, Controller
from litestar.response import ServerSentEvent
from litestar.logging.config import LoggingConfig
from litestar.channels import ChannelsPlugin
from litestar.channels.subscriber import Subscriber
from litestar.channels.backends.memory import MemoryChannelsBackend
from litestar.contrib.sqlalchemy.plugins import AsyncSessionConfig, SQLAlchemyAsyncConfig, SQLAlchemyInitPlugin
from litestar.contrib.sqlalchemy.base import UUIDBase

logger = LoggingConfig().configure()()


session_config = AsyncSessionConfig(expire_on_commit=False)
sqlalchemy_config = SQLAlchemyAsyncConfig(
    connection_string="sqlite+aiosqlite:///test.sqlite", session_config=session_config
)
sqlalchemy_plugin = SQLAlchemyInitPlugin(config=sqlalchemy_config)


async def db_on_startup() -> None:
    """Create all tables."""
    async with sqlalchemy_config.get_engine().begin() as conn:
        await conn.run_sync(UUIDBase.metadata.create_all)


class ChannelsEnum(str, Enum):
    public = "public"
    secret = "secret"


channels_plugin = ChannelsPlugin(
    backend=MemoryChannelsBackend(),
    channels=[ChannelsEnum.public, ChannelsEnum.secret]
)


async def notification_loop(channels: ChannelsPlugin, channels_to_subscribe: list[str]) -> AsyncGenerator:
    async with channels.start_subscription(channels_to_subscribe) as subscriber:
        subscriber: Subscriber
        logger.info("Subscribed! Iterating events")
        try:
            async for message in subscriber.iter_events():
                logger.info("Pushing SSE!")
                yield message
        except Exception as e:
            logger.error(f"Something happened. {e.__class__.__name__}: {e}")
            return
    logger.info("Notification loop ended")


@get("/sse/updates")
async def sse_handler(channels: ChannelsPlugin) -> ServerSentEvent:
    channels_to_subscribe = [ChannelsEnum.public]
    logger.info(f"Starting sse.")
    return ServerSentEvent(notification_loop(channels=channels, channels_to_subscribe=channels_to_subscribe))


async def notification_task():
    for _ in range(5):
        logger.info("publishing message")
        message = {
            "content": "test message",
            "type": "publish"
        }
        channels_plugin.publish(json.dumps(message), channels=[ChannelsEnum.public])
        await asyncio.sleep(1)


@get("/notify")
async def notify() -> None:
    """Starts a notification loop that sends 5 public or secret notifications via websocket."""
    asyncio.create_task(notification_task())


app = Litestar(
    route_handlers=[
        notify,
        sse_handler
    ],
    on_startup=[db_on_startup],
    plugins=[channels_plugin, sqlalchemy_plugin],
    debug=True
)

[provinzkraut]

I don't see anything that looks like it could cause something to hang but a few things:

1. Creating a task with asyncio.create_task is not reliably if you don't keep a reference to this task. It can be garbage collected at any point.
2. Doing a bare try/except is almost always a bad idea, especially within context managers. Here for example you're preventing the outer context manager from handling exceptions properly because it does not know about them.

[aranvir]

Thanks! Then it might just be a freak coincidence.

Both things I only did for "quick&dirty prototyping" 😅 E.g. having some way to start the notification loop and this was the fastest I could come up with. For a proper solution, I'd push the messages via a different approach. And the exception logging: yes, also just so I can see what exceptions would come up without it crashing the application (which probably wouldnt happen anyway because litestar probably has some handling for that)

[provinzkraut]

And the exception logging: yes, also just so I can see what exceptions would come up without it crashing the application (which probably wouldnt happen anyway because litestar probably has some handling for that)

Indeed. Exceptions won't crash your application, you'll get an error response instead.

[provinzkraut]

In general, something like this should work:

import asyncio
from contextlib import asynccontextmanager
from typing import AsyncGenerator

from litestar import Litestar, get
from litestar.channels import ChannelsPlugin
from litestar.channels.backends.memory import MemoryChannelsBackend
from litestar.response import ServerSentEvent


@asynccontextmanager
async def lifespan(app: Litestar) -> AsyncGenerator[None, None]:
    channels = app.plugins.get(ChannelsPlugin)
    stop = asyncio.Event()

    async def background_task():
        while not stop.is_set():
            channels.publish("hello", ["some_channel"])
            await asyncio.sleep(1)

    task = asyncio.create_task(background_task())
    yield
    stop.set()
    await task


@get("/")
async def sse_handler(channels: ChannelsPlugin) -> ServerSentEvent:
    async def generator() -> AsyncGenerator[bytes, None]:
        async with channels.start_subscription(["some_channel"]) as subscriber:
            async for event in subscriber.iter_events():
                yield event

    return ServerSentEvent(generator())


app = Litestar(
    [sse_handler],
    plugins=[ChannelsPlugin(MemoryChannelsBackend(), channels=["some_channel"])],
    lifespan=[lifespan],
)

I tested this locally and closing the client always correctly seems to terminate the connection and server shutdown does not hang for me.

[provinzkraut]

For this solution, does this not create a lot of async generators

No, this only creates one such generator. When the client connects to the sse_handler, it keeps this connection open and sends subsequent responses on there, so the generator() is only invoked once per connection.

When the connection terminates, the generator will be garbage collected eventually.

[aranvir]

Okay, clear! That's a more concise approach then.

[aranvir]

Just for completeness' sake: I used this example but with multiple lifespans (one per channel). During shutdown, a RuntimeError can occur on channels.publish("hello", ["some_channel"]) because the channels plugin is already stopped but the task is still running. I had a blocked shutdown once, but cannot not re-create it (so I doubt it has anything to do with this specifically). Nevertheless, added a try-except to catch this and stop the task properly during shutdown.

    async def background_task():
        while not stop.is_set():
            try:
                channels.publish("hello", ["some_channel"])
            except RuntimeError:
                break
            await asyncio.sleep(1)

[provinzkraut]

Just for completeness' sake: I used this example but with multiple lifespans (one per channel). During shutdown, a RuntimeError can occur on channels.publish("hello", ["some_channel"])

The solution would be to not use multiple lifespans here. Is there a specific reason for doing that?
I don't see what this accomplishes that can't be with a single lifespan manager.

[aranvir]

Was mostly just me having separate publication loops for "public" and "secret" data... and then that error happened and I got curious :) so no definite need for that.