From 5f286f38e1c37173b715285b570e0df634959d14 Mon Sep 17 00:00:00 2001 From: dawn <90008@gaze.systems> Date: Wed, 27 May 2026 15:39:45 +0300 Subject: [PATCH] [jetstream] introduce jetstream compatible endpoint for indexer and relay --- Cargo.lock | 1 + Cargo.toml | 2 + docs/api/README.md | 1 + docs/api/jetstream.md | 110 +++++ docs/build-features.md | 1 + docs/configuration.md | 4 +- src/api/debug.rs | 10 + src/api/jetstream.rs | 315 +++++++++++++ src/api/mod.rs | 4 + src/control/jetstream.rs | 266 +++++++++++ src/control/mod.rs | 11 + src/control/stream.rs | 511 ++++++++++++++++++++++ src/db/ephemeral.rs | 42 ++ src/db/keys/mod.rs | 26 ++ src/db/mod.rs | 74 +++- src/ingest/indexer.rs | 43 ++ src/ingest/relay.rs | 89 +++- src/jetstream.rs | 39 ++ src/lib.rs | 7 + src/ops.rs | 19 + src/types.rs | 154 ++++++- tests/common.nu | 10 + tests/run_all.nu | 38 +- tests/stream_jetstream_subscribe.nu | 261 +++++++++++ tests/stream_relay_jetstream_subscribe.nu | 212 +++++++++ 25 files changed, 2230 insertions(+), 20 deletions(-) create mode 100644 docs/api/jetstream.md create mode 100644 src/api/jetstream.rs create mode 100644 src/control/jetstream.rs create mode 100644 src/jetstream.rs create mode 100644 tests/stream_jetstream_subscribe.nu create mode 100644 tests/stream_relay_jetstream_subscribe.nu diff --git a/Cargo.lock b/Cargo.lock index 63481b1..601e029 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1543,6 +1543,7 @@ dependencies = [ "tracing", "tracing-subscriber", "url", + "zstd", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 99089ae..89ab044 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,7 @@ backlinks = [] relay = [] indexer = [] indexer_stream = ["indexer"] +jetstream = ["dep:zstd"] [dependencies] tokio = { version = "1.0", features = ["full"] } @@ -29,6 +30,7 @@ fjall = { git = "https://github.com/90-008/fjall.git", rev = "aee642660079ec0a0f lsm-tree = { git = "https://github.com/90-008/lsm-tree.git", features = ["lz4", "zstd"] } serde_ipld_dagcbor = "0.6" serde_urlencoded = "0.7" +zstd = { version = "0.13", optional = true } bytes = { version = "1.11.1", features = ["serde"] } url = { version = "2.5", features = ["serde"] } diff --git a/docs/api/README.md b/docs/api/README.md index 42aac59..5b818d7 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -7,6 +7,7 @@ hydrant's REST API is split into public endpoints (safe to expose) and managemen ## public - `GET /stream`: subscribe to the event stream. query params: `cursor` (optional, start from a specific event ID). slow consumers may receive a `{"type":"error","error":"ConsumerTooSlow",...}` frame before the connection closes. +- `GET /subscribe`: [jetstream API](jetstream.md), subscribe to a filtered, jetstream-compatible websocket event stream. - `GET /stats`: get stats about the database (counts of repos, records, events; sizes of keyspaces on disk). - `GET /health` / `GET /_health`: health check. diff --git a/docs/api/jetstream.md b/docs/api/jetstream.md new file mode 100644 index 0000000..ac5dc1e --- /dev/null +++ b/docs/api/jetstream.md @@ -0,0 +1,110 @@ +--- +title: jetstream stream +--- + +hydrant implements a jetstream-compatible websocket subscription endpoint when compiled with the `jetstream` feature. it allows clients to consume atproto repository updates (commits, identity updates, and account state transitions) via a simplified, JSON-based format. + +## GET /subscribe + +subscribe to the jetstream websocket stream. + +### query parameters + +| param | type | description | +| :--- | :--- | :--- | +| `wantedCollections` | string \| seq | list of collection NSIDs to receive (e.g. `app.bsky.feed.post`). supports namespace wildcards (e.g. `app.bsky.feed.*`). | +| `wantedDids` | string \| seq | list of DIDs to receive (e.g. `did:plc:abc123xyz`). | +| `maxMessageSizeBytes` | integer | filters out events whose serialized JSON size exceeds this value. | +| `cursor` | integer | unix microseconds timestamp (`time_us`) to replay historical events from. | +| `compress` | boolean | if `true` (or if header `Socket-Encoding` contains `zstd`), compresses frames using zstd and sends them as binary websocket frames. | +| `requireHello` | boolean | if `true`, the socket will not stream any events until the client sends a `hello` / `options_update` message first. | + +### in-stream options update + +clients can dynamically modify filtering criteria (`wantedCollections`, `wantedDids`, and `maxMessageSizeBytes`) without reconnecting by sending a text frame with the following JSON format: + +```json +{ + "type": "options_update", + "payload": { + "wantedCollections": ["app.bsky.feed.post", "app.bsky.like.*"], + "wantedDids": ["did:plc:abc123xyz"], + "maxMessageSizeBytes": 5000000 + } +} +``` + +### event payload formats + +all jetstream events are serialized as JSON objects containing: +* `did`: string (did of the repository) +* `time_us`: integer (unix microseconds timestamp) +* `kind`: string (`"commit"`, `"identity"`, or `"account"`) + +#### commit + +fired when a record is created, updated, or deleted: + +```json +{ + "did": "did:plc:abc123xyz", + "time_us": 1716823456789012, + "kind": "commit", + "commit": { + "rev": "3kpjxabc123", + "operation": "create", + "collection": "app.bsky.feed.post", + "rkey": "3kpjxabc123", + "record": { + "$type": "app.bsky.feed.post", + "text": "hello, world!", + "createdAt": "2026-05-27T12:00:00.000Z" + }, + "cid": "bafyreihy..." + } +} +``` + +if hydrant is running in indexer mode with `HYDRANT_ONLY_INDEX_LINKS=true`, record content blocks are not persisted. consequently, `record` and `cid` may be omitted from `commit` payloads. + +#### identity + +fired when a did document or handle updates: + +```json +{ + "did": "did:plc:abc123xyz", + "time_us": 1716823456789012, + "kind": "identity", + "identity": { + "did": "did:plc:abc123xyz", + "seq": 1234567, + "time": "2026-05-27T12:00:00Z", + "handle": "user.bsky.social" + } +} +``` + +#### account + +fired when a repository's status changes (active, deactivated, or deleted): + +```json +{ + "did": "did:plc:abc123xyz", + "time_us": 1716823456789012, + "kind": "account", + "account": { + "active": false, + "did": "did:plc:abc123xyz", + "seq": 1234567, + "time": "2026-05-27T12:00:00Z", + "status": "deactivated" + } +} +``` + +### additional details + +- **live stream scope**: jetstream subscribers only receive live firehose events. historical backfill and sync replay events are processed with `live: false` internally, and are never staged in the jetstream keyspace or broadcasted. +- **slow consumers**: if a socket's send buffer remains full for longer than the configured timeout, the server sends a `{"type":"error","error":"ConsumerTooSlow"}` message and drops the connection. diff --git a/docs/build-features.md b/docs/build-features.md index 6b98aec..9a8786b 100644 --- a/docs/build-features.md +++ b/docs/build-features.md @@ -10,6 +10,7 @@ title: build features | `indexer_stream` | yes | enables the event stream for the indexer. requires indexer feature. | | `relay` | no | makes hydrant act as a relay. incompatible with the indexer feature. | | `backlinks` | no | enables the backlinks indexer and XRPC endpoints (`blue.microcosm.links.*`). requires indexer feature. | +| `jetstream` | no | enables the jetstream-compatible `GET /subscribe` websocket stream for indexer, ephemeral indexer, and relay builds. requires `indexer_stream` or `relay`. hydrant stores live-event metadata only; replay is bounded by normal event ttl in ephemeral/relay mode. | to build with a specific feature: diff --git a/docs/configuration.md b/docs/configuration.md index 0d5a717..bfc5d1b 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -20,8 +20,8 @@ hydrant is configured via environment variables, all prefixed with `HYDRANT_` (e | :--- | :--- | :--- | | `FULL_NETWORK` | `false` (indexer), `true` (relay) | if `true`, discover and index all repos in the network | | `EPHEMERAL` | `false` (indexer), `true` (relay) | if enabled, no records are stored (in indexer mode). events are deleted after a certain duration (`EPHEMERAL_TTL`) | -| `EPHEMERAL_TTL` | `60min`, `3d` (relay) | how long to keep events before deletion | -| `ONLY_INDEX_LINKS` | `false` | indexer only. if enabled, record blocks are not stored, only the index (records, counts, events) is kept. `getRecord`, `listRecords`, and `getRepo` will return errors. the event stream still works but create/update events will not include record values | +| `EPHEMERAL_TTL` | `60min`, `3d` (relay) | how long to keep events before deletion. when built with `jetstream`, retained Jetstream replay metadata is pruned on the same schedule | +| `ONLY_INDEX_LINKS` | `false` | indexer only. if enabled, record blocks are not stored, only the index (records, counts, events) is kept. `getRecord`, `listRecords`, and `getRepo` will return errors. the event stream and Jetstream stream still work, but create/update events will not include record values | ## filter diff --git a/src/api/debug.rs b/src/api/debug.rs index 2f29ada..34d0f6c 100644 --- a/src/api/debug.rs +++ b/src/api/debug.rs @@ -112,6 +112,12 @@ fn deserialize_value(partition: &str, value: &[u8]) -> Value { return serde_json::to_value(event).unwrap_or(Value::Null); } } + #[cfg(feature = "jetstream")] + "jetstream_events" => { + if let Ok(event) = rmp_serde::from_slice::(value) { + return serde_json::to_value(event).unwrap_or(Value::Null); + } + } "records" => { if let Ok(s) = String::from_utf8(value.to_vec()) { match Cid::from_str(&s) { @@ -288,6 +294,8 @@ fn get_keyspace_by_name(db: &crate::db::Db, name: &str) -> Result Ok(db.resync.clone()), #[cfg(feature = "indexer_stream")] "events" => Ok(db.events.clone()), + #[cfg(feature = "jetstream")] + "jetstream_events" => Ok(db.jetstream_events.clone()), #[cfg(feature = "relay")] "relay_events" => Ok(db.relay_events.clone()), #[cfg(feature = "indexer")] @@ -330,6 +338,8 @@ pub async fn handle_debug_ephemeral_ttl_tick( crate::db::ephemeral::ephemeral_ttl_tick(&state.db, &state.ephemeral_ttl)?; #[cfg(feature = "relay")] crate::db::ephemeral::relay_events_ttl_tick(&state.db, &state.ephemeral_ttl)?; + #[cfg(feature = "jetstream")] + crate::db::ephemeral::jetstream_events_ttl_tick(&state.db, &state.ephemeral_ttl)?; Ok(()) }) .await diff --git a/src/api/jetstream.rs b/src/api/jetstream.rs new file mode 100644 index 0000000..1e5b6a8 --- /dev/null +++ b/src/api/jetstream.rs @@ -0,0 +1,315 @@ +use std::sync::Arc; +use std::time::Duration; + +use axum::extract::{Query, State}; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use axum_tws::{Message, WebSocket, WebSocketUpgrade}; +use futures::{SinkExt, StreamExt}; +use serde::{Deserialize, Deserializer}; +use tokio::time::{MissedTickBehavior, interval, timeout}; +use tracing::{debug, warn}; + +use crate::control::{Hydrant, JetstreamFilter, JetstreamSubscriberOptions}; + +const PING_INTERVAL: Duration = Duration::from_secs(30); +const CLOSE_TIMEOUT: Duration = Duration::from_secs(1); +const MAX_SUBSCRIBER_MESSAGE_BYTES: usize = 10_000_000; + +// serde_urlencoded (used by axum's query extractor) cannot deserialize a single +// query param value into Vec — it expects repeated keys for sequences. +// this visitor accepts both a bare string (single value) and a proper sequence. +fn deserialize_string_or_seq<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + struct StringOrSeq; + + impl<'de> serde::de::Visitor<'de> for StringOrSeq { + type Value = Vec; + + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "a string or sequence of strings") + } + + fn visit_str(self, v: &str) -> Result, E> { + Ok(vec![v.to_owned()]) + } + + fn visit_seq>( + self, + mut seq: A, + ) -> Result, A::Error> { + let mut out = Vec::new(); + while let Some(v) = seq.next_element()? { + out.push(v); + } + Ok(out) + } + } + + deserializer.deserialize_any(StringOrSeq) +} + +#[derive(Deserialize)] +pub struct JetstreamQuery { + #[serde( + default, + rename = "wantedCollections", + deserialize_with = "deserialize_string_or_seq" + )] + wanted_collections: Vec, + #[serde( + default, + rename = "wantedDids", + deserialize_with = "deserialize_string_or_seq" + )] + wanted_dids: Vec, + #[serde(default, rename = "maxMessageSizeBytes", alias = "maxSize")] + max_message_size_bytes: Option, + cursor: Option, + #[serde(default)] + compress: bool, + #[serde(default, rename = "requireHello")] + require_hello: bool, +} + +pub async fn handle_subscribe( + State(hydrant): State, + Query(query): Query, + headers: HeaderMap, + ws: WebSocketUpgrade, +) -> Response { + let options = match options_from_parts( + &query.wanted_collections, + &query.wanted_dids, + parse_max_message_size(query.max_message_size_bytes), + ) { + Ok(options) => JetstreamFilter::new(options), + Err(err) => return (StatusCode::BAD_REQUEST, err).into_response(), + }; + + let socket_encoding = headers + .get("Socket-Encoding") + .and_then(|v| v.to_str().ok()) + .unwrap_or_default(); + let compress = query.compress || socket_encoding.contains("zstd"); + let cursor = query + .cursor + .filter(|cursor| *cursor <= chrono::Utc::now().timestamp_micros()); + + ws.on_upgrade(move |socket| { + handle_socket( + socket, + hydrant, + cursor, + options, + compress, + query.require_hello, + ) + }) + .into_response() +} + +async fn handle_socket( + socket: WebSocket, + hydrant: Hydrant, + cursor: Option, + options: JetstreamFilter, + compress: bool, + require_hello: bool, +) { + let send_timeout = hydrant.stream_send_timeout(); + let (mut sink, mut ws_recv) = socket.split(); + + if require_hello { + loop { + match ws_recv.next().await { + Some(Ok(m)) if m.is_ping() => { + if sink.send(Message::pong(m.into_payload())).await.is_err() { + return; + } + } + Some(Ok(m)) if m.is_close() => return, + Some(Ok(m)) if m.is_text() => match handle_options_message(m, &options) { + Ok(true) => break, + Ok(false) => return, + Err(_) => return, + }, + Some(Ok(_)) => {} + Some(Err(err)) => { + warn!(err = %err, "Jetstream ws recv error while waiting for hello"); + return; + } + None => return, + } + } + } + + let mut events = hydrant.subscribe_jetstream(cursor, options.clone()); + let mut ping_timer = interval(PING_INTERVAL); + ping_timer.set_missed_tick_behavior(MissedTickBehavior::Delay); + ping_timer.tick().await; + + loop { + tokio::select! { + inbound = ws_recv.next() => match inbound { + Some(Ok(m)) if m.is_close() => break, + Some(Ok(m)) if m.is_ping() => { + if let Err(err) = sink.send(Message::pong(m.into_payload())).await { + warn!(err = %err, "Jetstream ws pong send error"); + break; + } + } + Some(Ok(m)) if m.is_text() => { + if let Err(err) = handle_options_message(m, &options) { + debug!(err = %err, "invalid Jetstream options update"); + break; + } + } + Some(Ok(m)) if m.is_binary() => { + debug!("Jetstream client sent binary frame, closing"); + break; + } + Some(Ok(_)) => {} + Some(Err(err)) => { + warn!(err = %err, "Jetstream ws recv error"); + break; + } + None => break, + }, + evt = events.next() => match evt { + Some(Ok(json)) => { + let msg = match frame_for_event(&json, compress, &options) { + Ok(Some(msg)) => msg, + Ok(None) => continue, + Err(err) => { + warn!(err = %err, "failed to encode Jetstream frame"); + break; + } + }; + match timeout(send_timeout, sink.send(msg)).await { + Ok(Ok(())) => {} + Ok(Err(err)) => { + warn!(err = %err, "Jetstream ws send error"); + break; + } + Err(_) => { + let _ = timeout( + CLOSE_TIMEOUT, + sink.send(error_message("ConsumerTooSlow", &format!( + "Jetstream socket send blocked for at least {} seconds", + send_timeout.as_secs() + ))), + ) + .await; + break; + } + } + } + Some(Err(err)) => { + let _ = timeout(send_timeout, sink.send(error_message(err.code(), &err.to_string()))).await; + break; + } + None => break, + }, + _ = ping_timer.tick() => { + if let Err(err) = sink.send(Message::ping(bytes::Bytes::new())).await { + warn!(err = %err, "Jetstream ws ping send error"); + break; + } + } + } + } + + let _ = timeout(CLOSE_TIMEOUT, sink.close()).await; +} + +fn handle_options_message(msg: Message, options: &JetstreamFilter) -> Result { + let payload: bytes::Bytes = msg.into_payload().into(); + if payload.len() > MAX_SUBSCRIBER_MESSAGE_BYTES { + return Err("subscriber message too large".into()); + } + let msg: SubscriberSourcedMessage = + serde_json::from_slice(&payload).map_err(|e| e.to_string())?; + if msg.kind != "options_update" { + return Ok(false); + } + let payload: SubscriberOptionsUpdatePayload = + serde_json::from_value(msg.payload).map_err(|e| e.to_string())?; + let next = options_from_parts( + &payload.wanted_collections, + &payload.wanted_dids, + parse_max_message_size(Some(payload.max_message_size_bytes)), + )?; + options.update(next); + Ok(true) +} + +fn frame_for_event( + json: &[u8], + compress: bool, + options: &JetstreamFilter, +) -> Result, String> { + if compress { + let compressed = zstd::bulk::compress(json, 3).map_err(|e| e.to_string())?; + if exceeds_max(compressed.len(), options) { + return Ok(None); + } + return Ok(Some(Message::binary(compressed))); + } + + if exceeds_max(json.len(), options) { + return Ok(None); + } + let text = std::str::from_utf8(json).map_err(|e| e.to_string())?; + Ok(Some(Message::text(text.to_string()))) +} + +fn exceeds_max(size: usize, options: &JetstreamFilter) -> bool { + let max = options.max_message_size_bytes(); + max > 0 && size > max as usize +} + +fn options_from_parts( + wanted_collections: &[String], + wanted_dids: &[String], + max_message_size_bytes: u32, +) -> Result { + JetstreamSubscriberOptions::parse(wanted_collections, wanted_dids, max_message_size_bytes) +} + +fn parse_max_message_size(value: Option) -> u32 { + value + .filter(|v| *v > 0) + .and_then(|v| u32::try_from(v).ok()) + .unwrap_or(0) +} + +fn error_message(code: &str, message: &str) -> Message { + Message::text( + serde_json::json!({ + "type": "error", + "error": code, + "message": message, + }) + .to_string(), + ) +} + +#[derive(Deserialize)] +struct SubscriberSourcedMessage { + #[serde(rename = "type")] + kind: String, + payload: serde_json::Value, +} + +#[derive(Deserialize)] +struct SubscriberOptionsUpdatePayload { + #[serde(default, rename = "wantedCollections")] + wanted_collections: Vec, + #[serde(default, rename = "wantedDids")] + wanted_dids: Vec, + #[serde(default, rename = "maxMessageSizeBytes", alias = "maxSize")] + max_message_size_bytes: i64, +} diff --git a/src/api/mod.rs b/src/api/mod.rs index 976328f..a3a6b5d 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -14,6 +14,8 @@ mod debug; mod filter; mod firehose; mod ingestion; +#[cfg(feature = "jetstream")] +mod jetstream; mod pds; mod repos; mod stats; @@ -45,6 +47,8 @@ pub async fn serve(hydrant: Hydrant, binds: ApiBinds) -> miette::Result<()> { .route("/stats", get(stats::get_stats)); #[cfg(feature = "indexer_stream")] let app = app.nest("/stream", stream::router()); + #[cfg(feature = "jetstream")] + let app = app.route("/subscribe", get(jetstream::handle_subscribe)); let app = app .merge(xrpc::router(blocks_available)) .merge(filter::router()) diff --git a/src/control/jetstream.rs b/src/control/jetstream.rs new file mode 100644 index 0000000..913e7e4 --- /dev/null +++ b/src/control/jetstream.rs @@ -0,0 +1,266 @@ +use super::*; +use crate::db::types::TrimmedDid; +use crate::types::StoredJetstreamEvent; +use jacquard_common::IntoStatic; +use jacquard_common::types::nsid::Nsid; +use jacquard_common::types::string::Did; +use smol_str::{SmolStr, ToSmolStr}; +use std::collections::HashSet; +use std::sync::Arc; + +pub struct JetstreamEventStream(mpsc::Receiver>); + +#[derive(Debug, Clone, thiserror::Error)] +pub enum JetstreamStreamError { + #[error("jetstream consumer too slow: {reason}")] + ConsumerTooSlow { reason: String }, +} + +impl JetstreamStreamError { + pub fn code(&self) -> &'static str { + match self { + Self::ConsumerTooSlow { .. } => "ConsumerTooSlow", + } + } +} + +impl Stream for JetstreamEventStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.0.poll_recv(cx) + } +} + +#[derive(Clone, Default)] +pub struct JetstreamFilter { + inner: Arc>, +} + +impl JetstreamFilter { + pub fn new(options: JetstreamSubscriberOptions) -> Self { + Self { + inner: Arc::new(arc_swap::ArcSwap::from_pointee(options)), + } + } + + pub fn update(&self, options: JetstreamSubscriberOptions) { + self.inner.store(Arc::new(options)); + } + + pub fn max_message_size_bytes(&self) -> u32 { + self.inner.load().max_message_size_bytes() + } + + pub(crate) fn wants(&self, event: &StoredJetstreamEvent<'_>) -> bool { + self.inner.load().wants(event) + } +} + +#[derive(Clone, Default)] +pub struct JetstreamSubscriberOptions { + wanted_collections: Option, + wanted_dids: Arc>>, + max_message_size_bytes: u32, +} + +#[derive(Clone)] +struct WantedCollections { + prefixes: Vec, + full_paths: HashSet, +} + +impl JetstreamSubscriberOptions { + pub fn parse( + wanted_collections: &[String], + wanted_dids: &[String], + max_message_size_bytes: u32, + ) -> std::result::Result { + let wanted_collections = parse_wanted_collections(wanted_collections)?; + let wanted_dids = parse_wanted_dids(wanted_dids)?; + Ok(Self { + wanted_collections, + wanted_dids: Arc::new(wanted_dids), + max_message_size_bytes, + }) + } + + pub fn max_message_size_bytes(&self) -> u32 { + self.max_message_size_bytes + } + + pub(crate) fn wants(&self, event: &StoredJetstreamEvent<'_>) -> bool { + if !self.wanted_dids.is_empty() { + let mut did = Vec::with_capacity(event.did().len()); + event.did().write_to_vec(&mut did); + if !self.wanted_dids.contains(&did) { + return false; + } + } + + let Some(wanted_collections) = &self.wanted_collections else { + return true; + }; + let Some(collection) = event.collection() else { + return true; + }; + + wanted_collections.full_paths.contains(collection) + || wanted_collections + .prefixes + .iter() + .any(|prefix| collection.starts_with(prefix.as_str())) + } +} + +impl Hydrant { + pub fn subscribe_jetstream( + &self, + cursor: Option, + filter: JetstreamFilter, + ) -> JetstreamEventStream { + let (tx, rx) = mpsc::channel(500); + let state = self.state.clone(); + let runtime = tokio::runtime::Handle::current(); + let opts = stream::StreamOptions::from_config(&self.config); + + std::thread::Builder::new() + .name("hydrant-jetstream".into()) + .spawn(move || { + let _g = runtime.enter(); + stream::jetstream_stream_thread(state, tx, cursor, filter, opts); + }) + .expect("failed to spawn Jetstream thread"); + + JetstreamEventStream(rx) + } +} + +fn parse_wanted_collections( + provided: &[String], +) -> std::result::Result, String> { + if provided.is_empty() { + return Ok(None); + } + if provided.len() > 100 { + return Err("too many wanted collections".into()); + } + + let mut prefixes = Vec::new(); + let mut full_paths = HashSet::new(); + for collection_raw in provided { + for collection in collection_raw.split(',') { + if collection.is_empty() { + continue; + } + if let Some(prefix) = collection.strip_suffix(".*") { + validate_collection_prefix(prefix) + .map_err(|_| format!("invalid collection prefix: {collection}"))?; + prefixes.push(format!("{prefix}.").to_smolstr()); + continue; + } + + let nsid = Nsid::new(collection) + .map_err(|_| format!("invalid collection: {collection}"))? + .into_static(); + full_paths.insert(nsid.as_str().to_smolstr()); + } + } + + Ok(Some(WantedCollections { + prefixes, + full_paths, + })) +} + +fn validate_collection_prefix(prefix: &str) -> std::result::Result<(), String> { + let synthetic = format!("{prefix}.x"); + Nsid::new(synthetic.as_str()) + .map(|_| ()) + .map_err(|e| e.to_string()) +} + +fn parse_wanted_dids(provided: &[String]) -> std::result::Result>, String> { + let mut wanted = HashSet::new(); + for did_raw in provided { + for did_str in did_raw.split(',') { + if did_str.is_empty() { + continue; + } + let did = Did::new(did_str).map_err(|_| "invalid wanted DID".to_string())?; + let trimmed = TrimmedDid::from(&did); + let mut buf = Vec::with_capacity(trimmed.len()); + trimmed.write_to_vec(&mut buf); + wanted.insert(buf); + } + } + if wanted.len() > 10_000 { + return Err("too many wanted DIDs".into()); + } + Ok(wanted) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn collection_prefixes_accept_domain_prefixes() { + let opts = JetstreamSubscriberOptions::parse(&["app.bsky.*".into()], &[], 0).unwrap(); + let wanted = opts.wanted_collections.unwrap(); + + assert!(wanted.prefixes.iter().any(|prefix| prefix == "app.bsky.")); + assert!(JetstreamSubscriberOptions::parse(&["app.bsky.feed.po*".into()], &[], 0).is_err()); + } + + #[test] + fn full_path_filter_matches_relay_commit_collection() { + use crate::db::types::TrimmedDid; + use crate::types::StoredJetstreamEvent; + use jacquard_common::{CowStr, IntoStatic}; + use smol_str::ToSmolStr; + + let opts = JetstreamSubscriberOptions::parse(&["app.bsky.feed.post".into()], &[], 0) + .expect("app.bsky.feed.post must be a valid collection"); + + let did = + TrimmedDid::from(&jacquard_common::types::string::Did::new("did:plc:abc123").unwrap()) + .into_static(); + + let matching = StoredJetstreamEvent::RelayCommit { + did: did.clone(), + collection: CowStr::Owned("app.bsky.feed.post".to_smolstr()), + relay_seq: 1, + op_index: 0, + }; + let non_matching = StoredJetstreamEvent::RelayCommit { + did: did.clone(), + collection: CowStr::Owned("app.bsky.actor.profile".to_smolstr()), + relay_seq: 2, + op_index: 0, + }; + let account = StoredJetstreamEvent::RelayAccount { did, relay_seq: 3 }; + + assert!( + opts.wants(&matching), + "app.bsky.feed.post commit should match" + ); + assert!( + !opts.wants(&non_matching), + "other collection must not match" + ); + assert!( + opts.wants(&account), + "account events must pass through collection filter" + ); + } + + #[test] + fn wanted_did_limit_counts_unique_dids() { + let dids = vec!["did:plc:abc123".to_string(); 10_001]; + + let opts = JetstreamSubscriberOptions::parse(&[], &dids, 0).unwrap(); + + assert_eq!(opts.wanted_dids.len(), 1); + } +} diff --git a/src/control/mod.rs b/src/control/mod.rs index c2ed413..5250003 100644 --- a/src/control/mod.rs +++ b/src/control/mod.rs @@ -18,6 +18,10 @@ pub use indexer::*; mod relay; #[cfg(feature = "relay")] pub use relay::*; +#[cfg(feature = "jetstream")] +mod jetstream; +#[cfg(feature = "jetstream")] +pub use jetstream::*; pub use filter::{FilterControl, FilterPatch, FilterSnapshot}; pub use firehose::{FirehoseHandle, FirehoseSourceInfo}; @@ -790,6 +794,11 @@ impl Hydrant { "relay_events", state.db.relay_events.approximate_len() as u64, ); + #[cfg(feature = "jetstream")] + counts.insert( + "jetstream_events", + state.db.jetstream_events.approximate_len() as u64, + ); let sizes = tokio::task::spawn_blocking(move || { let mut s = BTreeMap::new(); @@ -812,6 +821,8 @@ impl Hydrant { #[cfg(feature = "relay")] s.insert("relay_events", state.db.relay_events.disk_space()); + #[cfg(feature = "jetstream")] + s.insert("jetstream_events", state.db.jetstream_events.disk_space()); #[cfg(feature = "backlinks")] s.insert("backlinks", state.db.backlinks.disk_space()); diff --git a/src/control/stream.rs b/src/control/stream.rs index e188daf..0862577 100644 --- a/src/control/stream.rs +++ b/src/control/stream.rs @@ -28,6 +28,16 @@ use { #[cfg(feature = "relay")] use super::RelayStreamError; +#[cfg(feature = "jetstream")] +use { + super::{JetstreamStreamError, JetstreamSubscriberOptions}, + crate::types::{JetstreamBroadcast, StoredJetstreamEvent}, + bytes::Bytes, +}; + +#[cfg(all(feature = "relay", feature = "jetstream"))] +use crate::ingest::stream::{SubscribeReposMessage, decode_frame}; + #[cfg(any(feature = "indexer_stream", feature = "relay"))] const STREAM_SEND_RETRY_PAUSE: Duration = Duration::from_millis(10); @@ -103,6 +113,13 @@ impl From for RelayStreamError { } } +#[cfg(feature = "jetstream")] +impl From for JetstreamStreamError { + fn from(err: StreamTooSlow) -> Self { + Self::ConsumerTooSlow { reason: err.reason } + } +} + #[cfg(any(feature = "indexer_stream", feature = "relay"))] trait StreamBroadcast { fn sequence(&self) -> u64; @@ -137,6 +154,17 @@ impl StreamBroadcast for crate::types::RelayBroadcast { } } +#[cfg(feature = "jetstream")] +impl StreamBroadcast for JetstreamBroadcast { + fn sequence(&self) -> u64 { + self.id + } + + fn is_persisted_marker(&self) -> bool { + false + } +} + #[cfg(any(feature = "indexer_stream", feature = "relay"))] struct PendingLiveEvents { queue: VecDeque, @@ -648,6 +676,489 @@ fn relay_broadcast_to_frame(event: crate::types::RelayBroadcast) -> Option, + tx: mpsc::Sender>, + cursor: Option, + filter: crate::control::JetstreamFilter, + opts: StreamOptions, +) { + let mut event_rx = state.db.jetstream_tx.subscribe(); + let head = latest_jetstream_head(&state); + + let replay = cursor + .filter(|cursor| *cursor <= chrono::Utc::now().timestamp_micros()) + .map(|cursor| cursor.max(0) as u64); + + // for cursor replay: skip live events with id <= head to avoid duplicating + // events already covered by the replay window. + let mut current_id = replay.and(head).map(|(_, id)| id); + + if let (Some(cursor_us), Some((target_time_us, _))) = (replay, head) { + let start_key = keys::jetstream_event_key(cursor_us, 0); + match send_jetstream_replay_window( + &state, + &tx, + &mut event_rx, + start_key.to_vec(), + target_time_us, + &filter, + opts, + ) { + None => return, + Some(max_id) => { + if let Some(max_id) = max_id { + current_id = Some(max_id); + } + } + } + } + + // live tail: read directly from the broadcast channel without a buffered + // pending queue. this avoids the pending queue overflow that occurs when + // most events are filtered (e.g. wantedCollections) and send_stream_event + // is rarely called during replay, leaving pending unbounded. + loop { + match event_rx.blocking_recv() { + Ok(event) => { + let id = event.id; + if current_id.is_some_and(|c| id <= c) { + continue; + } + let out = jetstream_event_to_bytes(&state, event, &filter); + if let Some(bytes) = out { + match send_jetstream_live_event(&tx, &mut event_rx, bytes, opts) { + Ok(()) => current_id = Some(id), + Err(()) => return, + } + } else { + current_id = Some(id); + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => { + let err = StreamTooSlow::lagged(skipped); + warn!(%err, "closing slow Jetstream subscriber"); + send_stream_error(&tx, err.into()); + return; + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => break, + } + } +} + +#[cfg(feature = "jetstream")] +fn latest_jetstream_head(state: &AppState) -> Option<(u64, u64)> { + let guard = state.db.jetstream_events.iter().next_back()?; + let key = match guard.key() { + Ok(key) => key, + Err(e) => { + error!(err = %e, "failed to read latest Jetstream event key"); + return None; + } + }; + match keys::parse_jetstream_event_key(&key) { + Ok(parsed) => Some(parsed), + Err(e) => { + error!(err = %e, "failed to parse latest Jetstream event key"); + None + } + } +} + +// drain the Jetstream broadcast channel without buffering, to prevent the +// receiver from lagging during long replay windows. +#[cfg(feature = "jetstream")] +fn drain_jetstream_broadcast(event_rx: &mut broadcast::Receiver) { + loop { + match event_rx.try_recv() { + Ok(_) | Err(broadcast::error::TryRecvError::Lagged(_)) => {} + Err(broadcast::error::TryRecvError::Empty | broadcast::error::TryRecvError::Closed) => { + break; + } + } + } +} + +// send one live Jetstream event, retrying until the channel drains or timeout. +// drains the broadcast channel during retries to keep the receiver from lagging. +#[cfg(feature = "jetstream")] +fn send_jetstream_live_event( + tx: &mpsc::Sender>, + event_rx: &mut broadcast::Receiver, + bytes: Bytes, + opts: StreamOptions, +) -> Result<(), ()> { + let mut item = Ok(bytes); + let started = Instant::now(); + loop { + match tx.try_send(item) { + Ok(()) => return Ok(()), + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => return Err(()), + Err(tokio::sync::mpsc::error::TrySendError::Full(returned)) => { + item = returned; + drain_jetstream_broadcast(event_rx); + if started.elapsed() >= opts.send_timeout { + send_stream_error(tx, StreamTooSlow::send_timeout(opts.send_timeout).into()); + return Err(()); + } + std::thread::sleep(Duration::from_millis(10)); + } + } + } +} + +#[cfg(feature = "jetstream")] +fn send_jetstream_replay_window( + state: &AppState, + tx: &mpsc::Sender>, + event_rx: &mut broadcast::Receiver, + mut next_key: Vec, + target_time_us: u64, + filter: &crate::control::JetstreamFilter, + opts: StreamOptions, +) -> Option> { + let mut max_id_seen: Option = None; + loop { + // drain live events without buffering so the broadcast receiver never + // lags regardless of how many events arrive during the replay window. + drain_jetstream_broadcast(event_rx); + + let chunk = + read_jetstream_replay_chunk(state, &next_key, target_time_us, opts.replay_chunk_size); + for event in chunk.events { + let event_id = event.id; + next_key = keys::jetstream_event_key(event.time_us as u64, event_id.saturating_add(1)) + .to_vec(); + max_id_seen = Some(max_id_seen.unwrap_or(0).max(event_id)); + if let Some(bytes) = jetstream_event_to_bytes(state, event, filter) { + match send_jetstream_live_event(tx, event_rx, bytes, opts) { + Ok(()) => {} + Err(()) => return None, + } + } + } + + if chunk.exhausted { + return Some(max_id_seen); + } + if !opts.replay_chunk_pause.is_zero() { + std::thread::sleep(opts.replay_chunk_pause); + } + } +} + +#[cfg(feature = "jetstream")] +fn read_jetstream_replay_chunk( + state: &AppState, + start_key: &[u8], + target_time_us: u64, + chunk_size: usize, +) -> ReplayChunk { + let end_key = keys::jetstream_event_key(target_time_us, u64::MAX); + let mut events = Vec::with_capacity(chunk_size); + let mut exhausted = false; + let mut iter = state.db.jetstream_events.range(start_key..=&end_key[..]); + + while events.len() < chunk_size { + let Some(item) = iter.next() else { + exhausted = true; + break; + }; + + let (k, v) = match item.into_inner() { + Ok(kv) => kv, + Err(e) => { + error!(err = %e, "failed to read Jetstream event from db"); + exhausted = true; + break; + } + }; + let (time_us, id) = match keys::parse_jetstream_event_key(&k) { + Ok(parsed) => parsed, + Err(e) => { + error!(err = %e, "failed to parse Jetstream event key"); + continue; + } + }; + + let event: StoredJetstreamEvent = match rmp_serde::from_slice(&v) { + Ok(event) => event, + Err(e) => { + error!(err = %e, "failed to deserialize Jetstream event"); + continue; + } + }; + events.push(JetstreamBroadcast { + id, + time_us: time_us as i64, + event: event.into_static(), + }); + } + + ReplayChunk { + events, + last_seen_seq: None, + exhausted, + } +} + +#[cfg(feature = "jetstream")] +#[derive(serde::Serialize)] +struct JetstreamEvent<'a> { + did: &'a str, + time_us: i64, + #[serde(flatten)] + payload: JetstreamPayload<'a>, +} + +#[cfg(feature = "jetstream")] +#[derive(serde::Serialize)] +#[serde(tag = "kind", rename_all = "lowercase")] +enum JetstreamPayload<'a> { + Commit { + commit: JetstreamCommit<'a>, + }, + Identity { + #[serde(borrow)] + identity: JetstreamIdentity<'a>, + }, + Account { + #[serde(borrow)] + account: JetstreamAccount<'a>, + }, +} + +#[cfg(feature = "jetstream")] +#[derive(serde::Serialize)] +struct JetstreamCommit<'a> { + rev: &'a str, + operation: &'a str, + collection: &'a str, + rkey: &'a str, + #[serde(skip_serializing_if = "Option::is_none")] + record: Option<&'a serde_json::Value>, + #[serde(skip_serializing_if = "Option::is_none")] + cid: Option, +} + +#[cfg(feature = "jetstream")] +#[derive(serde::Serialize)] +struct JetstreamIdentity<'a> { + did: String, + seq: i64, + time: &'a crate::ingest::stream::Datetime, + #[serde(skip_serializing_if = "Option::is_none")] + handle: Option, +} + +#[cfg(feature = "jetstream")] +#[derive(serde::Serialize)] +struct JetstreamAccount<'a> { + active: bool, + did: String, + seq: i64, + time: &'a crate::ingest::stream::Datetime, + #[serde(skip_serializing_if = "Option::is_none")] + status: Option, +} + +#[cfg(feature = "jetstream")] +fn jetstream_event_to_bytes( + state: &AppState, + event: JetstreamBroadcast, + filter: &crate::control::JetstreamFilter, +) -> Option { + if !filter.wants(&event.event) { + return None; + } + + match &event.event { + #[cfg(feature = "indexer_stream")] + StoredJetstreamEvent::Commit { event_id, .. } => { + let bytes = state.db.events.get(keys::event_key(*event_id)).ok()??; + let stored: StoredEvent = rmp_serde::from_slice(&bytes).ok()?; + let evt = stored_to_event(state, *event_id, stored, None)?; + let rec = evt.record?; + let did_str = rec.did.as_str(); + + let json_event = JetstreamEvent { + did: did_str, + time_us: event.time_us, + payload: JetstreamPayload::Commit { + commit: JetstreamCommit { + rev: rec.rev.as_str(), + operation: rec.action.as_str(), + collection: rec.collection.as_str(), + rkey: rec.rkey.as_str(), + record: rec.record.as_ref(), + cid: rec.cid.map(|cid| cid.to_string()), + }, + }, + }; + serde_json::to_vec(&json_event).ok().map(Bytes::from) + } + #[cfg(feature = "relay")] + StoredJetstreamEvent::RelayCommit { + relay_seq, + op_index, + .. + } => { + let frame = state + .db + .relay_events + .get(keys::relay_event_key(*relay_seq)) + .ok()??; + let SubscribeReposMessage::Commit(commit) = decode_frame(frame.as_ref()).ok()? else { + return None; + }; + let op = commit.ops.get(*op_index as usize)?; + let (collection, rkey) = op.path.split_once('/')?; + let action = op.action.as_str(); + + let mut record_owned = None; + let cid = if matches!(action, "create" | "update") { + let cid = op.cid.as_ref()?; + let cid_ipld = cid.to_ipld().ok()?; + let parsed = tokio::runtime::Handle::current() + .block_on(jacquard_repo::car::reader::parse_car_bytes( + commit.blocks.as_ref(), + )) + .ok()?; + let block = parsed.blocks.get(&cid_ipld)?; + let val = serde_ipld_dagcbor::from_slice::(block).ok()?; + record_owned = Some(serde_json::to_value(val).ok()?); + Some(cid.to_string()) + } else { + None + }; + + let json_event = JetstreamEvent { + did: commit.repo.as_str(), + time_us: event.time_us, + payload: JetstreamPayload::Commit { + commit: JetstreamCommit { + rev: commit.rev.as_str(), + operation: action, + collection, + rkey, + record: record_owned.as_ref(), + cid, + }, + }, + }; + serde_json::to_vec(&json_event).ok().map(Bytes::from) + } + #[cfg(feature = "relay")] + StoredJetstreamEvent::RelayAccount { relay_seq, .. } => { + let frame = state + .db + .relay_events + .get(keys::relay_event_key(*relay_seq)) + .ok()??; + let SubscribeReposMessage::Account(account) = decode_frame(frame.as_ref()).ok()? else { + return None; + }; + let did_str = account.did.as_str(); + let status = account.status.as_ref().map(|s| { + use jacquard_common::{IntoStatic, cowstr::ToCowStr}; + s.to_cowstr().into_static().as_str().to_string() + }); + + let json_event = JetstreamEvent { + did: did_str, + time_us: event.time_us, + payload: JetstreamPayload::Account { + account: JetstreamAccount { + active: account.active, + did: did_str.to_string(), + seq: account.seq, + time: &account.time, + status, + }, + }, + }; + serde_json::to_vec(&json_event).ok().map(Bytes::from) + } + #[cfg(feature = "relay")] + StoredJetstreamEvent::RelayIdentity { relay_seq, .. } => { + let frame = state + .db + .relay_events + .get(keys::relay_event_key(*relay_seq)) + .ok()??; + let SubscribeReposMessage::Identity(identity) = decode_frame(frame.as_ref()).ok()? + else { + return None; + }; + let did_str = identity.did.as_str(); + let handle = identity.handle.as_ref().map(|h| h.as_str().to_string()); + + let json_event = JetstreamEvent { + did: did_str, + time_us: event.time_us, + payload: JetstreamPayload::Identity { + identity: JetstreamIdentity { + did: did_str.to_string(), + seq: identity.seq, + time: &identity.time, + handle, + }, + }, + }; + serde_json::to_vec(&json_event).ok().map(Bytes::from) + } + StoredJetstreamEvent::Account { + did, + active, + status, + seq, + time, + } => { + let did_str = did.to_did(); + let status = status.as_ref().map(|s| s.as_str().to_string()); + + let json_event = JetstreamEvent { + did: did_str.as_str(), + time_us: event.time_us, + payload: JetstreamPayload::Account { + account: JetstreamAccount { + active: *active, + did: did_str.as_str().to_string(), + seq: *seq, + time, + status, + }, + }, + }; + serde_json::to_vec(&json_event).ok().map(Bytes::from) + } + StoredJetstreamEvent::Identity { + did, + handle, + seq, + time, + } => { + let did_str = did.to_did(); + let handle = handle.as_ref().map(|h| h.as_str().to_string()); + + let json_event = JetstreamEvent { + did: did_str.as_str(), + time_us: event.time_us, + payload: JetstreamPayload::Identity { + identity: JetstreamIdentity { + did: did_str.as_str().to_string(), + seq: *seq, + time, + handle, + }, + }, + }; + serde_json::to_vec(&json_event).ok().map(Bytes::from) + } + } +} + #[cfg(all(test, any(feature = "indexer_stream", feature = "relay")))] mod tests { use super::*; diff --git a/src/db/ephemeral.rs b/src/db/ephemeral.rs index 7838bc1..e002922 100644 --- a/src/db/ephemeral.rs +++ b/src/db/ephemeral.rs @@ -17,6 +17,10 @@ pub fn ephemeral_ttl_worker(state: Arc) { if let Err(e) = ephemeral_ttl_tick(&state.db, &state.ephemeral_ttl) { error!(err = %e, "ephemeral TTL tick failed"); } + #[cfg(feature = "jetstream")] + if let Err(e) = jetstream_events_ttl_tick(&state.db, &state.ephemeral_ttl) { + error!(err = %e, "jetstream TTL tick failed"); + } } } @@ -28,6 +32,10 @@ pub fn relay_events_ttl_worker(state: Arc) { if let Err(e) = relay_events_ttl_tick(&state.db, &state.ephemeral_ttl) { error!(err = %e, "relay events TTL tick failed"); } + #[cfg(feature = "jetstream")] + if let Err(e) = jetstream_events_ttl_tick(&state.db, &state.ephemeral_ttl) { + error!(err = %e, "jetstream TTL tick failed"); + } } } @@ -57,6 +65,40 @@ pub fn relay_events_ttl_tick(db: &Db, ttl: &Duration) -> miette::Result<()> { ) } +#[cfg(feature = "jetstream")] +pub fn jetstream_events_ttl_tick(db: &Db, ttl: &Duration) -> miette::Result<()> { + let now = chrono::Utc::now().timestamp() as u64; + let cutoff_ts = now.saturating_sub(ttl.as_secs()); + let cutoff_us = cutoff_ts.saturating_mul(1_000_000); + + db.jetstream_events + .rotate_memtable_and_wait() + .into_diagnostic() + .wrap_err("failed to rotate memtable before Jetstream TTL range drop")?; + + let before_space = db.jetstream_events.disk_space(); + let before_tables = db.jetstream_events.table_count(); + db.jetstream_events + .drop_range(..keys::jetstream_event_key(cutoff_us, 0)) + .into_diagnostic() + .wrap_err("failed Jetstream TTL range drop for old events")?; + let after_space = db.jetstream_events.disk_space(); + let after_tables = db.jetstream_events.table_count(); + + info!( + cutoff_us, + reclaimed_bytes = before_space.saturating_sub(after_space), + dropped_tables = before_tables.saturating_sub(after_tables), + before_space, + after_space, + before_tables, + after_tables, + "dropped old Jetstream events for TTL" + ); + + Ok(()) +} + #[cfg(any(feature = "indexer_stream", feature = "relay"))] fn ttl_tick_inner( db: &Db, diff --git a/src/db/keys/mod.rs b/src/db/keys/mod.rs index c8e00b5..1f0c502 100644 --- a/src/db/keys/mod.rs +++ b/src/db/keys/mod.rs @@ -56,6 +56,32 @@ pub fn event_key(seq: u64) -> [u8; 8] { seq.to_be_bytes() } +#[cfg(feature = "jetstream")] +pub fn jetstream_event_key(time_us: u64, id: u64) -> [u8; 16] { + let mut key = [0u8; 16]; + key[..8].copy_from_slice(&time_us.to_be_bytes()); + key[8..].copy_from_slice(&id.to_be_bytes()); + key +} + +#[cfg(feature = "jetstream")] +pub fn parse_jetstream_event_key(key: &[u8]) -> miette::Result<(u64, u64)> { + if key.len() != 16 { + miette::bail!("jetstream event key is not 16 bytes"); + } + let time_us = u64::from_be_bytes( + key[..8] + .try_into() + .map_err(|e| miette::miette!("invalid jetstream time bytes: {e}"))?, + ); + let id = u64::from_be_bytes( + key[8..] + .try_into() + .map_err(|e| miette::miette!("invalid jetstream id bytes: {e}"))?, + ); + Ok((time_us, id)) +} + pub const COUNT_KS_PREFIX: &[u8] = &[b'k', SEP]; // count keys for the counts keyspace diff --git a/src/db/mod.rs b/src/db/mod.rs index 85d82ce..dd8f0a2 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -4,6 +4,8 @@ use crate::types::{RepoMetadata, RepoState}; #[cfg(feature = "indexer_stream")] use crate::types::BroadcastEvent; +#[cfg(feature = "jetstream")] +use crate::types::JetstreamBroadcast; #[cfg(feature = "relay")] use crate::types::RelayBroadcast; @@ -18,6 +20,8 @@ use smol_str::SmolStr; use std::cell::RefCell; use std::collections::{BTreeMap, BTreeSet, HashSet}; +#[cfg(feature = "jetstream")] +use std::sync::atomic::AtomicI64; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use url::Url; @@ -60,12 +64,22 @@ pub struct Db { pub resync_buffer: Keyspace, #[cfg(feature = "indexer_stream")] pub events: Keyspace, + #[cfg(feature = "jetstream")] + pub(crate) jetstream_events: Keyspace, #[cfg(feature = "backlinks")] pub backlinks: Keyspace, #[cfg(feature = "indexer_stream")] pub(crate) event_tx: broadcast::Sender, #[cfg(feature = "indexer_stream")] pub next_event_id: Arc, + #[cfg(feature = "jetstream")] + pub(crate) jetstream_tx: broadcast::Sender, + #[cfg(feature = "jetstream")] + pub(crate) next_jetstream_id: Arc, + #[cfg(feature = "jetstream")] + pub(crate) last_jetstream_time_us: Arc, + #[cfg(feature = "jetstream")] + pub(crate) jetstream_lock: Arc>, #[cfg(feature = "relay")] pub(crate) relay_events: Keyspace, #[cfg(feature = "relay")] @@ -246,16 +260,15 @@ impl Db { } None }; - let dicts = ["repos", "blocks", "events", "backlinks"].into_iter().fold( - std::collections::HashMap::new(), - |mut acc, name| { + let dicts = ["repos", "blocks", "events", "jetstream_events", "backlinks"] + .into_iter() + .fold(std::collections::HashMap::new(), |mut acc, name| { let Some(dict) = load_dict(name) else { return acc; }; acc.insert(name, dict); acc - }, - ); + }); let get_compression = |name: &str, level: i32| match cfg.data_compression { Compression::Lz4 => CompressionType::Lz4, Compression::Zstd => dicts @@ -433,6 +446,22 @@ impl Db { // ids are int, we can prefix truncate a lot .data_block_restart_interval_policy(RestartIntervalPolicy::new([64, 128])), )?; + + #[cfg(feature = "jetstream")] + let jetstream_events = open_ks( + "jetstream_events", + opts() + // time-ordered append-only stream metadata, only iterated for replay. + .expect_point_read_hits(true) + .max_memtable_size(mb(cfg.db_events_memtable_size_mb)) + .data_block_size_policy(BlockSizePolicy::new([kb(16), kb(64), kb(128)])) + .data_block_compression_policy(CompressionPolicy::new([ + CompressionType::None, + get_compression("jetstream_events", 3), + get_compression("jetstream_events", 5), + ])) + .data_block_restart_interval_policy(RestartIntervalPolicy::new([64, 128])), + )?; let counts = open_ks( "counts", opts() @@ -512,6 +541,9 @@ impl Db { #[cfg(feature = "relay")] let (relay_broadcast_tx, _) = broadcast::channel(10000); + #[cfg(feature = "jetstream")] + let (jetstream_tx, _) = broadcast::channel(10000); + let this = Self { inner: db, path: cfg.database_path.clone(), @@ -530,6 +562,8 @@ impl Db { resync_buffer, #[cfg(feature = "indexer_stream")] events, + #[cfg(feature = "jetstream")] + jetstream_events, counts, filter, crawler, @@ -540,6 +574,14 @@ impl Db { counts_map: HashMap::new(), #[cfg(feature = "indexer_stream")] next_event_id: Arc::new(AtomicU64::new(0)), + #[cfg(feature = "jetstream")] + jetstream_tx, + #[cfg(feature = "jetstream")] + next_jetstream_id: Arc::new(AtomicU64::new(0)), + #[cfg(feature = "jetstream")] + last_jetstream_time_us: Arc::new(AtomicI64::new(0)), + #[cfg(feature = "jetstream")] + jetstream_lock: Arc::new(parking_lot::Mutex::new(())), #[cfg(feature = "relay")] relay_events, #[cfg(feature = "relay")] @@ -587,6 +629,22 @@ impl Db { .store(last_id + 1, std::sync::atomic::Ordering::Relaxed); } + #[cfg(feature = "jetstream")] + { + let mut last_id = 0; + let mut last_time_us = 0; + if let Some(guard) = this.jetstream_events.iter().next_back() { + let k = guard.key().into_diagnostic()?; + let (time_us, id) = keys::parse_jetstream_event_key(&k)?; + last_id = id; + last_time_us = time_us; + } + this.next_jetstream_id + .store(last_id + 1, std::sync::atomic::Ordering::Relaxed); + this.last_jetstream_time_us + .store(last_time_us as i64, std::sync::atomic::Ordering::Relaxed); + } + // load counts into memory for guard in this.counts.prefix(keys::COUNT_KS_PREFIX) { let (k, v) = guard.into_inner().into_diagnostic()?; @@ -628,6 +686,8 @@ impl Db { "blocks" => &self.blocks, #[cfg(feature = "indexer_stream")] "events" => &self.events, + #[cfg(feature = "jetstream")] + "jetstream_events" => &self.jetstream_events, "repos" => &self.repos, #[cfg(feature = "backlinks")] "backlinks" => &self.backlinks, @@ -637,6 +697,7 @@ impl Db { let dict_size = match ks_name { "blocks" => kb(128), "events" => kb(64), + "jetstream_events" => kb(64), "repos" => kb(64), "backlinks" => kb(64), _ => kb(32), @@ -748,6 +809,9 @@ impl Db { #[cfg(feature = "indexer_stream")] tasks.push(compact(self.events.clone())); + #[cfg(feature = "jetstream")] + tasks.push(compact(self.jetstream_events.clone())); + #[cfg(feature = "relay")] tasks.push(compact(self.relay_events.clone())); diff --git a/src/ingest/indexer.rs b/src/ingest/indexer.rs index 4b54ae1..c060b2e 100644 --- a/src/ingest/indexer.rs +++ b/src/ingest/indexer.rs @@ -13,6 +13,11 @@ use fjall::OwnedWriteBatch; use jacquard_common::IntoStatic; use jacquard_common::types::did::Did; +#[cfg(feature = "jetstream")] +use crate::{ + db::types::TrimmedDid, + types::{JetstreamBroadcast, StoredJetstreamEvent}, +}; use jacquard_repo::error::CommitError; use miette::{Diagnostic, IntoDiagnostic, Result}; use std::sync::Arc; @@ -169,6 +174,8 @@ struct WorkerContext<'a> { count_deltas: &'a mut CountDeltas, #[cfg(feature = "indexer_stream")] broadcast_events: &'a mut Vec, + #[cfg(feature = "jetstream")] + jetstream_events: &'a mut Vec, } impl FirehoseWorker { @@ -213,11 +220,15 @@ impl FirehoseWorker { #[cfg(feature = "indexer_stream")] let mut broadcast_events = Vec::new(); + #[cfg(feature = "jetstream")] + let mut jetstream_events = Vec::new(); while let Some(msg) = rx.blocking_recv() { let batch = state.db.inner.batch(); #[cfg(feature = "indexer_stream")] broadcast_events.clear(); + #[cfg(feature = "jetstream")] + jetstream_events.clear(); let mut added_blocks = 0; let mut records_delta = 0; @@ -231,6 +242,8 @@ impl FirehoseWorker { count_deltas: &mut count_deltas, #[cfg(feature = "indexer_stream")] broadcast_events: &mut broadcast_events, + #[cfg(feature = "jetstream")] + jetstream_events: &mut jetstream_events, }; match msg { @@ -424,6 +437,10 @@ impl FirehoseWorker { for evt in broadcast_events.drain(..) { let _ = state.db.event_tx.send(evt); } + #[cfg(feature = "jetstream")] + for evt in jetstream_events.drain(..) { + let _ = state.db.jetstream_tx.send(evt); + } // state.db.inner.persist(fjall::PersistMode::Buffer).ok(); } @@ -515,6 +532,8 @@ impl FirehoseWorker { .push(BroadcastEvent::Persisted(last_id)); } } + #[cfg(feature = "jetstream")] + ctx.jetstream_events.extend(res.jetstream_events); Ok(RepoProcessResult::Ok(repo_state)) } @@ -530,6 +549,18 @@ impl FirehoseWorker { { let db = &ctx.state.db; let did = &identity.did; + #[cfg(feature = "jetstream")] + ctx.jetstream_events.push(crate::jetstream::stage_event( + &mut ctx.batch, + db, + StoredJetstreamEvent::Identity { + did: TrimmedDid::from(did).into_static(), + handle: identity.handle.clone().map(IntoStatic::into_static), + seq: identity.seq, + time: identity.time.clone(), + }, + )?); + if changed { let evt = IdentityEvt { did: did.clone().into_static(), @@ -558,6 +589,18 @@ impl FirehoseWorker { active: account.active, status: account.status.as_ref().map(|s| s.to_cowstr().into_static()), }; + #[cfg(feature = "jetstream")] + ctx.jetstream_events.push(crate::jetstream::stage_event( + &mut ctx.batch, + db, + StoredJetstreamEvent::Account { + did: TrimmedDid::from(did).into_static(), + active: account.active, + status: account.status.as_ref().map(|s| s.to_cowstr().into_static()), + seq: account.seq, + time: account.time.clone(), + }, + )?); if is_inactive { use crate::ingest::stream::AccountStatus; diff --git a/src/ingest/relay.rs b/src/ingest/relay.rs index 0216a5c..baca44d 100644 --- a/src/ingest/relay.rs +++ b/src/ingest/relay.rs @@ -18,6 +18,8 @@ use tracing::{debug, error, info, info_span, trace, warn}; use url::Url; use crate::db::keys::pds_account_count_key; +#[cfg(all(feature = "relay", feature = "jetstream"))] +use crate::db::types::TrimmedDid; use crate::db::{self, CountDeltas, keys}; use crate::ingest::stream::AccountStatus; #[cfg(feature = "relay")] @@ -31,6 +33,8 @@ use crate::ingest::{BufferRx, BufferTx, IngestMessage}; use crate::state::AppState; #[cfg(feature = "relay")] use crate::types::RelayBroadcast; +#[cfg(all(feature = "relay", feature = "jetstream"))] +use crate::types::{JetstreamBroadcast, StoredJetstreamEvent}; use crate::types::{RepoState, RepoStatus}; use crate::util; use smol_str::{SmolStr, ToSmolStr}; @@ -43,6 +47,8 @@ struct WorkerContext<'a> { count_deltas: CountDeltas, #[cfg(feature = "relay")] pending_broadcasts: Vec, + #[cfg(all(feature = "relay", feature = "jetstream"))] + pending_jetstream_events: Vec>, #[cfg(feature = "indexer")] pending_hook_messages: Vec, #[cfg(feature = "indexer")] @@ -162,6 +168,8 @@ impl RelayWorker { count_deltas: CountDeltas::default(), #[cfg(feature = "relay")] pending_broadcasts: Vec::with_capacity(2), + #[cfg(all(feature = "relay", feature = "jetstream"))] + pending_jetstream_events: Vec::with_capacity(2), #[cfg(feature = "indexer")] pending_hook_messages: Vec::with_capacity(2), #[cfg(feature = "indexer")] @@ -210,7 +218,29 @@ impl RelayWorker { .state .db .stage_count_deltas(&mut batch, &ctx.count_deltas); + + #[cfg(all(feature = "relay", feature = "jetstream"))] + let mut jetstream_broadcasts = Vec::new(); + + #[cfg(all(feature = "relay", feature = "jetstream"))] + let res = { + let _lock = ctx.state.db.jetstream_lock.lock(); + let mut stage_res = Ok(()); + for event in ctx.pending_jetstream_events.drain(..) { + match crate::jetstream::stage_event(&mut batch, &ctx.state.db, event) { + Ok(broadcast) => jetstream_broadcasts.push(broadcast), + Err(e) => { + stage_res = Err(e); + break; + } + } + } + stage_res.and_then(|_| batch.commit().into_diagnostic()) + }; + + #[cfg(not(all(feature = "relay", feature = "jetstream")))] let res = batch.commit(); + if let Err(e) = res { error!(shard = id, err = %e, "relay shard: failed to commit batch"); drop(reservation); @@ -223,6 +253,10 @@ impl RelayWorker { for broadcast in ctx.pending_broadcasts.drain(..) { let _ = state.db.relay_broadcast_tx.send(broadcast); } + #[cfg(all(feature = "relay", feature = "jetstream"))] + for broadcast in jetstream_broadcasts { + let _ = state.db.jetstream_tx.send(broadcast); + } #[cfg(feature = "indexer")] for msg in ctx.pending_hook_messages.drain(..) { let _ = ctx.hook.blocking_send(msg); @@ -330,10 +364,35 @@ impl RelayWorker { } #[cfg(feature = "relay")] { - ctx.queue_emit(|seq| { + #[cfg(feature = "jetstream")] + let jetstream_ops = commit + .ops + .iter() + .enumerate() + .filter_map(|(idx, op)| { + matches!(op.action.as_str(), "create" | "update" | "delete") + .then(|| split_collection(&op.path).map(|col| (idx as u32, col))) + .flatten() + }) + .collect::>(); + #[cfg(feature = "jetstream")] + let jetstream_did = TrimmedDid::from(&commit.repo).into_static(); + + let relay_seq = ctx.queue_emit(|seq| { commit.seq = seq; encode_frame("#commit", &commit) })?; + #[cfg(feature = "jetstream")] + for (op_index, collection) in jetstream_ops { + // todo: build live jetstream frames from this decoded commit and use stored frames only for replay. + ctx.pending_jetstream_events + .push(StoredJetstreamEvent::RelayCommit { + did: jetstream_did.clone(), + collection, + relay_seq, + op_index, + }); + } } repo_state.root = Some(commit_obj.into()); @@ -466,10 +525,16 @@ impl RelayWorker { } #[cfg(feature = "relay")] { - ctx.queue_emit(|seq| { + let relay_seq = ctx.queue_emit(|seq| { identity.seq = seq; encode_frame("#identity", &identity) })?; + #[cfg(feature = "jetstream")] + ctx.pending_jetstream_events + .push(StoredJetstreamEvent::RelayIdentity { + did: TrimmedDid::from(&identity.did).into_static(), + relay_seq, + }); } ctx.batch.insert( @@ -484,7 +549,7 @@ impl RelayWorker { fn handle_account( ctx: &mut WorkerContext, repo_state: &mut RepoState, - firehose: &Url, + #[allow(unused_variables)] firehose: &Url, #[allow(unused_mut)] mut account: Account<'static>, _is_pds: bool, ) -> Result<()> { @@ -558,10 +623,16 @@ impl RelayWorker { } #[cfg(feature = "relay")] { - ctx.queue_emit(|seq| { + let relay_seq = ctx.queue_emit(|seq| { account.seq = seq; encode_frame("#account", &account) })?; + #[cfg(feature = "jetstream")] + ctx.pending_jetstream_events + .push(StoredJetstreamEvent::RelayAccount { + did: TrimmedDid::from(&account.did).into_static(), + relay_seq, + }); } repo_state.touch(); @@ -931,7 +1002,7 @@ impl WorkerContext<'_> { } #[cfg(feature = "relay")] - fn queue_emit(&mut self, make_frame: impl FnOnce(i64) -> Result) -> Result<()> { + fn queue_emit(&mut self, make_frame: impl FnOnce(i64) -> Result) -> Result { let db = &self.state.db; let seq = db.next_relay_seq.fetch_add(1, Ordering::SeqCst); let frame = make_frame(seq as i64)?; @@ -940,10 +1011,16 @@ impl WorkerContext<'_> { self.pending_broadcasts .push(RelayBroadcast::Ephemeral(seq, frame)); self.pending_broadcasts.push(RelayBroadcast::Persisted(seq)); - Ok(()) + Ok(seq) } } +#[cfg(all(feature = "relay", feature = "jetstream"))] +fn split_collection(path: &str) -> Option> { + path.split_once('/') + .map(|(collection, _)| CowStr::Owned(collection.to_smolstr())) +} + /// outcome of a host authority check. enum AuthorityOutcome { /// stored pds matched the source host immediately. diff --git a/src/jetstream.rs b/src/jetstream.rs new file mode 100644 index 0000000..cda5678 --- /dev/null +++ b/src/jetstream.rs @@ -0,0 +1,39 @@ +use std::sync::atomic::Ordering; + +use fjall::OwnedWriteBatch; +use miette::{IntoDiagnostic, Result}; + +use crate::db::{Db, keys}; +use crate::types::{JetstreamBroadcast, StoredJetstreamEvent}; + +pub(crate) fn stage_event( + batch: &mut OwnedWriteBatch, + db: &Db, + event: StoredJetstreamEvent<'_>, +) -> Result { + let id = db.next_jetstream_id.fetch_add(1, Ordering::SeqCst); + let time_us = next_time_us(db); + let event = event.into_static(); + let bytes = rmp_serde::to_vec(&event).into_diagnostic()?; + batch.insert( + &db.jetstream_events, + keys::jetstream_event_key(time_us as u64, id), + bytes, + ); + Ok(JetstreamBroadcast { id, time_us, event }) +} + +fn next_time_us(db: &Db) -> i64 { + loop { + let last = db.last_jetstream_time_us.load(Ordering::SeqCst); + let now = chrono::Utc::now().timestamp_micros(); + let next = now.max(last.saturating_add(1)); + if db + .last_jetstream_time_us + .compare_exchange(last, next, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + { + return next; + } + } +} diff --git a/src/lib.rs b/src/lib.rs index dd792f1..d71f758 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,6 +26,11 @@ compile_error!("can't be relay and indexer at the same time"); any(feature = "indexer_stream", feature = "backlinks") ))] compile_error!("indexer dependent features (stream, backlinks) without indexer can't be enabled"); +#[cfg(all( + feature = "jetstream", + not(any(feature = "indexer_stream", feature = "relay")) +))] +compile_error!("jetstream requires either indexer_stream or relay"); pub(crate) mod api; #[cfg(feature = "indexer")] @@ -36,6 +41,8 @@ pub mod backlinks; pub(crate) mod crawler; pub(crate) mod db; pub(crate) mod ingest; +#[cfg(feature = "jetstream")] +pub(crate) mod jetstream; #[cfg(feature = "indexer")] pub(crate) mod ops; pub(crate) mod patch; diff --git a/src/ops.rs b/src/ops.rs index c2b3dfb..cec933d 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -26,6 +26,9 @@ use { std::sync::atomic::Ordering, }; +#[cfg(feature = "jetstream")] +use crate::types::{JetstreamBroadcast, StoredJetstreamEvent}; + pub fn persist_to_resync_buffer(db: &Db, did: &Did, commit: &Commit) -> Result<()> { let key = keys::resync_buffer_key(did, DbTid::from(&commit.rev)); let value = rmp_serde::to_vec_named(commit).into_diagnostic()?; @@ -212,6 +215,8 @@ pub struct ApplyCommitResults<'s> { pub live_events: Vec, #[cfg(feature = "indexer_stream")] pub last_event_id: Option, + #[cfg(feature = "jetstream")] + pub jetstream_events: Vec, } pub fn apply_commit<'s>( @@ -245,6 +250,8 @@ pub fn apply_commit<'s>( let mut live_events = Vec::new(); #[cfg(feature = "indexer_stream")] let mut last_event_id = None; + #[cfg(feature = "jetstream")] + let mut jetstream_events = Vec::new(); for op in &commit.ops { let (collection, rkey) = parse_path(&op.path)?; @@ -371,6 +378,16 @@ pub fn apply_commit<'s>( let bytes = rmp_serde::to_vec(&evt).into_diagnostic()?; batch.insert(&db.events, keys::event_key(event_id), bytes); + #[cfg(feature = "jetstream")] + { + let jetstream = StoredJetstreamEvent::Commit { + did: did_trimmed.clone().into_static(), + collection: collection.clone().into_static(), + event_id, + }; + jetstream_events.push(crate::jetstream::stage_event(batch, db, jetstream)?); + } + if should_broadcast_live { live_events.push(LiveRecordEvent { id: event_id, @@ -404,6 +421,8 @@ pub fn apply_commit<'s>( live_events, #[cfg(feature = "indexer_stream")] last_event_id, + #[cfg(feature = "jetstream")] + jetstream_events, }) } diff --git a/src/types.rs b/src/types.rs index a6a7546..0920f82 100644 --- a/src/types.rs +++ b/src/types.rs @@ -13,8 +13,12 @@ use smol_str::{SmolStr, ToSmolStr}; use crate::db::types::{DbTid, DidKey}; +#[cfg(any(feature = "indexer_stream", feature = "jetstream"))] +use crate::db::types::TrimmedDid; #[cfg(feature = "indexer_stream")] -use crate::db::types::{DbAction, DbRkey, TrimmedDid}; +use crate::db::types::{DbAction, DbRkey}; +#[cfg(feature = "jetstream")] +use crate::ingest::stream::Datetime; use crate::resolver::MiniDoc; pub(crate) mod v2 { @@ -432,6 +436,154 @@ pub(crate) struct LiveRecordEvent { pub inline_block: Option, } +#[cfg(feature = "jetstream")] +#[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(bound(deserialize = "'i: 'de"))] +pub(crate) enum StoredJetstreamEvent<'i> { + #[cfg(feature = "indexer_stream")] + Commit { + #[serde(borrow)] + did: TrimmedDid<'i>, + #[serde(borrow)] + collection: CowStr<'i>, + event_id: u64, + }, + #[cfg(feature = "relay")] + RelayCommit { + #[serde(borrow)] + did: TrimmedDid<'i>, + #[serde(borrow)] + collection: CowStr<'i>, + relay_seq: u64, + op_index: u32, + }, + #[cfg(feature = "relay")] + RelayAccount { + #[serde(borrow)] + did: TrimmedDid<'i>, + relay_seq: u64, + }, + #[cfg(feature = "relay")] + RelayIdentity { + #[serde(borrow)] + did: TrimmedDid<'i>, + relay_seq: u64, + }, + Account { + #[serde(borrow)] + did: TrimmedDid<'i>, + active: bool, + #[serde(borrow)] + status: Option>, + seq: i64, + time: Datetime, + }, + Identity { + #[serde(borrow)] + did: TrimmedDid<'i>, + #[serde(borrow)] + handle: Option>, + seq: i64, + time: Datetime, + }, +} + +#[cfg(feature = "jetstream")] +#[derive(Clone, Debug)] +pub(crate) struct JetstreamBroadcast { + pub id: u64, + pub time_us: i64, + pub event: StoredJetstreamEvent<'static>, +} + +#[cfg(feature = "jetstream")] +impl<'i> StoredJetstreamEvent<'i> { + pub(crate) fn did(&self) -> &TrimmedDid<'i> { + match self { + #[cfg(feature = "indexer_stream")] + Self::Commit { did, .. } => did, + #[cfg(feature = "relay")] + Self::RelayCommit { did, .. } + | Self::RelayAccount { did, .. } + | Self::RelayIdentity { did, .. } => did, + Self::Account { did, .. } | Self::Identity { did, .. } => did, + } + } + + pub(crate) fn collection(&self) -> Option<&str> { + match self { + #[cfg(feature = "indexer_stream")] + Self::Commit { collection, .. } => Some(collection.as_str()), + #[cfg(feature = "relay")] + Self::RelayCommit { collection, .. } => Some(collection.as_str()), + #[cfg(feature = "relay")] + Self::RelayAccount { .. } | Self::RelayIdentity { .. } => None, + Self::Account { .. } | Self::Identity { .. } => None, + } + } + + pub(crate) fn into_static(self) -> StoredJetstreamEvent<'static> { + match self { + #[cfg(feature = "indexer_stream")] + Self::Commit { + did, + collection, + event_id, + } => StoredJetstreamEvent::Commit { + did: did.into_static(), + collection: collection.into_static(), + event_id, + }, + #[cfg(feature = "relay")] + Self::RelayCommit { + did, + collection, + relay_seq, + op_index, + } => StoredJetstreamEvent::RelayCommit { + did: did.into_static(), + collection: collection.into_static(), + relay_seq, + op_index, + }, + #[cfg(feature = "relay")] + Self::RelayAccount { did, relay_seq } => StoredJetstreamEvent::RelayAccount { + did: did.into_static(), + relay_seq, + }, + #[cfg(feature = "relay")] + Self::RelayIdentity { did, relay_seq } => StoredJetstreamEvent::RelayIdentity { + did: did.into_static(), + relay_seq, + }, + Self::Account { + did, + active, + status, + seq, + time, + } => StoredJetstreamEvent::Account { + did: did.into_static(), + active, + status: status.map(IntoStatic::into_static), + seq, + time, + }, + Self::Identity { + did, + handle, + seq, + time, + } => StoredJetstreamEvent::Identity { + did: did.into_static(), + handle: handle.map(IntoStatic::into_static), + seq, + time, + }, + } + } +} + #[cfg(feature = "relay")] #[derive(Clone)] pub(crate) enum RelayBroadcast { diff --git a/tests/common.nu b/tests/common.nu index e62c014..dfe4088 100644 --- a/tests/common.nu +++ b/tests/common.nu @@ -145,6 +145,16 @@ export def build-hydrant-relay [] { parse-hydrant-executable $out.stdout } +# build the hydrant binary for relay with jetstream support (no default features) +export def build-hydrant-relay-jetstream [] { + if ($env | get --optional HYDRANT_BINARY | is-not-empty) { + return $env.HYDRANT_BINARY + } + print "building hydrant for relay with jetstream..." + let out = (^cargo build --no-default-features --features relay,jetstream --message-format json err> /dev/null | complete) + parse-hydrant-executable $out.stdout +} + # start hydrant in the background export def start-hydrant [binary: string, db_path: string, port: int] { let log_file = $"($db_path)/hydrant.log" diff --git a/tests/run_all.nu b/tests/run_all.nu index 5e5d680..36ccc07 100644 --- a/tests/run_all.nu +++ b/tests/run_all.nu @@ -4,7 +4,7 @@ # usage: # nu tests/run_all.nu # nu tests/run_all.nu --only [stream_live_backfill api_repos] -use common.nu [build-hydrant-features build-hydrant-relay] +use common.nu [build-hydrant-features build-hydrant-relay build-hydrant-relay-jetstream] def get_free_ports [count: int] { if $count == 0 { @@ -62,7 +62,17 @@ def test-needs-relay-binary [name: string] { # tests that build a relay-only binary must run last (and serially) to avoid racing on # `target/` artifacts while other tests are executing. try { - open --raw $"tests/($name).nu" | str contains "build-hydrant-relay" + let content = (open --raw $"tests/($name).nu") + # exclude relay-jetstream tests which use a different binary + ($content | str contains "build-hydrant-relay") and not ($content | str contains "build-hydrant-relay-jetstream") + } catch { + false + } +} + +def test-needs-relay-jetstream-binary [name: string] { + try { + open --raw $"tests/($name).nu" | str contains "build-hydrant-relay-jetstream" } catch { false } @@ -99,14 +109,19 @@ def main [--only: list = [], --skip-creds] { } let relay_tests = $tests | where {|t| test-needs-relay-binary $t } - let indexer_tests = $tests | where {|t| not ($relay_tests | any {$in == $t}) } + let relay_jetstream_tests = $tests | where {|t| test-needs-relay-jetstream-binary $t } + let indexer_tests = $tests | where {|t| + not ($relay_tests | any {$in == $t}) and not ($relay_jetstream_tests | any {$in == $t}) + } mut indexer_binary = null mut relay_binary = null + mut relay_jetstream_binary = null + let needs_snapshot = not ($relay_tests | is-empty) or not ($relay_jetstream_tests | is-empty) if not ($indexer_tests | is-empty) { - let built = build-hydrant-features "backlinks" - $indexer_binary = if not ($relay_tests | is-empty) { + let built = build-hydrant-features "backlinks,jetstream" + $indexer_binary = if $needs_snapshot { snapshot-binary $built "indexer" } else { $built @@ -118,6 +133,11 @@ def main [--only: list = [], --skip-creds] { $relay_binary = snapshot-binary $built "relay" } + if not ($relay_jetstream_tests | is-empty) { + let built = build-hydrant-relay-jetstream + $relay_jetstream_binary = snapshot-binary $built "relay_jetstream" + } + print "" let ports = get_free_ports (($tests | length) * 3) @@ -126,7 +146,13 @@ def main [--only: list = [], --skip-creds] { for test in ($tests | enumerate) { let p = {($test | get index) * 3 + $in} let name = ($test | get item) - let binary = if ($relay_tests | any {$in == $name}) { $relay_binary } else { $indexer_binary } + let binary = if ($relay_tests | any {$in == $name}) { + $relay_binary + } else if ($relay_jetstream_tests | any {$in == $name}) { + $relay_jetstream_binary + } else { + $indexer_binary + } let entry = { name: $name, api: ($ports | get (0 | do $p)), diff --git a/tests/stream_jetstream_subscribe.nu b/tests/stream_jetstream_subscribe.nu new file mode 100644 index 0000000..1ad6730 --- /dev/null +++ b/tests/stream_jetstream_subscribe.nu @@ -0,0 +1,261 @@ +#!/usr/bin/env nu +# Tests the Jetstream-compatible /subscribe WebSocket endpoint. +# +# Requires TEST_REPO and TEST_PASSWORD in .env (same credentials used by +# authenticated_stream tests). Skips gracefully when credentials are absent. +# +# Jetstream event format: +# {did, time_us, kind: "commit"|"identity"|"account", commit?|identity?|account?} +# commit: {rev, operation, collection, rkey, record?, cid?} +# identity: {did, seq, time, handle?} +# account: {did, active, seq, time, status?} + +use common.nu * + +# Collect newline-delimited JSON events from a WebSocket into a list. +def collect-ws-json [ws_url: string, output_file: string, duration: duration] { + let pid_file = $"($output_file).pid" + bash -c $"websocat -n '($ws_url)' > '($output_file)' 2>&1 & echo $! > '($pid_file)'" + sleep $duration + let pid = (open $pid_file | str trim | into int) + try { kill $pid } + sleep 200ms + if ($output_file | path exists) { + open $output_file + | str trim + | lines + | each { |line| try { $line | from json } catch { null } } + | compact + } else { + [] + } +} + +def assert-no-error-events [events: list, label: string, ...pids: int] { + let errors = ($events | where { |e| + let t = ($e | get -o type | default "") + let k = ($e | get -o kind | default "") + $t == "error" or $k == "error" + }) + if ($errors | length) > 0 { + fail $"($label): unexpected error event(s): ($errors | first)" ...$pids + } +} + +def assert-commit-structure [c: record, label: string, ...pids: int] { + for field in ["rev", "operation", "collection", "rkey"] { + if not ($field in $c) { + fail $"($label): commit missing field ($field)" ...$pids + } + } + if $c.operation == "create" or $c.operation == "update" { + if not ("cid" in $c) { + fail $"($label): create/update commit missing cid" ...$pids + } + if not ("record" in $c) { + fail $"($label): create/update commit missing record" ...$pids + } + } +} + +def main [] { + let env_vars = load-env-file + let did = ($env_vars | get --optional TEST_REPO) + let password = ($env_vars | get --optional TEST_PASSWORD) + + if ($did | is-empty) or ($password | is-empty) { + print "SKIP: TEST_REPO and TEST_PASSWORD not set in .env" + exit 0 + } + + let port = resolve-test-port 3016 + let url = $"http://localhost:($port)" + let ws_base = $"ws://127.0.0.1:($port)" + let db_path = (mktemp -d -t hydrant_stream_jetstream_subscribe.XXXXXX) + + print $"testing /subscribe \(Jetstream\) for ($did)..." + print $"database path: ($db_path)" + + let pds_url = resolve-pds $did + + # resolve session, activating if needed + let session = (do { + let s = authenticate $pds_url $did $password + let active = ($s | get -o active | default true) + if $active == false { + print "account deactivated; activating before test..." + activate-account $pds_url $s.accessJwt + sleep 2sec + authenticate $pds_url $did $password + } else { + $s + } + }) + mut jwt = $session.accessJwt + print "authentication successful" + + let binary = build-hydrant-features "jetstream" + let instance = (with-env { HYDRANT_RELAY_HOSTS: "wss://bsky.network" } { + start-hydrant $binary $db_path $port + }) + + if not (wait-for-api $url) { + fail "api failed to start" $instance.pid + } + + print $"adding repo ($did)..." + http put -t application/json $"($url)/repos" [{ did: $did }] + + if not (wait-for-backfill $url) { + fail "backfill timed out" $instance.pid + } + + # start live subscriber before creating records so we catch the events + let live_file = $"($db_path)/live.txt" + print "starting live /subscribe listener..." + let live_pid_file = $"($live_file).pid" + bash -c $"websocat -n '($ws_base)/subscribe' > '($live_file)' 2>&1 & echo $! > '($live_pid_file)'" + let live_pid = (open $live_pid_file | str trim | into int) + print $"live listener pid: ($live_pid)" + sleep 1sec + + # --- create / update / delete a record to generate live jetstream events --- + let collection = "app.bsky.feed.post" + let timestamp = (date now | format date "%Y-%m-%dT%H:%M:%SZ") + let record_data = { + "$type": "app.bsky.feed.post", + text: $"hydrant jetstream test ($timestamp)", + createdAt: $timestamp + } + + print "--- action: create ---" + let create_res = create-record $pds_url $jwt $did $collection $record_data + let rkey = ($create_res.uri | split row "/" | last) + print $"created rkey: ($rkey)" + + print "--- action: update ---" + try { + http post -t application/json -H ["Authorization" $"Bearer ($jwt)"] $"($pds_url)/xrpc/com.atproto.repo.putRecord" { + repo: $did, collection: $collection, rkey: $rkey, + record: ($record_data | update text $"updated text ($timestamp)") + } + print "updated record" + } catch { |e| print $"update skipped: ($e.msg)" } + + print "--- action: delete ---" + delete-record $pds_url $jwt $did $collection $rkey + print "deleted record" + + # wait for live events to arrive (up to 30s) + print "waiting for live commit events..." + mut live_events = [] + for i in 1..60 { + sleep 500ms + if ($live_file | path exists) { + $live_events = ( + open $live_file | str trim | lines + | each { |line| try { $line | from json } catch { null } } + | compact + ) + let commits = ($live_events | where { |e| ($e | get -o kind | default "") == "commit" and ($e | get -o did | default "") == $did }) + if ($commits | length) >= 2 { + break + } + } + } + try { kill $live_pid } + sleep 200ms + + print $"live events received: ($live_events | length)" + assert-no-error-events $live_events "live subscription" $instance.pid + + let our_commits = ($live_events | where { |e| ($e | get -o kind | default "") == "commit" and ($e | get -o did | default "") == $did }) + print $"our commit events: ($our_commits | length)" + if ($our_commits | length) < 2 { + fail $"expected >= 2 live commit events, got ($our_commits | length)" $instance.pid + } + + # validate commit event structure + for evt in $our_commits { + if not ("time_us" in $evt) { + fail "live commit event missing time_us" $instance.pid + } + assert-commit-structure $evt.commit "live commit" $instance.pid + } + print "live commit event structure is correct" + + # verify create event has record body and cid + let create_evts = ($our_commits | where { |e| ($e | get -o commit.operation | default "") == "create" }) + if ($create_evts | is-empty) { + fail "no create event found in live events" $instance.pid + } + let create_evt = ($create_evts | first) + let create_rkey = ($create_evt | get -o commit.rkey | default "") + if $create_rkey != $rkey { + fail $"create event rkey mismatch: expected ($rkey), got ($create_rkey)" $instance.pid + } + print "create event matched expected rkey and has record body" + + # --- scenario: cursor replay --- + print "--- scenario: cursor replay ---" + let replay_file = $"($db_path)/replay.txt" + + # cursor is time_us (microseconds since epoch); 60 seconds before now + # use %s (epoch seconds) then integer multiply to avoid nushell float division + let now_s = (date now | format date "%s" | into int) + let cursor_us = (($now_s - 60) * 1_000_000) + let replay_events = (collect-ws-json $"($ws_base)/subscribe?cursor=($cursor_us)" $replay_file 10sec) + print $"cursor replay: ($replay_events | length) events" + assert-no-error-events $replay_events "cursor replay" $instance.pid + + let replay_commits = ($replay_events | where { |e| ($e | get -o kind | default "") == "commit" and ($e | get -o did | default "") == $did }) + if ($replay_commits | length) < 2 { + fail $"cursor replay: expected >= 2 commit events for our DID, got ($replay_commits | length)" $instance.pid + } + print $"cursor replay: ($replay_commits | length) commit events for our DID" + + # verify monotonic time_us ordering + let time_us_vals = ($replay_events | get time_us) + let monotonic = ( + $time_us_vals + | zip ($time_us_vals | skip 1) + | all { |pair| $pair.0 <= $pair.1 } + ) + if not $monotonic { + fail "cursor replay events are not monotonically ordered by time_us" $instance.pid + } + print "time_us ordering is monotonically non-decreasing" + + # --- scenario: wantedCollections filter --- + print "--- scenario: wantedCollections filter ---" + let col_file = $"($db_path)/col_filter.txt" + let col_url = $"($ws_base)/subscribe?cursor=($cursor_us)&wantedCollections=($collection)" + let col_events = (collect-ws-json $col_url $col_file 8sec) + let col_commits = ($col_events | where { |e| ($e | get -o kind | default "") == "commit" }) + print $"wantedCollections=($collection): ($col_commits | length) commit events" + + let col_commits_with_data = ($col_commits | where { |e| ($e | get -o commit.collection | is-not-empty) }) + if ($col_commits_with_data | any { |e| ($e | get -o commit.collection | default "") != $collection }) { + fail "wantedCollections filter returned commits for unexpected collection" $instance.pid + } + # account/identity events always pass through even with collection filter + let non_commit = ($col_events | where { |e| ($e | get -o kind | default "") != "commit" }) + print $"non-commit events through collection filter: ($non_commit | length) \(expected: account/identity pass-through\)" + print "wantedCollections filter is correct" + + # --- scenario: wantedDids filter --- + print "--- scenario: wantedDids filter ---" + let did_file = $"($db_path)/did_filter.txt" + let did_url = $"($ws_base)/subscribe?cursor=($cursor_us)&wantedDids=($did)" + let did_events = (collect-ws-json $did_url $did_file 8sec) + print $"wantedDids=($did): ($did_events | length) events" + + let did_jetstream_events = ($did_events | where { |e| ($e | get -o did | is-not-empty) }) + if ($did_jetstream_events | any { |e| ($e | get -o did | default "") != $did }) { + fail "wantedDids filter returned events for wrong DID" $instance.pid + } + print "wantedDids filter is correct" + + try { kill $instance.pid } + print "=== jetstream /subscribe test PASSED ===" +} diff --git a/tests/stream_relay_jetstream_subscribe.nu b/tests/stream_relay_jetstream_subscribe.nu new file mode 100644 index 0000000..afe3c91 --- /dev/null +++ b/tests/stream_relay_jetstream_subscribe.nu @@ -0,0 +1,212 @@ +#!/usr/bin/env nu +# Tests the Jetstream-compatible /subscribe WebSocket endpoint. +# +# Starts hydrant with the jetstream feature, connects to wss://bsky.network +# as the upstream firehose with full-network mode, and verifies that events +# arrive and are correctly formatted. Crawler and backfill are disabled so +# only the single firehose connection is used (no PDS connections). +# +# Does not require credentials — relies on the public bsky.network firehose. + +use common.nu * + +def collect-ws-json [ws_url: string, output_file: string, duration: duration] { + let pid_file = $"($output_file).pid" + bash -c $"websocat -n '($ws_url)' > '($output_file)' 2>&1 & echo $! > '($pid_file)'" + sleep $duration + let pid = (open $pid_file | str trim | into int) + try { kill $pid } + sleep 200ms + if ($output_file | path exists) { + open $output_file + | str trim + | lines + | each { |line| try { $line | from json } catch { null } } + | compact + } else { + [] + } +} + +def assert-no-error-events [events: list, label: string, ...pids: int] { + let errors = ($events | where { |e| + let t = ($e | get -o type | default "") + let k = ($e | get -o kind | default "") + $t == "error" or $k == "error" + }) + if ($errors | length) > 0 { + fail $"($label): unexpected error event(s): ($errors | first)" ...$pids + } +} + +def assert-commit-structure [c: record, label: string, ...pids: int] { + for field in ["rev", "operation", "collection", "rkey"] { + if not ($field in $c) { + fail $"($label): commit missing field ($field)" ...$pids + } + } + if $c.operation == "create" or $c.operation == "update" { + if not ("cid" in $c) { + fail $"($label): create/update commit missing cid" ...$pids + } + if not ("record" in $c) { + fail $"($label): create/update commit missing record" ...$pids + } + } +} + +def main [] { + let port = resolve-test-port 3019 + let url = $"http://localhost:($port)" + let ws_base = $"ws://127.0.0.1:($port)" + let db_path = (mktemp -d -t hydrant_stream_relay_jetstream_subscribe.XXXXXX) + + print "testing Jetstream /subscribe..." + print $"database path: ($db_path)" + + let binary = build-hydrant-relay-jetstream + let instance = (with-env { + HYDRANT_RELAY_HOSTS: "wss://bsky.network", + HYDRANT_FULL_NETWORK: "true", + HYDRANT_SEED_HOSTS: "" + } { + start-hydrant $binary $db_path $port + }) + + if not (wait-for-api $url) { + fail "api failed to start" $instance.pid + } + + # --- scenario: basic live events --- + print "--- scenario: basic live events ---" + # capture cursor_us before the live window so replay is guaranteed to cover it + let cursor_us = ((date now | format date "%s" | into int) * 1_000_000) + let live_file = $"($db_path)/live.txt" + # bsky.network has high throughput; 15 seconds gives a large event sample + # and a reasonable chance of capturing identity/account events alongside commits + let events = (collect-ws-json $"($ws_base)/subscribe" $live_file 15sec) + print $"received ($events | length) events" + assert-no-error-events $events "live subscription" $instance.pid + + let commits = ($events | where { |e| ($e | get -o kind | default "") == "commit" }) + print $"commit events: ($commits | length)" + if ($commits | length) < 10 { + fail $"expected >= 10 live commit events, got ($commits | length)" $instance.pid + } + + # validate commit event structure for a sample + for evt in ($commits | first 10) { + if not ("time_us" in $evt) { + fail "commit event missing time_us" $instance.pid + } + if not ("did" in $evt) { + fail "commit event missing did" $instance.pid + } + assert-commit-structure $evt.commit "relay commit" $instance.pid + } + print "commit event structure is correct" + + # validate identity events if any arrived (lazy RelayIdentity inflation path) + let identity_events = ($events | where { |e| ($e | get -o kind | default "") == "identity" }) + print $"identity events: ($identity_events | length)" + for evt in $identity_events { + if not ("did" in $evt) { fail "identity event missing did" $instance.pid } + if not ("time_us" in $evt) { fail "identity event missing time_us" $instance.pid } + let identity = ($evt | get -o identity | default {}) + for field in ["did", "seq", "time"] { + if not ($field in $identity) { + fail $"identity event missing identity.($field)" $instance.pid + } + } + } + if ($identity_events | length) > 0 { + print "identity event structure is correct \(lazy RelayIdentity path\)" + } + + # validate account events if any arrived (lazy RelayAccount inflation path) + let account_events = ($events | where { |e| ($e | get -o kind | default "") == "account" }) + print $"account events: ($account_events | length)" + for evt in $account_events { + if not ("did" in $evt) { fail "account event missing did" $instance.pid } + if not ("time_us" in $evt) { fail "account event missing time_us" $instance.pid } + let account = ($evt | get -o account | default {}) + for field in ["active", "did", "seq", "time"] { + if not ($field in $account) { + fail $"account event missing account.($field)" $instance.pid + } + } + } + if ($account_events | length) > 0 { + print "account event structure is correct \(lazy RelayAccount path\)" + } + + # --- scenario: cursor replay --- + print "--- scenario: cursor replay ---" + # cursor_us was captured before the live window so all live events fall after it + print $"replaying from cursor_us: ($cursor_us)" + + let replay_file = $"($db_path)/replay.txt" + let replay_events = (collect-ws-json $"($ws_base)/subscribe?cursor=($cursor_us)" $replay_file 8sec) + print $"cursor replay: ($replay_events | length) events" + assert-no-error-events $replay_events "cursor replay" $instance.pid + + let replay_commits = ($replay_events | where { |e| ($e | get -o kind | default "") == "commit" }) + if ($replay_commits | length) < 1 { + fail $"cursor replay: expected >= 1 commit event, got ($replay_commits | length)" $instance.pid + } + print $"cursor replay: ($replay_commits | length) commit events" + + # verify monotonic time_us ordering + let replay_timed = ($replay_events | where { |e| ($e | get -o time_us | default 0) > 0 }) + if ($replay_timed | length) > 1 { + let time_us_vals = ($replay_timed | get time_us) + let monotonic = ( + $time_us_vals + | zip ($time_us_vals | skip 1) + | all { |pair| $pair.0 <= $pair.1 } + ) + if not $monotonic { + fail "cursor replay events are not monotonically ordered by time_us" $instance.pid + } + print "time_us ordering is monotonically non-decreasing" + } + + # --- scenario: wantedCollections filter --- + print "--- scenario: wantedCollections filter ---" + let collection = "app.bsky.feed.post" + let col_file = $"($db_path)/col_filter.txt" + let col_url = $"($ws_base)/subscribe?cursor=($cursor_us)&wantedCollections=($collection)" + let col_events = (collect-ws-json $col_url $col_file 8sec) + let col_commits = ($col_events | where { |e| ($e | get -o kind | default "") == "commit" }) + print $"wantedCollections=($collection): ($col_commits | length) commit events" + if ($col_commits | length) < 1 { + fail "wantedCollections filter returned no commits" $instance.pid + } + let col_commits_with_data = ($col_commits | where { |e| ($e | get -o commit.collection | is-not-empty) }) + if ($col_commits_with_data | any { |e| ($e | get -o commit.collection | default "") != $collection }) { + fail "wantedCollections filter returned commits for unexpected collection" $instance.pid + } + # account/identity events pass through even with collection filter + let non_commit = ($col_events | where { |e| ($e | get -o kind | default "") != "commit" }) + print $"non-commit events through collection filter: ($non_commit | length) \(expected: account/identity pass-through\)" + print "wantedCollections filter is correct" + + # --- scenario: wantedDids filter --- + print "--- scenario: wantedDids filter ---" + # pick a DID from the live commits and replay events for it; verify the filter + # excludes all other DIDs without requiring a specific minimum event count + let sample_did = ($commits | first | get did) + print $"filtering to DID: ($sample_did)" + let did_file = $"($db_path)/did_filter.txt" + let did_url = $"($ws_base)/subscribe?cursor=($cursor_us)&wantedDids=($sample_did)" + let did_events = (collect-ws-json $did_url $did_file 8sec) + print $"wantedDids=($sample_did): ($did_events | length) events" + let did_jetstream_events = ($did_events | where { |e| ($e | get -o did | is-not-empty) }) + if ($did_jetstream_events | any { |e| ($e | get -o did | default "") != $sample_did }) { + fail "wantedDids filter returned events for wrong DID" $instance.pid + } + print "wantedDids filter is correct" + + try { kill $instance.pid } + print "=== jetstream /subscribe test PASSED ===" +} -- 2.51.2