Skip to main content

armdb/fixed_replication/
protocol.rs

1//! Wire protocol for FixedStore replication.
2//!
3//! Frame format: `[type: u8][len: u32 LE][payload: len bytes]`.
4//! All integers are little-endian. See spec §6 for the full message table.
5
6use std::io::{self, Read, Write};
7
8use crate::fixed::slot::{STATUS_DELETED, STATUS_FREE, STATUS_OCCUPIED, status_of};
9
10pub const PROTOCOL_VERSION: u8 = 2;
11pub const BATCH_MAX_ENTRIES: usize = 256;
12pub const BATCH_MAX_BYTES: usize = 64 * 1024;
13pub const MAX_FRAME_PAYLOAD: usize = 256 * 1024;
14pub const TAIL_POLL_MS: u64 = 1;
15pub const HEARTBEAT_INTERVAL_SECS: u64 = 5;
16pub const ACK_INTERVAL: usize = 1000;
17
18/// Flags carried in SyncRequest.
19pub const FLAG_EMPTY_STATE: u8 = 1 << 0;
20
21#[repr(u8)]
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum FixedMessageType {
24    SyncRequest = 1,
25    ShardInfo = 2,
26    SlotBatch = 3,
27    CaughtUp = 4,
28    Ack = 5,
29    Heartbeat = 6,
30    Error = 255,
31}
32
33impl FixedMessageType {
34    pub fn from_u8(v: u8) -> Option<Self> {
35        Some(match v {
36            1 => Self::SyncRequest,
37            2 => Self::ShardInfo,
38            3 => Self::SlotBatch,
39            4 => Self::CaughtUp,
40            5 => Self::Ack,
41            6 => Self::Heartbeat,
42            255 => Self::Error,
43            _ => return None,
44        })
45    }
46}
47
48impl crate::frame_reader::FrameKind for FixedMessageType {
49    const MAX_PAYLOAD: usize = MAX_FRAME_PAYLOAD;
50    fn from_u8(b: u8) -> Option<Self> {
51        FixedMessageType::from_u8(b)
52    }
53}
54
55pub type Frame = crate::frame_reader::GenericFrame<FixedMessageType>;
56
57pub fn write_frame(w: &mut impl Write, frame: &Frame) -> io::Result<()> {
58    w.write_all(&[frame.msg_type as u8])?;
59    w.write_all(&(frame.payload.len() as u32).to_le_bytes())?;
60    w.write_all(&frame.payload)?;
61    w.flush()
62}
63
64pub fn read_frame(r: &mut impl Read) -> io::Result<Frame> {
65    let mut t = [0u8; 1];
66    r.read_exact(&mut t)?;
67    let msg_type = FixedMessageType::from_u8(t[0]).ok_or_else(|| {
68        io::Error::new(
69            io::ErrorKind::InvalidData,
70            format!("unknown message type: {}", t[0]),
71        )
72    })?;
73    let mut len_buf = [0u8; 4];
74    r.read_exact(&mut len_buf)?;
75    let len = u32::from_le_bytes(len_buf) as usize;
76    if len > MAX_FRAME_PAYLOAD {
77        return Err(io::Error::new(
78            io::ErrorKind::InvalidData,
79            format!("frame payload too large: {len} > {MAX_FRAME_PAYLOAD}"),
80        ));
81    }
82    let mut payload = vec![0u8; len];
83    if len > 0 {
84        r.read_exact(&mut payload)?;
85    }
86    Ok(Frame { msg_type, payload })
87}
88
89/// Stateful, resumable frame reader for non-blocking / read-timeout sockets
90/// (shared implementation: `crate::frame_reader`). On a `WouldBlock`/
91/// `TimedOut` mid-frame the buffered bytes are kept and the next call
92/// resumes exactly; after any terminal error discard the reader together
93/// with the connection.
94///
95/// Keep using the free [`read_frame`] on the blocking handshake paths.
96pub type FrameReader = crate::frame_reader::FrameReader<FixedMessageType>;
97
98/// F→L: ask leader to sync the given shard.
99#[derive(Debug)]
100pub struct SyncRequest {
101    pub shard_id: u8,
102    pub protocol_version: u8,
103    pub flags: u8, // FLAG_EMPTY_STATE etc.
104}
105
106impl SyncRequest {
107    pub fn encode(&self) -> Frame {
108        let payload = vec![self.shard_id, self.protocol_version, self.flags];
109        Frame {
110            msg_type: FixedMessageType::SyncRequest,
111            payload,
112        }
113    }
114    pub fn decode(payload: &[u8]) -> io::Result<Self> {
115        if payload.len() < 3 {
116            return Err(io::Error::new(
117                io::ErrorKind::InvalidData,
118                "SyncRequest too short",
119            ));
120        }
121        Ok(Self {
122            shard_id: payload[0],
123            protocol_version: payload[1],
124            flags: payload[2],
125        })
126    }
127}
128
129/// L→F: leader advertises shard topology after accepting a SyncRequest.
130///
131/// Wire layout (12 bytes):
132/// `[shard_count: u8][key_len: u16 LE][value_len: u16 LE][slot_size: u16 LE]`
133/// `[current_slot_count: u32 LE][shard_prefix_bits: u8]`
134pub struct ShardInfo {
135    pub shard_count: u8,
136    pub key_len: u16,
137    pub value_len: u16,
138    pub slot_size: u16,
139    pub current_slot_count: u32,
140    pub shard_prefix_bits: u8,
141}
142
143impl ShardInfo {
144    pub fn encode(&self) -> Frame {
145        let mut p = Vec::with_capacity(12);
146        p.push(self.shard_count);
147        p.extend_from_slice(&self.key_len.to_le_bytes());
148        p.extend_from_slice(&self.value_len.to_le_bytes());
149        p.extend_from_slice(&self.slot_size.to_le_bytes());
150        p.extend_from_slice(&self.current_slot_count.to_le_bytes());
151        p.push(self.shard_prefix_bits);
152        Frame {
153            msg_type: FixedMessageType::ShardInfo,
154            payload: p,
155        }
156    }
157    pub fn decode(payload: &[u8]) -> io::Result<Self> {
158        if payload.len() < 12 {
159            return Err(io::Error::new(
160                io::ErrorKind::InvalidData,
161                "ShardInfo too short",
162            ));
163        }
164        Ok(Self {
165            shard_count: payload[0],
166            key_len: u16::from_le_bytes(payload[1..3].try_into().expect("2 bytes")),
167            value_len: u16::from_le_bytes(payload[3..5].try_into().expect("2 bytes")),
168            slot_size: u16::from_le_bytes(payload[5..7].try_into().expect("2 bytes")),
169            current_slot_count: u32::from_le_bytes(payload[7..11].try_into().expect("4 bytes")),
170            shard_prefix_bits: payload[11],
171        })
172    }
173}
174
175/// Incremental encoder — caller adds events one at a time to avoid building
176/// an intermediate Vec of event structs.
177pub struct SlotBatchEncoder {
178    payload: Vec<u8>,
179    count: u32,
180    key_len: usize,
181    value_len: usize,
182    count_offset: usize,
183}
184
185impl SlotBatchEncoder {
186    pub fn new(shard_id: u8, key_len: usize, value_len: usize) -> Self {
187        let mut payload = Vec::with_capacity(BATCH_MAX_BYTES);
188        payload.push(shard_id);
189        let count_offset = payload.len();
190        payload.extend_from_slice(&0u32.to_le_bytes()); // placeholder count
191        Self {
192            payload,
193            count: 0,
194            key_len,
195            value_len,
196            count_offset,
197        }
198    }
199
200    pub fn add_occupied(&mut self, slot_id: u32, meta: u32, key: &[u8], value: &[u8]) {
201        debug_assert_eq!(key.len(), self.key_len);
202        debug_assert_eq!(value.len(), self.value_len);
203        debug_assert_eq!(status_of(meta), STATUS_OCCUPIED);
204        self.payload.extend_from_slice(&slot_id.to_le_bytes());
205        self.payload.extend_from_slice(&meta.to_le_bytes());
206        self.payload.extend_from_slice(key);
207        self.payload.extend_from_slice(value);
208        self.count += 1;
209    }
210
211    pub fn add_deleted(&mut self, slot_id: u32, meta: u32, key: &[u8]) {
212        debug_assert_eq!(key.len(), self.key_len);
213        self.payload.extend_from_slice(&slot_id.to_le_bytes());
214        self.payload.extend_from_slice(&meta.to_le_bytes());
215        self.payload.extend_from_slice(key);
216        self.count += 1;
217    }
218
219    pub fn add_reset(&mut self, slot_id: u32, meta: u32) {
220        debug_assert_eq!(status_of(meta), STATUS_FREE);
221        self.payload.extend_from_slice(&slot_id.to_le_bytes());
222        self.payload.extend_from_slice(&meta.to_le_bytes());
223        self.count += 1;
224    }
225
226    pub fn len(&self) -> u32 {
227        self.count
228    }
229    pub fn bytes(&self) -> usize {
230        self.payload.len()
231    }
232    pub fn is_empty(&self) -> bool {
233        self.count == 0
234    }
235
236    pub fn finish(mut self) -> Frame {
237        self.payload[self.count_offset..self.count_offset + 4]
238            .copy_from_slice(&self.count.to_le_bytes());
239        Frame {
240            msg_type: FixedMessageType::SlotBatch,
241            payload: self.payload,
242        }
243    }
244}
245
246/// Zero-copy decoder reading events out of a payload slice.
247pub struct SlotBatchDecoder<'a> {
248    pub shard_id: u8,
249    pub count: u32,
250    pub events_bytes: &'a [u8],
251    key_len: usize,
252    value_len: usize,
253}
254
255#[derive(Debug)]
256pub struct SlotEventRef<'a> {
257    pub slot_id: u32,
258    pub meta: u32,
259    pub key: &'a [u8],
260    pub value: &'a [u8], // empty if DELETED
261}
262
263impl<'a> SlotBatchDecoder<'a> {
264    pub fn new(payload: &'a [u8], key_len: usize, value_len: usize) -> io::Result<Self> {
265        if payload.len() < 5 {
266            return Err(io::Error::new(
267                io::ErrorKind::InvalidData,
268                "SlotBatch too short",
269            ));
270        }
271        let shard_id = payload[0];
272        let count = u32::from_le_bytes(payload[1..5].try_into().expect("4 bytes"));
273        Ok(Self {
274            shard_id,
275            count,
276            events_bytes: &payload[5..],
277            key_len,
278            value_len,
279        })
280    }
281
282    pub fn iter(&self) -> SlotEventIter<'a> {
283        SlotEventIter {
284            remaining: self.events_bytes,
285            left: self.count,
286            key_len: self.key_len,
287            value_len: self.value_len,
288        }
289    }
290}
291
292pub struct SlotEventIter<'a> {
293    remaining: &'a [u8],
294    left: u32,
295    key_len: usize,
296    value_len: usize,
297}
298
299impl<'a> Iterator for SlotEventIter<'a> {
300    type Item = io::Result<SlotEventRef<'a>>;
301
302    fn next(&mut self) -> Option<Self::Item> {
303        if self.left == 0 {
304            return None;
305        }
306        if self.remaining.len() < 8 {
307            return Some(Err(io::Error::new(
308                io::ErrorKind::InvalidData,
309                "SlotBatch event truncated at header",
310            )));
311        }
312        let slot_id = u32::from_le_bytes(self.remaining[0..4].try_into().expect("4 bytes"));
313        let meta = u32::from_le_bytes(self.remaining[4..8].try_into().expect("4 bytes"));
314
315        let (key_len, value_len) = match status_of(meta) {
316            STATUS_OCCUPIED => (self.key_len, self.value_len),
317            STATUS_DELETED => (self.key_len, 0),
318            STATUS_FREE => (0, 0),
319            other => {
320                return Some(Err(io::Error::new(
321                    io::ErrorKind::InvalidData,
322                    format!("unknown slot status: {other}"),
323                )));
324            }
325        };
326
327        let key_start = 8;
328        let key_end = key_start + key_len;
329        let value_end = key_end + value_len;
330        if self.remaining.len() < value_end {
331            return Some(Err(io::Error::new(
332                io::ErrorKind::InvalidData,
333                "SlotBatch event truncated",
334            )));
335        }
336        let key = &self.remaining[key_start..key_end];
337        let value = &self.remaining[key_end..value_end];
338        self.remaining = &self.remaining[value_end..];
339        self.left -= 1;
340        Some(Ok(SlotEventRef {
341            slot_id,
342            meta,
343            key,
344            value,
345        }))
346    }
347}
348
349pub struct CaughtUp {
350    pub shard_id: u8,
351    pub total_scanned: u64,
352}
353
354impl CaughtUp {
355    pub fn encode(&self) -> Frame {
356        let mut p = Vec::with_capacity(9);
357        p.push(self.shard_id);
358        p.extend_from_slice(&self.total_scanned.to_le_bytes());
359        Frame {
360            msg_type: FixedMessageType::CaughtUp,
361            payload: p,
362        }
363    }
364    pub fn decode(payload: &[u8]) -> io::Result<Self> {
365        if payload.len() < 9 {
366            return Err(io::Error::new(
367                io::ErrorKind::InvalidData,
368                "CaughtUp too short",
369            ));
370        }
371        Ok(Self {
372            shard_id: payload[0],
373            total_scanned: u64::from_le_bytes(payload[1..9].try_into().expect("8 bytes")),
374        })
375    }
376}
377
378pub struct Ack {
379    pub shard_id: u8,
380    pub applied_count: u64,
381    pub max_version_seen: u32,
382}
383
384impl Ack {
385    pub fn encode(&self) -> Frame {
386        let mut p = Vec::with_capacity(13);
387        p.push(self.shard_id);
388        p.extend_from_slice(&self.applied_count.to_le_bytes());
389        p.extend_from_slice(&self.max_version_seen.to_le_bytes());
390        Frame {
391            msg_type: FixedMessageType::Ack,
392            payload: p,
393        }
394    }
395    pub fn decode(payload: &[u8]) -> io::Result<Self> {
396        if payload.len() < 13 {
397            return Err(io::Error::new(io::ErrorKind::InvalidData, "Ack too short"));
398        }
399        Ok(Self {
400            shard_id: payload[0],
401            applied_count: u64::from_le_bytes(payload[1..9].try_into().expect("8 bytes")),
402            max_version_seen: u32::from_le_bytes(payload[9..13].try_into().expect("4 bytes")),
403        })
404    }
405}
406
407pub fn encode_heartbeat() -> Frame {
408    Frame {
409        msg_type: FixedMessageType::Heartbeat,
410        payload: Vec::new(),
411    }
412}
413
414pub fn encode_error(msg: &str) -> Frame {
415    Frame {
416        msg_type: FixedMessageType::Error,
417        payload: msg.as_bytes().to_vec(),
418    }
419}
420
421pub fn decode_error(payload: &[u8]) -> String {
422    String::from_utf8_lossy(payload).into_owned()
423}
424
425#[cfg(test)]
426mod tests {
427    use super::*;
428    use crate::fixed::slot::{STATUS_DELETED, STATUS_OCCUPIED, pack_meta};
429
430    #[test]
431    fn test_frame_roundtrip() {
432        let f = Frame {
433            msg_type: FixedMessageType::Heartbeat,
434            payload: vec![1, 2, 3],
435        };
436        let mut buf = Vec::new();
437        write_frame(&mut buf, &f).unwrap();
438        let parsed = read_frame(&mut &buf[..]).unwrap();
439        assert_eq!(parsed.msg_type, FixedMessageType::Heartbeat);
440        assert_eq!(parsed.payload, vec![1, 2, 3]);
441    }
442
443    #[test]
444    fn test_frame_unknown_type() {
445        let buf = [42u8, 0, 0, 0, 0];
446        let err = read_frame(&mut &buf[..]).unwrap_err();
447        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
448    }
449
450    #[test]
451    fn test_sync_request_roundtrip() {
452        let r = SyncRequest {
453            shard_id: 7,
454            protocol_version: PROTOCOL_VERSION,
455            flags: FLAG_EMPTY_STATE,
456        };
457        let frame = r.encode();
458        let d = SyncRequest::decode(&frame.payload).unwrap();
459        assert_eq!(d.shard_id, 7);
460        assert_eq!(d.protocol_version, PROTOCOL_VERSION);
461        assert_eq!(d.flags, FLAG_EMPTY_STATE);
462    }
463
464    #[test]
465    fn test_sync_request_truncated() {
466        assert_eq!(
467            SyncRequest::decode(&[1u8]).unwrap_err().kind(),
468            std::io::ErrorKind::InvalidData
469        );
470    }
471
472    #[test]
473    fn test_shard_info_roundtrip() {
474        let s = ShardInfo {
475            shard_count: 16,
476            key_len: 8,
477            value_len: 32,
478            slot_size: 48,
479            current_slot_count: 1_000_000,
480            shard_prefix_bits: 4,
481        };
482        let f = s.encode();
483        let d = ShardInfo::decode(&f.payload).unwrap();
484        assert_eq!(d.shard_count, 16);
485        assert_eq!(d.key_len, 8);
486        assert_eq!(d.value_len, 32);
487        assert_eq!(d.slot_size, 48);
488        assert_eq!(d.current_slot_count, 1_000_000);
489        assert_eq!(d.shard_prefix_bits, 4);
490    }
491
492    #[test]
493    fn test_slot_batch_mixed_roundtrip() {
494        let key_len = 4usize;
495        let value_len = 8usize;
496        let mut b = SlotBatchEncoder::new(5, key_len, value_len);
497        let occ_meta = pack_meta(STATUS_OCCUPIED, 42);
498        b.add_occupied(100, occ_meta, b"key0", b"12345678");
499        let del_meta = pack_meta(STATUS_DELETED, 43);
500        b.add_deleted(101, del_meta, b"key1");
501        let frame = b.finish();
502        assert_eq!(frame.msg_type, FixedMessageType::SlotBatch);
503
504        let d = SlotBatchDecoder::new(&frame.payload, key_len, value_len).unwrap();
505        assert_eq!(d.shard_id, 5);
506        let events: Vec<_> = d.iter().collect::<Result<_, _>>().unwrap();
507        assert_eq!(events.len(), 2);
508        assert_eq!(events[0].slot_id, 100);
509        assert_eq!(events[0].meta, occ_meta);
510        assert_eq!(events[0].key, b"key0");
511        assert_eq!(events[0].value, b"12345678");
512        assert_eq!(events[1].slot_id, 101);
513        assert_eq!(events[1].meta, del_meta);
514        assert_eq!(events[1].key, b"key1");
515        assert!(events[1].value.is_empty());
516    }
517
518    #[test]
519    fn test_slot_batch_truncated_event_fails() {
520        let key_len = 4usize;
521        let value_len = 8usize;
522        let mut b = SlotBatchEncoder::new(0, key_len, value_len);
523        b.add_occupied(1, pack_meta(STATUS_OCCUPIED, 1), b"key0", b"12345678");
524        let mut frame = b.finish();
525        // Strip last 4 bytes (half of value) → truncated event
526        frame.payload.truncate(frame.payload.len() - 4);
527        let d = SlotBatchDecoder::new(&frame.payload, key_len, value_len).unwrap();
528        let mut iter = d.iter();
529        let err = iter.next().unwrap().unwrap_err();
530        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
531    }
532
533    #[test]
534    fn test_caught_up_roundtrip() {
535        let c = CaughtUp {
536            shard_id: 3,
537            total_scanned: 123_456_789,
538        };
539        let f = c.encode();
540        let d = CaughtUp::decode(&f.payload).unwrap();
541        assert_eq!(d.shard_id, 3);
542        assert_eq!(d.total_scanned, 123_456_789);
543    }
544
545    #[test]
546    fn test_ack_roundtrip() {
547        let a = Ack {
548            shard_id: 4,
549            applied_count: 1000,
550            max_version_seen: 42,
551        };
552        let f = a.encode();
553        let d = Ack::decode(&f.payload).unwrap();
554        assert_eq!(d.shard_id, 4);
555        assert_eq!(d.applied_count, 1000);
556        assert_eq!(d.max_version_seen, 42);
557    }
558
559    #[test]
560    fn test_read_frame_rejects_oversized_payload_before_allocation() {
561        let mut buf = Vec::new();
562        buf.push(FixedMessageType::Heartbeat as u8);
563        buf.extend_from_slice(&((MAX_FRAME_PAYLOAD as u32) + 1).to_le_bytes());
564
565        let err = read_frame(&mut &buf[..]).unwrap_err();
566        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
567        assert!(err.to_string().contains("frame payload too large"));
568    }
569
570    #[test]
571    fn test_slot_batch_reset_roundtrip_without_key_or_value() {
572        let mut b = SlotBatchEncoder::new(5, 4, 8);
573        let reset_meta = pack_meta(STATUS_FREE, 77);
574        b.add_reset(100, reset_meta);
575
576        let frame = b.finish();
577        let d = SlotBatchDecoder::new(&frame.payload, 4, 8).unwrap();
578        let events: Vec<_> = d.iter().collect::<Result<_, _>>().unwrap();
579
580        assert_eq!(events.len(), 1);
581        assert_eq!(events[0].slot_id, 100);
582        assert_eq!(events[0].meta, reset_meta);
583        assert!(events[0].key.is_empty());
584        assert!(events[0].value.is_empty());
585    }
586
587    #[test]
588    fn test_slot_batch_mixed_v2_roundtrip() {
589        let mut b = SlotBatchEncoder::new(5, 4, 8);
590        let occ_meta = pack_meta(STATUS_OCCUPIED, 42);
591        let del_meta = pack_meta(STATUS_DELETED, 43);
592        let free_meta = pack_meta(STATUS_FREE, 44);
593
594        b.add_occupied(100, occ_meta, b"key0", b"12345678");
595        b.add_deleted(101, del_meta, b"key1");
596        b.add_reset(102, free_meta);
597
598        let frame = b.finish();
599        let d = SlotBatchDecoder::new(&frame.payload, 4, 8).unwrap();
600        let events: Vec<_> = d.iter().collect::<Result<_, _>>().unwrap();
601
602        assert_eq!(events.len(), 3);
603        assert_eq!(events[0].key, b"key0");
604        assert_eq!(events[0].value, b"12345678");
605        assert_eq!(events[1].key, b"key1");
606        assert!(events[1].value.is_empty());
607        assert!(events[2].key.is_empty());
608        assert!(events[2].value.is_empty());
609    }
610
611    #[test]
612    fn test_slot_batch_unknown_status_fails() {
613        let mut payload = vec![0u8];
614        payload.extend_from_slice(&1u32.to_le_bytes());
615        payload.extend_from_slice(&9u32.to_le_bytes());
616        payload.extend_from_slice(&0xC000_0003u32.to_le_bytes());
617
618        let d = SlotBatchDecoder::new(&payload, 4, 8).unwrap();
619        let err = d.iter().next().unwrap().unwrap_err();
620        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
621        assert!(err.to_string().contains("unknown slot status"));
622    }
623
624    /// Reader that yields exactly one byte per `read`, returning `WouldBlock`
625    /// between every byte and after the data is exhausted. Simulates a socket
626    /// that dribbles a frame across many partial reads with timeouts.
627    struct DribbleReader {
628        data: Vec<u8>,
629        pos: usize,
630        block_next: bool,
631    }
632
633    impl DribbleReader {
634        fn new(data: Vec<u8>) -> Self {
635            Self {
636                data,
637                pos: 0,
638                block_next: false,
639            }
640        }
641    }
642
643    impl std::io::Read for DribbleReader {
644        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
645            if self.block_next {
646                self.block_next = false;
647                return Err(std::io::Error::new(
648                    std::io::ErrorKind::WouldBlock,
649                    "dribble would block",
650                ));
651            }
652            if self.pos >= self.data.len() || buf.is_empty() {
653                return Err(std::io::Error::new(
654                    std::io::ErrorKind::WouldBlock,
655                    "dribble drained",
656                ));
657            }
658            buf[0] = self.data[self.pos];
659            self.pos += 1;
660            self.block_next = true; // WouldBlock after each byte
661            Ok(1)
662        }
663    }
664
665    #[test]
666    fn test_frame_reader_reassembles_across_would_block() {
667        // Encode two frames back-to-back; drip them one byte at a time with a
668        // WouldBlock between every byte. FrameReader must reassemble both
669        // exactly.
670        let mut wire = Vec::new();
671        let mut b = SlotBatchEncoder::new(2, 4, 8);
672        b.add_occupied(7, pack_meta(STATUS_OCCUPIED, 9), b"keyA", b"valueAAA");
673        b.add_deleted(8, pack_meta(STATUS_DELETED, 10), b"keyB");
674        write_frame(&mut wire, &b.finish()).unwrap();
675        write_frame(&mut wire, &encode_heartbeat()).unwrap();
676
677        let mut reader = DribbleReader::new(wire);
678        let mut fr = FrameReader::new();
679
680        // First frame: the SlotBatch.
681        let mut would_blocks = 0u32;
682        let f1 = loop {
683            match fr.read_frame(&mut reader) {
684                Ok(f) => break f,
685                Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
686                    would_blocks += 1;
687                    continue;
688                }
689                Err(e) => panic!("unexpected error: {e}"),
690            }
691        };
692        assert_eq!(f1.msg_type, FixedMessageType::SlotBatch);
693        let d = SlotBatchDecoder::new(&f1.payload, 4, 8).unwrap();
694        let events: Vec<_> = d.iter().collect::<Result<_, _>>().unwrap();
695        assert_eq!(events.len(), 2);
696        assert_eq!(events[0].slot_id, 7);
697        assert_eq!(events[0].key, b"keyA");
698        assert_eq!(events[0].value, b"valueAAA");
699        assert_eq!(events[1].slot_id, 8);
700        assert_eq!(events[1].key, b"keyB");
701        assert!(
702            would_blocks > 0,
703            "test should have exercised the WouldBlock resume path"
704        );
705
706        // Second frame: the Heartbeat, read from the same buffered stream.
707        let f2 = loop {
708            match fr.read_frame(&mut reader) {
709                Ok(f) => break f,
710                Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => continue,
711                Err(e) => panic!("unexpected error: {e}"),
712            }
713        };
714        assert_eq!(f2.msg_type, FixedMessageType::Heartbeat);
715        assert!(f2.payload.is_empty());
716    }
717
718    #[test]
719    fn test_frame_reader_matches_free_read_frame() {
720        // A frame delivered in one shot must decode identically to the free fn.
721        let mut wire = Vec::new();
722        let a = Ack {
723            shard_id: 3,
724            applied_count: 77,
725            max_version_seen: 5,
726        };
727        write_frame(&mut wire, &a.encode()).unwrap();
728
729        let mut fr = FrameReader::new();
730        let f = fr.read_frame(&mut &wire[..]).unwrap();
731        assert_eq!(f.msg_type, FixedMessageType::Ack);
732        let decoded = Ack::decode(&f.payload).unwrap();
733        assert_eq!(decoded.shard_id, 3);
734        assert_eq!(decoded.applied_count, 77);
735        assert_eq!(decoded.max_version_seen, 5);
736    }
737
738    #[test]
739    fn test_frame_reader_rejects_oversized_len() {
740        let mut buf = Vec::new();
741        buf.push(FixedMessageType::Heartbeat as u8);
742        buf.extend_from_slice(&((MAX_FRAME_PAYLOAD as u32) + 1).to_le_bytes());
743        let mut fr = FrameReader::new();
744        let err = fr.read_frame(&mut &buf[..]).unwrap_err();
745        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
746        assert!(err.to_string().contains("frame payload too large"));
747    }
748
749    #[test]
750    fn frame_reader_rejects_garbage_type_byte_immediately() {
751        // R3: early type-byte validation — an error with a single buffered
752        // byte, without waiting for the full 5-byte header. (The old
753        // implementation waited for 5 bytes; the generic validates earlier.)
754        let mut fr = FrameReader::new();
755        let err = fr.read_frame(&mut &[0xEEu8][..]).unwrap_err();
756        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
757        assert!(err.to_string().contains("unknown message type"));
758    }
759
760    #[test]
761    fn frame_reader_default_contract() {
762        // R5: the old FrameReader derived Default — the alias must keep
763        // that contract.
764        let fr = FrameReader::default();
765        let _ = fr;
766    }
767}