arc-malachitebft-engine 0.7.0-pre

use std::collections::{BTreeSet, HashMap};
use std::marker::PhantomData;

use async_trait::async_trait;
use derive_where::derive_where;
use eyre::eyre;
use libp2p::request_response;
use ractor::{Actor, ActorProcessingErr, ActorRef, RpcReplyPort};
use tokio::task::JoinHandle;
use tracing::{error, info, trace, warn};

use malachitebft_codec as codec;
use malachitebft_core_consensus::{LivenessMsg, SignedConsensusMsg};
use malachitebft_core_types::{
    Context, PolkaCertificate, RoundCertificate, SignedProposal, SignedVote, Validator,
    ValidatorSet,
};
use malachitebft_metrics::SharedRegistry;
use malachitebft_network::handle::CtrlHandle;
use malachitebft_network::{Channel, Config, Event, PeerId};

pub use malachitebft_network::{
    Multiaddr, NetworkIdentity, NetworkStateDump, PersistentPeerError, PersistentPeersOp,
};

use malachitebft_sync::{
    self as sync, InboundRequestId, OutboundRequestId, RawMessage, Request, Response,
};

use crate::consensus::ConsensusCodec;
use crate::sync::SyncCodec;
use crate::util::output_port::{OutputPort, OutputPortSubscriberTrait};
use crate::util::streaming::StreamMessage;

pub type NetworkRef<Ctx> = ActorRef<Msg<Ctx>>;
pub type NetworkMsg<Ctx> = Msg<Ctx>;

pub trait Subscriber<Msg>: OutputPortSubscriberTrait<Msg>
where
    Msg: Clone + ractor::Message,
{
    fn send(&self, msg: Msg);
}

impl<Msg, To> Subscriber<Msg> for ActorRef<To>
where
    Msg: Clone + ractor::Message,
    To: From<Msg> + ractor::Message,
{
    fn send(&self, msg: Msg) {
        if let Err(e) = self.cast(To::from(msg)) {
            error!("Failed to send message to subscriber: {e:?}");
        }
    }
}

pub struct Network<Ctx, Codec> {
    codec: Codec,
    span: tracing::Span,
    marker: PhantomData<Ctx>,
}

impl<Ctx, Codec> Network<Ctx, Codec> {
    pub fn new(codec: Codec, span: tracing::Span) -> Self {
        Self {
            codec,
            span,
            marker: PhantomData,
        }
    }
}

impl<Ctx, Codec> Network<Ctx, Codec>
where
    Ctx: Context,
    Codec: ConsensusCodec<Ctx>,
    Codec: SyncCodec<Ctx>,
    Codec: codec::HasEncodedLen<sync::Response<Ctx>>,
{
    pub async fn spawn(
        identity: NetworkIdentity,
        config: Config,
        metrics: SharedRegistry,
        codec: Codec,
        span: tracing::Span,
    ) -> Result<ActorRef<Msg<Ctx>>, ractor::SpawnErr> {
        let args = Args {
            identity,
            config: config.clone(),
            metrics,
        };

        let (actor_ref, _) = Actor::spawn(None, Self::new(codec, span), args).await?;
        Ok(actor_ref)
    }
}

pub struct Args {
    pub identity: NetworkIdentity,
    pub config: Config,
    pub metrics: SharedRegistry,
}

#[derive_where(Clone, Debug, PartialEq, Eq)]
pub enum NetworkEvent<Ctx: Context> {
    Listening(Multiaddr),

    PeerConnected(PeerId),
    PeerDisconnected(PeerId),

    Vote(PeerId, SignedVote<Ctx>),

    Proposal(PeerId, SignedProposal<Ctx>),
    ProposalPart(PeerId, StreamMessage<Ctx::ProposalPart>),

    PolkaCertificate(PeerId, PolkaCertificate<Ctx>),

    RoundCertificate(PeerId, RoundCertificate<Ctx>),

    Status(PeerId, Status<Ctx>),

