From ba4ba6a87da1c472b4ec369faf09db73aef12bc6 Mon Sep 17 00:00:00 2001 From: phil Date: Sat, 8 Mar 2025 16:51:55 -0500 Subject: [PATCH] cursor handling --- jetstream/README.md | 5 +- jetstream/src/events/mod.rs | 77 ++++++++++++++++++++++++- jetstream/src/lib.rs | 111 +++++++++++++++++++++++++----------- 3 files changed, 157 insertions(+), 36 deletions(-) diff --git a/jetstream/README.md b/jetstream/README.md index def4d4b..cd96072 100644 --- a/jetstream/README.md +++ b/jetstream/README.md @@ -1,7 +1,6 @@ -# jetstream-oxide +# fork of the awesome jetstream-oxide -[![Crate](https://img.shields.io/crates/v/jetstream-oxide.svg)](https://crates.io/crates/jetstream-oxide) -[![docs.rs](https://docs.rs/jetstream-oxide/badge.svg)](https://docs.rs/jetstream-oxide/latest/jetstream_oxide) +fork note: this readme is likely a bit out of date! i've been messing around with some apis. A typed Rust library for easily interacting with and consuming the Bluesky [Jetstream](https://github.com/bluesky-social/jetstream) diff --git a/jetstream/src/events/mod.rs b/jetstream/src/events/mod.rs index f8784a6..6db8e96 100644 --- a/jetstream/src/events/mod.rs +++ b/jetstream/src/events/mod.rs @@ -2,15 +2,27 @@ pub mod account; pub mod commit; pub mod identity; +use std::time::{ + Duration, + SystemTime, + UNIX_EPOCH, +}; + use serde::Deserialize; use crate::exports; +/// Opaque wrapper for the time_us cursor used by jetstream +/// +/// Generally, you should use a cursor +#[derive(Deserialize, Debug, Clone)] +pub struct Cursor(u64); + /// Basic data that is included with every event. #[derive(Deserialize, Debug)] pub struct EventInfo { pub did: exports::Did, - pub time_us: u64, + pub time_us: Cursor, pub kind: EventKind, } @@ -29,3 +41,66 @@ pub enum EventKind { Identity, Account, } + +impl JetstreamEvent { + pub fn cursor(&self) -> Cursor { + match self { + JetstreamEvent::Commit(commit::CommitEvent::Create { info, .. }) => { + info.time_us.clone() + } + JetstreamEvent::Commit(commit::CommitEvent::Update { info, .. }) => { + info.time_us.clone() + } + JetstreamEvent::Commit(commit::CommitEvent::Delete { info, .. }) => { + info.time_us.clone() + } + JetstreamEvent::Identity(e) => e.info.time_us.clone(), + JetstreamEvent::Account(e) => e.info.time_us.clone(), + } + } +} + +impl Cursor { + /// Get a cursor that will consume all available jetstream replay + /// + /// This sets the cursor to zero. + /// + /// Jetstream instances typically only have a few days of replay. + pub fn from_start() -> Self { + Self(0) + } + /// Get a cursor for a specific time + /// + /// Panics: if t is older than the unix epoch: Jan 1, 1970. + /// + /// If you want to receive all available jetstream replay (typically a few days), use + /// .from_start() + pub fn at(t: SystemTime) -> Self { + let unix_dt = t + .duration_since(UNIX_EPOCH) + .expect("cannot set jetstream cursor earlier than unix epoch"); + Self(unix_dt.as_micros() as u64) + } + /// Get a cursor rewound from now by this amount + /// + /// Panics: if d is greater than the time since the unix epoch: Jan 1, 1970. + /// + /// Jetstream instances typically only have a few days of replay. + pub fn back_by(d: Duration) -> Self { + Self::at(SystemTime::now() - d) + } + /// Get a Cursor from a raw u64 + /// + /// For example, from a jetstream event's `time_us` field. + pub fn from_raw_u64(time_us: u64) -> Self { + Self(time_us) + } + /// Get the raw u64 value from this cursor. + pub fn to_raw_u64(&self) -> u64 { + self.0 + } + /// Format the cursor value for use in a jetstream connection url querystring + pub fn to_jetstream(&self) -> String { + self.0.to_string() + } +} diff --git a/jetstream/src/lib.rs b/jetstream/src/lib.rs index 7d35a82..a87b3ef 100644 --- a/jetstream/src/lib.rs +++ b/jetstream/src/lib.rs @@ -4,7 +4,7 @@ pub mod exports; use std::{ io::{ - Cursor, + Cursor as IoCursor, Read, }, marker::PhantomData, @@ -16,7 +16,6 @@ use std::{ }; use atrium_api::record::KnownRecord; -use chrono::Utc; use futures_util::{ stream::StreamExt, SinkExt, @@ -49,7 +48,10 @@ use crate::{ ConnectionError, JetstreamEventError, }, - events::JetstreamEvent, + events::{ + Cursor, + JetstreamEvent, + }, }; /// The Jetstream endpoints officially provided by Bluesky themselves. @@ -167,13 +169,14 @@ pub struct JetstreamConfig { pub wanted_dids: Vec, /// The compression algorithm to request and use for the WebSocket connection (if any). pub compression: JetstreamCompression, - /// An optional timestamp to begin playback from. + /// Enable automatic cursor for auto-reconnect /// - /// An absent cursor or a cursor from the future will result in live-tail operation. + /// By default, reconnects will never set a cursor for the connection, so a small number of + /// events will always be dropped. /// - /// When reconnecting, use the time_us from your most recently processed event and maybe - /// provide a negative buffer (i.e. subtract a few seconds) to ensure gapless playback. - pub cursor: Option>, + /// If you want gapless playback across reconnects, set this to `true`. If you always want + /// the latest available events and can tolerate missing some: `false`. + pub replay_on_reconnect: bool, /// Maximum size of send channel for jetstream events. /// /// If your consuming task can't keep up with every new jetstream event in real-time, @@ -197,7 +200,7 @@ impl Default for JetstreamConfig { wanted_collections: Vec::new(), wanted_dids: Vec::new(), compression: JetstreamCompression::None, - cursor: None, + replay_on_reconnect: false, channel_size: 4096, // a few seconds of firehose buffer record_type: PhantomData, } @@ -225,14 +228,9 @@ impl JetstreamConfig { }, ); - let cursor = self - .cursor - .map(|c| ("cursor", c.timestamp_micros().to_string())); - let params = did_search_query .chain(collection_search_query) .chain(std::iter::once(compression)) - .chain(cursor) .collect::>(); Url::parse_with_params(endpoint, params) @@ -276,6 +274,29 @@ impl JetstreamConnector { /// A [JetstreamReceiver] is returned which can be used to respond to events. When all instances /// of this receiver are dropped, the connection and task are automatically closed. pub async fn connect(&self) -> Result, ConnectionError> { + self.base_connect(None).await + } + + /// Connects to a Jetstream instance as defined in the [JetstreamConfig] with playback from a + /// cursor + /// + /// A cursor from the future will result in live-tail operation. + /// + /// The cursor is only used for first successfull connection -- on auto-reconnect it will + /// live-tail by default. Set `replay_on_reconnect: true` in the config if you need to + /// receive every event, which will keep track of the last-seen cursor and reconnect from + /// there. + pub async fn connect_cursor( + &self, + cursor: Cursor, + ) -> Result, ConnectionError> { + self.base_connect(Some(cursor)).await + } + + async fn base_connect( + &self, + cursor: Option, + ) -> Result, ConnectionError> { // We validate the config again for good measure. Probably not necessary but it can't hurt. self.config .validate() @@ -288,6 +309,8 @@ impl JetstreamConnector { .construct_endpoint(&self.config.endpoint) .map_err(ConnectionError::InvalidEndpoint)?; + let replay_on_reconnect = self.config.replay_on_reconnect; + tokio::task::spawn(async move { let max_retries = 30; let base_delay_ms = 1_000; // 1 second @@ -295,36 +318,53 @@ impl JetstreamConnector { let success_threshold_s = 15; // 15 seconds, retry count is reset if we were connected at least this long let mut retry_attempt = 0; + let mut connect_cursor = cursor; loop { let dict = DecoderDictionary::copy(JETSTREAM_ZSTD_DICTIONARY); + let mut configured_endpoint = configured_endpoint.clone(); + if let Some(ref cursor) = connect_cursor { + configured_endpoint + .query_pairs_mut() + .append_pair("cursor", &cursor.to_jetstream()); + } + + let mut last_cursor = connect_cursor.clone(); + retry_attempt += 1; if let Ok((ws_stream, _)) = connect_async(&configured_endpoint).await { let t_connected = Instant::now(); - if let Err(e) = websocket_task(dict, ws_stream, send_channel.clone()).await { + if let Err(e) = + websocket_task(dict, ws_stream, send_channel.clone(), &mut last_cursor) + .await + { log::error!("Jetstream closed after encountering error: {e:?}"); } else { log::error!("Jetstream connection closed cleanly"); } if t_connected.elapsed() > Duration::from_secs(success_threshold_s) { retry_attempt = 0; - continue; } } if retry_attempt >= max_retries { - eprintln!("max retries, bye"); + log::error!("hit max retries, bye"); break; } - eprintln!("will try to reconnect"); - - // Exponential backoff - let delay_ms = base_delay_ms * (2_u64.pow(retry_attempt)); - - log::error!("Connection failed, retrying in {delay_ms}ms..."); - tokio::time::sleep(Duration::from_millis(delay_ms.min(max_delay_ms))).await; - log::info!("Attempting to reconnect...") + connect_cursor = if replay_on_reconnect { + last_cursor + } else { + None + }; + + if retry_attempt > 0 { + // Exponential backoff + let delay_ms = base_delay_ms * (2_u64.pow(retry_attempt)); + log::error!("Connection failed, retrying in {delay_ms}ms..."); + tokio::time::sleep(Duration::from_millis(delay_ms.min(max_delay_ms))).await; + log::info!("Attempting to reconnect..."); + } } log::error!("Connection retries exhausted. Jetstream is disconnected."); }); @@ -339,6 +379,7 @@ async fn websocket_task( dictionary: DecoderDictionary<'_>, ws: WebSocketStream>, send_channel: JetstreamSender, + last_cursor: &mut Option, ) -> Result<(), JetstreamEventError> { // TODO: Use the write half to allow the user to change configuration settings on the fly. let (socket_write, mut socket_read) = ws.split(); @@ -373,20 +414,23 @@ async fn websocket_task( Some(Ok(message)) => { match message { Message::Text(json) => { - let event = serde_json::from_str(&json) + let event: JetstreamEvent = serde_json::from_str(&json) .map_err(JetstreamEventError::ReceivedMalformedJSON)?; + let event_cursor = event.cursor(); if send_channel.send(event).await.is_err() { // We can assume that all receivers have been dropped, so we can close // the connection and exit the task. log::info!( - "All receivers for the Jetstream connection have been dropped, closing connection." - ); + "All receivers for the Jetstream connection have been dropped, closing connection." + ); closing_connection = true; + } else if let Some(v) = last_cursor.as_mut() { + *v = event_cursor; } } Message::Binary(zstd_json) => { - let mut cursor = Cursor::new(zstd_json); + let mut cursor = IoCursor::new(zstd_json); let mut decoder = zstd::stream::Decoder::with_prepared_dictionary( &mut cursor, &dictionary, @@ -398,16 +442,19 @@ async fn websocket_task( .read_to_string(&mut json) .map_err(JetstreamEventError::CompressionDecoderError)?; - let event = serde_json::from_str(&json) + let event: JetstreamEvent = serde_json::from_str(&json) .map_err(JetstreamEventError::ReceivedMalformedJSON)?; + let event_cursor = event.cursor(); if send_channel.send(event).await.is_err() { // We can assume that all receivers have been dropped, so we can close // the connection and exit the task. log::info!( - "All receivers for the Jetstream connection have been dropped, closing connection..." - ); + "All receivers for the Jetstream connection have been dropped, closing connection..." + ); closing_connection = true; + } else if let Some(v) = last_cursor.as_mut() { + *v = event_cursor; } } Message::Ping(vec) => { -- 2.51.2