actor.rs

use crate::{
    ApplicationConfig,
    ingress::{Mailbox, Message},
};
use anyhow::{Context, Result, anyhow};
use commonware_macros::select;
use commonware_runtime::{Clock, ContextCell, Handle, Metrics, Spawner, Storage, spawn_cell};
use commonware_utils::SystemTimeExt;
use commonware_utils::channel::mpsc;
use futures::{
    FutureExt,
    future::{self, Either, try_join},
};
use rand::Rng;
use tokio_util::sync::CancellationToken;

use commonware_consensus::simplex::scheme::Scheme;
use commonware_consensus::types::{Epoch, Epocher, Round, View};
use commonware_cryptography::bls12381::primitives::variant::Variant;
use commonware_cryptography::{PublicKey, Signer};
use std::marker::PhantomData;
#[cfg(feature = "permissioned")]
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use summit_finalizer::FinalizerMailbox;
use tracing::{debug, info, warn};

#[cfg(feature = "prom")]
use metrics::{counter, histogram};
use summit_syncer::ingress::mailbox::Mailbox as SyncerMailbox;
use summit_types::{Block, BlockAuxData, Digest, EngineClient};

pub struct Actor<
    R: Storage + Metrics + Clock + Spawner + governor::clock::Clock + Rng,
    C: EngineClient,
    S: Scheme<Digest>,
    P: PublicKey,
    K: Signer,
    V: Variant,
    ES: Epocher,
> {
    context: ContextCell<R>,
    mailbox: mpsc::Receiver<Message>,
    engine_client: C,
    built_block: Arc<Mutex<Option<(Block, Round)>>>,
    genesis_hash: [u8; 32],
    epocher: ES,
    cancellation_token: CancellationToken,
    #[cfg(feature = "permissioned")]
    paused: Arc<AtomicBool>,
    _scheme_marker: PhantomData<S>,
    _key_marker: PhantomData<P>,
    _signer_marker: PhantomData<K>,
    _variant_marker: PhantomData<V>,
}

impl<
    R: Storage + Metrics + Clock + Spawner + governor::clock::Clock + Rng,
    C: EngineClient,
    S: Scheme<Digest>,
    P: PublicKey,
    K: Signer,
    V: Variant,
    ES: Epocher,
