chadtmiller.com/knotless

//! Keeps mirrors warm and the index fed from Jetstream v2
//! (`network.bsky.jetstream.subscribeEvents`), the firehose re-encoded as
//! JSON and filtered server side. Asking only for `dev.knotless.*` commits,
//! a server receives its own collections' writes and every account's
//! identity, account and sync events, with each record decoded inline, so
//! indexing it needs no `getRecord` round trip.
//!
//! Jetstream frames are not signed: the record is what the Jetstream
//! operator says the commit held. That is the same trust a server places in
//! the relay it follows, and the index is a cache either way.

use crate::App;
use crate::firehose::{Event, Op, knotless};
use futures::StreamExt;
use serde_json::Value;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;

/// The collections asked for: every knotless and forge NSID.
pub const COLLECTIONS: &str = "dev.knotless.*";

/// The one framing the endpoint speaks.
const SUBPROTOCOL: &str = "xrpc.v1.json";

/// The payload `$type`s, after this prefix, name the event kind.
const KIND: &str = "network.bsky.jetstream.subscribeEvents#";

/// The subscribeEvents endpoint for a Jetstream host, from `cursor`, which
/// it replays inclusively.
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/network.bsky.jetstream.subscribeEvents");
    }
    url.push_str(&format!(
        "?collections={COLLECTIONS}&kinds=commit&kinds=identity&kinds=account&kinds=sync"
    ));
    if let Some(c) = cursor {
        url.push_str(&format!("&cursor={c}"));
    }
    url
}

/// Reads a text frame. `None` when it is not a message: an error frame, or
/// not Jetstream at all.
pub fn parse(frame: &str) -> Option<Event> {
    let mut frame: Value = serde_json::from_str(frame).ok()?;
    if frame["$type"] != "message" {
        return None;
    }
    let mut payload = frame["payload"].take();
    let kind = payload["$type"].as_str()?.strip_prefix(KIND)?.to_string();
    let mut event = Event {
        seq: payload["seq"].as_i64(),
        ..Event::default()
    };
    match kind.as_str() {
        "commit" => {
            let collection = text(&payload, "collection")?;
            if knotless(&collection) {
                let record = payload.get_mut("record").map(Value::take);
                event.ops.push(Op {
                    action: text(&payload, "operation")?,
                    collection,
                    rkey: text(&payload, "rkey")?,
                    record,
                });
                event.did = Some(text(&payload, "did")?);
            }
        }
        "identity" | "account" | "sync" => event.did = Some(text(&payload, "did")?),
        _ => {}
    }
    Some(event)
}

fn text(v: &Value, name: &str) -> Option<String> {
    v[name].as_str().map(str::to_string)
}

/// Follows Jetstream until the process ends, reconnecting after the last
/// sequence number seen. 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 request) = endpoint(&base, cursor).into_client_request() {
            request.headers_mut().insert(
                "Sec-WebSocket-Protocol",
                HeaderValue::from_static(SUBPROTOCOL),
            );
            if let Ok((mut ws, _)) = tokio_tungstenite::connect_async(request).await {
                app.watching.store(true, Ordering::Relaxed);
                while let Some(Ok(msg)) = ws.next().await {
                    if let Message::Text(frame) = msg
                        && let Some(event) = parse(&frame)
                    {
                        // The cursor is inclusive: resume after the last one.
                        cursor = event.seq.map(|s| s + 1).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 knotless::records::nsid;

    const TIME: &str = "1970-01-01T00:00:00.000001Z";

    fn commit(seq: i64, op: &str, collection: &str, record: &str) -> String {
        format!(
            r#"{{"$type":"message","payload":{{"$type":"{KIND}commit","seq":{seq},"did":"did:plc:a","time":"{TIME}","rev":"r","operation":"{op}","collection":"{collection}","rkey":"k"{record}}}}}"#
        )
    }

    #[test]
    fn the_filter_covers_every_collection() {
        let prefix = COLLECTIONS.trim_end_matches('*');
        for collection in nsid::ALL.iter().chain(nsid::FORGE.iter()) {
            assert!(collection.starts_with(prefix), "{collection}");
        }
    }

    #[test]
    fn frames() {
        let created = parse(&commit(
            42,
            "create",
            "dev.knotless.issue",
            r#","cid":"bafy","record":{"$type":"dev.knotless.issue","title":"hi"}"#,
        ))
        .unwrap();
        assert_eq!(created.seq, Some(42));
        assert_eq!(created.did.as_deref(), Some("did:plc:a"));
        assert_eq!(
            created.ops,
            [Op {
                action: "create".into(),
                collection: "dev.knotless.issue".into(),
                rkey: "k".into(),
                record: Some(serde_json::json!({"$type": "dev.knotless.issue", "title": "hi"})),
            }]
        );
        let deleted = parse(&commit(43, "delete", "dev.knotless.ref", "")).unwrap();
        assert_eq!(deleted.ops[0].record, None);
        let elsewhere = parse(&commit(44, "create", "app.bsky.feed.post", "")).unwrap();
        assert_eq!(
            (elsewhere.seq, elsewhere.did, elsewhere.ops.len()),
            (Some(44), None, 0)
        );

        let identity = format!(
            r#"{{"$type":"message","payload":{{"$type":"{KIND}identity","seq":5,"did":"did:plc:b","time":"{TIME}","identity":{{"did":"did:plc:b","handle":"b.test","seq":9,"time":"{TIME}"}}}}}}"#
        );
        let identity = parse(&identity).unwrap();
        assert_eq!(
            (identity.seq, identity.did.as_deref()),
            (Some(5), Some("did:plc:b"))
        );
        let info = format!(
            r#"{{"$type":"message","payload":{{"$type":"{KIND}info","name":"OutdatedCursor","message":"starting at seq 200"}}}}"#
        );
        assert_eq!(parse(&info), Some(Event::default()));

        assert_eq!(
            parse(r#"{"$type":"error","error":"ConsumerTooSlow"}"#),
            None
        );
        assert_eq!(
            parse(r#"{"$type":"message","payload":{"$type":"app.bsky.x"}}"#),
            None
        );
        assert_eq!(parse(r#"{"$type":"message","payload":{}}"#), None);
        assert_eq!(parse("not json"), None);
        let no_rkey = format!(
            r#"{{"$type":"message","payload":{{"$type":"{KIND}commit","seq":1,"did":"did:plc:a","operation":"create","collection":"dev.knotless.repo"}}}}"#
        );
        assert_eq!(parse(&no_rkey), None);
    }

    #[test]
    fn endpoints() {
        let filter =
            "collections=dev.knotless.*&kinds=commit&kinds=identity&kinds=account&kinds=sync";
        assert_eq!(
            endpoint("wss://jetstream.us-east.bsky.network/", None),
            format!(
                "wss://jetstream.us-east.bsky.network/xrpc/network.bsky.jetstream.subscribeEvents?{filter}"
            )
        );
        assert_eq!(
            endpoint(
                "ws://h/xrpc/network.bsky.jetstream.subscribeEvents",
                Some(3)
            ),
            format!("ws://h/xrpc/network.bsky.jetstream.subscribeEvents?{filter}&cursor=3")
        );
    }
}