    SyncRequest(InboundRequestId, PeerId, Request<Ctx>),
    SyncResponse(OutboundRequestId, PeerId, Option<Response<Ctx>>),
}

pub enum State<Ctx: Context> {
    Stopped,
    Running {
        listen_addrs: Vec<Multiaddr>,
        peers: BTreeSet<PeerId>,
        output_port: OutputPort<NetworkEvent<Ctx>>,
        ctrl_handle: Box<CtrlHandle>,
        recv_task: JoinHandle<()>,
        inbound_requests: HashMap<InboundRequestId, request_response::InboundRequestId>,
    },
}

#[derive_where(Clone, Debug, PartialEq, Eq)]
pub struct Status<Ctx: Context> {
    pub tip_height: Ctx::Height,
    pub history_min_height: Ctx::Height,
}

impl<Ctx: Context> Status<Ctx> {
    pub fn new(tip_height: Ctx::Height, history_min_height: Ctx::Height) -> Self {
        Self {
            tip_height,
            history_min_height,
        }
    }
}

pub enum Msg<Ctx: Context> {
    /// Subscribe this actor to receive gossip events
    Subscribe(Box<dyn Subscriber<NetworkEvent<Ctx>>>),

    /// Publish a signed consensus message
    PublishConsensusMsg(SignedConsensusMsg<Ctx>),

    /// Publish a liveness message
    PublishLivenessMsg(LivenessMsg<Ctx>),

    /// Publish a proposal part
    PublishProposalPart(StreamMessage<Ctx::ProposalPart>),

    /// Broadcast status to all direct peers
    BroadcastStatus(Status<Ctx>),

    /// Send a request to a peer, returning the outbound request ID
    OutgoingRequest(PeerId, Request<Ctx>, RpcReplyPort<OutboundRequestId>),

    /// Send a response for a request to a peer
    OutgoingResponse(InboundRequestId, Response<Ctx>),

    /// Request to dump the current network state
    DumpState(RpcReplyPort<Option<NetworkStateDump>>),

    /// Add or remove a persistent peer at runtime
    UpdatePersistentPeers(
        PersistentPeersOp,
        RpcReplyPort<Result<(), PersistentPeerError>>,
    ),

    /// Update the validator set for the current height
    UpdateValidatorSet(Ctx::ValidatorSet),

    // Event emitted by the gossip layer
    #[doc(hidden)]
    NewEvent(Event),
}

