From 2106ba23e6a79da0a40fdf319abe0919cc39180c Mon Sep 17 00:00:00 2001 From: Mia Date: Sat, 12 Apr 2025 12:52:18 +0100 Subject: [PATCH] multithreaded indexer also disables update_follow_stats as it kills perf --- Cargo.lock | 1 + consumer/Cargo.toml | 1 + consumer/src/config.rs | 6 +- consumer/src/indexer/mod.rs | 407 ++++++++++++++++++++---------------- consumer/src/main.rs | 2 +- 5 files changed, 236 insertions(+), 181 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d797223e..c4aec5d6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -648,6 +648,7 @@ dependencies = [ "eyre", "figment", "flume", + "foldhash", "futures", "ipld-core", "iroh-car", diff --git a/consumer/Cargo.toml b/consumer/Cargo.toml index 633ccc5e..59534be9 100644 --- a/consumer/Cargo.toml +++ b/consumer/Cargo.toml @@ -12,6 +12,7 @@ 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"] } +foldhash = "0.1.4" futures = "0.3.31" ipld-core = "0.4.1" iroh-car = "0.5.1" diff --git a/consumer/src/config.rs b/consumer/src/config.rs index 199e00d3..13ff721d 100644 --- a/consumer/src/config.rs +++ b/consumer/src/config.rs @@ -21,6 +21,8 @@ pub struct Config { pub ua_contact: Option, #[serde(default = "default_backfill_workers")] pub backfill_workers: u8, + #[serde(default = "default_indexer_workers")] + pub indexer_workers: u8, } #[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Deserialize)] @@ -32,4 +34,6 @@ pub enum HistoryMode { Realtime, } -fn default_backfill_workers() -> u8 { 4 } \ No newline at end of file +fn default_backfill_workers() -> u8 { 4 } + +fn default_indexer_workers() -> u8 { 4 } diff --git a/consumer/src/indexer/mod.rs b/consumer/src/indexer/mod.rs index 0c201c41..0f77a8dc 100644 --- a/consumer/src/indexer/mod.rs +++ b/consumer/src/indexer/mod.rs @@ -2,13 +2,15 @@ use crate::config::HistoryMode; use crate::firehose::{AtpAccountEvent, AtpCommitEvent, AtpIdentityEvent, CommitOp, FirehoseEvent}; use crate::indexer::types::{BackfillItem, BackfillItemInner, CollectionType, RecordTypes}; use did_resolver::Resolver; -use diesel_async::pooled_connection::deadpool::{Object, Pool}; +use diesel_async::pooled_connection::deadpool::Pool; use diesel_async::{AsyncConnection, AsyncPgConnection}; +use foldhash::quality::RandomState; use futures::StreamExt; use ipld_core::cid::Cid; use metrics::counter; use parakeet_db::types::{ActorStatus, ActorSyncState}; use std::collections::HashMap; +use std::hash::BuildHasher; use std::sync::Arc; use tokio::sync::mpsc::{channel, Receiver, Sender}; use tracing::instrument; @@ -17,14 +19,20 @@ pub mod db; pub mod records; pub mod types; -pub struct RelayIndexer { - conn: Object, - rx: Receiver, +#[derive(Clone)] +struct RelayIndexerState { backfill_tx: flume::Sender, resolver: Arc, do_backfill: bool, } +pub struct RelayIndexer { + pool: Pool, + state: RelayIndexerState, + rx: Receiver, + hasher: RandomState, +} + impl RelayIndexer { pub async fn new( pool: Pool, @@ -32,213 +40,251 @@ impl RelayIndexer { resolver: Arc, history_mode: HistoryMode, ) -> eyre::Result<(Self, Sender)> { - let conn = pool.get().await?; - let (tx, rx) = channel(128); let indexer = RelayIndexer { - conn, + pool, + state: RelayIndexerState { + backfill_tx, + resolver, + do_backfill: history_mode == HistoryMode::BackfillHistory, + }, rx, - backfill_tx, - resolver, - do_backfill: history_mode == HistoryMode::BackfillHistory, + hasher: RandomState::default(), }; Ok((indexer, tx)) } - pub async fn run(mut self) -> eyre::Result<()> { + pub async fn run(mut self, threads: u8) -> eyre::Result<()> { + let (submit, _handles) = (0..threads) + .map(|idx| { + let pool = self.pool.clone(); + let state = self.state.clone(); + let (tx, mut rx) = channel(16); + + let handle = tokio::spawn(async move { + let mut conn = pool.get().await.unwrap(); + + while let Some(event) = rx.recv().await { + counter!("ri_events_seen", "thread" => idx.to_string()).increment(1); + let res = match event { + FirehoseEvent::Identity(identity) => { + index_identity(&state, &mut conn, identity).await + } + FirehoseEvent::Account(account) => { + index_account(&state, &mut conn, account).await + } + FirehoseEvent::Commit(commit) => { + index_commit(&state, &mut conn, commit).await + } + FirehoseEvent::Label(_) => unreachable!(), + }; + + if let Err(e) = res { + tracing::error!("Indexing error: {e}"); + } + } + }); + + (tx, handle) + }) + .unzip::<_, _, Vec<_>, Vec<_>>(); + + let threads = threads as u64; while let Some(event) = self.rx.recv().await { - let res = match event { - FirehoseEvent::Identity(identity) => self.index_identity(identity).await, - FirehoseEvent::Account(account) => self.index_account(account).await, - FirehoseEvent::Commit(commit) => self.index_commit(commit).await, + let tgt_idx = match &event { + FirehoseEvent::Identity(identity) => self.hasher.hash_one(&identity.did) % threads, + FirehoseEvent::Account(account) => self.hasher.hash_one(&account.did) % threads, + FirehoseEvent::Commit(commit) => self.hasher.hash_one(&commit.repo) % threads, FirehoseEvent::Label(_) => { // We handle all labels through direct connections to labelers tracing::warn!("got #labels from the relay"); - Ok(()) + continue; } }; - if let Err(e) = res { - tracing::error!("Indexing error: {e}"); + if let Err(e) = submit[tgt_idx as usize].send(event).await { + tracing::error!("Error sending event: {e}"); } } Ok(()) } +} - #[instrument(skip_all, fields(seq = identity.seq, repo = identity.did))] - async fn index_identity(&mut self, identity: AtpIdentityEvent) -> eyre::Result<()> { - let Some(did_doc) = self.resolver.resolve_did(&identity.did).await? else { - eyre::bail!("missing did doc"); - }; +#[instrument(skip_all, fields(seq = identity.seq, repo = identity.did))] +async fn index_identity( + state: &RelayIndexerState, + conn: &mut AsyncPgConnection, + identity: AtpIdentityEvent, +) -> eyre::Result<()> { + let Some(did_doc) = state.resolver.resolve_did(&identity.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(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"); + }; - // 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}"); - }; + // 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) = state.resolver.resolve_handle(&handle).await? else { + eyre::bail!("Failed to resolve did for handle {handle}"); + }; - let new_handle = if handle_did != identity.did { - None - } else { - Some(identity.handle) - }; + let new_handle = if handle_did != identity.did { + None + } else { + Some(identity.handle) + }; - let sync_state = (!self.do_backfill).then_some(ActorSyncState::Synced); + let sync_state = (!state.do_backfill).then_some(ActorSyncState::Synced); - db::upsert_actor( - &mut self.conn, - &identity.did, - new_handle, - None, - sync_state, - identity.time, - ) - .await?; + db::upsert_actor( + conn, + &identity.did, + new_handle, + None, + sync_state, + identity.time, + ) + .await?; - Ok(()) - } + Ok(()) +} - #[instrument(skip_all, fields(seq = account.seq, repo = account.did))] - async fn index_account(&mut self, account: AtpAccountEvent) -> eyre::Result<()> { - let status = account - .status - .map(ActorStatus::from) - .unwrap_or(ActorStatus::Active); - - let sync_state = (!self.do_backfill).then_some(ActorSyncState::Synced); - - db::upsert_actor( - &mut self.conn, - &account.did, - None, - Some(status), - sync_state, - account.time, - ) - .await?; +#[instrument(skip_all, fields(seq = account.seq, repo = account.did))] +async fn index_account( + state: &RelayIndexerState, + conn: &mut AsyncPgConnection, + account: AtpAccountEvent, +) -> eyre::Result<()> { + let status = account + .status + .map(ActorStatus::from) + .unwrap_or(ActorStatus::Active); + + let sync_state = (!state.do_backfill).then_some(ActorSyncState::Synced); + + db::upsert_actor( + conn, + &account.did, + None, + Some(status), + sync_state, + account.time, + ) + .await?; - Ok(()) - } + Ok(()) +} - #[instrument(skip_all, fields(seq = commit.seq, repo = commit.repo, rev = commit.rev))] - async fn index_commit(&mut self, commit: AtpCommitEvent) -> eyre::Result<()> { - // what's the backfill status of this account? this respects locks held by the backfiller. - // we should drop events for 'dirty' and queue 'processing' - // if we can't find the actor, we should add them and then drop the event as with 'dirty' - let is_active = match db::account_sync_status(&mut self.conn, &commit.repo).await? { - Some(ActorSyncState::Synced) => true, - Some(ActorSyncState::Processing) => false, - Some(ActorSyncState::Dirty) => { - // in backfill mode, a new actor can be created by #account / #identity as 'dirty'. - // Completely new repos will always have since == None, so we don't need to trigger - // backfill for them and they can be marked active and indexed normally. - // 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_async(commit.repo.clone()).await?; - } - return Ok(()); +#[instrument(skip_all, fields(seq = commit.seq, repo = commit.repo, rev = commit.rev))] +async fn index_commit( + state: &RelayIndexerState, + conn: &mut AsyncPgConnection, + commit: AtpCommitEvent, +) -> eyre::Result<()> { + // what's the backfill status of this account? this respects locks held by the backfiller. + // we should drop events for 'dirty' and queue 'processing' + // if we can't find the actor, we should add them and then drop the event as with 'dirty' + let is_active = match db::account_sync_status(conn, &commit.repo).await? { + Some(ActorSyncState::Synced) => true, + Some(ActorSyncState::Processing) => false, + Some(ActorSyncState::Dirty) => { + // in backfill mode, a new actor can be created by #account / #identity as 'dirty'. + // Completely new repos will always have since == None, so we don't need to trigger + // backfill for them and they can be marked active and indexed normally. + // 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?; } - - // this is the first commit in an actor's repo - set them to Synced. - db::upsert_actor( - &mut self.conn, - &commit.repo, - None, - None, - Some(ActorSyncState::Synced), - commit.time, - ) - .await?; - - true + return Ok(()); } - None => { - tracing::debug!("found new repo from commit"); - let trigger_backfill = self.do_backfill && commit.since.is_some(); - - let sync_state = (!trigger_backfill).then_some(ActorSyncState::Synced); - - db::upsert_actor( - &mut self.conn, - &commit.repo, - None, - None, - sync_state, - commit.time, - ) - .await?; - - if trigger_backfill { - self.backfill_tx.send_async(commit.repo.clone()).await?; - return Ok(()); - } - true + // this is the first commit in an actor's repo - set them to Synced. + db::upsert_actor( + conn, + &commit.repo, + None, + None, + Some(ActorSyncState::Synced), + commit.time, + ) + .await?; + + true + } + None => { + tracing::debug!("found new repo from commit"); + let trigger_backfill = state.do_backfill && commit.since.is_some(); + + let sync_state = (!trigger_backfill).then_some(ActorSyncState::Synced); + + 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?; + return Ok(()); } - }; - // turn the car slice into a map of cid:block - let car_reader = iroh_car::CarReader::new(commit.blocks.as_slice()).await?; - let blocks = car_reader - .stream() - .filter_map(|car| async move { car.ok() }) - .collect::>>() - .await; - - self.conn - .transaction::<_, diesel::result::Error, _>(|t| { - Box::pin(async move { - if is_active { - // ensure new repo rev > current rev, but only when the repo is synced - // (if repo is processing, they'll end up on the queue anyway) - let updated = - db::update_repo_version(t, &commit.repo, &commit.rev, commit.commit) - .await?; - - if !updated { - tracing::debug!( - "Got a repo update older than the current rev. not processing ops." - ); - return Ok(false); - } + true + } + }; - for op in &commit.ops { - process_op(t, &commit.repo, op, &blocks).await?; - } - } else { - let items = commit - .ops - .iter() - .filter_map(|op| process_op_bf(&commit.repo, op, &blocks)) - .collect::>(); + // turn the car slice into a map of cid:block + let car_reader = iroh_car::CarReader::new(commit.blocks.as_slice()).await?; + let blocks = car_reader + .stream() + .filter_map(|car| async move { car.ok() }) + .collect::>>() + .await; + + conn.transaction::<_, diesel::result::Error, _>(|t| { + Box::pin(async move { + if is_active { + // ensure new repo rev > current rev, but only when the repo is synced + // (if repo is processing, they'll end up on the queue anyway) + let updated = + db::update_repo_version(t, &commit.repo, &commit.rev, commit.commit).await?; + + if !updated { + tracing::debug!( + "Got a repo update older than the current rev. not processing ops." + ); + return Ok(false); + } - let items = serde_json::to_value(items).unwrap_or_default(); + for op in &commit.ops { + process_op(t, &commit.repo, op, &blocks).await?; + } + } else { + let items = commit + .ops + .iter() + .filter_map(|op| process_op_bf(&commit.repo, op, &blocks)) + .collect::>(); - db::write_backfill_row(t, &commit.repo, &commit.rev, commit.commit, items) - .await?; - } + let items = serde_json::to_value(items).unwrap_or_default(); - Ok(true) - }) - }) - .await?; + db::write_backfill_row(t, &commit.repo, &commit.rev, commit.commit, items).await?; + } - Ok(()) - } + Ok(true) + }) + }) + .await?; + + Ok(()) } +#[inline(always)] fn process_op_bf( repo: &str, op: &CommitOp, @@ -284,6 +330,7 @@ fn process_op_bf( }) } +#[inline(always)] async fn process_op( conn: &mut AsyncPgConnection, repo: &str, @@ -375,13 +422,20 @@ pub async fn index_op( return Ok(()); } - let has_disable_rule = record.embedding_rules.contains(&records::PostgateEmbeddingRules::Disable); + let has_disable_rule = record + .embedding_rules + .contains(&records::PostgateEmbeddingRules::Disable); let disable_effective = has_disable_rule.then_some(record.created_at.naive_utc()); db::upsert_postgate(conn, at_uri, cid, &record).await?; - db::postgate_maintain_detaches(conn, &record.post, &record.detached_embedding_uris, disable_effective) - .await?; + db::postgate_maintain_detaches( + conn, + &record.post, + &record.detached_embedding_uris, + disable_effective, + ) + .await?; } RecordTypes::AppBskyFeedRepost(record) => { db::insert_repost(conn, repo, at_uri, record).await?; @@ -399,9 +453,7 @@ pub async fn index_op( db::insert_block(conn, repo, at_uri, record).await?; } RecordTypes::AppBskyGraphFollow(record) => { - let subj = record.subject.clone(); db::insert_follow(conn, repo, at_uri, record).await?; - db::update_follow_stats(conn, &[repo, &subj]).await?; } RecordTypes::AppBskyGraphList(record) => { db::upsert_list(conn, repo, at_uri, cid, record).await?; @@ -422,7 +474,7 @@ pub async fn index_op( } RecordTypes::AppBskyGraphStarterPack(record) => { db::upsert_starterpack(conn, repo, cid, at_uri, record).await?; - }, + } RecordTypes::ChatBskyActorDeclaration(record) => { if rkey == "self" { db::upsert_chat_decl(conn, repo, record).await?; @@ -449,11 +501,8 @@ pub async fn index_op_delete( CollectionType::BskyFeedRepost => db::delete_repost(conn, at_uri).await?, CollectionType::BskyFeedThreadgate => db::delete_threadgate(conn, at_uri).await?, CollectionType::BskyFollow => { - if let Some(subject) = db::delete_follow(conn, at_uri).await? { - db::update_follow_stats(conn, &[repo, &subject]).await? - } else { - 0 - } + db::delete_follow(conn, at_uri).await?; + 0 } CollectionType::BskyList => { db::delete_list(conn, at_uri).await? diff --git a/consumer/src/main.rs b/consumer/src/main.rs index 637d60d6..be7ff58f 100644 --- a/consumer/src/main.rs +++ b/consumer/src/main.rs @@ -44,7 +44,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(relay_indexer.run(conf.indexer_workers)), tokio::spawn(backfiller.run(conf.backfill_workers)) }?; -- 2.51.2