consensus_core/

transaction.rs

// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0
use std::{collections::BTreeMap, sync::Arc};

use consensus_config::Epoch;
use consensus_types::block::{
    BlockRef, NUM_RESERVED_TRANSACTION_INDICES, PING_TRANSACTION_INDEX, Round, TransactionIndex,
};
use mysten_common::debug_fatal;
use mysten_metrics::monitored_mpsc::{Receiver, Sender, channel};
use parking_lot::Mutex;
use tap::TapFallible;
use thiserror::Error;
use tokio::sync::oneshot;
use tracing::{error, warn};

use crate::{block::Transaction, context::Context};

/// The maximum number of transactions pending to the queue to be pulled for block proposal
const MAX_PENDING_TRANSACTIONS: usize = 2_000;

/// Priority of a submission to consensus. Consensus is agnostic to transaction type; the
/// submitter decides the priority. `High` submissions use a dedicated, reserved lane and
/// are pulled ahead of `Normal` ones for block proposal.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum Priority {
    Normal,
    High,
}

/// Reserved capacity for the local validator's own high-priority submissions, which are
/// low-volume and pulled ahead of normal-priority ones for block proposal.
const MAX_PENDING_PRIORITY_TRANSACTIONS: usize = 128;

/// The guard acts as an acknowledgment mechanism for the inclusion of the transactions to a block.
/// When its last transaction is included to a block then `included_in_block_ack` will be signalled.
/// If the guard is dropped without getting acknowledged that means the transactions have not been
/// included to a block and the consensus is shutting down.
pub(crate) struct TransactionsGuard {
    // Holds a list of transactions to be included in the block.
    // A TransactionsGuard may be partially consumed by `TransactionConsumer`, in which case, this holds the remaining transactions.
    transactions: Vec<Transaction>,

    // When the transactions are included in a block, this will be signalled with
    // the following information
    included_in_block_ack: oneshot::Sender<(
        // The block reference in which the transactions have been included
        BlockRef,
        // The indices of the transactions that have been included in the block
        Vec<TransactionIndex>,
        // A receiver to notify the submitter about the block status
        oneshot::Receiver<BlockStatus>,
    )>,
}

