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";
/// 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<AppState>) {
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<u64>,
}
struct ChunkFile {
tmp: PathBuf,
file: BufWriter<std::fs::File>,
hash: Sha256,
bytes: u64,
}
impl Write for ChunkFile {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
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<Self> {
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<u64>,
) -> 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<Self> {
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);
dasl_star::write_pair(&mut self.out, &id.to_be_bytes(), value).into_diagnostic()?;
self.events += 1;
Ok(())
}
fn finish(self) -> Result<(Chunk, PathBuf)> {
let mut out = self.out.finish().into_diagnostic()?;
out.write_all(&end_frame()?).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<Vec<u8>> {
let mut header = Vec::new();
dasl_star::write_header(
&mut header,
&Header {
kind: STAR_TYPE,
stream: state.db.stream.id(),
first,
last,
},
)
.into_diagnostic()?;
zstd::bulk::compress(&header, ZSTD_LEVEL).into_diagnostic()
}
/// a frame of its own, because a merge drops each chunk's end and writes one after the last
fn end_frame() -> Result<Vec<u8>> {
let mut end = Vec::new();
dasl_star::write_end(&mut end).into_diagnostic()?;
zstd::bulk::compress(&end, ZSTD_LEVEL).into_diagnostic()
}
fn pair_frames(chunk: &[u8]) -> Result<&[u8]> {
let mut starts = Vec::new();
let mut at = 0;
while at < chunk.len() {
starts.push(at);
at += zstd::zstd_safe::find_frame_compressed_size(&chunk[at..]).map_err(|code| {
miette::miette!(
"a chunk with a broken zstd frame: {}",
zstd::zstd_safe::get_error_name(code)
)
})?;
}
let (Some(&pairs), Some(&end)) = (starts.get(1), starts.last()) else {
miette::bail!("a chunk without a header frame and an end frame");
};
let ends = zstd::bulk::decompress(&chunk[end..], 1).is_ok_and(|end| end == [dasl_star::END]);
if !ends {
miette::bail!("a chunk whose last frame isn't the STAR end");
}
Ok(&chunk[pairs..end])
}
/// 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::<StoredEvent>(&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()?;
out.write_all(pair_frames(&bytes)?).into_diagnostic()?;
}
out.write_all(&end_frame()?).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<Chunk>) -> 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<AppState>) {
let tmp = tempfile::tempdir().unwrap();
let state = open(&tmp, span);
(tmp, state)
}
fn open(dir: &tempfile::TempDir, span: u64) -> Arc<AppState> {
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<jacquard_common::types::cid::IpldCid>,
did: String,
rev: String,
live: bool,
rkey: String,
action: String,
#[serde(default, with = "serde_bytes")]
record: Option<Vec<u8>>,
collection: String,
}
fn frame_json(id: u64, f: ChunkFrame) -> Vec<u8> {
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 read_chunk(path: &std::path::Path) -> (Header, Vec<Vec<u8>>) {
let bytes = zstd::decode_all(std::fs::File::open(path).unwrap()).unwrap();
let reader = dasl_star::Reader::new(&bytes[..]).unwrap();
let header: Header = serde_ipld_dagcbor::from_slice(reader.meta()).unwrap();
let frames = reader
.map(|pair| {
let (key, value) = pair.unwrap();
let id = u64::from_be_bytes(key.try_into().unwrap());
frame_json(id, serde_ipld_dagcbor::from_slice(&value).unwrap())
})
.collect();
(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<usize>)]) {
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<usize>) {
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<AppState>, cursor: u64) -> Vec<Vec<u8>> {
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::<super::super::Frame>(
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<Vec<u8>>, 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<AppState>) -> Vec<Vec<u8>> {
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 chunks_an_older_format_wrote_are_rebuilt() {
let tmp = tempfile::tempdir().unwrap();
let state = open(&tmp, 4);
rounds(&state, 0..6);
build(&state).unwrap();
let built = chunks(&state).manifest().chunks;
assert!(!built.is_empty());
let dir = tmp.path().join("stream-chunks");
drop(state);
for c in &built {
let path = dir.join(format!("{}.zst", c.sha256));
let bytes = std::fs::read(&path).unwrap();
let mut at = 0;
let mut last = 0;
while at < bytes.len() {
last = at;
at += zstd::zstd_safe::find_frame_compressed_size(&bytes[at..]).unwrap();
}
std::fs::write(&path, &bytes[..last]).unwrap();
}
let list = dir.join("state.json");
let mut saved: serde_json::Value =
serde_json::from_slice(&std::fs::read(&list).unwrap()).unwrap();
saved.as_object_mut().unwrap().remove("format").unwrap();
std::fs::write(&list, serde_json::to_vec(&saved).unwrap()).unwrap();
let state = open(&tmp, 4);
let chunks = chunks(&state);
assert!(chunks.manifest().chunks.is_empty());
assert!(built.iter().all(|c| chunks.servable(&c.sha256).is_none()));
build(&state).unwrap();
// >= because a restart moves the head on
let through = built.last().unwrap().last;
assert!(chunks.manifest().through.is_some_and(|t| t >= through));
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<u8>]) -> 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());
}
}