use anyhow::Result; use futures::StreamExt; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use tokio_tungstenite::tungstenite::Message; use tracing::{error, info, warn}; use crate::event::JetstreamEvent; pub struct JetstreamConsumer { pub url: String, pub collections: Vec, pub max_backoff_secs: u64, /// Optional handle the consumer toggles on every connect/disconnect. /// Useful for health endpoints that want a live "are we connected?" /// signal without polling. pub connected: Option>, /// Optional cursor (`time_us`) to resume from. 0 means "no cursor / start /// fresh" which is Jetstream's default. Set via `with_cursor(...)`. pub cursor_us: i64, } impl JetstreamConsumer { pub fn new(url: impl Into, collections: Vec) -> Self { Self { url: url.into(), collections, max_backoff_secs: 30, connected: None, cursor_us: 0, } } /// The URL actually dialled: base URL plus `wantedCollections` and /// `cursor` as query parameters. /// /// This used to connect to the bare URL and then send /// `{"type": "options", "wantedCollections": [...]}` as a text frame. /// Jetstream ignores that, and silently: filters are query parameters, /// and the only message-based path (`options_update`) requires the /// connection to have been opened with `requireHello=true`. So every /// deployment that thought it was subscribing to six collections was in /// fact taking the entire public firehose — measured against /// jetstream1.us-east: 3119 events in 8 s unfiltered versus 520 for a /// single collection. On the dev database that quietly grew to 3.3 M /// posts; on the production instance the AppView had to be switched off /// to stop it filling the disk. /// /// Note what fixing this does *not* solve: the collections this project /// wants (`app.bsky.feed.post` / `like` / `repost` / `graph.follow`) are /// ~97 % of the firehose by volume. Correct filtering is necessary, not /// sufficient — an instance that does not want the whole public network /// in its index wants `wantedDids`, or no Jetstream at all. pub fn subscribe_url(&self) -> String { let mut url = self.url.trim_end_matches('&').to_string(); let mut sep = if url.contains('?') { '&' } else { '?' }; for c in &self.collections { url.push(sep); url.push_str("wantedCollections="); url.push_str(&urlencode(c)); sep = '&'; } if self.cursor_us > 0 { url.push(sep); url.push_str("cursor="); url.push_str(&self.cursor_us.to_string()); } url } /// Build a consumer that shares a connection-state flag with the caller. pub fn with_connected_flag(mut self, flag: Arc) -> Self { self.connected = Some(flag); self } /// Build a consumer with an explicit reconnect backoff ceiling (seconds). pub fn with_max_backoff_secs(mut self, max: u64) -> Self { self.max_backoff_secs = max.max(1); self } /// Build a consumer that starts from a previously-persisted cursor /// (microseconds since epoch). Setting this avoids the post-restart /// gap where Jetstream's default backfill window might miss events. pub fn with_cursor(mut self, cursor_us: i64) -> Self { self.cursor_us = cursor_us; self } pub async fn run(&self, mut on_event: F) -> Result<()> where F: FnMut(JetstreamEvent) -> Fut + Send + 'static, Fut: std::future::Future> + Send, { let mut backoff_secs: u64 = 1; loop { let r = self.connect_and_consume(&mut on_event).await; // Any non-error return (clean disconnect, error) means we lost // the connection; flip the flag if we own one and back off. if let Some(flag) = &self.connected { flag.store(false, Ordering::Relaxed); } match r { Ok(()) => warn!("jetstream stream ended, reconnecting"), Err(e) => error!("jetstream error: {e:#}"), } warn!("reconnecting in {backoff_secs}s"); tokio::time::sleep(Duration::from_secs(backoff_secs)).await; backoff_secs = (backoff_secs * 2).min(self.max_backoff_secs); } } async fn connect_and_consume(&self, on_event: &mut F) -> Result<()> where F: FnMut(JetstreamEvent) -> Fut + Send, Fut: std::future::Future> + Send, { let url = self.subscribe_url(); let (mut ws, _) = tokio_tungstenite::connect_async(&url).await?; info!("connected to jetstream: {url}"); if let Some(flag) = &self.connected { flag.store(true, Ordering::Relaxed); } while let Some(msg) = ws.next().await { let msg = msg?; if let Message::Text(text) = msg { if let Ok(ev) = serde_json::from_str::(&text) { on_event(ev).await?; } } } Ok(()) } } /// Percent-encode everything outside the unreserved set. Collection NSIDs are /// dots and letters today, but a `wantedDids` value carries `:` — encoding /// unconditionally keeps this correct if the caller passes one. fn urlencode(s: &str) -> String { s.chars() .map(|c| match c { 'A'..='Z' | 'a'..='z' | '0'..='9' | '-' | '_' | '.' | '~' => c.to_string(), other => { let mut buf = [0u8; 4]; other .encode_utf8(&mut buf) .bytes() .map(|b| format!("%{b:02X}")) .collect() } }) .collect() } #[cfg(test)] mod url_tests { use super::*; #[test] fn collections_go_into_the_query_string() { let c = JetstreamConsumer::new( "wss://jetstream1.us-east.bsky.network/subscribe", vec!["app.twi.post".into(), "app.bsky.feed.like".into()], ); assert_eq!( c.subscribe_url(), "wss://jetstream1.us-east.bsky.network/subscribe\ ?wantedCollections=app.twi.post&wantedCollections=app.bsky.feed.like" .replace(' ', "") ); } #[test] fn cursor_is_appended_and_respects_an_existing_query() { let mut c = JetstreamConsumer::new("wss://host/subscribe?compress=false", vec![]); c.cursor_us = 1234; assert_eq!(c.subscribe_url(), "wss://host/subscribe?compress=false&cursor=1234"); } #[test] fn no_filters_leaves_the_url_alone() { let c = JetstreamConsumer::new("wss://host/subscribe", vec![]); assert_eq!(c.subscribe_url(), "wss://host/subscribe"); } /// A DID contains `:`, which has to survive as `%3A` in a query value. #[test] fn values_are_percent_encoded() { assert_eq!(urlencode("did:plc:abc"), "did%3Aplc%3Aabc"); assert_eq!(urlencode("app.bsky.feed.post"), "app.bsky.feed.post"); } }