feat(providers): pause heartbeat when no transactions are pending (#2800)

* feat(providers): pause heartbeat when no transactions are pending

* chore: downgrade log

* wip: tokio

* fix

* log

* Update crates/provider/src/blocks.rs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: DaniPopes

crates/provider/src/blocks.rs

@@ -5,7 +5,14 @@ use alloy_transport::RpcError;
 use async_stream::stream;
 use futures::{Stream, StreamExt};
 use lru::LruCache;
-use std::{marker::PhantomData, num::NonZeroUsize};
+use std::{
+    marker::PhantomData,
+    num::NonZeroUsize,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+};
 
 #[cfg(feature = "pubsub")]
 use futures::{future::Either, FutureExt};
@@ -19,6 +26,38 @@ const MAX_RETRIES: usize = 3;
 /// Default block number for when we don't have a block yet.
 const NO_BLOCK_NUMBER: BlockNumber = BlockNumber::MAX;
 
+#[derive(Default)]
+pub(crate) struct Paused {
+    is_paused: AtomicBool,
+    notify: tokio::sync::Notify,
+}
+
+impl Paused {
+    pub(crate) fn is_paused(&self) -> bool {
+        self.is_paused.load(Ordering::Acquire)
+    }
+
+    pub(crate) fn set_paused(&self, paused: bool) {
+        self.is_paused.store(paused, Ordering::Release);
+        if !paused {
+            self.notify.notify_waiters();
+        }
+    }
+
+    /// Waits until the paused state is changed to `false`.
+    ///
+    /// Returns `true` if the method actually waited for the paused state to become unpaused,
+    /// or `false` if it was already unpaused when called.
+    async fn wait(&self) -> bool {
+        if !self.is_paused() {
+            return false;
+        }
+        self.notify.notified().await;
+        debug_assert!(!self.is_paused());
+        true
+    }
+}
+
 /// Streams new blocks from the client.
 pub(crate) struct NewBlocks<N: Network = Ethereum> {
     client: WeakClient,
@@ -28,6 +67,7 @@ pub(crate) struct NewBlocks<N: Network = Ethereum> {
     next_yield: BlockNumber,
     /// LRU cache of known blocks. Only used by the polling task.
     known_blocks: LruCache<BlockNumber, N::BlockResponse>,
+    pub(crate) paused: Arc<Paused>,
     _phantom: PhantomData<N>,
 }
 
@@ -37,6 +77,7 @@ impl<N: Network> NewBlocks<N> {
             client,
             next_yield: NO_BLOCK_NUMBER,
             known_blocks: LruCache::new(BLOCK_CACHE_SIZE),
+            paused: Arc::default(),
             _phantom: PhantomData,
         }
     }
@@ -123,13 +164,17 @@ impl<N: Network> NewBlocks<N> {
                 yield known_block;
             }
 
+            // If we're paused, wait until we're unpaused.
+            // Once unpaused, reset `self.next_yield` to ignore the blocks that were included while we were paused.
+            let unpaused = self.paused.wait().await;
+
             // Get the tip.
             let Some(block_number) = numbers_stream.next().await else {
                 debug!("polling stream ended");
                 break 'task;
             };
             trace!(%block_number, "got block number");
-            if self.next_yield == NO_BLOCK_NUMBER {
+            if self.next_yield == NO_BLOCK_NUMBER || unpaused {
                 assert!(block_number < NO_BLOCK_NUMBER, "too many blocks");
                 // this stream can be initialized after the first tx was sent,
                 // to avoid the edge case where the tx is mined immediately, we should apply an

crates/provider/src/heart.rs

@@ -1,6 +1,6 @@
 //! Block heartbeat and pending transaction watcher.
 
-use crate::{Provider, RootProvider};
+use crate::{blocks::Paused, Provider, RootProvider};
 use alloy_consensus::BlockHeader;
 use alloy_json_rpc::RpcError;
 use alloy_network::{BlockResponse, Network};
@@ -14,6 +14,7 @@ use std::{
     collections::{BTreeMap, VecDeque},
     fmt,
     future::Future,
+    sync::Arc,
     time::Duration,
 };
 use tokio::{
@@ -463,18 +464,22 @@ pub(crate) struct Heartbeat<N, S> {
     /// Ordered map of transactions to reap at a certain time.
     reap_at: BTreeMap<Instant, B256>,
 
+    /// Whether the heartbeat is currently paused.
+    paused: Arc<Paused>,
+
     _network: std::marker::PhantomData<N>,
 }
 
 impl<N: Network, S: Stream<Item = N::BlockResponse> + Unpin + 'static> Heartbeat<N, S> {
     /// Create a new heartbeat task.
-    pub(crate) fn new(stream: S) -> Self {
+    pub(crate) fn new(stream: S, is_paused: Arc<Paused>) -> Self {
         Self {
             stream: stream.fuse(),
             past_blocks: Default::default(),
             unconfirmed: Default::default(),
             waiting_confs: Default::default(),
             reap_at: Default::default(),
+            paused: is_paused,
             _network: Default::default(),
         }
     }
@@ -531,6 +536,20 @@ impl<N: Network, S: Stream<Item = N::BlockResponse> + Unpin + 'static> Heartbeat
         }
     }
 
+    /// Check if we have any pending transactions.
+    fn has_pending_transactions(&self) -> bool {
+        !self.unconfirmed.is_empty() || !self.waiting_confs.is_empty()
+    }
+
+    /// Update the pause state based on whether we have pending transactions.
+    fn update_pause_state(&mut self) {
+        let should_pause = !self.has_pending_transactions();
+        if self.paused.is_paused() != should_pause {
+            debug!(paused = should_pause, "updating heartbeat pause state");
+            self.paused.set_paused(should_pause);
+        }
+    }
+
     /// Handle a watch instruction by adding it to the watch list, and
     /// potentially adding it to our `reap_at` list.
     fn handle_watch_ix(&mut self, to_watch: TxWatcher) {
@@ -597,7 +616,8 @@ impl<N: Network, S: Stream<Item = N::BlockResponse> + Unpin + 'static> Heartbeat
             // Check that the chain is continuous.
             if *last_height + 1 != block_height {
                 // Move all the transactions that were reset by the reorg to the unconfirmed list.
-                warn!(%block_height, last_height, "reorg detected");
+                // This can also happen if we unpaused the heartbeat after some time.
+                debug!(block_height, last_height, "reorg/unpause detected");
                 self.move_reorg_to_unconfirmed(block_height);
                 // Remove past blocks that are now invalid.
                 self.past_blocks.retain(|(h, _)| *h < block_height);
@@ -663,6 +683,8 @@ impl<N: Network, S: Stream<Item = N::BlockResponse> + Unpin + 'static> Heartbeat
     async fn into_future(mut self, mut ixns: mpsc::Receiver<TxWatcher>) {
         'shutdown: loop {
             {
+                self.update_pause_state();
+
                 let next_reap = self.next_reap();
                 let sleep = std::pin::pin!(sleep_until(next_reap.into()));
 

crates/provider/src/provider/root.rs

@@ -115,8 +115,9 @@ impl<N: Network> RootProvider<N> {
     pub(crate) fn get_heart(&self) -> &HeartbeatHandle {
         self.inner.heart.get_or_init(|| {
             let new_blocks = NewBlocks::<N>::new(self.inner.weak_client());
+            let paused = new_blocks.paused.clone();
             let stream = new_blocks.into_stream();
-            Heartbeat::<N, _>::new(Box::pin(stream)).spawn()
+            Heartbeat::<N, _>::new(Box::pin(stream), paused).spawn()
         })
     }
 }