chore: reorganise and add test infrastructure to make test readable

also fix some bugs along the way, of course

Signed-off-by: Roland Kuhn <rk@rkuhn.info>

Author: rkuhn

crates/amaru/src/stages/consensus/chain_forward.rs

@@ -12,12 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use crate::stages::PallasPoint;
 use acto::{AcTokio, ActoCell, ActoMsgSuper, ActoRef, ActoRuntime};
 use amaru_consensus::{consensus::store::ChainStore, IsHeader};
 use amaru_kernel::{Hash, Header};
 use amaru_ledger::BlockValidationResult;
 use client_protocol::{client_protocols, ClientProtocolMsg};
-use client_state::{to_pallas_point, ClientOp};
 use gasket::framework::*;
 use pallas_network::{
     facades::PeerServer,
@@ -40,9 +40,6 @@ pub const EVENT_TARGET: &str = "amaru::consensus::chain_forward";
 
 /// Forwarding stage of the consensus where blocks are stored and made
 /// available to downstream peers.
-///
-/// TODO: currently does nothing, should store block, update chain state, and
-/// forward new chain downstream
 
 #[derive(Stage)]
 #[stage(name = "consensus.forward", unit = "Unit", worker = "Worker")]
@@ -52,14 +49,16 @@ pub struct ForwardStage {
     pub network_magic: u64,
     pub runtime: AcTokio,
     pub listen_address: String,
-    pub downstream: Option<ActoRef<ForwardEvent>>,
+    pub downstream: ActoRef<ForwardEvent>,
     pub max_peers: usize,
+    pub our_tip: Tip,
 }
 
 #[derive(Debug, Clone)]
 pub enum ForwardEvent {
     Listening(u16),
     Forward(Point),
+    Backward(Point),
 }
 
 impl ForwardStage {
@@ -69,6 +68,7 @@ impl ForwardStage {
         network_magic: u64,
         listen_address: &str,
         max_peers: usize,
+        our_tip: Tip,
     ) -> Self {
         #[allow(clippy::expect_used)]
         let runtime =
@@ -79,8 +79,9 @@ impl ForwardStage {
             network_magic,
             runtime,
             listen_address: listen_address.to_string(),
-            downstream,
+            downstream: downstream.unwrap_or_else(ActoRef::blackhole),
             max_peers,
+            our_tip,
         }
     }
 }
@@ -93,7 +94,7 @@ pub enum Unit {
 pub struct Worker {
     server: JoinHandle<()>,
     incoming_peers: Receiver<PeerServer>,
-    tip: Tip,
+    our_tip: Tip,
     clients: ActoRef<ClientMsg>,
 }
 
@@ -103,16 +104,46 @@ impl Drop for Worker {
     }
 }
 
+#[allow(clippy::large_enum_variant)]
+#[derive(Debug, Clone)]
+enum ClientOp {
+    /// the tip to go back to
+    Backward(Tip),
+    /// the header to go forward to and the tip we will be at after sending this header
+    Forward(Header, Tip),
+}
+
+impl PartialEq for ClientOp {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Self::Backward(l0), Self::Backward(r0)) => (&l0.0, l0.1) == (&r0.0, r0.1),
+            (Self::Forward(l0, l1), Self::Forward(r0, r1)) => {
+                l0 == r0 && (&l1.0, l1.1) == (&r1.0, r1.1)
+            }
+            _ => false,
+        }
+    }
+}
+
+impl Eq for ClientOp {}
+
+impl ClientOp {
+    pub fn tip(&self) -> Tip {
+        match self {
+            ClientOp::Backward(tip) => tip.clone(),
+            ClientOp::Forward(_, tip) => tip.clone(),
+        }
+    }
+}
+
 #[async_trait::async_trait(?Send)]
 impl gasket::framework::Worker<ForwardStage> for Worker {
     async fn bootstrap(stage: &ForwardStage) -> Result<Self, WorkerError> {
         let server = TcpListener::bind(&stage.listen_address).await.or_panic()?;
-        if let Some(downstream) = &stage.downstream {
-            tracing::debug!("sending listening event");
-            downstream.send(ForwardEvent::Listening(
-                server.local_addr().or_panic()?.port(),
-            ));
-        }
+        tracing::debug!("sending listening event");
+        stage.downstream.send(ForwardEvent::Listening(
+            server.local_addr().or_panic()?.port(),
+        ));
 
         let (tx, incoming_peers) = mpsc::channel(10);
 
@@ -154,7 +185,7 @@ impl gasket::framework::Worker<ForwardStage> for Worker {
         Ok(Self {
             server,
             incoming_peers,
-            tip: Tip(Point::Origin, 0),
+            our_tip: stage.our_tip.clone(),
             clients,
         })
     }
@@ -191,12 +222,25 @@ impl gasket::framework::Worker<ForwardStage> for Worker {
                 // FIXME: block height should be part of BlockValidated message
                 let store = stage.store.lock().await;
                 if let Some(header) = store.load_header(&Hash::from(point)) {
-                    self.tip = Tip(to_pallas_point(point), header.block_height());
-                    self.clients.send(ClientMsg::Op(ClientOp::Forward(header)));
-                }
+                    // assert that the new tip is a direct successor of the old tip
+                    assert_eq!(header.block_height(), self.our_tip.1 + 1);
+                    match header.parent() {
+                        Some(parent) => assert_eq!(
+                            Point::new(self.our_tip.0.slot_or_default(), parent.as_ref().to_vec()),
+                            self.our_tip.0
+                        ),
+                        None => assert_eq!(self.our_tip.0, Point::Origin),
+                    }
 
-                if let Some(downstream) = &stage.downstream {
-                    downstream.send(ForwardEvent::Forward(to_pallas_point(point)));
+                    self.our_tip = Tip(point.pallas_point(), header.block_height());
+                    self.clients.send(ClientMsg::Op(ClientOp::Forward(
+                        header,
+                        self.our_tip.clone(),
+                    )));
+
+                    stage
+                        .downstream
+                        .send(ForwardEvent::Forward(point.pallas_point()));
                 }
 
                 Ok(())
@@ -213,9 +257,13 @@ impl gasket::framework::Worker<ForwardStage> for Worker {
                 // FIXME: block height should be part of BlockValidated message
                 let store = stage.store.lock().await;
                 if let Some(header) = store.load_header(&Hash::from(point)) {
-                    self.tip = Tip(to_pallas_point(point), header.block_height());
+                    self.our_tip = Tip(point.pallas_point(), header.block_height());
                     self.clients
-                        .send(ClientMsg::Op(ClientOp::Backward(to_pallas_point(point))));
+                        .send(ClientMsg::Op(ClientOp::Backward(self.our_tip.clone())));
+
+                    stage
+                        .downstream
+                        .send(ForwardEvent::Backward(point.pallas_point()));
                 }
 
                 Ok(())
@@ -235,7 +283,8 @@ impl gasket::framework::Worker<ForwardStage> for Worker {
                 // FIXME: gasket design bug that we only get &Unit and thus cannot take values from it without internal mutability
                 let peer = peer.borrow_mut().take();
                 if let Some(peer) = peer {
-                    self.clients.send(ClientMsg::Peer(peer, self.tip.clone()));
+                    self.clients
+                        .send(ClientMsg::Peer(peer, self.our_tip.clone()));
                 } else {
                     tracing::error!(target: EVENT_TARGET, "Unit::Peer was empty in execute");
                 }
@@ -247,29 +296,33 @@ impl gasket::framework::Worker<ForwardStage> for Worker {
 
 #[allow(clippy::large_enum_variant)]
 enum ClientMsg {
+    /// A new peer has connected to us.
+    ///
+    /// Our tip is included to get the connection handlers started correctly.
     Peer(PeerServer, Tip),
+    /// An operation to be executed on all clients.
     Op(ClientOp),
 }
 
 async fn client_supervisor(
-    mut cell: ActoCell<ClientMsg, impl ActoRuntime, Result<(), client_protocol::ClientError>>,
+    mut cell: ActoCell<ClientMsg, impl ActoRuntime, anyhow::Result<()>>,
     store: Arc<Mutex<dyn ChainStore<Header>>>,
     max_peers: usize,
 ) {
     let mut clients = HashMap::new();
     while let Some(msg) = cell.recv().await.has_senders() {
         match msg {
             ActoMsgSuper::Message(ClientMsg::Peer(peer, tip)) => {
-                if clients.len() >= max_peers {
-                    tracing::warn!(target: EVENT_TARGET, "max peers reached, dropping peer");
-                    continue;
-                }
-
                 let addr = peer
                     .accepted_address()
                     .map(|a| a.to_string())
                     .unwrap_or_default();
 
+                if clients.len() >= max_peers {
+                    tracing::warn!(target: EVENT_TARGET, "max peers reached, dropping peer from {addr}");
+                    continue;
+                }
+
                 let client = cell.spawn_supervised(&addr, {
                     let store = store.clone();
                     move |cell| client_protocols(cell, peer, store, tip)
@@ -294,3 +347,6 @@ mod client_state;
 
 #[cfg(test)]
 mod tests;
+
+#[cfg(test)]
+mod test_infra;