/// The TransactionConsumer is responsible for fetching the next transactions to be included for the block proposals.
/// The transactions are submitted to a channel which is shared between the TransactionConsumer and the TransactionClient
/// and are pulled every time the `next` method is called.
pub(crate) struct TransactionConsumer {
    tx_receiver: Receiver<TransactionsGuard>,
    // Reserved lane for the local validator's own high-priority submissions, drained
    // ahead of `tx_receiver` so they are never delayed behind a normal-priority backlog.
    priority_tx_receiver: Receiver<TransactionsGuard>,
    max_transactions_in_block_bytes: u64,
    max_num_transactions_in_block: u64,
    pending_transactions: Option<TransactionsGuard>,
    block_status_subscribers: Arc<Mutex<BTreeMap<BlockRef, Vec<oneshot::Sender<BlockStatus>>>>>,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum BlockStatus {
    /// The block has been sequenced as part of a committed sub dag. That means that any transaction that has been included in the block
    /// has been committed as well.
    Sequenced(BlockRef),
    /// The block has been garbage collected and will never be committed. Any transactions that have been included in the block should also
    /// be considered as impossible to be committed as part of this block and might need to be retried
    GarbageCollected(BlockRef),
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum LimitReached {
    // The maximum number of transactions have been included
    MaxNumOfTransactions,
    // The maximum number of bytes have been included
    MaxBytes,
    // All available transactions have been included
    AllTransactionsIncluded,
}

impl TransactionConsumer {
    pub(crate) fn new(
        tx_receiver: Receiver<TransactionsGuard>,
        priority_tx_receiver: Receiver<TransactionsGuard>,
        context: Arc<Context>,
    ) -> Self {
        // max_num_transactions_in_block - 1 is the max possible transaction index in a block.
        // TransactionIndex::MAX is reserved for the ping transaction.
        // Indexes down to TransactionIndex::MAX - 8 are also reserved for future use.
        // This check makes sure they do not overlap.
        assert!(
            context
                .protocol_config
                .max_num_transactions_in_block()
                .saturating_sub(1)
                < TransactionIndex::MAX.saturating_sub(NUM_RESERVED_TRANSACTION_INDICES) as u64,
            "Unsupported max_num_transactions_in_block: {}",
            context.protocol_config.max_num_transactions_in_block()
        );

        Self {
            tx_receiver,
            priority_tx_receiver,
            max_transactions_in_block_bytes: context
                .protocol_config
                .max_transactions_in_block_bytes(),
            max_num_transactions_in_block: context.protocol_config.max_num_transactions_in_block(),
            pending_transactions: None,
            block_status_subscribers: Arc::new(Mutex::new(BTreeMap::new())),
        }
    }

    // Attempts to fetch the next transactions that have been submitted for sequence. Respects the `max_transactions_in_block_bytes`
    // and `max_num_transactions_in_block` parameters specified via protocol config.
    // This returns one or more transactions to be included in the block and a callback to acknowledge the inclusion of those transactions.
    // Also returns a `LimitReached` enum to indicate which limit type has been reached.
    pub(crate) fn next(
        &mut self,
    ) -> (
        Vec<Transaction>,
        Box<dyn FnOnce(BlockRef) + Send>,
        LimitReached,
    ) {
        let mut transactions = Vec::new();
        let mut acks = Vec::new();
        let mut total_bytes = 0;
        let mut limit_reached = LimitReached::AllTransactionsIncluded;

        // Handle one batch of incoming transactions from TransactionGuard.
        // The method will return `None` if all the transactions can be included in the block. Otherwise none of the transactions will be
        // included in the block and the method will return the TransactionGuard.
        let mut handle_txs = |t: TransactionsGuard| -> Option<TransactionsGuard> {
            // If no transactions are submitted, it means that the transaction guard represents a ping transaction.
            // In this case, we need to push the `PING_TRANSACTION_INDEX` to the indices vector.
            let transactions_num = t.transactions.len() as u64;
            if transactions_num == 0 {
                acks.push((t.included_in_block_ack, vec![PING_TRANSACTION_INDEX]));
                return None;
            }

            // Check if the total bytes of the transactions exceed the max transactions in block bytes.
            let transactions_bytes =
                t.transactions.iter().map(|t| t.data().len()).sum::<usize>() as u64;
            if total_bytes + transactions_bytes > self.max_transactions_in_block_bytes {
                limit_reached = LimitReached::MaxBytes;
                return Some(t);
            }
            if transactions.len() as u64 + transactions_num > self.max_num_transactions_in_block {
                limit_reached = LimitReached::MaxNumOfTransactions;
                return Some(t);
            }

            total_bytes += transactions_bytes;

            // Calculate indices for this batch
            let start_idx = transactions.len() as TransactionIndex;
            let indices: Vec<TransactionIndex> =
                (start_idx..start_idx + t.transactions.len() as TransactionIndex).collect();

            // The transactions can be consumed, register its ack and transaction
            // indices to be sent with the ack.
            acks.push((t.included_in_block_ack, indices));
            transactions.extend(t.transactions);
            None
        };

        if let Some(t) = self.pending_transactions.take()
            && let Some(pending_transactions) = handle_txs(t)
        {
            debug_fatal!(
                "Previously pending transaction(s) should fit into an empty block! Dropping: {:?}",
                pending_transactions.transactions
            );
        }

        // Pull until we reach the block limit (which may already be reached above).
        // The reserved priority lane is drained first, so this validator's own
        // high-priority submissions get block space ahead of a normal-priority backlog.
        for receiver in [&mut self.priority_tx_receiver, &mut self.tx_receiver] {
            while self.pending_transactions.is_none() {
                if let Ok(t) = receiver.try_recv() {
                    self.pending_transactions = handle_txs(t);
                } else {
                    break;
                }
            }
        }

        let block_status_subscribers = self.block_status_subscribers.clone();
        (
            transactions,
            Box::new(move |block_ref: BlockRef| {
                let mut block_status_subscribers = block_status_subscribers.lock();

                for (ack, tx_indices) in acks {
                    let (status_tx, status_rx) = oneshot::channel();

                    block_status_subscribers
                        .entry(block_ref)
                        .or_default()
                        .push(status_tx);

                    let _ = ack.send((block_ref, tx_indices, status_rx));
                }
            }),
            limit_reached,
        )
    }

    /// Notifies all the transaction submitters who are waiting to receive an update on the status of the block.
    /// The `committed_blocks` are the blocks that have been committed and the `gc_round` is the round up to which the blocks have been garbage collected.
    /// First we'll notify for all the committed blocks, and then for all the blocks that have been garbage collected.
    pub(crate) fn notify_own_blocks_status(
        &self,
        committed_blocks: Vec<BlockRef>,
        gc_round: Round,
    ) {
        // Notify for all the committed blocks first
        let mut block_status_subscribers = self.block_status_subscribers.lock();
        for block_ref in committed_blocks {
            if let Some(subscribers) = block_status_subscribers.remove(&block_ref) {
                subscribers.into_iter().for_each(|s| {
                    let _ = s.send(BlockStatus::Sequenced(block_ref));
                });
            }
        }

        // Now notify everyone <= gc_round that their block has been garbage collected and clean up the entries
        while let Some((block_ref, subscribers)) = block_status_subscribers.pop_first() {
            if block_ref.round <= gc_round {
                subscribers.into_iter().for_each(|s| {
                    let _ = s.send(BlockStatus::GarbageCollected(block_ref));
                });
            } else {
                block_status_subscribers.insert(block_ref, subscribers);
                break;
            }
        }
    }

    #[cfg(test)]
    pub(crate) fn subscribe_for_block_status_testing(
        &self,
        block_ref: BlockRef,
    ) -> oneshot::Receiver<BlockStatus> {
        let (tx, rx) = oneshot::channel();
        let mut block_status_subscribers = self.block_status_subscribers.lock();
        block_status_subscribers
            .entry(block_ref)
            .or_default()
            .push(tx);
        rx
    }

    #[cfg(test)]
    fn is_empty(&mut self) -> bool {
        if self.pending_transactions.is_some() {
            return false;
        }
        for receiver in [&mut self.priority_tx_receiver, &mut self.tx_receiver] {
            if let Ok(t) = receiver.try_recv() {
                self.pending_transactions = Some(t);
                return false;
            }
        }
        true
    }
}

#[derive(Clone)]
pub struct TransactionClient {
    context: Arc<Context>,
    sender: Sender<TransactionsGuard>,
    // Reserved channel for the local validator's own high-priority submissions.
    priority_sender: Sender<TransactionsGuard>,
    max_transaction_size: u64,
    max_transactions_in_block_bytes: u64,
    max_transactions_in_block_count: u64,
}

#[derive(Debug, Error)]
pub enum ClientError {
    #[error("Failed to submit transaction, consensus is shutting down: {0}")]
    ConsensusShuttingDown(String),

    #[error("Transaction size ({0}B) is over limit ({1}B)")]
    OversizedTransaction(u64, u64),

    #[error("Transaction bundle size ({0}B) is over limit ({1}B)")]
    OversizedTransactionBundleBytes(u64, u64),

    #[error("Transaction bundle count ({0}) is over limit ({1})")]
    OversizedTransactionBundleCount(u64, u64),
}

impl TransactionClient {
    /// Returns the client and the receivers for normal and priority transactions, in that order.
    pub(crate) fn new(
        context: Arc<Context>,
    ) -> (
        Self,
        Receiver<TransactionsGuard>,
        Receiver<TransactionsGuard>,
    ) {
        Self::new_with_max_pending_transactions(
            context,
            MAX_PENDING_TRANSACTIONS,
            MAX_PENDING_PRIORITY_TRANSACTIONS,
        )
    }

    /// Returns the client and the receivers for normal and priority transactions, in that order.
    fn new_with_max_pending_transactions(
        context: Arc<Context>,
        max_pending_transactions: usize,
        max_pending_priority_transactions: usize,
    ) -> (
        Self,
        Receiver<TransactionsGuard>,
        Receiver<TransactionsGuard>,
    ) {
        let (sender, receiver) = channel("consensus_input", max_pending_transactions);
        let (priority_sender, priority_receiver) = channel(
            "consensus_input_priority",
            max_pending_priority_transactions,
        );
        (
            Self {
                sender,
                priority_sender,
                max_transaction_size: context.protocol_config.max_transaction_size_bytes(),

                max_transactions_in_block_bytes: context
                    .protocol_config
                    .max_transactions_in_block_bytes(),
                max_transactions_in_block_count: context
                    .protocol_config
                    .max_num_transactions_in_block(),
                context: context.clone(),
            },
            receiver,
            priority_receiver,
        )
    }

    /// Returns the current epoch of this client.
    pub fn epoch(&self) -> Epoch {
        self.context.committee.epoch()
    }

    /// Submits a list of transactions to be sequenced. The method returns when all the transactions have been successfully included
    /// to next proposed blocks.
    ///
    /// If `transactions` is empty, then this will be interpreted as a "ping" signal from the client in order to get information about the next
    /// block and simulate a transaction inclusion to the next block. In this an empty vector of the transaction index will be returned as response
    /// and the block status receiver.
    pub async fn submit(
        &self,
        transactions: Vec<Vec<u8>>,
        priority: Priority,
    ) -> Result<
        (
            BlockRef,
            Vec<TransactionIndex>,
            oneshot::Receiver<BlockStatus>,
        ),
        ClientError,
    > {
        let included_in_block = self.submit_no_wait(transactions, priority).await?;
        included_in_block
            .await
            .tap_err(|e| warn!("Transaction acknowledge failed with {:?}", e))
            .map_err(|e| ClientError::ConsensusShuttingDown(e.to_string()))
    }

    /// Submits a list of transactions to be sequenced.
    /// If any transaction's length exceeds `max_transaction_size`, no transaction will be submitted.
    /// That shouldn't be the common case as sizes should be aligned between consensus and client. The method returns
    /// a receiver to wait on until the transactions has been included in the next block to get proposed. The consumer should
    /// wait on it to consider as inclusion acknowledgement. If the receiver errors then consensus is shutting down and transaction
    /// has not been included to any block.
    /// If multiple transactions are submitted, the method will attempt to bundle them together in a single block. If the total size of
    /// the transactions exceeds `max_transactions_in_block_bytes`, no transaction will be submitted and an error will be returned instead.
    /// Similar if transactions exceed `max_transactions_in_block_count` an error will be returned.
    pub(crate) async fn submit_no_wait(
        &self,
        transactions: Vec<Vec<u8>>,
        priority: Priority,
    ) -> Result<
        oneshot::Receiver<(
            BlockRef,
            Vec<TransactionIndex>,
            oneshot::Receiver<BlockStatus>,
        )>,
        ClientError,
    > {
        let (included_in_block_ack_send, included_in_block_ack_receive) = oneshot::channel();

        let mut bundle_size = 0;

        if transactions.len() as u64 > self.max_transactions_in_block_count {
            return Err(ClientError::OversizedTransactionBundleCount(
                transactions.len() as u64,
                self.max_transactions_in_block_count,
            ));
        }

        for transaction in &transactions {
            if transaction.len() as u64 > self.max_transaction_size {
                return Err(ClientError::OversizedTransaction(
                    transaction.len() as u64,
                    self.max_transaction_size,
                ));
            }
            bundle_size += transaction.len() as u64;

            if bundle_size > self.max_transactions_in_block_bytes {
                return Err(ClientError::OversizedTransactionBundleBytes(
                    bundle_size,
                    self.max_transactions_in_block_bytes,
                ));
            }
        }

        let t = TransactionsGuard {
            transactions: transactions.into_iter().map(Transaction::new).collect(),
            included_in_block_ack: included_in_block_ack_send,
        };
        let sender = match priority {
            Priority::High => &self.priority_sender,
            Priority::Normal => &self.sender,
        };
        // A full reserved priority lane means a high-priority submission has to wait —
        // silently reintroducing the buffering the lane exists to avoid. Record it.
        if priority == Priority::High && sender.capacity() == 0 {
            self.context
                .metrics
                .node_metrics
                .priority_submission_backpressure
                .inc();
        }
        sender
            .send(t)
            .await
            .tap_err(|e| error!("Submit transactions failed with {:?}", e))
            .map_err(|e| ClientError::ConsensusShuttingDown(e.to_string()))?;
        Ok(included_in_block_ack_receive)
    }
}

/// `TransactionPool` supplies transactions for block proposals, as an alternative to
/// submitting transactions through `TransactionClient`. Like `TransactionVerifier`, the
/// implementation can be provided by Sui and passed into `ConsensusAuthority::start()`.
pub trait TransactionPool: Send + Sync + 'static {
    /// Called by the proposer while building a block. Takes transactions to include, up to
    /// `max_count` transactions and `max_bytes` total serialized bytes. Returns the
    /// transactions in block order, an ack callback the proposer invokes with the created
    /// block's reference after the block is durably created, and which limit stopped the
    /// take. Dropping the callback without invoking it means the transactions have not been
    /// included in a block, and the implementation may make them available to take again.
    fn take(
        &self,
        max_count: usize,
        max_bytes: usize,
    ) -> (
        Vec<Transaction>,
        Box<dyn FnOnce(BlockRef) + Send>,
        LimitReached,
    );

    /// Called from the commit path with this authority's own committed block refs and the
    /// current GC round. Own blocks at rounds <= `gc_round` that are not committed are GC'ed
    /// and will never commit.
    fn notify_committed(&self, own_committed_blocks: Vec<BlockRef>, gc_round: Round);
}

/// Adapts the channel-based `TransactionConsumer` to the `TransactionPool` interface, for
/// transactions submitted through `TransactionClient`.
pub(crate) struct TransactionConsumerPool {
    consumer: Mutex<TransactionConsumer>,
}

impl TransactionConsumerPool {
    pub(crate) fn new(consumer: TransactionConsumer) -> Self {
        Self {
            consumer: Mutex::new(consumer),
        }
    }

