mirror of
https://github.com/Pumpkin-MC/Pumpkin.git
synced 2026-08-30 20:14:23 +00:00
Fix File IO Lock Manager (#703)
* fix file locker * fix linear bounds check
This commit is contained in:
@@ -19,7 +19,8 @@ use tokio::{
|
||||
|
||||
use crate::{
|
||||
chunk::{
|
||||
ChunkData, ChunkReadingError, ChunkSerializingError, ChunkWritingError, CompressionError,
|
||||
ChunkData, ChunkParsingError, ChunkReadingError, ChunkSerializingError, ChunkWritingError,
|
||||
CompressionError,
|
||||
io::{ChunkSerializer, LoadedData},
|
||||
},
|
||||
generation::section_coords,
|
||||
@@ -239,7 +240,7 @@ impl AnvilChunkData {
|
||||
#[inline]
|
||||
fn raw_write_size(&self) -> usize {
|
||||
// 4 bytes for the *length* and 1 byte for the *compression* method
|
||||
self.compressed_data.remaining() + 4 + 1
|
||||
self.compressed_data.len() + 4 + 1
|
||||
}
|
||||
|
||||
/// Size of serialized chunk with padding
|
||||
@@ -260,6 +261,16 @@ impl AnvilChunkData {
|
||||
// Minus one for the compression byte
|
||||
let length = bytes.get_u32() as usize - 1;
|
||||
|
||||
if length > bytes.len() {
|
||||
return Err(ChunkReadingError::ParsingError(
|
||||
ChunkParsingError::ErrorDeserializingChunk(format!(
|
||||
"Chunk length is greater than available bytes ({} vs {})",
|
||||
length,
|
||||
bytes.len()
|
||||
)),
|
||||
));
|
||||
}
|
||||
|
||||
let compression_method = bytes.get_u8();
|
||||
let compression = Compression::from_byte(compression_method)
|
||||
.map_err(|_| ChunkReadingError::Compression(CompressionError::UnknownCompression))?;
|
||||
@@ -574,6 +585,17 @@ impl ChunkSerializer for AnvilChunkFile {
|
||||
let bytes_offset = (sector_offset - 2) * SECTOR_BYTES;
|
||||
let bytes_count = sector_count * SECTOR_BYTES;
|
||||
|
||||
if bytes_offset + bytes_count > raw_file_bytes.len() {
|
||||
return Err(ChunkReadingError::ParsingError(
|
||||
ChunkParsingError::ErrorDeserializingChunk(format!(
|
||||
"Not enough bytes available for the chunk {} ({} vs {})",
|
||||
i,
|
||||
bytes_count,
|
||||
raw_file_bytes.len().saturating_sub(bytes_offset)
|
||||
)),
|
||||
));
|
||||
}
|
||||
|
||||
let serialized_data = AnvilChunkData::from_bytes(
|
||||
raw_file_bytes.slice(bytes_offset..bytes_offset + bytes_count),
|
||||
)?;
|
||||
|
||||
@@ -296,7 +296,16 @@ impl ChunkSerializer for LinearFile {
|
||||
if header.size != 0 {
|
||||
let last_index = bytes_offset;
|
||||
bytes_offset += header.size as usize;
|
||||
chunks[i] = Some(buffer.slice(last_index..bytes_offset));
|
||||
if bytes_offset > buffer.len() {
|
||||
log::warn!(
|
||||
"Not enough bytes are available for chunk {} ({} vs {})",
|
||||
i,
|
||||
header.size,
|
||||
buffer.len() - last_index
|
||||
);
|
||||
} else {
|
||||
chunks[i] = Some(buffer.slice(last_index..bytes_offset));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -29,14 +29,91 @@ use super::{ChunkIO, ChunkSerializer, LoadedData};
|
||||
/// using parallelism and a cache for the files with ongoing IO operations.
|
||||
///
|
||||
/// It also avoid IO operations that could produce dataraces thanks to the
|
||||
/// custom *DashMap* like implementation.
|
||||
/// custom *DashMap*-like implementation.
|
||||
pub struct ChunkFileManager<S: ChunkSerializer<WriteBackend = PathBuf>> {
|
||||
// Dashmap has rw-locks on shards, but we want per-serializer
|
||||
file_locks: RwLock<BTreeMap<PathBuf, SerializerCacheEntry<S>>>,
|
||||
// Dashmap has rw-locks on shards, but we want per-serializer.
|
||||
//
|
||||
// Lock lookups need to be:
|
||||
// - Relatively quick insertion / lookup as to not extraneously block other concurrent hashmap
|
||||
// ops
|
||||
// - Guarantee that there is only one serializer per file at a time
|
||||
// - Lazily load files as to support point 1
|
||||
// - Allow for ease of usage (able to return the serializer from a function)
|
||||
file_locks: RwLock<BTreeMap<PathBuf, ChunkSerializerLazyLoader<S>>>,
|
||||
watchers: RwLock<BTreeMap<PathBuf, usize>>,
|
||||
}
|
||||
//to avoid clippy warnings we extract the type alias
|
||||
type SerializerCacheEntry<S> = OnceCell<Arc<RwLock<S>>>;
|
||||
|
||||
struct ChunkSerializerLazyLoader<S: ChunkSerializer<WriteBackend = PathBuf>> {
|
||||
path: PathBuf,
|
||||
internal: OnceCell<Arc<RwLock<S>>>,
|
||||
}
|
||||
|
||||
impl<S: ChunkSerializer<Data = ChunkData, WriteBackend = PathBuf>> ChunkSerializerLazyLoader<S> {
|
||||
fn new(path: PathBuf) -> Self {
|
||||
Self {
|
||||
path,
|
||||
internal: OnceCell::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// We can only remove this entry from the map if we are the only ones with a reference to it
|
||||
/// IMPORTANT: This must be called within the write lock of the parent map
|
||||
async fn can_remove(&self) -> bool {
|
||||
match self.internal.get() {
|
||||
Some(arc) => {
|
||||
let _write_lock = arc.write().await;
|
||||
Arc::strong_count(arc) == 1
|
||||
}
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
|
||||
async fn get(&self) -> Result<Arc<RwLock<S>>, ChunkReadingError> {
|
||||
self.internal
|
||||
.get_or_try_init(|| async {
|
||||
let serializer = self.read_from_disk().await?;
|
||||
Ok(Arc::new(RwLock::new(serializer)))
|
||||
})
|
||||
.await
|
||||
.cloned()
|
||||
}
|
||||
|
||||
async fn read_from_disk(&self) -> Result<S, ChunkReadingError> {
|
||||
trace!("Opening file from Disk: {:?}", self.path);
|
||||
let file = tokio::fs::OpenOptions::new()
|
||||
.read(true)
|
||||
.write(false)
|
||||
.create(false)
|
||||
.truncate(false)
|
||||
.open(&self.path)
|
||||
.await
|
||||
.map_err(|err| match err.kind() {
|
||||
ErrorKind::NotFound => ChunkReadingError::ChunkNotExist,
|
||||
kind => ChunkReadingError::IoError(kind),
|
||||
});
|
||||
|
||||
let value = match file {
|
||||
Ok(mut file) => {
|
||||
let capacity = match file.metadata().await {
|
||||
Ok(metadata) => metadata.len() as usize,
|
||||
Err(_) => 4096, // A sane default
|
||||
};
|
||||
|
||||
// TODO: Memmap?
|
||||
let mut file_bytes = Vec::with_capacity(capacity);
|
||||
file.read_to_end(&mut file_bytes)
|
||||
.await
|
||||
.map_err(|err| ChunkReadingError::IoError(err.kind()))?;
|
||||
S::read(file_bytes.into())?
|
||||
}
|
||||
Err(ChunkReadingError::ChunkNotExist) => S::default(),
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
|
||||
trace!("Successfully read file from Disk: {:?}", self.path);
|
||||
Ok(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: ChunkSerializer<WriteBackend = PathBuf>> Default for ChunkFileManager<S> {
|
||||
fn default() -> Self {
|
||||
@@ -51,71 +128,26 @@ impl<S: ChunkSerializer<WriteBackend = PathBuf>> ChunkFileManager<S> {
|
||||
fn map_key(folder: &LevelFolder, file_name: &str) -> PathBuf {
|
||||
folder.region_folder.join(file_name)
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_file(&self, path: &Path) -> Result<Arc<RwLock<S>>, ChunkReadingError> {
|
||||
impl<S: ChunkSerializer<Data = ChunkData, WriteBackend = PathBuf>> ChunkFileManager<S> {
|
||||
async fn get_serializer(&self, path: &Path) -> Result<Arc<RwLock<S>>, ChunkReadingError> {
|
||||
// We get the entry from the DashMap and try to insert a new lock if it doesn't exist
|
||||
// using dead-lock safe methods like `or_try_insert_with`
|
||||
|
||||
async fn read_from_disk<S: ChunkSerializer>(
|
||||
path: &Path,
|
||||
) -> Result<Arc<RwLock<S>>, ChunkReadingError> {
|
||||
trace!("Opening file from Disk: {:?}", path);
|
||||
let file = tokio::fs::OpenOptions::new()
|
||||
.read(true)
|
||||
.write(false)
|
||||
.create(false)
|
||||
.truncate(false)
|
||||
.open(path)
|
||||
.await
|
||||
.map_err(|err| match err.kind() {
|
||||
ErrorKind::NotFound => ChunkReadingError::ChunkNotExist,
|
||||
kind => ChunkReadingError::IoError(kind),
|
||||
});
|
||||
|
||||
let value = match file {
|
||||
Ok(mut file) => {
|
||||
let capacity = match file.metadata().await {
|
||||
Ok(metadata) => metadata.len() as usize,
|
||||
Err(_) => 4096, // A sane default
|
||||
};
|
||||
|
||||
let mut file_bytes = Vec::with_capacity(capacity);
|
||||
file.read_to_end(&mut file_bytes)
|
||||
.await
|
||||
.map_err(|err| ChunkReadingError::IoError(err.kind()))?;
|
||||
S::read(file_bytes.into())?
|
||||
}
|
||||
Err(ChunkReadingError::ChunkNotExist) => S::default(),
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
|
||||
trace!("Successfully read file from Disk: {:?}", path);
|
||||
Ok(Arc::new(RwLock::new(value)))
|
||||
}
|
||||
|
||||
// We use a once lock here to quickly make an insertion into the map without holding the
|
||||
// We use a lazy loader here to quickly make an insertion into the map without holding the
|
||||
// lock for too long starving other threads
|
||||
let serializer = if let Some(once_cell) = self.file_locks.read().await.get(path) {
|
||||
log::trace!("Loading file lock from cache: {:?}", path);
|
||||
once_cell
|
||||
.get_or_try_init(|| read_from_disk(path))
|
||||
.await?
|
||||
.clone()
|
||||
if let Some(serializer_loader) = self.file_locks.read().await.get(path) {
|
||||
serializer_loader.get().await
|
||||
} else {
|
||||
log::trace!("Cache miss loading file lock from cache: {:?}", path);
|
||||
let mut file_locks = self.file_locks.write().await;
|
||||
file_locks
|
||||
.entry(path.to_path_buf())
|
||||
.or_insert_with(OnceCell::new);
|
||||
let file_locks = file_locks.downgrade();
|
||||
let once_cell = file_locks.get(path).expect("We just inserted this!");
|
||||
once_cell
|
||||
.get_or_try_init(|| read_from_disk(path))
|
||||
.await?
|
||||
.clone()
|
||||
};
|
||||
|
||||
Ok(serializer)
|
||||
self.file_locks
|
||||
.write()
|
||||
.await
|
||||
.entry(path.into())
|
||||
.or_insert_with(|| ChunkSerializerLazyLoader::new(path.into()))
|
||||
.get()
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -185,7 +217,7 @@ where
|
||||
// Also improves File Cache utilizations.
|
||||
let region_read_tasks = regions_chunks.into_iter().map(async |(file_name, chunks)| {
|
||||
let path = Self::map_key(folder, &file_name);
|
||||
let chunk_serializer = match self.read_file(&path).await {
|
||||
let chunk_serializer = match self.get_serializer(&path).await {
|
||||
Ok(chunk_serializer) => chunk_serializer,
|
||||
Err(ChunkReadingError::ChunkNotExist) => {
|
||||
unreachable!("Default Serializer must be created")
|
||||
@@ -247,7 +279,7 @@ where
|
||||
let path = Self::map_key(folder, &file_name);
|
||||
log::trace!("Updating data for file {:?}", path);
|
||||
|
||||
let chunk_serializer = match self.read_file(&path).await {
|
||||
let chunk_serializer = match self.get_serializer(&path).await {
|
||||
Ok(file) => Ok(file),
|
||||
Err(ChunkReadingError::ChunkNotExist) => {
|
||||
unreachable!("Must be managed by the cache")
|
||||
@@ -297,7 +329,11 @@ where
|
||||
.write(path.clone())
|
||||
.await
|
||||
.map_err(|err| ChunkWritingError::IoError(err.kind()))?;
|
||||
|
||||
// Remove lock
|
||||
drop(serializer);
|
||||
// Decrement strong count
|
||||
drop(chunk_serializer);
|
||||
|
||||
// If there are still no watchers, drop from the locks
|
||||
let mut locks = self.file_locks.write().await;
|
||||
@@ -309,8 +345,18 @@ where
|
||||
.get(&path)
|
||||
.is_none_or(|count| count.is_zero())
|
||||
{
|
||||
locks.remove(&path);
|
||||
log::trace!("Removed lockfile cache {:?}", path);
|
||||
let can_remove = if let Some(loader) = locks.get(&path) {
|
||||
loader.can_remove().await
|
||||
} else {
|
||||
true
|
||||
};
|
||||
|
||||
if can_remove {
|
||||
locks.remove(&path);
|
||||
log::trace!("Removed lockfile cache {:?}", path);
|
||||
} else {
|
||||
log::trace!("Wanted to remove lockfile cache {:?} but someone still holds a reference to it!", path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -333,23 +379,10 @@ where
|
||||
//we need to block any other operation
|
||||
let serializer_cache = self.file_locks.write().await;
|
||||
|
||||
let locks: Vec<_> = serializer_cache
|
||||
.iter()
|
||||
.map(|(pos, value)| (pos, value.clone()))
|
||||
.collect();
|
||||
|
||||
// Acquire a write lock on all entries to verify they are complete
|
||||
let tasks = locks.iter().map(async |(pos, serializer)| {
|
||||
if let Some(lock) = serializer.get() {
|
||||
Some(lock.write().await)
|
||||
} else {
|
||||
log::warn!(
|
||||
"Closing FileManager while the File {} is being loaded",
|
||||
pos.display()
|
||||
);
|
||||
None
|
||||
}
|
||||
});
|
||||
let tasks = serializer_cache
|
||||
.values()
|
||||
.map(async |serializer| serializer.can_remove());
|
||||
|
||||
// We need to wait to ensure that all the locks are acquired
|
||||
// so there is no **operation** ongoing
|
||||
|
||||
Reference in New Issue
Block a user