diff --git a/src/control/mod.rs b/src/control/mod.rs --- a/src/control/mod.rs +++ b/src/control/mod.rs @@ -283,13 +283,23 @@ } let fut = async move { - // raw firehose events from pds/relay to RelayWorker - let (buffer_tx, buffer_rx) = mpsc::channel::(500); - - // validated IndexerMessages from RelayWorker/backfill to FirehoseWorker + // validated IndexerMessages from RelayWorker/backfill to FirehoseWorker. #[cfg(feature = "indexer")] - let (indexer_tx, indexer_rx) = - mpsc::channel::(500); + let (indexer_tx, firehose_worker) = + FirehoseWorker::new(state.clone(), config.firehose_workers); + + // raw firehose events from pds/relay to RelayWorker. + let (buffer_tx, relay_worker) = crate::ingest::relay::RelayWorker::new( + state.clone(), + #[cfg(feature = "indexer")] + indexer_tx.clone(), + matches!(config.verify_signatures, SignatureVerification::Full), + config.firehose_workers, + crate::ingest::validation::ValidationOptions { + verify_mst: config.verify_mst, + rev_clock_skew_secs: config.rev_clock_skew_secs, + }, + ); // 5. spawn the backfill worker (not used in relay mode) #[cfg(feature = "indexer")] @@ -673,28 +683,8 @@ // 12. spawn the relay worker let relay_worker = std::thread::spawn({ - let state = state.clone(); let handle = tokio::runtime::Handle::current(); - let config = config.clone(); - - #[cfg(feature = "indexer")] - let hook = indexer_tx.clone(); - - move || { - crate::ingest::relay::RelayWorker::new( - state, - buffer_rx, - #[cfg(feature = "indexer")] - hook, - matches!(config.verify_signatures, SignatureVerification::Full), - config.firehose_workers, - crate::ingest::validation::ValidationOptions { - verify_mst: config.verify_mst, - rev_clock_skew_secs: config.rev_clock_skew_secs, - }, - ) - .run(handle) - } + move || relay_worker.run(handle) }); let tx = Arc::clone(&fatal_tx); @@ -713,10 +703,8 @@ // 13. spawn the firehose worker (if enabled) #[cfg(feature = "indexer")] let firehose_worker = std::thread::spawn({ - let state = state.clone(); let handle = tokio::runtime::Handle::current(); - let config = config.clone(); - move || FirehoseWorker::new(state, indexer_rx, config.firehose_workers).run(handle) + move || firehose_worker.run(handle) }); #[cfg(feature = "indexer")] diff --git a/src/ingest/indexer.rs b/src/ingest/indexer.rs --- a/src/ingest/indexer.rs +++ b/src/ingest/indexer.rs @@ -1,5 +1,6 @@ use super::*; use crate::db::{self, CountDeltas, keys, ser_repo_meta}; +use crate::ingest::mailbox::{ShardedMessage, ShardedReceiver, ShardedSender}; use crate::ingest::stream::{Account, Commit, Identity}; use crate::ingest::validation; use crate::resolver::{NoSigningKeyError, ResolverError}; @@ -18,7 +19,6 @@ use std::sync::atomic::Ordering::SeqCst; use thiserror::Error; use tokio::runtime::Handle as TokioHandle; -use tokio::sync::mpsc; use tracing::{debug, error, info, warn}; #[cfg(feature = "indexer_stream")] use { @@ -80,8 +80,46 @@ BackfillFinished(Did<'static>), } -pub type IndexerTx = mpsc::Sender; -pub type IndexerRx = mpsc::Receiver; +#[derive(Clone, Debug)] +pub struct IndexerTx { + inner: ShardedSender, +} + +pub type IndexerRx = ShardedReceiver; + +impl IndexerTx { + pub async fn send( + &self, + msg: IndexerMessage, + ) -> Result<(), tokio::sync::mpsc::error::SendError> { + self.inner.send(msg).await + } + + pub fn blocking_send( + &self, + msg: IndexerMessage, + ) -> Result<(), tokio::sync::mpsc::error::SendError> { + self.inner.blocking_send(msg) + } +} + +impl ShardedMessage for IndexerMessage { + fn shard_idx(&self, num_shards: usize) -> usize { + // keep commit, new-repo, and backfill-finished messages for a did on one ordered path. + let did = match self { + IndexerMessage::Event(e) => match &e.data { + IndexerEventData::Commit(m) => &m.commit.repo, + IndexerEventData::Identity(m) => &m.identity.did, + IndexerEventData::Account(m) => &m.account.did, + IndexerEventData::Sync(did) => did, + }, + IndexerMessage::NewRepo(did) => did, + IndexerMessage::BackfillFinished(did) => did, + }; + + (util::hash(did) as usize) % num_shards + } +} #[derive(Debug, Diagnostic, Error)] enum IngestError { @@ -120,8 +158,7 @@ pub struct FirehoseWorker { state: Arc, - rx: IndexerRx, - num_shards: usize, + rxs: Vec, } struct WorkerContext<'a> { @@ -135,93 +172,42 @@ } impl FirehoseWorker { - pub fn new(state: Arc, rx: IndexerRx, num_shards: usize) -> Self { - Self { - state, - rx, - num_shards, - } + pub fn new(state: Arc, num_shards: usize) -> (IndexerTx, Self) { + let (inner, rxs) = ShardedSender::channel(num_shards); + (IndexerTx { inner }, Self { state, rxs }) } pub fn run(self, handle: TokioHandle) -> Result<()> { - use futures::{StreamExt, future::BoxFuture, stream::FuturesUnordered}; + let num_shards = self.rxs.len(); + let (exit_tx, exit_rx) = std::sync::mpsc::channel(); - let mut shards: Vec> = Vec::with_capacity(self.num_shards); - - for i in 0..self.num_shards { - let (tx, rx) = mpsc::channel(64); - shards.push(tx); - - let state = self.state.clone(); + for (i, rx) in self.rxs.into_iter().enumerate() { + let state = Arc::clone(&self.state); let handle = handle.clone(); + let exit_tx = exit_tx.clone(); std::thread::Builder::new() .name(format!("ingest-shard-{i}")) .spawn(move || { - Self::shard(i, rx, state, handle); + let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + Self::shard(i, rx, state, handle); + })); + let _ = exit_tx.send((i, res)); }) .into_diagnostic()?; } + drop(exit_tx); - info!(num = self.num_shards, "started shards"); + info!(num = num_shards, "started shards"); - let num_shards = self.num_shards; - let mut rx = self.rx; - - handle.block_on(async move { - let mut pending: FuturesUnordered< - BoxFuture<'_, Result<(), mpsc::error::SendError>>, - > = FuturesUnordered::new(); - - loop { - tokio::select! { - msg = rx.recv(), if pending.len() < num_shards => { - let Some(msg) = msg else { break; }; - let shard_idx = { - let did = match &msg { - IndexerMessage::Event(e) => match &e.data { - IndexerEventData::Commit(m) => &m.commit.repo, - IndexerEventData::Identity(m) => &m.identity.did, - IndexerEventData::Account(m) => &m.account.did, - IndexerEventData::Sync(did) => did, - }, - IndexerMessage::NewRepo(did) => did, - IndexerMessage::BackfillFinished(did) => did, - }; - (util::hash(did) as usize) % num_shards - }; - match shards[shard_idx].try_send(msg) { - Ok(()) => {} - Err(mpsc::error::TrySendError::Full(msg)) => { - pending.push(Box::pin(shards[shard_idx].send(msg))); - } - Err(mpsc::error::TrySendError::Closed(_)) => { - error!(shard = shard_idx, "shard closed unexpectedly"); - break; - } - } - } - Some(result) = pending.next(), if !pending.is_empty() => { - if let Err(e) = result { - error!(err = %e, "failed to send to shard, shard panicked?"); - break; - } - } - } - } - }); - - Err(miette::miette!( - "firehose worker dispatcher shutting down, shard died?" - )) + match exit_rx.recv() { + Ok((id, Ok(()))) => Err(miette::miette!("firehose worker shard {id} shut down")), + Ok((id, Err(_))) => Err(miette::miette!("firehose worker shard {id} panicked")), + Err(_) => Err(miette::miette!("firehose worker shards shut down")), + } } #[inline(always)] - fn shard( - id: usize, - mut rx: mpsc::Receiver, - state: Arc, - handle: TokioHandle, - ) { + fn shard(id: usize, mut rx: IndexerRx, state: Arc, handle: TokioHandle) { let _guard = handle.enter(); debug!(shard = id, "shard started"); diff --git a/src/ingest/mailbox.rs b/src/ingest/mailbox.rs new file mode 100644 --- /dev/null +++ b/src/ingest/mailbox.rs @@ -0,0 +1,56 @@ +use std::sync::Arc; + +use tokio::sync::mpsc; + +const SHARD_CHANNEL_CAPACITY: usize = 64; + +pub(crate) trait ShardedMessage { + fn shard_idx(&self, num_shards: usize) -> usize; +} + +#[derive(Debug)] +pub(crate) struct ShardedSender { + shards: Arc<[mpsc::Sender]>, +} + +pub(crate) type ShardedReceiver = mpsc::Receiver; + +impl Clone for ShardedSender { + fn clone(&self) -> Self { + Self { + shards: Arc::clone(&self.shards), + } + } +} + +impl ShardedSender { + pub(crate) fn channel(num_shards: usize) -> (Self, Vec>) { + assert!(num_shards > 0, "num_shards must be greater than zero"); + + let mut txs = Vec::with_capacity(num_shards); + let mut rxs = Vec::with_capacity(num_shards); + for _ in 0..num_shards { + let (tx, rx) = mpsc::channel(SHARD_CHANNEL_CAPACITY); + txs.push(tx); + rxs.push(rx); + } + + ( + Self { + shards: Arc::from(txs.into_boxed_slice()), + }, + rxs, + ) + } + + pub(crate) async fn send(&self, msg: T) -> Result<(), mpsc::error::SendError> { + let shard_idx = msg.shard_idx(self.shards.len()); + self.shards[shard_idx].send(msg).await + } + + #[cfg(feature = "indexer")] + pub(crate) fn blocking_send(&self, msg: T) -> Result<(), mpsc::error::SendError> { + let shard_idx = msg.shard_idx(self.shards.len()); + self.shards[shard_idx].blocking_send(msg) + } +} diff --git a/src/ingest/mod.rs b/src/ingest/mod.rs --- a/src/ingest/mod.rs +++ b/src/ingest/mod.rs @@ -1,13 +1,13 @@ -use tokio::sync::mpsc; - pub mod firehose; #[cfg(feature = "indexer")] pub mod indexer; +mod mailbox; pub mod relay; pub mod stream; pub mod validation; use url::Url; +use crate::ingest::mailbox::{ShardedMessage, ShardedReceiver, ShardedSender}; use crate::ingest::stream::SubscribeReposMessage; #[derive(Debug)] @@ -21,5 +21,34 @@ }, } -pub type BufferTx = mpsc::Sender; -pub type BufferRx = mpsc::Receiver; +#[derive(Clone, Debug)] +pub struct BufferTx { + inner: ShardedSender, +} + +pub type BufferRx = ShardedReceiver; + +impl BufferTx { + pub(crate) fn channel(num_shards: usize) -> (Self, Vec) { + let (inner, rxs) = ShardedSender::channel(num_shards); + (Self { inner }, rxs) + } + + pub async fn send( + &self, + msg: IngestMessage, + ) -> Result<(), tokio::sync::mpsc::error::SendError> { + self.inner.send(msg).await + } +} + +impl ShardedMessage for IngestMessage { + fn shard_idx(&self, num_shards: usize) -> usize { + let did = match self { + IngestMessage::Firehose { msg, .. } => msg.did(), + }; + + did.map(|did| (crate::util::hash(did) as usize) % num_shards) + .unwrap_or(0) + } +} diff --git a/src/ingest/relay.rs b/src/ingest/relay.rs --- a/src/ingest/relay.rs +++ b/src/ingest/relay.rs @@ -14,7 +14,6 @@ use jacquard_common::{CowStr, IntoStatic}; use miette::{IntoDiagnostic, Result}; use tokio::runtime::Handle; -use tokio::sync::mpsc; use tracing::{debug, error, info, info_span, trace, warn}; use url::Url; @@ -28,7 +27,7 @@ CommitValidationError, SyncValidationError, ValidatedCommit, ValidatedSync, ValidationContext, ValidationOptions, }; -use crate::ingest::{BufferRx, IngestMessage}; +use crate::ingest::{BufferRx, BufferTx, IngestMessage}; use crate::state::AppState; #[cfg(feature = "relay")] use crate::types::RelayBroadcast; @@ -60,7 +59,7 @@ pub struct RelayWorker { state: Arc, - rx: BufferRx, + rxs: Vec, #[cfg(feature = "indexer")] hook: crate::ingest::indexer::IndexerTx, verify_signatures: bool, @@ -72,128 +71,75 @@ impl RelayWorker { pub fn new( state: Arc, - rx: BufferRx, #[cfg(feature = "indexer")] hook: crate::ingest::indexer::IndexerTx, verify_signatures: bool, num_shards: usize, validation_opts: ValidationOptions, - ) -> Self { - Self { - state, - rx, - #[cfg(feature = "indexer")] - hook, - verify_signatures, - num_shards, - validation_opts: Arc::new(validation_opts), - http: reqwest::Client::new(), - } + ) -> (BufferTx, Self) { + let (tx, rxs) = BufferTx::channel(num_shards); + ( + tx, + Self { + state, + rxs, + #[cfg(feature = "indexer")] + hook, + verify_signatures, + num_shards, + validation_opts: Arc::new(validation_opts), + http: reqwest::Client::new(), + }, + ) } pub fn run(self, handle: Handle) -> Result<()> { - use futures::{StreamExt, future::BoxFuture, stream::FuturesUnordered}; + let (exit_tx, exit_rx) = std::sync::mpsc::channel(); - let mut shards: Vec> = Vec::with_capacity(self.num_shards); - - for i in 0..self.num_shards { - let (tx, rx) = mpsc::channel(64); - shards.push(tx); - - let state = self.state.clone(); + for (i, rx) in self.rxs.into_iter().enumerate() { + let state = Arc::clone(&self.state); #[cfg(feature = "indexer")] let hook = self.hook.clone(); let verify = self.verify_signatures; let h = handle.clone(); let opts = self.validation_opts.clone(); let http = self.http.clone(); + let exit_tx = exit_tx.clone(); std::thread::Builder::new() .name(format!("relay-shard-{i}")) .spawn(move || { - Self::shard( - i, - rx, - state, - #[cfg(feature = "indexer")] - hook, - verify, - h, - opts, - http, - ); + let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + Self::shard( + i, + rx, + state, + #[cfg(feature = "indexer")] + hook, + verify, + h, + opts, + http, + ); + })); + let _ = exit_tx.send((i, res)); }) .into_diagnostic()?; } + drop(exit_tx); info!(num = self.num_shards, "relay worker: started shards"); - let num_shards = self.num_shards; - let mut rx = self.rx; - - handle.block_on(async move { - let mut pending: FuturesUnordered< - BoxFuture<'_, Result<(), mpsc::error::SendError>>, - > = FuturesUnordered::new(); - - loop { - tokio::select! { - msg = rx.recv(), if pending.len() < num_shards => { - let Some(msg) = msg else { break; }; - let IngestMessage::Firehose { url, is_pds, msg } = msg; - - if let SubscribeReposMessage::Info(inf) = msg { - match inf.name { - InfoName::OutdatedCursor => {} - InfoName::Other(name) => { - let message = inf - .message - .unwrap_or(CowStr::Borrowed("")); - info!(name = %name, "relay sent info: {message}"); - } - } - continue; - } - - let shard_idx = { - let did = match &msg { - SubscribeReposMessage::Commit(c) => &c.repo, - SubscribeReposMessage::Identity(i) => &i.did, - SubscribeReposMessage::Account(a) => &a.did, - SubscribeReposMessage::Sync(s) => &s.did, - _ => continue, - }; - (util::hash(did) as usize) % num_shards - }; - - let worker_msg = WorkerMessage { firehose: url, is_pds, msg }; - match shards[shard_idx].try_send(worker_msg) { - Ok(()) => {} - Err(mpsc::error::TrySendError::Full(worker_msg)) => { - pending.push(Box::pin(shards[shard_idx].send(worker_msg))); - } - Err(mpsc::error::TrySendError::Closed(_)) => { - error!(shard = shard_idx, "relay shard closed unexpectedly"); - break; - } - } - } - Some(result) = pending.next(), if !pending.is_empty() => { - if let Err(e) = result { - error!(err = %e, "relay worker: failed to send to shard, shard panicked?"); - break; - } - } - } - } - }); - - Err(miette::miette!("relay worker dispatcher shutting down")) + match exit_rx.recv() { + Ok((id, Ok(()))) => Err(miette::miette!("relay worker shard {id} shut down")), + Ok((id, Err(_))) => Err(miette::miette!("relay worker shard {id} panicked")), + Err(_) => Err(miette::miette!("relay worker shards shut down")), + } } #[allow(clippy::too_many_arguments)] fn shard( id: usize, - mut rx: mpsc::Receiver, + mut rx: BufferRx, state: Arc, #[cfg(feature = "indexer")] hook: crate::ingest::indexer::IndexerTx, verify_signatures: bool, @@ -225,6 +171,24 @@ }; while let Some(msg) = rx.blocking_recv() { + let IngestMessage::Firehose { url, is_pds, msg } = msg; + if let SubscribeReposMessage::Info(inf) = msg { + match inf.name { + InfoName::OutdatedCursor => {} + InfoName::Other(name) => { + let message = inf.message.unwrap_or(CowStr::Borrowed("")); + info!(name = %name, "relay sent info: {message}"); + } + } + continue; + } + + let msg = WorkerMessage { + is_pds, + firehose: url, + msg, + }; + ctx.count_deltas = CountDeltas::default(); let (did, seq) = match &msg.msg { SubscribeReposMessage::Commit(c) => (c.repo.clone(), c.seq),