From fdeedc896cc61bd79fbf02c3c5dd3735d8423394 Mon Sep 17 00:00:00 2001 From: Mia Date: Sat, 5 Apr 2025 12:21:21 +0100 Subject: [PATCH] feat(consumer): multithreaded backfiller also some improvements to cut down on backfiller errors --- Cargo.lock | 29 ++++- consumer/Cargo.toml | 1 + consumer/src/backfill/db.rs | 31 ++++- consumer/src/backfill/mod.rs | 232 +++++++++++++++++++++------------ consumer/src/backfill/types.rs | 8 ++ consumer/src/config.rs | 6 +- consumer/src/indexer/mod.rs | 8 +- consumer/src/main.rs | 2 +- 8 files changed, 224 insertions(+), 93 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ab1af78a..d797223e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "addr2line" @@ -647,6 +647,7 @@ dependencies = [ "diesel-async", "eyre", "figment", + "flume", "futures", "ipld-core", "iroh-car", @@ -1037,6 +1038,18 @@ dependencies = [ "version_check", ] +[[package]] +name = "flume" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da0e4dd2a88388a1f4ccc7c9ce104604dab68d9f408dc34cd45823d5a9069095" +dependencies = [ + "futures-core", + "futures-sink", + "nanorand", + "spin", +] + [[package]] name = "fnv" version = "1.0.7" @@ -1198,8 +1211,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" dependencies = [ "cfg-if", + "js-sys", "libc", "wasi", + "wasm-bindgen", ] [[package]] @@ -1981,6 +1996,15 @@ dependencies = [ "unsigned-varint 0.8.0", ] +[[package]] +name = "nanorand" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a51313c5820b0b02bd422f4b44776fbf47961755c74ce64afc73bfad10226c3" +dependencies = [ + "getrandom", +] + [[package]] name = "native-tls" version = "0.2.12" @@ -2879,6 +2903,9 @@ name = "spin" version = "0.9.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +dependencies = [ + "lock_api", +] [[package]] name = "stable_deref_trait" diff --git a/consumer/Cargo.toml b/consumer/Cargo.toml index 9e927d8f..633ccc5e 100644 --- a/consumer/Cargo.toml +++ b/consumer/Cargo.toml @@ -11,6 +11,7 @@ diesel = { version = "2.2.6", features = ["chrono", "serde_json"] } diesel-async = { version = "0.5.2", features = ["deadpool", "postgres"] } eyre = "0.6.12" figment = { version = "0.10.19", features = ["env", "toml"] } +flume = { version = "0.11.1", features = ["async"] } futures = "0.3.31" ipld-core = "0.4.1" iroh-car = "0.5.1" diff --git a/consumer/src/backfill/db.rs b/consumer/src/backfill/db.rs index 8c246abb..71d444f0 100644 --- a/consumer/src/backfill/db.rs +++ b/consumer/src/backfill/db.rs @@ -25,13 +25,42 @@ pub async fn update_repo_sync_state( .await } +pub async fn update_handle( + conn: &mut AsyncPgConnection, + did: &str, + handle: Option, +) -> QueryResult { + diesel::update(schema::actors::table) + .set(schema::actors::handle.eq(handle)) + .filter(schema::actors::did.eq(did)) + .execute(conn) + .await +} + +pub async fn update_actor_status( + conn: &mut AsyncPgConnection, + did: &str, + status: types::ActorStatus, + sync_state: types::ActorSyncState, +) -> QueryResult { + diesel::update(schema::actors::table) + .set((schema::actors::status.eq(status), schema::actors::sync_state.eq(sync_state))) + .filter(schema::actors::did.eq(did)) + .execute(conn) + .await +} + pub async fn defer(conn: &mut AsyncPgConnection) -> QueryResult { diesel::sql_query("SET CONSTRAINTS ALL DEFERRED") .execute(conn) .await } -pub async fn pull_backfill_rows(conn: &mut AsyncPgConnection, repo: &str, rev: &str) -> QueryResult> { +pub async fn pull_backfill_rows( + conn: &mut AsyncPgConnection, + repo: &str, + rev: &str, +) -> QueryResult> { schema::backfill::table .select(models::BackfillRow::as_select()) .filter( diff --git a/consumer/src/backfill/mod.rs b/consumer/src/backfill/mod.rs index a2284f5c..fb156d3d 100644 --- a/consumer/src/backfill/mod.rs +++ b/consumer/src/backfill/mod.rs @@ -4,13 +4,13 @@ use crate::indexer::{self, db as indexer_db}; use did_resolver::Resolver; use diesel_async::pooled_connection::deadpool::Pool; use diesel_async::{AsyncConnection, AsyncPgConnection}; +use flume::{unbounded, Receiver, Sender}; use ipld_core::cid::Cid; use parakeet_db::types::{ActorStatus, ActorSyncState}; -use reqwest::Client; +use reqwest::{Client, StatusCode}; use std::str::FromStr; use std::sync::Arc; -use tokio::sync::mpsc::{channel, Receiver, Sender}; -use tracing::instrument; +use tracing::{instrument, Instrument}; mod db; mod repo; @@ -18,11 +18,16 @@ mod types; const PDS_SERVICE_ID: &str = "#atproto_pds"; -pub struct BackfillManager { +#[derive(Clone)] +pub struct BackfillManagerInner { pool: Pool, - rx: Receiver, resolver: Arc, client: Client, +} + +pub struct BackfillManager { + rx: Receiver, + inner: BackfillManagerInner, do_backfill: bool, } @@ -34,124 +39,162 @@ impl BackfillManager { ) -> eyre::Result<(Self, Sender)> { let client = Client::new(); - let (tx, rx) = channel(64); + let (tx, rx) = unbounded(); let mgr = BackfillManager { - pool: pool.clone(), + inner: BackfillManagerInner { + pool, + resolver, + client, + }, rx, - resolver, - client, do_backfill: history_mode == HistoryMode::BackfillHistory, }; Ok((mgr, tx)) } - pub async fn run(mut self) -> eyre::Result<()> { + pub async fn run(self, threads: u8) -> eyre::Result<()> { + let mut js = tokio::task::JoinSet::new(); + if self.do_backfill { - while let Some(did) = self.rx.recv().await { - tracing::trace!("backfilling {did}"); - if let Err(e) = self.backfill_actor(&did).await { - tracing::error!(did, "backfill failed: {e}"); - } + for idx in 0..threads { + let rx = self.rx.clone(); + let inner = self.inner.clone(); + + js.spawn( + async move { + while let Ok(did) = rx.recv_async().await { + tracing::trace!("backfilling {did}"); + if let Err(e) = backfill_actor(&inner, &did).await { + tracing::error!(did, "backfill failed: {e}"); + } + } + tracing::error!("backfill rx error?"); + } + .instrument(tracing::info_span!("backfill_worker", thread = idx)), + ); } } + js.join_all().await; + Ok(()) } +} - #[instrument(skip(self))] - async fn backfill_actor(&mut self, did: &str) -> eyre::Result<()> { - let mut conn = self.pool.get().await?; - - let (status, sync_state) = db::get_actor_status(&mut conn, did).await?; - - if sync_state != ActorSyncState::Dirty || status != ActorStatus::Active { - tracing::debug!("skipping non-dirty or inactive repo"); - return Ok(()); - } - - // resolve the did to a PDS (also validates the handle) - let Some(did_doc) = self.resolver.resolve_did(did).await? else { - eyre::bail!("missing did doc"); - }; +#[instrument(skip(inner))] +async fn backfill_actor(inner: &BackfillManagerInner, did: &str) -> eyre::Result<()> { + let mut conn = inner.pool.get().await?; - let Some(handle) = did_doc - .also_known_as - .and_then(|aka| aka.first().cloned()) - .and_then(|handle| handle.strip_prefix("at://").map(String::from)) - else { - eyre::bail!("DID doc contained no handle"); - }; - let Some(service) = did_doc - .service - .and_then(|services| services.into_iter().find(|svc| svc.id == PDS_SERVICE_ID)) - else { - eyre::bail!("DID doc contained no service endpoint"); - }; + let (status, sync_state) = db::get_actor_status(&mut conn, did).await?; - let pds_url = service.service_endpoint; + if sync_state != ActorSyncState::Dirty || status != ActorStatus::Active { + tracing::debug!("skipping non-dirty or inactive repo"); + return Ok(()); + } - // in theory, we can use com.atproto.identity.resolveHandle against a PDS, but that seems - // like a way to end up with really sus handles. - let Some(handle_did) = self.resolver.resolve_handle(&handle).await? else { - eyre::bail!("Failed to resolve did for handle {handle}"); - }; + // resolve the did to a PDS (also validates the handle) + let Some(did_doc) = inner.resolver.resolve_did(did).await? else { + eyre::bail!("missing did doc"); + }; + + let Some(handle) = did_doc + .also_known_as + .and_then(|aka| aka.first().cloned()) + .and_then(|handle| handle.strip_prefix("at://").map(String::from)) + else { + eyre::bail!("DID doc contained no handle"); + }; + let Some(service) = did_doc + .service + .and_then(|services| services.into_iter().find(|svc| svc.id == PDS_SERVICE_ID)) + else { + eyre::bail!("DID doc contained no service endpoint"); + }; + + let pds_url = service.service_endpoint; + + // in theory, we can use com.atproto.identity.resolveHandle against a PDS, but that seems + // like a way to end up with really sus handles. + let Some(handle_did) = inner.resolver.resolve_handle(&handle).await? else { + eyre::bail!("Failed to resolve did for handle {handle}"); + }; + + if handle_did != did { + eyre::bail!("requested DID doesn't match handle"); + } - if handle_did != did { - eyre::bail!("requested DID doesn't match handle"); - } + // set the handle from above + db::update_handle(&mut conn, did, Some(handle)).await?; + + // check the repo status + let Some(repo_status) = check_pds_repo_status(&inner.client, &pds_url, did).await? else { + // this repo can't be found - set dirty and assume deleted. + tracing::debug!("repo was deleted"); + db::update_actor_status(&mut conn, did, ActorStatus::Deleted, ActorSyncState::Dirty) + .await?; + return Ok(()); + }; + + if !repo_status.active { + tracing::debug!("repo is inactive"); + let status = repo_status + .status + .unwrap_or(crate::firehose::AtpAccountStatus::Deleted); + db::update_actor_status(&mut conn, did, status.into(), ActorSyncState::Dirty).await?; + return Ok(()); + } - // now we can start actually backfilling - db::update_repo_sync_state(&mut conn, did, ActorSyncState::Processing).await?; + // now we can start actually backfilling + db::update_repo_sync_state(&mut conn, did, ActorSyncState::Processing).await?; - tracing::trace!("pulling repo"); + tracing::trace!("pulling repo"); - let (rev, cid, records) = repo::pull_repo(&self.client, did, &pds_url).await?; + let (rev, cid, records) = repo::pull_repo(&inner.client, did, &pds_url).await?; - tracing::trace!("repo pulled - inserting"); + tracing::trace!("repo pulled - inserting"); - conn.transaction::<(), diesel::result::Error, _>(|t| { - Box::pin(async move { - db::defer(t).await?; + conn.transaction::<(), diesel::result::Error, _>(|t| { + Box::pin(async move { + db::defer(t).await?; - indexer_db::update_repo_version(t, did, &rev, cid).await?; + indexer_db::update_repo_version(t, did, &rev, cid).await?; - let mut follow_stats = vec![did.to_string()]; + let mut follow_stats = vec![did.to_string()]; - for (path, (cid, record)) in records { - let Some((_, rkey)) = path.split_once("/") else { - tracing::warn!("record contained invalid path {}", path); - return Err(diesel::result::Error::RollbackTransaction); - }; + for (path, (cid, record)) in records { + let Some((_, rkey)) = path.split_once("/") else { + tracing::warn!("record contained invalid path {}", path); + return Err(diesel::result::Error::RollbackTransaction); + }; - let full_path = format!("at://{did}/{path}"); + let full_path = format!("at://{did}/{path}"); - match record { - RecordTypes::AppBskyGraphFollow(record) => { - follow_stats.push(record.subject.clone()); - indexer_db::insert_follow(t, did, &full_path, record).await?; - } - _ => indexer::index_op(t, did, cid, record, &full_path, rkey).await?, + match record { + RecordTypes::AppBskyGraphFollow(record) => { + follow_stats.push(record.subject.clone()); + indexer_db::insert_follow(t, did, &full_path, record).await?; } + _ => indexer::index_op(t, did, cid, record, &full_path, rkey).await?, } + } - db::update_repo_sync_state(t, did, ActorSyncState::Synced).await?; + db::update_repo_sync_state(t, did, ActorSyncState::Synced).await?; - handle_backfill_rows(t, &mut follow_stats, did, &rev).await?; + handle_backfill_rows(t, &mut follow_stats, did, &rev).await?; - // on second thought, should this be done after the transaction? - // if we're loading a chunky repo, we might be a few seconds+ out of date? - indexer_db::update_follow_stats(t, &follow_stats).await?; + // on second thought, should this be done after the transaction? + // if we're loading a chunky repo, we might be a few seconds+ out of date? + indexer_db::update_follow_stats(t, &follow_stats).await?; - tracing::trace!("insertion finished"); - Ok(()) - }) + tracing::trace!("insertion finished"); + Ok(()) }) - .await?; + }) + .await?; - Ok(()) - } + Ok(()) } async fn handle_backfill_rows( @@ -210,3 +253,22 @@ async fn handle_backfill_rows( Ok(()) } + +async fn check_pds_repo_status( + client: &Client, + pds: &str, + repo: &str, +) -> eyre::Result> { + let res = client + .get(format!( + "{pds}/xrpc/com.atproto.sync.getRepoStatus?did={repo}" + )) + .send() + .await?; + + if res.status() == StatusCode::NOT_FOUND { + return Ok(None); + } + + Ok(res.json().await?) +} diff --git a/consumer/src/backfill/types.rs b/consumer/src/backfill/types.rs index 729a6671..51885a5e 100644 --- a/consumer/src/backfill/types.rs +++ b/consumer/src/backfill/types.rs @@ -34,3 +34,11 @@ pub struct CarCommitEntry { pub prev: Option, pub sig: ByteBuf, } + +#[derive(Debug, Deserialize)] +pub struct GetRepoStatusRes { + pub did: String, + pub active: bool, + pub status: Option, + pub rev: Option, +} \ No newline at end of file diff --git a/consumer/src/config.rs b/consumer/src/config.rs index 3a8c3841..199e00d3 100644 --- a/consumer/src/config.rs +++ b/consumer/src/config.rs @@ -19,6 +19,8 @@ pub struct Config { pub history_mode: HistoryMode, /// Adds contact details (email / bluesky handle / website) to the UA header. pub ua_contact: Option, + #[serde(default = "default_backfill_workers")] + pub backfill_workers: u8, } #[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Deserialize)] @@ -28,4 +30,6 @@ pub enum HistoryMode { BackfillHistory, /// Discover new accounts as they come and do not import history Realtime, -} \ No newline at end of file +} + +fn default_backfill_workers() -> u8 { 4 } \ No newline at end of file diff --git a/consumer/src/indexer/mod.rs b/consumer/src/indexer/mod.rs index 928edef3..6abff5a3 100644 --- a/consumer/src/indexer/mod.rs +++ b/consumer/src/indexer/mod.rs @@ -20,7 +20,7 @@ pub mod types; pub struct RelayIndexer { conn: Object, rx: Receiver, - backfill_tx: Sender, + backfill_tx: flume::Sender, resolver: Arc, do_backfill: bool, } @@ -28,7 +28,7 @@ pub struct RelayIndexer { impl RelayIndexer { pub async fn new( pool: Pool, - backfill_tx: Sender, + backfill_tx: flume::Sender, resolver: Arc, history_mode: HistoryMode, ) -> eyre::Result<(Self, Sender)> { @@ -146,7 +146,7 @@ impl RelayIndexer { // TODO: bridgy doesn't implement since atm - we need a special case if commit.since.is_some() { if self.do_backfill { - self.backfill_tx.send(commit.repo.clone()).await?; + self.backfill_tx.send_async(commit.repo.clone()).await?; } return Ok(()); } @@ -181,7 +181,7 @@ impl RelayIndexer { .await?; if trigger_backfill { - self.backfill_tx.send(commit.repo.clone()).await?; + self.backfill_tx.send_async(commit.repo.clone()).await?; return Ok(()); } diff --git a/consumer/src/main.rs b/consumer/src/main.rs index 4f558599..637d60d6 100644 --- a/consumer/src/main.rs +++ b/consumer/src/main.rs @@ -45,7 +45,7 @@ async fn main() -> eyre::Result<()> { let (firehose_res, indexer_res, backfill_res) = tokio::try_join! { tokio::spawn(relay_consumer(relay_firehose, tx)), tokio::spawn(relay_indexer.run()), - tokio::spawn(backfiller.run()) + tokio::spawn(backfiller.run(conf.backfill_workers)) }?; firehose_res.and(indexer_res).and(backfill_res) -- 2.51.2