/home/noah/src/realizar/src/memory.rs
Line | Count | Source |
1 | | //! Memory management for hot expert pinning |
2 | | //! |
3 | | //! Per Dean & Barroso (2013) "The Tail at Scale", page faults cause |
4 | | //! tail latency. This module provides mlock support for hot experts. |
5 | | |
6 | | /// Configuration for memory pinning |
7 | | #[derive(Debug, Clone, Default)] |
8 | | pub struct MlockConfig { |
9 | | /// Whether to attempt memory locking |
10 | | pub enabled: bool, |
11 | | /// Maximum bytes to lock (0 = unlimited) |
12 | | pub max_locked_bytes: usize, |
13 | | } |
14 | | |
15 | | /// Result of mlock operation |
16 | | #[derive(Debug, Clone, PartialEq, Eq)] |
17 | | pub enum MlockResult { |
18 | | /// Successfully locked memory |
19 | | Locked, |
20 | | /// Mlock disabled in config |
21 | | Disabled, |
22 | | /// Failed to lock (insufficient privileges) |
23 | | InsufficientPrivileges, |
24 | | /// Failed to lock (resource limit) |
25 | | ResourceLimit, |
26 | | /// Platform not supported |
27 | | Unsupported, |
28 | | } |
29 | | |
30 | | /// Memory region that can be pinned |
31 | | pub struct PinnedRegion { |
32 | | /// Pointer to locked memory (for tracking) |
33 | | ptr: *const u8, |
34 | | /// Size of locked region |
35 | | len: usize, |
36 | | /// Whether actually locked |
37 | | locked: bool, |
38 | | } |
39 | | |
40 | | // Safety: PinnedRegion only holds a pointer for tracking, doesn't access it |
41 | | unsafe impl Send for PinnedRegion {} |
42 | | unsafe impl Sync for PinnedRegion {} |
43 | | |
44 | | impl PinnedRegion { |
45 | | /// Create a pinned region (attempts mlock) |
46 | | /// |
47 | | /// # Safety |
48 | | /// |
49 | | /// Caller must ensure ptr points to valid memory for len bytes. |
50 | | #[must_use] |
51 | 5 | pub unsafe fn new(ptr: *const u8, len: usize, config: &MlockConfig) -> (Self, MlockResult) { |
52 | 5 | if !config.enabled { |
53 | 3 | return ( |
54 | 3 | Self { |
55 | 3 | ptr, |
56 | 3 | len, |
57 | 3 | locked: false, |
58 | 3 | }, |
59 | 3 | MlockResult::Disabled, |
60 | 3 | ); |
61 | 2 | } |
62 | | |
63 | 2 | if config.max_locked_bytes > 0 && len > config.max_locked_bytes1 { |
64 | 1 | return ( |
65 | 1 | Self { |
66 | 1 | ptr, |
67 | 1 | len, |
68 | 1 | locked: false, |
69 | 1 | }, |
70 | 1 | MlockResult::ResourceLimit, |
71 | 1 | ); |
72 | 1 | } |
73 | | |
74 | 1 | let result = Self::mlock_impl(ptr, len); |
75 | 1 | let locked = result == MlockResult::Locked; |
76 | 1 | (Self { ptr, len, locked }, result) |
77 | 5 | } |
78 | | |
79 | | /// Check if memory is locked |
80 | | #[must_use] |
81 | 2 | pub fn is_locked(&self) -> bool { |
82 | 2 | self.locked |
83 | 2 | } |
84 | | |
85 | | /// Get size of region |
86 | | #[must_use] |
87 | 3 | pub fn len(&self) -> usize { |
88 | 3 | self.len |
89 | 3 | } |
90 | | |
91 | | /// Check if region is empty |
92 | | #[must_use] |
93 | 2 | pub fn is_empty(&self) -> bool { |
94 | 2 | self.len == 0 |
95 | 2 | } |
96 | | |
97 | | #[cfg(target_family = "unix")] |
98 | 1 | fn mlock_impl(ptr: *const u8, len: usize) -> MlockResult { |
99 | | // Safety: ptr and len validated by caller |
100 | | // SAFETY: Memory safety ensured by bounds checking and alignment |
101 | 1 | let result = unsafe { libc::mlock(ptr.cast(), len) }; |
102 | 1 | if result == 0 { |
103 | 1 | MlockResult::Locked |
104 | | } else { |
105 | | // Check errno for specific error |
106 | 0 | let errno = std::io::Error::last_os_error().raw_os_error().unwrap_or(0); |
107 | 0 | if errno == libc::EPERM { |
108 | 0 | MlockResult::InsufficientPrivileges |
109 | | } else { |
110 | 0 | MlockResult::ResourceLimit |
111 | | } |
112 | | } |
113 | 1 | } |
114 | | |
115 | | #[cfg(not(target_family = "unix"))] |
116 | | fn mlock_impl(_ptr: *const u8, _len: usize) -> MlockResult { |
117 | | MlockResult::Unsupported |
118 | | } |
119 | | } |
120 | | |
121 | | impl Drop for PinnedRegion { |
122 | 5 | fn drop(&mut self) { |
123 | 5 | if self.locked { |
124 | | #[cfg(target_family = "unix")] |
125 | | // SAFETY: Memory safety ensured by bounds checking and alignment |
126 | 1 | unsafe { |
127 | 1 | libc::munlock(self.ptr.cast(), self.len); |
128 | 1 | } |
129 | 4 | } |
130 | 5 | } |
131 | | } |
132 | | |
133 | | /// Expert tier classification for memory management |
134 | | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
135 | | pub enum ExpertTier { |
136 | | /// Hot tier - pinned in memory, never swapped |
137 | | Hot, |
138 | | /// Warm tier - in memory but may be swapped |
139 | | Warm, |
140 | | /// Cold tier - on disk, loaded on demand |
141 | | Cold, |
142 | | } |
143 | | |
144 | | impl ExpertTier { |
145 | | /// Determine tier based on access frequency |
146 | | #[must_use] |
147 | 7 | pub fn from_access_count(count: usize, hot_threshold: usize, warm_threshold: usize) -> Self { |
148 | 7 | if count >= hot_threshold { |
149 | 2 | Self::Hot |
150 | 5 | } else if count >= warm_threshold { |
151 | 3 | Self::Warm |
152 | | } else { |
153 | 2 | Self::Cold |
154 | | } |
155 | 7 | } |
156 | | } |
157 | | |
158 | | #[cfg(test)] |
159 | | mod tests { |
160 | | use super::*; |
161 | | |
162 | | #[test] |
163 | 1 | fn test_mlock_disabled() { |
164 | 1 | let config = MlockConfig { |
165 | 1 | enabled: false, |
166 | 1 | max_locked_bytes: 0, |
167 | 1 | }; |
168 | 1 | let data = vec![0u8; 1024]; |
169 | | // SAFETY: Memory safety ensured by bounds checking and alignment |
170 | 1 | let (region, result) = unsafe { PinnedRegion::new(data.as_ptr(), data.len(), &config) }; |
171 | 1 | assert_eq!(result, MlockResult::Disabled); |
172 | 1 | assert!(!region.is_locked()); |
173 | 1 | } |
174 | | |
175 | | #[test] |
176 | 1 | fn test_mlock_resource_limit() { |
177 | 1 | let config = MlockConfig { |
178 | 1 | enabled: true, |
179 | 1 | max_locked_bytes: 100, |
180 | 1 | }; |
181 | 1 | let data = vec![0u8; 1024]; // Exceeds limit |
182 | | // SAFETY: Memory safety ensured by bounds checking and alignment |
183 | 1 | let (region, result) = unsafe { PinnedRegion::new(data.as_ptr(), data.len(), &config) }; |
184 | 1 | assert_eq!(result, MlockResult::ResourceLimit); |
185 | 1 | assert!(!region.is_locked()); |
186 | 1 | } |
187 | | |
188 | | #[test] |
189 | 1 | fn test_expert_tier_hot() { |
190 | 1 | let tier = ExpertTier::from_access_count(1000, 100, 10); |
191 | 1 | assert_eq!(tier, ExpertTier::Hot); |
192 | 1 | } |
193 | | |
194 | | #[test] |
195 | 1 | fn test_expert_tier_warm() { |
196 | 1 | let tier = ExpertTier::from_access_count(50, 100, 10); |
197 | 1 | assert_eq!(tier, ExpertTier::Warm); |
198 | 1 | } |
199 | | |
200 | | #[test] |
201 | 1 | fn test_expert_tier_cold() { |
202 | 1 | let tier = ExpertTier::from_access_count(5, 100, 10); |
203 | 1 | assert_eq!(tier, ExpertTier::Cold); |
204 | 1 | } |
205 | | |
206 | | #[test] |
207 | 1 | fn test_pinned_region_len() { |
208 | 1 | let config = MlockConfig::default(); |
209 | 1 | let data = vec![0u8; 1024]; |
210 | | // SAFETY: Memory safety ensured by bounds checking and alignment |
211 | 1 | let (region, _) = unsafe { PinnedRegion::new(data.as_ptr(), data.len(), &config) }; |
212 | 1 | assert_eq!(region.len(), 1024); |
213 | 1 | assert!(!region.is_empty()); |
214 | 1 | } |
215 | | |
216 | | #[test] |
217 | 1 | fn test_pinned_region_empty() { |
218 | 1 | let config = MlockConfig::default(); |
219 | 1 | let data: Vec<u8> = vec![]; |
220 | | // SAFETY: Memory safety ensured by bounds checking and alignment |
221 | 1 | let (region, _) = unsafe { PinnedRegion::new(data.as_ptr(), data.len(), &config) }; |
222 | 1 | assert_eq!(region.len(), 0); |
223 | 1 | assert!(region.is_empty()); |
224 | 1 | } |
225 | | |
226 | | #[test] |
227 | 1 | fn test_mlock_config_default() { |
228 | 1 | let config = MlockConfig::default(); |
229 | 1 | assert!(!config.enabled); |
230 | 1 | assert_eq!(config.max_locked_bytes, 0); |
231 | 1 | } |
232 | | |
233 | | #[test] |
234 | 1 | fn test_mlock_result_equality() { |
235 | 1 | assert_eq!(MlockResult::Locked, MlockResult::Locked); |
236 | 1 | assert_eq!(MlockResult::Disabled, MlockResult::Disabled); |
237 | 1 | assert_ne!(MlockResult::Locked, MlockResult::Disabled); |
238 | 1 | } |
239 | | |
240 | | #[test] |
241 | 1 | fn test_expert_tier_boundary() { |
242 | | // Test exact boundary conditions |
243 | 1 | assert_eq!(ExpertTier::from_access_count(100, 100, 10), ExpertTier::Hot); |
244 | 1 | assert_eq!(ExpertTier::from_access_count(99, 100, 10), ExpertTier::Warm); |
245 | 1 | assert_eq!(ExpertTier::from_access_count(10, 100, 10), ExpertTier::Warm); |
246 | 1 | assert_eq!(ExpertTier::from_access_count(9, 100, 10), ExpertTier::Cold); |
247 | 1 | } |
248 | | |
249 | | #[test] |
250 | 1 | fn test_mlock_enabled_within_limit() { |
251 | | // Test when enabled and within limit - will try actual mlock |
252 | 1 | let config = MlockConfig { |
253 | 1 | enabled: true, |
254 | 1 | max_locked_bytes: 0, // 0 = unlimited |
255 | 1 | }; |
256 | 1 | let data: [u8; 64] = [0u8; 64]; |
257 | | // SAFETY: Memory safety ensured by bounds checking and alignment |
258 | 1 | let (region, result) = unsafe { PinnedRegion::new(data.as_ptr(), data.len(), &config) }; |
259 | | // Result depends on system permissions, but API should work |
260 | 1 | assert!( |
261 | 1 | result == MlockResult::Locked |
262 | 0 | || result == MlockResult::InsufficientPrivileges |
263 | 0 | || result == MlockResult::ResourceLimit |
264 | 0 | || result == MlockResult::Unsupported |
265 | | ); |
266 | 1 | assert_eq!(region.len(), 64); |
267 | 1 | } |
268 | | |
269 | | // ======================================================================== |
270 | | // Coverage Tests |
271 | | // ======================================================================== |
272 | | |
273 | | #[test] |
274 | 1 | fn test_mlock_config_debug() { |
275 | 1 | let config = MlockConfig { |
276 | 1 | enabled: true, |
277 | 1 | max_locked_bytes: 1024, |
278 | 1 | }; |
279 | 1 | let debug = format!("{:?}", config); |
280 | 1 | assert!(debug.contains("MlockConfig")); |
281 | 1 | assert!(debug.contains("1024")); |
282 | 1 | } |
283 | | |
284 | | #[test] |
285 | 1 | fn test_mlock_config_clone() { |
286 | 1 | let config = MlockConfig { |
287 | 1 | enabled: true, |
288 | 1 | max_locked_bytes: 2048, |
289 | 1 | }; |
290 | 1 | let cloned = config.clone(); |
291 | 1 | assert_eq!(config.enabled, cloned.enabled); |
292 | 1 | assert_eq!(config.max_locked_bytes, cloned.max_locked_bytes); |
293 | 1 | } |
294 | | |
295 | | #[test] |
296 | 1 | fn test_mlock_result_debug() { |
297 | 1 | let result = MlockResult::Locked; |
298 | 1 | let debug = format!("{:?}", result); |
299 | 1 | assert!(debug.contains("Locked")); |
300 | | |
301 | 1 | let result2 = MlockResult::InsufficientPrivileges; |
302 | 1 | let debug2 = format!("{:?}", result2); |
303 | 1 | assert!(debug2.contains("InsufficientPrivileges")); |
304 | 1 | } |
305 | | |
306 | | #[test] |
307 | 1 | fn test_mlock_result_clone() { |
308 | 1 | let result = MlockResult::Locked; |
309 | 1 | let cloned = result.clone(); |
310 | 1 | assert_eq!(result, cloned); |
311 | 1 | } |
312 | | |
313 | | #[test] |
314 | 1 | fn test_expert_tier_debug() { |
315 | 1 | let tier = ExpertTier::Hot; |
316 | 1 | let debug = format!("{:?}", tier); |
317 | 1 | assert!(debug.contains("Hot")); |
318 | 1 | } |
319 | | |
320 | | #[test] |
321 | 1 | fn test_expert_tier_clone() { |
322 | 1 | let tier = ExpertTier::Warm; |
323 | 1 | let cloned = tier; |
324 | 1 | assert_eq!(tier, cloned); |
325 | 1 | } |
326 | | } |