use std::io::{self, Read, Seek, SeekFrom};
use std::path::Path;
pub const FRAME_HEADER_LEN: usize = 8;
pub const MAX_SEGMENT_BYTES: u64 = 64 * 1024 * 1024;
pub fn complete_frames(bytes: &[u8]) -> (usize, u64) {
let mut pos = 0usize;
let mut frames = 0u64;
while let Some(len) = frame_len_at(bytes, pos) {
let start = pos + FRAME_HEADER_LEN;
let Some(end) = start.checked_add(len).filter(|&e| e <= bytes.len()) else {
break;
};
let crc = u32::from_le_bytes(bytes[pos + 4..pos + 8].try_into().unwrap());
if crc32fast::hash(&bytes[start..end]) != crc {
break;
}
pos = end;
frames += 1;
}
(pos, frames)
}
#[derive(Debug, Default)]
pub struct FrameCounter {
header: [u8; FRAME_HEADER_LEN],
have: usize,
remaining: u64,
in_frame: bool,
frames: u64,
}
impl FrameCounter {
pub fn feed(&mut self, mut buf: &[u8]) {
while !buf.is_empty() {
if self.in_frame {
let take = (self.remaining as usize).min(buf.len());
self.remaining -= take as u64;
buf = &buf[take..];
if self.remaining == 0 {
self.in_frame = false;
self.frames += 1;
}
} else {
let take = (FRAME_HEADER_LEN - self.have).min(buf.len());
self.header[self.have..self.have + take].copy_from_slice(&buf[..take]);
self.have += take;
buf = &buf[take..];
if self.have == FRAME_HEADER_LEN {
self.have = 0;
self.remaining =
u32::from_le_bytes(self.header[..4].try_into().unwrap()) as u64;
if self.remaining == 0 {
self.frames += 1;
} else {
self.in_frame = true;
}
}
}
}
}
pub fn frames(&self) -> u64 {
self.frames
}
}
fn frame_len_at(bytes: &[u8], pos: usize) -> Option<usize> {
let header = bytes.get(pos..pos + FRAME_HEADER_LEN)?;
Some(u32::from_le_bytes(header[..4].try_into().unwrap()) as usize)
}
pub fn read_new_segment(
wal_path: &Path,
from_byte: u64,
durable_len: u64,
) -> io::Result<Option<(u64, u64, Vec<u8>)>> {
if durable_len <= from_byte {
return Ok(None);
}
let mut file = match std::fs::File::open(wal_path) {
Ok(f) => f,
Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(None),
Err(e) => return Err(e),
};
let available = durable_len.min(file.metadata()?.len());
if available <= from_byte {
return Ok(None);
}
let want = (available - from_byte).min(MAX_SEGMENT_BYTES);
file.seek(SeekFrom::Start(from_byte))?;
let mut bytes = read_up_to(&mut file, want)?;
if let Some(len) = frame_len_at(&bytes, 0) {
let frame = (FRAME_HEADER_LEN + len) as u64;
if frame > bytes.len() as u64 && from_byte + frame <= available {
let rest = read_up_to(&mut file, frame - bytes.len() as u64)?;
bytes.extend_from_slice(&rest);
}
}
let (whole, _) = complete_frames(&bytes);
if whole == 0 {
let first_frame_is_there = frame_len_at(&bytes, 0)
.is_some_and(|len| (FRAME_HEADER_LEN + len) as u64 <= bytes.len() as u64);
if first_frame_is_there {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("wal.bin at byte {from_byte} is not a valid frame; not shipping it"),
));
}
return Ok(None);
}
bytes.truncate(whole);
Ok(Some((from_byte, from_byte + whole as u64, bytes)))
}
fn read_up_to(file: &mut std::fs::File, n: u64) -> io::Result<Vec<u8>> {
let mut buf = Vec::with_capacity(n as usize);
file.take(n).read_to_end(&mut buf)?;
Ok(buf)
}
#[cfg(test)]
mod tests {
use super::*;
fn frame(payload: &[u8]) -> Vec<u8> {
let mut f = Vec::new();
f.extend((payload.len() as u32).to_le_bytes());
f.extend(crc32fast::hash(payload).to_le_bytes());
f.extend(payload);
f
}
#[test]
fn complete_frames_counts_whole_valid_frames() {
let mut b = frame(b"one");
b.extend(frame(b"two"));
let two = b.len();
b.extend(&frame(b"three")[..5]);
assert_eq!(complete_frames(&b), (two, 2));
assert_eq!(complete_frames(&[]), (0, 0));
}
#[test]
fn the_counter_finds_the_same_frames_in_any_chunking() {
let mut b = Vec::new();
for p in [
&b""[..],
b"one",
b"a much longer payload than the rest",
b"x",
] {
b.extend(frame(p));
}
let (_, whole) = complete_frames(&b);
assert_eq!(whole, 4);
for chunk in [1, 3, 7, 8, 9, 100] {
let mut c = FrameCounter::default();
for piece in b.chunks(chunk) {
c.feed(piece);
}
assert_eq!(c.frames(), whole, "chunk {chunk}");
}
let mut c = FrameCounter::default();
c.feed(&b[..b.len() - 1]);
assert_eq!(c.frames(), 3, "a partial last frame is not counted");
}
#[test]
fn a_bad_checksum_stops_the_scan() {
let mut b = frame(b"one");
let one = b.len();
let mut bad = frame(b"two");
let last = bad.len() - 1;
bad[last] ^= 0xff;
b.extend(bad);
assert_eq!(complete_frames(&b), (one, 1));
}
#[test]
fn nothing_past_the_durable_length_is_read() {
let dir = tempfile::tempdir().unwrap();
let p = dir.path().join("wal.bin");
let mut b = frame(b"one");
let one = b.len() as u64;
b.extend(frame(b"two"));
std::fs::write(&p, &b).unwrap();
let (s, e, bytes) = read_new_segment(&p, 0, one).unwrap().unwrap();
assert_eq!((s, e), (0, one));
assert_eq!(bytes, frame(b"one"));
assert!(read_new_segment(&p, one, one).unwrap().is_none());
}
#[test]
fn a_whole_frame_that_fails_its_checksum_is_an_error() {
let dir = tempfile::tempdir().unwrap();
let p = dir.path().join("wal.bin");
let mut b = frame(b"0123456789abcdef");
let last = b.len() - 1;
b[last] ^= 0xff;
std::fs::write(&p, &b).unwrap();
let err = read_new_segment(&p, 0, b.len() as u64).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn a_missing_file_has_nothing_to_ship() {
let dir = tempfile::tempdir().unwrap();
assert!(read_new_segment(&dir.path().join("wal.bin"), 0, 100)
.unwrap()
.is_none());
}
}