use std::fs::{File, OpenOptions};
use std::marker::PhantomData;
use std::ops;
use std::path::Path;
use anyhow::{Context, Result};
use log::error;
use memmap2::MmapMut;
use crate::merkle::Element;
use crate::store::{Store, StoreConfig};
#[derive(Debug)]
pub struct MmapStore<E: Element> {
map: Option<MmapMut>,
file: File,
len: usize,
_e: PhantomData<E>,
}
impl<E: Element> ops::Deref for MmapStore<E> {
type Target = [u8];
fn deref(&self) -> &Self::Target {
&self.map.as_ref().expect("couldn't dereference")[..]
}
}
impl<E: Element> Store<E> for MmapStore<E> {
#[allow(unsafe_code)]
fn new_with_config(size: usize, branches: usize, config: StoreConfig) -> Result<Self> {
let data_path = StoreConfig::data_path(&config.path, &config.id);
if Path::new(&data_path).exists() {
return Self::new_from_disk(size, branches, &config);
}
let file = OpenOptions::new()
.write(true)
.read(true)
.create_new(true)
.open(&data_path)?;
let store_size = E::byte_len() * size;
file.set_len(store_size as u64)?;
let map = unsafe { MmapMut::map_mut(&file)? };
Ok(MmapStore {
map: Some(map),
file,
len: 0,
_e: Default::default(),
})
}
#[allow(unsafe_code)]
fn new(size: usize) -> Result<Self> {
let store_size = E::byte_len() * size;
let file = tempfile::NamedTempFile::new()?;
file.as_file().set_len(store_size as u64)?;
let (file, _path) = file.into_parts();
let map = unsafe { MmapMut::map_mut(&file)? };
Ok(MmapStore {
map: Some(map),
file,
len: 0,
_e: Default::default(),
})
}
#[allow(unsafe_code)]
fn new_from_disk(size: usize, _branches: usize, config: &StoreConfig) -> Result<Self> {
let data_path = StoreConfig::data_path(&config.path, &config.id);
ensure!(Path::new(&data_path).exists(), "[MmapStore] new_from_disk constructor can be used only for instantiating already existing storages");
let file = match OpenOptions::new().write(true).read(true).open(&data_path) {
Ok(file) => file,
Err(e) => {
if e.kind() == std::io::ErrorKind::PermissionDenied {
error!("MmapStore doesn't support read-only storages");
}
panic!("{}", e)
}
};
let metadata = file.metadata()?;
let store_size = metadata.len() as usize;
ensure!(
store_size == size * E::byte_len(),
"Invalid formatted file provided. Expected {} bytes, found {} bytes",
size * E::byte_len(),
store_size
);
let map = unsafe { MmapMut::map_mut(&file)? };
Ok(MmapStore {
map: Some(map),
file,
len: size,
_e: Default::default(),
})
}
fn write_at(&mut self, el: E, index: usize) -> Result<()> {
let start = index * E::byte_len();
let end = start + E::byte_len();
if self.map.is_none() {
self.reinit()?;
}
self.map.as_mut().unwrap()[start..end].copy_from_slice(el.as_ref());
self.len = std::cmp::max(self.len, index + 1);
Ok(())
}
fn copy_from_slice(&mut self, buf: &[u8], start: usize) -> Result<()> {
ensure!(
buf.len() % E::byte_len() == 0,
"buf size must be a multiple of {}",
E::byte_len()
);
let map_start = start * E::byte_len();
let map_end = map_start + buf.len();
if self.map.is_none() {
self.reinit()?;
}
self.map.as_mut().unwrap()[map_start..map_end].copy_from_slice(buf);
self.len = std::cmp::max(self.len, start + (buf.len() / E::byte_len()));
Ok(())
}
fn new_from_slice_with_config(
size: usize,
branches: usize,
data: &[u8],
config: StoreConfig,
) -> Result<Self> {
ensure!(
data.len() % E::byte_len() == 0,
"data size must be a multiple of {}",
E::byte_len()
);
let mut store = Self::new_with_config(size, branches, config)?;
if !store.loaded_from_disk() {
if store.map.is_none() {
store.reinit()?;
}
let len = data.len();
store.map.as_mut().unwrap()[0..len].copy_from_slice(data);
store.len = len / E::byte_len();
}
Ok(store)
}
fn new_from_slice(size: usize, data: &[u8]) -> Result<Self> {
ensure!(
data.len() % E::byte_len() == 0,
"data size must be a multiple of {}",
E::byte_len()
);
let mut store = Self::new(size)?;
ensure!(store.map.is_some(), "Internal map needs to be initialized");
let len = data.len();
store.map.as_mut().unwrap()[0..len].copy_from_slice(data);
store.len = len / E::byte_len();
Ok(store)
}
fn read_at(&self, index: usize) -> Result<E> {
let start = index * E::byte_len();
let end = start + E::byte_len();
let len = self.len * E::byte_len();
ensure!(start < len, "start out of range {} >= {}", start, len);
ensure!(end <= len, "end out of range {} > {}", end, len);
let data = self
.map
.as_ref()
.context("Internal map needs to be initialized")?;
Ok(E::from_slice(&data[start..end]))
}
fn read_into(&self, index: usize, buf: &mut [u8]) -> Result<()> {
let start = index * E::byte_len();
let end = start + E::byte_len();
let len = self.len * E::byte_len();
ensure!(start < len, "start out of range {} >= {}", start, len);
ensure!(end <= len, "end out of range {} > {}", end, len);
let data = self
.map
.as_ref()
.context("Internal map needs to be initialized")?;
buf.copy_from_slice(&data[start..end]);
Ok(())
}
fn read_range_into(&self, _start: usize, _end: usize, _buf: &mut [u8]) -> Result<()> {
unimplemented!("Not required here");
}
fn read_range(&self, r: ops::Range<usize>) -> Result<Vec<E>> {
let start = r.start * E::byte_len();
let end = r.end * E::byte_len();
let len = self.len * E::byte_len();
ensure!(start < len, "start out of range {} >= {}", start, len);
ensure!(end <= len, "end out of range {} > {}", end, len);
let data = self
.map
.as_ref()
.context("Internal map needs to be initialized")?;
Ok(data[start..end]
.chunks(E::byte_len())
.map(E::from_slice)
.collect())
}
fn len(&self) -> usize {
self.len
}
fn loaded_from_disk(&self) -> bool {
false
}
fn compact(
&mut self,
_branches: usize,
_config: StoreConfig,
_store_version: u32,
) -> Result<bool> {
let map = self.map.take();
Ok(map.is_some())
}
#[allow(unsafe_code)]
fn reinit(&mut self) -> Result<()> {
self.map = unsafe { Some(MmapMut::map_mut(&self.file)?) };
ensure!(self.map.is_some(), "Re-init mapping failed");
Ok(())
}
fn delete(_config: StoreConfig) -> Result<()> {
Ok(())
}
fn is_empty(&self) -> bool {
self.len == 0
}
fn push(&mut self, el: E) -> Result<()> {
let l = self.len;
if self.map.is_none() {
self.reinit()?;
}
ensure!(
(l + 1) * E::byte_len() <= self.map.as_ref().unwrap().len(),
"not enough space"
);
self.write_at(el, l)
}
}