chore: add debug logs to chainsync server and fix bugs

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

Author: rkuhn

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

@@ -105,14 +105,53 @@ impl Drop for Worker {
 }
 
 #[allow(clippy::large_enum_variant)]
-#[derive(Debug, Clone)]
+#[derive(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 std::fmt::Debug for ClientOp {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Backward(tip) => f
+                .debug_struct("Backward")
+                .field("tip", &(tip.1, PrettyPoint(&tip.0)))
+                .finish(),
+            Self::Forward(header, tip) => f
+                .debug_struct("Forward")
+                .field(
+                    "header",
+                    &(header.block_height(), PrettyPoint(&header.pallas_point())),
+                )
+                .field("tip", &(tip.1, PrettyPoint(&tip.0)))
+                .finish(),
+        }
+    }
+}
+
+struct PrettyPoint<'a>(&'a Point);
+
+impl std::fmt::Debug for PrettyPoint<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "({}, {})",
+            self.0.slot_or_default(),
+            hex::encode(hash_point(&self.0))
+        )
+    }
+}
+
+fn hash_point(point: &Point) -> Hash<32> {
+    match point {
+        Point::Origin => Hash::from([0; 32]),
+        Point::Specific(_slot, hash) => Hash::from(hash.as_slice()),
+    }
+}
+
 impl PartialEq for ClientOp {
     fn eq(&self, other: &Self) -> bool {
         match (self, other) {

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

@@ -107,25 +107,34 @@ async fn chain_sync(
         let input = cell.recv().await;
         match input {
             ActoInput::Message(ChainSyncMsg::Op(op)) => {
+                tracing::debug!("got op {op:?}");
                 our_tip = op.tip();
                 state.add_op(op);
                 if waiting {
                     if let Some(op) = state.next_op() {
+                        tracing::debug!("sending op {op:?} to waiting handler");
                         waiting = false;
                         handler.send(Some((op, our_tip.clone())));
                     }
                 }
             }
             ActoInput::Message(ChainSyncMsg::ReqNext) => {
+                tracing::debug!("got req next");
                 if let Some(op) = state.next_op() {
+                    tracing::debug!("sending op {op:?} to handler");
                     handler.send(Some((op, our_tip.clone())));
                 } else {
+                    tracing::debug!("sending await reply");
                     handler.send(None);
                     waiting = true;
                 }
             }
-            ActoInput::NoMoreSenders => return Ok(()),
+            ActoInput::NoMoreSenders => {
+                tracing::debug!("no more senders");
+                return Ok(());
+            }
             ActoInput::Supervision { result, .. } => {
+                tracing::debug!("supervision result: {result:?}");
                 return result
                     .map_err(|e| anyhow::Error::from(ClientError::HandlerFailure(e.to_string())))
                     .and_then(|x| x);
@@ -144,24 +153,30 @@ async fn chain_sync_handler(
 ) -> anyhow::Result<()> {
     loop {
         let Some(req) = server.recv_while_idle().await? else {
+            tracing::debug!("client terminated");
             return Err(ClientError::ClientTerminated.into());
         };
         if !matches!(req, ClientRequest::RequestNext) {
+            tracing::debug!("late intersection");
             return Err(ClientError::LateIntersection.into());
         };
+        tracing::debug!("got req next");
         parent.send(ChainSyncMsg::ReqNext);
 
         if let ActoInput::Message(op) = cell.recv().await {
             match op {
                 Some((ClientOp::Forward(header, _), tip)) => {
+                    tracing::debug!("sending roll forward");
                     server
                         .send_roll_forward(to_header_content(header), tip)
                         .await?;
                 }
                 Some((ClientOp::Backward(point), tip)) => {
+                    tracing::debug!("sending roll backward");
                     server.send_roll_backward(point.0, tip).await?;
                 }
                 None => {
+                    tracing::debug!("sending await reply");
                     server.send_await_reply().await?;
                     let ActoInput::Message(Some((op, tip))) = cell.recv().await else {
                         return Ok(());
@@ -179,6 +194,7 @@ async fn chain_sync_handler(
                 }
             }
         } else {
+            tracing::debug!("parent terminated");
             // parent terminated
             return Ok(());
         }

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

@@ -1,7 +1,7 @@
-use super::ClientOp;
+use super::{hash_point, ClientOp};
 use crate::stages::AsTip;
 use amaru_consensus::{consensus::store::ChainStore, IsHeader};
-use amaru_kernel::{Hash, Header};
+use amaru_kernel::Header;
 use pallas_network::miniprotocols::{chainsync::Tip, Point};
 use std::collections::VecDeque;
 
@@ -19,24 +19,30 @@ impl ClientState {
     }
 
     pub fn next_op(&mut self) -> Option<ClientOp> {
+        tracing::debug!("next_op: {:?}", self.ops.front());
         self.ops.pop_front()
     }
 
     pub fn add_op(&mut self, op: ClientOp) {
+        tracing::debug!("add_op: {:?}", op);
         match op {
             ClientOp::Backward(tip) => {
-                if let Some(index) = self
-                    .ops
-                    .iter()
-                    .rposition(|op| matches!(op, ClientOp::Forward(_, tip) if &tip.0 == &tip.0))
+                if let Some((index, _)) =
+                    self.ops.iter().enumerate().rfind(
+                        |(_, op)| matches!(op, ClientOp::Forward(_, tip2) if &tip2.0 == &tip.0),
+                    )
                 {
+                    tracing::debug!("found backward op at index {index} in {:?}", self.ops);
                     self.ops.truncate(index + 1);
+                    tracing::debug!("last after truncate: {:?}", self.ops.back());
                 } else {
+                    tracing::debug!("clearing ops");
                     self.ops.clear();
                     self.ops.push_back(ClientOp::Backward(tip));
                 }
             }
             op @ ClientOp::Forward(..) => {
+                tracing::debug!("adding forward op");
                 self.ops.push_back(op);
             }
         }
@@ -85,10 +91,3 @@ pub(super) fn find_headers_between(
     headers.reverse();
     Some((headers, Tip(Point::Origin, 0)))
 }
-
-pub(super) fn hash_point(point: &Point) -> Hash<32> {
-    match point {
-        Point::Origin => Hash::from([0; 32]),
-        Point::Specific(_slot, hash) => Hash::from(hash.as_slice()),
-    }
-}

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

@@ -1,6 +1,6 @@
 #![allow(dead_code)]
 
-use super::{ForwardEvent, ForwardStage};
+use super::{ForwardEvent, ForwardStage, PrettyPoint};
 use crate::stages::PallasPoint;
 use acto::{AcTokio, AcTokioRuntime, ActoCell, ActoInput, ActoRuntime};
 use amaru_consensus::consensus::store::{ChainStore, Nonces, StoreError};
@@ -269,6 +269,21 @@ impl Client {
         ops
     }
 
+    pub fn recv_n<const N: usize>(&mut self) -> [ClientMsg; N] {
+        let mut ops = Vec::new();
+        for _ in 0..N {
+            let msg = block_on(&self.runtime, self.client.chainsync().request_next()).unwrap();
+            match msg {
+                NextResponse::RollForward(header, tip) => {
+                    ops.push(ClientMsg::Forward(from_cbor(&header.cbor).unwrap(), tip))
+                }
+                NextResponse::RollBackward(point, tip) => ops.push(ClientMsg::Backward(point, tip)),
+                NextResponse::Await => break,
+            }
+        }
+        ops.try_into().unwrap()
+    }
+
     pub fn recv_after_await(&mut self) -> ClientMsg {
         let msg = block_on(
             &self.runtime,
@@ -285,12 +300,32 @@ impl Client {
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub enum ClientMsg {
     Forward(Header, Tip),
     Backward(Point, Tip),
 }
 
+impl std::fmt::Debug for ClientMsg {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Forward(header, tip) => f
+                .debug_struct("Forward")
+                .field(
+                    "header",
+                    &(header.block_height(), PrettyPoint(&header.pallas_point())),
+                )
+                .field("tip", &(tip.1, PrettyPoint(&tip.0)))
+                .finish(),
+            Self::Backward(point, tip) => f
+                .debug_struct("Backward")
+                .field("point", &PrettyPoint(&point))
+                .field("tip", &(tip.1, PrettyPoint(&tip.0)))
+                .finish(),
+        }
+    }
+}
+
 impl PartialEq for ClientMsg {
     fn eq(&self, other: &Self) -> bool {
         match (self, other) {

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

@@ -131,3 +131,42 @@ fn test_chain_sync() {
 
     // Note: there’s no way to shut down the gasket stage without logging to ERRORs, sorry
 }
+
+#[test]
+fn test_sync_optimising_rollback() {
+    let mut setup = Setup::new(LOST_47);
+    let chain = setup.store.get_chain(TIP_47);
+    let lost = setup.store.get(&hash(LOST_47)).unwrap().clone();
+
+    let mut client = setup.connect();
+    client.find_intersect(vec![]).0.expect("no intersection");
+
+    let msgs = client.recv_n::<4>();
+    assert_eq!(
+        msgs,
+        [
+            ClientMsg::Forward(chain[0].clone(), lost.as_tip()),
+            ClientMsg::Forward(chain[1].clone(), lost.as_tip()),
+            ClientMsg::Forward(chain[2].clone(), lost.as_tip()),
+            ClientMsg::Forward(chain[3].clone(), lost.as_tip()),
+        ]
+    );
+
+    setup.send_backward(BRANCH_47);
+    setup.send_validated(&chain[7].hash().to_string());
+    setup.send_validated(&chain[8].hash().to_string());
+    setup.send_validated(&chain[9].hash().to_string());
+
+    let msgs = client.recv_until_await();
+    assert_eq!(
+        msgs,
+        [
+            ClientMsg::Forward(chain[4].clone(), chain[9].as_tip()),
+            ClientMsg::Forward(chain[5].clone(), chain[9].as_tip()),
+            ClientMsg::Forward(chain[6].clone(), chain[9].as_tip()),
+            ClientMsg::Forward(chain[7].clone(), chain[9].as_tip()),
+            ClientMsg::Forward(chain[8].clone(), chain[9].as_tip()),
+            ClientMsg::Forward(chain[9].clone(), chain[9].as_tip()),
+        ]
+    );
+}