hugetlb: Implemented `HugePage::compute_huge()`, `scan_hugepages()`, and a public lazy static `SYSTEM_HUGEPAGES`.

Fortune for mapped-file's current commit: Future small blessing − 末小吉

Cargo.toml

@@ -8,4 +8,6 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+lazy_static = "1.4.0"
 libc = "0.2.132"
+memchr = "2.5.0"

src/err.rs

@@ -5,6 +5,28 @@ use std::{
     fmt, error
 };
 
+macro_rules! opaque {
+    ($obj:expr => $msg:literal $(, $args:expr)*) => {
+	{
+            #[derive(Debug)]
+	    struct Opaque<'a, T: ?Sized + ::std::fmt::Debug>(::std::fmt::FormatArgs<'a>, T);
+	    impl<'a, T: ?Sized + ::std::fmt::Debug> ::std::fmt::Display for Opaque<T>
+	    {
+		#[inline] 
+		fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
+		{
+		    use ::std::fmt::Write;
+		    write!(f, "{}: {:?}", self.0, &self.1)
+		}
+	    }
+	    impl<'a, T: ?Sized + ::std::fmt::Debug> error::Error for Opaque<'a, T>{}
+
+	    Opaque(::std::format_args!($msg, $(, $args)*), $obj)
+	}
+    };
+}
+pub(crate) use opaque;
+
 /// Construct an ad-hoc error wrapping the last OS error.
 macro_rules! os_error {
     ($fmt:literal $(, $args:expr)*) => {

src/hugetlb.rs

@@ -4,6 +4,8 @@ use std::{
     mem,
     hash,
     num::NonZeroUsize,
+    fs,
+    path::{Path, PathBuf},
 };
 use libc::{
     c_int,
@@ -26,12 +28,8 @@ impl Default for MapHugeFlag
     }
 }
 #[inline(always)]
-//TODO: XXX: Check this implementation of `log2<usize>()`... It seems slightly wrong...
 const fn log2(n: usize) -> usize
 {
-    /*const fn num_bits<T>() -> usize {
-    mem::size_of::<T>() * (u8::BITS as usize)
-}*/
     usize::BITS as usize -  n.leading_zeros() as usize - 1
 }
 
@@ -61,6 +59,43 @@ impl MapHugeFlag
 	Self((log2(kilobytes.get()) << (MAP_HUGE_SHIFT as usize)) as c_int)
     }
 
+    /// Attempt to calculate `MAP_HUGE_*` flag from a size (in kB).
+    #[inline]
+    pub const fn try_calculate(kilobytes: usize) -> Option<Self>
+    {
+	match kilobytes {
+	    0 => None,
+	    kilobytes => {
+		if let Some(shift) = log2(kilobytes).checked_shl(MAP_HUGE_SHIFT as u32) {
+		    if shift <= c_int::MAX as usize {
+			return Some(Self(shift as c_int));
+		    }
+		}
+		None
+	    }
+	}
+    }
+
+    /// Attempt to calculate `MAP_HUGE_*`, or use `HUGE_DEFAULT` on failure.
+    ///
+    /// # Note
+    /// If `kilobytes` is `0`, or there is a calculation overflow, then `HUGE_DEFAULT` is returned.
+    #[inline] 
+    pub const fn calculate_or_default(kilobytes: usize) -> Self
+    {
+	match Self::try_calculate(kilobytes) {
+	    None => Self::HUGE_DEFAULT,
+	    Some(x) => x,
+	}
+    }
+
+    /// Check if this is the smallest huge-page size the kernel supports.
+    #[inline] 
+    pub const fn is_default(&self) -> bool
+    {
+	self.0 == Self::HUGE_DEFAULT.0
+    }
+
     /// Get the `MAP_HUGE_*` mask.
     #[inline(always)] 
     pub const fn get_mask(self) -> c_int
@@ -71,13 +106,13 @@ impl MapHugeFlag
 
 impl From<MapHugeFlag> for c_int
 {
+    #[inline] 
     fn from(from: MapHugeFlag) -> Self
     {
 	from.0
     }
 }
 