#[async_trait]
impl<Ctx, Codec> Actor for Network<Ctx, Codec>
where
    Ctx: Context,
    Codec: Send + Sync + 'static,
    Codec: codec::Codec<Ctx::ProposalPart>,
    Codec: codec::Codec<SignedConsensusMsg<Ctx>>,
    Codec: codec::Codec<StreamMessage<Ctx::ProposalPart>>,
    Codec: codec::Codec<LivenessMsg<Ctx>>,
    Codec: SyncCodec<Ctx>,
{
    type Msg = Msg<Ctx>;
    type State = State<Ctx>;
    type Arguments = Args;

    async fn pre_start(
        &self,
        myself: ActorRef<Msg<Ctx>>,
        args: Args,
    ) -> Result<Self::State, ActorProcessingErr> {
        let handle = malachitebft_network::spawn(args.identity, args.config, args.metrics).await?;

        let (mut recv_handle, ctrl_handle) = handle.split();

        let recv_task = tokio::spawn(async move {
            while let Some(event) = recv_handle.recv().await {
                if let Err(e) = myself.cast(Msg::NewEvent(event)) {
                    error!("Actor has died, stopping network: {e:?}");
                    break;
                }
            }
        });

        Ok(State::Running {
            listen_addrs: Vec::new(),
            peers: BTreeSet::new(),
            output_port: OutputPort::with_capacity(128),
            ctrl_handle: Box::new(ctrl_handle),
            recv_task,
            inbound_requests: HashMap::new(),
        })
    }

    async fn post_start(
        &self,
        _myself: ActorRef<Msg<Ctx>>,
        _state: &mut State<Ctx>,
    ) -> Result<(), ActorProcessingErr> {
        Ok(())
    }

    #[tracing::instrument(name = "network", parent = &self.span, skip_all)]
    async fn handle(
        &self,
        _myself: ActorRef<Msg<Ctx>>,
        msg: Msg<Ctx>,
        state: &mut State<Ctx>,
    ) -> Result<(), ActorProcessingErr> {
        // We need to handle before deconstructing `state` to always reply.
        if let Msg::DumpState(reply_to) = msg {
            handle_dump_state(state, reply_to).await;
            return Ok(());
        }

        if let Msg::UpdatePersistentPeers(op, reply_to) = msg {
            handle_update_persistent_peers(state, op, reply_to).await;
            return Ok(());
        }

        let State::Running {
            listen_addrs,
            peers,
            output_port,
            ctrl_handle,
            inbound_requests,
            ..
        } = state
        else {
            return Ok(());
        };

        match msg {
            Msg::Subscribe(subscriber) => {
                for addr in listen_addrs.iter() {
                    subscriber.send(NetworkEvent::Listening(addr.clone()));
                }

                for peer in peers.iter() {
                    subscriber.send(NetworkEvent::PeerConnected(*peer));
                }

                subscriber.subscribe_to_port(output_port);
            }

            Msg::PublishConsensusMsg(msg) => match self.codec.encode(&msg) {
                Ok(data) => ctrl_handle.publish(Channel::Consensus, data).await?,
                Err(e) => error!("Failed to encode consensus message: {e:?}"),
            },

            Msg::PublishLivenessMsg(msg) => match self.codec.encode(&msg) {
                Ok(data) => ctrl_handle.publish(Channel::Liveness, data).await?,
                Err(e) => error!("Failed to encode liveness message: {e:?}"),
            },

            Msg::PublishProposalPart(msg) => {
                trace!(
                    stream_id = %msg.stream_id,
                    sequence = %msg.sequence,
                    "Broadcasting proposal part"
                );

                let data = self.codec.encode(&msg);
                match data {
                    Ok(data) => ctrl_handle.publish(Channel::ProposalParts, data).await?,
                    Err(e) => error!("Failed to encode proposal part: {e:?}"),
                }
            }

            Msg::BroadcastStatus(status) => {
                let status = sync::Status {
                    peer_id: ctrl_handle.peer_id(),
                    tip_height: status.tip_height,
                    history_min_height: status.history_min_height,
                };

                let data = self.codec.encode(&status);
                match data {
                    Ok(data) => ctrl_handle.broadcast(Channel::Sync, data).await?,
                    Err(e) => error!("Failed to encode status message: {e:?}"),
                }
            }

            Msg::OutgoingRequest(peer_id, request, reply_to) => {
                let request = self.codec.encode(&request);

                match request {
                    Ok(data) => {
                        let p2p_request_id = ctrl_handle.sync_request(peer_id, data).await?;
                        reply_to.send(OutboundRequestId::new(p2p_request_id))?;
                    }
                    Err(e) => error!("Failed to encode request message: {e:?}"),
                }
            }

            Msg::OutgoingResponse(request_id, response) => {
                let response = self.codec.encode(&response);

                match response {
                    Ok(data) => {
                        let request_id = inbound_requests
                            .remove(&request_id)
                            .ok_or_else(|| eyre!("Unknown inbound request ID: {request_id}"))?;

                        ctrl_handle.sync_reply(request_id, data).await?
                    }
                    Err(e) => {
                        error!(%request_id, "Failed to encode response message: {e:?}");
                        return Ok(());
                    }
                };
            }

            Msg::NewEvent(Event::Listening(addr)) => {
                listen_addrs.push(addr.clone());
                output_port.send(NetworkEvent::Listening(addr));
            }

            Msg::NewEvent(Event::PeerConnected(peer_id)) => {
                peers.insert(peer_id);
                output_port.send(NetworkEvent::PeerConnected(peer_id));
            }

            Msg::NewEvent(Event::PeerDisconnected(peer_id)) => {
                peers.remove(&peer_id);
                output_port.send(NetworkEvent::PeerDisconnected(peer_id));
            }

            Msg::NewEvent(Event::LivenessMessage(Channel::Liveness, from, data)) => {
                let msg = match self.codec.decode(data) {
                    Ok(msg) => msg,
                    Err(e) => {
                        error!(%from, "Failed to decode liveness message: {e:?}");
                        return Ok(());
                    }
                };

                let event = match msg {
                    LivenessMsg::PolkaCertificate(polka_cert) => {
                        NetworkEvent::PolkaCertificate(from, polka_cert)
                    }
                    LivenessMsg::SkipRoundCertificate(round_cert) => {
                        NetworkEvent::RoundCertificate(from, round_cert)
                    }
                    LivenessMsg::Vote(vote) => NetworkEvent::Vote(from, vote),
                };

                output_port.send(event);
            }

            Msg::NewEvent(Event::LivenessMessage(channel, from, _)) => {
                error!(%from, "Unexpected liveness message on {channel} channel");
                return Ok(());
            }

            Msg::NewEvent(Event::ConsensusMessage(Channel::Consensus, from, data)) => {
                let msg = match self.codec.decode(data) {
                    Ok(msg) => msg,
                    Err(e) => {
                        error!(%from, "Failed to decode consensus message: {e:?}");
                        return Ok(());
                    }
                };

                let event = match msg {
                    SignedConsensusMsg::Vote(vote) => NetworkEvent::Vote(from, vote),
                    SignedConsensusMsg::Proposal(proposal) => {
                        NetworkEvent::Proposal(from, proposal)
                    }
                };

                output_port.send(event);
            }

            Msg::NewEvent(Event::ConsensusMessage(Channel::ProposalParts, from, data)) => {
                let msg: StreamMessage<Ctx::ProposalPart> = match self.codec.decode(data) {
                    Ok(stream_msg) => stream_msg,
                    Err(e) => {
                        error!(%from, "Failed to decode stream message: {e:?}");
                        return Ok(());
                    }
                };

                trace!(
                    %from,
                    stream_id = %msg.stream_id,
                    sequence = %msg.sequence,
                    "Received proposal part"
                );

                output_port.send(NetworkEvent::ProposalPart(from, msg));
            }

            Msg::NewEvent(Event::ConsensusMessage(Channel::Sync, from, data)) => {
                let status: sync::Status<Ctx> = match self.codec.decode(data) {
                    Ok(status) => status,
                    Err(e) => {
                        error!(%from, "Failed to decode status message: {e:?}");
                        return Ok(());
                    }
                };

                if from != status.peer_id {
                    error!(%from, %status.peer_id, "Mismatched peer ID in status message");
                    return Ok(());
                }

                trace!(%from, tip_height = %status.tip_height, "Received status");

                output_port.send(NetworkEvent::Status(
                    status.peer_id,
                    Status::new(status.tip_height, status.history_min_height),
                ));
            }

            Msg::NewEvent(Event::ConsensusMessage(channel, from, _)) => {
                error!(%from, "Unexpected consensus message on {channel} channel");
                return Ok(());
            }

            Msg::NewEvent(Event::Sync(raw_msg)) => match raw_msg {
                RawMessage::Request {
                    request_id,
                    peer,
                    body,
                } => {
                    let request = match self.codec.decode(body) {
                        Ok(request) => request,
                        Err(e) => {
                            error!(%peer, "Failed to decode sync request: {e:?}");
                            return Ok(());
                        }
                    };

                    inbound_requests.insert(InboundRequestId::new(request_id), request_id);

                    output_port.send(NetworkEvent::SyncRequest(
                        InboundRequestId::new(request_id),
                        peer,
                        request,
                    ));
                }

                RawMessage::Response {
                    request_id,
                    peer,
                    body,
                } => {
                    let response = match self.codec.decode(body) {
                        Ok(response) => Some(response),
                        Err(e) => {
                            error!(%peer, "Failed to decode sync response: {e:?}");
                            None
                        }
                    };

                    output_port.send(NetworkEvent::SyncResponse(
                        OutboundRequestId::new(request_id),
                        peer,
                        response,
                    ));
                }
            },

            Msg::UpdateValidatorSet(validator_set) => {
                info!(
                    "Updating validator set: {} validators",
                    validator_set.count()
                );
                // Convert ValidatorSet to Vec<ValidatorInfo>
                // Note: We don't pass the Ctx to the network layer
                let validators: Vec<_> = validator_set
                    .iter()
                    .map(|v| malachitebft_network::ValidatorInfo {
                        address: v.address().to_string(),
                        voting_power: v.voting_power(),
                    })
                    .collect();
                ctrl_handle.update_validator_set(validators).await?;
            }

            Msg::DumpState(_) => unreachable!("DumpState handled above to ensure a reply"),
            Msg::UpdatePersistentPeers(_, _) => {
                unreachable!("UpdatePersistentPeers handled above to ensure a reply")
            }
        }

        Ok(())
    }

    async fn post_stop(
        &self,
        _myself: ActorRef<Msg<Ctx>>,
        state: &mut State<Ctx>,
    ) -> Result<(), ActorProcessingErr> {
        let state = std::mem::replace(state, State::Stopped);

        if let State::Running {
            ctrl_handle,
            recv_task,
            ..
        } = state
        {
            ctrl_handle.wait_shutdown().await?;
            recv_task.await?;
        }

        Ok(())
    }
}

