refactor: roll custom canon state notification stream (#2486)

crates/storage/provider/Cargo.toml

@@ -26,6 +26,7 @@ tracing = "0.1"
 thiserror = "1.0.37"
 auto_impl = "1.0"
 itertools = "0.10"
+pin-project = "1.0"
 
 # test-utils
 reth-rlp = { path = "../../rlp", optional = true }

crates/storage/provider/src/traits/chain.rs

@@ -1,9 +1,14 @@
 //! Canonical chain state notification trait and types.
 use crate::{chain::BlockReceipts, Chain};
 use auto_impl::auto_impl;
-use std::sync::Arc;
+use std::{
+    pin::Pin,
+    sync::Arc,
+    task::{ready, Context, Poll},
+};
 use tokio::sync::broadcast;
-use tokio_stream::wrappers::BroadcastStream;
+use tokio_stream::{wrappers::BroadcastStream, Stream};
+use tracing::debug;
 
 /// Type alias for a receiver that receives [CanonStateNotification]
 pub type CanonStateNotifications = broadcast::Receiver<CanonStateNotification>;
@@ -20,8 +25,35 @@ pub trait CanonStateSubscriptions: Send + Sync {
     fn subscribe_to_canonical_state(&self) -> CanonStateNotifications;
 
     /// Convenience method to get a stream of [`CanonStateNotification`].
-    fn canonical_state_stream(&self) -> BroadcastStream<CanonStateNotification> {
-        BroadcastStream::new(self.subscribe_to_canonical_state())
+    fn canonical_state_stream(&self) -> CanonStateNotificationStream {
+        CanonStateNotificationStream {
+            st: BroadcastStream::new(self.subscribe_to_canonical_state()),
+        }
+    }
+}
+
+/// A Stream of [CanonStateNotification].
+#[derive(Debug)]
+#[pin_project::pin_project]
+pub struct CanonStateNotificationStream {
+    #[pin]
+    st: BroadcastStream<CanonStateNotification>,
+}
+
+impl Stream for CanonStateNotificationStream {
+    type Item = CanonStateNotification;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        loop {
+            return match ready!(self.as_mut().project().st.poll_next(cx)) {
+                Some(Ok(notification)) => Poll::Ready(Some(notification)),
+                Some(Err(err)) => {
+                    debug!(%err, "canonical state notification stream lagging behind");
+                    continue
+                }
+                None => Poll::Ready(None),
+            }
+        }
     }
 }