scrypt_opt/
memory.rs

1use core::ops::{Deref, DerefMut};
2
3#[repr(align(64))]
4#[derive(Default, Debug, PartialEq, Eq, Clone, Copy)]
5/// Align to 64 bytes
6pub struct Align64<T>(pub T);
7
8impl<T> Align64<T> {
9    /// Perform a reference cast to an [`Align64<Block<R>>`] without checking the alignment.
10    ///
11    /// # Safety
12    /// This function is unsafe because it does not check the alignment of the input pointer.
13    #[inline(always)]
14    pub const unsafe fn new_unchecked(input: &T) -> &Self {
15        unsafe { &*(input as *const T as *const Self) }
16    }
17
18    /// Perform a reference cast to an [`Align64<Block<R>>`].
19    ///
20    /// # Panics
21    ///
22    /// Panics if the input pointer is not aligned to 64 bytes.
23    pub fn new(input: &T) -> &Self {
24        let ptr = input as *const T;
25        assert_eq!(
26            ptr.align_offset(64),
27            0,
28            "Input pointer is not aligned to 64 bytes"
29        );
30        unsafe { Self::new_unchecked(input) }
31    }
32
33    /// Perform a mutable reference cast to an [`Align64<Block<R>>`] without checking the alignment.
34    ///
35    /// # Safety
36    /// This function is unsafe because it does not check the alignment of the input pointer.
37    #[inline(always)]
38    pub const unsafe fn new_mut_unchecked(input: &mut T) -> &mut Self {
39        unsafe { &mut *(input as *mut T as *mut Self) }
40    }
41
42    /// Perform a mutable reference cast to an [`Align64<Block<R>>`].
43    ///
44    /// # Panics
45    ///
46    /// Panics if the input pointer is not aligned to 64 bytes.
47    pub fn new_mut(input: &mut T) -> &mut Self {
48        let ptr = input as *mut T;
49        assert_eq!(
50            ptr.align_offset(64),
51            0,
52            "Input pointer is not aligned to 64 bytes"
53        );
54        unsafe { Self::new_mut_unchecked(input) }
55    }
56}
57
58impl<T> AsRef<T> for Align64<T> {
59    fn as_ref(&self) -> &T {
60        &self.0
61    }
62}
63
64impl<T> AsRef<Align64<T>> for Align64<T> {
65    fn as_ref(&self) -> &Align64<T> {
66        self
67    }
68}
69
70impl<T> AsMut<Align64<T>> for Align64<T> {
71    fn as_mut(&mut self) -> &mut Align64<T> {
72        self
73    }
74}
75
76impl<T> AsMut<T> for Align64<T> {
77    fn as_mut(&mut self) -> &mut T {
78        &mut self.0
79    }
80}
81
82impl<T> Deref for Align64<T> {
83    type Target = T;
84    fn deref(&self) -> &Self::Target {
85        &self.0
86    }
87}
88
89impl<T> DerefMut for Align64<T> {
90    fn deref_mut(&mut self) -> &mut Self::Target {
91        &mut self.0
92    }
93}
94
95/// A box for a hugepage-backed buffer
96#[cfg(feature = "huge-page")]
97pub struct HugeSlice<T> {
98    ptr: *mut T,
99    len: usize,
100    #[cfg(any(target_os = "android", target_os = "linux"))]
101    capacity: usize,
102    #[cfg(any(target_os = "android", target_os = "linux"))]
103    file: std::fs::File,
104}
105
106#[cfg(feature = "huge-page")]
107unsafe impl<T> Send for HugeSlice<T> where T: Send {}
108
109#[cfg(feature = "huge-page")]
110unsafe impl<T> Sync for HugeSlice<T> where T: Sync {}
111
112#[cfg(feature = "huge-page")]
113impl<T> core::convert::AsMut<T> for HugeSlice<T> {
114    fn as_mut(&mut self) -> &mut T {
115        unsafe { &mut *self.ptr }
116    }
117}
118
119#[cfg(feature = "huge-page")]
120impl<T> HugeSlice<T> {
121    /// Create a new hugepage-backed buffer
122    #[cfg(target_os = "windows")]
123    pub fn new(len: usize) -> Result<Self, std::io::Error> {
124        use windows_sys::Win32::{
125            Security::TOKEN_ADJUST_PRIVILEGES,
126            System::{
127                Memory::{
128                    GetLargePageMinimum, MEM_COMMIT, MEM_LARGE_PAGES, MEM_RESERVE, PAGE_READWRITE,
129                    VirtualAlloc,
130                },
131                Threading::{GetCurrentProcess, OpenProcessToken},
132            },
133        };
134        unsafe {
135            use windows_sys::Win32::Security::{
136                AdjustTokenPrivileges, LUID_AND_ATTRIBUTES, LookupPrivilegeValueA,
137                SE_PRIVILEGE_ENABLED, TOKEN_PRIVILEGES,
138            };
139
140            let large_page_minimum = GetLargePageMinimum();
141            if large_page_minimum == 0 {
142                return Err(std::io::Error::new(
143                    std::io::ErrorKind::Other,
144                    "Large page not supported by the system",
145                ));
146            }
147
148            if core::mem::align_of::<T>() > large_page_minimum {
149                return Err(std::io::Error::new(
150                    std::io::ErrorKind::Other,
151                    format!(
152                        "Alignment of the type is greater than the large page minimum: {} requires {} alignment, large page minimum is {}",
153                        core::any::type_name::<T>(),
154                        core::mem::align_of::<T>(),
155                        large_page_minimum
156                    ),
157                ));
158            }
159
160            let min_alloc_size = core::mem::size_of::<T>()
161                .checked_mul(len)
162                .and_then(|x| x.checked_next_multiple_of(large_page_minimum))
163                .and_then(|x| x.checked_next_multiple_of(8))
164                .ok_or_else(|| {
165                    std::io::Error::new(
166                        std::io::ErrorKind::Other,
167                        format!(
168                            "Allocation size overflow (requested: {}, max: {})",
169                            core::mem::size_of::<T>() as u128 * len as u128,
170                            usize::MAX
171                        ),
172                    )
173                })?;
174
175            if min_alloc_size == 0 {
176                return Err(std::io::Error::new(
177                    std::io::ErrorKind::Other,
178                    "Allocation size must be non-zero",
179                ));
180            }
181
182            let mut token_handle = Default::default();
183            if OpenProcessToken(
184                GetCurrentProcess(),
185                TOKEN_ADJUST_PRIVILEGES,
186                &mut token_handle,
187            ) == 0
188            {
189                return Err(std::io::Error::last_os_error());
190            }
191
192            let mut privs = [TOKEN_PRIVILEGES {
193                PrivilegeCount: 1,
194                Privileges: [LUID_AND_ATTRIBUTES {
195                    Luid: Default::default(),
196                    Attributes: SE_PRIVILEGE_ENABLED,
197                }],
198            }];
199
200            if LookupPrivilegeValueA(
201                core::ptr::null_mut(),
202                c"SeLockMemoryPrivilege".as_ptr().cast(),
203                &mut privs[0].Privileges[0].Luid,
204            ) == 0
205            {
206                return Err(std::io::Error::last_os_error());
207            }
208
209            if AdjustTokenPrivileges(
210                token_handle,
211                0,
212                privs.as_ptr(),
213                0,
214                core::ptr::null_mut(),
215                core::ptr::null_mut(),
216            ) == 0
217            {
218                return Err(std::io::Error::new(
219                    std::io::ErrorKind::Other,
220                    "Failed to adjust token privileges",
221                ));
222            }
223
224            let ptr = VirtualAlloc(
225                core::ptr::null_mut(),
226                min_alloc_size,
227                MEM_COMMIT | MEM_RESERVE | MEM_LARGE_PAGES,
228                PAGE_READWRITE,
229            );
230
231            if ptr == core::ptr::null_mut() {
232                return Err(std::io::Error::last_os_error());
233            }
234
235            Ok(Self {
236                ptr: ptr.cast::<T>(),
237                len,
238            })
239        }
240    }
241
242    /// Create a new hugepage-backed buffer
243    #[cfg(any(target_os = "android", target_os = "linux"))]
244    pub fn new(len: usize) -> Result<Self, std::io::Error> {
245        Self::new_unix(
246            len,
247            #[cfg(all(target_os = "linux", feature = "std"))]
248            None,
249        )
250    }
251
252    /// Create a new hugepage-backed buffer backed by a file
253    #[cfg(any(target_os = "android", target_os = "linux"))]
254    pub fn new_unix(
255        len: usize,
256        #[cfg(target_os = "linux")] file: Option<std::fs::File>,
257    ) -> Result<Self, std::io::Error> {
258        unsafe {
259            let pagesz = libc::sysconf(libc::_SC_PAGESIZE);
260            if pagesz == -1 {
261                return Err(std::io::Error::last_os_error());
262            }
263
264            let page_size = pagesz as usize;
265            if core::mem::align_of::<T>() > page_size {
266                return Err(std::io::Error::new(
267                    std::io::ErrorKind::Other,
268                    format!(
269                        "Alignment of the type is greater than the page size: {} requires {} alignment, page size is {}",
270                        core::any::type_name::<T>(),
271                        core::mem::align_of::<T>(),
272                        page_size
273                    ),
274                ));
275            }
276
277            let alloc_min_len = core::mem::size_of::<T>()
278                .checked_mul(len)
279                .and_then(|x| x.checked_next_multiple_of(page_size))
280                .and_then(|x| x.checked_next_multiple_of(8))
281                .ok_or(std::io::Error::new(
282                    std::io::ErrorKind::Other,
283                    format!(
284                        "Allocation length overflow (requested: {}, max: {})",
285                        core::mem::size_of::<T>() as u128 * len as u128,
286                        usize::MAX
287                    ),
288                ))?;
289
290            if alloc_min_len == 0 {
291                return Err(std::io::Error::new(
292                    std::io::ErrorKind::Other,
293                    "Allocation length must be non-zero",
294                ));
295            }
296
297            #[cfg(target_os = "linux")]
298            if let Some(file) = file {
299                let ptr = libc::mmap64(
300                    core::ptr::null_mut(),
301                    alloc_min_len,
302                    libc::PROT_READ | libc::PROT_WRITE,
303                    libc::MAP_PRIVATE,
304                    std::os::unix::io::AsRawFd::as_raw_fd(&file),
305                    0,
306                );
307                if ptr != libc::MAP_FAILED {
308                    return Ok(HugeSlice {
309                        ptr: ptr.cast::<T>(),
310                        len,
311                        capacity: alloc_min_len,
312                        file,
313                    });
314                }
315
316                return Err(std::io::Error::last_os_error());
317            }
318
319            for (try_page_size, try_flags) in [
320                #[cfg(target_os = "linux")]
321                ((1 << 30), libc::MFD_HUGE_1GB),
322                #[cfg(target_os = "linux")]
323                ((256 << 20), libc::MFD_HUGE_256MB),
324                #[cfg(target_os = "linux")]
325                ((32 << 20), libc::MFD_HUGE_32MB),
326                #[cfg(target_os = "linux")]
327                ((16 << 20), libc::MFD_HUGE_16MB),
328                #[cfg(target_os = "linux")]
329                ((8 << 20), libc::MFD_HUGE_8MB),
330                #[cfg(target_os = "linux")]
331                ((2 << 20), libc::MFD_HUGE_2MB),
332                ((page_size), 0),
333            ]
334            .into_iter()
335            .filter(|(try_page_size, _flags)| *try_page_size >= page_size)
336            // don't grossly over size by capping page size at 2x amount of memory needed
337            .filter(|(try_page_size, flags)| *flags == 0 || alloc_min_len * 2 > *try_page_size)
338            {
339                let fd = libc::memfd_create(
340                    c"scrypt-opt-huge-page-file".as_ptr().cast(),
341                    libc::MFD_CLOEXEC | libc::MFD_HUGETLB | try_flags,
342                );
343
344                if fd == -1 {
345                    continue;
346                }
347
348                let try_file = std::os::unix::io::FromRawFd::from_raw_fd(fd);
349
350                let try_size = alloc_min_len
351                    .checked_next_multiple_of(try_page_size)
352                    .ok_or(std::io::Error::new(
353                        std::io::ErrorKind::Other,
354                        format!(
355                            "Allocation length overflow (requested: {}, max: {})",
356                            alloc_min_len,
357                            usize::MAX
358                        ),
359                    ))?;
360
361                let ptr = libc::mmap64(
362                    core::ptr::null_mut(),
363                    try_size,
364                    libc::PROT_READ | libc::PROT_WRITE,
365                    libc::MAP_PRIVATE | libc::MAP_POPULATE,
366                    fd,
367                    0,
368                );
369                if ptr != libc::MAP_FAILED {
370                    return Ok(HugeSlice {
371                        ptr: ptr.cast::<T>(),
372                        len,
373                        capacity: try_size,
374                        file: try_file,
375                    });
376                }
377            }
378
379            Err(std::io::Error::last_os_error())
380        }
381    }
382
383    #[cfg(all(
384        not(target_os = "windows"),
385        not(target_os = "android"),
386        not(target_os = "linux")
387    ))]
388    /// Create a new hugepage-backed buffer
389    pub fn new(_len: usize) -> Result<Self, std::io::Error> {
390        Err(std::io::Error::new(
391            std::io::ErrorKind::Other,
392            "Huge page not supported on this platform",
393        ))
394    }
395}
396
397#[cfg(feature = "huge-page")]
398impl<T> HugeSlice<core::mem::MaybeUninit<T>> {
399    /// Assume the buffer is initialized
400    pub unsafe fn assume_init(self) -> HugeSlice<T> {
401        let forgotten = core::mem::ManuallyDrop::new(self);
402        HugeSlice {
403            ptr: forgotten.ptr.cast::<T>(),
404            len: forgotten.len,
405            #[cfg(any(target_os = "android", target_os = "linux"))]
406            capacity: forgotten.capacity,
407            #[cfg(any(target_os = "android", target_os = "linux"))]
408            file: unsafe {
409                std::os::unix::io::FromRawFd::from_raw_fd(std::os::unix::io::AsRawFd::as_raw_fd(
410                    &forgotten.file,
411                ))
412            },
413        }
414    }
415}
416
417#[cfg(feature = "huge-page")]
418impl<T> core::convert::AsRef<[T]> for HugeSlice<T> {
419    fn as_ref(&self) -> &[T] {
420        unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
421    }
422}
423
424#[cfg(feature = "huge-page")]
425impl<T> core::convert::AsMut<[T]> for HugeSlice<T> {
426    fn as_mut(&mut self) -> &mut [T] {
427        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.len) }
428    }
429}
430
431#[cfg(feature = "huge-page")]
432impl<T> core::ops::Deref for HugeSlice<T> {
433    type Target = [T];
434    fn deref(&self) -> &Self::Target {
435        unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
436    }
437}
438
439#[cfg(feature = "huge-page")]
440impl<T> core::ops::DerefMut for HugeSlice<T> {
441    fn deref_mut(&mut self) -> &mut Self::Target {
442        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.len) }
443    }
444}
445
446#[cfg(feature = "huge-page")]
447impl<T> Drop for HugeSlice<T> {
448    #[cfg(target_os = "windows")]
449    fn drop(&mut self) {
450        use windows_sys::Win32::System::Memory::{MEM_RELEASE, VirtualFree};
451        unsafe {
452            debug_assert!(
453                VirtualFree(self.ptr as *mut _, 0, MEM_RELEASE) != 0,
454                "Failed to free huge page"
455            );
456        }
457    }
458
459    #[cfg(any(target_os = "android", target_os = "linux"))]
460    fn drop(&mut self) {
461        unsafe {
462            libc::munmap(self.ptr as *mut libc::c_void, self.capacity);
463        }
464    }
465
466    #[cfg(all(
467        not(target_os = "windows"),
468        not(target_os = "android"),
469        not(target_os = "linux")
470    ))]
471    fn drop(&mut self) {
472        // Do nothing
473    }
474}
475
476#[cfg(feature = "alloc")]
477/// A box for a buffer that can be backed by a huge page or a normal box
478pub enum MaybeHugeSlice<T> {
479    #[cfg(feature = "huge-page")]
480    /// A huge page-backed buffer
481    Huge(HugeSlice<T>),
482    /// A normal buffer backed by heap
483    Normal(alloc::boxed::Box<[T]>),
484}
485
486impl<T> core::ops::Deref for MaybeHugeSlice<T> {
487    type Target = [T];
488    fn deref(&self) -> &Self::Target {
489        self.as_ref()
490    }
491}
492
493impl<T> core::ops::DerefMut for MaybeHugeSlice<T> {
494    fn deref_mut(&mut self) -> &mut Self::Target {
495        self.as_mut()
496    }
497}
498
499impl<T> AsRef<[T]> for MaybeHugeSlice<T> {
500    fn as_ref(&self) -> &[T] {
501        match self {
502            #[cfg(feature = "huge-page")]
503            MaybeHugeSlice::Huge(b) => b.as_ref(),
504            MaybeHugeSlice::Normal(b) => b,
505        }
506    }
507}
508
509impl<T> core::convert::AsMut<[T]> for MaybeHugeSlice<T> {
510    fn as_mut(&mut self) -> &mut [T] {
511        match self {
512            #[cfg(feature = "huge-page")]
513            MaybeHugeSlice::Huge(b) => b.as_mut(),
514            MaybeHugeSlice::Normal(b) => b.as_mut(),
515        }
516    }
517}
518
519impl<T> MaybeHugeSlice<T> {
520    /// Check if the buffer is backed by a huge page
521    pub fn is_huge_page(&self) -> bool {
522        match self {
523            #[cfg(feature = "huge-page")]
524            MaybeHugeSlice::Huge(_) => true,
525            MaybeHugeSlice::Normal(_) => false,
526        }
527    }
528
529    /// Create a new huge page-backed buffer
530    #[cfg(feature = "huge-page")]
531    pub fn new_huge_slice_zeroed(len: usize) -> Result<Self, std::io::Error> {
532        let b: HugeSlice<core::mem::MaybeUninit<T>> = HugeSlice::new(len)?;
533        unsafe {
534            core::ptr::write_bytes(b.ptr.cast::<T>(), 0, len);
535            Ok(MaybeHugeSlice::Huge(b.assume_init()))
536        }
537    }
538
539    /// Create a new normal buffer
540    #[cfg(feature = "alloc")]
541    pub fn new_slice_zeroed(len: usize) -> Self {
542        let mut b = alloc::vec::Vec::<T>::with_capacity(len);
543        unsafe {
544            b.set_len(len);
545            MaybeHugeSlice::Normal(b.into())
546        }
547    }
548
549    /// Create a new buffer
550    #[cfg(feature = "alloc")]
551    pub fn new_maybe(len: usize) -> Self {
552        #[cfg(feature = "huge-page")]
553        {
554            match Self::new_huge_slice_zeroed(len) {
555                Ok(huge) => huge,
556                Err(_) => Self::new_slice_zeroed(len),
557            }
558        }
559
560        #[cfg(not(feature = "huge-page"))]
561        Self::new_slice_zeroed(len)
562    }
563
564    /// Create a new huge page-backed buffer backed by a file
565    #[cfg(all(feature = "huge-page", target_os = "linux", feature = "std"))]
566    pub fn new_in(len: usize, file: std::fs::File) -> Result<Self, std::io::Error> {
567        let huge = HugeSlice::new_unix(len, Some(file))?;
568        Ok(Self::Huge(huge))
569    }
570
571    /// Create a new buffer
572    #[cfg(feature = "std")]
573    pub fn new(len: usize) -> (Self, Option<std::io::Error>) {
574        #[cfg(feature = "huge-page")]
575        {
576            match Self::new_huge_slice_zeroed(len) {
577                Ok(huge) => (huge, None),
578                Err(e) => (Self::new_slice_zeroed(len), Some(e.into())),
579            }
580        }
581
582        #[cfg(not(feature = "huge-page"))]
583        (Self::new_slice_zeroed(len), None)
584    }
585}
scrypt_opt/memory.rs

scrypt_opt/
memory.rs
