diff --git a/.gitignore b/.gitignore index 721a867..6fd4575 100644 --- a/.gitignore +++ b/.gitignore @@ -14,8 +14,5 @@ Cargo.lock *.pdb .lsm.data -.data -/old_* .test* -.block_index_test .bench diff --git a/Cargo.toml b/Cargo.toml index 9b99450..7396a05 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -101,3 +101,9 @@ name = "fd_table" harness = false path = "benches/fd_table.rs" required-features = [] + +[[bench]] +name = "partition_point" +harness = false +path = "benches/partition_point.rs" +required-features = [] diff --git a/README.md b/README.md index 150a32a..b136815 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@ A K.I.S.S. implementation of log-structured merge trees (LSM-trees/LSMTs) in Rus This is the most feature-rich LSM-tree implementation in Rust! It features: - Thread-safe BTreeMap-like API -- 100% safe & stable Rust +- [99.9% safe](./UNSAFE.md) & stable Rust - Block-based tables with compression support - Range & prefix searching with forward and reverse iteration - Size-tiered, (concurrent) Leveled and FIFO compaction diff --git a/UNSAFE.md b/UNSAFE.md new file mode 100644 index 0000000..5b7f603 --- /dev/null +++ b/UNSAFE.md @@ -0,0 +1,5 @@ +# Unsafe usage + +Currently, the project itself only uses one **1** unsafe block (ignoring dependencies which are tested themselves separately): + +- https://github.com/fjall-rs/lsm-tree/blob/2d8686e873369bd9c4ff2b562ed988c1cea38331/src/binary_search.rs#L23-L25 diff --git a/benches/partition_point.rs b/benches/partition_point.rs new file mode 100644 index 0000000..c528022 --- /dev/null +++ b/benches/partition_point.rs @@ -0,0 +1,23 @@ +use criterion::{criterion_group, criterion_main, Criterion}; +use lsm_tree::binary_search::partition_point; + +fn bench_partition_point(c: &mut Criterion) { + let mut group = c.benchmark_group("partition_point"); + + for item_count in [10, 100, 1_000, 10_000, 100_000, 1_000_000] { + let items = (0..item_count).collect::>(); + + // TODO: replace search key with random integer + + group.bench_function(format!("native {item_count}"), |b| { + b.iter(|| items.partition_point(|&x| x <= 5_000)) + }); + + group.bench_function(format!("rewrite {item_count}"), |b| { + b.iter(|| partition_point(&items, |&x| x <= 5_000)) + }); + } +} + +criterion_group!(benches, bench_partition_point); +criterion_main!(benches); diff --git a/fuzz/.gitignore b/fuzz/.gitignore index a4aa077..b400c27 100644 --- a/fuzz/.gitignore +++ b/fuzz/.gitignore @@ -1 +1,2 @@ corpus +artifacts diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml new file mode 100644 index 0000000..81c62b5 --- /dev/null +++ b/fuzz/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "lsm-tree-fuzz" +version = "0.0.0" +publish = false +edition = "2021" + +[package.metadata] +cargo-fuzz = true + +[dependencies] +libfuzzer-sys = "0.4" +lsm-tree = { path = ".." } + +[[bin]] +name = "partition_point" +path = "fuzz_targets/partition_point.rs" +test = false +doc = false +bench = false diff --git a/fuzz/fuzz_targets/partition_point.rs b/fuzz/fuzz_targets/partition_point.rs new file mode 100644 index 0000000..356e1c8 --- /dev/null +++ b/fuzz/fuzz_targets/partition_point.rs @@ -0,0 +1,19 @@ +#![no_main] +use libfuzzer_sys::{ + arbitrary::{Arbitrary, Unstructured}, + fuzz_target, +}; +use lsm_tree::binary_search::partition_point; + +fuzz_target!(|data: &[u8]| { + let mut unstructured = Unstructured::new(data); + + if let Ok(mut items) = as Arbitrary>::arbitrary(&mut unstructured) { + items.sort(); + items.dedup(); + + let idx = partition_point(&items, |&x| x < 128); + let std_pp_idx = items.partition_point(|&x| x < 128); + assert_eq!(std_pp_idx, idx); + } +}); diff --git a/src/binary_search.rs b/src/binary_search.rs new file mode 100644 index 0000000..2fd7329 --- /dev/null +++ b/src/binary_search.rs @@ -0,0 +1,91 @@ +// 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) + +/// Returns the index of the partition point according to the given predicate +/// (the index of the first element of the second partition). +/// +/// This seems to be faster than std's partition_point: https://github.com/rust-lang/rust/issues/138796 +pub fn partition_point(slice: &[T], pred: F) -> usize +where + F: Fn(&T) -> bool, +{ + let mut left = 0; + let mut right = slice.len(); + + if right == 0 { + return 0; + } + + while left < right { + let mid = (left + right) / 2; + + // SAFETY: See https://github.com/rust-lang/rust/blob/ebf0cf75d368c035f4c7e7246d203bd469ee4a51/library/core/src/slice/mod.rs#L2834-L2836 + #[warn(unsafe_code)] + let item = unsafe { slice.get_unchecked(mid) }; + + if pred(item) { + left = mid + 1; + } else { + right = mid; + } + } + + left +} + +#[cfg(test)] +mod tests { + use super::partition_point; + use test_log::test; + + #[test] + fn binary_search_first() { + let items = [1, 2, 3, 4, 5]; + let idx = partition_point(&items, |&x| x < 1); + assert_eq!(0, idx); + + let std_pp_idx = items.partition_point(|&x| x < 1); + assert_eq!(std_pp_idx, idx); + } + + #[test] + fn binary_search_last() { + let items = [1, 2, 3, 4, 5]; + let idx = partition_point(&items, |&x| x < 5); + assert_eq!(4, idx); + + let std_pp_idx = items.partition_point(|&x| x < 5); + assert_eq!(std_pp_idx, idx); + } + + #[test] + fn binary_search_middle() { + let items = [1, 2, 3, 4, 5]; + let idx = partition_point(&items, |&x| x < 3); + assert_eq!(2, idx); + + let std_pp_idx = items.partition_point(|&x| x < 3); + assert_eq!(std_pp_idx, idx); + } + + #[test] + fn binary_search_none() { + let items = [1, 2, 3, 4, 5]; + let idx = partition_point(&items, |&x| x < 10); + assert_eq!(5, idx); + + let std_pp_idx = items.partition_point(|&x| x < 10); + assert_eq!(std_pp_idx, idx); + } + + #[test] + fn binary_search_empty() { + let items: [i32; 0] = []; + let idx = partition_point(&items, |&x| x < 10); + assert_eq!(0, idx); + + let std_pp_idx = items.partition_point(|&x| x < 10); + assert_eq!(std_pp_idx, idx); + } +} diff --git a/src/level_manifest/level.rs b/src/level_manifest/level.rs index 32e854c..1648b11 100644 --- a/src/level_manifest/level.rs +++ b/src/level_manifest/level.rs @@ -2,7 +2,9 @@ // This source code is licensed under both the Apache 2.0 and MIT License // (found in the LICENSE-* files in the repository) -use crate::{segment::meta::SegmentId, HashSet, KeyRange, Segment, UserKey}; +use crate::{ + binary_search::partition_point, segment::meta::SegmentId, HashSet, KeyRange, Segment, UserKey, +}; use std::ops::Bound; /// Level of an LSM-tree @@ -175,13 +177,11 @@ pub struct DisjointLevel<'a>(&'a Level); impl<'a> DisjointLevel<'a> { /// Returns the segment that possibly contains the key. pub fn get_segment_containing_key(&self, key: &[u8]) -> Option { - let level = &self.0; - - let idx = level - .segments - .partition_point(|x| x.metadata.key_range.max() < &key); + let idx = partition_point(&self.0.segments, |segment| { + segment.metadata.key_range.max() < &key + }); - level + self.0 .segments .get(idx) .filter(|x| x.metadata.key_range.min() <= &key) @@ -197,12 +197,12 @@ impl<'a> DisjointLevel<'a> { let lo = match &key_range.0 { Bound::Unbounded => 0, - Bound::Included(start_key) => { - level.partition_point(|segment| segment.metadata.key_range.max() < start_key) - } - Bound::Excluded(start_key) => { - level.partition_point(|segment| segment.metadata.key_range.max() <= start_key) - } + Bound::Included(start_key) => partition_point(level, |segment| { + segment.metadata.key_range.max() < start_key + }), + Bound::Excluded(start_key) => partition_point(level, |segment| { + segment.metadata.key_range.max() <= start_key + }), }; if lo >= level.len() { @@ -213,7 +213,7 @@ impl<'a> DisjointLevel<'a> { Bound::Unbounded => level.len() - 1, Bound::Included(end_key) => { let idx = - level.partition_point(|segment| segment.metadata.key_range.min() <= end_key); + partition_point(level, |segment| segment.metadata.key_range.min() <= end_key); if idx == 0 { return None; @@ -223,7 +223,7 @@ impl<'a> DisjointLevel<'a> { } Bound::Excluded(end_key) => { let idx = - level.partition_point(|segment| segment.metadata.key_range.min() < end_key); + partition_point(level, |segment| segment.metadata.key_range.min() < end_key); if idx == 0 { return None; diff --git a/src/lib.rs b/src/lib.rs index f148aa6..aeea580 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -90,7 +90,7 @@ #![doc(html_logo_url = "https://raw.githubusercontent.com/fjall-rs/lsm-tree/main/logo.png")] #![doc(html_favicon_url = "https://raw.githubusercontent.com/fjall-rs/lsm-tree/main/logo.png")] -#![forbid(unsafe_code)] +#![deny(unsafe_code)] #![deny(clippy::all, missing_docs, clippy::cargo)] #![deny(clippy::unwrap_used)] #![deny(clippy::indexing_slicing)] @@ -124,6 +124,9 @@ mod any_tree; mod r#abstract; +#[doc(hidden)] +pub mod binary_search; + #[doc(hidden)] pub mod blob_tree; diff --git a/src/segment/block_index/mod.rs b/src/segment/block_index/mod.rs index 2ff8a04..97269d0 100644 --- a/src/segment/block_index/mod.rs +++ b/src/segment/block_index/mod.rs @@ -12,6 +12,7 @@ use super::{ block::{offset::BlockOffset, Block}, value_block::CachePolicy, }; +use crate::binary_search::partition_point; use block_handle::KeyedBlockHandle; use full_index::FullBlockIndex; use two_level_index::TwoLevelBlockIndex; @@ -44,7 +45,7 @@ impl KeyedBlockIndex for [KeyedBlockHandle] { key: &[u8], _: CachePolicy, ) -> crate::Result> { - let idx = self.partition_point(|x| &*x.end_key < key); + let idx = partition_point(self, |item| item.end_key < key); Ok(self.get(idx)) } @@ -53,7 +54,7 @@ impl KeyedBlockIndex for [KeyedBlockHandle] { key: &[u8], _: CachePolicy, ) -> crate::Result> { - let idx = self.partition_point(|x| &*x.end_key <= key); + let idx = partition_point(self, |x| &*x.end_key <= key); if idx == 0 { return Ok(self.first()); @@ -129,10 +130,10 @@ pub enum BlockIndexImpl { #[allow(clippy::expect_used)] mod tests { use super::*; - use crate::Slice; + use crate::{segment::block::offset::BlockOffset, UserKey}; use test_log::test; - fn bh>(end_key: K, offset: BlockOffset) -> KeyedBlockHandle { + fn bh>(end_key: K, offset: BlockOffset) -> KeyedBlockHandle { KeyedBlockHandle { end_key: end_key.into(), offset, diff --git a/src/segment/value_block.rs b/src/segment/value_block.rs index 7c04df6..7668a59 100644 --- a/src/segment/value_block.rs +++ b/src/segment/value_block.rs @@ -2,11 +2,11 @@ // This source code is licensed under both the Apache 2.0 and MIT License // (found in the LICENSE-* files in the repository) -use super::{ - block::{offset::BlockOffset, Block}, - id::GlobalSegmentId, +use super::{block::Block, id::GlobalSegmentId}; +use crate::{ + binary_search::partition_point, descriptor_table::FileDescriptorTable, + segment::block::offset::BlockOffset, value::InternalValue, Cache, }; -use crate::{cache::Cache, descriptor_table::FileDescriptorTable, value::InternalValue}; use std::sync::Arc; #[derive(Copy, Clone, Debug, PartialEq, Eq)] @@ -28,7 +28,7 @@ pub type ValueBlock = Block; impl ValueBlock { #[must_use] pub fn get_latest(&self, key: &[u8]) -> Option<&InternalValue> { - let idx = self.items.partition_point(|item| &*item.key.user_key < key); + let idx = partition_point(&self.items, |item| &*item.key.user_key < key); self.items .get(idx) diff --git a/src/segment/value_block_consumer.rs b/src/segment/value_block_consumer.rs index 3bb0506..f02fca0 100644 --- a/src/segment/value_block_consumer.rs +++ b/src/segment/value_block_consumer.rs @@ -3,7 +3,7 @@ // (found in the LICENSE-* files in the repository) use super::value_block::ValueBlock; -use crate::value::InternalValue; +use crate::{binary_search::partition_point, value::InternalValue}; use std::sync::Arc; pub struct ValueBlockConsumer { @@ -25,13 +25,13 @@ impl ValueBlockConsumer { end_key: Option<&[u8]>, ) -> Self { let mut lo = start_key.as_ref().map_or(0, |key| { - inner.items.partition_point(|x| &*x.key.user_key < *key) + partition_point(&inner.items, |x| &*x.key.user_key < *key) }); let hi = end_key.as_ref().map_or_else( || inner.items.len() - 1, |key| { - let idx = inner.items.partition_point(|x| &*x.key.user_key <= *key); + let idx = partition_point(&inner.items, |x| &*x.key.user_key <= *key); if idx == 0 { let first = inner