use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::Path;
use crate::crypto::KeyRing;
use crate::error::{Result, WalError};
use crate::preamble::{SegmentPreamble, WAL_PREAMBLE_MAGIC, read_leading_preamble};
use crate::reader::StopReason;
use crate::record::foreign_version::read_declared_payload;
use crate::record::{
HEADER_SIZE, MAX_WAL_PAYLOAD_SIZE, RecordHeader, RecordType, WalRecord, is_format_gap,
may_be_format_gap,
};
use crate::segment::SegmentDecryptor;
fn checked_offset_add(offset: u64, len: u64) -> Result<u64> {
offset.checked_add(len).ok_or_else(|| {
WalError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"WAL lazy-reader offset overflow",
))
})
}
pub struct LazyWalReader {
file: File,
offset: u64,
records_start: u64,
segment_preamble: Option<SegmentPreamble>,
double_write: Option<crate::double_write::DoubleWriteBuffer>,
stop_reason: Option<StopReason>,
decryptor: SegmentDecryptor,
}
impl LazyWalReader {
pub fn open(path: &Path, keys: Option<&KeyRing>) -> Result<Self> {
let mut file = File::open(path)?;
let dwb_path = path.with_extension("dwb");
let double_write = if dwb_path.exists() {
crate::double_write::DoubleWriteBuffer::open(
&dwb_path,
crate::double_write::DwbMode::Buffered,
)
.ok()
} else {
None
};
let (segment_preamble, start_offset) =
read_leading_preamble(&mut file, &WAL_PREAMBLE_MAGIC)?;
let decryptor = SegmentDecryptor::new(segment_preamble.as_ref(), keys);
Ok(Self {
file,
offset: start_offset,
records_start: start_offset,
segment_preamble,
double_write,
stop_reason: None,
decryptor,
})
}
pub fn segment_preamble(&self) -> Option<&SegmentPreamble> {
self.segment_preamble.as_ref()
}
pub fn stop_reason(&self) -> Option<StopReason> {
self.stop_reason
}
fn stop(&mut self, reason: StopReason) -> Result<Option<RecordHeader>> {
self.stop_reason = Some(reason);
Ok(None)
}
pub fn next_header(&mut self) -> Result<Option<RecordHeader>> {
loop {
let record_offset = self.offset;
let mut header_buf = [0u8; HEADER_SIZE];
match self.read_exact(&mut header_buf) {
Ok(()) => {}
Err(WalError::Io(e)) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
let reason = if record_offset == self.file.metadata()?.len() {
StopReason::Eof
} else {
StopReason::Corruption {
offset: record_offset,
}
};
return self.stop(reason);
}
Err(e) => return Err(e),
}
let header = RecordHeader::from_bytes(&header_buf);
match header.validate(record_offset) {
Ok(()) => {}
Err(error @ WalError::PayloadTooLarge { .. }) => return Err(error),
Err(error @ WalError::UnsupportedVersion { .. })
if may_be_format_gap(&header, record_offset == self.records_start) =>
{
let payload = read_declared_payload(&mut self.file, &header, self.offset)?;
if is_format_gap(&header, payload.as_deref()) {
return Err(error);
}
return self.stop(StopReason::Corruption {
offset: record_offset,
});
}
Err(_) => {
return self.stop(StopReason::Corruption {
offset: record_offset,
});
}
}
let logical_type = header.logical_record_type();
match RecordType::from_raw(logical_type) {
Some(RecordType::Noop) => {
self.skip_payload(&header)?;
continue;
}
Some(_) => {}
None if RecordType::is_required(logical_type) => {
return Err(WalError::UnknownRequiredRecordType {
record_type: header.record_type,
lsn: header.lsn,
});
}
None => {}
}
return Ok(Some(header));
}
}
pub fn read_payload(&mut self, header: &RecordHeader) -> Result<Vec<u8>> {
let raw = self.read_raw_payload(header)?;
self.decryptor.decrypt_payload(header, raw)
}
fn read_raw_payload(&mut self, header: &RecordHeader) -> Result<Vec<u8>> {
header.validate(self.offset)?;
let payload_len =
usize::try_from(header.payload_len).map_err(|_| WalError::PayloadTooLarge {
size: usize::MAX,
max: MAX_WAL_PAYLOAD_SIZE,
})?;
let mut payload = vec![0u8; payload_len];
if !payload.is_empty() {
match self.read_exact(&mut payload) {
Ok(()) => {}
Err(WalError::Io(e)) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
return Err(WalError::Io(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"torn write: incomplete payload",
)));
}
Err(e) => return Err(e),
}
}
let checksum_ok = header.crc32c == header.compute_checksum(&payload);
if !checksum_ok {
if let Some(dwb) = &mut self.double_write
&& let Ok(Some(recovered)) = dwb.recover_record(header.lsn)
{
let payload_start =
self.offset.checked_sub(payload_len as u64).ok_or_else(|| {
WalError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"WAL lazy-reader payload offset underflow",
))
})?;
let resume = checked_offset_add(payload_start, recovered.payload.len() as u64)?;
self.file.seek(SeekFrom::Start(resume))?;
self.offset = resume;
return Ok(recovered.payload);
}
return Err(WalError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"checksum mismatch",
)));
}
Ok(payload)
}
pub fn read_record(&mut self, header: &RecordHeader) -> Result<WalRecord> {
let raw = self.read_raw_payload(header)?;
let record = WalRecord {
header: *header,
payload: raw,
};
self.decryptor.decrypt_record(record)
}
pub fn skip_payload(&mut self, header: &RecordHeader) -> Result<()> {
header.validate(self.offset)?;
let len = u64::from(header.payload_len);
let target = checked_offset_add(self.offset, len)?;
let file_len = self.file.metadata()?.len();
if target > file_len {
return Err(WalError::Io(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"torn write: incomplete payload while skipping",
)));
}
if len > 0 {
let signed_len = i64::try_from(len).map_err(|_| {
WalError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"WAL lazy-reader skip length does not fit i64",
))
})?;
let position = self.file.seek(SeekFrom::Current(signed_len))?;
if position != target {
return Err(WalError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"WAL lazy-reader seek did not reach bounded target",
)));
}
self.offset = target;
}
Ok(())
}
pub fn offset(&self) -> u64 {
self.offset
}
fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
let len = u64::try_from(buf.len()).map_err(|_| {
WalError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"WAL lazy-reader read length does not fit u64",
))
})?;
let next_offset = checked_offset_add(self.offset, len)?;
self.file.read_exact(buf)?;
self.offset = next_offset;
Ok(())
}
}
pub fn replay_segment_lazy<F>(path: &Path, keys: Option<&KeyRing>, handler: F) -> Result<()>
where
F: FnMut(&mut LazyWalReader, &RecordHeader) -> Result<()>,
{
replay_one_segment(path, keys, handler).map(|_| ())
}
fn replay_one_segment<F>(path: &Path, keys: Option<&KeyRing>, mut handler: F) -> Result<u64>
where
F: FnMut(&mut LazyWalReader, &RecordHeader) -> Result<()>,
{
let mut reader = LazyWalReader::open(path, keys)?;
let mut last_lsn = 0u64;
while let Some(header) = reader
.next_header()
.map_err(|error| error.with_segment_path(path))?
{
last_lsn = header.lsn;
handler(&mut reader, &header)?;
}
crate::torn_tail::verify_committed_prefix(path, reader.stop_reason(), last_lsn)?;
Ok(last_lsn)
}
pub fn replay_all_segments_lazy<F>(
wal_dir: &Path,
keys: Option<&KeyRing>,
mut handler: F,
) -> Result<()>
where
F: FnMut(&mut LazyWalReader, &RecordHeader) -> Result<()>,
{
let segments = crate::segment::discover_segments(wal_dir)?;
let mut continuity = crate::segment::SegmentContinuity::new();
for seg in &segments {
continuity.check(seg)?;
let last_lsn = replay_one_segment(&seg.path, keys, &mut handler)?;
continuity.completed(seg, last_lsn);
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::record::RecordType;
use crate::writer::WalWriter;
#[test]
fn lazy_read_all_payloads() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test.wal");
{
let mut w = WalWriter::open_without_direct_io(&path).unwrap();
w.append(RecordType::Put as u32, 1, 0, 0, b"hello").unwrap();
w.append(RecordType::VectorPut as u32, 1, 0, 0, b"vector-data")
.unwrap();
w.append(RecordType::Put as u32, 2, 1, 0, b"world").unwrap();
w.sync().unwrap();
}
let mut reader = LazyWalReader::open(&path, None).unwrap();
let mut records = Vec::new();
while let Some(header) = reader.next_header().unwrap() {
let payload = reader.read_payload(&header).unwrap();
records.push((header, payload));
}
assert_eq!(records.len(), 3);
assert_eq!(records[0].1, b"hello");
assert_eq!(records[1].1, b"vector-data");
assert_eq!(records[2].1, b"world");
}
#[test]
fn lazy_skip_non_matching() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test.wal");
{
let mut w = WalWriter::open_without_direct_io(&path).unwrap();
w.append(RecordType::TimeseriesBatch as u32, 1, 0, 0, &[0u8; 10000])
.unwrap();
w.append(RecordType::TimeseriesBatch as u32, 1, 0, 0, &[0u8; 10000])
.unwrap();
w.append(RecordType::VectorPut as u32, 1, 0, 0, b"small-vec")
.unwrap();
w.append(RecordType::TimeseriesBatch as u32, 1, 0, 0, &[0u8; 10000])
.unwrap();
w.sync().unwrap();
}
let mut reader = LazyWalReader::open(&path, None).unwrap();
let mut vector_payloads = Vec::new();
let mut skipped = 0;
while let Some(header) = reader.next_header().unwrap() {
let rt = RecordType::from_raw(header.record_type);
if rt == Some(RecordType::VectorPut) {
let payload = reader.read_payload(&header).unwrap();
vector_payloads.push(payload);
} else {
reader.skip_payload(&header).unwrap();
skipped += 1;
}
}
assert_eq!(vector_payloads.len(), 1);
assert_eq!(vector_payloads[0], b"small-vec");
assert_eq!(skipped, 3);
}
#[test]
fn oversized_payload_header_is_a_typed_error_before_skip_or_allocation() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("oversized.wal");
let header = RecordHeader {
magic: crate::record::WAL_MAGIC,
format_version: crate::record::WAL_FORMAT_VERSION,
record_type: RecordType::Put as u32,
lsn: 1,
tenant_id: 1,
vshard_id: 0,
payload_len: (MAX_WAL_PAYLOAD_SIZE + 1) as u32,
database_id: 0,
apply_key: 0,
event_source: crate::record::NO_EVENT_SOURCE,
commit_hlc: 0,
crc32c: 0,
};
std::fs::write(&path, header.to_bytes()).unwrap();
let mut reader = LazyWalReader::open(&path, None).unwrap();
assert!(matches!(
reader.next_header(),
Err(WalError::PayloadTooLarge { .. })
));
}
#[test]
fn skip_payload_rejects_truncated_tail_without_seeking_beyond_eof() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("truncated.wal");
let mut writer = WalWriter::open_without_direct_io(&path).unwrap();
writer
.append(RecordType::Put as u32, 1, 0, 0, b"payload")
.unwrap();
writer.sync().unwrap();
drop(writer);
let length = std::fs::metadata(&path).unwrap().len();
std::fs::OpenOptions::new()
.write(true)
.open(&path)
.unwrap()
.set_len(length - 1)
.unwrap();
let mut reader = LazyWalReader::open(&path, None).unwrap();
let header = reader.next_header().unwrap().expect("header");
assert!(
matches!(reader.skip_payload(&header), Err(WalError::Io(error)) if error.kind() == std::io::ErrorKind::UnexpectedEof)
);
assert_eq!(reader.offset(), HEADER_SIZE as u64);
}
#[test]
fn checked_offset_add_rejects_overflow() {
assert!(checked_offset_add(u64::MAX, 1).is_err());
}
#[test]
fn replay_all_segments_lazy_works() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("wal-00000000000000000001.seg");
{
let mut w = WalWriter::open_without_direct_io(&path).unwrap();
w.append(RecordType::Put as u32, 1, 0, 0, b"a").unwrap();
w.append(RecordType::Put as u32, 1, 0, 0, b"b").unwrap();
w.sync().unwrap();
}
let mut count = 0;
replay_all_segments_lazy(dir.path(), None, |reader, header| {
reader.skip_payload(header)?;
count += 1;
Ok(())
})
.unwrap();
assert_eq!(count, 2);
}
#[test]
fn empty_wal_no_records() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("empty.wal");
{
let mut w = WalWriter::open_without_direct_io(&path).unwrap();
w.sync().unwrap();
}
let mut reader = LazyWalReader::open(&path, None).unwrap();
assert!(reader.next_header().unwrap().is_none());
}
#[test]
fn an_unknown_format_version_at_the_first_record_is_an_error() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("unknown.wal");
{
let mut writer = WalWriter::open_without_direct_io(&path).unwrap();
writer
.append(RecordType::Put as u32, 1, 0, 0, b"x")
.unwrap();
writer.sync().unwrap();
}
let mut bytes = std::fs::read(&path).unwrap();
crate::record::foreign_version::reseal_with_version(
&mut bytes,
0,
crate::record::WAL_FORMAT_VERSION + 1,
);
std::fs::write(&path, &bytes).unwrap();
let mut reader = LazyWalReader::open(&path, None).unwrap();
match reader.next_header() {
Err(WalError::UnsupportedVersion { version, supported }) => {
assert_eq!(version, crate::record::WAL_FORMAT_VERSION + 1);
assert_eq!(supported, crate::record::WAL_FORMAT_VERSION);
}
other => panic!("expected UnsupportedVersion, got {other:?}"),
}
}
#[test]
fn a_version_mismatch_after_the_first_record_is_damage() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("two-records.wal");
{
let mut writer = WalWriter::open_without_direct_io(&path).unwrap();
writer
.append(RecordType::Put as u32, 1, 0, 0, b"first")
.unwrap();
writer
.append(RecordType::Put as u32, 1, 0, 0, b"second")
.unwrap();
writer.sync().unwrap();
}
let mut reader = LazyWalReader::open(&path, None).unwrap();
let first = reader.next_header().unwrap().expect("the first header");
reader.skip_payload(&first).unwrap();
let second_at = reader.offset();
let mut bytes = std::fs::read(&path).unwrap();
let version_at = usize::try_from(second_at).unwrap() + 4;
bytes[version_at..version_at + 2]
.copy_from_slice(&(crate::record::WAL_FORMAT_VERSION + 1).to_le_bytes());
std::fs::write(&path, &bytes).unwrap();
let mut reader = LazyWalReader::open(&path, None).unwrap();
let first = reader
.next_header()
.unwrap()
.expect("the first record is still readable");
assert_eq!(first.lsn, 1);
reader.skip_payload(&first).unwrap();
assert!(
reader.next_header().unwrap().is_none(),
"the mismatch stops the stream instead of failing the read"
);
assert_eq!(
reader.stop_reason(),
Some(StopReason::Corruption { offset: second_at }),
"a mismatch after the first record is damage, not a format gap"
);
}
}