[consensus] Improve State Sync Ergonomics (#3828)

Author: patrick-ogrady

consensus/fuzz/src/lib.rs

@@ -18,7 +18,7 @@ use commonware_consensus::{
         config,
         mocks::{application, relay, reporter, twins},
         types::{Certificate, Vote},
-        Engine, ForwardingPolicy,
+        Engine, Floor, ForwardingPolicy,
     },
     types::{Delta, Epoch, View},
     Monitor, Viewable,
@@ -391,6 +391,7 @@ where
         partition: validator.to_string(),
         mailbox_size: NZUsize!(1024),
         epoch: Epoch::new(EPOCH),
+        floor: Floor::Genesis(application::genesis::<Sha256>(Epoch::new(EPOCH))),
         leader_timeout,
         certification_timeout,
         timeout_retry: Duration::from_secs(10),
@@ -634,6 +635,7 @@ fn run_with_twin_mutator<P: simplex::Simplex>(input: FuzzInput) {
                 partition: format!("twin_{idx}_primary"),
                 mailbox_size: NZUsize!(1024),
                 epoch: Epoch::new(EPOCH),
+                floor: Floor::Genesis(application::genesis::<Sha256>(Epoch::new(EPOCH))),
                 leader_timeout: Duration::from_secs(1),
                 certification_timeout: Duration::from_millis(1_500),
                 timeout_retry: Duration::from_secs(10),

consensus/src/aggregation/mocks/application.rs

@@ -1,7 +1,4 @@
-use crate::{
-    types::{Epoch, Height},
-    Automaton as A,
-};
+use crate::{types::Height, Automaton as A};
 use commonware_cryptography::{Hasher, Sha256};
 use commonware_utils::channel::oneshot;
 use tracing::trace;
@@ -33,12 +30,6 @@ impl A for Application {
     type Context = Height;
     type Digest = <Sha256 as Hasher>::Digest;
 
-    async fn genesis(&mut self, _epoch: Epoch) -> Self::Digest {
-        let mut hasher = Sha256::default();
-        hasher.update(b"genesis");
-        hasher.finalize()
-    }
-
     async fn propose(&mut self, context: Self::Context) -> oneshot::Receiver<Self::Digest> {
         let (sender, receiver) = oneshot::channel();
 

consensus/src/lib.rs

@@ -92,9 +92,6 @@ stability_scope!(BETA, cfg(not(target_arch = "wasm32")) {
         /// Hash of an arbitrary payload.
         type Digest: Digest;
 
-        /// Payload used to initialize the consensus engine.
-        fn genesis(&mut self, epoch: Epoch) -> impl Future<Output = Self::Digest> + Send;
-
         /// Generate a new payload for the given context.
         ///
         /// If it is possible to generate a payload, the Digest should be returned over the provided
@@ -258,12 +255,6 @@ stability_scope!(ALPHA, cfg(not(target_arch = "wasm32")) {
         /// The block type produced by the application's builder.
         type Block: Block;
 
-        /// Payload used to initialize the consensus engine in the first epoch.
-        ///
-        /// This future may be cancelled before it completes. Implementations must be
-        /// cancellation-safe.
-        fn genesis(&mut self) -> impl Future<Output = Self::Block> + Send;
-
         /// Build a new block on top of the provided parent ancestry. If the build job fails,
         /// or the proposer's slot should be skipped, the implementor should return [None].
         ///

consensus/src/marshal/ancestry.rs

@@ -1,7 +1,7 @@
 //! A stream that yields the ancestors of a block while prefetching parents.
 
 use crate::{types::Height, Block, Heightable};
-use commonware_cryptography::Digestible;
+use commonware_cryptography::{Digest, Digestible};
 use futures::{
     future::{BoxFuture, OptionFuture},
     FutureExt, Stream,
@@ -24,7 +24,7 @@ where
 }
 
 /// An interface for providing parent blocks.
-pub trait BlockProvider: Clone + Send + 'static {
+pub trait BlockProvider: Send + 'static {
     /// The block type the provider walks.
     type Block: Block;
 
@@ -41,29 +41,56 @@ pub trait BlockProvider: Clone + Send + 'static {
     ///
     /// The child block can carry variant-specific context needed to retrieve its parent.
     fn subscribe_parent(
-        self,
-        block: Self::Block,
-    ) -> impl Future<Output = Option<Self::Block>> + Send;
+        &self,
+        block: &Self::Block,
+    ) -> impl Future<Output = Option<Self::Block>> + Send + 'static;
 }
 
-/// Yields the ancestors of a block while prefetching parents, _not_ including the genesis block.
-///
-// TODO(<https://github.com/commonwarexyz/monorepo/issues/2982>): Once marshal can also yield the genesis block,
-// this stream should end at block height 0 rather than 1.
+// Expected parent height and digest for a pending fetch.
+struct ExpectedParent<D>(Height, D);
+
+// Pending parent fetch paired with the relationship it must satisfy.
+type PendingFetch<B> = BoxFuture<'static, Option<(ExpectedParent<<B as Digestible>::Digest>, B)>>;
+
+impl<D: Digest> ExpectedParent<D> {
+    fn from_child<B: Block<Digest = D>>(child: &B) -> Self {
+        Self(
+            child.height().previous().expect("child must have parent"),
+            child.parent(),
+        )
+    }
+
+    fn assert_matches<B: Block<Digest = D>>(self, parent: &B) {
+        let Self(parent_height, parent_digest) = self;
+        assert_eq!(
+            parent.height(),
+            parent_height,
+            "fetched parent must be contiguous in height"
+        );
+        assert_eq!(
+            parent.digest(),
+            parent_digest,
+            "fetched parent must be contiguous in ancestry"
+        );
+    }
+}
+
+/// Yields the ancestors of a block while prefetching parents, including the
+/// height-zero genesis block if it is available.
 #[pin_project]
 pub struct AncestorStream<M: BlockProvider> {
     buffered: Vec<M::Block>,
     marshal: M,
     #[pin]
-    pending: OptionFuture<BoxFuture<'static, Option<M::Block>>>,
+    pending: OptionFuture<PendingFetch<M::Block>>,
 }
 
 impl<M: BlockProvider> AncestorStream<M> {
     /// Creates a new [AncestorStream] starting from the given ancestry.
     ///
     /// # Panics
     ///
-    /// Panics if the initial blocks are not contiguous in height.
+    /// Panics if the initial blocks are not contiguous.
     pub(crate) fn new(marshal: M, initial: impl IntoIterator<Item = M::Block>) -> Self {
         let mut buffered = initial.into_iter().collect::<Vec<M::Block>>();
         buffered.sort_by_key(Heightable::height);
@@ -93,13 +120,11 @@ impl<M: BlockProvider> AncestorStream<M> {
 impl<M> Stream for AncestorStream<M>
 where
     M: BlockProvider,
-    M::Block: Clone,
 {
     type Item = M::Block;
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        // Because marshal cannot currently yield the genesis block, we stop at height 1.
-        const END_BOUND: Height = Height::new(1);
+        const END_BOUND: Height = Height::zero();
 
         let mut this = self.project();
 
@@ -109,14 +134,20 @@ where
             let should_walk_parent = height > END_BOUND;
             let end_of_buffered = this.buffered.is_empty();
             if should_walk_parent && end_of_buffered {
-                let future = this.marshal.clone().subscribe_parent(block.clone()).boxed();
+                let expected = ExpectedParent::from_child(&block);
+                let future = this
+                    .marshal
+                    .subscribe_parent(&block)
+                    .map(move |parent| parent.map(|parent| (expected, parent)))
+                    .boxed();
                 *this.pending.as_mut() = Some(future).into();
 
                 // Explicitly poll the next future to kick off the fetch. If it's already ready,
                 // buffer it for the next poll.
                 match this.pending.as_mut().poll(cx) {
-                    Poll::Ready(Some(Some(block))) => {
-                        this.buffered.push(block);
+                    Poll::Ready(Some(Some((expected, parent)))) => {
+                        expected.assert_matches(&parent);
+                        this.buffered.push(parent);
                     }
                     Poll::Ready(Some(None)) => {
                         *this.pending.as_mut() = None.into();
@@ -137,18 +168,25 @@ where
                 *this.pending.as_mut() = None.into();
                 Poll::Ready(None)
             }
-            Poll::Ready(Some(Some(block))) => {
+            Poll::Ready(Some(Some((expected, block)))) => {
+                expected.assert_matches(&block);
                 let height = block.height();
                 let should_walk_parent = height > END_BOUND;
                 if should_walk_parent {
-                    let future = this.marshal.clone().subscribe_parent(block.clone()).boxed();
+                    let expected = ExpectedParent::from_child(&block);
+                    let future = this
+                        .marshal
+                        .subscribe_parent(&block)
+                        .map(move |parent| parent.map(|parent| (expected, parent)))
+                        .boxed();
                     *this.pending.as_mut() = Some(future).into();
 
                     // Explicitly poll the next future to kick off the fetch. If it's already ready,
                     // buffer it for the next poll.
                     match this.pending.as_mut().poll(cx) {
-                        Poll::Ready(Some(Some(block))) => {
-                            this.buffered.push(block);
+                        Poll::Ready(Some(Some((expected, parent)))) => {
+                            expected.assert_matches(&parent);
+                            this.buffered.push(parent);
                         }
                         Poll::Ready(Some(None)) => {
                             *this.pending.as_mut() = None.into();
@@ -179,9 +217,25 @@ mod test {
     impl BlockProvider for MockProvider {
         type Block = Block<Sha256Digest, ()>;
 
-        async fn subscribe_parent(self, block: Self::Block) -> Option<Self::Block> {
+        fn subscribe_parent(
+            &self,
+            block: &Self::Block,
+        ) -> impl Future<Output = Option<Self::Block>> + Send + 'static {
             let parent = block.parent;
-            self.0.into_iter().find(|b| b.digest() == parent)
+            std::future::ready(self.0.iter().find(|b| b.digest() == parent).cloned())
+        }
+    }
+
+    #[derive(Clone)]
+    struct WrongParentProvider(Block<Sha256Digest, ()>);
+    impl BlockProvider for WrongParentProvider {
+        type Block = Block<Sha256Digest, ()>;
+
+        fn subscribe_parent(
+            &self,
+            _block: &Self::Block,
+        ) -> impl Future<Output = Option<Self::Block>> + Send + 'static {
+            std::future::ready(Some(self.0.clone()))
         }
     }
 
@@ -209,6 +263,45 @@ mod test {
         );
     }
 
+    #[test]
+    #[should_panic = "fetched parent must be contiguous in height"]
+    fn test_panics_on_non_contiguous_fetched_parent_height() {
+        let parent = Block::new::<Sha256>((), Sha256Digest::EMPTY, Height::zero(), 0);
+        let child = Block::new::<Sha256>((), parent.digest(), Height::new(3), 3);
+        let stream = AncestorStream::new(MockProvider(vec![parent]), [child]);
+        futures::pin_mut!(stream);
+
+        let waker = futures::task::noop_waker_ref();
+        let mut cx = std::task::Context::from_waker(waker);
+        let _ = futures::Stream::poll_next(stream.as_mut(), &mut cx);
+    }
+
+    #[test]
+    #[should_panic = "fetched parent must be contiguous in ancestry"]
+    fn test_panics_on_non_contiguous_fetched_parent_digest() {
+        let expected_parent = Block::new::<Sha256>((), Sha256Digest::EMPTY, Height::zero(), 0);
+        let fetched_parent = Block::new::<Sha256>((), Sha256Digest::EMPTY, Height::zero(), 1);
+        let child = Block::new::<Sha256>((), expected_parent.digest(), Height::new(1), 2);
+        let stream = AncestorStream::new(WrongParentProvider(fetched_parent), [child]);
+        futures::pin_mut!(stream);
+
+        let waker = futures::task::noop_waker_ref();
+        let mut cx = std::task::Context::from_waker(waker);
+        let _ = futures::Stream::poll_next(stream.as_mut(), &mut cx);
+    }
+
+    #[test_async]
+    async fn test_yields_genesis_and_stops() {
+        let genesis = Block::new::<Sha256>((), Sha256Digest::EMPTY, Height::zero(), 0);
+        let child = Block::new::<Sha256>((), genesis.digest(), Height::new(1), 1);
+
+        let provider = MockProvider(vec![genesis.clone()]);
+        let stream = AncestorStream::new(provider, [child.clone()]);
+
+        let results = stream.collect::<Vec<_>>().await;
+        assert_eq!(results, vec![child, genesis]);
+    }
+
     #[test_async]
     async fn test_empty_yields_none() {
         let mut stream: AncestorStream<MockProvider> =