use anyhow::{Result, anyhow};
use memmap2::{Mmap, MmapOptions};
use std::collections::{BTreeMap, HashMap};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::fs::File;
use std::sync::atomic::{AtomicUsize, Ordering};
#[derive(Debug, Clone, Copy)]
pub struct MapOptions {
pub read_only: bool,
pub shared: bool,
pub populate: bool,
}
impl Default for MapOptions {
fn default() -> Self {
Self {
read_only: true,
shared: false,
populate: false,
}
}
}
#[derive(Debug, Clone)]
pub struct MapperStats {
pub total_mappings: u64,
pub total_bytes_mapped: u64,
pub cache_hits: u64,
pub cache_misses: u64,
}
impl Default for MapperStats {
fn default() -> Self {
Self {
total_mappings: 0,
total_bytes_mapped: 0,
cache_hits: 0,
cache_misses: 0,
}
}
}
#[derive(Debug)]
pub struct MappedFile {
mmap: Arc<Mmap>,
path: PathBuf,
size: usize,
}
#[derive(Debug)]
struct MappingCache {
cache: lru::LruCache<PathBuf, Arc<Mmap>>,
stats: Arc<AtomicUsize>,
max_entries: usize,
}
impl MappingCache {
fn new(capacity: usize) -> Self {
Self {
cache: lru::LruCache::new(NonZeroUsize::new(capacity.max(1)).unwrap()),
stats: Arc::new(Mutex::new(AtomicUsize::new(0))),
max_entries: capacity,
}
}
fn get(&mut self, path: &Path) -> Option<&Arc<Mmap>> {
if let Some(mmap) = self.cache.get(path) {
self.stats.fetch_add(1, Ordering::Relaxed); Some(mmap)
} else {
self.stats.fetch_add(1, Ordering::Relaxed); None
}
}
fn insert(&mut self, path: PathBuf, mmap: Arc<Mmap>) {
if self.cache.len() >= self.max_entries {
if let Some((old_path, _)) = self.cache.pop_lru() {
}
}
self.cache.put(path, mmap);
}
fn clear(&mut self) {
self.cache.clear();
}
}
#[derive(Debug)]
pub struct MemoryMapper {
cache: Arc<tokio::sync::Mutex<MappingCache>>,
page_size: usize,
stats: Arc<AtomicUsize>,
}
impl MemoryMapper {
pub fn new(cache_size: usize) -> Result<Self> {
let page_size: _ = unsafe { libc::sysconf(libc::_SC_PAGESIZE) } as usize;
Ok(Self {
cache: Arc::new(Mutex::new(tokio::sync::Mutex::new(MappingCache::new(cache_size)))),
page_size,
stats: Arc::new(Mutex::new(AtomicUsize::new(0))),
})
}
pub async fn map_file(&self, path: &Path, options: MapOptions) -> Result<MappedFile> {
{
let mut cache = self.cache.lock().await;
if let Some(mmap) = cache.get(path) {
let mapped_file: _ = MappedFile {
mmap: Arc::clone(mmap),
path: path.to_path_buf(),
size: mmap.len(),
};
self.stats.fetch_add(mapped_file.size, Ordering::Relaxed);
return Ok(mapped_file);
}
}
let file: _ = File::open(path).await?;
let metadata: _ = file.metadata().await?;
let file_size: _ = metadata.len() as usize;
let mut mmap_options = MmapOptions::new();
if options.read_only {
} else {
unsafe { mmap_options.map_copy(&file)? };
}
if options.populate {
mmap_options.populate();
}
let mmap: _ = unsafe { mmap_options.map(&file)? };
let mmap_arc: _ = Arc::new(Mutex::new(mmap));
{
let mut cache = self.cache.lock().await;
cache.insert(path.to_path_buf(), Arc::clone(mmap_arc));
}
let mapped_file: _ = MappedFile {
mmap: mmap_arc,
path: path.to_path_buf(),
size: file_size,
};
self.stats.fetch_add(mapped_file.size, Ordering::Relaxed);
Ok(mapped_file)
}
pub fn align_to_page(&self, addr: usize) -> usize {
addr & !(self.page_size - 1)
}
pub fn align_size_to_page(&self, size: usize) -> usize {
(size + self.page_size - 1) & !(self.page_size - 1)
}
pub async fn create_copy_on_write(&self, base: &MappedFile) -> Result<MappedFile> {
let file: _ = File::open(&base.path).await?;
let mmap: _ = unsafe { MmapOptions::new().map(&file)? };
Ok(MappedFile {
mmap: Arc::new(Mutex::new(mmap)),
path: base.path.clone(),
size: base.size,
})
}
pub fn advise(&self, data: &[u8], advice: MemoryAdvice) -> Result<()> {
let addr: _ = data.as_ptr() as usize;
let len: _ = data.len();
let advice_value: _ = match advice {
MemoryAdvice::Normal => libc::POSIX_MADV_NORMAL,
MemoryAdvice::Sequential => libc::POSIX_MADV_SEQUENTIAL,
MemoryAdvice::Random => libc::POSIX_MADV_RANDOM,
MemoryAdvice::WillNeed => libc::POSIX_MADV_WILLNEED,
MemoryAdvice::DontNeed => libc::POSIX_MADV_DONTNEED,
};
unsafe {
libc::posix_madvise(
addr as *mut libc::c_void,
len,
advice_value,
);
}
Ok(())
}
pub async fn get_stats(&self) -> MapperStats {
let cache: _ = self.cache.lock().await;
MapperStats {
total_mappings: self.stats.load(Ordering::Relaxed) as u64,
total_bytes_mapped: self.stats.load(Ordering::Relaxed) as u64,
cache_hits: cache.stats.load(Ordering::Relaxed) as u64,
cache_misses: 0, }
}
pub async fn clear_cache(&self) {
let mut cache = self.cache.lock().await;
cache.clear();
}
}
#[derive(Debug, Clone, Copy)]
pub enum MemoryAdvice {
Normal,
Sequential,
Random,
WillNeed,
DontNeed,
}
impl MappedFile {
pub fn as_bytes(&self) -> &[u8] {
&self.mmap
}
pub fn as_mut_bytes(&mut self) -> Option<&mut [u8]> {
None
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn size(&self) -> usize {
self.size
}
pub fn clone(&self) -> Self {
MappedFile {
mmap: Arc::clone(&self.mmap),
path: self.path.clone(),
size: self.size,
}
}
}
#[cfg(test)]
mod tests {
#[tokio::test]
async fn test_map_file() {
let mapper: _ = MemoryMapper::new(16).unwrap();
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"Hello, World!").unwrap();
let mapped: _ = mapper.map_file(temp_file.path(), MapOptions::default()).await.unwrap();
assert_eq!(mapped.size(), 13);
assert_eq!(mapped.as_bytes(), b"Hello, World!");
}
#[tokio::test]
async fn test_align_to_page() {
let mapper: _ = MemoryMapper::new(16).unwrap();
let page_size: _ = mapper.page_size;
let addr: _ = 12345;
let aligned: _ = mapper.align_to_page(addr);
assert!(aligned <= addr);
assert_eq!(aligned % page_size, 0);
assert!(aligned + page_size > addr);
}
#[tokio::test]
async fn test_align_size_to_page() {
let mapper: _ = MemoryMapper::new(16).unwrap();
let page_size: _ = mapper.page_size;
let size: _ = 12345;
let aligned: _ = mapper.align_size_to_page(size);
assert!(aligned >= size);
assert_eq!(aligned % page_size, 0);
assert!(aligned < size + page_size);
}
#[tokio::test]
async fn test_copy_on_write() {
let mapper: _ = MemoryMapper::new(16).unwrap();
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"Original data").unwrap();
let base: _ = mapper.map_file(temp_file.path(), MapOptions::default()).await.unwrap();
let cow: _ = mapper.create_copy_on_write(&base).await.unwrap();
assert_eq!(cow.size(), base.size());
assert_eq!(cow.as_bytes(), base.as_bytes());
}
#[tokio::test]
async fn test_memory_advice() {
let mapper: _ = MemoryMapper::new(16).unwrap();
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"Test data").unwrap();
let mapped: _ = mapper.map_file(temp_file.path(), MapOptions::default()).await.unwrap();
let result: _ = mapper.advise(mapped.as_bytes(), MemoryAdvice::Sequential);
assert!(result.is_ok());
}
#[tokio::test]
async fn test_mapped_file_clone() {
let mapper: _ = MemoryMapper::new(16).unwrap();
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"Clone test").unwrap();
let mapped1: _ = mapper.map_file(temp_file.path(), MapOptions::default()).await.unwrap();
let mapped2: _ = mapped1.clone();
assert_eq!(mapped1.size(), mapped2.size());
assert_eq!(mapped1.as_bytes().as_ptr(), mapped2.as_bytes().as_ptr());
}
#[tokio::test]
async fn test_stats() {
let mapper: _ = MemoryMapper::new(16).unwrap();
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"Stats test").unwrap();
let _mapped: _ = mapper.map_file(temp_file.path(), MapOptions::default()).await.unwrap();
let stats: _ = mapper.get_stats().await;
assert!(stats.total_mappings > 0);
assert!(stats.total_bytes_mapped >= 10);
}
#[tokio::test]
async fn test_cache() {
let mapper: _ = MemoryMapper::new(16).unwrap();
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"Cache test").unwrap();
let mapped1: _ = mapper.map_file(temp_file.path(), MapOptions::default()).await.unwrap();
let _mapped2: _ = mapper.map_file(temp_file.path(), MapOptions::default()).await.unwrap();
let stats: _ = mapper.get_stats().await;
assert!(stats.cache_hits > 0);
}
}