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";
/// bump when the bytes of a chunk change. older chunks get rebuilt
const 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 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<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(())
}
}