use std::alloc::{self, Layout};
use std::cell::Cell;
use std::cmp::Ordering;
use std::mem::{align_of, size_of};
use std::ptr::{self, NonNull};
use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering::*};
use parking_lot::Mutex;
use crate::key;
pub const MAX_HEIGHT: usize = 16;
const BRANCHING: u64 = 4;
const BLOCK: usize = 32 << 10;
const SHARDS: usize = 16;
const LARGE: usize = BLOCK / 4;
const FIRST_CHUNK: usize = 64 << 10;
const MAX_CHUNK: usize = 4 << 20;
const POOL_BYTES: usize = 64 << 20;
static POOL: Mutex<Vec<Pooled>> = Mutex::new(Vec::new());
struct Pooled(NonNull<u8>);
unsafe impl Send for Pooled {}
fn chunk_layout(size: usize) -> Layout {
Layout::from_size_align(size, ALIGN).expect("chunk layout")
}
const ALIGN: usize = align_of::<AtomicPtr<Node>>();
const LINK: usize = size_of::<AtomicPtr<Node>>();
#[repr(C)]
struct Node {
key_len: u32,
value_len: u32,
}
const HEADER: usize = size_of::<Node>();
fn node_size(height: usize, key_len: usize, value_len: usize) -> usize {
(height * LINK + HEADER + key_len + value_len).next_multiple_of(ALIGN)
}
unsafe fn link<'a>(node: *const Node, level: usize) -> &'a AtomicPtr<Node> {
&*(node as *const u8)
.sub((level + 1) * LINK)
.cast::<AtomicPtr<Node>>()
}
unsafe fn key_of<'a>(node: *const Node) -> &'a [u8] {
std::slice::from_raw_parts((node as *const u8).add(HEADER), (*node).key_len as usize)
}
unsafe fn value_of<'a>(node: *const Node) -> &'a [u8] {
let at = (node as *const u8).add(HEADER + (*node).key_len as usize);
std::slice::from_raw_parts(at, (*node).value_len as usize)
}
static NEXT_SHARD: AtomicUsize = AtomicUsize::new(0);
thread_local! {
static SHARD: usize = NEXT_SHARD.fetch_add(1, Relaxed) % SHARDS;
}
struct Free {
at: *mut u8,
left: usize,
}
unsafe impl Send for Free {}
#[repr(align(128))]
struct Shard {
block: Mutex<Free>,
used: AtomicUsize,
}
struct Arena {
shards: [Shard; SHARDS],
chunks: Mutex<Chunks>,
used: AtomicUsize,
}
struct Chunks {
owned: Vec<(NonNull<u8>, Layout)>,
free: Free,
next_size: usize,
}
impl Chunks {
fn take(&mut self, size: usize) -> (*mut u8, usize) {
debug_assert_eq!(size % ALIGN, 0);
let mut lost = 0;
if self.free.left < size {
lost = self.free.left;
let chunk = size.max(self.next_size);
self.next_size = (self.next_size * 2).min(MAX_CHUNK);
let layout = chunk_layout(chunk);
let pooled = if chunk == MAX_CHUNK {
POOL.lock().pop().map(|c| c.0)
} else {
None
};
let p = pooled.unwrap_or_else(|| {
NonNull::new(unsafe { alloc::alloc(layout) })
.unwrap_or_else(|| alloc::handle_alloc_error(layout))
});
self.owned.push((p, layout));
self.free = Free {
at: p.as_ptr(),
left: chunk,
};
}
let p = self.free.at;
self.free.at = unsafe { self.free.at.add(size) };
self.free.left -= size;
(p, lost)
}
}
impl Arena {
fn new() -> Self {
Arena {
shards: std::array::from_fn(|_| Shard {
block: Mutex::new(Free {
at: ptr::null_mut(),
left: 0,
}),
used: AtomicUsize::new(0),
}),
chunks: Mutex::new(Chunks {
owned: Vec::new(),
free: Free {
at: ptr::null_mut(),
left: 0,
},
next_size: FIRST_CHUNK,
}),
used: AtomicUsize::new(0),
}
}
fn take_from_chunks(&self, size: usize) -> *mut u8 {
let (p, lost) = self.chunks.lock().take(size);
self.used.fetch_add(size + lost, Relaxed);
p
}
fn alloc(&self, size: usize) -> *mut u8 {
if size > LARGE {
return self.take_from_chunks(size);
}
let shard = &self.shards[SHARD.with(|s| *s)];
let mut block = shard.block.lock();
let mut handed = size;
if block.left < size {
handed += block.left;
*block = Free {
at: self.take_from_chunks(BLOCK),
left: BLOCK,
};
self.used.fetch_sub(BLOCK, Relaxed);
}
let p = block.at;
block.at = unsafe { block.at.add(size) };
block.left -= size;
shard.used.fetch_add(handed, Relaxed);
p
}
fn usage(&self) -> usize {
self.used.load(Relaxed)
+ self
.shards
.iter()
.map(|s| s.used.load(Relaxed))
.sum::<usize>()
}
}
impl Drop for Arena {
fn drop(&mut self) {
let mut pool = POOL.lock();
for (p, layout) in self.chunks.get_mut().owned.drain(..) {
if layout.size() == MAX_CHUNK && pool.len() * MAX_CHUNK < POOL_BYTES {
pool.push(Pooled(p));
} else {
unsafe { alloc::dealloc(p.as_ptr(), layout) };
}
}
}
}
thread_local! {
static RNG: Cell<u64> = Cell::new({
let local = 0u8;
(&local as *const u8 as u64) | 1
});
}
fn random_height() -> usize {
RNG.with(|rng| {
let mut x = rng.get();
let mut height = 1;
loop {
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
if height == MAX_HEIGHT || x % BRANCHING != 0 {
break;
}
height += 1;
}
rng.set(x);
height
})
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Cursor(NonNull<Node>);
unsafe impl Send for Cursor {}
unsafe impl Sync for Cursor {}
pub struct SkipList {
arena: Arena,
head: NonNull<Node>,
height: AtomicUsize,
}
unsafe impl Send for SkipList {}
unsafe impl Sync for SkipList {}
impl Default for SkipList {
fn default() -> Self {
Self::new()
}
}
impl SkipList {
pub fn new() -> Self {
let arena = Arena::new();
let (head, _) = new_node(&arena, MAX_HEIGHT, &[], &[]);
SkipList {
arena,
head,
height: AtomicUsize::new(1),
}
}
pub fn memory_usage(&self) -> usize {
self.arena.usage() - node_size(MAX_HEIGHT, 0, 0)
}
pub fn insert(&self, key_parts: &[&[u8]], value: &[u8]) -> bool {
let height = random_height();
let node = new_node(&self.arena, height, key_parts, value).0.as_ptr();
let key = unsafe { key_of(node) };
let mut top = self.height.load(Relaxed);
if height > top {
top = self.height.fetch_max(height, Relaxed).max(height);
}
let mut prev = [self.head.as_ptr(); MAX_HEIGHT];
let mut next = [ptr::null_mut(); MAX_HEIGHT];
let mut x = self.head.as_ptr();
for level in (0..top).rev() {
let (p, n) = self.find_splice(key, x, level);
prev[level] = p;
next[level] = n;
x = p;
}
for level in 0..height {
loop {
if level == 0 && self.equals(next[0], key) {
return false;
}
unsafe {
link(node, level).store(next[level], Relaxed);
match link(prev[level], level).compare_exchange(
next[level],
node,
Release,
Acquire,
) {
Ok(_) => break,
Err(_) => {
let (p, n) = self.find_splice(key, prev[level], level);
prev[level] = p;
next[level] = n;
}
}
}
}
}
true
}
fn equals(&self, node: *mut Node, key: &[u8]) -> bool {
!node.is_null() && key::compare(unsafe { key_of(node) }, key) == Ordering::Equal
}
fn find_splice(&self, key: &[u8], start: *mut Node, level: usize) -> (*mut Node, *mut Node) {
let mut x = start;
loop {
let n = unsafe { link(x, level).load(Acquire) };
if !n.is_null() && key::compare(unsafe { key_of(n) }, key) == Ordering::Less {
x = n;
} else {
return (x, n);
}
}
}
fn search(&self, before: impl Fn(&[u8]) -> bool) -> (*mut Node, *mut Node) {
let mut x = self.head.as_ptr();
let mut level = self.height.load(Relaxed) - 1;
loop {
let n = unsafe { link(x, level).load(Acquire) };
if !n.is_null() && before(unsafe { key_of(n) }) {
x = n;
} else if level == 0 {
return (x, n);
} else {
level -= 1;
}
}
}
fn cursor(&self, node: *mut Node) -> Option<Cursor> {
if node == self.head.as_ptr() {
return None;
}
NonNull::new(node).map(Cursor)
}
pub fn is_empty(&self) -> bool {
self.first().is_none()
}
pub fn first(&self) -> Option<Cursor> {
self.cursor(unsafe { link(self.head.as_ptr(), 0).load(Acquire) })
}
pub fn last(&self) -> Option<Cursor> {
self.cursor(self.search(|_| true).0)
}
pub fn seek_ge(&self, target: &[u8]) -> Option<Cursor> {
self.cursor(self.search(|k| key::compare(k, target) == Ordering::Less).1)
}
pub fn seek_le(&self, target: &[u8]) -> Option<Cursor> {
self.cursor(
self.search(|k| key::compare(k, target) != Ordering::Greater)
.0,
)
}
pub fn seek_lt(&self, target: &[u8]) -> Option<Cursor> {
self.cursor(self.search(|k| key::compare(k, target) == Ordering::Less).0)
}
pub unsafe fn next(&self, at: Cursor) -> Option<Cursor> {
self.cursor(unsafe { link(at.0.as_ptr(), 0).load(Acquire) })
}
pub unsafe fn key(&self, at: Cursor) -> &[u8] {
unsafe { key_of(at.0.as_ptr()) }
}
pub unsafe fn value(&self, at: Cursor) -> &[u8] {
unsafe { value_of(at.0.as_ptr()) }
}
pub fn iter(&self) -> impl Iterator<Item = (&[u8], &[u8])> + '_ {
std::iter::successors(self.first(), |&c| unsafe { self.next(c) })
.map(|c| unsafe { (self.key(c), self.value(c)) })
}
}
fn new_node(
arena: &Arena,
height: usize,
key_parts: &[&[u8]],
value: &[u8],
) -> (NonNull<Node>, usize) {
let key_len: usize = key_parts.iter().map(|p| p.len()).sum();
let key_len32 = u32::try_from(key_len).expect("a key under 4 GiB");
let value_len32 = u32::try_from(value.len()).expect("a value under 4 GiB");
let size = node_size(height, key_len, value.len());
let p = arena.alloc(size);
unsafe {
let links = p.cast::<AtomicPtr<Node>>();
for level in 0..height {
links.add(level).write(AtomicPtr::new(ptr::null_mut()));
}
let node = p.add(height * LINK).cast::<Node>();
ptr::addr_of_mut!((*node).key_len).write(key_len32);
ptr::addr_of_mut!((*node).value_len).write(value_len32);
let mut at = node.cast::<u8>().add(HEADER);
for part in key_parts {
ptr::copy_nonoverlapping(part.as_ptr(), at, part.len());
at = at.add(part.len());
}
ptr::copy_nonoverlapping(value.as_ptr(), at, value.len());
(NonNull::new_unchecked(node), size)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::key::{encode, Kind};
use std::collections::BTreeMap;
use std::sync::Arc;
fn scaled(full: u32, miri: u32) -> u32 {
if cfg!(miri) {
miri
} else {
full
}
}
fn key(list: &SkipList, c: Cursor) -> Vec<u8> {
unsafe { list.key(c) }.to_vec()
}
fn value(list: &SkipList, c: Cursor) -> Vec<u8> {
unsafe { list.value(c) }.to_vec()
}
fn ik(i: u32) -> Vec<u8> {
encode(
0,
format!("k{i:06}").as_bytes(),
u64::from(i) + 1,
Kind::Put,
)
}
#[test]
fn entries_come_back_in_key_order() {
let list = SkipList::new();
assert!(list.is_empty());
let n = scaled(2000, 300);
let mut keys: Vec<u32> = (0..n).map(|i| (i * 7919) % n).collect();
for &i in &keys {
assert!(list.insert(&[&ik(i)], format!("v{i}").as_bytes()));
}
keys.sort();
let got: Vec<_> = list.iter().map(|(k, v)| (k.to_vec(), v.to_vec())).collect();
let want: Vec<_> = keys
.iter()
.map(|&i| (ik(i), format!("v{i}").into_bytes()))
.collect();
assert_eq!(got, want);
assert_eq!(key(&list, list.last().unwrap()), ik(n - 1));
}
#[test]
fn a_key_given_in_parts_is_one_key_and_a_duplicate_is_refused() {
let list = SkipList::new();
let whole = encode(3, b"user", 9, Kind::Delete);
let (a, b) = whole.split_at(5);
assert!(list.insert(&[a, b], b""));
let before = list.memory_usage();
assert!(!list.insert(&[&whole], b"other"));
assert!(
list.memory_usage() > before,
"the refused node still counts"
);
let c = list.first().unwrap();
assert_eq!(key(&list, c), whole);
assert_eq!(value(&list, c), b"");
assert_eq!(unsafe { list.next(c) }, None);
}
#[test]
fn seeks_agree_with_a_sorted_map() {
let list = SkipList::new();
let mut model = BTreeMap::new();
let n = scaled(3000, 300);
for i in (0..n).step_by(3) {
list.insert(&[&ik(i)], &[]);
model.insert(i, ());
}
let key_at = |c: Option<Cursor>| c.map(|c| key(&list, c));
for probe in 0..n + 5 {
let t = ik(probe);
let ge = model.range(probe..).next().map(|(&i, _)| ik(i));
let le = model.range(..=probe).next_back().map(|(&i, _)| ik(i));
let lt = model.range(..probe).next_back().map(|(&i, _)| ik(i));
assert_eq!(key_at(list.seek_ge(&t)), ge, "ge {probe}");
assert_eq!(key_at(list.seek_le(&t)), le, "le {probe}");
assert_eq!(key_at(list.seek_lt(&t)), lt, "lt {probe}");
}
}
#[test]
fn values_of_every_size_survive_including_large_ones() {
let list = SkipList::new();
for (i, len) in [
0usize,
1,
7,
8,
9,
LARGE - 1,
LARGE + 1,
BLOCK,
scaled(3 * MAX_CHUNK as u32, MAX_CHUNK as u32 + 8) as usize,
]
.into_iter()
.enumerate()
{
let run: Vec<u8> = (0..251u32).map(|b| (b * 31 + i as u32) as u8).collect();
let mut value = Vec::with_capacity(len);
while value.len() < len {
let take = run.len().min(len - value.len());
value.extend_from_slice(&run[..take]);
}
list.insert(&[&ik(i as u32)], &value);
let c = list.seek_ge(&ik(i as u32)).unwrap();
assert_eq!(self::value(&list, c), value, "len {len}");
}
}
#[test]
fn the_usage_is_what_the_nodes_take() {
let list = SkipList::new();
let empty = list.memory_usage();
assert_eq!(empty, 0, "the head is not counted");
list.insert(&[&ik(1)], b"12345");
let one = list.memory_usage() - empty;
assert!(one >= node_size(1, ik(1).len(), 5), "{one}");
assert!(one <= node_size(MAX_HEIGHT, ik(1).len(), 5), "{one}");
assert_eq!(one % ALIGN, 0);
let list = SkipList::new();
let mut least = 0;
for i in 0..scaled(50_000, 2_000) {
list.insert(&[&ik(i)], &[7; 100]);
least += node_size(1, ik(i).len(), 100);
}
let usage = list.memory_usage();
assert!(usage >= least, "{usage} < {least}");
assert!(usage < least * 13 / 10, "{usage} vs {least}");
}
#[test]
fn heights_follow_a_branching_of_four() {
let n = scaled(100_000, 20_000) as usize;
let total: usize = (0..n).map(|_| random_height()).sum();
let mean = total as f64 / n as f64;
assert!((1.28..1.39).contains(&mean), "mean height {mean}");
}
#[test]
fn concurrent_inserts_and_readers() {
let list = Arc::new(SkipList::new());
let threads = scaled(8, 4);
let per = scaled(5000, 150);
let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
let readers: Vec<_> = (0..2)
.map(|_| {
let (list, stop) = (list.clone(), stop.clone());
std::thread::spawn(move || {
while !stop.load(Relaxed) {
let mut last: Option<Vec<u8>> = None;
for (k, _) in list.iter() {
if let Some(l) = &last {
assert_eq!(key::compare(l, k), Ordering::Less);
}
last = Some(k.to_vec());
}
}
})
})
.collect();
let writers: Vec<_> = (0..threads)
.map(|t| {
let list = list.clone();
std::thread::spawn(move || {
for i in 0..per {
let k = i * threads + t;
assert!(list.insert(&[&ik(k)], &k.to_le_bytes()));
}
})
})
.collect();
for w in writers {
w.join().unwrap();
}
stop.store(true, Relaxed);
for r in readers {
r.join().unwrap();
}
let got: Vec<u32> = list
.iter()
.map(|(_, v)| u32::from_le_bytes(v.try_into().unwrap()))
.collect();
assert_eq!(got, (0..threads * per).collect::<Vec<_>>());
}
#[test]
fn a_seek_is_not_fooled_by_a_neighbour_inserted_meanwhile() {
use std::sync::atomic::AtomicU32;
let list = Arc::new(SkipList::new());
let n = scaled(200_000, 1_500);
for i in 0..=n {
list.insert(&[&ik(i * 4)], b"");
}
let progress: Arc<Vec<AtomicU32>> = Arc::new((0..3).map(|_| AtomicU32::new(0)).collect());
let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
let readers: Vec<_> = (0..3usize)
.map(|r| {
let (list, stop, progress) = (list.clone(), stop.clone(), progress.clone());
std::thread::spawn(move || {
let mut seeks = 0u64;
while !stop.load(Relaxed) {
let p = progress[(seeks as usize + r) % 3].load(Relaxed);
let target = ik((p + 1).min(n) * 4);
let found = list.seek_ge(&target).map(|c| key(&list, c));
assert_eq!(found.as_deref(), Some(&target[..]), "seek_ge");
let found = list.seek_le(&target).map(|c| key(&list, c));
assert_eq!(found.as_deref(), Some(&target[..]), "seek_le");
seeks += 1;
}
seeks
})
})
.collect();
let writers: Vec<_> = (0..3u32)
.map(|w| {
let (list, progress) = (list.clone(), progress.clone());
std::thread::spawn(move || {
for i in 0..n {
progress[w as usize].store(i, Relaxed);
list.insert(&[&ik(i * 4 + 1 + w)], b"");
}
})
})
.collect();
for w in writers {
w.join().unwrap();
}
stop.store(true, Relaxed);
let seeks: u64 = readers.into_iter().map(|r| r.join().unwrap()).sum();
if !cfg!(miri) {
assert!(seeks > 1000, "{seeks} seeks");
}
assert_eq!(list.iter().count(), 4 * n as usize + 1);
}
#[test]
fn chunks_go_back_to_a_bounded_pool_and_are_reused() {
for round in 0..scaled(6, 2) {
let list = SkipList::new();
let value = vec![round as u8; 1000];
for i in 0..scaled(30_000, 9_000) {
list.insert(&[&ik(i)], &value);
}
assert!(list.iter().all(|(_, v)| v == &value[..]), "round {round}");
drop(list);
assert!(POOL.lock().len() * MAX_CHUNK <= POOL_BYTES);
}
}
#[test]
fn interleaved_lists_keep_their_own_blocks() {
let (a, b) = (SkipList::new(), SkipList::new());
let n = scaled(5000, 500);
for i in 0..n {
a.insert(&[&ik(i)], b"a");
b.insert(&[&ik(i)], b"b");
}
drop(b);
let c = SkipList::new();
for i in 0..n {
c.insert(&[&ik(i)], b"c");
}
assert!(a.iter().all(|(_, v)| v == b"a"));
assert!(a.memory_usage() < 1 << 20, "{}", a.memory_usage());
assert!(c.iter().all(|(_, v)| v == b"c"));
assert_eq!(a.iter().count(), n as usize);
}
}