maarcadetweet/crates/at-firehose/src/consumer.rs

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<String>,
    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<Arc<AtomicBool>>,
    /// 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<String>, collections: Vec<String>) -> 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<AtomicBool>) -> 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<F, Fut>(&self, mut on_event: F) -> Result<()>
    where
        F: FnMut(JetstreamEvent) -> Fut + Send + 'static,
        Fut: std::future::Future<Output = Result<()>> + 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<F, Fut>(&self, on_event: &mut F) -> Result<()>
    where
        F: FnMut(JetstreamEvent) -> Fut + Send,
        Fut: std::future::Future<Output = Result<()>> + 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::<JetstreamEvent>(&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");
    }
}