diff --git a/consumer/src/backfill/db.rs b/consumer/src/backfill/db.rs index 3e299ef5..3b83b9d8 100644 --- a/consumer/src/backfill/db.rs +++ b/consumer/src/backfill/db.rs @@ -2,6 +2,40 @@ use diesel::prelude::*; use diesel_async::{AsyncPgConnection, RunQueryDsl}; use parakeet_db::{models, schema, types}; +#[derive(Debug, QueryableByName)] +#[diesel(table_name = schema::backfill_jobs)] +pub struct JobInfo { + pub id: i32, + pub did: String, + pub since: Option, +} + +pub async fn get_job(conn: &mut AsyncPgConnection) -> QueryResult> { + diesel::sql_query(include_str!("../sql/backfill_job.sql")) + .get_result(conn) + .await + .optional() +} + +pub async fn update_job_status( + conn: &mut AsyncPgConnection, + job: i32, + status: &str, +) -> QueryResult { + diesel::update(schema::backfill_jobs::table) + .set(schema::backfill_jobs::status.eq(status)) + .filter(schema::backfill_jobs::id.eq(job)) + .execute(conn) + .await +} + +pub async fn clear_job(conn: &mut AsyncPgConnection, job: i32) -> QueryResult { + diesel::delete(schema::backfill_jobs::table) + .filter(schema::backfill_jobs::id.eq(job)) + .execute(conn) + .await +} + pub async fn get_actor_status( conn: &mut AsyncPgConnection, did: &str, diff --git a/consumer/src/backfill/mod.rs b/consumer/src/backfill/mod.rs index 4f9556a4..e584d851 100644 --- a/consumer/src/backfill/mod.rs +++ b/consumer/src/backfill/mod.rs @@ -1,10 +1,9 @@ -use crate::config::HistoryMode; use crate::indexer::types::{AggregateDeltaStore, BackfillItem, BackfillItemInner}; use crate::indexer::{self, db as indexer_db}; use did_resolver::Resolver; use diesel_async::pooled_connection::deadpool::Pool; +use diesel_async::scoped_futures::ScopedFutureExt; use diesel_async::{AsyncConnection, AsyncPgConnection}; -use flume::{unbounded, Receiver, Sender}; use ipld_core::cid::Cid; use metrics::counter; use parakeet_db::types::{ActorStatus, ActorSyncState}; @@ -12,7 +11,7 @@ use reqwest::{Client, StatusCode}; use std::collections::HashMap; use std::str::FromStr; use std::sync::Arc; -use tracing::{instrument, Instrument}; +use tracing::instrument; mod db; mod repo; @@ -25,69 +24,42 @@ const DELTA_BATCH_SIZE: usize = 32 * 1024; #[derive(Clone)] pub struct BackfillManagerInner { - pool: Pool, resolver: Arc, client: Client, index_client: parakeet_index::Client, } pub struct BackfillManager { - rx: Receiver, + pool: Pool, inner: BackfillManagerInner, - do_backfill: bool, } impl BackfillManager { pub async fn new( pool: Pool, - history_mode: HistoryMode, resolver: Arc, index_client: parakeet_index::Client, - ) -> eyre::Result<(Self, Sender)> { + ) -> eyre::Result { let client = Client::new(); - let (tx, rx) = unbounded(); - - let mgr = BackfillManager { + Ok(BackfillManager { + pool, inner: BackfillManagerInner { - pool, resolver, client, index_client, }, - rx, - do_backfill: history_mode == HistoryMode::BackfillHistory, - }; - - Ok((mgr, tx)) + }) } pub async fn run(self, threads: u8) -> eyre::Result<()> { let mut js = tokio::task::JoinSet::new(); - if self.do_backfill { - for idx in 0..threads { - let rx = self.rx.clone(); - let mut 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(&mut inner, &did).await { - tracing::error!(did, "backfill failed: {e}"); - counter!("backfill_failure").increment(1); - } else { - counter!("backfill_success").increment(1); - } - - counter!("backfill_pending").absolute(rx.len() as u64); - } - tracing::error!("backfill rx error?"); - } - .instrument(tracing::info_span!("backfill_worker", thread = idx)), - ); - } + for idx in 0..threads { + let inner = self.inner.clone(); + let pool = self.pool.clone(); + + js.spawn(backfill_worker(pool, inner, idx)); } js.join_all().await; @@ -96,11 +68,60 @@ impl BackfillManager { } } -#[instrument(skip(inner))] -async fn backfill_actor(inner: &mut BackfillManagerInner, did: &str) -> eyre::Result<()> { - let mut conn = inner.pool.get().await?; +#[instrument(skip_all, fields(thread = _thread))] +async fn backfill_worker( + pool: Pool, + mut inner: BackfillManagerInner, + _thread: u8, +) { + let mut conn = pool.get().await.unwrap(); + + loop { + let res = conn + .transaction::<_, eyre::Report, _>(|t| { + async { + let Some(job) = db::get_job(t).await? else { + return Ok(()); + }; + tracing::trace!("backfilling {}", &job.did); + + // savepoint for the actual backfill insert + let t_res = t + .transaction(|t2| { + async { backfill_actor(t2, &mut inner, &job.did).await }.scope_boxed() + }) + .await; + + if let Err(e) = t_res { + tracing::error!(did = &job.did, "backfill failed: {e}"); + counter!("backfill_failure").increment(1); - let (status, sync_state) = db::get_actor_status(&mut conn, did).await?; + db::update_job_status(t, job.id, "failed").await?; + } else { + counter!("backfill_success").increment(1); + + db::clear_job(t, job.id).await?; + } + + Ok(()) + } + .scope_boxed() + }) + .await; + + if let Err(e) = res { + tracing::error!("backfill process failure: {e}"); + } + } +} + +#[instrument(skip(conn, inner))] +async fn backfill_actor( + conn: &mut AsyncPgConnection, + inner: &mut BackfillManagerInner, + did: &str, +) -> eyre::Result<()> { + let (status, sync_state) = db::get_actor_status(conn, did).await?; if sync_state != ActorSyncState::Dirty || status != ActorStatus::Active { tracing::debug!("skipping non-dirty or inactive repo"); @@ -139,14 +160,13 @@ async fn backfill_actor(inner: &mut BackfillManagerInner, did: &str) -> eyre::Re } // set the handle from above - db::update_handle(&mut conn, did, Some(handle)).await?; + db::update_handle(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?; + db::update_actor_status(conn, did, ActorStatus::Deleted, ActorSyncState::Dirty).await?; return Ok(()); }; @@ -155,12 +175,12 @@ async fn backfill_actor(inner: &mut BackfillManagerInner, did: &str) -> eyre::Re let status = repo_status .status .unwrap_or(crate::firehose::AtpAccountStatus::Deleted); - db::update_actor_status(&mut conn, did, status.into(), ActorSyncState::Dirty).await?; + db::update_actor_status(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?; + db::update_repo_sync_state(conn, did, ActorSyncState::Processing).await?; tracing::trace!("pulling repo"); @@ -168,41 +188,32 @@ async fn backfill_actor(inner: &mut BackfillManagerInner, did: &str) -> eyre::Re tracing::trace!("repo pulled - inserting"); - let delta_store = conn - .transaction::<_, diesel::result::Error, _>(|t| { - Box::pin(async move { - let mut delta_store = HashMap::new(); - - db::defer(t).await?; + let mut delta_store = HashMap::new(); - indexer_db::update_repo_version(t, did, &rev, cid).await?; + db::defer(conn).await?; - // let mut follow_stats = vec![did.to_string()]; + indexer_db::update_repo_version(conn, did, &rev, cid).await?; - for (path, (cid, record)) in records { - let Some((collection, 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((collection, rkey)) = path.split_once("/") else { + tracing::warn!("record contained invalid path {}", path); + return Err(diesel::result::Error::RollbackTransaction.into()); + }; - counter!("backfilled_commits", "collection" => collection.to_string()) - .increment(1); + counter!("backfilled_commits", "collection" => collection.to_string()) + .increment(1); - let full_path = format!("at://{did}/{path}"); + let full_path = format!("at://{did}/{path}"); - indexer::index_op(t, &mut delta_store, did, cid, record, &full_path, rkey) - .await? - } + indexer::index_op(conn, &mut delta_store, did, cid, record, &full_path, rkey) + .await? + } - db::update_repo_sync_state(t, did, ActorSyncState::Synced).await?; + db::update_repo_sync_state(conn, did, ActorSyncState::Synced).await?; - handle_backfill_rows(t, &mut delta_store, did, &rev).await?; - tracing::trace!("insertion finished"); + handle_backfill_rows(conn, &mut delta_store, did, &rev).await?; - Ok(delta_store) - }) - }) - .await?; + tracing::trace!("insertion finished"); // submit the deltas let delta_store = delta_store diff --git a/consumer/src/indexer/db.rs b/consumer/src/indexer/db.rs index f283bfa2..95f29eaf 100644 --- a/consumer/src/indexer/db.rs +++ b/consumer/src/indexer/db.rs @@ -29,6 +29,16 @@ pub async fn write_backfill_row( .await } +pub async fn write_backfill_job(conn: &mut AsyncPgConnection, repo: &str) -> QueryResult { + diesel::insert_into(schema::backfill_jobs::table) + .values(( + schema::backfill_jobs::did.eq(repo), + schema::backfill_jobs::status.eq("pending"), + )) + .execute(conn) + .await +} + pub async fn upsert_actor( conn: &mut AsyncPgConnection, did: &str, diff --git a/consumer/src/indexer/mod.rs b/consumer/src/indexer/mod.rs index 970c9254..95aaac0d 100644 --- a/consumer/src/indexer/mod.rs +++ b/consumer/src/indexer/mod.rs @@ -24,7 +24,6 @@ pub mod types; #[derive(Clone)] struct RelayIndexerState { - backfill_tx: flume::Sender, idxc_tx: Sender, resolver: Arc, do_backfill: bool, @@ -40,7 +39,6 @@ pub struct RelayIndexer { impl RelayIndexer { pub async fn new( pool: Pool, - backfill_tx: flume::Sender, idxc_tx: Sender, resolver: Arc, history_mode: HistoryMode, @@ -50,7 +48,6 @@ impl RelayIndexer { let indexer = RelayIndexer { pool, state: RelayIndexerState { - backfill_tx, resolver, do_backfill: history_mode == HistoryMode::BackfillHistory, idxc_tx, @@ -209,7 +206,7 @@ async fn index_commit( // TODO: bridgy doesn't implement since atm - we need a special case if commit.since.is_some() { if state.do_backfill { - state.backfill_tx.send_async(commit.repo.clone()).await?; + db::write_backfill_job(conn, &commit.repo).await?; } return Ok(()); } @@ -236,7 +233,7 @@ async fn index_commit( db::upsert_actor(conn, &commit.repo, None, None, sync_state, commit.time).await?; if trigger_backfill { - state.backfill_tx.send_async(commit.repo.clone()).await?; + db::write_backfill_job(conn, &commit.repo).await?; return Ok(()); } diff --git a/consumer/src/main.rs b/consumer/src/main.rs index faf618ad..727082fe 100644 --- a/consumer/src/main.rs +++ b/consumer/src/main.rs @@ -41,24 +41,15 @@ async fn main() -> eyre::Result<()> { .await?; let relay_firehose = firehose::FirehoseConsumer::new_relay(&conf.relay_source, None, &user_agent).await?; - let (backfiller, backfill_tx) = backfill::BackfillManager::new( - pool.clone(), - conf.history_mode, - resolver.clone(), - index_client.clone(), - ) - .await?; + let backfiller = + backfill::BackfillManager::new(pool.clone(), resolver.clone(), index_client.clone()) + .await?; let (idxc_tx, idxc_rx) = tokio::sync::mpsc::channel(128); - let (relay_indexer, tx) = indexer::RelayIndexer::new( - pool.clone(), - backfill_tx, - idxc_tx, - resolver.clone(), - conf.history_mode, - ) - .await?; + let (relay_indexer, tx) = + indexer::RelayIndexer::new(pool.clone(), idxc_tx, resolver.clone(), conf.history_mode) + .await?; let (firehose_res, indexer_res, backfill_res, label_res, idxt_res) = tokio::try_join! { tokio::spawn(relay_consumer(relay_firehose, tx)), diff --git a/consumer/src/sql/backfill_job.sql b/consumer/src/sql/backfill_job.sql new file mode 100644 index 00000000..2432c9d6 --- /dev/null +++ b/consumer/src/sql/backfill_job.sql @@ -0,0 +1,5 @@ +UPDATE backfill_jobs +SET status='pending', + updated_at=NOW() +WHERE id = (SELECT id FROM backfill_jobs WHERE status = 'pending' ORDER BY id FOR UPDATE SKIP LOCKED LIMIT 1) +RETURNING id, did, since; \ No newline at end of file diff --git a/migrations/2025-02-12-211348_backfill/down.sql b/migrations/2025-02-12-211348_backfill/down.sql index 2248002f..3ff658e7 100644 --- a/migrations/2025-02-12-211348_backfill/down.sql +++ b/migrations/2025-02-12-211348_backfill/down.sql @@ -1 +1,2 @@ -drop table backfill; \ No newline at end of file +drop table backfill; +drop table backfill_jobs; \ No newline at end of file diff --git a/migrations/2025-02-12-211348_backfill/up.sql b/migrations/2025-02-12-211348_backfill/up.sql index ceb01342..4acd9446 100644 --- a/migrations/2025-02-12-211348_backfill/up.sql +++ b/migrations/2025-02-12-211348_backfill/up.sql @@ -4,7 +4,7 @@ create table backfill repo_ver text not null, cid text not null, - data jsonb not null, + data jsonb not null, indexed_at timestamp not null default now(), @@ -12,4 +12,18 @@ create table backfill ); create index backfill_repo_idx on backfill using hash (repo); -create index backfill_repover_idx on backfill (repo_ver); \ No newline at end of file +create index backfill_repover_idx on backfill (repo_ver); + +create table backfill_jobs +( + id serial primary key, + did text not null, + since text, -- for future use. + + status text not null, -- pending/processing/failed/completed (completed rows will be deleted) + created_at timestamp not null default now(), + updated_at timestamp not null default now() +); + +create index bfjobs_did on backfill_jobs (did); +create index bfjobs_status on backfill_jobs (status); \ No newline at end of file diff --git a/parakeet-db/src/schema.rs b/parakeet-db/src/schema.rs index 196e2601..cc92ae96 100644 --- a/parakeet-db/src/schema.rs +++ b/parakeet-db/src/schema.rs @@ -22,6 +22,17 @@ diesel::table! { } } +diesel::table! { + backfill_jobs (id) { + id -> Int4, + did -> Text, + since -> Nullable, + status -> Text, + created_at -> Timestamp, + updated_at -> Timestamp, + } +} + diesel::table! { blocks (at_uri) { at_uri -> Text, @@ -333,6 +344,7 @@ diesel::joinable!(verification -> actors (verifier)); diesel::allow_tables_to_appear_in_same_query!( actors, backfill, + backfill_jobs, blocks, chat_decls, feedgens,