feat: move the forward_chain stage to pure-stage

Signed-off-by: etorreborre <etorreborre@yahoo.com>

Author: etorreborre

crates/amaru-consensus/src/consensus/effects/mod.rs

@@ -13,4 +13,5 @@
 // limitations under the License.
 
 pub mod block_effects;
+pub mod network_effects;
 pub mod store_effects;

crates/amaru-consensus/src/consensus/effects/network_effects.rs

@@ -0,0 +1,102 @@
+// Copyright 2025 PRAGMA
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::consensus::errors::ProcessingFailed;
+use crate::consensus::tip::HeaderTip;
+use amaru_kernel::peer::Peer;
+use amaru_kernel::{Header, Point};
+use amaru_ouroboros_traits::IsHeader;
+use anyhow::anyhow;
+use async_trait::async_trait;
+use pure_stage::{ExternalEffect, ExternalEffectAPI, Resources};
+use serde::{Deserialize, Serialize};
+use std::fmt::Display;
+use std::sync::Arc;
+
+pub type ResourceForwardEventListener = Arc<dyn ForwardEventListener + Send + Sync>;
+
+/// A listener interface for forward events (new headers or rollbacks).
+/// These events are either caught for tests or forwarded to downstream peers (see the TcpForwardEventListener implementation).
+#[async_trait]
+pub trait ForwardEventListener {
+    async fn send(&self, event: ForwardEvent) -> anyhow::Result<()>;
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub enum ForwardEvent {
+    Forward(Header),
+    Backward(HeaderTip),
+}
+
+impl ForwardEvent {
+    pub fn point(&self) -> Point {
+        match self {
+            ForwardEvent::Forward(header) => header.point(),
+            ForwardEvent::Backward(tip) => tip.point(),
+        }
+    }
+}
+
+impl Display for ForwardEvent {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ForwardEvent::Forward(header) => write!(f, "Forward({})", header.point()),
+            ForwardEvent::Backward(tip) => write!(f, "Backward({})", tip),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ForwardEventEffect {
+    peer: Peer,
+    event: ForwardEvent,
+}
+
+impl ForwardEventEffect {
+    pub fn new(peer: &Peer, event: ForwardEvent) -> Self {
+        Self {
+            peer: peer.clone(),
+            event,
+        }
+    }
+}
+
+impl ExternalEffect for ForwardEventEffect {
+    #[expect(clippy::expect_used)]
+    fn run(
+        self: Box<Self>,
+        resources: Resources,
+    ) -> pure_stage::BoxFuture<'static, Box<dyn pure_stage::SendData>> {
+        Box::pin(async move {
+            let listener = resources
+                .get::<ResourceForwardEventListener>()
+                .expect("ForwardEventEffect requires a ForwardEventListener")
+                .clone();
+
+            let point = self.event.point();
+            let result: <Self as ExternalEffectAPI>::Response =
+                listener.send(self.event).await.map_err(|e| {
+                    ProcessingFailed::new(
+                        &self.peer,
+                        anyhow!("Cannot send the forward event {}: {e}", &point),
+                    )
+                });
+            Box::new(result) as Box<dyn pure_stage::SendData>
+        })
+    }
+}
+
+impl ExternalEffectAPI for ForwardEventEffect {
+    type Response = Result<(), ProcessingFailed>;
+}

crates/amaru-consensus/src/consensus/events.rs

@@ -257,6 +257,16 @@ pub enum BlockValidationResult {
     },
 }
 
