diff --git a/hubble-sync/src/config.rs b/hubble-sync/src/config.rs index 1655d8b..f40e409 100644 --- a/hubble-sync/src/config.rs +++ b/hubble-sync/src/config.rs @@ -114,6 +114,13 @@ pub struct SyncConfig { /// /// default: SyncScope::Everything pub sync_scope: SyncScope, + /// start a task to route a `ReconcileScope` message to every + /// [`SyncStatus::OutOfScope`] repo. + /// + /// one-shot at startup. + /// + /// default: false + pub rescope_sweep_on_start: bool, /// upstream (subscribeRepos source) config pub upstream: UpstreamConfig, /// deep crawl config @@ -141,6 +148,7 @@ impl Default for SyncConfig { reactive_permit_wait_timeout: Duration::from_secs(10), hubble_sync_storage_prefix: &[0x00], sync_scope: SyncScope::default(), + rescope_sweep_on_start: false, upstream: UpstreamConfig::default(), deep_crawl: None, firehose: FirehoseConfig::default(), diff --git a/hubble-sync/src/hubble_sync.rs b/hubble-sync/src/hubble_sync.rs index f201eab..d9ba2ce 100644 --- a/hubble-sync/src/hubble_sync.rs +++ b/hubble-sync/src/hubble_sync.rs @@ -19,6 +19,7 @@ use crate::metrics::{ REPOS_DESYNC_BY_REASON, REPOS_TOTAL, }; use crate::pending_identity_scheduler; +use crate::rescope_sweep; use crate::resync_scheduler; use crate::storage::repo::AccountStatusSource; @@ -46,6 +47,9 @@ pub enum HubbleSyncError { #[error("backfill repo discovery: {0}")] RepoDiscovery(#[source] CrawlError), + #[error("rescope sweep: {0}")] + RescopeSweep(#[source] rescope_sweep::RescopeSweepError), + #[error("background task panicked")] TaskPanic(#[from] tokio::task::JoinError), } @@ -272,6 +276,18 @@ where }); } + // one-shot rescope sweep (only widens) + if self.config.rescope_sweep_on_start { + let stor = self.storage.clone(); + let rr = self.repo_registry.clone(); + let c = self.cancel.clone(); + tasks.spawn(async move { + rescope_sweep::run(stor, rr, c) + .await + .map_err(HubbleSyncError::RescopeSweep) + }); + } + // storage engine maintenance { let storage = self.storage.clone(); diff --git a/hubble-sync/src/lib.rs b/hubble-sync/src/lib.rs index 8c6dce8..4c66cf5 100644 --- a/hubble-sync/src/lib.rs +++ b/hubble-sync/src/lib.rs @@ -11,6 +11,7 @@ mod metrics; mod pending_identity_scheduler; mod repo_actor; mod repo_slot; +mod rescope_sweep; mod resync; mod resync_scheduler; mod storage; diff --git a/hubble-sync/src/repo_actor/actor.rs b/hubble-sync/src/repo_actor/actor.rs index 8e946ac..3663545 100644 --- a/hubble-sync/src/repo_actor/actor.rs +++ b/hubble-sync/src/repo_actor/actor.rs @@ -100,6 +100,8 @@ pub enum RepoMessage { event: Box, reply: Option>, }, + /// re-evaluate this repo's sync scope against current policy + ReconcileScope, } /// repo actor's actual work (internal), executed by TaskProcessor @@ -123,6 +125,7 @@ pub(super) enum Task { event: ModerateEvent, reply: Option>, }, + ReconcileScope, } impl Task { @@ -138,6 +141,7 @@ impl Task { Task::ScheduledResync { .. } => "ScheduledResync", Task::RefreshIdentity { .. } => "RefreshIdentity", Task::Moderate { .. } => "Moderate", + Task::ReconcileScope => "ReconcileScope", } } } @@ -248,6 +252,14 @@ impl Intake { event: *event, reply, }, + M::ReconcileScope => { + // coalesce: one pending reconcile is enough + if self.queue.iter().any(|t| matches!(t, Task::ReconcileScope)) { + trace!("intake: dropping reconcile-scope: one already queued"); + return; + } + Task::ReconcileScope + } M::ScheduledIdentityResolve { resolve_permit } => { // TODO: right now we assume that this gets handled in wake(). // this is slightly sketch, and means a ScheduledIdentityResolve diff --git a/hubble-sync/src/repo_actor/task_processor/mod.rs b/hubble-sync/src/repo_actor/task_processor/mod.rs index 4184f44..7fdb2ea 100644 --- a/hubble-sync/src/repo_actor/task_processor/mod.rs +++ b/hubble-sync/src/repo_actor/task_processor/mod.rs @@ -302,6 +302,7 @@ impl, R: Resolve> TaskProcessor self.process_refresh(reply).await, Task::Moderate { event, reply } => self.process_moderate(event, reply).await, + Task::ReconcileScope => self.process_reconcile_scope().await, }; if self.repo.is_none() { error!("process handler failed to restore Some(repo)"); @@ -499,6 +500,24 @@ impl, R: Resolve> TaskProcessor PEResult { + let mut updating = self.repo().clone(); + let storage = self.storage.clone(); + let now = SystemTime::now(); + let updated = spawn_blocking(move || -> PEResult { + let mut batch = storage.batch(); + updating.reconcile_scope(now, &mut batch); + batch.commit().map_err(ProcessError::Storage)?; + Ok(updating) + }) + .await + .expect("task not to panic")?; + + self.repo = Some(Box::new(updated)); + Ok(()) + } + /// handle account events from subscribeRepos /// /// critical: right now this *trusts* any message received from upstream, so @@ -1480,6 +1499,45 @@ mod tests { .is_some() } + #[tokio::test] + async fn reconcile_scope_brings_out_of_scope_repo_into_scope() { + // an out-of-scope repo whose host is now in scope (the default registry + // is Everything) re-desyncs with CameIntoScope and re-enters the queue. + let did = plc_did("reconcilein"); + let (mut tp, _resolver, hosts) = processor( + &did, + SyncStatus::OutOfScope { + since: SystemTime::now(), + }, + ); + + tp.process_reconcile_scope().await.expect("processed"); + + match &tp.repo().info().sync_status { + SyncStatus::Desynchronized(d) => { + assert!(matches!(d.reason, DesyncReason::CameIntoScope)); + } + other => panic!("expected desynchronized, got {other:?}"), + } + assert!(queued_for(&tp, hosts.get(REPO_PDS).expect("interned"))); + } + + #[tokio::test] + async fn reconcile_scope_is_noop_for_in_scope_repo() { + // a synchronized (already in-scope) repo is left untouched: no spurious + // desync, nothing queued. + let did = plc_did("reconcilenoop"); + let (mut tp, _resolver, hosts) = processor(&did, SyncStatus::Synchronized); + + tp.process_reconcile_scope().await.expect("processed"); + + assert!(matches!( + tp.repo().info().sync_status, + SyncStatus::Synchronized + )); + assert!(!queued_for(&tp, hosts.get(REPO_PDS).expect("interned"))); + } + #[tokio::test] async fn commit_for_out_of_scope_repo_drops_without_reaching_identity() { let did = plc_did("ooscommit"); diff --git a/hubble-sync/src/rescope_sweep.rs b/hubble-sync/src/rescope_sweep.rs new file mode 100644 index 0000000..fd892fc --- /dev/null +++ b/hubble-sync/src/rescope_sweep.rs @@ -0,0 +1,106 @@ +//! route a `ReconcileScope` message to every [`SyncStatus::OutOfScope`] repo +//! +//! one-shot at startup + +use std::sync::Arc; +use std::time::Duration; + +use tokio::task::spawn_blocking; +use tokio_util::sync::CancellationToken; +use tracing::{info, warn}; + +use crate::cancel::CancelExt; +use crate::repo_actor::{RepoMessage, RepoRegistry}; +use crate::storage::repo::RepoInfo; +use crate::{ + Did, LoadError, PrefixedEngine, RepoSendError, Resolve, StorageEngine, StorageError, + SyncConsumer, +}; + +/// how long to wait before retrying a send when a repo's mailbox is full. +const SEND_BACKOFF: Duration = Duration::from_secs(1); + +/// log sweep progress every this many sent messages +const PROGRESS_EVERY: u64 = 50_000; + +#[derive(Debug, thiserror::Error)] +pub enum RescopeSweepError { + #[error("scanning repo info: {0}")] + Scan(#[source] LoadError), + + #[error("repo registry eviction stuck: cannot route reconcile-scope")] + EvictionStuck, + + #[error("rescope sweep scan task panicked")] + ScanPanic(#[source] tokio::task::JoinError), +} + +/// scan every stored repo and send `ReconcileScope` to every out-of-scope one +pub async fn run( + storage: PrefixedEngine, + repo_registry: Arc>, + cancel: CancellationToken, +) -> Result<(), RescopeSweepError> +where + S: StorageEngine, + A: SyncConsumer, + R: Resolve, +{ + // producer: blocking scan of the whole keyspace, streaming out-of-scope DIDs + // over a bounded channel so memory stays flat however many there are. + // `blocking_send` naturally backpressures the scan when the consumer stalls. + let (tx, mut rx) = tokio::sync::mpsc::channel::(1024); + let scan_storage = storage.clone(); + let scan_cancel = cancel.clone(); + let producer = spawn_blocking(move || -> Result> { + let mut found = 0u64; + for item in RepoInfo::scan_out_of_scope(&scan_storage) { + if scan_cancel.is_cancelled() { + break; + } + let did = item?; + found += 1; + if tx.blocking_send(did).is_err() { + break; // consumer gone + } + } + Ok(found) + }); + + // consumer: route each DID to its actor for an actor-owned reconcile. + let mut sent = 0u64; + while let Some(did) = rx.recv().await { + loop { + match repo_registry.try_send(&did, RepoMessage::ReconcileScope) { + Ok(()) => { + sent += 1; + if sent.is_multiple_of(PROGRESS_EVERY) { + info!(sent, "rescope sweep progress"); + } + break; + } + Err(RepoSendError::Draining(_)) => { + info!(sent, "rescope sweep stopping: repo registry draining"); + return Ok(()); + } + Err(RepoSendError::Backpressure(_)) => { + // mailbox full: back off and retry the same repo + if !cancel.sleep(SEND_BACKOFF).await { + return Ok(()); // cancelled + } + } + Err(RepoSendError::EvictionStuck(_)) => { + warn!(%did, "rescope sweep aborting: repo registry eviction stuck"); + return Err(RescopeSweepError::EvictionStuck); + } + } + } + } + + let found = producer + .await + .map_err(RescopeSweepError::ScanPanic)? + .map_err(RescopeSweepError::Scan)?; + info!(found, sent, "rescope sweep complete"); + Ok(()) +} diff --git a/hubble-sync/src/storage/repo/info.rs b/hubble-sync/src/storage/repo/info.rs index 656bf9b..ae9ef0b 100644 --- a/hubble-sync/src/storage/repo/info.rs +++ b/hubble-sync/src/storage/repo/info.rs @@ -15,7 +15,7 @@ use std::time::{Duration, SystemTime}; use dasl::drisl; use serde::{Deserialize, Serialize}; -use super::{LoadError, PREFIX_REPO_INFO, unix_ms_u64}; +use super::{DecodeError, LoadError, PREFIX_REPO_INFO, unix_ms_u64}; use crate::host::{Host, HostnameError}; use crate::{ Did, DidMethod, HostRegistry, SigningKey, StorageBatch, StorageEngine, StorageError, Tid, @@ -119,6 +119,31 @@ impl RepoInfo { .is_some()) } + /// scan the whole repo-info keyspace, yielding the DID of every repo + /// currently parked [`SyncStatus::OutOfScope`]. + /// + /// only the raw db shape is decoded (no host re-interning). + /// non-out-of-scope repos are skipped. + pub(crate) fn scan_out_of_scope( + storage: &S, + ) -> impl Iterator>> + '_ { + storage + .scan_from(PREFIX_REPO_INFO, b"") + .filter_map(|pair| match pair { + Err(e) => Some(Err(LoadError::Storage(e))), + Ok((key, value)) => { + let db: DbRepoInfo = match drisl::from_slice(&value) { + Ok(db) => db, + Err(e) => return Some(Err(e.into())), + }; + if !matches!(db.sync_status, SyncStatus::OutOfScope { .. }) { + return None; + } + Some(did_from_info_key(&key).map_err(LoadError::Decode)) + } + }) + } + pub fn store(&self, did: &Did, batch: &mut B) where E: StorageError, @@ -157,6 +182,14 @@ impl RepoInfo { } } +/// reconstruct a DID from a repo-info key (prefix-stripped) +fn did_from_info_key(key: &[u8]) -> Result { + std::str::from_utf8(key) + .map_err(|err| DecodeError::NotUtf8 { what: "did", err })? + .try_into() + .map_err(DecodeError::BadDid) +} + impl From for DbRepoInfo { fn from(info: RepoInfo) -> Self { let id = info.identity; @@ -805,6 +838,43 @@ mod tests { assert!(RepoInfo::exists(&eng, &did).unwrap(), "present after store"); } + #[test] + fn scan_out_of_scope_yields_only_parked_repos() { + // the sweep's enumerator: yields exactly the OutOfScope DIDs, skipping + // synchronized and gone repos, decoding only the raw db shape. DIDs are + // re-validated out of the key, so they must be well-formed (32 chars). + let eng = MemEngine::new(); + let host = h("example.com"); + let plc = |tag: char| Did::new(format!("did:plc:{}", tag.to_string().repeat(24))).unwrap(); + let store = |did: &Did, status: SyncStatus| { + let mut info = sample_info(host.clone()); + info.sync_status = status; + let mut b = eng.batch(); + info.store(did, &mut b); + b.commit().unwrap(); + }; + + store(&plc('a'), SyncStatus::OutOfScope { since: t(1) }); + store(&plc('b'), SyncStatus::Synchronized); + store(&plc('c'), SyncStatus::OutOfScope { since: t(2) }); + store( + &plc('d'), + SyncStatus::Gone { + since: t(3), + reason: "deleted".into(), + }, + ); + + let mut found: Vec = RepoInfo::scan_out_of_scope(&eng) + .map(|r| r.unwrap().as_str().to_string()) + .collect(); + found.sort(); + assert_eq!( + found, + vec![plc('a').as_str().to_string(), plc('c').as_str().to_string()] + ); + } + #[test] fn repo_info_round_trip_with_identity_interns_host() { let eng = MemEngine::new(); diff --git a/hubble-sync/src/storage/repo/mod.rs b/hubble-sync/src/storage/repo/mod.rs index a503687..013c188 100644 --- a/hubble-sync/src/storage/repo/mod.rs +++ b/hubble-sync/src/storage/repo/mod.rs @@ -575,7 +575,7 @@ impl Repo { } // TODO: follow app style, don't do an ad-hoc manual reconcile helper here - fn reconcile_scope>( + pub(crate) fn reconcile_scope>( &mut self, now: SystemTime, batch: &mut B, diff --git a/hubble/src/main.rs b/hubble/src/main.rs index 851cd1c..23d45b6 100644 --- a/hubble/src/main.rs +++ b/hubble/src/main.rs @@ -152,13 +152,13 @@ struct Args { /// bearer token for com.atproto.sync.getRepo requests to the upstream host /// - /// pds-direct upstreams only (conflicts with --relay-bsky). prefer the env - /// var over the flag so the token doesn't show up in `ps` + /// pds-direct upstreams only (conflicts with --upstream-mode-relay). prefer + /// the env var over the flag so the token doesn't show up in `ps` #[arg( long, env = "HUBBLE_UPSTREAM_GETREPO_TOKEN", hide_env_values = true, - conflicts_with = "relay_bsky" + conflicts_with = "upstream_mode_relay" )] upstream_getrepo_token: Option, @@ -179,13 +179,17 @@ struct Args { #[arg(long)] upstream_concurrency: Option, - /// use relay-mode upstream with is-bsky policy (near-full-network test) + /// run in relay mode instead of direct-pds #[arg(long, action)] - relay_bsky: bool, + upstream_mode_relay: bool, + + /// restrict sync scope to bsky-mushroom PDS hosts only + #[arg(long, action)] + host_scope_bsky: bool, /// enable the deep crawl: listHosts on the relay, then listRepos on each PDS /// directly, to discover repos the relay doesn't list. relay upstream only. - #[arg(long, action)] + #[arg(long, action, requires("upstream_mode_relay"))] deep_crawl: bool, /// don't wait for the relay listRepos crawl's first full pass before starting @@ -193,6 +197,10 @@ struct Args { #[arg(long, action, requires("deep_crawl"))] deep_crawl_no_wait: bool, + /// re-evaluate scope of all currently-out-of-scope repos (scans on startup) + #[arg(long, action)] + rescope_sweep: bool, + /// how many scheduled plc resolution requests can be in flight #[arg(long, default_value = "6")] max_plc_scheduled_requests: NonZeroUsize, @@ -336,10 +344,13 @@ async fn main() -> color_eyre::Result<()> { }, ..Default::default() }; - if args.relay_bsky { + if args.upstream_mode_relay { config.upstream.kind = UpstreamKind::Relay; + } + if args.host_scope_bsky { config.sync_scope = SyncScope::BskyPdsOnly; } + config.rescope_sweep_on_start = args.rescope_sweep; if args.deep_crawl { config.deep_crawl = Some(DeepCrawl { concurrency: NonZeroUsize::new(100).expect("nonzero"), @@ -551,13 +562,8 @@ mod tests { Args::command().debug_assert(); } - #[test] - fn getrepo_token_conflicts_with_relay_upstream() { - // the token is scoped to pds-direct upstreams. this covers the two- - // flags case; an env-sourced token also triggers the conflict (clap - // treats env values as explicit), backed by the assert in - // HubbleSync::with_resolver either way. - let err = Args::try_parse_from([ + fn base_args() -> Vec<&'static str> { + vec![ "hubble", "--data-dir", "/tmp/x", @@ -565,11 +571,44 @@ mod tests { "example.com", "--contact", "@bad-example.com", - "--relay-bsky", + ] + } + + #[test] + fn getrepo_token_conflicts_with_relay_upstream() { + // the token is scoped to pds-direct upstreams. this covers the two- + // flags case; an env-sourced token also triggers the conflict (clap + // treats env values as explicit), backed by the assert in + // HubbleSync::with_resolver either way. + let mut argv = base_args(); + argv.extend([ + "--upstream-mode-relay", "--upstream-getrepo-token", "sekrit", - ]) - .unwrap_err(); + ]); + let err = Args::try_parse_from(argv).unwrap_err(); assert_eq!(err.kind(), clap::error::ErrorKind::ArgumentConflict); } + + #[test] + fn getrepo_token_parses_with_default_pds_upstream() { + // control: the token is valid with the default (pds-direct) upstream + let mut argv = base_args(); + argv.extend(["--upstream-getrepo-token", "sekrit"]); + Args::try_parse_from(argv).expect("token valid with pds upstream"); + } + + #[test] + fn deep_crawl_requires_relay_upstream() { + // deep crawl only works against a relay: clap rejects it without + // --upstream-mode-relay, and accepts it with. + let mut without = base_args(); + without.push("--deep-crawl"); + let err = Args::try_parse_from(without).unwrap_err(); + assert_eq!(err.kind(), clap::error::ErrorKind::MissingRequiredArgument); + + let mut with = base_args(); + with.extend(["--deep-crawl", "--upstream-mode-relay"]); + Args::try_parse_from(with).expect("deep crawl valid with relay upstream"); + } }