diff --git a/src/storage/list_hosts_cursor.rs b/src/storage/list_hosts_cursor.rs new file mode 100644 index 0000000..57ac01c --- /dev/null +++ b/src/storage/list_hosts_cursor.rs @@ -0,0 +1,76 @@ +//! listHosts pagination cursor for the upstream relay. +//! +//! Tracks progress through `com.atproto.sync.listHosts` for deep-crawl mode. +//! +//! Key: `"lhs"`, value: plain string cursor + +use jacquard_common::url::Host; + +use crate::storage::{ + DbRef, PREFIX_LIST_HOSTS, + error::{StorageError, StorageResult}, +}; + +fn key(host: &Host) -> Vec { + let host = host.to_string(); + let mut k = Vec::with_capacity(PREFIX_LIST_HOSTS.len() + host.len()); + k.extend_from_slice(&PREFIX_LIST_HOSTS); + k.extend_from_slice(host.as_bytes()); + k +} + +/// Read the `listHosts` cursor for `host`. +/// +/// Returns `None` if no walk has been started yet. +pub fn get(db: &DbRef, host: &Host) -> StorageResult> { + let k = key(host); + match db.ks.get(k)? { + None => Ok(None), + Some(v) => { + let s = std::str::from_utf8(v.as_ref()).map_err(|_| StorageError::Corrupt { + key: host.to_string(), + reason: "listHosts cursor not valid UTF-8", + })?; + Ok(Some(s.to_owned())) + } + } +} + +/// Persist the `listHosts` cursor for `host`. +pub fn set(db: &DbRef, host: &Host, cursor: &str) -> StorageResult<()> { + db.ks.insert(key(host), cursor.as_bytes())?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::open_temporary; + + fn host(s: &str) -> Host { + Host::parse(s).unwrap() + } + + #[test] + fn none_when_missing() { + let db = open_temporary().unwrap(); + assert_eq!(get(&db, &host("relay.example.com")).unwrap(), None); + } + + #[test] + fn set_and_get() { + let db = open_temporary().unwrap(); + let h = host("relay.example.com"); + set(&db, &h, "page2cursor").unwrap(); + assert_eq!(get(&db, &h).unwrap(), Some("page2cursor".to_owned())); + } + + #[test] + fn overwrite_cursor() { + let db = open_temporary().unwrap(); + let h = host("relay.example.com"); + set(&db, &h, "first").unwrap(); + set(&db, &h, "second").unwrap(); + assert_eq!(get(&db, &h).unwrap(), Some("second".to_owned())); + } +} diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 4aef601..2c63a3b 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -2,6 +2,8 @@ pub mod backfill_progress; pub mod collection_index; pub mod error; pub mod firehose_cursor; +pub mod list_hosts_cursor; +pub mod pds_host; pub mod repo; pub mod resync_buffer; pub mod resync_queue; @@ -35,6 +37,10 @@ pub(crate) const PREFIX_LIST_REPOS: KeyPrefix = *b"lsr"; pub(crate) const PREFIX_RESYNC_QUEUE: KeyPrefix = *b"rsq"; /// Per-repo buffered firehose events during resync. See [`resync_buffer`]. pub(crate) const PREFIX_RESYNC_BUFFER: KeyPrefix = *b"rsb"; +/// Per-PDS host state (sync1.1 mode, trust, listRepos cursor/done). See [`pds_host`]. +pub(crate) const PREFIX_PDS_HOST: KeyPrefix = *b"pdh"; +/// listHosts walk cursor (per upstream relay host). See [`list_hosts_cursor`]. +pub(crate) const PREFIX_LIST_HOSTS: KeyPrefix = *b"lhs"; use std::path::Path; use std::sync::Arc; diff --git a/src/storage/pds_host.rs b/src/storage/pds_host.rs new file mode 100644 index 0000000..24aad7e --- /dev/null +++ b/src/storage/pds_host.rs @@ -0,0 +1,259 @@ +//! Per-PDS-host state storage +//! +//! Key: `"pdh"`, value: `\0\0\0` + +use jacquard_common::url::Host; + +use crate::storage::{ + DbRef, PREFIX_PDS_HOST, + error::{StorageError, StorageResult}, +}; + +// --------------------------------------------------------------------------- +// Key encoding +// --------------------------------------------------------------------------- + +fn key(host: &Host) -> Vec { + let host = host.to_string(); + let mut k = Vec::with_capacity(PREFIX_PDS_HOST.len() + host.len()); + k.extend_from_slice(&PREFIX_PDS_HOST); + k.extend_from_slice(host.as_bytes()); + k +} + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +/// Whether this host has been seen sending sync1.1-style events. +/// +/// Hosts start in `Lenient` mode and flip to `Strict` once any event with +/// `prevData` is observed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Sync11Mode { + Lenient, + Strict, +} + +impl Sync11Mode { + fn as_str(&self) -> &'static str { + match self { + Sync11Mode::Lenient => "lenient", + Sync11Mode::Strict => "strict", + } + } + + fn from_str(s: &str) -> Option { + Some(match s { + "lenient" => Sync11Mode::Lenient, + "strict" => Sync11Mode::Strict, + _ => return None, + }) + } +} + +/// Trust level for a PDS host. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PdsTrust { + Ok, + Cautious, + Blocked, +} + +impl PdsTrust { + fn as_str(&self) -> &'static str { + match self { + PdsTrust::Ok => "ok", + PdsTrust::Cautious => "cautious", + PdsTrust::Blocked => "blocked", + } + } + + fn from_str(s: &str) -> Option { + Some(match s { + "ok" => PdsTrust::Ok, + "cautious" => PdsTrust::Cautious, + "blocked" => PdsTrust::Blocked, + _ => return None, + }) + } +} + +/// Stored state for a PDS host. +#[derive(Debug, Clone)] +pub struct PdsHostInfo { + pub sync11_mode: Sync11Mode, + pub trust: PdsTrust, + /// Whether the `listRepos` walk for this PDS has completed. + pub list_repos_done: bool, + /// Current `listRepos` pagination cursor; empty when not yet started. + pub list_repos_cursor: String, +} + +impl Default for PdsHostInfo { + fn default() -> Self { + PdsHostInfo { + sync11_mode: Sync11Mode::Lenient, + trust: PdsTrust::Ok, + list_repos_done: false, + list_repos_cursor: String::new(), + } + } +} + +// --------------------------------------------------------------------------- +// Encode / decode +// --------------------------------------------------------------------------- + +/// Wire format: `\0\0\0` +fn encode(info: &PdsHostInfo) -> Vec { + let done = if info.list_repos_done { "1" } else { "0" }; + let s = format!( + "{}\0{}\0{}\0{}", + info.sync11_mode.as_str(), + info.trust.as_str(), + done, + info.list_repos_cursor, + ); + s.into_bytes() +} + +fn decode(bytes: &[u8], key: &str) -> StorageResult { + let s = std::str::from_utf8(bytes).map_err(|_| StorageError::Corrupt { + key: key.to_owned(), + reason: "pds host info not valid UTF-8", + })?; + let mut parts = s.splitn(4, '\0'); + let sync11_mode = parts + .next() + .and_then(Sync11Mode::from_str) + .ok_or(StorageError::Corrupt { + key: key.to_owned(), + reason: "invalid sync11_mode", + })?; + let trust = parts + .next() + .and_then(PdsTrust::from_str) + .ok_or(StorageError::Corrupt { + key: key.to_owned(), + reason: "invalid trust", + })?; + let list_repos_done = match parts.next() { + Some("0") => false, + Some("1") => true, + _ => { + return Err(StorageError::Corrupt { + key: key.to_owned(), + reason: "invalid list_repos_done", + }); + } + }; + let list_repos_cursor = parts.next().unwrap_or("").to_owned(); + Ok(PdsHostInfo { + sync11_mode, + trust, + list_repos_done, + list_repos_cursor, + }) +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +/// Read the [`PdsHostInfo`] for `host`. +/// +/// Returns the default (lenient, ok, not done, empty cursor) when the key is +/// absent — matching the implicit state for any host we haven't seen yet. +pub fn get(db: &DbRef, host: &Host) -> StorageResult { + let k = key(host); + match db.ks.get(&k)? { + None => Ok(PdsHostInfo::default()), + Some(v) => { + let key_str = String::from_utf8_lossy(&k).into_owned(); + decode(v.as_ref(), &key_str) + } + } +} + +/// Add a [`PdsHostInfo`] write to an existing batch. +pub fn put_into(batch: &mut fjall::OwnedWriteBatch, db: &DbRef, host: &Host, info: &PdsHostInfo) { + batch.insert(&db.ks, key(host), encode(info)); +} + +/// Write a [`PdsHostInfo`] for `host`. +pub fn put(db: &DbRef, host: &Host, info: &PdsHostInfo) -> StorageResult<()> { + let mut batch = db.database.batch(); + put_into(&mut batch, db, host, info); + batch.commit()?; + Ok(()) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::open_temporary; + + fn host(s: &str) -> Host { + Host::parse(s).unwrap() + } + + #[test] + fn default_when_missing() { + let db = open_temporary().unwrap(); + let info = get(&db, &host("pds.example.com")).unwrap(); + assert_eq!(info.sync11_mode, Sync11Mode::Lenient); + assert_eq!(info.trust, PdsTrust::Ok); + assert!(!info.list_repos_done); + assert!(info.list_repos_cursor.is_empty()); + } + + #[test] + fn roundtrip_default() { + let info = PdsHostInfo::default(); + let encoded = encode(&info); + let decoded = decode(&encoded, "test").unwrap(); + assert_eq!(decoded.sync11_mode, Sync11Mode::Lenient); + assert_eq!(decoded.trust, PdsTrust::Ok); + assert!(!decoded.list_repos_done); + assert!(decoded.list_repos_cursor.is_empty()); + } + + #[test] + fn roundtrip_strict_blocked_done_with_cursor() { + let info = PdsHostInfo { + sync11_mode: Sync11Mode::Strict, + trust: PdsTrust::Blocked, + list_repos_done: true, + list_repos_cursor: "abc123".to_owned(), + }; + let encoded = encode(&info); + let decoded = decode(&encoded, "test").unwrap(); + assert_eq!(decoded.sync11_mode, Sync11Mode::Strict); + assert_eq!(decoded.trust, PdsTrust::Blocked); + assert!(decoded.list_repos_done); + assert_eq!(decoded.list_repos_cursor, "abc123"); + } + + #[test] + fn put_and_get() { + let db = open_temporary().unwrap(); + let h = host("pds.example.com"); + let info = PdsHostInfo { + sync11_mode: Sync11Mode::Strict, + trust: PdsTrust::Cautious, + list_repos_done: false, + list_repos_cursor: "cursor42".to_owned(), + }; + put(&db, &h, &info).unwrap(); + let retrieved = get(&db, &h).unwrap(); + assert_eq!(retrieved.sync11_mode, Sync11Mode::Strict); + assert_eq!(retrieved.trust, PdsTrust::Cautious); + assert!(!retrieved.list_repos_done); + assert_eq!(retrieved.list_repos_cursor, "cursor42"); + } +}