From fe41ba1fc39233f504e19a726857705d1a6c84dc Mon Sep 17 00:00:00 2001 From: dawn <90008@gaze.systems> Date: Tue, 5 May 2026 06:30:03 +0300 Subject: [PATCH] [ingest] change the unbounded channels to be bounded --- src/ingest/indexer.rs | 76 ++++++++++++++++++++----------- src/ingest/relay.rs | 102 +++++++++++++++++++++++++----------------- 2 files changed, 111 insertions(+), 67 deletions(-) diff --git a/src/ingest/indexer.rs b/src/ingest/indexer.rs index c8b6ed0..12d4d3c 100644 --- a/src/ingest/indexer.rs +++ b/src/ingest/indexer.rs @@ -141,16 +141,13 @@ impl FirehoseWorker { } } - // starts the worker threads and the main dispatch loop - // the dispatch loop reads from the firehose channel and - // distributes messages to shards based on the hash of the DID - pub fn run(mut self, handle: TokioHandle) -> Result<()> { - let mut shards = Vec::with_capacity(self.num_shards); + pub fn run(self, handle: TokioHandle) -> Result<()> { + use futures::{StreamExt, future::BoxFuture, stream::FuturesUnordered}; + + let mut shards: Vec> = Vec::with_capacity(self.num_shards); for i in 0..self.num_shards { - // unbounded here so we dont block other shards potentially - // if one has a small lag or something - let (tx, rx) = mpsc::unbounded_channel(); + let (tx, rx) = mpsc::channel(64); shards.push(tx); let state = self.state.clone(); @@ -165,24 +162,51 @@ impl FirehoseWorker { info!(num = self.num_shards, "started shards"); - while let Some(msg) = self.rx.blocking_recv() { - 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, - }; - - let shard_idx = (util::hash(did) as usize) % self.num_shards; - if let Err(e) = shards[shard_idx].send(msg) { - error!(shard = shard_idx, err = %e, "failed to send message to shard, shard panicked?"); - break; + 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?" @@ -192,7 +216,7 @@ impl FirehoseWorker { #[inline(always)] fn shard( id: usize, - mut rx: mpsc::UnboundedReceiver, + mut rx: mpsc::Receiver, state: Arc, handle: TokioHandle, ) { diff --git a/src/ingest/relay.rs b/src/ingest/relay.rs index 488577c..9708f8e 100644 --- a/src/ingest/relay.rs +++ b/src/ingest/relay.rs @@ -90,11 +90,13 @@ impl RelayWorker { } } - pub fn run(mut self, handle: Handle) -> Result<()> { - let mut shards = Vec::with_capacity(self.num_shards); + pub fn run(self, handle: Handle) -> Result<()> { + use futures::{StreamExt, future::BoxFuture, stream::FuturesUnordered}; + + let mut shards: Vec> = Vec::with_capacity(self.num_shards); for i in 0..self.num_shards { - let (tx, rx) = mpsc::unbounded_channel(); + let (tx, rx) = mpsc::channel(64); shards.push(tx); let state = self.state.clone(); @@ -125,54 +127,72 @@ impl RelayWorker { info!(num = self.num_shards, "relay worker: started shards"); - let _g = handle.enter(); - - while let Some(msg) = self.rx.blocking_recv() { - let IngestMessage::Firehose { url, is_pds, msg } = msg; - - // #info only pertains to us, the direct consumer - if let SubscribeReposMessage::Info(inf) = msg { - match inf.name { - InfoName::OutdatedCursor => { - // todo: handle + 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_else(|| 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; + } + } } - InfoName::Other(name) => { - let message = inf - .message - .unwrap_or_else(|| CowStr::Borrowed("")); - info!(name = %name, "relay sent info: {message}"); + 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; + } } } - 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) % self.num_shards - }; - - if let Err(e) = shards[shard_idx].send(WorkerMessage { - firehose: url, - is_pds, - msg, - }) { - error!(shard = shard_idx, err = %e, "relay worker: failed to send to shard"); - break; } - } + }); Err(miette::miette!("relay worker dispatcher shutting down")) } fn shard( id: usize, - mut rx: mpsc::UnboundedReceiver, + mut rx: mpsc::Receiver, state: Arc, #[cfg(feature = "indexer")] hook: crate::ingest::indexer::IndexerTx, verify_signatures: bool, -- 2.51.2