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::{DecodedFrame, 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, routed: false }` 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            match self.sock.read(&mut chunk) {
35                Ok(0) => {
36                    if self.cursor < self.buf.len() {
37                        return Some(Err(ReplicaError::Truncated));
38                    }
39                    return None;
40                }
41                Ok(n) => self.buf.extend_from_slice(&chunk[..n]),
42                Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
43                Err(e) => return Some(Err(ReplicaError::Io(e))),
44            }
45        }
46    }
47
48    /// Try to decode one event from the buffered bytes. Returns
49    /// `None` when more bytes are needed (the loop in [`Self::next_event`]
50    /// will read more). Split out so the outer loop stays tiny + the
51    /// per-event dispatch fits the project's 50-LOC-fn rule.
52    fn try_decode_one_event(&mut self) -> Option<Result<ReplicaEvent, ReplicaError>> {
53        if self.cursor >= self.buf.len() {
54            return None;
55        }
56        let first = self.buf[self.cursor];
57        match first {
58            b'+' => self.try_decode_snapshot_marker(),
59            b'$' if self.in_snapshot => self.try_decode_snapshot_chunk(),
60            b'*' if self.in_snapshot => Some(Err(ReplicaError::UnexpectedInSnapshot)),
61            b'*' => self.try_decode_live_frame(),
62            _ => Some(Err(ReplicaError::Frame(WireError::BadEnvelope))),
63        }
64    }
65
66    fn try_decode_live_frame(&mut self) -> Option<Result<ReplicaEvent, ReplicaError>> {
67        match decode_frame(&self.buf[self.cursor..]) {
68            Ok((offset, argv, used)) => {
69                self.cursor += used;
70                self.maybe_compact_buf();
71                if offset != self.expected_offset {
72                    return Some(Err(ReplicaError::OffsetGap {
73                        expected: self.expected_offset,
74                        got: offset,
75                    }));
76                }
77                self.expected_offset = self.expected_offset.saturating_add(1);
78                Some(Ok(ReplicaEvent::Frame(DecodedFrame { offset, argv })))
79            }
80            Err(WireError::Truncated) => None,
81            Err(e) => Some(Err(ReplicaError::Frame(e))),
82        }
83    }
84
85    fn try_decode_snapshot_marker(&mut self) -> Option<Result<ReplicaEvent, ReplicaError>> {
86        match decode_snapshot_marker(&self.buf[self.cursor..]) {
87            Ok(Some((SnapshotMarker::Begin, used))) => {
88                self.cursor += used;
89                self.maybe_compact_buf();
90                self.in_snapshot = true;
91                Some(Ok(ReplicaEvent::SnapshotBegin))
92            }
93            Ok(Some((SnapshotMarker::Ping { generation, next_offset }, used))) => {
94                self.cursor += used;
95                self.maybe_compact_buf();
96                Some(Ok(ReplicaEvent::Ping { generation, primary_offset: next_offset }))
97            }
98            Ok(Some((SnapshotMarker::End(ack_offset), used))) => {
99                self.cursor += used;
100                self.maybe_compact_buf();
101                self.in_snapshot = false;
102                self.expected_offset = ack_offset;
103                Some(Ok(ReplicaEvent::SnapshotEnd { ack_offset }))
104            }
105            Ok(None) => Some(Err(ReplicaError::Frame(WireError::BadEnvelope))),
106            Err(WireError::Truncated) => None,
107            Err(e) => Some(Err(ReplicaError::Frame(e))),
108        }
109    }
110
111    fn try_decode_snapshot_chunk(&mut self) -> Option<Result<ReplicaEvent, ReplicaError>> {
112        match decode_snapshot_chunk(&self.buf[self.cursor..]) {
113            Ok((chunk, used)) => {
114                let owned = chunk.to_vec();
115                self.cursor += used;
116                self.maybe_compact_buf();
117                Some(Ok(ReplicaEvent::SnapshotChunk(owned)))
118            }
119            Err(WireError::Truncated) => None,
120            Err(e) => Some(Err(ReplicaError::Frame(e))),
121        }
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use crate::replica::{ReplicaClient, ReplicaError, ReplicaEvent};
128    use crate::wire::{
129        encode_frame, encode_snapshot_begin, encode_snapshot_chunk, encode_snapshot_end,
130    };
131    use kevy_resp::Argv;
132    use std::io::Write;
133    use std::net::{TcpListener, TcpStream};
134    use std::thread;
135
136    fn tcp_pair() -> (TcpStream, TcpStream) {
137        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
138        let addr = listener.local_addr().unwrap();
139        let client = TcpStream::connect(addr).unwrap();
140        let (server, _) = listener.accept().unwrap();
141        (server, client)
142    }
143
144    fn argv_for(args: &[&[u8]]) -> Argv {
145        let mut a = Argv::default();
146        for arg in args {
147            a.push(arg);
148        }
149        a
150    }
151
152    #[test]
153    fn next_event_snapshot_path_begin_chunks_end_then_frame() {
154        let (mut srv, cli) = tcp_pair();
155        thread::spawn(move || {
156            srv.write_all(&encode_snapshot_begin()).unwrap();
157            srv.write_all(&encode_snapshot_chunk(b"hello-snapshot")).unwrap();
158            srv.write_all(&encode_snapshot_chunk(b"more-snapshot-bytes")).unwrap();
159            srv.write_all(&encode_snapshot_end(42)).unwrap();
160            srv.write_all(&encode_frame(42, &argv_for(&[b"SET", b"k", b"v"]))).unwrap();
161            std::thread::sleep(std::time::Duration::from_millis(50));
162            drop(srv);
163        });
164        let mut client = ReplicaClient::from_socket_for_test(cli, 0);
165
166        assert!(matches!(client.next_event(), Some(Ok(ReplicaEvent::SnapshotBegin))));
167        match client.next_event() {
168            Some(Ok(ReplicaEvent::SnapshotChunk(bytes))) => {
169                assert_eq!(bytes, b"hello-snapshot");
170            }
171            other => panic!("expected SnapshotChunk, got {other:?}"),
172        }
173        match client.next_event() {
174            Some(Ok(ReplicaEvent::SnapshotChunk(bytes))) => {
175                assert_eq!(bytes, b"more-snapshot-bytes");
176            }
177            other => panic!("expected SnapshotChunk, got {other:?}"),
178        }
179        match client.next_event() {
180            Some(Ok(ReplicaEvent::SnapshotEnd { ack_offset })) => assert_eq!(ack_offset, 42),
181            other => panic!("expected SnapshotEnd, got {other:?}"),
182        }
183        assert_eq!(client.expected_offset(), 42);
184        match client.next_event() {
185            Some(Ok(ReplicaEvent::Frame(f))) => {
186                assert_eq!(f.offset, 42);
187                assert_eq!(f.argv, argv_for(&[b"SET", b"k", b"v"]));
188            }
189            other => panic!("expected Frame, got {other:?}"),
190        }
191    }
192
193    #[test]
194    fn next_event_live_frame_during_snapshot_is_unexpected() {
195        let (mut srv, cli) = tcp_pair();
196        thread::spawn(move || {
197            srv.write_all(&encode_snapshot_begin()).unwrap();
198            srv.write_all(&encode_snapshot_chunk(b"first")).unwrap();
199            srv.write_all(&encode_frame(0, &argv_for(&[b"PING"]))).unwrap();
200            std::thread::sleep(std::time::Duration::from_millis(50));
201            drop(srv);
202        });
203        let mut client = ReplicaClient::from_socket_for_test(cli, 0);
204        assert!(matches!(client.next_event(), Some(Ok(ReplicaEvent::SnapshotBegin))));
205        assert!(matches!(client.next_event(), Some(Ok(ReplicaEvent::SnapshotChunk(_)))));
206        assert!(matches!(client.next_event(), Some(Err(ReplicaError::UnexpectedInSnapshot))));
207    }
208
209    #[test]
210    fn next_frame_returns_snapshot_in_progress_when_snapshot_starts() {
211        let (mut srv, cli) = tcp_pair();
212        thread::spawn(move || {
213            srv.write_all(&encode_snapshot_begin()).unwrap();
214            std::thread::sleep(std::time::Duration::from_millis(50));
215            drop(srv);
216        });
217        let mut client = ReplicaClient::from_socket_for_test(cli, 0);
218        assert!(matches!(client.next_frame(), Some(Err(ReplicaError::SnapshotInProgress))));
219    }
220
221    #[test]
222    fn live_frame_path_via_next_event() {
223        let (mut srv, cli) = tcp_pair();
224        thread::spawn(move || {
225            srv.write_all(&encode_frame(0, &argv_for(&[b"SET", b"a", b"1"]))).unwrap();
226            srv.write_all(&encode_frame(1, &argv_for(&[b"SET", b"b", b"2"]))).unwrap();
227            std::thread::sleep(std::time::Duration::from_millis(50));
228            drop(srv);
229        });
230        let mut client = ReplicaClient::from_socket_for_test(cli, 0);
231        for expected_off in 0..2 {
232            match client.next_event() {
233                Some(Ok(ReplicaEvent::Frame(f))) => assert_eq!(f.offset, expected_off),
234                other => panic!("expected Frame {expected_off}, got {other:?}"),
235            }
236        }
237        assert_eq!(client.expected_offset(), 2);
238    }
239
240    #[test]
241    fn snapshot_end_with_zero_offset_handled() {
242        let (mut srv, cli) = tcp_pair();
243        thread::spawn(move || {
244            srv.write_all(&encode_snapshot_begin()).unwrap();
245            srv.write_all(&encode_snapshot_end(0)).unwrap();
246            std::thread::sleep(std::time::Duration::from_millis(50));
247            drop(srv);
248        });
249        let mut client = ReplicaClient::from_socket_for_test(cli, 0);
250        assert!(matches!(client.next_event(), Some(Ok(ReplicaEvent::SnapshotBegin))));
251        match client.next_event() {
252            Some(Ok(ReplicaEvent::SnapshotEnd { ack_offset })) => assert_eq!(ack_offset, 0),
253            other => panic!("expected SnapshotEnd, got {other:?}"),
254        }
255        assert_eq!(client.expected_offset(), 0);
256    }
257}
kevy_replicate/replica_decode.rs

kevy_replicate/
replica_decode.rs
