use std::{
collections::HashMap,
fs::{self, OpenOptions},
io::{Cursor, Read},
path::{Path, PathBuf},
sync::Arc,
};
use brk_core::{Error, Result};
use memmap2::MmapMut;
use parking_lot::{RwLock, RwLockWriteGuard};
use zerocopy::{FromBytes, IntoBytes};
use super::{
Identifier, PAGE_SIZE,
region::{Region, SIZE_OF_REGION},
};
#[derive(Debug)]
pub struct Regions {
id_to_index: HashMap<String, usize>,
id_to_index_path: PathBuf,
index_to_region: Vec<Option<Arc<RwLock<Region>>>>,
index_to_region_file: fs::File,
index_to_region_mmap: MmapMut,
}
impl Regions {
pub fn open(path: &Path) -> Result<Self> {
let path = path.join("regions");
fs::create_dir_all(&path)?;
let id_to_index_path = path.join("id_to_index");
let id_to_index: HashMap<String, usize> =
deserialize_hashmap_binary(&fs::read(&id_to_index_path).unwrap_or_default())
.unwrap_or_default();
let index_to_region_file = OpenOptions::new()
.read(true)
.create(true)
.write(true)
.truncate(false)
.open(path.join("index_to_region"))?;
let index_to_region_mmap = unsafe { MmapMut::map_mut(&index_to_region_file)? };
let mut index_to_region: Vec<Option<Arc<RwLock<Region>>>> = vec![];
id_to_index
.iter()
.try_for_each(|(_, &index)| -> Result<()> {
let start = index * SIZE_OF_REGION;
let end = start + SIZE_OF_REGION;
let region = Region::read_from_bytes(&index_to_region_mmap[start..end])?;
if index_to_region.len() < index + 1 {
index_to_region.resize_with(index + 1, Default::default);
}
index_to_region
.get_mut(index)
.unwrap()
.replace(Arc::new(RwLock::new(region)));
Ok(())
})?;
Ok(Self {
id_to_index,
id_to_index_path,
index_to_region,
index_to_region_file,
index_to_region_mmap,
})
}
pub fn set_min_len(&mut self, len: u64) -> Result<()> {
if self.index_to_region_mmap.len() < len as usize {
self.index_to_region_file.set_len(len)?;
self.index_to_region_mmap = unsafe { MmapMut::map_mut(&self.index_to_region_file)? };
}
Ok(())
}
pub fn create_region(
&mut self,
id: String,
start: u64,
) -> Result<(usize, Arc<RwLock<Region>>)> {
let index = self
.index_to_region
.iter()
.enumerate()
.find(|(_, opt)| opt.is_none())
.map(|(index, _)| index)
.unwrap_or_else(|| self.index_to_region.len());
let region = Region::new(start, 0, PAGE_SIZE);
self.set_min_len(((index + 1) * SIZE_OF_REGION) as u64)?;
let region_lock = RwLock::new(region);
self.write_to_mmap(®ion_lock.write(), index);
let region_arc = Arc::new(region_lock);
let region_opt = Some(region_arc.clone());
if index < self.index_to_region.len() {
self.index_to_region[index] = region_opt
} else {
self.index_to_region.push(region_opt);
self.index_to_region_mmap.flush()?;
}
if self.id_to_index.insert(id, index).is_some() {
return Err(Error::Str("Already exists"));
}
self.flush_id_to_index()?;
Ok((index, region_arc))
}
fn flush_id_to_index(&mut self) -> Result<()> {
fs::write(
&self.id_to_index_path,
serialize_hashmap_binary(&self.id_to_index),
)?;
Ok(())
}
#[inline]
pub fn get_region(&self, identifier: Identifier) -> Option<Arc<RwLock<Region>>> {
match identifier {
Identifier::Number(index) => self.get_region_from_index(index),
Identifier::String(id) => self.get_region_from_id(&id),
}
}
#[inline]
pub fn get_region_from_index(&self, index: usize) -> Option<Arc<RwLock<Region>>> {
self.index_to_region.get(index).cloned().flatten()
}
#[inline]
pub fn get_region_from_id(&self, id: &str) -> Option<Arc<RwLock<Region>>> {
self.get_region_index_from_id(id)
.and_then(|index| self.get_region_from_index(index))
}
#[inline]
pub fn get_region_index_from_id(&self, id: &str) -> Option<usize> {
self.id_to_index.get(id).copied()
}
fn find_id_from_index(&self, index: usize) -> Option<&String> {
Some(
self.id_to_index
.iter()
.find(|(_, v)| **v == index)
.unwrap()
.0,
)
}
#[inline]
pub fn index_to_region(&self) -> &[Option<Arc<RwLock<Region>>>] {
&self.index_to_region
}
#[inline]
pub fn id_to_index(&self) -> &HashMap<String, usize> {
&self.id_to_index
}
#[inline]
pub fn identifier_to_index(&self, identifier: Identifier) -> Option<usize> {
match identifier {
Identifier::Number(index) => Some(index),
Identifier::String(id) => self.get_region_index_from_id(&id),
}
}
pub fn remove_region(&mut self, identifier: Identifier) -> Result<Option<Arc<RwLock<Region>>>> {
match identifier {
Identifier::Number(index) => self.remove_region_from_index(index),
Identifier::String(id) => self.remove_region_from_id(&id),
}
}
pub fn remove_region_from_id(&mut self, id: &str) -> Result<Option<Arc<RwLock<Region>>>> {
let Some(index) = self.get_region_index_from_id(id) else {
return Ok(None);
};
self.remove_region_from_index(index)
}
pub fn remove_region_from_index(
&mut self,
index: usize,
) -> Result<Option<Arc<RwLock<Region>>>> {
let Some(region) = self.index_to_region.get_mut(index).and_then(Option::take) else {
return Ok(None);
};
self.index_to_region_mmap.flush()?;
self.id_to_index
.remove(&self.find_id_from_index(index).unwrap().to_owned());
self.flush_id_to_index()?;
Ok(Some(region))
}
pub fn write_to_mmap(&self, region: &RwLockWriteGuard<Region>, index: usize) {
let mmap = &self.index_to_region_mmap;
let start = index * SIZE_OF_REGION;
let end = start + SIZE_OF_REGION;
if end > mmap.len() {
unreachable!("Trying to write beyond mmap")
}
let slice = unsafe { std::slice::from_raw_parts_mut(mmap.as_ptr() as *mut u8, mmap.len()) };
slice[start..end].copy_from_slice(region.as_bytes());
}
pub fn flush(&self) -> Result<()> {
self.index_to_region_mmap.flush().map_err(|e| e.into())
}
}
fn serialize_hashmap_binary(map: &HashMap<String, usize>) -> Vec<u8> {
let mut buffer = Vec::new();
buffer.extend_from_slice(&map.len().to_ne_bytes());
for (key, value) in map {
buffer.extend_from_slice(&key.len().to_ne_bytes());
buffer.extend_from_slice(key.as_bytes());
buffer.extend_from_slice(&value.to_ne_bytes());
}
buffer
}
fn deserialize_hashmap_binary(data: &[u8]) -> Result<HashMap<String, usize>> {
let mut cursor = Cursor::new(data);
let mut buffer = [0u8; 8];
cursor
.read_exact(&mut buffer)
.map_err(|_| Error::Str("Failed to read entry count"))?;
let entry_count = usize::from_ne_bytes(buffer);
let mut map = HashMap::with_capacity(entry_count);
for _ in 0..entry_count {
cursor
.read_exact(&mut buffer)
.map_err(|_| Error::Str("Failed to read key length"))?;
let key_len = usize::from_ne_bytes(buffer);
let mut key_bytes = vec![0u8; key_len];
cursor
.read_exact(&mut key_bytes)
.map_err(|_| Error::Str("Failed to read key"))?;
let key = String::from_utf8(key_bytes).map_err(|_| Error::Str("Invalid UTF-8 in key"))?;
cursor
.read_exact(&mut buffer)
.map_err(|_| Error::Str("Failed to read value"))?;
let value = usize::from_ne_bytes(buffer);
map.insert(key, value);
}
Ok(map)
}