pulsify_sqlite/
hub.rs

1use crate::{relay_frame, validate_doc_id, Error, Result, SqliteStore};
2use std::collections::BTreeMap;
3
4/// Delivers frames to one session. Returns `false` once the peer is gone.
5pub trait Sink {
6    fn deliver(&self, session: u64, frame: Vec<u8>) -> bool;
7}
8impl<F: Fn(u64, Vec<u8>) -> bool> Sink for F {
9    fn deliver(&self, session: u64, frame: Vec<u8>) -> bool {
10        self(session, frame)
11    }
12}
13
14struct Session<S> {
15    doc_id: String,
16    owner: String,
17    sink: S,
18}
19
20/// One writer and its sessions: each session is pinned to a document, and rows
21/// stored through one session are forwarded to the document's other sessions.
22pub struct Hub<S> {
23    store: SqliteStore,
24    sessions: BTreeMap<u64, Session<S>>,
25    next_session: u64,
26}
27
28impl<S: Sink> Hub<S> {
29    pub fn new(store: SqliteStore) -> Self {
30        Self {
31            store,
32            sessions: BTreeMap::new(),
33            next_session: 1,
34        }
35    }
36    pub fn store(&mut self) -> &mut SqliteStore {
37        &mut self.store
38    }
39    /// `owner` groups sessions for `close_owner`, e.g. a webview label.
40    pub fn open(&mut self, doc_id: &str, owner: &str, sink: S) -> Result<u64> {
41        validate_doc_id(doc_id)?;
42        let session = self.next_session;
43        self.next_session += 1;
44        self.sessions.insert(
45            session,
46            Session {
47                doc_id: doc_id.to_owned(),
48                owner: owner.to_owned(),
49                sink,
50            },
51        );
52        Ok(session)
53    }
54    /// Relays one frame. Responses go to `session`, new rows to the document's other sessions.
55    pub fn send(&mut self, session: u64, bytes: &[u8]) -> Result<()> {
56        let peer = self
57            .sessions
58            .get(&session)
59            .ok_or_else(|| Error::new("CLOSED", "Unknown or closed session"))?;
60        let output = relay_frame(&mut self.store, bytes, Some(&peer.doc_id))?;
61        for frame in output.responses {
62            if !peer.sink.deliver(session, frame) {
63                self.sessions.remove(&session);
64                return Err(Error::new("CLOSED", "Session receiver is gone"));
65            }
66        }
67        if output.broadcast.is_empty() {
68            return Ok(());
69        }
70        let mut gone = Vec::new();
71        for (id, other) in &self.sessions {
72            if *id == session || other.doc_id != output.doc_id {
73                continue;
74            }
75            for frame in &output.broadcast {
76                if !other.sink.deliver(*id, frame.clone()) {
77                    gone.push(*id);
78                    break;
79                }
80            }
81        }
82        for id in gone {
83            self.sessions.remove(&id);
84        }
85        Ok(())
86    }
87    pub fn close(&mut self, session: u64) -> bool {
88        self.sessions.remove(&session).is_some()
89    }
90    /// Closes every session of an owner; returns how many were open.
91    pub fn close_owner(&mut self, owner: &str) -> usize {
92        let before = self.sessions.len();
93        self.sessions.retain(|_, session| session.owner != owner);
94        before - self.sessions.len()
95    }
96    pub fn session_count(&self) -> usize {
97        self.sessions.len()
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use crate::relay::tests::{connect, upload};
105    use crate::store::tests::{next, A};
106    use pulsify_protocol::{decode, Body};
107    use std::cell::RefCell;
108    use std::rc::Rc;
109
110    type Inbox = Rc<RefCell<Vec<(u64, Body)>>>;
111    type TestSink = Box<dyn Fn(u64, Vec<u8>) -> bool>;
112    fn sink(inbox: &Inbox, alive: bool) -> TestSink {
113        let inbox = inbox.clone();
114        Box::new(move |session, frame| {
115            inbox
116                .borrow_mut()
117                .push((session, decode(&frame).unwrap().body));
118            alive
119        })
120    }
121    fn hub() -> Hub<TestSink> {
122        Hub::new(SqliteStore::open_in_memory().unwrap())
123    }
124
125    #[test]
126    fn broadcasts_new_rows_to_other_sessions_of_the_document() {
127        let mut hub = hub();
128        let inbox: Inbox = Rc::default();
129        let first = hub.open("doc", "main", sink(&inbox, true)).unwrap();
130        let second = hub.open("doc", "other", sink(&inbox, true)).unwrap();
131        let unrelated = hub.open("elsewhere", "main", sink(&inbox, true)).unwrap();
132        hub.send(first, &connect(1, vec![])).unwrap();
133        assert_eq!(inbox.borrow().len(), 1);
134        inbox.borrow_mut().clear();
135        hub.send(first, &upload(2, vec![next(A, 1, 1)])).unwrap();
136        let received = inbox.take();
137        assert_eq!(received.len(), 2);
138        assert_eq!(received[0], (first, Body::Ack { vv: vec![(A, 2)] }));
139        assert_eq!(received[1].0, second);
140        assert!(matches!(received[1].1, Body::Update { .. }));
141        hub.send(second, &upload(3, vec![next(A, 1, 1)])).unwrap();
142        assert_eq!(inbox.take().len(), 1, "duplicates are not broadcast");
143        hub.send(unrelated, &connect(4, vec![])).unwrap();
144        let received = inbox.take();
145        let [(session, Body::Error { code, .. })] = &received[..] else {
146            panic!("a session only serves its own document");
147        };
148        assert_eq!((*session, code.as_str()), (unrelated, "INVALID_MESSAGE"));
149    }
150
151    #[test]
152    fn sessions_close_explicitly_by_owner_or_when_their_sink_is_gone() {
153        let mut hub = hub();
154        let inbox: Inbox = Rc::default();
155        let a = hub.open("doc", "main", sink(&inbox, true)).unwrap();
156        let b = hub.open("doc", "main", sink(&inbox, true)).unwrap();
157        let dead = hub.open("doc", "other", sink(&inbox, false)).unwrap();
158        hub.send(a, &upload(1, vec![next(A, 1, 1)])).unwrap();
159        assert_eq!(
160            hub.session_count(),
161            2,
162            "a failed broadcast drops the session"
163        );
164        assert_eq!(
165            hub.send(dead, &connect(2, vec![])).unwrap_err().code,
166            "CLOSED"
167        );
168        assert!(hub.close(b));
169        assert!(!hub.close(b));
170        assert_eq!(hub.send(b, &connect(3, vec![])).unwrap_err().code, "CLOSED");
171        hub.open("doc", "main", sink(&inbox, true)).unwrap();
172        assert_eq!(hub.close_owner("main"), 2);
173        assert_eq!(hub.session_count(), 0);
174        assert_eq!(
175            hub.open("", "main", sink(&inbox, true)).unwrap_err().code,
176            "INVALID_MESSAGE"
177        );
178    }
179
180    #[test]
181    fn invalid_frames_fail_the_call_but_keep_the_session() {
182        let mut hub = hub();
183        let inbox: Inbox = Rc::default();
184        let session = hub.open("doc", "main", sink(&inbox, true)).unwrap();
185        assert_eq!(
186            hub.send(session, &[9, 3, 1]).unwrap_err().code,
187            "UNSUPPORTED_VERSION"
188        );
189        hub.send(session, &connect(1, vec![])).unwrap();
190        assert_eq!(inbox.borrow().len(), 1);
191    }
192}
pulsify_sqlite/hub.rs

pulsify_sqlite/
hub.rs