+impl BlockValidationResult {
+    pub fn peer(&self) -> Peer {
+        match self {
+            BlockValidationResult::BlockValidated { peer, .. } => peer.clone(),
+            BlockValidationResult::BlockValidationFailed { peer, .. } => peer.clone(),
+            BlockValidationResult::RolledBackTo { peer, .. } => peer.clone(),
+        }
+    }
+}
+
 impl PartialEq for BlockValidationResult {
     fn eq(&self, other: &Self) -> bool {
         match (self, other) {

crates/amaru-consensus/src/consensus/stages/forward_chain.rs

@@ -0,0 +1,118 @@
+// Copyright 2025 PRAGMA
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::consensus::effects::network_effects::{ForwardEvent, ForwardEventEffect};
+use crate::consensus::errors::{ProcessingFailed, ValidationFailed};
+use crate::consensus::events::BlockValidationResult;
+use crate::consensus::span::adopt_current_span;
+use crate::consensus::tip::{AsHeaderTip, HeaderTip};
+use amaru_kernel::Point;
+use amaru_ouroboros_traits::IsHeader;
+use anyhow::anyhow;
+use pure_stage::{Effects, StageRef};
+use tracing::{Level, error, info, instrument, trace};
+
+pub const EVENT_TARGET: &str = "amaru::consensus::forward_chain";
+
+type State = (
+    HeaderTip,
+    StageRef<ValidationFailed>,
+    StageRef<ProcessingFailed>,
+);
+
+/// The forward chain stage forwards the headers of validated blocks to downstream peers, via the
+/// `ForwardEventEffect`. The current node tip is maintained in order to double check that the header
+/// we sent out is correct
+#[instrument(
+    level = Level::TRACE,
+    skip_all,
+    name = "stage.forward_chain",
+)]
+pub async fn stage(
+    state: State,
+    msg: BlockValidationResult,
+    eff: Effects<BlockValidationResult>,
+) -> State {
+    adopt_current_span(&msg);
+    let (mut our_tip, validation_errors, processing_errors) = state;
+    match msg {
+        BlockValidationResult::BlockValidated { peer, header, .. } => {
+            // assert that the new tip is a direct successor of the old tip
+            assert_eq!(header.block_height(), our_tip.block_height() + 1);
+            match header.parent() {
+                Some(parent) => assert_eq!(parent, our_tip.hash()),
+                None => assert_eq!(our_tip, HeaderTip::new(Point::Origin, 0)),
+            }
+            our_tip = header.as_header_tip();
+            trace!(
+                target: EVENT_TARGET,
+                tip = %header.point(),
+                "tip_changed"
+            );
+
+            if let Err(e) = eff
+                .external(ForwardEventEffect::new(
+                    &peer,
+                    ForwardEvent::Forward(header.clone()),
+                ))
+                .await
+            {
+                error!(
+                    target: EVENT_TARGET,
+                    %e,
+                    "failed to send forward event"
+                );
+                eff.send(&processing_errors, ProcessingFailed::new(&peer, anyhow!(e)))
+                    .await
+            }
+        }
+        BlockValidationResult::RolledBackTo {
+            peer,
+            rollback_header,
+            ..
+        } => {
+            info!(
+                target: EVENT_TARGET,
+                point = %rollback_header.point(),
+                "rolled_back_to"
+            );
+
+            our_tip = rollback_header.as_header_tip();
+            if let Err(e) = eff
+                .external(ForwardEventEffect::new(
+                    &peer,
+                    ForwardEvent::Backward(rollback_header.as_header_tip()),
+                ))
+                .await
+            {
+                error!(
+                    target: EVENT_TARGET,
+                    %e,
+                    "failed to send backward event"
+                );
+                eff.send(&processing_errors, ProcessingFailed::new(&peer, anyhow!(e)))
+                    .await
+            }
+        }
+        BlockValidationResult::BlockValidationFailed { point, .. } => {
+            error!(
+                target: EVENT_TARGET,
+                slot = %point.slot_or_default(),
+                hash = %point.hash(),
+                "block validation failed"
+            );
+        }
+    }
+    (our_tip, validation_errors, processing_errors)
+}

crates/amaru-consensus/src/consensus/stages/mod.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 pub mod fetch_block;
+pub mod forward_chain;
 pub mod receive_header;
 pub mod select_chain;
 pub mod store_block;

crates/amaru-consensus/src/consensus/tip.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use amaru_kernel::{HEADER_HASH_SIZE, ORIGIN_HASH, Point, cbor};
+use amaru_kernel::{HEADER_HASH_SIZE, Header, ORIGIN_HASH, Point, cbor};
 use amaru_ouroboros_traits::is_header::IsHeader;
 use pallas_crypto::hash::Hash;
 use serde::{Deserialize, Serialize};
@@ -50,14 +50,21 @@ impl<H: IsHeader + Display> Display for Tip<H> {
     }
 }
 
-impl<H> Tip<H> {
+impl<H: IsHeader> Tip<H> {
     /// Return the header for this Tip if the Tip doesn't represent the Genesis header.
     pub fn to_header(&self) -> Option<&H> {
         match self {
             Tip::Genesis => None,
             Tip::Hdr(h) => Some(h),
         }
     }
+
+    pub fn block_height(&self) -> u64 {
+        match self {
+            Tip::Genesis => 0,
+            Tip::Hdr(h) => h.block_height(),
+        }
+    }
 }
 
 impl<H: IsHeader, C> cbor::encode::Encode<C> for Tip<H>
@@ -137,3 +144,45 @@ impl<H: IsHeader> From<Option<H>> for Tip<H> {
         }
     }
 }
