use std::io::{self};
use std::path::{Path, PathBuf};
#[cfg(test)]
use std::fs;
#[cfg(test)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock};
use std::ops::ControlFlow;
use super::block::{Block, BlockBuilder, BlockHandle, RESTART_INTERVAL, decode_entry_at};
use super::block_cache::BlockCache;
use super::bloom::{BloomFilterBuilder, encode_bloom_block};
use super::checksum;
use super::filter_block::FilterBlock;
use super::index_block::IndexBlock;
use super::internal_key::{
VALUE_TYPE_DELETION, VALUE_TYPE_MERGE, decode_internal_key, user_key_of,
};
use super::lookup_key::LookupKey;
use super::range_tombstone::{RangeTombstone, RangeTombstoneSet};
use crate::DbSlice;
use crate::env::{BufferedWriter, Env, ReadFile, WriteMode};
use crate::options::{CompressionType, PrefixExtractor};
const MAGIC_V5: u64 = 0x5245474F_53535405;
const MAGIC_V6: u64 = 0x5245474F_53535406;
const FOOTER_SIZE: usize = 72;
const META_CHECKSUM_LEN: usize = 4;
const COMPRESSION_NONE: u8 = 0x00;
const COMPRESSION_LZ4: u8 = 0x01;
const COMPRESSION_SNAPPY: u8 = 0x02;
#[derive(Debug)]
struct Footer {
range_tombstone_offset: u64,
range_tombstone_size: u64,
bloom_offset: u64,
bloom_size: u64,
index_offset: u64,
index_size: u64,
num_entries: u64,
magic: u64,
}
impl Footer {
fn magic_is_partitioned(magic: u64) -> bool {
matches!(magic, MAGIC_V6)
}
fn size_for_magic(magic: u64) -> io::Result<usize> {
match magic {
MAGIC_V5 | MAGIC_V6 => Ok(FOOTER_SIZE),
other => Err(invalid_data(format!(
"invalid SSTable magic number: {other:#018x}"
))),
}
}
fn size(&self) -> usize {
FOOTER_SIZE
}
fn encode(&self) -> Vec<u8> {
let mut buf = vec![0u8; self.size()];
buf[0..8].copy_from_slice(&self.range_tombstone_offset.to_le_bytes());
buf[8..16].copy_from_slice(&self.range_tombstone_size.to_le_bytes());
buf[16..24].copy_from_slice(&self.bloom_offset.to_le_bytes());
buf[24..32].copy_from_slice(&self.bloom_size.to_le_bytes());
buf[32..40].copy_from_slice(&self.index_offset.to_le_bytes());
buf[40..48].copy_from_slice(&self.index_size.to_le_bytes());
buf[48..56].copy_from_slice(&self.num_entries.to_le_bytes());
let sum = checksum::sst_footer(&buf[0..56], self.magic).to_le_bytes();
buf[56..64].copy_from_slice(&sum);
buf[64..72].copy_from_slice(&self.magic.to_le_bytes());
buf
}
fn decode(buf: &[u8]) -> io::Result<Self> {
if buf.len() < FOOTER_SIZE {
return Err(invalid_data("SSTable footer is truncated"));
}
let magic = u64::from_le_bytes(buf[buf.len() - 8..].try_into().unwrap());
let expected = Self::size_for_magic(magic)?;
if buf.len() != expected {
return Err(invalid_data(format!(
"SSTable footer is {} bytes but its format version needs {expected}",
buf.len()
)));
}
let stored = u64::from_le_bytes(buf[56..64].try_into().unwrap());
if stored != checksum::sst_footer(&buf[0..56], magic) {
return Err(invalid_data("SSTable footer checksum mismatch"));
}
Ok(Self {
range_tombstone_offset: u64::from_le_bytes(buf[0..8].try_into().unwrap()),
range_tombstone_size: u64::from_le_bytes(buf[8..16].try_into().unwrap()),
bloom_offset: u64::from_le_bytes(buf[16..24].try_into().unwrap()),
bloom_size: u64::from_le_bytes(buf[24..32].try_into().unwrap()),
index_offset: u64::from_le_bytes(buf[32..40].try_into().unwrap()),
index_size: u64::from_le_bytes(buf[40..48].try_into().unwrap()),
num_entries: u64::from_le_bytes(buf[48..56].try_into().unwrap()),
magic,
})
}
}
fn read_footer(file: &dyn ReadFile, file_size: u64) -> io::Result<Footer> {
if file_size < 8 {
return Err(invalid_data(
"SSTable file too small to carry a magic number",
));
}
let mut magic_buf = [0u8; 8];
file.read_exact_at(file_size - 8, &mut magic_buf)?;
let footer_size = Footer::size_for_magic(u64::from_le_bytes(magic_buf))?;
if file_size < footer_size as u64 {
return Err(invalid_data("SSTable file too small for its footer"));
}
let mut footer_buf = vec![0u8; footer_size];
file.read_exact_at(file_size - footer_size as u64, &mut footer_buf)?;
Footer::decode(&footer_buf)
}
fn verify_meta_region<'a>(
region: &'a [u8],
kind: u8,
checksummed: bool,
name: &'static str,
) -> io::Result<&'a [u8]> {
if !checksummed {
return Ok(region);
}
if region.len() < META_CHECKSUM_LEN {
return Err(invalid_data(format!(
"{name} is too short to carry a checksum"
)));
}
let (payload, trailer) = region.split_at(region.len() - META_CHECKSUM_LEN);
let stored = u32::from_le_bytes(trailer.try_into().unwrap());
if stored != checksum::sst_meta(kind, payload) {
return Err(invalid_data(format!("{name} checksum mismatch")));
}
Ok(payload)
}
fn encode_range_tombstone_block(tombstones: &[RangeTombstone]) -> Vec<u8> {
let mut buf = Vec::new();
buf.extend_from_slice(&(tombstones.len() as u32).to_le_bytes());
for rt in tombstones {
buf.extend_from_slice(&(rt.start.len() as u32).to_le_bytes());
buf.extend_from_slice(&rt.start);
buf.extend_from_slice(&(rt.end.len() as u32).to_le_bytes());
buf.extend_from_slice(&rt.end);
buf.extend_from_slice(&rt.seq.to_le_bytes());
}
buf
}
pub(crate) fn decode_range_tombstone_block(data: &[u8]) -> io::Result<Vec<RangeTombstone>> {
if data.is_empty() {
return Ok(Vec::new());
}
if data.len() < 4 {
return Err(invalid_data("range tombstone block too short"));
}
let count = u32::from_le_bytes(data[0..4].try_into().unwrap()) as usize;
let max_records_by_size = (data.len() - 4) / 16;
let mut out = Vec::with_capacity(count.min(max_records_by_size));
let mut pos = 4;
for _ in 0..count {
let start_len = read_u32(data, &mut pos, "range tombstone start length")? as usize;
let start = read_bytes(data, &mut pos, start_len, "range tombstone start")?;
let end_len = read_u32(data, &mut pos, "range tombstone end length")? as usize;
let end = read_bytes(data, &mut pos, end_len, "range tombstone end")?;
let seq = read_u64(data, &mut pos, "range tombstone seq")?;
out.push(RangeTombstone::new(start, end, seq));
}
if pos != data.len() {
return Err(invalid_data("range tombstone block has trailing bytes"));
}
Ok(out)
}
#[derive(Debug, Clone)]
pub(crate) struct SsTableMeta {
pub(crate) file_id: u64,
pub(crate) smallest_key: Vec<u8>,
pub(crate) largest_key: Vec<u8>,
pub(crate) file_size: u64,
pub(crate) num_entries: u64,
}
pub(crate) struct LiveSst {
pub(crate) meta: SsTableMeta,
pub(crate) reader: Arc<SsTableReader>,
}
impl LiveSst {
pub(crate) fn new(meta: SsTableMeta, reader: Arc<SsTableReader>) -> Arc<Self> {
Arc::new(Self { meta, reader })
}
}
#[derive(Clone)]
pub(crate) enum SsTableBlockCursor {
Flat(usize),
Partitioned {
leaf_idx: usize,
entry_idx: usize,
leaf: Arc<IndexBlock>,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum MetadataPolicy {
Pinned,
Cached,
}
enum MetaSlot<T> {
Pinned(Arc<T>),
Cached(BlockHandle),
}
enum MetaRef<'a, T> {
Borrowed(&'a T),
Owned(Arc<T>),
}
impl<T> std::ops::Deref for MetaRef<'_, T> {
type Target = T;
fn deref(&self) -> &T {
match self {
Self::Borrowed(value) => value,
Self::Owned(value) => value,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum LookupResult<V = DbSlice> {
NotInTable,
Found { seq: u64, value: V },
FoundTombstone { seq: u64 },
}
impl<V> LookupResult<V> {
pub(crate) fn map_value<W>(self, f: impl FnOnce(V) -> W) -> LookupResult<W> {
match self {
Self::NotInTable => LookupResult::NotInTable,
Self::FoundTombstone { seq } => LookupResult::FoundTombstone { seq },
Self::Found { seq, value } => LookupResult::Found {
seq,
value: f(value),
},
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Materialize {
Value,
LengthOnly,
}
pub(crate) enum PointValue {
Value(DbSlice),
Length(usize),
}
impl PointValue {
pub(crate) fn of(value: DbSlice, materialize: Materialize) -> Self {
match materialize {
Materialize::Value => Self::Value(value),
Materialize::LengthOnly => Self::Length(value.len()),
}
}
pub(crate) fn len(&self) -> usize {
match self {
Self::Value(value) => value.len(),
Self::Length(len) => *len,
}
}
}
enum BlockHit {
Tombstone { seq: u64 },
Value {
seq: u64,
value_offset: usize,
value_len: usize,
},
}
#[derive(Debug)]
pub(crate) struct SsTableWriteSummary {
pub(crate) smallest_user_key: Vec<u8>,
pub(crate) largest_user_key: Vec<u8>,
pub(crate) num_entries: u64,
}
fn encode_index_block(entries: &[(Vec<u8>, BlockHandle)]) -> Vec<u8> {
let mut data = Vec::new();
data.extend_from_slice(&(entries.len() as u32).to_le_bytes());
for (key, handle) in entries {
data.extend_from_slice(&(key.len() as u32).to_le_bytes());
data.extend_from_slice(key);
data.extend_from_slice(&handle.offset.to_le_bytes());
data.extend_from_slice(&handle.size.to_le_bytes());
}
data
}
fn encoded_index_block_size(entries: &[(Vec<u8>, BlockHandle)]) -> usize {
4 + entries.iter().map(|(k, _)| 4 + k.len() + 16).sum::<usize>()
}
fn invalid_data(message: impl Into<String>) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, message.into())
}
fn read_u32(data: &[u8], pos: &mut usize, field: &'static str) -> io::Result<u32> {
let end = pos
.checked_add(4)
.ok_or_else(|| invalid_data(format!("{field} offset overflows")))?;
if end > data.len() {
return Err(invalid_data(format!("{field} is truncated")));
}
let value = u32::from_le_bytes(data[*pos..end].try_into().unwrap());
*pos = end;
Ok(value)
}
fn read_u64(data: &[u8], pos: &mut usize, field: &'static str) -> io::Result<u64> {
let end = pos
.checked_add(8)
.ok_or_else(|| invalid_data(format!("{field} offset overflows")))?;
if end > data.len() {
return Err(invalid_data(format!("{field} is truncated")));
}
let value = u64::from_le_bytes(data[*pos..end].try_into().unwrap());
*pos = end;
Ok(value)
}
fn read_bytes(
data: &[u8],
pos: &mut usize,
len: usize,
field: &'static str,
) -> io::Result<Vec<u8>> {
let end = pos
.checked_add(len)
.ok_or_else(|| invalid_data(format!("{field} length overflows")))?;
if end > data.len() {
return Err(invalid_data(format!("{field} is truncated")));
}
let bytes = data[*pos..end].to_vec();
*pos = end;
Ok(bytes)
}
fn validate_footer_regions(footer: &Footer, data_end: u64) -> io::Result<()> {
let trailer = META_CHECKSUM_LEN as u64;
if footer.bloom_size < 8 + trailer {
return Err(invalid_data("bloom region too short"));
}
if footer.index_size < 4 + trailer {
return Err(invalid_data("index block too short"));
}
validate_file_region(
footer.bloom_offset,
footer.bloom_size,
data_end,
"bloom region",
)?;
validate_file_region(
footer.index_offset,
footer.index_size,
data_end,
"index block",
)?;
if footer.range_tombstone_size > 0 {
if footer.range_tombstone_size < 4 + trailer {
return Err(invalid_data("range tombstone block too short"));
}
validate_file_region(
footer.range_tombstone_offset,
footer.range_tombstone_size,
data_end,
"range tombstone block",
)?;
}
Ok(())
}
fn write_meta_region(
writer: &mut BufferedWriter,
offset: &mut u64,
payload: &[u8],
kind: u8,
) -> io::Result<u64> {
let trailer = checksum::sst_meta(kind, payload).to_le_bytes();
writer.write_all(payload)?;
writer.write_all(&trailer)?;
let size = (payload.len() + trailer.len()) as u64;
*offset += size;
Ok(size)
}
fn read_file_region(
file: &dyn ReadFile,
offset: u64,
size: u64,
data_end: u64,
name: &'static str,
) -> io::Result<Vec<u8>> {
validate_file_region(offset, size, data_end, name)?;
let len = usize::try_from(size)
.map_err(|_| invalid_data(format!("{name} is too large to address")))?;
let mut buf = vec![0u8; len];
file.read_exact_at(offset, &mut buf)?;
Ok(buf)
}
fn validate_file_region(
offset: u64,
size: u64,
data_end: u64,
name: &'static str,
) -> io::Result<()> {
let end = offset
.checked_add(size)
.ok_or_else(|| invalid_data(format!("{name} range overflows")))?;
if end > data_end {
return Err(invalid_data(format!("{name} extends past SSTable data")));
}
if usize::try_from(size).is_err() {
return Err(invalid_data(format!("{name} is too large to address")));
}
Ok(())
}
pub(crate) struct SsTableWriter {
writer: BufferedWriter,
path: PathBuf,
env: Arc<dyn Env>,
block_builder: BlockBuilder,
index_entries: Vec<(Vec<u8>, BlockHandle)>,
bloom_builder: BloomFilterBuilder,
prefix_bloom_builder: Option<BloomFilterBuilder>,
prefix_extractor: Option<Arc<dyn PrefixExtractor>>,
last_prefix: Option<Vec<u8>>,
range_tombstones: Vec<RangeTombstone>,
block_size: usize,
current_offset: u64,
num_entries: u64,
last_internal_key: Vec<u8>,
smallest_user_key: Option<Vec<u8>>,
largest_user_key: Option<Vec<u8>>,
last_bloom_user_key: Vec<u8>,
compression: CompressionType,
partitioned_index: bool,
metadata_block_size: usize,
}
impl SsTableWriter {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new_in(
env: &Arc<dyn Env>,
path: &Path,
block_size: usize,
bloom_bits_per_key: usize,
compression: CompressionType,
prefix_extractor: Option<Arc<dyn PrefixExtractor>>,
partitioned_index: bool,
metadata_block_size: usize,
) -> io::Result<Self> {
let file = env.open_write(path, WriteMode::Truncate)?;
let prefix_bloom_builder = prefix_extractor
.as_ref()
.map(|_| BloomFilterBuilder::new(bloom_bits_per_key));
Ok(Self {
writer: BufferedWriter::new(file),
path: path.to_path_buf(),
env: Arc::clone(env),
block_builder: BlockBuilder::new(RESTART_INTERVAL),
index_entries: Vec::new(),
bloom_builder: BloomFilterBuilder::new(bloom_bits_per_key),
prefix_bloom_builder,
prefix_extractor,
last_prefix: None,
range_tombstones: Vec::new(),
block_size,
current_offset: 0,
num_entries: 0,
last_internal_key: Vec::new(),
smallest_user_key: None,
largest_user_key: None,
last_bloom_user_key: Vec::new(),
compression,
partitioned_index,
metadata_block_size,
})
}
#[cfg(test)]
pub(crate) fn new(
path: &Path,
block_size: usize,
bloom_bits_per_key: usize,
compression: CompressionType,
prefix_extractor: Option<Arc<dyn PrefixExtractor>>,
partitioned_index: bool,
metadata_block_size: usize,
) -> io::Result<Self> {
Self::new_in(
&crate::env::std_env(),
path,
block_size,
bloom_bits_per_key,
compression,
prefix_extractor,
partitioned_index,
metadata_block_size,
)
}
pub(crate) fn add_range_tombstone(&mut self, start: &[u8], end: &[u8], seq: u64) {
self.range_tombstones
.push(RangeTombstone::new(start.to_vec(), end.to_vec(), seq));
}
pub(crate) fn add(&mut self, internal_key: &[u8], value: &[u8]) -> io::Result<()> {
let user_key = user_key_of(internal_key);
if user_key != self.last_bloom_user_key.as_slice() {
self.bloom_builder.add_key(user_key);
self.last_bloom_user_key = user_key.to_vec();
if let (Some(extractor), Some(builder)) = (
self.prefix_extractor.as_ref(),
self.prefix_bloom_builder.as_mut(),
) && let Some(prefix) = extractor.extract(user_key)
{
let changed = match &self.last_prefix {
Some(p) => p.as_slice() != prefix,
None => true,
};
if changed {
builder.add_key(prefix);
self.last_prefix = Some(prefix.to_vec());
}
}
}
self.block_builder.add(internal_key, value);
self.last_internal_key = internal_key.to_vec();
self.num_entries += 1;
if self.smallest_user_key.is_none() {
self.smallest_user_key = Some(user_key.to_vec());
}
match &mut self.largest_user_key {
Some(existing) if existing.as_slice() == user_key => {}
slot => *slot = Some(user_key.to_vec()),
}
if self.block_builder.estimated_size() >= self.block_size {
self.flush_block()?;
}
Ok(())
}
pub(crate) fn finish(mut self) -> io::Result<Option<SsTableWriteSummary>> {
if self.num_entries == 0 && self.range_tombstones.is_empty() {
return Ok(None);
}
if !self.block_builder.is_empty() {
self.flush_block()?;
}
let range_tombstone_set =
RangeTombstoneSet::from_vec(std::mem::take(&mut self.range_tombstones));
let range_tombstone_data = if range_tombstone_set.is_empty() {
Vec::new()
} else {
encode_range_tombstone_block(range_tombstone_set.as_slice())
};
let range_tombstone_offset = self.current_offset;
let range_tombstone_size = if range_tombstone_data.is_empty() {
0
} else {
write_meta_region(
&mut self.writer,
&mut self.current_offset,
&range_tombstone_data,
checksum::META_KIND_RANGE_TOMBSTONE,
)?
};
let prefix_bloom_data = match self.prefix_bloom_builder.take() {
Some(builder) => encode_bloom_block(&builder.build()),
None => Vec::new(),
};
let user_bloom = self.bloom_builder.build();
let user_bloom_data = encode_bloom_block(&user_bloom);
let mut bloom_region =
Vec::with_capacity(8 + prefix_bloom_data.len() + user_bloom_data.len());
bloom_region.extend_from_slice(&(prefix_bloom_data.len() as u64).to_le_bytes());
bloom_region.extend_from_slice(&prefix_bloom_data);
bloom_region.extend_from_slice(&user_bloom_data);
let bloom_offset = self.current_offset;
let bloom_size = write_meta_region(
&mut self.writer,
&mut self.current_offset,
&bloom_region,
checksum::META_KIND_BLOOM,
)?;
let (index_offset, index_size, magic) = if self.partitioned_index {
let mut top_level: Vec<(Vec<u8>, BlockHandle)> = Vec::new();
let mut chunk_start = 0usize;
while chunk_start < self.index_entries.len() {
let mut chunk_end = chunk_start + 1;
while chunk_end < self.index_entries.len() {
let candidate = &self.index_entries[chunk_start..chunk_end + 1];
let size = encoded_index_block_size(candidate);
if size > self.metadata_block_size {
break;
}
chunk_end += 1;
}
let chunk = &self.index_entries[chunk_start..chunk_end];
let leaf_data = encode_index_block(chunk);
let last_key = match chunk.last() {
Some((key, _)) => key.clone(),
None => return Err(invalid_data("index partitioning produced an empty leaf")),
};
let leaf_offset = self.current_offset;
let leaf_size = write_meta_region(
&mut self.writer,
&mut self.current_offset,
&leaf_data,
checksum::META_KIND_INDEX_LEAF,
)?;
top_level.push((
last_key,
BlockHandle {
offset: leaf_offset,
size: leaf_size,
},
));
chunk_start = chunk_end;
}
let top_data = encode_index_block(&top_level);
let top_offset = self.current_offset;
let top_size = write_meta_region(
&mut self.writer,
&mut self.current_offset,
&top_data,
checksum::META_KIND_INDEX,
)?;
(top_offset, top_size, MAGIC_V6)
} else {
let index_data = encode_index_block(&self.index_entries);
let idx_offset = self.current_offset;
let idx_size = write_meta_region(
&mut self.writer,
&mut self.current_offset,
&index_data,
checksum::META_KIND_INDEX,
)?;
(idx_offset, idx_size, MAGIC_V5)
};
let footer = Footer {
range_tombstone_offset,
range_tombstone_size,
bloom_offset,
bloom_size,
index_offset,
index_size,
num_entries: self.num_entries,
magic,
};
self.writer.write_all(&footer.encode())?;
self.writer.sync_all()?;
crate::env::sync_parent_dir(&*self.env, &self.path)?;
let mut smallest_user_key = self.smallest_user_key.take();
let mut largest_user_key = self.largest_user_key.take();
for rt in range_tombstone_set.iter() {
if smallest_user_key
.as_ref()
.is_none_or(|smallest| rt.start.as_slice() < smallest.as_slice())
{
smallest_user_key = Some(rt.start.clone());
}
if largest_user_key
.as_ref()
.is_none_or(|largest| rt.end.as_slice() > largest.as_slice())
{
largest_user_key = Some(rt.end.clone());
}
}
let smallest_user_key = smallest_user_key.expect("checked non-empty above");
let largest_user_key = largest_user_key.expect("checked non-empty above");
Ok(Some(SsTableWriteSummary {
smallest_user_key,
largest_user_key,
num_entries: self.num_entries,
}))
}
fn flush_block(&mut self) -> io::Result<()> {
let last_internal = self.last_internal_key.clone();
let block_builder =
std::mem::replace(&mut self.block_builder, BlockBuilder::new(RESTART_INTERVAL));
let raw_data = block_builder.finish();
let block_offset = self.current_offset;
let (codec_byte, payload) = match self.compression {
CompressionType::None => (COMPRESSION_NONE, raw_data.clone()),
CompressionType::Lz4 => (COMPRESSION_LZ4, lz4_flex::compress_prepend_size(&raw_data)),
CompressionType::Snappy => {
let mut encoder = snap::raw::Encoder::new();
let compressed = encoder.compress_vec(&raw_data).map_err(|e| {
io::Error::new(io::ErrorKind::InvalidData, format!("snappy encode: {e}"))
})?;
let mut framed = Vec::with_capacity(4 + compressed.len());
framed.extend_from_slice(&(raw_data.len() as u32).to_le_bytes());
framed.extend_from_slice(&compressed);
(COMPRESSION_SNAPPY, framed)
}
};
let checksum = checksum::sst_block(codec_byte, &payload);
self.writer.write_all(&[codec_byte])?;
self.writer.write_all(&payload)?;
self.writer.write_all(&checksum.to_le_bytes())?;
let total = 1 + payload.len() + 4;
self.current_offset += total as u64;
let block_size = self.current_offset - block_offset;
self.index_entries.push((
last_internal,
BlockHandle {
offset: block_offset,
size: block_size,
},
));
Ok(())
}
}
pub(crate) struct SsTableReader {
file: Box<dyn ReadFile>,
data_end: u64,
pub(crate) file_id: u64,
index: MetaSlot<IndexBlock>,
filter: MetaSlot<FilterBlock>,
index_fallback: OnceLock<Arc<IndexBlock>>,
filter_fallback: OnceLock<Arc<FilterBlock>>,
range_tombstones: RangeTombstoneSet,
partitioned: bool,
#[cfg(test)]
index_leaf_reads: AtomicUsize,
}
pub(crate) struct SsTableInternalIter<'a> {
reader: &'a SsTableReader,
cache: &'a BlockCache,
next_cursor: Option<SsTableBlockCursor>,
current_block: Option<Arc<Block>>,
current_pos: usize,
current_key: Vec<u8>,
}
impl<'a> SsTableInternalIter<'a> {
pub(crate) fn next_entry(&mut self) -> io::Result<Option<(Vec<u8>, Vec<u8>)>> {
loop {
if let Some(block) = &self.current_block {
let data = block.entry_data();
if self.current_pos < data.len() {
let (consumed, value_offset, value_len) =
decode_entry_at(data, self.current_pos, &mut self.current_key);
self.current_pos += consumed;
let value = data[value_offset..value_offset + value_len].to_vec();
return Ok(Some((self.current_key.clone(), value)));
}
}
let Some(cursor) = self.next_cursor.take() else {
return Ok(None);
};
self.next_cursor = self.reader.next_block_cursor(&cursor, self.cache)?;
let block = self.reader.load_block_at_cursor(&cursor, self.cache)?;
self.current_block = Some(block);
self.current_pos = 0;
self.current_key.clear();
}
}
}
impl SsTableReader {
#[cfg(test)]
pub(crate) fn open(path: &Path, file_id: u64) -> io::Result<Self> {
Self::open_with(
&crate::env::std_env(),
path,
file_id,
MetadataPolicy::Pinned,
)
}
pub(crate) fn open_with(
env: &Arc<dyn Env>,
path: &Path,
file_id: u64,
policy: MetadataPolicy,
) -> io::Result<Self> {
let file = env.open_read(path)?;
let file_size = file.len()?;
let footer = read_footer(&*file, file_size)?;
let data_end = file_size - footer.size() as u64;
validate_footer_regions(&footer, data_end)?;
let filter_handle = BlockHandle {
offset: footer.bloom_offset,
size: footer.bloom_size,
};
let filter_region = read_file_region(
&*file,
filter_handle.offset,
filter_handle.size,
data_end,
"bloom region",
)?;
let filter_block = FilterBlock::decode(verify_meta_region(
&filter_region,
checksum::META_KIND_BLOOM,
true,
"bloom region",
)?)?;
let filter = match policy {
MetadataPolicy::Pinned => MetaSlot::Pinned(Arc::new(filter_block)),
MetadataPolicy::Cached => MetaSlot::Cached(filter_handle),
};
drop(filter_region);
let index_handle = BlockHandle {
offset: footer.index_offset,
size: footer.index_size,
};
let mut index_region = read_file_region(
&*file,
index_handle.offset,
index_handle.size,
data_end,
"index block",
)?;
let index_payload_len = verify_meta_region(
&index_region,
checksum::META_KIND_INDEX,
true,
"index block",
)?
.len();
index_region.truncate(index_payload_len);
let index_block = IndexBlock::decode(index_region)?;
let range_tombstones = if footer.range_tombstone_size == 0 {
Vec::new()
} else {
let rt_region = read_file_region(
&*file,
footer.range_tombstone_offset,
footer.range_tombstone_size,
data_end,
"range tombstone block",
)?;
decode_range_tombstone_block(verify_meta_region(
&rt_region,
checksum::META_KIND_RANGE_TOMBSTONE,
true,
"range tombstone block",
)?)?
};
let partitioned = Footer::magic_is_partitioned(footer.magic);
let index = if partitioned || policy == MetadataPolicy::Pinned {
MetaSlot::Pinned(Arc::new(index_block))
} else {
MetaSlot::Cached(index_handle)
};
Ok(Self {
file,
data_end,
file_id,
index,
filter,
index_fallback: OnceLock::new(),
filter_fallback: OnceLock::new(),
range_tombstones: RangeTombstoneSet::from_vec(range_tombstones),
partitioned,
#[cfg(test)]
index_leaf_reads: AtomicUsize::new(0),
})
}
pub(crate) fn pinned_metadata_bytes(&self) -> usize {
let index = match &self.index {
MetaSlot::Pinned(block) => block.charge(),
MetaSlot::Cached(_) => 0,
};
let filter = match &self.filter {
MetaSlot::Pinned(block) => block.charge(),
MetaSlot::Cached(_) => 0,
};
let index_refused = self.index_fallback.get().map_or(0, |block| block.charge());
let filter_refused = self.filter_fallback.get().map_or(0, |block| block.charge());
let tombstones: usize = self
.range_tombstones
.as_slice()
.iter()
.map(|rt| {
std::mem::size_of::<RangeTombstone>() + rt.start.capacity() + rt.end.capacity()
})
.sum();
index + filter + index_refused + filter_refused + tombstones
}
fn index(&self, cache: &BlockCache) -> io::Result<MetaRef<'_, IndexBlock>> {
let handle = match &self.index {
MetaSlot::Pinned(block) => return Ok(MetaRef::Borrowed(block)),
MetaSlot::Cached(handle) => *handle,
};
if let Some(block) = self.index_fallback.get() {
return Ok(MetaRef::Borrowed(block));
}
if let Some(block) = cache.get_index(self.file_id, handle.offset) {
return Ok(MetaRef::Owned(block));
}
let data = self.read_metadata_region(handle, checksum::META_KIND_INDEX, "index block")?;
let block = Arc::new(IndexBlock::decode(data)?);
if !cache.insert_index(self.file_id, handle.offset, Arc::clone(&block)) {
let _ = self.index_fallback.set(Arc::clone(&block));
}
Ok(MetaRef::Owned(block))
}
fn filter(&self, cache: &BlockCache) -> io::Result<MetaRef<'_, FilterBlock>> {
let handle = match &self.filter {
MetaSlot::Pinned(block) => return Ok(MetaRef::Borrowed(block)),
MetaSlot::Cached(handle) => *handle,
};
if let Some(block) = self.filter_fallback.get() {
return Ok(MetaRef::Borrowed(block));
}
if let Some(block) = cache.get_filter(self.file_id, handle.offset) {
return Ok(MetaRef::Owned(block));
}
let region =
self.read_metadata_region(handle, checksum::META_KIND_BLOOM, "bloom region")?;
let block = Arc::new(FilterBlock::decode(®ion)?);
if !cache.insert_filter(self.file_id, handle.offset, Arc::clone(&block)) {
let _ = self.filter_fallback.set(Arc::clone(&block));
}
Ok(MetaRef::Owned(block))
}
fn read_metadata_region(
&self,
handle: BlockHandle,
kind: u8,
name: &'static str,
) -> io::Result<Vec<u8>> {
let mut region =
{ read_file_region(&*self.file, handle.offset, handle.size, self.data_end, name)? };
let payload_len = verify_meta_region(®ion, kind, true, name)?.len();
region.truncate(payload_len);
Ok(region)
}
fn read_index_leaf(
&self,
handle: BlockHandle,
cache: &BlockCache,
) -> io::Result<Arc<IndexBlock>> {
if let Some(leaf) = cache.get_index(self.file_id, handle.offset) {
return Ok(leaf);
}
#[cfg(test)]
self.index_leaf_reads.fetch_add(1, Ordering::Relaxed);
let buf = self.read_metadata_region(
handle,
checksum::META_KIND_INDEX_LEAF,
"partitioned index leaf",
)?;
let leaf = Arc::new(IndexBlock::decode(buf)?);
cache.insert_index(self.file_id, handle.offset, Arc::clone(&leaf));
Ok(leaf)
}
#[cfg(test)]
fn index_leaf_read_count(&self) -> usize {
self.index_leaf_reads.load(Ordering::Relaxed)
}
#[cfg(test)]
fn reset_index_leaf_read_count(&self) {
self.index_leaf_reads.store(0, Ordering::Relaxed);
}
#[cfg(test)]
fn pinned_index_len(&self) -> Option<usize> {
match &self.index {
MetaSlot::Pinned(block) => Some(block.len()),
MetaSlot::Cached(_) => None,
}
}
fn cursor_from_leaf(
&self,
leaf_idx: usize,
entry_idx: usize,
leaf: Arc<IndexBlock>,
) -> Option<SsTableBlockCursor> {
if entry_idx >= leaf.len() {
return None;
}
Some(SsTableBlockCursor::Partitioned {
leaf_idx,
entry_idx,
leaf,
})
}
pub(crate) fn first_block_cursor(
&self,
cache: &BlockCache,
) -> io::Result<Option<SsTableBlockCursor>> {
let index = self.index(cache)?;
if !self.partitioned {
return Ok((!index.is_empty()).then_some(SsTableBlockCursor::Flat(0)));
}
for leaf_idx in 0..index.len() {
let Some(handle) = index.handle(leaf_idx) else {
continue;
};
let leaf = self.read_index_leaf(handle, cache)?;
if let Some(cursor) = self.cursor_from_leaf(leaf_idx, 0, leaf) {
return Ok(Some(cursor));
}
}
Ok(None)
}
pub(crate) fn last_block_cursor(
&self,
cache: &BlockCache,
) -> io::Result<Option<SsTableBlockCursor>> {
let index = self.index(cache)?;
if !self.partitioned {
return Ok(index.len().checked_sub(1).map(SsTableBlockCursor::Flat));
}
for leaf_idx in (0..index.len()).rev() {
let Some(handle) = index.handle(leaf_idx) else {
continue;
};
let leaf = self.read_index_leaf(handle, cache)?;
if !leaf.is_empty() {
let entry_idx = leaf.len() - 1;
return Ok(self.cursor_from_leaf(leaf_idx, entry_idx, leaf));
}
}
Ok(None)
}
pub(crate) fn seek_block_cursor(
&self,
target: &[u8],
cache: &BlockCache,
) -> io::Result<Option<SsTableBlockCursor>> {
let index = self.index(cache)?;
let mut leaf_idx = index.seek(target);
if leaf_idx >= index.len() {
return Ok(None);
}
if !self.partitioned {
return Ok(Some(SsTableBlockCursor::Flat(leaf_idx)));
}
while leaf_idx < index.len() {
let Some(handle) = index.handle(leaf_idx) else {
break;
};
let leaf = self.read_index_leaf(handle, cache)?;
let entry_idx = leaf.seek(target);
if let Some(cursor) = self.cursor_from_leaf(leaf_idx, entry_idx, leaf) {
return Ok(Some(cursor));
}
leaf_idx += 1;
}
Ok(None)
}
pub(crate) fn next_block_cursor(
&self,
cursor: &SsTableBlockCursor,
cache: &BlockCache,
) -> io::Result<Option<SsTableBlockCursor>> {
match cursor {
SsTableBlockCursor::Flat(idx) => {
let next = idx + 1;
let index = self.index(cache)?;
Ok((next < index.len()).then_some(SsTableBlockCursor::Flat(next)))
}
SsTableBlockCursor::Partitioned {
leaf_idx,
entry_idx,
leaf,
} => {
let next_entry = entry_idx + 1;
if next_entry < leaf.len() {
return Ok(Some(SsTableBlockCursor::Partitioned {
leaf_idx: *leaf_idx,
entry_idx: next_entry,
leaf: Arc::clone(leaf),
}));
}
let index = self.index(cache)?;
for next_leaf_idx in leaf_idx + 1..index.len() {
let Some(handle) = index.handle(next_leaf_idx) else {
continue;
};
let next_leaf = self.read_index_leaf(handle, cache)?;
if let Some(cursor) = self.cursor_from_leaf(next_leaf_idx, 0, next_leaf) {
return Ok(Some(cursor));
}
}
Ok(None)
}
}
}
pub(crate) fn prev_block_cursor(
&self,
cursor: &SsTableBlockCursor,
cache: &BlockCache,
) -> io::Result<Option<SsTableBlockCursor>> {
match cursor {
SsTableBlockCursor::Flat(idx) => {
if *idx == 0 {
Ok(None)
} else {
Ok(Some(SsTableBlockCursor::Flat(idx - 1)))
}
}
SsTableBlockCursor::Partitioned {
leaf_idx,
entry_idx,
leaf,
} => {
if *entry_idx > 0 {
return Ok(Some(SsTableBlockCursor::Partitioned {
leaf_idx: *leaf_idx,
entry_idx: entry_idx - 1,
leaf: Arc::clone(leaf),
}));
}
let index = self.index(cache)?;
for prev_leaf_idx in (0..*leaf_idx).rev() {
let Some(handle) = index.handle(prev_leaf_idx) else {
continue;
};
let prev_leaf = self.read_index_leaf(handle, cache)?;
if !prev_leaf.is_empty() {
let entry_idx = prev_leaf.len() - 1;
return Ok(self.cursor_from_leaf(prev_leaf_idx, entry_idx, prev_leaf));
}
}
Ok(None)
}
}
}
pub(crate) fn load_block_at_cursor(
&self,
cursor: &SsTableBlockCursor,
cache: &BlockCache,
) -> io::Result<Arc<Block>> {
let handle = self.cursor_handle(cursor, cache)?;
self.read_block(handle, cache)
}
pub(crate) fn cursor_handle(
&self,
cursor: &SsTableBlockCursor,
cache: &BlockCache,
) -> io::Result<BlockHandle> {
match cursor {
SsTableBlockCursor::Flat(idx) => self
.index(cache)?
.handle(*idx)
.ok_or_else(|| invalid_data("block index out of bounds")),
SsTableBlockCursor::Partitioned {
entry_idx, leaf, ..
} => leaf
.handle(*entry_idx)
.ok_or_else(|| invalid_data("partitioned block index out of bounds")),
}
}
fn find_block_handle(
&self,
search_key: &[u8],
cache: &BlockCache,
) -> io::Result<Option<BlockHandle>> {
self.seek_block_cursor(search_key, cache)?
.map(|cursor| self.cursor_handle(&cursor, cache))
.transpose()
}
pub(crate) fn may_have_prefix(&self, prefix: &[u8], cache: &BlockCache) -> io::Result<bool> {
Ok(self.filter(cache)?.may_have_prefix(prefix))
}
pub(crate) fn covering_range_tombstone_seq(&self, user_key: &[u8], snapshot_seq: u64) -> u64 {
if self.range_tombstones.is_empty() {
return 0;
}
self.range_tombstones
.max_covering_seq(user_key, snapshot_seq)
}
pub(crate) fn range_tombstones(&self) -> &[RangeTombstone] {
self.range_tombstones.as_slice()
}
pub(crate) fn get(
&self,
lk: &LookupKey,
key_buf: &mut Vec<u8>,
cache: &BlockCache,
) -> io::Result<LookupResult<DbSlice>> {
Ok(match self.probe(lk, key_buf, cache)? {
None => LookupResult::NotInTable,
Some((_, BlockHit::Tombstone { seq })) => LookupResult::FoundTombstone { seq },
Some((
block,
BlockHit::Value {
seq,
value_offset,
value_len,
},
)) => {
let value = DbSlice::from_block(block, value_offset, value_len)
.ok_or_else(|| invalid_data("block value extends past block"))?;
LookupResult::Found { seq, value }
}
})
}
pub(crate) fn get_size(
&self,
lk: &LookupKey,
key_buf: &mut Vec<u8>,
cache: &BlockCache,
) -> io::Result<LookupResult<usize>> {
Ok(match self.probe(lk, key_buf, cache)? {
None => LookupResult::NotInTable,
Some((_, BlockHit::Tombstone { seq })) => LookupResult::FoundTombstone { seq },
Some((_, BlockHit::Value { seq, value_len, .. })) => LookupResult::Found {
seq,
value: value_len,
},
})
}
fn probe(
&self,
lk: &LookupKey,
key_buf: &mut Vec<u8>,
cache: &BlockCache,
) -> io::Result<Option<(Arc<Block>, BlockHit)>> {
let user_key = lk.prefixed_user_key();
if !self.filter(cache)?.may_contain(user_key) {
cache.record_bloom_useful();
return Ok(None);
}
let search_key = lk.internal();
let Some(handle) = self.find_block_handle(search_key, cache)? else {
return Ok(None);
};
let block = self.read_block(handle, cache)?;
let hit = block.scan_from(search_key, key_buf, |ik, value_offset, value_len| {
let (uk, seq, vt) = decode_internal_key(ik);
if uk != user_key {
return ControlFlow::Break(None);
}
match vt {
VALUE_TYPE_MERGE => ControlFlow::Continue(()),
VALUE_TYPE_DELETION => ControlFlow::Break(Some(BlockHit::Tombstone { seq })),
_ => ControlFlow::Break(Some(BlockHit::Value {
seq,
value_offset,
value_len,
})),
}
});
match hit.flatten() {
None => Ok(None),
Some(hit) => {
cache.record_bloom_full_positive();
Ok(Some((block, hit)))
}
}
}
pub(crate) fn collect_merge_chain(
&self,
lk: &LookupKey,
key_buf: &mut Vec<u8>,
cache: &BlockCache,
out: &mut Vec<(u64, u8, DbSlice)>,
) -> io::Result<bool> {
let user_key = lk.prefixed_user_key();
if !self.filter(cache)?.may_contain(user_key) {
return Ok(false);
}
let search_key = lk.internal();
let handle = match self.find_block_handle(search_key, cache)? {
Some(h) => h,
None => return Ok(false),
};
let block = self.read_block(handle, cache)?;
let mut spans: Vec<(u64, u8, usize, usize)> = Vec::new();
let terminated = block
.scan_from(search_key, key_buf, |ik, value_offset, value_len| {
let (uk, seq, vt) = decode_internal_key(ik);
if uk != user_key {
return ControlFlow::Break(false);
}
spans.push((seq, vt, value_offset, value_len));
if vt != VALUE_TYPE_MERGE {
ControlFlow::Break(true)
} else {
ControlFlow::Continue(())
}
})
.unwrap_or(false);
for (seq, vt, value_offset, value_len) in spans {
let value = DbSlice::from_block(Arc::clone(&block), value_offset, value_len)
.ok_or_else(|| invalid_data("block value extends past block"))?;
out.push((seq, vt, value));
}
Ok(terminated)
}
#[cfg(test)]
pub(crate) fn iter_internal(&self, cache: &BlockCache) -> io::Result<Vec<(Vec<u8>, Vec<u8>)>> {
let mut stream = self.iter_internal_stream(cache)?;
let mut result = Vec::new();
while let Some(entry) = stream.next_entry()? {
result.push(entry);
}
Ok(result)
}
pub(crate) fn iter_internal_stream<'a>(
&'a self,
cache: &'a BlockCache,
) -> io::Result<SsTableInternalIter<'a>> {
Ok(SsTableInternalIter {
reader: self,
cache,
next_cursor: self.first_block_cursor(cache)?,
current_block: None,
current_pos: 0,
current_key: Vec::new(),
})
}
pub(crate) fn approximate_size_in_range(
&self,
start: &[u8],
end: &[u8],
cache: &BlockCache,
) -> u64 {
if start >= end {
return 0;
}
let Ok(index) = self.index(cache) else {
return 0;
};
if index.is_empty() {
return 0;
}
let lo = LookupKey::from_prefixed(start, u64::MAX);
let hi = LookupKey::from_prefixed(end, u64::MAX);
let (lo_probe, hi_probe) = (lo.internal(), hi.internal());
if !self.partitioned {
return index.approximate_size_in_range(lo_probe, hi_probe);
}
let first_leaf = index.seek(lo_probe);
let end_leaf = index.seek(hi_probe).min(index.len() - 1);
if first_leaf > end_leaf {
return 0;
}
let mut total = 0;
for leaf_idx in first_leaf..=end_leaf {
let Some(handle) = index.handle(leaf_idx) else {
continue;
};
let leaf = match self.read_index_leaf(handle, cache) {
Ok(leaf) => leaf,
Err(_) => return 0,
};
total += leaf.approximate_size_in_range(lo_probe, hi_probe);
}
total
}
pub(crate) fn read_block(
&self,
handle: BlockHandle,
cache: &BlockCache,
) -> io::Result<Arc<Block>> {
if let Some(block) = cache.get(self.file_id, handle.offset) {
return Ok(block);
}
if handle.size < 5 {
return Err(invalid_data("block frame too short"));
}
let block_data = read_file_region(
&*self.file,
handle.offset,
handle.size,
self.data_end,
"data block",
)?;
self.decode_block_frame(handle, &block_data, cache)
}
pub(crate) fn read_span(&self, offset: u64, len: u64) -> io::Result<Vec<u8>> {
let len = len.min(self.data_end.saturating_sub(offset));
read_file_region(&*self.file, offset, len, self.data_end, "data span")
}
pub(crate) fn decode_block_frame(
&self,
handle: BlockHandle,
block_data: &[u8],
cache: &BlockCache,
) -> io::Result<Arc<Block>> {
if block_data.len() < 5 {
return Err(invalid_data("block frame too short"));
}
let compression_type = block_data[0];
let checksum_offset = block_data.len() - 4;
let stored_checksum = u32::from_le_bytes(block_data[checksum_offset..].try_into().unwrap());
let compressed_data = &block_data[1..checksum_offset];
let computed_checksum = checksum::sst_block(compression_type, compressed_data);
if stored_checksum != computed_checksum {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"block checksum mismatch",
));
}
let raw_data = match compression_type {
COMPRESSION_NONE => compressed_data.to_vec(),
COMPRESSION_LZ4 => lz4_flex::decompress_size_prepended(compressed_data)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?,
COMPRESSION_SNAPPY => {
if compressed_data.len() < 4 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"snappy block too short",
));
}
let raw_len =
u32::from_le_bytes(compressed_data[0..4].try_into().unwrap()) as usize;
let mut decoder = snap::raw::Decoder::new();
let mut out = vec![0u8; raw_len];
let n = decoder
.decompress(&compressed_data[4..], &mut out)
.map_err(|e| {
io::Error::new(io::ErrorKind::InvalidData, format!("snappy decode: {e}"))
})?;
if n != raw_len {
return Err(invalid_data("snappy decoded length mismatch"));
}
out.truncate(n);
out
}
_ => {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("unknown compression type: {}", compression_type),
));
}
};
let block = Arc::new(Block::decode_data_block(raw_data)?);
cache.insert(self.file_id, handle.offset, Arc::clone(&block));
Ok(block)
}
}
pub(crate) fn table_carries_data(env: &dyn Env, path: &Path) -> io::Result<bool> {
let file = env.open_read(path)?;
let file_size = file.len()?;
if file_size < 8 {
return Err(invalid_data(
"SSTable file too small to carry a magic number",
));
}
let mut magic_buf = [0u8; 8];
file.read_exact_at(file_size - 8, &mut magic_buf)?;
let footer_size = Footer::size_for_magic(u64::from_le_bytes(magic_buf))?;
let data_end = file_size
.checked_sub(footer_size as u64)
.ok_or_else(|| invalid_data("SSTable file too small for its footer"))?;
let mut footer_buf = vec![0u8; footer_size];
file.read_exact_at(data_end, &mut footer_buf)?;
let footer = Footer::decode(&footer_buf)?;
validate_footer_regions(&footer, data_end)?;
if footer.num_entries > 0 || footer.range_tombstone_size > 0 {
return Ok(true);
}
let mut index_region = read_file_region(
&*file,
footer.index_offset,
footer.index_size,
data_end,
"index block",
)?;
let payload_len = verify_meta_region(
&index_region,
checksum::META_KIND_INDEX,
true,
"index block",
)?
.len();
index_region.truncate(payload_len);
Ok(!IndexBlock::decode(index_region)?.is_empty())
}
pub(crate) fn sst_filename(id: u64) -> String {
format!("{:06}.sst", id)
}
pub(crate) fn remove_sst_in(env: &dyn Env, path: &Path) -> io::Result<()> {
env.remove_file(path)
}
#[cfg(test)]
mod tests {
use super::*;
fn probe_get(
reader: &SsTableReader,
key: &[u8],
snapshot_seq: u64,
cache: &BlockCache,
) -> io::Result<LookupResult> {
let lk = LookupKey::from_prefixed(key, snapshot_seq);
let mut key_buf = Vec::new();
reader.get(&lk, &mut key_buf, cache)
}
fn probe_get_size(
reader: &SsTableReader,
key: &[u8],
snapshot_seq: u64,
cache: &BlockCache,
) -> io::Result<LookupResult<usize>> {
let lk = LookupKey::from_prefixed(key, snapshot_seq);
let mut key_buf = Vec::new();
reader.get_size(&lk, &mut key_buf, cache)
}
use crate::engine::internal_key::{VALUE_TYPE_VALUE, encode_internal_key};
use tempfile::TempDir;
fn ik(key: &[u8], seq: u64) -> Vec<u8> {
encode_internal_key(key, seq, VALUE_TYPE_VALUE)
}
fn tombstone(key: &[u8], seq: u64) -> Vec<u8> {
encode_internal_key(key, seq, VALUE_TYPE_DELETION)
}
#[test]
fn test_sstable_write_read() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("test.sst");
let cache = BlockCache::new(64 * 1024 * 1024);
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::Lz4, None, false, 4096)
.unwrap();
for i in 0..100 {
let user_key = format!("key_{:04}", i);
let value = format!("value_{}", i);
writer
.add(&ik(user_key.as_bytes(), 1), value.as_bytes())
.unwrap();
}
let summary = writer.finish().unwrap().unwrap();
assert_eq!(summary.num_entries, 100);
assert_eq!(summary.smallest_user_key, b"key_0000");
assert_eq!(summary.largest_user_key, b"key_0099");
}
let reader = SsTableReader::open(&path, 1).unwrap();
assert_eq!(
probe_get(&reader, b"key_0042", u64::MAX, &cache).unwrap(),
LookupResult::Found {
seq: 1,
value: DbSlice::from_vec(b"value_42".to_vec())
}
);
assert_eq!(
probe_get(&reader, b"nonexistent", u64::MAX, &cache).unwrap(),
LookupResult::NotInTable
);
}
#[test]
fn test_sstable_mvcc_visibility() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("mvcc.sst");
let cache = BlockCache::new(1024 * 1024);
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096)
.unwrap();
writer.add(&tombstone(b"k", 5), b"").unwrap();
writer.add(&ik(b"k", 3), b"v3").unwrap();
writer.add(&ik(b"k", 1), b"v1").unwrap();
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 7).unwrap();
assert_eq!(
probe_get(&reader, b"k", 6, &cache).unwrap(),
LookupResult::FoundTombstone { seq: 5 }
);
assert_eq!(
probe_get(&reader, b"k", 4, &cache).unwrap(),
LookupResult::Found {
seq: 3,
value: DbSlice::from_vec(b"v3".to_vec())
}
);
assert_eq!(
probe_get(&reader, b"k", 2, &cache).unwrap(),
LookupResult::Found {
seq: 1,
value: DbSlice::from_vec(b"v1".to_vec())
}
);
assert_eq!(
probe_get(&reader, b"k", 0, &cache).unwrap(),
LookupResult::NotInTable
);
}
#[test]
fn test_sstable_no_compression() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("nocomp.sst");
let cache = BlockCache::new(1024 * 1024);
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096)
.unwrap();
writer.add(&ik(b"hello", 1), b"world").unwrap();
writer.add(&ik(b"test", 1), b"data").unwrap();
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 2).unwrap();
assert_eq!(
probe_get(&reader, b"hello", u64::MAX, &cache).unwrap(),
LookupResult::Found {
seq: 1,
value: DbSlice::from_vec(b"world".to_vec())
}
);
assert_eq!(
probe_get(&reader, b"test", u64::MAX, &cache).unwrap(),
LookupResult::Found {
seq: 1,
value: DbSlice::from_vec(b"data".to_vec())
}
);
}
#[test]
fn test_sstable_range_tombstones_roundtrip() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("rt.sst");
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096)
.unwrap();
writer.add(&ik(b"a", 1), b"v_a").unwrap();
writer.add(&ik(b"m", 2), b"v_m").unwrap();
writer.add(&ik(b"z", 3), b"v_z").unwrap();
writer.add_range_tombstone(b"b", b"n", 10);
writer.add_range_tombstone(b"p", b"y", 15);
let summary = writer.finish().unwrap().unwrap();
assert_eq!(summary.num_entries, 3);
}
let reader = SsTableReader::open(&path, 1).unwrap();
assert_eq!(reader.range_tombstones().len(), 2);
assert_eq!(reader.covering_range_tombstone_seq(b"a", 100), 0);
assert_eq!(reader.covering_range_tombstone_seq(b"b", 100), 10);
assert_eq!(reader.covering_range_tombstone_seq(b"m", 100), 10);
assert_eq!(reader.covering_range_tombstone_seq(b"n", 100), 0);
assert_eq!(reader.covering_range_tombstone_seq(b"p", 100), 15);
assert_eq!(reader.covering_range_tombstone_seq(b"z", 100), 0);
assert_eq!(reader.covering_range_tombstone_seq(b"p", 14), 0);
}
#[test]
fn test_sstable_rt_only_no_points() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("rt_only.sst");
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096)
.unwrap();
writer.add_range_tombstone(b"aa", b"kk", 5);
writer.add_range_tombstone(b"mm", b"zz", 7);
let summary = writer.finish().unwrap().unwrap();
assert_eq!(summary.num_entries, 0);
assert_eq!(summary.smallest_user_key, b"aa");
assert_eq!(summary.largest_user_key, b"zz");
}
let reader = SsTableReader::open(&path, 2).unwrap();
assert_eq!(reader.range_tombstones().len(), 2);
assert_eq!(reader.covering_range_tombstone_seq(b"bb", 100), 5);
assert_eq!(reader.covering_range_tombstone_seq(b"nn", 100), 7);
}
#[test]
fn test_sstable_prefix_bloom_roundtrip() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("prefix.sst");
let cache = BlockCache::new(1024 * 1024);
let extractor: Arc<dyn PrefixExtractor> = Arc::new(crate::options::FixedLengthPrefix(4));
{
let mut writer = SsTableWriter::new(
&path,
4096,
10,
CompressionType::None,
Some(extractor.clone()),
false,
4096,
)
.unwrap();
for tenant in &["aaaa", "bbbb", "cccc"] {
for i in 0..4 {
let user_key = format!("{}:key{}", tenant, i);
writer.add(&ik(user_key.as_bytes(), 1), b"v").unwrap();
}
}
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 1).unwrap();
assert!(reader.may_have_prefix(b"aaaa", &cache).unwrap());
assert!(reader.may_have_prefix(b"bbbb", &cache).unwrap());
assert!(reader.may_have_prefix(b"cccc", &cache).unwrap());
let mut false_positives = 0;
for i in 0..200u32 {
let p = format!("zz{:02}", i);
if reader.may_have_prefix(p.as_bytes(), &cache).unwrap() {
false_positives += 1;
}
}
assert!(
false_positives < 20,
"too many prefix bloom false positives: {}",
false_positives
);
assert_eq!(
probe_get(&reader, b"bbbb:key2", u64::MAX, &cache).unwrap(),
LookupResult::Found {
seq: 1,
value: DbSlice::from_vec(b"v".to_vec())
}
);
}
#[test]
fn test_sstable_without_prefix_bloom_is_superset() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("noprefix.sst");
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096)
.unwrap();
writer.add(&ik(b"aaaa:1", 1), b"v").unwrap();
writer.add(&ik(b"bbbb:1", 1), b"v").unwrap();
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 1).unwrap();
let cache = BlockCache::new(1024 * 1024);
assert!(reader.may_have_prefix(b"aaaa", &cache).unwrap());
assert!(reader.may_have_prefix(b"zzzz", &cache).unwrap());
assert!(reader.may_have_prefix(b"anything", &cache).unwrap());
}
#[test]
fn sst_filename_pads_six_digits() {
assert_eq!(sst_filename(0), "000000.sst");
assert_eq!(sst_filename(1), "000001.sst");
assert_eq!(sst_filename(999_999), "999999.sst");
}
#[test]
fn footer_decode_rejects_bad_magic() {
let mut buf = [0u8; FOOTER_SIZE];
buf[56..64].copy_from_slice(&0xDEAD_BEEF_u64.to_le_bytes());
let err = Footer::decode(&buf).expect_err("should reject");
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
assert!(err.to_string().contains("magic"), "got: {err}");
}
#[test]
fn footer_decode_catches_a_flip_in_every_checksummed_field() {
let f = Footer {
range_tombstone_offset: 7,
range_tombstone_size: 11,
bloom_offset: 13,
bloom_size: 17,
index_offset: 19,
index_size: 23,
num_entries: 29,
magic: MAGIC_V5,
};
let clean = f.encode();
for byte in 0..clean.len() {
for bit in 0..8u8 {
let mut damaged = clean.clone();
damaged[byte] ^= 1 << bit;
assert!(
Footer::decode(&damaged).is_err(),
"byte {byte} bit {bit} of a V3 footer was not caught"
);
}
}
}
#[test]
fn metadata_checksums_are_domain_separated_by_kind() {
let payload = b"the same bytes in two different regions";
let index = checksum::sst_meta(checksum::META_KIND_INDEX, payload);
assert_ne!(
index,
checksum::sst_meta(checksum::META_KIND_BLOOM, payload)
);
assert_ne!(
index,
checksum::sst_meta(checksum::META_KIND_INDEX_LEAF, payload)
);
assert_ne!(
index,
checksum::sst_meta(checksum::META_KIND_RANGE_TOMBSTONE, payload)
);
}
#[test]
fn a_freshly_written_table_carries_the_checksummed_format() {
for partitioned in [false, true] {
let dir = TempDir::new().unwrap();
let path = dir.path().join("fresh.sst");
let mut writer = SsTableWriter::new(
&path,
256,
10,
CompressionType::None,
None,
partitioned,
256,
)
.unwrap();
for i in 0..64u32 {
let key = super::super::internal_key::encode_internal_key(
format!("k{i:04}").as_bytes(),
1,
VALUE_TYPE_VALUE,
);
writer.add(&key, format!("v{i}").as_bytes()).unwrap();
}
writer.finish().unwrap().unwrap();
let on_disk = fs::read(&path).unwrap();
let magic = u64::from_le_bytes(on_disk[on_disk.len() - 8..].try_into().unwrap());
assert_eq!(
magic,
if partitioned { MAGIC_V6 } else { MAGIC_V5 },
"a fresh table must carry the REGOSST magic"
);
assert_eq!(
&magic.to_be_bytes()[..7],
b"REGOSST",
"the identifier itself must read as REGOSST"
);
let reader = SsTableReader::open(&path, 1).unwrap();
assert_eq!(reader.partitioned, partitioned);
assert!(table_carries_data(&*crate::env::std_env(), &path).unwrap());
}
}
#[test]
fn a_flip_in_an_index_leaf_or_a_range_tombstone_block_is_caught() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("partitioned-rt.sst");
let mut writer =
SsTableWriter::new(&path, 64, 10, CompressionType::None, None, true, 96).unwrap();
for i in 0..40 {
writer
.add(&ik(format!("key_{i:04}").as_bytes(), 1), b"value")
.unwrap();
}
writer.add_range_tombstone(b"key_0000", b"key_0005", 3);
writer.finish().unwrap().unwrap();
let clean = fs::read(&path).unwrap();
let f = &clean[clean.len() - FOOTER_SIZE..];
let word = |i: usize| u64::from_le_bytes(f[i * 8..i * 8 + 8].try_into().unwrap());
let range_tombstones = word(0)..word(0) + word(1);
let leaves = word(2) + word(3)..word(4);
assert!(
!range_tombstones.is_empty(),
"the fixture has no range tombstones"
);
assert!(
leaves.end - leaves.start > 1,
"the fixture did not partition its index"
);
for offset in range_tombstones.chain(leaves) {
for bit in [0u8, 5] {
let mut damaged = clean.clone();
damaged[offset as usize] ^= 1 << bit;
fs::write(&path, &damaged).unwrap();
let cache = BlockCache::new(1024 * 1024);
let caught = match SsTableReader::open(&path, 1) {
Err(_) => true,
Ok(reader) => reader.iter_internal(&cache).is_err(),
};
assert!(caught, "byte {offset} bit {bit} was not caught");
}
}
}
fn expect_reader_open_err(path: &Path) -> io::Error {
match SsTableReader::open(path, 1) {
Err(e) => e,
Ok(_) => panic!("expected invalid SSTable"),
}
}
#[test]
fn open_rejects_file_smaller_than_footer() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("tiny.sst");
fs::write(&path, vec![0u8; 10]).unwrap();
let kind = match SsTableReader::open(&path, 1) {
Err(e) => e.kind(),
Ok(_) => panic!("expected InvalidData"),
};
assert_eq!(kind, io::ErrorKind::InvalidData);
}
#[test]
fn open_rejects_file_with_bogus_magic() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("bogus.sst");
fs::write(&path, vec![0u8; FOOTER_SIZE]).unwrap();
let kind = match SsTableReader::open(&path, 1) {
Err(e) => e.kind(),
Ok(_) => panic!("expected InvalidData"),
};
assert_eq!(kind, io::ErrorKind::InvalidData);
}
#[test]
fn open_rejects_footer_region_past_file_data() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("bad-region.sst");
let footer = Footer {
range_tombstone_offset: 0,
range_tombstone_size: 0,
bloom_offset: 0,
bloom_size: 8,
index_offset: 8,
index_size: 4,
num_entries: 0,
magic: MAGIC_V5,
};
fs::write(&path, footer.encode()).unwrap();
let err = expect_reader_open_err(&path);
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn open_missing_file_returns_not_found() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("never.sst");
let kind = match SsTableReader::open(&path, 1) {
Err(e) => e.kind(),
Ok(_) => panic!("expected NotFound"),
};
assert_eq!(kind, io::ErrorKind::NotFound);
}
#[test]
fn read_block_rejects_short_frame() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("short-block.bin");
fs::write(&path, [0u8; 4]).unwrap();
let reader = SsTableReader {
file: crate::env::std_env().open_read(&path).unwrap(),
data_end: 4,
file_id: 1,
index: MetaSlot::Pinned(Arc::new(
IndexBlock::decode(0u32.to_le_bytes().to_vec()).unwrap(),
)),
filter: MetaSlot::Cached(BlockHandle { offset: 0, size: 0 }),
index_fallback: OnceLock::new(),
filter_fallback: OnceLock::new(),
range_tombstones: RangeTombstoneSet::default(),
partitioned: false,
index_leaf_reads: AtomicUsize::new(0),
};
let cache = BlockCache::new(1024);
let err = match reader.read_block(BlockHandle { offset: 0, size: 4 }, &cache) {
Err(e) => e,
Ok(_) => panic!("expected invalid block frame"),
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn block_checksum_covers_compression_type() {
let payload = b"payload";
let baseline = checksum::sst_block(COMPRESSION_NONE, payload);
assert_ne!(baseline, checksum::sst_block(COMPRESSION_LZ4, payload));
}
#[test]
fn read_block_rejects_compression_header_flip() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("header-flip.sst");
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096)
.unwrap();
writer.add(&ik(b"k", 1), b"v").unwrap();
writer.finish().unwrap().unwrap();
}
let mut bytes = fs::read(&path).unwrap();
bytes[0] = COMPRESSION_LZ4;
fs::write(&path, bytes).unwrap();
let reader = SsTableReader::open(&path, 1).unwrap();
let cache = BlockCache::new(1024);
let err = match probe_get(&reader, b"k", u64::MAX, &cache) {
Err(e) => e,
Ok(v) => panic!("expected checksum error, got {v:?}"),
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn finish_on_empty_writer_returns_none() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("empty.sst");
let writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096).unwrap();
assert!(writer.finish().unwrap().is_none());
}
#[test]
fn finish_with_only_range_tombstones_returns_summary() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("rt_only_summary.sst");
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096).unwrap();
writer.add_range_tombstone(b"a", b"z", 7);
let summary = writer.finish().unwrap().expect("should produce summary");
assert_eq!(summary.num_entries, 0);
assert_eq!(summary.smallest_user_key, b"a");
assert_eq!(summary.largest_user_key, b"z");
}
#[test]
fn summary_records_first_and_last_user_keys() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("bounds.sst");
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096).unwrap();
writer.add(&ik(b"alpha", 1), b"1").unwrap();
writer.add(&ik(b"bravo", 1), b"2").unwrap();
writer.add(&ik(b"delta", 1), b"3").unwrap();
let summary = writer.finish().unwrap().unwrap();
assert_eq!(summary.smallest_user_key, b"alpha");
assert_eq!(summary.largest_user_key, b"delta");
assert_eq!(summary.num_entries, 3);
}
#[test]
fn summary_includes_range_tombstone_bounds_with_point_entries() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("mixed_rt_bounds.sst");
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096).unwrap();
writer.add(&ik(b"m", 1), b"1").unwrap();
writer.add(&ik(b"n", 1), b"2").unwrap();
writer.add_range_tombstone(b"a", b"c", 7);
writer.add_range_tombstone(b"x", b"z", 8);
let summary = writer.finish().unwrap().unwrap();
assert_eq!(summary.smallest_user_key, b"a");
assert_eq!(summary.largest_user_key, b"z");
assert_eq!(summary.num_entries, 2);
}
#[test]
fn multi_block_lookup_walks_index_correctly() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("multi.sst");
let cache = BlockCache::new(1024 * 1024);
{
let mut writer =
SsTableWriter::new(&path, 64, 10, CompressionType::None, None, false, 4096)
.unwrap();
for i in 0..200 {
let k = format!("k_{:04}", i);
let v = format!("v_{}", i);
writer.add(&ik(k.as_bytes(), 1), v.as_bytes()).unwrap();
}
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 1).unwrap();
for i in 0..200 {
let k = format!("k_{:04}", i);
let expected = format!("v_{}", i);
assert_eq!(
probe_get(&reader, k.as_bytes(), u64::MAX, &cache).unwrap(),
LookupResult::Found {
seq: 1,
value: DbSlice::from_vec(expected.into_bytes())
},
);
}
}
#[test]
fn iter_internal_returns_every_entry_in_order() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("iter.sst");
let cache = BlockCache::new(1024 * 1024);
let mut writer =
SsTableWriter::new(&path, 64, 10, CompressionType::None, None, false, 4096).unwrap();
writer.add(&ik(b"a", 3), b"a3").unwrap();
writer.add(&ik(b"a", 1), b"a1").unwrap();
writer.add(&ik(b"b", 2), b"b2").unwrap();
writer.finish().unwrap().unwrap();
let reader = SsTableReader::open(&path, 1).unwrap();
let pairs = reader.iter_internal(&cache).unwrap();
assert_eq!(pairs.len(), 3);
let user_keys: Vec<&[u8]> = pairs.iter().map(|(k, _)| user_key_of(k)).collect();
assert_eq!(user_keys, vec![&b"a"[..], &b"a"[..], &b"b"[..]]);
}
#[test]
fn iter_internal_stream_matches_vec_iterator() {
for partitioned in [false, true] {
let dir = TempDir::new().unwrap();
let path = dir.path().join("stream.sst");
let cache = BlockCache::new(1024 * 1024);
let mut writer =
SsTableWriter::new(&path, 64, 10, CompressionType::None, None, partitioned, 96)
.unwrap();
for i in 0..128 {
let key = format!("k{i:04}");
let value = format!("v{i}");
writer
.add(&ik(key.as_bytes(), 1), value.as_bytes())
.unwrap();
}
writer.finish().unwrap().unwrap();
let reader = SsTableReader::open(&path, 1).unwrap();
let expected = reader.iter_internal(&cache).unwrap();
let mut stream = reader.iter_internal_stream(&cache).unwrap();
let mut actual = Vec::new();
while let Some(entry) = stream.next_entry().unwrap() {
actual.push(entry);
}
assert_eq!(actual, expected);
}
}
#[test]
fn approximate_size_empty_range_is_zero() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("approx.sst");
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096).unwrap();
writer.add(&ik(b"m", 1), b"v").unwrap();
writer.finish().unwrap().unwrap();
let reader = SsTableReader::open(&path, 1).unwrap();
let cache = BlockCache::new(1024 * 1024);
assert_eq!(reader.approximate_size_in_range(b"z", b"a", &cache), 0);
assert_eq!(reader.approximate_size_in_range(b"m", b"m", &cache), 0);
}
#[test]
fn approximate_size_grows_with_range_width() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("approx_grow.sst");
let mut writer =
SsTableWriter::new(&path, 128, 10, CompressionType::None, None, false, 4096).unwrap();
for i in 0..500 {
writer
.add(&ik(format!("k_{:04}", i).as_bytes(), 1), b"value")
.unwrap();
}
writer.finish().unwrap().unwrap();
let reader = SsTableReader::open(&path, 1).unwrap();
let cache = BlockCache::new(1024 * 1024);
let narrow = reader.approximate_size_in_range(b"k_0000", b"k_0001", &cache);
let wide = reader.approximate_size_in_range(b"k_0000", b"k_0499", &cache);
assert!(
wide > narrow,
"wide ({wide}) should exceed narrow ({narrow})"
);
}
#[test]
fn lz4_compression_round_trips() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("lz4.sst");
let cache = BlockCache::new(1024 * 1024);
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::Lz4, None, false, 4096)
.unwrap();
for i in 0..50 {
writer
.add(
&ik(format!("key_{:03}", i).as_bytes(), 1),
b"AAAAAAAAAAAAAAAAAAAAAAAAAA",
)
.unwrap();
}
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 1).unwrap();
assert_eq!(
probe_get(&reader, b"key_025", u64::MAX, &cache).unwrap(),
LookupResult::Found {
seq: 1,
value: DbSlice::from_vec(b"AAAAAAAAAAAAAAAAAAAAAAAAAA".to_vec())
}
);
}
#[test]
fn snappy_compression_round_trips() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("snappy.sst");
let cache = BlockCache::new(1024 * 1024);
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::Snappy, None, false, 4096)
.unwrap();
for i in 0..50 {
writer
.add(
&ik(format!("key_{:03}", i).as_bytes(), 1),
b"BBBBBBBBBBBBBBBBBBBBBBBBBB",
)
.unwrap();
}
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 1).unwrap();
assert_eq!(
probe_get(&reader, b"key_007", u64::MAX, &cache).unwrap(),
LookupResult::Found {
seq: 1,
value: DbSlice::from_vec(b"BBBBBBBBBBBBBBBBBBBBBBBBBB".to_vec())
}
);
}
#[test]
fn range_tombstone_block_round_trip() {
let input = vec![
RangeTombstone::new(b"a".to_vec(), b"c".to_vec(), 5),
RangeTombstone::new(b"f".to_vec(), b"k".to_vec(), 7),
];
let bytes = encode_range_tombstone_block(&input);
let got = decode_range_tombstone_block(&bytes).unwrap();
assert_eq!(got.len(), 2);
assert_eq!(got[0].start, b"a");
assert_eq!(got[0].end, b"c");
assert_eq!(got[0].seq, 5);
assert_eq!(got[1].seq, 7);
}
#[test]
fn range_tombstone_block_empty_input() {
assert!(decode_range_tombstone_block(&[]).unwrap().is_empty());
}
#[test]
fn range_tombstone_block_rejects_tiny_header() {
let kind = match decode_range_tombstone_block(&[0, 1]) {
Err(e) => e.kind(),
Ok(v) => panic!("expected error, got {} tombstones", v.len()),
};
assert_eq!(kind, io::ErrorKind::InvalidData);
}
#[test]
fn range_tombstone_block_rejects_truncated_mid_record() {
let mut bytes = vec![];
bytes.extend_from_slice(&2u32.to_le_bytes());
bytes.extend_from_slice(&2u32.to_le_bytes());
bytes.extend_from_slice(b"aa");
bytes.extend_from_slice(&2u32.to_le_bytes());
bytes.extend_from_slice(b"bb");
bytes.extend_from_slice(&1u64.to_le_bytes());
bytes.extend_from_slice(&2u32.to_le_bytes());
bytes.push(b'c');
let err = decode_range_tombstone_block(&bytes).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn range_tombstone_block_rejects_trailing_bytes() {
let mut bytes =
encode_range_tombstone_block(&[RangeTombstone::new(b"a".to_vec(), b"b".to_vec(), 1)]);
bytes.push(0xAA);
let err = decode_range_tombstone_block(&bytes).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn index_block_rejects_tiny_header() {
let err = IndexBlock::decode(vec![1, 2]).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn index_block_rejects_truncated_entry() {
let mut bytes = Vec::new();
bytes.extend_from_slice(&1u32.to_le_bytes());
bytes.extend_from_slice(&4u32.to_le_bytes());
bytes.extend_from_slice(b"ab");
let err = IndexBlock::decode(bytes).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn index_block_rejects_trailing_bytes() {
let mut bytes = encode_index_block(&[]);
bytes.push(0xAA);
let err = IndexBlock::decode(bytes).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn reader_with_no_range_tombstones_reports_empty_list() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("no_rt.sst");
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096)
.unwrap();
writer.add(&ik(b"k", 1), b"v").unwrap();
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 1).unwrap();
assert!(reader.range_tombstones().is_empty());
assert_eq!(reader.covering_range_tombstone_seq(b"k", u64::MAX), 0);
}
#[test]
fn negative_lookup_on_many_keys_mostly_short_circuits() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("bloomy.sst");
let cache = BlockCache::new(1024 * 1024);
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096)
.unwrap();
for i in 0..1000 {
writer
.add(&ik(format!("hit_{:04}", i).as_bytes(), 1), b"v")
.unwrap();
}
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 1).unwrap();
for i in 0..5000 {
let absent = format!("miss_{:05}", i);
assert_eq!(
probe_get(&reader, absent.as_bytes(), u64::MAX, &cache).unwrap(),
LookupResult::NotInTable
);
}
}
fn write_partitioned_fixture(path: &Path) {
let mut writer =
SsTableWriter::new(path, 96, 10, CompressionType::None, None, true, 96).unwrap();
for i in 0..300 {
writer
.add(&ik(format!("k_{:04}", i).as_bytes(), 1), b"value")
.unwrap();
}
writer.finish().unwrap().unwrap();
}
#[test]
fn partitioned_index_round_trip() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("partitioned.sst");
let cache = BlockCache::new(1024 * 1024);
write_partitioned_fixture(&path);
let reader = SsTableReader::open(&path, 1).unwrap();
for i in [0usize, 75, 150, 225, 299] {
let k = format!("k_{:04}", i);
assert_eq!(
probe_get(&reader, k.as_bytes(), u64::MAX, &cache).unwrap(),
LookupResult::Found {
seq: 1,
value: DbSlice::from_vec(b"value".to_vec()),
}
);
}
}
#[test]
fn partitioned_point_lookup_reads_one_index_leaf() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("partitioned_point.sst");
let cache = BlockCache::new(1024 * 1024);
write_partitioned_fixture(&path);
let reader = SsTableReader::open(&path, 1).unwrap();
assert!(
reader.pinned_index_len().unwrap() > 3,
"fixture must contain several index leaves"
);
reader.reset_index_leaf_read_count();
assert_eq!(
probe_get(&reader, b"k_0150", u64::MAX, &cache).unwrap(),
LookupResult::Found {
seq: 1,
value: DbSlice::from_vec(b"value".to_vec()),
}
);
assert_eq!(
reader.index_leaf_read_count(),
1,
"point lookup must load only the selected leaf index"
);
}
#[test]
fn partitioned_internal_stream_starts_without_expanding_all_leaves() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("partitioned_stream.sst");
let cache = BlockCache::new(1024 * 1024);
write_partitioned_fixture(&path);
let reader = SsTableReader::open(&path, 1).unwrap();
let leaf_count = reader.pinned_index_len().unwrap();
assert!(leaf_count > 3, "fixture must contain several index leaves");
reader.reset_index_leaf_read_count();
let mut stream = reader.iter_internal_stream(&cache).unwrap();
assert_eq!(
reader.index_leaf_read_count(),
1,
"stream construction should open only the first leaf"
);
let first = stream.next_entry().unwrap().expect("fixture has entries");
assert_eq!(user_key_of(&first.0), b"k_0000");
assert!(
reader.index_leaf_read_count() < leaf_count,
"reading the first entry must not expand every leaf"
);
}
#[test]
fn partitioned_approximate_size_reads_only_overlapped_leaves() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("partitioned_approx.sst");
write_partitioned_fixture(&path);
let reader = SsTableReader::open(&path, 1).unwrap();
let cache = BlockCache::new(1024 * 1024);
let leaf_count = reader.pinned_index_len().unwrap();
assert!(leaf_count > 3, "fixture must contain several index leaves");
reader.reset_index_leaf_read_count();
assert!(reader.approximate_size_in_range(b"k_0100", b"k_0101", &cache) > 0);
assert!(
reader.index_leaf_read_count() <= 2,
"narrow range should load at most the boundary leaves"
);
assert!(
reader.index_leaf_read_count() < leaf_count,
"narrow range must not expand every leaf"
);
}
fn write_flat_fixture(path: &Path) {
let mut writer =
SsTableWriter::new(path, 256, 10, CompressionType::None, None, false, 4096).unwrap();
for i in 0..500 {
writer
.add(&ik(format!("k_{:04}", i).as_bytes(), 1), b"value")
.unwrap();
}
writer.finish().unwrap().unwrap();
}
#[test]
fn partitioned_leaf_is_read_from_disk_once_then_served_from_the_cache() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("leaf_cached.sst");
let cache = BlockCache::new(1024 * 1024);
write_partitioned_fixture(&path);
let reader = SsTableReader::open(&path, 1).unwrap();
reader.reset_index_leaf_read_count();
assert!(
probe_get(&reader, b"k_0150", u64::MAX, &cache).unwrap() != LookupResult::NotInTable
);
let after_first = reader.index_leaf_read_count();
assert_eq!(after_first, 1, "first lookup must read its leaf from disk");
for _ in 0..8 {
assert!(
probe_get(&reader, b"k_0150", u64::MAX, &cache).unwrap()
!= LookupResult::NotInTable
);
}
assert_eq!(
reader.index_leaf_read_count(),
after_first,
"repeat lookups must be served from the cached leaf"
);
}
#[test]
fn cached_leaf_survives_a_second_reader_over_the_same_file_id() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("leaf_shared.sst");
let cache = BlockCache::new(1024 * 1024);
write_partitioned_fixture(&path);
let warm = SsTableReader::open(&path, 9).unwrap();
probe_get(&warm, b"k_0150", u64::MAX, &cache).unwrap();
let fresh = SsTableReader::open(&path, 9).unwrap();
fresh.reset_index_leaf_read_count();
probe_get(&fresh, b"k_0150", u64::MAX, &cache).unwrap();
assert_eq!(
fresh.index_leaf_read_count(),
0,
"a leaf cached under the same file id is shared across readers"
);
}
#[test]
fn cached_policy_answers_the_same_as_pinned() {
let dir = TempDir::new().unwrap();
for (name, partitioned) in [("cached_flat.sst", false), ("cached_part.sst", true)] {
let path = dir.path().join(name);
if partitioned {
write_partitioned_fixture(&path);
} else {
write_flat_fixture(&path);
}
let pinned_cache = BlockCache::new(1024 * 1024);
let cached_cache = BlockCache::new(1024 * 1024);
let pinned =
SsTableReader::open_with(&crate::env::std_env(), &path, 1, MetadataPolicy::Pinned)
.unwrap();
let cached =
SsTableReader::open_with(&crate::env::std_env(), &path, 1, MetadataPolicy::Cached)
.unwrap();
for probe in ["k_0000", "k_0123", "k_0299", "k_9999"] {
let a = probe_get(&pinned, probe.as_bytes(), u64::MAX, &pinned_cache).unwrap();
let b = probe_get(&cached, probe.as_bytes(), u64::MAX, &cached_cache).unwrap();
assert_eq!(a, b, "{name} disagreed on {probe}");
}
let pinned_all = pinned.iter_internal(&pinned_cache).unwrap();
let cached_all = cached.iter_internal(&cached_cache).unwrap();
assert_eq!(pinned_all, cached_all, "{name} full scans disagreed");
}
}
#[test]
fn cached_policy_moves_flat_index_and_filter_into_the_cache() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("charge.sst");
write_flat_fixture(&path);
let pinned =
SsTableReader::open_with(&crate::env::std_env(), &path, 1, MetadataPolicy::Pinned)
.unwrap();
let cached =
SsTableReader::open_with(&crate::env::std_env(), &path, 2, MetadataPolicy::Cached)
.unwrap();
assert!(
pinned.pinned_metadata_bytes() > cached.pinned_metadata_bytes(),
"the cached reader must hold less outside the budget"
);
let cache = BlockCache::new(1024 * 1024);
assert_eq!(cache.usage(), 0);
probe_get(&cached, b"k_0100", u64::MAX, &cache).unwrap();
assert!(
cache.usage() > 0,
"index and filter bytes must be charged to the cache"
);
}
#[test]
fn cached_policy_still_rejects_a_corrupt_index_at_open() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("corrupt_index.sst");
write_flat_fixture(&path);
let mut bytes = fs::read(&path).unwrap();
let footer_at = bytes.len() - FOOTER_SIZE;
let index_size =
u64::from_le_bytes(bytes[footer_at + 40..footer_at + 48].try_into().unwrap());
bytes[footer_at + 40..footer_at + 48].copy_from_slice(&(index_size - 1).to_le_bytes());
fs::write(&path, &bytes).unwrap();
for policy in [MetadataPolicy::Pinned, MetadataPolicy::Cached] {
let err = match SsTableReader::open_with(&crate::env::std_env(), &path, 1, policy) {
Err(e) => e,
Ok(_) => panic!("a corrupt index must fail the open under either policy"),
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
}
#[test]
fn strict_cache_refusal_pins_the_metadata_instead_of_re_reading_it() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("strict.sst");
write_flat_fixture(&path);
let cache = BlockCache::with_config(512, 0, true);
let reader =
SsTableReader::open_with(&crate::env::std_env(), &path, 1, MetadataPolicy::Cached)
.unwrap();
assert_eq!(reader.pinned_metadata_bytes(), 0);
for i in 0..64 {
let key = format!("k_{:04}", i);
assert!(
probe_get(&reader, key.as_bytes(), u64::MAX, &cache).unwrap()
!= LookupResult::NotInTable,
"strict-cache reads must still be correct"
);
}
assert!(
reader.pinned_metadata_bytes() > 0,
"a refused metadata insert must fall back to pinning"
);
}
#[test]
fn index_filter_and_data_offsets_do_not_collide() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("offsets.sst");
write_partitioned_fixture(&path);
let mut bytes = fs::read(&path).unwrap();
let magic = u64::from_le_bytes(bytes[bytes.len() - 8..].try_into().unwrap());
let footer_size = Footer::size_for_magic(magic).unwrap();
let footer_at = bytes.len() - footer_size;
let footer = Footer::decode(&bytes[footer_at..]).unwrap();
bytes.truncate(footer_at);
let cache = BlockCache::new(1024 * 1024);
let reader = SsTableReader::open(&path, 1).unwrap();
let mut offsets = std::collections::HashSet::new();
assert!(offsets.insert(footer.bloom_offset), "filter offset");
assert!(offsets.insert(footer.index_offset), "top index offset");
let index = reader.index(&cache).unwrap();
for leaf_idx in 0..index.len() {
let leaf_handle = index.handle(leaf_idx).unwrap();
assert!(
offsets.insert(leaf_handle.offset),
"leaf {leaf_idx} offset collides"
);
let leaf = reader.read_index_leaf(leaf_handle, &cache).unwrap();
for entry_idx in 0..leaf.len() {
let handle = leaf.handle(entry_idx).unwrap();
assert!(
offsets.insert(handle.offset),
"data block offset {} collides",
handle.offset
);
}
}
}
#[test]
fn pinned_metadata_bytes_counts_index_filter_and_tombstones() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("pinned_bytes.sst");
{
let mut writer =
SsTableWriter::new(&path, 4096, 10, CompressionType::None, None, false, 4096)
.unwrap();
writer.add(&ik(b"a", 1), b"v").unwrap();
writer.add_range_tombstone(b"m", b"z", 5);
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 1).unwrap();
let reported = reader.pinned_metadata_bytes();
assert!(
reported > std::mem::size_of::<RangeTombstone>(),
"pinned bytes {reported} must cover more than one tombstone"
);
}
fn reader_over(
entries: &std::collections::BTreeMap<Vec<u8>, Vec<u8>>,
) -> (TempDir, SsTableReader, BlockCache) {
let dir = TempDir::new().expect("tempdir");
let path = dir.path().join("agree.sst");
let cache = BlockCache::new(4 * 1024 * 1024);
{
let mut writer =
SsTableWriter::new(&path, 512, 10, CompressionType::None, None, false, 512)
.expect("writer");
for (key, value) in entries {
writer.add(&ik(key, 1), value).expect("add");
}
writer.finish().expect("finish").expect("non-empty");
}
let reader = SsTableReader::open(&path, 7).expect("open");
(dir, reader, cache)
}
fn assert_size_agrees(reader: &SsTableReader, key: &[u8], seq: u64, cache: &BlockCache) {
let full = probe_get(reader, key, seq, cache).expect("get");
let sized = probe_get_size(reader, key, seq, cache).expect("get_size");
match (full, sized) {
(LookupResult::NotInTable, LookupResult::NotInTable) => {}
(LookupResult::FoundTombstone { seq: a }, LookupResult::FoundTombstone { seq: b }) => {
assert_eq!(a, b);
}
(LookupResult::Found { seq: a, value }, LookupResult::Found { seq: b, value: len }) => {
assert_eq!(a, b);
assert_eq!(value.len(), len);
}
(a, b) => panic!("get and get_size disagreed: {a:?} vs {b:?}"),
}
}
#[test]
fn get_size_agrees_with_get_on_tombstones_and_misses() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("mixed.sst");
let cache = BlockCache::new(1024 * 1024);
{
let mut writer =
SsTableWriter::new(&path, 512, 10, CompressionType::None, None, false, 512)
.unwrap();
writer.add(&ik(b"alive", 2), b"payload").unwrap();
writer.add(&tombstone(b"dead", 3), b"").unwrap();
writer.add(&ik(b"empty", 4), b"").unwrap();
writer.finish().unwrap().unwrap();
}
let reader = SsTableReader::open(&path, 3).unwrap();
for key in [b"alive".as_ref(), b"dead", b"empty", b"absent", b"zzz"] {
assert_size_agrees(&reader, key, u64::MAX, &cache);
}
assert_eq!(
probe_get_size(&reader, b"alive", u64::MAX, &cache).unwrap(),
LookupResult::Found { seq: 2, value: 7 }
);
assert_eq!(
probe_get_size(&reader, b"empty", u64::MAX, &cache).unwrap(),
LookupResult::Found { seq: 4, value: 0 },
"a zero-length value is present with length 0, not absent"
);
assert_eq!(
probe_get_size(&reader, b"dead", u64::MAX, &cache).unwrap(),
LookupResult::FoundTombstone { seq: 3 }
);
assert_eq!(
probe_get_size(&reader, b"absent", u64::MAX, &cache).unwrap(),
LookupResult::NotInTable
);
}
proptest::proptest! {
#[test]
fn get_size_agrees_with_get(
entries in proptest::collection::btree_map(
proptest::collection::vec(proptest::prelude::any::<u8>(), 1..12),
proptest::collection::vec(proptest::prelude::any::<u8>(), 0..40),
1..40,
),
probes in proptest::collection::vec(
proptest::collection::vec(proptest::prelude::any::<u8>(), 1..12), 1..12),
) {
let (_dir, reader, cache) = reader_over(&entries);
for key in entries.keys() {
assert_size_agrees(&reader, key, u64::MAX, &cache);
}
for probe in &probes {
assert_size_agrees(&reader, probe, u64::MAX, &cache);
}
}
}
}