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 crate::compression::{CompressionDict, CompressionLevel};
31use crate::directories::{FileHandle, OwnedBytes};
32
33pub const SSTABLE_MAGIC: u32 = 0x53544234; pub const BLOCK_SIZE: usize = 16 * 1024;
39
40pub const DEFAULT_DICT_SIZE: usize = 64 * 1024;
42
43pub const BLOOM_BITS_PER_KEY: usize = 10;
45
46pub const BLOOM_HASH_COUNT: usize = 7;
48const BLOOM_FILTER_HEADER_SIZE: usize = 16;
49
50const MAX_SSTABLE_BLOCK_BYTES: usize = 64 * 1024 * 1024;
51const MAX_SSTABLE_DICTIONARY_BYTES: u64 = 16 * 1024 * 1024;
52
53pub type PrefixScanResult<V> = (Vec<(Vec<u8>, V)>, bool);
55
56#[derive(Debug, Clone)]
62pub struct BloomFilter {
63 bits: BloomBits,
64 num_bits: usize,
65 num_hashes: usize,
66}
67
68#[derive(Debug, Clone)]
70enum BloomBits {
71 Vec(Vec<u64>),
73 Bytes(OwnedBytes),
75}
76
77impl BloomBits {
78 #[inline]
79 fn len(&self) -> usize {
80 match self {
81 BloomBits::Vec(v) => v.len(),
82 BloomBits::Bytes(b) => b.len() / 8,
83 }
84 }
85
86 #[inline]
87 fn get(&self, word_idx: usize) -> u64 {
88 match self {
89 BloomBits::Vec(v) => v[word_idx],
90 BloomBits::Bytes(b) => {
91 let off = word_idx * 8;
92 u64::from_le_bytes([
93 b[off],
94 b[off + 1],
95 b[off + 2],
96 b[off + 3],
97 b[off + 4],
98 b[off + 5],
99 b[off + 6],
100 b[off + 7],
101 ])
102 }
103 }
104 }
105
106 #[inline]
107 fn set_bit(&mut self, word_idx: usize, bit_idx: usize) {
108 match self {
109 BloomBits::Vec(v) => v[word_idx] |= 1u64 << bit_idx,
110 BloomBits::Bytes(_) => panic!("cannot mutate read-only bloom filter"),
111 }
112 }
113
114 fn size_bytes(&self) -> usize {
115 match self {
116 BloomBits::Vec(v) => v.len() * 8,
117 BloomBits::Bytes(b) => b.len(),
118 }
119 }
120
121 fn write_to(&self, writer: &mut (impl Write + ?Sized)) -> io::Result<()> {
122 match self {
123 BloomBits::Vec(words) => {
124 #[cfg(target_endian = "little")]
125 {
126 let bytes = unsafe {
129 std::slice::from_raw_parts(
130 words.as_ptr().cast::<u8>(),
131 words.len().saturating_mul(8),
132 )
133 };
134 writer.write_all(bytes)
135 }
136 #[cfg(target_endian = "big")]
137 {
138 for &word in words {
139 writer.write_u64::<LittleEndian>(word)?;
140 }
141 Ok(())
142 }
143 }
144 BloomBits::Bytes(bytes) => writer.write_all(bytes.as_slice()),
145 }
146 }
147}
148
149impl BloomFilter {
150 pub(crate) const SERIALIZED_HEADER_SIZE: usize = BLOOM_FILTER_HEADER_SIZE;
151
152 pub fn new(expected_keys: usize, bits_per_key: usize) -> Self {
154 let num_bits = expected_keys.saturating_mul(bits_per_key).max(64);
155 let num_words = num_bits.div_ceil(64);
156 Self {
157 bits: BloomBits::Vec(vec![0u64; num_words]),
158 num_bits,
159 num_hashes: BLOOM_HASH_COUNT,
160 }
161 }
162
163 pub fn from_bytes_mutable(data: &[u8]) -> io::Result<Self> {
167 if data.len() < BLOOM_FILTER_HEADER_SIZE {
168 return Err(io::Error::new(
169 io::ErrorKind::InvalidData,
170 "Bloom filter data too short",
171 ));
172 }
173 let num_bits = usize::try_from(u64::from_le_bytes(data[0..8].try_into().unwrap()))
174 .map_err(|_| {
175 io::Error::new(
176 io::ErrorKind::InvalidData,
177 "Bloom filter bit count exceeds addressable memory",
178 )
179 })?;
180 let num_hashes = u32::from_le_bytes(data[8..12].try_into().unwrap()) as usize;
181 let num_words = u32::from_le_bytes(data[12..16].try_into().unwrap()) as usize;
182
183 validate_bloom_header(data.len(), num_bits, num_hashes, num_words)?;
184 let expected_len = BLOOM_FILTER_HEADER_SIZE + num_words * 8;
185 if data.len() != expected_len {
186 return Err(io::Error::new(
187 io::ErrorKind::InvalidData,
188 "Bloom filter data truncated",
189 ));
190 }
191
192 let mut vec = vec![0u64; num_words];
193 for (i, v) in vec.iter_mut().enumerate() {
194 let off = BLOOM_FILTER_HEADER_SIZE + i * 8;
195 *v = u64::from_le_bytes(data[off..off + 8].try_into().unwrap());
196 }
197
198 Ok(Self {
199 bits: BloomBits::Vec(vec),
200 num_bits,
201 num_hashes,
202 })
203 }
204
205 pub fn from_owned_bytes(data: OwnedBytes) -> io::Result<Self> {
207 if data.len() < BLOOM_FILTER_HEADER_SIZE {
208 return Err(io::Error::new(
209 io::ErrorKind::InvalidData,
210 "Bloom filter data too short",
211 ));
212 }
213 let d = data.as_slice();
214 let num_bits =
215 usize::try_from(u64::from_le_bytes(d[0..8].try_into().unwrap())).map_err(|_| {
216 io::Error::new(
217 io::ErrorKind::InvalidData,
218 "Bloom filter bit count exceeds addressable memory",
219 )
220 })?;
221 let num_hashes = u32::from_le_bytes(d[8..12].try_into().unwrap()) as usize;
222 let num_words = u32::from_le_bytes(d[12..16].try_into().unwrap()) as usize;
223
224 validate_bloom_header(d.len(), num_bits, num_hashes, num_words)?;
225 let expected_len = BLOOM_FILTER_HEADER_SIZE + num_words * 8;
226 if d.len() != expected_len {
227 return Err(io::Error::new(
228 io::ErrorKind::InvalidData,
229 "Bloom filter data truncated",
230 ));
231 }
232
233 let bits_bytes =
235 data.slice(BLOOM_FILTER_HEADER_SIZE..BLOOM_FILTER_HEADER_SIZE + num_words * 8);
236
237 Ok(Self {
238 bits: BloomBits::Bytes(bits_bytes),
239 num_bits,
240 num_hashes,
241 })
242 }
243
244 pub fn serialized_len(&self) -> usize {
246 BLOOM_FILTER_HEADER_SIZE + self.bits.len() * 8
247 }
248
249 pub fn write_to(&self, writer: &mut (impl Write + ?Sized)) -> io::Result<()> {
251 let num_words = self.bits.len();
252 write_bloom_header(writer, self.num_bits, self.num_hashes, num_words)?;
253 self.bits.write_to(writer)
254 }
255
256 pub fn to_bytes(&self) -> Vec<u8> {
258 let mut data = Vec::with_capacity(self.serialized_len());
259 self.write_to(&mut data)
260 .expect("writing a bloom filter to Vec cannot fail");
261 data
262 }
263
264 pub fn insert(&mut self, key: &[u8]) {
266 let (h1, h2) = self.hash_pair(key);
267 for i in 0..self.num_hashes {
268 let bit_pos = self.get_bit_pos(h1, h2, i);
269 let word_idx = bit_pos / 64;
270 let bit_idx = bit_pos % 64;
271 if word_idx < self.bits.len() {
272 self.bits.set_bit(word_idx, bit_idx);
273 }
274 }
275 }
276
277 pub fn may_contain(&self, key: &[u8]) -> bool {
280 let (h1, h2) = self.hash_pair(key);
281 for i in 0..self.num_hashes {
282 let bit_pos = self.get_bit_pos(h1, h2, i);
283 let word_idx = bit_pos / 64;
284 let bit_idx = bit_pos % 64;
285 if word_idx >= self.bits.len() || (self.bits.get(word_idx) & (1u64 << bit_idx)) == 0 {
286 return false;
287 }
288 }
289 true
290 }
291
292 pub fn size_bytes(&self) -> usize {
294 BLOOM_FILTER_HEADER_SIZE + self.bits.size_bytes()
295 }
296
297 pub fn insert_hashed(&mut self, h1: u64, h2: u64) {
299 for i in 0..self.num_hashes {
300 let bit_pos = self.get_bit_pos(h1, h2, i);
301 let word_idx = bit_pos / 64;
302 let bit_idx = bit_pos % 64;
303 if word_idx < self.bits.len() {
304 self.bits.set_bit(word_idx, bit_idx);
305 }
306 }
307 }
308
309 #[inline]
311 fn hash_pair(&self, key: &[u8]) -> (u64, u64) {
312 let mut h1: u64 = 0xcbf29ce484222325;
313 let mut h2: u64 = 0x84222325cbf29ce4;
314 for &byte in key {
315 h1 ^= byte as u64;
316 h1 = h1.wrapping_mul(0x100000001b3);
317 h2 = h2.wrapping_mul(0x100000001b3);
318 h2 ^= byte as u64;
319 }
320 (h1, h2)
321 }
322
323 #[inline]
325 fn get_bit_pos(&self, h1: u64, h2: u64, i: usize) -> usize {
326 (h1.wrapping_add((i as u64).wrapping_mul(h2)) % (self.num_bits as u64)) as usize
327 }
328}
329
330fn write_bloom_header(
331 writer: &mut (impl Write + ?Sized),
332 num_bits: usize,
333 num_hashes: usize,
334 num_words: usize,
335) -> io::Result<()> {
336 writer.write_u64::<LittleEndian>(u64::try_from(num_bits).map_err(|_| {
337 io::Error::new(
338 io::ErrorKind::InvalidInput,
339 "Bloom filter bit count exceeds u64",
340 )
341 })?)?;
342 writer.write_u32::<LittleEndian>(u32::try_from(num_hashes).map_err(|_| {
343 io::Error::new(
344 io::ErrorKind::InvalidInput,
345 "Bloom filter hash count exceeds u32",
346 )
347 })?)?;
348 writer.write_u32::<LittleEndian>(u32::try_from(num_words).map_err(|_| {
349 io::Error::new(
350 io::ErrorKind::InvalidInput,
351 "Bloom filter word count exceeds u32",
352 )
353 })?)?;
354 Ok(())
355}
356
357fn validate_bloom_header(
358 data_len: usize,
359 num_bits: usize,
360 num_hashes: usize,
361 num_words: usize,
362) -> io::Result<()> {
363 if num_bits == 0 || num_hashes == 0 || num_hashes > 32 || num_words == 0 {
364 return Err(io::Error::new(
365 io::ErrorKind::InvalidData,
366 "invalid bloom filter parameters",
367 ));
368 }
369 let word_bytes = num_words
370 .checked_mul(8)
371 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
372 let expected_len = BLOOM_FILTER_HEADER_SIZE
373 .checked_add(word_bytes)
374 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
375 let capacity_bits = num_words
376 .checked_mul(64)
377 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
378 if expected_len > data_len
379 || num_bits > capacity_bits
380 || num_bits <= capacity_bits.saturating_sub(64)
381 {
382 return Err(io::Error::new(
383 io::ErrorKind::InvalidData,
384 "inconsistent bloom filter dimensions",
385 ));
386 }
387 Ok(())
388}
389
390#[inline]
393fn bloom_hash_pair(key: &[u8]) -> (u64, u64) {
394 let mut h1: u64 = 0xcbf29ce484222325;
395 let mut h2: u64 = 0x84222325cbf29ce4;
396 for &byte in key {
397 h1 ^= byte as u64;
398 h1 = h1.wrapping_mul(0x100000001b3);
399 h2 = h2.wrapping_mul(0x100000001b3);
400 h2 ^= byte as u64;
401 }
402 (h1, h2)
403}
404
405pub trait SSTableValue: Clone + Send + Sync {
407 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()>;
408 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self>;
409}
410
411impl SSTableValue for u64 {
413 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
414 write_vint(writer, *self)
415 }
416
417 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
418 read_vint(reader)
419 }
420}
421
422impl SSTableValue for Vec<u8> {
424 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
425 write_vint(writer, self.len() as u64)?;
426 writer.write_all(self)
427 }
428
429 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
430 let len = usize::try_from(read_vint(reader)?).map_err(|_| {
431 io::Error::new(io::ErrorKind::InvalidData, "SSTable value length too large")
432 })?;
433 if len > MAX_SSTABLE_BLOCK_BYTES {
434 return Err(io::Error::new(
435 io::ErrorKind::InvalidData,
436 "SSTable value exceeds block safety limit",
437 ));
438 }
439 let mut data = vec![0u8; len];
440 reader.read_exact(&mut data)?;
441 Ok(data)
442 }
443}
444
445#[derive(Debug, Clone, Copy, PartialEq, Eq)]
448pub struct SparseDimInfo {
449 pub offset: u64,
451 pub length: u32,
453}
454
455impl SparseDimInfo {
456 pub fn new(offset: u64, length: u32) -> Self {
457 Self { offset, length }
458 }
459}
460
461impl SSTableValue for SparseDimInfo {
462 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
463 write_vint(writer, self.offset)?;
464 write_vint(writer, self.length as u64)
465 }
466
467 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
468 let offset = read_vint(reader)?;
469 let length = u32::try_from(read_vint(reader)?).map_err(|_| {
470 io::Error::new(
471 io::ErrorKind::InvalidData,
472 "sparse posting length exceeds u32",
473 )
474 })?;
475 Ok(Self { offset, length })
476 }
477}
478
479pub const MAX_INLINE_POSTINGS: usize = 3;
481
482#[derive(Debug, Clone, PartialEq, Eq)]
490pub enum TermInfo {
491 Inline {
494 doc_freq: u8,
496 data: [u8; 16],
499 data_len: u8,
501 },
502 External {
504 posting_offset: u64,
505 posting_len: u64,
506 doc_freq: u32,
507 position_offset: u64,
509 position_len: u64,
511 },
512}
513
514impl TermInfo {
515 pub fn external(posting_offset: u64, posting_len: u64, doc_freq: u32) -> Self {
517 TermInfo::External {
518 posting_offset,
519 posting_len,
520 doc_freq,
521 position_offset: 0,
522 position_len: 0,
523 }
524 }
525
526 pub fn external_with_positions(
528 posting_offset: u64,
529 posting_len: u64,
530 doc_freq: u32,
531 position_offset: u64,
532 position_len: u64,
533 ) -> Self {
534 TermInfo::External {
535 posting_offset,
536 posting_len,
537 doc_freq,
538 position_offset,
539 position_len,
540 }
541 }
542
543 pub fn try_inline(doc_ids: &[u32], term_freqs: &[u32]) -> Option<Self> {
546 if doc_ids.len() > MAX_INLINE_POSTINGS
547 || doc_ids.is_empty()
548 || doc_ids.len() != term_freqs.len()
549 {
550 return None;
551 }
552
553 let mut data = [0u8; 16];
554 let mut cursor = std::io::Cursor::new(&mut data[..]);
555 let mut prev_doc_id = 0u32;
556
557 for (i, &doc_id) in doc_ids.iter().enumerate() {
558 let delta = doc_id.checked_sub(prev_doc_id)?;
559 if write_vint(&mut cursor, delta as u64).is_err() {
560 return None;
561 }
562 if write_vint(&mut cursor, term_freqs[i] as u64).is_err() {
563 return None;
564 }
565 prev_doc_id = doc_id;
566 }
567
568 let data_len = cursor.position() as u8;
569 if data_len > 16 {
570 return None;
571 }
572
573 Some(TermInfo::Inline {
574 doc_freq: doc_ids.len() as u8,
575 data,
576 data_len,
577 })
578 }
579
580 pub fn try_inline_iter(count: usize, iter: impl Iterator<Item = (u32, u32)>) -> Option<Self> {
584 if count > MAX_INLINE_POSTINGS || count == 0 {
585 return None;
586 }
587
588 let mut data = [0u8; 16];
589 let mut cursor = std::io::Cursor::new(&mut data[..]);
590 let mut prev_doc_id = 0u32;
591
592 let mut actual_count = 0usize;
593 for (doc_id, tf) in iter {
594 if actual_count >= count {
595 return None;
596 }
597 let delta = doc_id.checked_sub(prev_doc_id)?;
598 if write_vint(&mut cursor, delta as u64).is_err() {
599 return None;
600 }
601 if write_vint(&mut cursor, tf as u64).is_err() {
602 return None;
603 }
604 prev_doc_id = doc_id;
605 actual_count += 1;
606 }
607
608 if actual_count != count {
609 return None;
610 }
611
612 let data_len = cursor.position() as u8;
613
614 Some(TermInfo::Inline {
615 doc_freq: count as u8,
616 data,
617 data_len,
618 })
619 }
620
621 pub fn doc_freq(&self) -> u32 {
623 match self {
624 TermInfo::Inline { doc_freq, .. } => *doc_freq as u32,
625 TermInfo::External { doc_freq, .. } => *doc_freq,
626 }
627 }
628
629 pub fn is_inline(&self) -> bool {
631 matches!(self, TermInfo::Inline { .. })
632 }
633
634 pub fn external_info(&self) -> Option<(u64, u64)> {
636 match self {
637 TermInfo::External {
638 posting_offset,
639 posting_len,
640 ..
641 } => Some((*posting_offset, *posting_len)),
642 TermInfo::Inline { .. } => None,
643 }
644 }
645
646 pub fn position_info(&self) -> Option<(u64, u64)> {
648 match self {
649 TermInfo::External {
650 position_offset,
651 position_len,
652 ..
653 } if *position_len > 0 => Some((*position_offset, *position_len)),
654 _ => None,
655 }
656 }
657
658 pub fn decode_inline(&self) -> Option<(Vec<u32>, Vec<u32>)> {
661 match self {
662 TermInfo::Inline {
663 doc_freq,
664 data,
665 data_len,
666 } => {
667 if *doc_freq == 0
668 || *doc_freq as usize > MAX_INLINE_POSTINGS
669 || *data_len as usize > data.len()
670 {
671 return None;
672 }
673 let mut doc_ids = Vec::with_capacity(*doc_freq as usize);
674 let mut term_freqs = Vec::with_capacity(*doc_freq as usize);
675 let mut reader = &data[..*data_len as usize];
676 let mut prev_doc_id = 0u32;
677
678 for _ in 0..*doc_freq {
679 let delta = read_vint(&mut reader).ok()? as u32;
680 let tf = read_vint(&mut reader).ok()? as u32;
681 let doc_id = prev_doc_id.checked_add(delta)?;
682 doc_ids.push(doc_id);
683 term_freqs.push(tf);
684 prev_doc_id = doc_id;
685 }
686
687 Some((doc_ids, term_freqs))
688 }
689 TermInfo::External { .. } => None,
690 }
691 }
692}
693
694impl SSTableValue for TermInfo {
695 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
696 match self {
697 TermInfo::Inline {
698 doc_freq,
699 data,
700 data_len,
701 } => {
702 if *doc_freq == 0
703 || *doc_freq as usize > MAX_INLINE_POSTINGS
704 || *data_len as usize > data.len()
705 {
706 return Err(io::Error::new(
707 io::ErrorKind::InvalidInput,
708 "invalid inline TermInfo",
709 ));
710 }
711 writer.write_u8(0xFF)?;
713 writer.write_u8(*doc_freq)?;
714 writer.write_u8(*data_len)?;
715 writer.write_all(&data[..*data_len as usize])?;
716 }
717 TermInfo::External {
718 posting_offset,
719 posting_len,
720 doc_freq,
721 position_offset,
722 position_len,
723 } => {
724 if *position_len > 0 {
727 writer.write_u8(0x01)?;
728 write_vint(writer, *doc_freq as u64)?;
729 write_vint(writer, *posting_offset)?;
730 write_vint(writer, *posting_len)?;
731 write_vint(writer, *position_offset)?;
732 write_vint(writer, *position_len)?;
733 } else {
734 writer.write_u8(0x00)?;
735 write_vint(writer, *doc_freq as u64)?;
736 write_vint(writer, *posting_offset)?;
737 write_vint(writer, *posting_len)?;
738 }
739 }
740 }
741 Ok(())
742 }
743
744 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
745 let tag = reader.read_u8()?;
746
747 if tag == 0xFF {
748 let doc_freq = reader.read_u8()?;
750 let data_len = reader.read_u8()?;
751 if doc_freq == 0 || doc_freq as usize > MAX_INLINE_POSTINGS || data_len as usize > 16 {
752 return Err(io::Error::new(
753 io::ErrorKind::InvalidData,
754 "invalid inline TermInfo lengths",
755 ));
756 }
757 let mut data = [0u8; 16];
758 reader.read_exact(&mut data[..data_len as usize])?;
759 Ok(TermInfo::Inline {
760 doc_freq,
761 data,
762 data_len,
763 })
764 } else if tag == 0x00 {
765 let doc_freq = read_vint(reader)? as u32;
767 let posting_offset = read_vint(reader)?;
768 let posting_len = read_vint(reader)?;
769 Ok(TermInfo::External {
770 posting_offset,
771 posting_len,
772 doc_freq,
773 position_offset: 0,
774 position_len: 0,
775 })
776 } else if tag == 0x01 {
777 let doc_freq = read_vint(reader)? as u32;
779 let posting_offset = read_vint(reader)?;
780 let posting_len = read_vint(reader)?;
781 let position_offset = read_vint(reader)?;
782 let position_len = read_vint(reader)?;
783 Ok(TermInfo::External {
784 posting_offset,
785 posting_len,
786 doc_freq,
787 position_offset,
788 position_len,
789 })
790 } else {
791 Err(io::Error::new(
792 io::ErrorKind::InvalidData,
793 format!("Invalid TermInfo tag: {}", tag),
794 ))
795 }
796 }
797}
798
799pub fn write_vint<W: Write + ?Sized>(writer: &mut W, mut value: u64) -> io::Result<()> {
801 loop {
802 let byte = (value & 0x7F) as u8;
803 value >>= 7;
804 if value == 0 {
805 writer.write_u8(byte)?;
806 return Ok(());
807 } else {
808 writer.write_u8(byte | 0x80)?;
809 }
810 }
811}
812
813pub fn read_vint<R: Read>(reader: &mut R) -> io::Result<u64> {
815 let mut result = 0u64;
816 let mut shift = 0;
817
818 loop {
819 let byte = reader.read_u8()?;
820 result |= ((byte & 0x7F) as u64) << shift;
821 if byte & 0x80 == 0 {
822 return Ok(result);
823 }
824 shift += 7;
825 if shift >= 64 {
826 return Err(io::Error::new(
827 io::ErrorKind::InvalidData,
828 "varint too long",
829 ));
830 }
831 }
832}
833
834pub fn common_prefix_len(a: &[u8], b: &[u8]) -> usize {
836 a.iter().zip(b.iter()).take_while(|(x, y)| x == y).count()
837}
838
839#[derive(Debug, Clone)]
841pub struct SSTableStats {
842 pub num_blocks: usize,
843 pub num_sparse_entries: usize,
844 pub num_entries: u64,
845 pub has_bloom_filter: bool,
846 pub has_dictionary: bool,
847 pub bloom_filter_size: usize,
848 pub dictionary_size: usize,
849}
850
851#[derive(Debug, Clone)]
853pub struct SSTableWriterConfig {
854 pub compression_level: CompressionLevel,
856 pub use_dictionary: bool,
858 pub dict_size: usize,
860 pub use_bloom_filter: bool,
862 pub bloom_bits_per_key: usize,
864}
865
866impl Default for SSTableWriterConfig {
867 fn default() -> Self {
868 Self::from_optimization(crate::structures::IndexOptimization::default())
869 }
870}
871
872impl SSTableWriterConfig {
873 pub fn from_optimization(optimization: crate::structures::IndexOptimization) -> Self {
875 use crate::structures::IndexOptimization;
876 match optimization {
877 IndexOptimization::Adaptive => Self {
878 compression_level: CompressionLevel::BETTER, use_dictionary: false,
880 dict_size: DEFAULT_DICT_SIZE,
881 use_bloom_filter: true, bloom_bits_per_key: BLOOM_BITS_PER_KEY,
883 },
884 IndexOptimization::SizeOptimized => Self {
885 compression_level: CompressionLevel::MAX, use_dictionary: true,
887 dict_size: DEFAULT_DICT_SIZE,
888 use_bloom_filter: true,
889 bloom_bits_per_key: BLOOM_BITS_PER_KEY,
890 },
891 IndexOptimization::PerformanceOptimized => Self {
892 compression_level: CompressionLevel::FAST, use_dictionary: false,
894 dict_size: DEFAULT_DICT_SIZE,
895 use_bloom_filter: true, bloom_bits_per_key: BLOOM_BITS_PER_KEY,
897 },
898 }
899 }
900
901 pub fn fast() -> Self {
903 Self::from_optimization(crate::structures::IndexOptimization::PerformanceOptimized)
904 }
905
906 pub fn max_compression() -> Self {
908 Self::from_optimization(crate::structures::IndexOptimization::SizeOptimized)
909 }
910}
911
912pub struct SSTableWriter<W: Write, V: SSTableValue> {
918 writer: W,
919 block_buffer: Vec<u8>,
920 prev_key: Vec<u8>,
921 index: Vec<BlockIndexEntry>,
922 current_offset: u64,
923 num_entries: u64,
924 block_first_key: Option<Vec<u8>>,
925 config: SSTableWriterConfig,
926 dictionary: Option<CompressionDict>,
928 bloom_hashes: Vec<(u64, u64)>,
931 _phantom: std::marker::PhantomData<V>,
932}
933
934impl<W: Write, V: SSTableValue> SSTableWriter<W, V> {
935 pub fn new(writer: W) -> Self {
937 Self::with_config(writer, SSTableWriterConfig::default())
938 }
939
940 pub fn with_config(writer: W, config: SSTableWriterConfig) -> Self {
942 Self {
943 writer,
944 block_buffer: Vec::with_capacity(BLOCK_SIZE),
945 prev_key: Vec::new(),
946 index: Vec::new(),
947 current_offset: 0,
948 num_entries: 0,
949 block_first_key: None,
950 config,
951 dictionary: None,
952 bloom_hashes: Vec::new(),
953 _phantom: std::marker::PhantomData,
954 }
955 }
956
957 pub fn with_dictionary(
959 writer: W,
960 config: SSTableWriterConfig,
961 dictionary: CompressionDict,
962 ) -> Self {
963 Self {
964 writer,
965 block_buffer: Vec::with_capacity(BLOCK_SIZE),
966 prev_key: Vec::new(),
967 index: Vec::new(),
968 current_offset: 0,
969 num_entries: 0,
970 block_first_key: None,
971 config,
972 dictionary: Some(dictionary),
973 bloom_hashes: Vec::new(),
974 _phantom: std::marker::PhantomData,
975 }
976 }
977
978 pub fn insert(&mut self, key: &[u8], value: &V) -> io::Result<()> {
979 if self.block_first_key.is_none() {
980 self.block_first_key = Some(key.to_vec());
981 }
982
983 if self.config.use_bloom_filter {
985 self.bloom_hashes.push(bloom_hash_pair(key));
986 }
987
988 let prefix_len = common_prefix_len(&self.prev_key, key);
989 let suffix = &key[prefix_len..];
990
991 write_vint(&mut self.block_buffer, prefix_len as u64)?;
992 write_vint(&mut self.block_buffer, suffix.len() as u64)?;
993 self.block_buffer.extend_from_slice(suffix);
994 value.serialize(&mut self.block_buffer)?;
995
996 self.prev_key.clear();
997 self.prev_key.extend_from_slice(key);
998 self.num_entries += 1;
999
1000 if self.block_buffer.len() >= BLOCK_SIZE {
1001 self.flush_block()?;
1002 }
1003
1004 Ok(())
1005 }
1006
1007 fn flush_block(&mut self) -> io::Result<()> {
1009 if self.block_buffer.is_empty() {
1010 return Ok(());
1011 }
1012
1013 let compressed = if let Some(ref dict) = self.dictionary {
1015 crate::compression::compress_with_dict(
1016 &self.block_buffer,
1017 self.config.compression_level,
1018 dict,
1019 )?
1020 } else {
1021 crate::compression::compress(&self.block_buffer, self.config.compression_level)?
1022 };
1023
1024 if let Some(first_key) = self.block_first_key.take() {
1025 self.index.push(BlockIndexEntry {
1026 first_key,
1027 offset: self.current_offset,
1028 length: compressed.len() as u32,
1029 });
1030 }
1031
1032 self.writer.write_all(&compressed)?;
1033 self.current_offset += compressed.len() as u64;
1034 self.block_buffer.clear();
1035 self.prev_key.clear();
1036
1037 Ok(())
1038 }
1039
1040 pub fn finish(mut self) -> io::Result<W> {
1041 self.flush_block()?;
1043
1044 let bloom_filter = if self.config.use_bloom_filter && !self.bloom_hashes.is_empty() {
1046 let mut bloom =
1047 BloomFilter::new(self.bloom_hashes.len(), self.config.bloom_bits_per_key);
1048 for (h1, h2) in &self.bloom_hashes {
1049 bloom.insert_hashed(*h1, *h2);
1050 }
1051 Some(bloom)
1052 } else {
1053 None
1054 };
1055
1056 let data_end_offset = self.current_offset;
1057
1058 let entries: Vec<(Vec<u8>, BlockAddr)> = self
1061 .index
1062 .iter()
1063 .map(|e| {
1064 (
1065 e.first_key.clone(),
1066 BlockAddr {
1067 offset: e.offset,
1068 length: e.length,
1069 },
1070 )
1071 })
1072 .collect();
1073
1074 #[cfg(feature = "native")]
1076 let index_bytes = FstBlockIndex::build(&entries)?;
1077 #[cfg(not(feature = "native"))]
1078 let index_bytes = MmapBlockIndex::build(&entries)?;
1079
1080 self.writer
1082 .write_u32::<LittleEndian>(index_bytes.len() as u32)?;
1083 self.writer.write_all(&index_bytes)?;
1084 self.current_offset += 4 + index_bytes.len() as u64;
1085
1086 let bloom_offset = if let Some(ref bloom) = bloom_filter {
1088 let offset = self.current_offset;
1089 bloom.write_to(&mut self.writer)?;
1090 self.current_offset += bloom.serialized_len() as u64;
1091 offset
1092 } else {
1093 0
1094 };
1095
1096 let dict_offset = if let Some(ref dict) = self.dictionary {
1098 let dict_bytes = dict.as_bytes();
1099 let offset = self.current_offset;
1100 self.writer
1101 .write_u32::<LittleEndian>(dict_bytes.len() as u32)?;
1102 self.writer.write_all(dict_bytes)?;
1103 self.current_offset += 4 + dict_bytes.len() as u64;
1104 offset
1105 } else {
1106 0
1107 };
1108
1109 self.writer.write_u64::<LittleEndian>(data_end_offset)?;
1111 self.writer.write_u64::<LittleEndian>(self.num_entries)?;
1112 self.writer.write_u64::<LittleEndian>(bloom_offset)?; self.writer.write_u64::<LittleEndian>(dict_offset)?; self.writer
1115 .write_u8(self.config.compression_level.0 as u8)?;
1116 self.writer.write_u32::<LittleEndian>(SSTABLE_MAGIC)?;
1117
1118 Ok(self.writer)
1119 }
1120}
1121
1122#[derive(Debug, Clone)]
1124struct BlockIndexEntry {
1125 first_key: Vec<u8>,
1126 offset: u64,
1127 length: u32,
1128}
1129
1130pub struct AsyncSSTableReader<V: SSTableValue> {
1137 data_slice: FileHandle,
1139 block_index: BlockIndex,
1141 num_entries: u64,
1142 cache: RwLock<BlockCache>,
1144 bloom_filter: Option<BloomFilter>,
1146 dictionary: Option<CompressionDict>,
1148 #[allow(dead_code)]
1150 compression_level: CompressionLevel,
1151 _phantom: std::marker::PhantomData<V>,
1152}
1153
1154struct BlockCache {
1160 blocks: FxHashMap<u64, Arc<[u8]>>,
1161 lru_order: std::collections::VecDeque<u64>,
1162 max_blocks: usize,
1163}
1164
1165impl BlockCache {
1166 fn new(max_blocks: usize) -> Self {
1167 Self {
1168 blocks: FxHashMap::default(),
1169 lru_order: std::collections::VecDeque::with_capacity(max_blocks),
1170 max_blocks,
1171 }
1172 }
1173
1174 fn get(&mut self, offset: u64) -> Option<Arc<[u8]>> {
1175 if self.blocks.contains_key(&offset) {
1176 self.promote(offset);
1177 self.blocks.get(&offset).map(Arc::clone)
1178 } else {
1179 None
1180 }
1181 }
1182
1183 fn peek(&self, offset: u64) -> Option<Arc<[u8]>> {
1185 self.blocks.get(&offset).map(Arc::clone)
1186 }
1187
1188 fn insert(&mut self, offset: u64, block: Arc<[u8]>) {
1189 if self.max_blocks == 0 {
1190 return;
1191 }
1192 if self.blocks.contains_key(&offset) {
1193 self.promote(offset);
1194 return;
1195 }
1196 while self.blocks.len() >= self.max_blocks {
1197 if let Some(evict_offset) = self.lru_order.pop_front() {
1198 self.blocks.remove(&evict_offset);
1199 } else {
1200 break;
1201 }
1202 }
1203 self.blocks.insert(offset, block);
1204 self.lru_order.push_back(offset);
1205 }
1206
1207 fn promote(&mut self, offset: u64) {
1209 if let Some(pos) = self.lru_order.iter().position(|&k| k == offset) {
1210 self.lru_order.remove(pos);
1211 self.lru_order.push_back(offset);
1212 }
1213 }
1214}
1215
1216impl<V: SSTableValue> AsyncSSTableReader<V> {
1217 pub async fn open(file_handle: FileHandle, cache_blocks: usize) -> io::Result<Self> {
1222 let file_len = file_handle.len();
1223 if file_len < 37 {
1224 return Err(io::Error::new(
1225 io::ErrorKind::InvalidData,
1226 "SSTable too small",
1227 ));
1228 }
1229
1230 let footer_bytes = file_handle
1233 .read_bytes_range(file_len - 37..file_len)
1234 .await?;
1235
1236 let mut reader = footer_bytes.as_slice();
1237 let data_end_offset = reader.read_u64::<LittleEndian>()?;
1238 let num_entries = reader.read_u64::<LittleEndian>()?;
1239 let bloom_offset = reader.read_u64::<LittleEndian>()?;
1240 let dict_offset = reader.read_u64::<LittleEndian>()?;
1241 let compression_level = CompressionLevel(reader.read_u8()? as i32);
1242 let magic = reader.read_u32::<LittleEndian>()?;
1243
1244 if magic != SSTABLE_MAGIC {
1245 return Err(io::Error::new(
1246 io::ErrorKind::InvalidData,
1247 format!("Invalid SSTable magic: 0x{:08X}", magic),
1248 ));
1249 }
1250
1251 let footer_start = file_len - 37;
1252 if data_end_offset > footer_start {
1253 return Err(io::Error::new(
1254 io::ErrorKind::InvalidData,
1255 "SSTable data section extends past its footer",
1256 ));
1257 }
1258 if bloom_offset != 0 && (bloom_offset < data_end_offset || bloom_offset >= footer_start) {
1259 return Err(io::Error::new(
1260 io::ErrorKind::InvalidData,
1261 "SSTable bloom filter offset is out of bounds",
1262 ));
1263 }
1264 if dict_offset != 0
1265 && (dict_offset < data_end_offset
1266 || dict_offset >= footer_start
1267 || (bloom_offset != 0 && dict_offset <= bloom_offset))
1268 {
1269 return Err(io::Error::new(
1270 io::ErrorKind::InvalidData,
1271 "SSTable dictionary offset is out of bounds",
1272 ));
1273 }
1274
1275 let index_start = data_end_offset;
1277 let index_end = if bloom_offset != 0 {
1278 bloom_offset
1279 } else if dict_offset != 0 {
1280 dict_offset
1281 } else {
1282 footer_start
1283 };
1284 let index_bytes = file_handle.read_bytes_range(index_start..index_end).await?;
1285
1286 let mut idx_reader = index_bytes.as_slice();
1288 let index_len = idx_reader.read_u32::<LittleEndian>()? as usize;
1289
1290 if index_len != idx_reader.len() {
1291 return Err(io::Error::new(
1292 io::ErrorKind::InvalidData,
1293 "Index data truncated",
1294 ));
1295 }
1296
1297 let index_data = index_bytes.slice(4..4 + index_len);
1298
1299 #[cfg(feature = "fst-index")]
1301 let block_index = match FstBlockIndex::load(index_data.clone()) {
1302 Ok(fst_idx) => BlockIndex::Fst(fst_idx),
1303 Err(_) => BlockIndex::Mmap(MmapBlockIndex::load(index_data)?),
1304 };
1305 #[cfg(not(feature = "fst-index"))]
1306 let block_index = BlockIndex::Mmap(MmapBlockIndex::load(index_data)?);
1307
1308 let mut expected_offset = 0u64;
1309 for addr in block_index.all_addrs() {
1310 let end = addr.offset.checked_add(addr.length as u64).ok_or_else(|| {
1311 io::Error::new(io::ErrorKind::InvalidData, "SSTable block range overflow")
1312 })?;
1313 if addr.length == 0 || addr.offset != expected_offset || end > data_end_offset {
1314 return Err(io::Error::new(
1315 io::ErrorKind::InvalidData,
1316 "SSTable block addresses are inconsistent",
1317 ));
1318 }
1319 expected_offset = end;
1320 }
1321 if expected_offset != data_end_offset {
1322 return Err(io::Error::new(
1323 io::ErrorKind::InvalidData,
1324 "SSTable block addresses do not cover the data section",
1325 ));
1326 }
1327
1328 let bloom_filter = if bloom_offset > 0 {
1330 let bloom_start = bloom_offset;
1331 let bloom_end = if dict_offset != 0 {
1332 dict_offset
1333 } else {
1334 footer_start
1335 };
1336 let header_size = BloomFilter::SERIALIZED_HEADER_SIZE as u64;
1338 let bloom_header_end = bloom_start.checked_add(header_size).ok_or_else(|| {
1339 io::Error::new(io::ErrorKind::InvalidData, "bloom filter range overflow")
1340 })?;
1341 if bloom_header_end > bloom_end {
1342 return Err(io::Error::new(
1343 io::ErrorKind::UnexpectedEof,
1344 "bloom filter header is truncated",
1345 ));
1346 }
1347 let bloom_header = file_handle
1348 .read_bytes_range(bloom_start..bloom_header_end)
1349 .await?;
1350 let num_words = u32::from_le_bytes([
1351 bloom_header[12],
1352 bloom_header[13],
1353 bloom_header[14],
1354 bloom_header[15],
1355 ]) as u64;
1356 let bloom_size = num_words
1357 .checked_mul(8)
1358 .and_then(|bytes| bytes.checked_add(header_size))
1359 .ok_or_else(|| {
1360 io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow")
1361 })?;
1362 let actual_bloom_size = bloom_end - bloom_start;
1363 if bloom_size != actual_bloom_size {
1364 return Err(io::Error::new(
1365 io::ErrorKind::InvalidData,
1366 "bloom filter length is inconsistent",
1367 ));
1368 }
1369 let bloom_data = file_handle.read_bytes_range(bloom_start..bloom_end).await?;
1370 Some(BloomFilter::from_owned_bytes(bloom_data)?)
1371 } else {
1372 None
1373 };
1374
1375 let dictionary = if dict_offset > 0 {
1377 let dict_start = dict_offset;
1378 let dict_header_end = dict_start.checked_add(4).ok_or_else(|| {
1380 io::Error::new(io::ErrorKind::InvalidData, "dictionary range overflow")
1381 })?;
1382 if dict_header_end > footer_start {
1383 return Err(io::Error::new(
1384 io::ErrorKind::UnexpectedEof,
1385 "dictionary header is truncated",
1386 ));
1387 }
1388 let dict_len_bytes = file_handle
1389 .read_bytes_range(dict_start..dict_header_end)
1390 .await?;
1391 let dict_len = u32::from_le_bytes([
1392 dict_len_bytes[0],
1393 dict_len_bytes[1],
1394 dict_len_bytes[2],
1395 dict_len_bytes[3],
1396 ]) as u64;
1397 if dict_len > MAX_SSTABLE_DICTIONARY_BYTES {
1398 return Err(io::Error::new(
1399 io::ErrorKind::InvalidData,
1400 "SSTable dictionary exceeds safety limit",
1401 ));
1402 }
1403 let dict_end = dict_header_end.checked_add(dict_len).ok_or_else(|| {
1404 io::Error::new(io::ErrorKind::InvalidData, "dictionary range overflow")
1405 })?;
1406 if dict_end != footer_start {
1407 return Err(io::Error::new(
1408 io::ErrorKind::InvalidData,
1409 "SSTable dictionary length is inconsistent",
1410 ));
1411 }
1412 let dict_data = file_handle
1413 .read_bytes_range(dict_header_end..dict_end)
1414 .await?;
1415 Some(CompressionDict::from_owned_bytes(dict_data))
1416 } else {
1417 None
1418 };
1419
1420 let data_slice = file_handle.slice(0..data_end_offset);
1422
1423 Ok(Self {
1424 data_slice,
1425 block_index,
1426 num_entries,
1427 cache: RwLock::new(BlockCache::new(cache_blocks)),
1428 bloom_filter,
1429 dictionary,
1430 compression_level,
1431 _phantom: std::marker::PhantomData,
1432 })
1433 }
1434
1435 pub fn num_entries(&self) -> u64 {
1437 self.num_entries
1438 }
1439
1440 pub fn stats(&self) -> SSTableStats {
1442 SSTableStats {
1443 num_blocks: self.block_index.len(),
1444 num_sparse_entries: 0, num_entries: self.num_entries,
1446 has_bloom_filter: self.bloom_filter.is_some(),
1447 has_dictionary: self.dictionary.is_some(),
1448 bloom_filter_size: self
1449 .bloom_filter
1450 .as_ref()
1451 .map(|b| b.size_bytes())
1452 .unwrap_or(0),
1453 dictionary_size: self.dictionary.as_ref().map(|d| d.len()).unwrap_or(0),
1454 }
1455 }
1456
1457 pub fn cached_blocks(&self) -> usize {
1459 self.cache.read().blocks.len()
1460 }
1461
1462 pub fn cached_bytes(&self) -> usize {
1468 self.cache
1469 .read()
1470 .blocks
1471 .values()
1472 .map(|block| block.len())
1473 .sum()
1474 }
1475
1476 pub async fn get(&self, key: &[u8]) -> io::Result<Option<V>> {
1481 log::debug!(
1482 "SSTable::get called, key_len={}, total_blocks={}",
1483 key.len(),
1484 self.block_index.len()
1485 );
1486
1487 if let Some(ref bloom) = self.bloom_filter
1489 && !bloom.may_contain(key)
1490 {
1491 log::debug!("SSTable::get bloom filter negative");
1492 return Ok(None);
1493 }
1494
1495 let block_idx = match self.block_index.locate(key) {
1497 Some(idx) => idx,
1498 None => {
1499 log::debug!("SSTable::get key not found (before first block)");
1500 return Ok(None);
1501 }
1502 };
1503
1504 log::debug!("SSTable::get loading block_idx={}", block_idx);
1505
1506 let block_data = self.load_block(block_idx).await?;
1508 self.search_block(&block_data, key)
1509 }
1510
1511 pub async fn get_batch(&self, keys: &[&[u8]]) -> io::Result<Vec<Option<V>>> {
1517 if keys.is_empty() {
1518 return Ok(Vec::new());
1519 }
1520
1521 let mut key_to_block: Vec<(usize, usize)> = Vec::with_capacity(keys.len());
1523 for (key_idx, key) in keys.iter().enumerate() {
1524 if let Some(ref bloom) = self.bloom_filter
1526 && !bloom.may_contain(key)
1527 {
1528 key_to_block.push((key_idx, usize::MAX)); continue;
1530 }
1531
1532 match self.block_index.locate(key) {
1533 Some(block_idx) => key_to_block.push((key_idx, block_idx)),
1534 None => key_to_block.push((key_idx, usize::MAX)), }
1536 }
1537
1538 let mut blocks_to_load: Vec<usize> = key_to_block
1540 .iter()
1541 .filter(|(_, b)| *b != usize::MAX)
1542 .map(|(_, b)| *b)
1543 .collect();
1544 blocks_to_load.sort_unstable();
1545 blocks_to_load.dedup();
1546
1547 for &block_idx in &blocks_to_load {
1549 let _ = self.load_block(block_idx).await?;
1550 }
1551
1552 let mut results = vec![None; keys.len()];
1554 for (key_idx, block_idx) in key_to_block {
1555 if block_idx == usize::MAX {
1556 continue;
1557 }
1558 let block_data = self.load_block(block_idx).await?; results[key_idx] = self.search_block(&block_data, keys[key_idx])?;
1560 }
1561
1562 Ok(results)
1563 }
1564
1565 pub async fn preload_all_blocks(&self) -> io::Result<()> {
1570 for block_idx in 0..self.block_index.len() {
1571 self.load_block(block_idx).await?;
1572 }
1573 Ok(())
1574 }
1575
1576 pub async fn prefetch_all_data_bulk(&self) -> io::Result<()> {
1582 let num_blocks = self.block_index.len();
1583 if num_blocks == 0 {
1584 return Ok(());
1585 }
1586
1587 let mut max_end: u64 = 0;
1589 for i in 0..num_blocks {
1590 if let Some(addr) = self.block_index.get_addr(i) {
1591 let end = addr.offset.checked_add(addr.length as u64).ok_or_else(|| {
1592 io::Error::new(io::ErrorKind::InvalidData, "SSTable block range overflow")
1593 })?;
1594 max_end = max_end.max(end);
1595 }
1596 }
1597
1598 let all_data = self.data_slice.read_bytes_range(0..max_end).await?;
1600 let buf = all_data.as_slice();
1601
1602 let mut cache = self.cache.write();
1604 cache.max_blocks = cache.max_blocks.max(num_blocks);
1605 for i in 0..num_blocks {
1606 let addr = self.block_index.get_addr(i).unwrap();
1607 if cache.get(addr.offset).is_some() {
1608 continue;
1609 }
1610 let start = usize::try_from(addr.offset).map_err(|_| {
1611 io::Error::new(io::ErrorKind::InvalidData, "SSTable block offset too large")
1612 })?;
1613 let end = addr
1614 .offset
1615 .checked_add(addr.length as u64)
1616 .and_then(|end| usize::try_from(end).ok())
1617 .ok_or_else(|| {
1618 io::Error::new(io::ErrorKind::InvalidData, "SSTable block range overflow")
1619 })?;
1620 let compressed = buf.get(start..end).ok_or_else(|| {
1621 io::Error::new(io::ErrorKind::UnexpectedEof, "SSTable block is truncated")
1622 })?;
1623 let decompressed = if let Some(ref dict) = self.dictionary {
1624 crate::compression::decompress_with_dict_limited(
1625 compressed,
1626 dict,
1627 MAX_SSTABLE_BLOCK_BYTES,
1628 )?
1629 } else {
1630 crate::compression::decompress_limited(compressed, MAX_SSTABLE_BLOCK_BYTES)?
1631 };
1632 cache.insert(addr.offset, Arc::from(decompressed));
1633 }
1634
1635 Ok(())
1636 }
1637
1638 async fn load_block(&self, block_idx: usize) -> io::Result<Arc<[u8]>> {
1641 let addr = self.block_index.get_addr(block_idx).ok_or_else(|| {
1642 io::Error::new(io::ErrorKind::InvalidInput, "Block index out of range")
1643 })?;
1644
1645 {
1647 if let Some(block) = self.cache.read().peek(addr.offset) {
1648 return Ok(block);
1649 }
1650 }
1651
1652 log::debug!(
1653 "SSTable::load_block idx={} CACHE MISS, reading bytes [{}-{}]",
1654 block_idx,
1655 addr.offset,
1656 addr.offset + addr.length as u64
1657 );
1658
1659 let range = addr.byte_range();
1661 let compressed = self.data_slice.read_bytes_range(range).await?;
1662
1663 let decompressed = if let Some(ref dict) = self.dictionary {
1665 crate::compression::decompress_with_dict_limited(
1666 compressed.as_slice(),
1667 dict,
1668 MAX_SSTABLE_BLOCK_BYTES,
1669 )?
1670 } else {
1671 crate::compression::decompress_limited(compressed.as_slice(), MAX_SSTABLE_BLOCK_BYTES)?
1672 };
1673
1674 let block: Arc<[u8]> = Arc::from(decompressed);
1675
1676 {
1678 let mut cache = self.cache.write();
1679 cache.insert(addr.offset, Arc::clone(&block));
1680 }
1681
1682 Ok(block)
1683 }
1684
1685 #[cfg(feature = "sync")]
1687 fn load_block_sync(&self, block_idx: usize) -> io::Result<Arc<[u8]>> {
1688 let addr = self.block_index.get_addr(block_idx).ok_or_else(|| {
1689 io::Error::new(io::ErrorKind::InvalidInput, "Block index out of range")
1690 })?;
1691
1692 {
1694 if let Some(block) = self.cache.read().peek(addr.offset) {
1695 return Ok(block);
1696 }
1697 }
1698
1699 let range = addr.byte_range();
1701 let compressed = self.data_slice.read_bytes_range_sync(range)?;
1702
1703 let decompressed = if let Some(ref dict) = self.dictionary {
1705 crate::compression::decompress_with_dict_limited(
1706 compressed.as_slice(),
1707 dict,
1708 MAX_SSTABLE_BLOCK_BYTES,
1709 )?
1710 } else {
1711 crate::compression::decompress_limited(compressed.as_slice(), MAX_SSTABLE_BLOCK_BYTES)?
1712 };
1713
1714 let block: Arc<[u8]> = Arc::from(decompressed);
1715
1716 {
1718 let mut cache = self.cache.write();
1719 cache.insert(addr.offset, Arc::clone(&block));
1720 }
1721
1722 Ok(block)
1723 }
1724
1725 #[cfg(feature = "sync")]
1727 pub fn get_sync(&self, key: &[u8]) -> io::Result<Option<V>> {
1728 if let Some(ref bloom) = self.bloom_filter
1730 && !bloom.may_contain(key)
1731 {
1732 return Ok(None);
1733 }
1734
1735 let block_idx = match self.block_index.locate(key) {
1737 Some(idx) => idx,
1738 None => {
1739 return Ok(None);
1740 }
1741 };
1742
1743 let block_data = self.load_block_sync(block_idx)?;
1744 self.search_block(&block_data, key)
1745 }
1746
1747 fn search_block(&self, block_data: &[u8], target_key: &[u8]) -> io::Result<Option<V>> {
1748 let mut reader = block_data;
1749 let mut current_key = Vec::new();
1750
1751 while !reader.is_empty() {
1752 let common_prefix_len = read_vint(&mut reader)? as usize;
1753 let suffix_len = read_vint(&mut reader)? as usize;
1754
1755 if suffix_len > reader.len() {
1756 return Err(io::Error::new(
1757 io::ErrorKind::UnexpectedEof,
1758 "SSTable block suffix truncated",
1759 ));
1760 }
1761 current_key.truncate(common_prefix_len);
1762 current_key.extend_from_slice(&reader[..suffix_len]);
1763 reader = &reader[suffix_len..];
1764
1765 let value = V::deserialize(&mut reader)?;
1766
1767 match current_key.as_slice().cmp(target_key) {
1768 std::cmp::Ordering::Equal => return Ok(Some(value)),
1769 std::cmp::Ordering::Greater => return Ok(None),
1770 std::cmp::Ordering::Less => continue,
1771 }
1772 }
1773
1774 Ok(None)
1775 }
1776
1777 pub async fn prefetch_range(&self, start_key: &[u8], end_key: &[u8]) -> io::Result<()> {
1779 let start_block = self.block_index.locate(start_key).unwrap_or(0);
1780 let end_block = self
1781 .block_index
1782 .locate(end_key)
1783 .unwrap_or(self.block_index.len().saturating_sub(1));
1784
1785 for block_idx in start_block..=end_block.min(self.block_index.len().saturating_sub(1)) {
1786 let _ = self.load_block(block_idx).await?;
1787 }
1788
1789 Ok(())
1790 }
1791
1792 pub fn iter(&self) -> AsyncSSTableIterator<'_, V> {
1794 AsyncSSTableIterator::new(self)
1795 }
1796
1797 pub async fn all_entries(&self) -> io::Result<Vec<(Vec<u8>, V)>> {
1799 let mut results = Vec::new();
1800
1801 for block_idx in 0..self.block_index.len() {
1802 let block_data = self.load_block(block_idx).await?;
1803 let mut reader = &block_data[..];
1804 let mut current_key = Vec::new();
1805
1806 while !reader.is_empty() {
1807 let common_prefix_len = read_vint(&mut reader)? as usize;
1808 let suffix_len = read_vint(&mut reader)? as usize;
1809
1810 if suffix_len > reader.len() {
1811 return Err(io::Error::new(
1812 io::ErrorKind::UnexpectedEof,
1813 "SSTable block suffix truncated",
1814 ));
1815 }
1816 current_key.truncate(common_prefix_len);
1817 current_key.extend_from_slice(&reader[..suffix_len]);
1818 reader = &reader[suffix_len..];
1819
1820 let value = V::deserialize(&mut reader)?;
1821 results.push((current_key.clone(), value));
1822 }
1823 }
1824
1825 Ok(results)
1826 }
1827
1828 pub async fn prefix_scan(&self, prefix: &[u8]) -> io::Result<Vec<(Vec<u8>, V)>> {
1834 let (results, _) = self.prefix_scan_limited(prefix, usize::MAX).await?;
1835 Ok(results)
1836 }
1837
1838 pub async fn prefix_scan_limited(
1841 &self,
1842 prefix: &[u8],
1843 max_results: usize,
1844 ) -> io::Result<PrefixScanResult<V>> {
1845 if self.block_index.is_empty() || prefix.is_empty() {
1846 return Ok((Vec::new(), false));
1847 }
1848
1849 let start_block = match self.block_index.locate(prefix) {
1850 Some(idx) => idx,
1851 None => return Ok((Vec::new(), false)),
1852 };
1853
1854 let mut results = Vec::new();
1855
1856 for block_idx in start_block..self.block_index.len() {
1857 let block_data = self.load_block(block_idx).await?;
1858 let mut reader = &block_data[..];
1859 let mut current_key = Vec::new();
1860
1861 while !reader.is_empty() {
1862 let common_prefix_len = read_vint(&mut reader)? as usize;
1863 let suffix_len = read_vint(&mut reader)? as usize;
1864
1865 if suffix_len > reader.len() {
1866 return Err(io::Error::new(
1867 io::ErrorKind::UnexpectedEof,
1868 "SSTable block suffix truncated",
1869 ));
1870 }
1871 current_key.truncate(common_prefix_len);
1872 current_key.extend_from_slice(&reader[..suffix_len]);
1873 reader = &reader[suffix_len..];
1874
1875 let value = V::deserialize(&mut reader)?;
1876
1877 if current_key.starts_with(prefix) {
1878 if results.len() >= max_results {
1879 return Ok((results, true));
1880 }
1881 results.push((current_key.clone(), value));
1882 } else if current_key.as_slice() > prefix {
1883 return Ok((results, false));
1885 }
1886 }
1887 }
1888
1889 Ok((results, false))
1890 }
1891
1892 #[cfg(feature = "sync")]
1894 pub fn prefix_scan_sync(&self, prefix: &[u8]) -> io::Result<Vec<(Vec<u8>, V)>> {
1895 let (results, _) = self.prefix_scan_limited_sync(prefix, usize::MAX)?;
1896 Ok(results)
1897 }
1898
1899 #[cfg(feature = "sync")]
1901 pub fn prefix_scan_limited_sync(
1902 &self,
1903 prefix: &[u8],
1904 max_results: usize,
1905 ) -> io::Result<PrefixScanResult<V>> {
1906 if self.block_index.is_empty() || prefix.is_empty() {
1907 return Ok((Vec::new(), false));
1908 }
1909
1910 let start_block = match self.block_index.locate(prefix) {
1911 Some(idx) => idx,
1912 None => return Ok((Vec::new(), false)),
1913 };
1914
1915 let mut results = Vec::new();
1916
1917 for block_idx in start_block..self.block_index.len() {
1918 let block_data = self.load_block_sync(block_idx)?;
1919 let mut reader = &block_data[..];
1920 let mut current_key = Vec::new();
1921
1922 while !reader.is_empty() {
1923 let common_prefix_len = read_vint(&mut reader)? as usize;
1924 let suffix_len = read_vint(&mut reader)? as usize;
1925
1926 if suffix_len > reader.len() {
1927 return Err(io::Error::new(
1928 io::ErrorKind::UnexpectedEof,
1929 "SSTable block suffix truncated",
1930 ));
1931 }
1932 current_key.truncate(common_prefix_len);
1933 current_key.extend_from_slice(&reader[..suffix_len]);
1934 reader = &reader[suffix_len..];
1935
1936 let value = V::deserialize(&mut reader)?;
1937
1938 if current_key.starts_with(prefix) {
1939 if results.len() >= max_results {
1940 return Ok((results, true));
1941 }
1942 results.push((current_key.clone(), value));
1943 } else if current_key.as_slice() > prefix {
1944 return Ok((results, false));
1945 }
1946 }
1947 }
1948
1949 Ok((results, false))
1950 }
1951}
1952
1953pub struct AsyncSSTableIterator<'a, V: SSTableValue> {
1955 reader: &'a AsyncSSTableReader<V>,
1956 current_block: usize,
1957 block_data: Option<Arc<[u8]>>,
1958 block_offset: usize,
1959 current_key: Vec<u8>,
1960 finished: bool,
1961}
1962
1963impl<'a, V: SSTableValue> AsyncSSTableIterator<'a, V> {
1964 fn new(reader: &'a AsyncSSTableReader<V>) -> Self {
1965 Self {
1966 reader,
1967 current_block: 0,
1968 block_data: None,
1969 block_offset: 0,
1970 current_key: Vec::new(),
1971 finished: reader.block_index.is_empty(),
1972 }
1973 }
1974
1975 async fn load_next_block(&mut self) -> io::Result<bool> {
1976 if self.current_block >= self.reader.block_index.len() {
1977 self.finished = true;
1978 return Ok(false);
1979 }
1980
1981 self.block_data = Some(self.reader.load_block(self.current_block).await?);
1982 self.block_offset = 0;
1983 self.current_key.clear();
1984 self.current_block += 1;
1985 Ok(true)
1986 }
1987
1988 pub async fn next(&mut self) -> io::Result<Option<(Vec<u8>, V)>> {
1990 if self.finished {
1991 return Ok(None);
1992 }
1993
1994 if self.block_data.is_none() && !self.load_next_block().await? {
1995 return Ok(None);
1996 }
1997
1998 loop {
1999 let block = self.block_data.as_ref().unwrap();
2000 if self.block_offset >= block.len() {
2001 if !self.load_next_block().await? {
2002 return Ok(None);
2003 }
2004 continue;
2005 }
2006
2007 let mut reader = &block[self.block_offset..];
2008 let start_len = reader.len();
2009
2010 let common_prefix_len = read_vint(&mut reader)? as usize;
2011 let suffix_len = read_vint(&mut reader)? as usize;
2012
2013 if suffix_len > reader.len() {
2014 return Err(io::Error::new(
2015 io::ErrorKind::UnexpectedEof,
2016 "SSTable block suffix truncated",
2017 ));
2018 }
2019 self.current_key.truncate(common_prefix_len);
2020 self.current_key.extend_from_slice(&reader[..suffix_len]);
2021 reader = &reader[suffix_len..];
2022
2023 let value = V::deserialize(&mut reader)?;
2024
2025 self.block_offset += start_len - reader.len();
2026
2027 return Ok(Some((self.current_key.clone(), value)));
2028 }
2029 }
2030}
2031
2032#[cfg(test)]
2033mod tests {
2034 use super::*;
2035
2036 #[test]
2037 fn test_bloom_filter_basic() {
2038 let mut bloom = BloomFilter::new(100, 10);
2039
2040 bloom.insert(b"hello");
2041 bloom.insert(b"world");
2042 bloom.insert(b"test");
2043
2044 assert!(bloom.may_contain(b"hello"));
2045 assert!(bloom.may_contain(b"world"));
2046 assert!(bloom.may_contain(b"test"));
2047
2048 assert!(!bloom.may_contain(b"notfound"));
2050 assert!(!bloom.may_contain(b"missing"));
2051 }
2052
2053 #[test]
2054 fn test_bloom_filter_serialization() {
2055 let mut bloom = BloomFilter::new(100, 10);
2056 bloom.insert(b"key1");
2057 bloom.insert(b"key2");
2058
2059 let bytes = bloom.to_bytes();
2060 let restored = BloomFilter::from_owned_bytes(OwnedBytes::new(bytes)).unwrap();
2061
2062 assert!(restored.may_contain(b"key1"));
2063 assert!(restored.may_contain(b"key2"));
2064 assert!(!restored.may_contain(b"key3"));
2065 }
2066
2067 #[test]
2068 fn bloom_header_preserves_bit_counts_above_u32() {
2069 let num_bits = u32::MAX as usize + 1;
2070 let mut header = Vec::new();
2071 write_bloom_header(&mut header, num_bits, BLOOM_HASH_COUNT, 1).unwrap();
2072 assert_eq!(header.len(), BLOOM_FILTER_HEADER_SIZE);
2073 assert_eq!(
2074 u64::from_le_bytes(header[0..8].try_into().unwrap()),
2075 num_bits as u64
2076 );
2077 }
2078
2079 #[test]
2080 fn test_bloom_filter_false_positive_rate() {
2081 let num_keys = 10000;
2082 let mut bloom = BloomFilter::new(num_keys, BLOOM_BITS_PER_KEY);
2083
2084 for i in 0..num_keys {
2086 let key = format!("key_{}", i);
2087 bloom.insert(key.as_bytes());
2088 }
2089
2090 for i in 0..num_keys {
2092 let key = format!("key_{}", i);
2093 assert!(bloom.may_contain(key.as_bytes()));
2094 }
2095
2096 let mut false_positives = 0;
2098 let test_count = 10000;
2099 for i in 0..test_count {
2100 let key = format!("nonexistent_{}", i);
2101 if bloom.may_contain(key.as_bytes()) {
2102 false_positives += 1;
2103 }
2104 }
2105
2106 let fp_rate = false_positives as f64 / test_count as f64;
2109 assert!(
2110 fp_rate < 0.03,
2111 "False positive rate {} is too high",
2112 fp_rate
2113 );
2114 }
2115
2116 #[test]
2117 fn test_sstable_writer_config() {
2118 use crate::structures::IndexOptimization;
2119
2120 let config = SSTableWriterConfig::default();
2122 assert_eq!(config.compression_level.0, 9); assert!(config.use_bloom_filter); assert!(!config.use_dictionary);
2125
2126 let adaptive = SSTableWriterConfig::from_optimization(IndexOptimization::Adaptive);
2128 assert_eq!(adaptive.compression_level.0, 9);
2129 assert!(adaptive.use_bloom_filter);
2130 assert!(!adaptive.use_dictionary);
2131
2132 let size = SSTableWriterConfig::from_optimization(IndexOptimization::SizeOptimized);
2134 assert_eq!(size.compression_level.0, 22); assert!(size.use_bloom_filter);
2136 assert!(size.use_dictionary);
2137
2138 let perf = SSTableWriterConfig::from_optimization(IndexOptimization::PerformanceOptimized);
2140 assert_eq!(perf.compression_level.0, 1); assert!(perf.use_bloom_filter); assert!(!perf.use_dictionary);
2143
2144 let fast = SSTableWriterConfig::fast();
2146 assert_eq!(fast.compression_level.0, 1);
2147
2148 let max = SSTableWriterConfig::max_compression();
2149 assert_eq!(max.compression_level.0, 22);
2150 }
2151
2152 #[test]
2153 fn test_vint_roundtrip() {
2154 let test_values = [0u64, 1, 127, 128, 255, 256, 16383, 16384, u64::MAX];
2155
2156 for &val in &test_values {
2157 let mut buf = Vec::new();
2158 write_vint(&mut buf, val).unwrap();
2159 let mut reader = buf.as_slice();
2160 let decoded = read_vint(&mut reader).unwrap();
2161 assert_eq!(val, decoded, "Failed for value {}", val);
2162 }
2163 }
2164
2165 #[test]
2166 fn test_common_prefix_len() {
2167 assert_eq!(common_prefix_len(b"hello", b"hello"), 5);
2168 assert_eq!(common_prefix_len(b"hello", b"help"), 3);
2169 assert_eq!(common_prefix_len(b"hello", b"world"), 0);
2170 assert_eq!(common_prefix_len(b"", b"hello"), 0);
2171 assert_eq!(common_prefix_len(b"hello", b""), 0);
2172 }
2173}