    #[cfg(test)]
    pub(crate) fn subscribe_for_block_status_testing(
        &self,
        block_ref: BlockRef,
    ) -> oneshot::Receiver<BlockStatus> {
        self.consumer
            .lock()
            .subscribe_for_block_status_testing(block_ref)
    }
}

impl TransactionPool for TransactionConsumerPool {
    // TransactionConsumer enforces the same max limits internally via protocol config.
    fn take(
        &self,
        _max_count: usize,
        _max_bytes: usize,
    ) -> (
        Vec<Transaction>,
        Box<dyn FnOnce(BlockRef) + Send>,
        LimitReached,
    ) {
        self.consumer.lock().next()
    }

    fn notify_committed(&self, own_committed_blocks: Vec<BlockRef>, gc_round: Round) {
        self.consumer
            .lock()
            .notify_own_blocks_status(own_committed_blocks, gc_round);
    }
}

/// `TransactionVerifier` implementation is supplied by Sui to validate transactions in a block,
/// before acceptance of the block.
pub trait TransactionVerifier: Send + Sync + 'static {
    /// Determines if this batch of transactions is valid.
    /// Fails if any one of the transactions is invalid.
    fn verify_batch(&self, batch: &[&[u8]]) -> Result<(), ValidationError>;

    /// Returns indices of transactions to reject, or a transaction validation error.
    /// Currently only uncertified user transactions can be voted to reject, which are created
    /// by Mysticeti fastpath client.
    /// Honest validators may disagree on voting for uncertified user transactions.
    /// The other types of transactions are implicitly voted to be accepted if they pass validation.
    ///
    /// Honest validators should produce the same validation outcome on the same batch of
    /// transactions. So if a batch from a peer fails validation, the peer is equivocating.
    fn verify_and_vote_batch(
        &self,
        block_ref: &BlockRef,
        batch: &[&[u8]],
    ) -> Result<Vec<TransactionIndex>, ValidationError>;
}

