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 consumed(&self) -> u64 {
27        self.consumed.load(Ordering::Acquire)
28    }
29
30    pub fn reserve(&self, bytes: u64) -> std::result::Result<(), EntryBytesReserveError> {
31        let mut current = self.consumed.load(Ordering::Acquire);
32        loop {
33            let next = current.saturating_add(bytes);
34            if next > self.slot_limit {
35                return Err(EntryBytesReserveError::ExceedsSlotLimit);
36            }
37
38            match self.consumed.compare_exchange_weak(
39                current,
40                next,
41                Ordering::AcqRel,
42                Ordering::Acquire,
43            ) {
44                Ok(_) => return Ok(()),
45                Err(actual) => current = actual,
46            }
47        }
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54    const TEST_SLOT_LIMIT: u64 = 1_000;
55
56    #[test]
57    fn test_load_new() {
58        let budget = EntryBytesBudget::new(TEST_SLOT_LIMIT);
59        assert_eq!(budget.consumed(), 0);
60        assert_eq!(budget.slot_limit(), TEST_SLOT_LIMIT);
61    }
62
63    #[test]
64    fn test_reserve() {
65        let budget = EntryBytesBudget::new(TEST_SLOT_LIMIT);
66
67        assert!(budget.reserve(100).is_ok());
68        assert_eq!(budget.consumed(), 100);
69    }
70
71    #[test]
72    fn test_reserve_rejects_over_limit() {
73        let budget = EntryBytesBudget::new(TEST_SLOT_LIMIT);
74
75        assert!(budget.reserve(TEST_SLOT_LIMIT - 1).is_ok());
76        assert_eq!(
77            budget.reserve(2),
78            Err(EntryBytesReserveError::ExceedsSlotLimit)
79        );
80        assert_eq!(budget.consumed(), TEST_SLOT_LIMIT - 1);
81    }
82}
solana_runtime/bank/entry_bytes_budget.rs

solana_runtime/bank/
entry_bytes_budget.rs
