feat: dynamic Jetstream subscription driven by lexicon registry · gamesgamesgamesgames.games/happyview@09bf19e

+18 -1

src/admin.rs

··· 6 6 use serde_json::Value; 7 7 8 8 use crate::error::AppError; 9 - use crate::lexicon::ParsedLexicon; 9 + use crate::lexicon::{LexiconType, ParsedLexicon}; 10 10 use crate::AppState; 11 11 12 12 // --------------------------------------------------------------------------- ··· 59 59 } 60 60 61 61 // --------------------------------------------------------------------------- 62 + // Jetstream notification 63 + // --------------------------------------------------------------------------- 64 + 65 + /// Send the current record collection list to the Jetstream task so it 66 + /// reconnects with the updated filter. 67 + async fn notify_jetstream(state: &AppState) { 68 + let collections = state.lexicons.get_record_collections().await; 69 + let _ = state.collections_tx.send(collections); 70 + } 71 + 72 + // --------------------------------------------------------------------------- 62 73 // Request / response types 63 74 // --------------------------------------------------------------------------- 64 75 ··· 154 165 // Update in-memory registry with correct revision 155 166 let parsed = ParsedLexicon::parse(body.lexicon_json, revision) 156 167 .map_err(|e| AppError::Internal(format!("failed to re-parse lexicon: {e}")))?; 168 + let is_record = parsed.lexicon_type == LexiconType::Record; 157 169 state.lexicons.upsert(parsed).await; 158 170 171 + if is_record { 172 + notify_jetstream(&state).await; 173 + } 174 + 159 175 let status = if revision == 1 { 160 176 StatusCode::CREATED 161 177 } else { ··· 250 266 } 251 267 252 268 state.lexicons.remove(&id).await; 269 + notify_jetstream(&state).await; 253 270 254 271 Ok(StatusCode::NO_CONTENT) 255 272 }

+116 -74

src/jetstream.rs

