diff --git a/src/admin.rs b/src/admin.rs index 51637c2..7296154 100644 --- a/src/admin.rs +++ b/src/admin.rs @@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::error::AppError; -use crate::lexicon::ParsedLexicon; +use crate::lexicon::{LexiconType, ParsedLexicon}; use crate::AppState; // --------------------------------------------------------------------------- @@ -58,6 +58,17 @@ impl axum::extract::FromRequestParts for AdminAuth { } } +// --------------------------------------------------------------------------- +// Jetstream notification +// --------------------------------------------------------------------------- + +/// Send the current record collection list to the Jetstream task so it +/// reconnects with the updated filter. +async fn notify_jetstream(state: &AppState) { + let collections = state.lexicons.get_record_collections().await; + let _ = state.collections_tx.send(collections); +} + // --------------------------------------------------------------------------- // Request / response types // --------------------------------------------------------------------------- @@ -154,8 +165,13 @@ async fn upload_lexicon( // Update in-memory registry with correct revision let parsed = ParsedLexicon::parse(body.lexicon_json, revision) .map_err(|e| AppError::Internal(format!("failed to re-parse lexicon: {e}")))?; + let is_record = parsed.lexicon_type == LexiconType::Record; state.lexicons.upsert(parsed).await; + if is_record { + notify_jetstream(&state).await; + } + let status = if revision == 1 { StatusCode::CREATED } else { @@ -250,6 +266,7 @@ async fn delete_lexicon( } state.lexicons.remove(&id).await; + notify_jetstream(&state).await; Ok(StatusCode::NO_CONTENT) } diff --git a/src/jetstream.rs b/src/jetstream.rs index 3ef6552..131ed91 100644 --- a/src/jetstream.rs +++ b/src/jetstream.rs @@ -4,6 +4,7 @@ use serde_json::Value; use sqlx::PgPool; use std::sync::atomic::{AtomicI64, Ordering}; use std::sync::Arc; +use tokio::sync::watch; use tokio_tungstenite::tungstenite::Message; // --------------------------------------------------------------------------- @@ -32,19 +33,40 @@ struct JetstreamCommit { // --------------------------------------------------------------------------- /// Spawn a background task that subscribes to the Jetstream firehose and -/// indexes `games.gamesgamesgamesgames.*` records into PostgreSQL. -pub fn spawn(db: PgPool, jetstream_url: String) { +/// indexes records for collections specified by the watch channel. +/// +/// When the collection list is empty, the task idles without connecting. +/// When collections change, it disconnects and reconnects with the new filter. +pub fn spawn(db: PgPool, jetstream_url: String, mut collections_rx: watch::Receiver>) { tokio::spawn(async move { let cursor: Arc = Arc::new(AtomicI64::new(0)); loop { - if let Err(e) = run(&db, &jetstream_url, &cursor).await { - tracing::warn!("jetstream disconnected: {e}"); + // Wait until we have at least one collection to subscribe to. + let collections = collections_rx.borrow_and_update().clone(); + if collections.is_empty() { + tracing::info!("no collections configured, jetstream idle"); + // Block until the collection list changes. + if collections_rx.changed().await.is_err() { + // Sender dropped — shut down. + tracing::info!("jetstream watch channel closed, shutting down"); + return; + } + continue; } - // Back off before reconnecting. - tokio::time::sleep(std::time::Duration::from_secs(2)).await; - tracing::info!("reconnecting to jetstream..."); + // Connect and process events. If the collection list changes + // mid-stream, `run` returns so we can reconnect with new filters. + match run(&db, &jetstream_url, &cursor, &collections, &mut collections_rx).await { + Ok(()) => { + tracing::info!("jetstream reconnecting due to collection change"); + } + Err(e) => { + tracing::warn!("jetstream disconnected: {e}"); + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + tracing::info!("reconnecting to jetstream..."); + } + } } }); } @@ -57,11 +79,16 @@ async fn run( db: &PgPool, jetstream_url: &str, cursor: &Arc, + collections: &[String], + collections_rx: &mut watch::Receiver>, ) -> Result<(), Box> { - let mut url = format!( - "{}?wantedCollections=games.gamesgamesgamesgames.*", - jetstream_url, - ); + let wanted: String = collections + .iter() + .map(|c| format!("wantedCollections={c}")) + .collect::>() + .join("&"); + + let mut url = format!("{jetstream_url}?{wanted}"); let last = cursor.load(Ordering::Relaxed); if last > 0 { @@ -71,85 +98,100 @@ async fn run( tracing::info!(cursor = rewound, "resuming jetstream with cursor"); } + tracing::info!(collections = ?collections, "connecting to jetstream"); + let (ws, _) = tokio_tungstenite::connect_async(&url).await?; tracing::info!("connected to jetstream"); let (_, mut read) = ws.split(); - while let Some(msg) = read.next().await { - let msg = msg?; - - let text = match msg { - Message::Text(t) => t, - Message::Close(_) => break, - _ => continue, - }; - - let event: JetstreamEvent = match serde_json::from_str(&text) { - Ok(e) => e, - Err(e) => { - tracing::debug!("skipping unparseable event: {e}"); - continue; - } - }; + loop { + tokio::select! { + msg = read.next() => { + let msg = match msg { + Some(Ok(m)) => m, + Some(Err(e)) => return Err(e.into()), + None => break, + }; - // Update cursor. - cursor.store(event.time_us, Ordering::Relaxed); + let text = match msg { + Message::Text(t) => t, + Message::Close(_) => break, + _ => continue, + }; - if event.kind != "commit" { - continue; - } + let event: JetstreamEvent = match serde_json::from_str(&text) { + Ok(e) => e, + Err(e) => { + tracing::debug!("skipping unparseable event: {e}"); + continue; + } + }; - let commit = match event.commit { - Some(c) => c, - None => continue, - }; + // Update cursor. + cursor.store(event.time_us, Ordering::Relaxed); - let uri = format!( - "at://{}/{}/{}", - event.did, commit.collection, commit.rkey, - ); + if event.kind != "commit" { + continue; + } - match commit.operation.as_str() { - "create" | "update" => { - let record = match commit.record { - Some(r) => r, + let commit = match event.commit { + Some(c) => c, None => continue, }; - let cid = commit.cid.unwrap_or_default(); - - if let Err(e) = sqlx::query( - r#" - INSERT INTO records (uri, did, collection, rkey, record, cid, indexed_at) - VALUES ($1, $2, $3, $4, $5, $6, NOW()) - ON CONFLICT (uri) DO UPDATE - SET record = EXCLUDED.record, - cid = EXCLUDED.cid, - indexed_at = NOW() - "#, - ) - .bind(&uri) - .bind(&event.did) - .bind(&commit.collection) - .bind(&commit.rkey) - .bind(&record) - .bind(&cid) - .execute(db) - .await - { - tracing::warn!(uri = %uri, "failed to upsert record: {e}"); + + let uri = format!( + "at://{}/{}/{}", + event.did, commit.collection, commit.rkey, + ); + + match commit.operation.as_str() { + "create" | "update" => { + let record = match commit.record { + Some(r) => r, + None => continue, + }; + let cid = commit.cid.unwrap_or_default(); + + if let Err(e) = sqlx::query( + r#" + INSERT INTO records (uri, did, collection, rkey, record, cid, indexed_at) + VALUES ($1, $2, $3, $4, $5, $6, NOW()) + ON CONFLICT (uri) DO UPDATE + SET record = EXCLUDED.record, + cid = EXCLUDED.cid, + indexed_at = NOW() + "#, + ) + .bind(&uri) + .bind(&event.did) + .bind(&commit.collection) + .bind(&commit.rkey) + .bind(&record) + .bind(&cid) + .execute(db) + .await + { + tracing::warn!(uri = %uri, "failed to upsert record: {e}"); + } + } + "delete" => { + if let Err(e) = sqlx::query("DELETE FROM records WHERE uri = $1") + .bind(&uri) + .execute(db) + .await + { + tracing::warn!(uri = %uri, "failed to delete record: {e}"); + } + } + _ => {} } } - "delete" => { - if let Err(e) = sqlx::query("DELETE FROM records WHERE uri = $1") - .bind(&uri) - .execute(db) - .await - { - tracing::warn!(uri = %uri, "failed to delete record: {e}"); - } + // If the collection list changes, break out to reconnect. + _ = collections_rx.changed() => { + tracing::info!("collection filter changed, will reconnect"); + return Ok(()); } - _ => {} } } diff --git a/src/main.rs b/src/main.rs index ad91323..4217a30 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,6 +10,7 @@ mod server; use config::Config; use lexicon::LexiconRegistry; +use tokio::sync::watch; use tracing::info; #[derive(Clone)] @@ -18,6 +19,7 @@ pub struct AppState { pub http: reqwest::Client, pub db: sqlx::PgPool, pub lexicons: LexiconRegistry, + pub collections_tx: watch::Sender>, } #[tokio::main] @@ -51,14 +53,18 @@ async fn main() { .await .expect("failed to load lexicons"); + let initial_collections = lexicons.get_record_collections().await; + let (collections_tx, collections_rx) = watch::channel(initial_collections); + let state = AppState { config: config.clone(), http: reqwest::Client::new(), db, lexicons, + collections_tx, }; - jetstream::spawn(state.db.clone(), config.jetstream_url.clone()); + jetstream::spawn(state.db.clone(), config.jetstream_url.clone(), collections_rx); let app = server::router(state); let addr = config.listen_addr();