#[derive(Debug, Error)]
pub enum ValidationError {
    #[error("Invalid transaction: {0}")]
    InvalidTransaction(String),
}

/// `NoopTransactionVerifier` accepts all transactions.
#[cfg(any(test, msim))]
pub struct NoopTransactionVerifier;

#[cfg(any(test, msim))]
impl TransactionVerifier for NoopTransactionVerifier {
    fn verify_batch(&self, _batch: &[&[u8]]) -> Result<(), ValidationError> {
        Ok(())
    }

    fn verify_and_vote_batch(
        &self,
        _block_ref: &BlockRef,
        _batch: &[&[u8]],
    ) -> Result<Vec<TransactionIndex>, ValidationError> {
        Ok(vec![])
    }
}

#[cfg(test)]
mod tests {
    use std::{sync::Arc, time::Duration};

    use consensus_config::AuthorityIndex;
    use consensus_types::block::{
        BlockDigest, BlockRef, NUM_RESERVED_TRANSACTION_INDICES, PING_TRANSACTION_INDEX,
        TransactionIndex,
    };
    use futures::{StreamExt, stream::FuturesUnordered};
    use tokio::time::timeout;

    use crate::transaction::NoopTransactionVerifier;
    use crate::{
        block_verifier::SignedBlockVerifier,
        context::Context,
        transaction::{
            BlockStatus, LimitReached, MAX_PENDING_PRIORITY_TRANSACTIONS, Priority,
            TransactionClient, TransactionConsumer,
        },
    };

    #[tokio::test(flavor = "current_thread", start_paused = true)]
    async fn basic_submit_and_consume() {
        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(2_000); // 2KB
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(2_000);
        let context = Arc::new(context);
        let (client, tx_receiver, priority_tx_receiver) = TransactionClient::new(context.clone());
        let mut consumer =
            TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());

        // submit asynchronously the transactions and keep the waiters
        let mut included_in_block_waiters = FuturesUnordered::new();
        for i in 0..3 {
            let transaction =
                bcs::to_bytes(&format!("transaction {i}")).expect("Serialization should not fail.");
            let w = client
                .submit_no_wait(vec![transaction], Priority::Normal)
                .await
                .expect("Shouldn't submit successfully transaction");
            included_in_block_waiters.push(w);
        }

        // now pull the transactions from the consumer
        let (transactions, ack_transactions, _limit_reached) = consumer.next();
        assert_eq!(transactions.len(), 3);

        for (i, t) in transactions.iter().enumerate() {
            let t: String = bcs::from_bytes(t.data()).unwrap();
            assert_eq!(format!("transaction {i}").to_string(), t);
        }

        assert!(
            timeout(Duration::from_secs(1), included_in_block_waiters.next())
                .await
                .is_err(),
            "We should expect to timeout as none of the transactions have been acknowledged yet"
        );

        // Now acknowledge the inclusion of transactions
        ack_transactions(BlockRef::MIN);

        // Now make sure that all the waiters have returned
        while let Some(result) = included_in_block_waiters.next().await {
            assert!(result.is_ok());
        }

