use crate::wal::frame::append_frame_bytes;
use crate::wal::rotation::RotationReason;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Write};
use std::path::PathBuf;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use thiserror::Error;
pub const SEGMENT_MAGIC: u32 = 0x4145_4442;
pub const SEGMENT_HEADER_SIZE: usize = 64;
#[derive(Debug, Error)]
pub enum SegmentError {
#[error("io error: {0}")]
Io(#[from] std::io::Error),
#[error("invalid segment magic")]
InvalidMagic,
#[error("segment is not open")]
NotOpen,
}
#[derive(Debug, Clone)]
pub struct SegmentConfig {
pub max_segment_bytes: u64,
pub max_segment_age: Duration,
}
impl Default for SegmentConfig {
fn default() -> Self {
Self {
max_segment_bytes: 64 * 1024 * 1024,
max_segment_age: Duration::from_secs(60),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SegmentHeader {
pub magic: u32,
pub format_version: u16,
pub instance_id: u64,
pub segment_seq: u64,
pub created_at_micros: u64,
pub prev_segment_hash: [u8; 32],
}
impl SegmentHeader {
pub fn new(instance_id: u64, segment_seq: u64, prev_segment_hash: [u8; 32]) -> Self {
let created_at_micros = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_micros() as u64;
Self {
magic: SEGMENT_MAGIC,
format_version: 1,
instance_id,
segment_seq,
created_at_micros,
prev_segment_hash,
}
}
pub fn to_bytes(&self) -> [u8; SEGMENT_HEADER_SIZE] {
let mut out = [0u8; SEGMENT_HEADER_SIZE];
out[0..4].copy_from_slice(&self.magic.to_be_bytes());
out[4..6].copy_from_slice(&self.format_version.to_be_bytes());
out[6..8].copy_from_slice(&0u16.to_be_bytes());
out[8..16].copy_from_slice(&self.instance_id.to_be_bytes());
out[16..24].copy_from_slice(&self.segment_seq.to_be_bytes());
out[24..32].copy_from_slice(&self.created_at_micros.to_be_bytes());
out[32..64].copy_from_slice(&self.prev_segment_hash);
out
}
pub fn from_bytes(bytes: &[u8; SEGMENT_HEADER_SIZE]) -> Result<Self, SegmentError> {
let magic = u32::from_be_bytes(copy_header_bytes::<4>(bytes, 0));
if magic != SEGMENT_MAGIC {
return Err(SegmentError::InvalidMagic);
}
Ok(Self {
magic,
format_version: u16::from_be_bytes(copy_header_bytes::<2>(bytes, 4)),
instance_id: u64::from_be_bytes(copy_header_bytes::<8>(bytes, 8)),
segment_seq: u64::from_be_bytes(copy_header_bytes::<8>(bytes, 16)),
created_at_micros: u64::from_be_bytes(copy_header_bytes::<8>(bytes, 24)),
prev_segment_hash: copy_header_bytes::<32>(bytes, 32),
})
}
}
fn copy_header_bytes<const N: usize>(bytes: &[u8; SEGMENT_HEADER_SIZE], offset: usize) -> [u8; N] {
let mut out = [0u8; N];
out.copy_from_slice(&bytes[offset..offset + N]);
out
}
#[derive(Debug, Clone)]
pub struct ClosedSegment {
pub path: PathBuf,
pub segment_seq: u64,
pub hash: [u8; 32],
pub size_bytes: u64,
}
struct ActiveSegment {
file: File,
path: PathBuf,
segment_seq: u64,
opened_at: Instant,
size_bytes: u64,
commit_count: u64,
hasher: blake3::Hasher,
}
pub(crate) struct PendingFrame<'a> {
pub seq: u64,
pub timestamp_micros: u64,
pub payload_type: u8,
pub payload: &'a [u8],
}
pub struct SegmentManager {
dir: PathBuf,
config: SegmentConfig,
instance_id: u64,
prev_hash: [u8; 32],
active: Option<ActiveSegment>,
force_rotate: bool,
encode_scratch: Vec<u8>,
}
impl SegmentManager {
const FRAME_FIXED_BYTES: usize = 4 + 8 + 8 + 1 + 4;
pub fn new(dir: impl Into<PathBuf>, config: SegmentConfig, instance_id: u64) -> Self {
Self {
dir: dir.into(),
config,
instance_id,
prev_hash: [0u8; 32],
active: None,
force_rotate: false,
encode_scratch: Vec::new(),
}
}
pub fn open_active(&mut self, segment_seq: u64) -> Result<(), SegmentError> {
fs::create_dir_all(&self.dir)?;
let path = self.dir.join(format!("segment_{segment_seq:016}.aedbwal"));
let mut file = OpenOptions::new()
.create(true)
.append(true)
.read(true)
.open(&path)?;
let mut hasher = blake3::Hasher::new();
if file.metadata()?.len() == 0 {
let header = SegmentHeader::new(self.instance_id, segment_seq, self.prev_hash);
let header_bytes = header.to_bytes();
file.write_all(&header_bytes)?;
hasher.update(&header_bytes);
file.flush()?;
} else {
let mut reader = File::open(&path)?;
let mut buf = [0u8; 64 * 1024];
loop {
let n = reader.read(&mut buf)?;
if n == 0 {
break;
}
hasher.update(&buf[..n]);
}
}
let size_bytes = file.metadata()?.len();
self.active = Some(ActiveSegment {
file,
path,
segment_seq,
opened_at: Instant::now(),
size_bytes,
commit_count: 0,
hasher,
});
Ok(())
}
#[cfg(test)]
pub(crate) fn append_frame(
&mut self,
seq: u64,
timestamp_micros: u64,
payload_type: u8,
payload: &[u8],
) -> Result<(), SegmentError> {
let single = [PendingFrame {
seq,
timestamp_micros,
payload_type,
payload,
}];
self.append_frames_with_sync(&single, true)
}
pub(crate) fn append_frames_with_sync(
&mut self,
frames: &[PendingFrame<'_>],
sync: bool,
) -> Result<(), SegmentError> {
if frames.is_empty() {
return Ok(());
}
let active = self.active.as_mut().ok_or(SegmentError::NotOpen)?;
let estimated_encoded_bytes = frames.iter().fold(0usize, |acc, frame| {
acc.saturating_add(Self::FRAME_FIXED_BYTES.saturating_add(frame.payload.len()))
});
let mut encoded_frames = std::mem::take(&mut self.encode_scratch);
encoded_frames.clear();
if encoded_frames.capacity() < estimated_encoded_bytes {
encoded_frames.reserve(estimated_encoded_bytes - encoded_frames.capacity());
}
for frame in frames {
append_frame_bytes(
&mut encoded_frames,
frame.seq,
frame.timestamp_micros,
frame.payload_type,
frame.payload,
)
.map_err(|e| SegmentError::Io(std::io::Error::other(e.to_string())))?;
}
let encoded_size_bytes = encoded_frames.len() as u64;
if crate::faults::any_armed() {
crate::faults::trip("wal_append")
.map_err(|e| SegmentError::Io(std::io::Error::other(e.to_string())))?;
}
active.file.write_all(&encoded_frames)?;
active.hasher.update(&encoded_frames);
active.size_bytes = active.size_bytes.saturating_add(encoded_size_bytes);
if sync {
if crate::faults::any_armed() {
crate::faults::trip("wal_sync")
.map_err(|e| SegmentError::Io(std::io::Error::other(e.to_string())))?;
}
active.file.flush()?;
active.file.sync_data()?;
}
active.commit_count = active.commit_count.saturating_add(frames.len() as u64);
encoded_frames.clear();
self.encode_scratch = encoded_frames;
Ok(())
}
pub fn sync_active(&mut self) -> Result<(), SegmentError> {
let active = self.active.as_mut().ok_or(SegmentError::NotOpen)?;
if crate::faults::any_armed() {
crate::faults::trip("wal_sync")
.map_err(|e| SegmentError::Io(std::io::Error::other(e.to_string())))?;
}
active.file.sync_data()?;
Ok(())
}
pub(crate) fn try_clone_active_file(&self) -> Result<Option<File>, SegmentError> {
match &self.active {
Some(active) => Ok(Some(active.file.try_clone()?)),
None => Ok(None),
}
}
pub fn should_rotate(&self) -> Option<RotationReason> {
self.should_rotate_at(Instant::now())
}
pub fn should_rotate_at(&self, now: Instant) -> Option<RotationReason> {
let active = self.active.as_ref()?;
if self.force_rotate {
return Some(RotationReason::Forced);
}
if active.size_bytes >= self.config.max_segment_bytes {
return Some(RotationReason::Size);
}
if now.duration_since(active.opened_at) >= self.config.max_segment_age {
return Some(RotationReason::Time);
}
None
}
pub fn force_rotate(&mut self) {
self.force_rotate = true;
}
pub fn rotate(&mut self) -> Result<ClosedSegment, SegmentError> {
let closed = self.close_active()?;
self.prev_hash = closed.hash;
self.force_rotate = false;
self.open_active(closed.segment_seq + 1)?;
Ok(closed)
}
pub fn close_active(&mut self) -> Result<ClosedSegment, SegmentError> {
let mut active = self.active.take().ok_or(SegmentError::NotOpen)?;
active.file.flush()?;
let hash = *blake3::Hasher::finalize(&active.hasher).as_bytes();
let size_bytes = active.file.metadata()?.len();
Ok(ClosedSegment {
path: active.path,
segment_seq: active.segment_seq,
hash,
size_bytes,
})
}
}
#[cfg(test)]
mod tests;