use crate::wire::manifest::MetadataRow;
use serde::{Deserialize, Serialize};
use std::io::Read;
use std::num::NonZeroUsize;
use thiserror::Error;
use xxhash_rust::xxh64::xxh64;
pub const DEFAULT_TARGET_BLOCK_BYTES: usize = 64 * 1024;
pub const DEFAULT_MAX_DELTA_RUNS: usize = 8;
pub const DEFAULT_MAX_ROWS_PER_SEGMENT: usize = 65_536;
pub const DEFAULT_TARGET_SEGMENT_BYTES: usize = 8 * 1024 * 1024;
pub const DEFAULT_MAX_REORGANIZATION_INPUT_RUNS: usize = 8;
pub const DEFAULT_MAX_REORGANIZATION_INPUT_ROWS: usize = 131_072;
pub const DEFAULT_MAX_REORGANIZATION_INPUT_BYTES: usize = 64 * 1024 * 1024;
pub const DEFAULT_INLINE_FILTER_MAX_BYTES: u32 = 1024;
pub const RESTART_INTERVAL: usize = 16;
pub const FILTER_BITS_PER_KEY: usize = 10;
pub const FILTER_HASH_COUNT: u32 = 7;
const FILTER_HASH_SEED_ONE: u64 = 0;
const FILTER_HASH_SEED_TWO: u64 = 0x9e37_79b9_7f4a_7c15;
pub(crate) const ZSTD_LEVEL: i32 = 3;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct BlockHandle {
pub offset: u64,
pub stored_len: u32,
pub decoded_len: u32,
pub crc32c: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct SegmentIndexEntry {
pub last_row_key: String,
pub block: BlockHandle,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BuiltSegmentBlocks {
pub bytes: Vec<u8>,
pub index: BlockHandle,
pub filter: BlockHandle,
pub row_count: u64,
pub min_row_key: String,
pub max_row_key: String,
}
impl BuiltSegmentBlocks {
pub fn inline_filter_hex(&self) -> Option<String> {
(self.filter.stored_len <= DEFAULT_INLINE_FILTER_MAX_BYTES).then(|| {
let start = self.filter.offset as usize;
crate::wire::hex::hex_encode_bytes(
&self.bytes[start..start + self.filter.stored_len as usize],
)
})
}
}
pub fn key_range_may_intersect(
min_row_key: &str,
max_row_key: &str,
row_count: u64,
lower_bound: &str,
upper_bound: Option<&str>,
) -> bool {
row_count != 0
&& max_row_key >= lower_bound
&& upper_bound.is_none_or(|upper_bound| min_row_key < upper_bound)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DecodedDataBlock<R = MetadataRow> {
pub row_keys: Vec<String>,
pub rows: Vec<R>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SegmentFilter {
n_hashes: u32,
bit_len: u64,
bits: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq, Error)]
#[non_exhaustive]
pub enum SstBlockCodecError {
#[error("segment must contain at least one row")]
EmptySegment,
#[error("row key `{offered}` is not in ascending order after `{previous}`")]
RowKeysOutOfOrder {
previous: String,
offered: String,
},
#[error("stored bytes length {actual} does not match handle length {expected}")]
StoredLengthMismatch {
expected: u32,
actual: usize,
},
#[error("block checksum mismatch: expected {expected:#010x}, actual {actual:#010x}")]
ChecksumMismatch {
expected: u32,
actual: u32,
},
#[error("decoded length {actual} does not match handle length {expected}")]
DecodedLengthMismatch {
expected: u32,
actual: usize,
},
#[error("malformed block: {0}")]
Malformed(String),
#[error("block codec error: {0}")]
Codec(String),
}
#[derive(Debug)]
#[must_use]
pub struct SegmentBlocksBuilder {
target_block_bytes: usize,
entries: Vec<u8>,
restarts: Vec<u32>,
entry_count: usize,
previous_key: String,
bytes: Vec<u8>,
index: Vec<SegmentIndexEntry>,
decoded_data_bytes: usize,
filter_hashes: Vec<(u64, u64)>,
row_count: u64,
min_row_key: String,
}
impl Default for SegmentBlocksBuilder {
fn default() -> Self {
Self::new(const { NonZeroUsize::new(DEFAULT_TARGET_BLOCK_BYTES).unwrap() })
}
}
impl SegmentBlocksBuilder {
pub fn new(target_block_bytes: NonZeroUsize) -> Self {
Self {
target_block_bytes: target_block_bytes.get(),
entries: Vec::new(),
restarts: Vec::new(),
entry_count: 0,
previous_key: String::new(),
bytes: Vec::new(),
index: Vec::new(),
decoded_data_bytes: 0,
filter_hashes: Vec::new(),
row_count: 0,
min_row_key: String::new(),
}
}
pub fn push<R: Serialize>(
&mut self,
row_key: &str,
filter_key: &str,
row: &R,
) -> Result<(), SstBlockCodecError> {
self.push_with_encoded_row(row_key, filter_key, row)
.map(drop)
}
pub fn push_with_encoded_row<R: Serialize>(
&mut self,
row_key: &str,
filter_key: &str,
row: &R,
) -> Result<Vec<u8>, SstBlockCodecError> {
if self.row_count > 0 && row_key < self.previous_key.as_str() {
return Err(SstBlockCodecError::RowKeysOutOfOrder {
previous: self.previous_key.clone(),
offered: row_key.to_owned(),
});
}
if self.row_count == 0 {
self.min_row_key = row_key.to_owned();
}
self.filter_hashes.push(filter_key_hashes(filter_key));
let restart = self.entry_count % RESTART_INTERVAL == 0;
if restart {
self.restarts.push(self.entries.len() as u32);
}
let shared_len = if restart {
0
} else {
shared_prefix_len(&self.previous_key, row_key)
};
let suffix = &row_key.as_bytes()[shared_len..];
let mut row_bytes = Vec::new();
ciborium::ser::into_writer(row, &mut row_bytes)
.map_err(|error| SstBlockCodecError::Codec(error.to_string()))?;
let previous_len = self.entries.len();
write_varint(&mut self.entries, shared_len as u64);
write_varint(&mut self.entries, suffix.len() as u64);
self.entries.extend_from_slice(suffix);
write_varint(&mut self.entries, row_bytes.len() as u64);
self.entries.extend_from_slice(&row_bytes);
self.decoded_data_bytes += self.entries.len() - previous_len;
self.entry_count += 1;
self.row_count += 1;
self.previous_key.clear();
self.previous_key.push_str(row_key);
if self.entries.len() >= self.target_block_bytes {
self.finish_data_block()?;
}
Ok(row_bytes)
}
fn finish_data_block(&mut self) -> Result<(), SstBlockCodecError> {
if self.entries.is_empty() {
return Ok(());
}
let mut payload = std::mem::take(&mut self.entries);
for restart in &self.restarts {
payload.extend_from_slice(&restart.to_le_bytes());
}
payload.extend_from_slice(&(self.restarts.len() as u32).to_le_bytes());
self.decoded_data_bytes += (self.restarts.len() + 1) * 4;
self.restarts.clear();
self.entry_count = 0;
let block = append_section(&mut self.bytes, &payload, true)?;
self.index.push(SegmentIndexEntry {
last_row_key: self.previous_key.clone(),
block,
});
Ok(())
}
pub fn row_count(&self) -> u64 {
self.row_count
}
pub fn decoded_data_bytes(&self) -> usize {
self.decoded_data_bytes
+ if self.entries.is_empty() {
0
} else {
(self.restarts.len() + 1) * 4
}
}
pub fn finish(mut self) -> Result<BuiltSegmentBlocks, SstBlockCodecError> {
if self.row_count == 0 {
return Err(SstBlockCodecError::EmptySegment);
}
self.finish_data_block()?;
let mut bytes = self.bytes;
let index = self.index;
let filter_payload = build_filter_payload(&self.filter_hashes);
let filter = append_section(&mut bytes, &filter_payload, false)?;
let mut index_payload = Vec::new();
ciborium::ser::into_writer(&index, &mut index_payload)
.map_err(|error| SstBlockCodecError::Codec(error.to_string()))?;
let index = append_section(&mut bytes, &index_payload, true)?;
Ok(BuiltSegmentBlocks {
bytes,
index,
filter,
row_count: self.row_count,
min_row_key: self.min_row_key,
max_row_key: self.previous_key,
})
}
}
pub fn decode_index_block(
stored: &[u8],
handle: &BlockHandle,
) -> Result<Vec<SegmentIndexEntry>, SstBlockCodecError> {
let payload = decode_section(stored, handle, true)?;
let entries: Vec<SegmentIndexEntry> = ciborium::de::from_reader(payload.as_slice())
.map_err(|error| SstBlockCodecError::Codec(error.to_string()))?;
let mut previous: Option<(&String, u64)> = None;
for entry in &entries {
let end = entry
.block
.offset
.checked_add(u64::from(entry.block.stored_len))
.ok_or_else(|| {
SstBlockCodecError::Malformed(format!(
"index block `{}` byte range overflows",
entry.last_row_key
))
})?;
if let Some((previous_key, previous_end)) = previous {
if previous_key > &entry.last_row_key {
return Err(SstBlockCodecError::Malformed(format!(
"index blocks out of key order: `{}` follows `{previous_key}`",
entry.last_row_key
)));
}
if entry.block.offset != previous_end {
return Err(SstBlockCodecError::Malformed(format!(
"index block `{}` does not start where `{previous_key}` ends",
entry.last_row_key
)));
}
}
previous = Some((&entry.last_row_key, end));
}
Ok(entries)
}
pub fn decode_data_block(
stored: &[u8],
handle: &BlockHandle,
) -> Result<DecodedDataBlock, SstBlockCodecError> {
decode_data_block_rows::<MetadataRow>(stored, handle)
}
pub fn decode_data_block_rows<R: serde::de::DeserializeOwned>(
stored: &[u8],
handle: &BlockHandle,
) -> Result<DecodedDataBlock<R>, SstBlockCodecError> {
let payload = decode_section(stored, handle, true)?;
if payload.len() < 4 {
return Err(SstBlockCodecError::Malformed(
"data block shorter than its restart count".to_owned(),
));
}
let (body, restart_count_bytes) = payload.split_at(payload.len() - 4);
let restart_count = u32::from_le_bytes(
restart_count_bytes
.try_into()
.expect("split_at should leave exactly four bytes"),
) as usize;
let restarts_len = restart_count
.checked_mul(4)
.filter(|len| *len <= body.len())
.ok_or_else(|| SstBlockCodecError::Malformed("restart array exceeds block".to_owned()))?;
let (entries, restarts) = body.split_at(body.len() - restarts_len);
let mut restarts = restarts.chunks_exact(4);
let mut row_keys = Vec::new();
let mut rows = Vec::new();
let mut cursor = 0usize;
let mut previous_key = String::new();
while cursor < entries.len() {
let restart = rows.len() % RESTART_INTERVAL == 0;
if restart {
let offset = restarts.next().map(|bytes| {
u32::from_le_bytes(
bytes
.try_into()
.expect("restart offsets should have four bytes"),
)
});
if offset.map(u64::from) != Some(cursor as u64) {
return Err(SstBlockCodecError::Malformed(
"restart offset disagrees with its entry".to_owned(),
));
}
}
let shared_len = read_entry_length(entries, &mut cursor)?;
let suffix_len = read_entry_length(entries, &mut cursor)?;
if restart && shared_len != 0 {
return Err(SstBlockCodecError::Malformed(
"restart entry has a shared prefix".to_owned(),
));
}
if shared_len > previous_key.len() || !previous_key.is_char_boundary(shared_len) {
return Err(SstBlockCodecError::Malformed(
"shared prefix exceeds previous key".to_owned(),
));
}
let suffix = take_slice(entries, &mut cursor, suffix_len)?;
let suffix = std::str::from_utf8(suffix)
.map_err(|_| SstBlockCodecError::Malformed("row key is not utf-8".to_owned()))?;
let key_len = shared_len.checked_add(suffix.len()).ok_or_else(|| {
SstBlockCodecError::Malformed("row key length exceeds address space".to_owned())
})?;
let mut key = String::with_capacity(key_len);
key.push_str(&previous_key[..shared_len]);
key.push_str(suffix);
let row_len = read_entry_length(entries, &mut cursor)?;
let row_bytes = take_slice(entries, &mut cursor, row_len)?;
let row: R = ciborium::de::from_reader(row_bytes)
.map_err(|error| SstBlockCodecError::Codec(error.to_string()))?;
if key.as_str() < previous_key.as_str() {
return Err(SstBlockCodecError::Malformed(format!(
"rows out of row-key order: `{key}` follows `{previous_key}`"
)));
}
previous_key.clear();
previous_key.push_str(&key);
row_keys.push(key);
rows.push(row);
}
if restarts.next().is_some() {
return Err(SstBlockCodecError::Malformed(
"restart array contains unused offsets".to_owned(),
));
}
Ok(DecodedDataBlock { row_keys, rows })
}
pub fn decode_filter_block(
stored: &[u8],
handle: &BlockHandle,
) -> Result<SegmentFilter, SstBlockCodecError> {
let payload = decode_section(stored, handle, false)?;
if payload.len() < 12 {
return Err(SstBlockCodecError::Malformed(
"filter block shorter than its header".to_owned(),
));
}
let n_hashes = u32::from_le_bytes(
payload[0..4]
.try_into()
.expect("header length should be checked above"),
);
let bit_len = u64::from_le_bytes(
payload[4..12]
.try_into()
.expect("header length should be checked above"),
);
if n_hashes != FILTER_HASH_COUNT {
return Err(SstBlockCodecError::Malformed(
"filter hash count must be seven".to_owned(),
));
}
if bit_len < 64 {
return Err(SstBlockCodecError::Malformed(
"filter bit length must be at least 64".to_owned(),
));
}
let bits = &payload[12..];
if bit_len.div_ceil(8) != bits.len() as u64 {
return Err(SstBlockCodecError::Malformed(
"filter bit length disagrees with its bytes".to_owned(),
));
}
Ok(SegmentFilter {
n_hashes,
bit_len,
bits: bits.to_vec(),
})
}
impl SegmentFilter {
pub fn may_contain(&self, filter_key: &str) -> bool {
let (h1, h2) = filter_key_hashes(filter_key);
for probe in 0..u64::from(self.n_hashes) {
let bit = h1.wrapping_add(probe.wrapping_mul(h2)) % self.bit_len;
let byte = self.bits[(bit / 8) as usize];
if byte & (1 << (bit % 8)) == 0 {
return false;
}
}
true
}
}
pub fn string_prefix_upper_bound(prefix: &str) -> Option<String> {
let mut bytes = prefix.as_bytes().to_vec();
for index in (0..bytes.len()).rev() {
if bytes[index] != u8::MAX {
bytes[index] += 1;
bytes.truncate(index + 1);
return String::from_utf8(bytes).ok();
}
}
None
}
pub fn index_blocks_for_key_range(
index: &[SegmentIndexEntry],
lower_bound: &str,
upper_bound: Option<&str>,
) -> std::ops::Range<usize> {
let start = index.partition_point(|entry| entry.last_row_key.as_str() < lower_bound);
let end = upper_bound.map_or(index.len(), |upper_bound| {
index
.partition_point(|entry| entry.last_row_key.as_str() < upper_bound)
.saturating_add(1)
.min(index.len())
});
start..end.max(start)
}
fn append_section(
bytes: &mut Vec<u8>,
payload: &[u8],
compress: bool,
) -> Result<BlockHandle, SstBlockCodecError> {
let stored = if compress {
zstd::bulk::compress(payload, ZSTD_LEVEL)
.map_err(|error| SstBlockCodecError::Codec(error.to_string()))?
} else {
payload.to_vec()
};
let handle = BlockHandle {
offset: bytes.len() as u64,
stored_len: stored.len() as u32,
decoded_len: payload.len() as u32,
crc32c: crc32c::crc32c(&stored),
};
bytes.extend_from_slice(&stored);
Ok(handle)
}
fn decode_section(
stored: &[u8],
handle: &BlockHandle,
compressed: bool,
) -> Result<Vec<u8>, SstBlockCodecError> {
if stored.len() != handle.stored_len as usize {
return Err(SstBlockCodecError::StoredLengthMismatch {
expected: handle.stored_len,
actual: stored.len(),
});
}
let actual = crc32c::crc32c(stored);
if actual != handle.crc32c {
return Err(SstBlockCodecError::ChecksumMismatch {
expected: handle.crc32c,
actual,
});
}
let payload = if compressed {
let mut payload =
Vec::with_capacity((handle.decoded_len as usize).min(DEFAULT_TARGET_BLOCK_BYTES));
zstd::Decoder::new(stored)
.and_then(|decoder| {
decoder
.take(u64::from(handle.decoded_len) + 1)
.read_to_end(&mut payload)
})
.map_err(|error| SstBlockCodecError::Codec(error.to_string()))?;
payload
} else {
if stored.len() != handle.decoded_len as usize {
return Err(SstBlockCodecError::DecodedLengthMismatch {
expected: handle.decoded_len,
actual: stored.len(),
});
}
stored.to_vec()
};
if payload.len() != handle.decoded_len as usize {
return Err(SstBlockCodecError::DecodedLengthMismatch {
expected: handle.decoded_len,
actual: payload.len(),
});
}
Ok(payload)
}
fn build_filter_payload(hashes: &[(u64, u64)]) -> Vec<u8> {
let bit_len = (hashes.len() * FILTER_BITS_PER_KEY).max(64) as u64;
let mut bits = vec![0u8; bit_len.div_ceil(8) as usize];
for (h1, h2) in hashes {
for probe in 0..u64::from(FILTER_HASH_COUNT) {
let bit = h1.wrapping_add(probe.wrapping_mul(*h2)) % bit_len;
bits[(bit / 8) as usize] |= 1 << (bit % 8);
}
}
let mut payload = Vec::with_capacity(12 + bits.len());
payload.extend_from_slice(&FILTER_HASH_COUNT.to_le_bytes());
payload.extend_from_slice(&bit_len.to_le_bytes());
payload.extend_from_slice(&bits);
payload
}
fn filter_key_hashes(filter_key: &str) -> (u64, u64) {
(
xxh64(filter_key.as_bytes(), FILTER_HASH_SEED_ONE),
xxh64(filter_key.as_bytes(), FILTER_HASH_SEED_TWO),
)
}
fn shared_prefix_len(previous: &str, current: &str) -> usize {
let mut len = previous
.as_bytes()
.iter()
.zip(current.as_bytes())
.take_while(|(a, b)| a == b)
.count();
while !current.is_char_boundary(len) {
len -= 1;
}
len
}
pub fn write_varint(bytes: &mut Vec<u8>, mut value: u64) {
loop {
let byte = (value & 0x7f) as u8;
value >>= 7;
if value == 0 {
bytes.push(byte);
return;
}
bytes.push(byte | 0x80);
}
}
pub fn read_varint(bytes: &[u8], cursor: &mut usize) -> Result<u64, SstBlockCodecError> {
let mut value = 0u64;
let mut shift = 0u32;
loop {
let byte = *bytes.get(*cursor).ok_or_else(|| {
SstBlockCodecError::Malformed("varint runs past the block".to_owned())
})?;
*cursor += 1;
if shift == 63 && byte > 1 {
return Err(SstBlockCodecError::Malformed(
"varint exceeds 64 bits".to_owned(),
));
}
value |= u64::from(byte & 0x7f) << shift;
if byte & 0x80 == 0 {
return Ok(value);
}
shift += 7;
}
}
fn read_entry_length(bytes: &[u8], cursor: &mut usize) -> Result<usize, SstBlockCodecError> {
usize::try_from(read_varint(bytes, cursor)?)
.map_err(|_| SstBlockCodecError::Malformed("entry length exceeds address space".to_owned()))
}
fn take_slice<'a>(
bytes: &'a [u8],
cursor: &mut usize,
len: usize,
) -> Result<&'a [u8], SstBlockCodecError> {
let end = cursor.checked_add(len).filter(|end| *end <= bytes.len());
match end {
Some(end) => {
let slice = &bytes[*cursor..end];
*cursor = end;
Ok(slice)
}
None => Err(SstBlockCodecError::Malformed(
"entry runs past the block".to_owned(),
)),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{ChangeSeq, InodeId, InodeKind};
#[test]
fn decoded_byte_accounting_includes_closed_and_open_blocks() {
let mut builder =
SegmentBlocksBuilder::new(NonZeroUsize::new(512).expect("nonzero block target"));
for index in 0..17 {
let (key, filter, row) = inode_row(index);
builder.push(&key, &filter, &row).expect("encode row");
}
let decoded_bytes = builder.decoded_data_bytes();
assert_eq!(builder.row_count(), 17);
assert!(
!builder.index.is_empty(),
"completed blocks are already compressed"
);
let built = builder.finish().expect("finish segment");
let start = built.index.offset as usize;
let end = start + built.index.stored_len as usize;
let index = decode_index_block(&built.bytes[start..end], &built.index).expect("index");
assert_eq!(
decoded_bytes,
index
.iter()
.map(|entry| entry.block.decoded_len as usize)
.sum::<usize>(),
);
}
fn inode_row(inode_id: u64) -> (String, String, MetadataRow) {
let row = MetadataRow::Inode(crate::manifest::InodeRecord {
inode_id: InodeId(inode_id),
inode_kind: InodeKind::File,
created_seq: ChangeSeq(inode_id),
commit_id: crate::CommitId::parse(format!("c_row_{inode_id}"))
.expect("valid commit id"),
created_by: crate::ActorId::loonfs(),
created_at_ms: inode_id,
});
let key = row.row_key();
(key.clone(), key, row)
}
fn build_segment(rows: usize) -> BuiltSegmentBlocks {
let mut builder = SegmentBlocksBuilder::default();
for index in 0..rows {
let (key, filter_key, row) = inode_row(index as u64);
builder.push(&key, &filter_key, &row).expect("push row");
}
builder.finish().expect("finish segment")
}
const SHARED_SEGMENT_ROWS: usize = 5_000;
fn shared_segment() -> &'static BuiltSegmentBlocks {
static SEGMENT: std::sync::OnceLock<BuiltSegmentBlocks> = std::sync::OnceLock::new();
SEGMENT.get_or_init(|| build_segment(SHARED_SEGMENT_ROWS))
}
fn section<'a>(bytes: &'a [u8], handle: &BlockHandle) -> &'a [u8] {
&bytes[handle.offset as usize..handle.offset as usize + handle.stored_len as usize]
}
fn encode_index(entries: &[SegmentIndexEntry]) -> (Vec<u8>, BlockHandle) {
let mut payload = Vec::new();
ciborium::ser::into_writer(entries, &mut payload).expect("encode index");
let mut bytes = Vec::new();
let handle = append_section(&mut bytes, &payload, true).expect("append section");
(bytes, handle)
}
fn index_entry(last_row_key: &str, offset: u64, stored_len: u32) -> SegmentIndexEntry {
SegmentIndexEntry {
last_row_key: last_row_key.to_owned(),
block: BlockHandle {
offset,
stored_len,
decoded_len: stored_len,
crc32c: 0,
},
}
}
#[test]
fn segment_round_trips_every_row_through_index_and_blocks() {
let rows = SHARED_SEGMENT_ROWS;
let built = shared_segment();
let index =
decode_index_block(section(&built.bytes, &built.index), &built.index).expect("index");
assert!(index.len() > 1, "5k inode rows should span several blocks");
let mut recovered = Vec::new();
for entry in &index {
let block = decode_data_block(section(&built.bytes, &entry.block), &entry.block)
.expect("data block");
assert_eq!(block.row_keys.len(), block.rows.len());
assert_eq!(
block.row_keys.last().expect("blocks are never empty"),
&entry.last_row_key
);
recovered.extend(block.row_keys.iter().cloned());
}
let expected: Vec<String> = (0..rows).map(|i| inode_row(i as u64).0).collect();
assert_eq!(recovered, expected);
assert_eq!(built.row_count, rows as u64);
assert_eq!(built.min_row_key, expected[0]);
assert_eq!(&built.max_row_key, expected.last().expect("rows"));
}
#[test]
fn index_narrows_point_lookups_to_one_block() {
let built = shared_segment();
let index =
decode_index_block(section(&built.bytes, &built.index), &built.index).expect("index");
let (key, _, row) = inode_row(3_217);
let upper = format!("{key}\0");
let range = index_blocks_for_key_range(&index, &key, Some(&upper));
assert_eq!(range.len(), 1, "a point lookup should touch one block");
let entry = &index[range.start];
let block =
decode_data_block(section(&built.bytes, &entry.block), &entry.block).expect("block");
let position = block
.row_keys
.binary_search_by(|candidate| candidate.as_str().cmp(key.as_str()))
.expect("row should be present");
assert_eq!(block.rows[position], row);
}
#[test]
fn key_range_scan_covers_exactly_the_matching_blocks() {
let built = shared_segment();
let index =
decode_index_block(section(&built.bytes, &built.index), &built.index).expect("index");
let lower = inode_row(1_000).0;
let upper = inode_row(1_500).0;
let range = index_blocks_for_key_range(&index, &lower, Some(&upper));
let mut keys = Vec::new();
for entry in &index[range] {
let block = decode_data_block(section(&built.bytes, &entry.block), &entry.block)
.expect("block");
keys.extend(block.row_keys);
}
let keys: Vec<&String> = keys
.iter()
.filter(|key| key.as_str() >= lower.as_str() && key.as_str() < upper.as_str())
.collect();
assert_eq!(keys.len(), 500);
}
#[test]
fn out_of_order_and_empty_segments_are_rejected() {
let mut builder = SegmentBlocksBuilder::default();
let (key_b, filter_b, row_b) = inode_row(2);
let (key_a, filter_a, row_a) = inode_row(1);
builder.push(&key_b, &filter_b, &row_b).expect("first row");
let error = builder
.push(&key_a, &filter_a, &row_a)
.expect_err("descending key should be rejected");
assert!(matches!(
error,
SstBlockCodecError::RowKeysOutOfOrder { .. }
));
let error = SegmentBlocksBuilder::default()
.finish()
.expect_err("empty segment should be rejected");
assert!(matches!(error, SstBlockCodecError::EmptySegment));
}
#[test]
fn max_key_survives_a_last_row_that_closes_its_block() {
let rows = 100usize;
let single_block = build_segment(rows);
let calibration = decode_index_block(
section(&single_block.bytes, &single_block.index),
&single_block.index,
)
.expect("index");
assert_eq!(calibration.len(), 1, "the calibration segment is one block");
let restarts = rows.div_ceil(RESTART_INTERVAL);
let entry_bytes = calibration[0].block.decoded_len as usize - 4 * restarts - 4;
let mut builder =
SegmentBlocksBuilder::new(NonZeroUsize::new(entry_bytes).expect("positive target"));
for index in 0..rows {
let (key, filter_key, row) = inode_row(index as u64);
builder.push(&key, &filter_key, &row).expect("push row");
}
let built = builder.finish().expect("finish segment");
let expected_max = inode_row((rows - 1) as u64).0;
assert_eq!(built.min_row_key, inode_row(0).0);
assert_eq!(built.max_row_key, expected_max);
assert_eq!(built.row_count, rows as u64);
assert_eq!(built.bytes, single_block.bytes);
let index =
decode_index_block(section(&built.bytes, &built.index), &built.index).expect("index");
assert_eq!(index.len(), 1);
assert_eq!(index[0].last_row_key, expected_max);
}
#[test]
fn block_geometry_does_not_change_the_segment_key_range() {
let rows = 400usize;
let expected: Vec<String> = (0..rows).map(|index| inode_row(index as u64).0).collect();
for target in [1usize, 64, 257, 1_024, 4_096, 65_536] {
let mut builder =
SegmentBlocksBuilder::new(NonZeroUsize::new(target).expect("positive target"));
for index in 0..rows {
let (key, filter_key, row) = inode_row(index as u64);
builder.push(&key, &filter_key, &row).expect("push row");
}
let built = builder.finish().expect("finish segment");
assert_eq!(built.min_row_key, expected[0], "target {target}");
assert_eq!(
&built.max_row_key,
expected.last().expect("rows"),
"target {target}"
);
assert_eq!(built.row_count, rows as u64, "target {target}");
let index = decode_index_block(section(&built.bytes, &built.index), &built.index)
.expect("index");
let mut recovered = Vec::new();
for entry in &index {
let block = decode_data_block(section(&built.bytes, &entry.block), &entry.block)
.expect("data block");
assert_eq!(
block.row_keys.last().expect("blocks are never empty"),
&entry.last_row_key,
"target {target}"
);
recovered.extend(block.row_keys);
}
assert_eq!(recovered, expected, "target {target}");
assert_eq!(
index.last().expect("blocks").last_row_key,
built.max_row_key,
"target {target}"
);
}
}
#[test]
fn a_descending_row_after_a_block_boundary_is_rejected() {
let mut builder = SegmentBlocksBuilder::new(NonZeroUsize::MIN);
let (key_high, filter_high, row_high) = inode_row(9);
builder
.push(&key_high, &filter_high, &row_high)
.expect("first row");
let (key_low, filter_low, row_low) = inode_row(3);
let error = builder
.push(&key_low, &filter_low, &row_low)
.expect_err("a descending key across a block boundary should be rejected");
assert!(
matches!(
&error,
SstBlockCodecError::RowKeysOutOfOrder { previous, offered }
if previous == &key_high && offered == &key_low
),
"unexpected error: {error}"
);
}
#[test]
fn adjacent_equal_keys_are_permitted() {
let mut builder = SegmentBlocksBuilder::default();
let (key, filter_key, row) = inode_row(7);
builder.push(&key, &filter_key, &row).expect("first copy");
builder.push(&key, &filter_key, &row).expect("second copy");
let built = builder.finish().expect("finish");
assert_eq!(built.row_count, 2);
}
#[test]
fn corrupted_sections_fail_their_checksums() {
let built = build_segment(200);
let index =
decode_index_block(section(&built.bytes, &built.index), &built.index).expect("index");
let mut corrupted = built.bytes.clone();
let target = index[0].block.offset as usize + 3;
corrupted[target] ^= 0xff;
let error = decode_data_block(section(&corrupted, &index[0].block), &index[0].block)
.expect_err("corrupted data block should fail");
assert!(matches!(error, SstBlockCodecError::ChecksumMismatch { .. }));
let mut corrupted = built.bytes.clone();
let target = built.index.offset as usize + 3;
corrupted[target] ^= 0xff;
let error = decode_index_block(section(&corrupted, &built.index), &built.index)
.expect_err("corrupted index should fail");
assert!(matches!(error, SstBlockCodecError::ChecksumMismatch { .. }));
let mut corrupted = built.bytes.clone();
let target = built.filter.offset as usize + 12;
corrupted[target] ^= 0xff;
let error = decode_filter_block(section(&corrupted, &built.filter), &built.filter)
.expect_err("corrupted filter should fail");
assert!(matches!(error, SstBlockCodecError::ChecksumMismatch { .. }));
}
#[test]
fn block_expansion_stops_after_the_first_byte_beyond_the_declared_length() {
let mut bytes = Vec::new();
let mut handle = append_section(&mut bytes, &[0; 1024], true).expect("compress block");
handle.decoded_len = 8;
assert_eq!(
decode_data_block(&bytes, &handle).expect_err("expansion exceeds the handle"),
SstBlockCodecError::DecodedLengthMismatch {
expected: 8,
actual: 9,
}
);
}
#[test]
fn a_large_declared_length_does_not_require_a_large_initial_allocation() {
let mut bytes = Vec::new();
let mut handle = append_section(&mut bytes, &[0; 4], true).expect("compress block");
handle.decoded_len = u32::MAX;
assert_eq!(
decode_index_block(&bytes, &handle).expect_err("declared length exceeds expansion"),
SstBlockCodecError::DecodedLengthMismatch {
expected: u32::MAX,
actual: 4,
}
);
}
#[test]
fn truncated_index_and_filter_sections_are_rejected_before_decoding() {
let built = shared_segment();
let index = section(&built.bytes, &built.index);
let filter = section(&built.bytes, &built.filter);
for (handle, error) in [
(
built.index,
decode_index_block(&index[..index.len() - 1], &built.index)
.expect_err("truncated index"),
),
(
built.filter,
decode_filter_block(&filter[..filter.len() - 1], &built.filter)
.expect_err("truncated filter"),
),
] {
assert_eq!(
error,
SstBlockCodecError::StoredLengthMismatch {
expected: handle.stored_len,
actual: handle.stored_len as usize - 1,
}
);
}
}
#[test]
fn filter_headers_reject_zero_bits_and_invalid_hash_counts() {
for (n_hashes, bit_len) in [(FILTER_HASH_COUNT, 0u64), (u32::MAX, 64), (0, 64)] {
let mut payload = Vec::new();
payload.extend_from_slice(&n_hashes.to_le_bytes());
payload.extend_from_slice(&bit_len.to_le_bytes());
payload.resize(12 + bit_len.div_ceil(8) as usize, 0xff);
let mut bytes = Vec::new();
let handle = append_section(&mut bytes, &payload, false).expect("filter section");
assert!(matches!(
decode_filter_block(&bytes, &handle),
Err(SstBlockCodecError::Malformed(_))
));
}
}
#[test]
fn restart_entries_bound_shared_prefix_reconstruction() {
let mut payload = Vec::new();
let mut restart_offsets = Vec::new();
for index in 0..=RESTART_INTERVAL {
if index % RESTART_INTERVAL == 0 {
restart_offsets.push(payload.len() as u32);
}
write_varint(&mut payload, u64::from(index != 0));
write_varint(&mut payload, u64::from(index == 0));
if index == 0 {
payload.push(b'a');
}
write_varint(&mut payload, 1);
payload.push(0);
}
for offset in restart_offsets {
payload.extend_from_slice(&offset.to_le_bytes());
}
payload.extend_from_slice(&2u32.to_le_bytes());
let mut bytes = Vec::new();
let handle = append_section(&mut bytes, &payload, true).expect("data section");
assert_eq!(
decode_data_block_rows::<u8>(&bytes, &handle).expect_err("shared restart prefix"),
SstBlockCodecError::Malformed("restart entry has a shared prefix".to_owned())
);
}
#[test]
fn restart_offsets_must_name_their_entries() {
let payload = [0, 1, b'a', 1, 0, 1, 0, 0, 0, 1, 0, 0, 0];
let mut bytes = Vec::new();
let handle = append_section(&mut bytes, &payload, true).expect("data section");
assert_eq!(
decode_data_block_rows::<u8>(&bytes, &handle).expect_err("wrong restart offset"),
SstBlockCodecError::Malformed("restart offset disagrees with its entry".to_owned())
);
}
#[test]
fn varints_reject_overflow_in_the_tenth_byte() {
let mut bytes = [0xff; 10];
for last in [2, 0x80] {
bytes[9] = last;
assert_eq!(
read_varint(&bytes, &mut 0).expect_err("overflowing varint"),
SstBlockCodecError::Malformed("varint exceeds 64 bits".to_owned())
);
}
bytes[9] = 1;
assert_eq!(
read_varint(&bytes, &mut 0).expect("largest varint"),
u64::MAX
);
}
#[test]
fn filter_has_no_false_negatives_and_few_false_positives() {
let rows = 2_000;
let built = build_segment(rows);
let filter = decode_filter_block(section(&built.bytes, &built.filter), &built.filter)
.expect("filter");
for index in 0..rows {
let (key, _, _) = inode_row(index as u64);
assert!(filter.may_contain(&key), "inserted key must stay positive");
}
let mut false_positives = 0usize;
let probes = 10_000usize;
for index in 0..probes {
let (absent, _, _) = inode_row((rows + 10_000 + index) as u64);
if filter.may_contain(&absent) {
false_positives += 1;
}
}
let rate = false_positives as f64 / probes as f64;
assert!(rate < 0.02, "false positive rate {rate} exceeds 2%");
}
#[test]
fn decoding_rejects_out_of_order_rows_in_a_block() {
let mut entries = Vec::new();
for inode in [9u64, 3u64] {
let (key, _, row) = inode_row(inode);
let mut row_bytes = Vec::new();
ciborium::ser::into_writer(&row, &mut row_bytes).expect("encode row");
write_varint(&mut entries, 0);
write_varint(&mut entries, key.len() as u64);
entries.extend_from_slice(key.as_bytes());
write_varint(&mut entries, row_bytes.len() as u64);
entries.extend_from_slice(&row_bytes);
}
let mut payload = entries;
payload.extend_from_slice(&0u32.to_le_bytes());
payload.extend_from_slice(&1u32.to_le_bytes());
let mut bytes = Vec::new();
let handle = append_section(&mut bytes, &payload, true).expect("append section");
let error =
decode_data_block(&bytes, &handle).expect_err("descending rows should be rejected");
assert!(
matches!(&error, SstBlockCodecError::Malformed(message) if message.contains("row-key order")),
"unexpected error: {error}"
);
}
#[test]
fn decoding_rejects_a_shared_prefix_inside_a_utf8_code_point() {
let (_, _, row) = inode_row(1);
let mut row_bytes = Vec::new();
ciborium::ser::into_writer(&row, &mut row_bytes).expect("encode row");
let mut payload = Vec::new();
write_varint(&mut payload, 0);
write_varint(&mut payload, "é".len() as u64);
payload.extend_from_slice("é".as_bytes());
write_varint(&mut payload, row_bytes.len() as u64);
payload.extend_from_slice(&row_bytes);
write_varint(&mut payload, 1);
write_varint(&mut payload, 0);
write_varint(&mut payload, row_bytes.len() as u64);
payload.extend_from_slice(&row_bytes);
payload.extend_from_slice(&0u32.to_le_bytes());
payload.extend_from_slice(&1u32.to_le_bytes());
let mut bytes = Vec::new();
let handle = append_section(&mut bytes, &payload, true).expect("append section");
let error = decode_data_block(&bytes, &handle)
.expect_err("a partial utf-8 prefix should be rejected");
assert!(matches!(
&error,
SstBlockCodecError::Malformed(message)
if message == "shared prefix exceeds previous key"
));
}
#[test]
fn decoding_rejects_out_of_order_index_entries() {
let entries = vec![
index_entry("inode-00000000000000000009", 0, 1),
index_entry("inode-00000000000000000003", 0, 1),
];
let (bytes, handle) = encode_index(&entries);
let error =
decode_index_block(&bytes, &handle).expect_err("descending index should be rejected");
assert!(
matches!(&error, SstBlockCodecError::Malformed(message) if message.contains("key order")),
"unexpected error: {error}"
);
}
#[test]
fn decoding_rejects_out_of_order_index_offsets() {
let entries = [index_entry("a", 10, 1), index_entry("b", 5, 1)];
let (bytes, handle) = encode_index(&entries);
let error = decode_index_block(&bytes, &handle)
.expect_err("descending block offsets should be rejected");
assert!(
matches!(&error, SstBlockCodecError::Malformed(message) if message.contains("does not start where")),
"unexpected error: {error}"
);
}
#[test]
fn decoding_rejects_overlapping_index_ranges() {
let entries = [index_entry("a", 10, 5), index_entry("b", 14, 1)];
let (bytes, handle) = encode_index(&entries);
let error = decode_index_block(&bytes, &handle)
.expect_err("overlapping block ranges should be rejected");
assert!(
matches!(&error, SstBlockCodecError::Malformed(message) if message.contains("does not start where")),
"unexpected error: {error}"
);
}
#[test]
fn decoding_rejects_a_gap_between_index_blocks() {
let entries = [index_entry("a", 0, 10), index_entry("b", 20, 1)];
let (bytes, handle) = encode_index(&entries);
let error = decode_index_block(&bytes, &handle)
.expect_err("a gap between block ranges should be rejected");
assert!(
matches!(&error, SstBlockCodecError::Malformed(message) if message.contains("does not start where")),
"unexpected error: {error}"
);
}
#[test]
fn decoding_rejects_a_single_block_range_past_the_integer_edge() {
let entries = [index_entry("a", u64::MAX - 10, 100)];
let (bytes, handle) = encode_index(&entries);
let error = decode_index_block(&bytes, &handle)
.expect_err("an overflowing single range should be rejected");
assert!(
matches!(&error, SstBlockCodecError::Malformed(message) if message.contains("byte range overflows")),
"unexpected error: {error}"
);
}
#[test]
fn decoding_rejects_a_final_block_range_past_the_integer_edge() {
let entries = [
index_entry("a", u64::MAX - 110, 100),
index_entry("b", u64::MAX - 10, 100),
];
let (bytes, handle) = encode_index(&entries);
let error = decode_index_block(&bytes, &handle)
.expect_err("a trailing overflowing range should be rejected");
assert!(
matches!(&error, SstBlockCodecError::Malformed(message) if message.contains("byte range overflows")),
"unexpected error: {error}"
);
}
}