From 27d1b2bfa77dabe835ae34f364e93c915e369727 Mon Sep 17 00:00:00 2001 From: "@cinnabar-cavy.test.pds.did.bot" Date: Sun, 25 Oct 2026 19:55:44 -0400 Subject: [PATCH] fix: unstall the global crawl and split it into hash-sharded tasks Checkpoints no longer copy the database, every repository scan has a time limit, and the global pass can run as N shards whose counts merge without double counting. The stalled run spent its time in checkpoints: VACUUM INTO plus the per-day queries re-read a 3.8 GB index on a 4 GB task, and once the task's storage throttled one checkpoint took 94 minutes with scanning paused. A checkpoint now folds the WAL in and uploads the file as it is, and the summary comes from one pass that exports each counted change with its earliest day per count. Shards (TALLY_SHARD_INDEX/COUNT or --shard i/N) crawl repositories by FNV-1a of the DID under their own state key, seed themselves from the single-task database, and publish counted changes; global-merge takes the earliest day per change across shards. One shard is the old single task. Co-Authored-By: Claude Opus 5.5 Change-Id: Ie842fed00510fe6495bcf82db0c3f4b103297e27 --- src/config.rs | 51 ++++- src/crawl.rs | 113 ++++++++-- src/db.rs | 387 +++++++++++++++++++++++++++++++-- src/global_summary.rs | 493 +++++++++++++++++++++++++++++++++++++++--- src/main.rs | 137 ++++++++++-- src/publish.rs | 101 ++++++--- src/shard.rs | 205 ++++++++++++++++++ 7 files changed, 1359 insertions(+), 128 deletions(-) create mode 100644 src/shard.rs diff --git a/src/config.rs b/src/config.rs index f935e2e..a7f5da7 100644 --- a/src/config.rs +++ b/src/config.rs @@ -108,6 +108,13 @@ pub struct Config { pub global_checkpoint_secs: u64, pub global_checkpoint_repos: usize, pub global_min_coverage_pct: f64, + /// Which part of the global crawl this task owns (TALLY_SHARD_INDEX of + /// TALLY_SHARD_COUNT, or `--shard i/N`). One shard, the default, is the + /// whole crawl in one task. + pub shard: crate::shard::Shard, + /// Give up on one repository after this long, so a slow knot or a huge + /// history can't hold the run (TALLY_REPO_TIMEOUT_SECS; 0 for no limit). + pub repo_timeout: Option, /// Circuit breaker for a host that keeps returning 5xx or transport errors. pub circuit_failures: u32, pub circuit_cooldown_secs: u64, @@ -132,8 +139,14 @@ pub struct Config { pub state_bucket: Option, #[cfg_attr(not(feature = "aws"), allow(dead_code))] pub site_bucket: Option, + /// The state object this task reads and writes: TALLY_STATE_KEY, or for + /// one shard of a sharded global crawl the shard's own key beside it. #[cfg_attr(not(feature = "aws"), allow(dead_code))] pub state_key: String, + /// TALLY_STATE_KEY as given: the single-task key a shard derives its own + /// keys from, and seeds a missing database from. + #[cfg_attr(not(feature = "aws"), allow(dead_code))] + pub base_state_key: String, #[cfg_attr(not(feature = "aws"), allow(dead_code))] pub s3_prefix: String, pub publish: bool, @@ -221,7 +234,9 @@ impl Config { let global_progress_secs = env_parse("TALLY_GLOBAL_PROGRESS_SECS", 600u64)?; let global_progress_repos = env_parse("TALLY_GLOBAL_PROGRESS_REPOS", 250usize)?; let global_checkpoint_secs = env_parse("TALLY_GLOBAL_CHECKPOINT_SECS", 1_800u64)?; - let global_checkpoint_repos = env_parse("TALLY_GLOBAL_CHECKPOINT_REPOS", 1_000usize)?; + // Time alone paces checkpoints by default: a count could fire again + // right after a slow checkpoint, and a run would do little else. + let global_checkpoint_repos = env_parse("TALLY_GLOBAL_CHECKPOINT_REPOS", 0usize)?; let global_min_coverage_pct = env_parse("TALLY_GLOBAL_MIN_COVERAGE_PCT", 90.0f64)?; if !(0.0..=100.0).contains(&global_min_coverage_pct) { anyhow::bail!("TALLY_GLOBAL_MIN_COVERAGE_PCT must be between 0 and 100"); @@ -231,8 +246,13 @@ impl Config { anyhow::bail!("TALLY_CIRCUIT_FAILURES must be at least 1"); } let circuit_cooldown_secs = env_parse("TALLY_CIRCUIT_COOLDOWN_SECS", 60u64)?; + let shard = crate::shard::Shard::new( + env_parse("TALLY_SHARD_INDEX", 0usize)?, + env_parse("TALLY_SHARD_COUNT", 1usize)?, + ) + .context("TALLY_SHARD_INDEX and TALLY_SHARD_COUNT")?; - Ok(Self { + let mut config = Self { personal, personal_emails, didbot_pds, @@ -265,6 +285,11 @@ impl Config { global_checkpoint_secs, global_checkpoint_repos, global_min_coverage_pct, + shard: crate::shard::Shard::WHOLE, + repo_timeout: match env_parse("TALLY_REPO_TIMEOUT_SECS", 900u64)? { + 0 => None, + secs => Some(Duration::from_secs(secs)), + }, circuit_failures, circuit_cooldown_secs, concurrency, @@ -291,11 +316,29 @@ impl Config { site_bucket: env::var("TALLY_SITE_BUCKET") .ok() .filter(|value| !value.trim().is_empty()), - state_key: env::var("TALLY_STATE_KEY") + state_key: String::new(), + base_state_key: env::var("TALLY_STATE_KEY") .unwrap_or_else(|_| "state/tally.sqlite3".to_string()), s3_prefix: env::var("TALLY_S3_PREFIX").unwrap_or_default(), publish: env_parse("TALLY_PUBLISH", false)?, - }) + }; + config.set_shard(shard); + Ok(config) + } + + /// Choose the shard this task crawls and the state keys that go with it. + pub fn set_shard(&mut self, shard: crate::shard::Shard) { + self.shard = shard; + self.state_key = shard.state_key(&self.base_state_key); + } + + /// The state object a shard with no database of its own starts from: the + /// single-task database, which it trims to its own repositories. + #[cfg_attr(not(feature = "aws"), allow(dead_code))] + pub fn seed_state_key(&self) -> Option<&str> { + self.shard + .is_sharded() + .then_some(self.base_state_key.as_str()) } pub fn github_configured(&self) -> bool { diff --git a/src/crawl.rs b/src/crawl.rs index f52601d..977f897 100644 --- a/src/crawl.rs +++ b/src/crawl.rs @@ -366,7 +366,13 @@ pub async fn run(config: &Config, db: &Db, options: &CrawlOptions) -> Result = Vec::new(); let mut global_due: Vec = Vec::new(); let mut owned_count = 0usize; + let mut shard_count = 0usize; for repo in all_repos { + // A shard of the global crawl reads only its own repositories. + if options.global && !config.shard.contains(&repo.repo_did) { + continue; + } + shard_count += 1; let is_owned = owned_dids.contains(&repo.owner_did); if is_owned { owned_count += 1; @@ -444,8 +450,16 @@ pub async fn run(config: &Config, db: &Db, options: &CrawlOptions) -> Result Result Result result, + Err(_) => { + let message = format!("scan timed out after {}s", limit.as_secs()); + let _ = resolver + .db + .mark_repo_scanned(&repo.repo_did, Some(&message)); + tracing::warn!( + repo = %repo.repo_did, + timeout_secs = limit.as_secs(), + "repository scan timed out" + ); + RepoResult { + report: ScanReport { + repos_seen: 1, + repos_failed: 1, + ..Default::default() + }, + budget_exhausted: false, + } + } + } } })) - .buffer_unordered(config.concurrency); + .buffer_unordered(concurrency); futures::pin_mut!(results); let mut last_progress = Instant::now(); @@ -577,24 +620,54 @@ pub async fn run(config: &Config, db: &Db, options: &CrawlOptions) -> Result Result<()> { - let path = config.db_path.with_extension("checkpoint.sqlite3"); - db.snapshot_to(&path)?; + let started = Instant::now(); + publish_global(config, db, true).await?; + tracing::info!( + elapsed_secs = started.elapsed().as_secs(), + "global checkpoint complete" + ); + Ok(()) +} + +/// Publish a global run's results: the database, then the summary, or for one +/// shard its counted changes for the merge. +pub async fn publish_global(config: &Config, db: &Db, partial: bool) -> Result<()> { + db.flush_wal()?; + #[cfg(feature = "aws")] + if config.publish { + crate::publish::publish_state(config).await?; + } + let counts = config.db_path.with_extension("counts.sqlite3"); + if config.shard.is_sharded() { + let changes = global_summary::export_counts( + db, + &counts, + config.global_repo_age_grace_days, + config.shard, + partial, + )?; + tracing::info!(shard = %config.shard, changes, path = %counts.display(), "wrote counted changes"); + #[cfg(feature = "aws")] + if config.publish { + let key = config.shard.counts_key(&config.base_state_key); + crate::publish::publish_file(config, &counts, &key).await?; + } + return Ok(()); + } let summary = global_summary::from_db( db, + &counts, &config.global_imported_before, config.global_repo_age_grace_days, - true, + partial, )?; - if let Err(err) = global_summary::write(config, &summary).await { - tracing::warn!(error = %err, "could not write partial global summary"); - } - #[cfg(feature = "aws")] - if config.publish { - crate::publish::publish_state_path(config, &path).await?; - } - let _ = tokio::fs::remove_file(&path).await; - Ok(()) + global_summary::write(config, &summary).await } #[derive(Debug, Default)] @@ -680,6 +753,7 @@ async fn scan_repo( if resolver.global && resolver.review_repo_commits > 0 && let Some(default) = branches.first() + && !branch_is_current(resolver.db, &repo.repo_did, default) { match resolver .tangled @@ -978,6 +1052,15 @@ async fn scan_branch( Ok((report, false, false)) } +/// Whether the stored branch is indexed at the head Tangled reports now: a +/// walk would add nothing, so neither does asking for its size. +fn branch_is_current(db: &Db, repo_did: &str, branch: &BranchHead) -> bool { + db.get_branch(repo_did, &branch.name) + .ok() + .flatten() + .is_some_and(|stored| stored.indexed && stored.head_sha == branch.head_sha) +} + fn all_repos_count(db: &Db) -> Result { Ok(db.list_active_repos()?.len()) } diff --git a/src/db.rs b/src/db.rs index 2e337f2..a7ad006 100644 --- a/src/db.rs +++ b/src/db.rs @@ -16,6 +16,7 @@ use chrono::{DateTime, Utc}; use rusqlite::{Connection, OptionalExtension, params}; use crate::model::{Branch, Commit, CommitRef, Identity, IdentityKind, Repo}; +use crate::shard::Shard; /// Each repository's own rows are found through `repos`' primary key: /// `p.repo_did = substr(r.repo_key, 9)` strips the `tangled:` prefix from the @@ -23,6 +24,10 @@ use crate::model::{Branch, Commit, CommitRef, Identity, IdentityKind, Repo}; /// can't use an index (the comparison takes `repo_key`'s text affinity, the /// expression has none), so SQLite scanned every repository for every /// reference: hours on a full global index instead of minutes. +/// +/// The per-day count the global summary used before counted changes were +/// exported; kept as the reference the export is tested against. +#[cfg(test)] const TANGLED_DAILY_COUNTS_SQL: &str = r#" select day, count(*) from ( select substr(min(c.authored_at), 1, 10) as day @@ -48,6 +53,47 @@ const TANGLED_DAILY_COUNTS_SQL: &str = r#" order by day "#; +/// The six columns of a counted-changes export: the earliest UTC author day a +/// logical change counts on, for everyone, the owner's identities and their +/// agents, counting the default branch only and then every branch. +pub const COUNTED_COLUMNS: [&str; 6] = [ + "all_default", + "owner_default", + "agent_default", + "all_any", + "owner_any", + "agent_any", +]; + +/// Every counted Tangled logical change, keyed by Change-Id or SHA, in one +/// pass over the references, under the same rules as the per-day counts: a +/// reference in an excluded repository doesn't count, nor does a commit +/// authored before its repository existed (less the grace period). A day is +/// the earliest author date's first ten characters, and the minimum of +/// prefixes is the prefix of the minimum, so this matches grouping first. +/// `{shard}` narrows the references to one shard's repositories. +const COUNTED_CHANGES_SQL: &str = r#" + select case when c.change_id != '' then c.change_id else c.sha end as k, + min(case when r.on_default = 1 then substr(c.authored_at, 1, 10) end), + min(case when r.on_default = 1 and c.author_did in (select did from temp.counted_owner) + then substr(c.authored_at, 1, 10) end), + min(case when r.on_default = 1 and c.author_did in (select did from temp.counted_agent) + then substr(c.authored_at, 1, 10) end), + min(substr(c.authored_at, 1, 10)), + min(case when c.author_did in (select did from temp.counted_owner) + then substr(c.authored_at, 1, 10) end), + min(case when c.author_did in (select did from temp.counted_agent) + then substr(c.authored_at, 1, 10) end) + from commit_refs r + join commits c on c.sha = r.sha + left join repos p on p.repo_did = substr(r.repo_key, 9) + where r.source = 'tangled' + and coalesce(p.excluded, 0) = 0 + and coalesce(date(c.authored_at) < date(p.created_at, ?1), 0) = 0 + {shard} + group by k +"#; + const REPO_LOGICAL_COUNTS_SQL: &str = " select r.repo_key, count(distinct case when c.change_id != '' then c.change_id else c.sha end) @@ -140,19 +186,50 @@ impl Db { Ok(db) } - /// Write a consistent copy for a checkpoint. `VACUUM INTO` is safe while - /// the crawler is still writing; copying the live file is not. - pub fn snapshot_to(&self, path: &Path) -> Result<()> { - if path.exists() { - std::fs::remove_file(path) - .with_context(|| format!("replacing checkpoint {}", path.display()))?; - } + /// Fold the write-ahead log into the database file, so the file alone is + /// the whole database. Uploading it is then safe for as long as nothing + /// writes: the crawl's checkpoint holds the only connection while it + /// uploads, so no copy of a multi-gigabyte file is needed first. + pub fn flush_wal(&self) -> Result<()> { let conn = self.conn.lock().unwrap(); - conn.execute("vacuum into ?1", [path.to_string_lossy().as_ref()]) - .with_context(|| format!("snapshotting database to {}", path.display()))?; + let (busy, _, _): (i64, i64, i64) = + conn.query_row("pragma wal_checkpoint(truncate)", [], |row| { + Ok((row.get(0)?, row.get(1)?, row.get(2)?)) + })?; + if busy != 0 { + anyhow::bail!("the write-ahead log could not be checkpointed"); + } Ok(()) } + /// Create the database at `path` for one shard of the global crawl from + /// the single-task database at `seed`: every repository's metadata, and + /// only this shard's branches, references and commits. A shard then + /// resumes where the single task left off instead of reading Tangled + /// from the start, from a file a fraction of the size. + pub fn seed_shard(path: &Path, seed: &Path, shard: Shard) -> Result { + // Bring the seed to this schema first, so the copy can name columns. + drop(Self::open(seed).with_context(|| format!("opening seed {}", seed.display()))?); + let db = Self::open(path)?; + { + let conn = db.conn.lock().unwrap(); + // One large transaction into a new file: a rollback journal has + // nothing to save, where the write-ahead log would hold a second + // copy of every page until the next checkpoint. + conn.pragma_update(None, "journal_mode", "DELETE")?; + conn.execute( + "attach database ?1 as seed", + [seed.to_string_lossy().as_ref()], + )?; + let copied = seed_shard_rows(&conn, shard); + let detached = conn.execute("detach database seed", []); + copied?; + detached?; + conn.pragma_update(None, "journal_mode", "WAL")?; + } + Ok(db) + } + #[cfg(test)] pub fn in_memory() -> Result { let db = Self { @@ -910,6 +987,7 @@ impl Db { Ok(()) } + #[cfg(test)] pub fn branch_scope(&self) -> BranchScope { *self.scope.lock().unwrap() } @@ -954,6 +1032,7 @@ impl Db { Ok(conn.query_row(&sql, refs.as_slice(), |row| row.get::<_, i64>(0))? as u64) } + #[cfg(test)] /// Distinct logical Tangled changes per UTC author day by `dids`, under the /// same exclusions and repo-age filter as `tangled_daily_counts`, so the /// global summary can carry the owner's side of the share from the same @@ -1103,6 +1182,7 @@ impl Db { Ok(rows.collect::>>()?) } + #[cfg(test)] /// Distinct logical Tangled changes per UTC author day. This is the /// denominator the global task publishes; the hourly task never /// derives it from its own owned-only database. @@ -1498,6 +1578,73 @@ impl Db { }) } + /// [`Db::repo_coverage`] over one shard's repositories. Every shard + /// keeps the whole repository list; only its own count toward it. + pub fn repo_coverage_in(&self, shard: Shard) -> Result { + if !shard.is_sharded() { + return self.repo_coverage(); + } + let conn = self.conn.lock().unwrap(); + let mut coverage = Coverage { + repos_total: 0, + repos_scanned: 0, + repos_failed: 0, + branches_total: 0, + branches_indexed: 0, + }; + let mut stmt = conn.prepare( + "select repo_did, last_scanned is not null and last_error is null, + last_error is not null + from repos where active = 1 and excluded = 0", + )?; + let mut rows = stmt.query([])?; + while let Some(row) = rows.next()? { + if shard.contains(&row.get::<_, String>(0)?) { + coverage.repos_total += 1; + coverage.repos_scanned += u64::from(row.get::<_, bool>(1)?); + coverage.repos_failed += u64::from(row.get::<_, bool>(2)?); + } + } + let mut stmt = conn.prepare("select repo_did, indexed from branches")?; + let mut rows = stmt.query([])?; + while let Some(row) = rows.next()? { + if shard.contains(&row.get::<_, String>(0)?) { + coverage.branches_total += 1; + coverage.branches_indexed += u64::from(row.get::<_, bool>(1)?); + } + } + Ok(coverage) + } + + /// Write this database's counted Tangled changes to a new SQLite file at + /// `out` (see [`COUNTED_COLUMNS`]), with `meta` beside them. `owner` and + /// `agents` are the DIDs whose own columns are filled; `shard` keeps only + /// references to its repositories. Returns the number of changes. + pub fn write_counted( + &self, + out: &Path, + owner: &[String], + agents: &[String], + repo_age_grace_days: i64, + shard: Shard, + meta: &[(&str, String)], + ) -> Result { + if out.exists() { + std::fs::remove_file(out) + .with_context(|| format!("replacing counted changes at {}", out.display()))?; + } + let conn = self.conn.lock().unwrap(); + conn.execute( + "attach database ?1 as counts", + [out.to_string_lossy().as_ref()], + )?; + let written = write_counted_rows(&conn, owner, agents, repo_age_grace_days, shard, meta); + let detached = conn.execute("detach database counts", []); + let written = written?; + detached?; + Ok(written) + } + pub fn error_summary(&self, limit: usize) -> Result> { let conn = self.conn.lock().unwrap(); let mut stmt = conn.prepare( @@ -1615,6 +1762,204 @@ pub struct Coverage { pub branches_indexed: u64, } +/// Per day, how many logical changes each of [`COUNTED_COLUMNS`] counts, +/// after merging one or more exports, with each export's `meta`. +#[derive(Debug, Default)] +pub struct MergedCounts { + pub days: [Vec<(String, u64)>; 6], + pub meta: Vec>, +} + +fn write_counted_rows( + conn: &Connection, + owner: &[String], + agents: &[String], + repo_age_grace_days: i64, + shard: Shard, + meta: &[(&str, String)], +) -> Result { + let columns = COUNTED_COLUMNS + .map(|column| format!("{column} text")) + .join(", "); + conn.execute_batch(&format!( + "create table counts.meta (key text primary key, value text not null); + create table counts.counted (key text primary key, {columns}) without rowid; + drop table if exists temp.counted_owner; + drop table if exists temp.counted_agent; + drop table if exists temp.counted_shard; + create temp table counted_owner (did text primary key); + create temp table counted_agent (did text primary key); + create temp table counted_shard (repo_key text primary key);" + ))?; + for (table, dids) in [("counted_owner", owner), ("counted_agent", agents)] { + let mut stmt = conn.prepare(&format!("insert or ignore into temp.{table} values (?1)"))?; + for did in dids { + stmt.execute([did])?; + } + } + let filter = if shard.is_sharded() { + let keys: Vec = conn + .prepare("select distinct repo_key from commit_refs where source = 'tangled'")? + .query_map([], |row| row.get(0))? + .collect::>()?; + let mut stmt = conn.prepare("insert into temp.counted_shard values (?1)")?; + for key in keys.iter().filter(|key| shard.contains_repo_key(key)) { + stmt.execute([key])?; + } + "and r.repo_key in (select repo_key from temp.counted_shard)" + } else { + "" + }; + let written = conn.execute( + &format!( + "insert into counts.counted {}", + COUNTED_CHANGES_SQL.replace("{shard}", filter) + ), + [format!("-{repo_age_grace_days} days")], + )?; + let mut stmt = conn.prepare("insert into counts.meta (key, value) values (?1, ?2)")?; + for (key, value) in meta { + stmt.execute(params![key, value])?; + } + conn.execute_batch( + "drop table temp.counted_owner; + drop table temp.counted_agent; + drop table temp.counted_shard;", + )?; + Ok(written as u64) +} + +/// Merge counted-change exports written by [`Db::write_counted`]. A change +/// in several exports (its repositories on different shards) counts once, on +/// the earliest day any of them gives it, column by column. +pub fn merge_counted(paths: &[impl AsRef]) -> Result { + let conn = Connection::open_in_memory()?; + let columns = COUNTED_COLUMNS + .map(|column| format!("{column} text")) + .join(", "); + conn.execute_batch(&format!( + "create table merged (key text primary key, {columns}) without rowid;" + ))?; + let earliest = COUNTED_COLUMNS + .map(|column| { + // min() of two values is null if either is; keep whichever exists. + format!( + "{column} = coalesce(min({column}, excluded.{column}), {column}, excluded.{column})" + ) + }) + .join(", "); + let names = COUNTED_COLUMNS.join(", "); + let mut merged = MergedCounts::default(); + for path in paths { + let path = path.as_ref(); + if !path.exists() { + anyhow::bail!("no counted changes at {}", path.display()); + } + conn.execute( + "attach database ?1 as export", + [path.to_string_lossy().as_ref()], + )?; + let read = (|| -> Result> { + conn.execute( + &format!( + "insert into merged (key, {names}) + select key, {names} from export.counted where true + on conflict (key) do update set {earliest}" + ), + [], + )?; + let meta = conn + .prepare("select key, value from export.meta")? + .query_map([], |row| Ok((row.get(0)?, row.get(1)?)))? + .collect::>()?; + Ok(meta) + })(); + let detached = conn.execute("detach database export", []); + let meta = read.with_context(|| format!("merging {}", path.display()))?; + detached?; + merged.meta.push(meta); + } + for (index, column) in COUNTED_COLUMNS.iter().enumerate() { + merged.days[index] = conn + .prepare(&format!( + "select {column}, count(*) from merged + where {column} is not null group by {column} order by {column}" + ))? + .query_map([], |row| Ok((row.get(0)?, row.get::<_, i64>(1)? as u64)))? + .collect::>()?; + } + Ok(merged) +} + +/// Copy the seed's rows that belong to `shard` into the main database. +fn seed_shard_rows(conn: &Connection, shard: Shard) -> Result<()> { + conn.execute_batch( + "drop table if exists temp.seed_shard; + create temp table seed_shard (repo_did text primary key);", + )?; + let repos: Vec = conn + .prepare( + "select repo_did from seed.repos + union + select substr(repo_key, 9) from seed.commit_refs where source = 'tangled'", + )? + .query_map([], |row| row.get(0))? + .collect::>()?; + let mut stmt = conn.prepare("insert or ignore into temp.seed_shard values (?1)")?; + for repo in repos.iter().filter(|repo| shard.contains(repo)) { + stmt.execute([repo])?; + } + drop(stmt); + let in_shard = "(select repo_did from temp.seed_shard)"; + conn.execute_batch("begin")?; + let copied = (|| -> Result<()> { + for table in ["meta", "identities", "repos", "github_repos"] { + copy_seed_table(conn, table, "")?; + } + copy_seed_table(conn, "branches", &format!("where repo_did in {in_shard}"))?; + copy_seed_table( + conn, + "commit_refs", + &format!("where source = 'tangled' and substr(repo_key, 9) in {in_shard}"), + )?; + copy_seed_table( + conn, + "commits", + "where sha in (select sha from main.commit_refs)", + )?; + Ok(()) + })(); + match copied { + Ok(()) => conn.execute_batch("commit; drop table temp.seed_shard;")?, + Err(err) => { + let _ = conn.execute_batch("rollback"); + return Err(err); + } + } + Ok(()) +} + +/// `insert or ignore` the seed's rows of `table` into the main database by +/// column name: the two may have gained columns in a different order. +fn copy_seed_table(conn: &Connection, table: &str, filter: &str) -> Result<()> { + let columns: Vec = conn + .prepare(&format!( + "select name from pragma_table_info('{table}', 'main')" + ))? + .query_map([], |row| row.get(0))? + .collect::>()?; + let columns = columns.join(", "); + conn.execute( + &format!( + "insert or ignore into main.{table} ({columns}) + select {columns} from seed.{table} {filter}" + ), + [], + ) + .with_context(|| format!("copying {table} from the seed"))?; + Ok(()) +} + pub fn now_string() -> String { Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true) } @@ -1750,7 +2095,13 @@ mod tests { // on a real global index; each lookup must be a primary-key search. let db = Db::in_memory().unwrap(); let conn = db.conn.lock().unwrap(); - for sql in [TANGLED_DAILY_COUNTS_SQL, REPO_LOGICAL_COUNTS_SQL] { + conn.execute_batch( + "create temp table counted_owner (did text primary key); + create temp table counted_agent (did text primary key);", + ) + .unwrap(); + let counted = COUNTED_CHANGES_SQL.replace("{shard}", ""); + for sql in [TANGLED_DAILY_COUNTS_SQL, REPO_LOGICAL_COUNTS_SQL, &counted] { let mut stmt = conn.prepare(&format!("explain query plan {sql}")).unwrap(); let params = std::iter::repeat_n("-30 days", stmt.parameter_count()); let plan: Vec = stmt @@ -1958,15 +2309,17 @@ mod tests { } #[test] - fn snapshot_is_a_consistent_copy() { - let db = Db::in_memory().unwrap(); + fn a_flushed_database_file_is_whole_without_its_log() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("live.sqlite3"); + let db = Db::open(&path).unwrap(); db.insert_commit(&commit("a", "did:plc:me"), &reference("a", "tangled:r1")) .unwrap(); - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("snapshot.sqlite3"); - db.snapshot_to(&path).unwrap(); - let copy = Db::open(&path).unwrap(); - assert_eq!(copy.tangled_total().unwrap(), 1); + db.flush_wal().unwrap(); + // What an upload of the file alone carries. + let copy = dir.path().join("uploaded.sqlite3"); + std::fs::copy(&path, ©).unwrap(); + assert_eq!(Db::open(©).unwrap().tangled_total().unwrap(), 1); } #[test] diff --git a/src/global_summary.rs b/src/global_summary.rs index 6b7b3b3..a7a0626 100644 --- a/src/global_summary.rs +++ b/src/global_summary.rs @@ -5,15 +5,17 @@ //! trust the small owned-only database for an ecosystem-wide number. use std::collections::BTreeMap; +use std::path::Path; use anyhow::{Context, Result, bail}; use chrono::{DateTime, Duration, NaiveDate, Utc}; use serde::{Deserialize, Serialize}; use crate::config::Config; -use crate::db::{BranchScope, Db, now_string}; +use crate::db::{self, BranchScope, Db, now_string}; use crate::http::HttpClient; use crate::model::IdentityKind; +use crate::shard::Shard; pub const SUMMARY_VERSION: u32 = 1; @@ -190,22 +192,18 @@ fn sum_since(days: &BTreeMap, since: Option<&str>) -> u64 { .sum() } -pub fn from_db( +/// Write `db`'s counted Tangled changes to `out` for [`from_counts`]. The +/// identities are the ones this run just refreshed: the configured personal +/// DIDs, and every bot the didbot PDSes have shown as the owner's, including +/// inactive ones. A shard writes only its own repositories' changes. +pub fn export_counts( db: &Db, - imported_before: &str, + out: &Path, repo_age_grace_days: i64, + shard: Shard, partial: bool, -) -> Result { - let coverage = db.repo_coverage()?; - let since_cutoff = |counts: Vec<(String, u64)>| -> BTreeMap { - counts - .into_iter() - .filter(|(day, _)| day.as_str() >= imported_before) - .collect() - }; - // Both sides of the share from this one snapshot. The identities are the - // ones this run just refreshed: the configured personal DIDs, and every - // bot the didbot PDSes have shown as the owner's, including inactive ones. +) -> Result { + let coverage = db.repo_coverage_in(shard)?; let identities = db.list_identities()?; let dids_of = |kind: IdentityKind| -> Vec { identities @@ -214,29 +212,73 @@ pub fn from_db( .map(|identity| identity.did.clone()) .collect() }; - let owner = dids_of(IdentityKind::Personal); - let agents = dids_of(IdentityKind::Bot); - // Both branch scopes from the same snapshot: storage keeps every branch, - // and the view the counts read is switched for each, then put back. - let restore = db.branch_scope(); - let counts = |scope: BranchScope| -> Result<[BTreeMap; 3]> { - db.set_branch_scope(scope)?; - Ok([ - since_cutoff(db.tangled_daily_counts(repo_age_grace_days)?), - since_cutoff(db.tangled_daily_counts_by(&owner, repo_age_grace_days)?), - since_cutoff(db.tangled_daily_counts_by(&agents, repo_age_grace_days)?), - ]) - }; - let default = counts(BranchScope::Default); - let all = counts(BranchScope::All); - db.set_branch_scope(restore)?; - let [days, owner_days, agent_days] = default?; - let [days_all, owner_all, agent_all] = all?; + let meta = [ + ("shard", shard.to_string()), + ("generated_at", now_string()), + ("partial", if partial { "1" } else { "0" }.to_string()), + ("repos_reached", coverage.repos_scanned.to_string()), + ("repos_total", coverage.repos_total.to_string()), + ]; + db.write_counted( + out, + &dids_of(IdentityKind::Personal), + &dids_of(IdentityKind::Bot), + repo_age_grace_days, + shard, + &meta, + ) +} + +/// The summary from counted-change exports: the whole crawl's one, or one +/// from each of `shards` shards, which must all be present. A change several +/// shards counted counts once, on its earliest day; summing each shard's +/// per-day counts instead would count it once per shard. The summary is as +/// old as its oldest export and partial if any export is. +pub fn from_counts( + paths: &[impl AsRef], + imported_before: &str, + shards: usize, +) -> Result { + let merged = db::merge_counted(paths)?; + let mut seen = vec![false; shards]; + let (mut repos_reached, mut repos_total, mut partial) = (0u64, 0u64, false); + let mut generated_at: Option = None; + for (path, meta) in paths.iter().zip(&merged.meta) { + let path = path.as_ref().display(); + let field = |key: &str| { + meta.get(key) + .with_context(|| format!("{path} has no {key}")) + }; + let shard = Shard::parse(field("shard")?)?; + if shard.count != shards || std::mem::replace(&mut seen[shard.index], true) { + bail!("{path} is shard {shard}, not one of {shards} expected shards"); + } + repos_reached += field("repos_reached")?.parse::()?; + repos_total += field("repos_total")?.parse::()?; + partial |= field("partial")? == "1"; + let at = field("generated_at")?; + if generated_at + .as_deref() + .is_none_or(|oldest| at.as_str() < oldest) + { + generated_at = Some(at.clone()); + } + } + if let Some(missing) = seen.iter().position(|seen| !seen) { + bail!("the counted changes of shard {missing}/{shards} are missing"); + } + let [days, owner_days, agent_days, days_all, owner_all, agent_all] = + merged.days.map(|counts| { + counts + .into_iter() + .filter(|(day, _)| day.as_str() >= imported_before) + .collect::>() + }); Ok(GlobalSummary { version: SUMMARY_VERSION, - generated_at: now_string(), - repos_reached: coverage.repos_scanned, - repos_total: coverage.repos_total, + generated_at: generated_at.unwrap_or_else(now_string), + repos_reached, + repos_total, partial, days, owner_days: Some(owner_days), @@ -248,6 +290,21 @@ pub fn from_db( }) } +/// The summary of one whole database, through the same export the shards +/// write, using `scratch` for it. +pub fn from_db( + db: &Db, + scratch: &Path, + imported_before: &str, + repo_age_grace_days: i64, + partial: bool, +) -> Result { + export_counts(db, scratch, repo_age_grace_days, Shard::WHOLE, partial)?; + let summary = from_counts(&[scratch], imported_before, 1); + let _ = std::fs::remove_file(scratch); + summary +} + pub async fn read(config: &Config) -> Result> { let Some(source) = config.global_summary.as_deref() else { return Ok(None); @@ -551,7 +608,15 @@ mod tests { message: String::new(), }; db.insert_commit(&f, &side).unwrap(); - let value = from_db(&db, "2025-01-01", 30, false).unwrap(); + let scratch = tempfile::tempdir().unwrap(); + let value = from_db( + &db, + &scratch.path().join("counts.sqlite3"), + "2025-01-01", + 30, + false, + ) + .unwrap(); assert_eq!(value.branch_scope.as_deref(), Some("default")); assert_eq!(db.branch_scope(), BranchScope::Default); let all = value.clone().in_scope(BranchScope::All); @@ -596,4 +661,360 @@ mod tests { value.days.insert("not-a-day".to_string(), 1); assert!(value.validate().is_err()); } + + mod shards { + use super::super::*; + use crate::model::{Commit, CommitRef, Identity, Repo}; + + const OWNER: &str = "did:plc:me"; + const AGENT: &str = "did:web:agent.test"; + + /// A repository DID that falls in shard `index` of `count`. + fn repo_in(index: usize, count: usize, skip: usize) -> String { + (0..) + .map(|i| format!("did:plc:repo{i}")) + .filter(|did| Shard::of(did, count) == index) + .nth(skip) + .unwrap() + } + + fn identities(db: &Db) { + for (did, kind) in [(OWNER, IdentityKind::Personal), (AGENT, IdentityKind::Bot)] { + db.upsert_identity(&Identity { + did: did.into(), + handle: did.into(), + kind, + operator_did: None, + color: String::new(), + active: true, + pds: None, + collections: String::new(), + checked_at: None, + }) + .unwrap(); + } + } + + fn repo(db: &Db, did: &str, created: &str) { + db.upsert_repo(&Repo { + repo_did: did.into(), + owner_did: "did:plc:someone".into(), + name: did.into(), + knot: "knot.test".into(), + url: String::new(), + website: None, + fork: false, + created_at: Some(created.into()), + active: true, + last_scanned: Some("2026-10-25T00:00:00Z".into()), + last_error: None, + }) + .unwrap(); + db.mark_repo_scanned(did, None).unwrap(); + } + + fn add( + db: &Db, + sha: &str, + change_id: &str, + author: &str, + at: &str, + repo: &str, + on_default: bool, + ) { + let commit = Commit { + sha: sha.into(), + change_id: change_id.into(), + author_did: author.into(), + author_email: String::new(), + author_name: String::new(), + authored_at: at.into(), + committed_at: at.into(), + message: String::new(), + }; + let reference = CommitRef { + sha: sha.into(), + source: "tangled".into(), + repo_key: format!("tangled:{repo}"), + url: String::new(), + on_default, + }; + db.insert_commit(&commit, &reference).unwrap(); + } + + /// A database whose repositories span both of two shards, with the + /// cases the count has to get right on each side. + fn fixture() -> Db { + let db = Db::in_memory().unwrap(); + fill(&db); + db + } + + /// The repositories `fill` writes: two in each shard of two, the + /// last one excluded. + fn repos() -> [String; 4] { + [ + repo_in(0, 2, 0), + repo_in(1, 2, 0), + repo_in(0, 2, 1), + repo_in(1, 2, 1), + ] + } + + fn fill(db: &Db) { + identities(db); + let [a, b, a2, gone] = repos(); + for did in [&a, &b, &a2, &gone] { + repo(db, did, "2026-01-01T00:00:00Z"); + } + db.exclude_repo(&gone, "fork").unwrap(); + // One change, rebased: a different SHA in each shard, and the + // later shard has the earlier author date. + add( + db, + "c1", + "Ichange1", + OWNER, + "2026-10-21T09:00:00Z", + &a, + true, + ); + add( + db, + "c1b", + "Ichange1", + OWNER, + "2026-10-20T09:00:00Z", + &b, + true, + ); + // The same commit in a repository on each shard. + add(db, "c2", "", AGENT, "2026-10-21T10:00:00Z", &a, true); + add(db, "c2", "", AGENT, "2026-10-21T10:00:00Z", &b, true); + // Merged in one shard, only on a side branch in the other. + add( + db, + "c3", + "", + "did:plc:other", + "2026-10-22T00:00:00Z", + &a2, + false, + ); + add( + db, + "c3", + "", + "did:plc:other", + "2026-10-22T00:00:00Z", + &b, + true, + ); + // Only on a side branch anywhere. + add(db, "c4", "", OWNER, "2026-10-22T01:00:00Z", &a, false); + // An excluded repository's commit, and the same change elsewhere. + add( + db, + "c5", + "Ichange5", + OWNER, + "2026-10-19T00:00:00Z", + &gone, + true, + ); + add( + db, + "c5b", + "Ichange5", + OWNER, + "2026-10-23T00:00:00Z", + &a, + true, + ); + // Predates its repository, but counts through another one that + // has no creation date to hold against it. + add(db, "c6", "", AGENT, "2025-06-01T00:00:00Z", &a, true); + add( + db, + "c6", + "", + AGENT, + "2025-06-01T00:00:00Z", + "did:plc:unlisted", + true, + ); + // An offset author date: the day is the stored text's. + add( + db, + "c7", + "", + "did:plc:other", + "2026-10-23T23:30:00-05:00", + &b, + true, + ); + } + + /// The per-day queries the summary used before counted changes were + /// exported, as the reference: three counts in each branch scope. + fn reference_counts(db: &Db, imported_before: &str) -> Vec> { + let cut = |counts: Vec<(String, u64)>| -> BTreeMap { + counts + .into_iter() + .filter(|(day, _)| day.as_str() >= imported_before) + .collect() + }; + let mut out = Vec::new(); + for scope in [BranchScope::Default, BranchScope::All] { + db.set_branch_scope(scope).unwrap(); + out.push(cut(db.tangled_daily_counts(30).unwrap())); + out.push(cut(db + .tangled_daily_counts_by(&[OWNER.into()], 30) + .unwrap())); + out.push(cut(db + .tangled_daily_counts_by(&[AGENT.into()], 30) + .unwrap())); + } + db.set_branch_scope(BranchScope::Default).unwrap(); + out + } + + fn maps(summary: &GlobalSummary) -> Vec> { + vec![ + summary.days.clone(), + summary.owner_days.clone().unwrap(), + summary.agent_days.clone().unwrap(), + summary.days_all_branches.clone().unwrap(), + summary.owner_days_all_branches.clone().unwrap(), + summary.agent_days_all_branches.clone().unwrap(), + ] + } + + #[test] + fn the_export_counts_what_the_per_day_queries_counted() { + let db = fixture(); + let dir = tempfile::tempdir().unwrap(); + for cutoff in ["2000-01-01", "2026-10-21"] { + let summary = + from_db(&db, &dir.path().join("counts.sqlite3"), cutoff, 30, false).unwrap(); + assert_eq!( + maps(&summary), + reference_counts(&db, cutoff), + "cutoff {cutoff}" + ); + } + let summary = + from_db(&db, &dir.path().join("c.sqlite3"), "2000-01-01", 30, true).unwrap(); + assert!(summary.partial); + assert_eq!((summary.repos_reached, summary.repos_total), (3, 3)); + // Spot checks: change 1 once, on its earliest day; c2 once; c4 + // only with every branch; the offset date stays on its own day. + assert_eq!( + summary.owner_days.as_ref().unwrap().get("2026-10-20"), + Some(&1) + ); + assert_eq!(summary.owner_total(None), 2); + assert_eq!( + summary.clone().in_scope(BranchScope::All).owner_total(None), + 3 + ); + assert_eq!(summary.days.get("2026-10-23"), Some(&2)); + } + + #[test] + fn a_change_in_repositories_on_two_shards_counts_once() { + let db = fixture(); + let dir = tempfile::tempdir().unwrap(); + let whole = from_db( + &db, + &dir.path().join("whole.sqlite3"), + "2000-01-01", + 30, + false, + ) + .unwrap(); + let paths: Vec<_> = (0..2) + .map(|index| { + let path = dir.path().join(format!("shard-{index}.sqlite3")); + let shard = Shard::new(index, 2).unwrap(); + export_counts(&db, &path, 30, shard, false).unwrap(); + path + }) + .collect(); + let merged = from_counts(&paths, "2000-01-01", 2).unwrap(); + assert_eq!(maps(&merged), maps(&whole)); + assert_eq!((merged.repos_reached, merged.repos_total), (3, 3)); + + // Adding up each shard's own per-day counts instead counts + // change 1, c2 and c3 once per shard. + let added: u64 = paths + .iter() + .flat_map(|path| db::merge_counted(&[path]).unwrap().days[3].clone()) + .map(|(_, count)| count) + .sum(); + let all = whole.clone().in_scope(BranchScope::All).total(None); + assert_eq!(added, all + 3); + } + + #[test] + fn a_merge_needs_every_shard_once() { + let db = fixture(); + let dir = tempfile::tempdir().unwrap(); + let path = |index: usize| { + let path = dir.path().join(format!("s{index}.sqlite3")); + export_counts(&db, &path, 30, Shard::new(index, 2).unwrap(), true).unwrap(); + path + }; + let (zero, one) = (path(0), path(1)); + assert!(from_counts(&[&zero], "2000-01-01", 2).is_err()); + assert!(from_counts(&[&zero, &zero], "2000-01-01", 2).is_err()); + assert!(from_counts(&[&zero, &one], "2000-01-01", 3).is_err()); + let merged = from_counts(&[&one, &zero], "2000-01-01", 2).unwrap(); + assert!(merged.partial); + } + + #[test] + fn a_seeded_shard_keeps_only_its_repositories() { + let dir = tempfile::tempdir().unwrap(); + let seed = dir.path().join("global.sqlite3"); + let whole_db = Db::open(&seed).unwrap(); + fill(&whole_db); + let whole = from_db( + &whole_db, + &dir.path().join("w.sqlite3"), + "2000-01-01", + 30, + false, + ) + .unwrap(); + drop(whole_db); + let mut paths = Vec::new(); + for index in 0..2 { + let shard = Shard::new(index, 2).unwrap(); + let path = dir.path().join(format!("shard-{index}.sqlite3")); + let db = Db::seed_shard(&path, &seed, shard).unwrap(); + for did in repos() + .iter() + .map(String::as_str) + .chain(["did:plc:unlisted"]) + { + let held = !db.repo_commit_shas(did).unwrap().is_empty(); + assert_eq!(held, shard.contains(did), "{did} in shard {shard}"); + } + // Every repository is still listed; only its own count. + assert_eq!(db.list_active_repos().unwrap().len(), 3); + assert_eq!(db.repo_coverage().unwrap().repos_total, 3); + assert_eq!( + db.repo_coverage_in(shard).unwrap().repos_total, + 2 - index as u64 + ); + assert_eq!(db.list_identities().unwrap().len(), 2); + let counts = dir.path().join(format!("counts-{index}.sqlite3")); + export_counts(&db, &counts, 30, shard, false).unwrap(); + paths.push(counts); + } + let merged = from_counts(&paths, "2000-01-01", 2).unwrap(); + assert_eq!(maps(&merged), maps(&whole)); + } + } } diff --git a/src/main.rs b/src/main.rs index 733d870..46f1bba 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,6 +14,7 @@ mod land_gap; mod model; mod project_timelines; mod serve; +mod shard; mod site; mod stats; mod tangled; @@ -22,7 +23,7 @@ mod work_inbox; #[cfg(feature = "aws")] mod publish; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use anyhow::{Context, Result}; use clap::{Args, Parser, Subcommand}; @@ -51,6 +52,8 @@ enum Command { Serve(ServeArgs), /// Crawl on a timer and publish to S3 when configured. Daemon(DaemonArgs), + /// Merge the shards' counted changes into the global summary. + GlobalMerge(MergeArgs), /// Print the didbot identities currently known to the database. Bots, /// Upload an existing build to S3. @@ -85,6 +88,15 @@ struct CrawlArgs { filter_only: bool, #[arg(long)] max_repos: Option, + /// Crawl one shard of the global pass, as index/count (for example 2/8). + /// Overrides TALLY_SHARD_INDEX and TALLY_SHARD_COUNT. + #[arg(long, value_parser = shard::Shard::parse, requires = "global")] + shard: Option, + /// Start a shard with no database of its own from this single-task + /// database, keeping only the shard's repositories. With publishing on, + /// the single-task state object is the seed. + #[arg(long, requires = "shard")] + seed: Option, } impl CrawlArgs { @@ -98,6 +110,9 @@ impl CrawlArgs { if self.max_repos.is_some() { config.max_repos_per_run = self.max_repos; } + if let Some(shard) = self.shard { + config.set_shard(shard); + } } fn options(&self) -> CrawlOptions { @@ -113,6 +128,18 @@ impl CrawlArgs { } } +#[derive(Debug, Args)] +struct MergeArgs { + /// Counted-change files to merge. Without them, each shard's file is + /// downloaded from the state bucket. + #[arg(long)] + counts: Vec, + /// How many shards there are (default TALLY_SHARD_COUNT, or one per + /// --counts file). + #[arg(long)] + shards: Option, +} + #[derive(Debug, Args)] struct ServeArgs { #[arg(long, default_value = "dist")] @@ -159,7 +186,7 @@ async fn main() -> Result<()> { Command::Crawl(args) => { let mut config = Config::from_env()?; args.apply(&mut config); - crawl_once(&config, &args.options()).await?; + crawl_once(&config, &args.options(), args.seed.as_deref()).await?; } Command::Serve(args) => serve::run(args.dir, &args.bind).await?, Command::Daemon(args) => { @@ -167,6 +194,10 @@ async fn main() -> Result<()> { args.crawl.apply(&mut config); daemon(&config, &args).await?; } + Command::GlobalMerge(args) => { + let config = Config::from_env()?; + global_merge(&config, &args).await?; + } Command::Bots => print_bots().await?, #[cfg(feature = "aws")] Command::Publish(args) => { @@ -180,42 +211,58 @@ async fn main() -> Result<()> { Ok(()) } -async fn crawl_once(config: &Config, options: &CrawlOptions) -> Result<()> { +async fn crawl_once(config: &Config, options: &CrawlOptions, seed: Option<&Path>) -> Result<()> { + // A seed downloaded for this run is removed once used; a --seed file is + // the caller's. + #[allow(unused_mut)] + let (mut seed, mut downloaded) = (seed.map(Path::to_path_buf), false); #[cfg(feature = "aws")] - if config.publish { - publish::restore_state_if_missing(config).await?; + if config.publish + && let Some(restored) = publish::restore_state_if_missing(config).await? + { + (seed, downloaded) = (Some(restored), true); } - let db = Db::open(&config.db_path) - .with_context(|| format!("opening {}", config.db_path.display()))?; + let db = match seed { + Some(seed) if config.shard.is_sharded() && !config.db_path.exists() => { + let started = std::time::Instant::now(); + let db = Db::seed_shard(&config.db_path, &seed, config.shard).with_context(|| { + format!( + "seeding {} from {}", + config.db_path.display(), + seed.display() + ) + })?; + tracing::info!( + shard = %config.shard, + elapsed_secs = started.elapsed().as_secs(), + "seeded the shard's database from the single-task one" + ); + if downloaded { + let _ = std::fs::remove_file(&seed); + } + db + } + _ => Db::open(&config.db_path) + .with_context(|| format!("opening {}", config.db_path.display()))?, + }; db.set_branch_scope(config.branch_scope)?; let report = crawl::run(config, &db, options).await?; if options.global { - let summary = global_summary::from_db( - &db, - &config.global_imported_before, - config.global_repo_age_grace_days, - options.filter_only, - )?; - global_summary::write(config, &summary).await?; - #[cfg(feature = "aws")] - if config.publish { - publish::publish_state(config).await?; - } #[cfg(not(feature = "aws"))] if config.publish { tracing::warn!( "TALLY_PUBLISH is set but this binary was built without the aws feature" ); } + crawl::publish_global(config, &db, options.filter_only).await?; tracing::info!( + shard = %config.shard, repos_total = report.repos_total, repos_queued = report.repos_queued, commits_seen = report.scan.commits_seen, commits_new = report.scan.commits_new, - days = summary.days.len(), - repos_reached = summary.repos_reached, "global crawl complete" ); return Ok(()); @@ -247,6 +294,9 @@ async fn crawl_once(config: &Config, options: &CrawlOptions) -> Result<()> { #[cfg(feature = "aws")] if config.publish { + // The upload is the database file: fold the log in first, or the + // restored copy would lack the run's last writes. + db.flush_wal()?; publish::publish(config, &config.out_dir).await?; } #[cfg(not(feature = "aws"))] @@ -257,6 +307,51 @@ async fn crawl_once(config: &Config, options: &CrawlOptions) -> Result<()> { Ok(()) } +/// Download every shard's counted changes from the state bucket. +#[cfg(feature = "aws")] +async fn download_counts(config: &Config, shards: usize) -> Result> { + let mut paths = Vec::with_capacity(shards); + for index in 0..shards { + let key = shard::Shard::new(index, shards)?.counts_key(&config.base_state_key); + let path = config + .db_path + .with_extension(format!("shard-{index}.counts.sqlite3")); + publish::download_file(config, &key, &path).await?; + paths.push(path); + } + Ok(paths) +} + +#[cfg(not(feature = "aws"))] +async fn download_counts(_config: &Config, _shards: usize) -> Result> { + anyhow::bail!("pass --counts files; this build cannot read them from S3") +} + +/// Build the global summary from every shard's counted changes and write it. +async fn global_merge(config: &Config, args: &MergeArgs) -> Result<()> { + let shards = args.shards.unwrap_or(if args.counts.is_empty() { + config.shard.count + } else { + args.counts.len() + }); + let paths = if args.counts.is_empty() { + download_counts(config, shards).await? + } else { + args.counts.clone() + }; + let summary = global_summary::from_counts(&paths, &config.global_imported_before, shards)?; + global_summary::write(config, &summary).await?; + tracing::info!( + shards, + partial = summary.partial, + generated_at = %summary.generated_at, + days = summary.days.len(), + changes = summary.total(None), + "merged the global shards" + ); + Ok(()) +} + async fn daemon(config: &Config, args: &DaemonArgs) -> Result<()> { if let Some(bind) = &args.health_bind { let dir = config.out_dir.clone(); @@ -269,7 +364,7 @@ async fn daemon(config: &Config, args: &DaemonArgs) -> Result<()> { } loop { - match crawl_once(config, &args.crawl.options()).await { + match crawl_once(config, &args.crawl.options(), args.crawl.seed.as_deref()).await { Ok(()) => {} Err(err) => tracing::error!(error = %err, "crawl failed; keeping the previous site"), } diff --git a/src/publish.rs b/src/publish.rs index 35c63b0..ffd0e21 100644 --- a/src/publish.rs +++ b/src/publish.rs @@ -16,46 +16,69 @@ async fn client() -> aws_sdk_s3::Client { /// Download the last published index when this task has no local copy. A /// container restart is the normal case; a warm task keeps its database. -pub async fn restore_state_if_missing(config: &Config) -> Result<()> { +/// +/// A shard of the global crawl with no state of its own yet downloads the +/// single-task database instead and returns its path, for the caller to trim +/// to the shard with `Db::seed_shard`. +pub async fn restore_state_if_missing(config: &Config) -> Result> { let Some(bucket) = config.state_bucket.as_deref() else { - return Ok(()); + return Ok(None); }; if tokio::fs::metadata(&config.db_path).await.is_ok() { - return Ok(()); + return Ok(None); } let client = client().await; - let output = client - .get_object() - .bucket(bucket) - .key(&config.state_key) - .send() - .await; - let output = match output { + if download(&client, bucket, &config.state_key, &config.db_path).await? { + tracing::info!(bucket, key = %config.state_key, "restored index state"); + return Ok(None); + } + if let Some(seed_key) = config.seed_state_key() { + let seed = config.db_path.with_extension("seed.sqlite3"); + if download(&client, bucket, seed_key, &seed).await? { + tracing::info!(bucket, key = seed_key, shard = %config.shard, "restored seed state for a new shard"); + return Ok(Some(seed)); + } + } + tracing::info!(bucket, key = %config.state_key, "no previous state to restore"); + Ok(None) +} + +/// Download one shard's counted changes for the merge. +pub async fn download_file(config: &Config, key: &str, path: &Path) -> Result<()> { + let Some(bucket) = config.state_bucket.as_deref() else { + bail!("TALLY_STATE_BUCKET is not set"); + }; + if !download(&client().await, bucket, key, path).await? { + bail!("s3://{bucket}/{key} does not exist"); + } + Ok(()) +} + +/// Stream `s3://bucket/key` to `path`. False when there is no such object: a +/// first run has no state yet, and that is not an error. Any other failure +/// must stop the run: starting from an empty index and then uploading it +/// would replace good state with a partial one. +async fn download( + client: &aws_sdk_s3::Client, + bucket: &str, + key: &str, + path: &Path, +) -> Result { + let output = match client.get_object().bucket(bucket).key(key).send().await { Ok(output) => output, Err(err) => { - // A first run has no state object yet. That is not an error. Any - // other failure must stop the run: starting from an empty index - // and then uploading it would replace good state with a partial - // one. let not_found = err .as_service_error() .map(|service| service.is_no_such_key()) .unwrap_or(false); if not_found { - tracing::info!(bucket, key = %config.state_key, "no previous state to restore"); - return Ok(()); + return Ok(false); } - return Err(err).with_context(|| { - format!( - "reading previous state from s3://{bucket}/{}", - config.state_key - ) - }); + return Err(err).with_context(|| format!("reading s3://{bucket}/{key}")); } }; - write_body_to(output.body, &config.db_path).await?; - tracing::info!(bucket, key = %config.state_key, "restored index state"); - Ok(()) + write_body_to(output.body, path).await?; + Ok(true) } /// Stream an object body to `path` through a sibling file, then rename it into @@ -135,28 +158,36 @@ pub async fn publish_state(config: &Config) -> Result<()> { publish_state_path(config, &config.db_path).await } -/// Upload a snapshot made by `Db::snapshot_to`, used by the periodic global -/// checkpoint while the live database is still being written. -pub async fn publish_state_path(config: &Config, path: &Path) -> Result<()> { - let Some(state_bucket) = config.state_bucket.as_deref() else { +/// Upload the database at `path` to the state key. The caller makes sure the +/// file is whole (`Db::flush_wal`) and that nothing writes until this returns. +async fn publish_state_path(config: &Config, path: &Path) -> Result<()> { + if config.state_bucket.is_none() { tracing::warn!("TALLY_STATE_BUCKET is not set; not uploading index state"); return Ok(()); - }; + } if tokio::fs::metadata(path).await.is_err() { return Ok(()); } - let client = client().await; - client + publish_file(config, path, &config.state_key).await +} + +/// Upload a file to the state bucket under `key`. +pub async fn publish_file(config: &Config, path: &Path, key: &str) -> Result<()> { + let Some(state_bucket) = config.state_bucket.as_deref() else { + bail!("TALLY_STATE_BUCKET is not set"); + }; + client() + .await .put_object() .bucket(state_bucket) - .key(&config.state_key) + .key(key) .body(ByteStream::from_path(path).await?) .content_type("application/x-sqlite3") .cache_control("no-store") .send() .await - .with_context(|| format!("uploading {}", config.state_key))?; - tracing::info!(bucket = state_bucket, key = %config.state_key, "published index state"); + .with_context(|| format!("uploading {key}"))?; + tracing::info!(bucket = state_bucket, key, "published index state"); Ok(()) } diff --git a/src/shard.rs b/src/shard.rs new file mode 100644 index 0000000..36839ef --- /dev/null +++ b/src/shard.rs @@ -0,0 +1,205 @@ +//! Splitting the global crawl into shards. +//! +//! Repositories are independent, so the global pass can run as N tasks, each +//! owning the repositories whose DID hashes to its index. The hash is FNV-1a, +//! written out here rather than taken from `std`, whose hasher may change +//! between Rust releases: a repository has to stay in the same shard from one +//! month to the next, or its history would be read again by a shard whose +//! database has never seen it. +//! +//! Counts can't simply be added across shards: a logical change reachable from +//! repositories on two shards would count twice. Each shard exports its counted +//! changes by key instead, and the merge takes the earliest day per key (see +//! `global_summary`). + +use std::fmt; + +use anyhow::{Context, Result, bail}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Shard { + pub index: usize, + pub count: usize, +} + +impl Default for Shard { + fn default() -> Self { + Self::WHOLE + } +} + +impl Shard { + /// The single-task crawl: one shard holding every repository. + pub const WHOLE: Self = Self { index: 0, count: 1 }; + + pub fn new(index: usize, count: usize) -> Result { + if count == 0 { + bail!("the shard count must be at least 1"); + } + if index >= count { + bail!("shard index {index} is out of range for {count} shards"); + } + Ok(Self { index, count }) + } + + /// Parse `i/N`, as `--shard` takes it. + pub fn parse(value: &str) -> Result { + let (index, count) = value.split_once('/').with_context(|| { + format!("a shard is written index/count, such as 0/8, not {value:?}") + })?; + let index = index + .trim() + .parse() + .with_context(|| format!("invalid shard index in {value:?}"))?; + let count = count + .trim() + .parse() + .with_context(|| format!("invalid shard count in {value:?}"))?; + Self::new(index, count) + } + + pub fn is_sharded(self) -> bool { + self.count > 1 + } + + /// The shard a repository belongs to under `count` shards. + pub fn of(repo_did: &str, count: usize) -> usize { + (fnv1a(repo_did.as_bytes()) % count.max(1) as u64) as usize + } + + pub fn contains(self, repo_did: &str) -> bool { + self.count <= 1 || Self::of(repo_did, self.count) == self.index + } + + /// The same for a `tangled:` reference key. Keys from other sources + /// aren't Tangled repositories and belong to no global shard but the + /// whole one. + pub fn contains_repo_key(self, repo_key: &str) -> bool { + match repo_key.strip_prefix("tangled:") { + Some(repo_did) => self.contains(repo_did), + None => self.count <= 1, + } + } + + /// The object key this shard keeps its database under, derived from the + /// single-task key: `state/global.sqlite3` becomes + /// `state/global-shard-2-of-8.sqlite3`. The whole crawl keeps the key. + pub fn state_key(self, base: &str) -> String { + if !self.is_sharded() { + return base.to_string(); + } + let (stem, extension) = split_extension(base); + format!("{stem}-shard-{}-of-{}{extension}", self.index, self.count) + } + + /// Where this shard publishes its counted changes for the merge: + /// `state/global-shard-2-of-8.counts.sqlite3`. + #[cfg_attr(not(feature = "aws"), allow(dead_code))] + pub fn counts_key(self, base: &str) -> String { + let key = self.state_key(base); + let (stem, _) = split_extension(&key); + format!("{stem}.counts.sqlite3") + } +} + +impl fmt::Display for Shard { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}/{}", self.index, self.count) + } +} + +fn split_extension(key: &str) -> (&str, &str) { + let name_start = key.rfind('/').map_or(0, |slash| slash + 1); + match key[name_start..].rfind('.') { + Some(dot) if dot > 0 => key.split_at(name_start + dot), + _ => (key, ""), + } +} + +fn fnv1a(bytes: &[u8]) -> u64 { + let mut hash: u64 = 0xcbf2_9ce4_8422_2325; + for byte in bytes { + hash ^= u64::from(*byte); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + hash +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_and_rejects_shards() { + assert_eq!(Shard::parse("2/8").unwrap(), Shard { index: 2, count: 8 }); + assert_eq!(Shard::parse("0/1").unwrap(), Shard::WHOLE); + for bad in ["8/8", "1/0", "x/2", "3", ""] { + assert!(Shard::parse(bad).is_err(), "{bad} should not parse"); + } + } + + #[test] + fn assignment_is_pinned_so_repositories_never_move() { + // These values are part of the stored state's layout: changing the + // hash would hand every repository to a shard that has never read it. + assert_eq!(fnv1a(b""), 0xcbf2_9ce4_8422_2325); + assert_eq!(fnv1a(b"a"), 0xaf63_dc4c_8601_ec8c); + let did = "did:plc:ts5zhqjndasud5h5etjzks62"; + assert_eq!(Shard::of(did, 8), (fnv1a(did.as_bytes()) % 8) as usize); + assert_eq!(Shard::of(did, 8), Shard::of(did, 8)); + assert_eq!(Shard::of(did, 1), 0); + } + + #[test] + fn every_repository_is_in_exactly_one_shard() { + let dids: Vec = (0..5_000).map(|i| format!("did:plc:repo{i:05}")).collect(); + for count in [1, 2, 3, 8, 13] { + let shards: Vec = (0..count).map(|i| Shard::new(i, count).unwrap()).collect(); + let mut sizes = vec![0usize; count]; + for did in &dids { + let owners: Vec = shards + .iter() + .filter(|shard| shard.contains(did)) + .map(|shard| shard.index) + .collect(); + assert_eq!(owners.len(), 1, "{did} under {count} shards: {owners:?}"); + assert!(shards[owners[0]].contains_repo_key(&format!("tangled:{did}"))); + sizes[owners[0]] += 1; + } + // Roughly even: no shard more than half again the mean. + let mean = dids.len() / count; + assert!( + sizes.iter().all(|size| *size < mean + mean / 2), + "{sizes:?}" + ); + } + } + + #[test] + fn keys_derive_from_the_single_task_key() { + let shard = Shard::new(2, 8).unwrap(); + assert_eq!( + shard.state_key("state/global.sqlite3"), + "state/global-shard-2-of-8.sqlite3" + ); + assert_eq!( + shard.counts_key("state/global.sqlite3"), + "state/global-shard-2-of-8.counts.sqlite3" + ); + assert_eq!( + Shard::WHOLE.state_key("state/global.sqlite3"), + "state/global.sqlite3" + ); + assert_eq!( + Shard::WHOLE.counts_key("state/global.sqlite3"), + "state/global.counts.sqlite3" + ); + assert_eq!(shard.state_key("a.b/global"), "a.b/global-shard-2-of-8"); + assert!( + !Shard::new(0, 2) + .unwrap() + .contains_repo_key("github:someone/repo") + ); + assert!(Shard::WHOLE.contains_repo_key("github:someone/repo")); + } +} -- knotmirror 0.0.1