cratefield_runtime_native/ports/
realtime.rs1#![allow(clippy::disallowed_types)]
15
16use std::collections::HashMap;
17use std::sync::{Arc, Mutex};
18use std::time::Duration;
19
20use async_trait::async_trait;
21use cratefield_core::{Member, Realtime, RealtimeError, RoomContext, RoomHandler};
22use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
23
24type MemberSender = (String, UnboundedSender<Vec<u8>>);
26
27#[derive(Clone)]
29pub struct InProcessRealtime {
30 handler: Arc<dyn RoomHandler>,
31 rooms: Arc<Mutex<HashMap<String, Room>>>,
32}
33
34#[derive(Default)]
35struct Room {
36 members: HashMap<String, UnboundedSender<Vec<u8>>>,
39 alarm: Option<tokio::task::JoinHandle<()>>,
41}
42
43impl InProcessRealtime {
44 #[must_use]
45 pub fn new(handler: Arc<dyn RoomHandler>) -> Self {
46 Self {
47 handler,
48 rooms: Arc::new(Mutex::new(HashMap::new())),
49 }
50 }
51
52 fn context(&self, room_id: &str) -> NativeContext {
53 NativeContext {
54 rooms: Arc::clone(&self.rooms),
55 handler: Arc::clone(&self.handler),
56 room_id: room_id.to_owned(),
57 }
58 }
59
60 pub async fn connect(
73 &self,
74 room_id: &str,
75 member: Member,
76 ) -> Result<Connection, RealtimeError> {
77 let (tx, rx) = unbounded_channel();
78 {
79 let mut rooms = self.rooms.lock().expect("rooms lock");
80 rooms
81 .entry(room_id.to_owned())
82 .or_default()
83 .members
84 .insert(member.id.clone(), tx);
85 }
86 let ctx = self.context(room_id);
87 self.handler.on_join(&ctx, &member).await?;
88 Ok(Connection {
89 rx,
90 member,
91 room_id: room_id.to_owned(),
92 realtime: self.clone(),
93 })
94 }
95}
96
97pub struct Connection {
100 rx: UnboundedReceiver<Vec<u8>>,
101 member: Member,
102 room_id: String,
103 realtime: InProcessRealtime,
104}
105
106impl Connection {
107 pub async fn next_outgoing(&mut self) -> Option<Vec<u8>> {
110 self.rx.recv().await
111 }
112
113 pub async fn deliver(&self, message: &[u8]) -> Result<(), RealtimeError> {
124 let present = {
125 let rooms = self.realtime.rooms.lock().expect("rooms lock");
126 rooms
127 .get(&self.room_id)
128 .is_some_and(|room| room.members.contains_key(&self.member.id))
129 };
130 if !present {
131 return Err(RealtimeError::NotAMember);
132 }
133 let ctx = self.realtime.context(&self.room_id);
134 self.realtime
135 .handler
136 .on_message(&ctx, &self.member, message)
137 .await
138 }
139
140 pub async fn close(self) -> Result<(), RealtimeError> {
151 {
152 let mut rooms = self.realtime.rooms.lock().expect("rooms lock");
153 if let Some(room) = rooms.get_mut(&self.room_id) {
154 room.members.remove(&self.member.id);
155 if room.members.is_empty() {
156 if let Some(alarm) = room.alarm.take() {
157 alarm.abort();
158 }
159 rooms.remove(&self.room_id);
160 }
161 }
162 }
163 let ctx = self.realtime.context(&self.room_id);
164 self.realtime.handler.on_leave(&ctx, &self.member).await
165 }
166}
167
168struct NativeContext {
170 rooms: Arc<Mutex<HashMap<String, Room>>>,
171 handler: Arc<dyn RoomHandler>,
172 room_id: String,
173}
174
175impl NativeContext {
176 fn senders(&self) -> Vec<MemberSender> {
180 let rooms = self.rooms.lock().expect("rooms lock");
181 rooms
182 .get(&self.room_id)
183 .map(|room| {
184 room.members
185 .iter()
186 .map(|(id, tx)| (id.clone(), tx.clone()))
187 .collect()
188 })
189 .unwrap_or_default()
190 }
191}
192
193#[async_trait]
194impl RoomContext for NativeContext {
195 fn room_id(&self) -> &str {
196 &self.room_id
197 }
198
199 async fn members(&self) -> Vec<Member> {
200 self.senders()
201 .into_iter()
202 .map(|(id, _)| Member::new(id))
203 .collect()
204 }
205
206 async fn broadcast(&self, message: &[u8]) -> Result<(), RealtimeError> {
207 for (_, tx) in self.senders() {
208 let _ = tx.send(message.to_vec());
210 }
211 Ok(())
212 }
213
214 async fn send(&self, member_id: &str, message: &[u8]) -> Result<(), RealtimeError> {
215 let sender = self
216 .senders()
217 .into_iter()
218 .find(|(id, _)| id == member_id)
219 .map(|(_, tx)| tx);
220 if let Some(tx) = sender {
221 let _ = tx.send(message.to_vec());
222 }
223 Ok(())
224 }
225
226 async fn set_alarm(&self, delay: Duration) -> Result<(), RealtimeError> {
227 let rooms = Arc::clone(&self.rooms);
228 let handler = Arc::clone(&self.handler);
229 let room_id = self.room_id.clone();
230 let handle = tokio::spawn(async move {
231 tokio::time::sleep(delay).await;
232 let ctx = NativeContext {
233 rooms,
234 handler: Arc::clone(&handler),
235 room_id,
236 };
237 let _ = handler.on_alarm(&ctx).await;
238 });
239 let mut rooms = self.rooms.lock().expect("rooms lock");
240 if let Some(room) = rooms.get_mut(&self.room_id) {
241 if let Some(previous) = room.alarm.replace(handle) {
242 previous.abort();
243 }
244 } else {
245 handle.abort();
247 }
248 Ok(())
249 }
250}
251
252#[async_trait]
253impl Realtime for InProcessRealtime {
254 async fn broadcast(&self, room_id: &str, message: &[u8]) -> Result<(), RealtimeError> {
255 self.context(room_id).broadcast(message).await
256 }
257
258 async fn members(&self, room_id: &str) -> Result<Vec<Member>, RealtimeError> {
259 Ok(self.context(room_id).members().await)
260 }
261}
262
263#[cfg(test)]
264mod tests {
265 use super::*;
266 use std::sync::atomic::{AtomicUsize, Ordering};
267
268 #[derive(Default)]
271 struct Recorder {
272 joins: AtomicUsize,
273 leaves: AtomicUsize,
274 alarms: AtomicUsize,
275 }
276
277 #[async_trait]
278 impl RoomHandler for Recorder {
279 async fn on_join(
280 &self,
281 ctx: &dyn RoomContext,
282 member: &Member,
283 ) -> Result<(), RealtimeError> {
284 self.joins.fetch_add(1, Ordering::SeqCst);
285 ctx.broadcast(format!("join:{}", member.id).as_bytes())
286 .await
287 }
288
289 async fn on_message(
290 &self,
291 ctx: &dyn RoomContext,
292 member: &Member,
293 message: &[u8],
294 ) -> Result<(), RealtimeError> {
295 let echo = format!("{}:{}", member.id, String::from_utf8_lossy(message));
296 ctx.broadcast(echo.as_bytes()).await
297 }
298
299 async fn on_leave(
300 &self,
301 _ctx: &dyn RoomContext,
302 _member: &Member,
303 ) -> Result<(), RealtimeError> {
304 self.leaves.fetch_add(1, Ordering::SeqCst);
305 Ok(())
306 }
307
308 async fn on_alarm(&self, ctx: &dyn RoomContext) -> Result<(), RealtimeError> {
309 self.alarms.fetch_add(1, Ordering::SeqCst);
310 ctx.broadcast(b"tick").await
311 }
312 }
313
314 fn recorder() -> (InProcessRealtime, Arc<Recorder>) {
315 let recorder = Arc::new(Recorder::default());
316 (InProcessRealtime::new(recorder.clone()), recorder)
317 }
318
319 #[tokio::test]
320 async fn join_broadcasts_to_every_member() {
321 let (rt, rec) = recorder();
322 let mut a = rt.connect("room", Member::new("a")).await.unwrap();
323 let mut b = rt.connect("room", Member::new("b")).await.unwrap();
324
325 rt.broadcast("room", b"hello").await.unwrap();
327 assert_eq!(rt.members("room").await.unwrap().len(), 2);
328 assert_eq!(rec.joins.load(Ordering::SeqCst), 2);
329
330 assert_eq!(a.next_outgoing().await.unwrap(), b"join:a");
332 assert_eq!(a.next_outgoing().await.unwrap(), b"join:b");
333 assert_eq!(a.next_outgoing().await.unwrap(), b"hello");
334 assert_eq!(b.next_outgoing().await.unwrap(), b"join:b");
336 assert_eq!(b.next_outgoing().await.unwrap(), b"hello");
337 }
338
339 #[tokio::test]
340 async fn a_message_reaches_the_other_members() {
341 let (rt, _rec) = recorder();
342 let a = rt.connect("room", Member::new("a")).await.unwrap();
343 let mut b = rt.connect("room", Member::new("b")).await.unwrap();
344 assert_eq!(b.next_outgoing().await.unwrap(), b"join:b");
347
348 a.deliver(b"hi").await.unwrap();
349 assert_eq!(b.next_outgoing().await.unwrap(), b"a:hi");
350 }
351
352 #[tokio::test]
353 async fn leaving_removes_the_member_and_runs_on_leave() {
354 let (rt, rec) = recorder();
355 let a = rt.connect("room", Member::new("a")).await.unwrap();
356 rt.connect("room", Member::new("b")).await.unwrap();
357 a.close().await.unwrap();
358 assert_eq!(rec.leaves.load(Ordering::SeqCst), 1);
359 let ids: Vec<String> = rt
360 .members("room")
361 .await
362 .unwrap()
363 .into_iter()
364 .map(|m| m.id)
365 .collect();
366 assert_eq!(ids, ["b"]);
367 }
368
369 #[tokio::test]
370 async fn a_message_from_a_departed_member_is_refused() {
371 let (rt, _rec) = recorder();
372 rt.connect("room", Member::new("keep")).await.unwrap();
373 let gone = rt.connect("room", Member::new("gone")).await.unwrap();
374 let realtime = rt.clone();
376 let member = gone.member.clone();
377 gone.close().await.unwrap();
378 let orphan = Connection {
379 rx: unbounded_channel().1,
380 member,
381 room_id: "room".to_owned(),
382 realtime,
383 };
384 assert!(matches!(
385 orphan.deliver(b"late").await,
386 Err(RealtimeError::NotAMember)
387 ));
388 }
389
390 #[tokio::test(start_paused = true)]
391 async fn an_alarm_fires_exactly_once() {
392 let (rt, rec) = recorder();
393 let mut a = rt.connect("room", Member::new("a")).await.unwrap();
394 let _ = a.next_outgoing().await; rt.context("room")
398 .set_alarm(Duration::from_millis(50))
399 .await
400 .unwrap();
401 tokio::time::advance(Duration::from_millis(60)).await;
402 tokio::task::yield_now().await;
404
405 assert_eq!(a.next_outgoing().await.unwrap(), b"tick");
406 assert_eq!(rec.alarms.load(Ordering::SeqCst), 1);
407 }
408}