···
···
use atrium_api::record::KnownRecord;
···
52
-
events::JetstreamEvent,
/// The Jetstream endpoints officially provided by Bluesky themselves.
···
pub wanted_dids: Vec<exports::Did>,
/// The compression algorithm to request and use for the WebSocket connection (if any).
pub compression: JetstreamCompression,
170
-
/// An optional timestamp to begin playback from.
172
+
/// Enable automatic cursor for auto-reconnect
172
-
/// An absent cursor or a cursor from the future will result in live-tail operation.
174
+
/// By default, reconnects will never set a cursor for the connection, so a small number of
175
+
/// events will always be dropped.
174
-
/// When reconnecting, use the time_us from your most recently processed event and maybe
175
-
/// provide a negative buffer (i.e. subtract a few seconds) to ensure gapless playback.
176
-
pub cursor: Option<chrono::DateTime<Utc>>,
177
+
/// If you want gapless playback across reconnects, set this to `true`. If you always want
178
+
/// the latest available events and can tolerate missing some: `false`.
179
+
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,
···
wanted_collections: Vec::new(),
compression: JetstreamCompression::None,
203
+
replay_on_reconnect: false,
channel_size: 4096, // a few seconds of firehose buffer
record_type: PhantomData,
···
230
-
.map(|c| ("cursor", c.timestamp_micros().to_string()));
let params = did_search_query
.chain(collection_search_query)
.chain(std::iter::once(compression))
.collect::<Vec<(&str, String)>>();
Url::parse_with_params(endpoint, params)
···
/// 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<JetstreamReceiver<R>, ConnectionError> {
277
+
self.base_connect(None).await
280
+
/// Connects to a Jetstream instance as defined in the [JetstreamConfig] with playback from a
283
+
/// A cursor from the future will result in live-tail operation.
285
+
/// The cursor is only used for first successfull connection -- on auto-reconnect it will
286
+
/// live-tail by default. Set `replay_on_reconnect: true` in the config if you need to
287
+
/// receive every event, which will keep track of the last-seen cursor and reconnect from
289
+
pub async fn connect_cursor(
292
+
) -> Result<JetstreamReceiver<R>, ConnectionError> {
293
+
self.base_connect(Some(cursor)).await
296
+
async fn base_connect(
298
+
cursor: Option<Cursor>,
299
+
) -> Result<JetstreamReceiver<R>, ConnectionError> {
// We validate the config again for good measure. Probably not necessary but it can't hurt.
···
.construct_endpoint(&self.config.endpoint)
.map_err(ConnectionError::InvalidEndpoint)?;
312
+
let replay_on_reconnect = self.config.replay_on_reconnect;
tokio::task::spawn(async move {
let base_delay_ms = 1_000; // 1 second
···
let success_threshold_s = 15; // 15 seconds, retry count is reset if we were connected at least this long
let mut retry_attempt = 0;
321
+
let mut connect_cursor = cursor;
let dict = DecoderDictionary::copy(JETSTREAM_ZSTD_DICTIONARY);
325
+
let mut configured_endpoint = configured_endpoint.clone();
326
+
if let Some(ref cursor) = connect_cursor {
327
+
configured_endpoint
329
+
.append_pair("cursor", &cursor.to_jetstream());
332
+
let mut last_cursor = connect_cursor.clone();
if let Ok((ws_stream, _)) = connect_async(&configured_endpoint).await {
let t_connected = Instant::now();
304
-
if let Err(e) = websocket_task(dict, ws_stream, send_channel.clone()).await {
338
+
websocket_task(dict, ws_stream, send_channel.clone(), &mut last_cursor)
log::error!("Jetstream closed after encountering error: {e:?}");
log::error!("Jetstream connection closed cleanly");
if t_connected.elapsed() > Duration::from_secs(success_threshold_s) {
if retry_attempt >= max_retries {
316
-
eprintln!("max retries, bye");
351
+
log::error!("hit max retries, bye");
320
-
eprintln!("will try to reconnect");
355
+
connect_cursor = if replay_on_reconnect {
322
-
// Exponential backoff
323
-
let delay_ms = base_delay_ms * (2_u64.pow(retry_attempt));
325
-
log::error!("Connection failed, retrying in {delay_ms}ms...");
326
-
tokio::time::sleep(Duration::from_millis(delay_ms.min(max_delay_ms))).await;
327
-
log::info!("Attempting to reconnect...")
361
+
if retry_attempt > 0 {
362
+
// Exponential backoff
363
+
let delay_ms = base_delay_ms * (2_u64.pow(retry_attempt));
364
+
log::error!("Connection failed, retrying in {delay_ms}ms...");
365
+
tokio::time::sleep(Duration::from_millis(delay_ms.min(max_delay_ms))).await;
366
+
log::info!("Attempting to reconnect...");
log::error!("Connection retries exhausted. Jetstream is disconnected.");
···
dictionary: DecoderDictionary<'_>,
ws: WebSocketStream<MaybeTlsStream<TcpStream>>,
send_channel: JetstreamSender<R>,
382
+
last_cursor: &mut Option<Cursor>,
) -> 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();
···
376
-
let event = serde_json::from_str(&json)
417
+
let event: JetstreamEvent<R> = serde_json::from_str(&json)
.map_err(JetstreamEventError::ReceivedMalformedJSON)?;
419
+
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.
383
-
"All receivers for the Jetstream connection have been dropped, closing connection."
425
+
"All receivers for the Jetstream connection have been dropped, closing connection."
closing_connection = true;
428
+
} else if let Some(v) = last_cursor.as_mut() {
Message::Binary(zstd_json) => {
389
-
let mut cursor = Cursor::new(zstd_json);
433
+
let mut cursor = IoCursor::new(zstd_json);
let mut decoder = zstd::stream::Decoder::with_prepared_dictionary(
···
.read_to_string(&mut json)
.map_err(JetstreamEventError::CompressionDecoderError)?;
401
-
let event = serde_json::from_str(&json)
445
+
let event: JetstreamEvent<R> = serde_json::from_str(&json)
.map_err(JetstreamEventError::ReceivedMalformedJSON)?;
447
+
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.
408
-
"All receivers for the Jetstream connection have been dropped, closing connection..."
453
+
"All receivers for the Jetstream connection have been dropped, closing connection..."
closing_connection = true;
456
+
} else if let Some(v) = last_cursor.as_mut() {