#[derive(Debug, Clone, Default)]
pub struct MlockConfig {
pub enabled: bool,
pub max_locked_bytes: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MlockResult {
Locked,
Disabled,
InsufficientPrivileges,
ResourceLimit,
Unsupported,
}
pub struct PinnedRegion {
ptr: *const u8,
len: usize,
locked: bool,
}
unsafe impl Send for PinnedRegion {}
unsafe impl Sync for PinnedRegion {}
impl PinnedRegion {
#[must_use]
pub unsafe fn new(ptr: *const u8, len: usize, config: &MlockConfig) -> (Self, MlockResult) {
if !config.enabled {
return (
Self {
ptr,
len,
locked: false,
},
MlockResult::Disabled,
);
}
if config.max_locked_bytes > 0 && len > config.max_locked_bytes {
return (
Self {
ptr,
len,
locked: false,
},
MlockResult::ResourceLimit,
);
}
let result = Self::mlock_impl(ptr, len);
let locked = result == MlockResult::Locked;
(Self { ptr, len, locked }, result)
}
#[must_use]
pub fn is_locked(&self) -> bool {
self.locked
}
#[must_use]
pub fn len(&self) -> usize {
self.len
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[cfg(target_family = "unix")]
fn mlock_impl(ptr: *const u8, len: usize) -> MlockResult {
let result = unsafe { libc::mlock(ptr.cast(), len) };
if result == 0 {
MlockResult::Locked
} else {
let errno = std::io::Error::last_os_error().raw_os_error().unwrap_or(0);
if errno == libc::EPERM {
MlockResult::InsufficientPrivileges
} else {
MlockResult::ResourceLimit
}
}
}
#[cfg(not(target_family = "unix"))]
fn mlock_impl(_ptr: *const u8, _len: usize) -> MlockResult {
MlockResult::Unsupported
}
}
impl Drop for PinnedRegion {
fn drop(&mut self) {
if self.locked {
#[cfg(target_family = "unix")]
unsafe {
libc::munlock(self.ptr.cast(), self.len);
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExpertTier {
Hot,
Warm,
Cold,
}
impl ExpertTier {
#[must_use]
pub fn from_access_count(count: usize, hot_threshold: usize, warm_threshold: usize) -> Self {
if count >= hot_threshold {
Self::Hot
} else if count >= warm_threshold {
Self::Warm
} else {
Self::Cold
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_mlock_disabled() {
let config = MlockConfig {
enabled: false,
max_locked_bytes: 0,
};
let data = vec![0u8; 1024];
let (region, result) = unsafe { PinnedRegion::new(data.as_ptr(), data.len(), &config) };
assert_eq!(result, MlockResult::Disabled);
assert!(!region.is_locked());
}
#[test]
fn test_mlock_resource_limit() {
let config = MlockConfig {
enabled: true,
max_locked_bytes: 100,
};
let data = vec![0u8; 1024]; let (region, result) = unsafe { PinnedRegion::new(data.as_ptr(), data.len(), &config) };
assert_eq!(result, MlockResult::ResourceLimit);
assert!(!region.is_locked());
}
#[test]
fn test_expert_tier_hot() {
let tier = ExpertTier::from_access_count(1000, 100, 10);
assert_eq!(tier, ExpertTier::Hot);
}
#[test]
fn test_expert_tier_warm() {
let tier = ExpertTier::from_access_count(50, 100, 10);
assert_eq!(tier, ExpertTier::Warm);
}
#[test]
fn test_expert_tier_cold() {
let tier = ExpertTier::from_access_count(5, 100, 10);
assert_eq!(tier, ExpertTier::Cold);
}
#[test]
fn test_pinned_region_len() {
let config = MlockConfig::default();
let data = vec![0u8; 1024];
let (region, _) = unsafe { PinnedRegion::new(data.as_ptr(), data.len(), &config) };
assert_eq!(region.len(), 1024);
assert!(!region.is_empty());
}
#[test]
fn test_pinned_region_empty() {
let config = MlockConfig::default();
let data: Vec<u8> = vec![];
let (region, _) = unsafe { PinnedRegion::new(data.as_ptr(), data.len(), &config) };
assert_eq!(region.len(), 0);
assert!(region.is_empty());
}
#[test]
fn test_mlock_config_default() {
let config = MlockConfig::default();
assert!(!config.enabled);
assert_eq!(config.max_locked_bytes, 0);
}
#[test]
fn test_mlock_result_equality() {
assert_eq!(MlockResult::Locked, MlockResult::Locked);
assert_eq!(MlockResult::Disabled, MlockResult::Disabled);
assert_ne!(MlockResult::Locked, MlockResult::Disabled);
}
#[test]
fn test_expert_tier_boundary() {
assert_eq!(ExpertTier::from_access_count(100, 100, 10), ExpertTier::Hot);
assert_eq!(ExpertTier::from_access_count(99, 100, 10), ExpertTier::Warm);
assert_eq!(ExpertTier::from_access_count(10, 100, 10), ExpertTier::Warm);
assert_eq!(ExpertTier::from_access_count(9, 100, 10), ExpertTier::Cold);
}
#[test]
fn test_mlock_enabled_within_limit() {
let config = MlockConfig {
enabled: true,
max_locked_bytes: 0, };
let data: [u8; 64] = [0u8; 64];
let (region, result) = unsafe { PinnedRegion::new(data.as_ptr(), data.len(), &config) };
assert!(
result == MlockResult::Locked
|| result == MlockResult::InsufficientPrivileges
|| result == MlockResult::ResourceLimit
|| result == MlockResult::Unsupported
);
assert_eq!(region.len(), 64);
}
}