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")); + } +}