1use std::{
2 cmp::Ordering,
3 hash::BuildHasherDefault,
4 path::Path,
5 rc::Rc,
6 sync::{
7 Arc,
8 atomic::{AtomicU64, Ordering as AtomicOrdering},
9 },
10};
11
12use anyhow::{Context, Result, bail, ensure};
13use fs_err::File;
14use memmap2::Mmap;
15use quick_cache::{Lifecycle, sync::GuardResult};
16use rustc_hash::FxHasher;
17use smallvec::SmallVec;
18
19use crate::{
20 QueryKey,
21 arc_bytes::ArcBytes,
22 be,
23 compression::checksum_block,
24 constants::MAX_INLINE_VALUE_SIZE,
25 lookup_entry::{IterValue, LookupEntry, LookupValue},
26 mmap_helper::advise_mmap_for_persistence,
27 rc_bytes::RcBytes,
28 shared_bytes::SharedBytes,
29 static_sorted_file_builder::{BLOCK_HEADER_SIZE, INDEX_BLOCK_ENTRY_SIZE},
30};
31
32pub const BLOCK_TYPE_INDEX: u8 = 0;
34pub const BLOCK_TYPE_KEY_WITH_HASH: u8 = 1;
36pub const BLOCK_TYPE_KEY_NO_HASH: u8 = 2;
38pub const BLOCK_TYPE_FIXED_KEY_WITH_HASH: u8 = 3;
40pub const BLOCK_TYPE_FIXED_KEY_NO_HASH: u8 = 4;
42
43pub const KEY_BLOCK_ENTRY_TYPE_SMALL: u8 = 0;
45pub const KEY_BLOCK_ENTRY_TYPE_BLOB: u8 = 1;
47pub const KEY_BLOCK_ENTRY_TYPE_DELETED: u8 = 2;
49pub const KEY_BLOCK_ENTRY_TYPE_MEDIUM: u8 = 3;
51pub const KEY_BLOCK_ENTRY_TYPE_INLINE_MIN: u8 = 8;
53
54pub(crate) const SMALL_VALUE_REF_SIZE: usize = 8;
56pub(crate) const MEDIUM_VALUE_REF_SIZE: usize = 2;
58pub(crate) const BLOB_VALUE_REF_SIZE: usize = 4;
60pub(crate) const DELETED_VALUE_REF_SIZE: usize = 0;
62
63const _: () = assert!(
66 MAX_INLINE_VALUE_SIZE <= (u8::MAX - KEY_BLOCK_ENTRY_TYPE_INLINE_MIN) as usize,
67 "MAX_INLINE_VALUE_SIZE exceeds what can be encoded in key type byte"
68);
69
70pub enum SstLookupResult {
72 Found(SmallVec<[LookupValue; 1]>),
74 NotFound,
76}
77
78impl From<LookupValue> for SstLookupResult {
79 fn from(value: LookupValue) -> Self {
80 SstLookupResult::Found(smallvec::smallvec![value])
81 }
82}
83
84#[derive(Clone, Default)]
85pub struct BlockWeighter;
86
87impl quick_cache::Weighter<(u32, u16), ArcBytes> for BlockWeighter {
88 fn weight(&self, _key: &(u32, u16), val: &ArcBytes) -> u64 {
89 if val.is_mmap_backed() {
90 debug_assert!(
93 !val.is_mmap_backed(),
94 "mmap-backed block should not be inserted into BlockCache"
95 );
96 64
97 } else {
98 val.len() as u64 + 8
99 }
100 }
101}
102
103#[derive(Clone, Default)]
106pub struct BlockCacheLifecycle;
107
108impl Lifecycle<(u32, u16), ArcBytes> for BlockCacheLifecycle {
109 type RequestState = ();
110
111 #[inline]
112 fn is_pinned(&self, _key: &(u32, u16), val: &ArcBytes) -> bool {
113 val.is_shared_arc()
114 }
115
116 #[inline]
117 fn begin_request(&self) -> Self::RequestState {}
118
119 #[inline]
120 fn on_evict(&self, _state: &mut Self::RequestState, _key: (u32, u16), _val: ArcBytes) {}
121}
122
123pub type BlockCache = quick_cache::sync::Cache<
124 (u32, u16),
125 ArcBytes,
126 BlockWeighter,
127 BuildHasherDefault<FxHasher>,
128 BlockCacheLifecycle,
129>;
130
131trait ValueBlockCache<B: SharedBytes> {
138 fn get_or_read(
139 self,
140 mmap: &B::MmapHandle,
141 meta: &StaticSortedFileMetaData,
142 block_index: u16,
143 ) -> Result<B>;
144}
145
146#[derive(Clone, Copy)]
149struct ArcBlockCacheReader<'a> {
150 cache: &'a BlockCache,
151 verified_blocks: &'a [AtomicU64],
152}
153
154impl ValueBlockCache<ArcBytes> for ArcBlockCacheReader<'_> {
157 fn get_or_read(
158 self,
159 mmap: &Arc<Mmap>,
160 meta: &StaticSortedFileMetaData,
161 block_index: u16,
162 ) -> Result<ArcBytes> {
163 get_or_cache_block(mmap, meta, block_index, self.cache, self.verified_blocks)
164 }
165}
166
167impl ValueBlockCache<RcBytes> for &mut Option<(u16, RcBytes)> {
169 fn get_or_read(
170 self,
171 mmap: &Rc<Mmap>,
172 meta: &StaticSortedFileMetaData,
173 block_index: u16,
174 ) -> Result<RcBytes> {
175 if let Some((idx, block)) = self.as_ref()
176 && *idx == block_index
177 {
178 return Ok(block.clone());
179 }
180 let block: RcBytes = read_block_generic(mmap, meta, block_index)?;
181 *self = Some((block_index, block.clone()));
182 Ok(block)
183 }
184}
185
186#[derive(Clone, Copy, Debug)]
187pub struct StaticSortedFileMetaData {
188 pub sequence_number: u32,
190 pub block_count: u16,
192}
193
194impl StaticSortedFileMetaData {
195 pub fn block_offsets_start(&self, sst_len: usize) -> usize {
196 let bc: usize = self.block_count.into();
197 sst_len - (bc * size_of::<u32>())
198 }
199}
200
201pub struct StaticSortedFile {
203 meta: StaticSortedFileMetaData,
205 mmap: Arc<Mmap>,
209 verified_blocks: Box<[AtomicU64]>,
214}
215
216impl StaticSortedFile {
217 pub fn open(db_path: &Path, meta: StaticSortedFileMetaData) -> Result<Self> {
220 let filename = format!("{:08}.sst", meta.sequence_number);
221 let path = db_path.join(&filename);
222 let file = File::open(&path)?;
223 let mmap = unsafe { Mmap::map(file.file()) }.with_context(|| {
224 format!(
225 "Failed to mmap SST file {} ({} bytes)",
226 path.display(),
227 file.metadata().map(|m| m.len()).unwrap_or(0)
228 )
229 })?;
230 #[cfg(unix)]
231 {
232 mmap.advise(memmap2::Advice::Random)?;
233 let offset = meta.block_offsets_start(mmap.len());
234 let _ = mmap.advise_range(memmap2::Advice::Sequential, offset, mmap.len() - offset);
235 }
236 advise_mmap_for_persistence(&mmap)?;
237 let bitmap_words = (meta.block_count as usize).div_ceil(u64::BITS as usize);
238 let verified_blocks = (0..bitmap_words)
239 .map(|_| AtomicU64::new(0))
240 .collect::<Box<[_]>>();
241 Ok(Self {
242 meta,
243 mmap: Arc::new(mmap),
244 verified_blocks,
245 })
246 }
247
248 pub fn lookup<K: QueryKey, const FIND_ALL: bool>(
254 &self,
255 key_hash: u64,
256 key: &K,
257 key_block_cache: &BlockCache,
258 value_block_cache: &BlockCache,
259 ) -> Result<SstLookupResult> {
260 let index_block_index = self.meta.block_count - 1;
263 let index_block = get_or_cache_block(
264 &self.mmap,
265 &self.meta,
266 index_block_index,
267 key_block_cache,
268 &self.verified_blocks,
269 )?;
270 let key_block_index = self.lookup_index_block(&index_block, key_hash)?;
271
272 let key_block_arc = get_or_cache_block(
273 &self.mmap,
274 &self.meta,
275 key_block_index,
276 key_block_cache,
277 &self.verified_blocks,
278 )?;
279 let reader = ArcBlockCacheReader {
280 cache: value_block_cache,
281 verified_blocks: &self.verified_blocks,
282 };
283 let block_type = be::read_u8(&key_block_arc);
284 match block_type {
285 BLOCK_TYPE_KEY_WITH_HASH | BLOCK_TYPE_KEY_NO_HASH => {
286 let has_hash = block_type == BLOCK_TYPE_KEY_WITH_HASH;
287 self.lookup_key_block::<K, FIND_ALL>(key_block_arc, key_hash, key, has_hash, reader)
288 }
289
290 BLOCK_TYPE_FIXED_KEY_WITH_HASH | BLOCK_TYPE_FIXED_KEY_NO_HASH => {
291 let has_hash = block_type == BLOCK_TYPE_FIXED_KEY_WITH_HASH;
292 self.lookup_fixed_key_block::<K, FIND_ALL>(
293 key_block_arc,
294 key_hash,
295 key,
296 has_hash,
297 reader,
298 )
299 }
300 _ => {
301 bail!("Invalid block type");
302 }
303 }
304 }
305
306 fn lookup_index_block(&self, block: &[u8], hash: u64) -> Result<u16> {
308 ensure!(block.len() >= 3, "index block too short");
309 debug_assert!(
310 be::read_u8(block) == BLOCK_TYPE_INDEX,
311 "expected index block as last block"
312 );
313 let first_block = be::read_u16(&block[1..]);
314 let (entries, remainder) = block[3..].as_chunks::<INDEX_BLOCK_ENTRY_SIZE>();
315 if entries.is_empty() {
316 return Ok(first_block);
317 }
318 if !remainder.is_empty() {
319 bail!("invalid index block, {} extra bytes", remainder.len())
320 }
321 match entries.binary_search_by(|entry| be::read_u64(entry).cmp(&hash)) {
322 Ok(i) => Ok(be::read_u16(&entries[i][8..])),
323 Err(0) => Ok(first_block),
324 Err(i) => Ok(be::read_u16(&entries[i - 1][8..])),
325 }
326 }
327
328 fn lookup_key_block<K: QueryKey, const FIND_ALL: bool>(
333 &self,
334 block: ArcBytes,
335 key_hash: u64,
336 key: &K,
337 has_hash: bool,
338 reader: ArcBlockCacheReader<'_>,
339 ) -> Result<SstLookupResult> {
340 let hash_len: u8 = if has_hash { 8 } else { 0 };
341 ensure!(block.len() >= 4, "key block too short");
342 let entry_count = be::read_u24(&block[1..]) as usize;
343 let data = &block[4..];
344 ensure!(
345 data.len() >= entry_count * 4,
346 "key block too short for {entry_count} entries"
347 );
348 let offsets = &data[..entry_count * 4];
349 let entries = &data[entry_count * 4..];
350
351 self.lookup_block_inner::<K, FIND_ALL>(&block, entry_count, key_hash, key, reader, |i| {
352 get_key_entry(offsets, entries, entry_count, i, hash_len)
353 })
354 }
355
356 fn lookup_fixed_key_block<K: QueryKey, const FIND_ALL: bool>(
361 &self,
362 block: ArcBytes,
363 key_hash: u64,
364 key: &K,
365 has_hash: bool,
366 reader: ArcBlockCacheReader<'_>,
367 ) -> Result<SstLookupResult> {
368 let hash_len: u8 = if has_hash { 8 } else { 0 };
369 ensure!(block.len() >= 6, "fixed key block too short");
370 let entry_count = be::read_u24(&block[1..]) as usize;
371 let key_size = be::read_u8(&block[4..]) as usize;
372 let value_type = be::read_u8(&block[5..]);
373 let val_size = entry_val_size(value_type)?;
374 let stride = hash_len as usize + key_size + val_size;
375 let entries = &block[6..];
376 ensure!(
377 entries.len() == entry_count * stride,
378 "fixed key block for {entry_count} entries must is the wrong size"
379 );
380
381 self.lookup_block_inner::<K, FIND_ALL>(&block, entry_count, key_hash, key, reader, |i| {
382 Ok(get_fixed_key_entry(
383 entries, i, hash_len, key_size, value_type, stride,
384 ))
385 })
386 }
387
388 fn lookup_block_inner<'a, K: QueryKey, const FIND_ALL: bool>(
393 &self,
394 block: &ArcBytes,
395 entry_count: usize,
396 key_hash: u64,
397 key: &K,
398 reader: ArcBlockCacheReader<'_>,
399 get_entry: impl Fn(usize) -> Result<GetKeyEntryResult<'a>>,
400 ) -> Result<SstLookupResult> {
401 let mut l = 0;
402 let mut r = entry_count;
403 while l < r {
405 let m = (l + r) / 2;
406 let GetKeyEntryResult {
407 hash: mid_hash,
408 key: mid_key,
409 ty,
410 val,
411 } = get_entry(m)?;
412
413 let comparison = compare_hash_key(mid_hash, mid_key, key_hash, key);
414
415 match comparison {
416 Ordering::Less => r = m,
417 Ordering::Equal => {
418 if !FIND_ALL {
419 let result = self.handle_key_match(ty, val, block, reader)?;
422 return Ok(SstLookupResult::Found(SmallVec::from_buf([result])));
423 }
424 let mut results = SmallVec::new();
430 for i in (l..m).rev() {
431 let GetKeyEntryResult {
432 hash,
433 key: entry_key,
434 ty,
435 val,
436 } = get_entry(i)?;
437 if !entry_matches_key(hash, entry_key, key_hash, key) {
438 break;
439 }
440 results.push(self.handle_key_match(ty, val, block, reader)?);
441 }
442 results.push(self.handle_key_match(ty, val, block, reader)?);
451 for i in (m + 1)..r {
452 let GetKeyEntryResult {
453 hash,
454 key: entry_key,
455 ty,
456 val,
457 } = get_entry(i)?;
458 if !entry_matches_key(hash, entry_key, key_hash, key) {
459 break;
460 }
461 results.push(self.handle_key_match(ty, val, block, reader)?);
462 }
463 return Ok(SstLookupResult::Found(results));
464 }
465 Ordering::Greater => l = m + 1,
466 }
467 }
468
469 Ok(SstLookupResult::NotFound)
470 }
471
472 fn handle_key_match(
474 &self,
475 ty: u8,
476 val: &[u8],
477 key_block_arc: &ArcBytes,
478 reader: ArcBlockCacheReader<'_>,
479 ) -> Result<LookupValue> {
480 handle_key_match_generic(&self.mmap, &self.meta, ty, val, key_block_arc, reader)
481 }
482}
483
484fn get_or_cache_block(
493 mmap: &Arc<Mmap>,
494 meta: &StaticSortedFileMetaData,
495 block_index: u16,
496 cache: &BlockCache,
497 verified_blocks: &[AtomicU64],
498) -> Result<ArcBytes> {
499 let (uncompressed_length, checksum, block_data) = get_raw_block_slice(mmap, meta, block_index)
500 .with_context(|| {
501 format!(
502 "Failed to read raw block {} from {:08}.sst",
503 block_index, meta.sequence_number
504 )
505 })?;
506
507 if uncompressed_length == 0 {
508 verify_checksum_once(meta, block_data, checksum, block_index, verified_blocks)?;
510 return Ok(unsafe { ArcBytes::from_mmap(mmap, block_data) });
512 }
513
514 Ok(
516 match cache.get_value_or_guard(&(meta.sequence_number, block_index), None) {
517 GuardResult::Value(block) => block,
518 GuardResult::Guard(guard) => {
519 verify_checksum_once(meta, block_data, checksum, block_index, verified_blocks)?;
522 let block = ArcBytes::from_decompressed(uncompressed_length, block_data)
523 .with_context(|| {
524 format!(
525 "Failed to decompress block {} from {:08}.sst ({} bytes uncompressed)",
526 block_index, meta.sequence_number, uncompressed_length
527 )
528 })?;
529 let _ = guard.insert(block.clone());
530 block
531 }
532 GuardResult::Timeout => unreachable!(),
533 },
534 )
535}
536
537fn get_raw_block_slice<'a>(
540 mmap: &'a Mmap,
541 meta: &StaticSortedFileMetaData,
542 block_index: u16,
543) -> Result<(u32, u32, &'a [u8])> {
544 #[cfg(feature = "strict_checks")]
545 if block_index >= meta.block_count {
546 bail!(
547 "Corrupted file seq:{} block:{} > number of blocks {} (block_offsets: {:x})",
548 meta.sequence_number,
549 block_index,
550 meta.block_count,
551 meta.block_offsets_start(mmap.len()),
552 );
553 }
554 let offset = meta.block_offsets_start(mmap.len()) + block_index as usize * 4;
555 #[cfg(feature = "strict_checks")]
556 if offset + 4 > mmap.len() {
557 bail!(
558 "Corrupted file seq:{} block:{} block offset locations {} + 4 bytes > file end {} \
559 (block_offsets: {:x})",
560 meta.sequence_number,
561 block_index,
562 offset,
563 mmap.len(),
564 meta.block_offsets_start(mmap.len()),
565 );
566 }
567 let block_start = if block_index == 0 {
568 0
569 } else {
570 be::read_u32(&mmap[offset - 4..]) as usize
571 };
572 let block_end = be::read_u32(&mmap[offset..]) as usize;
573 #[cfg(feature = "strict_checks")]
574 if block_end > mmap.len() || block_start > mmap.len() {
575 bail!(
576 "Corrupted file seq:{} block:{} block {} - {} > file end {} (block_offsets: {:x})",
577 meta.sequence_number,
578 block_index,
579 block_start,
580 block_end,
581 mmap.len(),
582 meta.block_offsets_start(mmap.len()),
583 );
584 }
585 ensure!(
586 block_start + BLOCK_HEADER_SIZE <= block_end,
587 "block {} header truncated in {:08}.sst",
588 block_index,
589 meta.sequence_number
590 );
591 let uncompressed_length = be::read_u32(&mmap[block_start..]);
592 let checksum = be::read_u32(&mmap[block_start + 4..]);
593 let block = &mmap[block_start + BLOCK_HEADER_SIZE..block_end];
594 Ok((uncompressed_length, checksum, block))
595}
596
597fn verify_checksum(
599 meta: &StaticSortedFileMetaData,
600 data: &[u8],
601 expected: u32,
602 block_index: u16,
603) -> Result<()> {
604 let actual = checksum_block(data);
605 if actual != expected {
606 bail!(
607 "Cache corruption detected: checksum mismatch in block {} of {:08}.sst (expected \
608 {:08x}, got {:08x})",
609 block_index,
610 meta.sequence_number,
611 expected,
612 actual
613 );
614 }
615 Ok(())
616}
617
618fn verify_checksum_once(
625 meta: &StaticSortedFileMetaData,
626 data: &[u8],
627 expected: u32,
628 block_index: u16,
629 verified_blocks: &[AtomicU64],
630) -> Result<()> {
631 let word_idx = block_index as usize / u64::BITS as usize;
632 let bit = 1u64 << (block_index as usize % u64::BITS as usize);
633 if verified_blocks[word_idx].load(AtomicOrdering::Relaxed) & bit != 0 {
634 return Ok(());
635 }
636 verify_checksum(meta, data, expected, block_index)?;
637 verified_blocks[word_idx].fetch_or(bit, AtomicOrdering::Relaxed);
638 Ok(())
639}
640
641fn get_raw_block_generic<B: SharedBytes>(
644 mmap: &B::MmapHandle,
645 meta: &StaticSortedFileMetaData,
646 block_index: u16,
647) -> Result<(u32, u32, B)> {
648 let (uncompressed_length, checksum, block) = get_raw_block_slice(mmap, meta, block_index)?;
649 Ok((uncompressed_length, checksum, unsafe {
651 B::from_mmap(mmap, block)
652 }))
653}
654
655#[tracing::instrument(level = "info", name = "reading database block", skip_all)]
658fn read_block_generic<B: SharedBytes>(
659 mmap: &B::MmapHandle,
660 meta: &StaticSortedFileMetaData,
661 block_index: u16,
662) -> Result<B> {
663 let (uncompressed_length, expected_checksum, block) =
664 get_raw_block_slice(mmap, meta, block_index).with_context(|| {
665 format!(
666 "Failed to read raw block {} from {:08}.sst",
667 block_index, meta.sequence_number
668 )
669 })?;
670
671 verify_checksum(meta, block, expected_checksum, block_index)?;
672
673 if uncompressed_length == 0 {
674 return Ok(unsafe { B::from_mmap(mmap, block) });
676 }
677
678 let buffer = B::from_decompressed(uncompressed_length, block).with_context(|| {
679 format!(
680 "Failed to decompress block {} from {:08}.sst ({} bytes uncompressed)",
681 block_index, meta.sequence_number, uncompressed_length
682 )
683 })?;
684 Ok(buffer)
685}
686
687fn handle_key_match_generic<B: SharedBytes>(
689 mmap: &B::MmapHandle,
690 meta: &StaticSortedFileMetaData,
691 ty: u8,
692 val: &[u8],
693 key_block: &B,
694 reader: impl ValueBlockCache<B>,
695) -> Result<LookupValue<B>> {
696 Ok(match ty {
697 KEY_BLOCK_ENTRY_TYPE_SMALL => {
698 let block = be::read_u16(val);
699 let size = be::read_u16(&val[2..]) as usize;
700 let position = be::read_u32(&val[4..]) as usize;
701 let value = reader
702 .get_or_read(mmap, meta, block)?
703 .slice(position..position + size);
704 LookupValue::Slice { value }
705 }
706 KEY_BLOCK_ENTRY_TYPE_MEDIUM => {
707 let block = be::read_u16(val);
708 let value = read_block_generic(mmap, meta, block)?;
709 LookupValue::Slice { value }
710 }
711 KEY_BLOCK_ENTRY_TYPE_BLOB => {
712 let sequence_number = be::read_u32(val);
713 LookupValue::Blob { sequence_number }
714 }
715 KEY_BLOCK_ENTRY_TYPE_DELETED => LookupValue::Deleted,
716 _ => {
717 let value = unsafe { key_block.slice_from_subslice(val) };
720 LookupValue::Slice { value }
721 }
722 })
723}
724
725pub struct StaticSortedFileIter {
727 mmap: Rc<Mmap>,
730 meta: StaticSortedFileMetaData,
732
733 index_entries: RcBytes,
736 num_index_entries: usize,
738 index_pos: usize,
740 current_key_block: CurrentKeyBlock,
741 value_block_cache: Option<(u16, RcBytes)>,
745}
746
747enum CurrentKeyBlockKind {
748 Variable { offsets: RcBytes, hash_len: u8 },
750 Fixed {
752 hash_len: u8,
753 key_size: usize,
754 value_type: u8,
755 stride: usize,
756 },
757}
758
759struct CurrentKeyBlock {
760 kind: CurrentKeyBlockKind,
761 entries: RcBytes,
762 entry_count: u32,
764 index: u32,
766}
767
768impl Iterator for StaticSortedFileIter {
769 type Item = Result<LookupEntry>;
770
771 fn next(&mut self) -> Option<Self::Item> {
772 self.next_internal().transpose()
773 }
774}
775
776impl StaticSortedFileIter {
777 pub fn open(db_path: &Path, meta: StaticSortedFileMetaData) -> Result<Self> {
781 let filename = format!("{:08}.sst", meta.sequence_number);
782 let path = db_path.join(&filename);
783 let file = File::open(&path)?;
784 let mmap = unsafe { Mmap::map(file.file()) }.with_context(|| {
785 format!(
786 "Failed to mmap SST file {} ({} bytes)",
787 path.display(),
788 file.metadata().map(|m| m.len()).unwrap_or(0)
789 )
790 })?;
791 #[cfg(unix)]
792 mmap.advise(memmap2::Advice::Sequential)?;
793 advise_mmap_for_persistence(&mmap)?;
794 Self::new(Rc::new(mmap), meta)
795 .with_context(|| format!("Unable to open static sorted file {filename}"))
796 }
797
798 fn new(mmap: Rc<Mmap>, meta: StaticSortedFileMetaData) -> Result<Self> {
799 let root_block_index = meta.block_count - 1;
800 let block: RcBytes = read_block_generic(&mmap, &meta, root_block_index)?;
801 let block_type = block[0];
802
803 if block_type != BLOCK_TYPE_INDEX {
805 bail!("Root block must be an index block");
806 }
807 let block_len = block.len();
808 ensure!(block_len >= 3, "index block too short");
809 let index_entries = block.slice(1..block_len);
810 let first_child = be::read_u16(&index_entries);
811 let num_index_entries: usize = (index_entries.len() + INDEX_BLOCK_ENTRY_SIZE
816 - size_of::<u16>())
817 / INDEX_BLOCK_ENTRY_SIZE;
818
819 let current_key_block = Self::parse_key_block(&mmap, &meta, first_child)?;
820 Ok(StaticSortedFileIter {
821 mmap,
822 meta,
823 index_entries,
824 num_index_entries,
825 index_pos: 1,
826 current_key_block,
827 value_block_cache: None,
828 })
829 }
830
831 fn parse_key_block(
833 mmap: &Rc<Mmap>,
834 meta: &StaticSortedFileMetaData,
835 block_index: u16,
836 ) -> Result<CurrentKeyBlock> {
837 let block: RcBytes = read_block_generic(mmap, meta, block_index)?;
838 let data = &*block;
839 ensure!(data.len() >= 4, "key block too short");
840 let block_type = data[0];
841 let entry_count = be::read_u24(&data[1..]);
842 match block_type {
843 BLOCK_TYPE_KEY_WITH_HASH | BLOCK_TYPE_KEY_NO_HASH => {
844 let hash_len = if block_type == BLOCK_TYPE_KEY_WITH_HASH {
845 8
846 } else {
847 0
848 };
849 let n = entry_count as usize;
850 let offsets_range = 4..4 + n * 4;
851 let entries_range = 4 + n * 4..block.len();
852 let offsets = block.clone().slice(offsets_range);
853 let entries = block.slice(entries_range);
854 Ok(CurrentKeyBlock {
855 kind: CurrentKeyBlockKind::Variable { offsets, hash_len },
856 entries,
857 entry_count,
858 index: 0,
859 })
860 }
861 BLOCK_TYPE_FIXED_KEY_WITH_HASH | BLOCK_TYPE_FIXED_KEY_NO_HASH => {
862 let hash_len = if block_type == BLOCK_TYPE_FIXED_KEY_WITH_HASH {
863 8
864 } else {
865 0
866 };
867 let key_size = data[4] as usize;
868 let value_type = data[5];
869 let val_size = entry_val_size(value_type)?;
870 let stride = hash_len as usize + key_size + val_size;
871 let entries_range = 6..block.len();
873 let entries = block.slice(entries_range);
874 Ok(CurrentKeyBlock {
875 kind: CurrentKeyBlockKind::Fixed {
876 hash_len,
877 key_size,
878 value_type,
879 stride,
880 },
881 entries,
882 entry_count,
883 index: 0,
884 })
885 }
886 _ => {
887 bail!("Invalid key block type: {block_type}");
888 }
889 }
890 }
891
892 fn next_internal(&mut self) -> Result<Option<LookupEntry>> {
894 loop {
895 let kb = &mut self.current_key_block;
896 if kb.index < kb.entry_count {
897 let index = kb.index as usize;
898 let entry_count = kb.entry_count as usize;
899 let GetKeyEntryResult { hash, key, ty, val } = match &kb.kind {
900 CurrentKeyBlockKind::Variable { offsets, hash_len } => {
901 get_key_entry(offsets, &kb.entries, entry_count, index, *hash_len)?
902 }
903 CurrentKeyBlockKind::Fixed {
904 hash_len,
905 key_size,
906 value_type,
907 stride,
908 } => get_fixed_key_entry(
909 &kb.entries,
910 index,
911 *hash_len,
912 *key_size,
913 *value_type,
914 *stride,
915 ),
916 };
917 let full_hash = if hash.is_empty() {
918 crate::key::hash_key(&key)
919 } else {
920 be::read_u64(hash)
921 };
922 let value = if ty == KEY_BLOCK_ENTRY_TYPE_MEDIUM {
923 let block = be::read_u16(val);
924 let (uncompressed_size, checksum, block) =
925 get_raw_block_generic(&self.mmap, &self.meta, block)?;
926 IterValue::Medium {
927 uncompressed_size,
928 checksum,
929 block,
930 }
931 } else {
932 handle_key_match_generic(
933 &self.mmap,
934 &self.meta,
935 ty,
936 val,
937 &kb.entries,
938 &mut self.value_block_cache,
939 )?
940 .into()
941 };
942 let entry = LookupEntry {
943 hash: full_hash,
944 key: unsafe { kb.entries.slice_from_subslice(key) },
945 value,
946 };
947 kb.index += 1;
948 return Ok(Some(entry));
949 }
950 if self.index_pos < self.num_index_entries {
951 let base = self.index_pos * INDEX_BLOCK_ENTRY_SIZE;
952 let block_index = be::read_u16(&self.index_entries[base..]);
953 self.index_pos += 1;
954 self.current_key_block =
955 Self::parse_key_block(&self.mmap, &self.meta, block_index)?;
956 } else {
957 return Ok(None);
958 }
959 }
960 }
961}
962
963struct GetKeyEntryResult<'l> {
964 hash: &'l [u8],
965 key: &'l [u8],
966 ty: u8,
967 val: &'l [u8],
968}
969
970fn compare_hash_key<K: QueryKey>(
974 entry_hash: &[u8],
975 entry_key: &[u8],
976 full_hash: u64,
977 query_key: &K,
978) -> Ordering {
979 if entry_hash.is_empty() {
980 let entry_full_hash = crate::key::hash_key(&entry_key);
982 match full_hash.cmp(&entry_full_hash) {
983 Ordering::Equal => query_key.cmp(entry_key),
984 ord => ord,
985 }
986 } else {
987 let full_hash_bytes = full_hash.to_be_bytes();
989 match full_hash_bytes[..].cmp(entry_hash) {
990 Ordering::Equal => query_key.cmp(entry_key),
991 ord => ord,
992 }
993 }
994}
995
996fn entry_matches_key<K: QueryKey>(
1000 entry_hash: &[u8],
1001 entry_key: &[u8],
1002 full_hash: u64,
1003 query_key: &K,
1004) -> bool {
1005 if entry_hash.is_empty() {
1006 query_key.cmp(entry_key) == Ordering::Equal
1008 } else {
1009 full_hash.to_be_bytes()[..] == *entry_hash && query_key.cmp(entry_key) == Ordering::Equal
1011 }
1012}
1013
1014fn entry_val_size(ty: u8) -> Result<usize> {
1016 match ty {
1017 KEY_BLOCK_ENTRY_TYPE_SMALL => Ok(SMALL_VALUE_REF_SIZE),
1018 KEY_BLOCK_ENTRY_TYPE_MEDIUM => Ok(MEDIUM_VALUE_REF_SIZE),
1019 KEY_BLOCK_ENTRY_TYPE_BLOB => Ok(BLOB_VALUE_REF_SIZE),
1020 KEY_BLOCK_ENTRY_TYPE_DELETED => Ok(DELETED_VALUE_REF_SIZE),
1021 ty if ty >= KEY_BLOCK_ENTRY_TYPE_INLINE_MIN => {
1022 Ok((ty - KEY_BLOCK_ENTRY_TYPE_INLINE_MIN) as usize)
1023 }
1024 _ => bail!("Invalid key block entry type: {ty}"),
1025 }
1026}
1027
1028#[inline(always)]
1031fn read_offset_entry(offsets: &[u8], index: usize) -> (u8, usize) {
1032 let base = index * 4;
1033 let word = be::read_u32(&offsets[base..]);
1034 let ty = (word >> 24) as u8;
1035 let offset = (word & 0x00FF_FFFF) as usize;
1036 (ty, offset)
1037}
1038
1039fn get_key_entry<'l>(
1041 offsets: &[u8],
1042 entries: &'l [u8],
1043 entry_count: usize,
1044 index: usize,
1045 hash_len: u8,
1046) -> Result<GetKeyEntryResult<'l>> {
1047 let hash_len_usize = hash_len as usize;
1048 let (ty, start) = read_offset_entry(offsets, index);
1049 let end = if index == entry_count - 1 {
1050 entries.len()
1051 } else {
1052 let (_, next_start) = read_offset_entry(offsets, index + 1);
1053 next_start
1054 };
1055 let hash = &entries[start..start + hash_len_usize];
1057 let val_size = entry_val_size(ty)?;
1058 Ok(GetKeyEntryResult {
1059 hash,
1060 key: &entries[start + hash_len_usize..end - val_size],
1061 ty,
1062 val: &entries[end - val_size..end],
1063 })
1064}
1065
1066fn get_fixed_key_entry<'l>(
1071 entries: &'l [u8],
1072 index: usize,
1073 hash_len: u8,
1074 key_size: usize,
1075 value_type: u8,
1076 stride: usize,
1077) -> GetKeyEntryResult<'l> {
1078 let hash_len_usize = hash_len as usize;
1079 let start = index * stride;
1080 GetKeyEntryResult {
1081 hash: &entries[start..start + hash_len_usize],
1082 key: &entries[start + hash_len_usize..start + hash_len_usize + key_size],
1083 ty: value_type,
1084 val: &entries[start + hash_len_usize + key_size..(index + 1) * stride],
1085 }
1086}