        // try to pull again transactions, result should be empty
        assert!(consumer.is_empty());
    }

    #[tokio::test(flavor = "current_thread", start_paused = true)]
    async fn high_priority_transactions_included_first() {
        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(2_000);
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(2_000);
        let context = Arc::new(context);
        let (client, tx_receiver, priority_tx_receiver) = TransactionClient::new(context.clone());
        let mut consumer =
            TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());

        // Submit normal transactions first, then a high-priority one.
        for i in 0..3 {
            let t = bcs::to_bytes(&format!("normal {i}")).unwrap();
            client
                .submit_no_wait(vec![t], Priority::Normal)
                .await
                .unwrap();
        }
        let high = bcs::to_bytes(&"high".to_string()).unwrap();
        client
            .submit_no_wait(vec![high], Priority::High)
            .await
            .unwrap();

        // The high-priority transaction is pulled ahead of the normal backlog.
        let (transactions, _ack, _limit) = consumer.next();
        let decoded: Vec<String> = transactions
            .iter()
            .map(|t| bcs::from_bytes(t.data()).unwrap())
            .collect();
        assert_eq!(decoded, vec!["high", "normal 0", "normal 1", "normal 2"]);
    }

    #[tokio::test(flavor = "current_thread", start_paused = true)]
    async fn high_priority_bypasses_full_normal_lane() {
        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(2_000);
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(2_000);
        let context = Arc::new(context);
        // Each lane holds a single entry. The receivers are held (never drained) so
        // the lanes stay full once an entry is buffered.
        let (client, _tx_receiver, _priority_tx_receiver) =
            TransactionClient::new_with_max_pending_transactions(context.clone(), 1, 1);

        // Fill the normal lane.
        let n0 = bcs::to_bytes(&"n0".to_string()).unwrap();
        client
            .submit_no_wait(vec![n0], Priority::Normal)
            .await
            .unwrap();

        // A further normal submission blocks on the full normal lane.
        let n1 = bcs::to_bytes(&"n1".to_string()).unwrap();
        assert!(
            timeout(
                Duration::from_millis(100),
                client.submit_no_wait(vec![n1], Priority::Normal)
            )
            .await
            .is_err(),
            "normal lane is full, so this submission must block"
        );

        // A high-priority submission uses the reserved lane and is not blocked.
        let h = bcs::to_bytes(&"h".to_string()).unwrap();
        timeout(
            Duration::from_millis(100),
            client.submit_no_wait(vec![h], Priority::High),
        )
        .await
        .expect("high-priority submission must not block on a full normal lane")
        .expect("high-priority submission should succeed");
    }

    #[tokio::test(flavor = "current_thread", start_paused = true)]
    async fn priority_overflow_held_and_drained_before_normal() {
        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(2_000);
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(2_000);
        // Only two transactions fit per block.
        context
            .protocol_config
            .set_max_num_transactions_in_block_for_testing(2);
        let context = Arc::new(context);
        let (client, tx_receiver, priority_tx_receiver) = TransactionClient::new(context.clone());
        let mut consumer =
            TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());

        for i in 0..3 {
            let t = bcs::to_bytes(&format!("p{i}")).unwrap();
            client
                .submit_no_wait(vec![t], Priority::High)
                .await
                .unwrap();
        }
        let n = bcs::to_bytes(&"n0".to_string()).unwrap();
        client
            .submit_no_wait(vec![n], Priority::Normal)
            .await
            .unwrap();

        // Block 1: only the first two priority txns fit; the rest is held.
        let (txs, ack, limit) = consumer.next();
        let got: Vec<String> = txs
            .iter()
            .map(|t| bcs::from_bytes(t.data()).unwrap())
            .collect();
        assert_eq!(got, vec!["p0", "p1"]);
        assert_eq!(limit, LimitReached::MaxNumOfTransactions);
        ack(BlockRef::MIN);

        // Block 2: the held priority txn is drained ahead of the normal one.
        let (txs, ack, _) = consumer.next();
        let got: Vec<String> = txs
            .iter()
            .map(|t| bcs::from_bytes(t.data()).unwrap())
            .collect();
        assert_eq!(got, vec!["p2", "n0"]);
        ack(BlockRef::MIN);
        assert!(consumer.is_empty());
    }

    #[tokio::test(flavor = "current_thread", start_paused = true)]
    async fn priority_overflow_by_bytes_held_and_re_offered() {
        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(2_000);
        // Each payload below is 11 bytes; only one fits in a 15-byte block.
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(15);
        let context = Arc::new(context);
        let (client, tx_receiver, priority_tx_receiver) = TransactionClient::new(context.clone());
        let mut consumer =
            TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());

        for p in ["AAAAAAAAAA", "BBBBBBBBBB"] {
            let t = bcs::to_bytes(&p.to_string()).unwrap();
            assert_eq!(t.len(), 11);
            client
                .submit_no_wait(vec![t], Priority::High)
                .await
                .unwrap();
        }

        // Block 1: only the first priority txn fits (byte limit); the second is held.
        let (txs, ack, limit) = consumer.next();
        let got: Vec<String> = txs
            .iter()
            .map(|t| bcs::from_bytes(t.data()).unwrap())
            .collect();
        assert_eq!(got, vec!["AAAAAAAAAA"]);
        assert_eq!(limit, LimitReached::MaxBytes);
        ack(BlockRef::MIN);

        // Block 2: the held priority txn is re-offered first.
        let (txs, ack, _) = consumer.next();
        let got: Vec<String> = txs
            .iter()
            .map(|t| bcs::from_bytes(t.data()).unwrap())
            .collect();
        assert_eq!(got, vec!["BBBBBBBBBB"]);
        ack(BlockRef::MIN);
        assert!(consumer.is_empty());
    }

    #[tokio::test(flavor = "current_thread", start_paused = true)]
    async fn interleaved_submissions_ordered_priority_first_then_fifo() {
        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(2_000);
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(2_000);
        let context = Arc::new(context);
        let (client, tx_receiver, priority_tx_receiver) = TransactionClient::new(context.clone());
        let mut consumer =
            TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());

        // Interleave the two lanes.
        for (payload, priority) in [
            ("n0", Priority::Normal),
            ("p0", Priority::High),
            ("n1", Priority::Normal),
            ("p1", Priority::High),
        ] {
            let t = bcs::to_bytes(&payload.to_string()).unwrap();
            client.submit_no_wait(vec![t], priority).await.unwrap();
        }

        // All priority txns first (in submission order), then all normal txns.
        let (txs, _ack, _) = consumer.next();
        let got: Vec<String> = txs
            .iter()
            .map(|t| bcs::from_bytes(t.data()).unwrap())
            .collect();
        assert_eq!(got, vec!["p0", "p1", "n0", "n1"]);
    }

    #[tokio::test(flavor = "current_thread", start_paused = true)]
    async fn block_status_update() {
        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(2_000); // 2KB
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(2_000);
        context.protocol_config.set_gc_depth_for_testing(10);
        let context = Arc::new(context);
        let (client, tx_receiver, priority_tx_receiver) = TransactionClient::new(context.clone());
        let mut consumer =
            TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());

        // submit the transactions and include 2 of each on a new block
        let mut included_in_block_waiters = FuturesUnordered::new();
        for i in 1..=10 {
            let transaction =
                bcs::to_bytes(&format!("transaction {i}")).expect("Serialization should not fail.");
            let w = client
                .submit_no_wait(vec![transaction], Priority::Normal)
                .await
                .expect("Shouldn't submit successfully transaction");
            included_in_block_waiters.push(w);

            // Every 2 transactions simulate the creation of a new block and acknowledge the inclusion of the transactions
            if i % 2 == 0 {
                let (transactions, ack_transactions, _limit_reached) = consumer.next();
                assert_eq!(transactions.len(), 2);
                ack_transactions(BlockRef::new(
                    i,
                    AuthorityIndex::new_for_test(0),
                    BlockDigest::MIN,
                ));
            }
        }

        let mut transaction_count = 0;
        // Now iterate over all the waiters. Everyone should have been acknowledged.
        let mut block_status_waiters = Vec::new();
        while let Some(result) = included_in_block_waiters.next().await {
            let (block_ref, tx_indices, block_status_waiter) =
                result.expect("Block inclusion waiter shouldn't fail");
            // tx is submitted one at a time so tx acks should only return one tx index
            assert_eq!(tx_indices.len(), 1);
            // The first transaction in the block should have index 0, the second one 1, etc.
            // because we submit 2 transactions per block, the index should be 0 then 1 and then
            // reset back to 0 for the next block.
            assert_eq!(tx_indices[0], transaction_count % 2);
            transaction_count += 1;

            block_status_waiters.push((block_ref, block_status_waiter));
        }

        // Now acknowledge the commit of the blocks 6, 8, 10 and set gc_round = 5, which should trigger the garbage collection of blocks 1..=5
        let gc_round = 5;
        consumer.notify_own_blocks_status(
            vec![
                BlockRef::new(6, AuthorityIndex::new_for_test(0), BlockDigest::MIN),
                BlockRef::new(8, AuthorityIndex::new_for_test(0), BlockDigest::MIN),
                BlockRef::new(10, AuthorityIndex::new_for_test(0), BlockDigest::MIN),
            ],
            gc_round,
        );

        // Now iterate over all the block status waiters. Everyone should have been notified.
        for (block_ref, waiter) in block_status_waiters {
            let block_status = waiter.await.expect("Block status waiter shouldn't fail");

            if block_ref.round <= gc_round {
                assert!(matches!(block_status, BlockStatus::GarbageCollected(_)))
            } else {
                assert!(matches!(block_status, BlockStatus::Sequenced(_)));
            }
        }

        // Ensure internal structure is clear
        assert!(consumer.block_status_subscribers.lock().is_empty());
    }

    #[tokio::test]
    async fn submit_over_max_fetch_size_and_consume() {
        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(100);
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(100);
        let context = Arc::new(context);
        let (client, tx_receiver, priority_tx_receiver) = TransactionClient::new(context.clone());
        let mut consumer =
            TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());

        // submit some transactions
        for i in 0..10 {
            let transaction =
                bcs::to_bytes(&format!("transaction {i}")).expect("Serialization should not fail.");
            let _w = client
                .submit_no_wait(vec![transaction], Priority::Normal)
                .await
                .expect("Shouldn't submit successfully transaction");
        }

        // now pull the transactions from the consumer
        let mut all_transactions = Vec::new();
        let (transactions, _ack_transactions, _limit_reached) = consumer.next();
        assert_eq!(transactions.len(), 7);

        // ensure their total size is less than `max_bytes_to_fetch`
        let total_size: u64 = transactions.iter().map(|t| t.data().len() as u64).sum();
        assert!(
            total_size <= context.protocol_config.max_transactions_in_block_bytes(),
            "Should have fetched transactions up to {}",
            context.protocol_config.max_transactions_in_block_bytes()
        );
        all_transactions.extend(transactions);

        // try to pull again transactions, next should be provided
        let (transactions, _ack_transactions, _limit_reached) = consumer.next();
        assert_eq!(transactions.len(), 3);

        // ensure their total size is less than `max_bytes_to_fetch`
        let total_size: u64 = transactions.iter().map(|t| t.data().len() as u64).sum();
        assert!(
            total_size <= context.protocol_config.max_transactions_in_block_bytes(),
            "Should have fetched transactions up to {}",
            context.protocol_config.max_transactions_in_block_bytes()
        );
        all_transactions.extend(transactions);

        // try to pull again transactions, result should be empty
        assert!(consumer.is_empty());

        for (i, t) in all_transactions.iter().enumerate() {
            let t: String = bcs::from_bytes(t.data()).unwrap();
            assert_eq!(format!("transaction {i}").to_string(), t);
        }
    }

    #[tokio::test]
    async fn submit_large_batch_and_ack() {
        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(15);
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(200);
        let context = Arc::new(context);
        let (client, tx_receiver, priority_tx_receiver) = TransactionClient::new(context.clone());
        let mut consumer =
            TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());
        let mut all_receivers = Vec::new();
        // submit a few transactions individually.
        for i in 0..10 {
            let transaction =
                bcs::to_bytes(&format!("transaction {i}")).expect("Serialization should not fail.");
            let w = client
                .submit_no_wait(vec![transaction], Priority::Normal)
                .await
                .expect("Should submit successfully transaction");
            all_receivers.push(w);
        }

        // construct an acceptable batch and submit, it should be accepted
        {
            let transactions: Vec<_> = (10..15)
                .map(|i| {
                    bcs::to_bytes(&format!("transaction {i}"))
                        .expect("Serialization should not fail.")
                })
                .collect();
            let w = client
                .submit_no_wait(transactions, Priority::Normal)
                .await
                .expect("Should submit successfully transaction");
            all_receivers.push(w);
        }

        // submit another individual transaction.
        {
            let i = 15;
            let transaction =
                bcs::to_bytes(&format!("transaction {i}")).expect("Serialization should not fail.");
            let w = client
                .submit_no_wait(vec![transaction], Priority::Normal)
                .await
                .expect("Shouldn't submit successfully transaction");
            all_receivers.push(w);
        }

        // construct a over-size-limit batch and submit, it should not be accepted
        {
            let transactions: Vec<_> = (16..32)
                .map(|i| {
                    bcs::to_bytes(&format!("transaction {i}"))
                        .expect("Serialization should not fail.")
                })
                .collect();
            let result = client
                .submit_no_wait(transactions, Priority::Normal)
                .await
                .unwrap_err();
            assert_eq!(
                result.to_string(),
                "Transaction bundle size (210B) is over limit (200B)"
            );
        }

        // now pull the transactions from the consumer.
        // we expect all transactions are fetched in order, not missing any, and not exceeding the size limit.
        let mut all_acks: Vec<Box<dyn FnOnce(BlockRef)>> = Vec::new();
        let mut batch_index = 0;
        while !consumer.is_empty() {
            let (transactions, ack_transactions, _limit_reached) = consumer.next();

            assert!(
                transactions.len() as u64
                    <= context.protocol_config.max_num_transactions_in_block(),
                "Should have fetched transactions up to {}",
                context.protocol_config.max_num_transactions_in_block()
            );

            let total_size: u64 = transactions.iter().map(|t| t.data().len() as u64).sum();
            assert!(
                total_size <= context.protocol_config.max_transactions_in_block_bytes(),
                "Should have fetched transactions up to {}",
                context.protocol_config.max_transactions_in_block_bytes()
            );

            // first batch should contain all transactions from 0..10. The softbundle it is to big to fit as well, so it's parked.
            if batch_index == 0 {
                assert_eq!(transactions.len(), 10);
                for (i, transaction) in transactions.iter().enumerate() {
                    let t: String = bcs::from_bytes(transaction.data()).unwrap();
                    assert_eq!(format!("transaction {}", i).to_string(), t);
                }
            // second batch will contain the soft bundle and the additional last transaction.
            } else if batch_index == 1 {
                assert_eq!(transactions.len(), 6);
                for (i, transaction) in transactions.iter().enumerate() {
                    let t: String = bcs::from_bytes(transaction.data()).unwrap();
                    assert_eq!(format!("transaction {}", i + 10).to_string(), t);
                }
            } else {
                panic!("Unexpected batch index");
            }

            batch_index += 1;

            all_acks.push(ack_transactions);
        }

        // now acknowledge the inclusion of all transactions.
        for ack in all_acks {
            ack(BlockRef::MIN);
        }

        // expect all receivers to be resolved.
        for w in all_receivers {
            let r = w.await;
            assert!(r.is_ok());
        }
    }

    #[tokio::test]
    async fn test_submit_over_max_block_size_and_validate_block_size() {
        // submit transactions individually so we make sure that we have reached the block size limit of 10
        {
            let (mut context, _) = Context::new_for_test(4);
            context
                .protocol_config
                .set_max_transaction_size_bytes_for_testing(100);
            context
                .protocol_config
                .set_max_num_transactions_in_block_for_testing(10);
            context
                .protocol_config
                .set_max_transactions_in_block_bytes_for_testing(300);
            let context = Arc::new(context);
            let (client, tx_receiver, priority_tx_receiver) =
                TransactionClient::new(context.clone());
            let mut consumer =
                TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());
            let mut all_receivers = Vec::new();

            // create enough transactions
            let max_num_transactions_in_block =
                context.protocol_config.max_num_transactions_in_block();
            for i in 0..2 * max_num_transactions_in_block {
                let transaction = bcs::to_bytes(&format!("transaction {i}"))
                    .expect("Serialization should not fail.");
                let w = client
                    .submit_no_wait(vec![transaction], Priority::Normal)
                    .await
                    .expect("Should submit successfully transaction");
                all_receivers.push(w);
            }

            // Fetch the next transactions to be included in a block
            let (transactions, _ack_transactions, limit) = consumer.next();
            assert_eq!(limit, LimitReached::MaxNumOfTransactions);
            assert_eq!(transactions.len() as u64, max_num_transactions_in_block);

            // Now create a block and verify that transactions are within the size limits
            let block_verifier =
                SignedBlockVerifier::new(context.clone(), Arc::new(NoopTransactionVerifier {}));

            let batch: Vec<_> = transactions.iter().map(|t| t.data()).collect();
            assert!(
                block_verifier.check_transactions(&batch).is_ok(),
                "Number of transactions limit verification failed"
            );
        }

        // submit transactions individually so we make sure that we have reached the block size bytes 300
        {
            let (mut context, _) = Context::new_for_test(4);
            context
                .protocol_config
                .set_max_transaction_size_bytes_for_testing(100);
            context
                .protocol_config
                .set_max_num_transactions_in_block_for_testing(1_000);
            context
                .protocol_config
                .set_max_transactions_in_block_bytes_for_testing(300);
            let context = Arc::new(context);
            let (client, tx_receiver, priority_tx_receiver) =
                TransactionClient::new(context.clone());
            let mut consumer =
                TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());
            let mut all_receivers = Vec::new();

            let max_transactions_in_block_bytes =
                context.protocol_config.max_transactions_in_block_bytes();
            let mut total_size = 0;
            loop {
                let transaction = bcs::to_bytes(&"transaction".to_string())
                    .expect("Serialization should not fail.");
                total_size += transaction.len() as u64;
                let w = client
                    .submit_no_wait(vec![transaction], Priority::Normal)
                    .await
                    .expect("Should submit successfully transaction");
                all_receivers.push(w);

                // create enough transactions to reach the block size limit
                if total_size >= 2 * max_transactions_in_block_bytes {
                    break;
                }
            }

            // Fetch the next transactions to be included in a block
            let (transactions, _ack_transactions, limit) = consumer.next();
            let batch: Vec<_> = transactions.iter().map(|t| t.data()).collect();
            let size = batch.iter().map(|t| t.len() as u64).sum::<u64>();

            assert_eq!(limit, LimitReached::MaxBytes);
            assert!(
                batch.len() < context.protocol_config.max_num_transactions_in_block() as usize,
                "Should have submitted less than the max number of transactions in a block"
            );
            assert!(size <= max_transactions_in_block_bytes);

            // Now create a block and verify that transactions are within the size limits
            let block_verifier =
                SignedBlockVerifier::new(context.clone(), Arc::new(NoopTransactionVerifier {}));

            assert!(
                block_verifier.check_transactions(&batch).is_ok(),
                "Total size of transactions limit verification failed"
            );
        }
    }

    // This is the case where the client submits a "ping" signal to the consensus to get information about the next block and simulate a transaction inclusion to the next block.
    #[tokio::test]
    async fn submit_with_no_transactions() {
        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(15);
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(200);
        let context = Arc::new(context);
        let (client, tx_receiver, priority_tx_receiver) = TransactionClient::new(context.clone());
        let mut consumer =
            TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());

        let w_no_transactions = client
            .submit_no_wait(vec![], Priority::Normal)
            .await
            .expect("Should submit successfully empty array of transactions");

        let transaction =
            bcs::to_bytes(&"transaction".to_string()).expect("Serialization should not fail.");
        let w_with_transactions = client
            .submit_no_wait(vec![transaction], Priority::Normal)
            .await
            .expect("Should submit successfully transaction");

        let (transactions, ack_transactions, _limit_reached) = consumer.next();
        assert_eq!(transactions.len(), 1);

        // Acknowledge the inclusion of the transactions
        ack_transactions(BlockRef::MIN);

        {
            let r = w_no_transactions.await;
            let (block_ref, indices, _status) = r.unwrap();
            assert_eq!(block_ref, BlockRef::MIN);
            assert_eq!(indices, vec![PING_TRANSACTION_INDEX]);
        }

        {
            let r = w_with_transactions.await;
            let (block_ref, indices, _status) = r.unwrap();
            assert_eq!(block_ref, BlockRef::MIN);
            assert_eq!(indices, vec![0]);
        }
    }

    #[tokio::test]
    async fn ping_transaction_index_never_reached() {
        // Set the max number of transactions in a block to the max value of u16.
        static MAX_NUM_TRANSACTIONS_IN_BLOCK: u64 =
            (TransactionIndex::MAX - NUM_RESERVED_TRANSACTION_INDICES) as u64;

        // Ensure that enough space is allocated in the channel for the pending transactions, so we don't end up consuming the transactions in chunks.
        static MAX_PENDING_TRANSACTIONS: usize = 2 * MAX_NUM_TRANSACTIONS_IN_BLOCK as usize;

        let (mut context, _) = Context::new_for_test(4);
        context
            .protocol_config
            .set_max_transaction_size_bytes_for_testing(200_000);
        context
            .protocol_config
            .set_max_transactions_in_block_bytes_for_testing(1_000_000);
        context
            .protocol_config
            .set_max_num_transactions_in_block_for_testing(MAX_NUM_TRANSACTIONS_IN_BLOCK);
        let context = Arc::new(context);
        let (client, tx_receiver, priority_tx_receiver) =
            TransactionClient::new_with_max_pending_transactions(
                context.clone(),
                MAX_PENDING_TRANSACTIONS,
                MAX_PENDING_PRIORITY_TRANSACTIONS,
            );
        let mut consumer =
            TransactionConsumer::new(tx_receiver, priority_tx_receiver, context.clone());

        // Add 10 more transactions than the max number of transactions in a block.
        for i in 0..MAX_NUM_TRANSACTIONS_IN_BLOCK + 10 {
            println!("Submitting transaction {i}");
            let transaction =
                bcs::to_bytes(&format!("t {i}")).expect("Serialization should not fail.");
            let _w = client
                .submit_no_wait(vec![transaction], Priority::Normal)
                .await
                .expect("Shouldn't submit successfully transaction");
        }

        // now pull the transactions from the consumer
        let (transactions, _ack_transactions, _limit_reached) = consumer.next();
        assert_eq!(transactions.len() as u64, MAX_NUM_TRANSACTIONS_IN_BLOCK);

        let t: String = bcs::from_bytes(transactions.last().unwrap().data()).unwrap();
        assert_eq!(
            t,
            format!(
                "t {}",
                PING_TRANSACTION_INDEX - NUM_RESERVED_TRANSACTION_INDICES - 1
            )
        );
    }
}