diff --git a/Cargo.lock b/Cargo.lock index 41a926a..552d98c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1438,6 +1438,7 @@ dependencies = [ "clap", "futures-util", "governor", + "hubble-pds", "ipld-core", "jacquard-api", "jacquard-common", @@ -1451,6 +1452,7 @@ dependencies = [ "serde_json", "star-lite", "star-repo", + "tempfile", "thiserror 2.0.18", "tokio", "tokio-util", diff --git a/mini/Cargo.toml b/mini/Cargo.toml index 6ffce5c..7ace9b4 100644 --- a/mini/Cargo.toml +++ b/mini/Cargo.toml @@ -9,6 +9,7 @@ async-compression = { version = "0.4", features = ["tokio", "gzip", "brotli", "z clap = { workspace = true } futures-util = "0.3" governor = "0.10.4" +hubble-pds = { workspace = true } ipld-core = "0.4" jacquard-api = { workspace = true } jacquard-common = { workspace = true, features = ["streaming"] } @@ -27,3 +28,6 @@ tokio = { workspace = true, features = ["full"] } tokio-util = { version = "0.7", features = ["io"] } tracing = { workspace = true } tracing-subscriber = { workspace = true } + +[dev-dependencies] +tempfile = "3" diff --git a/mini/src/check_repo.rs b/mini/src/check_repo.rs new file mode 100644 index 0000000..189b61e --- /dev/null +++ b/mini/src/check_repo.rs @@ -0,0 +1,96 @@ +//! `check-repo` subcommand: assert MST integrity of a stored repo. +//! +//! Reads `C\0` for the expected `data` CID. Range-scans +//! `r\0\0...` records in rocksdb order (which equals MST lex +//! order thanks to the `/` separator we picked in `storage/record.rs`). +//! Feeds each `(mst_key, record_bytes)` into +//! `star_lite::mst::MstStack`. Compares the computed +//! root to the stored commit's `data` CID. +//! +//! Authenticity (signature) verification is out of scope here — that +//! needs DID resolution to recover the signing key, which mini +//! deliberately doesn't do. The stored `StoredCommit.sig` field is +//! preserved for whatever does want to do that later. + +use std::str::FromStr; + +use hubble_pds::Hostname; +use ipld_core::cid::Cid; +use star_lite::mst::{BuilderOptions, MstStack, VerifyBackend}; + +use crate::storage::{Db, DbError, commit, record}; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("storage: {0}")] + Storage(#[from] DbError), + + #[error("rocksdb iter: {0}")] + Rocks(#[from] rocksdb::Error), + + #[error("no commit stored for {pds}\\0{did} — run snapshot-repos first")] + NoCommit { pds: String, did: String }, + + #[error("stored commit data CID is not a valid CID: {raw:?}: {reason}")] + BadStoredCid { raw: String, reason: String }, + + #[error("mst build: {0}")] + MstBuild(#[from] star_lite::Error), + + /// Iterator key didn't begin with `r\0\0` — shouldn't + /// happen given the upper/lower bounds, but checked anyway so a + /// malformed key surfaces here instead of corrupting verification. + #[error("malformed record key (bad prefix)")] + BadRecordKey, +} + +#[derive(Debug)] +pub struct Outcome { + pub matched: bool, + pub expected_root: Cid, + pub computed_root: Cid, + pub record_count: u64, +} + +pub fn run(db: &Db, pds: &Hostname, did: &str) -> Result { + let stored = commit::get(db, pds, did)?.ok_or_else(|| Error::NoCommit { + pds: pds.as_ref().to_string(), + did: did.to_string(), + })?; + let expected_root: Cid = Cid::from_str(&stored.data).map_err(|e| Error::BadStoredCid { + raw: stored.data.clone(), + reason: e.to_string(), + })?; + + let cf = db.cf_records()?; + let (lo, hi) = record::range_for_repo(pds, did); + // We need to strip the `r\0\0` prefix from each key to get + // the bare MST key. The lower bound is exactly that prefix. + let prefix_len = lo.len(); + let mut read_opts = rocksdb::ReadOptions::default(); + read_opts.set_iterate_lower_bound(lo); + read_opts.set_iterate_upper_bound(hi); + let iter = db + .inner + .iterator_cf_opt(cf, read_opts, rocksdb::IteratorMode::Start); + + let mut stack = MstStack::new(VerifyBackend, BuilderOptions::default()); + let mut record_count: u64 = 0; + for item in iter { + let (k, v) = item?; + if k.len() < prefix_len { + return Err(Error::BadRecordKey); + } + let mst_key = &k[prefix_len..]; + stack.insert(mst_key, &v)?; + record_count += 1; + } + + let (computed_root, _backend) = stack.finish()?; + Ok(Outcome { + matched: computed_root == expected_root, + expected_root, + computed_root, + record_count, + }) +} diff --git a/mini/src/find_pdses.rs b/mini/src/find_pdses.rs new file mode 100644 index 0000000..91a6dcb --- /dev/null +++ b/mini/src/find_pdses.rs @@ -0,0 +1,176 @@ +//! Discover PDS hosts: relay `listHosts` plus (optionally) mary's +//! atproto-scraping `instances.json`. Filters out `*.host.bsky.network` +//! at the boundary — those hosts are never inserted. + +use std::collections::BTreeMap; + +use hubble_pds::Hostname; +use jacquard_api::com_atproto::sync::list_hosts::ListHosts; +use jacquard_common::deps::fluent_uri::Uri; +use jacquard_common::deps::smol_str::SmolStr; +use jacquard_common::xrpc::XrpcExt; +use reqwest::Url; +use serde::Deserialize; + +use crate::storage::{Db, host}; + +/// listHosts page size. The lexicon caps this at 1000. +const LIST_HOSTS_PAGE: i64 = 1000; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("storage: {0}")] + Storage(#[from] crate::storage::DbError), + + #[error("http: {0}")] + Http(#[from] reqwest::Error), + + /// `jacquard`'s `ClientError` is large; box it. + #[error("xrpc: {0}")] + Xrpc(Box), + + /// Response decode error from `parse::()`. The inner `XrpcError` is + /// generic-on-error-type, so we collapse to a string at this boundary. + #[error("xrpc decode: {0}")] + XrpcDecode(String), + + #[error("scraping json: {0}")] + ScrapingJson(#[from] serde_json::Error), +} + +impl From for Error { + fn from(e: jacquard_common::error::ClientError) -> Self { + Self::Xrpc(Box::new(e)) + } +} + +pub async fn run( + db: &Db, + relay: &Url, + scraping_url: Option<&Url>, + ua: &str, +) -> Result<(), Error> { + let client = reqwest::Client::builder() + .user_agent(ua) + .pool_max_idle_per_host(2) + .pool_idle_timeout(std::time::Duration::from_secs(30)) + .build()?; + + get_from_relay(db, &client, relay).await?; + + if let Some(url) = scraping_url { + get_from_scraping(db, &client, url).await?; + } + + Ok(()) +} + +async fn get_from_relay(db: &Db, client: &reqwest::Client, relay: &Url) -> Result<(), Error> { + let mut cursor: Option = None; + let mut total: u64 = 0; + let mut skipped_bsky: u64 = 0; + + tracing::info!(%relay, "listHosts: starting"); + loop { + // url::Url is always a valid RFC 3986 URI, so this never fails. + let base = Uri::parse(relay.as_str()).expect("url::Url parses as fluent_uri::Uri"); + + let request = ListHosts::::new() + .maybe_cursor(cursor.clone()) + .limit(Some(LIST_HOSTS_PAGE)) + .build(); + + let response = client.xrpc(base).send(&request).await?; + let output = response + .parse::() + .map_err(|e| Error::XrpcDecode(e.to_string()))?; + + for entry in &output.hosts { + let hostname = match Hostname::new(entry.hostname.as_ref()) { + Ok(h) => h, + Err(e) => { + tracing::warn!(%e, "listHosts: skipping unparseable hostname"); + continue; + } + }; + if hostname.is_bsky() { + skipped_bsky += 1; + continue; + } + let status = entry.status.as_ref().map(|s| s.as_str()); + let count = entry + .account_count + .map(|c| u64::try_from(c).expect("non-negative pds account count from relay")); + host::upsert_from_listhosts(db, &hostname, status, count)?; + total += 1; + } + + match output.cursor { + Some(c) if !c.is_empty() => { + tracing::debug!(total, "listHosts page"); + cursor = Some(c); + } + _ => break, + } + } + + tracing::info!(total, skipped_bsky, "listHosts: done"); + Ok(()) +} + +/// Top-level shape of mary's `instances.json` (and any compatible scrape). +#[derive(Deserialize)] +struct ScrapingHosts { + pdses: BTreeMap, +} + +#[derive(Deserialize)] +#[allow(dead_code)] +struct ScrapingPds { + status: Option, + version: Option, + invite_code_required: Option, +} + +async fn get_from_scraping(db: &Db, client: &reqwest::Client, url: &Url) -> Result<(), Error> { + tracing::info!(%url, "scraping: fetching"); + let bytes = client + .get(url.clone()) + .send() + .await? + .error_for_status()? + .bytes() + .await?; + let doc: ScrapingHosts = serde_json::from_slice(&bytes)?; + + let mut total: u64 = 0; + let mut skipped_unparseable: u64 = 0; + let mut skipped_bsky: u64 = 0; + for raw in doc.pdses.keys() { + let parsed = match Url::parse(raw) { + Ok(u) => u, + Err(e) => { + tracing::warn!(%raw, %e, "scraping: skipping unparseable url"); + skipped_unparseable += 1; + continue; + } + }; + let hostname = match Hostname::from_url(&parsed) { + Ok(h) => h, + Err(e) => { + tracing::warn!(%raw, %e, "scraping: skipping url with no usable host"); + skipped_unparseable += 1; + continue; + } + }; + if hostname.is_bsky() { + skipped_bsky += 1; + continue; + } + host::upsert_from_scraping(db, &hostname)?; + total += 1; + } + + tracing::info!(total, skipped_unparseable, skipped_bsky, "scraping: done"); + Ok(()) +} diff --git a/mini/src/main.rs b/mini/src/main.rs index f328e4d..98989c5 100644 --- a/mini/src/main.rs +++ b/mini/src/main.rs @@ -1 +1,245 @@ -fn main() {} +//! `hubble-mini` — rolling snapshots of non-bluesky PDSes into RocksDB. +//! +//! See `local/mini-proposal.md` for the design. + +mod check_repo; +mod find_pdses; +mod pds; +mod repo; +mod snapshot_repos; +mod storage; + +use std::path::PathBuf; +use std::sync::Arc; + +use clap::{Args, Parser, Subcommand}; +use hubble_pds::Hostname; +use reqwest::Url; +use tracing_subscriber::EnvFilter; + +const DEFAULT_RELAY: &str = "https://bsky.network"; +const DEFAULT_MARY_URL: &str = "https://raw.githubusercontent.com/mary-ext/atproto-scraping/refs/heads/trunk/dist/instances.json"; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("storage: {0}")] + Storage(#[from] storage::DbError), + + #[error("find_pdses: {0}")] + FindPdses(#[from] find_pdses::Error), + + #[error("snapshot_repos: {0}")] + SnapshotRepos(#[from] snapshot_repos::Error), + + #[error("check_repo: {0}")] + CheckRepo(#[from] check_repo::Error), +} + +#[derive(Parser, Debug)] +#[command(name = "hubble-mini", version, about)] +struct Cli { + /// Path to the rocksdb database directory. + #[arg(long, global = true, default_value = "./local/mini-db")] + db: PathBuf, + + /// You're about to send a lot of requests: who are you? (appears in + /// User-Agent header of outbound requests) + /// + /// Must be passed before the subcommand name. (Not `global = true` + /// because clap's global+required combo requires the value at every + /// subcommand level — passing it once at the top is enough.) + #[arg(long)] + contact: String, + + #[command(subcommand)] + command: Command, +} + +#[derive(Subcommand, Debug)] +enum Command { + /// Find PDS hosts via relay listHosts and (optionally) mary's + /// atproto-scraping instances.json. Filters out `*.host.bsky.network`. + FindPdses(FindPdsesArgs), + + /// Add PDS hosts explicitly (no network calls). + AddHosts(AddHostsArgs), + + /// Walk known hosts and snapshot every repo into rocksdb. + SnapshotRepos(SnapshotReposArgs), + + /// Re-verify the MST of one stored repo against its stored commit. + /// Exits 0 on match, 1 on mismatch. + CheckRepo(CheckRepoArgs), + + /// `find-pdses` then `snapshot-repos` in one run. + Auto { + #[command(flatten)] + find_pdses: FindPdsesArgs, + #[command(flatten)] + snapshot_repos: SnapshotReposArgs, + }, +} + +#[derive(Args, Debug)] +struct FindPdsesArgs { + /// target for com.atproto.sync.listHosts + #[arg(long, default_value = DEFAULT_RELAY)] + relay: Url, + + /// also pull pds hosts from an `atproto-scraping`-dist-style json file. + /// optionally provide a url to override the default source (mary's). + #[arg(long, num_args = 0..=1, default_missing_value = DEFAULT_MARY_URL)] + with_scraping: Option, +} + +#[derive(Args, Debug)] +struct AddHostsArgs { + /// hostname of a pds to track (repeatable) + #[arg(long, num_args = 1..)] + host: Vec, +} + +#[derive(Args, Debug)] +struct CheckRepoArgs { + /// PDS hostname + #[arg(long)] + pds: Hostname, + + /// DID of the repo to check + #[arg(long)] + did: String, +} + +#[derive(Args, Debug)] +struct SnapshotReposArgs { + /// number of pds hosts to work on at once + #[arg(long, default_value_t = 100)] + host_workers_limit: usize, + + /// number of requests-in-flight per pds + #[arg(long, default_value_t = 10)] + per_pds_limit: usize, + + /// concurrency limit for big repos (prevent memory spikes during + /// tier-2 in-memory CAR loads, ~150 MiB each) + #[arg(long, default_value_t = 5)] + big_repos_limit: usize, + + /// directory for tier-3 disk spilling. each spill creates a fresh + /// sub-directory here (cleaned up on drop). avoid `/tmp` on linux — + /// it's often tmpfs (RAM-backed), which defeats the purpose. + #[arg(long, default_value = "./local/spill")] + spill_dir: PathBuf, + + /// skip these PDS hosts entirely — never picked as candidates. + /// Repeatable: `--skip foo.example --skip bar.example`. + #[arg(long)] + skip: Vec, +} + +#[tokio::main] +async fn main() -> Result<(), Error> { + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")), + ) + .init(); + let cli = Cli::parse(); + tracing::debug!(db = ?cli.db, command = ?cli.command, "parsed cli"); + + let db = Arc::new(storage::Db::open(&cli.db)?); + + let ua = format!( + "hubble-mini v{}, from @microcosm.blue. contact: {}", + env!("CARGO_PKG_VERSION"), + cli.contact, + ); + + match cli.command { + Command::FindPdses(args) => { + find_pdses::run(&db, &args.relay, args.with_scraping.as_ref(), &ua).await?; + } + Command::AddHosts(args) => { + run_add_hosts(&db, &args)?; + } + Command::SnapshotRepos(args) => { + snapshot_repos::run(db, snapshot_repos_args(args), &ua).await?; + } + Command::CheckRepo(args) => { + run_check_repo(&db, &args)?; + } + Command::Auto { + find_pdses: fp, + snapshot_repos: sr, + } => { + tracing::info!("1. Running command: find-pdses"); + find_pdses::run(&db, &fp.relay, fp.with_scraping.as_ref(), &ua).await?; + + tracing::info!("2. Running command: snapshot-repos"); + snapshot_repos::run(db, snapshot_repos_args(sr), &ua).await?; + } + } + Ok(()) +} + +fn run_add_hosts(db: &storage::Db, args: &AddHostsArgs) -> Result<(), Error> { + let mut added = 0u32; + let mut existing = 0u32; + let mut skipped_bsky = 0u32; + for host in &args.host { + if host.is_bsky() { + tracing::info!(%host, "add-hosts: skipping bsky.network host"); + skipped_bsky += 1; + continue; + } + if storage::host::upsert_manual(db, host)? { + tracing::info!(%host, "add-hosts: inserted"); + added += 1; + } else { + tracing::info!(%host, "add-hosts: already present"); + existing += 1; + } + } + tracing::info!( + added, + existing, + skipped_bsky, + total = args.host.len(), + "add-hosts: done" + ); + Ok(()) +} + +fn run_check_repo(db: &storage::Db, args: &CheckRepoArgs) -> Result<(), Error> { + let outcome = check_repo::run(db, &args.pds, &args.did)?; + if outcome.matched { + tracing::info!( + pds = %args.pds, + did = %args.did, + records = outcome.record_count, + root = %outcome.computed_root, + "check-repo: ok" + ); + Ok(()) + } else { + tracing::error!( + pds = %args.pds, + did = %args.did, + records = outcome.record_count, + expected = %outcome.expected_root, + computed = %outcome.computed_root, + "check-repo: mst root mismatch" + ); + std::process::exit(1); + } +} + +fn snapshot_repos_args(a: SnapshotReposArgs) -> snapshot_repos::Args { + snapshot_repos::Args { + host_workers_limit: a.host_workers_limit, + per_pds_limit: a.per_pds_limit, + big_repos_limit: a.big_repos_limit, + spill_dir: a.spill_dir, + skip_hosts: a.skip, + } +} diff --git a/mini/src/pds.rs b/mini/src/pds.rs new file mode 100644 index 0000000..cce1fd9 --- /dev/null +++ b/mini/src/pds.rs @@ -0,0 +1,308 @@ +//! Per-host PDS operations: probe, listRepos pagination, getRepo. +//! +//! The [`Pds`] struct owns a [`Hostname`] + a `reqwest::Client` + a +//! per-host governor rate limiter, and exposes them as methods so callers +//! don't have to thread the trio through every signature. +//! +//! Adapted from `stats-backfill/src/pds.rs`. Differences: +//! - No DB-side state: `list_repos` yields tuples through a channel and +//! leaves persistence to the orchestrator (so we don't bake in +//! pagination-resume / cursor recovery in v0). +//! - `get_repo` returns mini's own [`crate::repo::FetchOk`]. + +use std::num::NonZeroU32; +use std::sync::Arc; + +use governor::{DefaultDirectRateLimiter, Quota, RateLimiter}; +use hubble_pds::Hostname; +use jacquard_api::com_atproto::sync::list_repos::ListRepos; +use jacquard_common::deps::fluent_uri::Uri; +use jacquard_common::deps::smol_str::SmolStr; +use jacquard_common::xrpc::XrpcExt; +use reqwest::Url; +use tokio::sync::{Semaphore, mpsc}; + +use crate::repo; +use crate::storage::Db; + +const PROBE_BODY_LIMIT: usize = 256 * 1024; +const LIST_REPOS_PAGE: i64 = 1000; + +const SELF_RATE_LIMIT_RPS: NonZeroU32 = NonZeroU32::new(3).unwrap(); +const SELF_RATE_LIMIT_BRIDGY_RPS: NonZeroU32 = NonZeroU32::new(1).unwrap(); +const SELF_RATE_LIMIT_BSKY_RPS: NonZeroU32 = NonZeroU32::new(10).unwrap(); + +#[derive(Debug, thiserror::Error)] +pub enum Error { + /// `jacquard`'s `ClientError` is large; box it. + #[error("xrpc: {0}")] + Xrpc(Box), + + /// Response decode error from `parse::()`. + #[error("xrpc decode: {0}")] + XrpcDecode(String), + + #[error("channel closed (consumer dropped)")] + ChannelClosed, +} + +impl From for Error { + fn from(e: jacquard_common::error::ClientError) -> Self { + Self::Xrpc(Box::new(e)) + } +} + +#[derive(Debug, Clone)] +enum ProbeResult { + /// 2xx with a lil json body + OkJson(String), + /// non-2xx, empty body, body-too-big + OtherHttp(String), + /// no http response at all (DNS / connect / TLS / timeout / read / ...) + FailedTransport(String), +} + +impl ProbeResult { + fn is_reachable(&self) -> bool { + match self { + Self::OkJson(_) => true, + Self::OtherHttp(_) => true, + Self::FailedTransport(_) => false, + } + } + fn ok(self) -> Option { + match self { + Self::OkJson(s) => Some(s), + _ => None, + } + } + fn err(&self) -> Option { + match self { + Self::OkJson(_) => None, + Self::OtherHttp(e) => Some(e.clone()), + Self::FailedTransport(e) => Some(e.clone()), + } + } +} + +#[derive(Debug)] +pub struct ProbeOutcome { + pub health_json: Option, + pub describe_json: Option, + pub probe_error: Option, + pub reachable: bool, +} + +/// One entry emitted by [`Pds::list_repos`]. +#[derive(Debug, Clone)] +pub struct ListedRepo { + pub did: String, + pub rev: String, + pub active: Option, + pub status: Option, +} + +/// Operations against a single PDS host. Cheap to clone (`reqwest::Client` +/// is internally `Arc`-shared, the limiter is `Arc`-wrapped), so this can +/// be moved into spawned tasks. All clones for one host share the same +/// rate budget. +#[derive(Debug, Clone)] +pub struct Pds { + host: Hostname, + client: reqwest::Client, + limiter: Arc, +} + +impl Pds { + pub fn new(host: Hostname, client: reqwest::Client) -> Self { + let qps = if host.is_bsky() { + SELF_RATE_LIMIT_BSKY_RPS + } else if host.is_bridgy() { + SELF_RATE_LIMIT_BRIDGY_RPS + } else { + SELF_RATE_LIMIT_RPS + }; + + let limiter = Arc::new(RateLimiter::direct(Quota::per_second(qps))); + Self { + host, + client, + limiter, + } + } + + pub fn host(&self) -> &Hostname { + &self.host + } + + /// Wait for a token in this host's rate-limit bucket. Acquires before + /// every outbound HTTP call (probe, listRepos page, getRepo). + async fn wait(&self) { + self.limiter.until_ready().await; + } + + /// Get `/xrpc/_health` + `/xrpc/com.atproto.server.describeServer` and + /// stash response json. + /// + /// A host is considered reachable if at least one of the two probes + /// returned any http response. (Fails for DNS / connect / TLS / + /// timeout.) + pub async fn probe(&self) -> ProbeOutcome { + let mut reachable = false; + + let base: Url = format!("https://{}", self.host) + .parse() + .expect("the hostname to parse into a url"); + + self.wait().await; + let health = get_a_lil_json(&self.client, base.clone(), "/xrpc/_health").await; + reachable |= health.is_reachable(); + + self.wait().await; + let describe = get_a_lil_json( + &self.client, + base, + "/xrpc/com.atproto.server.describeServer", + ) + .await; + reachable |= describe.is_reachable(); + + let probe_error = match (health.err(), describe.err()) { + (Some(h), Some(d)) => Some([h, d].join("; ")), + (Some(e), _) | (_, Some(e)) => Some(e), + _ => None, + }; + + ProbeOutcome { + health_json: health.ok(), + describe_json: describe.ok(), + probe_error, + reachable, + } + } + + /// Paginate `com.atproto.sync.listRepos` and emit every observed repo + /// through `tx` as a [`ListedRepo`]. No state is persisted here — + /// the orchestrator decides what to do with each tuple. + pub async fn list_repos(&self, tx: mpsc::Sender) -> Result<(), Error> { + let mut cursor: Option = None; + let mut total: u64 = 0; + + tracing::info!(host = %self.host, "list_repos: starting"); + let base_url = self.host.xrpc_base_url(); + loop { + let request = ListRepos::::new() + .maybe_cursor(cursor.clone()) + .limit(Some(LIST_REPOS_PAGE)) + .build(); + let base = Uri::parse(base_url.as_str()) + .expect("https://{hostname} parses as fluent_uri::Uri"); + + self.wait().await; + let response = self.client.xrpc(base).send(&request).await?; + let output = response + .parse::() + .map_err(|e| Error::XrpcDecode(e.to_string()))?; + + for repo in &output.repos { + let listed = ListedRepo { + did: repo.did.as_ref().to_string(), + rev: repo.rev.as_ref().to_string(), + active: repo.active, + status: repo.status.as_ref().map(|s| s.as_str().to_string()), + }; + tx.send(listed).await.map_err(|_| Error::ChannelClosed)?; + total += 1; + } + + match output.cursor { + Some(c) if !c.is_empty() => { + // Defend against stuck-cursor PDSes: if the server hands + // back the same cursor we just sent, it's looping the + // same page forever. Bail (one duplicate page of DIDs + // already went through tx — downstream is idempotent + // on rev comparison). + if cursor.as_deref() == Some(c.as_str()) { + tracing::warn!( + host = %self.host, + cursor = %c, + total, + "list_repos: server returned same cursor as request; suspect fake-cursor loop, stopping" + ); + break; + } + tracing::debug!(host = %self.host, total, "list_repos page"); + cursor = Some(c); + } + _ => break, + } + } + + tracing::info!(host = %self.host, total, "list_repos: done"); + Ok(()) + } + + /// Fetch one repo. Thin rate-limited wrapper over [`crate::repo::get_repo`]. + /// `listed` carries the `(did, rev, active, status)` tuple from + /// `listRepos`; the new `R\0` state takes its `active` / + /// `status` fields from there. + pub async fn get_repo( + &self, + listed: &ListedRepo, + db: Arc, + big_mem_tokens: &Arc, + spill_dir: &std::path::Path, + ) -> Result { + self.wait().await; + repo::get_repo( + &self.client, + &self.host, + &listed.did, + listed, + db, + big_mem_tokens, + spill_dir, + ) + .await + } +} + +async fn get_a_lil_json(client: &reqwest::Client, mut base: Url, path: &str) -> ProbeResult { + base.set_path(path); + let mut response = match client.get(base).send().await { + Ok(r) => r, + Err(e) => { + if let Some(status) = e.status() { + return ProbeResult::OtherHttp(format!("status: {status} ({path})")); + } + return ProbeResult::FailedTransport(format!("send: {e} ({path})")); + } + }; + let status = response.status(); + if !status.is_success() { + return ProbeResult::OtherHttp(format!("status: {status} ({path})")); + } + + let mut body = Vec::new(); + loop { + let bytes = match response.chunk().await { + Ok(Some(c)) => c, + Ok(None) => break, + Err(e) => { + return ProbeResult::FailedTransport(format!("body: failed read: {e} ({path})")); + } + }; + if body.len() + bytes.len() > PROBE_BODY_LIMIT { + return ProbeResult::OtherHttp(format!("body: too big ({path})")); + } + body.extend_from_slice(&bytes); + } + + let Ok(value) = serde_json::from_slice::(&body) else { + return ProbeResult::OtherHttp(format!("body: not json ({path})")); + }; + + let normalized = serde_json::to_string(&value).expect("serde_json to be round-trip"); + + ProbeResult::OkJson(normalized) +} diff --git a/mini/src/repo.rs b/mini/src/repo.rs new file mode 100644 index 0000000..58ef3eb --- /dev/null +++ b/mini/src/repo.rs @@ -0,0 +1,658 @@ +//! Streaming `getRepo` fetch + repo-stream walk + record writes. +//! +//! Adapted from `stats-backfill/src/repo/{mod,measure}.rs`. Key +//! differences: +//! +//! - No STAR-lite + zstd-3 re-encoding pipe. We're storing the raw +//! per-record bytes as they appeared in the source CAR; star-lite +//! sizing would be redundant work. +//! - The walker writes records into rocksdb directly through a chain of +//! [`rocksdb::WriteBatch`]es with an 8 MiB soft cap, with the first +//! batch range-deleting any previous records for this `(pds, did)` +//! and the final batch atomically committing the trailing records + +//! the new `R\0` state + `s\0` stats. +//! - On any error, partial record writes may already be committed; the +//! next pass detects the rev mismatch + range-deletes + retries, so +//! no GC pass is needed. +//! +//! The three-tier load policy (small in-memory → semaphore-gated large +//! in-memory → per-repo disk spill) and its FD budget come over +//! unchanged from stats-backfill. + +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use futures_util::TryStreamExt; +use hubble_pds::{ACCEPT_ENCODING, CountingRead, Hostname, KnownEncoding, RequestEncodingError}; +use ipld_core::cid::Cid; +use ipld_core::ipld::Ipld; +use jacquard_api::com_atproto::sync::get_repo::GetRepo; +use jacquard_common::deps::fluent_uri::Uri; +use jacquard_common::deps::smol_str::SmolStr; +use jacquard_common::stream::{ByteStream, StreamError}; +use jacquard_common::types::string::{AtStrError, Did}; +use jacquard_common::xrpc::XrpcExt; +use microcosm_links::parse_any_link; +use repo_stream::{DiskBuilder, DiskError, DriveError, DriverBuilder, LoadError}; +use reqwest::StatusCode; +use reqwest::header::{self, HeaderValue}; +use rocksdb::WriteBatch; +use tokio::io::{AsyncRead, AsyncReadExt}; +use tokio::sync::Semaphore; +use tokio_util::io::StreamReader; + +use crate::pds::ListedRepo; +use crate::storage::{ + Db, DbError, commit, commit::StoredCommit, now_unix, record, repo_state, + repo_state::RepoState, repo_stats, repo_stats::RepoStats, +}; + +pub const FETCH_TIMEOUT: Duration = Duration::from_secs(300); + +// free pass for repos under this +const MEM_LIMIT_SMALL_MB: usize = 2; +// concurrency-gated (semaphore) for repos over this +const MEM_LIMIT_LARGE_MB: usize = 150; + +// per-spill fjall config; modest cache since this is a single-pass walk +const SPILL_CACHE_MB: usize = 16; +const SPILL_MAX_STORED_MB: usize = 10 * 1024; // 10 GiB + +/// Commit a write batch when it crosses this many bytes. The trailing +/// records + R/ + s/ commit together in the final batch regardless of +/// size. +const BATCH_SOFT_CAP_BYTES: usize = 8 * 1024 * 1024; + +/// Per-process counter for naming spill subdirectories. +static SPILL_COUNTER: AtomicU64 = AtomicU64::new(0); + +pub fn spill_count() -> u64 { + SPILL_COUNTER.load(Ordering::Relaxed) +} + +#[derive(Debug, thiserror::Error)] +pub enum FetchError { + #[error("invalid request: {0}")] + InvalidRequest(#[from] AtStrError), + + /// no HTTP response (DNS / connect / TLS / read error). + #[error("transport: {0}")] + Transport(#[from] StreamError), + + #[error("http {status}")] + HttpStatus { + status: StatusCode, + body: Option, + retry_after_secs: Option, + }, + + #[error(transparent)] + UnsupportedEncoding(#[from] RequestEncodingError), + + /// repo-stream couldn't parse the CAR (malformed / truncated). + #[error("car parse: {0}")] + CarParse(String), + + /// Couldn't open the per-spill fjall db (disk full, permissions, fd + /// exhaustion, ...). Local-system failure, not the host's fault. + #[error("spill open: {0}")] + SpillOpen(#[source] DiskError), + + #[error("spill drive: {0}")] + SpillDrive(#[source] DriveError), + + /// CAR exceeded `SPILL_MAX_STORED_MB` after spilling. + #[error("car too large for any tier")] + CarTooLarge, + + /// Post-load walk failed (CBOR decode, repo-stream walk error, or + /// rocksdb write error). + #[error("walk: {0}")] + Walk(#[from] WalkError), + + /// 300s overall budget exhausted. + #[error("timeout")] + Timeout, +} + +#[derive(Debug, thiserror::Error)] +pub enum WalkError { + #[error("repo-stream walk: {0}")] + StreamWalk(#[from] repo_stream::walk::WalkError), + + #[error("disk walk: {0}")] + DiskWalk(#[from] repo_stream::disk::DriveError), + + /// One record's bytes didn't decode as DAG-CBOR. We treat this as a + /// terminal repo failure: the bytes are what the PDS published, and + /// atproto records are required to be valid DAG-CBOR. + #[error("record cbor decode at key {key:?}: {reason}")] + RecordCbor { key: String, reason: String }, + + /// Record key didn't have the expected `collection/rkey` shape. + /// A protocol violation; same disposition as CBOR errors. + #[error("malformed record key {key:?}")] + BadKey { key: String }, + + #[error("storage: {0}")] + Storage(#[from] DbError), + + #[error("rocksdb write: {0}")] + Rocks(#[from] rocksdb::Error), +} + +#[derive(Debug)] +pub struct FetchOk { + pub rev: String, + pub wire_bytes: u64, + pub car_bytes: u64, + /// `Content-Encoding` the server used (or `None` for identity). + /// Captured for diagnostic logging even when no caller reads it. + #[allow(dead_code)] + pub content_encoding: Option, + pub stats: RepoStats, +} + +pub async fn get_repo( + client: &reqwest::Client, + host: &Hostname, + did: &str, + listed: &ListedRepo, + db: Arc, + big_mem_tokens: &Arc, + spill_dir: &Path, +) -> Result { + let base_url = format!("https://{}", host); + let did_typed = Did::::new(SmolStr::from(did))?; + let request = GetRepo:: { + did: did_typed, + since: None, + }; + + let host_owned = host.clone(); + let did_owned = did.to_string(); + let listed_owned = listed.clone(); + + let go = async { + let base = + Uri::parse(base_url.as_str()).expect("https://{hostname} parses as fluent_uri::Uri"); + + // reqwest's gzip/br/zstd/deflate features are not used because + // they strip Content-Encoding from the response. + let response = client + .xrpc(base) + .header(header::ACCEPT_ENCODING, HeaderValue::from_static(ACCEPT_ENCODING)) + .download(&request) + .await?; + + let status = response.status(); + let encoding = KnownEncoding::from_response_headers(response.headers())?; + let encoding_wire = encoding.map(|e| e.as_str().to_string()); + let retry_after_secs = parse_retry_after_secs(response.headers()); + let (_parts, body) = response.into_parts(); + + if !status.is_success() { + let body_text = drain_body_to_string(body).await; + return Err(FetchError::HttpStatus { + status, + body: Some(body_text), + retry_after_secs, + }); + } + + let mapped = body.into_inner().map_err(std::io::Error::other); + let stream_reader = StreamReader::new(mapped); + let (counting_wire, wire_count) = CountingRead::new(stream_reader); + + let decoded_reader: Box = match encoding { + None => Box::new(counting_wire), + Some(enc) => enc.decode(counting_wire), + }; + + let (counting_decoded, decoded_count) = CountingRead::new(decoded_reader); + + // Phase 1: small in-memory limit. + // + // Byte counters are read just before each `walk_*` call: by the + // time a tier-load `await` returns, the entire byte stream has + // flowed through both `CountingRead` wrappers, so the totals are + // stable. The walker stamps them onto `RepoStats` itself so the + // persisted s/ row carries the wire / decoded sizes. + let walk_out = match DriverBuilder::new() + .with_mem_limit_mb(MEM_LIMIT_SMALL_MB) + .load_car(counting_decoded) + .await + { + Ok(car) => { + let wire_b = wire_count.get(); + let car_b = decoded_count.get(); + walk_mem( + car, + db.clone(), + host_owned.clone(), + did_owned.clone(), + listed_owned.clone(), + encoding_wire.clone(), + wire_b, + car_b, + ) + .await? + } + Err(LoadError::MemoryLimitReached(partial)) => { + // Phase 2: gated by the global big-memory semaphore. + #[expect(unused_variables, reason = "semaphore drop guard")] + let permit = big_mem_tokens + .clone() + .acquire_owned() + .await + .expect("big_mem_tokens semaphore not closed"); + tracing::debug!(%host, %did, "phase-2 large-mem load"); + match partial.continue_loading(MEM_LIMIT_LARGE_MB).await { + Ok(car) => { + let wire_b = wire_count.get(); + let car_b = decoded_count.get(); + walk_mem( + car, + db.clone(), + host_owned.clone(), + did_owned.clone(), + listed_owned.clone(), + encoding_wire.clone(), + wire_b, + car_b, + ) + .await? + } + Err(LoadError::MemoryLimitReached(partial)) => { + // Phase 3: spill to disk. Each spill gets its own + // fjall db dir; `temporary(true)` (set inside + // repo-stream) cleans it on drop. + let counter = SPILL_COUNTER.fetch_add(1, Ordering::Relaxed); + let path = spill_dir.join(format!("spill-{counter}")); + tracing::debug!(%host, %did, ?path, "phase-3 disk spill"); + let store = DiskBuilder::new() + .with_cache_size_mb(SPILL_CACHE_MB) + .with_max_stored_mb(SPILL_MAX_STORED_MB) + .open(path) + .await + .map_err(FetchError::SpillOpen)?; + match partial.finish_loading(store).await { + Ok((commit, _root, driver)) => { + let wire_b = wire_count.get(); + let car_b = decoded_count.get(); + walk_disk( + commit, + driver, + db.clone(), + host_owned.clone(), + did_owned.clone(), + listed_owned.clone(), + encoding_wire.clone(), + wire_b, + car_b, + ) + .await? + } + Err(DriveError::StorageError(DiskError::MaxSizeExceeded)) => { + return Err(FetchError::CarTooLarge); + } + Err(e) => return Err(FetchError::SpillDrive(e)), + } + } + Err(e) => return Err(FetchError::CarParse(e.to_string())), + } + } + Err(e) => return Err(FetchError::CarParse(e.to_string())), + }; + + Ok(FetchOk { + rev: walk_out.rev, + wire_bytes: walk_out.stats.wire_bytes, + car_bytes: walk_out.stats.car_bytes, + content_encoding: encoding, + stats: walk_out.stats, + }) + }; + + tokio::time::timeout(FETCH_TIMEOUT, go) + .await + .map_err(|_| FetchError::Timeout)? +} + +struct WalkOut { + rev: String, + stats: RepoStats, +} + +/// Convert a `repo_stream::Commit` (parsed signed-commit fields) into +/// the `StoredCommit` we persist at `C\0`. CIDs are rendered +/// as base32-multibase strings; `sig` is captured byte-for-byte. +fn stored_commit_from(c: &repo_stream::Commit) -> StoredCommit { + let data: Cid = c.data.into(); + let prev: Option = c.prev.map(Into::into); + StoredCommit { + did: c.did.clone(), + version: c.version, + data: data.to_string(), + rev: c.rev.clone(), + prev: prev.map(|p| p.to_string()), + sig: c.sig.to_vec(), + } +} + +/// Walk an in-memory `MemCar` and write records on the blocking pool. +#[allow(clippy::too_many_arguments)] +async fn walk_mem( + mut car: repo_stream::MemCar, + db: Arc, + host: Hostname, + did: String, + listed: ListedRepo, + content_encoding: Option, + wire_bytes: u64, + car_bytes: u64, +) -> Result { + tokio::task::spawn_blocking(move || -> Result { + let stored_commit = stored_commit_from(&car.commit); + let rev = car.commit.rev.clone(); + let mut writer = RecordWriter::start( + &db, + &host, + &did, + content_encoding, + wire_bytes, + car_bytes, + stored_commit, + )?; + while let Some(out) = car.next_strict()? { + writer.process_one(&out.key, &out.data)?; + } + let stats = writer.finish(&listed, &rev)?; + Ok(WalkOut { rev, stats }) + }) + .await + .expect("walk_mem spawn_blocking task panic") +} + +/// Walk a `DiskDriver` (phase-3 disk-spilled CAR) and write records on +/// the blocking pool. +#[allow(clippy::too_many_arguments)] +async fn walk_disk( + commit: repo_stream::mst::Commit, + mut driver: repo_stream::DiskDriver, + db: Arc, + host: Hostname, + did: String, + listed: ListedRepo, + content_encoding: Option, + wire_bytes: u64, + car_bytes: u64, +) -> Result { + tokio::task::spawn_blocking(move || -> Result { + let stored_commit = stored_commit_from(&commit); + let rev = commit.rev.clone(); + let mut writer = RecordWriter::start( + &db, + &host, + &did, + content_encoding, + wire_bytes, + car_bytes, + stored_commit, + )?; + while let Some((key, data)) = driver.next_blocking()? { + writer.process_one(&key, &data)?; + } + let stats = writer.finish(&listed, &rev)?; + Ok(WalkOut { rev, stats }) + }) + .await + .expect("walk_disk spawn_blocking task panic") +} + +/// Per-repo write batcher. Holds: +/// +/// - a [`WriteBatch`] that accumulates record puts +/// - an `&Arc` for intermediate commits +/// - the [`Hostname`] / `did` strings every key encoder needs +/// - the parsed signed commit, written to `C\0` in the final +/// batch alongside R/ and s/ +/// +/// MST integrity verification is **not** done during snapshot — see the +/// `check-repo` subcommand. Snapshot just stores everything we received +/// and lets a separate read-back path confirm consistency on demand. +/// +/// Drops everything (and commits the final batch) via [`finish`]. +struct RecordWriter<'a> { + db: &'a Arc, + host: &'a Hostname, + did: &'a str, + batch: WriteBatch, + stats: RepoStats, + commit: StoredCommit, +} + +impl<'a> RecordWriter<'a> { + /// Create a new writer with batch-1 already containing the + /// `delete_range_cf` over this `(pds, did)`'s record range. The + /// `content_encoding` (literal wire name, or `None` for identity) + /// and final `wire_bytes` / `car_bytes` totals from the + /// `CountingRead`s are stamped onto `RepoStats` up-front so the + /// final batch carries them regardless of how the walk terminates. + /// `commit` is the signed commit block parsed from the CAR; it + /// rides through to [`finish`] and lands at `C\0`. + #[allow(clippy::too_many_arguments)] + fn start( + db: &'a Arc, + host: &'a Hostname, + did: &'a str, + content_encoding: Option, + wire_bytes: u64, + car_bytes: u64, + commit: StoredCommit, + ) -> Result { + let mut batch = WriteBatch::default(); + record::delete_range_into(&mut batch, db, host, did)?; + let stats = RepoStats { + content_encoding, + wire_bytes, + car_bytes, + ..RepoStats::default() + }; + Ok(Self { + db, + host, + did, + batch, + stats, + commit, + }) + } + + fn process_one(&mut self, mst_key: &str, data: &[u8]) -> Result<(), WalkError> { + // Repo-stream MST keys are `/`; we don't split, we + // store byte-for-byte under `r\0\0` so + // rocksdb byte-order matches MST lex order. + if !mst_key.contains('/') { + return Err(WalkError::BadKey { + key: mst_key.to_string(), + }); + } + + // CBOR-decode once for the blob/backlink walk. Records must be + // DAG-CBOR per atproto; a decode error is a terminal repo + // failure (the bytes are not records we want to store). + let value: Ipld = + serde_ipld_dagcbor::from_slice(data).map_err(|e| WalkError::RecordCbor { + key: mst_key.to_string(), + reason: e.to_string(), + })?; + + record::put_into(&mut self.batch, self.db, self.host, self.did, mst_key, data)?; + self.stats.record_count += 1; + walk_value(&value, &mut self.stats); + + if self.batch.size_in_bytes() > BATCH_SOFT_CAP_BYTES { + let to_commit = std::mem::take(&mut self.batch); + self.db.inner.write(to_commit)?; + } + + Ok(()) + } + + /// Commit the final batch atomically: any remaining records + the + /// new `R\0` + `s\0` + `C\0` entries. + fn finish(self, listed: &ListedRepo, rev: &str) -> Result { + let Self { + db, + host, + did, + mut batch, + stats, + commit: stored_commit, + } = self; + + let state = RepoState { + last_seen_rev: rev.to_string(), + last_fetched_at_unix: now_unix(), + listrepos_active: listed.active.unwrap_or(true), + listrepos_status: listed.status.clone(), + tombstoned: false, + last_error: None, + }; + repo_state::put_into(&mut batch, db, host, did, &state)?; + repo_stats::put_into(&mut batch, db, host, did, &stats)?; + commit::put_into(&mut batch, db, host, did, &stored_commit)?; + db.inner.write(batch)?; + Ok(stats) + } +} + +/// Recursive walk over a record's CBOR value, accumulating both: +/// - blob refs: `{$type: "blob", ref: , size?: int, mimeType?: str, ...}` +/// - backlinks: string values that `microcosm_links::parse_any_link` +/// recognises (at-uri, did, or generic URI). +fn walk_value(value: &Ipld, stats: &mut RepoStats) { + match value { + Ipld::Map(map) => { + if matches!(map.get("$type"), Some(Ipld::String(t)) if t == "blob") + && matches!(map.get("ref"), Some(Ipld::Link(_))) + { + stats.blob_ref_count += 1; + let size = match map.get("size") { + Some(Ipld::Integer(n)) if *n > 0 => u64::try_from(*n).unwrap_or(0), + _ => 0, + }; + if size == 0 { + stats.unsized_blob_count += 1; + } else { + stats.blob_total_size = stats.blob_total_size.saturating_add(size); + } + match map.get("mimeType") { + Some(Ipld::String(s)) if !s.is_empty() => { + *stats.blob_mime_counts.entry(s.clone()).or_default() += 1; + } + _ => { + stats.blob_no_mime_count += 1; + } + } + // Don't recurse into the matched blob — its remaining fields + // are scalars and would cause false positives. + return; + } + for v in map.values() { + walk_value(v, stats); + } + } + Ipld::List(items) => { + for item in items { + walk_value(item, stats); + } + } + Ipld::String(s) if parse_any_link(s).is_some() => { + stats.backlink_count += 1; + } + _ => {} + } +} + +fn parse_retry_after_secs(headers: &reqwest::header::HeaderMap) -> Option { + headers + .get(reqwest::header::RETRY_AFTER)? + .to_str() + .ok()? + .trim() + .parse::() + .ok() +} + +async fn drain_body_to_string(body: ByteStream) -> String { + let mapped = body.into_inner().map_err(std::io::Error::other); + let mut reader = StreamReader::new(mapped); + let mut buf = [0; 2048]; + let Ok(n) = reader.read(&mut buf).await else { + return "[could not read error response body]".to_string(); + }; + String::from_utf8_lossy(&buf[..n]).to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + use ipld_core::cid::Cid; + use std::collections::BTreeMap; + + fn cid_str(s: &str) -> Cid { + s.parse().expect("test cid parses") + } + + #[test] + fn walk_value_counts_blob_with_mime() { + let cid = cid_str("bafkreidkjxqphmt3yhvdgmwx5p64q7s7tcmw3rj7ynvtu7tnp7nb6yurxe"); + let mut map = BTreeMap::new(); + map.insert("$type".into(), Ipld::String("blob".into())); + map.insert("ref".into(), Ipld::Link(cid)); + map.insert("mimeType".into(), Ipld::String("image/jpeg".into())); + map.insert("size".into(), Ipld::Integer(12345)); + + let mut stats = RepoStats::default(); + walk_value(&Ipld::Map(map), &mut stats); + assert_eq!(stats.blob_ref_count, 1); + assert_eq!(stats.blob_total_size, 12345); + assert_eq!(stats.unsized_blob_count, 0); + assert_eq!(stats.blob_mime_counts.get("image/jpeg"), Some(&1)); + assert_eq!(stats.blob_no_mime_count, 0); + } + + #[test] + fn walk_value_counts_blob_missing_mime() { + let cid = cid_str("bafkreidkjxqphmt3yhvdgmwx5p64q7s7tcmw3rj7ynvtu7tnp7nb6yurxe"); + let mut map = BTreeMap::new(); + map.insert("$type".into(), Ipld::String("blob".into())); + map.insert("ref".into(), Ipld::Link(cid)); + map.insert("size".into(), Ipld::Integer(100)); + // no mimeType + + let mut stats = RepoStats::default(); + walk_value(&Ipld::Map(map), &mut stats); + assert_eq!(stats.blob_ref_count, 1); + assert!(stats.blob_mime_counts.is_empty()); + assert_eq!(stats.blob_no_mime_count, 1); + } + + #[test] + fn walk_value_counts_at_uri_backlink() { + let mut map = BTreeMap::new(); + map.insert( + "uri".into(), + Ipld::String("at://did:plc:foo/app.bsky.feed.post/3kabc".into()), + ); + map.insert("text".into(), Ipld::String("hello".into())); + + let mut stats = RepoStats::default(); + walk_value(&Ipld::Map(map), &mut stats); + assert_eq!(stats.backlink_count, 1); + } +} diff --git a/mini/src/snapshot_repos.rs b/mini/src/snapshot_repos.rs new file mode 100644 index 0000000..e39d43b --- /dev/null +++ b/mini/src/snapshot_repos.rs @@ -0,0 +1,699 @@ +//! `snapshot` orchestrator: walk every known PDS host, probe it, paginate +//! listRepos, fetch every repo whose `rev` advanced (or is new) into +//! rocksdb. Tombstone repos that disappeared from `listRepos`. +//! +//! Adapted from `stats-backfill/src/backfill_repos.rs`. Mini-specific +//! differences: +//! +//! - No SQLite-backed `eligible_hosts` / `fetch_state` plumbing — the host +//! list comes from a single `h` prefix scan and each pass runs each +//! host to completion exactly once. +//! - Per-repo rev gating happens in-process by reading the existing +//! `R\0` state; only fetches when `listrepos.rev` strictly +//! exceeds the stored rev. +//! - Inactive repos (`listrepos.active == Some(false)`) get an R/ write +//! inline and skip the record fetch. +//! - Disappeared repos (in stored R/ but absent from this run's +//! listRepos) are tombstoned, not deleted. +//! - No `--max-rss-mb` self-restart ceiling. + +use std::collections::{HashMap, HashSet}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use hubble_pds::Hostname; +use reqwest::StatusCode; +use tokio::sync::{Semaphore, mpsc}; +use tokio::task::JoinSet; + +use crate::pds::{ListedRepo, Pds}; +use crate::repo::{self, FetchError}; +use crate::storage::{ + Db, DbError, host, now_unix, probe, probe::Probe, repo_state, repo_state::RepoState, +}; + +/// How many *consecutive* transient outcomes (429 / 5xx / connect / timeout) +/// trip a host into cooldown. Resets on any success. +const COOLDOWN_TRANSIENT_THRESHOLD: u32 = 3; + +/// Floor for cooldown duration. Server-provided `Retry-After` above this +/// takes precedence (capped by `COOLDOWN_MAX_SECS`). +const COOLDOWN_DURATION_SECS: u64 = 60; + +/// Hard cap on cooldown duration. Defends against a server pinning a +/// host out of rotation indefinitely via `Retry-After: 99999999`. +const COOLDOWN_MAX_SECS: u64 = 3600; + +/// Rough generous fd budget per concurrent disk-spilling fetch. +const NOFILE_PER_BIG_REPO: u64 = 1024; +/// fds for rest of process: rocksdb SST handles + reqwest sockets + +/// stdio + small reserves. Higher than stats-backfill's 256 because +/// rocks keeps many SSTs open by default. +const NOFILE_SLACK: u64 = 2048; + +const PROGRESS_TICK_SECS: u64 = 5; + +const LISTREPOS_CHANNEL_DEPTH: usize = 1024; + +// --------------------------------------------------------------------------- +// Progress + classification +// --------------------------------------------------------------------------- + +#[derive(Default)] +struct Progress { + hosts_active: AtomicUsize, + hosts_done: AtomicU64, + fetched_ok: AtomicU64, + fetched_err: AtomicU64, + bytes_wire: AtomicU64, + records_written: AtomicU64, + tombstoned: AtomicU64, +} + +#[derive(Default, Clone, Copy)] +struct ProgressSnapshot { + hosts_active: usize, + hosts_done: u64, + fetched_ok: u64, + fetched_err: u64, + bytes_wire: u64, + records_written: u64, + tombstoned: u64, + disk_spills_started: u64, +} + +impl Progress { + fn snapshot(&self) -> ProgressSnapshot { + ProgressSnapshot { + hosts_active: self.hosts_active.load(Ordering::Relaxed), + hosts_done: self.hosts_done.load(Ordering::Relaxed), + fetched_ok: self.fetched_ok.load(Ordering::Relaxed), + fetched_err: self.fetched_err.load(Ordering::Relaxed), + bytes_wire: self.bytes_wire.load(Ordering::Relaxed), + records_written: self.records_written.load(Ordering::Relaxed), + tombstoned: self.tombstoned.load(Ordering::Relaxed), + disk_spills_started: repo::spill_count(), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Classification { + Ok, + RateLimited, + Server5xx, + ConnectError, + Timeout, + ParseError, + TerminalXrpcError, +} + +// --------------------------------------------------------------------------- +// Cooldown +// --------------------------------------------------------------------------- + +/// Per-host transient-outcome tracker shared across spawned fetch tasks. +/// Tracks the maximum `Retry-After` seen so the cooldown honours the +/// strongest server hint. +#[derive(Default)] +struct CooldownState { + consecutive_transient: AtomicU32, + triggered: AtomicBool, + max_retry_after_secs: AtomicU64, +} + +impl CooldownState { + fn record(&self, class: Classification, retry_after_secs: Option) { + match class { + Classification::Ok => { + self.consecutive_transient.store(0, Ordering::Relaxed); + } + Classification::RateLimited + | Classification::Server5xx + | Classification::ConnectError + | Classification::Timeout => { + let n = self.consecutive_transient.fetch_add(1, Ordering::Relaxed) + 1; + if n >= COOLDOWN_TRANSIENT_THRESHOLD { + self.triggered.store(true, Ordering::Relaxed); + } + } + Classification::TerminalXrpcError | Classification::ParseError => { + // Terminal outcomes don't move the cooldown needle either way. + } + } + + if let Some(secs) = retry_after_secs { + let bounded = secs.min(COOLDOWN_MAX_SECS); + let mut current = self.max_retry_after_secs.load(Ordering::Relaxed); + while bounded > current { + match self.max_retry_after_secs.compare_exchange_weak( + current, + bounded, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => break, + Err(c) => current = c, + } + } + } + } + + fn triggered(&self) -> bool { + self.triggered.load(Ordering::Relaxed) + } + + fn cooldown_duration_secs(&self) -> u64 { + self.max_retry_after_secs + .load(Ordering::Relaxed) + .max(COOLDOWN_DURATION_SECS) + } +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("storage: {0}")] + Storage(#[from] DbError), + + #[error("http client: {0}")] + HttpClient(#[from] reqwest::Error), + + #[error("rlimit query/set failed: {0}")] + Rlimit(#[source] std::io::Error), + + #[error("could not create spill_dir {path:?}: {source}")] + SpillDirCreate { + path: PathBuf, + #[source] + source: std::io::Error, + }, +} + +pub struct Args { + pub host_workers_limit: usize, + pub per_pds_limit: usize, + /// Cap on concurrent in-flight tier-2 CAR loads (~150 MiB each). + pub big_repos_limit: usize, + /// Directory under which each tier-3 disk spill creates a fresh subdir. + pub spill_dir: PathBuf, + /// Hosts to skip entirely (never picked as candidates). + pub skip_hosts: Vec, +} + +pub async fn run(db: Arc, args: Args, ua: &str) -> Result<(), Error> { + ensure_nofile_for_big_repos(args.big_repos_limit)?; + std::fs::create_dir_all(&args.spill_dir).map_err(|source| Error::SpillDirCreate { + path: args.spill_dir.clone(), + source, + })?; + + let client = reqwest::Client::builder() + .user_agent(ua) + .connect_timeout(Duration::from_secs(10)) + .pool_max_idle_per_host(2) + .pool_idle_timeout(Duration::from_secs(30)) + .build()?; + + let big_mem_tokens = Arc::new(Semaphore::new(args.big_repos_limit)); + let spill_dir = Arc::new(args.spill_dir); + let progress = Arc::new(Progress::default()); + + let skip_hosts: HashSet = args.skip_hosts.into_iter().collect(); + if !skip_hosts.is_empty() { + tracing::info!( + count = skip_hosts.len(), + hosts = ?skip_hosts.iter().collect::>(), + "skipping hosts (--skip)" + ); + } + + // Snapshot host set up-front. The discover step writes h rows; + // snapshot doesn't add new hosts. + let mut hosts: Vec = host::list_all(&db)? + .into_iter() + .map(|(h, _)| h) + .filter(|h| !h.is_bsky() && !skip_hosts.contains(h)) + .collect(); + // Stable order so reproducible runs / logs are easier to follow. + hosts.sort(); + tracing::info!(host_count = hosts.len(), "snapshot: starting"); + let mut next_host = hosts.into_iter(); + + let ticker = spawn_progress_ticker(progress.clone()); + + let mut active: JoinSet<()> = JoinSet::new(); + loop { + // Top up workers up to the cap. + while active.len() < args.host_workers_limit { + let Some(host) = next_host.next() else { break }; + let db = db.clone(); + let pds = Pds::new(host, client.clone()); + let per_pds = args.per_pds_limit; + let big_mem = big_mem_tokens.clone(); + let spill = spill_dir.clone(); + let prog = progress.clone(); + progress.hosts_active.fetch_add(1, Ordering::Relaxed); + active.spawn(async move { + if let Err(e) = host_worker(db, pds.clone(), per_pds, big_mem, spill, prog.clone()) + .await + { + tracing::warn!(host = %pds.host(), %e, "host_worker error"); + } + prog.hosts_active.fetch_sub(1, Ordering::Relaxed); + prog.hosts_done.fetch_add(1, Ordering::Relaxed); + }); + } + + let Some(joined) = active.join_next().await else { + break; + }; + if let Err(e) = joined { + tracing::warn!(%e, "host_worker join error (panic?)"); + } + } + + ticker.abort(); + let snap = progress.snapshot(); + tracing::info!( + hosts_done = snap.hosts_done, + fetched_ok = snap.fetched_ok, + fetched_err = snap.fetched_err, + records_written = snap.records_written, + wire_mb = snap.bytes_wire / (1024 * 1024), + tombstoned = snap.tombstoned, + disk_spills = snap.disk_spills_started, + "snapshot: done" + ); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Per-host worker +// --------------------------------------------------------------------------- + +async fn host_worker( + db: Arc, + pds: Pds, + per_pds_limit: usize, + big_mem_tokens: Arc, + spill_dir: Arc, + progress: Arc, +) -> Result<(), Error> { + // Step 1: probe. + let info = pds.probe().await; + let probe_value = Probe { + probed_at_unix: now_unix(), + health_json: info.health_json.clone(), + describe_json: info.describe_json.clone(), + probe_error: info.probe_error.clone(), + }; + probe::put(&db, pds.host(), &probe_value)?; + if !info.reachable { + tracing::warn!(host = %pds.host(), error = ?info.probe_error, "host unreachable; skipping"); + return Ok(()); + } + + // Step 2: load the existing R\0... DID set so we can detect + // disappearance after listRepos completes. + let seen_initial: HashMap = repo_state::dids_for_pds(&db, pds.host())? + .into_iter() + .collect(); + tracing::debug!(host = %pds.host(), known = seen_initial.len(), "host_worker: loaded prior repo set"); + + // Step 3: listRepos paginator → channel → fetch loop. + let (tx, rx) = mpsc::channel::(LISTREPOS_CHANNEL_DEPTH); + + let lr_pds = pds.clone(); + let lr_handle = tokio::spawn(async move { lr_pds.list_repos(tx).await }); + + let fr_pds = pds.clone(); + let fr_db = db.clone(); + let fr_big_mem = big_mem_tokens.clone(); + let fr_spill = spill_dir.clone(); + let fr_progress = progress.clone(); + let fr_handle = tokio::spawn(async move { + fetch_repos_loop( + fr_db, + fr_pds, + rx, + per_pds_limit, + fr_big_mem, + fr_spill, + fr_progress, + seen_initial, + ) + .await + }); + + let fr_outcome = fr_handle.await.expect("fetch_repos_loop task panic"); + let lr_res = lr_handle.await.expect("list_repos task panic"); + + if let Err(e) = lr_res { + tracing::warn!(host = %pds.host(), %e, "list_repos failed"); + } + + match fr_outcome { + Ok(FetchLoopOutcome::Done) => { + tracing::info!(host = %pds.host(), "host done"); + } + Ok(FetchLoopOutcome::CoolDown { until_unix }) => { + tracing::warn!(host = %pds.host(), until_unix, "host cooled down on transient errors"); + } + Err(e) => { + tracing::warn!(host = %pds.host(), %e, "fetch_repos failed"); + } + } + + Ok(()) +} + +enum FetchLoopOutcome { + Done, + CoolDown { until_unix: u64 }, +} + +#[allow(clippy::too_many_arguments)] +async fn fetch_repos_loop( + db: Arc, + pds: Pds, + mut rx: mpsc::Receiver, + per_pds_limit: usize, + big_mem_tokens: Arc, + spill_dir: Arc, + progress: Arc, + mut seen: HashMap, +) -> Result { + let semaphore = Arc::new(Semaphore::new(per_pds_limit)); + let cooldown = Arc::new(CooldownState::default()); + let mut in_flight: JoinSet<()> = JoinSet::new(); + + while let Some(listed) = rx.recv().await { + if cooldown.triggered() { + tracing::warn!( + host = %pds.host(), + "cooldown threshold reached; ceasing new fetches and draining in-flight", + ); + break; + } + + let prev = seen.remove(&listed.did); + + // Inactive: write R/ inline; never fetch records. + if listed.active == Some(false) { + let state = RepoState { + last_seen_rev: listed.rev.clone(), + last_fetched_at_unix: prev + .as_ref() + .map(|p| p.last_fetched_at_unix) + .unwrap_or(0), + listrepos_active: false, + listrepos_status: listed.status.clone(), + tombstoned: false, + last_error: None, + }; + if let Err(e) = repo_state::put(&db, pds.host(), &listed.did, &state) { + tracing::warn!(host = %pds.host(), did = %listed.did, %e, "write inactive R/ failed"); + } + continue; + } + + // Active: fetch only if rev strictly advanced. + let needs_fetch = match prev.as_ref() { + None => true, // brand new on this host + Some(p) => listed.rev > p.last_seen_rev, // TID lex-compare + }; + if !needs_fetch { + continue; + } + + let permit = semaphore + .clone() + .acquire_owned() + .await + .expect("semaphore not closed"); + let db_task = db.clone(); + let pds_task = pds.clone(); + let big_mem = big_mem_tokens.clone(); + let spill = spill_dir.clone(); + let cooldown_task = cooldown.clone(); + let progress_task = progress.clone(); + in_flight.spawn(async move { + #[expect(unused_variables, reason = "semaphore drop-guard")] + let permit = permit; + do_one_fetch( + &db_task, + &pds_task, + &listed, + &big_mem, + &spill, + &progress_task, + &cooldown_task, + ) + .await; + }); + } + + drop(rx); + while in_flight.join_next().await.is_some() {} + + // Tombstone disappeared repos: anything left in `seen` was in R/ + // before but absent from this run's listRepos. + for (did, prev_state) in seen { + // Already tombstoned? Skip re-write to save IO. + if prev_state.tombstoned { + continue; + } + let next = RepoState { + tombstoned: true, + ..prev_state + }; + match repo_state::put(&db, pds.host(), &did, &next) { + Ok(()) => { + progress.tombstoned.fetch_add(1, Ordering::Relaxed); + } + Err(e) => tracing::warn!(host = %pds.host(), %did, %e, "tombstone write failed"), + } + } + + if cooldown.triggered() { + let secs = cooldown.cooldown_duration_secs(); + tracing::info!( + host = %pds.host(), + cooldown_secs = secs, + "applying cooldown (with Retry-After hint if present)" + ); + Ok(FetchLoopOutcome::CoolDown { + until_unix: now_unix() + secs, + }) + } else { + Ok(FetchLoopOutcome::Done) + } +} + +async fn do_one_fetch( + db: &Arc, + pds: &Pds, + listed: &ListedRepo, + big_mem_tokens: &Arc, + spill_dir: &Path, + progress: &Progress, + cooldown: &CooldownState, +) { + let started = Instant::now(); + let result = pds + .get_repo(listed, db.clone(), big_mem_tokens, spill_dir) + .await; + let duration_ms = started.elapsed().as_millis() as u64; + + match result { + Ok(ok) => { + progress.fetched_ok.fetch_add(1, Ordering::Relaxed); + progress.bytes_wire.fetch_add(ok.wire_bytes, Ordering::Relaxed); + progress + .records_written + .fetch_add(ok.stats.record_count, Ordering::Relaxed); + cooldown.record(Classification::Ok, None); + tracing::debug!( + host = %pds.host(), + did = %listed.did, + rev = %ok.rev, + records = ok.stats.record_count, + wire_b = ok.wire_bytes, + car_b = ok.car_bytes, + ms = duration_ms, + "fetch ok" + ); + } + Err(e) => { + progress.fetched_err.fetch_add(1, Ordering::Relaxed); + let (class, retry_after) = classify(&e); + cooldown.record(class, retry_after); + tracing::warn!( + host = %pds.host(), + did = %listed.did, + ms = duration_ms, + class = ?class, + %e, + "fetch err" + ); + } + } +} + +fn classify(e: &FetchError) -> (Classification, Option) { + match e { + FetchError::Transport(_) => (Classification::ConnectError, None), + FetchError::Timeout => (Classification::Timeout, None), + FetchError::CarParse(_) | FetchError::CarTooLarge => (Classification::ParseError, None), + FetchError::SpillOpen(_) | FetchError::SpillDrive(_) => (Classification::ParseError, None), + FetchError::Walk(_) => (Classification::ParseError, None), + FetchError::InvalidRequest(_) => (Classification::TerminalXrpcError, None), + FetchError::UnsupportedEncoding(_) => (Classification::ParseError, None), + FetchError::HttpStatus { + status, + retry_after_secs, + .. + } => (classify_status(*status), *retry_after_secs), + } +} + +fn classify_status(s: StatusCode) -> Classification { + match s.as_u16() { + 429 => Classification::RateLimited, + 500..=599 => Classification::Server5xx, + _ => Classification::TerminalXrpcError, + } +} + +// --------------------------------------------------------------------------- +// Progress ticker +// --------------------------------------------------------------------------- + +fn spawn_progress_ticker(progress: Arc) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + let started = Instant::now(); + let mut last = progress.snapshot(); + let mut interval = tokio::time::interval(Duration::from_secs(PROGRESS_TICK_SECS)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + interval.tick().await; + let now = progress.snapshot(); + let dt = PROGRESS_TICK_SECS as f64; + let d_ok = now.fetched_ok.saturating_sub(last.fetched_ok); + let d_wire = now.bytes_wire.saturating_sub(last.bytes_wire); + tracing::info!( + hosts_active = now.hosts_active, + hosts_done = now.hosts_done, + fetched = now.fetched_ok, + failed = now.fetched_err, + records = now.records_written, + rate_per_s = format!("{:.1}", d_ok as f64 / dt), + wire_mb_per_s = format!("{:.2}", d_wire as f64 / dt / (1024.0 * 1024.0)), + wire_mb = now.bytes_wire / (1024 * 1024), + tombstoned = now.tombstoned, + spills = now.disk_spills_started, + uptime_s = started.elapsed().as_secs(), + "progress" + ); + last = now; + } + }) +} + +// --------------------------------------------------------------------------- +// rlimit +// --------------------------------------------------------------------------- + +/// Check that RLIMIT_NOFILE is high enough for `big_repos_limit` concurrent +/// disk-spilling fjall instances plus rocksdb SSTs and reqwest sockets. +/// Raises the soft limit toward the hard limit if there's headroom. +fn ensure_nofile_for_big_repos(big_repos_limit: usize) -> Result<(), Error> { + let needed = big_repos_limit as u64 * NOFILE_PER_BIG_REPO + NOFILE_SLACK; + let (soft, hard) = rlimit::Resource::NOFILE.get().map_err(Error::Rlimit)?; + if soft >= needed { + tracing::debug!(soft, hard, needed, "RLIMIT_NOFILE ok"); + return Ok(()); + } + let target = needed.min(hard); + if target > soft { + rlimit::Resource::NOFILE + .set(target, hard) + .map_err(Error::Rlimit)?; + tracing::info!( + from = soft, + to = target, + hard, + "raised RLIMIT_NOFILE soft cap" + ); + } + if target < needed { + tracing::warn!( + hard, + needed, + big_repos_limit, + "RLIMIT_NOFILE hard cap below required; concurrent spills may fail. raise ulimit -n" + ); + } + Ok(()) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cooldown_trips_after_threshold_of_consecutive_transients() { + let s = CooldownState::default(); + for _ in 0..(COOLDOWN_TRANSIENT_THRESHOLD - 1) { + s.record(Classification::ConnectError, None); + assert!(!s.triggered(), "below threshold should not trigger"); + } + s.record(Classification::Server5xx, None); + assert!(s.triggered(), "Nth transient should trigger"); + } + + #[test] + fn cooldown_resets_on_success() { + let s = CooldownState::default(); + s.record(Classification::ConnectError, None); + s.record(Classification::Timeout, None); + s.record(Classification::Ok, None); + for _ in 0..(COOLDOWN_TRANSIENT_THRESHOLD - 1) { + s.record(Classification::RateLimited, None); + assert!(!s.triggered()); + } + s.record(Classification::ConnectError, None); + assert!(s.triggered()); + } + + #[test] + fn cooldown_duration_uses_retry_after_when_larger() { + let s = CooldownState::default(); + assert_eq!(s.cooldown_duration_secs(), COOLDOWN_DURATION_SECS); + s.record(Classification::RateLimited, Some(10)); + assert_eq!(s.cooldown_duration_secs(), COOLDOWN_DURATION_SECS); + s.record(Classification::RateLimited, Some(180)); + assert_eq!(s.cooldown_duration_secs(), 180); + s.record(Classification::RateLimited, Some(420)); + assert_eq!(s.cooldown_duration_secs(), 420); + s.record(Classification::RateLimited, Some(30)); + assert_eq!(s.cooldown_duration_secs(), 420); + } + + #[test] + fn cooldown_duration_caps_pathological_retry_after() { + let s = CooldownState::default(); + s.record(Classification::RateLimited, Some(99_999_999)); + assert_eq!(s.cooldown_duration_secs(), COOLDOWN_MAX_SECS); + } +} diff --git a/mini/src/storage/commit.rs b/mini/src/storage/commit.rs new file mode 100644 index 0000000..c8d1e5b --- /dev/null +++ b/mini/src/storage/commit.rs @@ -0,0 +1,163 @@ +//! The signed repo commit at `C\0` in the `default` CF. +//! +//! Written once per successful fetch from `getRepo` (in the same final +//! batch as the new `R\0` and `s\0` rows). The +//! parsed fields cover everything an offline authenticity check needs — +//! `did`, `version`, `data`, `rev`, `prev`, `sig` — with the spec's +//! canonical-DAG-CBOR contract (§2.7) doing the byte-fidelity work +//! implicitly: a verifier re-canonicalizes commit-minus-sig from these +//! fields, hashes, and checks the signature against the recovered key. +//! If a PDS published non-canonical signed bytes, signature +//! verification *should* fail — that's the spec working as designed. + +use hubble_pds::Hostname; +use rocksdb::WriteBatch; +use serde::{Deserialize, Serialize}; + +use super::{Db, DbError, PREFIX_COMMIT, Result, SEP, assert_no_nul}; + +// --------------------------------------------------------------------------- +// Key encoding +// --------------------------------------------------------------------------- + +pub fn key(pds: &Hostname, did: &str) -> Vec { + assert_no_nul(did); + let pds = pds.as_ref(); + let mut k = Vec::with_capacity(1 + pds.len() + 1 + did.len()); + k.push(PREFIX_COMMIT); + k.extend_from_slice(pds.as_bytes()); + k.push(SEP); + k.extend_from_slice(did.as_bytes()); + k +} + +// --------------------------------------------------------------------------- +// Type +// --------------------------------------------------------------------------- + +/// All six signed-commit fields preserved from the last successful +/// fetch. `data` and `prev` are stored as base32-multibase CID +/// strings; `sig` is the raw signature bytes. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct StoredCommit { + pub did: String, + pub version: u64, + /// MST root CID, base32-multibase. + pub data: String, + /// TID-shaped revision; same lex-monotonic semantics as + /// `RepoState.last_seen_rev`. + pub rev: String, + /// Previous commit CID, base32-multibase. Spec §2.4: required field + /// on wire, typically null in v3. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prev: Option, + /// Raw signature bytes, as captured from the wire commit block. + pub sig: Vec, +} + +// --------------------------------------------------------------------------- +// Encode / decode +// --------------------------------------------------------------------------- + +fn encode(c: &StoredCommit) -> Result> { + serde_ipld_dagcbor::to_vec(c).map_err(|e| DbError::CborEncode(e.to_string())) +} + +fn decode(pds: &Hostname, did: &str, bytes: &[u8]) -> Result { + serde_ipld_dagcbor::from_slice(bytes).map_err(|e| DbError::CborDecode { + key: format!("C{}\\0{}", pds.as_ref(), did), + reason: e.to_string(), + }) +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +pub fn get(db: &Db, pds: &Hostname, did: &str) -> Result> { + let cf = db.cf_default()?; + match db.inner.get_cf(cf, key(pds, did))? { + None => Ok(None), + Some(bytes) => Ok(Some(decode(pds, did, &bytes)?)), + } +} + +/// Stage a write into `batch`. +pub fn put_into( + batch: &mut WriteBatch, + db: &Db, + pds: &Hostname, + did: &str, + commit: &StoredCommit, +) -> Result<()> { + let cf = db.cf_default()?; + batch.put_cf(cf, key(pds, did), encode(commit)?); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::open_temporary; + + fn h(s: &str) -> Hostname { + Hostname::new(s).unwrap() + } + + fn sample() -> StoredCommit { + StoredCommit { + did: "did:plc:abc".into(), + version: 3, + data: "bafyreidkjxqphmt3yhvdgmwx5p64q7s7tcmw3rj7ynvtu7tnp7nb6yurxe".into(), + rev: "3kabc123".into(), + prev: None, + sig: vec![0xde, 0xad, 0xbe, 0xef, 0x00, 0x01, 0x02, 0x03], + } + } + + #[test] + fn key_uses_c_prefix() { + let k = key(&h("example.com"), "did:plc:abc"); + assert_eq!(k[0], b'C'); + assert_eq!(&k[1..12], b"example.com"); + assert_eq!(k[12], 0); + assert_eq!(&k[13..], b"did:plc:abc"); + } + + #[test] + fn roundtrip_via_batch() { + let (db, _g) = open_temporary(); + let pds = h("example.com"); + let did = "did:plc:abc"; + let c = sample(); + let mut batch = WriteBatch::default(); + put_into(&mut batch, &db, &pds, did, &c).unwrap(); + db.inner.write(batch).unwrap(); + let got = get(&db, &pds, did).unwrap().unwrap(); + assert_eq!(got, c); + } + + #[test] + fn missing_is_none() { + let (db, _g) = open_temporary(); + assert!(get(&db, &h("example.com"), "did:plc:nope").unwrap().is_none()); + } + + #[test] + fn prev_round_trips_when_present() { + let (db, _g) = open_temporary(); + let pds = h("example.com"); + let did = "did:plc:abc"; + let mut c = sample(); + c.prev = Some("bafyreichainparentexamplecidstringalpha2345678901234567".into()); + let mut batch = WriteBatch::default(); + put_into(&mut batch, &db, &pds, did, &c).unwrap(); + db.inner.write(batch).unwrap(); + let got = get(&db, &pds, did).unwrap().unwrap(); + assert_eq!(got.prev, c.prev); + } +} diff --git a/mini/src/storage/host.rs b/mini/src/storage/host.rs new file mode 100644 index 0000000..44e414a --- /dev/null +++ b/mini/src/storage/host.rs @@ -0,0 +1,197 @@ +//! Per-PDS-host state at `h` in the `default` CF. +//! +//! Value is a dag-cbor-encoded [`HostState`]. Fields are added +//! incrementally; existing values round-trip via `#[serde(default)]` on +//! every Optional field. + +use hubble_pds::Hostname; +use serde::{Deserialize, Serialize}; + +use super::{Db, DbError, PREFIX_HOST, Result, now_unix}; + +// --------------------------------------------------------------------------- +// Key encoding +// --------------------------------------------------------------------------- + +pub(crate) fn key(host: &Hostname) -> Vec { + let host = host.as_ref(); + let mut k = Vec::with_capacity(1 + host.len()); + k.push(PREFIX_HOST); + k.extend_from_slice(host.as_bytes()); + k +} + +// --------------------------------------------------------------------------- +// Type +// --------------------------------------------------------------------------- + +/// Discovery- and snapshot-related state recorded against a PDS host. +/// +/// Optional fields are populated incrementally as signals arrive. New +/// fields added later land here with `#[serde(default)]` so old values +/// round-trip. +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct HostState { + pub discovered_at_unix: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub listhosts_seen_at_unix: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub listhosts_status: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub listhosts_account_count: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scraping_seen_at_unix: Option, +} + +// --------------------------------------------------------------------------- +// Encode / decode +// --------------------------------------------------------------------------- + +fn encode(state: &HostState) -> Result> { + serde_ipld_dagcbor::to_vec(state).map_err(|e| DbError::CborEncode(e.to_string())) +} + +fn decode(host: &Hostname, bytes: &[u8]) -> Result { + serde_ipld_dagcbor::from_slice(bytes).map_err(|e| DbError::CborDecode { + key: format!("h{}", host.as_ref()), + reason: e.to_string(), + }) +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +/// Read the current `HostState` for `host`, if any. +pub fn get(db: &Db, host: &Hostname) -> Result> { + let cf = db.cf_default()?; + match db.inner.get_cf(cf, key(host))? { + None => Ok(None), + Some(bytes) => Ok(Some(decode(host, &bytes)?)), + } +} + +fn put(db: &Db, host: &Hostname, state: &HostState) -> Result<()> { + let cf = db.cf_default()?; + db.inner.put_cf(cf, key(host), encode(state)?)?; + Ok(()) +} + +/// Apply listHosts-derived fields. Creates the row if absent (with +/// `discovered_at_unix = now`); otherwise updates only the listhosts fields. +pub fn upsert_from_listhosts( + db: &Db, + host: &Hostname, + status: Option<&str>, + account_count: Option, +) -> Result<()> { + let now = now_unix(); + let mut state = get(db, host)?.unwrap_or(HostState { + discovered_at_unix: now, + ..Default::default() + }); + state.listhosts_seen_at_unix = Some(now); + state.listhosts_status = status.map(str::to_string); + state.listhosts_account_count = account_count; + put(db, host, &state) +} + +/// Apply scraping-derived fields. Creates the row if absent. +pub fn upsert_from_scraping(db: &Db, host: &Hostname) -> Result<()> { + let now = now_unix(); + let mut state = get(db, host)?.unwrap_or(HostState { + discovered_at_unix: now, + ..Default::default() + }); + state.scraping_seen_at_unix = Some(now); + put(db, host, &state) +} + +/// Read every `(Hostname, HostState)` pair currently stored. Hostnames +/// returned in lex order; safe to filter / shuffle at the call site. +pub fn list_all(db: &Db) -> Result> { + let cf = db.cf_default()?; + let mut read_opts = rocksdb::ReadOptions::default(); + read_opts.set_iterate_lower_bound([PREFIX_HOST]); + read_opts.set_iterate_upper_bound([PREFIX_HOST + 1]); + let iter = db + .inner + .iterator_cf_opt(cf, read_opts, rocksdb::IteratorMode::Start); + + let mut out = Vec::new(); + for item in iter { + let (k, v) = item?; + let host_bytes = &k[1..]; // strip prefix + let host_str = std::str::from_utf8(host_bytes).map_err(|e| DbError::CborDecode { + key: "h".to_string(), + reason: format!("hostname not utf-8: {e}"), + })?; + let hostname = Hostname::new(host_str).map_err(|e| DbError::CborDecode { + key: format!("h{}", host_str), + reason: format!("hostname parse: {e}"), + })?; + let state = decode(&hostname, &v)?; + out.push((hostname, state)); + } + Ok(out) +} + +/// Insert a host with a fresh `discovered_at_unix` if absent; otherwise +/// no-op. Returns `true` iff a row was inserted. +pub fn upsert_manual(db: &Db, host: &Hostname) -> Result { + if get(db, host)?.is_some() { + return Ok(false); + } + let state = HostState { + discovered_at_unix: now_unix(), + ..Default::default() + }; + put(db, host, &state)?; + Ok(true) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::open_temporary; + + fn h(s: &str) -> Hostname { + Hostname::new(s).unwrap() + } + + #[test] + fn key_uses_h_prefix() { + assert_eq!(key(&h("example.com")), b"hexample.com"); + } + + #[test] + fn missing_is_none() { + let (db, _g) = open_temporary(); + assert!(get(&db, &h("missing.example")).unwrap().is_none()); + } + + #[test] + fn upsert_manual_is_idempotent() { + let (db, _g) = open_temporary(); + let host = h("example.com"); + assert!(upsert_manual(&db, &host).unwrap(), "first insert"); + assert!(!upsert_manual(&db, &host).unwrap(), "second is noop"); + } + + #[test] + fn upsert_listhosts_then_scraping_merges_fields() { + let (db, _g) = open_temporary(); + let host = h("example.com"); + upsert_from_listhosts(&db, &host, Some("active"), Some(123)).unwrap(); + upsert_from_scraping(&db, &host).unwrap(); + let s = get(&db, &host).unwrap().unwrap(); + assert_eq!(s.listhosts_status.as_deref(), Some("active")); + assert_eq!(s.listhosts_account_count, Some(123)); + assert!(s.listhosts_seen_at_unix.is_some()); + assert!(s.scraping_seen_at_unix.is_some()); + } +} diff --git a/mini/src/storage/mod.rs b/mini/src/storage/mod.rs new file mode 100644 index 0000000..c9fa4f9 --- /dev/null +++ b/mini/src/storage/mod.rs @@ -0,0 +1,181 @@ +//! RocksDB storage for hubble-mini. +//! +//! Two column families, with the layout pinned in `local/mini-proposal.md`: +//! - `records` CF: lz4 + zstd-3 bottommost, 32 KiB blocks, no dict +//! - `default` CF: lz4 + zstd-3 bottommost, default 4 KiB blocks, no dict +//! +//! Both use 4 subcompactions. Single-byte ASCII prefixes; composite keys +//! separate caller-supplied components with `\x00`. +//! +//! Each typed value lives in its own submodule that owns the key +//! encoding, the value type, encode/decode, and the get/put surface — the +//! pattern from `lightrail/src/storage/`. Submodules are added as new +//! typed values come online; the prefix registry below documents all of +//! them. + +pub mod commit; +pub mod host; +pub mod probe; +pub mod record; +pub mod repo_state; +pub mod repo_stats; + +use std::path::Path; +use std::sync::Arc; + +use rocksdb::{ + BlockBasedOptions, ColumnFamily, ColumnFamilyDescriptor, DB, DBCompressionType, Options, +}; + +#[derive(Debug, thiserror::Error)] +pub enum DbError { + #[error("rocksdb: {0}")] + Rocks(#[from] rocksdb::Error), + + #[error("cbor encode: {0}")] + CborEncode(String), + + #[error("cbor decode at key {key:?}: {reason}")] + CborDecode { key: String, reason: String }, + + #[error("column family `{0}` missing")] + MissingCf(&'static str), +} + +pub type Result = std::result::Result; + +pub(crate) const CF_RECORDS: &str = "records"; +pub(crate) const CF_DEFAULT: &str = "default"; + +/// Prefix registry. Each prefix below is one byte, used as the first byte +/// of every key in its key family. Submodules under `storage::` own the +/// key encoder/decoder, value type, and CRUD for each prefix; this list +/// is the index. +/// +/// | prefix | cf | typed by | +/// |--------|----------|--------------------------| +/// | `r` | records | [`record`] | +/// | `h` | default | [`host`] | +/// | `p` | default | [`probe`] | +/// | `R` | default | [`repo_state`] | +/// | `s` | default | [`repo_stats`] | +/// | `C` | default | [`commit`] | +/// +/// (No `m`/meta prefix yet — re-introduce when there's a concrete +/// process-level value to store, e.g. schema_version, last_run_started_at.) +pub(crate) const PREFIX_RECORD: u8 = b'r'; +pub(crate) const PREFIX_HOST: u8 = b'h'; +pub(crate) const PREFIX_PROBE: u8 = b'p'; +pub(crate) const PREFIX_REPO_STATE: u8 = b'R'; +pub(crate) const PREFIX_REPO_STATS: u8 = b's'; +pub(crate) const PREFIX_COMMIT: u8 = b'C'; + +pub(crate) const SEP: u8 = 0x00; + +pub(crate) fn assert_no_nul(part: &str) { + assert!( + !part.as_bytes().contains(&SEP), + "rocksdb key component contains NUL: {part:?}" + ); +} + +/// Shared handle to the rocksdb database. Cheap to clone (it's an `Arc`). +/// +/// `SingleThreaded` mode is fine here: mini opens all column families +/// once at startup and never creates them at runtime, and rocksdb's get +/// /put surface is `Send + Sync` regardless of the thread-mode marker. +/// `MultiThreaded` only buys lock-free dynamic CF management, which we +/// don't need. +pub struct Db { + pub(crate) inner: Arc, +} + +#[allow(dead_code)] // used by snapshot workers once they land +pub type DbRef = Arc; + +impl Db { + pub fn open(path: &Path) -> Result { + let mut db_opts = Options::default(); + db_opts.create_if_missing(true); + db_opts.create_missing_column_families(true); + + let cfs = vec![ + ColumnFamilyDescriptor::new(CF_DEFAULT, default_cf_opts()), + ColumnFamilyDescriptor::new(CF_RECORDS, records_cf_opts()), + ]; + + let inner = DB::open_cf_descriptors(&db_opts, path, cfs)?; + Ok(Self { inner: Arc::new(inner) }) + } + + pub(crate) fn cf_default(&self) -> Result<&ColumnFamily> { + self.inner + .cf_handle(CF_DEFAULT) + .ok_or(DbError::MissingCf(CF_DEFAULT)) + } + + #[allow(dead_code)] // used by record.rs once it lands + pub(crate) fn cf_records(&self) -> Result<&ColumnFamily> { + self.inner + .cf_handle(CF_RECORDS) + .ok_or(DbError::MissingCf(CF_RECORDS)) + } +} + +/// Current Unix time in seconds. Saturates to 0 if the system clock is +/// before the epoch (shouldn't happen, but `SystemTime` is allowed to be). +pub(crate) fn now_unix() -> u64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +fn default_cf_opts() -> Options { + let mut o = Options::default(); + o.set_compression_type(DBCompressionType::Lz4); + o.set_bottommost_compression_type(DBCompressionType::Zstd); + // window_bits=-14 is rocksdb's documented zstd default; level=3; + // strategy=0; max_dict_bytes=0 (no compression dictionary); enabled=true. + o.set_bottommost_compression_options(-14, 3, 0, 0, true); + o.set_max_subcompactions(4); + o +} + +fn records_cf_opts() -> Options { + let mut o = default_cf_opts(); + let mut block_opts = BlockBasedOptions::default(); + block_opts.set_block_size(32 * 1024); + o.set_block_based_table_factory(&block_opts); + o +} + +#[cfg(test)] +pub(crate) fn open_temporary() -> (Db, tempfile::TempDir) { + let dir = tempfile::tempdir().expect("tempdir"); + let db = Db::open(dir.path()).expect("open rocks"); + (db, dir) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn open_creates_and_reopens_with_both_cfs() { + let dir = tempfile::tempdir().unwrap(); + { + let _db = Db::open(dir.path()).unwrap(); + } + // Reopen must succeed without recreating CFs. + let _db = Db::open(dir.path()).unwrap(); + } + + #[test] + fn cf_handles_resolve() { + let (db, _g) = open_temporary(); + let _ = db.cf_default().unwrap(); + let _ = db.cf_records().unwrap(); + } +} diff --git a/mini/src/storage/probe.rs b/mini/src/storage/probe.rs new file mode 100644 index 0000000..e1cf0fe --- /dev/null +++ b/mini/src/storage/probe.rs @@ -0,0 +1,99 @@ +//! Per-PDS probe bodies at `p` in the `default` CF. +//! +//! We capture `_health` + `describeServer` responses verbatim at each +//! snapshot pass. Classification (is-this-PDS-trustworthy?) is deferred — +//! the bodies are kept for later analysis. + +use hubble_pds::Hostname; +use serde::{Deserialize, Serialize}; + +use super::{Db, DbError, PREFIX_PROBE, Result}; + +// --------------------------------------------------------------------------- +// Key encoding +// --------------------------------------------------------------------------- + +pub fn key(pds: &Hostname) -> Vec { + let pds = pds.as_ref(); + let mut k = Vec::with_capacity(1 + pds.len()); + k.push(PREFIX_PROBE); + k.extend_from_slice(pds.as_bytes()); + k +} + +// --------------------------------------------------------------------------- +// Type +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct Probe { + pub probed_at_unix: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub health_json: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub describe_json: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub probe_error: Option, +} + +// --------------------------------------------------------------------------- +// Encode / decode +// --------------------------------------------------------------------------- + +fn encode(probe: &Probe) -> Result> { + serde_ipld_dagcbor::to_vec(probe).map_err(|e| DbError::CborEncode(e.to_string())) +} + +#[cfg(test)] +fn decode(pds: &Hostname, bytes: &[u8]) -> Result { + serde_ipld_dagcbor::from_slice(bytes).map_err(|e| DbError::CborDecode { + key: format!("p{}", pds.as_ref()), + reason: e.to_string(), + }) +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +pub fn put(db: &Db, pds: &Hostname, probe: &Probe) -> Result<()> { + let cf = db.cf_default()?; + db.inner.put_cf(cf, key(pds), encode(probe)?)?; + Ok(()) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::open_temporary; + + fn h(s: &str) -> Hostname { + Hostname::new(s).unwrap() + } + + #[test] + fn key_uses_p_prefix() { + assert_eq!(key(&h("example.com")), b"pexample.com"); + } + + #[test] + fn roundtrip() { + let (db, _g) = open_temporary(); + let pds = h("example.com"); + let p = Probe { + probed_at_unix: 1_700_000_000, + health_json: Some(r#"{"version":"0.4.0"}"#.into()), + describe_json: Some(r#"{"availableUserDomains":[".bsky.social"]}"#.into()), + probe_error: None, + }; + put(&db, &pds, &p).unwrap(); + let cf = db.cf_default().unwrap(); + let raw = db.inner.get_cf(cf, key(&pds)).unwrap().unwrap(); + let got = decode(&pds, &raw).unwrap(); + assert_eq!(got, p); + } +} diff --git a/mini/src/storage/record.rs b/mini/src/storage/record.rs new file mode 100644 index 0000000..c54322e --- /dev/null +++ b/mini/src/storage/record.rs @@ -0,0 +1,238 @@ +//! Per-record bytes in the `records` CF at `r\0\0`. +//! +//! `` is the literal MST key from the source repo — +//! `/` byte-for-byte. Using `/` (not `\0`) between +//! collection and rkey means rocksdb's byte order over the records range +//! for one `(pds, did)` matches MST lex order, so MST verification +//! and any future stream-ordered export can iterate directly without +//! re-sorting. +//! +//! Values are the raw DAG-CBOR record bytes as they appeared in the +//! source CAR — no struct, no wrapper. Block compression handles +//! redundancy on disk. +//! +//! Records are written through a [`rocksdb::WriteBatch`] so a repo +//! fetch can atomically commit a chunk of records + the trailing +//! `R\0` and `s\0` state in one shot. Use +//! [`delete_range_into`] before writing a new set of records for a +//! repo — range-delete then puts in the same batch is safe (puts win +//! for keys covered by the delete-range tombstone, applied in batch +//! order by rocksdb). + +use hubble_pds::Hostname; +use rocksdb::WriteBatch; + +use super::{Db, PREFIX_RECORD, Result, SEP, assert_no_nul}; + +// --------------------------------------------------------------------------- +// Key encoding +// --------------------------------------------------------------------------- + +/// Build the `records` CF key for one record. +/// +/// `mst_key` is the MST key as repo-stream emits it: `/`, +/// containing exactly one `/` between an NSID-shaped collection and the +/// record key. We don't split it — the `/` is part of the stored key so +/// rocksdb byte-order = MST lex order within a single `(pds, did)`. +pub fn key(pds: &Hostname, did: &str, mst_key: &str) -> Vec { + assert_no_nul(did); + assert_no_nul(mst_key); + let pds = pds.as_ref(); + let mut k = Vec::with_capacity(1 + pds.len() + 1 + did.len() + 1 + mst_key.len()); + k.push(PREFIX_RECORD); + k.extend_from_slice(pds.as_bytes()); + k.push(SEP); + k.extend_from_slice(did.as_bytes()); + k.push(SEP); + k.extend_from_slice(mst_key.as_bytes()); + k +} + +/// Inclusive lower and exclusive upper bounds covering every record key +/// for `(pds, did)`. Use with `delete_range_cf` to wipe a repo. +pub fn range_for_repo(pds: &Hostname, did: &str) -> (Vec, Vec) { + assert_no_nul(did); + let pds = pds.as_ref(); + let mut lo = Vec::with_capacity(1 + pds.len() + 1 + did.len() + 1); + lo.push(PREFIX_RECORD); + lo.extend_from_slice(pds.as_bytes()); + lo.push(SEP); + lo.extend_from_slice(did.as_bytes()); + lo.push(SEP); + let mut hi = lo.clone(); + *hi.last_mut().expect("non-empty") = SEP + 1; + (lo, hi) +} + +// --------------------------------------------------------------------------- +// Batch helpers +// --------------------------------------------------------------------------- + +/// Stage one record write into `batch`. +pub fn put_into( + batch: &mut WriteBatch, + db: &Db, + pds: &Hostname, + did: &str, + mst_key: &str, + value: &[u8], +) -> Result<()> { + let cf = db.cf_records()?; + batch.put_cf(cf, key(pds, did, mst_key), value); + Ok(()) +} + +/// Stage a `delete_range_cf` over every record key for `(pds, did)`. +/// Subsequent `put_into` calls for the same repo within the same batch +/// overwrite the tombstone for those specific keys. +pub fn delete_range_into( + batch: &mut WriteBatch, + db: &Db, + pds: &Hostname, + did: &str, +) -> Result<()> { + let cf = db.cf_records()?; + let (lo, hi) = range_for_repo(pds, did); + batch.delete_range_cf(cf, lo, hi); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::open_temporary; + + fn h(s: &str) -> Hostname { + Hostname::new(s).unwrap() + } + + #[test] + fn key_layout_uses_slash_inside_mst_key() { + let k = key(&h("example.com"), "did:plc:abc", "app.bsky.feed.post/3k"); + assert_eq!(k[0], b'r'); + let rest = &k[1..]; + // r example.com \0 did:plc:abc \0 app.bsky.feed.post/3k + let nul_count = rest.iter().filter(|&&b| b == 0).count(); + assert_eq!(nul_count, 2, "exactly two \\0 separators (pds|did, did|mst_key)"); + assert!(rest.windows(b"app.bsky.feed.post/3k".len()) + .any(|w| w == b"app.bsky.feed.post/3k")); + } + + #[test] + fn rocksdb_iter_order_matches_mst_lex_order() { + // Within one (pds, did), byte-order over the `r..` range must equal + // MST lex order over the (coll, rkey) keys. The two collection + // strings here would have swapped order under the old `\0` + // separator: `\0` (0x00) < `.` (0x2E) put `app.bsky.feed.post` + // before `app.bsky.feed.post.x`. With `/` (0x2F) > `.`, the order + // flips to MST-correct: `app.bsky.feed.post.x/...` comes first. + let (db, _g) = open_temporary(); + let pds = h("example.com"); + let did = "did:plc:abc"; + + let mut batch = WriteBatch::default(); + put_into(&mut batch, &db, &pds, did, "app.bsky.feed.post/z", b"A").unwrap(); + put_into(&mut batch, &db, &pds, did, "app.bsky.feed.post.x/a", b"B").unwrap(); + db.inner.write(batch).unwrap(); + + let cf = db.cf_records().unwrap(); + let (lo, hi) = range_for_repo(&pds, did); + let mut read_opts = rocksdb::ReadOptions::default(); + read_opts.set_iterate_lower_bound(lo); + read_opts.set_iterate_upper_bound(hi); + let iter = db + .inner + .iterator_cf_opt(cf, read_opts, rocksdb::IteratorMode::Start); + let mut mst_keys: Vec = Vec::new(); + for item in iter { + let (k, _v) = item.unwrap(); + // Extract the mst_key suffix: everything after the second \0. + let mut splits = k.splitn(3, |b| *b == 0); + let _prefix_and_pds = splits.next().unwrap(); + let _did = splits.next().unwrap(); + let mst_key = splits.next().unwrap(); + mst_keys.push(String::from_utf8(mst_key.to_vec()).unwrap()); + } + // Pure-string lex sort over the MST keys should match the order + // rocksdb emitted them in. + let mut sorted = mst_keys.clone(); + sorted.sort(); + assert_eq!(mst_keys, sorted); + assert_eq!( + mst_keys, + vec![ + "app.bsky.feed.post.x/a".to_string(), + "app.bsky.feed.post/z".to_string(), + ] + ); + } + + #[test] + fn range_bounds_cover_only_one_repo() { + let pds = h("example.com"); + let (lo, hi) = range_for_repo(&pds, "did:plc:abc"); + let same = key(&pds, "did:plc:abc", "app.bsky.feed.post/x"); + assert!(lo.as_slice() <= same.as_slice() && same.as_slice() < hi.as_slice()); + // A different did is outside the range. + let other = key(&pds, "did:plc:abd", "app.bsky.feed.post/x"); + assert!(other.as_slice() >= hi.as_slice()); + } + + #[test] + #[should_panic(expected = "NUL")] + fn key_panics_on_nul_in_did() { + let _ = key(&h("example.com"), "did:\0bad", "app.bsky.feed.post/x"); + } + + #[test] + fn write_then_range_delete_in_one_batch() { + let (db, _g) = open_temporary(); + let pds = h("example.com"); + let did = "did:plc:abc"; + + let mut b = WriteBatch::default(); + put_into(&mut b, &db, &pds, did, "app.bsky.feed.post/x", b"AAA").unwrap(); + put_into(&mut b, &db, &pds, did, "app.bsky.feed.post/y", b"BBB").unwrap(); + put_into(&mut b, &db, &pds, "did:plc:other", "app.bsky.feed.post/z", b"CCC") + .unwrap(); + db.inner.write(b).unwrap(); + + let cf = db.cf_records().unwrap(); + assert!( + db.inner + .get_cf(cf, key(&pds, did, "app.bsky.feed.post/x")) + .unwrap() + .is_some() + ); + + let mut b = WriteBatch::default(); + delete_range_into(&mut b, &db, &pds, did).unwrap(); + put_into(&mut b, &db, &pds, did, "app.bsky.feed.post/x", b"NEW").unwrap(); + db.inner.write(b).unwrap(); + + assert!( + db.inner + .get_cf(cf, key(&pds, did, "app.bsky.feed.post/y")) + .unwrap() + .is_none() + ); + assert_eq!( + db.inner + .get_cf(cf, key(&pds, did, "app.bsky.feed.post/x")) + .unwrap() + .as_deref(), + Some(&b"NEW"[..]) + ); + assert_eq!( + db.inner + .get_cf(cf, key(&pds, "did:plc:other", "app.bsky.feed.post/z")) + .unwrap() + .as_deref(), + Some(&b"CCC"[..]) + ); + } +} diff --git a/mini/src/storage/repo_state.rs b/mini/src/storage/repo_state.rs new file mode 100644 index 0000000..4376703 --- /dev/null +++ b/mini/src/storage/repo_state.rs @@ -0,0 +1,232 @@ +//! Per-repo state at `R\0` in the `default` CF. +//! +//! Value is a dag-cbor-encoded [`RepoState`]. Keyed by `(pds, did)` so the +//! same DID can carry independent state under multiple PDSes — we never +//! cross-host de-dup. + +use hubble_pds::Hostname; +use rocksdb::WriteBatch; +use serde::{Deserialize, Serialize}; + +use super::{Db, DbError, PREFIX_REPO_STATE, Result, SEP, assert_no_nul}; + +// --------------------------------------------------------------------------- +// Key encoding +// --------------------------------------------------------------------------- + +pub fn key(pds: &Hostname, did: &str) -> Vec { + assert_no_nul(did); + let pds = pds.as_ref(); + let mut k = Vec::with_capacity(1 + pds.len() + 1 + did.len()); + k.push(PREFIX_REPO_STATE); + k.extend_from_slice(pds.as_bytes()); + k.push(SEP); + k.extend_from_slice(did.as_bytes()); + k +} + +/// Inclusive lower and exclusive upper bounds covering every `R\0...` +/// key for one host. Used by the tombstone-detection prefix scan. +pub fn range_for_pds(pds: &Hostname) -> (Vec, Vec) { + let pds = pds.as_ref(); + let mut lo = Vec::with_capacity(1 + pds.len() + 1); + lo.push(PREFIX_REPO_STATE); + lo.extend_from_slice(pds.as_bytes()); + lo.push(SEP); + let mut hi = lo.clone(); + *hi.last_mut().expect("non-empty") = SEP + 1; + (lo, hi) +} + +// --------------------------------------------------------------------------- +// Type +// --------------------------------------------------------------------------- + +/// Per-repo state under a single PDS host. +/// +/// Fields are added as the snapshot loop matures. Optional fields use +/// `#[serde(default, skip_serializing_if = "Option::is_none")]` so old +/// values round-trip after schema additions. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct RepoState { + /// `listRepos.rev` last observed. TID-shaped; compared as a string. + /// Distinct from `StoredCommit.rev` (the commit-block's own rev): + /// the listRepos-reported value is what scheduling decisions key + /// on; the commit-block rev is what was actually fetched. + pub last_seen_rev: String, + /// Unix seconds when we last successfully fetched + wrote this repo. + pub last_fetched_at_unix: u64, + /// `listRepos.active` flag at last observation. + pub listrepos_active: bool, + /// `listRepos.status` string (set when `listrepos_active = false`). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub listrepos_status: Option, + /// `true` once this `(pds, did)` stopped appearing in `listRepos`. + /// Records are NOT deleted in v0; tombstoned repos are skipped on + /// future passes. + #[serde(default)] + pub tombstoned: bool, + /// Short summary of the last terminal fetch error, if any. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_error: Option, +} + +// --------------------------------------------------------------------------- +// Encode / decode +// --------------------------------------------------------------------------- + +fn encode(state: &RepoState) -> Result> { + serde_ipld_dagcbor::to_vec(state).map_err(|e| DbError::CborEncode(e.to_string())) +} + +fn decode(pds: &Hostname, did: &str, bytes: &[u8]) -> Result { + serde_ipld_dagcbor::from_slice(bytes).map_err(|e| DbError::CborDecode { + key: format!("R{}\\0{}", pds.as_ref(), did), + reason: e.to_string(), + }) +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +#[allow(dead_code)] // useful read API; consumers will land alongside future tooling +pub fn get(db: &Db, pds: &Hostname, did: &str) -> Result> { + let cf = db.cf_default()?; + match db.inner.get_cf(cf, key(pds, did))? { + None => Ok(None), + Some(bytes) => Ok(Some(decode(pds, did, &bytes)?)), + } +} + +/// Stage a write into `batch`. +pub fn put_into( + batch: &mut WriteBatch, + db: &Db, + pds: &Hostname, + did: &str, + state: &RepoState, +) -> Result<()> { + let cf = db.cf_default()?; + batch.put_cf(cf, key(pds, did), encode(state)?); + Ok(()) +} + +/// Standalone write (no batch). Used when we're updating just one repo's +/// state, e.g. tombstoning a disappeared repo. +pub fn put(db: &Db, pds: &Hostname, did: &str, state: &RepoState) -> Result<()> { + let cf = db.cf_default()?; + db.inner.put_cf(cf, key(pds, did), encode(state)?)?; + Ok(()) +} + +/// All `(did, RepoState)` pairs currently stored under `pds`. Used to +/// build the pre-pass DID set for tombstone detection. +pub fn dids_for_pds(db: &Db, pds: &Hostname) -> Result> { + let cf = db.cf_default()?; + let (lo, hi) = range_for_pds(pds); + let mut read_opts = rocksdb::ReadOptions::default(); + read_opts.set_iterate_lower_bound(lo); + read_opts.set_iterate_upper_bound(hi.clone()); + let iter = db + .inner + .iterator_cf_opt(cf, read_opts, rocksdb::IteratorMode::Start); + + let prefix_len = 1 + pds.as_ref().len() + 1; // 'R' + pds + '\0' + let mut out = Vec::new(); + for item in iter { + let (k, v) = item?; + let key_bytes: &[u8] = &k; + debug_assert!(key_bytes.starts_with(&hi[..hi.len() - 1])); // 'R'+pds+SEP + let did_bytes = &key_bytes[prefix_len..]; + let did = std::str::from_utf8(did_bytes) + .map_err(|e| DbError::CborDecode { + key: format!("R{}\\0", pds.as_ref()), + reason: format!("did not utf-8: {e}"), + })? + .to_string(); + let state = decode(pds, &did, &v)?; + out.push((did, state)); + } + Ok(out) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::open_temporary; + + fn h(s: &str) -> Hostname { + Hostname::new(s).unwrap() + } + + fn state(rev: &str) -> RepoState { + RepoState { + last_seen_rev: rev.into(), + last_fetched_at_unix: 1_700_000_000, + listrepos_active: true, + listrepos_status: None, + tombstoned: false, + last_error: None, + } + } + + #[test] + fn key_layout() { + let k = key(&h("example.com"), "did:plc:abc"); + assert_eq!(k[0], b'R'); + assert_eq!(&k[1..12], b"example.com"); + assert_eq!(k[12], 0); + assert_eq!(&k[13..], b"did:plc:abc"); + } + + #[test] + fn roundtrip_get_put() { + let (db, _g) = open_temporary(); + let pds = h("example.com"); + let s = state("rev1"); + put(&db, &pds, "did:plc:abc", &s).unwrap(); + assert_eq!(get(&db, &pds, "did:plc:abc").unwrap(), Some(s)); + } + + #[test] + fn put_into_then_write_commits() { + let (db, _g) = open_temporary(); + let pds = h("example.com"); + let mut batch = WriteBatch::default(); + put_into(&mut batch, &db, &pds, "did:plc:a", &state("ra")).unwrap(); + put_into(&mut batch, &db, &pds, "did:plc:b", &state("rb")).unwrap(); + db.inner.write(batch).unwrap(); + assert!(get(&db, &pds, "did:plc:a").unwrap().is_some()); + assert!(get(&db, &pds, "did:plc:b").unwrap().is_some()); + } + + #[test] + fn dids_for_pds_scoped_by_host() { + let (db, _g) = open_temporary(); + let a = h("a.example"); + let b = h("b.example"); + put(&db, &a, "did:plc:1", &state("r1")).unwrap(); + put(&db, &a, "did:plc:2", &state("r2")).unwrap(); + put(&db, &b, "did:plc:3", &state("r3")).unwrap(); + + let mut for_a: Vec = dids_for_pds(&db, &a) + .unwrap() + .into_iter() + .map(|(d, _)| d) + .collect(); + for_a.sort(); + assert_eq!(for_a, vec!["did:plc:1", "did:plc:2"]); + + let for_b: Vec = dids_for_pds(&db, &b) + .unwrap() + .into_iter() + .map(|(d, _)| d) + .collect(); + assert_eq!(for_b, vec!["did:plc:3"]); + } +} diff --git a/mini/src/storage/repo_stats.rs b/mini/src/storage/repo_stats.rs new file mode 100644 index 0000000..37d3b32 --- /dev/null +++ b/mini/src/storage/repo_stats.rs @@ -0,0 +1,155 @@ +//! Per-repo measurement stats at `s\0` in the `default` CF. +//! +//! Written once per successful repo fetch (in the same final batch as +//! the new `R\0` state). + +use std::collections::BTreeMap; + +use hubble_pds::Hostname; +use rocksdb::WriteBatch; +use serde::{Deserialize, Serialize}; + +use super::{Db, DbError, PREFIX_REPO_STATS, Result, SEP, assert_no_nul}; + +// --------------------------------------------------------------------------- +// Key encoding +// --------------------------------------------------------------------------- + +pub fn key(pds: &Hostname, did: &str) -> Vec { + assert_no_nul(did); + let pds = pds.as_ref(); + let mut k = Vec::with_capacity(1 + pds.len() + 1 + did.len()); + k.push(PREFIX_REPO_STATS); + k.extend_from_slice(pds.as_bytes()); + k.push(SEP); + k.extend_from_slice(did.as_bytes()); + k +} + +// --------------------------------------------------------------------------- +// Type +// --------------------------------------------------------------------------- + +/// Per-repo measurement totals from the most recent successful fetch. +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct RepoStats { + /// `Content-Encoding` the server actually used on the response + /// (literal wire name: `"gzip"` / `"br"` / `"zstd"` / `"deflate"`), + /// or `None` for identity. Stored as a string rather than the + /// `KnownEncoding` enum so the schema is stable across future + /// enum changes. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub content_encoding: Option, + /// Raw HTTP body bytes received (before any `Content-Encoding` + /// decode — the on-the-wire CAR size). + pub wire_bytes: u64, + /// Decoded CAR bytes (after gzip/br/zstd decompression). + pub car_bytes: u64, + /// Number of records walked from the MST. + pub record_count: u64, + /// Total link references found inside records (per + /// `microcosm-links`). + pub backlink_count: u64, + /// Total blob references inside records. + pub blob_ref_count: u64, + /// Sum of `size` fields across blob refs that declared one. + pub blob_total_size: u64, + /// Blob refs that omitted a numeric `size`. + pub unsized_blob_count: u64, + /// Counts of blob refs grouped by `mimeType`. Refs with no + /// `mimeType` are not represented here; see `blob_no_mime_count`. + pub blob_mime_counts: BTreeMap, + /// Count of blob refs whose `mimeType` was missing or null. + pub blob_no_mime_count: u64, +} + +// --------------------------------------------------------------------------- +// Encode / decode +// --------------------------------------------------------------------------- + +fn encode(stats: &RepoStats) -> Result> { + serde_ipld_dagcbor::to_vec(stats).map_err(|e| DbError::CborEncode(e.to_string())) +} + +#[cfg(test)] +fn decode(pds: &Hostname, did: &str, bytes: &[u8]) -> Result { + serde_ipld_dagcbor::from_slice(bytes).map_err(|e| DbError::CborDecode { + key: format!("s{}\\0{}", pds.as_ref(), did), + reason: e.to_string(), + }) +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +/// Stage a write into `batch`. +pub fn put_into( + batch: &mut WriteBatch, + db: &Db, + pds: &Hostname, + did: &str, + stats: &RepoStats, +) -> Result<()> { + let cf = db.cf_default()?; + batch.put_cf(cf, key(pds, did), encode(stats)?); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::open_temporary; + + fn h(s: &str) -> Hostname { + Hostname::new(s).unwrap() + } + + fn sample() -> RepoStats { + let mut by_mime = BTreeMap::new(); + by_mime.insert("image/jpeg".to_string(), 5); + by_mime.insert("video/mp4".to_string(), 1); + RepoStats { + content_encoding: Some("zstd".into()), + wire_bytes: 1_234_567, + car_bytes: 2_345_678, + record_count: 42, + backlink_count: 17, + blob_ref_count: 8, + blob_total_size: 9_000_000, + unsized_blob_count: 1, + blob_mime_counts: by_mime, + blob_no_mime_count: 2, + } + } + + #[test] + fn key_layout() { + let k = key(&h("example.com"), "did:plc:abc"); + assert_eq!(k[0], b's'); + assert_eq!(&k[1..12], b"example.com"); + assert_eq!(k[12], 0); + assert_eq!(&k[13..], b"did:plc:abc"); + } + + #[test] + fn roundtrip_via_batch() { + let (db, _g) = open_temporary(); + let pds = h("example.com"); + let did = "did:plc:abc"; + + let stats = sample(); + let mut batch = WriteBatch::default(); + put_into(&mut batch, &db, &pds, did, &stats).unwrap(); + db.inner.write(batch).unwrap(); + + let cf = db.cf_default().unwrap(); + let raw = db.inner.get_cf(cf, key(&pds, did)).unwrap().unwrap(); + let got = decode(&pds, did, &raw).unwrap(); + assert_eq!(got, stats); + } +}