Skip to main content

doido_cable/
channel.rs

1use crate::protocol::ServerMessage;
2use crate::pubsub::PubSub;
3use doido_core::Result;
4use std::sync::Arc;
5use tokio::sync::{mpsc, Mutex};
6use tokio::task::JoinHandle;
7
8/// Per-subscription context handed to a [`Channel`]'s lifecycle callbacks. It
9/// carries the ActionCable `identifier`, an outbound sink to *this* connection,
10/// and the pub/sub backend, so a channel can `transmit`, `stream_from`, read
11/// `params`, and `stop_all_streams` (Rails `ActionCable::Channel`).
12pub struct ChannelContext {
13    /// The raw ActionCable subscription identifier (a JSON string).
14    pub identifier: String,
15    tx: mpsc::UnboundedSender<String>,
16    pubsub: Arc<dyn PubSub>,
17    /// Forwarding tasks for the streams this connection is subscribed to.
18    streams: Arc<Mutex<Vec<JoinHandle<()>>>>,
19}
20
21impl ChannelContext {
22    /// Build a context bound to a connection's outbound sink and pub/sub backend.
23    /// Used by the cable server; apps get one via the channel callbacks.
24    pub(crate) fn new(
25        identifier: String,
26        tx: mpsc::UnboundedSender<String>,
27        pubsub: Arc<dyn PubSub>,
28    ) -> Self {
29        Self {
30            identifier,
31            tx,
32            pubsub,
33            streams: Arc::new(Mutex::new(Vec::new())),
34        }
35    }
36
37    /// A context wired to a fresh channel and in-memory pub/sub, returning the
38    /// receiving end so tests can observe `transmit`/`stream_from` output.
39    pub fn for_test(identifier: impl Into<String>) -> (Self, mpsc::UnboundedReceiver<String>) {
40        let (tx, rx) = mpsc::unbounded_channel();
41        let ctx = Self::new(
42            identifier.into(),
43            tx,
44            Arc::new(crate::pubsub::MemoryPubSub::new()),
45        );
46        (ctx, rx)
47    }
48
49    /// The subscription params: the parsed ActionCable `identifier` JSON (e.g.
50    /// `{ "channel": "ChatChannel", "room": "42" }`), or `Null` if unparseable.
51    pub fn params(&self) -> serde_json::Value {
52        serde_json::from_str(&self.identifier).unwrap_or(serde_json::Value::Null)
53    }
54
55    /// Send a message to just this connection (Rails `transmit`). Wrapped as an
56    /// ActionCable [`ServerMessage`] carrying this subscription's identifier.
57    pub fn transmit(&self, message: serde_json::Value) {
58        let frame = ServerMessage::new(self.identifier.clone(), message);
59        if let Ok(json) = frame.to_json() {
60            let _ = self.tx.send(json);
61        }
62    }
63
64    /// Subscribe this connection to `stream`: anything published to it (e.g. via
65    /// [`Cable::broadcast`]) is forwarded to the client (Rails `stream_from`).
66    ///
67    /// [`Cable::broadcast`]: crate::cable::Cable::broadcast
68    pub async fn stream_from(&self, stream: &str) {
69        if let Ok(mut rx) = self.pubsub.subscribe(stream).await {
70            let tx = self.tx.clone();
71            let handle = tokio::spawn(async move {
72                while let Ok(msg) = rx.recv().await {
73                    if tx.send(msg).is_err() {
74                        break;
75                    }
76                }
77            });
78            self.streams.lock().await.push(handle);
79        }
80    }
81
82    /// Stop forwarding every stream this connection subscribed to (Rails
83    /// `stop_all_streams`).
84    pub async fn stop_all_streams(&self) {
85        let mut streams = self.streams.lock().await;
86        for handle in streams.drain(..) {
87            handle.abort();
88        }
89    }
90}
91
92#[async_trait::async_trait]
93pub trait Channel: Send + Sync {
94    async fn subscribed(&self, ctx: &ChannelContext) -> Result<()>;
95    async fn unsubscribed(&self, ctx: &ChannelContext) -> Result<()>;
96    async fn received(&self, ctx: &ChannelContext, data: serde_json::Value) -> Result<()>;
97}
98
99/// A channel's registration name — its struct name (`ChatChannel`), used to
100/// route ActionCable `identifier` messages to the right channel. Generated by
101/// the `#[channel]` macro.
102pub trait ChannelName {
103    fn channel_name() -> &'static str;
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[tokio::test]
111    async fn params_parses_the_identifier_json() {
112        let (ctx, _rx) = ChannelContext::for_test(r#"{"channel":"ChatChannel","room":"42"}"#);
113        assert_eq!(ctx.params()["channel"], "ChatChannel");
114        assert_eq!(ctx.params()["room"], "42");
115    }
116
117    #[tokio::test]
118    async fn transmit_sends_a_server_message_to_this_connection() {
119        let (ctx, mut rx) = ChannelContext::for_test(r#"{"channel":"ChatChannel"}"#);
120        ctx.transmit(serde_json::json!({ "hello": "world" }));
121        let raw = rx.recv().await.unwrap();
122        let msg = ServerMessage::parse(&raw).unwrap();
123        assert_eq!(msg.identifier, r#"{"channel":"ChatChannel"}"#);
124        assert_eq!(msg.message["hello"], "world");
125    }
126
127    #[tokio::test]
128    async fn stream_from_forwards_published_messages_and_stops() {
129        let pubsub: Arc<dyn PubSub> = Arc::new(crate::pubsub::MemoryPubSub::new());
130        let (tx, mut rx) = mpsc::unbounded_channel();
131        let ctx = ChannelContext::new("id".into(), tx, pubsub.clone());
132
133        ctx.stream_from("room:1").await;
134        // Give the forwarder task a moment to subscribe.
135        tokio::task::yield_now().await;
136        pubsub.publish("room:1", "hi").await.unwrap();
137        assert_eq!(rx.recv().await.unwrap(), "hi");
138
139        ctx.stop_all_streams().await;
140        tokio::task::yield_now().await;
141        pubsub.publish("room:1", "after-stop").await.unwrap();
142        // The forwarder was aborted, so nothing more arrives.
143        assert!(rx.try_recv().is_err());
144    }
145}