memtable/skiplist: add a purpose-built skiplist

Author: ajwerner

Cargo.toml

@@ -47,6 +47,7 @@ fs_extra = "1.3.0"
 nanoid = "0.4.0"
 rand = "0.9.0"
 test-log = "0.2.16"
+quickcheck = "1.0.3"
 
 # half 2.5.0 has MSRV 1.81
 half = "=2.4.0"

src/memtable/mod.rs

@@ -5,18 +5,20 @@
 use crate::key::InternalKey;
 use crate::segment::block::ItemSize;
 use crate::value::{InternalValue, SeqNo, UserValue, ValueType};
-use crossbeam_skiplist::SkipMap;
 use std::ops::RangeBounds;
 use std::sync::atomic::{AtomicU32, AtomicU64};
 
+#[allow(unsafe_code)]
+mod skiplist;
+
 /// The memtable serves as an intermediary, ephemeral, sorted storage for new items
 ///
 /// When the Memtable exceeds some size, it should be flushed to a disk segment.
 #[derive(Default)]
 pub struct Memtable {
     /// The actual content, stored in a lock-free skiplist.
     #[doc(hidden)]
-    pub items: SkipMap<InternalKey, UserValue>,
+    pub items: skiplist::SkipMap<InternalKey, UserValue>,
 
     /// Approximate active memtable size.
     ///
@@ -32,7 +34,7 @@ pub struct Memtable {
 impl Memtable {
     /// Clears the memtable.
     pub fn clear(&mut self) {
-        self.items.clear();
+        self.items = Default::default();
         self.highest_seqno = AtomicU64::new(0);
         self.approximate_size
             .store(0, std::sync::atomic::Ordering::Release);
@@ -131,7 +133,11 @@ impl Memtable {
             .fetch_add(item_size, std::sync::atomic::Ordering::AcqRel);
 
         let key = InternalKey::new(item.key.user_key, item.key.seqno, item.key.value_type);
-        self.items.insert(key, item.value);
+        // TODO(ajwerner): Decide what we want to do here. The panic is sort of
+        // extreme, but also seems right given the invariants.
+        if let Err((key, _value)) = self.items.insert(key, item.value) {
+            panic!("duplicate insert of {key:?} into memtable")
+        }
 
         self.highest_seqno
             .fetch_max(item.key.seqno, std::sync::atomic::Ordering::AcqRel);

src/memtable/skiplist/arena.rs

@@ -0,0 +1,116 @@
+// Copyright (c) 2024-present, fjall-rs 
+// This source code is licensed under both the Apache 2.0 and MIT License
+// (found in the LICENSE-* files in the repository)
+
+use std::{
+    alloc::Layout,
+    mem::offset_of,
+    sync::{
+        atomic::{AtomicPtr, AtomicUsize, Ordering},
+        Mutex,
+    },
+};
+
+// DEFAULT_BUFFER_SIZE needs to be at least big enough for one fullly-aligned node
+// for the crate to work correctly. Anything larger than that will work.
+//
+// TODO: Justify this size.
+const DEFAULT_BUFFER_SIZE: usize = (32 << 10) - size_of::<AtomicUsize>();
+
+impl<const BUFFER_SIZE: usize> Default for Arenas<BUFFER_SIZE> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+unsafe impl<const N: usize> Send for Arenas<N> {}
+unsafe impl<const N: usize> Sync for Arenas<N> {}
+
+pub(crate) struct Arenas<const BUFFER_SIZE: usize = DEFAULT_BUFFER_SIZE> {
+    // The current set of Arenas
+    arenas: Mutex<Vec<*mut Buffer<BUFFER_SIZE>>>,
+    // Cache of the currently open Arena. It'll be the last item in the buffers
+    // vec. This atomic is only ever written while holding the buffers Mutex.
+    open_arena: AtomicPtr<Buffer<BUFFER_SIZE>>,
+}
+
+impl<const BUFFER_SIZE: usize> Arenas<BUFFER_SIZE> {
+    pub(crate) fn new() -> Self {
+        Self {
+            arenas: Default::default(),
+            open_arena: AtomicPtr::default(),
+        }
+    }
+}
+
+impl<const BUFFER_SIZE: usize> Arenas<BUFFER_SIZE> {
+    pub(crate) fn alloc(&self, layout: Layout) -> *mut u8 {
+        loop {
+            let buffer_tail = self.open_arena.load(Ordering::Acquire);
+            if !buffer_tail.is_null() {
+                if let Some(offset) = try_alloc(buffer_tail, layout) {
+                    return offset;
+                }
+            }
+            let mut buffers = self.arenas.lock().unwrap();
+            let buffer = buffers.last().unwrap_or(&std::ptr::null_mut());
+            if *buffer != buffer_tail {
+                // Lost the race with somebody else.
+                continue;
+            }
+            let new_buffer: Box<Buffer<BUFFER_SIZE>> = Box::new(Buffer::default());
+            let new_buffer = Box::into_raw(new_buffer);
+            self.open_arena.store(new_buffer, Ordering::Release);
+            buffers.push(new_buffer);
+        }
+    }
+}
+
+struct Buffer<const N: usize> {
+    offset: AtomicUsize,
+    data: [u8; N],
+}
+
+impl<const N: usize> Default for Buffer<N> {
+    fn default() -> Self {
+        Self {
+            offset: Default::default(),
+            data: [0; N],
+        }
+    }
+}
+
+impl<const N: usize> Drop for Arenas<N> {
+    fn drop(&mut self) {
+        let mut buffers = self.arenas.lock().unwrap();
+        for buffer in buffers.drain(..) {
+            drop(unsafe { Box::from_raw(buffer) })
+        }
+    }
+}
+
+fn try_alloc<const N: usize>(buf: *mut Buffer<N>, layout: Layout) -> Option<*mut u8> {
+    let mut cur_offset = unsafe { &(*buf).offset }.load(Ordering::Relaxed);
+    loop {
+        let buf_start = unsafe { buf.byte_add(offset_of!(Buffer<N>, data)) as *mut u8 };
+        let free_start = unsafe { buf_start.byte_add(cur_offset) };
+        let start_addr = unsafe { free_start.byte_add(free_start.align_offset(layout.align())) };
+        let new_offset = ((start_addr as usize) + layout.size()) - (buf_start as usize);
+        if new_offset > N {
+            return None;
+        }
+
+        // Note that we can get away with using relaxed ordering here because we're not
+        // asserting anything about the contents of the buffer. We're just trying to
+        // allocate a new node.
+        match unsafe { &(*buf).offset }.compare_exchange(
+            cur_offset,
+            new_offset,
+            Ordering::Relaxed,
+            Ordering::Relaxed,
+        ) {
+            Ok(_offset) => return Some(start_addr),
+            Err(offset) => cur_offset = offset,
+        }
+    }
+}

src/memtable/skiplist/mod.rs

@@ -0,0 +1,30 @@
+// Copyright (c) 2024-present, fjall-rs
+// This source code is licensed under both the Apache 2.0 and MIT License
+// (found in the LICENSE-* files in the repository)
+
+// This implementation was heavily inspired by:
+//  * https://github.com/andy-kimball/arenaskl/tree/f7010085
+//  * https://github.com/crossbeam-rs/crossbeam/tree/983d56b6/crossbeam-skiplist
+
+//! This mod is a purpose-built concurrent skiplist intended for use
+//! by the memtable.
+//!
+//! Due to the requirements of memtable, there are a number of notable in the
+//! features it lacks:
+//!     * Updates
+//!     * Deletes
+//!     * Overwrites
+//!
+//! The main reasons for its existence are that it
+//!     * provides concurrent reads and inserts, and
+//!     * batches memory allocations.
+//!
+//! Prior to this implementation, crossbeam_skiplist was used.
+
+mod arena;
+mod skipmap;
+
+pub use skipmap::SkipMap;
+
+#[cfg(test)]
+mod test;