Skip to main content

cratestack_axum/rpc/
subscription_bridge.rs

1//! Bridges a `CratestackEventBus`-style push callback into a bounded,
2//! backpressure-aware `Stream` for the SSE encoder ([`super::sse`]). See
3//! `docs/design/rpc-transport.md` §3.4/§3.4a: "bounded per-subscription
4//! send buffer; on overflow, emit Error{unavailable} ... and end the
5//! stream."
6//!
7//! `CratestackEventBus::subscribe` handlers must never block or fail the
8//! emitting transaction just because one particular SSE client is slow
9//! to drain, so the push side here is always non-blocking (`try_send`)
10//! and infallible from the bus's point of view — overflow is signaled to
11//! the *consumer* by closing the channel, never back to the caller of
12//! `emit()`. See [`super::sse`]'s module doc for why "the stream ends"
13//! is therefore an unambiguous overflow signal, never confused with an
14//! ordinary client disconnect (which just drops the whole future
15//! instead of closing this channel).
16
17use std::sync::{Arc, Mutex};
18
19use futures_util::Stream;
20use futures_util::stream;
21use tokio::sync::mpsc;
22
23// One SSE subscription's outbox capacity is the backplane's
24// (`SUBSCRIPTION_BUFFER_CAPACITY`) before it's considered lagged. A slow
25// consumer past this many buffered, undelivered events gets disconnected with
26// `Error{unavailable}` (see [`super::sse`]) rather than growing memory
27// unboundedly.
28use cratestack_core::backplane::SUBSCRIPTION_BUFFER_CAPACITY;
29
30/// Handed to one or more `CratestackEventBus::subscribe` callbacks registered
31/// for the same logical subscription (e.g. one per `@@emit`ted operation
32/// on a model). Every clone shares the same underlying sender slot, so
33/// the *first* overflow observed by any of them permanently closes the
34/// channel — subsequent pushes from any clone become silent no-ops.
35pub struct SubscriptionPush<T> {
36    slot: Arc<Mutex<Option<mpsc::Sender<T>>>>,
37}
38
39impl<T> Clone for SubscriptionPush<T> {
40    fn clone(&self) -> Self {
41        Self {
42            slot: Arc::clone(&self.slot),
43        }
44    }
45}
46
47impl<T: Send + 'static> SubscriptionPush<T> {
48    /// Non-blocking push. Never fails from the caller's perspective —
49    /// overflow just closes the slot so this and every other clone
50    /// becomes a no-op from then on; the *consumer* observes that as the
51    /// stream ending (see [`guarded_receiver_stream`]).
52    pub fn push(&self, value: T) {
53        let sender = {
54            let guard = self.slot.lock().expect("subscription sender slot poisoned");
55            guard.clone()
56        };
57        let Some(sender) = sender else {
58            return;
59        };
60        if sender.try_send(value).is_err() {
61            // Either full (backpressure) or the receiver already
62            // dropped (client disconnected without ever overflowing —
63            // harmless to also close here, `guarded_receiver_stream`'s
64            // future was already being torn down for that case anyway).
65            *self.slot.lock().expect("subscription sender slot poisoned") = None;
66        }
67    }
68}
69
70/// Builds a fresh bounded channel plus the [`SubscriptionPush`] handle
71/// callers clone into every `CratestackEventBus::subscribe` closure that
72/// should feed it.
73pub fn subscription_channel<T: Send + 'static>() -> (SubscriptionPush<T>, mpsc::Receiver<T>) {
74    let (tx, rx) = mpsc::channel(SUBSCRIPTION_BUFFER_CAPACITY);
75    (
76        SubscriptionPush {
77            slot: Arc::new(Mutex::new(Some(tx))),
78        },
79        rx,
80    )
81}
82
83/// Wraps a raw `mpsc::Receiver` into a `Stream`, keeping `guard` alive
84/// for exactly as long as the stream is — dropped together whether the
85/// stream ends normally (overflow, see [`SubscriptionPush`]) or is
86/// cancelled mid-poll (an ordinary client disconnect just drops this
87/// whole future). This is what lets a `cratestack_core::SubscriptionGuard`
88/// passed as `guard` unsubscribe cleanly in either case without the
89/// caller needing to distinguish which one happened.
90pub fn guarded_receiver_stream<T, G>(
91    rx: mpsc::Receiver<T>,
92    guard: G,
93) -> impl Stream<Item = T> + Send + 'static
94where
95    T: Send + 'static,
96    G: Send + 'static,
97{
98    stream::unfold((rx, guard), |(mut rx, guard)| async move {
99        let item = rx.recv().await?;
100        Some((item, (rx, guard)))
101    })
102}
103
104#[cfg(test)]
105mod tests;