Something went wrong. Try again.
Something went wrong. Try again.
Rust
at batch-writes
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430use std::collections::HashSet;use std::str::FromStr;
use anyhow::{anyhow, Context, Result};use chrono::DateTime;use futures_util::SinkExt;use futures_util::StreamExt;use http::HeaderValue;use http::Uri;use tokio::sync::mpsc::Sender;use tokio::time::{sleep, Instant};use tokio_util::sync::CancellationToken;use tokio_websockets::{ClientBuilder, Message};
use crate::config;use crate::matcher::FeedMatchers;use crate::matcher::Match;use crate::storage;use crate::storage::consumer_control_get;use crate::storage::denylist_all;use crate::storage::StoragePool;use crate::writer::WriteCommand;
const MAX_MESSAGE_SIZE: usize = 25000;
#[derive(Clone)]pub struct ConsumerTaskConfig { pub user_agent: String, pub compression: bool, pub zstd_dictionary_location: String, pub jetstream_hostname: String, pub feeds: config::Feeds, pub collections: Vec<String>,}
pub struct ConsumerTask { cancellation_token: CancellationToken, pool: StoragePool, config: ConsumerTaskConfig, feed_matchers: FeedMatchers, write_tx: Sender<WriteCommand>,}
impl ConsumerTask { pub fn new( pool: StoragePool, config: ConsumerTaskConfig, cancellation_token: CancellationToken, write_tx: Sender<WriteCommand>, ) -> Result<Self> { let feed_matchers = FeedMatchers::from_config(&config.feeds)?;
Ok(Self { pool, cancellation_token, config, feed_matchers, write_tx, }) }
pub async fn run_background(&self) -> Result<()> { tracing::debug!("ConsumerTask started");
let last_time_us = consumer_control_get(&self.pool, &self.config.jetstream_hostname).await?;
tracing::info!(cursor = ?last_time_us, "loaded cursor from database");
let mut denylist: HashSet<String> = denylist_all(&self.pool).await?; tracing::info!(count = denylist.len(), "loaded denylist from database");
let cursor_param = if let Some(cursor) = last_time_us { format!("&cursor={}", cursor) } else { String::new() };
let uri = Uri::from_str(&format!( "wss://{}/subscribe?compress={}&requireHello=true{}", self.config.jetstream_hostname, self.config.compression, cursor_param )) .context("invalid jetstream URL")?;
tracing::info!(uri = %uri, "connecting to jetstream");
let (mut client, _) = ClientBuilder::from_uri(uri) .add_header( http::header::USER_AGENT, HeaderValue::from_str(&self.config.user_agent)?, ) .connect() .await .map_err(|err| anyhow::Error::new(err).context("cannot connect to jetstream"))?;
let update = model::SubscriberSourcedMessage::Update { wanted_collections: self.config.collections.clone(), wanted_dids: vec![], max_message_size_bytes: MAX_MESSAGE_SIZE as u64, cursor: None, }; let serialized_update = serde_json::to_string(&update) .map_err(|err| anyhow::Error::msg(err).context("cannot serialize update"))?;
tracing::info!(message = %serialized_update, "sending subscription update to jetstream");
client .send(Message::text(serialized_update)) .await .map_err(|err| anyhow::Error::msg(err).context("cannot send update"))?;
let mut decompressor = if self.config.compression { // mkdir -p data/ && curl -o data/zstd_dictionary https://github.com/bluesky-social/jetstream/raw/refs/heads/main/pkg/models/zstd_dictionary let data: Vec<u8> = std::fs::read(self.config.zstd_dictionary_location.clone()) .context("unable to load zstd dictionary")?; zstd::bulk::Decompressor::with_dictionary(&data) .map_err(|err| anyhow::Error::msg(err).context("cannot create decompressor"))? } else { zstd::bulk::Decompressor::new() .map_err(|err| anyhow::Error::msg(err).context("cannot create decompressor"))? };
let interval = std::time::Duration::from_secs(120); let sleeper = sleep(interval); tokio::pin!(sleeper);
let heartbeat_interval = std::time::Duration::from_secs(15); let heartbeat_sleeper = sleep(heartbeat_interval); tokio::pin!(heartbeat_sleeper);
let denylist_interval = std::time::Duration::from_secs(60); let denylist_sleeper = sleep(denylist_interval); tokio::pin!(denylist_sleeper);
let mut time_usec = 0i64;
// Diagnostics: socket drain rate and read-stall detection. let mut messages_read: u64 = 0; let mut messages_since_heartbeat: u64 = 0; let mut last_message_at: Option<Instant> = None;
loop { tokio::select! { () = self.cancellation_token.cancelled() => { break; }, () = &mut sleeper => { // Hand the cursor to the writer task, which flushes any // buffered writes before persisting it. self.write_tx .send(WriteCommand::Checkpoint(time_usec)) .await .map_err(|_| anyhow!("write channel closed"))?; sleeper.as_mut().reset(Instant::now() + interval); }, () = &mut denylist_sleeper => { denylist = denylist_all(&self.pool).await?; denylist_sleeper.as_mut().reset(Instant::now() + denylist_interval); }, () = &mut heartbeat_sleeper => { let timestamp = (time_usec > 0) .then(|| DateTime::from_timestamp_micros(time_usec).map(|dt| dt.to_rfc3339())) .flatten(); tracing::info!( time_us = time_usec, timestamp = ?timestamp, messages_since_last = messages_since_heartbeat, write_channel_available = self.write_tx.capacity(), "consumer heartbeat" ); messages_since_heartbeat = 0; heartbeat_sleeper.as_mut().reset(Instant::now() + heartbeat_interval); }, item = client.next() => { if item.is_none() { let since_last_ms = last_message_at.map(|t| t.elapsed().as_millis()); tracing::warn!( messages_read, since_last_message_ms = ?since_last_ms, "jetstream connection closed" ); break; } let item = item.unwrap();
if let Err(err) = item { let since_last_ms = last_message_at.map(|t| t.elapsed().as_millis()); tracing::error!(error = ?err, messages_read, since_last_message_ms = ?since_last_ms, "error reading from jetstream"); continue; } let item = item.unwrap();
messages_read += 1; messages_since_heartbeat += 1; last_message_at = Some(Instant::now());
// Handle control frames once, independent of payload // compression (control frames are never compressed). if item.is_ping() { // tokio-websockets queues a pong when it reads a ping // but only sends it on flush. We otherwise never write // after the initial subscription, so flush now or // jetstream never sees our keepalive reply. if let Err(err) = client.flush().await { tracing::warn!(error = ?err, "failed to flush pong response"); } continue; } if item.is_pong() { continue; } if item.is_close() { match item.as_close() { Some((code, reason)) => { tracing::warn!(?code, reason = %reason, "jetstream sent close frame"); } None => tracing::warn!("jetstream sent close frame"), } continue; }
let event = if self.config.compression { if !item.is_binary() { // Log unexpected non-binary message types tracing::warn!("received unexpected non-binary message from jetstream (not ping/pong/close)"); continue; } let payload = item.into_payload();
let decoded = decompressor.decompress(&payload, MAX_MESSAGE_SIZE * 3); if let Err(err) = decoded { tracing::debug!(err = ?err, "cannot decompress message"); continue; } let decoded = decoded.unwrap(); serde_json::from_slice::<model::Event>(&decoded) .context(anyhow!("cannot deserialize message")) } else { if !item.is_text() { // Log unexpected non-text message types tracing::warn!("received unexpected non-text message from jetstream (not ping/pong/close)"); continue; } item.as_text() .ok_or(anyhow!("cannot convert message to text")) .and_then(|value| { serde_json::from_str::<model::Event>(value) .context(anyhow!("cannot deserialize message")) }) }; if let Err(err) = event { tracing::error!(error = ?err, "error processing jetstream message");
continue; } let event = event.unwrap();
let previous_time_usec = time_usec; time_usec = std::cmp::max(time_usec, event.time_us);
if previous_time_usec == 0 { let datetime = DateTime::from_timestamp_micros(event.time_us) .map(|dt| dt.to_rfc3339()) .unwrap_or_else(|| format!("{} microseconds", event.time_us)); tracing::info!(time_us = event.time_us, timestamp = %datetime, "received first event from jetstream"); }
if event.kind != "commit" { continue; }
let event_value = serde_json::to_value(&event); if let Err(err) = event_value { tracing::error!(error = ?err, "error processing jetstream message"); continue; } let event_value = event_value.unwrap();
// Assumption: Performing a query for each event will cost more in the // long-term than evaluating each event against all matchers and if there's a // match, then checking both the event DID and the AT-URI DID. 'matchers_loop: for feed_matcher in self.feed_matchers.0.iter() { if let Some(Match(op, aturi)) = feed_matcher.matches(&event_value) { tracing::debug!(feed_id = ?feed_matcher.feed, "matched event");
let aturi_did = did_from_aturi(&aturi); if denylist.contains(event.did.as_str()) || denylist.contains(aturi_did.as_str()) { break 'matchers_loop; }
let feed_content = storage::model::FeedContent{ feed_id: feed_matcher.feed.clone(), uri: aturi, indexed_at: event.time_us, score: 1, }; self.write_tx .send(WriteCommand::Content(feed_content, op)) .await .map_err(|_| anyhow!("write channel closed"))?; } } } } }
// Persist the cursor on the way out (including on a jetstream // disconnect) so progress survives a restart. The writer flushes any // remaining buffered writes before persisting it. let _ = self .write_tx .send(WriteCommand::Checkpoint(time_usec)) .await;
tracing::debug!("ConsumerTask stopped");
Ok(()) }}
fn did_from_aturi(aturi: &str) -> String { let aturi_len = aturi.len(); if aturi_len < 6 { return "".to_string(); } let collection_start = aturi[5..] .find("/") .map(|value| value + 5) .unwrap_or(aturi_len); aturi[5..collection_start].to_string()}
pub(crate) mod model {
use std::collections::HashMap;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", content = "payload")] pub(crate) enum SubscriberSourcedMessage { #[serde(rename = "options_update")] Update { #[serde(rename = "wantedCollections")] wanted_collections: Vec<String>,
#[serde(rename = "wantedDids", skip_serializing_if = "Vec::is_empty", default)] wanted_dids: Vec<String>,
#[serde(rename = "maxMessageSizeBytes")] max_message_size_bytes: u64,
#[serde(skip_serializing_if = "Option::is_none")] cursor: Option<i64>, }, }
#[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct Facet { pub(crate) features: Vec<HashMap<String, String>>, }
#[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct StrongRef { pub(crate) uri: String, }
#[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct Reply { pub(crate) root: Option<StrongRef>, pub(crate) parent: Option<StrongRef>, }
#[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "$type")] pub(crate) enum Record { #[serde(rename = "app.bsky.feed.post")] Post { #[serde(flatten)] extra: HashMap<String, serde_json::Value>, }, #[serde(rename = "app.bsky.feed.like")] Like { #[serde(flatten)] extra: HashMap<String, serde_json::Value>, },
#[serde(untagged)] Other { #[serde(flatten)] extra: HashMap<String, serde_json::Value>, }, }
#[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "operation")] pub(crate) enum CommitOp { #[serde(rename = "create")] Create { rev: String, collection: String, rkey: String, record: Record, cid: String, }, #[serde(rename = "update")] Update { rev: String, collection: String, rkey: String, record: Record, cid: String, }, #[serde(rename = "delete")] Delete { rev: String, collection: String, rkey: String, }, }
#[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct Event { pub(crate) did: String, pub(crate) kind: String, pub(crate) time_us: i64, pub(crate) commit: Option<CommitOp>, }}