near-primitives 0.35.1

use crate::bandwidth_scheduler::{
    Bandwidth, BandwidthRequest, BandwidthRequestValues, BandwidthRequests,
    BandwidthSchedulerParams, BlockBandwidthRequests,
};
use borsh::{BorshDeserialize, BorshSerialize};
use near_primitives_core::types::{Balance, BlockHeight, Gas, ShardId};
use near_schema_checker_lib::ProtocolSchema;
use std::collections::BTreeMap;

/// Information gathered during chunk application.
/// Provides insight into what happened when the chunk was applied.
/// How many transactions and receipts were processed, buffered, forwarded, etc.
/// Useful for debugging, metrics and sanity checks.
#[derive(Debug, Clone, BorshSerialize, BorshDeserialize, ProtocolSchema)]
#[borsh(use_discriminant = true)]
#[repr(u8)]
pub enum ChunkApplyStats {
    V0(ChunkApplyStatsV0) = 0,
}

/// Information gathered during chunk application.
/// This feature is still in development. Consider V0 as unstable, fields might be added or removed
/// from it at any time. We will do proper versioning after stabilization when there will be other
/// services depending on this structure.
#[derive(Debug, Clone, BorshSerialize, BorshDeserialize, ProtocolSchema)]
pub struct ChunkApplyStatsV0 {
    /// Height at which the chunk was applied
    pub height: BlockHeight,
    /// Shard ID of the chunk
    pub shard_id: ShardId,
    /// Was this chunk applied as a new (non-missing) chunk or a missing one (apply_old_chunk)?
    pub is_new_chunk: bool,
    /// Number of new transactions in this chunk
    pub transactions_num: u64,
    /// Number of incoming receipts to this chunk
    pub incoming_receipts_num: u64,

    /// Receipt sink stats - forwarded receipts, buffered receipts, outgoing limits
    pub receipt_sink: ReceiptSinkStats,
    /// Bandwidth scheduler stats
    pub bandwidth_scheduler: BandwidthSchedulerStats,
    /// Balance stats - used in balance checker.
    pub balance: BalanceStats,
}

impl ChunkApplyStatsV0 {
    pub fn new(height: BlockHeight, shard_id: ShardId) -> ChunkApplyStatsV0 {
        ChunkApplyStatsV0 {
            height: height,
            shard_id: shard_id,
            is_new_chunk: false,
            transactions_num: 0,
            incoming_receipts_num: 0,
            bandwidth_scheduler: Default::default(),
            balance: Default::default(),
            receipt_sink: Default::default(),
        }
    }

    pub fn set_new_bandwidth_requests(
        &mut self,
        requests: &BandwidthRequests,
        params: &BandwidthSchedulerParams,
    ) {
        self.bandwidth_scheduler.set_new_bandwidth_requests(self.shard_id, requests, params);
    }

    /// Dummy data for tests.
    pub fn dummy() -> ChunkApplyStatsV0 {
        ChunkApplyStatsV0 {
            height: 0,
            shard_id: ShardId::new(0),
            is_new_chunk: false,
            transactions_num: 0,
            incoming_receipts_num: 0,
            bandwidth_scheduler: Default::default(),
            balance: Default::default(),
            receipt_sink: Default::default(),
        }
    }
}

#[derive(Debug, Clone, Default, BorshSerialize, BorshDeserialize, ProtocolSchema)]
pub struct BandwidthSchedulerStats {
    /// Scheduler params, should always be Some but there is no Default impl.
    pub params: Option<BandwidthSchedulerParams>,
    /// Bandwidth requests generated by previous chunks, used as input to bandwidth scheduler.
    pub prev_bandwidth_requests: BTreeMap<(ShardId, ShardId), Vec<Bandwidth>>,
    /// Number of previous bandwidth requests (prev_bandwidth_requests.len()).
    pub prev_bandwidth_requests_num: u64,
    /// How long it took to run the bandwidth scheduler (in milliseconds).
    pub time_to_run_ms: u128,
    /// Bandwidth granted by the scheduler.
    pub granted_bandwidth: BTreeMap<(ShardId, ShardId), Bandwidth>,
    /// Bandwidth requests generated at the end of chunk application.
    pub new_bandwidth_requests: BTreeMap<(ShardId, ShardId), Vec<Bandwidth>>,
}

impl BandwidthSchedulerStats {
    /// Set `prev_bandwidth_requests` and `prev_bandwidth_requests_num`. Automatically converts to
    /// the target representation.
    pub fn set_prev_bandwidth_requests(
        &mut self,
        requests: &BlockBandwidthRequests,
        params: &BandwidthSchedulerParams,
    ) {
        for (from_shard, shard_requests) in &requests.shards_bandwidth_requests {
            Self::add_requests_to_map(
                shard_requests,
                &mut self.prev_bandwidth_requests,
                *from_shard,
                params,
            );
        }
        self.prev_bandwidth_requests_num = self.prev_bandwidth_requests.len().try_into().unwrap();
    }

