chadtmiller.com/knotless

//! 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<i64>,
    /// The account, when it wrote to a knotless collection or its identity
    /// or repo changed under it.
    pub did: Option<String>,
    /// Writes to the knotless collections.
    pub ops: Vec<Op>,
}

#[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<serde_json::Value>,
}

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<Event> {
    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<ciborium::Value> {
        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<i64>) -> 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<App>, 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<u8> {
        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<V>| {
            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"
        );
    }
}