+
+/// This is equivalent to pallas_network::miniprotocols::chainsync::protocol::Tip
+/// but does not incur the dependency on pallas_network.
+///
+/// This can also replace Tip<Header> in places where we only need the Point and block height,
+/// without needing the full Header.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct HeaderTip(Point, u64);
+
+impl Display for HeaderTip {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}.{}", self.1, self.0.hash())
+    }
+}
+
+impl HeaderTip {
+    pub fn new(point: Point, block_height: u64) -> Self {
+        Self(point, block_height)
+    }
+
+    pub fn hash(&self) -> Hash<HEADER_HASH_SIZE> {
+        self.0.hash()
+    }
+
+    pub fn point(&self) -> Point {
+        self.0.clone()
+    }
+
+    pub fn block_height(&self) -> u64 {
+        self.1
+    }
+}
+
+pub trait AsHeaderTip {
+    fn as_header_tip(&self) -> HeaderTip;
+}
+
+impl AsHeaderTip for Header {
+    fn as_header_tip(&self) -> HeaderTip {
+        HeaderTip::new(self.point(), self.block_height())
+    }
+}

crates/amaru/Cargo.toml

@@ -36,6 +36,7 @@ sysinfo.workspace = true
 thiserror.workspace = true
 tokio = { workspace = true, features = [
   "fs",
+  "net",
   "rt",
   "rt-multi-thread",
   "signal",

crates/amaru/src/stages/build_stage_graph.rs

@@ -14,13 +14,14 @@
 
 use amaru_consensus::consensus::effects::store_effects::ResourceHeaderValidation;
 use amaru_consensus::consensus::errors::{ProcessingFailed, ValidationFailed};
-use amaru_consensus::consensus::events::{BlockValidationResult, ChainSyncEvent};
+use amaru_consensus::consensus::events::ChainSyncEvent;
 use amaru_consensus::consensus::stages::select_chain::SelectChain;
 use amaru_consensus::consensus::stages::validate_header::ValidateHeader;
 use amaru_consensus::consensus::stages::{
-    fetch_block, receive_header, select_chain, store_block, store_header, validate_block,
-    validate_header,
+    fetch_block, forward_chain, receive_header, select_chain, store_block, store_header,
+    validate_block, validate_header,
 };
+use amaru_consensus::consensus::tip::HeaderTip;
 use amaru_kernel::protocol_parameters::GlobalParameters;
 use pure_stage::{StageGraph, StageRef};
 
@@ -31,8 +32,8 @@ pub fn build_stage_graph(
     global_parameters: &GlobalParameters,
     header_validation: ResourceHeaderValidation,
     chain_selector: SelectChain,
+    our_tip: HeaderTip,
     network: &mut impl StageGraph,
-    outputs: StageRef<BlockValidationResult>,
 ) -> StageRef<ChainSyncEvent> {
     let receive_header_stage = network.stage("receive_header", receive_header::stage);
     let store_header_stage = network.stage("store_header", store_header::stage);
@@ -41,6 +42,7 @@ pub fn build_stage_graph(
     let fetch_block_stage = network.stage("fetch_block", fetch_block::stage);
     let store_block_stage = network.stage("store_block", store_block::stage);
     let validate_block_stage = network.stage("validate_block", validate_block::stage);
+    let forward_chain_stage = network.stage("forward_chain", forward_chain::stage);
 
     // TODO: currently only validate_header errors, will need to grow into all error handling
     let validation_errors_stage = network.stage(
@@ -65,10 +67,18 @@ pub fn build_stage_graph(
     let validation_errors_stage = network.wire_up(validation_errors_stage, ());
     let processing_errors_stage = network.wire_up(processing_errors_stage, ());
 
+    let forward_chain_stage = network.wire_up(
+        forward_chain_stage,
+        (
+            our_tip,
+            validation_errors_stage.clone().without_state(),
+            processing_errors_stage.clone().without_state(),
+        ),
+    );
     let validate_block_stage = network.wire_up(
         validate_block_stage,
         (
-            outputs,
+            forward_chain_stage.without_state(),
             validation_errors_stage.clone().without_state(),
             processing_errors_stage.clone().without_state(),
         ),