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
8pub struct ChannelContext {
13 pub identifier: String,
15 tx: mpsc::UnboundedSender<String>,
16 pubsub: Arc<dyn PubSub>,
17 streams: Arc<Mutex<Vec<JoinHandle<()>>>>,
19}
20
21impl ChannelContext {
22 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 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 pub fn params(&self) -> serde_json::Value {
52 serde_json::from_str(&self.identifier).unwrap_or(serde_json::Value::Null)
53 }
54
55 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 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 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
99pub 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 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 assert!(rx.try_recv().is_err());
144 }
145}