//! Keeps mirrors warm from `com.atproto.sync.subscribeRepos`: a commit that //! touches a knotless collection marks its account's mirrors stale and //! refreshes them, so the next clone needs no round trip to the PDS. use crate::App; use futures::StreamExt; use knotless::records::nsid; use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::Duration; use tokio_tungstenite::tungstenite::Message; /// What a firehose frame says happened. #[derive(Debug, Default, PartialEq)] pub struct Event { pub seq: Option, /// The account, when it wrote to a knotless collection or its identity /// or repo changed under it. pub did: Option, /// Writes to the knotless collections. pub ops: Vec, } #[derive(Debug, PartialEq)] pub struct Op { /// `create`, `update` or `delete`. pub action: String, pub collection: String, pub rkey: String, /// The record as written, when the source carries it decoded (Jetstream /// does; the firehose's is CBOR in a CAR, so it is fetched instead). pub record: Option, } pub(crate) fn knotless(collection: &str) -> bool { nsid::ALL.contains(&collection) || nsid::FORGE.contains(&collection) } /// Reads a frame. `None` when it is not one. pub fn parse(frame: &[u8]) -> Option { let mut reader = frame; let header: ciborium::Value = ciborium::from_reader(&mut reader).ok()?; let body: ciborium::Value = ciborium::from_reader(&mut reader).ok()?; let field = |v: &ciborium::Value, name: &str| -> Option { v.as_map()? .iter() .find(|(k, _)| k.as_text() == Some(name)) .map(|(_, v)| v.clone()) }; let text = |v: &ciborium::Value, name: &str| field(v, name)?.as_text().map(str::to_string); let kind = text(&header, "t")?; let mut event = Event { seq: field(&body, "seq") .and_then(|s| s.as_integer()) .and_then(|i| i64::try_from(i).ok()), ..Event::default() }; match kind.as_str() { "#commit" => { for op in field(&body, "ops")?.as_array()? { let (Some(action), Some(path)) = (text(op, "action"), text(op, "path")) else { continue; }; let Some((collection, rkey)) = path.split_once('/') else { continue; }; if knotless(collection) { event.ops.push(Op { action, collection: collection.into(), rkey: rkey.into(), record: None, }); } } if !event.ops.is_empty() { event.did = Some(text(&body, "repo")?); } } "#sync" | "#account" | "#identity" => event.did = Some(text(&body, "did")?), _ => {} } Some(event) } /// The subscribeRepos endpoint for a relay or PDS base URL. pub fn endpoint(base: &str, cursor: Option) -> String { let mut url = base.trim_end_matches('/').to_string(); if !url.contains("/xrpc/") { url.push_str("/xrpc/com.atproto.sync.subscribeRepos"); } if let Some(c) = cursor { url.push_str(&format!("?cursor={c}")); } url } /// Follows the firehose until the process ends, reconnecting from the last /// sequence number. While disconnected, mirrors fall back to the TTL. pub async fn watch(app: Arc, base: String) { let mut cursor = None; loop { if let Ok((mut ws, _)) = tokio_tungstenite::connect_async(endpoint(&base, cursor)).await { app.watching.store(true, Ordering::Relaxed); while let Some(Ok(msg)) = ws.next().await { if let Message::Binary(frame) = msg && let Some(event) = parse(&frame) { cursor = event.seq.or(cursor); app.observe(event); } } } app.watching.store(false, Ordering::Relaxed); tokio::time::sleep(Duration::from_millis(500)).await; } } #[cfg(test)] mod tests { use super::*; use ciborium::Value as V; fn frame(header: V, body: V) -> Vec { let mut out = Vec::new(); ciborium::into_writer(&header, &mut out).unwrap(); ciborium::into_writer(&body, &mut out).unwrap(); out } fn map(pairs: Vec<(&str, V)>) -> V { V::Map( pairs .into_iter() .map(|(k, v)| (V::Text(k.into()), v)) .collect(), ) } #[test] fn frames() { let commit = |ops: Vec| { frame( map(vec![ ("op", V::Integer(1.into())), ("t", V::Text("#commit".into())), ]), map(vec![ ("seq", V::Integer(7.into())), ("repo", V::Text("did:plc:a".into())), ("ops", V::Array(ops)), ]), ) }; let op = |action: &str, path: &str| { map(vec![ ("action", V::Text(action.into())), ("path", V::Text(path.into())), ]) }; let event = parse(&commit(vec![ op("create", "dev.knotless.ref/k.x"), op("delete", "dev.knotless.issue/3"), op("create", "app.bsky.feed.post/x"), op("create", "no-slash"), map(vec![("path", V::Text("dev.knotless.ref/y".into()))]), ])) .unwrap(); assert_eq!(event.seq, Some(7)); assert_eq!(event.did.as_deref(), Some("did:plc:a")); assert_eq!( event.ops, [ Op { action: "create".into(), collection: "dev.knotless.ref".into(), rkey: "k.x".into(), record: None }, Op { action: "delete".into(), collection: "dev.knotless.issue".into(), rkey: "3".into(), record: None }, ] ); let elsewhere = parse(&commit(vec![op("create", "app.bsky.feed.post/x")])).unwrap(); assert_eq!((elsewhere.did, elsewhere.ops.len()), (None, 0)); let identity = frame( map(vec![("t", V::Text("#identity".into()))]), map(vec![("did", V::Text("did:plc:b".into()))]), ); assert_eq!(parse(&identity).unwrap().did.as_deref(), Some("did:plc:b")); let info = frame(map(vec![("t", V::Text("#info".into()))]), map(vec![])); assert_eq!(parse(&info), Some(Event::default())); assert_eq!(parse(b"\xff"), None); assert_eq!(parse(&frame(map(vec![]), map(vec![]))), None); } #[test] fn endpoints() { assert_eq!( endpoint("wss://relay/", None), "wss://relay/xrpc/com.atproto.sync.subscribeRepos" ); assert_eq!( endpoint("ws://h/xrpc/com.atproto.sync.subscribeRepos", Some(3)), "ws://h/xrpc/com.atproto.sync.subscribeRepos?cursor=3" ); } }