async fn handle_dump_state<Ctx>(
    state: &mut State<Ctx>,
    reply_to: RpcReplyPort<Option<NetworkStateDump>>,
) where
    Ctx: Context,
{
    let dump = match state {
        State::Stopped => {
            info!("Dumping network state: not started");
            None
        }
        State::Running { ctrl_handle, .. } => match ctrl_handle.dump_state().await {
            Ok(snapshot) => Some(snapshot),
            Err(error) => {
                error!(%error, "Failed to obtain network dump");
                None
            }
        },
    };

    if let Err(error) = reply_to.send(dump) {
        error!(%error, "Failed to reply with network state dump");
    }
}

async fn handle_update_persistent_peers<Ctx>(
    state: &mut State<Ctx>,
    op: PersistentPeersOp,
    reply_to: RpcReplyPort<Result<(), PersistentPeerError>>,
) where
    Ctx: Context,
{
    fn log_result(result: &Result<(), PersistentPeerError>, op: &PersistentPeersOp) {
        match result {
            Ok(_) => match op {
                PersistentPeersOp::Add(addr) => {
                    info!("Successfully added persistent peer: {addr}");
                }
                PersistentPeersOp::Remove(addr) => {
                    info!("Successfully removed persistent peer: {addr}");
                }
            },
            Err(error) => {
                error!(%error, "Failed to update persistent peers");
            }
        }
    }

    let result = match state {
        State::Stopped => {
            warn!("Cannot update persistent peers: network not started");
            Err(PersistentPeerError::NetworkStopped)
        }
        State::Running { ctrl_handle, .. } => {
            let op_result = match &op {
                PersistentPeersOp::Add(addr) => ctrl_handle.add_persistent_peer(addr.clone()).await,
                PersistentPeersOp::Remove(addr) => {
                    ctrl_handle.remove_persistent_peer(addr.clone()).await
                }
            };

            op_result
                .inspect(|res| log_result(res, &op))
                .unwrap_or_else(|error| {
                    error!(%error, "Internal error: failed to update persistent peers");
                    Err(PersistentPeerError::InternalError(error.to_string()))
                })
        }
    };

    if let Err(error) = reply_to.send(result) {
        error!(%error, "Failed to reply to UpdatePersistentPeers");
    }
}