use std::collections::HashMap;
use std::fs::File;
use std::os::unix::fs::FileExt;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use anyhow::{Result, anyhow, bail};
use arrow::ipc::reader::StreamReader;
use arrow::ipc::writer::StreamWriter;
use crate::archive::ZnippyReader;
use crate::codec::CompressCtx;
use crate::index::{ChunkLoc, data_subindex_schema};
use crate::meta_sink::{ArchiveMetaSink, ArrowIpcSink, GroupKey, ReservedSectionBuilder};
use crate::meta_sink_append::{base_batch_from_rows, decode_base_rows, write_blobs};
const JOURNAL_INFIX: &str = ".hot.";
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct HotAppendReport {
pub rows_added: u64,
pub rows_replaced: u64,
pub blob_append_offset: u64,
pub blob_bytes_added: u64,
pub journal_bytes_added: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HotSealReport {
pub archive: PathBuf,
pub rows_sealed: u64,
pub segments_removed: usize,
pub journal_bytes_reclaimed: u64,
pub sealed_total_bytes: u64,
}
pub struct HotArchive {
archive: PathBuf,
blobs: Arc<File>,
blob_end: u64,
paths: Vec<String>,
locs: Vec<ChunkLoc>,
by_path: HashMap<String, usize>,
journal: StreamWriter<File>,
journal_path: PathBuf,
journal_len: u64,
ctx: std::sync::Mutex<CompressCtx>,
policy: crate::SkipPolicy,
}
impl HotArchive {
pub fn open(archive: &Path, level: i32, policy: crate::SkipPolicy) -> Result<Self> {
let archive = archive.to_path_buf();
let segments = journal_segments(&archive)?;
let (mut paths, mut locs) = (Vec::new(), Vec::new());
let mut reheated = false;
if !segments.is_empty() {
for (_, seg) in &segments {
let (p, l) = read_segment(seg)?;
paths.extend(p);
locs.extend(l);
}
} else if archive.is_file() && crate::index::read_znippy_full_manifest(&archive).is_ok() {
let (entries, _) = crate::index::read_znippy_full_manifest(&archive)?;
let (p, l) = crate::meta_sink_append::recover_rows(&archive, &entries)?;
paths = p;
locs = l;
reheated = true;
}
let (paths, locs) = live_rows(paths, locs);
let blob_end = locs
.iter()
.map(|l| l.blob_offset + l.blob_size)
.max()
.unwrap_or(0);
let next = segments.last().map(|(n, _)| n + 1).unwrap_or(0);
let journal_path = segment_path(&archive, next);
let mut journal = new_segment(&journal_path)?;
if reheated {
let batch = base_batch_from_rows(&paths, &locs)?;
journal
.write(&batch)
.map_err(|e| anyhow!("hot: seeding journal from sealed index: {e}"))?;
journal.flush().map_err(|e| anyhow!("hot: journal flush: {e}"))?;
journal
.get_ref()
.sync_all()
.map_err(|e| anyhow!("hot: journal fsync: {e}"))?;
sync_parent_dir(&journal_path);
}
let journal_len = std::fs::metadata(&journal_path).map(|m| m.len()).unwrap_or(0);
let blobs = Arc::new(
std::fs::OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&archive)
.map_err(|e| anyhow!("hot: open {} : {e}", archive.display()))?,
);
if blobs.metadata()?.len() > blob_end {
blobs.set_len(blob_end)?;
blobs.sync_all()?;
}
let mut by_path = HashMap::with_capacity(paths.len());
for (i, p) in paths.iter().enumerate() {
by_path.insert(p.clone(), i);
}
Ok(Self {
archive,
blobs,
blob_end,
paths,
locs,
by_path,
journal,
journal_path,
journal_len,
ctx: std::sync::Mutex::new(CompressCtx::new(level)?),
policy,
})
}
pub fn is_hot(archive: &Path) -> bool {
journal_segments(archive).map(|s| !s.is_empty()).unwrap_or(false)
}
pub fn archive(&self) -> &Path {
&self.archive
}
pub fn rows(&self) -> u64 {
self.paths.len() as u64
}
pub fn blob_end(&self) -> u64 {
self.blob_end
}
pub fn append(&mut self, files: &[(String, Vec<u8>)]) -> Result<HotAppendReport> {
if files.is_empty() {
return Ok(HotAppendReport {
blob_append_offset: self.blob_end,
..Default::default()
});
}
let blob_append_offset = self.blob_end;
let (new_paths, new_locs, cursor) =
write_blobs(&self.blobs, self.blob_end, files, &mut self.ctx.lock().unwrap(), self.policy)?;
let blob_bytes_added = cursor - blob_append_offset;
self.blobs.sync_all()?;
self.blob_end = cursor;
let batch = base_batch_from_rows(&new_paths, &new_locs)?;
self.journal
.write(&batch)
.map_err(|e| anyhow!("hot: journal write: {e}"))?;
self.journal.flush().map_err(|e| anyhow!("hot: journal flush: {e}"))?;
self.journal
.get_ref()
.sync_all()
.map_err(|e| anyhow!("hot: journal fsync: {e}"))?;
let journal_len = std::fs::metadata(&self.journal_path)?.len();
let journal_bytes_added = journal_len - self.journal_len;
self.journal_len = journal_len;
let mut rows_replaced = 0u64;
for (p, l) in new_paths.into_iter().zip(new_locs) {
match self.by_path.get(&p).copied() {
Some(i) => {
rows_replaced += 1;
self.locs[i] = l;
}
None => {
self.by_path.insert(p.clone(), self.paths.len());
self.paths.push(p);
self.locs.push(l);
}
}
}
Ok(HotAppendReport {
rows_added: batch.num_rows() as u64,
rows_replaced,
blob_append_offset,
blob_bytes_added,
journal_bytes_added,
})
}
pub fn seal(self, reserved: Option<ReservedSectionBuilder>) -> Result<HotSealReport> {
let Self {
archive,
blobs,
blob_end,
paths,
locs,
journal,
journal_path,
..
} = self;
drop(journal);
let rows_sealed = paths.len() as u64;
let mut sink = ArrowIpcSink::new(blobs, blob_end);
if let Some(b) = reserved {
sink = sink.with_reserved_builder(b);
}
if !paths.is_empty() {
let batch = base_batch_from_rows(&paths, &locs)?;
let schema = data_subindex_schema();
sink.push_subindex(schema.as_ref(), &[batch], GroupKey {
pkg_type: 0,
repo: String::new(),
module_name: String::new(),
})?;
}
let sealed_total_bytes = Box::new(sink).finish()?;
let segments = journal_segments(&archive)?;
let mut journal_bytes_reclaimed = 0u64;
for (_, seg) in &segments {
journal_bytes_reclaimed += std::fs::metadata(seg).map(|m| m.len()).unwrap_or(0);
std::fs::remove_file(seg)
.map_err(|e| anyhow!("hot: removing journal {}: {e}", seg.display()))?;
}
sync_parent_dir(&journal_path);
Ok(HotSealReport {
archive,
rows_sealed,
segments_removed: segments.len(),
journal_bytes_reclaimed,
sealed_total_bytes,
})
}
}
impl ZnippyReader for HotArchive {
fn list_files(&self) -> Result<Vec<String>> {
Ok(self.paths.clone())
}
fn contains(&self, relative_path: &str) -> bool {
self.by_path.contains_key(relative_path)
}
fn file_size(&self, relative_path: &str) -> Option<u64> {
self.by_path
.get(relative_path)
.map(|&i| self.locs[i].uncompressed_size)
}
fn extract_file(&self, relative_path: &str) -> Result<Vec<u8>> {
let &i = self
.by_path
.get(relative_path)
.ok_or_else(|| anyhow!("file not found in archive: {relative_path}"))?;
let loc = &self.locs[i];
let mut blob = vec![0u8; loc.blob_size as usize];
self.blobs.read_exact_at(&mut blob, loc.blob_offset)?;
if loc.compressed {
let mut out = Vec::new();
crate::codec::decompress_into(&blob, &mut out)?;
Ok(out)
} else {
Ok(blob)
}
}
}
fn live_rows(paths: Vec<String>, locs: Vec<ChunkLoc>) -> (Vec<String>, Vec<ChunkLoc>) {
let mut at: HashMap<&str, usize> = HashMap::with_capacity(paths.len());
let mut order: Vec<usize> = Vec::with_capacity(paths.len());
for (i, p) in paths.iter().enumerate() {
match at.get(p.as_str()).copied() {
Some(slot) => order[slot] = i,
None => {
at.insert(p.as_str(), order.len());
order.push(i);
}
}
}
let order: Vec<usize> = order;
let mut out_p = Vec::with_capacity(order.len());
let mut out_l = Vec::with_capacity(order.len());
for &i in &order {
out_p.push(paths[i].clone());
out_l.push(locs[i].clone());
}
(out_p, out_l)
}
fn segment_path(archive: &Path, n: u32) -> PathBuf {
let mut s = archive.as_os_str().to_os_string();
s.push(format!("{JOURNAL_INFIX}{n:05}"));
PathBuf::from(s)
}
fn journal_segments(archive: &Path) -> Result<Vec<(u32, PathBuf)>> {
let dir = archive.parent().filter(|p| !p.as_os_str().is_empty());
let dir = dir.map(|d| d.to_path_buf()).unwrap_or_else(|| PathBuf::from("."));
let stem = archive
.file_name()
.and_then(|n| n.to_str())
.ok_or_else(|| anyhow!("hot: {} has no usable file name", archive.display()))?;
let prefix = format!("{stem}{JOURNAL_INFIX}");
let mut out = Vec::new();
let rd = match std::fs::read_dir(&dir) {
Ok(rd) => rd,
Err(_) => return Ok(out),
};
for e in rd.flatten() {
let name = e.file_name();
let Some(name) = name.to_str() else { continue };
let Some(tail) = name.strip_prefix(&prefix) else { continue };
let Ok(n) = tail.parse::<u32>() else { continue };
out.push((n, e.path()));
}
out.sort_by_key(|(n, _)| *n);
Ok(out)
}
const JOURNAL_ALIGNMENT: usize = 8;
fn new_segment(path: &Path) -> Result<StreamWriter<File>> {
let f = File::create(path)
.map_err(|e| anyhow!("hot: create journal segment {}: {e}", path.display()))?;
let schema = data_subindex_schema();
let opts = arrow::ipc::writer::IpcWriteOptions::try_new(
JOURNAL_ALIGNMENT,
false,
arrow::ipc::MetadataVersion::V5,
)
.map_err(|e| anyhow!("hot: journal write options: {e}"))?;
let w = StreamWriter::try_new_with_options(f, schema.as_ref(), opts)
.map_err(|e| anyhow!("hot: journal schema: {e}"))?;
sync_parent_dir(path);
Ok(w)
}
fn read_segment(path: &Path) -> Result<(Vec<String>, Vec<ChunkLoc>)> {
let bytes = std::fs::read(path)
.map_err(|e| anyhow!("hot: reading journal {}: {e}", path.display()))?;
if bytes.is_empty() {
return Ok((Vec::new(), Vec::new()));
}
let reader = match StreamReader::try_new(std::io::Cursor::new(&bytes), None) {
Ok(r) => r,
Err(_) => return Ok((Vec::new(), Vec::new())),
};
let mut paths = Vec::new();
let mut locs = Vec::new();
let mut complete: Vec<arrow::record_batch::RecordBatch> = Vec::new();
for batch in reader {
match batch {
Ok(b) => complete.push(b),
Err(_) => break, }
}
for b in &complete {
let mut ipc: Vec<u8> = Vec::new();
{
let mut w = StreamWriter::try_new(&mut ipc, b.schema().as_ref())
.map_err(|e| anyhow!("hot: re-encode: {e}"))?;
w.write(b).map_err(|e| anyhow!("hot: re-encode write: {e}"))?;
w.finish().map_err(|e| anyhow!("hot: re-encode finish: {e}"))?;
}
let (p, l) = decode_base_rows(&ipc)?;
paths.extend(p);
locs.extend(l);
}
Ok((paths, locs))
}
fn sync_parent_dir(path: &Path) {
let parent = match path.parent() {
Some(p) if !p.as_os_str().is_empty() => p.to_path_buf(),
_ => PathBuf::from("."),
};
if let Ok(f) = File::open(parent) {
let _ = f.sync_all();
}
}
pub fn seal_hot(archive: &Path, level: i32, policy: crate::SkipPolicy) -> Result<HotSealReport> {
if !HotArchive::is_hot(archive) {
bail!("{} has no journal; it is not hot", archive.display());
}
HotArchive::open(archive, level, policy)?.seal(None)
}
#[cfg(all(test, feature = "openzl"))]
mod tests {
use super::*;
use crate::{ZnippyArchive, create_archive};
use std::time::{SystemTime, UNIX_EPOCH};
fn unique_dir(tag: &str) -> PathBuf {
let ns = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos();
let d = std::env::temp_dir()
.join(format!("znippy_hot_{tag}_{ns}_{:?}", std::thread::current().id()));
std::fs::create_dir_all(&d).unwrap();
d
}
fn synth(n: usize, salt: u64) -> Vec<(String, Vec<u8>)> {
(0..n)
.map(|i| {
let g = (i.wrapping_mul(2_654_435_761) ^ salt as usize) % 997;
let p = format!("repo/grp{g:03}/file{i:08}_{salt}.bin");
let body = format!("payload {i} salt {salt} {}\n", "z".repeat(8 + (i % 40)));
(p, body.into_bytes())
})
.collect()
}
#[test]
fn seal_matches_create_archive_byte_for_byte() {
let dir = unique_dir("identity");
let files = synth(400, 11);
let reference = dir.join("ref.znippy");
create_archive(&reference, &files, 3).unwrap();
let hot = dir.join("hot.znippy");
{
let mut h = HotArchive::open(&hot, 3, crate::SkipPolicy::resolve()).unwrap();
for f in &files {
h.append(std::slice::from_ref(f)).unwrap();
}
let report = h.seal(None).unwrap();
assert_eq!(report.rows_sealed, files.len() as u64);
assert_eq!(report.segments_removed, 1, "one process, one segment");
}
let a = std::fs::read(&reference).unwrap();
let b = std::fs::read(&hot).unwrap();
assert_eq!(
a.len(),
b.len(),
"sealed length differs: create_archive {} vs hot-then-seal {}",
a.len(),
b.len()
);
let first_diff = a.iter().zip(&b).position(|(x, y)| x != y);
assert!(
first_diff.is_none(),
"hot-then-seal is NOT byte-identical to create_archive; first difference at byte {:?}",
first_diff
);
assert!(!HotArchive::is_hot(&hot), "seal must remove the journal");
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn a_hot_archive_serves_reads_before_the_seal() {
let dir = unique_dir("read");
let path = dir.join("a.znippy");
let files = synth(120, 3);
let mut h = HotArchive::open(&path, 3, crate::SkipPolicy::resolve()).unwrap();
for f in &files {
h.append(std::slice::from_ref(f)).unwrap();
}
for (p, bytes) in &files {
assert!(h.contains(p), "hot archive must contain {p}");
assert_eq!(h.file_size(p), Some(bytes.len() as u64));
assert_eq!(&h.extract_file(p).unwrap(), bytes, "hot read mismatch for {p}");
}
assert_eq!(h.list_files().unwrap().len(), files.len());
h.seal(None).unwrap();
let ar = ZnippyArchive::open(&path).unwrap();
for (p, bytes) in &files {
assert_eq!(&ar.extract_file(p).unwrap(), bytes, "sealed read mismatch for {p}");
}
assert!(!crate::locate_file(&path, &files[77].0).unwrap().is_empty());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn reopening_recovers_every_segment() {
let dir = unique_dir("reopen");
let path = dir.join("a.znippy");
let mut all: Vec<(String, Vec<u8>)> = Vec::new();
for round in 0..3u64 {
let files = synth(40, round + 1);
let mut h = HotArchive::open(&path, 3, crate::SkipPolicy::resolve()).unwrap();
assert_eq!(h.rows(), all.len() as u64, "reopen must recover prior rows");
for f in &files {
h.append(std::slice::from_ref(f)).unwrap();
}
all.extend(files);
}
let h = HotArchive::open(&path, 3, crate::SkipPolicy::resolve()).unwrap();
assert_eq!(h.rows(), all.len() as u64);
let report = h.seal(None).unwrap();
assert_eq!(report.segments_removed, 4, "three appending opens + the final one");
let ar = ZnippyArchive::open(&path).unwrap();
for (p, bytes) in &all {
assert_eq!(&ar.extract_file(p).unwrap(), bytes, "byte mismatch after restarts for {p}");
}
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn a_sealed_archive_reheats_and_keeps_its_rows() {
let dir = unique_dir("reheat");
let path = dir.join("a.znippy");
let sealed = synth(60, 5);
create_archive(&path, &sealed, 3).unwrap();
let cold_len = std::fs::metadata(&path).unwrap().len();
let mut h = HotArchive::open(&path, 3, crate::SkipPolicy::resolve()).unwrap();
assert_eq!(h.rows(), sealed.len() as u64, "re-heat must recover the sealed rows");
assert!(
h.blob_end() < cold_len,
"the metadata tail must be dropped: blob_end {} vs sealed length {cold_len}",
h.blob_end()
);
let extra = synth(20, 6);
for f in &extra {
h.append(std::slice::from_ref(f)).unwrap();
}
h.seal(None).unwrap();
let ar = ZnippyArchive::open(&path).unwrap();
for (p, bytes) in sealed.iter().chain(extra.iter()) {
assert_eq!(&ar.extract_file(p).unwrap(), bytes, "byte mismatch after re-heat for {p}");
}
let mut listed = ar.list_files().unwrap();
listed.sort();
let mut want: Vec<String> =
sealed.iter().chain(extra.iter()).map(|(p, _)| p.clone()).collect();
want.sort();
assert_eq!(listed, want);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn re_appending_a_path_replaces_it() {
let dir = unique_dir("replace");
let path = dir.join("a.znippy");
let mut h = HotArchive::open(&path, 3, crate::SkipPolicy::resolve()).unwrap();
h.append(&[("x.bin".into(), b"first version".to_vec())]).unwrap();
h.append(&[("y.bin".into(), b"other".to_vec())]).unwrap();
let r = h.append(&[("x.bin".into(), b"SECOND version, longer".to_vec())]).unwrap();
assert_eq!(r.rows_replaced, 1, "re-appending x.bin must replace, not duplicate");
assert_eq!(h.rows(), 2);
assert_eq!(h.extract_file("x.bin").unwrap(), b"SECOND version, longer");
h.seal(None).unwrap();
let ar = ZnippyArchive::open(&path).unwrap();
assert_eq!(ar.list_files().unwrap().len(), 2, "the seal must carry ONE row per path");
assert_eq!(ar.extract_file("x.bin").unwrap(), b"SECOND version, longer");
assert_eq!(ar.extract_file("y.bin").unwrap(), b"other");
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn a_torn_journal_tail_is_dropped_not_fatal() {
let dir = unique_dir("torn");
let path = dir.join("a.znippy");
let files = synth(30, 2);
{
let mut h = HotArchive::open(&path, 3, crate::SkipPolicy::resolve()).unwrap();
for f in &files {
h.append(std::slice::from_ref(f)).unwrap();
}
}
let seg = segment_path(&path, 0);
let len = std::fs::metadata(&seg).unwrap().len();
std::fs::OpenOptions::new().write(true).open(&seg).unwrap().set_len(len - 40).unwrap();
let h = HotArchive::open(&path, 3, crate::SkipPolicy::resolve()).unwrap();
let rows = h.rows();
assert!(
rows > 0 && rows < files.len() as u64,
"a torn tail must cost SOME rows and not all of them; got {rows} of {}",
files.len()
);
for p in h.list_files().unwrap() {
let want = &files.iter().find(|(q, _)| *q == p).unwrap().1;
assert_eq!(&h.extract_file(&p).unwrap(), want, "torn-journal read mismatch for {p}");
}
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn the_index_cost_of_an_append_does_not_grow_with_the_archive() {
let dir = unique_dir("flat");
let path = dir.join("a.znippy");
let mut h = HotArchive::open(&path, 3, crate::SkipPolicy::resolve()).unwrap();
let mut early = 0u64;
let mut late = 0u64;
for i in 0..2_000usize {
let f = vec![(format!("pack-{i:040x}.pack"), vec![b'q'; 64])];
let r = h.append(&f).unwrap();
if i == 10 {
early = r.journal_bytes_added;
}
if i == 1_999 {
late = r.journal_bytes_added;
}
}
assert!(early > 0, "an append must cost SOME journal bytes; the probe read 0");
assert_eq!(
early, late,
"the index cost of an append grew from {early} B at row 10 to {late} B at row 2000 — \
the append is paying for the archive again"
);
let _ = std::fs::remove_dir_all(&dir);
}
}