solana_cost_model/
cost_tracker.rs

1//! `cost_tracker` keeps tracking transaction cost per chained accounts as well as for entire block
2//! The main functions are:
3//! - would_fit(&tx_cost), immutable function to test if tx with tx_cost would fit into current block
4//! - add_transaction_cost(&tx_cost), mutable function to accumulate tx_cost to tracker.
5//!
6use {
7    crate::{block_cost_limits::*, transaction_cost::TransactionCost},
8    solana_metrics::datapoint_info,
9    solana_sdk::{
10        clock::Slot, pubkey::Pubkey, saturating_add_assign, transaction::TransactionError,
11    },
12    solana_svm_transaction::svm_message::SVMMessage,
13    std::{cmp::Ordering, collections::HashMap},
14};
15
16const WRITABLE_ACCOUNTS_PER_BLOCK: usize = 4096;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum CostTrackerError {
20    /// would exceed block max limit
21    WouldExceedBlockMaxLimit,
22
23    /// would exceed vote max limit
24    WouldExceedVoteMaxLimit,
25
26    /// would exceed account max limit
27    WouldExceedAccountMaxLimit,
28
29    /// would exceed account data block limit
30    WouldExceedAccountDataBlockLimit,
31
32    /// would exceed account data total limit
33    WouldExceedAccountDataTotalLimit,
34}
35
36impl From<CostTrackerError> for TransactionError {
37    fn from(err: CostTrackerError) -> Self {
38        match err {
39            CostTrackerError::WouldExceedBlockMaxLimit => Self::WouldExceedMaxBlockCostLimit,
40            CostTrackerError::WouldExceedVoteMaxLimit => Self::WouldExceedMaxVoteCostLimit,
41            CostTrackerError::WouldExceedAccountMaxLimit => Self::WouldExceedMaxAccountCostLimit,
42            CostTrackerError::WouldExceedAccountDataBlockLimit => {
43                Self::WouldExceedAccountDataBlockLimit
44            }
45            CostTrackerError::WouldExceedAccountDataTotalLimit => {
46                Self::WouldExceedAccountDataTotalLimit
47            }
48        }
49    }
50}
51
52/// Relevant block costs that were updated after successful `try_add()`
53#[derive(Debug, Default)]
54pub struct UpdatedCosts {
55    pub updated_block_cost: u64,
56    // for all write-locked accounts `try_add()` successfully updated, the highest account cost
57    // can be useful info.
58    pub updated_costliest_account_cost: u64,
59}
60
61#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
62#[derive(Debug)]
63pub struct CostTracker {
64    account_cost_limit: u64,
65    block_cost_limit: u64,
66    vote_cost_limit: u64,
67    cost_by_writable_accounts: HashMap<Pubkey, u64, ahash::RandomState>,
68    block_cost: u64,
69    vote_cost: u64,
70    transaction_count: u64,
71    allocated_accounts_data_size: u64,
72    transaction_signature_count: u64,
73    secp256k1_instruction_signature_count: u64,
74    ed25519_instruction_signature_count: u64,
75    /// The number of transactions that have had their estimated cost added to
76    /// the tracker, but are still waiting for an update with actual usage or
77    /// removal if the transaction does not end up getting committed.
78    in_flight_transaction_count: usize,
79}
80
81impl Default for CostTracker {
82    fn default() -> Self {
83        // Clippy doesn't like asserts in const contexts, so need to explicitly allow them.  For
84        // more info, see this issue: https://github.com/rust-lang/rust-clippy/issues/8159
85        #![allow(clippy::assertions_on_constants)]
86        const _: () = assert!(MAX_WRITABLE_ACCOUNT_UNITS <= MAX_BLOCK_UNITS);
87        const _: () = assert!(MAX_VOTE_UNITS <= MAX_BLOCK_UNITS);
88
89        Self {
90            account_cost_limit: MAX_WRITABLE_ACCOUNT_UNITS,
91            block_cost_limit: MAX_BLOCK_UNITS,
92            vote_cost_limit: MAX_VOTE_UNITS,
93            cost_by_writable_accounts: HashMap::with_capacity_and_hasher(
94                WRITABLE_ACCOUNTS_PER_BLOCK,
95                ahash::RandomState::new(),
96            ),
97            block_cost: 0,
98            vote_cost: 0,
99            transaction_count: 0,
100            allocated_accounts_data_size: 0,
101            transaction_signature_count: 0,
102            secp256k1_instruction_signature_count: 0,
103            ed25519_instruction_signature_count: 0,
104            in_flight_transaction_count: 0,
105        }
106    }
107}
108
109impl CostTracker {
110    pub fn new_from_parent_limits(&self) -> Self {
111        let mut new = Self::default();
112        new.set_limits(
113            self.account_cost_limit,
114            self.block_cost_limit,
115            self.vote_cost_limit,
116        );
117        new
118    }
119
120    pub fn reset(&mut self) {
121        self.cost_by_writable_accounts.clear();
122        self.block_cost = 0;
123        self.vote_cost = 0;
124        self.transaction_count = 0;
125        self.allocated_accounts_data_size = 0;
126        self.transaction_signature_count = 0;
127        self.secp256k1_instruction_signature_count = 0;
128        self.ed25519_instruction_signature_count = 0;
129        self.in_flight_transaction_count = 0;
130    }
131
132    /// Get the overall block limit.
133    pub fn get_block_limit(&self) -> u64 {
134        self.block_cost_limit
135    }
136
137    /// allows to adjust limits initiated during construction
138    pub fn set_limits(
139        &mut self,
140        account_cost_limit: u64,
141        block_cost_limit: u64,
142        vote_cost_limit: u64,
143    ) {
144        self.account_cost_limit = account_cost_limit;
145        self.block_cost_limit = block_cost_limit;
146        self.vote_cost_limit = vote_cost_limit;
147    }
148
149    pub fn in_flight_transaction_count(&self) -> usize {
150        self.in_flight_transaction_count
151    }
152
153    pub fn add_transactions_in_flight(&mut self, in_flight_transaction_count: usize) {
154        saturating_add_assign!(
155            self.in_flight_transaction_count,
156            in_flight_transaction_count
157        );
158    }
159
160    pub fn sub_transactions_in_flight(&mut self, in_flight_transaction_count: usize) {
161        self.in_flight_transaction_count = self
162            .in_flight_transaction_count
163            .saturating_sub(in_flight_transaction_count);
164    }
165
166    pub fn try_add(
167        &mut self,
168        tx_cost: &TransactionCost<impl SVMMessage>,
169    ) -> Result<UpdatedCosts, CostTrackerError> {
170        self.would_fit(tx_cost)?;
171        let updated_costliest_account_cost = self.add_transaction_cost(tx_cost);
172        Ok(UpdatedCosts {
173            updated_block_cost: self.block_cost,
174            updated_costliest_account_cost,
175        })
176    }
177
178    pub fn update_execution_cost(
179        &mut self,
180        estimated_tx_cost: &TransactionCost<impl SVMMessage>,
181        actual_execution_units: u64,
182        actual_loaded_accounts_data_size_cost: u64,
183    ) {
184        let actual_load_and_execution_units =
185            actual_execution_units.saturating_add(actual_loaded_accounts_data_size_cost);
186        let estimated_load_and_execution_units = estimated_tx_cost
187            .programs_execution_cost()
188            .saturating_add(estimated_tx_cost.loaded_accounts_data_size_cost());
189        match actual_load_and_execution_units.cmp(&estimated_load_and_execution_units) {
190            Ordering::Equal => (),
191            Ordering::Greater => {
192                self.add_transaction_execution_cost(
193                    estimated_tx_cost,
194                    actual_load_and_execution_units - estimated_load_and_execution_units,
195                );
196            }
197            Ordering::Less => {
198                self.sub_transaction_execution_cost(
199                    estimated_tx_cost,
200                    estimated_load_and_execution_units - actual_load_and_execution_units,
201                );
202            }
203        }
204    }
205
206    pub fn remove(&mut self, tx_cost: &TransactionCost<impl SVMMessage>) {
207        self.remove_transaction_cost(tx_cost);
208    }
209
210    pub fn block_cost(&self) -> u64 {
211        self.block_cost
212    }
213
214    pub fn vote_cost(&self) -> u64 {
215        self.vote_cost
216    }
217
218    pub fn transaction_count(&self) -> u64 {
219        self.transaction_count
220    }
221
222    pub fn report_stats(&self, bank_slot: Slot) {
223        // skip reporting if block is empty
224        if self.transaction_count == 0 {
225            return;
226        }
227
228        let (costliest_account, costliest_account_cost) = self.find_costliest_account();
229
230        datapoint_info!(
231            "cost_tracker_stats",
232            ("bank_slot", bank_slot as i64, i64),
233            ("block_cost", self.block_cost as i64, i64),
234            ("vote_cost", self.vote_cost as i64, i64),
235            ("transaction_count", self.transaction_count as i64, i64),
236            ("number_of_accounts", self.number_of_accounts() as i64, i64),
237            ("costliest_account", costliest_account.to_string(), String),
238            ("costliest_account_cost", costliest_account_cost as i64, i64),
239            (
240                "allocated_accounts_data_size",
241                self.allocated_accounts_data_size,
242                i64
243            ),
244            (
245                "transaction_signature_count",
246                self.transaction_signature_count,
247                i64
248            ),
249            (
250                "secp256k1_instruction_signature_count",
251                self.secp256k1_instruction_signature_count,
252                i64
253            ),
254            (
255                "ed25519_instruction_signature_count",
256                self.ed25519_instruction_signature_count,
257                i64
258            ),
259            (
260                "inflight_transaction_count",
261                self.in_flight_transaction_count,
262                i64
263            ),
264        );
265    }
266
267    fn find_costliest_account(&self) -> (Pubkey, u64) {
268        self.cost_by_writable_accounts
269            .iter()
270            .max_by_key(|(_, &cost)| cost)
271            .map(|(&pubkey, &cost)| (pubkey, cost))
272            .unwrap_or_default()
273    }
274
275    fn would_fit(
276        &self,
277        tx_cost: &TransactionCost<impl SVMMessage>,
278    ) -> Result<(), CostTrackerError> {
279        let cost: u64 = tx_cost.sum();
280
281        if tx_cost.is_simple_vote() {
282            // if vote transaction, check if it exceeds vote_transaction_limit
283            if self.vote_cost.saturating_add(cost) > self.vote_cost_limit {
284                return Err(CostTrackerError::WouldExceedVoteMaxLimit);
285            }
286        }
287
288        if self.block_cost.saturating_add(cost) > self.block_cost_limit {
289            // check against the total package cost
290            return Err(CostTrackerError::WouldExceedBlockMaxLimit);
291        }
292
293        // check if the transaction itself is more costly than the account_cost_limit
294        if cost > self.account_cost_limit {
295            return Err(CostTrackerError::WouldExceedAccountMaxLimit);
296        }
297
298        let allocated_accounts_data_size = self
299            .allocated_accounts_data_size
300            .saturating_add(tx_cost.allocated_accounts_data_size());
301
302        if allocated_accounts_data_size > MAX_BLOCK_ACCOUNTS_DATA_SIZE_DELTA {
303            return Err(CostTrackerError::WouldExceedAccountDataBlockLimit);
304        }
305
306        // check each account against account_cost_limit,
307        for account_key in tx_cost.writable_accounts() {
308            match self.cost_by_writable_accounts.get(account_key) {
309                Some(chained_cost) => {
310                    if chained_cost.saturating_add(cost) > self.account_cost_limit {
311                        return Err(CostTrackerError::WouldExceedAccountMaxLimit);
312                    } else {
313                        continue;
314                    }
315                }
316                None => continue,
317            }
318        }
319
320        Ok(())
321    }
322
323    // Returns the highest account cost for all write-lock accounts `TransactionCost` updated
324    fn add_transaction_cost(&mut self, tx_cost: &TransactionCost<impl SVMMessage>) -> u64 {
325        saturating_add_assign!(
326            self.allocated_accounts_data_size,
327            tx_cost.allocated_accounts_data_size()
328        );
329        saturating_add_assign!(self.transaction_count, 1);
330        saturating_add_assign!(
331            self.transaction_signature_count,
332            tx_cost.num_transaction_signatures()
333        );
334        saturating_add_assign!(
335            self.secp256k1_instruction_signature_count,
336            tx_cost.num_secp256k1_instruction_signatures()
337        );
338        saturating_add_assign!(
339            self.ed25519_instruction_signature_count,
340            tx_cost.num_ed25519_instruction_signatures()
341        );
342        self.add_transaction_execution_cost(tx_cost, tx_cost.sum())
343    }
344
345    fn remove_transaction_cost(&mut self, tx_cost: &TransactionCost<impl SVMMessage>) {
346        let cost = tx_cost.sum();
347        self.sub_transaction_execution_cost(tx_cost, cost);
348        self.allocated_accounts_data_size = self
349            .allocated_accounts_data_size
350            .saturating_sub(tx_cost.allocated_accounts_data_size());
351        self.transaction_count = self.transaction_count.saturating_sub(1);
352        self.transaction_signature_count = self
353            .transaction_signature_count
354            .saturating_sub(tx_cost.num_transaction_signatures());
355        self.secp256k1_instruction_signature_count = self
356            .secp256k1_instruction_signature_count
357            .saturating_sub(tx_cost.num_secp256k1_instruction_signatures());
358        self.ed25519_instruction_signature_count = self
359            .ed25519_instruction_signature_count
360            .saturating_sub(tx_cost.num_ed25519_instruction_signatures());
361    }
362
363    /// Apply additional actual execution units to cost_tracker
364    /// Return the costliest account cost that were updated by `TransactionCost`
365    fn add_transaction_execution_cost(
366        &mut self,
367        tx_cost: &TransactionCost<impl SVMMessage>,
368        adjustment: u64,
369    ) -> u64 {
370        let mut costliest_account_cost = 0;
371        for account_key in tx_cost.writable_accounts() {
372            let account_cost = self
373                .cost_by_writable_accounts
374                .entry(*account_key)
375                .or_insert(0);
376            *account_cost = account_cost.saturating_add(adjustment);
377            costliest_account_cost = costliest_account_cost.max(*account_cost);
378        }
379        self.block_cost = self.block_cost.saturating_add(adjustment);
380        if tx_cost.is_simple_vote() {
381            self.vote_cost = self.vote_cost.saturating_add(adjustment);
382        }
383
384        costliest_account_cost
385    }
386
387    /// Subtract extra execution units from cost_tracker
388    fn sub_transaction_execution_cost(
389        &mut self,
390        tx_cost: &TransactionCost<impl SVMMessage>,
391        adjustment: u64,
392    ) {
393        for account_key in tx_cost.writable_accounts() {
394            let account_cost = self
395                .cost_by_writable_accounts
396                .entry(*account_key)
397                .or_insert(0);
398            *account_cost = account_cost.saturating_sub(adjustment);
399        }
400        self.block_cost = self.block_cost.saturating_sub(adjustment);
401        if tx_cost.is_simple_vote() {
402            self.vote_cost = self.vote_cost.saturating_sub(adjustment);
403        }
404    }
405
406    /// count number of none-zero CU accounts
407    fn number_of_accounts(&self) -> usize {
408        self.cost_by_writable_accounts
409            .values()
410            .filter(|units| **units > 0)
411            .count()
412    }
413}
414
415#[cfg(test)]
416mod tests {
417    use {
418        super::*,
419        crate::transaction_cost::{WritableKeysTransaction, *},
420        solana_sdk::{
421            message::TransactionSignatureDetails,
422            signature::{Keypair, Signer},
423        },
424        std::cmp,
425    };
426
427    impl CostTracker {
428        fn new(account_cost_limit: u64, block_cost_limit: u64, vote_cost_limit: u64) -> Self {
429            assert!(account_cost_limit <= block_cost_limit);
430            assert!(vote_cost_limit <= block_cost_limit);
431            Self {
432                account_cost_limit,
433                block_cost_limit,
434                vote_cost_limit,
435                ..Self::default()
436            }
437        }
438    }
439
440    fn test_setup() -> Keypair {
441        solana_logger::setup();
442        Keypair::new()
443    }
444
445    fn build_simple_transaction(mint_keypair: &Keypair) -> WritableKeysTransaction {
446        WritableKeysTransaction(vec![mint_keypair.pubkey()])
447    }
448
449    fn simple_usage_cost_details(
450        transaction: &WritableKeysTransaction,
451        programs_execution_cost: u64,
452    ) -> UsageCostDetails<WritableKeysTransaction> {
453        UsageCostDetails {
454            transaction,
455            signature_cost: 0,
456            write_lock_cost: 0,
457            data_bytes_cost: 0,
458            programs_execution_cost,
459            loaded_accounts_data_size_cost: 0,
460            allocated_accounts_data_size: 0,
461            signature_details: TransactionSignatureDetails::new(0, 0, 0),
462        }
463    }
464
465    fn simple_transaction_cost(
466        transaction: &WritableKeysTransaction,
467        programs_execution_cost: u64,
468    ) -> TransactionCost<WritableKeysTransaction> {
469        TransactionCost::Transaction(simple_usage_cost_details(
470            transaction,
471            programs_execution_cost,
472        ))
473    }
474
475    fn simple_vote_transaction_cost(
476        transaction: &WritableKeysTransaction,
477    ) -> TransactionCost<WritableKeysTransaction> {
478        TransactionCost::SimpleVote { transaction }
479    }
480
481    #[test]
482    fn test_cost_tracker_initialization() {
483        let testee = CostTracker::new(10, 11, 8);
484        assert_eq!(10, testee.account_cost_limit);
485        assert_eq!(11, testee.block_cost_limit);
486        assert_eq!(8, testee.vote_cost_limit);
487        assert_eq!(0, testee.cost_by_writable_accounts.len());
488        assert_eq!(0, testee.block_cost);
489    }
490
491    #[test]
492    fn test_cost_tracker_ok_add_one() {
493        let mint_keypair = test_setup();
494        let tx = build_simple_transaction(&mint_keypair);
495        let tx_cost = simple_transaction_cost(&tx, 5);
496        let cost = tx_cost.sum();
497
498        // build testee to have capacity for one simple transaction
499        let mut testee = CostTracker::new(cost, cost, cost);
500        assert!(testee.would_fit(&tx_cost).is_ok());
501        testee.add_transaction_cost(&tx_cost);
502        assert_eq!(cost, testee.block_cost);
503        assert_eq!(0, testee.vote_cost);
504        let (_costliest_account, costliest_account_cost) = testee.find_costliest_account();
505        assert_eq!(cost, costliest_account_cost);
506    }
507
508    #[test]
509    fn test_cost_tracker_ok_add_one_vote() {
510        let mint_keypair = test_setup();
511        let tx = build_simple_transaction(&mint_keypair);
512        let tx_cost = simple_vote_transaction_cost(&tx);
513        let cost = tx_cost.sum();
514
515        // build testee to have capacity for one simple transaction
516        let mut testee = CostTracker::new(cost, cost, cost);
517        assert!(testee.would_fit(&tx_cost).is_ok());
518        testee.add_transaction_cost(&tx_cost);
519        assert_eq!(cost, testee.block_cost);
520        assert_eq!(cost, testee.vote_cost);
521        let (_costliest_account, costliest_account_cost) = testee.find_costliest_account();
522        assert_eq!(cost, costliest_account_cost);
523    }
524
525    #[test]
526    fn test_cost_tracker_add_data() {
527        let mint_keypair = test_setup();
528        let tx = build_simple_transaction(&mint_keypair);
529        let mut tx_cost = simple_transaction_cost(&tx, 5);
530        if let TransactionCost::Transaction(ref mut usage_cost) = tx_cost {
531            usage_cost.allocated_accounts_data_size = 1;
532        } else {
533            unreachable!();
534        }
535        let cost = tx_cost.sum();
536
537        // build testee to have capacity for one simple transaction
538        let mut testee = CostTracker::new(cost, cost, cost);
539        assert!(testee.would_fit(&tx_cost).is_ok());
540        let old = testee.allocated_accounts_data_size;
541        testee.add_transaction_cost(&tx_cost);
542        assert_eq!(old + 1, testee.allocated_accounts_data_size);
543    }
544
545    #[test]
546    fn test_cost_tracker_ok_add_two_same_accounts() {
547        let mint_keypair = test_setup();
548        // build two transactions with same signed account
549        let tx1 = build_simple_transaction(&mint_keypair);
550        let tx_cost1 = simple_transaction_cost(&tx1, 5);
551        let cost1 = tx_cost1.sum();
552        let tx2 = build_simple_transaction(&mint_keypair);
553        let tx_cost2 = simple_transaction_cost(&tx2, 5);
554        let cost2 = tx_cost2.sum();
555
556        // build testee to have capacity for two simple transactions, with same accounts
557        let mut testee = CostTracker::new(cost1 + cost2, cost1 + cost2, cost1 + cost2);
558        {
559            assert!(testee.would_fit(&tx_cost1).is_ok());
560            testee.add_transaction_cost(&tx_cost1);
561        }
562        {
563            assert!(testee.would_fit(&tx_cost2).is_ok());
564            testee.add_transaction_cost(&tx_cost2);
565        }
566        assert_eq!(cost1 + cost2, testee.block_cost);
567        assert_eq!(1, testee.cost_by_writable_accounts.len());
568        let (_ccostliest_account, costliest_account_cost) = testee.find_costliest_account();
569        assert_eq!(cost1 + cost2, costliest_account_cost);
570    }
571
572    #[test]
573    fn test_cost_tracker_ok_add_two_diff_accounts() {
574        let mint_keypair = test_setup();
575        // build two transactions with diff accounts
576        let second_account = Keypair::new();
577        let tx1 = build_simple_transaction(&mint_keypair);
578        let tx_cost1 = simple_transaction_cost(&tx1, 5);
579        let cost1 = tx_cost1.sum();
580
581        let tx2 = build_simple_transaction(&second_account);
582        let tx_cost2 = simple_transaction_cost(&tx2, 5);
583        let cost2 = tx_cost2.sum();
584
585        // build testee to have capacity for two simple transactions, with same accounts
586        let mut testee = CostTracker::new(cmp::max(cost1, cost2), cost1 + cost2, cost1 + cost2);
587        {
588            assert!(testee.would_fit(&tx_cost1).is_ok());
589            testee.add_transaction_cost(&tx_cost1);
590        }
591        {
592            assert!(testee.would_fit(&tx_cost2).is_ok());
593            testee.add_transaction_cost(&tx_cost2);
594        }
595        assert_eq!(cost1 + cost2, testee.block_cost);
596        assert_eq!(2, testee.cost_by_writable_accounts.len());
597        let (_ccostliest_account, costliest_account_cost) = testee.find_costliest_account();
598        assert_eq!(std::cmp::max(cost1, cost2), costliest_account_cost);
599    }
600
601    #[test]
602    fn test_cost_tracker_chain_reach_limit() {
603        let mint_keypair = test_setup();
604        // build two transactions with same signed account
605        let tx1 = build_simple_transaction(&mint_keypair);
606        let tx_cost1 = simple_transaction_cost(&tx1, 5);
607        let cost1 = tx_cost1.sum();
608        let tx2 = build_simple_transaction(&mint_keypair);
609        let tx_cost2 = simple_transaction_cost(&tx2, 5);
610        let cost2 = tx_cost2.sum();
611
612        // build testee to have capacity for two simple transactions, but not for same accounts
613        let mut testee = CostTracker::new(cmp::min(cost1, cost2), cost1 + cost2, cost1 + cost2);
614        // should have room for first transaction
615        {
616            assert!(testee.would_fit(&tx_cost1).is_ok());
617            testee.add_transaction_cost(&tx_cost1);
618        }
619        // but no more sapce on the same chain (same signer account)
620        {
621            assert!(testee.would_fit(&tx_cost2).is_err());
622        }
623    }
624
625    #[test]
626    fn test_cost_tracker_reach_limit() {
627        let mint_keypair = test_setup();
628        // build two transactions with diff accounts
629        let second_account = Keypair::new();
630        let tx1 = build_simple_transaction(&mint_keypair);
631        let tx_cost1 = simple_transaction_cost(&tx1, 5);
632        let cost1 = tx_cost1.sum();
633        let tx2 = build_simple_transaction(&second_account);
634        let tx_cost2 = simple_transaction_cost(&tx2, 5);
635        let cost2 = tx_cost2.sum();
636
637        // build testee to have capacity for each chain, but not enough room for both transactions
638        let mut testee =
639            CostTracker::new(cmp::max(cost1, cost2), cost1 + cost2 - 1, cost1 + cost2 - 1);
640        // should have room for first transaction
641        {
642            assert!(testee.would_fit(&tx_cost1).is_ok());
643            testee.add_transaction_cost(&tx_cost1);
644        }
645        // but no more room for package as whole
646        {
647            assert!(testee.would_fit(&tx_cost2).is_err());
648        }
649    }
650
651    #[test]
652    fn test_cost_tracker_reach_vote_limit() {
653        let mint_keypair = test_setup();
654        // build two mocking vote transactions with diff accounts
655        let second_account = Keypair::new();
656        let tx1 = build_simple_transaction(&mint_keypair);
657        let tx_cost1 = simple_vote_transaction_cost(&tx1);
658        let cost1 = tx_cost1.sum();
659        let tx2 = build_simple_transaction(&second_account);
660        let tx_cost2 = simple_vote_transaction_cost(&tx2);
661        let cost2 = tx_cost2.sum();
662
663        // build testee to have capacity for each chain, but not enough room for both votes
664        let mut testee = CostTracker::new(cmp::max(cost1, cost2), cost1 + cost2, cost1 + cost2 - 1);
665        // should have room for first vote
666        {
667            assert!(testee.would_fit(&tx_cost1).is_ok());
668            testee.add_transaction_cost(&tx_cost1);
669        }
670        // but no more room for package as whole
671        {
672            assert!(testee.would_fit(&tx_cost2).is_err());
673        }
674        // however there is room for none-vote tx3
675        {
676            let third_account = Keypair::new();
677            let tx3 = build_simple_transaction(&third_account);
678            let tx_cost3 = simple_transaction_cost(&tx3, 5);
679            assert!(testee.would_fit(&tx_cost3).is_ok());
680        }
681    }
682
683    #[test]
684    fn test_cost_tracker_reach_data_block_limit() {
685        let mint_keypair = test_setup();
686        // build two transactions with diff accounts
687        let second_account = Keypair::new();
688        let tx1 = build_simple_transaction(&mint_keypair);
689        let mut tx_cost1 = simple_transaction_cost(&tx1, 5);
690        let tx2 = build_simple_transaction(&second_account);
691        let mut tx_cost2 = simple_transaction_cost(&tx2, 5);
692        if let TransactionCost::Transaction(ref mut usage_cost) = tx_cost1 {
693            usage_cost.allocated_accounts_data_size = MAX_BLOCK_ACCOUNTS_DATA_SIZE_DELTA;
694        } else {
695            unreachable!();
696        }
697        if let TransactionCost::Transaction(ref mut usage_cost) = tx_cost2 {
698            usage_cost.allocated_accounts_data_size = MAX_BLOCK_ACCOUNTS_DATA_SIZE_DELTA + 1;
699        } else {
700            unreachable!();
701        }
702        let cost1 = tx_cost1.sum();
703        let cost2 = tx_cost2.sum();
704
705        // build testee that passes
706        let testee = CostTracker::new(cmp::max(cost1, cost2), cost1 + cost2 - 1, cost1 + cost2 - 1);
707        assert!(testee.would_fit(&tx_cost1).is_ok());
708        // data is too big
709        assert_eq!(
710            testee.would_fit(&tx_cost2),
711            Err(CostTrackerError::WouldExceedAccountDataBlockLimit),
712        );
713    }
714
715    #[test]
716    fn test_cost_tracker_remove() {
717        let mint_keypair = test_setup();
718        // build two transactions with diff accounts
719        let second_account = Keypair::new();
720        let tx1 = build_simple_transaction(&mint_keypair);
721        let tx_cost1 = simple_transaction_cost(&tx1, 5);
722        let tx2 = build_simple_transaction(&second_account);
723        let tx_cost2 = simple_transaction_cost(&tx2, 5);
724        let cost1 = tx_cost1.sum();
725        let cost2 = tx_cost2.sum();
726
727        // build testee
728        let mut testee = CostTracker::new(cost1 + cost2, cost1 + cost2, cost1 + cost2);
729
730        assert!(testee.try_add(&tx_cost1).is_ok());
731        assert!(testee.try_add(&tx_cost2).is_ok());
732        assert_eq!(cost1 + cost2, testee.block_cost);
733
734        // removing a tx_cost affects block_cost
735        testee.remove(&tx_cost1);
736        assert_eq!(cost2, testee.block_cost);
737
738        // add back tx1
739        assert!(testee.try_add(&tx_cost1).is_ok());
740        assert_eq!(cost1 + cost2, testee.block_cost);
741
742        // cannot add tx1 again, cost limit would be exceeded
743        assert!(testee.try_add(&tx_cost1).is_err());
744    }
745
746    #[test]
747    fn test_cost_tracker_try_add_is_atomic() {
748        let acct1 = Pubkey::new_unique();
749        let acct2 = Pubkey::new_unique();
750        let acct3 = Pubkey::new_unique();
751        let cost = 100;
752        let account_max = cost * 2;
753        let block_max = account_max * 3; // for three accts
754
755        let mut testee = CostTracker::new(account_max, block_max, block_max);
756
757        // case 1: a tx writes to 3 accounts, should success, we will have:
758        // | acct1 | $cost |
759        // | acct2 | $cost |
760        // | acct3 | $cost |
761        // and block_cost = $cost
762        {
763            let transaction = WritableKeysTransaction(vec![acct1, acct2, acct3]);
764            let tx_cost = simple_transaction_cost(&transaction, cost);
765            assert!(testee.try_add(&tx_cost).is_ok());
766            let (_costliest_account, costliest_account_cost) = testee.find_costliest_account();
767            assert_eq!(cost, testee.block_cost);
768            assert_eq!(3, testee.cost_by_writable_accounts.len());
769            assert_eq!(cost, costliest_account_cost);
770        }
771
772        // case 2: add tx writes to acct2 with $cost, should succeed, result to
773        // | acct1 | $cost |
774        // | acct2 | $cost * 2 |
775        // | acct3 | $cost |
776        // and block_cost = $cost * 2
777        {
778            let transaction = WritableKeysTransaction(vec![acct2]);
779            let tx_cost = simple_transaction_cost(&transaction, cost);
780            assert!(testee.try_add(&tx_cost).is_ok());
781            let (costliest_account, costliest_account_cost) = testee.find_costliest_account();
782            assert_eq!(cost * 2, testee.block_cost);
783            assert_eq!(3, testee.cost_by_writable_accounts.len());
784            assert_eq!(cost * 2, costliest_account_cost);
785            assert_eq!(acct2, costliest_account);
786        }
787
788        // case 3: add tx writes to [acct1, acct2], acct2 exceeds limit, should failed atomically,
789        // we should still have:
790        // | acct1 | $cost |
791        // | acct2 | $cost * 2 |
792        // | acct3 | $cost |
793        // and block_cost = $cost * 2
794        {
795            let transaction = WritableKeysTransaction(vec![acct1, acct2]);
796            let tx_cost = simple_transaction_cost(&transaction, cost);
797            assert!(testee.try_add(&tx_cost).is_err());
798            let (costliest_account, costliest_account_cost) = testee.find_costliest_account();
799            assert_eq!(cost * 2, testee.block_cost);
800            assert_eq!(3, testee.cost_by_writable_accounts.len());
801            assert_eq!(cost * 2, costliest_account_cost);
802            assert_eq!(acct2, costliest_account);
803        }
804    }
805
806    #[test]
807    fn test_adjust_transaction_execution_cost() {
808        let acct1 = Pubkey::new_unique();
809        let acct2 = Pubkey::new_unique();
810        let acct3 = Pubkey::new_unique();
811        let cost = 100;
812        let account_max = cost * 2;
813        let block_max = account_max * 3; // for three accts
814
815        let mut testee = CostTracker::new(account_max, block_max, block_max);
816        let transaction = WritableKeysTransaction(vec![acct1, acct2, acct3]);
817        let tx_cost = simple_transaction_cost(&transaction, cost);
818        let mut expected_block_cost = tx_cost.sum();
819        let expected_tx_count = 1;
820        assert!(testee.try_add(&tx_cost).is_ok());
821        assert_eq!(expected_block_cost, testee.block_cost());
822        assert_eq!(expected_tx_count, testee.transaction_count());
823        testee
824            .cost_by_writable_accounts
825            .iter()
826            .for_each(|(_key, units)| {
827                assert_eq!(expected_block_cost, *units);
828            });
829
830        // adjust up
831        {
832            let adjustment = 50u64;
833            testee.add_transaction_execution_cost(&tx_cost, adjustment);
834            expected_block_cost += 50;
835            assert_eq!(expected_block_cost, testee.block_cost());
836            assert_eq!(expected_tx_count, testee.transaction_count());
837            testee
838                .cost_by_writable_accounts
839                .iter()
840                .for_each(|(_key, units)| {
841                    assert_eq!(expected_block_cost, *units);
842                });
843        }
844
845        // adjust down
846        {
847            let adjustment = 50u64;
848            testee.sub_transaction_execution_cost(&tx_cost, adjustment);
849            expected_block_cost -= 50;
850            assert_eq!(expected_block_cost, testee.block_cost());
851            assert_eq!(expected_tx_count, testee.transaction_count());
852            testee
853                .cost_by_writable_accounts
854                .iter()
855                .for_each(|(_key, units)| {
856                    assert_eq!(expected_block_cost, *units);
857                });
858        }
859
860        // adjust overflow
861        {
862            testee.add_transaction_execution_cost(&tx_cost, u64::MAX);
863            // expect block cost set to limit
864            assert_eq!(u64::MAX, testee.block_cost());
865            assert_eq!(expected_tx_count, testee.transaction_count());
866            testee
867                .cost_by_writable_accounts
868                .iter()
869                .for_each(|(_key, units)| {
870                    assert_eq!(u64::MAX, *units);
871                });
872        }
873
874        // adjust underflow
875        {
876            testee.sub_transaction_execution_cost(&tx_cost, u64::MAX);
877            // expect block cost set to limit
878            assert_eq!(u64::MIN, testee.block_cost());
879            assert_eq!(expected_tx_count, testee.transaction_count());
880            testee
881                .cost_by_writable_accounts
882                .iter()
883                .for_each(|(_key, units)| {
884                    assert_eq!(u64::MIN, *units);
885                });
886            // assert the number of non-empty accounts is zero, but map
887            // still contains 3 account
888            assert_eq!(0, testee.number_of_accounts());
889            assert_eq!(3, testee.cost_by_writable_accounts.len());
890        }
891    }
892
893    #[test]
894    fn test_update_execution_cost() {
895        let estimated_programs_execution_cost = 100;
896        let estimated_loaded_accounts_data_size_cost = 200;
897        let number_writeble_accounts = 3;
898        let transaction = WritableKeysTransaction(
899            std::iter::repeat_with(Pubkey::new_unique)
900                .take(number_writeble_accounts)
901                .collect(),
902        );
903
904        let mut usage_cost =
905            simple_usage_cost_details(&transaction, estimated_programs_execution_cost);
906        usage_cost.loaded_accounts_data_size_cost = estimated_loaded_accounts_data_size_cost;
907        let tx_cost = TransactionCost::Transaction(usage_cost);
908        // confirm tx_cost is only made up by programs_execution_cost and
909        // loaded_accounts_data_size_cost
910        let estimated_tx_cost = tx_cost.sum();
911        assert_eq!(
912            estimated_tx_cost,
913            estimated_programs_execution_cost + estimated_loaded_accounts_data_size_cost
914        );
915
916        let test_update_cost_tracker =
917            |execution_cost_adjust: i64, loaded_acounts_data_size_cost_adjust: i64| {
918                let mut cost_tracker = CostTracker::default();
919                assert!(cost_tracker.try_add(&tx_cost).is_ok());
920
921                let actual_programs_execution_cost =
922                    (estimated_programs_execution_cost as i64 + execution_cost_adjust) as u64;
923                let actual_loaded_accounts_data_size_cost =
924                    (estimated_loaded_accounts_data_size_cost as i64
925                        + loaded_acounts_data_size_cost_adjust) as u64;
926                let expected_cost = (estimated_tx_cost as i64
927                    + execution_cost_adjust
928                    + loaded_acounts_data_size_cost_adjust)
929                    as u64;
930
931                cost_tracker.update_execution_cost(
932                    &tx_cost,
933                    actual_programs_execution_cost,
934                    actual_loaded_accounts_data_size_cost,
935                );
936
937                assert_eq!(expected_cost, cost_tracker.block_cost);
938                assert_eq!(0, cost_tracker.vote_cost);
939                assert_eq!(
940                    number_writeble_accounts,
941                    cost_tracker.cost_by_writable_accounts.len()
942                );
943                for writable_account_cost in cost_tracker.cost_by_writable_accounts.values() {
944                    assert_eq!(expected_cost, *writable_account_cost);
945                }
946                assert_eq!(1, cost_tracker.transaction_count);
947            };
948
949        test_update_cost_tracker(0, 0);
950        test_update_cost_tracker(0, 9);
951        test_update_cost_tracker(0, -9);
952        test_update_cost_tracker(9, 0);
953        test_update_cost_tracker(9, 9);
954        test_update_cost_tracker(9, -9);
955        test_update_cost_tracker(-9, 0);
956        test_update_cost_tracker(-9, 9);
957        test_update_cost_tracker(-9, -9);
958    }
959
960    #[test]
961    fn test_remove_transaction_cost() {
962        let mut cost_tracker = CostTracker::default();
963
964        let cost = 100u64;
965        let transaction = WritableKeysTransaction(vec![Pubkey::new_unique()]);
966        let tx_cost = simple_transaction_cost(&transaction, cost);
967        cost_tracker.add_transaction_cost(&tx_cost);
968        // assert cost_tracker is reverted to default
969        assert_eq!(1, cost_tracker.transaction_count);
970        assert_eq!(1, cost_tracker.number_of_accounts());
971        assert_eq!(cost, cost_tracker.block_cost);
972        assert_eq!(0, cost_tracker.vote_cost);
973        assert_eq!(0, cost_tracker.allocated_accounts_data_size);
974
975        cost_tracker.remove_transaction_cost(&tx_cost);
976        // assert cost_tracker is reverted to default
977        assert_eq!(0, cost_tracker.transaction_count);
978        assert_eq!(0, cost_tracker.number_of_accounts());
979        assert_eq!(0, cost_tracker.block_cost);
980        assert_eq!(0, cost_tracker.vote_cost);
981        assert_eq!(0, cost_tracker.allocated_accounts_data_size);
982    }
983}