··· 4 4 use sqlx::PgPool; 5 5 use std::sync::atomic::{AtomicI64, Ordering}; 6 6 use std::sync::Arc; 7 + use tokio::sync::watch; 7 8 use tokio_tungstenite::tungstenite::Message; 8 9 9 10 // --------------------------------------------------------------------------- ··· 32 33 // --------------------------------------------------------------------------- 33 34 34 35 /// Spawn a background task that subscribes to the Jetstream firehose and 35 - /// indexes `games.gamesgamesgamesgames.*` records into PostgreSQL. 36 - pub fn spawn(db: PgPool, jetstream_url: String) { 36 + /// indexes records for collections specified by the watch channel. 37 + /// 38 + /// When the collection list is empty, the task idles without connecting. 39 + /// When collections change, it disconnects and reconnects with the new filter. 40 + pub fn spawn(db: PgPool, jetstream_url: String, mut collections_rx: watch::Receiver<Vec<String>>) { 37 41 tokio::spawn(async move { 38 42 let cursor: Arc<AtomicI64> = Arc::new(AtomicI64::new(0)); 39 43 40 44 loop { 41 - if let Err(e) = run(&db, &jetstream_url, &cursor).await { 42 - tracing::warn!("jetstream disconnected: {e}"); 45 + // Wait until we have at least one collection to subscribe to. 46 + let collections = collections_rx.borrow_and_update().clone(); 47 + if collections.is_empty() { 48 + tracing::info!("no collections configured, jetstream idle"); 49 + // Block until the collection list changes. 50 + if collections_rx.changed().await.is_err() { 51 + // Sender dropped — shut down. 52 + tracing::info!("jetstream watch channel closed, shutting down"); 53 + return; 54 + } 55 + continue; 43 56 } 44 57 45 - // Back off before reconnecting. 46 - tokio::time::sleep(std::time::Duration::from_secs(2)).await; 47 - tracing::info!("reconnecting to jetstream..."); 58 + // Connect and process events. If the collection list changes 59 + // mid-stream, `run` returns so we can reconnect with new filters. 60 + match run(&db, &jetstream_url, &cursor, &collections, &mut collections_rx).await { 61 + Ok(()) => { 62 + tracing::info!("jetstream reconnecting due to collection change"); 63 + } 64 + Err(e) => { 65 + tracing::warn!("jetstream disconnected: {e}"); 66 + tokio::time::sleep(std::time::Duration::from_secs(2)).await; 67 + tracing::info!("reconnecting to jetstream..."); 68 + } 69 + } 48 70 } 49 71 }); 50 72 } ··· 57 79 db: &PgPool, 58 80 jetstream_url: &str, 59 81 cursor: &Arc<AtomicI64>, 82 + collections: &[String], 83 + collections_rx: &mut watch::Receiver<Vec<String>>, 60 84 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> { 61 - let mut url = format!( 62 - "{}?wantedCollections=games.gamesgamesgamesgames.*", 63 - jetstream_url, 64 - ); 85 + let wanted: String = collections 86 + .iter() 87 + .map(|c| format!("wantedCollections={c}")) 88 + .collect::<Vec<_>>() 89 + .join("&"); 90 + 91 + let mut url = format!("{jetstream_url}?{wanted}"); 65 92 66 93 let last = cursor.load(Ordering::Relaxed); 67 94 if last > 0 { ··· 71 98 tracing::info!(cursor = rewound, "resuming jetstream with cursor"); 72 99 } 73 100 101 + tracing::info!(collections = ?collections, "connecting to jetstream"); 102 + 74 103 let (ws, _) = tokio_tungstenite::connect_async(&url).await?; 75 104 tracing::info!("connected to jetstream"); 76 105 77 106 let (_, mut read) = ws.split(); 78 107 79 - while let Some(msg) = read.next().await { 80 - let msg = msg?; 108 + loop { 109 + tokio::select! { 110 + msg = read.next() => { 111 + let msg = match msg { 112 + Some(Ok(m)) => m, 113 + Some(Err(e)) => return Err(e.into()), 114 + None => break, 115 + }; 81 116 82 - let text = match msg { 83 - Message::Text(t) => t, 84 - Message::Close(_) => break, 85 - _ => continue, 86 - }; 117 + let text = match msg { 118 + Message::Text(t) => t, 119 + Message::Close(_) => break, 120 + _ => continue, 121 + }; 87 122 88 - let event: JetstreamEvent = match serde_json::from_str(&text) { 89 - Ok(e) => e, 90 - Err(e) => { 91 - tracing::debug!("skipping unparseable event: {e}"); 92 - continue; 93 - } 94 - }; 123 + let event: JetstreamEvent = match serde_json::from_str(&text) { 124 + Ok(e) => e, 125 + Err(e) => { 126 + tracing::debug!("skipping unparseable event: {e}"); 127 + continue; 128 + } 129 + }; 95 130 96 - // Update cursor. 97 - cursor.store(event.time_us, Ordering::Relaxed); 98 - 99 - if event.kind != "commit" { 100 - continue; 101 - } 102 - 103 - let commit = match event.commit { 104 - Some(c) => c, 105 - None => continue, 106 - }; 131 + // Update cursor. 132 + cursor.store(event.time_us, Ordering::Relaxed); 107 133 108 - let uri = format!( 109 - "at://{}/{}/{}", 110 - event.did, commit.collection, commit.rkey, 111 - ); 134 + if event.kind != "commit" { 135 + continue; 136 + } 112 137 113 - match commit.operation.as_str() { 114 - "create" | "update" => { 115 - let record = match commit.record { 116 - Some(r) => r, 138 + let commit = match event.commit { 139 + Some(c) => c, 117 140 None => continue, 118 141 }; 119 - let cid = commit.cid.unwrap_or_default(); 120 142 121 - if let Err(e) = sqlx::query( 122 - r#" 123 - INSERT INTO records (uri, did, collection, rkey, record, cid, indexed_at) 124 - VALUES ($1, $2, $3, $4, $5, $6, NOW()) 125 - ON CONFLICT (uri) DO UPDATE 126 - SET record = EXCLUDED.record, 127 - cid = EXCLUDED.cid, 128 - indexed_at = NOW() 129 - "#, 130 - ) 131 - .bind(&uri) 132 - .bind(&event.did) 133 - .bind(&commit.collection) 134 - .bind(&commit.rkey) 135 - .bind(&record) 136 - .bind(&cid) 137 - .execute(db) 138 - .await 139 - { 140 - tracing::warn!(uri = %uri, "failed to upsert record: {e}"); 143 + let uri = format!( 144 + "at://{}/{}/{}", 145 + event.did, commit.collection, commit.rkey, 146 + ); 147 + 148 + match commit.operation.as_str() { 149 + "create" | "update" => { 150 + let record = match commit.record { 151 + Some(r) => r, 152 + None => continue, 153 + }; 154 + let cid = commit.cid.unwrap_or_default(); 155 + 156 + if let Err(e) = sqlx::query( 157 + r#" 158 + INSERT INTO records (uri, did, collection, rkey, record, cid, indexed_at) 159 + VALUES ($1, $2, $3, $4, $5, $6, NOW()) 160 + ON CONFLICT (uri) DO UPDATE 161 + SET record = EXCLUDED.record, 162 + cid = EXCLUDED.cid, 163 + indexed_at = NOW() 164 + "#, 165 + ) 166 + .bind(&uri) 167 + .bind(&event.did) 168 + .bind(&commit.collection) 169 + .bind(&commit.rkey) 170 + .bind(&record) 171 + .bind(&cid) 172 + .execute(db) 173 + .await 174 + { 175 + tracing::warn!(uri = %uri, "failed to upsert record: {e}"); 176 + } 177 + } 178 + "delete" => { 179 + if let Err(e) = sqlx::query("DELETE FROM records WHERE uri = $1") 180 + .bind(&uri) 181 + .execute(db) 182 + .await 183 + { 184 + tracing::warn!(uri = %uri, "failed to delete record: {e}"); 185 + } 186 + } 187 + _ => {} 141 188 } 142 189 } 143 - "delete" => { 144 - if let Err(e) = sqlx::query("DELETE FROM records WHERE uri = $1") 145 - .bind(&uri) 146 - .execute(db) 147 - .await 148 - { 149 - tracing::warn!(uri = %uri, "failed to delete record: {e}"); 150 - } 190 + // If the collection list changes, break out to reconnect. 191 + _ = collections_rx.changed() => { 192 + tracing::info!("collection filter changed, will reconnect"); 193 + return Ok(()); 151 194 } 152 - _ => {} 153 195 } 154 196 } 155 197

+7 -1

src/main.rs

··· 10 10 11 11 use config::Config; 12 12 use lexicon::LexiconRegistry; 13 + use tokio::sync::watch; 13 14 use tracing::info; 14 15 15 16 #[derive(Clone)] ··· 18 19 pub http: reqwest::Client, 19 20 pub db: sqlx::PgPool, 20 21 pub lexicons: LexiconRegistry, 22 + pub collections_tx: watch::Sender<Vec<String>>, 21 23 } 22 24 23 25 #[tokio::main] ··· 51 53 .await 52 54 .expect("failed to load lexicons"); 53 55 56 + let initial_collections = lexicons.get_record_collections().await; 57 + let (collections_tx, collections_rx) = watch::channel(initial_collections); 58 + 54 59 let state = AppState { 55 60 config: config.clone(), 56 61 http: reqwest::Client::new(), 57 62 db, 58 63 lexicons, 64 + collections_tx, 59 65 }; 60 66 61 - jetstream::spawn(state.db.clone(), config.jetstream_url.clone()); 67 + jetstream::spawn(state.db.clone(), config.jetstream_url.clone(), collections_rx); 62 68 63 69 let app = server::router(state); 64 70 let addr = config.listen_addr();

Configure Feed

Configure Feed