Lock memory map when reading from file

admin/src/bench/mod.rs

@@ -5,7 +5,7 @@ use super::*;
 
 mod sizes;
 
-pub use parity_db::{CompressionType, Db, Key, Value};
+pub use parity_db::{Db, Key};
 
 use rand::{RngCore, SeedableRng};
 use std::{

src/btree/mod.rs

@@ -16,7 +16,7 @@ use crate::{
 	},
 	Operation,
 };
-pub use iter::{BTreeIterator, LastIndex, LastKey};
+pub use iter::{BTreeIterator, LastKey};
 use node::SeparatorInner;
 
 #[allow(clippy::module_inception)]

src/file.rs

@@ -124,15 +124,22 @@ impl TableFile {
 		Ok(())
 	}
 
-	pub fn slice_at(&self, offset: u64, len: usize) -> &[u8] {
+	#[cfg(not(feature = "loom"))]
+	pub fn slice_at(&self, offset: u64, len: usize) -> MappedBytesGuard {
+		let offset = offset as usize;
+		let map = self.map.read();
+		parking_lot::RwLockReadGuard::map(map, |map| {
+			let (map, _) = map.as_ref().unwrap();
+			&map[offset..offset + len]
+		})
+	}
+
+	#[cfg(feature = "loom")]
+	pub fn slice_at(&self, offset: u64, len: usize) -> MappedBytesGuard {
 		let offset = offset as usize;
 		let map = self.map.read();
 		let (map, _) = map.as_ref().unwrap();
-		let data: &[u8] = unsafe {
-			let ptr = map.as_ptr().add(offset);
-			std::slice::from_raw_parts(ptr, len)
-		};
-		data
+		MappedBytesGuard::new(map[offset..offset + len].to_vec())
 	}
 
 	pub fn write_at(&self, buf: &[u8], offset: u64) -> Result<()> {
@@ -169,8 +176,6 @@ impl TableFile {
 					let new_map = mmap(&file, new_len as usize)?;
 					let old_map = std::mem::replace(map, new_map);
 					try_io!(old_map.flush());
-					// Leak the old mapping as there might be concurrent readers.
-					std::mem::forget(old_map);
 				}
 				new_len
 			},
@@ -197,3 +202,29 @@ impl TableFile {
 		Ok(())
 	}
 }
+
+// Loom is missing support for guard projection, so we copy the data as a workaround.
+#[cfg(feature = "loom")]
+pub struct MappedBytesGuard<'a> {
+	_phantom: std::marker::PhantomData<&'a ()>,
+	data: Vec<u8>,
+}
+
+#[cfg(feature = "loom")]
+impl<'a> MappedBytesGuard<'a> {
+	pub fn new(data: Vec<u8>) -> Self {
+		Self { _phantom: std::marker::PhantomData, data }
+	}
+}
+
+#[cfg(feature = "loom")]
+impl<'a> std::ops::Deref for MappedBytesGuard<'a> {
+	type Target = [u8];
+
+	fn deref(&self) -> &Self::Target {
+		self.data.as_slice()
+	}
+}
+
+#[cfg(not(feature = "loom"))]
+pub type MappedBytesGuard<'a> = parking_lot::MappedRwLockReadGuard<'a, [u8]>;

src/log.rs

@@ -4,6 +4,7 @@
 use crate::{
 	column::ColId,
 	error::{try_io, Error, Result},
+	file::MappedBytesGuard,
 	index::{Chunk as IndexChunk, TableId as IndexTableId, ENTRY_BYTES},
 	options::Options,
 	parking_lot::{RwLock, RwLockWriteGuard},
@@ -92,32 +93,6 @@ impl LogOverlays {
 	}
 }
 
-// Loom is missing support for guard projection, so we copy the data as a workaround.
-#[cfg(feature = "loom")]
-pub struct MappedBytesGuard<'a> {
-	_phantom: std::marker::PhantomData<&'a ()>,
-	data: Vec<u8>,
-}
-
-#[cfg(feature = "loom")]
-impl<'a> MappedBytesGuard<'a> {
-	fn new(data: Vec<u8>) -> Self {
-		Self { _phantom: std::marker::PhantomData, data }
-	}
-}
-
-#[cfg(feature = "loom")]
-impl<'a> std::ops::Deref for MappedBytesGuard<'a> {
-	type Target = [u8];
-
-	fn deref(&self) -> &Self::Target {
-		self.data.as_slice()
-	}
-}
-
-#[cfg(not(feature = "loom"))]
-type MappedBytesGuard<'a> = parking_lot::MappedRwLockReadGuard<'a, [u8]>;
-
 impl LogQuery for RwLock<LogOverlays> {
 	type ValueRef<'a> = MappedBytesGuard<'a>;
 

src/table.rs

@@ -382,6 +382,20 @@ impl<B: AsRef<[u8]> + AsMut<[u8]>> std::ops::IndexMut<std::ops::Range<usize>> fo
 	}
 }
 
+enum LockedSlice<O: std::ops::Deref<Target = [u8]>, F: std::ops::Deref<Target = [u8]>> {
+	FromOverlay(O),
+	FromFile(F),
+}
+
+impl<O: std::ops::Deref<Target = [u8]>, F: std::ops::Deref<Target = [u8]>> LockedSlice<O, F> {
+	fn as_slice(&self) -> &[u8] {
+		match self {
+			LockedSlice::FromOverlay(slice) => &*slice,
+			LockedSlice::FromFile(slice) => &*slice,
+		}
+	}
+}
+
 impl ValueTable {
 	pub fn open(
 		path: Arc<std::path::PathBuf>,
@@ -458,26 +472,25 @@ impl ValueTable {
 		let entry_size = self.entry_size as usize;
 		loop {
 			let vbuf = log.value_ref(self.id, index);
-			let buf: &[u8] = if let Some(buf) = vbuf.as_deref() {
+			let buf: LockedSlice<_, _> = if let Some(buf) = vbuf.as_deref() {
 				log::trace!(
 					target: "parity-db",
 					"{}: Found in overlay {}",
 					self.id,
 					index,
 				);
-				buf
+				LockedSlice::FromOverlay(buf)
 			} else {
 				log::trace!(
 					target: "parity-db",
 					"{}: Query slot {}",
 					self.id,
 					index,
 				);
-				self.file.slice_at(index * self.entry_size as u64, entry_size)
+				let vbuf = self.file.slice_at(index * self.entry_size as u64, entry_size);
+				LockedSlice::FromFile(vbuf)
 			};
-			let mut buf = EntryRef::new(buf);
-
-			buf.set_offset(0);
+			let mut buf = EntryRef::new(buf.as_slice());
 
 			if buf.is_tombstone() {
 				return Ok((0, false))