> Actor<R, C, S, P, K, V, ES>
{
    pub async fn new(context: R, cfg: ApplicationConfig<C, ES>) -> (Self, Mailbox<P>) {
        let (tx, rx) = mpsc::channel(cfg.mailbox_size);

        let genesis_hash = cfg.genesis_hash;

        (
            Self {
                context: ContextCell::new(context),
                mailbox: rx,
                engine_client: cfg.engine_client,
                built_block: Arc::new(Mutex::new(None)),
                genesis_hash,
                epocher: cfg.epocher,
                cancellation_token: cfg.cancellation_token,
                #[cfg(feature = "permissioned")]
                paused: cfg.paused,
                _scheme_marker: PhantomData,
                _key_marker: PhantomData,
                _signer_marker: PhantomData,
                _variant_marker: PhantomData,
            },
            Mailbox::new(tx),
        )
    }

    pub fn start(
        mut self,
        syncer: SyncerMailbox<S, Block>,
        finalizer: FinalizerMailbox<S, Block>,
    ) -> Handle<()> {
        spawn_cell!(self.context, self.run(syncer, finalizer).await)
    }

    pub async fn run(
        mut self,
        mut syncer: SyncerMailbox<S, Block>,
        mut finalizer: FinalizerMailbox<S, Block>,
    ) {
        let rand_id: u8 = rand::random();
        let mut signal = self.context.stopped().fuse();
        let cancellation_token = self.cancellation_token.clone();
        loop {
            select! {
                message = self.mailbox.recv() => {
                    let Some(message) = message else {
                        break;
                    };
                    match message {
                        Message::Genesis { response, epoch } => {
                            if epoch.get() == 0 {
                                let _ = response.send(self.genesis_hash.into());
                            } else {
                                let epoch_genesis_hash = finalizer
                                    .get_epoch_genesis_hash(epoch.get())
                                    .await
                                    .await
                                    .expect("failed to get epoch genesis hash from finalizer");
                                let _ = response.send(epoch_genesis_hash.into());
                            }
                        }
                        Message::Propose {
                            round,
                            parent,
                            mut response,
                        } => {
                            #[cfg(feature = "permissioned")]
                            if self.paused.load(Ordering::Relaxed) {
                                warn!("consensus paused, skipping proposal for round {round}");
                                continue;
                            }

                            debug!("{rand_id} application: Handling message Propose for round {} (epoch {}, view {}), parent view: {}",
                                round, round.epoch(), round.view(), parent.0);

                            let built = self.built_block.clone();
                            #[cfg(feature = "prom")]
                            let proposal_start = std::time::Instant::now();

                            select! {
                                    res = self.handle_proposal(parent, &mut syncer, &mut finalizer, round) => {
                                        match res {
                                            Ok(block) => {
                                                // store block
                                                let digest = block.digest();
                                                let height = block.height();
                                                let tx_count = block.payload.payload_inner.payload_inner.transactions.len();
                                                {
                                                    let mut built = built.lock().expect("locked poisoned");
                                                    *built = Some((block.clone(), round));
                                                }

                                                info!(
                                                    height,
                                                    epoch = round.epoch().get(),
                                                    view = round.view().get(),
                                                    tx_count,
                                                    "proposed block"
                                                );

                                                // send block to syncer for caching and broadcasting
                                                syncer.proposed(round, block).await;

                                                // send digest to consensus
                                                let _ = response.send(digest);
                                            },
                                            Err(e) => warn!("Failed to create a block for round {round}: {e}")
                                        }
                                    },
                                    _ = response.closed() => {
                                        // simplex dropped receiver
                                        #[cfg(feature = "prom")]
                                        {
                                            let elapsed = proposal_start.elapsed();
                                            warn!(
                                                round = ?round,
                                                parent_view = parent.0.get(),
                                                parent_digest = ?parent.1,
                                                elapsed_ms = elapsed.as_millis(),
                                                "proposal aborted - consensus timed out waiting for block (possible notarize-nullify race)"
                                            );
                                            counter!("proposal_timeout_total").increment(1);
                                            histogram!("proposal_timeout_elapsed_ms").record(elapsed.as_millis() as f64);
                                        }

                                        #[cfg(not(feature = "prom"))]
                                        warn!(
                                            round = ?round,
                                            parent_view = parent.0.get(),
                                            parent_digest = ?parent.1,
                                            "proposal aborted - consensus timed out waiting for block (possible notarize-nullify race)"
                                        );
                                    }
                            }
                        }
                        Message::Broadcast { payload: _ } => {
                            #[cfg(feature = "permissioned")]
                            if self.paused.load(Ordering::Relaxed) {
                                warn!("consensus paused, skipping broadcast");
                                continue;
                            }

                            info!("{rand_id} Handling message Broadcast");

                            let built_block = self.built_block.lock().expect("poisoned lock").take();

                            if let Some((block, round)) = built_block {
                                syncer.proposed(round, block).await;
                            } else {
                                warn!("Asked to broadcast a block without one built");
                            }
                        }

                        Message::Verify {
                            round,
                            parent,
                            payload,
                            mut response,
                        } => {
                            #[cfg(feature = "permissioned")]
                            if self.paused.load(Ordering::Relaxed) {
                                warn!("consensus paused, rejecting verify for round {round}");
                                let _ = response.send(false);
                                continue;
                            }

                            debug!("{rand_id} application: Handling message Verify for round {} (epoch {}, view {}), parent view: {}",
                                round, round.epoch(), round.view(), parent.0);

                            // Subscribe to blocks (will wait for them if not available)
                            let parent_request = if parent.1 == self.genesis_hash.into() {
                                Either::Left(future::ready(Ok(Block::genesis(self.genesis_hash))))
                            } else {
                                let parent_round = if parent.0.get() == 0 {
                                    // Parent view is 0, which means that this is the first block of the epoch
                                    None
                                } else {
                                    Some(Round::new(round.epoch(), parent.0))
                                };
                                Either::Right(
                                    syncer
                                        .subscribe(parent_round, parent.1)
                                        .await,
                                )
                            };
                            let block_request = syncer.subscribe(Some(round), payload).await;

                            // Wait for the blocks to be available or the request to be canceled in a separate task (to
                            // continue processing other messages)
                            self.context.with_label("verify").spawn({
                                let mut syncer = syncer.clone();
                                let mut finalizer_clone = finalizer.clone();
                                let epocher = self.epocher.clone();
                                move |context| async move {
                                    let requester = try_join(parent_request, block_request);
                                    select! {
                                        result = requester => {
                                            let (parent, block) = result.unwrap();

                                            let parent_digest = parent.digest();
                                            let parent_height = parent.height();

                                            // Wait for parent block to be executed by finalizer
                                            // This ensures the parent's state is available for aux_data
                                            let parent_executed = finalizer_clone
                                                .notify_at_height(parent_height, parent_digest)
                                                .await
                                                .await
                                                .unwrap_or(false);

                                            if !parent_executed {
                                                warn!(
                                                    ?round,
                                                    parent_height,
                                                    ?parent_digest,
                                                    "parent block not executed by finalizer"
                                                );
                                                let _ = response.send(false);
                                                return;
                                            }

                                            // Request aux data for the block we're verifying
                                            #[cfg(feature = "prom")]
                                            let aux_data_start = std::time::Instant::now();
                                            let maybe_aux_data = finalizer_clone
                                                .get_aux_data(parent_height + 1, parent_digest)
                                                .await
                                                .await
                                                .expect("Finalizer dropped");

                                            if let Some(aux_data) = maybe_aux_data {
                                                #[cfg(feature = "prom")]
                                                {
                                                    let aux_data_duration = aux_data_start.elapsed().as_millis() as f64;
                                                    histogram!("handle_verify_aux_data_duration_millis").record(aux_data_duration);
                                                }

                                                let now_millis = context.current().epoch_millis();
                                                if handle_verify(&block, parent, &epocher, &aux_data, now_millis) {
                                                    // persist valid block
                                                    syncer.verified(round, block).await;

                                                    // respond
                                                    let _ = response.send(true);
                                                } else {
                                                    info!("Unsuccessful vote for round {round} because the block is invalid");
                                                    let _ = response.send(false);
                                                }
                                            } else {
                                                info!(
                                                    "Unsuccessful vote for round {round} because of an outdated height notification",
                                                );
                                                let _ = response.send(false);
                                            }
                                        },
                                        _ = response.closed() => {
                                            warn!("verify aborted for round {round}");
                                        }
                                    }
                                }
                            });
                        }
                    }
                },
                _ = cancellation_token.cancelled() => {
                    info!("application received cancellation signal, exiting");
                    break;
                },
                sig = &mut signal => {
                    info!("runtime terminated, shutting down application: {}", sig.unwrap());
                    break;
                }
            }
        }
    }

    async fn handle_proposal(
        &mut self,
        parent: (View, Digest),
        syncer: &mut SyncerMailbox<S, Block>,
        finalizer: &mut FinalizerMailbox<S, Block>,
        round: Round,
    ) -> Result<Block> {
        #[cfg(feature = "prom")]
        let proposal_start = std::time::Instant::now();

        // STEP 1: Get the parent block
        debug!(
            ?round,
            parent_view = parent.0.get(),
            parent_digest = ?parent.1,
            "proposal step 1: fetching parent block"
        );
        #[cfg(feature = "prom")]
        let parent_fetch_start = std::time::Instant::now();
        let parent_block = if parent.1 == self.genesis_hash.into() {
            Either::Left(future::ready(Ok(Block::genesis(self.genesis_hash))))
        } else {
            let parent_round = if parent.0.get() == 0 {
                // Parent view is 0, which means that this is the first block of the epoch
                None
            } else {
                Some(Round::new(round.epoch(), parent.0))
            };
            Either::Right(
                syncer
                    .subscribe(parent_round, parent.1)
                    .await
                    .map(|x| x.context("")),
            )
        };
        let parent_block = parent_block.await.expect("sender dropped");

        #[cfg(feature = "prom")]
        {
            let parent_fetch_duration = parent_fetch_start.elapsed().as_millis() as f64;
            histogram!("handle_proposal_parent_fetch_duration_millis")
                .record(parent_fetch_duration);
        }

        // STEP 2: Wait for finalizer notification
        debug!(
            ?round,
            parent_height = parent_block.height(),
            parent_digest = ?parent_block.digest(),
            "proposal step 2: waiting for finalizer notification"
        );
        #[cfg(feature = "prom")]
        let finalizer_wait_start = std::time::Instant::now();
        // now that we have the parent additionally await for that to be executed by the finalizer
        let parent_height = parent_block.height();
        let parent_digest = parent_block.digest();
        let rx = finalizer
            .notify_at_height(parent_height, parent_digest)
            .await;
        // await for notification
        if !rx.await.expect("Finalizer dropped") {
            debug!(
                "Aborting block proposal for epoch {} and height {} because of an outdated height notification",
                round.epoch().get(),
                parent_height + 1,
            );
            return Err(anyhow!(
                "Aborting block proposal for epoch {} and height {} because of an outdated height notification",
                round.epoch().get(),
                parent_height + 1,
            ));
        }
        #[cfg(feature = "prom")]
        {
            let finalizer_wait_duration = finalizer_wait_start.elapsed().as_millis() as f64;
            histogram!("handle_proposal_finalizer_wait_duration_millis")
                .record(finalizer_wait_duration);
        }

        // STEP 3: Request aux data (withdrawals, checkpoint hash, header hash)
        debug!(
            ?round,
            parent_height, "proposal step 3: requesting aux data"
        );
        #[cfg(feature = "prom")]
        let aux_data_start = std::time::Instant::now();
        let maybe_aux_data = finalizer
            .get_aux_data(parent_height + 1, parent_digest)
            .await
            .await
            .expect("Finalizer dropped");

        let Some(aux_data) = maybe_aux_data else {
            debug!(
                "Aborting block proposal for epoch {} and height {} because of an outdated aux data request",
                round.epoch().get(),
                parent_height + 1,
            );
            return Err(anyhow!(
                "Aborting block proposal for epoch {} and height {} because of an outdated aux data request",
                round.epoch().get(),
                parent_height + 1,
            ));
        };

        #[cfg(feature = "prom")]
        {
            let aux_data_duration = aux_data_start.elapsed().as_millis() as f64;
            histogram!("handle_proposal_aux_data_duration_millis").record(aux_data_duration);
        }

        if aux_data.epoch != round.epoch().get() {
            // This might happen because the finalizer notifies the orchestrator at the end of an
            // epoch to shut down Simplex. While Simplex is being shutdown, it will still continue to produce blocks.
            return Err(anyhow!(
                "Aborting block proposal for height {} and epoch {}. Current epoch is {}",
                parent_height + 1,
                round.epoch().get(),
                aux_data.epoch,
            ));
        }

        // Special case: If the parent block is the last block in the epoch,
        // re-propose it as to not produce any blocks that will be cut out
        // by the epoch transition.
        let last_in_epoch = self
            .epocher
            .last(Epoch::new(aux_data.epoch))
            .expect("epoch should exist");
        if parent_block.height() == last_in_epoch.get() {
            debug!(round = ?round, digest = ?parent_block.digest(), "re-proposed parent block at epoch boundary");
            return Ok(parent_block);
        }

        let pending_withdrawals = aux_data.withdrawals;
        let checkpoint_hash = aux_data.checkpoint_hash;

        let mut current = self.context.current().epoch_millis();
        if current <= parent_block.timestamp() {
            current = parent_block.timestamp() + 1;
        }

        // STEP 4: Start building block (Engine Client)
        debug!(
            ?round,
            parent_height,
            epoch = aux_data.epoch,
            "proposal step 4: building block via engine client"
        );
        #[cfg(feature = "prom")]
        let start_building_start = std::time::Instant::now();

        //  aux_data.forkchoice.head_block_hash = parent_block.eth_block_hash().into();

        // Add pending withdrawals to the block
        let withdrawals = pending_withdrawals.into_iter().map(|w| w.inner).collect();
        let payload_id = {
            #[cfg(feature = "bench")]
            {
                self.engine_client
                    .start_building_block(
                        aux_data.forkchoice,
                        current,
                        withdrawals,
                        aux_data.suggested_fee_recipient,
                        None,
                        parent_block.height(),
                    )
                    .await
            }
            #[cfg(not(feature = "bench"))]
            {
                self.engine_client
                    .start_building_block(
                        aux_data.forkchoice,
                        current,
                        withdrawals,
                        aux_data.suggested_fee_recipient,
                        Some(aux_data.state_root.into()),
                    )
                    .await
            }
        }
        .ok_or(anyhow!("Unable to build payload"))?;

        #[cfg(feature = "prom")]
        {
            let start_building_duration = start_building_start.elapsed().as_millis() as f64;
            histogram!("handle_proposal_start_building_duration_millis")
                .record(start_building_duration);
        }

        self.context.sleep(Duration::from_millis(50)).await;

        // STEP 5: Get payload (Engine Client)
        #[cfg(feature = "prom")]
        let get_payload_start = std::time::Instant::now();
        let payload_envelope = self.engine_client.get_payload(payload_id).await;
        #[cfg(feature = "prom")]
        {
            let get_payload_duration = get_payload_start.elapsed().as_millis() as f64;
            histogram!("handle_proposal_get_payload_duration_millis").record(get_payload_duration);
        }

        // STEP 6: Compute block digest
        #[cfg(feature = "prom")]
        let compute_digest_start = std::time::Instant::now();

        let block = Block::compute_digest(
            parent_block.digest(),
            parent_block.height() + 1,
            current,
            payload_envelope.envelope_inner.execution_payload,
            payload_envelope.execution_requests.to_vec(),
            payload_envelope.envelope_inner.block_value,
            round.epoch().get(),
            round.view().get(),
            checkpoint_hash,
            aux_data.header_hash,
            aux_data.added_validators,
            aux_data.removed_validators,
            aux_data.state_root,
        );

        #[cfg(feature = "prom")]
        {
            let compute_digest_duration = compute_digest_start.elapsed().as_millis() as f64;
            histogram!("handle_proposal_compute_digest_duration_millis")
                .record(compute_digest_duration);
        }

        #[cfg(feature = "prom")]
        {
            let proposal_duration = proposal_start.elapsed().as_millis() as f64;
            histogram!("handle_proposal_duration_millis").record(proposal_duration);
        }
        Ok(block)
    }
}

