open_geocode/
memory.rs

1//! Build memory accounting.
2//!
3//! One [`MemoryBudget`] is shared by every large buffer of a build: sort
4//! buffers (counted by capacity, so `Vec` growth is included), merge read
5//! buffers and the text indexing buffers. A buffer reserves bytes before it
6//! grows; when the pool is full its owner spills to disk instead. The budget
7//! is the `--memory-budget-mb` limit, and the build report records the peak.
8
9use std::sync::{
10    Arc,
11    atomic::{AtomicUsize, Ordering},
12};
13
14#[derive(Debug)]
15pub struct MemoryBudget {
16    limit: usize,
17    used: AtomicUsize,
18    peak: AtomicUsize,
19}
20
21impl MemoryBudget {
22    pub fn new(limit: usize) -> Arc<Self> {
23        Arc::new(Self {
24            limit,
25            used: AtomicUsize::new(0),
26            peak: AtomicUsize::new(0),
27        })
28    }
29
30    /// A budget that never refuses, for writers used outside a build.
31    pub fn unlimited() -> Arc<Self> {
32        Self::new(usize::MAX)
33    }
34
35    pub fn limit(&self) -> usize {
36        self.limit
37    }
38
39    pub fn used(&self) -> usize {
40        self.used.load(Ordering::Relaxed)
41    }
42
43    /// Most bytes reserved at once.
44    pub fn peak(&self) -> usize {
45        self.peak.load(Ordering::Relaxed)
46    }
47
48    /// Reserve `bytes` if they fit.
49    pub fn try_reserve(&self, bytes: usize) -> bool {
50        let mut used = self.used.load(Ordering::Relaxed);
51        loop {
52            let Some(next) = used.checked_add(bytes).filter(|next| *next <= self.limit) else {
53                return false;
54            };
55            match self
56                .used
57                .compare_exchange_weak(used, next, Ordering::Relaxed, Ordering::Relaxed)
58            {
59                Ok(_) => {
60                    self.peak.fetch_max(next, Ordering::Relaxed);
61                    return true;
62                }
63                Err(current) => used = current,
64            }
65        }
66    }
67
68    /// Reserve `bytes` even past the limit: the minimum a caller needs to make
69    /// progress, such as a single item after spilling everything else.
70    pub fn force_reserve(&self, bytes: usize) {
71        let next = self
72            .used
73            .fetch_add(bytes, Ordering::Relaxed)
74            .saturating_add(bytes);
75        self.peak.fetch_max(next, Ordering::Relaxed);
76    }
77
78    pub fn release(&self, bytes: usize) {
79        self.used.fetch_sub(bytes, Ordering::Relaxed);
80    }
81}
82
83/// Bytes held from a budget, returned when dropped.
84#[derive(Debug)]
85pub struct Reservation {
86    budget: Arc<MemoryBudget>,
87    bytes: usize,
88}
89
90impl Reservation {
91    pub fn empty(budget: &Arc<MemoryBudget>) -> Self {
92        Self {
93            budget: Arc::clone(budget),
94            bytes: 0,
95        }
96    }
97
98    /// A reservation that is granted regardless of the limit.
99    pub fn forced(budget: &Arc<MemoryBudget>, bytes: usize) -> Self {
100        budget.force_reserve(bytes);
101        Self {
102            budget: Arc::clone(budget),
103            bytes,
104        }
105    }
106
107    pub fn bytes(&self) -> usize {
108        self.bytes
109    }
110
111    pub fn try_grow(&mut self, bytes: usize) -> bool {
112        let granted = self.budget.try_reserve(bytes);
113        if granted {
114            self.bytes += bytes;
115        }
116        granted
117    }
118
119    pub fn force_grow(&mut self, bytes: usize) {
120        self.budget.force_reserve(bytes);
121        self.bytes += bytes;
122    }
123
124    pub fn release_all(&mut self) {
125        self.budget.release(self.bytes);
126        self.bytes = 0;
127    }
128}
129
130impl Drop for Reservation {
131    fn drop(&mut self) {
132        self.release_all();
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn reservations_respect_the_limit_and_track_the_peak() {
142        let budget = MemoryBudget::new(100);
143        let mut first = Reservation::empty(&budget);
144        assert!(first.try_grow(60));
145        let mut second = Reservation::empty(&budget);
146        assert!(!second.try_grow(50));
147        assert!(second.try_grow(40));
148        assert_eq!(budget.used(), 100);
149        drop(first);
150        assert_eq!(budget.used(), 40);
151        second.force_grow(80);
152        assert_eq!(budget.used(), 120);
153        assert_eq!(budget.peak(), 120);
154        drop(second);
155        assert_eq!(budget.used(), 0);
156    }
157}
open_geocode/memory.rs

open_geocode/
memory.rs
