From 95dd793835c129ffa048f311b90a9bd7b12cfee2 Mon Sep 17 00:00:00 2001 From: Al Liu Date: Sun, 30 Jun 2024 02:51:41 +0800 Subject: [PATCH 1/4] Optimize value callback --- CHANGELOG.md | 5 ++ Cargo.toml | 2 +- src/map.rs | 78 ++++++++++++++++++++++-------- src/map/api.rs | 12 ++--- src/map/tests.rs | 2 +- src/types.rs | 120 ++++++++++++++++++++++++++++++++++++++++++++++- 6 files changed, 191 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index de2b3eff..5b231df5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ # CHANGELOG +## 0.12.0 + +- Bump `rarena-allocator`'s version +- Give users back the key offset in the ARENA and the key in value callback + ## 0.11.0 - Refactor and extract lock-free ARENA allocator implementation to [`rarena-allocator`](https://github.com/al8n/rarena) crate. diff --git a/Cargo.toml b/Cargo.toml index fa55e65d..978cc8fc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "skl" -version = "0.11.4" +version = "0.12.0" edition = "2021" rust-version = "1.56.0" repository = "https://github.com/al8n/skl" diff --git a/src/map.rs b/src/map.rs index adc9e6b3..1b9d4622 100644 --- a/src/map.rs +++ b/src/map.rs @@ -399,8 +399,10 @@ impl Node { &self, arena: &'a Arena, trailer: T, + key_offset: u32, + key: &'a [u8], value_size: u32, - f: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, + f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result<(), Either> { let mut bytes = arena .alloc_aligned_bytes::(value_size) @@ -409,9 +411,10 @@ impl Node { let trailer_offset = bytes.offset(); let value_offset = trailer_offset + mem::size_of::(); - let mut oval = VacantBuffer::new(value_size as usize, value_offset as u32, unsafe { + let buffer = VacantBuffer::new(value_size as usize, value_offset as u32, unsafe { arena.get_bytes_mut(value_offset, value_size as usize) }); + let mut oval = VacantValue::new(key_offset, key, buffer); f(&mut oval).map_err(Either::Left)?; let remaining = oval.remaining(); @@ -715,7 +718,7 @@ impl SkipMap { key_size: u32, kf: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, value_size: u32, - vf: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, + vf: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result<(NodePtr, Deallocator), Either> { self .check_node_size(height, key_size, value_size) @@ -759,6 +762,8 @@ impl SkipMap { .fill_vacant_value( trailer_offset as u32, trailer_and_value.capacity() as u32, + key_offset as u32, + self.arena.get_bytes(key_offset, key_cap), value_size, value_offset, vf, @@ -897,7 +902,7 @@ impl SkipMap { key_size: u32, key_offset: u32, value_size: u32, - vf: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, + vf: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result<(NodePtr, Deallocator), Either> { self .check_node_size(height, key_size, value_size) @@ -931,6 +936,8 @@ impl SkipMap { .fill_vacant_value( trailer_offset as u32, trailer_and_value.capacity() as u32, + key_offset, + self.arena.get_bytes(key_offset as usize, key_size as usize), value_size, value_offset, vf, @@ -1026,19 +1033,23 @@ impl SkipMap { Ok((oval.len() as u32, Pointer::new(offset, size))) } + #[allow(clippy::too_many_arguments)] #[inline] unsafe fn fill_vacant_value<'a, E>( &'a self, offset: u32, size: u32, + key_offset: u32, + key: &'a [u8], value_size: u32, value_offset: u32, - f: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, + f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result<(u32, Pointer), E> { let buf = self .arena .get_bytes_mut(value_offset as usize, value_size as usize); - let mut oval = VacantBuffer::new(value_size as usize, value_offset, buf); + let buffer = VacantBuffer::new(value_size as usize, value_offset, buf); + let mut oval = VacantValue::new(key_offset, key, buffer); if let Err(e) = f(&mut oval) { self.arena.dealloc(offset, size); return Err(e); @@ -1151,7 +1162,7 @@ impl SkipMap { key: &Key<'a, 'b>, trailer: T, value_size: u32, - f: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, + f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result<(NodePtr, u32, Deallocator), Either> { let height = super::random_height(self.opts.max_height().into()); let (nd, deallocator) = match key { @@ -1732,7 +1743,7 @@ impl SkipMap { trailer: T, key: Key<'a, 'b>, value_size: u32, - f: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E> + Copy, + f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E> + Copy, success: Ordering, failure: Ordering, ins: &mut Inserter, @@ -1746,12 +1757,21 @@ impl SkipMap { if found { let node_ptr = ptr.expect("the NodePtr cannot be `None` when we found"); let old = VersionedEntryRef::from_node(node_ptr, self); - - key.on_fail(&self.arena); + let node = node_ptr.as_ref(); + key.release(&self.arena); if upsert { return self.upsert( - old, node_ptr, &key, trailer, value_size, f, success, failure, + old, + node_ptr, + node.key_offset, + node.key_size(), + &key, + trailer, + value_size, + f, + success, + failure, ); } @@ -1793,7 +1813,7 @@ impl SkipMap { }; let (nd, height, mut deallocator) = self.new_node(&k, trailer, value_size, f).map_err(|e| { - k.on_fail(&self.arena); + k.release(&self.arena); e })?; @@ -1908,12 +1928,23 @@ impl SkipMap { .curr .expect("the current should not be `None` when we found"); let old = VersionedEntryRef::from_node(node_ptr, self); - - k.on_fail(&self.arena); + let node = node_ptr.as_ref(); + k.release(&self.arena); if upsert { deallocator.dealloc(&self.arena); - return self.upsert(old, node_ptr, &k, trailer, value_size, f, success, failure); + return self.upsert( + old, + node_ptr, + node.key_offset, + node.key_size(), + &k, + trailer, + value_size, + f, + success, + failure, + ); } deallocator.dealloc(&self.arena); @@ -1925,7 +1956,7 @@ impl SkipMap { } if let Some(p) = fr.found_key { - k.on_fail(&self.arena); + k.release(&self.arena); let node = nd.as_mut(); node.key_offset = p.offset; node.key_size_and_height = encode_key_size_and_height(p.size, p.height.unwrap()); @@ -1970,17 +2001,26 @@ impl SkipMap { &'a self, old: VersionedEntryRef<'a, T, C>, node_ptr: NodePtr, + key_offset: u32, + key_len: u32, key: &Key<'a, 'b>, trailer: T, value_size: u32, - f: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, + f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, success: Ordering, failure: Ordering, ) -> Result, Either> { match key { Key::Occupied(_) | Key::Vacant(_) | Key::Pointer { .. } => node_ptr .as_ref() - .set_value(&self.arena, trailer, value_size, f) + .set_value( + &self.arena, + trailer, + key_offset, + self.arena.get_bytes(key_offset as usize, key_len as usize), + value_size, + f, + ) .map(|_| Either::Left(if old.is_removed() { None } else { Some(old) })), Key::Remove(_) | Key::RemoveVacant(_) | Key::RemovePointer { .. } => { let node = node_ptr.as_ref(); @@ -2116,6 +2156,6 @@ const fn decode_key_size_and_height(size: u32) -> (u32, u8) { #[cold] #[inline(never)] -fn noop(_: &mut VacantBuffer<'_>) -> Result<(), E> { +fn noop(_: &mut VacantValue<'_>) -> Result<(), E> { Ok(()) } diff --git a/src/map/api.rs b/src/map/api.rs index 43a18551..98d06712 100644 --- a/src/map/api.rs +++ b/src/map/api.rs @@ -457,7 +457,7 @@ impl SkipMap { return Err(Error::read_only()); } - let copy = |buf: &mut VacantBuffer| { + let copy = |buf: &mut VacantValue| { let _ = buf.write(value); Ok(()) }; @@ -536,7 +536,7 @@ impl SkipMap { trailer: T, key: &'b [u8], value_size: u32, - f: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E> + Copy, + f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E> + Copy, ) -> Result>, Either> { if self.arena.read_only() { return Err(Either::Right(Error::read_only())); @@ -580,7 +580,7 @@ impl SkipMap { return Err(Error::read_only()); } - let copy = |buf: &mut VacantBuffer| { + let copy = |buf: &mut VacantValue| { let _ = buf.write(value); Ok(()) }; @@ -660,7 +660,7 @@ impl SkipMap { trailer: T, key: &'b [u8], value_size: u32, - f: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E> + Copy, + f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E> + Copy, ) -> Result>, Either> { if self.arena.read_only() { return Err(Either::Right(Error::read_only())); @@ -742,7 +742,7 @@ impl SkipMap { key_size: u27, key: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, val_size: u32, - val: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E> + Copy, + val: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E> + Copy, ) -> Result>, Either> { let vk = self.fetch_vacant_key(u32::from(key_size), key)?; @@ -820,7 +820,7 @@ impl SkipMap { key_size: u27, key: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, val_size: u32, - val: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E> + Copy, + val: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E> + Copy, ) -> Result>, Either> { let vk = self.fetch_vacant_key(u32::from(key_size), key)?; diff --git a/src/map/tests.rs b/src/map/tests.rs index 1d317fc0..0478c04e 100644 --- a/src/map/tests.rs +++ b/src/map/tests.rs @@ -170,7 +170,7 @@ fn full_in(l: impl FnOnce(usize) -> SkipMap) { assert!(found_arena_full); let e = l - .get_or_insert(0, &make_int_key(full_at + 1), &make_value(full_at + 1)) + .get_or_insert(0, &make_int_key(full_at), &make_value(full_at)) .unwrap_err(); assert!(matches!( diff --git a/src/types.rs b/src/types.rs index b6485379..b5ec2fb2 100644 --- a/src/types.rs +++ b/src/types.rs @@ -18,6 +18,124 @@ impl core::fmt::Display for TooLarge { #[cfg(feature = "std")] impl std::error::Error for TooLarge {} +/// A vacant value in the skiplist. +#[must_use = "vacant value must be filled with bytes."] +#[derive(Debug)] +pub struct VacantValue<'a> { + buffer: VacantBuffer<'a>, + key_offset: u32, + key: &'a [u8], +} + +impl<'a> core::ops::Deref for VacantValue<'a> { + type Target = VacantBuffer<'a>; + + fn deref(&self) -> &Self::Target { + &self.buffer + } +} + +impl<'a> core::ops::DerefMut for VacantValue<'a> { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.buffer + } +} + +impl<'a> AsRef<[u8]> for VacantValue<'a> { + fn as_ref(&self) -> &[u8] { + &self.buffer + } +} + +impl<'a> AsMut<[u8]> for VacantValue<'a> { + fn as_mut(&mut self) -> &mut [u8] { + &mut self.buffer + } +} + +impl<'a> PartialEq<[u8]> for VacantValue<'a> { + fn eq(&self, other: &[u8]) -> bool { + self.buffer.eq(other) + } +} + +impl<'a> PartialEq> for [u8] { + fn eq(&self, other: &VacantValue<'a>) -> bool { + self.eq(&other.buffer) + } +} + +impl<'a> PartialEq<[u8]> for &VacantValue<'a> { + fn eq(&self, other: &[u8]) -> bool { + self.buffer.eq(other) + } +} + +impl<'a> PartialEq<&VacantValue<'a>> for [u8] { + fn eq(&self, other: &&VacantValue<'a>) -> bool { + self.eq(&other.buffer) + } +} + +impl<'a, const N: usize> PartialEq<[u8; N]> for VacantValue<'a> { + fn eq(&self, other: &[u8; N]) -> bool { + self.buffer.eq(other) + } +} + +impl<'a, const N: usize> PartialEq> for [u8; N] { + fn eq(&self, other: &VacantValue<'a>) -> bool { + self.as_ref().eq(&other.buffer) + } +} + +impl<'a, const N: usize> PartialEq<&VacantValue<'a>> for [u8; N] { + fn eq(&self, other: &&VacantValue<'a>) -> bool { + self.as_ref().eq(&other.buffer) + } +} + +impl<'a, const N: usize> PartialEq<[u8; N]> for &VacantValue<'a> { + fn eq(&self, other: &[u8; N]) -> bool { + self.buffer.eq(other) + } +} + +impl<'a, const N: usize> PartialEq<&mut VacantValue<'a>> for [u8; N] { + fn eq(&self, other: &&mut VacantValue<'a>) -> bool { + self.eq(&other.buffer) + } +} + +impl<'a, const N: usize> PartialEq<[u8; N]> for &mut VacantValue<'a> { + fn eq(&self, other: &[u8; N]) -> bool { + self.buffer.eq(other) + } +} + +impl<'a> VacantValue<'a> { + #[inline] + pub(crate) const fn new(key_offset: u32, key: &'a [u8], buffer: VacantBuffer<'a>) -> Self { + Self { + buffer, + key_offset, + key, + } + } + + /// Returns the offset of the key (corresponding to the value) in the underlying ARENA. + #[inline] + pub const fn key_offset(&self) -> u32 { + self.key_offset + } + + /// Returns the key of the vacant value. + #[inline] + pub const fn key(&self) -> &'a [u8] { + self.key + } +} + /// A vacant buffer in the skiplist. #[must_use = "vacant buffer must be filled with bytes."] #[derive(Debug)] @@ -202,7 +320,7 @@ pub(crate) enum Key<'a, 'b: 'a> { impl<'a, 'b: 'a> Key<'a, 'b> { #[inline] - pub(crate) fn on_fail(&self, arena: &super::Arena) { + pub(crate) fn release(&self, arena: &super::Arena) { match self { Self::Occupied(_) | Self::Remove(_) | Self::Pointer { .. } | Self::RemovePointer { .. } => {} Self::Vacant(key) | Self::RemoveVacant(key) => unsafe { From fe0339d8b70164771b80394a80b0e1cf2857be76 Mon Sep 17 00:00:00 2001 From: Al Liu Date: Sun, 30 Jun 2024 02:55:51 +0800 Subject: [PATCH 2/4] Optimize value callback --- CHANGELOG.md | 2 +- src/map.rs | 27 ++++++++++++++------------- src/map/api.rs | 8 ++++---- 3 files changed, 19 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b231df5..f38333ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ ## 0.12.0 - Bump `rarena-allocator`'s version +- Change value callback from `impl FnOnce + Copy` to `impl Fn` - Give users back the key offset in the ARENA and the key in value callback ## 0.11.0 @@ -84,4 +85,3 @@ ## UNRELEASED - diff --git a/src/map.rs b/src/map.rs index 1b9d4622..8da6aeb9 100644 --- a/src/map.rs +++ b/src/map.rs @@ -402,7 +402,7 @@ impl Node { key_offset: u32, key: &'a [u8], value_size: u32, - f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, + f: impl Fn(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result<(), Either> { let mut bytes = arena .alloc_aligned_bytes::(value_size) @@ -718,7 +718,7 @@ impl SkipMap { key_size: u32, kf: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, value_size: u32, - vf: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, + vf: impl Fn(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result<(NodePtr, Deallocator), Either> { self .check_node_size(height, key_size, value_size) @@ -902,7 +902,7 @@ impl SkipMap { key_size: u32, key_offset: u32, value_size: u32, - vf: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, + vf: impl Fn(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result<(NodePtr, Deallocator), Either> { self .check_node_size(height, key_size, value_size) @@ -1043,7 +1043,7 @@ impl SkipMap { key: &'a [u8], value_size: u32, value_offset: u32, - f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, + f: impl Fn(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result<(u32, Pointer), E> { let buf = self .arena @@ -1162,7 +1162,7 @@ impl SkipMap { key: &Key<'a, 'b>, trailer: T, value_size: u32, - f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, + f: &impl Fn(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result<(NodePtr, u32, Deallocator), Either> { let height = super::random_height(self.opts.max_height().into()); let (nd, deallocator) = match key { @@ -1743,7 +1743,7 @@ impl SkipMap { trailer: T, key: Key<'a, 'b>, value_size: u32, - f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E> + Copy, + f: impl Fn(&mut VacantValue<'a>) -> Result<(), E>, success: Ordering, failure: Ordering, ins: &mut Inserter, @@ -1769,7 +1769,7 @@ impl SkipMap { &key, trailer, value_size, - f, + &f, success, failure, ); @@ -1812,10 +1812,11 @@ impl SkipMap { } }; - let (nd, height, mut deallocator) = self.new_node(&k, trailer, value_size, f).map_err(|e| { - k.release(&self.arena); - e - })?; + let (nd, height, mut deallocator) = + self.new_node(&k, trailer, value_size, &f).map_err(|e| { + k.release(&self.arena); + e + })?; // We always insert from the base level and up. After you add a node in base // level, we cannot create a node in the level above because it would have @@ -1941,7 +1942,7 @@ impl SkipMap { &k, trailer, value_size, - f, + &f, success, failure, ); @@ -2006,7 +2007,7 @@ impl SkipMap { key: &Key<'a, 'b>, trailer: T, value_size: u32, - f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E>, + f: &impl Fn(&mut VacantValue<'a>) -> Result<(), E>, success: Ordering, failure: Ordering, ) -> Result, Either> { diff --git a/src/map/api.rs b/src/map/api.rs index 98d06712..05821276 100644 --- a/src/map/api.rs +++ b/src/map/api.rs @@ -536,7 +536,7 @@ impl SkipMap { trailer: T, key: &'b [u8], value_size: u32, - f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E> + Copy, + f: impl Fn(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result>, Either> { if self.arena.read_only() { return Err(Either::Right(Error::read_only())); @@ -660,7 +660,7 @@ impl SkipMap { trailer: T, key: &'b [u8], value_size: u32, - f: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E> + Copy, + f: impl Fn(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result>, Either> { if self.arena.read_only() { return Err(Either::Right(Error::read_only())); @@ -742,7 +742,7 @@ impl SkipMap { key_size: u27, key: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, val_size: u32, - val: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E> + Copy, + val: impl Fn(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result>, Either> { let vk = self.fetch_vacant_key(u32::from(key_size), key)?; @@ -820,7 +820,7 @@ impl SkipMap { key_size: u27, key: impl FnOnce(&mut VacantBuffer<'a>) -> Result<(), E>, val_size: u32, - val: impl FnOnce(&mut VacantValue<'a>) -> Result<(), E> + Copy, + val: impl Fn(&mut VacantValue<'a>) -> Result<(), E>, ) -> Result>, Either> { let vk = self.fetch_vacant_key(u32::from(key_size), key)?; From 75c2142c510709edb18cfb6d6f12396206be1659 Mon Sep 17 00:00:00 2001 From: Al Liu Date: Sun, 30 Jun 2024 13:52:18 +0800 Subject: [PATCH 3/4] Update tests.rs --- src/map/tests.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/map/tests.rs b/src/map/tests.rs index 0478c04e..4fbae622 100644 --- a/src/map/tests.rs +++ b/src/map/tests.rs @@ -1502,8 +1502,7 @@ fn test_concurrent_one_key_map_mut() { let open_options = OpenOptions::default() .create(Some(ARENA_SIZE as u32)) .read(true) - .write(true) - .shrink_on_drop(true); + .write(true); let map_options = MmapOptions::default(); concurrent_one_key(Arc::new( SkipMap::map_mut(p, open_options, map_options) @@ -2301,7 +2300,7 @@ fn test_reopen_mmap() { l.flush().unwrap(); } - let open_options = OpenOptions::default().read(true).shrink_on_drop(true); + let open_options = OpenOptions::default().read(true); let map_options = MmapOptions::default(); let l = SkipMap::::map(&p, open_options, map_options, 0).unwrap(); assert_eq!(1000, l.len()); From a1c70b5ce857717c6ed44dd159bb38149968d56c Mon Sep 17 00:00:00 2001 From: Al Liu Date: Mon, 1 Jul 2024 02:15:00 +0800 Subject: [PATCH 4/4] bumpup version --- src/map/tests.rs | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/src/map/tests.rs b/src/map/tests.rs index 4fbae622..9129a370 100644 --- a/src/map/tests.rs +++ b/src/map/tests.rs @@ -153,8 +153,6 @@ fn test_empty_map_anon_unify() { fn full_in(l: impl FnOnce(usize) -> SkipMap) { let l = l(1000); let mut found_arena_full = false; - - let mut full_at = 0; for i in 0..100 { if let Err(e) = l.get_or_insert(0, &make_int_key(i), &make_value(i)) { assert!(matches!( @@ -162,21 +160,11 @@ fn full_in(l: impl FnOnce(usize) -> SkipMap) { Error::Arena(ArenaError::InsufficientSpace { .. }) )); found_arena_full = true; - full_at = i; break; } } assert!(found_arena_full); - - let e = l - .get_or_insert(0, &make_int_key(full_at), &make_value(full_at)) - .unwrap_err(); - - assert!(matches!( - e, - Error::Arena(ArenaError::InsufficientSpace { .. }) - )); } #[test]