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            if let Some(file) = file {
298                let ptr = libc::mmap64(
299                    core::ptr::null_mut(),
300                    alloc_min_len,
301                    libc::PROT_READ | libc::PROT_WRITE,
302                    libc::MAP_PRIVATE,
303                    std::os::unix::io::AsRawFd::as_raw_fd(&file),
304                    0,
305                );
306                if ptr != libc::MAP_FAILED {
307                    return Ok(HugeSlice {
308                        ptr: ptr.cast::<T>(),
309                        len,
310                        capacity: alloc_min_len,
311                        file,
312                    });
313                }
314
315                return Err(std::io::Error::last_os_error());
316            }
317
318            for (try_page_size, try_flags) in [
319                #[cfg(target_os = "linux")]
320                ((1 << 30), libc::MFD_HUGE_1GB),
321                #[cfg(target_os = "linux")]
322                ((256 << 20), libc::MFD_HUGE_256MB),
323                #[cfg(target_os = "linux")]
324                ((32 << 20), libc::MFD_HUGE_32MB),
325                #[cfg(target_os = "linux")]
326                ((16 << 20), libc::MFD_HUGE_16MB),
327                #[cfg(target_os = "linux")]
328                ((8 << 20), libc::MFD_HUGE_8MB),
329                #[cfg(target_os = "linux")]
330                ((2 << 20), libc::MFD_HUGE_2MB),
331                ((page_size), 0),
332            ]
333            .into_iter()
334            .filter(|(try_page_size, _flags)| *try_page_size >= page_size)
335            // don't grossly over size by capping page size at 2x amount of memory needed
336            .filter(|(try_page_size, flags)| *flags == 0 || alloc_min_len * 2 > *try_page_size)
337            {
338                let fd = libc::memfd_create(
339                    c"scrypt-opt-huge-page-file".as_ptr().cast(),
340                    libc::MFD_CLOEXEC | libc::MFD_HUGETLB | try_flags,
341                );
342
343                if fd == -1 {
344                    continue;
345                }
346
347                let try_file = std::os::unix::io::FromRawFd::from_raw_fd(fd);
348
349                let try_size = alloc_min_len
350                    .checked_next_multiple_of(try_page_size)
351                    .ok_or(std::io::Error::new(
352                        std::io::ErrorKind::Other,
353                        format!(
354                            "Allocation length overflow (requested: {}, max: {})",
355                            alloc_min_len,
356                            usize::MAX
357                        ),
358                    ))?;
359
360                let ptr = libc::mmap64(
361                    core::ptr::null_mut(),
362                    try_size,
363                    libc::PROT_READ | libc::PROT_WRITE,
364                    libc::MAP_PRIVATE | libc::MAP_POPULATE,
365                    fd,
366                    0,
367                );
368                if ptr != libc::MAP_FAILED {
369                    return Ok(HugeSlice {
370                        ptr: ptr.cast::<T>(),
371                        len,
372                        capacity: try_size,
373                        file: try_file,
374                    });
375                }
376            }
377
378            Err(std::io::Error::last_os_error())
379        }
380    }
381
382    #[cfg(all(
383        not(target_os = "windows"),
384        not(target_os = "android"),
385        not(target_os = "linux")
386    ))]
387    pub fn new(len: usize) -> Result<Self, std::io::Error> {
388        Err(std::io::Error::new(
389            std::io::ErrorKind::Other,
390            "Huge page not supported on this platform",
391        ))
392    }
393}
394
395#[cfg(feature = "huge-page")]
396impl<T> HugeSlice<core::mem::MaybeUninit<T>> {
397    /// Assume the buffer is initialized
398    pub unsafe fn assume_init(self) -> HugeSlice<T> {
399        let forgotten = core::mem::ManuallyDrop::new(self);
400        HugeSlice {
401            ptr: forgotten.ptr.cast::<T>(),
402            len: forgotten.len,
403            #[cfg(any(target_os = "android", target_os = "linux"))]
404            capacity: forgotten.capacity,
405            #[cfg(any(target_os = "android", target_os = "linux"))]
406            file: unsafe {
407                std::os::unix::io::FromRawFd::from_raw_fd(std::os::unix::io::AsRawFd::as_raw_fd(
408                    &forgotten.file,
409                ))
410            },
411        }
412    }
413}
414
415#[cfg(feature = "huge-page")]
416impl<T> core::convert::AsRef<[T]> for HugeSlice<T> {
417    fn as_ref(&self) -> &[T] {
418        unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
419    }
420}
421
422#[cfg(feature = "huge-page")]
423impl<T> core::convert::AsMut<[T]> for HugeSlice<T> {
424    fn as_mut(&mut self) -> &mut [T] {
425        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.len) }
426    }
427}
428
429#[cfg(feature = "huge-page")]
430impl<T> core::ops::Deref for HugeSlice<T> {
431    type Target = [T];
432    fn deref(&self) -> &Self::Target {
433        unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
434    }
435}
436
437#[cfg(feature = "huge-page")]
438impl<T> core::ops::DerefMut for HugeSlice<T> {
439    fn deref_mut(&mut self) -> &mut Self::Target {
440        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.len) }
441    }
442}
443
444#[cfg(feature = "huge-page")]
445impl<T> Drop for HugeSlice<T> {
446    #[cfg(target_os = "windows")]
447    fn drop(&mut self) {
448        use windows_sys::Win32::System::Memory::{MEM_RELEASE, VirtualFree};
449        unsafe {
450            debug_assert!(
451                VirtualFree(self.ptr as *mut _, 0, MEM_RELEASE) != 0,
452                "Failed to free huge page"
453            );
454        }
455    }
456
457    #[cfg(any(target_os = "android", target_os = "linux"))]
458    fn drop(&mut self) {
459        unsafe {
460            libc::munmap(self.ptr as *mut libc::c_void, self.capacity);
461        }
462    }
463
464    #[cfg(all(
465        not(target_os = "windows"),
466        not(target_os = "android"),
467        not(target_os = "linux")
468    ))]
469    fn drop(&mut self) {
470        // Do nothing
471    }
472}
473
474#[cfg(feature = "alloc")]
475/// A box for a buffer that can be backed by a huge page or a normal box
476pub enum MaybeHugeSlice<T> {
477    #[cfg(feature = "huge-page")]
478    /// A huge page-backed buffer
479    Huge(HugeSlice<T>),
480    /// A normal buffer backed by heap
481    Normal(alloc::boxed::Box<[T]>),
482}
483
484impl<T> core::ops::Deref for MaybeHugeSlice<T> {
485    type Target = [T];
486    fn deref(&self) -> &Self::Target {
487        self.as_ref()
488    }
489}
490
491impl<T> core::ops::DerefMut for MaybeHugeSlice<T> {
492    fn deref_mut(&mut self) -> &mut Self::Target {
493        self.as_mut()
494    }
495}
496
497impl<T> AsRef<[T]> for MaybeHugeSlice<T> {
498    fn as_ref(&self) -> &[T] {
499        match self {
500            #[cfg(feature = "huge-page")]
501            MaybeHugeSlice::Huge(b) => b.as_ref(),
502            MaybeHugeSlice::Normal(b) => b,
503        }
504    }
505}
506
507impl<T> core::convert::AsMut<[T]> for MaybeHugeSlice<T> {
508    fn as_mut(&mut self) -> &mut [T] {
509        match self {
510            #[cfg(feature = "huge-page")]
511            MaybeHugeSlice::Huge(b) => b.as_mut(),
512            MaybeHugeSlice::Normal(b) => b.as_mut(),
513        }
514    }
515}
516
517impl<T> MaybeHugeSlice<T> {
518    /// Check if the buffer is backed by a huge page
519    pub fn is_huge_page(&self) -> bool {
520        match self {
521            #[cfg(feature = "huge-page")]
522            MaybeHugeSlice::Huge(_) => true,
523            MaybeHugeSlice::Normal(_) => false,
524        }
525    }
526
527    /// Create a new huge page-backed buffer
528    #[cfg(feature = "huge-page")]
529    pub fn new_huge_slice_zeroed(len: usize) -> Result<Self, std::io::Error> {
530        let b: HugeSlice<core::mem::MaybeUninit<T>> = HugeSlice::new(len)?;
531        unsafe {
532            core::ptr::write_bytes(b.ptr.cast::<T>(), 0, len);
533            Ok(MaybeHugeSlice::Huge(b.assume_init()))
534        }
535    }
536
537    /// Create a new normal buffer
538    #[cfg(feature = "alloc")]
539    pub fn new_slice_zeroed(len: usize) -> Self {
540        let mut b = alloc::vec::Vec::<T>::with_capacity(len);
541        unsafe {
542            b.set_len(len);
543            MaybeHugeSlice::Normal(b.into())
544        }
545    }
546
547    /// Create a new buffer
548    #[cfg(feature = "alloc")]
549    pub fn new_maybe(len: usize) -> Self {
550        #[cfg(feature = "huge-page")]
551        {
552            match Self::new_huge_slice_zeroed(len) {
553                Ok(huge) => huge,
554                Err(_) => Self::new_slice_zeroed(len),
555            }
556        }
557
558        #[cfg(not(feature = "huge-page"))]
559        Self::new_slice_zeroed(len)
560    }
561
562    /// Create a new huge page-backed buffer backed by a file
563    #[cfg(all(feature = "huge-page", target_os = "linux", feature = "std"))]
564    pub fn new_in(len: usize, file: std::fs::File) -> Result<Self, std::io::Error> {
565        let huge = HugeSlice::new_unix(len, Some(file))?;
566        Ok(Self::Huge(huge))
567    }
568
569    /// Create a new buffer
570    #[cfg(feature = "std")]
571    pub fn new(len: usize) -> (Self, Option<std::io::Error>) {
572        #[cfg(feature = "huge-page")]
573        {
574            match Self::new_huge_slice_zeroed(len) {
575                Ok(huge) => (huge, None),
576                Err(e) => (Self::new_slice_zeroed(len), Some(e.into())),
577            }
578        }
579
580        #[cfg(not(feature = "huge-page"))]
581        (Self::new_slice_zeroed(len), None)
582    }
583}
scrypt_opt/memory.rs

scrypt_opt/
memory.rs
