iota_core/epoch/

randomness.rs

// Copyright (c) Mysten Labs, Inc.
// Modifications Copyright (c) 2024 IOTA Stiftung
// SPDX-License-Identifier: Apache-2.0

use std::{
    collections::{BTreeMap, HashMap},
    sync::{Arc, Weak},
    time::Instant,
};

use anemo::PeerId;
use fastcrypto::{
    encoding::{Encoding, Hex},
    error::{FastCryptoError, FastCryptoResult},
    groups::bls12381,
    serde_helpers::ToFromByteArray,
    traits::{KeyPair, ToFromBytes},
};
use fastcrypto_tbls::{dkg_v1, dkg_v1::Output, nodes, nodes::PartyId};
use futures::{StreamExt, stream::FuturesUnordered};
use iota_common::{debug_fatal, fatal};
use iota_macros::{fail_point, fail_point_if};
use iota_network::randomness;
use iota_sdk_types::RandomnessRound;
use iota_types::{
    base_types::{AuthorityName, CommitRound},
    committee::{Committee, EpochId, StakeUnit},
    crypto::{AuthorityKeyPair, RandomnessSignature},
    error::{IotaError, IotaResult},
    iota_system_state::epoch_start_iota_system_state::EpochStartSystemStateTrait,
    messages_consensus::{ConsensusTransaction, VersionedDkgConfirmation, VersionedDkgMessage},
};
use parking_lot::Mutex;
use rand08::{
    SeedableRng,
    rngs::{OsRng, StdRng},
};
use serde::{Deserialize, Serialize};
use tokio::{sync::OnceCell, task::JoinHandle};
use tracing::{debug, error, info, warn};
use typed_store::Map;

use crate::{
    authority::{
        authority_per_epoch_store::{
            AuthorityPerEpochStore, consensus_quarantine::ConsensusCommitOutput,
        },
        epoch_start_configuration::EpochStartConfigTrait,
    },
    consensus_adapter::SubmitToConsensus,
};

/// The epoch UNIX timestamp in milliseconds
pub type CommitTimestampMs = u64;

type PkG = bls12381::G2Element;
type EncG = bls12381::G2Element;

pub const SINGLETON_KEY: u64 = 0;

// Wrappers for DKG messages (to simplify upgrades).

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum VersionedProcessedMessage {
    V1(dkg_v1::ProcessedMessage<PkG, EncG>),
}

impl VersionedProcessedMessage {
    pub fn sender(&self) -> PartyId {
        match self {
            VersionedProcessedMessage::V1(msg) => msg.message.sender,
        }
    }

    pub fn unwrap_v1(self) -> dkg_v1::ProcessedMessage<PkG, EncG> {
        match self {
            VersionedProcessedMessage::V1(msg) => msg,
        }
    }

    pub fn process(
        party: Arc<dkg_v1::Party<PkG, EncG>>,
        message: VersionedDkgMessage,
    ) -> FastCryptoResult<VersionedProcessedMessage> {
        // All inputs are verified in add_message, so we can assume they are of the
        // correct version.
        let processed = party.process_message(message.unwrap_v1(), &mut rand08::thread_rng())?;
        Ok(VersionedProcessedMessage::V1(processed))
    }

