cratestack_axum/rpc/
sse_subscription.rs

1//! The SSE binding as an adapter over a [`SubscriptionStream`]
2//! (`docs/design/subscriptions-websocket.md` §6.1, slice S5a): the
3//! generated `GET /rpc/subscribe/{op_id}` handler opens the subscription
4//! through the schema's `SubscriptionDispatcher` and hands the stream here.
5//!
6//! First `event: ready` (`{"id": 1, "ready": {}}`, design §7.1): the
7//! subscription is live, and nothing committed after this point will be
8//! missed (§2.1). Then [`super::sse`]'s frames: one `event: message` per item
9//! (`{"id": <n>, "next": <item>}`), and the stream's terminal error as one
10//! `event: error` (`{"id": <n>, "err": RpcErrorBody}`). A model feed's
11//! overflow arrives as that error (`unavailable` / `"subscription lagged"`).
12//! A stream that simply ends sends nothing more. While no item is due, a
13//! `: keepalive` comment goes out every 15 s (§7.1), so proxies and load
14//! balancers that drop idle connections leave a quiet subscription open;
15//! `EventSource` ignores comments.
16
17use axum::body::{Body, Bytes};
18use axum::response::Response;
19use cratestack_core::SubscriptionStream;
20use cratestack_core::rpc::RpcErrorBody;
21use futures_util::future::{Either, select};
22use futures_util::stream::{self, StreamExt};
23
24use super::sse::{format_error_event, format_message_event, sse_response};
25
26/// How long a subscription may stay quiet before a keepalive comment.
27pub(crate) const KEEPALIVE: std::time::Duration = std::time::Duration::from_secs(15);
28
29/// Encode a subscription's items as a `text/event-stream` response.
30pub fn encode_subscription_sse_response(items: SubscriptionStream<serde_json::Value>) -> Response {
31    let ready = stream::once(async {
32        Ok::<_, std::convert::Infallible>(Bytes::from_static(
33            b"event: ready\ndata: {\"id\":1,\"ready\":{}}\n\n",
34        ))
35    });
36    let items = stream::unfold((items, 1u64, false), |(mut items, id, ended)| async move {
37        if ended {
38            return None;
39        }
40        let next = match select(items.next(), Box::pin(tokio::time::sleep(KEEPALIVE))).await {
41            Either::Left((next, _)) => next?,
42            Either::Right(((), _)) => {
43                let keepalive = Bytes::from_static(b": keepalive\n\n");
44                return Some((Ok(keepalive), (items, id, false)));
45            }
46        };
47        let id = id + 1;
48        let (bytes, ended): (Bytes, bool) = match next {
49            Ok(item) => (format_message_event(id, &item), false),
50            Err(error) => (
51                format_error_event(id, &RpcErrorBody::from_cratestack(&error)),
52                true,
53            ),
54        };
55        Some((Ok::<_, std::convert::Infallible>(bytes), (items, id, ended)))
56    });
57    sse_response(Body::from_stream(ready.chain(items)))
58}
59
60#[cfg(test)]
61mod tests;
cratestack_axum/rpc/sse_subscription.rs

cratestack_axum/rpc/
sse_subscription.rs