-
 #[derive(Default, Clone, Copy)]
 pub enum HugePage {
     /// A staticly presented `MAP_HUGE_*` flag. See `MapHugeFlag` for details.
@@ -86,6 +121,8 @@ pub enum HugePage {
     ///
     /// # Safety
     /// The kernel must actually support huge-pages of this size.
+    ///
+    /// If `kilobytes` is 0, or an overflow in calculation happens, then this is identical to `Smallest`.
     Dynamic{ kilobytes: usize },
     /// The smallest huge-page size on the system
     #[default]
@@ -94,7 +131,7 @@ pub enum HugePage {
     Largest,
     /// Use a callback function to select the huge-page size (*in kB*) from an *ordered* (lowest to highest) enumeration of all available on the system.
     //TODO: Remember to order the HUGEPAGE_LOCATION parsing results before passing them to this!
-    Selected(for<'r> fn (&'r [usize]) -> &'r usize),
+    Selected(for<'r> fn (&'r [usize]) -> Option<&'r usize>),
 }
 
 impl hash::Hash for HugePage {
@@ -110,7 +147,6 @@ impl hash::Hash for HugePage {
     }
 }
 
-
 impl fmt::Debug for HugePage
 {
     #[inline] 
@@ -138,7 +174,7 @@ impl PartialEq for HugePage
     #[inline] 
     fn eq(&self, other: &Self) -> bool
     {
-	 match (self, other) {
+	match (self, other) {
 	    (Self::Static(hpf), Self::Static(hpf2)) => hpf == hpf2,
 	    (Self::Dynamic { kilobytes }, Self::Dynamic { kilobytes: kilobytes2 }) => kilobytes == kilobytes2,
 	    (Self::Selected(func), Self::Selected(func2)) => ptr::eq(func, func2),
@@ -149,11 +185,133 @@ impl PartialEq for HugePage
 
 impl HugePage
 {
-    pub fn compute_huge(&self) -> Option<MapHugeFlag>
+    /// Compute the `MapHugeFlag` from this huge-page specification.
+    ///
+    /// # Returns
+    /// * `None` - If there was an error in computing the correct flag.
+    /// * `Some` - If the computation was successful.
+    /// 
+    /// # Panics
+    /// In debug builds, if scanning the system for huge-pages fails after `SYSTEM_HUGEPAGES` has already failed.
+    #[inline]  // This call is recursive, but also can be large for variant `Selected`, which we have factored out into a non-inline local function. All other variants are small enough for this to be okay.
+    pub fn compute_huge(self) -> Option<MapHugeFlag>
     {
-	todo!("TODO: copy `utf8encode`'s `compute_huge_flag()` -> pub fn compute_flag(&self) -> Option<MapHugeFlag>;")
+	use HugePage::*;
+	match self {
+	    Dynamic { kilobytes: 0 } |
+	    Smallest |
+	    Static(MapHugeFlag::HUGE_DEFAULT) => Some(MapHugeFlag::HUGE_DEFAULT),
+	    Static(mask) => Some(mask),
+	    Dynamic { kilobytes } => {
+		MapHugeFlag::try_calculate(kilobytes) //XXX: Should we use `calculate_or_default()` here?
+	    },
+	    Largest => Self::Selected(|sizes| sizes.iter().max()).compute_huge(),
+	    Selected(func) => {
+		// Factored out into a non-`inline` function since it's the only one doing actual work, and allows the parent function to be `inline` without bloating to much
+		fn compute_selected(func: for<'r> fn (&'r [usize]) -> Option<&'r usize>) -> Option<MapHugeFlag>
+		{
+		    use std::borrow::Cow;
+		    let mask = match SYSTEM_HUGEPAGE_SIZES.as_ref() {
+			Ok(avail) => Cow::Borrowed(&avail[..]),
+			Err(_) => {
+			    // Attempt to re-scan the system. Fail if scan fails.
+			    #[cold]
+			    fn rescan() -> io::Result<Vec<usize>>
+			    {
+				scan_hugepages().and_then(|x| x.into_iter().collect())
+			    }
+			    let v = rescan();
+			    let mut v = if cfg!(debug_assertions) {
+				v.expect("Failed to compute available hugetlb sizes")
+			    } else {
+				v.ok()?
+			    };
+			    v.sort_unstable();
+			    Cow::Owned(v)
+			},
+		    };
+
+		    match func(mask.as_ref()) {
+			Some(mask) => Dynamic { kilobytes: *mask }.compute_huge(),
+			None => Some(MapHugeFlag::HUGE_DEFAULT),
+		    }
+		}
+		compute_selected(func)
+	    },
+	}
     }
-    //TODO: ^
 }
 
-//TODO: implement `memfd`'s hugetlb interface from `utf8encode` here.
+lazy_static! {
+    /// A persistent invocation of `scan_hugepages()`.
+    pub(crate) static ref SYSTEM_HUGEPAGE_SIZES: io::Result<Vec<usize>> = {
+	let mut val: io::Result<Vec<usize>> = scan_hugepages().and_then(|x| x.into_iter().collect());
+	if let Ok(ref mut arr) = val.as_mut() {
+	    arr.sort_unstable();
+	};
+	val
+    };
+
+    /// A list of all availble huge-page flags if enumeration of them is possible.
+    ///
+    /// This is created from a persistent invocation of `scan_hugepages()`.
+    pub static ref SYSTEM_HUGEPAGES: io::Result<Vec<MapHugeFlag>> =
+	SYSTEM_HUGEPAGE_SIZES.as_ref()
+	.map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, format!("SYSTEM_HUGEPAGES failed with error {err}")))
+	.map(|vec| vec.iter().map(|&size| MapHugeFlag::calculate_or_default(size)).collect());
+}
+
+/// Scan the system for available huge-page sizes (in kB).
+///
+/// # Returns
+/// If reading the directory `HUGEPAGE_LOCATION` fails, then the error is returned.
+/// Otherwise, an iterator over each item in this location, parsed for its size, is returned.
+/// If reading an entry fails, an error is returned.
+///
+/// If an entry is not parsed correctly, then it is skipped.
+pub fn scan_hugepages() -> io::Result<impl IntoIterator<Item=io::Result<usize>> + Send + Sync + 'static>
+{
+    let path = Path::new(HUGEPAGE_LOCATION);
+    let dir = fs::read_dir(path)?;
+
+    #[derive(Debug)]
+    struct FilteredIterator(fs::ReadDir);
+    
+    impl Iterator for FilteredIterator
+    {
+	type Item = io::Result<usize>;
+	fn next(&mut self) -> Option<Self::Item> {
+	    loop {
+		break if let Some(next) = self.0.next() {
+		    let path = match next {
+			Ok(next) => next.file_name(),
+			Err(err) => return Some(Err(err)),
+		    };
+		    let kbs = if let Some(dash) = memchr::memchr(b'-', path.as_bytes()) {
+			let name = &path.as_bytes()[(dash+1)..];
+			if let Some(k_loc) = memchr::memrchr(b'k', &name) {
+			    &name[..k_loc]
+			} else {
+			    continue
+			}
+		    } else {
+			continue
+		    };
+		    let kb = if let Ok(kbs) = std::str::from_utf8(kbs) {
+			kbs.parse::<usize>().ok()
+		    } else {
+			continue
+		    };
+		    match kb {
+			None => continue,
+			valid => valid.map(Ok)
+		    }
+		} else {
+		    None
+		}
+	    }
+	}
+    }
+    
+    Ok(FilteredIterator(dir))
+}

src/lib.rs

@@ -1,4 +1,6 @@
 
+#[macro_use] extern crate lazy_static;
+
 use libc::{
     mmap,
     MAP_FAILED,
@@ -38,7 +40,10 @@ mod flags;
 pub use flags::*;
 
 pub mod err;
-use err::os_error;
+use err::{
+    os_error,
+    opaque,
+};
 
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]