pulsify_sqlite/
relay.rs

1use crate::{Error, Result, SqliteStore};
2use pulsify_protocol::{
3    decode, encode, Body, Message, PatchRow, MAX_MESSAGE_ROWS, MAX_PATCH_BYTES, MAX_SCHEMA_NODES,
4    MAX_VECTOR_ENTRIES,
5};
6
7/// The frames produced by one incoming frame.
8#[derive(Debug, Default, PartialEq, Eq)]
9pub struct RelayOutput {
10    pub doc_id: String,
11    /// Frames for the requesting peer, in order.
12    pub responses: Vec<Vec<u8>>,
13    /// Update frames (request ID 0) for the document's other peers, when rows were inserted.
14    pub broadcast: Vec<Vec<u8>>,
15}
16
17/// Mirror of `chunkRows`: about 2 MiB, 4096 rows and the relay node budget per page.
18pub fn chunk_rows(rows: Vec<PatchRow>) -> Vec<Vec<PatchRow>> {
19    let max_nodes = MAX_SCHEMA_NODES - (3 * MAX_VECTOR_ENTRIES + 16);
20    let mut chunks = Vec::new();
21    let mut chunk: Vec<PatchRow> = Vec::new();
22    let (mut bytes, mut nodes) = (0, 0);
23    for row in rows {
24        let size = row.operations.len()
25            + row.dependencies.len() * 16
26            + row.doc_id.encode_utf16().count() * 4
27            + 64;
28        let row_nodes = 7 + 3 * row.dependencies.len();
29        if !chunk.is_empty()
30            && (bytes + size > 2 * MAX_PATCH_BYTES
31                || chunk.len() >= MAX_MESSAGE_ROWS
32                || nodes + row_nodes > max_nodes)
33        {
34            chunks.push(std::mem::take(&mut chunk));
35            bytes = 0;
36            nodes = 0;
37        }
38        chunk.push(row);
39        bytes += size;
40        nodes += row_nodes;
41    }
42    if !chunk.is_empty() || chunks.is_empty() {
43        chunks.push(chunk);
44    }
45    chunks
46}
47
48/// Serve one protocol frame from the store, like `relayStoreMessage`. Rows are
49/// committed before the acknowledgement. A frame which cannot be decoded is an
50/// `Err`; every other failure becomes an error frame for the peer.
51pub fn relay_frame(
52    store: &mut SqliteStore,
53    bytes: &[u8],
54    session_doc: Option<&str>,
55) -> Result<RelayOutput> {
56    let Message {
57        doc_id,
58        request_id,
59        body,
60    } = decode(bytes)?;
61    let mut output = RelayOutput {
62        doc_id,
63        ..RelayOutput::default()
64    };
65    if let Err(error) = serve(store, &mut output, request_id, body, session_doc) {
66        output.broadcast.clear();
67        output.responses = vec![encode(&Message {
68            doc_id: output.doc_id.clone(),
69            request_id,
70            body: Body::Error {
71                code: error.code.to_owned(),
72                message: error.message,
73            },
74        })?];
75    }
76    Ok(output)
77}
78
79fn serve(
80    store: &mut SqliteStore,
81    output: &mut RelayOutput,
82    request_id: u64,
83    body: Body,
84    session_doc: Option<&str>,
85) -> Result<()> {
86    let doc_id = output.doc_id.clone();
87    if session_doc.is_some_and(|pinned| pinned != doc_id) {
88        return Err(Error::new("INVALID_MESSAGE", "One document per session"));
89    }
90    let frame = |request_id: u64, body: Body| {
91        encode(&Message {
92            doc_id: doc_id.clone(),
93            request_id,
94            body,
95        })
96    };
97    let rows = match body {
98        Body::Connect { vv } => {
99            let (rows, vv) = store.sync_state(&doc_id, &vv)?;
100            let chunks = chunk_rows(rows);
101            let last = chunks.len() - 1;
102            for (i, rows) in chunks.into_iter().enumerate() {
103                output.responses.push(frame(
104                    request_id,
105                    Body::Sync {
106                        found: !vv.is_empty(),
107                        complete: i == last,
108                        vv: vv.clone(),
109                        rows,
110                    },
111                )?);
112            }
113            return Ok(());
114        }
115        Body::SyncBack { rows, .. } | Body::Update { rows, .. } => rows,
116        _ => {
117            return Err(Error::new(
118                "INVALID_MESSAGE",
119                "Expected connect, syncBack or update",
120            ))
121        }
122    };
123    let inserted = store.append(&rows)?.inserted;
124    let vv = store.vector(&doc_id)?;
125    output
126        .responses
127        .push(frame(request_id, Body::Ack { vv: vv.clone() })?);
128    if inserted > 0 {
129        for rows in chunk_rows(rows) {
130            let vv = vv.clone();
131            output.broadcast.push(frame(0, Body::Update { vv, rows })?);
132        }
133    }
134    Ok(())
135}
136
137#[cfg(test)]
138pub(crate) mod tests {
139    use super::*;
140    use crate::store::tests::{next, row, A, B};
141    use pulsify_protocol::VersionVector;
142
143    pub fn frame(request_id: u64, body: Body) -> Vec<u8> {
144        encode(&Message {
145            doc_id: "doc".into(),
146            request_id,
147            body,
148        })
149        .unwrap()
150    }
151    pub fn connect(request_id: u64, vv: VersionVector) -> Vec<u8> {
152        frame(request_id, Body::Connect { vv })
153    }
154    pub fn upload(request_id: u64, rows: Vec<PatchRow>) -> Vec<u8> {
155        frame(request_id, Body::SyncBack { vv: vec![], rows })
156    }
157    fn bodies(frames: &[Vec<u8>]) -> Vec<(u64, Body)> {
158        frames
159            .iter()
160            .map(|bytes| {
161                let message = decode(bytes).unwrap();
162                (message.request_id, message.body)
163            })
164            .collect()
165    }
166
167    #[test]
168    fn connect_to_an_empty_document_returns_one_complete_page() {
169        let mut store = SqliteStore::open_in_memory().unwrap();
170        let output = relay_frame(&mut store, &connect(1, vec![]), Some("doc")).unwrap();
171        assert_eq!(output.doc_id, "doc");
172        assert!(output.broadcast.is_empty());
173        assert_eq!(
174            bodies(&output.responses),
175            vec![(
176                1,
177                Body::Sync {
178                    found: false,
179                    complete: true,
180                    vv: vec![],
181                    rows: vec![]
182                }
183            )]
184        );
185    }
186
187    #[test]
188    fn uploads_are_acknowledged_broadcast_and_served() {
189        let mut store = SqliteStore::open_in_memory().unwrap();
190        let rows = vec![next(A, 1, 2), next(A, 3, 1)];
191        let output = relay_frame(&mut store, &upload(2, rows.clone()), None).unwrap();
192        assert_eq!(
193            bodies(&output.responses),
194            vec![(2, Body::Ack { vv: vec![(A, 4)] })]
195        );
196        assert_eq!(
197            bodies(&output.broadcast),
198            vec![(
199                0,
200                Body::Update {
201                    vv: vec![(A, 4)],
202                    rows: rows.clone()
203                }
204            )]
205        );
206        let update = frame(3, Body::Update { vv: vec![], rows });
207        let again = relay_frame(&mut store, &update, None).unwrap();
208        assert_eq!(
209            bodies(&again.responses),
210            vec![(3, Body::Ack { vv: vec![(A, 4)] })]
211        );
212        assert!(again.broadcast.is_empty());
213        let sync = relay_frame(&mut store, &connect(4, vec![(A, 3)]), None).unwrap();
214        assert_eq!(
215            bodies(&sync.responses),
216            vec![(
217                4,
218                Body::Sync {
219                    found: true,
220                    complete: true,
221                    vv: vec![(A, 4)],
222                    rows: vec![next(A, 3, 1)]
223                }
224            )]
225        );
226    }
227
228    #[test]
229    fn pending_rows_are_acknowledged_without_advancing_the_vector() {
230        let mut store = SqliteStore::open_in_memory().unwrap();
231        let output = relay_frame(&mut store, &upload(1, vec![next(A, 2, 1)]), None).unwrap();
232        assert_eq!(
233            bodies(&output.responses),
234            vec![(1, Body::Ack { vv: vec![] })]
235        );
236        assert_eq!(output.broadcast.len(), 1);
237        let output = relay_frame(&mut store, &upload(2, vec![next(A, 1, 1)]), None).unwrap();
238        assert_eq!(
239            bodies(&output.responses),
240            vec![(2, Body::Ack { vv: vec![(A, 3)] })]
241        );
242    }
243
244    #[test]
245    fn failures_become_error_frames() {
246        let mut store = SqliteStore::open_in_memory().unwrap();
247        relay_frame(&mut store, &upload(1, vec![next(A, 1, 2)]), None).unwrap();
248        let mut changed = next(A, 1, 2);
249        changed.operations = vec![7];
250        let cases = [
251            (upload(2, vec![next(B, 1, 1), changed]), None, "CONFLICT"),
252            (connect(3, vec![]), Some("other"), "INVALID_MESSAGE"),
253            (
254                frame(4, Body::Ack { vv: vec![] }),
255                Some("doc"),
256                "INVALID_MESSAGE",
257            ),
258        ];
259        for (bytes, pinned, expected) in cases {
260            let output = relay_frame(&mut store, &bytes, pinned).unwrap();
261            assert!(output.broadcast.is_empty());
262            let [(_, Body::Error { code, .. })] = &bodies(&output.responses)[..] else {
263                panic!("expected one error frame");
264            };
265            assert_eq!(code, expected);
266        }
267        assert_eq!(store.vector("doc").unwrap(), vec![(A, 3)]);
268        let error = relay_frame(&mut store, &[1, 3, 9], None).unwrap_err();
269        assert_eq!(error.code, "INVALID_MESSAGE");
270        let error = relay_frame(&mut store, &[9, 3, 1], None).unwrap_err();
271        assert_eq!(error.code, "UNSUPPORTED_VERSION");
272    }
273
274    #[test]
275    fn large_documents_are_paged() {
276        let mut store = SqliteStore::open_in_memory().unwrap();
277        let mut clock = 1;
278        for _ in 0..3 {
279            let rows: Vec<_> = (0..2)
280                .map(|_| {
281                    let mut segment = next(A, clock, 1);
282                    segment.operations = vec![1; 700 * 1024];
283                    clock += 1;
284                    segment
285                })
286                .collect();
287            relay_frame(&mut store, &upload(1, rows), None).unwrap();
288        }
289        let tiny: Vec<_> = (0..MAX_VECTOR_ENTRIES as u64 - 1)
290            .map(|i| row(B + i, 1, 1, &[]))
291            .collect();
292        relay_frame(&mut store, &upload(2, tiny), None).unwrap();
293        let output = relay_frame(&mut store, &connect(3, vec![]), None).unwrap();
294        let pages = bodies(&output.responses);
295        assert!(pages.len() >= 3);
296        let mut total = 0;
297        for (i, (request_id, body)) in pages.iter().enumerate() {
298            let Body::Sync {
299                found,
300                complete,
301                vv,
302                rows,
303            } = body
304            else {
305                panic!("expected sync pages");
306            };
307            assert_eq!(*request_id, 3);
308            assert!(found);
309            assert_eq!(*complete, i == pages.len() - 1);
310            assert_eq!(vv.len(), MAX_VECTOR_ENTRIES);
311            total += rows.len();
312        }
313        assert_eq!(total, 6 + MAX_VECTOR_ENTRIES - 1);
314    }
315
316    #[test]
317    fn chunks_respect_row_count_and_size_budgets() {
318        let rows: Vec<_> = (1..=MAX_MESSAGE_ROWS as u64 * 2 + 1)
319            .map(|clock| next(A, clock, 1))
320            .collect();
321        let chunks = chunk_rows(rows);
322        assert_eq!(
323            chunks.iter().map(Vec::len).collect::<Vec<_>>(),
324            vec![MAX_MESSAGE_ROWS, MAX_MESSAGE_ROWS, 1]
325        );
326        let mut big = next(A, 1, 1);
327        big.operations = vec![0; MAX_PATCH_BYTES];
328        assert_eq!(chunk_rows(vec![big.clone(), big.clone(), big]).len(), 3);
329        assert_eq!(chunk_rows(vec![]), vec![Vec::<PatchRow>::new()]);
330    }
331}
pulsify_sqlite/relay.rs

pulsify_sqlite/
relay.rs
