solana_runtime/bank/
entry_bytes_budget.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum EntryBytesReserveError {
5    ExceedsSlotLimit,
6}
7
8#[derive(Debug)]
9pub struct EntryBytesBudget {
10    consumed: AtomicU64,
11    slot_limit: u64,
12}
13
14impl EntryBytesBudget {
15    pub const fn new(slot_limit: u64) -> Self {
16        Self {
17            consumed: AtomicU64::new(0),
18            slot_limit,
19        }
20    }
21
22    pub const fn slot_limit(&self) -> u64 {
23        self.slot_limit
24    }
25
26    pub fn reserve(&self, bytes: u64) -> std::result::Result<(), EntryBytesReserveError> {
27        loop {
28            let current = self.consumed.load(Ordering::Acquire);
29            let next = current.saturating_add(bytes);
30            if next > self.slot_limit {
31                return Err(EntryBytesReserveError::ExceedsSlotLimit);
32            }
33
34            if self
35                .consumed
36                .compare_exchange(current, next, Ordering::AcqRel, Ordering::Acquire)
37                .is_ok()
38            {
39                return Ok(());
40            }
41        }
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48    const TEST_SLOT_LIMIT: u64 = 1_000;
49
50    #[test]
51    fn test_load_new() {
52        let budget = EntryBytesBudget::new(TEST_SLOT_LIMIT);
53        assert_eq!(budget.consumed.load(Ordering::Acquire), 0);
54        assert_eq!(budget.slot_limit(), TEST_SLOT_LIMIT);
55    }
56
57    #[test]
58    fn test_reserve() {
59        let budget = EntryBytesBudget::new(TEST_SLOT_LIMIT);
60
61        assert!(budget.reserve(100).is_ok());
62        assert_eq!(budget.consumed.load(Ordering::Acquire), 100);
63    }
64
65    #[test]
66    fn test_reserve_rejects_over_limit() {
67        let budget = EntryBytesBudget::new(TEST_SLOT_LIMIT);
68
69        assert!(budget.reserve(TEST_SLOT_LIMIT - 1).is_ok());
70        assert_eq!(
71            budget.reserve(2),
72            Err(EntryBytesReserveError::ExceedsSlotLimit)
73        );
74        assert_eq!(budget.consumed.load(Ordering::Acquire), TEST_SLOT_LIMIT - 1);
75    }
76}
solana_runtime/bank/entry_bytes_budget.rs

solana_runtime/bank/
entry_bytes_budget.rs
