use std::collections::BTreeSet; use std::io::{BufWriter, Write}; use std::path::PathBuf; use std::sync::Arc; use std::time::{Duration, Instant}; use miette::{IntoDiagnostic, Result}; use sha2::{Digest, Sha256}; use tracing::{error, info, warn}; use super::indexer::stored_to_chunk_value; use crate::db::keys; use crate::db::stream_chunks::{Chunk, FANIN, did_hash}; use crate::state::AppState; use crate::types::StoredEvent; const INTERVAL: Duration = Duration::from_secs(30); const ZSTD_LEVEL: i32 = 3; const STAR_TYPE: &str = "systems.gaze.hydrant.streamChunk"; const STAR_MAGIC: [u8; 4] = [0x2a, 0x6c, 0x01, 0x00]; /// once this many firehose frames are read but not committed, ingest counts as busy const BACKLOG: usize = 64; pub(crate) fn chunk_worker(state: Arc) { lower_priority(); loop { if let Err(e) = build(&state) { error!(err = %e, "can't build stream chunks"); } if state.stop.wait(INTERVAL) { break; } } } fn lower_priority() { // linux nice is per thread and this has to be the tid #[cfg(target_os = "linux")] if unsafe { libc::setpriority(libc::PRIO_PROCESS, libc::gettid() as libc::id_t, 19) } != 0 { warn!(err = %std::io::Error::last_os_error(), "can't lower the stream chunk worker's priority"); } } /// under a firehose backlog a round seals only one span pub(crate) fn build(state: &AppState) -> Result<()> { let Some(chunks) = state.db.stream.chunks() else { return Ok(()); }; chunks.settle()?; for dirty in chunks.dirty() { write( state, dirty.first, dirty.last, dirty.level, std::slice::from_ref(&dirty), )?; } merge_all(state)?; loop { let first = chunks.sealed_next(); let last = first + chunks.span - 1; if state.db.stream.ids.head().is_none_or(|head| head < last) { break; } write(state, first, last, 0, &[])?; merge_all(state)?; if state.firehose_backlog() >= BACKLOG { break; } } chunks.collect_garbage() } /// panics when chunks are off fn chunks(state: &AppState) -> &crate::db::stream_chunks::StreamChunks { state.db.stream.chunks().expect("stream chunks are on") } fn merge_all(state: &AppState) -> Result<()> { while let Some(run) = chunks(state).mergeable() { merge(state, run)?; } Ok(()) } struct Writer { out: zstd::Encoder<'static, ChunkFile>, first: u64, last: u64, level: u32, events: u64, dids: BTreeSet, } struct ChunkFile { tmp: PathBuf, file: BufWriter, hash: Sha256, bytes: u64, } impl Write for ChunkFile { fn write(&mut self, buf: &[u8]) -> std::io::Result { let n = self.file.write(buf)?; self.hash.update(&buf[..n]); self.bytes += n as u64; Ok(n) } fn flush(&mut self) -> std::io::Result<()> { self.file.flush() } } impl ChunkFile { fn create(state: &AppState, first: u64, last: u64) -> Result { let tmp = chunks(state).tmp_path(); let mut this = Self { file: BufWriter::new(std::fs::File::create(&tmp).into_diagnostic()?), tmp, hash: Sha256::new(), bytes: 0, }; this.write_all(&header_frame(state, first, last)?) .into_diagnostic()?; Ok(this) } fn finish( self, first: u64, last: u64, level: u32, events: u64, dids: BTreeSet, ) -> Result<(Chunk, PathBuf)> { self.file .into_inner() .map_err(|e| e.into_error()) .into_diagnostic()? .sync_all() .into_diagnostic()?; let chunk = Chunk { first, last, level, events, bytes: self.bytes, sha256: crate::util::hex::encode(self.hash.finalize()), dids: dids.into_iter().collect(), dirty: false, }; Ok((chunk, self.tmp)) } } impl Writer { fn new(state: &AppState, first: u64, last: u64, level: u32) -> Result { let out = ChunkFile::create(state, first, last)?; Ok(Self { out: zstd::Encoder::new(out, ZSTD_LEVEL).into_diagnostic()?, first, last, level, events: 0, dids: BTreeSet::new(), }) } fn push(&mut self, did: u64, id: u64, value: &[u8]) -> Result<()> { self.dids.insert(did); let mut lens = Vec::with_capacity(20); varint(&mut lens, 8); lens.extend_from_slice(&id.to_be_bytes()); varint(&mut lens, value.len() as u64); self.out.write_all(&lens).into_diagnostic()?; self.out.write_all(value).into_diagnostic()?; self.events += 1; Ok(()) } fn finish(self) -> Result<(Chunk, PathBuf)> { let out = self.out.finish().into_diagnostic()?; out.finish(self.first, self.last, self.level, self.events, self.dids) } } #[derive(serde::Serialize)] struct Header<'a> { #[serde(rename = "$type")] kind: &'a str, stream: &'a str, first: u64, last: u64, } /// its own zstd frame because a merge cuts it off and writes a new one fn header_frame(state: &AppState, first: u64, last: u64) -> Result> { let meta = serde_ipld_dagcbor::to_vec(&Header { kind: STAR_TYPE, stream: state.db.stream.id(), first, last, }) .into_diagnostic()?; let mut header = STAR_MAGIC.to_vec(); varint(&mut header, meta.len() as u64); header.extend_from_slice(&meta); zstd::bulk::compress(&header, ZSTD_LEVEL).into_diagnostic() } fn varint(out: &mut Vec, mut n: u64) { while n >= 0x80 { out.push(n as u8 | 0x80); n >>= 7; } out.push(n as u8); } /// built from the event rows, not copied from any chunk it replaces fn write(state: &AppState, first: u64, last: u64, level: u32, replaces: &[Chunk]) -> Result<()> { let started = Instant::now(); let chunks = chunks(state); let mut writer = Writer::new(state, first, last, level)?; let rows = state .db .stream .event_range(keys::event_key(first)..=keys::event_key(last)); for guard in rows { let (key, value) = guard.into_inner().into_diagnostic()?; let id = u64::from_be_bytes(key.as_ref().try_into().into_diagnostic()?); // replay skips undecodable rows and chunks have to match it let Ok(stored) = rmp_serde::from_slice::(&value) else { continue; }; let did = did_hash(stored.did.to_did().as_str()); writer.push(did, id, &stored_to_chunk_value(state, id, stored))?; } let (chunk, tmp) = writer.finish()?; let (events, bytes) = (chunk.events, chunk.bytes); chunks.publish(replaces, chunk, &tmp)?; let rebuilt = !replaces.is_empty(); info!(first, last, events, bytes, rebuilt, elapsed = ?started.elapsed(), "wrote a stream chunk"); Ok(()) } fn concat(state: &AppState, run: &[Chunk]) -> Result<(Chunk, PathBuf)> { let chunks = chunks(state); let (first, last) = (run[0].first, run[run.len() - 1].last); let mut out = ChunkFile::create(state, first, last)?; for chunk in run { let bytes = std::fs::read(chunks.path(&chunk.sha256)).into_diagnostic()?; let header = zstd::zstd_safe::find_frame_compressed_size(&bytes).map_err(|code| { miette::miette!( "a chunk without a header frame: {}", zstd::zstd_safe::get_error_name(code) ) })?; out.write_all(&bytes[header..]).into_diagnostic()?; } let events = run.iter().map(|c| c.events).sum(); let dids = run.iter().flat_map(|c| c.dids.iter().copied()).collect(); out.finish(first, last, run[0].level + 1, events, dids) } fn merge(state: &AppState, run: Vec) -> Result<()> { let started = Instant::now(); let chunks = chunks(state); let (chunk, tmp) = concat(state, &run)?; let (first, last, bytes) = (chunk.first, chunk.last, chunk.bytes); chunks.publish(&run, chunk, &tmp)?; info!(first, last, bytes, elapsed = ?started.elapsed(), "merged stream chunks"); Ok(()) } #[cfg(test)] mod tests { use jacquard_common::types::did::Did; use jacquard_common::types::string::Tid; use tokio::sync::mpsc; use super::*; use crate::config::Config; use crate::db::Txn; use crate::db::types::{DbAction, DbRkey, DbTid}; const ALICE: &str = "did:plc:aaaaaaaaaaaaaaaaaaaaaaaa"; const BOB: &str = "did:plc:bbbbbbbbbbbbbbbbbbbbbbbb"; const COLLECTION: &str = "app.bsky.feed.post"; fn state(span: u64) -> (tempfile::TempDir, Arc) { let tmp = tempfile::tempdir().unwrap(); let state = open(&tmp, span); (tmp, state) } fn open(dir: &tempfile::TempDir, span: u64) -> Arc { let config = Config { database_path: dir.path().to_path_buf(), stream_chunk_span: span, ..Default::default() }; Arc::new(AppState::new(&config).unwrap()) } fn frame(line: &[u8]) -> serde_json::Value { serde_json::from_slice(line).unwrap() } #[derive(serde::Deserialize)] struct ChunkFrame { cid: Option, did: String, rev: String, live: bool, rkey: String, action: String, #[serde(default, with = "serde_bytes")] record: Option>, collection: String, } fn frame_json(id: u64, f: ChunkFrame) -> Vec { let json = |s: &str| serde_json::to_string(s).unwrap(); let mut out = format!( r#"{{"id":{},"type":"record","record":{{"live":{},"did":{},"rev":{},"collection":{},"rkey":{},"action":{}"#, id, f.live, json(&f.did), json(&f.rev), json(&f.collection), json(&f.rkey), json(&f.action), ) .into_bytes(); let cid = f .cid .or_else(|| f.record.as_deref().map(crate::car::compute_cid)); if let Some(body) = f.record { let mark = out.len(); out.extend_from_slice(br#","record":"#); if crate::types::write_record_json(&body, &mut out).is_err() { out.truncate(mark); } } if let Some(cid) = cid { out.extend_from_slice(format!(r#","cid":"{cid}""#).as_bytes()); } out.extend_from_slice(b"}}"); out } #[derive(serde::Deserialize)] struct Header { #[serde(rename = "$type")] kind: String, stream: String, first: u64, last: u64, } fn varint(bytes: &[u8], at: &mut usize) -> usize { let (mut n, mut shift) = (0, 0); loop { let b = bytes[*at]; *at += 1; n |= usize::from(b & 0x7f) << shift; if b & 0x80 == 0 { return n; } shift += 7; } } fn read_chunk(path: &std::path::Path) -> (Header, Vec>) { let bytes = zstd::decode_all(std::fs::File::open(path).unwrap()).unwrap(); assert_eq!(bytes[..4], STAR_MAGIC); let mut at = 4; let len = varint(&bytes, &mut at); let header: Header = serde_ipld_dagcbor::from_slice(&bytes[at..at + len]).unwrap(); at += len; let mut frames = Vec::new(); while at < bytes.len() { let len = varint(&bytes, &mut at); let id = u64::from_be_bytes(bytes[at..at + len].try_into().unwrap()); at += len; let len = varint(&bytes, &mut at); let value = serde_ipld_dagcbor::from_slice(&bytes[at..at + len]).unwrap(); at += len; frames.push(frame_json(id, value)); } (header, frames) } fn rkey(n: usize) -> DbRkey { DbRkey::Str(smol_str::format_smolstr!("r{n}")) } fn commit(state: &AppState, did: &'static str, ops: &[(usize, Option)]) { let did = Did::new_static(did).unwrap(); let mut txn = Txn::new(&state.db); let mut records = txn .records(state, &DbTid::from(&Tid::now_0()), &did) .unwrap(); for (n, version) in ops { match version { Some(v) => { let body = serde_json::json!({ "$type": COLLECTION, "n": n, "v": v }); let block = crate::car::CarBlock::from_body(serde_ipld_dagcbor::to_vec(&body).unwrap()); let action = if *v == 0 { DbAction::Create } else { DbAction::Update }; records .put_record(COLLECTION, &rkey(*n), &block, action) .unwrap(); } None => records.delete_record(COLLECTION, &rkey(*n)).unwrap(), } } txn.commit().unwrap(); } fn rounds(state: &AppState, rounds: std::ops::Range) { for round in rounds { for n in 0..5 { let op = match (round + n) % 4 { 3 => None, v => Some(round * 4 + v), }; commit(state, ALICE, &[(n, op)]); commit(state, BOB, &[(n, Some(round)), (n + 10, op)]); } } } /// what the real stream thread sends, which chunks have to match fn replay(state: &Arc, cursor: u64) -> Vec> { let head = state.db.stream.event_head().unwrap(); if head.is_none_or(|head| head < cursor) { return Vec::new(); } let (tx, mut rx) = mpsc::channel(1 << 16); let thread_state = state.clone(); let opts = crate::control::stream::StreamOptions::from_config(&Config::default()); std::thread::spawn(move || { super::super::event_stream_thread::( thread_state, tx, Some(cursor), opts, ) }); let mut frames = Vec::new(); loop { let line = rx.blocking_recv().unwrap().unwrap().to_vec(); let id = frame(&line)["id"].as_u64(); frames.push(line); if id == head { return frames; } } } fn chunk_frames(state: &AppState) -> (Vec>, u64) { let chunks = chunks(state); let manifest = chunks.manifest(); let mut frames = Vec::new(); for c in &manifest.chunks { let path = chunks .servable(&c.sha256) .expect("a listed chunk is served"); let file = std::fs::read(&path).unwrap(); assert_eq!(crate::util::hex::encode(Sha256::digest(&file)), c.sha256); let (header, chunk) = read_chunk(&path); assert_eq!( ( header.kind.as_str(), header.stream.as_str(), header.first, header.last ), (STAR_TYPE, state.db.stream.id(), c.first, c.last) ); frames.extend(chunk); } (frames, manifest.through.map_or(0, |through| through + 1)) } fn with_tail(state: &Arc) -> Vec> { let (mut frames, next) = chunk_frames(state); frames.extend(replay(state, next)); frames } #[test] fn chunks_and_the_tail_are_the_replay() { let (_tmp, state) = state(8); rounds(&state, 0..3); build(&state).unwrap(); // later writes hit records already in sealed chunks rounds(&state, 3..6); build(&state).unwrap(); let (frames, next) = chunk_frames(&state); assert!(next > 64 && !frames.is_empty(), "nothing sealed"); assert_eq!(with_tail(&state), replay(&state, 0)); assert!(chunks(&state).manifest().chunks.len() < 8); } #[test] fn the_worker_builds_until_hydrant_stops() { let (_tmp, state) = state(4); rounds(&state, 0..6); let worker = { let state = state.clone(); std::thread::spawn(move || chunk_worker(state)) }; while chunks(&state).manifest().chunks.is_empty() { std::thread::sleep(Duration::from_millis(10)); } state.stop.trigger(); worker.join().unwrap(); assert_eq!(with_tail(&state), replay(&state, 0)); } #[test] fn merges_stop_at_sixteen_spans() { let (_tmp, state) = state(1); rounds(&state, 0..6); build(&state).unwrap(); let manifest = chunks(&state).manifest(); let widest = manifest.chunks.iter().map(|c| c.last + 1 - c.first).max(); assert_eq!(widest, Some(16)); assert_eq!(with_tail(&state), replay(&state, 0)); } #[test] fn a_restart_keeps_the_chunks_it_built() { let tmp = tempfile::tempdir().unwrap(); let state = open(&tmp, 4); rounds(&state, 0..6); build(&state).unwrap(); let before = serde_json::to_value(chunks(&state).manifest()).unwrap(); assert!(!before["chunks"].as_array().unwrap().is_empty()); drop(state); let state = open(&tmp, 4); let chunks = chunks(&state); assert_eq!(serde_json::to_value(chunks.manifest()).unwrap(), before); let through = before["through"].as_u64().unwrap(); assert_eq!(chunks.sealed_next(), through + 1); rounds(&state, 6..8); build(&state).unwrap(); assert_eq!(with_tail(&state), replay(&state, 0)); } #[test] fn a_busy_ingest_gets_one_span_a_round() { let (_tmp, state) = state(4); rounds(&state, 0..6); let firehose = url::Url::parse("wss://relay.invalid").unwrap(); let tracker = futures::executor::block_on(state.firehose_tracker(&firehose)); let tickets: Vec<_> = (0..BACKLOG as i64) .map(|seq| tracker.ticket(seq, 0)) .collect(); let chunks = chunks(&state); build(&state).unwrap(); assert_eq!(chunks.sealed_next(), 4); build(&state).unwrap(); assert_eq!(chunks.sealed_next(), 8); for ticket in tickets { ticket.done(|_| {}); } build(&state).unwrap(); let head = state.db.stream.ids.head().unwrap(); assert!( chunks.sealed_next() + 4 > head, "an idle ingest left spans unsealed" ); } fn bodies(frames: &[Vec]) -> usize { frames .iter() .filter(|line| frame(line)["record"].get("record").is_some()) .count() } fn holds(state: &AppState, sha256: &str, did: &str) -> bool { read_chunk(&chunks(state).path(sha256)) .1 .iter() .any(|line| frame(line)["record"]["did"] == did) } #[test] fn a_redaction_reaches_every_chunk_holding_the_record() { let (_tmp, state) = state(4); rounds(&state, 0..6); build(&state).unwrap(); let before = chunks(&state).manifest(); crate::db::redact_record_bodies( &state.db, &Did::new_static(BOB).unwrap(), COLLECTION, &rkey(1), crate::types::DeleteBodyTarget::All, ) .unwrap(); let chunks = chunks(&state); for chunk in &before.chunks { let served = chunks.servable(&chunk.sha256).is_some(); assert_eq!(served, !holds(&state, &chunk.sha256, BOB)); } let dirty = chunks.manifest(); assert!(dirty.through < before.through, "no chunk went dirty"); assert!(dirty.generation > before.generation); build(&state).unwrap(); assert_eq!(with_tail(&state), replay(&state, 0)); assert_eq!(chunks.manifest().through, before.through); } #[test] fn an_erase_reaches_every_chunk_holding_the_repo() { let (_tmp, state) = state(4); rounds(&state, 0..6); build(&state).unwrap(); let with_bodies = bodies(&chunk_frames(&state).0); crate::db::erase_repo(&state.db, &Did::new_static(BOB).unwrap()).unwrap(); build(&state).unwrap(); let replayed = replay(&state, 0); assert!(bodies(&replayed) < with_bodies, "the erase took no body"); assert_eq!(with_tail(&state), replayed); } fn change(state: &AppState, did: &str) { let mut batch = state.db.inner.batch(); chunks(state).mark(&mut batch, did); batch.commit().unwrap(); chunks(state).touch(did); } /// crashes right after change, before the chunk list gets written fn crash_after(change: impl FnOnce(&AppState)) { let tmp = tempfile::tempdir().unwrap(); let state = open(&tmp, 4); rounds(&state, 0..6); build(&state).unwrap(); let list = tmp.path().join("stream-chunks/state.json"); let unaware = std::fs::read(&list).unwrap(); change(&state); state.db.persist().unwrap(); drop(state); std::fs::write(&list, unaware).unwrap(); let state = open(&tmp, 4); let (served, _) = chunk_frames(&state); assert_eq!(served[..], replay(&state, 0)[..served.len()]); build(&state).unwrap(); assert_eq!(with_tail(&state), replay(&state, 0)); } #[test] fn turning_chunks_off_forgets_them() { let tmp = tempfile::tempdir().unwrap(); let state = open(&tmp, 4); rounds(&state, 0..2); build(&state).unwrap(); drop(state); drop(open(&tmp, 0)); assert!(!tmp.path().join("stream-chunks").exists()); } #[test] fn a_crash_before_a_redaction_reaches_the_chunks_leaves_none_holding_the_body() { use crate::types::DeleteBodyTarget::{Head, History}; for target in [Head, History] { crash_after(|state| { let bob = Did::new_static(BOB).unwrap(); crate::db::redact_record_bodies(&state.db, &bob, COLLECTION, &rkey(1), target) .unwrap(); }); } } #[test] fn a_crash_before_an_erase_reaches_the_chunks_leaves_none_holding_the_bodies() { crash_after(|state| { crate::db::erase_repo(&state.db, &Did::new_static(BOB).unwrap()).unwrap(); }); } #[test] fn a_touch_while_chunks_merge_serves_none_of_them() { let (_tmp, state) = state(4); rounds(&state, 0..6); build(&state).unwrap(); rounds(&state, 6..12); while chunks(&state).sealed_next() + 4 <= state.db.stream.ids.head().unwrap() { let first = chunks(&state).sealed_next(); write(&state, first, first + 3, 0, &[]).unwrap(); } let chunks = chunks(&state); let run = chunks.mergeable().expect("a run to merge"); change(&state, BOB); let (chunk, tmp) = concat(&state, &run).unwrap(); chunks.publish(&run, chunk, &tmp).unwrap(); for replaced in &run { assert!(chunks.servable(&replaced.sha256).is_none()); } } #[test] fn a_touch_while_a_chunk_builds_leaves_it_dirty() { let (_tmp, state) = state(4); rounds(&state, 0..6); let chunks = chunks(&state); let mut writer = Writer::new(&state, 0, 3, 0).unwrap(); writer.push(did_hash(BOB), 0, &[0xa0]).unwrap(); change(&state, BOB); let (chunk, tmp) = writer.finish().unwrap(); chunks.publish(&[], chunk, &tmp).unwrap(); assert!(chunks.manifest().chunks.is_empty()); } }