    /// Set `new_bandwidth_requests`. Automatically converts to the target representation.
    pub fn set_new_bandwidth_requests(
        &mut self,
        from_shard: ShardId,
        requests: &BandwidthRequests,
        params: &BandwidthSchedulerParams,
    ) {
        Self::add_requests_to_map(requests, &mut self.new_bandwidth_requests, from_shard, params);
    }

    /// Convert bandwidth requests to the target representation and add to the given map.
    fn add_requests_to_map(
        requests: &BandwidthRequests,
        map: &mut BTreeMap<(ShardId, ShardId), Vec<u64>>,
        from_shard: ShardId,
        params: &BandwidthSchedulerParams,
    ) {
        match requests {
            BandwidthRequests::V1(requests_v1) => {
                for request in &requests_v1.requests {
                    map.insert(
                        (from_shard, request.to_shard.into()),
                        get_requested_values(request, params),
                    );
                }
            }
        }
    }
}

#[derive(Debug, Clone, Default, BorshSerialize, BorshDeserialize, ProtocolSchema)]
pub struct ReceiptSinkStats {
    /// Outgoing size and gas limits to every shard.
    pub outgoing_limits: BTreeMap<ShardId, OutgoingLimitStats>,
    /// New outgoing receipts generated during this chunk application.
    pub forwarded_receipts: BTreeMap<ShardId, ReceiptsStats>,
    /// New buffered receipts that couldn't be forwarded because of the outgoing limits.
    pub buffered_receipts: BTreeMap<ShardId, ReceiptsStats>,
    /// Final state of the outgoing buffers after all new receipts have been forwarded or buffered.
    /// Used to generate new bandwidth requests.
    pub final_outgoing_buffers: BTreeMap<ShardId, ReceiptsStats>,
    /// Whether the `ReceiptGroupsQueue` is fully initialized. Can only be false during the protocol
    /// upgrade that enables bandwidth scheduler.
    pub is_outgoing_metadata_ready: BTreeMap<ShardId, bool>,
    /// Whether `is_outgoing_metadata_ready` is true for all shards. This must be true before
    /// resharding can start.
    pub all_outgoing_metadatas_ready: bool,
}

impl ReceiptSinkStats {
    pub fn set_outgoing_limits(&mut self, limits: impl Iterator<Item = (ShardId, (u64, Gas))>) {
        for (shard_id, (size, gas)) in limits {
            self.outgoing_limits.insert(shard_id, OutgoingLimitStats { size, gas });
        }
    }
}

#[derive(Debug, Clone, Default, BorshSerialize, BorshDeserialize, ProtocolSchema)]
pub struct OutgoingLimitStats {
    pub size: u64,
    pub gas: Gas,
}

/// Stats about a set of receipts
#[derive(Debug, Clone, Default, BorshSerialize, BorshDeserialize, ProtocolSchema)]
pub struct ReceiptsStats {
    /// Number of receipts
    pub num: u64,
    /// Total size of receipts, as calculated by `congestion_control::compute_receipt_size`.
    pub total_size: u64,
    /// Total gas of receipts, as calculated by `compute_receipt_congestion_gas`.
    pub total_gas: u128,
}

impl ReceiptsStats {
    pub fn add_receipt(&mut self, size: u64, gas: Gas) {
        self.num += 1;
        self.total_size += size;
        let gas_u128: u128 = gas.as_gas().into();
        self.total_gas += gas_u128;
    }
}

/// Stats about token balance, used in balance checker.
#[derive(Debug, Clone, Default, BorshSerialize, BorshDeserialize, ProtocolSchema)]
pub struct BalanceStats {
    pub tx_burnt_amount: Balance,
    pub slashed_burnt_amount: Balance,
    pub other_burnt_amount: Balance,
    /// This is a negative amount. This amount was not charged from the account that issued
    /// the transaction. It's likely due to the delayed queue of the receipts.
    pub gas_deficit_amount: Balance,
    /// No longer used, keeping to preserve borsh deserialization of the old data in the db.
    pub _deprecated_global_actions_burnt_amount: Balance,
}

/// Convert a bandwidth request from the bitmap representation to a list of requested values.
fn get_requested_values(
    bandwidth_request: &BandwidthRequest,
    params: &BandwidthSchedulerParams,
) -> Vec<Bandwidth> {
    let values = BandwidthRequestValues::new(params);
    let mut res = Vec::new();
    for i in 0..bandwidth_request.requested_values_bitmap.len() {
        if bandwidth_request.requested_values_bitmap.get_bit(i) {
            res.push(values.values[i]);
        }
    }
    res
}