impl<
    R: Storage + Metrics + Clock + Spawner + governor::clock::Clock + Rng,
    C: EngineClient,
    S: Scheme<Digest>,
    P: PublicKey,
    K: Signer,
    V: Variant,
    ES: Epocher,
> Drop for Actor<R, C, S, P, K, V, ES>
{
    fn drop(&mut self) {
        self.cancellation_token.cancel();
    }
}

fn handle_verify<ES: Epocher>(
    block: &Block,
    parent: Block,
    epocher: &ES,
    aux_data: &BlockAuxData,
    now_millis: u64,
) -> bool {
    // You can only re-propose the same block if it's the last height in the epoch.
    if parent.digest() == block.digest() {
        let last_in_epoch = epocher
            .last(Epoch::new(aux_data.epoch))
            .expect("epoch should exist");
        return block.height() == last_in_epoch.get();
    }

    // Basic structural validation
    if block.parent() != parent.digest() {
        warn!("block parent mismatch");
        return false;
    }
    if block.height() != parent.height() + 1 {
        warn!("block height mismatch");
        return false;
    }
    if block.timestamp() <= parent.timestamp() {
        warn!("block timestamp not increasing");
        return false;
    }
    if block.timestamp() > now_millis + aux_data.allowed_timestamp_future_ms {
        warn!(
            block_timestamp = block.timestamp(),
            now_millis,
            allowed_timestamp_future_ms = aux_data.allowed_timestamp_future_ms,
            "block timestamp too far in the future"
        );
        return false;
    }

    // Validate consensus trie state root
    if block.header.parent_beacon_block_root != aux_data.state_root {
        warn!(
            expected = ?aux_data.state_root,
            actual = ?block.header.parent_beacon_block_root,
            "parent_beacon_block_root mismatch"
        );
        return false;
    }

    // Validate checkpoint_hash (None means [0; 32], matching Block::compute_digest)
    let expected_checkpoint_hash: Digest =
        aux_data.checkpoint_hash.unwrap_or_else(|| [0; 32].into());
    if block.header.checkpoint_hash != expected_checkpoint_hash {
        warn!(
            expected = ?expected_checkpoint_hash,
            actual = ?block.header.checkpoint_hash,
            "checkpoint_hash mismatch"
        );
        return false;
    }

    // Validate added_validators
    if block.header.added_validators != aux_data.added_validators {
        warn!(
            expected_count = aux_data.added_validators.len(),
            actual_count = block.header.added_validators.len(),
            "added_validators mismatch"
        );
        return false;
    }

    // Validate removed_validators
    if block.header.removed_validators != aux_data.removed_validators {
        warn!(
            expected_count = aux_data.removed_validators.len(),
            actual_count = block.header.removed_validators.len(),
            "removed_validators mismatch"
        );
        return false;
    }

    // Validate withdrawals
    let expected_withdrawals: Vec<_> = aux_data.withdrawals.iter().map(|w| w.inner).collect();
    let actual_withdrawals: &[_] = &block.payload.payload_inner.withdrawals;
    if actual_withdrawals != expected_withdrawals.as_slice() {
        warn!(
            expected_count = expected_withdrawals.len(),
            actual_count = actual_withdrawals.len(),
            "withdrawals mismatch"
        );
        return false;
    }

    true
}