Something went wrong. Try again.
very fast at protocol indexer with flexible filtering, xrpc queries, cursor-backed event stream, and more, built on fjall
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361use 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 oncepub(crate) const TOP_LEVEL: u32 = 2;/// readers with an older manifest still get replaced chunks for this longconst GRACE: Duration = Duration::from_secs(15 * 60);const STATE_FILE: &str = "state.json";/// bump when the bytes of a chunk change. older chunks get rebuiltconst FORMAT: u32 = 1;
#[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<u64>, /// 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<u64>,}
#[derive(Debug, Default, Serialize, Deserialize)]struct State { stream: String, /// reads as 0 when the list was saved without one #[serde(default)] format: u32, generation: u64, chunks: Vec<Chunk>, retired: Vec<Retired>,}
impl State { fn upgrade(&mut self) { if self.format == FORMAT { return; } for chunk in &mut self.chunks { chunk.dirty = true; } self.retired.clear(); self.format = FORMAT; self.generation += 1; }
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<u64>, pub chunks: Vec<ManifestChunk>,}
#[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<State>, marks: Ks<Cursors>,}
/// saved chunk lists hold these, changing it breaks thempub(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<u8> { [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<Cursors>) -> Result<Self> { 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::<State>(&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 mut state = match state { Some(state) if state.stream == stream => state, _ => State { stream: stream.to_owned(), format: FORMAT, ..State::default() }, }; state.upgrade(); 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::<u64>())) }
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<ManifestChunk> = 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<PathBuf> { 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<Chunk> { self.state .lock() .chunks .iter() .filter(|c| c.dirty) .cloned() .collect() }
pub(crate) fn mergeable(&self) -> Option<Vec<Chunk>> { 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<Vec<u64>> { 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<String> = { 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(()) }}