use std::fs::{File, OpenOptions};
use std::io::{Cursor, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use anyhow::{Result, anyhow};
use znippy_common::arrow::datatypes::Schema;
use znippy_common::arrow::ipc::reader::StreamReader;
use znippy_common::arrow::ipc::writer::StreamWriter;
use znippy_common::arrow::record_batch::RecordBatch;
use znippy_common::{ReservedSection, read_reserved_section_bytes};
#[derive(Debug, Default)]
pub struct PushLogScan {
pub pushes: Vec<RecordBatch>,
pub torn_tail_bytes: u64,
}
impl PushLogScan {
pub fn is_clean(&self) -> bool {
self.torn_tail_bytes == 0
}
}
pub fn encode_frame(schema: &Arc<Schema>, batch: &RecordBatch) -> Result<Vec<u8>> {
let mut buf = Vec::new();
{
let mut w = StreamWriter::try_new(&mut buf, schema)
.map_err(|e| anyhow!("push log: opening a frame: {e}"))?;
w.write(batch).map_err(|e| anyhow!("push log: writing a frame: {e}"))?;
w.finish().map_err(|e| anyhow!("push log: closing a frame: {e}"))?;
}
Ok(buf)
}
const IPC_EOS: [u8; 8] = [0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x00];
pub fn scan_frames(bytes: &[u8]) -> PushLogScan {
let mut out = PushLogScan::default();
let mut off = 0usize;
while off < bytes.len() {
let rest = &bytes[off..];
let Ok(mut reader) = StreamReader::try_new(Cursor::new(rest), None) else {
break;
};
let mut batches = Vec::new();
let mut errored = false;
loop {
match reader.next() {
Some(Ok(b)) => batches.push(b),
Some(Err(_)) => {
errored = true;
break;
}
None => break,
}
}
if errored {
break;
}
let consumed = reader.get_ref().position() as usize;
if consumed < IPC_EOS.len()
|| batches.len() != 1
|| rest[consumed - IPC_EOS.len()..consumed] != IPC_EOS
{
break;
}
out.pushes.append(&mut batches);
off += consumed;
}
out.torn_tail_bytes = (bytes.len() - off) as u64;
out
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CompactionPolicy {
pub after_frames: usize,
}
impl CompactionPolicy {
pub const fn never() -> Self {
Self { after_frames: 0 }
}
pub fn should_compact(&self, frames: usize) -> bool {
self.after_frames != 0 && frames >= self.after_frames
}
}
impl Default for CompactionPolicy {
fn default() -> Self {
Self { after_frames: 1024 }
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CompactionReport {
pub frames_before: usize,
pub frames_after: usize,
pub rows: usize,
pub bytes_before: u64,
pub bytes_after: u64,
}
fn compacting_path(path: &Path) -> PathBuf {
let name = path
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_else(|| "log".to_owned());
path.with_file_name(format!("{name}.compacting"))
}
fn sync_parent_dir(path: &Path) -> Result<()> {
let parent = match path.parent() {
Some(p) if !p.as_os_str().is_empty() => p.to_path_buf(),
_ => PathBuf::from("."),
};
File::open(parent)?.sync_all()?;
Ok(())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Finish {
Swap,
StopBeforeSwap,
}
pub struct PushLog {
path: PathBuf,
schema: Arc<Schema>,
}
impl PushLog {
pub fn new(path: impl Into<PathBuf>, schema: Arc<Schema>) -> Self {
Self { path: path.into(), schema }
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn schema(&self) -> &Arc<Schema> {
&self.schema
}
pub fn append(&self, batch: &RecordBatch) -> Result<u64> {
if batch.schema() != self.schema {
return Err(anyhow!(
"push log {}: batch schema does not match the log schema",
self.path.display()
));
}
let frame = encode_frame(&self.schema, batch)?;
let mut f = OpenOptions::new()
.create(true)
.append(true)
.open(&self.path)
.map_err(|e| anyhow!("push log {}: {e}", self.path.display()))?;
let offset = f.metadata()?.len();
f.write_all(&frame)?;
f.sync_all()?;
Ok(offset)
}
pub fn scan(&self) -> Result<PushLogScan> {
match std::fs::read(&self.path) {
Ok(b) => Ok(scan_frames(&b)),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(PushLogScan::default()),
Err(e) => Err(anyhow!("push log {}: {e}", self.path.display())),
}
}
pub fn frame_count(&self) -> Result<usize> {
Ok(self.scan()?.pushes.len())
}
pub fn compact(&self) -> Result<CompactionReport> {
self.compact_with(Finish::Swap)
}
pub fn compact_with(&self, finish: Finish) -> Result<CompactionReport> {
let scan = self.scan()?;
let frames_before = scan.pushes.len();
let bytes_before = std::fs::metadata(&self.path).map(|m| m.len()).unwrap_or(0);
if frames_before == 0 {
return Ok(CompactionReport {
frames_before,
frames_after: frames_before,
rows: 0,
bytes_before,
bytes_after: bytes_before,
});
}
let merged = znippy_common::arrow::compute::concat_batches(&self.schema, scan.pushes.iter())
.map_err(|e| anyhow!("compacting {}: {e}", self.path.display()))?;
let rows = merged.num_rows();
let bytes = encode_frame(&self.schema, &merged)?;
let staged = compacting_path(&self.path);
{
let mut f = File::create(&staged)?;
f.write_all(&bytes)?;
f.sync_all()?;
}
if finish == Finish::StopBeforeSwap {
return Ok(CompactionReport {
frames_before,
frames_after: 1,
rows,
bytes_before,
bytes_after: bytes.len() as u64,
});
}
std::fs::rename(&staged, &self.path)?;
sync_parent_dir(&self.path)?;
Ok(CompactionReport {
frames_before,
frames_after: 1,
rows,
bytes_before,
bytes_after: bytes.len() as u64,
})
}
pub fn maybe_compact(&self, policy: CompactionPolicy) -> Result<Option<CompactionReport>> {
if policy == CompactionPolicy::never() {
return Ok(None);
}
let frames = self.frame_count()?;
if !policy.should_compact(frames) {
return Ok(None);
}
self.compact().map(Some)
}
pub fn seal_section(&self, module_name: &str) -> Result<ReservedSection> {
let scan = self.scan()?;
Ok(ReservedSection::arrow(module_name, self.schema.clone(), scan.pushes))
}
}
pub fn read_sealed(archive: &Path, module_name: &str) -> Result<Option<Vec<RecordBatch>>> {
let Some(bytes) = read_reserved_section_bytes(archive, module_name)? else {
return Ok(None);
};
let reader = StreamReader::try_new(Cursor::new(&bytes[..]), None)
.map_err(|e| anyhow!("{module_name}: {e}"))?;
let mut out = Vec::new();
for b in reader {
out.push(b.map_err(|e| anyhow!("{module_name}: {e}"))?);
}
Ok(Some(out))
}
#[doc(hidden)]
pub fn truncate_for_test(path: &Path, len: u64) -> Result<()> {
let f = File::options().write(true).open(path)?;
f.set_len(len)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use znippy_common::arrow::array::{StringArray, UInt64Array};
use znippy_common::arrow::datatypes::{DataType, Field};
fn schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("name", DataType::Utf8, false),
Field::new("seq", DataType::UInt64, false),
]))
}
fn batch(names: &[&str], seq: u64) -> RecordBatch {
RecordBatch::try_new(
schema(),
vec![
Arc::new(StringArray::from(names.to_vec())),
Arc::new(UInt64Array::from(vec![seq; names.len()])),
],
)
.unwrap()
}
fn tmpdir(tag: &str) -> PathBuf {
let ns = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos();
let d = std::env::temp_dir().join(format!("znippy_pushlog_{tag}_{ns}"));
std::fs::create_dir_all(&d).unwrap();
d
}
#[test]
fn one_push_is_one_frame_is_one_batch() {
let dir = tmpdir("frames");
let log = PushLog::new(dir.join("refs.log"), schema());
log.append(&batch(&["a", "b"], 1)).unwrap();
log.append(&batch(&["c"], 2)).unwrap();
log.append(&batch(&["d", "e", "f"], 3)).unwrap();
let scan = log.scan().unwrap();
assert!(scan.is_clean(), "no crash happened, tail must be clean");
assert_eq!(scan.pushes.len(), 3, "one RecordBatch per push");
assert_eq!(scan.pushes[0].num_rows(), 2);
assert_eq!(scan.pushes[1].num_rows(), 1);
assert_eq!(scan.pushes[2].num_rows(), 3);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn a_torn_final_frame_loses_only_that_push() {
let dir = tmpdir("torn");
let path = dir.join("refs.log");
let log = PushLog::new(&path, schema());
log.append(&batch(&["a", "b"], 1)).unwrap();
log.append(&batch(&["c"], 2)).unwrap();
let third_start = log.append(&batch(&["doomed"], 3)).unwrap();
let full_len = std::fs::metadata(&path).unwrap().len();
assert!(full_len > third_start, "third frame must have real bytes");
let intact = std::fs::read(&path).unwrap();
for cut in (third_start + 1)..full_len {
std::fs::write(&path, &intact).unwrap();
truncate_for_test(&path, cut).unwrap();
let scan = log.scan().unwrap();
assert_eq!(
scan.pushes.len(),
2,
"cut at {cut}: the torn push must vanish whole, leaving exactly the two \
that completed (got {} batches)",
scan.pushes.len()
);
assert_eq!(scan.pushes[0].num_rows(), 2, "cut at {cut}: push 1 damaged");
assert_eq!(scan.pushes[1].num_rows(), 1, "cut at {cut}: push 2 damaged");
assert_eq!(
scan.torn_tail_bytes,
cut - third_start,
"cut at {cut}: the torn tail must be reported, not silently swallowed"
);
let names = scan.pushes[0]
.column(0)
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
assert_eq!(names.value(0), "a", "cut at {cut}: push 1 content damaged");
assert_eq!(names.value(1), "b");
}
std::fs::write(&path, &intact).unwrap();
let scan = log.scan().unwrap();
assert_eq!(scan.pushes.len(), 3);
assert!(scan.is_clean());
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn trailing_garbage_is_a_torn_tail_not_an_error() {
let dir = tmpdir("garbage");
let path = dir.join("refs.log");
let log = PushLog::new(&path, schema());
log.append(&batch(&["a"], 1)).unwrap();
for junk in [
&b"\x00"[..],
&b"\xff\xff\xff\xff"[..],
&b"\xff\xff\xff\xff\x10\x00\x00\x00partial"[..],
&[0xAB; 4096][..],
] {
let mut bytes = std::fs::read(&path).unwrap();
let clean_len = bytes.len();
bytes.extend_from_slice(junk);
let scan = scan_frames(&bytes);
assert_eq!(scan.pushes.len(), 1, "the complete push must survive {junk:?}");
assert_eq!(scan.torn_tail_bytes, (bytes.len() - clean_len) as u64);
}
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn a_missing_log_is_empty_and_clean() {
let dir = tmpdir("missing");
let log = PushLog::new(dir.join("nope.log"), schema());
let scan = log.scan().unwrap();
assert!(scan.pushes.is_empty());
assert!(scan.is_clean());
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn compaction_preserves_every_row_and_its_order() {
let dir = tmpdir("compact");
let path = dir.join("refs.log");
let log = PushLog::new(&path, schema());
for i in 0..200u64 {
log.append(&batch(&["a", "b", "c", "d"], i)).unwrap();
}
let before = log.scan().unwrap();
let before_rows: usize = before.pushes.iter().map(|b| b.num_rows()).sum();
assert_eq!(before.pushes.len(), 200);
let report = log.compact().unwrap();
assert_eq!(report.frames_before, 200);
assert_eq!(report.frames_after, 1, "everything must fold into one frame");
assert_eq!(report.rows, before_rows, "compaction must not drop a row");
let after = log.scan().unwrap();
assert!(after.is_clean(), "a compacted log must scan clean");
assert_eq!(after.pushes.len(), 1);
assert_eq!(
after.pushes[0].num_rows(),
before_rows,
"the merged batch must carry every row the frames did"
);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn compaction_actually_shrinks_the_log() {
let dir = tmpdir("compact_size");
let log = PushLog::new(dir.join("refs.log"), schema());
for i in 0..200u64 {
log.append(&batch(&["a", "b", "c", "d"], i)).unwrap();
}
let report = log.compact().unwrap();
assert!(
report.bytes_after * 2 < report.bytes_before,
"compaction saved almost nothing: {} -> {} bytes. It costs a rewrite and the \
frame boundaries; if it does not pay for them it should not run.",
report.bytes_before,
report.bytes_after
);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn an_interrupted_compaction_leaves_the_old_log_serving_at_every_cut() {
let dir = tmpdir("compact_crash");
let path = dir.join("refs.log");
let log = PushLog::new(&path, schema());
for i in 0..40u64 {
log.append(&batch(&["x", "y"], i)).unwrap();
}
let intact = std::fs::read(&path).unwrap();
let expected_rows: usize = log.scan().unwrap().pushes.iter().map(|b| b.num_rows()).sum();
let report = log.compact_with(Finish::StopBeforeSwap).unwrap();
let staged = compacting_path(&path);
assert!(staged.exists(), "the staged replacement must exist to cut into");
let staged_bytes = std::fs::read(&staged).unwrap();
assert!(staged_bytes.len() > 32);
for cut in 0..staged_bytes.len() {
std::fs::write(&staged, &staged_bytes[..cut]).unwrap();
assert_eq!(
std::fs::read(&path).unwrap(),
intact,
"cut at {cut}: the LIVE log was modified by a compaction that never \
completed — it must never be written in place"
);
let scan = log.scan().unwrap();
assert!(scan.is_clean(), "cut at {cut}: the live log stopped scanning clean");
assert_eq!(
scan.pushes.len(),
40,
"cut at {cut}: the live log lost pushes to an interrupted compaction"
);
let rows: usize = scan.pushes.iter().map(|b| b.num_rows()).sum();
assert_eq!(rows, expected_rows, "cut at {cut}: rows went missing");
}
std::fs::remove_file(&staged).ok();
let done = log.compact().unwrap();
assert_eq!(done.frames_after, 1);
assert_eq!(done.rows, expected_rows);
assert_eq!(report.rows, expected_rows);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn compacted_and_uncompacted_logs_are_the_same_format() {
let dir = tmpdir("compat");
let path = dir.join("refs.log");
let log = PushLog::new(&path, schema());
for i in 0..8u64 {
log.append(&batch(&["a"], i)).unwrap();
}
let uncompacted = log.scan().unwrap();
log.compact().unwrap();
let compacted = log.scan().unwrap();
let flat = |s: &PushLogScan| -> Vec<u64> {
let mut v = Vec::new();
for b in &s.pushes {
let c = b.column(1).as_any().downcast_ref::<UInt64Array>().unwrap();
v.extend((0..c.len()).map(|i| c.value(i)));
}
v
};
assert_eq!(flat(&uncompacted), flat(&compacted), "compaction changed the row sequence");
log.append(&batch(&["z"], 99)).unwrap();
let after = log.scan().unwrap();
assert!(after.is_clean());
assert_eq!(after.pushes.len(), 2, "a compacted log must still accept appends");
assert_eq!(*flat(&after).last().unwrap(), 99);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn a_foreign_schema_is_refused_at_append() {
let dir = tmpdir("schema");
let log = PushLog::new(dir.join("refs.log"), schema());
let other = Arc::new(Schema::new(vec![Field::new("x", DataType::Utf8, false)]));
let foreign =
RecordBatch::try_new(other, vec![Arc::new(StringArray::from(vec!["v"]))]).unwrap();
let err = log.append(&foreign).unwrap_err().to_string();
assert!(err.contains("schema"), "expected a schema error, got: {err}");
std::fs::remove_dir_all(&dir).ok();
}
}