diff --git a/Cargo.lock b/Cargo.lock index 08fbc863d..7df9b4fac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -877,6 +877,7 @@ dependencies = [ "serde", "serde_ipld_dagcbor", "serde_json", + "sha2 0.11.0", "thiserror 2.0.18", "tokio", "tokio-stream", @@ -886,6 +887,7 @@ dependencies = [ "url", "uuid", "wiremock", + "zstd", ] [[package]] @@ -960,6 +962,7 @@ dependencies = [ "scc", "serde", "serde_json", + "sha2 0.11.0", "thiserror 2.0.18", "tokio", "tokio-util", diff --git a/Cargo.toml b/Cargo.toml index 49b7842a4..85d245cad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -190,6 +190,7 @@ k256 = { version = "0.13", features = ["ecdsa"] } aes-gcm = "0.11.0-rc.4" hkdf = "0.13" sha2 = "0.11" +zstd = "0.13" base32 = "0.5" base64 = "0.22" bs58 = "0.5" diff --git a/bobbin/crates/bobbin-sim/src/runtime.rs b/bobbin/crates/bobbin-sim/src/runtime.rs index 0629fded4..60f2ccba5 100644 --- a/bobbin/crates/bobbin-sim/src/runtime.rs +++ b/bobbin/crates/bobbin-sim/src/runtime.rs @@ -136,6 +136,7 @@ impl Sim { start_cursor: bobbin_edge_index::HydrantCursor::new(0), parallelism, idle_promote_min_events: DEFAULT_IDLE_PROMOTE_MIN_EVENTS, + chunks: false, }; let mut ingest_handle = tokio::spawn(async move { let _ = run_ingest(ingest_config, ingest_runtime).await; diff --git a/bobbin/crates/bobbin/src/config.rs b/bobbin/crates/bobbin/src/config.rs index 114e4ccfb..648005b70 100644 --- a/bobbin/crates/bobbin/src/config.rs +++ b/bobbin/crates/bobbin/src/config.rs @@ -25,6 +25,7 @@ const KNOWN_KEYS: &[&str] = &[ "server.cors_origins", "hydrant.url", "hydrant.start_cursor", + "hydrant.chunks", "ingest.parallelism", "ingest.idle_promote_min_events", "ingest.author_fill_concurrency", @@ -66,6 +67,7 @@ const KNOWN_ENVS: &[&str] = &[ "BOBBIN_CORS_ORIGINS", "BOBBIN_HYDRANT_URL", "BOBBIN_START_CURSOR", + "BOBBIN_HYDRANT_CHUNKS", "BOBBIN_INGEST_PARALLELISM", "BOBBIN_INGEST_IDLE_PROMOTE_MIN_EVENTS", "BOBBIN_INGEST_AUTHOR_FILL_CONCURRENCY", @@ -231,6 +233,10 @@ pub struct HydrantConfig { /// strictly after the last cursor seen internally. #[config(env = "BOBBIN_START_CURSOR", default = 0)] pub start_cursor: u64, + + /// cold start loads hydrant's sealed /stream chunks, or replays /stream from 0 when there are none + #[config(env = "BOBBIN_HYDRANT_CHUNKS", default = true)] + pub chunks: bool, } #[derive(Debug, Config)] diff --git a/bobbin/crates/bobbin/src/main.rs b/bobbin/crates/bobbin/src/main.rs index f33aea3bc..7ee8fc799 100644 --- a/bobbin/crates/bobbin/src/main.rs +++ b/bobbin/crates/bobbin/src/main.rs @@ -336,6 +336,7 @@ async fn run(cfg: BobbinConfig) -> anyhow::Result<()> { start_cursor: HydrantCursor::new(cfg.hydrant.start_cursor), parallelism, idle_promote_min_events, + chunks: cfg.hydrant.chunks, }; let cancel = CancellationToken::new(); let analytics = if cfg.posthog.api_key.is_empty() { diff --git a/bobbin/crates/ingest/Cargo.toml b/bobbin/crates/ingest/Cargo.toml index c3aba69cf..e7fa7202a 100644 --- a/bobbin/crates/ingest/Cargo.toml +++ b/bobbin/crates/ingest/Cargo.toml @@ -25,6 +25,7 @@ scc = { workspace = true } serde = { workspace = true } serde_ipld_dagcbor = { workspace = true } serde_json = { workspace = true } +sha2 = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true } tokio-stream = { workspace = true } @@ -32,6 +33,7 @@ tokio-util = { workspace = true } tracing = { workspace = true } url = { workspace = true } uuid = { workspace = true } +zstd = { workspace = true } [dev-dependencies] http = { workspace = true } diff --git a/bobbin/crates/ingest/src/chunks.rs b/bobbin/crates/ingest/src/chunks.rs new file mode 100644 index 000000000..bd483cced --- /dev/null +++ b/bobbin/crates/ingest/src/chunks.rs @@ -0,0 +1,422 @@ +use std::collections::VecDeque; +use std::io::{BufRead, Read as _}; +use std::sync::Arc; +use std::time::Duration; + +use bobbin_edge_index::HydrantCursor; +use bobbin_resolver::{HydrantClient, StreamChunk}; +use bobbin_runtime::Clock; +use bobbin_types::search::SearchSink; +use bytes::Bytes; +use jacquard_common::types::cid::{Cid, IpldCid}; +use jacquard_common::types::did::Did; +use jacquard_common::types::nsid::Nsid; +use jacquard_common::types::recordkey::Rkey; +use jacquard_common::types::tid::Tid; +use jacquard_common::{DefaultStr, RawData}; +use serde::Deserialize; +use sha2::{Digest as _, Sha256}; +use tracing::{info, warn}; + +use crate::{ + FrameKind, HydrantFrame, IngestConfig, IngestRuntime, RecordAction, RecordFrame, + promote_at_head, spawn_pipeline, +}; + +const IN_FLIGHT: usize = 6; +/// the first chunk goes in even past the window +const WINDOW_BYTES: u64 = 64 << 20; +const ATTEMPTS: u32 = 3; +const RETRY_DELAY: Duration = Duration::from_millis(500); + +/// stopping partway is fine because /stream resumes from what got applied +pub(crate) async fn load( + config: &IngestConfig, + runtime: &IngestRuntime, + replay_head: Option, +) -> Option { + let client = runtime.hydrant.as_ref()?; + let manifest = match client.stream_manifest().await { + Ok(Some(manifest)) if !manifest.chunks.is_empty() => manifest, + Ok(_) => { + info!("hydrant has no stream chunks, replaying /stream"); + return None; + } + Err(err) => { + warn!( + ?err, + "couldn't read hydrant's stream chunks, replaying /stream" + ); + return None; + } + }; + let through = manifest.through?; + info!( + chunks = manifest.chunks.len(), + bytes = manifest.chunks.iter().map(|c| c.bytes).sum::(), + through, + "loading hydrant's stream chunks" + ); + let (frame_tx, processor) = spawn_pipeline(config, runtime, replay_head); + let mut queue = manifest.chunks.iter().cloned().peekable(); + let mut downloads = VecDeque::new(); + let mut window = 0; + let mut complete = true; + loop { + while downloads.len() < IN_FLIGHT + && let Some(chunk) = + queue.next_if(|chunk| downloads.is_empty() || window + chunk.bytes <= WINDOW_BYTES) + { + window += chunk.bytes; + let download = + tokio::spawn(fetch(client.clone(), runtime.clock.clone(), chunk.clone())); + downloads.push_back((chunk, download)); + } + let Some((chunk, download)) = downloads.front_mut() else { + break; + }; + let body = tokio::select! { + biased; + _ = runtime.cancel.cancelled() => break, + body = download => body.map_err(|e| e.to_string()).and_then(|body| body), + }; + let applied = match body { + Ok(body) => apply(body, &manifest.stream, chunk, frame_tx.clone()).await, + Err(err) => Err(err), + }; + if let Err(err) = applied { + warn!(first = chunk.first, last = chunk.last, %err, "stopped loading stream chunks, replaying the rest from /stream"); + complete = false; + break; + } + window -= chunk.bytes; + downloads.pop_front(); + } + for (_, download) in &downloads { + download.abort(); + } + drop(frame_tx); + if let Err(join) = processor.await { + warn!(?join, "frame processor task panicked"); + } + if complete && !runtime.cancel.is_cancelled() { + // through, not the last event id, because the range can end on ids with no event + runtime + .coverage + .update(|c| c.advance(HydrantCursor::new(through))); + if let Some(head) = replay_head { + promote_at_head(runtime, head); + } + } + let loaded = runtime.coverage.snapshot(); + info!( + events = loaded.events_processed(), + last_cursor = loaded.last_cursor().raw(), + complete, + "loaded hydrant's stream chunks" + ); + (loaded.events_processed() > 0).then_some(manifest.stream) +} + +async fn fetch( + client: HydrantClient, + clock: Arc, + chunk: StreamChunk, +) -> Result { + let mut attempt = 1; + loop { + match client.stream_chunk(&chunk).await { + Ok(body) => return Ok(body), + Err(err) if attempt < ATTEMPTS => { + warn!(first = chunk.first, attempt, %err, "retrying a stream chunk"); + clock.sleep(RETRY_DELAY * attempt).await; + attempt += 1; + } + Err(err) => return Err(err.to_string()), + } + } +} + +async fn apply( + body: Bytes, + stream: &str, + chunk: &StreamChunk, + frame_tx: tokio::sync::mpsc::Sender, +) -> Result<(), String> { + let (stream, first, last) = (stream.to_owned(), chunk.first, chunk.last); + tokio::task::spawn_blocking(move || { + for frame in chunk_frames(body, &stream, first, last)? { + if frame_tx.blocking_send(frame?).is_err() { + return Err("the frame processor stopped".to_owned()); + } + } + Ok(()) + }) + .await + .map_err(|e| e.to_string())? +} + +/// the $type of the dag-cbor header after the STAR magic +const STAR_TYPE: &str = "systems.gaze.hydrant.streamChunk"; +const STAR_MAGIC: [u8; 4] = [0x2a, 0x6c, 0x01, 0x00]; + +#[derive(Deserialize)] +struct Header<'a> { + #[serde(rename = "$type")] + kind: &'a str, + stream: &'a str, + first: u64, + last: u64, +} + +#[derive(Deserialize)] +struct ChunkEvent<'a> { + #[serde(default)] + cid: Option, + did: Did, + rev: Tid, + live: bool, + rkey: Rkey, + action: RecordAction, + #[serde(default, borrow)] + record: Option<&'a [u8]>, + collection: Nsid, +} + +pub fn chunk_frames( + body: Bytes, + stream: &str, + first: u64, + last: u64, +) -> Result>, String> { + let decoder = + zstd::stream::read::Decoder::new(std::io::Cursor::new(body)).map_err(|e| e.to_string())?; + let mut star = std::io::BufReader::new(decoder); + let mut magic = [0; 4]; + star.read_exact(&mut magic).map_err(|e| e.to_string())?; + if magic != STAR_MAGIC { + return Err("not a STAR".to_owned()); + } + let len = varint(&mut star)?.ok_or("a STAR without a header")?; + let meta = read_len(&mut star, len)?; + let header: Header = serde_ipld_dagcbor::from_slice(&meta).map_err(|e| e.to_string())?; + if (header.kind, header.stream, header.first, header.last) != (STAR_TYPE, stream, first, last) { + return Err(format!( + "a chunk of {} {} {}..={}, not the one listed", + header.kind, header.stream, header.first, header.last + )); + } + let mut next = first; + Ok(std::iter::from_fn(move || { + let pair = (|| { + let Some(len) = varint(&mut star)? else { + return Ok(None); + }; + let key = read_len(&mut star, len)?; + let id = <[u8; 8]>::try_from(key.as_slice()) + .map(u64::from_be_bytes) + .map_err(|_| "a key that isn't an id")?; + if !(next..=last).contains(&id) { + return Err(format!("id {id} out of order")); + } + next = id + 1; + let len = varint(&mut star)?.ok_or("a key without a value")?; + let value = read_len(&mut star, len)?; + let event = serde_ipld_dagcbor::from_slice(&value).map_err(|e| e.to_string())?; + Ok(Some(stream_frame(id, event))) + })(); + pair.transpose() + })) +} + +fn varint(r: &mut impl BufRead) -> Result, String> { + if r.fill_buf().map_err(|e| e.to_string())?.is_empty() { + return Ok(None); + } + let mut n = 0; + for shift in (0..63).step_by(7) { + let mut byte = [0]; + r.read_exact(&mut byte).map_err(|e| e.to_string())?; + n |= u64::from(byte[0] & 0x7f) << shift; + if byte[0] & 0x80 == 0 { + if byte[0] == 0 && shift > 0 { + return Err("a varint with a trailing zero byte".to_owned()); + } + return Ok(Some(n)); + } + } + Err("a varint past 9 bytes".to_owned()) +} + +fn read_len(r: &mut impl BufRead, len: u64) -> Result, String> { + let mut out = Vec::new(); + r.take(len) + .read_to_end(&mut out) + .map_err(|e| e.to_string())?; + if out.len() as u64 != len { + return Err("a STAR cut short".to_owned()); + } + Ok(out) +} + +fn stream_frame(id: u64, frame: ChunkEvent<'_>) -> HydrantFrame { + let cid = frame.cid.or_else(|| frame.record.map(body_cid)); + let record = frame.record.and_then(|body| { + let value = serde_ipld_dagcbor::from_slice::(body).ok()?; + serde_json::value::to_raw_value(&value).ok() + }); + HydrantFrame { + id, + kind: FrameKind::Record, + record: Some(RecordFrame { + live: frame.live, + did: frame.did, + rev: frame.rev, + collection: frame.collection, + rkey: frame.rkey, + action: frame.action, + record, + cid: cid.map(|cid| Cid::from(cid.to_string())), + }), + identity: None, + account: None, + } +} + +/// cidv1 of the record body, dag-cbor (0x71) and sha2-256 (0x12) +fn body_cid(body: &[u8]) -> IpldCid { + let mut bytes = vec![0x01, 0x71, 0x12, 0x20]; + bytes.extend_from_slice(&Sha256::digest(body)); + IpldCid::try_from(bytes.as_slice()).expect("a sha2-256 cidv1 is a cid") +} + +#[cfg(test)] +pub(crate) fn star_chunk(stream: &str, first: u64, last: u64, events: &[(u64, Vec)]) -> Bytes { + let varint = |out: &mut Vec, mut n: u64| { + while n >= 0x80 { + out.push(n as u8 | 0x80); + n >>= 7; + } + out.push(n as u8); + }; + let meta = serde_ipld_dagcbor::to_vec(&ipld_core::ipld::Ipld::Map( + [ + ("$type", ipld_core::ipld::Ipld::String(STAR_TYPE.to_owned())), + ("stream", ipld_core::ipld::Ipld::String(stream.to_owned())), + ("first", ipld_core::ipld::Ipld::Integer(first.into())), + ("last", ipld_core::ipld::Ipld::Integer(last.into())), + ] + .into_iter() + .map(|(k, v)| (k.to_owned(), v)) + .collect(), + )) + .unwrap(); + let mut star = STAR_MAGIC.to_vec(); + varint(&mut star, meta.len() as u64); + star.extend(meta); + for (id, value) in events { + varint(&mut star, 8); + star.extend(id.to_be_bytes()); + varint(&mut star, value.len() as u64); + star.extend(value); + } + Bytes::from(zstd::encode_all(star.as_slice(), 3).unwrap()) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use ipld_core::ipld::Ipld; + + use super::*; + + fn event(body: &[u8], cid: Option) -> Vec { + let text = |s: &str| Ipld::String(s.to_owned()); + let mut map = BTreeMap::from([ + ("did".to_owned(), text("did:plc:abalone")), + ("rev".to_owned(), text("3mabcdefghijk")), + ("live".to_owned(), Ipld::Bool(true)), + ("rkey".to_owned(), text("3mabcdefghijk")), + ("action".to_owned(), text("create")), + ("record".to_owned(), Ipld::Bytes(body.to_vec())), + ("collection".to_owned(), text("org.tangled.graph.follow")), + ]); + if let Some(cid) = cid { + map.insert("cid".to_owned(), Ipld::Link(cid)); + } + serde_ipld_dagcbor::to_vec(&Ipld::Map(map)).unwrap() + } + + fn frames(body: Bytes, first: u64, last: u64) -> Result, String> { + chunk_frames(body, "s1", first, last)?.collect() + } + + #[test] + fn a_chunk_event_is_the_frame_stream_sends() { + let body = serde_ipld_dagcbor::to_vec(&Ipld::Map(BTreeMap::from([ + ( + "$type".to_owned(), + Ipld::String("org.tangled.graph.follow".to_owned()), + ), + ("subject".to_owned(), Ipld::String("did:plc:uni".to_owned())), + ( + "createdAt".to_owned(), + Ipld::String("2026-05-01T00:00:00Z".to_owned()), + ), + ]))) + .unwrap(); + let mut other = vec![0x01, 0x71, 0x12, 0x20]; + other.extend([7; 32]); + let other = IpldCid::try_from(other.as_slice()).unwrap(); + let chunk = star_chunk( + "s1", + 0, + 3, + &[(1, event(&body, None)), (2, event(&body, Some(other)))], + ); + let [derived, kept] = <[_; 2]>::try_from(frames(chunk, 0, 3).unwrap()).unwrap(); + assert_eq!((derived.id, kept.id), (1, 2)); + let record = derived.record.unwrap(); + let json: serde_json::Value = serde_json::from_str(record.record.unwrap().get()).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "$type": "org.tangled.graph.follow", + "subject": "did:plc:uni", + "createdAt": "2026-05-01T00:00:00Z", + }) + ); + assert_eq!( + record.cid.unwrap().as_str(), + Cid::from(body_cid(&body).to_string()).as_str() + ); + assert_eq!( + kept.record.unwrap().cid.unwrap().as_str(), + other.to_string() + ); + } + + #[test] + fn a_chunk_that_isnt_the_one_listed_is_refused() { + let body = serde_ipld_dagcbor::to_vec(&Ipld::Map(BTreeMap::new())).unwrap(); + let chunk = |events: &[(u64, Vec)]| star_chunk("s1", 0, 3, events); + assert!(frames(chunk(&[]), 0, 3).unwrap().is_empty()); + assert!(chunk_frames(chunk(&[]), "s2", 0, 3).is_err()); + assert!(frames(chunk(&[]), 0, 7).is_err()); + let two = [(2, event(&body, None)), (1, event(&body, None))]; + assert!(frames(chunk(&two), 0, 3).is_err()); + assert!(frames(chunk(&[(4, event(&body, None))]), 0, 3).is_err()); + } + + #[test] + fn varints_are_read_as_star_says() { + let read = |bytes: &[u8]| varint(&mut std::io::Cursor::new(bytes)); + assert_eq!(read(&[]), Ok(None)); + assert_eq!(read(&[0x00]), Ok(Some(0))); + assert_eq!(read(&[0xac, 0x02]), Ok(Some(300))); + assert!(read(&[0x80, 0x00]).is_err()); + assert!(read(&[0x80]).is_err()); + assert!(read(&[0xff; 10]).is_err()); + } +} diff --git a/bobbin/crates/ingest/src/lib.rs b/bobbin/crates/ingest/src/lib.rs index 05d59f981..7ff655bf9 100644 --- a/bobbin/crates/ingest/src/lib.rs +++ b/bobbin/crates/ingest/src/lib.rs @@ -45,6 +45,7 @@ use url::Url; mod analytics; mod authorship; +mod chunks; mod frame; mod metrics; pub mod moderation; @@ -87,6 +88,7 @@ pub const DEFAULT_IDLE_PROMOTE_MIN_EVENTS: NonZeroU64 = match NonZeroU64::new(25 pub use analytics::{Analytics, PosthogConfig}; pub use authorship::TicketAuthorship; +pub use chunks::chunk_frames; #[derive(Clone, Debug)] pub struct IngestConfig { @@ -94,6 +96,7 @@ pub struct IngestConfig { pub start_cursor: HydrantCursor, pub parallelism: NonZeroUsize, pub idle_promote_min_events: NonZeroU64, + pub chunks: bool, } impl IngestConfig { @@ -103,15 +106,20 @@ impl IngestConfig { start_cursor: HydrantCursor::new(0), parallelism: DEFAULT_INGEST_PARALLELISM, idle_promote_min_events: DEFAULT_IDLE_PROMOTE_MIN_EVENTS, + chunks: false, } } - fn stream_url(&self, cursor: HydrantCursor) -> Result { + /// hydrant refuses a cursor from another stream + fn stream_url(&self, cursor: HydrantCursor, stream: Option<&str>) -> Result { let mut url = websocket_url(&self.hydrant_base)?; url.set_path("/stream"); url.query_pairs_mut() .clear() .append_pair("cursor", &cursor.raw().to_string()); + if let Some(stream) = stream { + url.query_pairs_mut().append_pair("stream", stream); + } Ok(url) } } @@ -373,6 +381,14 @@ async fn run_inner( let mut backoff = RECONNECT_INITIAL_DELAY; let mut connect_timeout = CONNECT_TIMEOUT; let mut replay_head = None; + let mut stream = None; + if config.chunks + && runtime.coverage.snapshot().events_processed() == 0 + && config.start_cursor.raw() == 0 + { + replay_head = fetch_replay_head(runtime).await; + stream = chunks::load(&config, runtime, replay_head).await; + } loop { let cursor = next_connect_cursor(runtime.coverage.snapshot(), config.start_cursor); let opened = runtime.clock.now_instant(); @@ -382,8 +398,16 @@ async fn run_inner( if replay_head.is_none() && !runtime.coverage.snapshot().is_ready() { replay_head = fetch_replay_head(runtime).await; } - let SessionEnd { outcome, error } = - run_session(&config, cursor, replay_head, connect_timeout, runtime, frames).await; + let SessionEnd { outcome, error } = run_session( + &config, + cursor, + stream.as_deref(), + replay_head, + connect_timeout, + runtime, + frames, + ) + .await; let open_for = runtime.clock.now_instant().duration_since(opened); if runtime.cancel.is_cancelled() { info!( @@ -427,6 +451,12 @@ async fn run_inner( runtime.clock.now_unix_micros(), error.as_ref().map(|err| err.to_string()).as_deref(), ); + // only a cold start gets past a stream mismatch + if let Some(err) = error.filter( + |err| matches!(err, IngestError::HydrantStream { code, .. } if code == "StreamMismatch"), + ) { + return Err(err); + } connect_timeout = CONNECT_TIMEOUT; match outcome { SessionOutcome::Progressed => backoff = RECONNECT_INITIAL_DELAY, @@ -688,12 +718,13 @@ fn jittered(base: Duration, entropy: &dyn Entropy) -> Duration { async fn run_session( config: &IngestConfig, cursor: HydrantCursor, + stream: Option<&str>, replay_head: Option, connect_timeout: Duration, runtime: &IngestRuntime, frames: &FrameCounter, ) -> SessionEnd { - let url = match config.stream_url(cursor) { + let url = match config.stream_url(cursor, stream) { Ok(u) => u, Err(e) => { return SessionEnd { @@ -729,33 +760,7 @@ async fn run_session( .stream_health .connected(runtime.clock.now_unix_micros()); - let (frame_tx, frame_rx) = tokio::sync::mpsc::channel::(FRAME_CHANNEL_DEPTH); - let parallelism = config.parallelism.get(); - let processor_runtime = runtime.clone(); - let processor = tokio::spawn(async move { - let cancel = processor_runtime.cancel.clone(); - let prep_rt = processor_runtime.clone(); - let claim_rt = processor_runtime.clone(); - let resolve_rt = processor_runtime.clone(); - let commit_rt = processor_runtime; - - let pipeline = ReceiverStream::new(frame_rx) - .map(move |frame| prep_stage(frame, prep_rt.clone())) - .buffered(parallelism) - // then not buffered, resolve reads what this stage writes - .then(move |staged| claim_stage(staged, claim_rt.clone())) - .map(move |staged| resolve_stage(staged, resolve_rt.clone(), replay_head)) - .buffered(parallelism) - .for_each(move |staged| { - commit_stage(staged, commit_rt.clone(), parallelism, replay_head) - }); - - tokio::select! { - biased; - _ = cancel.cancelled() => {}, - _ = pipeline => {}, - } - }); + let (frame_tx, processor) = spawn_pipeline(config, runtime, replay_head); let (control_tx, mut control_rx) = tokio::sync::mpsc::channel::(CONTROL_CHANNEL_DEPTH); let session_cancel = runtime.cancel.child_token(); @@ -831,6 +836,44 @@ async fn run_session( } } +/// frames prep in parallel and buffered keeps them in cursor order +fn spawn_pipeline( + config: &IngestConfig, + runtime: &IngestRuntime, + replay_head: Option, +) -> ( + tokio::sync::mpsc::Sender, + tokio::task::JoinHandle<()>, +) { + let (frame_tx, frame_rx) = tokio::sync::mpsc::channel::(FRAME_CHANNEL_DEPTH); + let parallelism = config.parallelism.get(); + let processor_runtime = runtime.clone(); + let processor = tokio::spawn(async move { + let cancel = processor_runtime.cancel.clone(); + let prep_rt = processor_runtime.clone(); + let claim_rt = processor_runtime.clone(); + let resolve_rt = processor_runtime.clone(); + let commit_rt = processor_runtime; + + let pipeline = ReceiverStream::new(frame_rx) + .map(move |frame| prep_stage(frame, prep_rt.clone())) + .buffered(parallelism) + .then(move |staged| claim_stage(staged, claim_rt.clone())) + .map(move |staged| resolve_stage(staged, resolve_rt.clone(), replay_head)) + .buffered(parallelism) + .for_each(move |staged| { + commit_stage(staged, commit_rt.clone(), parallelism, replay_head) + }); + + tokio::select! { + biased; + _ = cancel.cancelled() => {}, + _ = pipeline => {}, + } + }); + (frame_tx, processor) +} + async fn timed_send( sink: &mut Box, msg: WsMessage, @@ -3651,6 +3694,7 @@ mod tests { &cfg, HydrantCursor::new(0), None, + None, CONNECT_TIMEOUT, &runtime, &FrameCounter::default(), @@ -3778,6 +3822,252 @@ mod tests { assert!(matches!(outcome, Ok(())), "got {outcome:?}"); } + struct TestChunk { + first: u64, + last: u64, + body: Bytes, + sha256: String, + latency: Duration, + } + + fn test_chunk(first: u64, last: u64, ids: std::ops::RangeInclusive) -> TestChunk { + let events: Vec<(u64, Vec)> = ids + .map(|id| { + let event = json!({ + "did": "did:plc:abalone", "rev": "3mabcdefghijk", "live": false, + "rkey": "3mabcdefghijk", "action": "delete", "collection": "app.bsky.feed.post", + }); + (id, serde_ipld_dagcbor::to_vec(&event).unwrap()) + }) + .collect(); + let body = crate::chunks::star_chunk("s1", first, last, &events); + let sha256 = { + use sha2::Digest as _; + sha2::Sha256::digest(&body) + .iter() + .map(|b| format!("{b:02x}")) + .collect::() + }; + TestChunk { + first, + last, + body, + sha256, + latency: Duration::ZERO, + } + } + + struct ChunkResponder { + head: u64, + chunks: Vec, + /// per chunk, not total + failures: usize, + fetches: std::sync::Mutex>, + } + + impl bobbin_runtime::MemHttpResponder for ChunkResponder { + fn respond( + &self, + request: &bobbin_runtime::HttpRequest, + ) -> bobbin_runtime::MemHttpResponse { + let path = request.url.path(); + let json = |value: serde_json::Value| { + bobbin_runtime::MemHttpBody::ok_json(value.to_string().into()) + }; + let mut latency = Duration::ZERO; + let body = match path { + "/stream/head" => json(json!({ "id": self.head })), + "/stream/manifest" => json(json!({ + "stream": "s1", + "generation": 1, + "through": self.chunks.last().map(|c| c.last), + "chunks": self.chunks.iter().map(|c| json!({ + "first": c.first, + "last": c.last, + "events": 0, + "bytes": c.body.len(), + "sha256": c.sha256, + })).collect::>(), + })), + _ => match path + .strip_prefix("/stream/chunk/") + .and_then(|sha| self.chunks.iter().find(|c| c.sha256 == sha)) + { + Some(chunk) => { + latency = chunk.latency; + let mut fetches = self.fetches.lock().unwrap(); + let fetched = fetches.entry(chunk.sha256.clone()).or_default(); + *fetched += 1; + if *fetched > self.failures { + bobbin_runtime::MemHttpBody::ok_json(chunk.body.clone()) + } else { + bobbin_runtime::MemHttpBody::status_only( + http::StatusCode::SERVICE_UNAVAILABLE, + ) + } + } + // a wrong hash still serves the first chunk's bytes and the ingest has to catch it + None if path.starts_with("/stream/chunk/") => { + bobbin_runtime::MemHttpBody::ok_json(self.chunks[0].body.clone()) + } + None => bobbin_runtime::MemHttpBody::status_only(http::StatusCode::NOT_FOUND), + }, + }; + bobbin_runtime::MemHttpResponse { + latency, + result: Ok(body), + } + } + } + + struct RecordingWs { + urls: std::sync::Mutex>, + inner: Arc, + } + + impl WsTransport for RecordingWs { + fn connect(&self, url: Url) -> WsConnectFuture { + self.urls.lock().unwrap().push(url.clone()); + self.inner.connect(url) + } + } + + struct ChunkedStart { + coverage: Arc, + ws: Arc, + cancel: CancellationToken, + task: tokio::task::JoinHandle>, + } + + fn chunked_start( + chunks: Vec, + failures: usize, + head: u64, + live: Vec>, + ) -> ChunkedStart { + let (ws, ws_rx) = tokio::sync::mpsc::channel(16); + live.into_iter() + .for_each(|msg| ws.try_send(msg).expect("room for the stream")); + let cancel = CancellationToken::new(); + let mut runtime = fresh_runtime(cancel.clone()); + let recording = Arc::new(RecordingWs { + urls: std::sync::Mutex::new(Vec::new()), + inner: ScriptedWs::dialing(WsConn { + sink: Box::new(PongingSink { + ws, + backlog: Backlog::Drained, + }), + stream: Box::new(ChannelWsStream { rx: ws_rx }), + }), + }); + runtime.ws = recording.clone(); + let responder = Arc::new(ChunkResponder { + head, + chunks, + failures, + fetches: std::sync::Mutex::default(), + }); + let http = bobbin_runtime::MemHttpTransport::shared(responder, runtime.clock.clone()); + runtime.hydrant = Some( + HydrantClient::new(Url::parse("ws://127.0.0.1:1").expect("hydrant url"), http) + .expect("hydrant client"), + ); + let coverage = runtime.coverage.clone(); + let mut cfg = IngestConfig::new(Url::parse("ws://127.0.0.1:1").expect("hydrant url")); + cfg.chunks = true; + let task = tokio::spawn(async move { run(cfg, runtime).await }); + ChunkedStart { + coverage, + ws: recording, + cancel, + task, + } + } + + impl ChunkedStart { + async fn ready_at(self, cursor: u64) -> (u64, Vec) { + self.coverage + .subscribe() + .wait_for(|c| c.is_ready() && c.last_cursor() == HydrantCursor::new(cursor)) + .await + .expect("coverage watch outlives the ingest"); + let events = self.coverage.snapshot().events_processed(); + let queries = self + .ws + .urls + .lock() + .unwrap() + .iter() + .map(|url| url.query().unwrap_or_default().to_owned()) + .collect(); + self.cancel.cancel(); + assert!(matches!(self.task.await.unwrap(), Ok(()))); + (events, queries) + } + } + + #[tokio::test(start_paused = true)] + async fn a_cold_start_streams_on_from_where_the_chunks_end() { + let chunks = vec![test_chunk(0, 5, 1..=3)]; + let start = chunked_start(chunks, 0, 7, vec![other_frame(6), other_frame(7)]); + let (_, connected) = start.ready_at(7).await; + assert_eq!(connected, vec!["cursor=6&stream=s1"]); + } + + #[tokio::test(start_paused = true)] + async fn a_chunk_that_fails_its_hash_leaves_the_replay_to_the_stream() { + let live = (1..=7).map(other_frame).collect(); + let mut chunk = test_chunk(0, 5, 1..=3); + chunk.sha256 = "not the hash".to_owned(); + let start = chunked_start(vec![chunk], 0, 7, live); + assert_eq!(start.ready_at(7).await, (7, vec!["cursor=0".to_owned()])); + } + + #[tokio::test(start_paused = true)] + async fn chunks_fetched_at_once_apply_in_order() { + let chunks = (0..8u64) + .map(|n| TestChunk { + latency: Duration::from_millis(10 * (8 - n)), + ..test_chunk(n * 4, n * 4 + 3, n * 4 + 1..=n * 4 + 3) + }) + .collect(); + let started = tokio::time::Instant::now(); + let start = chunked_start(chunks, 0, 33, vec![other_frame(32), other_frame(33)]); + let (events, connected) = start.ready_at(33).await; + // one at a time would take 360ms, over the 150ms bound + assert!(started.elapsed() < Duration::from_millis(150)); + // 3 per chunk, 2 live, and 1 step to where the chunks end + assert_eq!(events, 8 * 3 + 2 + 1); + assert_eq!(connected, vec!["cursor=32&stream=s1"]); + } + + #[tokio::test(start_paused = true)] + async fn a_chunk_that_fails_a_fetch_is_fetched_again() { + let chunks = vec![test_chunk(0, 3, 1..=3), test_chunk(4, 5, 4..=5)]; + let start = chunked_start(chunks, 2, 7, vec![other_frame(6), other_frame(7)]); + let (_, connected) = start.ready_at(7).await; + assert_eq!(connected, vec!["cursor=6&stream=s1"]); + } + + #[tokio::test(start_paused = true)] + async fn a_stream_that_isnt_the_chunks_ends_the_ingest() { + let refusal = json!({ + "type": "error", + "error": "StreamMismatch", + "message": "cursor is from stream s1, this is s2", + }); + let chunks = vec![test_chunk(0, 5, 1..=3)]; + let start = chunked_start(chunks, 0, 7, vec![Ok(WsMessage::Text(refusal.to_string()))]); + let outcome = tokio::time::timeout(RECONNECT_INITIAL_DELAY * 64, start.task) + .await + .expect("the ingest kept reconnecting to a stream that isn't its own") + .expect("ingest panicked"); + assert!( + matches!(&outcome, Err(IngestError::HydrantStream { code, .. }) if code == "StreamMismatch"), + "got {outcome:?}" + ); + } + #[tokio::test(start_paused = true)] async fn a_replay_reaching_the_hydrant_head_promotes_without_waiting_for_idle() { let replay = head_replay(http::StatusCode::OK, json!(3), 3); @@ -4350,6 +4640,7 @@ mod tests { &cfg, HydrantCursor::new(0), None, + None, CONNECT_TIMEOUT, &runtime, &FrameCounter::default(), diff --git a/bobbin/crates/xrpc/tests/await_record_e2e.rs b/bobbin/crates/xrpc/tests/await_record_e2e.rs index 1ad45841d..4c87cbcaf 100644 --- a/bobbin/crates/xrpc/tests/await_record_e2e.rs +++ b/bobbin/crates/xrpc/tests/await_record_e2e.rs @@ -145,6 +145,7 @@ async fn a_record_off_the_stream_answers_an_await() { start_cursor: HydrantCursor::new(0), parallelism: bobbin_ingest::DEFAULT_INGEST_PARALLELISM, idle_promote_min_events: bobbin_ingest::DEFAULT_IDLE_PROMOTE_MIN_EVENTS, + chunks: false, }; let ingesting = tokio::spawn(run_ingest(config, ingest)); let answer = tokio::time::timeout(Duration::from_secs(10), router(state).oneshot(request))