1use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
22use parking_lot::RwLock;
23use rustc_hash::FxHashMap;
24use std::io::{self, Read, Write};
25use std::sync::Arc;
26
27#[cfg(feature = "fst-index")]
28use super::sstable_index::FstBlockIndex;
29use super::sstable_index::{BlockAddr, BlockIndex, MmapBlockIndex};
30use super::vint::{read_vint, write_vint};
31use crate::compression::{CompressionDict, CompressionLevel};
32use crate::directories::{FileHandle, OwnedBytes};
33
34pub const SSTABLE_MAGIC: u32 = 0x53544235; pub const RESTART_INTERVAL: usize = 16;
42
43pub const BLOCK_SIZE: usize = 16 * 1024;
45
46pub const DEFAULT_DICT_SIZE: usize = 64 * 1024;
48
49pub const BLOOM_BITS_PER_KEY: usize = 10;
51
52pub const BLOOM_HASH_COUNT: usize = 7;
54const BLOOM_FILTER_HEADER_SIZE: usize = 16;
55
56const MAX_SSTABLE_BLOCK_BYTES: usize = 64 * 1024 * 1024;
57const MAX_SSTABLE_DICTIONARY_BYTES: u64 = 16 * 1024 * 1024;
58
59pub type PrefixScanResult<V> = (Vec<(Vec<u8>, V)>, bool);
61
62#[derive(Debug, Clone)]
68pub struct BloomFilter {
69 bits: BloomBits,
70 num_bits: usize,
71 num_hashes: usize,
72}
73
74#[derive(Debug, Clone)]
76enum BloomBits {
77 Vec(Vec<u64>),
79 Bytes(OwnedBytes),
81}
82
83impl BloomBits {
84 #[inline]
85 fn len(&self) -> usize {
86 match self {
87 BloomBits::Vec(v) => v.len(),
88 BloomBits::Bytes(b) => b.len() / 8,
89 }
90 }
91
92 #[inline]
93 fn get(&self, word_idx: usize) -> u64 {
94 match self {
95 BloomBits::Vec(v) => v[word_idx],
96 BloomBits::Bytes(b) => {
97 let off = word_idx * 8;
98 u64::from_le_bytes([
99 b[off],
100 b[off + 1],
101 b[off + 2],
102 b[off + 3],
103 b[off + 4],
104 b[off + 5],
105 b[off + 6],
106 b[off + 7],
107 ])
108 }
109 }
110 }
111
112 #[inline]
113 fn set_bit(&mut self, word_idx: usize, bit_idx: usize) {
114 match self {
115 BloomBits::Vec(v) => v[word_idx] |= 1u64 << bit_idx,
116 BloomBits::Bytes(_) => panic!("cannot mutate read-only bloom filter"),
117 }
118 }
119
120 fn size_bytes(&self) -> usize {
121 match self {
122 BloomBits::Vec(v) => v.len() * 8,
123 BloomBits::Bytes(b) => b.len(),
124 }
125 }
126
127 fn write_to(&self, writer: &mut (impl Write + ?Sized)) -> io::Result<()> {
128 match self {
129 BloomBits::Vec(words) => {
130 #[cfg(target_endian = "little")]
131 {
132 let bytes = unsafe {
135 std::slice::from_raw_parts(
136 words.as_ptr().cast::<u8>(),
137 words.len().saturating_mul(8),
138 )
139 };
140 writer.write_all(bytes)
141 }
142 #[cfg(target_endian = "big")]
143 {
144 for &word in words {
145 writer.write_u64::<LittleEndian>(word)?;
146 }
147 Ok(())
148 }
149 }
150 BloomBits::Bytes(bytes) => writer.write_all(bytes.as_slice()),
151 }
152 }
153}
154
155impl BloomFilter {
156 pub(crate) const SERIALIZED_HEADER_SIZE: usize = BLOOM_FILTER_HEADER_SIZE;
157
158 pub fn new(expected_keys: usize, bits_per_key: usize) -> Self {
160 let num_bits = expected_keys.saturating_mul(bits_per_key).max(64);
161 let num_words = num_bits.div_ceil(64);
162 Self {
163 bits: BloomBits::Vec(vec![0u64; num_words]),
164 num_bits,
165 num_hashes: BLOOM_HASH_COUNT,
166 }
167 }
168
169 pub fn from_bytes_mutable(data: &[u8]) -> io::Result<Self> {
173 if data.len() < BLOOM_FILTER_HEADER_SIZE {
174 return Err(io::Error::new(
175 io::ErrorKind::InvalidData,
176 "Bloom filter data too short",
177 ));
178 }
179 let num_bits = usize::try_from(u64::from_le_bytes(data[0..8].try_into().unwrap()))
180 .map_err(|_| {
181 io::Error::new(
182 io::ErrorKind::InvalidData,
183 "Bloom filter bit count exceeds addressable memory",
184 )
185 })?;
186 let num_hashes = u32::from_le_bytes(data[8..12].try_into().unwrap()) as usize;
187 let num_words = u32::from_le_bytes(data[12..16].try_into().unwrap()) as usize;
188
189 validate_bloom_header(data.len(), num_bits, num_hashes, num_words)?;
190 let expected_len = BLOOM_FILTER_HEADER_SIZE + num_words * 8;
191 if data.len() != expected_len {
192 return Err(io::Error::new(
193 io::ErrorKind::InvalidData,
194 "Bloom filter data truncated",
195 ));
196 }
197
198 let mut vec = vec![0u64; num_words];
199 for (i, v) in vec.iter_mut().enumerate() {
200 let off = BLOOM_FILTER_HEADER_SIZE + i * 8;
201 *v = u64::from_le_bytes(data[off..off + 8].try_into().unwrap());
202 }
203
204 Ok(Self {
205 bits: BloomBits::Vec(vec),
206 num_bits,
207 num_hashes,
208 })
209 }
210
211 pub fn from_owned_bytes(data: OwnedBytes) -> io::Result<Self> {
213 if data.len() < BLOOM_FILTER_HEADER_SIZE {
214 return Err(io::Error::new(
215 io::ErrorKind::InvalidData,
216 "Bloom filter data too short",
217 ));
218 }
219 let d = data.as_slice();
220 let num_bits =
221 usize::try_from(u64::from_le_bytes(d[0..8].try_into().unwrap())).map_err(|_| {
222 io::Error::new(
223 io::ErrorKind::InvalidData,
224 "Bloom filter bit count exceeds addressable memory",
225 )
226 })?;
227 let num_hashes = u32::from_le_bytes(d[8..12].try_into().unwrap()) as usize;
228 let num_words = u32::from_le_bytes(d[12..16].try_into().unwrap()) as usize;
229
230 validate_bloom_header(d.len(), num_bits, num_hashes, num_words)?;
231 let expected_len = BLOOM_FILTER_HEADER_SIZE + num_words * 8;
232 if d.len() != expected_len {
233 return Err(io::Error::new(
234 io::ErrorKind::InvalidData,
235 "Bloom filter data truncated",
236 ));
237 }
238
239 let bits_bytes =
241 data.slice(BLOOM_FILTER_HEADER_SIZE..BLOOM_FILTER_HEADER_SIZE + num_words * 8);
242
243 Ok(Self {
244 bits: BloomBits::Bytes(bits_bytes),
245 num_bits,
246 num_hashes,
247 })
248 }
249
250 pub fn serialized_len(&self) -> usize {
252 BLOOM_FILTER_HEADER_SIZE + self.bits.len() * 8
253 }
254
255 pub fn write_to(&self, writer: &mut (impl Write + ?Sized)) -> io::Result<()> {
257 let num_words = self.bits.len();
258 write_bloom_header(writer, self.num_bits, self.num_hashes, num_words)?;
259 self.bits.write_to(writer)
260 }
261
262 pub fn to_bytes(&self) -> Vec<u8> {
264 let mut data = Vec::with_capacity(self.serialized_len());
265 self.write_to(&mut data)
266 .expect("writing a bloom filter to Vec cannot fail");
267 data
268 }
269
270 pub fn insert(&mut self, key: &[u8]) {
272 let (h1, h2) = self.hash_pair(key);
273 for i in 0..self.num_hashes {
274 let bit_pos = self.get_bit_pos(h1, h2, i);
275 let word_idx = bit_pos / 64;
276 let bit_idx = bit_pos % 64;
277 if word_idx < self.bits.len() {
278 self.bits.set_bit(word_idx, bit_idx);
279 }
280 }
281 }
282
283 pub fn may_contain(&self, key: &[u8]) -> bool {
286 let (h1, h2) = self.hash_pair(key);
287 for i in 0..self.num_hashes {
288 let bit_pos = self.get_bit_pos(h1, h2, i);
289 let word_idx = bit_pos / 64;
290 let bit_idx = bit_pos % 64;
291 if word_idx >= self.bits.len() || (self.bits.get(word_idx) & (1u64 << bit_idx)) == 0 {
292 return false;
293 }
294 }
295 true
296 }
297
298 pub fn size_bytes(&self) -> usize {
300 BLOOM_FILTER_HEADER_SIZE + self.bits.size_bytes()
301 }
302
303 pub fn insert_hashed(&mut self, h1: u64, h2: u64) {
305 for i in 0..self.num_hashes {
306 let bit_pos = self.get_bit_pos(h1, h2, i);
307 let word_idx = bit_pos / 64;
308 let bit_idx = bit_pos % 64;
309 if word_idx < self.bits.len() {
310 self.bits.set_bit(word_idx, bit_idx);
311 }
312 }
313 }
314
315 #[inline]
317 fn hash_pair(&self, key: &[u8]) -> (u64, u64) {
318 let mut h1: u64 = 0xcbf29ce484222325;
319 let mut h2: u64 = 0x84222325cbf29ce4;
320 for &byte in key {
321 h1 ^= byte as u64;
322 h1 = h1.wrapping_mul(0x100000001b3);
323 h2 = h2.wrapping_mul(0x100000001b3);
324 h2 ^= byte as u64;
325 }
326 (h1, h2)
327 }
328
329 #[inline]
331 fn get_bit_pos(&self, h1: u64, h2: u64, i: usize) -> usize {
332 (h1.wrapping_add((i as u64).wrapping_mul(h2)) % (self.num_bits as u64)) as usize
333 }
334}
335
336fn write_bloom_header(
337 writer: &mut (impl Write + ?Sized),
338 num_bits: usize,
339 num_hashes: usize,
340 num_words: usize,
341) -> io::Result<()> {
342 writer.write_u64::<LittleEndian>(u64::try_from(num_bits).map_err(|_| {
343 io::Error::new(
344 io::ErrorKind::InvalidInput,
345 "Bloom filter bit count exceeds u64",
346 )
347 })?)?;
348 writer.write_u32::<LittleEndian>(u32::try_from(num_hashes).map_err(|_| {
349 io::Error::new(
350 io::ErrorKind::InvalidInput,
351 "Bloom filter hash count exceeds u32",
352 )
353 })?)?;
354 writer.write_u32::<LittleEndian>(u32::try_from(num_words).map_err(|_| {
355 io::Error::new(
356 io::ErrorKind::InvalidInput,
357 "Bloom filter word count exceeds u32",
358 )
359 })?)?;
360 Ok(())
361}
362
363fn validate_bloom_header(
364 data_len: usize,
365 num_bits: usize,
366 num_hashes: usize,
367 num_words: usize,
368) -> io::Result<()> {
369 if num_bits == 0 || num_hashes == 0 || num_hashes > 32 || num_words == 0 {
370 return Err(io::Error::new(
371 io::ErrorKind::InvalidData,
372 "invalid bloom filter parameters",
373 ));
374 }
375 let word_bytes = num_words
376 .checked_mul(8)
377 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
378 let expected_len = BLOOM_FILTER_HEADER_SIZE
379 .checked_add(word_bytes)
380 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
381 let capacity_bits = num_words
382 .checked_mul(64)
383 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
384 if expected_len > data_len
385 || num_bits > capacity_bits
386 || num_bits <= capacity_bits.saturating_sub(64)
387 {
388 return Err(io::Error::new(
389 io::ErrorKind::InvalidData,
390 "inconsistent bloom filter dimensions",
391 ));
392 }
393 Ok(())
394}
395
396#[inline]
399fn bloom_hash_pair(key: &[u8]) -> (u64, u64) {
400 let mut h1: u64 = 0xcbf29ce484222325;
401 let mut h2: u64 = 0x84222325cbf29ce4;
402 for &byte in key {
403 h1 ^= byte as u64;
404 h1 = h1.wrapping_mul(0x100000001b3);
405 h2 = h2.wrapping_mul(0x100000001b3);
406 h2 ^= byte as u64;
407 }
408 (h1, h2)
409}
410
411pub trait SSTableValue: Clone + Send + Sync {
418 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()>;
420
421 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self>;
423}
424
425impl SSTableValue for u64 {
427 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
428 write_vint(writer, *self)
429 }
430
431 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
432 read_vint(reader)
433 }
434}
435
436impl SSTableValue for Vec<u8> {
438 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
439 write_vint(writer, self.len() as u64)?;
440 writer.write_all(self)
441 }
442
443 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
444 let len = usize::try_from(read_vint(reader)?).map_err(|_| {
445 io::Error::new(io::ErrorKind::InvalidData, "SSTable value length too large")
446 })?;
447 if len > MAX_SSTABLE_BLOCK_BYTES {
448 return Err(io::Error::new(
449 io::ErrorKind::InvalidData,
450 "SSTable value exceeds block safety limit",
451 ));
452 }
453 let mut data = vec![0u8; len];
454 reader.read_exact(&mut data)?;
455 Ok(data)
456 }
457}
458
459#[derive(Debug, Clone, Copy, PartialEq, Eq)]
462pub struct SparseDimInfo {
463 pub offset: u64,
465 pub length: u32,
467}
468
469impl SparseDimInfo {
470 pub fn new(offset: u64, length: u32) -> Self {
471 Self { offset, length }
472 }
473}
474
475impl SSTableValue for SparseDimInfo {
476 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
477 write_vint(writer, self.offset)?;
478 write_vint(writer, self.length as u64)
479 }
480
481 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
482 let offset = read_vint(reader)?;
483 let length = u32::try_from(read_vint(reader)?).map_err(|_| {
484 io::Error::new(
485 io::ErrorKind::InvalidData,
486 "sparse posting length exceeds u32",
487 )
488 })?;
489 Ok(Self { offset, length })
490 }
491}
492
493pub const MAX_INLINE_POSTINGS: usize = 3;
495
496#[derive(Debug, Clone, PartialEq, Eq)]
504pub enum TermInfo {
505 Inline {
508 doc_freq: u8,
510 data: [u8; 16],
513 data_len: u8,
515 },
516 External {
518 posting_offset: u64,
519 posting_len: u64,
520 doc_freq: u32,
521 position_offset: u64,
523 position_len: u64,
525 },
526}
527
528impl TermInfo {
529 pub fn external(posting_offset: u64, posting_len: u64, doc_freq: u32) -> Self {
531 TermInfo::External {
532 posting_offset,
533 posting_len,
534 doc_freq,
535 position_offset: 0,
536 position_len: 0,
537 }
538 }
539
540 pub fn external_with_positions(
542 posting_offset: u64,
543 posting_len: u64,
544 doc_freq: u32,
545 position_offset: u64,
546 position_len: u64,
547 ) -> Self {
548 TermInfo::External {
549 posting_offset,
550 posting_len,
551 doc_freq,
552 position_offset,
553 position_len,
554 }
555 }
556
557 pub fn try_inline(doc_ids: &[u32], term_freqs: &[u32]) -> Option<Self> {
560 if doc_ids.len() > MAX_INLINE_POSTINGS
561 || doc_ids.is_empty()
562 || doc_ids.len() != term_freqs.len()
563 {
564 return None;
565 }
566
567 let mut data = [0u8; 16];
568 let mut cursor = std::io::Cursor::new(&mut data[..]);
569 let mut prev_doc_id = 0u32;
570
571 for (i, &doc_id) in doc_ids.iter().enumerate() {
572 let delta = doc_id.checked_sub(prev_doc_id)?;
573 if write_vint(&mut cursor, delta as u64).is_err() {
574 return None;
575 }
576 if write_vint(&mut cursor, term_freqs[i] as u64).is_err() {
577 return None;
578 }
579 prev_doc_id = doc_id;
580 }
581
582 let data_len = cursor.position() as u8;
583 if data_len > 16 {
584 return None;
585 }
586
587 Some(TermInfo::Inline {
588 doc_freq: doc_ids.len() as u8,
589 data,
590 data_len,
591 })
592 }
593
594 pub fn try_inline_iter(count: usize, iter: impl Iterator<Item = (u32, u32)>) -> Option<Self> {
598 if count > MAX_INLINE_POSTINGS || count == 0 {
599 return None;
600 }
601
602 let mut data = [0u8; 16];
603 let mut cursor = std::io::Cursor::new(&mut data[..]);
604 let mut prev_doc_id = 0u32;
605
606 let mut actual_count = 0usize;
607 for (doc_id, tf) in iter {
608 if actual_count >= count {
609 return None;
610 }
611 let delta = doc_id.checked_sub(prev_doc_id)?;
612 if write_vint(&mut cursor, delta as u64).is_err() {
613 return None;
614 }
615 if write_vint(&mut cursor, tf as u64).is_err() {
616 return None;
617 }
618 prev_doc_id = doc_id;
619 actual_count += 1;
620 }
621
622 if actual_count != count {
623 return None;
624 }
625
626 let data_len = cursor.position() as u8;
627
628 Some(TermInfo::Inline {
629 doc_freq: count as u8,
630 data,
631 data_len,
632 })
633 }
634
635 pub fn doc_freq(&self) -> u32 {
637 match self {
638 TermInfo::Inline { doc_freq, .. } => *doc_freq as u32,
639 TermInfo::External { doc_freq, .. } => *doc_freq,
640 }
641 }
642
643 pub fn is_inline(&self) -> bool {
645 matches!(self, TermInfo::Inline { .. })
646 }
647
648 pub fn external_info(&self) -> Option<(u64, u64)> {
650 match self {
651 TermInfo::External {
652 posting_offset,
653 posting_len,
654 ..
655 } => Some((*posting_offset, *posting_len)),
656 TermInfo::Inline { .. } => None,
657 }
658 }
659
660 pub fn position_info(&self) -> Option<(u64, u64)> {
662 match self {
663 TermInfo::External {
664 position_offset,
665 position_len,
666 ..
667 } if *position_len > 0 => Some((*position_offset, *position_len)),
668 _ => None,
669 }
670 }
671
672 pub fn decode_inline(&self) -> Option<(Vec<u32>, Vec<u32>)> {
675 match self {
676 TermInfo::Inline {
677 doc_freq,
678 data,
679 data_len,
680 } => {
681 if *doc_freq == 0
682 || *doc_freq as usize > MAX_INLINE_POSTINGS
683 || *data_len as usize > data.len()
684 {
685 return None;
686 }
687 let mut doc_ids = Vec::with_capacity(*doc_freq as usize);
688 let mut term_freqs = Vec::with_capacity(*doc_freq as usize);
689 let mut reader = &data[..*data_len as usize];
690 let mut prev_doc_id = 0u32;
691
692 for _ in 0..*doc_freq {
693 let delta = read_vint(&mut reader).ok()? as u32;
694 let tf = read_vint(&mut reader).ok()? as u32;
695 let doc_id = prev_doc_id.checked_add(delta)?;
696 doc_ids.push(doc_id);
697 term_freqs.push(tf);
698 prev_doc_id = doc_id;
699 }
700
701 Some((doc_ids, term_freqs))
702 }
703 TermInfo::External { .. } => None,
704 }
705 }
706}
707
708impl SSTableValue for TermInfo {
709 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
710 match self {
711 TermInfo::Inline {
712 doc_freq,
713 data,
714 data_len,
715 } => {
716 if *doc_freq == 0
717 || *doc_freq as usize > MAX_INLINE_POSTINGS
718 || *data_len as usize > data.len()
719 {
720 return Err(io::Error::new(
721 io::ErrorKind::InvalidInput,
722 "invalid inline TermInfo",
723 ));
724 }
725 writer.write_u8(0xFF)?;
727 writer.write_u8(*doc_freq)?;
728 writer.write_u8(*data_len)?;
729 writer.write_all(&data[..*data_len as usize])?;
730 }
731 TermInfo::External {
732 posting_offset,
733 posting_len,
734 doc_freq,
735 position_offset,
736 position_len,
737 } => {
738 if *position_len > 0 {
741 writer.write_u8(0x01)?;
742 write_vint(writer, *doc_freq as u64)?;
743 write_vint(writer, *posting_offset)?;
744 write_vint(writer, *posting_len)?;
745 write_vint(writer, *position_offset)?;
746 write_vint(writer, *position_len)?;
747 } else {
748 writer.write_u8(0x00)?;
749 write_vint(writer, *doc_freq as u64)?;
750 write_vint(writer, *posting_offset)?;
751 write_vint(writer, *posting_len)?;
752 }
753 }
754 }
755 Ok(())
756 }
757
758 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
759 let tag = reader.read_u8()?;
760
761 if tag == 0xFF {
762 let doc_freq = reader.read_u8()?;
764 let data_len = reader.read_u8()?;
765 if doc_freq == 0 || doc_freq as usize > MAX_INLINE_POSTINGS || data_len as usize > 16 {
766 return Err(io::Error::new(
767 io::ErrorKind::InvalidData,
768 "invalid inline TermInfo lengths",
769 ));
770 }
771 let mut data = [0u8; 16];
772 reader.read_exact(&mut data[..data_len as usize])?;
773 Ok(TermInfo::Inline {
774 doc_freq,
775 data,
776 data_len,
777 })
778 } else if tag == 0x00 {
779 let doc_freq = read_vint(reader)? as u32;
781 let posting_offset = read_vint(reader)?;
782 let posting_len = read_vint(reader)?;
783 Ok(TermInfo::External {
784 posting_offset,
785 posting_len,
786 doc_freq,
787 position_offset: 0,
788 position_len: 0,
789 })
790 } else if tag == 0x01 {
791 let doc_freq = read_vint(reader)? as u32;
793 let posting_offset = read_vint(reader)?;
794 let posting_len = read_vint(reader)?;
795 let position_offset = read_vint(reader)?;
796 let position_len = read_vint(reader)?;
797 Ok(TermInfo::External {
798 posting_offset,
799 posting_len,
800 doc_freq,
801 position_offset,
802 position_len,
803 })
804 } else {
805 Err(io::Error::new(
806 io::ErrorKind::InvalidData,
807 format!("Invalid TermInfo tag: {}", tag),
808 ))
809 }
810 }
811}
812
813pub fn common_prefix_len(a: &[u8], b: &[u8]) -> usize {
815 a.iter().zip(b.iter()).take_while(|(x, y)| x == y).count()
816}
817
818fn decode_block_entry<V: SSTableValue>(
824 reader: &mut &[u8],
825 current_key: &mut Vec<u8>,
826) -> io::Result<V> {
827 let common_prefix_len = read_vint(reader)? as usize;
828 let suffix_len = read_vint(reader)? as usize;
829
830 if suffix_len > reader.len() {
831 return Err(io::Error::new(
832 io::ErrorKind::UnexpectedEof,
833 "SSTable block suffix truncated",
834 ));
835 }
836
837 current_key.truncate(common_prefix_len);
838 current_key.extend_from_slice(&reader[..suffix_len]);
839 *reader = &reader[suffix_len..];
840
841 V::deserialize(reader)
842}
843
844#[derive(Debug, Clone)]
846pub struct SSTableStats {
847 pub num_blocks: usize,
849 pub num_sparse_entries: usize,
851 pub num_entries: u64,
853 pub has_bloom_filter: bool,
855 pub has_dictionary: bool,
857 pub bloom_filter_size: usize,
859 pub dictionary_size: usize,
861}
862
863#[derive(Debug, Clone)]
865pub struct SSTableWriterConfig {
866 pub compression_level: CompressionLevel,
868 pub use_dictionary: bool,
870 pub dict_size: usize,
872 pub use_bloom_filter: bool,
874 pub bloom_bits_per_key: usize,
876}
877
878impl Default for SSTableWriterConfig {
879 fn default() -> Self {
880 Self::from_optimization(crate::structures::IndexOptimization::default())
881 }
882}
883
884impl SSTableWriterConfig {
885 pub fn from_optimization(optimization: crate::structures::IndexOptimization) -> Self {
887 use crate::structures::IndexOptimization;
888 match optimization {
889 IndexOptimization::Adaptive => Self {
890 compression_level: CompressionLevel::BETTER, use_dictionary: false,
892 dict_size: DEFAULT_DICT_SIZE,
893 use_bloom_filter: true, bloom_bits_per_key: BLOOM_BITS_PER_KEY,
895 },
896 IndexOptimization::SizeOptimized => Self {
897 compression_level: CompressionLevel::MAX, use_dictionary: true,
899 dict_size: DEFAULT_DICT_SIZE,
900 use_bloom_filter: true,
901 bloom_bits_per_key: BLOOM_BITS_PER_KEY,
902 },
903 IndexOptimization::PerformanceOptimized => Self {
904 compression_level: CompressionLevel::FAST, use_dictionary: false,
906 dict_size: DEFAULT_DICT_SIZE,
907 use_bloom_filter: true, bloom_bits_per_key: BLOOM_BITS_PER_KEY,
909 },
910 }
911 }
912
913 pub fn fast() -> Self {
915 Self::from_optimization(crate::structures::IndexOptimization::PerformanceOptimized)
916 }
917
918 pub fn max_compression() -> Self {
920 Self::from_optimization(crate::structures::IndexOptimization::SizeOptimized)
921 }
922}
923
924pub struct SSTableWriter<W: Write, V: SSTableValue> {
930 writer: W,
931 block_buffer: Vec<u8>,
932 prev_key: Vec<u8>,
933 index: Vec<BlockIndexEntry>,
934 current_offset: u64,
935 num_entries: u64,
936 block_first_key: Option<Vec<u8>>,
937 config: SSTableWriterConfig,
938 dictionary: Option<CompressionDict>,
940 bloom_hashes: Vec<(u64, u64)>,
943 block_restarts: Vec<u32>,
946 block_entry_count: usize,
948 _phantom: std::marker::PhantomData<V>,
949}
950
951impl<W: Write, V: SSTableValue> SSTableWriter<W, V> {
952 pub fn new(writer: W) -> Self {
954 Self::with_config(writer, SSTableWriterConfig::default())
955 }
956
957 pub fn with_config(writer: W, config: SSTableWriterConfig) -> Self {
959 Self {
960 writer,
961 block_buffer: Vec::with_capacity(BLOCK_SIZE),
962 prev_key: Vec::new(),
963 index: Vec::new(),
964 current_offset: 0,
965 num_entries: 0,
966 block_first_key: None,
967 config,
968 dictionary: None,
969 bloom_hashes: Vec::new(),
970 block_restarts: Vec::new(),
971 block_entry_count: 0,
972 _phantom: std::marker::PhantomData,
973 }
974 }
975
976 pub fn with_dictionary(
978 writer: W,
979 config: SSTableWriterConfig,
980 dictionary: CompressionDict,
981 ) -> Self {
982 Self {
983 writer,
984 block_buffer: Vec::with_capacity(BLOCK_SIZE),
985 prev_key: Vec::new(),
986 index: Vec::new(),
987 current_offset: 0,
988 num_entries: 0,
989 block_first_key: None,
990 config,
991 dictionary: Some(dictionary),
992 bloom_hashes: Vec::new(),
993 block_restarts: Vec::new(),
994 block_entry_count: 0,
995 _phantom: std::marker::PhantomData,
996 }
997 }
998
999 pub fn insert(&mut self, key: &[u8], value: &V) -> io::Result<()> {
1000 if self.block_first_key.is_none() {
1001 self.block_first_key = Some(key.to_vec());
1002 }
1003
1004 if self.config.use_bloom_filter {
1006 self.bloom_hashes.push(bloom_hash_pair(key));
1007 }
1008
1009 let restart = self.block_entry_count.is_multiple_of(RESTART_INTERVAL);
1012 let prefix_len = if restart {
1013 self.block_restarts.push(self.block_buffer.len() as u32);
1014 0
1015 } else {
1016 common_prefix_len(&self.prev_key, key)
1017 };
1018 let suffix = &key[prefix_len..];
1019
1020 write_vint(&mut self.block_buffer, prefix_len as u64)?;
1021 write_vint(&mut self.block_buffer, suffix.len() as u64)?;
1022 self.block_buffer.extend_from_slice(suffix);
1023 value.serialize(&mut self.block_buffer)?;
1024
1025 self.prev_key.clear();
1026 self.prev_key.extend_from_slice(key);
1027 self.num_entries += 1;
1028 self.block_entry_count += 1;
1029
1030 if self.block_buffer.len() >= BLOCK_SIZE {
1031 self.flush_block()?;
1032 }
1033
1034 Ok(())
1035 }
1036
1037 fn flush_block(&mut self) -> io::Result<()> {
1039 if self.block_buffer.is_empty() {
1040 return Ok(());
1041 }
1042
1043 for offset in &self.block_restarts {
1045 self.block_buffer.extend_from_slice(&offset.to_le_bytes());
1046 }
1047 self.block_buffer
1048 .extend_from_slice(&(self.block_restarts.len() as u32).to_le_bytes());
1049 self.block_restarts.clear();
1050 self.block_entry_count = 0;
1051
1052 let compressed = if let Some(ref dict) = self.dictionary {
1054 crate::compression::compress_with_dict(
1055 &self.block_buffer,
1056 self.config.compression_level,
1057 dict,
1058 )?
1059 } else {
1060 crate::compression::compress(&self.block_buffer, self.config.compression_level)?
1061 };
1062
1063 if let Some(first_key) = self.block_first_key.take() {
1064 self.index.push(BlockIndexEntry {
1065 first_key,
1066 offset: self.current_offset,
1067 length: compressed.len() as u32,
1068 });
1069 }
1070
1071 self.writer.write_all(&compressed)?;
1072 self.current_offset += compressed.len() as u64;
1073 self.block_buffer.clear();
1074 self.prev_key.clear();
1075
1076 Ok(())
1077 }
1078
1079 pub fn finish(mut self) -> io::Result<W> {
1080 self.flush_block()?;
1082
1083 let bloom_filter = if self.config.use_bloom_filter && !self.bloom_hashes.is_empty() {
1085 let mut bloom =
1086 BloomFilter::new(self.bloom_hashes.len(), self.config.bloom_bits_per_key);
1087 for (h1, h2) in &self.bloom_hashes {
1088 bloom.insert_hashed(*h1, *h2);
1089 }
1090 Some(bloom)
1091 } else {
1092 None
1093 };
1094
1095 let data_end_offset = self.current_offset;
1096
1097 let entries: Vec<(Vec<u8>, BlockAddr)> = self
1100 .index
1101 .iter()
1102 .map(|e| {
1103 (
1104 e.first_key.clone(),
1105 BlockAddr {
1106 offset: e.offset,
1107 length: e.length,
1108 },
1109 )
1110 })
1111 .collect();
1112
1113 #[cfg(feature = "native")]
1115 let index_bytes = FstBlockIndex::build(&entries)?;
1116 #[cfg(not(feature = "native"))]
1117 let index_bytes = MmapBlockIndex::build(&entries)?;
1118
1119 self.writer
1121 .write_u32::<LittleEndian>(index_bytes.len() as u32)?;
1122 self.writer.write_all(&index_bytes)?;
1123 self.current_offset += 4 + index_bytes.len() as u64;
1124
1125 let bloom_offset = if let Some(ref bloom) = bloom_filter {
1127 let offset = self.current_offset;
1128 bloom.write_to(&mut self.writer)?;
1129 self.current_offset += bloom.serialized_len() as u64;
1130 offset
1131 } else {
1132 0
1133 };
1134
1135 let dict_offset = if let Some(ref dict) = self.dictionary {
1137 let dict_bytes = dict.as_bytes();
1138 let offset = self.current_offset;
1139 self.writer
1140 .write_u32::<LittleEndian>(dict_bytes.len() as u32)?;
1141 self.writer.write_all(dict_bytes)?;
1142 self.current_offset += 4 + dict_bytes.len() as u64;
1143 offset
1144 } else {
1145 0
1146 };
1147
1148 self.writer.write_u64::<LittleEndian>(data_end_offset)?;
1150 self.writer.write_u64::<LittleEndian>(self.num_entries)?;
1151 self.writer.write_u64::<LittleEndian>(bloom_offset)?; self.writer.write_u64::<LittleEndian>(dict_offset)?; self.writer
1154 .write_u8(self.config.compression_level.0 as u8)?;
1155 self.writer.write_u32::<LittleEndian>(SSTABLE_MAGIC)?;
1156
1157 Ok(self.writer)
1158 }
1159}
1160
1161#[derive(Debug, Clone)]
1163struct BlockIndexEntry {
1164 first_key: Vec<u8>,
1165 offset: u64,
1166 length: u32,
1167}
1168
1169pub struct AsyncSSTableReader<V: SSTableValue> {
1176 data_slice: FileHandle,
1178 block_index: BlockIndex,
1180 num_entries: u64,
1181 cache: RwLock<BlockCache>,
1183 bloom_filter: Option<BloomFilter>,
1185 dictionary: Option<CompressionDict>,
1187 #[allow(dead_code)]
1189 compression_level: CompressionLevel,
1190 _phantom: std::marker::PhantomData<V>,
1191}
1192
1193struct BlockParts<'b> {
1195 entries: &'b [u8],
1196 restart_table: &'b [u8],
1198}
1199
1200impl<'b> BlockParts<'b> {
1201 fn split(block: &'b [u8]) -> io::Result<Self> {
1202 if block.len() < 4 {
1203 return Err(io::Error::new(
1204 io::ErrorKind::InvalidData,
1205 "SSTable block shorter than its restart trailer",
1206 ));
1207 }
1208 let count_at = block.len() - 4;
1209 let count = u32::from_le_bytes(block[count_at..].try_into().unwrap()) as usize;
1210 let table_len = count.checked_mul(4).ok_or_else(|| {
1211 io::Error::new(io::ErrorKind::InvalidData, "SSTable restart table overflow")
1212 })?;
1213 let table_at = count_at.checked_sub(table_len).ok_or_else(|| {
1214 io::Error::new(
1215 io::ErrorKind::InvalidData,
1216 "SSTable restart table exceeds its block",
1217 )
1218 })?;
1219 Ok(Self {
1220 entries: &block[..table_at],
1221 restart_table: &block[table_at..count_at],
1222 })
1223 }
1224
1225 #[inline]
1226 fn num_restarts(&self) -> usize {
1227 self.restart_table.len() / 4
1228 }
1229
1230 #[inline]
1231 fn restart_offset(&self, i: usize) -> io::Result<usize> {
1232 let at = i * 4;
1233 let offset =
1234 u32::from_le_bytes(self.restart_table[at..at + 4].try_into().unwrap()) as usize;
1235 if offset >= self.entries.len() {
1236 return Err(io::Error::new(
1237 io::ErrorKind::InvalidData,
1238 "SSTable restart offset outside the block",
1239 ));
1240 }
1241 Ok(offset)
1242 }
1243
1244 fn restart_key(&self, i: usize) -> io::Result<&'b [u8]> {
1246 let offset = self.restart_offset(i)?;
1247 let mut reader = &self.entries[offset..];
1248 let prefix_len = read_vint(&mut reader)?;
1249 if prefix_len != 0 {
1250 return Err(io::Error::new(
1251 io::ErrorKind::InvalidData,
1252 "SSTable restart entry has a non-zero key prefix",
1253 ));
1254 }
1255 let suffix_len = read_vint(&mut reader)? as usize;
1256 if suffix_len > reader.len() {
1257 return Err(io::Error::new(
1258 io::ErrorKind::UnexpectedEof,
1259 "SSTable restart key truncated",
1260 ));
1261 }
1262 Ok(&reader[..suffix_len])
1263 }
1264}
1265
1266struct BlockCache {
1272 blocks: FxHashMap<u64, Arc<[u8]>>,
1273 lru_order: std::collections::VecDeque<u64>,
1274 max_blocks: usize,
1275}
1276
1277impl BlockCache {
1278 fn new(max_blocks: usize) -> Self {
1279 Self {
1280 blocks: FxHashMap::default(),
1281 lru_order: std::collections::VecDeque::with_capacity(max_blocks),
1282 max_blocks,
1283 }
1284 }
1285
1286 fn get(&mut self, offset: u64) -> Option<Arc<[u8]>> {
1287 if self.blocks.contains_key(&offset) {
1288 self.promote(offset);
1289 self.blocks.get(&offset).map(Arc::clone)
1290 } else {
1291 None
1292 }
1293 }
1294
1295 fn peek(&self, offset: u64) -> Option<Arc<[u8]>> {
1297 self.blocks.get(&offset).map(Arc::clone)
1298 }
1299
1300 fn insert(&mut self, offset: u64, block: Arc<[u8]>) {
1301 if self.max_blocks == 0 {
1302 return;
1303 }
1304 if self.blocks.contains_key(&offset) {
1305 self.promote(offset);
1306 return;
1307 }
1308 while self.blocks.len() >= self.max_blocks {
1309 if let Some(evict_offset) = self.lru_order.pop_front() {
1310 self.blocks.remove(&evict_offset);
1311 } else {
1312 break;
1313 }
1314 }
1315 self.blocks.insert(offset, block);
1316 self.lru_order.push_back(offset);
1317 }
1318
1319 fn promote(&mut self, offset: u64) {
1321 if let Some(pos) = self.lru_order.iter().position(|&k| k == offset) {
1322 self.lru_order.remove(pos);
1323 self.lru_order.push_back(offset);
1324 }
1325 }
1326}
1327
1328impl<V: SSTableValue> AsyncSSTableReader<V> {
1329 pub async fn open(file_handle: FileHandle, cache_blocks: usize) -> io::Result<Self> {
1334 let file_len = file_handle.len();
1335 if file_len < 37 {
1336 return Err(io::Error::new(
1337 io::ErrorKind::InvalidData,
1338 "SSTable too small",
1339 ));
1340 }
1341
1342 let footer_bytes = file_handle
1345 .read_bytes_range(file_len - 37..file_len)
1346 .await?;
1347
1348 let mut reader = footer_bytes.as_slice();
1349 let data_end_offset = reader.read_u64::<LittleEndian>()?;
1350 let num_entries = reader.read_u64::<LittleEndian>()?;
1351 let bloom_offset = reader.read_u64::<LittleEndian>()?;
1352 let dict_offset = reader.read_u64::<LittleEndian>()?;
1353 let compression_level = CompressionLevel(reader.read_u8()? as i32);
1354 let magic = reader.read_u32::<LittleEndian>()?;
1355
1356 if magic != SSTABLE_MAGIC {
1357 return Err(io::Error::new(
1358 io::ErrorKind::InvalidData,
1359 format!(
1360 "Invalid SSTable magic: 0x{magic:08X} (required STB5); the term dictionary \
1361 was written by an incompatible Hermes"
1362 ),
1363 ));
1364 }
1365
1366 let footer_start = file_len - 37;
1367 if data_end_offset > footer_start {
1368 return Err(io::Error::new(
1369 io::ErrorKind::InvalidData,
1370 "SSTable data section extends past its footer",
1371 ));
1372 }
1373 if bloom_offset != 0 && (bloom_offset < data_end_offset || bloom_offset >= footer_start) {
1374 return Err(io::Error::new(
1375 io::ErrorKind::InvalidData,
1376 "SSTable bloom filter offset is out of bounds",
1377 ));
1378 }
1379 if dict_offset != 0
1380 && (dict_offset < data_end_offset
1381 || dict_offset >= footer_start
1382 || (bloom_offset != 0 && dict_offset <= bloom_offset))
1383 {
1384 return Err(io::Error::new(
1385 io::ErrorKind::InvalidData,
1386 "SSTable dictionary offset is out of bounds",
1387 ));
1388 }
1389
1390 let index_start = data_end_offset;
1392 let index_end = if bloom_offset != 0 {
1393 bloom_offset
1394 } else if dict_offset != 0 {
1395 dict_offset
1396 } else {
1397 footer_start
1398 };
1399 let index_bytes = file_handle.read_bytes_range(index_start..index_end).await?;
1400
1401 let mut idx_reader = index_bytes.as_slice();
1403 let index_len = idx_reader.read_u32::<LittleEndian>()? as usize;
1404
1405 if index_len != idx_reader.len() {
1406 return Err(io::Error::new(
1407 io::ErrorKind::InvalidData,
1408 "Index data truncated",
1409 ));
1410 }
1411
1412 let index_data = index_bytes.slice(4..4 + index_len);
1413
1414 #[cfg(feature = "fst-index")]
1416 let block_index = match FstBlockIndex::load(index_data.clone()) {
1417 Ok(fst_idx) => BlockIndex::Fst(fst_idx),
1418 Err(_) => BlockIndex::Mmap(MmapBlockIndex::load(index_data)?),
1419 };
1420 #[cfg(not(feature = "fst-index"))]
1421 let block_index = BlockIndex::Mmap(MmapBlockIndex::load(index_data)?);
1422
1423 let mut expected_offset = 0u64;
1424 for addr in block_index.all_addrs() {
1425 let end = addr.offset.checked_add(addr.length as u64).ok_or_else(|| {
1426 io::Error::new(io::ErrorKind::InvalidData, "SSTable block range overflow")
1427 })?;
1428 if addr.length == 0 || addr.offset != expected_offset || end > data_end_offset {
1429 return Err(io::Error::new(
1430 io::ErrorKind::InvalidData,
1431 "SSTable block addresses are inconsistent",
1432 ));
1433 }
1434 expected_offset = end;
1435 }
1436 if expected_offset != data_end_offset {
1437 return Err(io::Error::new(
1438 io::ErrorKind::InvalidData,
1439 "SSTable block addresses do not cover the data section",
1440 ));
1441 }
1442
1443 let bloom_filter = if bloom_offset > 0 {
1445 let bloom_start = bloom_offset;
1446 let bloom_end = if dict_offset != 0 {
1447 dict_offset
1448 } else {
1449 footer_start
1450 };
1451 let header_size = BloomFilter::SERIALIZED_HEADER_SIZE as u64;
1453 let bloom_header_end = bloom_start.checked_add(header_size).ok_or_else(|| {
1454 io::Error::new(io::ErrorKind::InvalidData, "bloom filter range overflow")
1455 })?;
1456 if bloom_header_end > bloom_end {
1457 return Err(io::Error::new(
1458 io::ErrorKind::UnexpectedEof,
1459 "bloom filter header is truncated",
1460 ));
1461 }
1462 let bloom_header = file_handle
1463 .read_bytes_range(bloom_start..bloom_header_end)
1464 .await?;
1465 let num_words = u32::from_le_bytes([
1466 bloom_header[12],
1467 bloom_header[13],
1468 bloom_header[14],
1469 bloom_header[15],
1470 ]) as u64;
1471 let bloom_size = num_words
1472 .checked_mul(8)
1473 .and_then(|bytes| bytes.checked_add(header_size))
1474 .ok_or_else(|| {
1475 io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow")
1476 })?;
1477 let actual_bloom_size = bloom_end - bloom_start;
1478 if bloom_size != actual_bloom_size {
1479 return Err(io::Error::new(
1480 io::ErrorKind::InvalidData,
1481 "bloom filter length is inconsistent",
1482 ));
1483 }
1484 let bloom_data = file_handle.read_bytes_range(bloom_start..bloom_end).await?;
1485 Some(BloomFilter::from_owned_bytes(bloom_data)?)
1486 } else {
1487 None
1488 };
1489
1490 let dictionary = if dict_offset > 0 {
1492 let dict_start = dict_offset;
1493 let dict_header_end = dict_start.checked_add(4).ok_or_else(|| {
1495 io::Error::new(io::ErrorKind::InvalidData, "dictionary range overflow")
1496 })?;
1497 if dict_header_end > footer_start {
1498 return Err(io::Error::new(
1499 io::ErrorKind::UnexpectedEof,
1500 "dictionary header is truncated",
1501 ));
1502 }
1503 let dict_len_bytes = file_handle
1504 .read_bytes_range(dict_start..dict_header_end)
1505 .await?;
1506 let dict_len = u32::from_le_bytes([
1507 dict_len_bytes[0],
1508 dict_len_bytes[1],
1509 dict_len_bytes[2],
1510 dict_len_bytes[3],
1511 ]) as u64;
1512 if dict_len > MAX_SSTABLE_DICTIONARY_BYTES {
1513 return Err(io::Error::new(
1514 io::ErrorKind::InvalidData,
1515 "SSTable dictionary exceeds safety limit",
1516 ));
1517 }
1518 let dict_end = dict_header_end.checked_add(dict_len).ok_or_else(|| {
1519 io::Error::new(io::ErrorKind::InvalidData, "dictionary range overflow")
1520 })?;
1521 if dict_end != footer_start {
1522 return Err(io::Error::new(
1523 io::ErrorKind::InvalidData,
1524 "SSTable dictionary length is inconsistent",
1525 ));
1526 }
1527 let dict_data = file_handle
1528 .read_bytes_range(dict_header_end..dict_end)
1529 .await?;
1530 Some(CompressionDict::from_owned_bytes(dict_data))
1531 } else {
1532 None
1533 };
1534
1535 let data_slice = file_handle.slice(0..data_end_offset);
1537
1538 Ok(Self {
1539 data_slice,
1540 block_index,
1541 num_entries,
1542 cache: RwLock::new(BlockCache::new(cache_blocks)),
1543 bloom_filter,
1544 dictionary,
1545 compression_level,
1546 _phantom: std::marker::PhantomData,
1547 })
1548 }
1549
1550 pub fn num_entries(&self) -> u64 {
1552 self.num_entries
1553 }
1554
1555 pub fn stats(&self) -> SSTableStats {
1557 SSTableStats {
1558 num_blocks: self.block_index.len(),
1559 num_sparse_entries: 0, num_entries: self.num_entries,
1561 has_bloom_filter: self.bloom_filter.is_some(),
1562 has_dictionary: self.dictionary.is_some(),
1563 bloom_filter_size: self
1564 .bloom_filter
1565 .as_ref()
1566 .map(|b| b.size_bytes())
1567 .unwrap_or(0),
1568 dictionary_size: self.dictionary.as_ref().map(|d| d.len()).unwrap_or(0),
1569 }
1570 }
1571
1572 pub fn cached_blocks(&self) -> usize {
1574 self.cache.read().blocks.len()
1575 }
1576
1577 pub fn cached_bytes(&self) -> usize {
1583 self.cache
1584 .read()
1585 .blocks
1586 .values()
1587 .map(|block| block.len())
1588 .sum()
1589 }
1590
1591 pub async fn get(&self, key: &[u8]) -> io::Result<Option<V>> {
1596 log::debug!(
1597 "SSTable::get called, key_len={}, total_blocks={}",
1598 key.len(),
1599 self.block_index.len()
1600 );
1601
1602 if let Some(ref bloom) = self.bloom_filter
1604 && !bloom.may_contain(key)
1605 {
1606 log::debug!("SSTable::get bloom filter negative");
1607 return Ok(None);
1608 }
1609
1610 let block_idx = match self.block_index.locate(key) {
1612 Some(idx) => idx,
1613 None => {
1614 log::debug!("SSTable::get key not found (before first block)");
1615 return Ok(None);
1616 }
1617 };
1618
1619 log::debug!("SSTable::get loading block_idx={}", block_idx);
1620
1621 let block_data = self.load_block(block_idx).await?;
1623 self.search_block(&block_data, key)
1624 }
1625
1626 pub async fn get_batch(&self, keys: &[&[u8]]) -> io::Result<Vec<Option<V>>> {
1632 if keys.is_empty() {
1633 return Ok(Vec::new());
1634 }
1635
1636 let mut key_to_block: Vec<(usize, usize)> = Vec::with_capacity(keys.len());
1638 for (key_idx, key) in keys.iter().enumerate() {
1639 if let Some(ref bloom) = self.bloom_filter
1641 && !bloom.may_contain(key)
1642 {
1643 key_to_block.push((key_idx, usize::MAX)); continue;
1645 }
1646
1647 match self.block_index.locate(key) {
1648 Some(block_idx) => key_to_block.push((key_idx, block_idx)),
1649 None => key_to_block.push((key_idx, usize::MAX)), }
1651 }
1652
1653 let mut blocks_to_load: Vec<usize> = key_to_block
1655 .iter()
1656 .filter(|(_, b)| *b != usize::MAX)
1657 .map(|(_, b)| *b)
1658 .collect();
1659 blocks_to_load.sort_unstable();
1660 blocks_to_load.dedup();
1661
1662 for &block_idx in &blocks_to_load {
1664 let _ = self.load_block(block_idx).await?;
1665 }
1666
1667 let mut results = vec![None; keys.len()];
1669 for (key_idx, block_idx) in key_to_block {
1670 if block_idx == usize::MAX {
1671 continue;
1672 }
1673 let block_data = self.load_block(block_idx).await?; results[key_idx] = self.search_block(&block_data, keys[key_idx])?;
1675 }
1676
1677 Ok(results)
1678 }
1679
1680 pub async fn preload_all_blocks(&self) -> io::Result<()> {
1685 for block_idx in 0..self.block_index.len() {
1686 self.load_block(block_idx).await?;
1687 }
1688 Ok(())
1689 }
1690
1691 pub async fn prefetch_all_data_bulk(&self) -> io::Result<()> {
1697 let num_blocks = self.block_index.len();
1698 if num_blocks == 0 {
1699 return Ok(());
1700 }
1701
1702 let mut max_end: u64 = 0;
1704 for i in 0..num_blocks {
1705 if let Some(addr) = self.block_index.get_addr(i) {
1706 let end = addr.offset.checked_add(addr.length as u64).ok_or_else(|| {
1707 io::Error::new(io::ErrorKind::InvalidData, "SSTable block range overflow")
1708 })?;
1709 max_end = max_end.max(end);
1710 }
1711 }
1712
1713 let all_data = self.data_slice.read_bytes_range(0..max_end).await?;
1715 let buf = all_data.as_slice();
1716
1717 let mut cache = self.cache.write();
1719 cache.max_blocks = cache.max_blocks.max(num_blocks);
1720 for i in 0..num_blocks {
1721 let addr = self.block_index.get_addr(i).unwrap();
1722 if cache.get(addr.offset).is_some() {
1723 continue;
1724 }
1725 let start = usize::try_from(addr.offset).map_err(|_| {
1726 io::Error::new(io::ErrorKind::InvalidData, "SSTable block offset too large")
1727 })?;
1728 let end = addr
1729 .offset
1730 .checked_add(addr.length as u64)
1731 .and_then(|end| usize::try_from(end).ok())
1732 .ok_or_else(|| {
1733 io::Error::new(io::ErrorKind::InvalidData, "SSTable block range overflow")
1734 })?;
1735 let compressed = buf.get(start..end).ok_or_else(|| {
1736 io::Error::new(io::ErrorKind::UnexpectedEof, "SSTable block is truncated")
1737 })?;
1738 let decompressed = if let Some(ref dict) = self.dictionary {
1739 crate::compression::decompress_with_dict_limited(
1740 compressed,
1741 dict,
1742 MAX_SSTABLE_BLOCK_BYTES,
1743 )?
1744 } else {
1745 crate::compression::decompress_limited(compressed, MAX_SSTABLE_BLOCK_BYTES)?
1746 };
1747 cache.insert(addr.offset, Arc::from(decompressed));
1748 }
1749
1750 Ok(())
1751 }
1752
1753 async fn load_block(&self, block_idx: usize) -> io::Result<Arc<[u8]>> {
1756 let addr = self.block_index.get_addr(block_idx).ok_or_else(|| {
1757 io::Error::new(io::ErrorKind::InvalidInput, "Block index out of range")
1758 })?;
1759
1760 {
1762 if let Some(block) = self.cache.read().peek(addr.offset) {
1763 return Ok(block);
1764 }
1765 }
1766
1767 log::debug!(
1768 "SSTable::load_block idx={} CACHE MISS, reading bytes [{}-{}]",
1769 block_idx,
1770 addr.offset,
1771 addr.offset + addr.length as u64
1772 );
1773
1774 let range = addr.byte_range();
1776 let compressed = self.data_slice.read_bytes_range(range).await?;
1777
1778 let decompressed = if let Some(ref dict) = self.dictionary {
1780 crate::compression::decompress_with_dict_limited(
1781 compressed.as_slice(),
1782 dict,
1783 MAX_SSTABLE_BLOCK_BYTES,
1784 )?
1785 } else {
1786 crate::compression::decompress_limited(compressed.as_slice(), MAX_SSTABLE_BLOCK_BYTES)?
1787 };
1788
1789 let block: Arc<[u8]> = Arc::from(decompressed);
1790
1791 {
1793 let mut cache = self.cache.write();
1794 cache.insert(addr.offset, Arc::clone(&block));
1795 }
1796
1797 Ok(block)
1798 }
1799
1800 #[cfg(feature = "sync")]
1802 fn load_block_sync(&self, block_idx: usize) -> io::Result<Arc<[u8]>> {
1803 let addr = self.block_index.get_addr(block_idx).ok_or_else(|| {
1804 io::Error::new(io::ErrorKind::InvalidInput, "Block index out of range")
1805 })?;
1806
1807 {
1809 if let Some(block) = self.cache.read().peek(addr.offset) {
1810 return Ok(block);
1811 }
1812 }
1813
1814 let range = addr.byte_range();
1816 let compressed = self.data_slice.read_bytes_range_sync(range)?;
1817
1818 let decompressed = if let Some(ref dict) = self.dictionary {
1820 crate::compression::decompress_with_dict_limited(
1821 compressed.as_slice(),
1822 dict,
1823 MAX_SSTABLE_BLOCK_BYTES,
1824 )?
1825 } else {
1826 crate::compression::decompress_limited(compressed.as_slice(), MAX_SSTABLE_BLOCK_BYTES)?
1827 };
1828
1829 let block: Arc<[u8]> = Arc::from(decompressed);
1830
1831 {
1833 let mut cache = self.cache.write();
1834 cache.insert(addr.offset, Arc::clone(&block));
1835 }
1836
1837 Ok(block)
1838 }
1839
1840 #[cfg(feature = "sync")]
1842 pub fn get_sync(&self, key: &[u8]) -> io::Result<Option<V>> {
1843 if let Some(ref bloom) = self.bloom_filter
1845 && !bloom.may_contain(key)
1846 {
1847 return Ok(None);
1848 }
1849
1850 let block_idx = match self.block_index.locate(key) {
1852 Some(idx) => idx,
1853 None => {
1854 return Ok(None);
1855 }
1856 };
1857
1858 let block_data = self.load_block_sync(block_idx)?;
1859 self.search_block(&block_data, key)
1860 }
1861
1862 fn block_entries<'b>(&self, block_data: &'b [u8]) -> io::Result<&'b [u8]> {
1864 Ok(BlockParts::split(block_data)?.entries)
1865 }
1866
1867 fn search_block(&self, block_data: &[u8], target_key: &[u8]) -> io::Result<Option<V>> {
1868 let parts = BlockParts::split(block_data)?;
1869
1870 let start = if parts.num_restarts() > 0 {
1873 let (mut lo, mut hi) = (0usize, parts.num_restarts());
1874 while lo < hi {
1875 let mid = lo + (hi - lo) / 2;
1876 if parts.restart_key(mid)? <= target_key {
1877 lo = mid + 1;
1878 } else {
1879 hi = mid;
1880 }
1881 }
1882 if lo == 0 {
1883 return Ok(None);
1885 }
1886 parts.restart_offset(lo - 1)?
1887 } else {
1888 0
1889 };
1890
1891 let mut reader = &parts.entries[start..];
1892 let mut current_key = Vec::new();
1893
1894 while !reader.is_empty() {
1895 let value = decode_block_entry(&mut reader, &mut current_key)?;
1896
1897 match current_key.as_slice().cmp(target_key) {
1898 std::cmp::Ordering::Equal => return Ok(Some(value)),
1899 std::cmp::Ordering::Greater => return Ok(None),
1900 std::cmp::Ordering::Less => continue,
1901 }
1902 }
1903
1904 Ok(None)
1905 }
1906
1907 pub async fn prefetch_range(&self, start_key: &[u8], end_key: &[u8]) -> io::Result<()> {
1909 let start_block = self.block_index.locate(start_key).unwrap_or(0);
1910 let end_block = self
1911 .block_index
1912 .locate(end_key)
1913 .unwrap_or(self.block_index.len().saturating_sub(1));
1914
1915 for block_idx in start_block..=end_block.min(self.block_index.len().saturating_sub(1)) {
1916 let _ = self.load_block(block_idx).await?;
1917 }
1918
1919 Ok(())
1920 }
1921
1922 pub fn iter(&self) -> AsyncSSTableIterator<'_, V> {
1924 AsyncSSTableIterator::new(self)
1925 }
1926
1927 pub async fn all_entries(&self) -> io::Result<Vec<(Vec<u8>, V)>> {
1929 let mut results = Vec::new();
1930
1931 for block_idx in 0..self.block_index.len() {
1932 let block_data = self.load_block(block_idx).await?;
1933 let mut reader = self.block_entries(&block_data)?;
1934 let mut current_key = Vec::new();
1935
1936 while !reader.is_empty() {
1937 let value = decode_block_entry(&mut reader, &mut current_key)?;
1938 results.push((current_key.clone(), value));
1939 }
1940 }
1941
1942 Ok(results)
1943 }
1944
1945 pub async fn prefix_scan(&self, prefix: &[u8]) -> io::Result<Vec<(Vec<u8>, V)>> {
1951 let (results, _) = self.prefix_scan_limited(prefix, usize::MAX).await?;
1952 Ok(results)
1953 }
1954
1955 pub async fn prefix_scan_limited(
1958 &self,
1959 prefix: &[u8],
1960 max_results: usize,
1961 ) -> io::Result<PrefixScanResult<V>> {
1962 if self.block_index.is_empty() || prefix.is_empty() {
1963 return Ok((Vec::new(), false));
1964 }
1965
1966 let start_block = self.block_index.locate(prefix).unwrap_or(0);
1970
1971 let mut results = Vec::new();
1972
1973 for block_idx in start_block..self.block_index.len() {
1974 let block_data = self.load_block(block_idx).await?;
1975 let mut reader = self.block_entries(&block_data)?;
1976 let mut current_key = Vec::new();
1977
1978 while !reader.is_empty() {
1979 let value = decode_block_entry(&mut reader, &mut current_key)?;
1980
1981 if current_key.starts_with(prefix) {
1982 if results.len() >= max_results {
1983 return Ok((results, true));
1984 }
1985 results.push((current_key.clone(), value));
1986 } else if current_key.as_slice() > prefix {
1987 return Ok((results, false));
1989 }
1990 }
1991 }
1992
1993 Ok((results, false))
1994 }
1995
1996 #[cfg(feature = "sync")]
1998 pub fn prefix_scan_sync(&self, prefix: &[u8]) -> io::Result<Vec<(Vec<u8>, V)>> {
1999 let (results, _) = self.prefix_scan_limited_sync(prefix, usize::MAX)?;
2000 Ok(results)
2001 }
2002
2003 #[cfg(feature = "sync")]
2005 pub fn prefix_scan_limited_sync(
2006 &self,
2007 prefix: &[u8],
2008 max_results: usize,
2009 ) -> io::Result<PrefixScanResult<V>> {
2010 if self.block_index.is_empty() || prefix.is_empty() {
2011 return Ok((Vec::new(), false));
2012 }
2013
2014 let start_block = self.block_index.locate(prefix).unwrap_or(0);
2016
2017 let mut results = Vec::new();
2018
2019 for block_idx in start_block..self.block_index.len() {
2020 let block_data = self.load_block_sync(block_idx)?;
2021 let mut reader = self.block_entries(&block_data)?;
2022 let mut current_key = Vec::new();
2023
2024 while !reader.is_empty() {
2025 let value = decode_block_entry(&mut reader, &mut current_key)?;
2026
2027 if current_key.starts_with(prefix) {
2028 if results.len() >= max_results {
2029 return Ok((results, true));
2030 }
2031 results.push((current_key.clone(), value));
2032 } else if current_key.as_slice() > prefix {
2033 return Ok((results, false));
2034 }
2035 }
2036 }
2037
2038 Ok((results, false))
2039 }
2040}
2041
2042pub struct AsyncSSTableIterator<'a, V: SSTableValue> {
2044 reader: &'a AsyncSSTableReader<V>,
2045 current_block: usize,
2046 block_data: Option<Arc<[u8]>>,
2047 block_offset: usize,
2048 block_entries_end: usize,
2050 current_key: Vec<u8>,
2051 finished: bool,
2052}
2053
2054impl<'a, V: SSTableValue> AsyncSSTableIterator<'a, V> {
2055 fn new(reader: &'a AsyncSSTableReader<V>) -> Self {
2056 Self {
2057 reader,
2058 current_block: 0,
2059 block_data: None,
2060 block_offset: 0,
2061 block_entries_end: 0,
2062 current_key: Vec::new(),
2063 finished: reader.block_index.is_empty(),
2064 }
2065 }
2066
2067 async fn load_next_block(&mut self) -> io::Result<bool> {
2068 if self.current_block >= self.reader.block_index.len() {
2069 self.finished = true;
2070 return Ok(false);
2071 }
2072
2073 let block = self.reader.load_block(self.current_block).await?;
2074 self.block_entries_end = self.reader.block_entries(&block)?.len();
2075 self.block_data = Some(block);
2076 self.block_offset = 0;
2077 self.current_key.clear();
2078 self.current_block += 1;
2079 Ok(true)
2080 }
2081
2082 pub async fn next(&mut self) -> io::Result<Option<(Vec<u8>, V)>> {
2084 if self.finished {
2085 return Ok(None);
2086 }
2087
2088 if self.block_data.is_none() && !self.load_next_block().await? {
2089 return Ok(None);
2090 }
2091
2092 loop {
2093 let block = self.block_data.as_ref().unwrap();
2094 if self.block_offset >= self.block_entries_end {
2095 if !self.load_next_block().await? {
2096 return Ok(None);
2097 }
2098 continue;
2099 }
2100
2101 let mut reader = &block[self.block_offset..self.block_entries_end];
2102 let start_len = reader.len();
2103
2104 let value = decode_block_entry(&mut reader, &mut self.current_key)?;
2105
2106 self.block_offset += start_len - reader.len();
2107
2108 return Ok(Some((self.current_key.clone(), value)));
2109 }
2110 }
2111}
2112
2113#[cfg(test)]
2114mod tests {
2115 use super::*;
2116
2117 #[test]
2118 fn test_bloom_filter_basic() {
2119 let mut bloom = BloomFilter::new(100, 10);
2120
2121 bloom.insert(b"hello");
2122 bloom.insert(b"world");
2123 bloom.insert(b"test");
2124
2125 assert!(bloom.may_contain(b"hello"));
2126 assert!(bloom.may_contain(b"world"));
2127 assert!(bloom.may_contain(b"test"));
2128
2129 assert!(!bloom.may_contain(b"notfound"));
2131 assert!(!bloom.may_contain(b"missing"));
2132 }
2133
2134 #[test]
2135 fn test_bloom_filter_serialization() {
2136 let mut bloom = BloomFilter::new(100, 10);
2137 bloom.insert(b"key1");
2138 bloom.insert(b"key2");
2139
2140 let bytes = bloom.to_bytes();
2141 let restored = BloomFilter::from_owned_bytes(OwnedBytes::new(bytes)).unwrap();
2142
2143 assert!(restored.may_contain(b"key1"));
2144 assert!(restored.may_contain(b"key2"));
2145 assert!(!restored.may_contain(b"key3"));
2146 }
2147
2148 #[test]
2149 fn bloom_header_preserves_bit_counts_above_u32() {
2150 let num_bits = u32::MAX as usize + 1;
2151 let mut header = Vec::new();
2152 write_bloom_header(&mut header, num_bits, BLOOM_HASH_COUNT, 1).unwrap();
2153 assert_eq!(header.len(), BLOOM_FILTER_HEADER_SIZE);
2154 assert_eq!(
2155 u64::from_le_bytes(header[0..8].try_into().unwrap()),
2156 num_bits as u64
2157 );
2158 }
2159
2160 #[test]
2161 fn test_bloom_filter_false_positive_rate() {
2162 let num_keys = 10000;
2163 let mut bloom = BloomFilter::new(num_keys, BLOOM_BITS_PER_KEY);
2164
2165 for i in 0..num_keys {
2167 let key = format!("key_{}", i);
2168 bloom.insert(key.as_bytes());
2169 }
2170
2171 for i in 0..num_keys {
2173 let key = format!("key_{}", i);
2174 assert!(bloom.may_contain(key.as_bytes()));
2175 }
2176
2177 let mut false_positives = 0;
2179 let test_count = 10000;
2180 for i in 0..test_count {
2181 let key = format!("nonexistent_{}", i);
2182 if bloom.may_contain(key.as_bytes()) {
2183 false_positives += 1;
2184 }
2185 }
2186
2187 let fp_rate = false_positives as f64 / test_count as f64;
2190 assert!(
2191 fp_rate < 0.03,
2192 "False positive rate {} is too high",
2193 fp_rate
2194 );
2195 }
2196
2197 #[test]
2198 fn test_sstable_writer_config() {
2199 use crate::structures::IndexOptimization;
2200
2201 let config = SSTableWriterConfig::default();
2203 assert_eq!(config.compression_level.0, 9); assert!(config.use_bloom_filter); assert!(!config.use_dictionary);
2206
2207 let adaptive = SSTableWriterConfig::from_optimization(IndexOptimization::Adaptive);
2209 assert_eq!(adaptive.compression_level.0, 9);
2210 assert!(adaptive.use_bloom_filter);
2211 assert!(!adaptive.use_dictionary);
2212
2213 let size = SSTableWriterConfig::from_optimization(IndexOptimization::SizeOptimized);
2215 assert_eq!(size.compression_level.0, 22); assert!(size.use_bloom_filter);
2217 assert!(size.use_dictionary);
2218
2219 let perf = SSTableWriterConfig::from_optimization(IndexOptimization::PerformanceOptimized);
2221 assert_eq!(perf.compression_level.0, 1); assert!(perf.use_bloom_filter); assert!(!perf.use_dictionary);
2224
2225 let fast = SSTableWriterConfig::fast();
2227 assert_eq!(fast.compression_level.0, 1);
2228
2229 let max = SSTableWriterConfig::max_compression();
2230 assert_eq!(max.compression_level.0, 22);
2231 }
2232
2233 #[test]
2234 fn test_vint_roundtrip() {
2235 let test_values = [0u64, 1, 127, 128, 255, 256, 16383, 16384, u64::MAX];
2236
2237 for &val in &test_values {
2238 let mut buf = Vec::new();
2239 write_vint(&mut buf, val).unwrap();
2240 let mut reader = buf.as_slice();
2241 let decoded = read_vint(&mut reader).unwrap();
2242 assert_eq!(val, decoded, "Failed for value {}", val);
2243 }
2244 }
2245
2246 fn keyed_table(num_keys: usize) -> (Vec<u8>, Vec<Vec<u8>>) {
2247 let mut keys: Vec<Vec<u8>> = (0..num_keys)
2248 .map(|i| format!("field{:02}/term{:07}", i % 7, i * 7919 % 100_003).into_bytes())
2249 .collect();
2250 keys.sort();
2251 keys.dedup();
2252 let mut writer = SSTableWriter::<_, u64>::new(Vec::new());
2253 for (i, key) in keys.iter().enumerate() {
2254 writer.insert(key, &(i as u64)).unwrap();
2255 }
2256 (writer.finish().unwrap(), keys)
2257 }
2258
2259 async fn check_every_key_and_scan(bytes: Vec<u8>, keys: &[Vec<u8>]) {
2260 let handle = FileHandle::from_bytes(OwnedBytes::new(bytes));
2261 let reader = AsyncSSTableReader::<u64>::open(handle, 8).await.unwrap();
2262 assert!(reader.block_index.len() > 3, "test needs several blocks");
2263
2264 for (i, key) in keys.iter().enumerate() {
2266 assert_eq!(reader.get(key).await.unwrap(), Some(i as u64), "key {i}");
2267 let mut before = key.clone();
2268 *before.last_mut().unwrap() -= 1;
2269 let mut after = key.clone();
2270 after.push(0);
2271 assert!(
2272 reader.get(&before).await.unwrap().is_none() || keys.binary_search(&before).is_ok()
2273 );
2274 assert!(
2275 reader.get(&after).await.unwrap().is_none() || keys.binary_search(&after).is_ok()
2276 );
2277 }
2278 assert!(reader.get(b"").await.unwrap().is_none());
2279 assert!(reader.get(b"zzz").await.unwrap().is_none());
2280
2281 let mut it = reader.iter();
2283 let mut seen = Vec::new();
2284 while let Some((k, v)) = it.next().await.unwrap() {
2285 assert_eq!(v as usize, seen.len());
2286 seen.push(k);
2287 }
2288 assert_eq!(seen, keys);
2289 let scanned = reader.prefix_scan(b"field03/").await.unwrap();
2290 let expected: Vec<&Vec<u8>> = keys.iter().filter(|k| k.starts_with(b"field03/")).collect();
2291 assert_eq!(scanned.len(), expected.len());
2292 assert!(scanned.iter().zip(expected).all(|((k, _), e)| k == e));
2293 assert_eq!(reader.all_entries().await.unwrap().len(), keys.len());
2294 let batch: Vec<&[u8]> = keys.iter().step_by(97).map(|k| k.as_slice()).collect();
2295 let got = reader.get_batch(&batch).await.unwrap();
2296 assert!(got.iter().all(|v| v.is_some()));
2297 }
2298
2299 #[cfg(feature = "native")]
2302 #[tokio::test]
2303 async fn v5_restart_points_find_every_key_across_blocks() {
2304 let (bytes, keys) = keyed_table(20_000);
2305 check_every_key_and_scan(bytes, &keys).await;
2306 }
2307
2308 #[cfg(feature = "native")]
2310 #[tokio::test]
2311 async fn unknown_sstable_magic_is_refused() {
2312 let (mut bytes, _) = keyed_table(100);
2313 let n = bytes.len();
2314 bytes[n - 4..].copy_from_slice(&0x5354_4236u32.to_le_bytes()); let handle = FileHandle::from_bytes(OwnedBytes::new(bytes));
2316 let err = match AsyncSSTableReader::<u64>::open(handle, 8).await {
2317 Ok(_) => panic!("unknown magic must be refused"),
2318 Err(err) => err,
2319 };
2320 assert!(err.to_string().contains("incompatible Hermes"), "{err}");
2321 }
2322
2323 #[cfg(feature = "native")]
2326 #[tokio::test]
2327 async fn prefix_scan_matches_prefix_of_smallest_key() {
2328 let mut writer = SSTableWriter::<_, u64>::new(Vec::new());
2329 writer.insert(b"apple", &1).unwrap();
2330 writer.insert(b"apricot", &2).unwrap();
2331 writer.insert(b"banana", &3).unwrap();
2332 let bytes = writer.finish().unwrap();
2333 let handle = FileHandle::from_bytes(OwnedBytes::new(bytes));
2334 let reader = AsyncSSTableReader::<u64>::open(handle, 4).await.unwrap();
2335
2336 let keys = |entries: Vec<(Vec<u8>, u64)>| -> Vec<Vec<u8>> {
2337 entries.into_iter().map(|(k, _)| k).collect()
2338 };
2339
2340 assert_eq!(
2343 keys(reader.prefix_scan(b"ap").await.unwrap()),
2344 vec![b"apple".to_vec(), b"apricot".to_vec()]
2345 );
2346 assert_eq!(
2347 keys(reader.prefix_scan(b"a").await.unwrap()),
2348 vec![b"apple".to_vec(), b"apricot".to_vec()]
2349 );
2350 assert_eq!(
2351 keys(reader.prefix_scan(b"ba").await.unwrap()),
2352 vec![b"banana".to_vec()]
2353 );
2354 assert!(reader.prefix_scan(b"0").await.unwrap().is_empty());
2356
2357 #[cfg(feature = "sync")]
2358 {
2359 assert_eq!(
2360 keys(reader.prefix_scan_sync(b"ap").unwrap()),
2361 vec![b"apple".to_vec(), b"apricot".to_vec()]
2362 );
2363 assert!(reader.prefix_scan_sync(b"0").unwrap().is_empty());
2364 }
2365 }
2366
2367 #[test]
2368 fn test_common_prefix_len() {
2369 assert_eq!(common_prefix_len(b"hello", b"hello"), 5);
2370 assert_eq!(common_prefix_len(b"hello", b"help"), 3);
2371 assert_eq!(common_prefix_len(b"hello", b"world"), 0);
2372 assert_eq!(common_prefix_len(b"", b"hello"), 0);
2373 assert_eq!(common_prefix_len(b"hello", b""), 0);
2374 }
2375
2376 #[test]
2377 fn decode_block_entry_reconstructs_keys_and_consumes_one_entry() {
2378 let mut encoded = Vec::new();
2379
2380 write_vint(&mut encoded, 0).unwrap();
2381 write_vint(&mut encoded, 5).unwrap();
2382 encoded.extend_from_slice(b"alpha");
2383 7_u64.serialize(&mut encoded).unwrap();
2384
2385 write_vint(&mut encoded, 3).unwrap();
2386 write_vint(&mut encoded, 3).unwrap();
2387 encoded.extend_from_slice(b"ine");
2388 11_u64.serialize(&mut encoded).unwrap();
2389
2390 let mut reader = encoded.as_slice();
2391 let mut key = Vec::new();
2392
2393 assert_eq!(decode_block_entry::<u64>(&mut reader, &mut key).unwrap(), 7);
2394 assert_eq!(key, b"alpha");
2395 assert!(!reader.is_empty(), "the second entry must remain unread");
2396
2397 assert_eq!(
2398 decode_block_entry::<u64>(&mut reader, &mut key).unwrap(),
2399 11
2400 );
2401 assert_eq!(key, b"alpine");
2402 assert!(reader.is_empty());
2403 }
2404
2405 #[test]
2406 fn decode_block_entry_rejects_truncated_suffix() {
2407 let encoded = [0, 4, b'o', b'n'];
2408 let mut reader = encoded.as_slice();
2409 let mut key = Vec::new();
2410
2411 let error = decode_block_entry::<u64>(&mut reader, &mut key).unwrap_err();
2412
2413 assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
2414 assert_eq!(error.to_string(), "SSTable block suffix truncated");
2415 }
2416}