use anyhow::{anyhow, Context, Result};
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::thread;
use tracing::{debug, error, trace, warn};
use crate::fileio::{sync_dir, FileIo};
use crate::logfile::{
append_footer, encode_records, log_file_path, now_ms, write_header, LogfileError,
FILE_HEADER_SIZE, FOOTER_SIZE, RECORD_HEADER_SIZE,
};
use crate::record::RecordID;
use crate::write_ahead::{WriteAheadError, WriteAheadOptions};
pub(crate) enum WriterCommand {
Write(flume::Sender<Result<WriteAck>>, Vec<Vec<u8>>),
Trim(flume::Sender<TrimStats>, u64),
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TrimStats {
pub files_deleted: u64,
pub bytes_reclaimed: u64,
}
#[derive(Debug)]
pub(crate) struct WriteAck {
pub file_id: u64,
pub offsets: Vec<u64>,
}
#[derive(Debug, Clone)]
pub(crate) enum WriterEvent {
Created(u64, PathBuf),
Sealed(u64),
Deleted(u64),
}
#[derive(Debug, Clone)]
pub(crate) struct FileMeta {
pub path: PathBuf,
pub size: u64,
pub seal_timestamp_ms: Option<u64>,
}
#[derive(Debug, Clone)]
pub struct WriteHandle {
pub(crate) tx: flume::Sender<WriterCommand>,
}
impl WriteHandle {
pub async fn write(&self, record: Vec<u8>) -> Result<RecordID> {
let mut ids = self.write_batch(vec![record]).await?;
Ok(ids.remove(0))
}
pub async fn write_batch(&self, records: Vec<Vec<u8>>) -> Result<Vec<RecordID>> {
let (tx, rx) = flume::bounded(1);
self.tx
.send(WriterCommand::Write(tx, records))
.map_err(|_| WriteAheadError::WriterClosed)?;
let ack = rx
.recv_async()
.await
.map_err(|_| WriteAheadError::WriterClosed)??;
Ok(ack
.offsets
.iter()
.map(|offset| RecordID::new(ack.file_id, *offset))
.collect())
}
pub async fn trim_before(&self, file_id: u64) -> Result<TrimStats> {
let (tx, rx) = flume::bounded(1);
self.tx
.send(WriterCommand::Trim(tx, file_id))
.map_err(|_| WriteAheadError::WriterClosed)?;
Ok(rx
.recv_async()
.await
.map_err(|_| WriteAheadError::WriterClosed)?)
}
}
type ReplyTx = flume::Sender<Result<WriteAck>>;
const MAX_GROUP_COMMIT_BYTES: usize = 4 * 1024 * 1024;
fn batch_bytes(records: &[Vec<u8>]) -> usize {
records
.iter()
.map(|r| RECORD_HEADER_SIZE as usize + r.len())
.sum()
}
pub(crate) struct WalWriter<F: FileIo> {
options: WriteAheadOptions,
fio: F,
file_id: u64,
file_length: u64,
files: BTreeMap<u64, FileMeta>,
recv: flume::Receiver<WriterCommand>,
events: flume::Sender<WriterEvent>,
}
impl<F: FileIo + 'static> WalWriter<F> {
pub(crate) fn launch(
options: WriteAheadOptions,
active_id: u64,
active_path: PathBuf,
mut files: BTreeMap<u64, FileMeta>,
) -> Result<(flume::Sender<WriterCommand>, flume::Receiver<WriterEvent>)> {
let (tx, rx) = flume::unbounded();
let (event_tx, event_rx) = flume::unbounded();
let mut fio = F::open(&active_path)?;
let len = fio.len()?;
let file_length = if len == 0 {
write_header(&mut fio)?;
fio.sync()?;
sync_dir(&options.log_dir)?;
FILE_HEADER_SIZE
} else if len < FILE_HEADER_SIZE {
return Err(anyhow!(LogfileError::InvalidHeader));
} else {
len
};
files.insert(
active_id,
FileMeta {
path: active_path,
size: file_length,
seal_timestamp_ms: None,
},
);
let mut writer = Self {
options,
fio,
file_id: active_id,
file_length,
files,
recv: rx,
events: event_tx,
};
writer.apply_retention();
thread::Builder::new()
.name("wal-writer".into())
.spawn(move || writer.actor_loop())
.context("Failed to spawn writer thread")?;
Ok((tx, event_rx))
}
fn actor_loop(mut self) {
while let Ok(first) = self.recv.recv() {
let mut writes: Vec<(ReplyTx, Vec<Vec<u8>>)> = Vec::new();
let mut trims: Vec<(flume::Sender<TrimStats>, u64)> = Vec::new();
let mut queued_bytes = 0usize;
match first {
WriterCommand::Write(reply, data) => {
queued_bytes += batch_bytes(&data);
writes.push((reply, data));
}
WriterCommand::Trim(reply, upto) => trims.push((reply, upto)),
}
while queued_bytes < MAX_GROUP_COMMIT_BYTES {
match self.recv.try_recv() {
Ok(WriterCommand::Write(reply, data)) => {
queued_bytes += batch_bytes(&data);
writes.push((reply, data));
}
Ok(WriterCommand::Trim(reply, upto)) => trims.push((reply, upto)),
Err(_) => break,
}
}
if !writes.is_empty() {
self.commit_group(writes);
}
for (reply, upto) in trims {
let _ = reply.send(self.trim_before(upto));
}
if let Err(e) = self.maybe_rotate() {
error!("log rotation failed: {e:#}");
}
}
trace!("wal writer actor shut down");
}
fn commit_group(&mut self, group: Vec<(ReplyTx, Vec<Vec<u8>>)>) {
let mut valid: Vec<(ReplyTx, Vec<Vec<u8>>)> = Vec::with_capacity(group.len());
let mut total = 0usize;
for (reply, records) in group {
if records.iter().any(|r| r.len() > u32::MAX as usize) {
let _ = reply.send(Err(anyhow!(LogfileError::RecordTooLarge)));
continue;
}
total += batch_bytes(&records);
valid.push((reply, records));
}
if valid.is_empty() {
return;
}
if total == 0 {
for (reply, _) in valid {
let _ = reply.send(Ok(WriteAck {
file_id: self.file_id,
offsets: Vec::new(),
}));
}
return;
}
let mut buf = Vec::with_capacity(total);
let mut current_offset = self.file_length;
let mut per_caller_offsets: Vec<Vec<u64>> = Vec::with_capacity(valid.len());
for (_, records) in &valid {
let offsets = encode_records(&mut buf, records, current_offset);
current_offset += batch_bytes(records) as u64;
per_caller_offsets.push(offsets);
}
let result = self
.fio
.write_at(self.file_length, &buf)
.context("Failed to write records")
.and_then(|_| self.fio.sync().context("Failed to sync records"));
match result {
Ok(()) => {
self.file_length = current_offset;
if let Some(meta) = self.files.get_mut(&self.file_id) {
meta.size = current_offset;
}
for ((reply, _), offsets) in valid.into_iter().zip(per_caller_offsets) {
let _ = reply.send(Ok(WriteAck {
file_id: self.file_id,
offsets,
}));
}
}
Err(e) => {
let msg = format!("{e:#}");
for (reply, _) in valid {
let _ = reply.send(Err(anyhow!("{}", msg.clone())));
}
}
}
}
fn maybe_rotate(&mut self) -> Result<()> {
if self.file_length <= self.options.max_file_size {
return Ok(());
}
let next_id = self.file_id + 1;
let path = log_file_path(&self.options.log_dir, next_id);
let mut fio = F::open(&path)?;
write_header(&mut fio)?;
fio.sync()?;
sync_dir(&self.options.log_dir)?;
self.files.insert(
next_id,
FileMeta {
path: path.clone(),
size: FILE_HEADER_SIZE,
seal_timestamp_ms: None,
},
);
let _ = self.events.send(WriterEvent::Created(next_id, path));
let ts = now_ms();
match append_footer(&mut self.fio, self.file_length, ts) {
Ok(()) => {
if let Some(meta) = self.files.get_mut(&self.file_id) {
meta.size = self.file_length + FOOTER_SIZE;
meta.seal_timestamp_ms = Some(ts);
}
let _ = self.events.send(WriterEvent::Sealed(self.file_id));
debug!("sealed log file {}", self.file_id);
}
Err(e) => warn!("failed to seal log file {}: {e:#}", self.file_id),
}
self.fio = fio;
self.file_id = next_id;
self.file_length = FILE_HEADER_SIZE;
self.apply_retention();
Ok(())
}
fn apply_retention(&mut self) {
if let Some(meta) = self.files.get_mut(&self.file_id) {
meta.size = self.file_length;
}
let mut doomed: Vec<u64> = Vec::new();
let ttl_ms = self.options.retention.ttl.as_millis() as u64;
if ttl_ms > 0 {
let cutoff = now_ms().saturating_sub(ttl_ms);
doomed.extend(
self.files
.iter()
.filter(|(id, meta)| {
**id != self.file_id && meta.seal_timestamp_ms.is_some_and(|ts| ts < cutoff)
})
.map(|(id, _)| *id),
);
}
let max_total = self.options.retention.max_total_size;
if max_total > 0 {
let doomed_size: u64 = doomed
.iter()
.filter_map(|id| self.files.get(id))
.map(|m| m.size)
.sum();
let mut total: u64 = self.files.values().map(|m| m.size).sum::<u64>() - doomed_size;
for (id, meta) in &self.files {
if total <= max_total {
break;
}
if *id == self.file_id || meta.seal_timestamp_ms.is_none() || doomed.contains(id) {
continue;
}
doomed.push(*id);
total -= meta.size;
}
}
self.delete_files(doomed, "retention");
}
fn trim_before(&mut self, upto: u64) -> TrimStats {
let doomed: Vec<u64> = self
.files
.iter()
.filter(|(id, meta)| {
**id < upto && **id != self.file_id && meta.seal_timestamp_ms.is_some()
})
.map(|(id, _)| *id)
.collect();
self.delete_files(doomed, "trim")
}
fn delete_files(&mut self, doomed: Vec<u64>, why: &str) -> TrimStats {
let mut stats = TrimStats::default();
if doomed.is_empty() {
return stats;
}
for id in doomed {
if let Some(meta) = self.files.remove(&id) {
debug!("{}: deleting logfile {}", why, id);
if let Err(e) = std::fs::remove_file(&meta.path) {
warn!("{}: failed to delete {}: {e}", why, meta.path.display());
self.files.insert(id, meta);
continue;
}
stats.files_deleted += 1;
stats.bytes_reclaimed += meta.size;
let _ = self.events.send(WriterEvent::Deleted(id));
}
}
if let Err(e) = sync_dir(&self.options.log_dir) {
warn!("{}: failed to sync log dir: {e:#}", why);
}
stats
}
}