use std::collections::HashSet; use std::io::Write; use std::path::{Path, PathBuf}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use fjall::OwnedWriteBatch; use miette::{IntoDiagnostic, Result, WrapErr}; use parking_lot::Mutex; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use super::keys::STREAM_CHUNKS_DIRTY_PREFIX; use super::schema::{Cursors, Ks}; pub(crate) const FANIN: usize = 4; /// 16 spans, small enough for a reader to fetch several at once pub(crate) const TOP_LEVEL: u32 = 2; /// readers with an older manifest still get replaced chunks for this long const GRACE: Duration = Duration::from_secs(15 * 60); const STATE_FILE: &str = "state.json"; #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct Chunk { /// ids with no event count too pub first: u64, pub last: u64, pub level: u32, pub events: u64, pub bytes: u64, pub sha256: String, /// lookups binary search this, keep it sorted pub dids: Vec, /// holds a redacted or erased body and isn't served until rebuilt pub dirty: bool, } #[derive(Clone, Debug, Serialize, Deserialize)] struct Retired { sha256: String, at: u64, dids: Vec, } #[derive(Debug, Default, Serialize, Deserialize)] struct State { stream: String, generation: u64, chunks: Vec, retired: Vec, } impl State { fn reach(&mut self, hash: u64) -> bool { let mut changed = false; for chunk in &mut self.chunks { if !chunk.dirty && chunk.dids.binary_search(&hash).is_ok() { chunk.dirty = true; changed = true; } } let before = self.retired.len(); self.retired .retain(|r| r.dids.binary_search(&hash).is_err()); changed || self.retired.len() != before } } /// stops at the first dirty chunk, even if later ones are clean #[derive(Debug, Serialize)] pub struct Manifest { pub stream: String, pub generation: u64, /// inclusive, `/stream` picks up at through + 1 pub through: Option, pub chunks: Vec, } #[derive(Debug, Serialize)] pub struct ManifestChunk { pub first: u64, pub last: u64, pub events: u64, pub bytes: u64, pub sha256: String, } pub struct StreamChunks { dir: PathBuf, pub(crate) span: u64, state: Mutex, marks: Ks, } /// saved chunk lists hold these, changing it breaks them pub(crate) fn did_hash(did: &str) -> u64 { let digest = Sha256::digest(did.as_bytes()); u64::from_be_bytes(digest[..8].try_into().expect("8 bytes")) } fn mark_key(hash: u64) -> Vec { [STREAM_CHUNKS_DIRTY_PREFIX, &hash.to_be_bytes()].concat() } fn now_secs() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .map_or(0, |d| d.as_secs()) } impl StreamChunks { pub(crate) fn open(dir: PathBuf, span: u64, stream: &str, marks: Ks) -> Result { std::fs::create_dir_all(&dir) .into_diagnostic() .wrap_err_with(|| format!("can't create {}", dir.display()))?; let state = match std::fs::read(dir.join(STATE_FILE)) { Ok(bytes) => serde_json::from_slice::(&bytes) .inspect_err( |e| tracing::warn!(err = %e, "unreadable stream chunk list, starting over"), ) .ok(), Err(e) if e.kind() == std::io::ErrorKind::NotFound => None, Err(e) => return Err(e).into_diagnostic(), }; let state = match state { Some(state) if state.stream == stream => state, _ => State { stream: stream.to_owned(), ..State::default() }, }; let this = Self { dir, span, state: Mutex::new(state), marks, }; this.persist(&this.state.lock())?; this.settle()?; this.collect_garbage()?; Ok(this) } pub(crate) fn path(&self, sha256: &str) -> PathBuf { self.dir.join(format!("{sha256}.zst")) } pub(crate) fn tmp_path(&self) -> PathBuf { self.dir .join(format!("build-{}.tmp", rand::random::())) } fn persist(&self, state: &State) -> Result<()> { let tmp = self.dir.join(format!("{STATE_FILE}.tmp")); let mut file = std::fs::File::create(&tmp).into_diagnostic()?; file.write_all(&serde_json::to_vec(state).into_diagnostic()?) .into_diagnostic()?; file.sync_all().into_diagnostic()?; std::fs::rename(&tmp, self.dir.join(STATE_FILE)).into_diagnostic() } pub fn manifest(&self) -> Manifest { let state = self.state.lock(); let chunks: Vec = state .chunks .iter() .take_while(|c| !c.dirty) .map(|c| ManifestChunk { first: c.first, last: c.last, events: c.events, bytes: c.bytes, sha256: c.sha256.clone(), }) .collect(); Manifest { stream: state.stream.clone(), generation: state.generation, through: chunks.last().map(|c| c.last), chunks, } } pub fn servable(&self, sha256: &str) -> Option { let state = self.state.lock(); let listed = state.chunks.iter().any(|c| c.sha256 == sha256 && !c.dirty); let cutoff = now_secs().saturating_sub(GRACE.as_secs()); let retired = state .retired .iter() .any(|r| r.sha256 == sha256 && r.at >= cutoff); (listed || retired).then(|| self.path(sha256)) } pub(crate) fn sealed_next(&self) -> u64 { self.state.lock().chunks.last().map_or(0, |c| c.last + 1) } pub(crate) fn dirty(&self) -> Vec { self.state .lock() .chunks .iter() .filter(|c| c.dirty) .cloned() .collect() } pub(crate) fn mergeable(&self) -> Option> { let state = self.state.lock(); let list = &state.chunks; if list.iter().any(|c| c.dirty) { return None; } let run = list.get(list.len().checked_sub(FANIN)?..)?; let level = run[0].level; (level < TOP_LEVEL && run.iter().all(|c| c.level == level)).then(|| run.to_vec()) } /// goes in the same batch as the body change so a crash can't split them pub(crate) fn mark(&self, batch: &mut OwnedWriteBatch, did: &str) { batch.insert(&self.marks, mark_key(did_hash(did)), []); } fn marked(&self) -> Result> { self.marks .prefix(STREAM_CHUNKS_DIRTY_PREFIX) .map(|guard| { let key = guard.key().into_diagnostic()?; let hash = key[STREAM_CHUNKS_DIRTY_PREFIX.len()..].try_into(); Ok(u64::from_be_bytes(hash.into_diagnostic()?)) }) .collect() } /// a build round calls this before reading dirty() pub(crate) fn settle(&self) -> Result<()> { let marked = self.marked()?; if marked.is_empty() { return Ok(()); } { let mut state = self.state.lock(); let mut changed = false; for &hash in &marked { changed |= state.reach(hash); } if changed { state.generation += 1; } self.persist(&state)?; } for hash in marked { self.marks.remove(mark_key(hash)).into_diagnostic()?; } Ok(()) } /// not durable. anything that has to survive a crash needs mark pub(crate) fn touch(&self, did: &str) { let mut state = self.state.lock(); if state.reach(did_hash(did)) { state.generation += 1; if let Err(e) = self.persist(&state) { tracing::error!(err = %e, "can't persist the stream chunk list"); } } } pub(crate) fn publish(&self, replaces: &[Chunk], mut chunk: Chunk, tmp: &Path) -> Result<()> { std::fs::rename(tmp, self.path(&chunk.sha256)).into_diagnostic()?; let mut state = self.state.lock(); // marks since settle are changes this build may have missed chunk.dirty = self .marked()? .iter() .any(|hash| chunk.dids.binary_search(hash).is_ok()); let list = &mut state.chunks; let mut gone = Vec::new(); let at = match replaces.first() { Some(first) => { let at = list.iter().position(|c| c.sha256 == first.sha256); let end = at.map_or(0, |at| at + replaces.len()); let same = at.and_then(|at| list.get(at..end)).is_some_and(|current| { current .iter() .zip(replaces) .all(|(a, b)| a.sha256 == b.sha256) }); let Some(at) = at.filter(|_| same) else { miette::bail!("the chunks a build replaces changed under it"); }; gone.extend(list.drain(at..end)); at } None => list.len(), }; list.insert(at, chunk); let now = now_secs(); // a chunk dirtied during the build holds an old body and can't be served as retired let retired = gone.into_iter().filter(|c| !c.dirty).map(|c| Retired { sha256: c.sha256, at: now, dids: c.dids, }); state.retired.extend(retired); state.generation += 1; self.persist(&state) } pub(crate) fn collect_garbage(&self) -> Result<()> { let keep: HashSet = { let mut state = self.state.lock(); let cutoff = now_secs().saturating_sub(GRACE.as_secs()); let before = state.retired.len(); state.retired.retain(|r| r.at >= cutoff); if state.retired.len() != before { self.persist(&state)?; } state .chunks .iter() .map(|c| c.sha256.clone()) .chain(state.retired.iter().map(|r| r.sha256.clone())) .collect() }; for entry in std::fs::read_dir(&self.dir).into_diagnostic()? { let path = entry.into_diagnostic()?.path(); let Some(name) = path.file_name().and_then(|n| n.to_str()) else { continue; }; let stale = match name.strip_suffix(".zst") { Some(sha256) => !keep.contains(sha256), None => name.starts_with("build-"), }; if stale { let _ = std::fs::remove_file(&path); } } Ok(()) } }