kevy_replicate/
replica_decode.rs

1//! Snapshot-aware event decoding for [`crate::replica::ReplicaClient`]
2//! — split from `replica.rs` to keep that file under the 500-LOC
3//! project ceiling. The state-machine helpers live here; the type
4//! definitions, `connect`, and `next_frame` stay in `replica.rs`.
5
6use crate::replica::{ReplicaClient, ReplicaError, ReplicaEvent};
7use crate::wire::{
8    SnapshotMarker, WireError, decode_frame, decode_snapshot_chunk, decode_snapshot_marker,
9};
10use std::io::{self, Read};
11
12impl ReplicaClient {
13    /// Snapshot-aware iterator step. Returns one [`ReplicaEvent`] per
14    /// call — a live `Frame`, a `SnapshotBegin`/`SnapshotChunk`/
15    /// `SnapshotEnd`, or one of the [`ReplicaError`] variants.
16    /// Returns `None` on clean peer EOF.
17    ///
18    /// Snapshot bookkeeping:
19    /// - Entering `SnapshotBegin` sets `in_snapshot = true`; chunk
20    ///   bytes are valid until `SnapshotEnd`.
21    /// - `SnapshotEnd { ack_offset }` sets `expected_offset =
22    ///   ack_offset` (so the next live `Frame` has no gap) and
23    ///   clears `in_snapshot`.
24    /// - Live `*2\r\n` bytes during a snapshot return
25    ///   [`ReplicaError::UnexpectedInSnapshot`] (interleaving is
26    ///   forbidden — see `docs/snapshot.md`).
27    pub fn next_event(&mut self) -> Option<Result<ReplicaEvent, ReplicaError>> {
28        loop {
29            if let Some(result) = self.try_decode_one_event() {
30                return Some(result);
31            }
32            // Need more bytes off the socket.
33            let mut chunk = [0u8; 4096];
34            let read = match self.noise.as_mut() {
35                Some(noise) => noise.read(&mut self.sock, &mut chunk, &mut self.buf),
36                None => {
37                    self.sock.read(&mut chunk).inspect(|&n| self.buf.extend_from_slice(&chunk[..n]))
38                }
39            };
40            match read {
41                Ok(0) => {
42                    if self.cursor < self.buf.len() {
43                        return Some(Err(ReplicaError::Truncated));
44                    }
45                    return None;
46                }
47                Ok(_) => {}
48                Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
49                Err(e) => return Some(Err(ReplicaError::Io(e))),
50            }
51        }
52    }
53
54    /// Try to decode one event from the buffered bytes. Returns
55    /// `None` when more bytes are needed (the loop in [`Self::next_event`]
56    /// will read more). Split out so the outer loop stays tiny + the
57    /// per-event dispatch fits the project's 50-LOC-fn rule.
58    fn try_decode_one_event(&mut self) -> Option<Result<ReplicaEvent, ReplicaError>> {
59        if self.cursor >= self.buf.len() {
60            return None;
61        }
62        let first = self.buf[self.cursor];
63        match first {
64            b'+' => self.try_decode_snapshot_marker(),
65            b'$' if self.in_snapshot => self.try_decode_snapshot_chunk(),
66            b'*' if self.in_snapshot => Some(Err(ReplicaError::UnexpectedInSnapshot)),
67            b'*' => self.try_decode_live_frame(),
68            _ => Some(Err(ReplicaError::Frame(WireError::BadEnvelope))),
69        }
70    }
71
72    fn try_decode_live_frame(&mut self) -> Option<Result<ReplicaEvent, ReplicaError>> {
73        match decode_frame(&self.buf[self.cursor..]) {
74            Ok((frame, used)) => {
75                self.cursor += used;
76                self.maybe_compact_buf();
77                if frame.offset != self.expected_offset {
78                    return Some(Err(ReplicaError::OffsetGap {
79                        expected: self.expected_offset,
80                        got: frame.offset,
81                    }));
82                }
83                self.expected_offset = self.expected_offset.saturating_add(1);
84                Some(Ok(ReplicaEvent::Frame(frame)))
85            }
86            Err(WireError::Truncated) => None,
87            Err(e) => Some(Err(ReplicaError::Frame(e))),
88        }
89    }
90
91    fn try_decode_snapshot_marker(&mut self) -> Option<Result<ReplicaEvent, ReplicaError>> {
92        match decode_snapshot_marker(&self.buf[self.cursor..]) {
93            Ok(Some((SnapshotMarker::Begin, used))) => {
94                self.cursor += used;
95                self.maybe_compact_buf();
96                self.in_snapshot = true;
97                Some(Ok(ReplicaEvent::SnapshotBegin))
98            }
99            Ok(Some((SnapshotMarker::Ping(tail), used))) => {
100                self.cursor += used;
101                self.maybe_compact_buf();
102                Some(Ok(ReplicaEvent::Ping(tail)))
103            }
104            Ok(Some((SnapshotMarker::End(ack_offset), used))) => {
105                self.cursor += used;
106                self.maybe_compact_buf();
107                self.in_snapshot = false;
108                self.expected_offset = ack_offset;
109                Some(Ok(ReplicaEvent::SnapshotEnd { ack_offset }))
110            }
111            Ok(None) => Some(Err(ReplicaError::Frame(WireError::BadEnvelope))),
112            Err(WireError::Truncated) => None,
113            Err(e) => Some(Err(ReplicaError::Frame(e))),
114        }
115    }
116
117    fn try_decode_snapshot_chunk(&mut self) -> Option<Result<ReplicaEvent, ReplicaError>> {
118        match decode_snapshot_chunk(&self.buf[self.cursor..]) {
119            Ok((chunk, used)) => {
120                let owned = chunk.to_vec();
121                self.cursor += used;
122                self.maybe_compact_buf();
123                Some(Ok(ReplicaEvent::SnapshotChunk(owned)))
124            }
125            Err(WireError::Truncated) => None,
126            Err(e) => Some(Err(ReplicaError::Frame(e))),
127        }
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use crate::replica::{ReplicaClient, ReplicaError, ReplicaEvent};
134    use crate::wire::{
135        encode_frame, encode_snapshot_begin, encode_snapshot_chunk, encode_snapshot_end,
136    };
137    use kevy_resp::Argv;
138    use std::io::Write;
139    use std::net::{TcpListener, TcpStream};
140    use std::thread;
141
142    fn tcp_pair() -> (TcpStream, TcpStream) {
143        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
144        let addr = listener.local_addr().unwrap();
145        let client = TcpStream::connect(addr).unwrap();
146        let (server, _) = listener.accept().unwrap();
147        (server, client)
148    }
149
150    fn argv_for(args: &[&[u8]]) -> Argv {
151        let mut a = Argv::default();
152        for arg in args {
153            a.push(arg);
154        }
155        a
156    }
157
158    #[test]
159    fn next_event_snapshot_path_begin_chunks_end_then_frame() {
160        let (mut srv, cli) = tcp_pair();
161        thread::spawn(move || {
162            srv.write_all(&encode_snapshot_begin()).unwrap();
163            srv.write_all(&encode_snapshot_chunk(b"hello-snapshot")).unwrap();
164            srv.write_all(&encode_snapshot_chunk(b"more-snapshot-bytes")).unwrap();
165            srv.write_all(&encode_snapshot_end(42)).unwrap();
166            srv.write_all(&encode_frame(42, &argv_for(&[b"SET", b"k", b"v"]))).unwrap();
167            std::thread::sleep(std::time::Duration::from_millis(50));
168            drop(srv);
169        });
170        let mut client = ReplicaClient::from_socket_for_test(cli, 0);
171
172        assert!(matches!(client.next_event(), Some(Ok(ReplicaEvent::SnapshotBegin))));
173        match client.next_event() {
174            Some(Ok(ReplicaEvent::SnapshotChunk(bytes))) => {
175                assert_eq!(bytes, b"hello-snapshot");
176            }
177            other => panic!("expected SnapshotChunk, got {other:?}"),
178        }
179        match client.next_event() {
180            Some(Ok(ReplicaEvent::SnapshotChunk(bytes))) => {
181                assert_eq!(bytes, b"more-snapshot-bytes");
182            }
183            other => panic!("expected SnapshotChunk, got {other:?}"),
184        }
185        match client.next_event() {
186            Some(Ok(ReplicaEvent::SnapshotEnd { ack_offset })) => assert_eq!(ack_offset, 42),
187            other => panic!("expected SnapshotEnd, got {other:?}"),
188        }
189        assert_eq!(client.expected_offset(), 42);
190        match client.next_event() {
191            Some(Ok(ReplicaEvent::Frame(f))) => {
192                assert_eq!(f.offset, 42);
193                assert_eq!(f.argv, argv_for(&[b"SET", b"k", b"v"]));
194            }
195            other => panic!("expected Frame, got {other:?}"),
196        }
197    }
198
199    #[test]
200    fn next_event_live_frame_during_snapshot_is_unexpected() {
201        let (mut srv, cli) = tcp_pair();
202        thread::spawn(move || {
203            srv.write_all(&encode_snapshot_begin()).unwrap();
204            srv.write_all(&encode_snapshot_chunk(b"first")).unwrap();
205            srv.write_all(&encode_frame(0, &argv_for(&[b"PING"]))).unwrap();
206            std::thread::sleep(std::time::Duration::from_millis(50));
207            drop(srv);
208        });
209        let mut client = ReplicaClient::from_socket_for_test(cli, 0);
210        assert!(matches!(client.next_event(), Some(Ok(ReplicaEvent::SnapshotBegin))));
211        assert!(matches!(client.next_event(), Some(Ok(ReplicaEvent::SnapshotChunk(_)))));
212        assert!(matches!(client.next_event(), Some(Err(ReplicaError::UnexpectedInSnapshot))));
213    }
214
215    #[test]
216    fn next_frame_returns_snapshot_in_progress_when_snapshot_starts() {
217        let (mut srv, cli) = tcp_pair();
218        thread::spawn(move || {
219            srv.write_all(&encode_snapshot_begin()).unwrap();
220            std::thread::sleep(std::time::Duration::from_millis(50));
221            drop(srv);
222        });
223        let mut client = ReplicaClient::from_socket_for_test(cli, 0);
224        assert!(matches!(client.next_frame(), Some(Err(ReplicaError::SnapshotInProgress))));
225    }
226
227    #[test]
228    fn live_frame_path_via_next_event() {
229        let (mut srv, cli) = tcp_pair();
230        thread::spawn(move || {
231            srv.write_all(&encode_frame(0, &argv_for(&[b"SET", b"a", b"1"]))).unwrap();
232            srv.write_all(&encode_frame(1, &argv_for(&[b"SET", b"b", b"2"]))).unwrap();
233            std::thread::sleep(std::time::Duration::from_millis(50));
234            drop(srv);
235        });
236        let mut client = ReplicaClient::from_socket_for_test(cli, 0);
237        for expected_off in 0..2 {
238            match client.next_event() {
239                Some(Ok(ReplicaEvent::Frame(f))) => assert_eq!(f.offset, expected_off),
240                other => panic!("expected Frame {expected_off}, got {other:?}"),
241            }
242        }
243        assert_eq!(client.expected_offset(), 2);
244    }
245
246    #[test]
247    fn snapshot_end_with_zero_offset_handled() {
248        let (mut srv, cli) = tcp_pair();
249        thread::spawn(move || {
250            srv.write_all(&encode_snapshot_begin()).unwrap();
251            srv.write_all(&encode_snapshot_end(0)).unwrap();
252            std::thread::sleep(std::time::Duration::from_millis(50));
253            drop(srv);
254        });
255        let mut client = ReplicaClient::from_socket_for_test(cli, 0);
256        assert!(matches!(client.next_event(), Some(Ok(ReplicaEvent::SnapshotBegin))));
257        match client.next_event() {
258            Some(Ok(ReplicaEvent::SnapshotEnd { ack_offset })) => assert_eq!(ack_offset, 0),
259            other => panic!("expected SnapshotEnd, got {other:?}"),
260        }
261        assert_eq!(client.expected_offset(), 0);
262    }
263}
kevy_replicate/replica_decode.rs

kevy_replicate/
replica_decode.rs
