virtual_fs/
limiter.rs

1use std::sync::Arc;
2
3use crate::FsError;
4
5pub use self::tracked_vec::TrackedVec;
6
7/// Allows tracking and limiting the memory usage of a memfs [`FileSystem`](crate::FileSystem).
8pub trait FsMemoryLimiter: Send + Sync + std::fmt::Debug {
9    fn on_grow(&self, grown_bytes: usize) -> std::result::Result<(), FsError>;
10    fn on_shrink(&self, shrunk_bytes: usize);
11}
12
13pub type DynFsMemoryLimiter = Arc<dyn FsMemoryLimiter + Send + Sync>;
14
15#[cfg(feature = "tracking")]
16mod tracked_vec {
17    use crate::FsError;
18
19    use super::DynFsMemoryLimiter;
20
21    #[derive(Debug, Clone)]
22    pub struct TrackedVec {
23        data: Vec<u8>,
24        pub(super) limiter: Option<DynFsMemoryLimiter>,
25    }
26
27    impl TrackedVec {
28        pub fn new(limiter: Option<DynFsMemoryLimiter>) -> Self {
29            Self {
30                data: Vec::new(),
31                limiter,
32            }
33        }
34
35        pub fn limiter(&self) -> Option<&DynFsMemoryLimiter> {
36            self.limiter.as_ref()
37        }
38
39        pub fn with_capacity(
40            capacity: usize,
41            limiter: Option<DynFsMemoryLimiter>,
42        ) -> Result<Self, FsError> {
43            if let Some(limiter) = &limiter {
44                limiter.on_grow(capacity)?;
45            }
46            Ok(Self {
47                data: Vec::with_capacity(capacity),
48                limiter,
49            })
50        }
51
52        pub fn clear(&mut self) {
53            self.data.clear();
54        }
55
56        pub fn append(&mut self, other: &mut Self) -> Result<(), FsError> {
57            let old_capacity = self.data.capacity();
58            self.data.append(&mut other.data);
59
60            if let Some(limiter) = &self.limiter {
61                let new = self.data.capacity() - old_capacity;
62                limiter.on_grow(new)?;
63            }
64
65            Ok(())
66        }
67
68        pub fn split_off(&mut self, at: usize) -> Result<Self, FsError> {
69            let other = self.data.split_off(at);
70
71            if let Some(limiter) = &self.limiter {
72                // NOTE: split_off leaves the original vector capacity intact, so
73                // we can just add the new length.
74                let new_len = other.capacity();
75                limiter.on_grow(new_len)?;
76            }
77
78            Ok(Self {
79                data: other,
80                limiter: self.limiter.clone(),
81            })
82        }
83
84        pub fn resize(&mut self, new_len: usize, value: u8) -> Result<(), FsError> {
85            let old_capacity = self.data.capacity();
86            self.data.resize(new_len, value);
87            if let Some(limiter) = &self.limiter {
88                let new = self.data.capacity() - old_capacity;
89                limiter.on_grow(new)?;
90            }
91            Ok(())
92        }
93
94        pub fn extend_from_slice(&mut self, other: &[u8]) -> Result<(), FsError> {
95            let old_capacity = self.data.capacity();
96            self.data.extend_from_slice(other);
97            if let Some(limiter) = &self.limiter {
98                let new = self.data.capacity() - old_capacity;
99                limiter.on_grow(new)?;
100            }
101            Ok(())
102        }
103
104        pub fn reserve_exact(&mut self, additional: usize) -> Result<(), FsError> {
105            let old_capacity = self.data.capacity();
106            self.data.reserve_exact(additional);
107            if let Some(limiter) = &self.limiter {
108                let new = self.data.capacity() - old_capacity;
109                limiter.on_grow(new)?;
110            }
111            Ok(())
112        }
113    }
114
115    impl Drop for TrackedVec {
116        fn drop(&mut self) {
117            if let Some(limiter) = &self.limiter {
118                limiter.on_shrink(self.data.capacity());
119            }
120        }
121    }
122
123    impl std::ops::Deref for TrackedVec {
124        type Target = [u8];
125
126        fn deref(&self) -> &Self::Target {
127            &self.data
128        }
129    }
130
131    impl std::ops::DerefMut for TrackedVec {
132        fn deref_mut(&mut self) -> &mut Self::Target {
133            &mut self.data
134        }
135    }
136}
137
138#[cfg(not(feature = "tracking"))]
139mod tracked_vec {
140    use crate::FsError;
141
142    use super::DynFsMemoryLimiter;
143
144    #[derive(Debug)]
145    pub struct TrackedVec {
146        data: Vec<u8>,
147    }
148
149    impl TrackedVec {
150        pub fn new(_limiter: Option<DynFsMemoryLimiter>) -> Self {
151            Self { data: Vec::new() }
152        }
153
154        pub fn limiter(&self) -> Option<&DynFsMemoryLimiter> {
155            None
156        }
157
158        pub fn with_capacity(
159            capacity: usize,
160            _limiter: Option<DynFsMemoryLimiter>,
161        ) -> Result<Self, FsError> {
162            Ok(Self {
163                data: Vec::with_capacity(capacity),
164            })
165        }
166
167        pub fn clear(&mut self) {
168            self.data.clear();
169        }
170
171        pub fn append(&mut self, other: &mut Self) -> Result<(), FsError> {
172            self.data.append(&mut other.data);
173            Ok(())
174        }
175
176        pub fn split_off(&mut self, at: usize) -> Result<Self, FsError> {
177            let other = self.data.split_off(at);
178            Ok(Self { data: other })
179        }
180
181        pub fn resize(&mut self, new_len: usize, value: u8) -> Result<(), FsError> {
182            self.data.resize(new_len, value);
183            Ok(())
184        }
185
186        pub fn extend_from_slice(&mut self, other: &[u8]) -> Result<(), FsError> {
187            self.data.extend_from_slice(other);
188            Ok(())
189        }
190
191        pub fn reserve_exact(&mut self, additional: usize) -> Result<(), FsError> {
192            self.data.reserve_exact(additional);
193            Ok(())
194        }
195    }
196
197    impl std::ops::Deref for TrackedVec {
198        type Target = Vec<u8>;
199
200        fn deref(&self) -> &Self::Target {
201            &self.data
202        }
203    }
204
205    impl std::ops::DerefMut for TrackedVec {
206        fn deref_mut(&mut self) -> &mut Self::Target {
207            &mut self.data
208        }
209    }
210}
virtual_fs/limiter.rs

virtual_fs/
limiter.rs
