feat: add helper stream type for finalized notifications

crates/chain-state/Cargo.toml

@@ -53,6 +53,7 @@ alloy-signer.workspace = true
 alloy-signer-local.workspace = true
 alloy-consensus.workspace = true
 rand.workspace = true
+futures.workspace = true
 
 [features]
 test-utils = [

crates/chain-state/src/notifications.rs

@@ -2,6 +2,7 @@
 
 use alloy_eips::eip2718::Encodable2718;
 use derive_more::{Deref, DerefMut};
+use pin_project::pin_project;
 use reth_execution_types::{BlockReceipts, Chain};
 use reth_primitives::{NodePrimitives, RecoveredBlock, SealedHeader};
 use reth_storage_api::NodePrimitivesProvider;
@@ -211,6 +212,86 @@ impl<T: Clone + Sync + Send + 'static> Stream for ForkChoiceStream<T> {
     }
 }
 
+/// A stream that emits canonical chain notifications only when blocks are finalized.
+#[pin_project]
+pub(crate) struct FinalizedNotificationStream<N: NodePrimitives> {
+    #[pin]
+    canon_stream: CanonStateNotificationStream<N>,
+    #[pin]
+    finalized_stream: ForkChoiceStream<SealedHeader<alloy_consensus::Header>>,
+    // Buffer for the most recent canonical notification
+    buffered_notification: Option<CanonStateNotification<N>>,
+    // Buffer for the most recent finalization event
+    buffered_finalization: Option<SealedHeader<alloy_consensus::Header>>,
+}
+
+impl<N: NodePrimitives> FinalizedNotificationStream<N>
+where
+    N::Block: Clone + Send + Sync + 'static,
+{
+    // Helper method to check if we have a matching pair of notifications
+    fn try_match(
+        buffered_notification: &Option<CanonStateNotification<N>>,
+        buffered_finalization: &Option<SealedHeader<alloy_consensus::Header>>,
+    ) -> Option<CanonStateNotification<N>> {
+        match (buffered_notification, buffered_finalization) {
+            (Some(notification), Some(finalized_header))
+                if notification.tip().hash() == finalized_header.hash() =>
+            {
+                Some(notification.clone())
+            }
+            _ => None,
+        }
+    }
+}
+
+impl<N: NodePrimitives> Stream for FinalizedNotificationStream<N>
+where
+    N::Block: Send + Sync + 'static,
+{
+    type Item = CanonStateNotification<N>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+        loop {
+            let mut progress = false;
+
+            // Poll both streams and update buffers as needed
+            match this.canon_stream.as_mut().poll_next(cx) {
+                Poll::Ready(Some(notification)) => {
+                    progress = true;
+                    *this.buffered_notification = Some(notification);
+                }
+                Poll::Ready(None) => return Poll::Ready(None),
+                Poll::Pending => {}
+            }
+
+            match this.finalized_stream.as_mut().poll_next(cx) {
+                Poll::Ready(Some(finalized_header)) => {
+                    progress = true;
+                    *this.buffered_finalization = Some(finalized_header);
+                }
+                Poll::Ready(None) => return Poll::Ready(None),
+                Poll::Pending => {}
+            }
+
+            // Check if we have a matching pair using the projected fields
+            if let Some(notification) =
+                Self::try_match(this.buffered_notification, this.buffered_finalization)
+            {
+                return Poll::Ready(Some(notification));
+            }
+
+            // If we made no progress and both streams are pending, break the loop
+            if !progress {
+                break;
+            }
+        }
+
+        Poll::Pending
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -472,4 +553,66 @@ mod tests {
         // Confirm this is from the committed segment.
         assert!(!block_receipts[1].1);
     }
+
+    #[test]
+    fn test_finalized_notification_stream() {
+        // Create a default block for testing
+        let block: RecoveredBlock<reth_primitives::Block> = Default::default();
+
+        // Create the header first so we can use its hash
+        let header = alloy_consensus::Header { number: 1u64, ..Default::default() };
+        let sealed_header = SealedHeader::seal_slow(header);
+        let block_hash = sealed_header.hash();
+
+        // Now create the block with the matching hash
+        let mut test_block = block;
+        test_block.set_block_number(1);
+        test_block.set_hash(block_hash);
+
+        // Create the chain with our test block
+        let chain: Arc<Chain> =
+            Arc::new(Chain::new(vec![test_block.clone()], ExecutionOutcome::default(), None));
+
+        // Create channels for both streams
+        let (canon_tx, canon_rx) = broadcast::channel(10);
+        let (finalized_tx, finalized_rx) = watch::channel(None);
+
+        // Create the notification stream
+        let mut stream = FinalizedNotificationStream {
+            canon_stream: CanonStateNotificationStream { st: BroadcastStream::new(canon_rx) },
+            finalized_stream: ForkChoiceStream::new(finalized_rx),
+            buffered_notification: None,
+            buffered_finalization: None,
+        };
+
+        // Create a notification
+        let notification = CanonStateNotification::Commit { new: chain };
+
+        // Send the canonical notification
+        canon_tx.send(notification).unwrap();
+
+        // Create a waker and context for polling
+        let waker = futures::task::noop_waker();
+        let mut cx = Context::from_waker(&waker);
+
+        // First poll should be pending since we haven't received finalization yet
+        assert!(matches!(Pin::new(&mut stream).poll_next(&mut cx), Poll::Pending));
+
+        // Send the finalization notification
+        finalized_tx.send(Some(sealed_header)).unwrap();
+
+        // Now polling should return the notification since we have both events
+        match Pin::new(&mut stream).poll_next(&mut cx) {
+            Poll::Ready(Some(received)) => {
+                assert_eq!(received.tip().hash(), block_hash);
+                match received {
+                    CanonStateNotification::Commit { new } => {
+                        assert_eq!(new.tip().hash(), block_hash);
+                    }
+                    _ => panic!("Expected Commit notification"),
+                }
+            }
+            other => panic!("Expected Ready(Some), got: {:?}", other),
+        }
+    }
 }