    pub fn merge(
        party: Arc<dkg_v1::Party<PkG, EncG>>,
        messages: Vec<Self>,
    ) -> FastCryptoResult<(VersionedDkgConfirmation, VersionedUsedProcessedMessages)> {
        // All inputs were created by this validator, so we can assume they are of the
        // correct version.
        let (conf, msgs) = party.merge(
            &messages
                .into_iter()
                .map(|vm| vm.unwrap_v1())
                .collect::<Vec<_>>(),
        )?;
        Ok((
            VersionedDkgConfirmation::V1(conf),
            VersionedUsedProcessedMessages::V1(msgs),
        ))
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum VersionedUsedProcessedMessages {
    V1(dkg_v1::UsedProcessedMessages<PkG, EncG>),
}

impl VersionedUsedProcessedMessages {
    fn complete_dkg<'a, Iter: Iterator<Item = &'a VersionedDkgConfirmation>>(
        &self,
        party: Arc<dkg_v1::Party<PkG, EncG>>,
        confirmations: Iter,
    ) -> FastCryptoResult<Output<PkG, EncG>> {
        // All inputs are verified in add_confirmation, so we can assume they are of the
        // correct version.
        let rng = &mut StdRng::from_rng(OsRng).expect("RNG construction should not fail");
        let VersionedUsedProcessedMessages::V1(msg) = self;
        party.complete(
            msg,
            &confirmations
                .map(|vm| vm.unwrap_v1())
                .cloned()
                .collect::<Vec<_>>(),
            rng,
        )
    }
}

/// Error from advancing DKG. Runtime shutdown is a separate variant so that a
/// cancelled blocking task is never handled like a failed DKG step.
#[derive(Debug)]
pub(crate) enum AdvanceDkgError {
    Iota(IotaError),
    /// The runtime is shutting down and cancelled DKG message processing. The
    /// commit being processed must be abandoned without persisting its
    /// output.
    RuntimeShutdown,
}

impl From<IotaError> for AdvanceDkgError {
    fn from(e: IotaError) -> Self {
        Self::Iota(e)
    }
}

// State machine for randomness DKG and generation.
//
// DKG protocol:
// 1. This validator sends out a `Message` to all other validators.
// 2. Once sufficient valid `Message`s are received from other validators via consensus and
//    processed, this validator sends out a `Confirmation` to all other validators.
// 3. Once sufficient `Confirmation`s are received from other validators via consensus and
//    processed, they are combined to form a public VSS key and local private key shares.
// 4. Randomness generation begins.
//
// Randomness generation:
// 1. For each new round, AuthorityPerEpochStore eventually calls `generate_randomness`.
// 2. This kicks off a process in RandomnessEventLoop to send partial signatures for the new round
//    to all other validators.
// 3. Once enough partial signatures for the round are collected, a RandomnessStateUpdate
//    transaction is generated and injected into the TransactionManager.
// 4. Once the RandomnessStateUpdate transaction is seen in a certified checkpoint,
//    `notify_randomness_in_checkpoint` is called to complete the round and stop sending partial
//    signatures for it.
pub struct RandomnessManager {
    epoch_store: Weak<AuthorityPerEpochStore>,
    epoch: EpochId,
    consensus_adapter: Box<dyn SubmitToConsensus>,
    network_handle: randomness::Handle,
    authority_info: HashMap<AuthorityName, (PeerId, PartyId)>,

    // State for DKG.
    dkg_start_time: OnceCell<Instant>,
    party: Arc<dkg_v1::Party<PkG, EncG>>,
    enqueued_messages: BTreeMap<PartyId, JoinHandle<Option<VersionedProcessedMessage>>>,
    processed_messages: BTreeMap<PartyId, VersionedProcessedMessage>,
    used_messages: OnceCell<VersionedUsedProcessedMessages>,
    confirmations: BTreeMap<PartyId, VersionedDkgConfirmation>,
    dkg_output: OnceCell<Option<dkg_v1::Output<PkG, EncG>>>,
    dkg_timeout_reported: bool,

    // State for randomness generation.
    next_randomness_round: RandomnessRound,
    highest_completed_round: Arc<Mutex<Option<RandomnessRound>>>,
}

impl RandomnessManager {
    /// Returns an error in case of invalid input or other failure to
    /// initialize DKG.
    pub async fn try_new(
        epoch_store_weak: Weak<AuthorityPerEpochStore>,
        consensus_adapter: Box<dyn SubmitToConsensus>,
        network_handle: randomness::Handle,
        authority_key_pair: &AuthorityKeyPair,
    ) -> IotaResult<Self> {
        let epoch_store = epoch_store_weak
            .upgrade()
            .ok_or_else(|| IotaError::Unknown("AuthorityPerEpochStore already gone".to_string()))?;
        let tables = epoch_store.tables()?;
        let protocol_config = epoch_store.protocol_config();

        let name: AuthorityName = authority_key_pair.public().into();
        let committee = epoch_store.committee();
        let info = RandomnessManager::randomness_dkg_info_from_committee(committee);
        if tracing::enabled!(tracing::Level::DEBUG) {
            // Log first few entries in DKG info for debugging.
            for (id, name, pk, stake) in info.iter().filter(|(id, _, _, _)| *id < 3) {
                let pk_bytes = pk.as_element().to_byte_array();
                debug!(
                    "random beacon: DKG info: id={id}, stake={stake}, name={name}, pk={pk_bytes:x?}"
                );
            }
        }
        let authority_ids: HashMap<_, _> =
            info.iter().map(|(id, name, _, _)| (*name, *id)).collect();
        let authority_peer_ids = epoch_store
            .epoch_start_config()
            .epoch_start_state()
            .get_authority_names_to_peer_ids();
        let authority_info: HashMap<AuthorityName, (PeerId, PartyId)> = authority_ids
            .into_iter()
            .map(|(name, id)| {
                authority_peer_ids
                    .get(&name)
                    .map(|peer_id| (name, (*peer_id, id)))
            })
            .collect::<Option<_>>()
            .ok_or_else(|| {
                IotaError::Unknown("authority name missing from peer_ids".to_string())
            })?;
        let nodes = info
            .iter()
            .map(|(id, _, pk, stake)| {
                Some(nodes::Node::<EncG> {
                    id: *id,
                    pk: pk.clone(),
                    weight: (*stake).try_into().ok()?,
                })
            })
            .collect::<Option<Vec<_>>>()
            .ok_or_else(|| IotaError::Unknown("validator stake does not fit in u16".to_string()))?;
        let validity_threshold: u16 = committee.validity_threshold().try_into().map_err(|_| {
            IotaError::Unknown(format!(
                "committee validity threshold {} does not fit in u16",
                committee.validity_threshold()
            ))
        })?;
        let reduction_lower_bound: u16 = protocol_config
            .random_beacon_reduction_lower_bound()
            .try_into()
            .map_err(|_| {
                IotaError::Unknown(format!(
                    "random beacon reduction lower bound {} does not fit in u16",
                    protocol_config.random_beacon_reduction_lower_bound()
                ))
            })?;
        let (nodes, t) = nodes::Nodes::new_reduced(
            nodes,
            validity_threshold,
            protocol_config.random_beacon_reduction_allowed_delta(),
            reduction_lower_bound,
        )
        .map_err(|err| IotaError::Unknown(format!("error while initializing Nodes: {err:?}")))?;
        let total_weight = nodes.total_weight();
        let num_nodes = nodes.num_nodes();
        let prefix_str = format!(
            "dkg {} {}",
            Hex::encode(epoch_store.get_chain_identifier().as_bytes()),
            committee.epoch()
        );
        let randomness_private_key = bls12381::Scalar::from_byte_array(
            authority_key_pair
                .copy()
                .private()
                .as_bytes()
                .try_into()
                .expect("key length should match"),
        )
        .expect("should work to convert BLS key to Scalar");
        let party = dkg_v1::Party::<PkG, EncG>::new(
            fastcrypto_tbls::ecies_v1::PrivateKey::<bls12381::G2Element>::from(
                randomness_private_key,
            ),
            nodes,
            t,
            fastcrypto_tbls::random_oracle::RandomOracle::new(prefix_str.as_str()),
            &mut rand08::thread_rng(),
        )
        .map_err(|err| IotaError::Unknown(format!("error while initializing Party: {err:?}")))?;
        info!(
            "random beacon: state initialized with authority={name}, total_weight={total_weight}, t={t}, num_nodes={num_nodes}, oracle initial_prefix={prefix_str:?}",
        );

        // Reset DKG metrics to a clean baseline for this epoch. The success/failure
        // branches below restore the right value if a persisted verdict exists;
        // pending starts at 0 across the board.
        epoch_store.metrics.epoch_random_beacon_dkg_failed.set(0);
        epoch_store
            .metrics
            .epoch_random_beacon_dkg_completed_after_timeout
            .set(0);
        epoch_store
            .metrics
            .epoch_random_beacon_dkg_num_shares
            .set(0);

        // Load existing data from store.
        let highest_completed_round = tables
            .randomness_highest_completed_round
            .get(&SINGLETON_KEY)
            .expect("typed_store should not fail");
        let mut rm = RandomnessManager {
            epoch_store: epoch_store_weak,
            epoch: committee.epoch(),
            consensus_adapter,
            network_handle: network_handle.clone(),
            authority_info,
            dkg_start_time: OnceCell::new(),
            party: Arc::new(party),
            enqueued_messages: BTreeMap::new(),
            processed_messages: BTreeMap::new(),
            used_messages: OnceCell::new(),
            confirmations: BTreeMap::new(),
            dkg_output: OnceCell::new(),
            dkg_timeout_reported: false,
            next_randomness_round: RandomnessRound::new(0),
            highest_completed_round: Arc::new(Mutex::new(highest_completed_round)),
        };
        // Load the persisted DKG verdict. Read the v2 table first so a terminal
        // failure (`Some(None)`) is restored across restarts; fall back to the
        // legacy `dkg_output` table for stores written before this fix.
        let dkg_output = match tables
            .dkg_output_v2
            .get(&SINGLETON_KEY)
            .expect("typed_store should not fail")
        {
            Some(persisted) => Some(persisted),
            None => tables
                .dkg_output
                .get(&SINGLETON_KEY)
                .expect("typed_store should not fail")
                .map(Some),
        };
        match dkg_output {
            Some(Some(dkg_output)) => {
                info!(
                    "random beacon: loaded existing DKG output for epoch {}",
                    committee.epoch()
                );
                epoch_store
                    .metrics
                    .epoch_random_beacon_dkg_num_shares
                    .set(dkg_output.shares.as_ref().map_or(0, |shares| shares.len()) as i64);
                rm.dkg_output
                    .set(Some(dkg_output.clone()))
                    .expect("setting new OnceCell should succeed");
                network_handle.update_epoch(
                    committee.epoch(),
                    rm.authority_info.clone(),
                    dkg_output,
                    rm.party.t(),
                    highest_completed_round,
                );
            }
            Some(None) => {
                // Terminal-failure verdict was persisted to `dkg_output_v2`. This is
                // only written when late DKG completion is disabled. Restore it so DKG
                // isn't re-run and randomness stays disabled for the epoch.
                error!(
                    "random beacon: loaded failed DKG for epoch {}. Randomness disabled for this epoch. All randomness-using transactions will fail.",
                    committee.epoch()
                );
                epoch_store.metrics.epoch_random_beacon_dkg_failed.set(1);
                rm.dkg_output
                    .set(None)
                    .expect("setting new OnceCell should succeed");
            }
            None => {
                info!(
                    "random beacon: no existing DKG output found for epoch {}",
                    committee.epoch()
                );

                // Load intermediate data.
                assert!(
                    epoch_store.protocol_config().dkg_version() > 0,
                    "BUG: DKG version 0 is deprecated"
                );
                rm.processed_messages.extend(
                    tables
                        .dkg_processed_messages
                        .safe_iter()
                        .map(|result| result.expect("typed_store should not fail")),
                );
                if let Some(used_messages) = tables
                    .dkg_used_messages
                    .get(&SINGLETON_KEY)
                    .expect("typed_store should not fail")
                {
                    rm.used_messages
                        .set(used_messages)
                        .expect("setting new OnceCell should succeed");
                }
                rm.confirmations.extend(
                    tables
                        .dkg_confirmations
                        .safe_iter()
                        .map(|result| result.expect("typed_store should not fail")),
                );
            }
        }

        // Resume randomness generation from where we left off.
        // This must be loaded regardless of whether DKG has finished yet, since the
        // RandomnessEventLoop and commit-handling logic in AuthorityPerEpochStore both
        // depend on this state.
        rm.next_randomness_round = tables
            .randomness_next_round
            .get(&SINGLETON_KEY)
            .expect("typed_store should not fail")
            .unwrap_or(RandomnessRound::new(0));
        info!(
            "random beacon: starting from next_randomness_round={}",
            rm.next_randomness_round
        );
        let first_incomplete_round = highest_completed_round
            .map(|r| r + 1)
            .unwrap_or(RandomnessRound::new(0));
        if first_incomplete_round < rm.next_randomness_round {
            info!(
                "random beacon: resuming generation for randomness rounds from {first_incomplete_round} to {}",
                rm.next_randomness_round - 1,
            );
            for r in first_incomplete_round.value()..rm.next_randomness_round.value() {
                network_handle.send_partial_signatures(committee.epoch(), RandomnessRound::new(r));
            }
        }

        Ok(rm)
    }

    /// Sends the initial dkg::Message to begin the randomness DKG protocol.
    pub async fn start_dkg(&mut self) -> IotaResult {
        if self.used_messages.initialized() || self.dkg_output.initialized() {
            // DKG already started (or completed or failed).
            return Ok(());
        }

        let _ = self.dkg_start_time.set(Instant::now());

        let epoch_store = self.epoch_store()?;
        let dkg_version = epoch_store.protocol_config().dkg_version();
        info!("random beacon: starting DKG, version {dkg_version}");

        let msg = match VersionedDkgMessage::create(dkg_version, self.party.clone()) {
            Ok(msg) => msg,
            Err(FastCryptoError::IgnoredMessage) => {
                info!(
                    "random beacon: no DKG Message for party id={} (zero weight)",
                    self.party.id
                );
                return Ok(());
            }
            Err(e) => {
                error!("random beacon: error while creating a DKG Message: {e:?}");
                return Ok(());
            }
        };

        info!("random beacon: created {msg:?} with dkg version {dkg_version}");
        let transaction = ConsensusTransaction::new_randomness_dkg_message(epoch_store.name, &msg);

        #[cfg_attr(not(any(msim, fail_points)), expect(unused_mut))]
        let mut fail_point_skip_sending = false;
        fail_point_if!("rb-dkg", || {
            // maybe skip sending in simtests
            fail_point_skip_sending = true;
        });
        if !fail_point_skip_sending {
            self.consensus_adapter
                .submit_to_consensus(&[transaction], &epoch_store)?;
        }

        epoch_store
            .metrics
            .epoch_random_beacon_dkg_message_time_ms
            .set(
                self.dkg_start_time
                    .get()
                    .unwrap() // already set above
                    .elapsed()
                    .as_millis() as i64,
            );
        Ok(())
    }

    /// Processes all received messages and advances the randomness DKG state
    /// machine when possible, sending out a dkg::Confirmation and
    /// generating final output.
    pub(crate) async fn advance_dkg(
        &mut self,
        consensus_output: &mut ConsensusCommitOutput,
        round: CommitRound,
    ) -> Result<(), AdvanceDkgError> {
        let epoch_store = self.epoch_store()?;

        // Once we have enough Messages, send a Confirmation.
        if !self.dkg_output.initialized() && !self.used_messages.initialized() {
            // Process all enqueued messages.
            let mut handles: FuturesUnordered<_> = std::mem::take(&mut self.enqueued_messages)
                .into_values()
                .collect();
            while let Some(res) = handles.next().await {
                let processed = match res {
                    Ok(processed) => processed,
                    Err(e) if e.is_panic() => std::panic::resume_unwind(e.into_panic()),
                    // Only runtime shutdown cancels the task. Treating it like an invalid
                    // message would persist DKG state that is missing this message.
                    Err(_) => return Err(AdvanceDkgError::RuntimeShutdown),
                };
                if let Some(processed) = processed {
                    self.processed_messages
                        .insert(processed.sender(), processed.clone());
                    consensus_output.insert_dkg_processed_message(processed);
                }
            }

            // Attempt to generate the Confirmation.
            match VersionedProcessedMessage::merge(
                self.party.clone(),
                self.processed_messages
                    .values()
                    .cloned()
                    .collect::<Vec<_>>(),
            ) {
                Ok((conf, used_msgs)) => {
                    info!(
                        "random beacon: sending DKG Confirmation with {} complaints",
                        conf.num_of_complaints()
                    );
                    if self.used_messages.set(used_msgs.clone()).is_err() {
                        error!("BUG: used_messages should only ever be set once");
                    }
                    consensus_output.insert_dkg_used_messages(used_msgs);

                    let transaction = ConsensusTransaction::new_randomness_dkg_confirmation(
                        epoch_store.name,
                        &conf,
                    );

                    #[cfg_attr(not(any(msim, fail_points)), expect(unused_mut))]
                    let mut fail_point_skip_sending = false;
                    fail_point_if!("rb-dkg", || {
                        // maybe skip sending in simtests
                        fail_point_skip_sending = true;
                    });
                    if !fail_point_skip_sending {
                        self.consensus_adapter
                            .submit_to_consensus(&[transaction], &epoch_store)?;
                    }

                    let elapsed = self.dkg_start_time.get().map(|t| t.elapsed().as_millis());
                    if let Some(elapsed) = elapsed {
                        epoch_store
                            .metrics
                            .epoch_random_beacon_dkg_confirmation_time_ms
                            .set(elapsed as i64);
                    }
                }
                Err(FastCryptoError::NotEnoughInputs) => (), // wait for more input
                Err(e) => debug!("random beacon: error while merging DKG Messages: {e:?}"),
            }
        }

        // Once we have enough Confirmations, process them and update shares.
        if !self.dkg_output.initialized() && self.used_messages.initialized() {
            match self
                .used_messages
                .get()
                .expect("checked above that `used_messages` is initialized")
                .complete_dkg(self.party.clone(), self.confirmations.values())
            {
                Ok(output) => {
                    let num_shares = output.shares.as_ref().map_or(0, |shares| shares.len());
                    let epoch_elapsed = epoch_store.epoch_open_time.elapsed().as_millis();
                    let elapsed = self.dkg_start_time.get().map(|t| t.elapsed().as_millis());
                    let completed_after_timeout =
                        Self::is_past_dkg_timeout_round(round, epoch_store.protocol_config());
                    info!(
                        "random beacon: DKG complete in {epoch_elapsed}ms since epoch start, {elapsed:?}ms since DKG start, with {num_shares} shares for this node, completed_after_timeout={completed_after_timeout}"
                    );
                    epoch_store
                        .metrics
                        .epoch_random_beacon_dkg_num_shares
                        .set(num_shares as i64);
                    epoch_store
                        .metrics
                        .epoch_random_beacon_dkg_epoch_start_completion_time_ms
                        .set(epoch_elapsed as i64);
                    epoch_store.metrics.epoch_random_beacon_dkg_failed.set(0);
                    epoch_store
                        .metrics
                        .epoch_random_beacon_dkg_completed_after_timeout
                        .set(i64::from(completed_after_timeout));
                    if completed_after_timeout {
                        fail_point!("rb-dkg-completed-after-timeout");
                    }
                    if let Some(elapsed) = elapsed {
                        epoch_store
                            .metrics
                            .epoch_random_beacon_dkg_completion_time_ms
                            .set(elapsed as i64);
                    }
                    self.dkg_output
                        .set(Some(output.clone()))
                        .expect("checked above that `dkg_output` is uninitialized");
                    self.network_handle.update_epoch(
                        epoch_store.committee().epoch(),
                        self.authority_info.clone(),
                        output.clone(),
                        self.party.t(),
                        None,
                    );
                    consensus_output.set_dkg_output(Some(output));
                }
                Err(FastCryptoError::NotEnoughInputs) => (), // wait for more input
                // The confirmations are sequenced by consensus and the used
                // messages are derived from sequenced messages, so every
                // validator computes this step from the same input, and
                // honest input cannot make it fail. A failure has one of two
                // causes. Either the DKG state this node persisted no longer
                // matches what consensus delivered, and only this node stops.
                // Or a library bug or more than a third of the stake dealing
                // bad shares makes the step fail everywhere, and every
                // validator stops on the same commit. A restart replays the
                // same input, so it does not clear either cause. Carrying on
                // would let the node defer or cancel randomness-using
                // transactions that its peers execute, and nothing would
                // notice until its checkpoint differs from the certified one.
                Err(e) => {
                    fatal!(
                        "random beacon: error while processing DKG Confirmations: {e:?}. A \
                            restart replays the same input. Restore the epoch database or \
                            state-sync before rejoining; if other validators stopped too, the \
                            cause is shared and needs a fix"
                    )
                }
            }
        }

        // If we ran out of time, mark DKG as failed.
        if !epoch_store
            .protocol_config()
            .allow_dkg_completion_after_timeout()
            && !self.dkg_output.initialized()
            && Self::is_past_dkg_timeout_round(round, epoch_store.protocol_config())
        {
            error!(
                "random beacon: DKG timed out. Randomness disabled for this epoch. All randomness-using transactions will fail."
            );
            epoch_store.metrics.epoch_random_beacon_dkg_failed.set(1);
            self.dkg_output
                .set(None)
                .expect("checked above that `dkg_output` is uninitialized");
            // Persist the terminal-failure verdict so a restart past the
            // timeout round doesn't drop back to a "pending" view of DKG and
            // defer randomness-using transactions indefinitely.
            consensus_output.set_dkg_output(None);
        }

        Ok(())
    }

    /// Adds a received VersionedDkgMessage to the randomness DKG state machine.
    pub fn add_message(
        &mut self,
        authority: &AuthorityName,
        msg: VersionedDkgMessage,
    ) -> IotaResult {
        // message was received from other validators, so we need to ensure it uses a
        // supported version before we call other functions that assume the
        // version is correct
        let dkg_version = self.epoch_store()?.protocol_config().dkg_version();
        if !msg.is_valid_version(dkg_version) {
            warn!("ignoring DKG Message from authority {authority:?} with unsupported version");
            return Ok(());
        }

        if self.used_messages.initialized() || self.dkg_output.initialized() {
            // We've already sent a `Confirmation`, so we can't add any more messages.
            return Ok(());
        }
        let Some((_, party_id)) = self.authority_info.get(authority) else {
            debug_fatal!(
                "random beacon: received DKG Message from unknown authority: {authority:?}"
            );
            return Ok(());
        };
        if *party_id != msg.sender() {
            warn!(
                "ignoring equivocating DKG Message from authority {authority:?} pretending to be PartyId {party_id}"
            );
            return Ok(());
        }
        if self.enqueued_messages.contains_key(&msg.sender())
            || self.processed_messages.contains_key(&msg.sender())
        {
            info!("ignoring duplicate DKG Message from authority {authority:?}");
            return Ok(());
        }

        let party = self.party.clone();
        // TODO: Could save some CPU by not processing messages if we already have
        // enough to merge.
        self.enqueued_messages.insert(
            msg.sender(),
            tokio::task::spawn_blocking(move || {
                match VersionedProcessedMessage::process(party, msg) {
                    Ok(processed) => Some(processed),
                    Err(err) => {
                        debug!("random beacon: error while processing DKG Message: {err:?}");
                        None
                    }
                }
            }),
        );
        Ok(())
    }

    /// Adds a received dkg::Confirmation to the randomness DKG state machine.
    pub(crate) fn add_confirmation(
        &mut self,
        output: &mut ConsensusCommitOutput,
        authority: &AuthorityName,
        conf: VersionedDkgConfirmation,
    ) -> IotaResult {
        // confirmation was received from other validators, so we need to ensure it uses
        // a supported version before we call other functions that assume the
        // version is correct
        let dkg_version = self.epoch_store()?.protocol_config().dkg_version();
        if !conf.is_valid_version(dkg_version) {
            warn!(
                "ignoring DKG Confirmation from authority {authority:?} with unsupported version"
            );
            return Ok(());
        }

        if self.dkg_output.initialized() {
            // Once we have completed DKG, no more `Confirmation`s are needed.
            return Ok(());
        }
        let Some((_, party_id)) = self.authority_info.get(authority) else {
            error!(
                "random beacon: received DKG Confirmation from unknown authority: {authority:?}"
            );
            return Ok(());
        };
        if *party_id != conf.sender() {
            warn!(
                "ignoring equivocating DKG Confirmation from authority {authority:?} pretending to be PartyId {party_id}"
            );
            return Ok(());
        }
        self.confirmations.insert(conf.sender(), conf.clone());
        output.insert_dkg_confirmation(conf);
        Ok(())
    }

    /// Reserves the next available round number for randomness generation if
    /// enough time has elapsed, or returns None if not yet ready (based on
    /// ProtocolConfig setting). Once the given batch is written,
    /// `generate_randomness` must be called to start the process. On restart,
    /// any reserved rounds for which the batch was written will automatically
    /// be resumed.
    pub(crate) fn reserve_next_randomness(
        &mut self,
        commit_timestamp: CommitTimestampMs,
        output: &mut ConsensusCommitOutput,
    ) -> IotaResult<Option<RandomnessRound>> {
        let epoch_store = self.epoch_store()?;

        let last_round_timestamp = epoch_store
            .get_randomness_last_round_timestamp()
            .expect("read should not fail");

        if let Some(last_round_timestamp) = last_round_timestamp {
            if commit_timestamp - last_round_timestamp
                < epoch_store
                    .protocol_config()
                    .random_beacon_min_round_interval_ms()
            {
                return Ok(None);
            }
        }

        let randomness_round = self.next_randomness_round;
        self.next_randomness_round = self
            .next_randomness_round
            .checked_add(1)
            .ok_or_else(|| IotaError::Unknown("RandomnessRound overflow".to_string()))?;

        output.reserve_next_randomness_round(self.next_randomness_round, commit_timestamp);

        Ok(Some(randomness_round))
    }

    /// Starts the process of generating the given RandomnessRound.
    pub fn generate_randomness(&self, epoch: EpochId, randomness_round: RandomnessRound) {
        self.network_handle
            .send_partial_signatures(epoch, randomness_round);
    }

    pub fn dkg_status(&self) -> DkgStatus {
        match self.dkg_output.get() {
            Some(Some(_)) => DkgStatus::Successful,
            Some(None) => DkgStatus::TimedOut,
            None => DkgStatus::Pending,
        }
    }

    /// Returns true if `round` is past the DKG timeout round. This comparison
    /// is consensus-critical: all validators must agree on the exact commit
    /// round at which DKG times out.
    fn is_past_dkg_timeout_round(
        round: CommitRound,
        protocol_config: &iota_protocol_config::ProtocolConfig,
    ) -> bool {
        round > protocol_config.random_beacon_dkg_timeout_round() as u64
    }

    /// Returns the DKG status as observed at the given consensus commit round.
    /// When `allow_dkg_completion_after_timeout` is enabled, an incomplete DKG
    /// reports `TimedOut` for commits past the timeout round while the protocol
    /// keeps running, so a later commit may observe `Successful`.
    ///
    /// Returns `IotaError::EpochEnded` if the epoch store is already gone.
    pub fn dkg_status_for_commit_round(&mut self, round: CommitRound) -> IotaResult<DkgStatus> {
        let status = self.dkg_status();
        if status != DkgStatus::Pending {
            return Ok(status);
        }

        let epoch_store = self.epoch_store()?;
        if !epoch_store
            .protocol_config()
            .allow_dkg_completion_after_timeout()
            || !Self::is_past_dkg_timeout_round(round, epoch_store.protocol_config())
        {
            return Ok(status);
        }

        if !self.dkg_timeout_reported {
            error!(
                "random beacon: DKG timed out; randomness-using transactions will be canceled unless/until DKG completes"
            );
            epoch_store.metrics.epoch_random_beacon_dkg_failed.set(1);
            self.dkg_timeout_reported = true;
        }
        Ok(DkgStatus::TimedOut)
    }

    /// Generates a new RandomnessReporter for reporting observed rounds to this
    /// RandomnessManager.
    pub fn reporter(&self) -> RandomnessReporter {
        RandomnessReporter {
            epoch_store: self.epoch_store.clone(),
            epoch: self.epoch,
            network_handle: self.network_handle.clone(),
            highest_completed_round: self.highest_completed_round.clone(),
        }
    }

    fn epoch_store(&self) -> IotaResult<Arc<AuthorityPerEpochStore>> {
        self.epoch_store
            .upgrade()
            .ok_or(IotaError::EpochEnded(self.epoch))
    }

    fn randomness_dkg_info_from_committee(
        committee: &Committee,
    ) -> Vec<(
        u16,
        AuthorityName,
        fastcrypto_tbls::ecies_v1::PublicKey<bls12381::G2Element>,
        StakeUnit,
    )> {
        committee
            .members()
            .map(|(name, stake)| {
                let index: u16 = committee
                    .authority_index(name)
                    .expect("lookup of known committee member should succeed")
                    .try_into()
                    .expect("authority index should fit in u16");
                let pk = bls12381::G2Element::from_byte_array(
                    committee
                        .public_key(name)
                        .expect("lookup of known committee member should succeed")
                        .as_bytes()
                        .try_into()
                        .expect("key length should match"),
                )
                .expect("should work to convert BLS key to G2Element");
                (
                    index,
                    *name,
                    fastcrypto_tbls::ecies_v1::PublicKey::from(pk),
                    *stake,
                )
            })
            .collect()
    }
}

// Used by other components to notify the randomness system of observed
// randomness.
#[derive(Clone)]
pub struct RandomnessReporter {
    epoch_store: Weak<AuthorityPerEpochStore>,
    epoch: EpochId,
    network_handle: randomness::Handle,
    highest_completed_round: Arc<Mutex<Option<RandomnessRound>>>,
}

impl RandomnessReporter {
    /// Notifies the associated randomness manager that randomness for the given
    /// round has been durably committed in a checkpoint. This completes the
    /// process of generating randomness for the round.
    pub fn notify_randomness_in_checkpoint(&self, round: RandomnessRound) -> IotaResult {
        let epoch_store = self
            .epoch_store
            .upgrade()
            .ok_or(IotaError::EpochEnded(self.epoch))?;
        let mut highest_completed_round = self.highest_completed_round.lock();
        if Some(round) > *highest_completed_round {
            *highest_completed_round = Some(round);
            epoch_store
                .tables()?
                .randomness_highest_completed_round
                .insert(&SINGLETON_KEY, &round)?;
            self.network_handle
                .complete_round(epoch_store.committee().epoch(), round);
        }
        Ok(())
    }

    /// Passes the randomness for `round` found in a certified checkpoint to
    /// the randomness event loop, so a round this node missed completes
    /// without peers resending it.
    pub fn forward_randomness_from_checkpoint(&self, round: RandomnessRound, random_bytes: &[u8]) {
        match bcs::from_bytes::<RandomnessSignature>(random_bytes) {
            Ok(sig) => self
                .network_handle
                .receive_full_signature_from_checkpoint(self.epoch, round, sig),
            Err(e) => error!(
                "random beacon: cannot decode randomness for round {round} from a certified checkpoint: {e:?}"
            ),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DkgStatus {
    Pending,
    /// DKG did not complete by the timeout round. If `allow_dkg_completion_after_timeout`
    /// is enabled, DKG may still complete later and reach `Successful`.
    TimedOut,
    Successful,
}

#[cfg(test)]
mod tests {
    use std::num::NonZeroUsize;

    use iota_protocol_config::{Chain, ProtocolConfig, ProtocolVersion};
    use iota_types::{crypto::AuthorityKeyPair, messages_consensus::ConsensusTransactionKind};
    use starfish_core::{BlockRef, BlockStatus};
    use tokio::sync::mpsc;

    use crate::{
        authority::{
            authority_per_epoch_store::{
                AuthorityPerEpochStore, ExecutionIndices, ExecutionIndicesWithStats,
            },
            test_authority_builder::TestAuthorityBuilder,
        },
        checkpoints::CheckpointStore,
        consensus_adapter::{
            ConnectionMonitorStatusForTests, ConsensusAdapter, ConsensusAdapterMetrics,
            MockConsensusClient,
        },
        epoch::randomness::*,
        mock_consensus::with_block_status,
    };

    #[tokio::test]
    async fn test_dkg_v1() {
        test_dkg(1).await;
    }

    async fn test_dkg(version: u64) {
        telemetry_subscribers::init_for_testing();

        let network_config =
            iota_swarm_config::network_config_builder::ConfigBuilder::new_with_temp_dir()
                .committee_size(NonZeroUsize::new(4).unwrap())
                .with_reference_gas_price(500)
                .build();

        let mut protocol_config =
            ProtocolConfig::get_for_version(ProtocolVersion::max(), Chain::Unknown);
        protocol_config.set_random_beacon_dkg_version_for_testing(version);

        let mut epoch_stores = Vec::new();
        let mut randomness_managers = Vec::new();
        let (tx_consensus, mut rx_consensus) = mpsc::channel(100);

        for validator in network_config.validator_configs.iter() {
            // Send consensus messages to channel.
            let mut mock_consensus_client = MockConsensusClient::new();
            let tx_consensus = tx_consensus.clone();
            mock_consensus_client
                .expect_submit()
                .withf(move |transactions: &[ConsensusTransaction], _epoch_store| {
                    tx_consensus.try_send(transactions.to_vec()).unwrap();
                    true
                })
                .returning(|_, _| {
                    Ok(with_block_status(BlockStatus::Sequenced(
                        starfish_core::GenericTransactionRef::BlockRef(BlockRef::MIN),
                    )))
                });

            let state = TestAuthorityBuilder::new()
                .with_protocol_config(protocol_config.clone())
                .with_genesis_and_keypair(&network_config.genesis, validator.authority_key_pair())
                .build()
                .await;
            let consensus_adapter = Arc::new(ConsensusAdapter::new(
                Arc::new(mock_consensus_client),
                CheckpointStore::new_for_tests(),
                state.name,
                Arc::new(ConnectionMonitorStatusForTests {}),
                100_000,
                100_000,
                None,
                None,
                ConsensusAdapterMetrics::new_test(),
                50,
            ));
            let epoch_store = state.epoch_store_for_testing();
            let randomness_manager = RandomnessManager::try_new(
                Arc::downgrade(&epoch_store),
                Box::new(consensus_adapter.clone()),
                iota_network::randomness::Handle::new_stub(),
                validator.authority_key_pair(),
            )
            .await
            .unwrap();

            epoch_stores.push(epoch_store);
            randomness_managers.push(randomness_manager);
        }

        // Generate and distribute Messages.
        let mut dkg_messages = Vec::new();
        for randomness_manager in randomness_managers.iter_mut() {
            randomness_manager.start_dkg().await.unwrap();

            let mut dkg_message = rx_consensus.recv().await.unwrap();
            assert!(dkg_message.len() == 1);
            match dkg_message.remove(0).kind {
                ConsensusTransactionKind::RandomnessDkgMessage(_, bytes) => {
                    let msg: VersionedDkgMessage = bcs::from_bytes(&bytes)
                        .expect("DKG message deserialization should not fail");
                    dkg_messages.push(msg);
                }
                _ => panic!("wrong type of message sent"),
            }
        }

        // A DKG message task cancelled by runtime shutdown must abandon the
        // commit instead of being skipped like an invalid message.
        let shutdown_runtime = tokio::runtime::Builder::new_current_thread()
            .build()
            .unwrap();
        let shutdown_handle = shutdown_runtime.handle().clone();
        shutdown_runtime.shutdown_background();
        randomness_managers[0]
            .enqueued_messages
            .insert(0, shutdown_handle.spawn_blocking(|| None));
        let result = randomness_managers[0]
            .advance_dkg(&mut ConsensusCommitOutput::new(0, 0), 0)
            .await;
        assert!(matches!(result, Err(AdvanceDkgError::RuntimeShutdown)));
        assert!(randomness_managers[0].processed_messages.is_empty());

        for i in 0..randomness_managers.len() {
            let mut output = ConsensusCommitOutput::new(0, 0);
            output.record_consensus_commit_stats(ExecutionIndicesWithStats {
                index: ExecutionIndices {
                    last_committed_round: 0,
                    ..Default::default()
                },
                ..Default::default()
            });
            for (j, dkg_message) in dkg_messages.iter().cloned().enumerate() {
                randomness_managers[i]
                    .add_message(&epoch_stores[j].name, dkg_message)
                    .unwrap();
            }
            randomness_managers[i]
                .advance_dkg(&mut output, 0)
                .await
                .unwrap();
            let mut batch = epoch_stores[i].db_batch_for_test();
            output.write_to_batch(&epoch_stores[i], &mut batch).unwrap();
            batch.write().unwrap();
        }

        // Generate and distribute Confirmations.
        let mut dkg_confirmations = Vec::new();
        for _ in 0..randomness_managers.len() {
            let mut dkg_confirmation = rx_consensus.recv().await.unwrap();
            assert!(dkg_confirmation.len() == 1);
            match dkg_confirmation.remove(0).kind {
                ConsensusTransactionKind::RandomnessDkgConfirmation(_, bytes) => {
                    let msg: VersionedDkgConfirmation = bcs::from_bytes(&bytes)
                        .expect("DKG message deserialization should not fail");
                    dkg_confirmations.push(msg);
                }
                _ => panic!("wrong type of message sent"),
            }
        }
        for i in 0..randomness_managers.len() {
            let mut output = ConsensusCommitOutput::new(0, 0);
            output.record_consensus_commit_stats(ExecutionIndicesWithStats {
                index: ExecutionIndices {
                    last_committed_round: 1,
                    ..Default::default()
                },
                ..Default::default()
            });
            for (j, dkg_confirmation) in dkg_confirmations.iter().cloned().enumerate() {
                randomness_managers[i]
                    .add_confirmation(&mut output, &epoch_stores[j].name, dkg_confirmation)
                    .unwrap();
            }
            randomness_managers[i]
                .advance_dkg(&mut output, 0)
                .await
                .unwrap();
            let mut batch = epoch_stores[i].db_batch_for_test();
            output.write_to_batch(&epoch_stores[i], &mut batch).unwrap();
            batch.write().unwrap();
        }

        // Verify DKG completed.
        for randomness_manager in &randomness_managers {
            assert_eq!(DkgStatus::Successful, randomness_manager.dkg_status());
        }
    }

    #[tokio::test]
    async fn test_dkg_expiration_v1() {
        test_dkg_expiration(1).await;
    }

    async fn test_dkg_expiration(version: u64) {
        telemetry_subscribers::init_for_testing();

        let network_config =
            iota_swarm_config::network_config_builder::ConfigBuilder::new_with_temp_dir()
                .committee_size(NonZeroUsize::new(4).unwrap())
                .with_reference_gas_price(500)
                .build();

        let mut epoch_stores = Vec::new();
        let mut randomness_managers = Vec::new();
        let (tx_consensus, mut rx_consensus) = mpsc::channel(100);

        let mut protocol_config =
            ProtocolConfig::get_for_version(ProtocolVersion::max(), Chain::Unknown);
        protocol_config.set_random_beacon_dkg_version_for_testing(version);
        protocol_config.set_allow_dkg_completion_after_timeout_for_testing(false);

        for validator in network_config.validator_configs.iter() {
            // Send consensus messages to channel.
            let mut mock_consensus_client = MockConsensusClient::new();
            let tx_consensus = tx_consensus.clone();
            mock_consensus_client
                .expect_submit()
                .withf(move |transactions: &[ConsensusTransaction], _epoch_store| {
                    tx_consensus.try_send(transactions.to_vec()).unwrap();
                    true
                })
                .returning(|_, _| {
                    Ok(with_block_status(starfish_core::BlockStatus::Sequenced(
                        starfish_core::GenericTransactionRef::BlockRef(BlockRef::MIN),
                    )))
                });

            let state = TestAuthorityBuilder::new()
                .with_protocol_config(protocol_config.clone())
                .with_genesis_and_keypair(&network_config.genesis, validator.authority_key_pair())
                .build()
                .await;
            let consensus_adapter = Arc::new(ConsensusAdapter::new(
                Arc::new(mock_consensus_client),
                CheckpointStore::new_for_tests(),
                state.name,
                Arc::new(ConnectionMonitorStatusForTests {}),
                100_000,
                100_000,
                None,
                None,
                ConsensusAdapterMetrics::new_test(),
                50,
            ));
            let epoch_store = state.epoch_store_for_testing();
            let randomness_manager = RandomnessManager::try_new(
                Arc::downgrade(&epoch_store),
                Box::new(consensus_adapter.clone()),
                iota_network::randomness::Handle::new_stub(),
                validator.authority_key_pair(),
            )
            .await
            .unwrap();

            epoch_stores.push(epoch_store);
            randomness_managers.push(randomness_manager);
        }

        // Generate and distribute Messages.
        let mut dkg_messages = Vec::new();
        for randomness_manager in randomness_managers.iter_mut() {
            randomness_manager.start_dkg().await.unwrap();

            let mut dkg_message = rx_consensus.recv().await.unwrap();
            assert!(dkg_message.len() == 1);
            match dkg_message.remove(0).kind {
                ConsensusTransactionKind::RandomnessDkgMessage(_, bytes) => {
                    let msg: VersionedDkgMessage = bcs::from_bytes(&bytes)
                        .expect("DKG message deserialization should not fail");
                    dkg_messages.push(msg);
                }
                _ => panic!("wrong type of message sent"),
            }
        }

        // With late completion disabled, `dkg_status_for_commit_round` passes
        // through the raw status; timeout failure is only marked by `advance_dkg`.
        assert_eq!(
            DkgStatus::Pending,
            randomness_managers[0]
                .dkg_status_for_commit_round(u64::MAX)
                .unwrap()
        );

        for i in 0..randomness_managers.len() {
            let mut output = ConsensusCommitOutput::new(0, 0);
            output.record_consensus_commit_stats(ExecutionIndicesWithStats {
                index: ExecutionIndices {
                    last_committed_round: 0,
                    ..Default::default()
                },
                ..Default::default()
            });
            for (j, dkg_message) in dkg_messages.iter().cloned().enumerate() {
                randomness_managers[i]
                    .add_message(&epoch_stores[j].name, dkg_message)
                    .unwrap();
            }
            randomness_managers[i]
                .advance_dkg(&mut output, u64::MAX)
                .await
                .unwrap();
            let mut batch = epoch_stores[i].db_batch_for_test();
            output.write_to_batch(&epoch_stores[i], &mut batch).unwrap();
            batch.write().unwrap();
        }

        // Verify DKG failed.
        for randomness_manager in &randomness_managers {
            assert_eq!(DkgStatus::TimedOut, randomness_manager.dkg_status());
        }

        // Simulate a restart: rebuild each manager from its persisted epoch store
        // and verify the failure verdict is durably restored (not Pending). This
        // exercises the `dkg_output_v2` `Some(None)` failure-load path in
        // `try_new`; before the fix the verdict lived only in the in-memory
        // `OnceCell`, so a restart re-loaded the manager as Pending and
        // randomness-using transactions were deferred forever.
        for (epoch_store, validator) in epoch_stores
            .iter()
            .zip(network_config.validator_configs.iter())
        {
            let recovered = recover_randomness_manager(
                epoch_store,
                validator.authority_key_pair(),
                tx_consensus.clone(),
            )
            .await;
            assert_eq!(DkgStatus::TimedOut, recovered.dkg_status());
        }
    }

    /// Rebuilds a `RandomnessManager` from an existing (already-persisted)
    /// epoch store, simulating a validator restart. Only on-disk DKG state
    /// is carried over: the in-memory state machine starts fresh and must
    /// be restored from the persisted tables by `try_new`. Like node startup,
    /// it then re-sends the consensus transactions still pending in the epoch
    /// store through `ConsensusAdapter::submit_recovered`.
    async fn recover_randomness_manager(
        epoch_store: &Arc<AuthorityPerEpochStore>,
        authority_key_pair: &AuthorityKeyPair,
        tx_consensus: mpsc::Sender<Vec<ConsensusTransaction>>,
    ) -> RandomnessManager {
        let mut mock_consensus_client = MockConsensusClient::new();
        mock_consensus_client
            .expect_submit()
            .withf(move |transactions: &[ConsensusTransaction], _epoch_store| {
                tx_consensus.try_send(transactions.to_vec()).unwrap();
                true
            })
            .returning(|_, _| {
                Ok(with_block_status(BlockStatus::Sequenced(
                    starfish_core::GenericTransactionRef::BlockRef(BlockRef::MIN),
                )))
            });
        let consensus_adapter = Arc::new(ConsensusAdapter::new(
            Arc::new(mock_consensus_client),
            CheckpointStore::new_for_tests(),
            epoch_store.name,
            Arc::new(ConnectionMonitorStatusForTests {}),
            100_000,
            100_000,
            None,
            None,
            ConsensusAdapterMetrics::new_test(),
            50,
        ));
        let randomness_manager = RandomnessManager::try_new(
            Arc::downgrade(epoch_store),
            Box::new(consensus_adapter.clone()),
            iota_network::randomness::Handle::new_stub(),
            authority_key_pair,
        )
        .await
        .unwrap();
        consensus_adapter.submit_recovered(epoch_store);
        randomness_manager
    }

    /// Distributes DKG messages (but no confirmations) so DKG is still pending,
    /// then reconstructs every `RandomnessManager` from the same epoch stores
    /// and asserts the in-progress state (processed and used messages) is
    /// restored and DKG is still reported as pending. This exercises the
    /// no-output (`None`) load branch in `try_new`. The messages are applied
    /// past the timeout round. The restart sends each validator's confirmation
    /// again, and the restarted managers then complete DKG from them.
    #[tokio::test]
    async fn test_dkg_recovers_processed_messages_after_restart() {
        telemetry_subscribers::init_for_testing();

        let network_config =
            iota_swarm_config::network_config_builder::ConfigBuilder::new_with_temp_dir()
                .committee_size(NonZeroUsize::new(4).unwrap())
                .with_reference_gas_price(500)
                .build();

        let mut protocol_config =
            ProtocolConfig::get_for_version(ProtocolVersion::max(), Chain::Unknown);
        protocol_config.set_random_beacon_dkg_version_for_testing(1);

        let mut epoch_stores = Vec::new();
        let mut randomness_managers = Vec::new();
        let (tx_consensus, mut rx_consensus) = mpsc::channel(100);

        for validator in network_config.validator_configs.iter() {
            let mut mock_consensus_client = MockConsensusClient::new();
            let tx_consensus = tx_consensus.clone();
            mock_consensus_client
                .expect_submit()
                .withf(move |transactions: &[ConsensusTransaction], _epoch_store| {
                    tx_consensus.try_send(transactions.to_vec()).unwrap();
                    true
                })
                .returning(|_, _| {
                    Ok(with_block_status(BlockStatus::Sequenced(
                        starfish_core::GenericTransactionRef::BlockRef(BlockRef::MIN),
                    )))
                });

            let state = TestAuthorityBuilder::new()
                .with_protocol_config(protocol_config.clone())
                .with_genesis_and_keypair(&network_config.genesis, validator.authority_key_pair())
                .build()
                .await;
            let consensus_adapter = Arc::new(ConsensusAdapter::new(
                Arc::new(mock_consensus_client),
                CheckpointStore::new_for_tests(),
                state.name,
                Arc::new(ConnectionMonitorStatusForTests {}),
                100_000,
                100_000,
                None,
                None,
                ConsensusAdapterMetrics::new_test(),
                50,
            ));
            let epoch_store = state.epoch_store_for_testing();
            let randomness_manager = RandomnessManager::try_new(
                Arc::downgrade(&epoch_store),
                Box::new(consensus_adapter.clone()),
                iota_network::randomness::Handle::new_stub(),
                validator.authority_key_pair(),
            )
            .await
            .unwrap();

            epoch_stores.push(epoch_store);
            randomness_managers.push(randomness_manager);
        }

        let timeout = epoch_stores[0]
            .protocol_config()
            .random_beacon_dkg_timeout_round() as u64;

        // Generate and distribute Messages.
        let mut dkg_messages = Vec::new();
        for randomness_manager in randomness_managers.iter_mut() {
            randomness_manager.start_dkg().await.unwrap();

            let mut dkg_message = rx_consensus.recv().await.unwrap();
            assert!(dkg_message.len() == 1);
            match dkg_message.remove(0).kind {
                ConsensusTransactionKind::RandomnessDkgMessage(_, bytes) => {
                    let msg: VersionedDkgMessage = bcs::from_bytes(&bytes)
                        .expect("DKG message deserialization should not fail");
                    dkg_messages.push(msg);
                }
                _ => panic!("wrong type of message sent"),
            }
        }
        for i in 0..randomness_managers.len() {
            let mut output = ConsensusCommitOutput::new(0, 0);
            output.record_consensus_commit_stats(ExecutionIndicesWithStats {
                index: ExecutionIndices {
                    last_committed_round: 0,
                    ..Default::default()
                },
                ..Default::default()
            });
            for (j, dkg_message) in dkg_messages.iter().cloned().enumerate() {
                randomness_managers[i]
                    .add_message(&epoch_stores[j].name, dkg_message)
                    .unwrap();
            }
            // Past the timeout round, DKG keeps running when late completion is
            // allowed.
            randomness_managers[i]
                .advance_dkg(&mut output, timeout + 1)
                .await
                .unwrap();
            let mut batch = epoch_stores[i].db_batch_for_test();
            output.write_to_batch(&epoch_stores[i], &mut batch).unwrap();
            batch.write().unwrap();
        }

        // All messages processed and a confirmation emitted, but no confirmations
        // applied yet: DKG is still pending.
        for randomness_manager in &randomness_managers {
            assert_eq!(DkgStatus::Pending, randomness_manager.dkg_status());
            assert_eq!(
                network_config.validator_configs.len(),
                randomness_manager.processed_messages.len()
            );
            assert!(randomness_manager.used_messages.initialized());
        }
        let dkg_confirmations =
            collect_dkg_confirmations(&mut rx_consensus, network_config.validator_configs.len())
                .await;

        // Simulate a restart: rebuild each manager and verify the in-progress DKG
        // state is restored and still pending.
        let mut recovered_managers = Vec::new();
        for (epoch_store, validator) in epoch_stores
            .iter()
            .zip(network_config.validator_configs.iter())
        {
            let mut recovered = recover_randomness_manager(
                epoch_store,
                validator.authority_key_pair(),
                tx_consensus.clone(),
            )
            .await;
            assert_eq!(DkgStatus::Pending, recovered.dkg_status());
            assert_eq!(
                network_config.validator_configs.len(),
                recovered.processed_messages.len()
            );
            assert!(recovered.used_messages.initialized());
            assert_eq!(
                DkgStatus::Pending,
                recovered.dkg_status_for_commit_round(timeout).unwrap()
            );
            assert_eq!(
                DkgStatus::TimedOut,
                recovered.dkg_status_for_commit_round(timeout + 1).unwrap()
            );
            recovered_managers.push(recovered);
        }

        // The confirmations were never processed by consensus, so they are still
        // pending in each epoch store and are sent again on restart.
        let mut resent_confirmations = tokio::time::timeout(
            std::time::Duration::from_secs(30),
            collect_dkg_confirmations(&mut rx_consensus, network_config.validator_configs.len()),
        )
        .await
        .expect("restart should send the pending DKG confirmations again");
        let mut sent_confirmations = dkg_confirmations.clone();
        resent_confirmations.sort_by_key(|conf| conf.sender());
        sent_confirmations.sort_by_key(|conf| conf.sender());
        assert_eq!(sent_confirmations, resent_confirmations);

        // DKG completes after its timeout once the confirmations are sequenced.
        for i in 0..recovered_managers.len() {
            let mut output = ConsensusCommitOutput::new(0, 0);
            output.record_consensus_commit_stats(ExecutionIndicesWithStats {
                index: ExecutionIndices {
                    last_committed_round: 1,
                    ..Default::default()
                },
                ..Default::default()
            });
            for (j, dkg_confirmation) in dkg_confirmations.iter().cloned().enumerate() {
                recovered_managers[i]
                    .add_confirmation(&mut output, &epoch_stores[j].name, dkg_confirmation)
                    .unwrap();
            }
            recovered_managers[i]
                .advance_dkg(&mut output, timeout + 2)
                .await
                .unwrap();
            let mut batch = epoch_stores[i].db_batch_for_test();
            output.write_to_batch(&epoch_stores[i], &mut batch).unwrap();
            batch.write().unwrap();
        }
        for recovered in &recovered_managers {
            assert_eq!(DkgStatus::Successful, recovered.dkg_status());
        }
    }

    /// Receives one DKG confirmation from each of `num_validators` validators,
    /// skipping any DKG messages sent in between.
    async fn collect_dkg_confirmations(
        rx_consensus: &mut mpsc::Receiver<Vec<ConsensusTransaction>>,
        num_validators: usize,
    ) -> Vec<VersionedDkgConfirmation> {
        let mut dkg_confirmations = Vec::new();
        while dkg_confirmations.len() < num_validators {
            let mut transactions = rx_consensus.recv().await.unwrap();
            assert!(transactions.len() == 1);
            match transactions.remove(0).kind {
                ConsensusTransactionKind::RandomnessDkgConfirmation(_, bytes) => {
                    let msg: VersionedDkgConfirmation = bcs::from_bytes(&bytes)
                        .expect("DKG confirmation deserialization should not fail");
                    dkg_confirmations.push(msg);
                }
                ConsensusTransactionKind::RandomnessDkgMessage(..) => {}
                _ => panic!("wrong type of message sent"),
            }
        }
        dkg_confirmations
    }

    #[tokio::test]
    async fn try_new_fails_when_reduction_lower_bound_exceeds_u16() {
        telemetry_subscribers::init_for_testing();

        let network_config =
            iota_swarm_config::network_config_builder::ConfigBuilder::new_with_temp_dir()
                .committee_size(NonZeroUsize::new(4).unwrap())
                .with_reference_gas_price(500)
                .build();

        let mut protocol_config =
            ProtocolConfig::get_for_version(ProtocolVersion::max(), Chain::Unknown);
        protocol_config
            .set_random_beacon_reduction_lower_bound_for_testing(u32::from(u16::MAX) + 1);

        let validator = &network_config.validator_configs[0];
        let state = TestAuthorityBuilder::new()
            .with_protocol_config(protocol_config)
            .with_genesis_and_keypair(&network_config.genesis, validator.authority_key_pair())
            .build()
            .await;
        let consensus_adapter = Arc::new(ConsensusAdapter::new(
            Arc::new(MockConsensusClient::new()),
            CheckpointStore::new_for_tests(),
            state.name,
            Arc::new(ConnectionMonitorStatusForTests {}),
            100_000,
            100_000,
            None,
            None,
            ConsensusAdapterMetrics::new_test(),
            50,
        ));
        let epoch_store = state.epoch_store_for_testing();

        let randomness_manager = RandomnessManager::try_new(
            Arc::downgrade(&epoch_store),
            Box::new(consensus_adapter),
            iota_network::randomness::Handle::new_stub(),
            validator.authority_key_pair(),
        )
        .await;

        let err = randomness_manager.err().expect("construction should fail");
        assert!(err.to_string().contains("reduction lower bound"));
    }

    /// A failure to compute the DKG output halts the node: every validator
    /// computes it from the same sequenced input, so a failure is a local
    /// fault, and a node that carried on would defer or cancel
    /// randomness-using transactions its peers execute.
    ///
    /// The failure is forced on the first manager by seeding it with an empty
    /// set of used messages: once enough confirmations arrive, completion has
    /// no message left to aggregate and fails.
    #[tokio::test]
    #[should_panic(expected = "error while processing DKG Confirmations")]
    async fn test_dkg_completion_error_is_fatal() {
        telemetry_subscribers::init_for_testing();

        let network_config =
            iota_swarm_config::network_config_builder::ConfigBuilder::new_with_temp_dir()
                .committee_size(NonZeroUsize::new(4).unwrap())
                .with_reference_gas_price(500)
                .build();

        let mut protocol_config =
            ProtocolConfig::get_for_version(ProtocolVersion::max(), Chain::Unknown);
        protocol_config.set_random_beacon_dkg_version_for_testing(1);

        let mut epoch_stores = Vec::new();
        let mut randomness_managers = Vec::new();
        let (tx_consensus, mut rx_consensus) = mpsc::channel(100);

        for validator in network_config.validator_configs.iter() {
            let mut mock_consensus_client = MockConsensusClient::new();
            let tx_consensus = tx_consensus.clone();
            mock_consensus_client
                .expect_submit()
                .withf(move |transactions: &[ConsensusTransaction], _epoch_store| {
                    tx_consensus.try_send(transactions.to_vec()).unwrap();
                    true
                })
                .returning(|_, _| {
                    Ok(with_block_status(BlockStatus::Sequenced(
                        starfish_core::GenericTransactionRef::BlockRef(BlockRef::MIN),
                    )))
                });

            let state = TestAuthorityBuilder::new()
                .with_protocol_config(protocol_config.clone())
                .with_genesis_and_keypair(&network_config.genesis, validator.authority_key_pair())
                .build()
                .await;
            let consensus_adapter = Arc::new(ConsensusAdapter::new(
                Arc::new(mock_consensus_client),
                CheckpointStore::new_for_tests(),
                state.name,
                Arc::new(ConnectionMonitorStatusForTests {}),
                100_000,
                100_000,
                None,
                None,
                ConsensusAdapterMetrics::new_test(),
                50,
            ));
            let epoch_store = state.epoch_store_for_testing();
            let randomness_manager = RandomnessManager::try_new(
                Arc::downgrade(&epoch_store),
                Box::new(consensus_adapter.clone()),
                iota_network::randomness::Handle::new_stub(),
                validator.authority_key_pair(),
            )
            .await
            .unwrap();

            epoch_stores.push(epoch_store);
            randomness_managers.push(randomness_manager);
        }

        // Every manager sends its DKG Message.
        let mut dkg_messages = Vec::new();
        for randomness_manager in randomness_managers.iter_mut() {
            randomness_manager.start_dkg().await.unwrap();

            let mut dkg_message = rx_consensus.recv().await.unwrap();
            assert_eq!(dkg_message.len(), 1);
            match dkg_message.remove(0).kind {
                ConsensusTransactionKind::RandomnessDkgMessage(_, bytes) => {
                    let msg: VersionedDkgMessage = bcs::from_bytes(&bytes)
                        .expect("DKG message deserialization should not fail");
                    dkg_messages.push(msg);
                }
                _ => panic!("wrong type of message sent"),
            }
        }

        // The first manager gets an empty set of used messages instead of
        // merging the received ones, so it sends no Confirmation of its own.
        let (faulty_manager, peer_managers) = randomness_managers.split_first_mut().unwrap();
        faulty_manager
            .used_messages
            .set(VersionedUsedProcessedMessages::V1(
                dkg_v1::UsedProcessedMessages(vec![]),
            ))
            .unwrap();

        // The other managers process the Messages and send their Confirmations.
        for randomness_manager in peer_managers.iter_mut() {
            let mut output = ConsensusCommitOutput::new(0, 0);
            for (j, dkg_message) in dkg_messages.iter().cloned().enumerate() {
                randomness_manager
                    .add_message(&epoch_stores[j].name, dkg_message)
                    .unwrap();
            }
            randomness_manager
                .advance_dkg(&mut output, 0)
                .await
                .unwrap();
        }
        let mut dkg_confirmations = Vec::new();
        for _ in 0..peer_managers.len() {
            let mut dkg_confirmation = rx_consensus.recv().await.unwrap();
            assert_eq!(dkg_confirmation.len(), 1);
            match dkg_confirmation.remove(0).kind {
                ConsensusTransactionKind::RandomnessDkgConfirmation(_, bytes) => {
                    let msg: VersionedDkgConfirmation = bcs::from_bytes(&bytes)
                        .expect("DKG confirmation deserialization should not fail");
                    dkg_confirmations.push(msg);
                }
                _ => panic!("wrong type of message sent"),
            }
        }

        // Three of four Confirmations are enough for completion to run, and
        // the empty set of used messages leaves it nothing to aggregate.
        let mut output = ConsensusCommitOutput::new(0, 0);
        for (dkg_confirmation, epoch_store) in dkg_confirmations.into_iter().zip(&epoch_stores[1..])
        {
            faulty_manager
                .add_confirmation(&mut output, &epoch_store.name, dkg_confirmation)
                .unwrap();
        }
        assert_eq!(DkgStatus::Pending, faulty_manager.dkg_status());
        let _ = faulty_manager.advance_dkg(&mut output, 0).await;
    }
}