From a21adcf931e07bb49eafe4632e1d6e5529cb55de Mon Sep 17 00:00:00 2001 From: dawn <90008@gaze.systems> Date: Fri, 13 Mar 2026 02:27:12 +0300 Subject: [PATCH] fjall: implement /export/stream --- Cargo.lock | 60 ++++++++++++++++++- Cargo.toml | 2 +- src/mirror/fjall.rs | 139 ++++++++++++++++++++++++++++++++++++++++---- src/plc_fjall.rs | 66 ++++++++++++++++----- 4 files changed, 239 insertions(+), 28 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f9ce0a6..a16d63a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -60,7 +60,7 @@ dependencies = [ "tokio", "tokio-postgres", "tokio-stream", - "tokio-tungstenite", + "tokio-tungstenite 0.26.2", "tokio-util", "tracing", "tracing-opentelemetry", @@ -1104,6 +1104,19 @@ dependencies = [ "slab", ] +[[package]] +name = "futures_codec" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad004dd81205978a2bba6c566ed70535ccf88c0be34649e628186474603f43ca" +dependencies = [ + "bytes", + "futures-sink", + "futures-util", + "memchr", + "pin-project-lite", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -2329,10 +2342,13 @@ dependencies = [ "serde_json", "serde_urlencoded", "smallvec", + "sse-codec", "sync_wrapper", "thiserror 2.0.18", "tokio", "tokio-rustls", + "tokio-stream", + "tokio-tungstenite 0.27.0", "tokio-util", "tracing", "wildmatch", @@ -3252,6 +3268,17 @@ dependencies = [ "der", ] +[[package]] +name = "sse-codec" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a395a858c7ff5c4b42aeab0501e07c978ac5e1ae5059f301884dab3fa405f47" +dependencies = [ + "futures-io", + "futures_codec", + "memchr", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -3557,7 +3584,19 @@ dependencies = [ "native-tls", "tokio", "tokio-native-tls", - "tungstenite", + "tungstenite 0.26.2", +] + +[[package]] +name = "tokio-tungstenite" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "489a59b6730eda1b0171fcfda8b121f4bee2b35cba8645ca35c5f7ba3eb736c1" +dependencies = [ + "futures-util", + "log", + "tokio", + "tungstenite 0.27.0", ] [[package]] @@ -3778,6 +3817,23 @@ dependencies = [ "utf-8", ] +[[package]] +name = "tungstenite" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eadc29d668c91fcc564941132e17b28a7ceb2f3ebf0b9dae3e03fd7a6748eb0d" +dependencies = [ + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.9.2", + "sha1", + "thiserror 2.0.18", + "utf-8", +] + [[package]] name = "twox-hash" version = "2.1.2" diff --git a/Cargo.toml b/Cargo.toml index a075d5d..71cfce4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,7 +27,7 @@ native-tls = "0.2.14" opentelemetry = "0.30.0" opentelemetry-otlp = { version = "0.30.0" } opentelemetry_sdk = { version = "0.30.0", features = ["rt-tokio"] } -poem = { version = "3.1.12", features = ["acme", "compression"] } +poem = { version = "3.1.12", features = ["acme", "compression", "test", "websocket"] } postgres-native-tls = "0.5.1" reqwest = { version = "0.12.23", features = ["stream", "json", "gzip"] } reqwest-middleware = "0.4.2" diff --git a/src/mirror/fjall.rs b/src/mirror/fjall.rs index 37864f6..420ee5b 100644 --- a/src/mirror/fjall.rs +++ b/src/mirror/fjall.rs @@ -1,8 +1,23 @@ +use std::sync::Arc; + use super::*; -use futures::StreamExt; +use futures::{SinkExt as _, StreamExt as _}; +use poem::IntoResponse; use poem::web::Query; use serde::Deserialize; +async fn spawn_blocking(f: F) -> poem::Result +where + R: Send + 'static, + F: FnOnce() -> anyhow::Result + Send + 'static, +{ + tokio::task::spawn_blocking(f) + .await + .map_err(anyhow::Error::from) + .flatten() + .map_err(|e| Error::from_string(e.to_string(), StatusCode::INTERNAL_SERVER_ERROR)) +} + #[derive(Clone)] struct FjallState { client: Client, @@ -145,13 +160,11 @@ async fn fjall_resolve(req: &Request, Data(state): Data<&FjallState>) -> Result< let did = did.to_string(); let db = state.fjall.clone(); - let ops = tokio::task::spawn_blocking(move || { + let ops = spawn_blocking(move || { let iter = db.ops_for_did(&did)?; iter.collect::>>() }) - .await - .map_err(|e| Error::from_string(e.to_string(), StatusCode::INTERNAL_SERVER_ERROR))? - .map_err(|e| Error::from_string(e.to_string(), StatusCode::INTERNAL_SERVER_ERROR))?; + .await?; if ops.is_empty() { return Err(Error::from_string( @@ -256,7 +269,7 @@ struct ExportQuery { } #[handler] -async fn fjall_export( +async fn export( _req: &Request, Query(query): Query, Data(FjallState { fjall, .. }): Data<&FjallState>, @@ -265,13 +278,11 @@ async fn fjall_export( let limit = 1000; let db = fjall.clone(); - let ops = tokio::task::spawn_blocking(move || { + let ops = spawn_blocking(move || { let iter = db.export_ops(after..)?; iter.take(limit).collect::>>() }) - .await - .map_err(|e| Error::from_string(e.to_string(), StatusCode::INTERNAL_SERVER_ERROR))? - .map_err(|e| Error::from_string(e.to_string(), StatusCode::INTERNAL_SERVER_ERROR))?; + .await?; let stream = futures::stream::iter(ops).map(|op| { let mut json = serde_json::to_string(&op).unwrap(); @@ -282,6 +293,111 @@ async fn fjall_export( Ok(Body::from_bytes_stream(stream)) } +#[derive(Deserialize)] +struct StreamQuery { + cursor: Option, +} + +#[handler] +async fn export_stream( + _req: &Request, + Query(query): Query, + ws: poem::web::websocket::WebSocket, + Data(FjallState { fjall, .. }): Data<&FjallState>, +) -> poem::Result { + use poem::web::websocket::Message; + use tokio::sync::Notify; + + let db = fjall.clone(); + + let latest_cursor = spawn_blocking({ + let db = db.clone(); + move || db.get_latest().map(|res| res.map(|(c, _)| c)) + }) + .await? + .unwrap_or(0); + + let mut cursor = match query.cursor { + Some(cursor) => { + if cursor > latest_cursor { + return Err(Error::from_string( + format!("cursor {cursor} is in the future"), + StatusCode::BAD_REQUEST, + )); + } + + let created_at = spawn_blocking({ + let db = db.clone(); + move || { + db.get_op_at_or_after(cursor) + .map(|res| res.map(|op| op.created_at)) + } + }) + .await?; + + match created_at { + Some(created_at) => { + // check that the provided cursor is not stale + if (chrono::Utc::now() - created_at).num_days() > 1 { + return Err(Error::from_string( + format!("cursor {cursor} is stale, catch up using /export first"), + StatusCode::BAD_REQUEST, + )); + } + cursor + } + None => latest_cursor, + } + } + None => { + // if cursor is not provided, start at the latest op + latest_cursor + } + }; + + Ok(ws.on_upgrade(move |mut socket| async move { + let errored = Arc::new(Notify::new()); + + loop { + let (tx, mut op_rx) = tokio::sync::mpsc::channel(64); + + tokio::task::spawn_blocking({ + let db = db.clone(); + let errored = errored.clone(); + move || { + let iter = match db.export_ops(cursor..) { + Ok(it) => it, + Err(e) => { + log::error!("read failed: {e}"); + errored.notify_one(); + return; + } + }; + for op in iter.flatten() { + if tx.blocking_send(op).is_err() { + return; + } + } + } + }); + + while let Some(op) = op_rx.recv().await { + cursor = op.seq; + let json = serde_json::to_string(&op).unwrap(); + if let Err(e) = socket.send(Message::Text(json)).await { + log::warn!("closing export stream: {e}"); + return; + } + } + + tokio::select! { + _ = db.subscribe() => {}, + _ = errored.notified() => return, + } + } + })) +} + #[handler] async fn fjall_nope(Data(FjallState { upstream, .. }): Data<&FjallState>) -> (StatusCode, String) { ( @@ -334,7 +450,8 @@ pub async fn serve_fjall( .at("/", get(fjall_hello)) .at("/favicon.ico", get(favicon)) .at("/_health", get(fjall_health)) - .at("/export", get(fjall_export)); + .at("/export", get(export)) + .at("/export/stream", get(export_stream)); if experimental.write_upstream { log::info!("enabling experimental write forwarding to upstream"); diff --git a/src/plc_fjall.rs b/src/plc_fjall.rs index df4479b..21a56d2 100644 --- a/src/plc_fjall.rs +++ b/src/plc_fjall.rs @@ -12,7 +12,7 @@ use std::fmt; use std::path::Path; use std::sync::Arc; use std::time::Instant; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{Notify, futures::Notified, mpsc, oneshot}; const SEP: u8 = 0; @@ -829,8 +829,9 @@ struct DbOp { } // we have our own Op struct for fjall since we dont want to have to convert Value back to RawValue -#[derive(Debug, Serialize)] +#[derive(Debug, Serialize, Deserialize, Clone)] pub struct Op { + pub seq: u64, pub did: String, pub cid: String, pub created_at: Dt, @@ -849,6 +850,7 @@ struct FjallInner { ops: Keyspace, /// secondary index: [encoded_did][SEP][seq_varint] -> [] by_did: Keyspace, + notify_stream: Notify, } impl FjallDb { @@ -880,12 +882,17 @@ impl FjallDb { let by_did = db.keyspace("by_did", || { opts() .max_memtable_size(mb(64)) - // this isn't gonna compress well anyway, since its just keys (did + timestamp + cid) + // this isn't gonna compress well anyway, since its just keys (did + seq) // and dids dont have many operations in the first place, so we can use small blocks .data_block_size_policy(BlockSizePolicy::all(kb(2))) })?; Ok(Self { - inner: Arc::new(FjallInner { db, ops, by_did }), + inner: Arc::new(FjallInner { + db, + ops, + by_did, + notify_stream: Notify::new(), + }), }) } @@ -905,6 +912,10 @@ impl FjallDb { Ok(()) } + pub fn subscribe(&self) -> Notified<'_> { + self.inner.notify_stream.notified() + } + /// Returns `(seq, created_at)` for the last stored op, or `None` if empty. pub fn get_latest(&self) -> anyhow::Result> { let Some(guard) = self.inner.ops.last_key_value() else { @@ -950,7 +961,7 @@ impl FjallDb { self._ops_for_did(&op.did) .map(|ops| { ops.rev() - .find(|r| r.as_ref().map_or(true, |(_, cid, _)| cid == prev_cid)) + .find(|r| r.as_ref().map_or(true, |(_, _, cid, _)| cid == prev_cid)) .transpose() }) .flatten() @@ -958,7 +969,7 @@ impl FjallDb { .transpose()? .flatten(); - let prev_stored = prev_op.as_ref().map(|(_, _, p)| &p.operation); + let prev_stored = prev_op.as_ref().map(|(_, _, _, p)| &p.operation); match verify_op_sig(&operation, prev_stored) { Ok(results) => { @@ -1002,18 +1013,42 @@ impl FjallDb { batch.insert(&self.inner.by_did, by_did_key_bytes, &[]); batch.commit()?; + self.inner.notify_stream.notify_waiters(); + Ok(1) } -} -impl FjallDb { + pub(crate) fn get_op_at_or_after(&self, seq: u64) -> anyhow::Result> { + self.inner + .ops + .range(seq_key(seq)..) + .next() + .map(|v| { + rmp_serde::from_slice::(&v.value()?) + .context("failed to decode op") + .map(|op| { + Ok(Op { + seq, + did: decode_did(&op.did), + cid: decode_cid(&op.cid)?, + created_at: Dt::from_timestamp_micros(op.created_at as i64) + .ok_or_else(|| anyhow::anyhow!("invalid created_at in op"))?, + nullified: op.nullified, + operation: op.operation.to_json_value(), + }) + }) + .flatten() + }) + .transpose() + } + /// Decode a `by_did` entry: extract the seq from the key suffix, then /// look up the full `DbOp` in the `ops` keyspace. fn decode_by_did_entry( &self, by_did_key_bytes: &[u8], prefix_len: usize, - ) -> anyhow::Result<(Dt, PlcCid, DbOp)> { + ) -> anyhow::Result<(u64, Dt, PlcCid, DbOp)> { let key_suffix = by_did_key_bytes .get(prefix_len..) .ok_or_else(|| anyhow::anyhow!("invalid by_did key {by_did_key_bytes:?}"))?; @@ -1032,13 +1067,13 @@ impl FjallDb { .ok_or_else(|| anyhow::anyhow!("invalid created_at_micros {}", op.created_at))?; let cid = PlcCid(op.cid.clone()); - Ok((ts, cid, op)) + Ok((seq, ts, cid, op)) } fn _ops_for_did( &self, did: &str, - ) -> anyhow::Result> + '_> + ) -> anyhow::Result> + '_> { let prefix = by_did_prefix(did)?; @@ -1055,10 +1090,11 @@ impl FjallDb { did: &str, ) -> anyhow::Result> + '_> { Ok(self._ops_for_did(did)?.map(|res| { - let (ts, cid, op) = res?; + let (seq, ts, cid, op) = res?; let cid = decode_cid(&cid.0)?; let did = decode_did(&op.did); Ok(Op { + seq, did, cid, created_at: ts, @@ -1087,9 +1123,10 @@ impl FjallDb { .ops .range(range) .map(|item| -> anyhow::Result { - let (_, value) = item + let (key, value) = item .into_inner() .map_err(|e: fjall::Error| anyhow::anyhow!("fjall read error: {e}"))?; + let seq = decode_seq_key(&key)?; let db_op: DbOp = rmp_serde::from_slice(&value)?; let created_at = Dt::from_timestamp_micros(db_op.created_at as i64).ok_or_else(|| { @@ -1098,6 +1135,7 @@ impl FjallDb { let cid = decode_cid(&db_op.cid)?; let did = decode_did(&db_op.did); Ok(Op { + seq, did, cid, created_at, @@ -1257,7 +1295,7 @@ impl FjallDb { current_prefix = Some(prefix_array); } - did_ops.push(op); + did_ops.push((op.1, op.2, op.3)); processed_ops += 1; } -- 2.51.2