Skip to main content

sparse_mmap/
lib.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Memory-related abstractions.
5
6// UNSAFETY: Manual pointer manipulation, dealing with mmap, and a signal handler.
7#![expect(unsafe_code)]
8#![expect(missing_docs)]
9#![expect(clippy::undocumented_unsafe_blocks, clippy::missing_safety_doc)]
10
11pub mod alloc;
12pub mod unix;
13pub mod windows;
14
15pub use sys::AsMappableRef;
16pub use sys::Mappable;
17pub use sys::MappableRef;
18pub use sys::SparseMapping;
19pub use sys::alloc_shared_memory;
20pub use sys::alloc_shared_memory_hugetlb;
21pub use sys::new_mappable_from_file;
22
23use std::mem::MaybeUninit;
24use std::sync::atomic::AtomicU8;
25use thiserror::Error;
26#[cfg(unix)]
27use unix as sys;
28#[cfg(windows)]
29use windows as sys;
30use zerocopy::FromBytes;
31use zerocopy::Immutable;
32use zerocopy::IntoBytes;
33use zerocopy::KnownLayout;
34
35#[derive(Debug, Error)]
36pub enum SparseMappingError {
37    #[error("out of bounds")]
38    OutOfBounds,
39    #[error(transparent)]
40    Memory(trycopy::MemoryError),
41}
42
43/// Computes the reservation alignment for a mapping of `len` bytes.
44///
45/// Larger mappings are aligned to large-page boundaries (2 MB, then 1 GB) so
46/// that they can back large pages without the caller having to ask. The result
47/// is always at least `minimum_alignment` and at least the system page size.
48///
49/// Returns an error if `minimum_alignment` is not a power of two.
50fn reservation_alignment(len: usize, minimum_alignment: usize) -> std::io::Result<usize> {
51    if !minimum_alignment.is_power_of_two() {
52        return Err(std::io::Error::new(
53            std::io::ErrorKind::InvalidInput,
54            "alignment must be a power of two",
55        ));
56    }
57    const SIZE_2M: usize = 0x200000;
58    const SIZE_1G: usize = 0x40000000;
59    let default_alignment = if len < SIZE_2M {
60        SparseMapping::page_size()
61    } else if len < SIZE_1G {
62        SIZE_2M
63    } else {
64        SIZE_1G
65    };
66    Ok(default_alignment.max(minimum_alignment))
67}
68
69impl SparseMapping {
70    /// Gets the supported page size for sparse mappings.
71    pub fn page_size() -> usize {
72        sys::page_size()
73    }
74
75    fn check(&self, offset: usize, len: usize) -> Result<(), SparseMappingError> {
76        if self.len() < offset || self.len() - offset < len {
77            return Err(SparseMappingError::OutOfBounds);
78        }
79        Ok(())
80    }
81
82    /// Reads a type `T` from `offset` in the sparse mapping using a single read instruction.
83    ///
84    /// Panics if `T` is not 1, 2, 4, or 8 bytes in size.
85    pub fn read_volatile<T: FromBytes + Immutable + KnownLayout>(
86        &self,
87        offset: usize,
88    ) -> Result<T, SparseMappingError> {
89        assert!(self.is_local(), "cannot read from remote mappings");
90
91        self.check(offset, size_of::<T>())?;
92        // SAFETY: the bounds have been checked above.
93        unsafe { trycopy::try_read_volatile(self.as_ptr().byte_add(offset).cast()) }
94            .map_err(SparseMappingError::Memory)
95    }
96
97    /// Writes a type `T` at `offset` in the sparse mapping using a single write instruciton.
98    ///
99    /// Panics if `T` is not 1, 2, 4, or 8 bytes in size.
100    pub fn write_volatile<T: IntoBytes + Immutable + KnownLayout>(
101        &self,
102        offset: usize,
103        value: &T,
104    ) -> Result<(), SparseMappingError> {
105        assert!(self.is_local(), "cannot write to remote mappings");
106
107        self.check(offset, size_of::<T>())?;
108        // SAFETY: the bounds have been checked above.
109        unsafe { trycopy::try_write_volatile(self.as_ptr().byte_add(offset).cast(), value) }
110            .map_err(SparseMappingError::Memory)
111    }
112
113    /// Tries to write into the sparse mapping.
114    pub fn write_at(&self, offset: usize, data: &[u8]) -> Result<(), SparseMappingError> {
115        assert!(self.is_local(), "cannot write to remote mappings");
116
117        self.check(offset, data.len())?;
118        // SAFETY: the bounds have been checked above.
119        unsafe {
120            let dest = self.as_ptr().cast::<u8>().add(offset);
121            trycopy::try_copy(data.as_ptr(), dest, data.len()).map_err(SparseMappingError::Memory)
122        }
123    }
124
125    /// Tries to read from the sparse mapping.
126    pub fn read_at(&self, offset: usize, data: &mut [u8]) -> Result<(), SparseMappingError> {
127        assert!(self.is_local(), "cannot read from remote mappings");
128
129        self.check(offset, data.len())?;
130        // SAFETY: the bounds have been checked above.
131        unsafe {
132            let src = (self.as_ptr() as *const u8).add(offset);
133            trycopy::try_copy(src, data.as_mut_ptr(), data.len())
134                .map_err(SparseMappingError::Memory)
135        }
136    }
137
138    /// Tries to read a type `T` from `offset`.
139    pub fn read_plain<T: FromBytes + Immutable + KnownLayout>(
140        &self,
141        offset: usize,
142    ) -> Result<T, SparseMappingError> {
143        if matches!(size_of::<T>(), 1 | 2 | 4 | 8) {
144            self.read_volatile(offset)
145        } else {
146            let mut obj = MaybeUninit::<T>::uninit();
147            // SAFETY: `obj` is a valid target for writes.
148            unsafe {
149                self.read_at(
150                    offset,
151                    std::slice::from_raw_parts_mut(obj.as_mut_ptr().cast::<u8>(), size_of::<T>()),
152                )?;
153            }
154            // SAFETY: `obj` was fully initialized by `read_at`.
155            Ok(unsafe { obj.assume_init() })
156        }
157    }
158
159    /// Tries to fill a region of the sparse mapping with `val`.
160    pub fn fill_at(&self, offset: usize, val: u8, len: usize) -> Result<(), SparseMappingError> {
161        assert!(self.is_local(), "cannot fill remote mappings");
162
163        self.check(offset, len)?;
164        // SAFETY: the bounds have been checked above.
165        unsafe {
166            let dest = self.as_ptr().cast::<u8>().add(offset);
167            trycopy::try_write_bytes(dest, val, len).map_err(SparseMappingError::Memory)
168        }
169    }
170
171    /// Gets a slice for accessing the mapped data directly.
172    ///
173    /// This is safe from a Rust memory model perspective, since the underlying
174    /// VA is either mapped and is owned in a shared state by this object (in
175    /// which case &[AtomicU8] access from multiple threads is fine), or the VA
176    /// is not mapped but is reserved and so will not be mapped by another Rust
177    /// object.
178    ///
179    /// In the latter case, actually accessing the data may cause a fault, which
180    /// will likely lead to a process crash, so care must nonetheless be taken
181    /// when using this method.
182    pub fn atomic_slice(&self, start: usize, len: usize) -> &[AtomicU8] {
183        assert!(self.len() >= start && self.len() - start >= len);
184        // SAFETY: slice is within the mapped range
185        unsafe { std::slice::from_raw_parts((self.as_ptr() as *const AtomicU8).add(start), len) }
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    static BUF: [u8; 65536] = [0xcc; 65536];
194
195    fn test_with(range_size: usize) {
196        let page_size = SparseMapping::page_size();
197
198        let mapping = SparseMapping::new(range_size).unwrap();
199        mapping.alloc(page_size, page_size).unwrap();
200        let slice = unsafe {
201            std::slice::from_raw_parts_mut(mapping.as_ptr().add(page_size).cast::<u8>(), page_size)
202        };
203        slice.copy_from_slice(&BUF[..page_size]);
204        mapping.unmap(page_size, page_size).unwrap();
205
206        mapping.alloc(range_size - page_size, page_size).unwrap();
207        let slice = unsafe {
208            std::slice::from_raw_parts_mut(
209                mapping.as_ptr().add(range_size - page_size).cast::<u8>(),
210                page_size,
211            )
212        };
213        slice.copy_from_slice(&BUF[..page_size]);
214        mapping.unmap(range_size - page_size, page_size).unwrap();
215        drop(mapping);
216    }
217
218    #[test]
219    fn test_sparse_mapping() {
220        test_with(0x100000);
221        test_with(0x200000);
222        test_with(0x200000 + SparseMapping::page_size());
223        test_with(0x40000000);
224        test_with(0x40000000 + SparseMapping::page_size());
225    }
226
227    #[test]
228    fn test_sparse_mapping_minimum_alignment() {
229        SparseMapping::new_with_minimum_alignment(SparseMapping::page_size(), 0).unwrap_err();
230
231        let mapping =
232            SparseMapping::new_with_minimum_alignment(SparseMapping::page_size(), 1).unwrap();
233        assert_eq!(mapping.as_ptr() as usize % SparseMapping::page_size(), 0);
234
235        let alignment = 0x10000;
236        let mapping =
237            SparseMapping::new_with_minimum_alignment(SparseMapping::page_size(), alignment)
238                .unwrap();
239        assert_eq!(mapping.as_ptr() as usize % alignment, 0);
240
241        // Alignments larger than the allocation granularity are honored on
242        // both platforms.
243        let alignment = 0x200000;
244        let mapping =
245            SparseMapping::new_with_minimum_alignment(SparseMapping::page_size(), alignment)
246                .unwrap();
247        assert_eq!(mapping.as_ptr() as usize % alignment, 0);
248    }
249
250    #[test]
251    fn test_sparse_mapping_default_alignment() {
252        // A small mapping only needs page alignment.
253        let mapping = SparseMapping::new(SparseMapping::page_size()).unwrap();
254        assert_eq!(mapping.as_ptr() as usize % SparseMapping::page_size(), 0);
255
256        // Mappings of at least 2 MB are aligned to a 2 MB boundary so that they
257        // can back large pages.
258        let mapping = SparseMapping::new(0x200000).unwrap();
259        assert_eq!(mapping.as_ptr() as usize % 0x200000, 0);
260
261        // Mappings of at least 1 GB are aligned to a 1 GB boundary.
262        let mapping = SparseMapping::new(0x40000000).unwrap();
263        assert_eq!(mapping.as_ptr() as usize % 0x40000000, 0);
264    }
265
266    #[test]
267    fn test_overlapping_mappings() {
268        #![expect(clippy::identity_op)]
269
270        let page_size = SparseMapping::page_size();
271        let mapping = SparseMapping::new(0x10 * page_size).unwrap();
272        mapping.alloc(0x1 * page_size, 0x4 * page_size).unwrap();
273        mapping.alloc(0x1 * page_size, 0x2 * page_size).unwrap();
274        mapping.alloc(0x2 * page_size, 0x3 * page_size).unwrap();
275        mapping.alloc(0, 0x10 * page_size).unwrap();
276        mapping.alloc(0x8 * page_size, 0x8 * page_size).unwrap();
277        mapping.unmap(0xc * page_size, 0x2 * page_size).unwrap();
278        mapping.alloc(0x9 * page_size, 0x4 * page_size).unwrap();
279        mapping.unmap(0x3 * page_size, 0xb * page_size).unwrap();
280
281        mapping.alloc(0x5 * page_size, 0x4 * page_size).unwrap();
282        mapping.alloc(0x6 * page_size, 0x2 * page_size).unwrap();
283        mapping.alloc(0x6 * page_size, 0x1 * page_size).unwrap();
284        mapping.alloc(0x4 * page_size, 0x3 * page_size).unwrap();
285
286        let shmem = alloc_shared_memory(0x4 * page_size, "test").unwrap();
287        mapping
288            .map_file(0x5 * page_size, 0x4 * page_size, &shmem, 0, true)
289            .unwrap();
290        mapping
291            .map_file(0x6 * page_size, 0x2 * page_size, &shmem, 0, true)
292            .unwrap();
293        mapping
294            .map_file(0x6 * page_size, 0x1 * page_size, &shmem, 0, true)
295            .unwrap();
296        mapping
297            .map_file(0x4 * page_size, 0x3 * page_size, &shmem, 0, true)
298            .unwrap();
299
300        drop(mapping);
301    }
302
303    #[test]
304    fn test_decommit_zeros_pages() {
305        let page_size = SparseMapping::page_size();
306        let mapping = SparseMapping::new(4 * page_size).unwrap();
307
308        // Allocate and write a pattern.
309        mapping.alloc(0, 4 * page_size).unwrap();
310        let pattern = vec![0xABu8; page_size];
311        mapping.write_at(0, &pattern).unwrap();
312        mapping.write_at(page_size, &pattern).unwrap();
313
314        // Verify data is present.
315        let mut buf = vec![0u8; page_size];
316        mapping.read_at(0, &mut buf).unwrap();
317        assert_eq!(buf, pattern);
318
319        // Decommit the first page.
320        mapping.decommit(0, page_size).unwrap();
321
322        // Read it back — should be zeros (on Linux, kernel gives zero pages;
323        // on Windows, the page is decommitted so we skip this read there).
324        #[cfg(unix)]
325        {
326            let mut buf = vec![0xFFu8; page_size];
327            mapping.read_at(0, &mut buf).unwrap();
328            assert!(
329                buf.iter().all(|&b| b == 0),
330                "decommitted page should be zeros"
331            );
332        }
333
334        // Second page should still have its data.
335        let mut buf2 = vec![0u8; page_size];
336        mapping.read_at(page_size, &mut buf2).unwrap();
337        assert_eq!(buf2, pattern);
338    }
339
340    #[test]
341    fn test_commit_after_decommit() {
342        let page_size = SparseMapping::page_size();
343        let mapping = SparseMapping::new(4 * page_size).unwrap();
344
345        // Allocate and write data.
346        mapping.alloc(0, 4 * page_size).unwrap();
347        let pattern = vec![0xCDu8; page_size];
348        mapping.write_at(0, &pattern).unwrap();
349
350        // Decommit then recommit.
351        mapping.decommit(0, page_size).unwrap();
352        mapping.commit(0, page_size).unwrap();
353
354        // After recommit, the page should be accessible and zeroed.
355        let mut buf = vec![0xFFu8; page_size];
356        mapping.read_at(0, &mut buf).unwrap();
357        assert!(
358            buf.iter().all(|&b| b == 0),
359            "recommitted page should be zeros"
360        );
361    }
362
363    #[test]
364    fn test_commit_idempotent() {
365        let page_size = SparseMapping::page_size();
366        let mapping = SparseMapping::new(4 * page_size).unwrap();
367
368        // Allocate (commit) pages.
369        mapping.alloc(0, 4 * page_size).unwrap();
370
371        // Commit the same range again — should be a no-op, no error.
372        mapping.commit(0, 4 * page_size).unwrap();
373        mapping.commit(0, page_size).unwrap();
374        mapping.commit(page_size, page_size).unwrap();
375
376        // Write and read to verify pages still work.
377        let pattern = vec![0xEFu8; page_size];
378        mapping.write_at(0, &pattern).unwrap();
379        let mut buf = vec![0u8; page_size];
380        mapping.read_at(0, &mut buf).unwrap();
381        assert_eq!(buf, pattern);
382    }
383
384    #[test]
385    #[cfg(target_os = "linux")]
386    fn test_madvise_hugepage() {
387        let page_size = SparseMapping::page_size();
388        let size = 2 * 1024 * 1024;
389        let mapping = SparseMapping::new(size).unwrap();
390        mapping.alloc(0, size).unwrap();
391
392        mapping.madvise_hugepage(0, size).unwrap();
393
394        // Memory should still work after the madvise.
395        let pattern = vec![0xABu8; page_size];
396        mapping.write_at(0, &pattern).unwrap();
397        let mut buf = vec![0u8; page_size];
398        mapping.read_at(0, &mut buf).unwrap();
399        assert_eq!(buf, pattern);
400
401        // Decommit should still zero pages with THP enabled.
402        mapping.decommit(0, page_size).unwrap();
403        #[cfg(unix)]
404        {
405            let mut buf = vec![0xFFu8; page_size];
406            mapping.read_at(0, &mut buf).unwrap();
407            assert!(
408                buf.iter().all(|&b| b == 0),
409                "decommitted page should be zeros even with THP"
410            );
411        }
412    }
413
414    #[test]
415    #[cfg(target_os = "linux")]
416    fn test_madvise_hugepage_shared() {
417        let page_size = SparseMapping::page_size();
418        let size = 2 * 1024 * 1024;
419        let shmem = alloc_shared_memory(size, "test-thp").unwrap();
420        let mapping = SparseMapping::new(size).unwrap();
421        mapping.map_file(0, size, &shmem, 0, true).unwrap();
422
423        mapping.madvise_hugepage(0, size).unwrap();
424
425        // Memory should still work after the madvise.
426        let pattern = vec![0xABu8; page_size];
427        mapping.write_at(0, &pattern).unwrap();
428        let mut buf = vec![0u8; page_size];
429        mapping.read_at(0, &mut buf).unwrap();
430        assert_eq!(buf, pattern);
431    }
432
433    #[test]
434    #[cfg(any(target_os = "linux", windows))]
435    fn test_alloc_numa_node0() {
436        let page_size = SparseMapping::page_size();
437        let size = 4 * page_size;
438        let mapping = SparseMapping::new(size).unwrap();
439
440        // Allocate with NUMA node 0 (always present).
441        #[cfg(unix)]
442        {
443            mapping.alloc(0, size).unwrap();
444            mapping.mbind_at(0, size, 0).unwrap();
445        }
446        #[cfg(windows)]
447        mapping.alloc_numa(0, size, Some(0)).unwrap();
448
449        // Memory should be accessible and writable.
450        let pattern = vec![0xABu8; page_size];
451        mapping.write_at(0, &pattern).unwrap();
452        let mut buf = vec![0u8; page_size];
453        mapping.read_at(0, &mut buf).unwrap();
454        assert_eq!(buf, pattern);
455    }
456
457    #[test]
458    #[cfg(any(target_os = "linux", windows))]
459    fn test_map_file_numa_node0() {
460        let page_size = SparseMapping::page_size();
461        let size = 4 * page_size;
462        let mapping = SparseMapping::new(size).unwrap();
463        let shmem = alloc_shared_memory(size, "test-numa").unwrap();
464
465        // Map with NUMA node 0 (always present).
466        #[cfg(unix)]
467        {
468            mapping.map_file(0, size, &shmem, 0, true).unwrap();
469            mapping.mbind_at(0, size, 0).unwrap();
470        }
471        #[cfg(windows)]
472        mapping
473            .map_file_numa(0, size, &shmem, 0, true, Some(0))
474            .unwrap();
475
476        // Memory should be accessible and writable.
477        let pattern = vec![0xCDu8; page_size];
478        mapping.write_at(0, &pattern).unwrap();
479        let mut buf = vec![0u8; page_size];
480        mapping.read_at(0, &mut buf).unwrap();
481        assert_eq!(buf, pattern);
482    }
483
484    #[test]
485    #[cfg(any(target_os = "linux", windows))]
486    fn test_alloc_numa_invalid_node() {
487        let page_size = SparseMapping::page_size();
488        let mapping = SparseMapping::new(page_size).unwrap();
489
490        // A very large NUMA node number should fail with an error (not panic).
491        #[cfg(unix)]
492        {
493            mapping.alloc(0, page_size).unwrap();
494            let result = mapping.mbind_at(0, page_size, 99999);
495            assert!(result.is_err());
496        }
497        #[cfg(windows)]
498        {
499            let result = mapping.alloc_numa(0, page_size, Some(99999));
500            assert!(result.is_err());
501        }
502    }
503}