fix: disable overflow in head change broadcast channel

Cargo.toml

@@ -37,6 +37,7 @@ anes = "0.2"
 anyhow = { workspace = true }
 arc-swap = "1"
 argon2 = "0.5"
+async-broadcast = "0.7"
 async-compression = { version = "0.4", features = ["tokio", "zstd"] }
 async-fs = "2"
 async-trait = "0.1"

src/chain/store/chain_store.rs

@@ -39,11 +39,11 @@ use serde::{Serialize, de::DeserializeOwned};
 use std::{
     num::NonZeroUsize,
     sync::atomic::{self, AtomicI64},
+    time::Duration,
 };
-use tokio::sync::broadcast;
 
-// A cap on the size of the future_sink
-const SINK_CAP: usize = 200;
+// Capacity of the head change broadcast channel
+const HEAD_CHANGE_BROADCAST_CHANNEL_CAP: usize = 1000;
 
 // Assume a tipset has 5 blocks on average, we cache 1-day-worth of validated blocks. (5 * 2 * 60 * 24 = 14400)
 const VALIDATED_BLOCKS_CACHE_SIZE: NonZeroUsize = nonzero!(14400usize);
@@ -81,8 +81,11 @@ pub type HeadChanges = PathChanges<Tipset>;
 /// epoch. This structure is thread-safe, and all caches are wrapped in a mutex
 /// to allow a consistent `ChainStore` to be shared across tasks.
 pub struct ChainStore {
-    /// Publisher for head change events
-    head_changes_tx: broadcast::Sender<HeadChanges>,
+    /// The non-blocking bridge sender for head change events. This is used to send head changes to the `head_changes_tx` channel.
+    head_changes_tx_bridge: flume::Sender<HeadChanges>,
+
+    /// Inactive receiver for head change events. This is used to keep the channel alive even if there are no active subscribers.
+    head_changes_rx_inactive: Arc<async_broadcast::InactiveReceiver<HeadChanges>>,
 
     /// Heaviest tipset cache
     heaviest_tipset: Arc<ArcSwap<Tipset>>,
@@ -118,7 +121,8 @@ pub struct ChainStore {
 impl ShallowClone for ChainStore {
     fn shallow_clone(&self) -> Self {
         Self {
-            head_changes_tx: self.head_changes_tx.clone(),
+            head_changes_tx_bridge: self.head_changes_tx_bridge.clone(),
+            head_changes_rx_inactive: self.head_changes_rx_inactive.shallow_clone(),
             heaviest_tipset: self.heaviest_tipset.shallow_clone(),
             f3_finalized_tipset: self.f3_finalized_tipset.shallow_clone(),
             ec_calculator_finalized_epoch: self.ec_calculator_finalized_epoch.shallow_clone(),
@@ -142,7 +146,39 @@ impl ChainStore {
         let db = db.into();
         let genesis = genesis.into();
         anyhow::ensure!(genesis.epoch() == 0, "genesis tipset must be at epoch 0");
-        let (publisher, _) = broadcast::channel(SINK_CAP);
+        let (mut head_changes_tx, head_changes_rx) =
+            async_broadcast::broadcast(HEAD_CHANGE_BROADCAST_CHANNEL_CAP);
+        head_changes_tx.set_await_active(false);
+        head_changes_tx.set_overflow(false); // Disable overflow to not drop head changes
+        // Bridge the flume channel to the async_broadcast channel in a background task,
+        // it's unbounded to take the back pressure and not block `set_heaviest_head`
+        // in case `head_changes_tx` is unexpectedly full and blocked.
+        let (head_changes_tx_bridge, head_changes_rx_bridge) = flume::unbounded();
+        // Warn if the broadcast channel is blocked (timed out after 1 second)
+        if tokio::runtime::Handle::try_current().is_ok() {
+            tokio::spawn(async move {
+                // The loop breaks when the flume channel is closed, which happens when the `ChainStore` is dropped.
+                while let Ok(m) = head_changes_rx_bridge.recv_async().await {
+                    const TIMEOUT: Duration = Duration::from_secs(1);
+                    if tokio::time::timeout(TIMEOUT, head_changes_tx.broadcast_direct(m))
+                        .await
+                        .is_err()
+                    {
+                        error!(
+                            "Head change broadcast channel is full. This indicates some consumers are not processing head changes fast enough."
+                        );
+                    }
+                }
+            });
+        } else {
+            cfg_if::cfg_if! {
+                if #[cfg(test)] {
+                    warn!("ChainStore::new() is called outside of a Tokio runtime, head change broadcast channel is not working in this test");
+                } else {
+                    anyhow::bail!("ChainStore::new() must be called from within a Tokio runtime");
+                }
+            }
+        }
         let head = if let Some(head_tsk) = db
             .heaviest_tipset_key()
             .context("failed to load head tipset key")?
@@ -166,7 +202,8 @@ impl ChainStore {
             }
         }));
         Ok(Self {
-            head_changes_tx: publisher,
+            head_changes_tx_bridge,
+            head_changes_rx_inactive: head_changes_rx.deactivate().into(),
             chain_index,
             tipset_tracker: TipsetTracker::new(db, chain_config.clone()),
             heaviest_tipset,
@@ -254,7 +291,8 @@ impl ChainStore {
         }
 
         let old_head = self.heaviest_tipset.swap(head.shallow_clone().into());
-        if crate::utils::broadcast::has_subscribers(&self.head_changes_tx) {
+        // Only publish head changes when there are active subscribers and head is changed.
+        if self.head_changes_rx_inactive.receiver_count() > 0 && old_head.key() != head.key() {
             let changes = match crate::rpc::chain::chain_get_path(self, old_head.key(), head.key())
             {
                 Ok(changes) => changes,
@@ -270,8 +308,17 @@ impl ChainStore {
                     }
                 }
             };
-            if self.head_changes_tx.send(changes).is_err() {
-                debug!("did not publish changes, no active receivers");
+            // Do not publish empty change and check active receivers again
+            if !changes.is_empty()
+                && self.head_changes_rx_inactive.receiver_count() > 0
+                // Use an unbounded bridge channel to avoid blocking `set_heaviest_tipset`.
+                // Consider refactoring `set_heaviest_tipset` to be async and move the timeout logic here instead.
+                // Note: head change is only published after tipset validation and the 30s block delay
+                // should be sufficient for any consumer to catch up. If this blocks, the consumer logic
+                // needs to be fixed, e.g. spawning a non-blocking task the process the head changes.                
+                && self.head_changes_tx_bridge.send(changes).is_err()
+            {
+                debug!("no active receivers");
             }
         }
 
@@ -345,8 +392,8 @@ impl ChainStore {
     }
 
     /// Subscribes head changes.
-    pub fn subscribe_head_changes(&self) -> broadcast::Receiver<HeadChanges> {
-        self.head_changes_tx.subscribe()
+    pub fn subscribe_head_changes(&self) -> async_broadcast::Receiver<HeadChanges> {
+        self.head_changes_rx_inactive.activate_cloned()
     }
 
     /// Returns a borrowed key-value store instance.
@@ -843,9 +890,14 @@ pub fn get_parent_receipt(
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::utils::multihash::prelude::*;
-    use crate::{blocks::RawBlockHeader, shim::address::Address};
+    use crate::{
+        blocks::{Chain4U, RawBlockHeader, chain4u},
+        shim::address::Address,
+        utils::multihash::prelude::*,
+    };
     use fvm_ipld_encoding::DAG_CBOR;
+    use std::time::Duration;
+    use tokio_util::task::AbortOnDropHandle;
 
     #[test]
     fn genesis_test() {
@@ -960,4 +1012,49 @@ mod tests {
         );
         assert!(inserter_executed.load(std::sync::atomic::Ordering::Relaxed));
     }
+
+    #[tokio::test]
+    async fn test_head_changes() {
+        let c4u = Chain4U::new();
+        chain4u! {
+            in c4u;
+            t0 @ [genesis]
+            -> t1 @ [_b1_0]
+            -> t2 @ [_b2_0, _b2_1]
+            -> t3 @ [_b3_0]
+            -> t4 @ [_b4_1]
+        };
+
+        let db = DbImpl::from(Arc::new(crate::db::MemoryDB::default()));
+        let chain_config = Arc::new(ChainConfig::default());
+        let cs = ChainStore::new(db, chain_config, genesis).unwrap();
+        let mut rx = cs.subscribe_head_changes();
+
+        let handle = AbortOnDropHandle::new(tokio::spawn({
+            let tipsets = vec![
+                // This duplicate head should not be published
+                t0.shallow_clone(),
+                t1.shallow_clone(),
+                t2.shallow_clone(),
+                t3.shallow_clone(),
+                t4.shallow_clone(),
+            ];
+            async move {
+                for ts in tipsets {
+                    cs.set_heaviest_tipset(ts).unwrap();
+                    tokio::time::sleep(Duration::from_millis(100)).await;
+                }
+            }
+        }));
+
+        // t0 is set as head in `ChainStore::new`, so the first published change is t1.
+        for ts in [&t1, &t2, &t3, &t4] {
+            let changes = rx.recv().await.unwrap();
+            assert_eq!(changes.applies, vec![ts.shallow_clone()]);
+        }
+
+        rx.try_recv().unwrap_err(); // no more messages
+
+        handle.await.unwrap();
+    }
 }

src/daemon/db_util.rs

@@ -17,8 +17,8 @@ use crate::state_manager::StateManager;
 use crate::utils::db::car_stream::CarStream;
 use crate::utils::io::EitherMmapOrRandomAccessFile;
 use crate::utils::net::{DownloadFileOption, download_to};
-use anyhow::{Context, bail};
-use futures::TryStreamExt;
+use anyhow::bail;
+use futures::TryStreamExt as _;
 use serde::{Deserialize, Serialize};
 use std::sync::LazyLock;
 use std::sync::atomic::{AtomicI64, Ordering};
@@ -29,8 +29,7 @@ use std::{
     time,
 };
 use tokio::io::AsyncWriteExt;
-use tokio::sync::broadcast::error::TryRecvError;
-use tokio_util::sync::CancellationToken;
+use tokio_util::{sync::CancellationToken, task::AbortOnDropHandle};
 use tracing::{debug, info, warn};
 use url::Url;
 use walkdir::WalkDir;
@@ -726,7 +725,18 @@ pub async fn run_backfill(
     let cancel = guard.cancellation_token();
 
     // Subscribe before the walk so applies/reverts that happen during it are observed.
-    let mut head_rx = state_manager.chain_store().subscribe_head_changes();
+    // We should consume the channel ASAP to avoid blocking the sending part or dropping
+    // tipsets when overflow is enabled.
+    let mut head_changes_rx = state_manager.chain_store().subscribe_head_changes();
+    let (tx, rx) = flume::unbounded();
+    let bridge_handle = AbortOnDropHandle::new(tokio::spawn(async move {
+        while let Some(changes) = head_changes_rx.next().await {
+            if let Err(e) = tx.send_async(changes).await {
+                error!("failed to send head changes to backfill: {e}");
+                break;
+            }
+        }
+    }));
 
     // Optionally clamp the start below finality to avoid indexing revert-prone near-head tipsets.
     let start_ts = if options.allow_near_head {
@@ -807,33 +817,22 @@ pub async fn run_backfill(
 
     // Re-index tipsets applied during the walk so the canonical mapping wins.
     if !report.cancelled {
+        // drop the handle and the underlying channels
+        drop(bridge_handle);
         let mut extra: Vec<(SignedMessage, u64)> = vec![];
-        loop {
-            match head_rx.try_recv() {
-                Ok(changes) => {
-                    for ts in changes.applies {
-                        if ts.epoch() >= lowest_epoch && ts.epoch() <= start_ts.epoch() {
-                            tracing::debug!(
-                                "re-indexing tipset @{} applied during backfill",
-                                ts.epoch()
-                            );
-                            if let Err(e) =
-                                process_ts(&ts, state_manager, &mut extra, options.allow_recompute)
-                                    .await
-                            {
-                                tracing::warn!(
-                                    "failed to re-index applied tipset @{}: {e:#}",
-                                    ts.epoch()
-                                );
-                            }
-                        }
+        // Not using `rx.drain` to make sure tx is dropped and the channel is closed.
+        // Otherwise the loop below will block forever and timeout the CI tests.
+        let mut rx_stream = rx.into_stream();
+        while let Some(changes) = rx_stream.next().await {
+            for ts in changes.applies {
+                if ts.epoch() >= lowest_epoch && ts.epoch() <= start_ts.epoch() {
+                    tracing::debug!("re-indexing tipset @{} applied during backfill", ts.epoch());
+                    if let Err(e) =
+                        process_ts(&ts, state_manager, &mut extra, options.allow_recompute).await
+                    {
+                        tracing::warn!("failed to re-index applied tipset @{}: {e:#}", ts.epoch());
                     }
                 }
-                Err(TryRecvError::Empty) | Err(TryRecvError::Closed) => break,
-                Err(TryRecvError::Lagged(n)) => {
-                    tracing::warn!("backfill head-change listener lagged: skipped {n} events");
-                    continue;
-                }
             }
         }
         if !extra.is_empty() {

src/daemon/mod.rs

@@ -750,26 +750,18 @@ fn maybe_start_indexer_service(
         let chain_store = ctx.state_manager.chain_store().shallow_clone();
         services.spawn(async move {
             tracing::info!("Starting indexer service");
-
             // Continuously listen for head changes
-            loop {
-                match head_changes_rx.recv().await {
-                    Ok(changes) => {
-                        for ts in changes.applies {
-                            tracing::debug!("Indexing tipset {}", ts.key());
-                            let delegated_messages = chain_store
-                                .headers_delegated_messages(ts.block_headers().iter())?;
-                            // Head indexing writes the newest tipset, so use the blind-write
-                            // fast path (no read-before-write timestamp comparison).
-                            chain_store.process_signed_messages(&delegated_messages, false)?;
-                        }
-                    }
-                    Err(RecvError::Lagged(n)) => {
-                        warn!("indexer service lagged: skipping {n} events")
-                    }
-                    Err(RecvError::Closed) => break Ok(()),
+            while let Some(changes) = head_changes_rx.next().await {
+                for ts in changes.applies {
+                    tracing::debug!("Indexing tipset {}", ts.key());
+                    let delegated_messages =
+                        chain_store.headers_delegated_messages(ts.block_headers().iter())?;
+                    // Head indexing writes the newest tipset, so use the blind-write
+                    // fast path (no read-before-write timestamp comparison).
+                    chain_store.process_signed_messages(&delegated_messages, false)?;
                 }
             }
+            Ok(())
         });
 
         // Run the collector only if chain indexer is enabled

src/lib.rs

@@ -84,7 +84,7 @@ mod prelude {
     pub use ahash::{HashMapExt as _, HashSetExt as _};
     pub use anyhow::Context as _;
     pub use cid::Cid;
-    pub use futures::FutureExt as _;
+    pub use futures::{FutureExt as _, StreamExt as _};
     pub use itertools::Itertools as _;
     pub use std::{ops::Deref as _, sync::Arc};
     pub use tracing::{debug, error, info, trace, warn};