use super::*; use crate::db::types::TrimmedDid; use crate::types::StoredJetstreamEvent; use jacquard_common::IntoStatic; use jacquard_common::types::nsid::Nsid; use jacquard_common::types::string::Did; use smol_str::{SmolStr, ToSmolStr}; use std::collections::HashSet; use std::sync::Arc; pub struct JetstreamEventStream(mpsc::Receiver>); #[derive(Debug, Clone, thiserror::Error)] pub enum JetstreamStreamError { #[error("jetstream consumer too slow: {reason}")] ConsumerTooSlow { reason: String }, } impl JetstreamStreamError { pub fn code(&self) -> &'static str { match self { Self::ConsumerTooSlow { .. } => "ConsumerTooSlow", } } } impl Stream for JetstreamEventStream { type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.0.poll_recv(cx) } } #[derive(Clone, Default)] pub struct JetstreamFilter { inner: Arc>, } impl JetstreamFilter { pub fn new(options: JetstreamSubscriberOptions) -> Self { Self { inner: Arc::new(arc_swap::ArcSwap::from_pointee(options)), } } pub fn update(&self, options: JetstreamSubscriberOptions) { self.inner.store(Arc::new(options)); } pub fn max_message_size_bytes(&self) -> u32 { self.inner.load().max_message_size_bytes() } pub(crate) fn wants(&self, event: &StoredJetstreamEvent<'_>) -> bool { self.inner.load().wants(event) } /// whether backfilled (`live: false`) commits that finish after /// connecting are sent from the live tail. only when `historical` is /// asked for explicitly; cursor replay sends them by default. pub(crate) fn tails_historical(&self) -> bool { self.inner .load() .wanted_event_types .as_ref() .is_some_and(|wanted| wanted.historical) } } #[derive(Clone, Default)] pub struct JetstreamSubscriberOptions { wanted_collections: Option, wanted_dids: Arc>>, max_message_size_bytes: u32, wanted_event_types: Option, } #[derive(Clone)] struct WantedCollections { prefixes: Vec, full_paths: HashSet, } #[derive(Clone)] struct WantedEventTypes { live: bool, historical: bool, } impl JetstreamSubscriberOptions { pub fn parse( wanted_collections: &[String], wanted_dids: &[String], max_message_size_bytes: u32, wanted_event_types: &[String], ) -> std::result::Result { let wanted_collections = parse_wanted_collections(wanted_collections)?; let wanted_dids = parse_wanted_dids(wanted_dids)?; let wanted_event_types = parse_wanted_event_types(wanted_event_types)?; Ok(Self { wanted_collections, wanted_dids: Arc::new(wanted_dids), max_message_size_bytes, wanted_event_types, }) } pub fn max_message_size_bytes(&self) -> u32 { self.max_message_size_bytes } pub(crate) fn wants(&self, event: &StoredJetstreamEvent<'_>) -> bool { if let Some(wanted) = &self.wanted_event_types { let wanted_type = match event.is_live() { true => wanted.live, false => wanted.historical, }; if !wanted_type { return false; } } if !self.wanted_dids.is_empty() { let mut did = Vec::with_capacity(event.did().len()); event.did().write_to_vec(&mut did); if !self.wanted_dids.contains(&did) { return false; } } let Some(wanted_collections) = &self.wanted_collections else { return true; }; let Some(collection) = event.collection() else { return true; }; wanted_collections.full_paths.contains(collection) || wanted_collections .prefixes .iter() .any(|prefix| collection.starts_with(prefix.as_str())) } } impl Hydrant { /// subscribe to the jetstream compatible stream. the replay runs as fast as you read it. /// only `/subscribe` readers on the network are capped, when `HYDRANT_THROTTLED_STREAMS` /// lists it. pub fn subscribe_jetstream( &self, cursor: Option, filter: JetstreamFilter, ) -> JetstreamEventStream { let opts = stream::StreamOptions::from_config(&self.config); self.subscribe_jetstream_with(cursor, filter, opts) } /// [`Self::subscribe_jetstream`] for a reader on the network, whose replay is capped. pub(crate) fn subscribe_jetstream_network( &self, cursor: Option, filter: JetstreamFilter, ) -> JetstreamEventStream { let opts = stream::StreamOptions::network(&self.config, crate::config::StreamKind::Jetstream); self.subscribe_jetstream_with(cursor, filter, opts) } fn subscribe_jetstream_with( &self, cursor: Option, filter: JetstreamFilter, opts: stream::StreamOptions, ) -> JetstreamEventStream { let (tx, rx) = mpsc::channel(stream::STREAM_CHANNEL_CAPACITY); let state = self.state.clone(); let runtime = tokio::runtime::Handle::current(); std::thread::Builder::new() .name("hydrant-jetstream".into()) .spawn(move || { let _g = runtime.enter(); stream::jetstream_stream_thread(state, tx, cursor, filter, opts); }) .expect("failed to spawn Jetstream thread"); JetstreamEventStream(rx) } } fn parse_wanted_collections( provided: &[String], ) -> std::result::Result, String> { if provided.is_empty() { return Ok(None); } if provided.len() > 100 { return Err("too many wanted collections".into()); } let mut prefixes = Vec::new(); let mut full_paths = HashSet::new(); for collection_raw in provided { for collection in collection_raw.split(',') { if collection.is_empty() { continue; } if let Some(prefix) = collection.strip_suffix(".*") { validate_collection_prefix(prefix) .map_err(|_| format!("invalid collection prefix: {collection}"))?; prefixes.push(format!("{prefix}.").to_smolstr()); continue; } let nsid = Nsid::new(collection) .map_err(|_| format!("invalid collection: {collection}"))? .into_static(); full_paths.insert(nsid.as_str().to_smolstr()); } } Ok(Some(WantedCollections { prefixes, full_paths, })) } fn validate_collection_prefix(prefix: &str) -> std::result::Result<(), String> { let synthetic = format!("{prefix}.x"); Nsid::new(synthetic.as_str()) .map(|_| ()) .map_err(|e| e.to_string()) } fn parse_wanted_event_types( provided: &[String], ) -> std::result::Result, String> { if provided.is_empty() { return Ok(None); } let mut live = false; let mut historical = false; for ev in provided { match ev.as_str() { "live" => live = true, "historical" => historical = true, _ => return Err(format!("unknown wantedEventTypes value: {}", ev)), } } Ok(Some(WantedEventTypes { live, historical })) } fn parse_wanted_dids(provided: &[String]) -> std::result::Result>, String> { let mut wanted = HashSet::new(); for did_raw in provided { for did_str in did_raw.split(',') { if did_str.is_empty() { continue; } let did = Did::new_owned(did_str).map_err(|_| "invalid wanted DID".to_string())?; let trimmed = TrimmedDid::from(&did); let mut buf = Vec::with_capacity(trimmed.len()); trimmed.write_to_vec(&mut buf); wanted.insert(buf); } } if wanted.len() > 10_000 { return Err("too many wanted DIDs".into()); } Ok(wanted) } #[cfg(test)] mod tests { use super::*; #[test] fn collection_prefixes_accept_domain_prefixes() { let opts = JetstreamSubscriberOptions::parse(&["app.bsky.*".into()], &[], 0, &[]).unwrap(); let wanted = opts.wanted_collections.unwrap(); assert!(wanted.prefixes.iter().any(|prefix| prefix == "app.bsky.")); assert!( JetstreamSubscriberOptions::parse(&["app.bsky.feed.po*".into()], &[], 0, &[]).is_err() ); } #[test] #[cfg(feature = "indexer_stream")] fn full_path_filter_matches_commit_collection() { use crate::db::types::TrimmedDid; use crate::types::StoredJetstreamEvent; use jacquard_common::{CowStr, IntoStatic}; use smol_str::ToSmolStr; let opts = JetstreamSubscriberOptions::parse(&["app.bsky.feed.post".into()], &[], 0, &[]) .expect("app.bsky.feed.post must be a valid collection"); let did = TrimmedDid::from( &jacquard_common::types::string::Did::new_static("did:plc:abc123").unwrap(), ) .into_static(); let matching = StoredJetstreamEvent::Commit { did: did.clone(), collection: CowStr::Owned("app.bsky.feed.post".to_smolstr()), event_id: 1, live: true, }; let non_matching = StoredJetstreamEvent::Commit { did: did.clone(), collection: CowStr::Owned("app.bsky.actor.profile".to_smolstr()), event_id: 2, live: true, }; let account = StoredJetstreamEvent::Account { did, active: true, status: None, seq: 1, time: crate::ingest::stream::Datetime(chrono::Utc::now().into()), }; assert!( opts.wants(&matching), "app.bsky.feed.post commit should match" ); assert!( !opts.wants(&non_matching), "other collection must not match" ); assert!( opts.wants(&account), "account events must pass through collection filter" ); } #[test] #[cfg(feature = "indexer_stream")] fn backfilled_commits_tail_only_when_asked_for() { use jacquard_common::CowStr; let did = TrimmedDid::from(&Did::new_static("did:plc:abc123").unwrap()).into_static(); let commit = |live| StoredJetstreamEvent::Commit { did: did.clone(), collection: CowStr::Borrowed("app.bsky.feed.post"), event_id: 1, live, }; let filter = |types: &[&str]| { let types: Vec<_> = types.iter().map(|t| t.to_string()).collect(); JetstreamFilter::new(JetstreamSubscriberOptions::parse(&[], &[], 0, &types).unwrap()) }; // (types, passes live, passes historical, tails historical) for (types, live, historical, tails) in [ (&[][..], true, true, false), (&["live"][..], true, false, false), (&["historical"][..], false, true, true), (&["live", "historical"][..], true, true, true), ] { let filter = filter(types); assert_eq!(filter.wants(&commit(true)), live, "{types:?}"); assert_eq!(filter.wants(&commit(false)), historical, "{types:?}"); assert_eq!(filter.tails_historical(), tails, "{types:?}"); } } #[test] fn wanted_did_limit_counts_unique_dids() { let dids = vec!["did:plc:abc123".to_string(); 10_001]; let opts = JetstreamSubscriberOptions::parse(&[], &dids, 0, &[]).unwrap(); assert_eq!(opts.wanted_dids.len(), 1); } }