1use std::{
2 borrow::Cow,
3 collections::VecDeque,
4 io::{BufWriter, Write},
5 path::{Path, PathBuf},
6};
7
8use anyhow::{Context, Result};
9use byteorder::{BE, ByteOrder, WriteBytesExt};
10use either::Either;
11use fs_err::File;
12
13use crate::{
14 Compression,
15 compression::{Compressor, checksum_block},
16 constants::{MAX_INLINE_VALUE_SIZE, MAX_SMALL_VALUE_SIZE, MIN_SMALL_VALUE_BLOCK_SIZE},
17 meta_file::MetaEntryFlags,
18 static_sorted_file::{
19 BLOB_VALUE_REF_SIZE, BLOCK_TYPE_INDEX, FIXED_KEY_BLOCK_MIXED_VALUE_TYPE, FixedRegions,
20 KEY_BLOCK_ENTRY_TYPE_BLOB, KEY_BLOCK_ENTRY_TYPE_INLINE_MIN,
21 KEY_BLOCK_ENTRY_TYPE_KEY_DELETED, KEY_BLOCK_ENTRY_TYPE_KEY_VALUE_DELETED_MIN,
22 KEY_BLOCK_ENTRY_TYPE_MEDIUM, KEY_BLOCK_ENTRY_TYPE_SMALL,
23 KEY_BLOCK_TABLE_ENTRY_SIZE_NO_HASH, KEY_DELETED_REF_SIZE, KeyBlockLayout,
24 MEDIUM_VALUE_REF_SIZE, SMALL_VALUE_REF_SIZE, key_block_table_stride,
25 },
26};
27
28pub const BLOCK_HEADER_SIZE: usize = 8;
30
31const MAX_KEY_BLOCK_ENTRIES: usize = MAX_KEY_BLOCK_SIZE / KEY_BLOCK_ENTRY_META_OVERHEAD;
33const MAX_KEY_BLOCK_SIZE: usize = 16 * 1024;
36const KEY_BLOCK_ENTRY_META_OVERHEAD: usize = 20;
45const AMQF_FALSE_POSITIVE_RATE: f64 = 0.01;
47const AVG_SMALL_VALUE_SIZE: usize = 64;
52
53const BLOCK_INDEX_CAPACITY_BUFFER: usize = 16;
57
58const MIN_KEY_SIZE_FOR_COMPRESSION: usize = 16;
63
64const MAX_FIXED_KEY_LEN: usize = u8::MAX as usize;
69
70const MAX_FIXED_VAL_SIZE: usize = u8::MAX as usize;
75
76#[derive(Clone, Copy, PartialEq, Eq, Debug)]
81struct EntryType(u8);
82
83#[derive(Clone, Copy)]
97enum KeyBlockFormat {
98 Unknown,
100 Fixed {
102 key_len: u8,
103 val_size: u8,
104 value_type: Option<EntryType>,
106 },
107 Variable,
109}
110
111impl KeyBlockFormat {
112 fn update(&mut self, key_len: usize, value_type: EntryType) {
117 let val_size = value_type_val_size(value_type);
118 *self = match *self {
119 KeyBlockFormat::Unknown => {
120 if key_len <= MAX_FIXED_KEY_LEN && val_size <= MAX_FIXED_VAL_SIZE {
121 KeyBlockFormat::Fixed {
122 key_len: key_len as u8,
123 val_size: val_size as u8,
124 value_type: Some(value_type),
125 }
126 } else {
127 KeyBlockFormat::Variable
128 }
129 }
130 KeyBlockFormat::Fixed {
131 key_len: k,
132 val_size: s,
133 value_type: v,
134 } if k as usize == key_len && s as usize == val_size => KeyBlockFormat::Fixed {
135 key_len: k,
136 val_size: s,
137 value_type: v.filter(|v| *v == value_type),
139 },
140 KeyBlockFormat::Fixed { .. } | KeyBlockFormat::Variable => KeyBlockFormat::Variable,
141 };
142 }
143}
144
145#[derive(Clone, Copy)]
147struct KeyBlockFlushInfo {
148 max_key_len: usize,
149 min_key_len: usize,
150 format: KeyBlockFormat,
151}
152
153impl KeyBlockFlushInfo {
154 fn uniform_key_len(&self) -> Option<usize> {
156 (self.min_key_len == self.max_key_len).then_some(self.max_key_len)
157 }
158}
159
160struct KeyBlockAccumulator {
166 size: usize,
168 entry_count: usize,
170 max_key_len: usize,
172 min_key_len: usize,
173 last_hash: u64,
176 format: KeyBlockFormat,
178}
179
180impl KeyBlockAccumulator {
181 fn new() -> Self {
182 Self {
183 size: 0,
184 entry_count: 0,
185 max_key_len: 0,
186 min_key_len: usize::MAX,
187 last_hash: 0,
188 format: KeyBlockFormat::Unknown,
189 }
190 }
191
192 fn add(&mut self, key_len: usize, key_hash: u64, value_type: EntryType) {
194 self.size += key_len + KEY_BLOCK_ENTRY_META_OVERHEAD;
195 self.max_key_len = self.max_key_len.max(key_len);
196 self.min_key_len = self.min_key_len.min(key_len);
197 self.entry_count += 1;
198 self.last_hash = key_hash;
199 self.format.update(key_len, value_type);
200 }
201
202 fn flush_info(&self) -> KeyBlockFlushInfo {
204 KeyBlockFlushInfo {
205 max_key_len: self.max_key_len,
206 min_key_len: self.min_key_len,
207 format: self.format,
208 }
209 }
210
211 fn should_flush(&self, next_key_len: usize, next_key_hash: u64) -> bool {
215 if self.entry_count == 0 {
216 return false;
217 }
218 let would_exceed_size =
219 self.size + next_key_len + KEY_BLOCK_ENTRY_META_OVERHEAD > MAX_KEY_BLOCK_SIZE;
220 let would_exceed_entries = self.entry_count >= MAX_KEY_BLOCK_ENTRIES;
221 (would_exceed_size || would_exceed_entries) && self.last_hash != next_key_hash
223 }
224
225 fn reset(&mut self) {
227 self.size = 0;
228 self.entry_count = 0;
229 self.max_key_len = 0;
230 self.min_key_len = usize::MAX;
231 self.format = KeyBlockFormat::Unknown;
232 }
234}
235
236fn choose_layout(max_key_len: usize) -> KeyBlockLayout {
238 if max_key_len > 32 {
241 KeyBlockLayout::HashThenKey
242 } else {
243 KeyBlockLayout::KeyOnly
244 }
245}
246
247#[inline]
248fn be_key_u32(key: &[u8]) -> u32 {
249 u32::from_be_bytes(key.try_into().expect("4-byte key"))
250}
251
252#[inline]
253fn be_key_u64(key: &[u8]) -> u64 {
254 u64::from_be_bytes(key.try_into().expect("8-byte key"))
255}
256
257pub trait Entry {
259 fn key_hash(&self) -> u64;
261 fn key_len(&self) -> usize;
263 fn key_bytes(&self) -> &[u8];
265 fn write_key_to(&self, buf: &mut Vec<u8>) {
267 buf.extend_from_slice(self.key_bytes());
268 }
269
270 fn value(&self) -> EntryValue<'_>;
272}
273
274impl<E: Entry> Entry for &E {
275 fn key_hash(&self) -> u64 {
276 (*self).key_hash()
277 }
278 fn key_len(&self) -> usize {
279 (*self).key_len()
280 }
281 fn key_bytes(&self) -> &[u8] {
282 (*self).key_bytes()
283 }
284 fn value(&self) -> EntryValue<'_> {
285 (*self).value()
286 }
287}
288
289#[derive(Copy, Clone)]
291pub enum EntryValue<'l> {
292 Inline { value: &'l [u8] },
294 Small { value: &'l [u8] },
296 Medium { value: &'l [u8] },
298 MediumRaw {
300 uncompressed_size: u32,
303 checksum: u32,
305 block: &'l [u8],
306 },
307 Large { blob: u32 },
309 KeyDeleted,
311 KeyValueDeleted { value: &'l [u8] },
315}
316
317#[derive(Debug, Clone)]
318pub struct StaticSortedFileBuilderMeta<'a> {
319 pub min_hash: u64,
321 pub max_hash: u64,
323 pub amqf: Cow<'a, [u8]>,
325 pub block_count: u16,
327 pub size: u64,
329 pub flags: MetaEntryFlags,
331 pub entries: u64,
333}
334
335pub fn write_static_stored_file<E: Entry>(
345 entries: &[E],
346 file: &Path,
347 flags: MetaEntryFlags,
348 compression: Compression,
349) -> Result<(StaticSortedFileBuilderMeta<'static>, File)> {
350 debug_assert!(
351 entries
352 .iter()
353 .map(|e| (e.key_hash(), e.key_bytes()))
354 .is_sorted()
355 );
356 let mut writer = StreamingSstWriter::new(file, flags, entries.len() as u64, compression)?;
357 for entry in entries {
358 if let Err(err) = writer.add(entry) {
359 writer.cancel();
360 return Err(err);
361 }
362 }
363 writer.close()
364}
365
366fn write_raw_block_to_file(
375 file: &mut BufWriter<File>,
376 block_offsets: &mut Vec<u32>,
377 uncompressed_size: u32,
378 checksum: u32,
379 block: &[u8],
380) -> Result<u16> {
381 let block_index: u16 = block_offsets
382 .len()
383 .try_into()
384 .expect("Block index overflow");
385
386 let len: u32 = (block.len() + BLOCK_HEADER_SIZE).try_into().unwrap();
387 let offset = block_offsets
388 .last()
389 .copied()
390 .unwrap_or_default()
391 .checked_add(len)
392 .expect("Block offset overflow");
393 block_offsets.push(offset);
394
395 file.write_u32::<BE>(uncompressed_size)
396 .context("Failed to write uncompressed size")?;
397 file.write_u32::<BE>(checksum)
398 .context("Failed to write checksum")?;
399 file.write_all(block)
400 .context("Failed to write block data")?;
401 Ok(block_index)
402}
403
404fn write_block_to_file(
406 file: &mut BufWriter<File>,
407 compress_buffer: &mut Vec<u8>,
408 block_offsets: &mut Vec<u32>,
409 block: &[u8],
410 try_compress: bool,
411 compressor: &mut Compressor,
412) -> Result<u16> {
413 let (uncompressed_size, data_to_write): (u32, &[u8]) = if try_compress {
414 compressor.compress_into_buffer(block, compress_buffer)?;
415 if compress_buffer.len() < block.len() - (block.len() / 8) {
417 (block.len().try_into().unwrap(), compress_buffer)
418 } else {
419 (0, block)
420 }
421 } else {
422 (0, block)
423 };
424
425 let checksum = checksum_block(data_to_write);
427
428 write_raw_block_to_file(
429 file,
430 block_offsets,
431 uncompressed_size,
432 checksum,
433 data_to_write,
434 )
435}
436
437enum ValueRef {
443 Small {
445 block_index: u16,
446 offset: u32,
447 size: u16,
448 },
449 PendingSmall {
452 #[cfg(debug_assertions)]
453 small_block_id: u16,
454 offset: u32,
455 size: u16,
456 },
457 Medium { block_index: u16 },
459 Inline {
461 data: [u8; MAX_INLINE_VALUE_SIZE],
462 len: u8,
463 },
464 Blob { blob_id: u32 },
466 KeyDeleted,
468 KeyValueDeleted {
471 data: [u8; MAX_INLINE_VALUE_SIZE],
472 len: u8,
473 },
474}
475
476impl ValueRef {
477 fn entry_type(&self) -> EntryType {
479 EntryType(match self {
480 ValueRef::Small { .. } | ValueRef::PendingSmall { .. } => KEY_BLOCK_ENTRY_TYPE_SMALL,
481 ValueRef::Medium { .. } => KEY_BLOCK_ENTRY_TYPE_MEDIUM,
482 ValueRef::Inline { len, .. } => KEY_BLOCK_ENTRY_TYPE_INLINE_MIN + *len,
483 ValueRef::Blob { .. } => KEY_BLOCK_ENTRY_TYPE_BLOB,
484 ValueRef::KeyDeleted => KEY_BLOCK_ENTRY_TYPE_KEY_DELETED,
485 ValueRef::KeyValueDeleted { len, .. } => {
486 KEY_BLOCK_ENTRY_TYPE_KEY_VALUE_DELETED_MIN + *len
487 }
488 })
489 }
490
491 fn write_value_to(&self, buffer: &mut Vec<u8>) {
496 match self {
497 ValueRef::Small {
498 block_index,
499 offset,
500 size,
501 } => {
502 let mut scratch = [0u8; 8];
503 BE::write_u16(&mut scratch, *block_index);
504 BE::write_u16(&mut scratch[2..], *size);
505 BE::write_u32(&mut scratch[4..], *offset);
506 buffer.extend(&scratch);
507 }
508 ValueRef::Medium { block_index } => {
509 let mut scratch = [0u8; 2];
510 BE::write_u16(&mut scratch, *block_index);
511 buffer.extend(scratch);
512 }
513 ValueRef::Inline { data, len } => {
514 buffer.extend(&data[..*len as usize]);
515 }
516 ValueRef::Blob { blob_id } => {
517 let mut scratch = [0u8; 4];
518 BE::write_u32(&mut scratch, *blob_id);
519 buffer.extend(scratch);
520 }
521 ValueRef::KeyDeleted => { }
522 ValueRef::KeyValueDeleted { data, len } => {
523 buffer.extend(&data[..*len as usize]);
524 }
525 ValueRef::PendingSmall { .. } => {
526 unreachable!("PendingSmall should have been resolved");
527 }
528 }
529 }
530}
531
532struct PendingEntry<E> {
533 entry: E,
534 value_ref: ValueRef,
535}
536
537pub struct StreamingSstWriter<E: Entry> {
547 file: Option<BufWriter<File>>,
550 file_path: PathBuf,
551 preserve_file: bool,
553 compress_buffer: Vec<u8>,
554 block_offsets: Vec<u32>,
555 compressor: Compressor,
556
557 pending_keys: VecDeque<PendingEntry<E>>,
580
581 first_pending_small_index: usize,
586
587 #[cfg(debug_assertions)]
589 current_small_block_id: u16,
590
591 pending_small_value_block: Vec<u8>,
593
594 key_buffer: Vec<u8>,
596
597 key_value_buffer: Vec<u8>,
600
601 collected_fingerprints: Vec<u32>,
604
605 key_block_boundaries: Vec<(u64, u16)>,
607
608 min_hash: u64,
610 max_hash: u64,
611 entry_count: u64,
612 flags: MetaEntryFlags,
613
614 total_key_size: usize,
616 total_value_size: usize,
617
618 pending_key_total_size: usize,
620
621 current_key_block: KeyBlockAccumulator,
623
624 #[cfg(debug_assertions)]
627 finished: bool,
628}
629
630impl<E: Entry> StreamingSstWriter<E> {
631 pub fn new(
635 file: &Path,
636 flags: MetaEntryFlags,
637 max_entry_count: u64,
638 compression: Compression,
639 ) -> Result<Self> {
640 let file_path = file.to_owned();
641 let file = BufWriter::new(File::create(file)?);
642 let compressor = Compressor::new(compression)?;
643
644 let estimated_key_blocks = (max_entry_count as usize)
647 .div_ceil(MAX_KEY_BLOCK_ENTRIES)
648 .max(1);
649 let entries_per_value_block = MIN_SMALL_VALUE_BLOCK_SIZE / AVG_SMALL_VALUE_SIZE;
652 let estimated_value_blocks = (max_entry_count as usize)
653 .div_ceil(entries_per_value_block)
654 .max(1);
655 let estimated_total_blocks = estimated_key_blocks + estimated_value_blocks + 1;
656
657 Ok(Self {
658 file: Some(file),
659 file_path,
660 preserve_file: false,
661 compress_buffer: Vec::with_capacity(MIN_SMALL_VALUE_BLOCK_SIZE + MAX_SMALL_VALUE_SIZE),
662 block_offsets: Vec::with_capacity(estimated_total_blocks),
663 compressor,
664 pending_keys: VecDeque::with_capacity(entries_per_value_block),
665 first_pending_small_index: 0,
666 #[cfg(debug_assertions)]
667 current_small_block_id: 0,
668 pending_small_value_block: Vec::with_capacity(
669 MIN_SMALL_VALUE_BLOCK_SIZE + MAX_SMALL_VALUE_SIZE,
670 ),
671 key_buffer: Vec::with_capacity(MAX_KEY_BLOCK_SIZE),
675 key_value_buffer: Vec::new(),
676 collected_fingerprints: Vec::with_capacity(max_entry_count as usize),
677 key_block_boundaries: Vec::with_capacity(estimated_key_blocks),
678 min_hash: u64::MAX,
679 max_hash: 0,
680 entry_count: 0,
681 flags,
682 total_key_size: 0,
683 total_value_size: 0,
684 pending_key_total_size: 0,
685 current_key_block: KeyBlockAccumulator::new(),
686 #[cfg(debug_assertions)]
687 finished: false,
688 })
689 }
690
691 pub fn is_full(&self, max_entries: usize, max_data_size: usize) -> bool {
695 self.entry_count as usize >= max_entries
696 || self.total_key_size + self.total_value_size >= max_data_size
697 || !self.has_block_index_capacity()
698 }
699
700 fn has_block_index_capacity(&self) -> bool {
704 let blocks_written = self.block_offsets.len();
705 let pending_small_block = usize::from(!self.pending_small_value_block.is_empty());
710 let pending_key_blocks = self
711 .pending_keys
712 .len()
713 .div_ceil(MAX_KEY_BLOCK_ENTRIES)
714 .max(self.pending_key_total_size.div_ceil(MAX_KEY_BLOCK_SIZE))
715 .max(1);
716 let index_block = 1;
717 let buffer = BLOCK_INDEX_CAPACITY_BUFFER;
718 blocks_written + pending_small_block + pending_key_blocks + index_block + buffer
719 < u16::MAX as usize
720 }
721
722 pub fn add(&mut self, entry: E) -> Result<()> {
724 let key_hash = entry.key_hash();
725 let key_len = entry.key_len();
726
727 if self.entry_count == 0 {
729 self.min_hash = key_hash;
730 }
731 self.max_hash = key_hash;
732 self.entry_count += 1;
733
734 self.collected_fingerprints.push(key_hash as u32);
736
737 self.total_key_size += key_len;
739 self.pending_key_total_size += key_len;
740
741 let value_ref = match entry.value() {
743 EntryValue::Medium { value } => {
744 self.total_value_size += value.len();
745 let block_index = write_block_to_file(
746 self.file.as_mut().unwrap(),
747 &mut self.compress_buffer,
748 &mut self.block_offsets,
749 value,
750 true,
751 &mut self.compressor,
752 )
753 .context("Failed to write value block")?;
754 ValueRef::Medium { block_index }
755 }
756 EntryValue::MediumRaw {
757 uncompressed_size,
758 checksum,
759 block,
760 } => {
761 self.total_value_size += block.len();
764 let block_index = write_raw_block_to_file(
765 self.file.as_mut().unwrap(),
766 &mut self.block_offsets,
767 uncompressed_size,
768 checksum,
769 block,
770 )
771 .context("Failed to write compressed value block")?;
772 ValueRef::Medium { block_index }
773 }
774 EntryValue::Small { value } => {
775 self.total_value_size += value.len();
776
777 let offset = self.pending_small_value_block.len() as u32;
778 let size: u16 = value.len().try_into().unwrap();
779 self.pending_small_value_block.extend_from_slice(value);
780
781 if self.first_pending_small_index >= self.pending_keys.len() {
783 self.first_pending_small_index = self.pending_keys.len();
784 }
785
786 let value_ref = ValueRef::PendingSmall {
787 #[cfg(debug_assertions)]
788 small_block_id: self.current_small_block_id,
789 offset,
790 size,
791 };
792
793 self.push_pending_key_entry(entry, value_ref);
794
795 if self.pending_small_value_block.len() >= MIN_SMALL_VALUE_BLOCK_SIZE {
799 self.flush_small_value_block()?;
800 }
801
802 return Ok(());
803 }
804 EntryValue::Inline { value } => {
805 debug_assert!(value.len() <= MAX_INLINE_VALUE_SIZE);
806 let mut data = [0u8; MAX_INLINE_VALUE_SIZE];
807 data[..value.len()].copy_from_slice(value);
808 ValueRef::Inline {
809 data,
810 len: value.len() as u8,
811 }
812 }
813 EntryValue::Large { blob } => ValueRef::Blob { blob_id: blob },
814 EntryValue::KeyDeleted => ValueRef::KeyDeleted,
815 EntryValue::KeyValueDeleted { value } => {
816 debug_assert!(value.len() <= MAX_INLINE_VALUE_SIZE);
818 let mut data = [0u8; MAX_INLINE_VALUE_SIZE];
819 data[..value.len()].copy_from_slice(value);
820 ValueRef::KeyValueDeleted {
821 data,
822 len: value.len() as u8,
823 }
824 }
825 };
826
827 self.push_pending_key_entry(entry, value_ref);
828 self.try_flush_key_blocks()
829 }
830
831 pub fn cancel(mut self) {
833 self.discard_partial_file();
834 #[cfg(debug_assertions)]
835 {
836 self.finished = true;
837 }
838 }
839
840 fn discard_partial_file(&mut self) {
842 if let Some(file) = self.file.take() {
843 let (file, _) = file.into_parts();
845 drop(file);
846 }
847 let _ = fs_err::remove_file(&self.file_path);
850 }
851
852 fn push_pending_key_entry(&mut self, entry: E, value_ref: ValueRef) {
854 self.pending_keys
855 .push_back(PendingEntry { entry, value_ref });
856 }
857
858 fn try_flush_key_blocks(&mut self) -> Result<()> {
864 debug_assert!(!matches!(
865 self.pending_keys.back().unwrap().value_ref,
866 ValueRef::PendingSmall { .. }
867 ));
868 if self.first_pending_small_index != self.pending_keys.len() - 1 {
869 return Ok(());
871 }
872 self.advance_boundary_to(self.pending_keys.len())
873 }
874
875 fn advance_boundary_to(&mut self, new_boundary: usize) -> Result<()> {
881 let mut last_flushed_end = 0usize;
882 let mut cumulative_key_size = 0usize;
885 let mut flushed_key_size = 0usize;
886
887 for i in self.first_pending_small_index..new_boundary {
888 let entry = &self.pending_keys[i];
889 let key_len = entry.entry.key_len();
890 let key_hash = entry.entry.key_hash();
891 let value_type = entry.value_ref.entry_type();
892
893 if self.current_key_block.should_flush(key_len, key_hash) {
894 let block_end = last_flushed_end + self.current_key_block.entry_count;
895 let info = self.current_key_block.flush_info();
896 self.flush_key_block(last_flushed_end, block_end, info)?;
897 flushed_key_size = cumulative_key_size;
898 last_flushed_end = block_end;
899 self.current_key_block.reset();
900 }
901
902 cumulative_key_size += key_len;
903 self.current_key_block.add(key_len, key_hash, value_type);
904 }
905
906 if last_flushed_end > 0 {
907 self.pending_key_total_size -= flushed_key_size;
908 self.pending_keys.drain(..last_flushed_end);
909 }
910
911 self.first_pending_small_index = new_boundary - last_flushed_end;
912 Ok(())
913 }
914
915 fn flush_small_value_block(&mut self) -> Result<()> {
918 if self.pending_small_value_block.is_empty() {
921 return Ok(());
922 }
923
924 let block_index = write_block_to_file(
925 self.file.as_mut().unwrap(),
926 &mut self.compress_buffer,
927 &mut self.block_offsets,
928 &self.pending_small_value_block,
929 true,
930 &mut self.compressor,
931 )
932 .context("Failed to write small value block")?;
933
934 #[cfg(debug_assertions)]
938 let flushed_id = self.current_small_block_id;
939 for i in self.first_pending_small_index..self.pending_keys.len() {
940 let entry = &mut self.pending_keys[i];
941 if let ValueRef::PendingSmall {
942 #[cfg(debug_assertions)]
943 small_block_id,
944 offset,
945 size,
946 } = entry.value_ref
947 {
948 #[cfg(debug_assertions)]
949 debug_assert_eq!(
950 small_block_id, flushed_id,
951 "all pending small entries must reference the small value block that was just \
952 written"
953 );
954 entry.value_ref = ValueRef::Small {
955 block_index,
956 offset,
957 size,
958 };
959 }
960 }
961
962 self.advance_boundary_to(self.pending_keys.len())?;
965
966 #[cfg(debug_assertions)]
968 {
969 self.current_small_block_id += 1;
970 }
971 self.pending_small_value_block.clear();
972
973 Ok(())
974 }
975
976 fn flush_key_block(&mut self, start: usize, end: usize, info: KeyBlockFlushInfo) -> Result<()> {
980 let entry_count = end - start;
981 let layout = choose_layout(info.max_key_len);
982 let try_compress = info.max_key_len >= MIN_KEY_SIZE_FOR_COMPRESSION;
983
984 let first_hash = self.pending_keys[start].entry.key_hash();
987 let Self {
990 key_buffer,
991 key_value_buffer,
992 pending_keys,
993 ..
994 } = self;
995 key_buffer.clear();
996 let block_entries = |start: usize, end: usize| {
1002 if layout == KeyBlockLayout::HashThenKey {
1003 return Either::Left(pending_keys.range(start..end));
1004 }
1005 let mut entries: Vec<&PendingEntry<E>> = pending_keys.range(start..end).collect();
1006 match info.uniform_key_len() {
1008 Some(4) => entries.sort_by_key(|&e| be_key_u32(e.entry.key_bytes())),
1009 Some(8) => entries.sort_by_key(|&e| be_key_u64(e.entry.key_bytes())),
1010 _ => entries.sort_by_key(|&e| e.entry.key_bytes()),
1011 }
1012 Either::Right(entries.into_iter())
1013 };
1014
1015 if let KeyBlockFormat::Fixed {
1016 key_len: key_size,
1017 val_size,
1018 value_type,
1019 } = info.format
1020 {
1021 let mut builder = FixedKeyBlockBuilder::new(
1022 key_buffer,
1023 key_value_buffer,
1024 entry_count as u32,
1025 layout,
1026 key_size,
1027 val_size,
1028 value_type,
1029 );
1030 for pending in block_entries(start, end) {
1031 builder.put(&pending.entry, &pending.value_ref);
1032 }
1033 builder.finish();
1034 } else {
1035 let mut builder = KeyBlockBuilder::new(key_buffer, entry_count as u32, layout);
1036 for pending in block_entries(start, end) {
1037 builder.put(&pending.entry, &pending.value_ref);
1038 }
1039 builder.finish();
1040 }
1041
1042 let block_index = write_block_to_file(
1043 self.file.as_mut().unwrap(),
1044 &mut self.compress_buffer,
1045 &mut self.block_offsets,
1046 &self.key_buffer,
1047 try_compress,
1048 &mut self.compressor,
1049 )
1050 .context("Failed to write key block")?;
1051 self.key_block_boundaries.push((first_hash, block_index));
1052
1053 Ok(())
1054 }
1055
1056 pub fn close(mut self) -> Result<(StaticSortedFileBuilderMeta<'static>, File)> {
1059 #[cfg(debug_assertions)]
1060 {
1061 self.finished = true;
1062 }
1063
1064 self.flush_small_value_block()?;
1066
1067 self.flush_remaining_key_blocks()?;
1069
1070 assert!(
1071 !self.key_block_boundaries.is_empty(),
1072 "StreamingSstWriter::close() called with no entries"
1073 );
1074
1075 let mut file = self.file.take().unwrap();
1076
1077 let index_entry_count: u16 = (self.key_block_boundaries.len() - 1)
1080 .try_into()
1081 .expect("Index entries count overflow");
1082 let index_block_size: usize =
1083 INDEX_BLOCK_HEADER_SIZE + index_entry_count as usize * INDEX_BLOCK_ENTRY_SIZE;
1084 let mut index_buf = Vec::with_capacity(index_block_size);
1085 {
1086 let first_block = self.key_block_boundaries[0].1;
1087 let mut index_block = IndexBlockBuilder::new(&mut index_buf, first_block);
1088 for &(hash, block) in &self.key_block_boundaries[1..] {
1089 index_block.put(hash, block);
1090 }
1091 }
1092 let index_checksum = checksum_block(&index_buf);
1093 write_raw_block_to_file(
1094 &mut file,
1095 &mut self.block_offsets,
1096 0,
1097 index_checksum,
1098 &index_buf,
1099 )
1100 .context("Failed to write index block")?;
1101
1102 for offset in &self.block_offsets {
1104 file.write_u32::<BE>(*offset)
1105 .context("Failed to write block offset")?;
1106 }
1107
1108 let block_count: u16 = self
1109 .block_offsets
1110 .len()
1111 .try_into()
1112 .expect("Block count overflow");
1113
1114 let actual_count = self.collected_fingerprints.len() as u64;
1118 let mut builder = qfilter::Builder::new(
1119 qfilter::Filter::new(actual_count.max(1), AMQF_FALSE_POSITIVE_RATE)
1120 .expect("Filter can't be constructed"),
1121 );
1122
1123 let fp_size = builder.fingerprint_size();
1124 assert!(fp_size < 32, "fp_size {fp_size} exceeds u32");
1125 let fp_mask = (1u32 << fp_size) - 1;
1126 self.collected_fingerprints
1128 .sort_unstable_by_key(|&h| h & fp_mask);
1129 for &h in &self.collected_fingerprints {
1130 builder
1131 .insert_fingerprint(false, h as u64)
1132 .expect("AMQF insert failed");
1133 }
1134 let filter = builder.into_filter();
1135
1136 let amqf = postcard::to_allocvec(&filter).expect("AMQF serialization failed");
1138
1139 let last_block_end = self.block_offsets.last().copied().unwrap_or_default() as u64;
1142 let offset_table_size = block_count as u64 * size_of::<u32>() as u64;
1143 let file_size = last_block_end + offset_table_size;
1144
1145 let meta = StaticSortedFileBuilderMeta {
1146 min_hash: self.min_hash,
1147 max_hash: self.max_hash,
1148 amqf: Cow::Owned(amqf),
1149 block_count,
1150 size: file_size,
1151 flags: self.flags,
1152 entries: self.entry_count,
1153 };
1154
1155 let file = file.into_inner()?;
1156 self.preserve_file = true;
1157 Ok((meta, file))
1158 }
1159
1160 fn flush_remaining_key_blocks(&mut self) -> Result<()> {
1167 if self.pending_keys.is_empty() {
1168 return Ok(());
1169 }
1170
1171 debug_assert_eq!(
1176 self.first_pending_small_index,
1177 self.pending_keys.len(),
1178 "expected no unresolved PendingSmall entries after flush_small_value_block"
1179 );
1180
1181 let total = self.pending_keys.len();
1182 let mut block_start = 0;
1183 let mut acc = KeyBlockAccumulator::new();
1184
1185 for i in 0..total {
1186 let entry = &self.pending_keys[i];
1187 let key_len = entry.entry.key_len();
1188 let key_hash = entry.entry.key_hash();
1189 let value_type = entry.value_ref.entry_type();
1190
1191 if acc.should_flush(key_len, key_hash) {
1192 self.flush_key_block(block_start, i, acc.flush_info())?;
1193 block_start = i;
1194 acc.reset();
1195 }
1196
1197 acc.add(key_len, key_hash, value_type);
1198 }
1199
1200 if block_start < total {
1202 self.flush_key_block(block_start, total, acc.flush_info())?;
1203 }
1204
1205 self.pending_keys.clear();
1207 Ok(())
1208 }
1209}
1210
1211impl<E: Entry> Drop for StreamingSstWriter<E> {
1212 fn drop(&mut self) {
1213 if !self.preserve_file {
1214 self.discard_partial_file();
1215 }
1216
1217 #[cfg(debug_assertions)]
1219 if !std::thread::panicking() {
1220 assert!(
1221 self.finished || self.entry_count == 0,
1222 "StreamingSstWriter dropped without calling close() or cancel()"
1223 );
1224 }
1225 }
1226}
1227
1228struct KeyBlockBuilder<'l> {
1237 current_entry: usize,
1238 header_size: usize,
1239 layout: KeyBlockLayout,
1242 table_stride: usize,
1244 buffer: &'l mut Vec<u8>,
1245}
1246
1247const KEY_BLOCK_HEADER_SIZE: usize = 4;
1249
1250impl<'l> KeyBlockBuilder<'l> {
1251 fn new(buffer: &'l mut Vec<u8>, entry_count: u32, layout: KeyBlockLayout) -> Self {
1253 debug_assert!(entry_count < (1 << 24));
1254
1255 const ESTIMATED_KEY_SIZE: usize = 16;
1256 let table_stride = key_block_table_stride(layout.hash_len());
1257 buffer.reserve(entry_count as usize * (ESTIMATED_KEY_SIZE + table_stride));
1258 let block_type = layout.block_type(false);
1259 buffer.write_u8(block_type).unwrap();
1260 buffer.write_u24::<BE>(entry_count).unwrap();
1261 buffer.resize(buffer.len() + entry_count as usize * table_stride, 0);
1263 Self {
1264 current_entry: 0,
1265 header_size: buffer.len(),
1266 layout,
1267 table_stride,
1268 buffer,
1269 }
1270 }
1271
1272 fn write_entry_header(&mut self, entry_type: EntryType) {
1276 let pos = self.buffer.len() - self.header_size;
1277 let slot = KEY_BLOCK_HEADER_SIZE + self.current_entry * self.table_stride;
1278 let word_offset = slot + self.table_stride - KEY_BLOCK_TABLE_ENTRY_SIZE_NO_HASH;
1279 let header = (pos as u32) | ((entry_type.0 as u32) << 24);
1280 BE::write_u32(&mut self.buffer[word_offset..word_offset + 4], header);
1281 }
1282
1283 fn put<E: Entry>(&mut self, entry: &E, value_ref: &ValueRef) {
1285 self.write_entry_header(value_ref.entry_type());
1286 if self.layout == KeyBlockLayout::HashThenKey {
1287 let slot = KEY_BLOCK_HEADER_SIZE + self.current_entry * self.table_stride;
1288 self.buffer[slot..slot + size_of::<u64>()]
1289 .copy_from_slice(&entry.key_hash().to_be_bytes());
1290 }
1291 entry.write_key_to(self.buffer);
1292 value_ref.write_value_to(self.buffer);
1293 self.current_entry += 1;
1294 }
1295 fn finish(self) -> &'l mut Vec<u8> {
1297 self.buffer
1298 }
1299}
1300
1301const FIXED_KEY_BLOCK_HEADER_SIZE: usize = 6;
1308
1309struct FixedKeyBlockBuilder<'l> {
1322 buffer: &'l mut Vec<u8>,
1324 tail: &'l mut Vec<u8>,
1326 per_entry_type: bool,
1328 layout: KeyBlockLayout,
1331}
1332
1333impl<'l> FixedKeyBlockBuilder<'l> {
1334 fn new(
1335 buffer: &'l mut Vec<u8>,
1336 tail: &'l mut Vec<u8>,
1337 entry_count: u32,
1338 layout: KeyBlockLayout,
1339 key_size: u8,
1340 val_size: u8,
1341 value_type: Option<EntryType>,
1342 ) -> Self {
1343 let per_entry_type = value_type.is_none();
1344 let FixedRegions {
1349 search_stride,
1350 tail_stride,
1351 ..
1352 } = FixedRegions::new(
1353 entry_count as usize,
1354 layout,
1355 key_size as usize,
1356 val_size as usize + usize::from(per_entry_type),
1357 );
1358 buffer.reserve(
1361 FIXED_KEY_BLOCK_HEADER_SIZE + entry_count as usize * (search_stride + tail_stride),
1362 );
1363 tail.clear();
1364 tail.reserve(entry_count as usize * tail_stride);
1365
1366 let block_type = layout.block_type(true);
1367 buffer.extend_from_slice(&[
1368 block_type,
1369 (entry_count >> 16) as u8,
1370 (entry_count >> 8) as u8,
1371 entry_count as u8,
1372 key_size,
1373 value_type.map_or(FIXED_KEY_BLOCK_MIXED_VALUE_TYPE, |ty| ty.0),
1374 ]);
1375 if per_entry_type {
1378 buffer.push(val_size);
1379 }
1380
1381 Self {
1382 buffer,
1383 tail,
1384 per_entry_type,
1385 layout,
1386 }
1387 }
1388
1389 fn put<E: Entry>(&mut self, entry: &E, value_ref: &ValueRef) {
1393 match self.layout {
1394 KeyBlockLayout::HashThenKey => {
1395 self.buffer
1396 .extend_from_slice(&entry.key_hash().to_be_bytes());
1397 entry.write_key_to(self.tail);
1398 }
1399 KeyBlockLayout::KeyOnly => entry.write_key_to(self.buffer),
1400 }
1401 self.put_tail(value_ref);
1402 }
1403
1404 fn put_tail(&mut self, value_ref: &ValueRef) {
1406 if self.per_entry_type {
1407 self.tail.push(value_ref.entry_type().0);
1408 }
1409 value_ref.write_value_to(self.tail);
1410 }
1411
1412 fn finish(self) -> &'l mut Vec<u8> {
1413 self.buffer.extend_from_slice(self.tail);
1414 self.buffer
1415 }
1416}
1417
1418fn value_type_val_size(ty: EntryType) -> usize {
1423 match ty.0 {
1424 KEY_BLOCK_ENTRY_TYPE_SMALL => SMALL_VALUE_REF_SIZE,
1425 KEY_BLOCK_ENTRY_TYPE_MEDIUM => MEDIUM_VALUE_REF_SIZE,
1426 KEY_BLOCK_ENTRY_TYPE_BLOB => BLOB_VALUE_REF_SIZE,
1427 KEY_BLOCK_ENTRY_TYPE_KEY_DELETED => KEY_DELETED_REF_SIZE,
1428 ty if ty >= KEY_BLOCK_ENTRY_TYPE_KEY_VALUE_DELETED_MIN => {
1430 (ty - KEY_BLOCK_ENTRY_TYPE_KEY_VALUE_DELETED_MIN) as usize
1431 }
1432 ty if ty >= KEY_BLOCK_ENTRY_TYPE_INLINE_MIN => {
1433 (ty - KEY_BLOCK_ENTRY_TYPE_INLINE_MIN) as usize
1434 }
1435 _ => panic!("Invalid key block entry type: {:?}", ty),
1436 }
1437}
1438
1439struct IndexBlockBuilder<W: Write> {
1445 writer: W,
1446}
1447
1448pub(crate) const INDEX_BLOCK_ENTRY_SIZE: usize = size_of::<u64>() + size_of::<u16>();
1450
1451pub(crate) const INDEX_BLOCK_HEADER_SIZE: usize = size_of::<u8>() + size_of::<u16>();
1453
1454impl<W: Write> IndexBlockBuilder<W> {
1455 fn new(mut writer: W, first_block: u16) -> Self {
1458 writer.write_u8(BLOCK_TYPE_INDEX).unwrap();
1459 writer.write_u16::<BE>(first_block).unwrap();
1460 Self { writer }
1461 }
1462
1463 fn put(&mut self, hash: u64, block: u16) {
1465 self.writer.write_u64::<BE>(hash).unwrap();
1466 self.writer.write_u16::<BE>(block).unwrap();
1467 }
1468}
1469
1470#[cfg(test)]
1471mod tests {
1472 use super::*;
1473 use crate::{
1474 key::hash_key,
1475 lookup_entry::LookupValue,
1476 static_sorted_file::{
1477 BlockCache, SstLookupResult, StaticSortedFile, StaticSortedFileMetaData,
1478 },
1479 };
1480
1481 fn make_cache() -> BlockCache {
1482 BlockCache::with(
1483 100,
1484 4 * 1024 * 1024,
1485 Default::default(),
1486 Default::default(),
1487 Default::default(),
1488 )
1489 }
1490
1491 struct TestEntry {
1493 key: Vec<u8>,
1494 hash: u64,
1495 value_kind: TestValueKind,
1496 }
1497
1498 enum TestValueKind {
1499 Inline(Vec<u8>),
1500 Small(Vec<u8>),
1501 Medium(Vec<u8>),
1502 MediumRaw(Vec<u8>),
1504 Blob(u32),
1505 KeyDeleted,
1506 KeyValueDeleted(Vec<u8>),
1507 }
1508
1509 impl TestEntry {
1510 fn new(key: &[u8], value_kind: TestValueKind) -> Self {
1511 let key = key.to_vec();
1512 let hash = hash_key(&key);
1513 Self {
1514 key,
1515 hash,
1516 value_kind,
1517 }
1518 }
1519
1520 fn small(key: &[u8], value: &[u8]) -> Self {
1521 Self::new(key, TestValueKind::Small(value.to_vec()))
1522 }
1523
1524 fn inline(key: &[u8], value: &[u8]) -> Self {
1525 debug_assert!(value.len() <= MAX_INLINE_VALUE_SIZE);
1526 Self::new(key, TestValueKind::Inline(value.to_vec()))
1527 }
1528
1529 fn medium(key: &[u8], value: &[u8]) -> Self {
1530 Self::new(key, TestValueKind::Medium(value.to_vec()))
1531 }
1532
1533 fn blob(key: &[u8], blob_id: u32) -> Self {
1534 Self::new(key, TestValueKind::Blob(blob_id))
1535 }
1536
1537 fn deleted(key: &[u8]) -> Self {
1538 Self::new(key, TestValueKind::KeyDeleted)
1539 }
1540
1541 fn medium_raw(key: &[u8], value: &[u8]) -> Self {
1542 Self::new(key, TestValueKind::MediumRaw(value.to_vec()))
1543 }
1544
1545 fn expected_value(&self) -> Option<&[u8]> {
1546 match &self.value_kind {
1547 TestValueKind::Inline(v)
1548 | TestValueKind::Small(v)
1549 | TestValueKind::Medium(v)
1550 | TestValueKind::MediumRaw(v) => Some(v),
1551 _ => None,
1552 }
1553 }
1554 }
1555
1556 impl Entry for TestEntry {
1557 fn key_hash(&self) -> u64 {
1558 self.hash
1559 }
1560
1561 fn key_len(&self) -> usize {
1562 self.key.len()
1563 }
1564
1565 fn key_bytes(&self) -> &[u8] {
1566 &self.key
1567 }
1568
1569 fn value(&self) -> EntryValue<'_> {
1570 match &self.value_kind {
1571 TestValueKind::Inline(v) => EntryValue::Inline { value: v },
1572 TestValueKind::Small(v) => EntryValue::Small { value: v },
1573 TestValueKind::Medium(v) => EntryValue::Medium { value: v },
1574 TestValueKind::MediumRaw(v) => EntryValue::MediumRaw {
1575 uncompressed_size: 0,
1577 checksum: checksum_block(v),
1578 block: v,
1579 },
1580 TestValueKind::Blob(id) => EntryValue::Large { blob: *id },
1581 TestValueKind::KeyDeleted => EntryValue::KeyDeleted,
1582 TestValueKind::KeyValueDeleted(v) => EntryValue::KeyValueDeleted { value: v },
1583 }
1584 }
1585 }
1586
1587 fn sort_entries(entries: &mut [TestEntry]) {
1589 entries.sort_by(|a, b| a.hash.cmp(&b.hash).then_with(|| a.key.cmp(&b.key)));
1590 }
1591
1592 fn open_sst(
1594 dir: &Path,
1595 seq: u32,
1596 meta: &StaticSortedFileBuilderMeta<'_>,
1597 ) -> Result<StaticSortedFile> {
1598 StaticSortedFile::open(
1599 dir,
1600 StaticSortedFileMetaData {
1601 sequence_number: seq,
1602 block_count: meta.block_count,
1603 },
1604 Compression::Lz4,
1605 crate::mmap_access_mode(),
1606 )
1607 }
1608
1609 fn write_sst(
1611 dir: &Path,
1612 seq: u32,
1613 entries: &[TestEntry],
1614 flags: MetaEntryFlags,
1615 ) -> Result<StaticSortedFileBuilderMeta<'static>> {
1616 let sst_path = dir.join(format!("{seq:08}.sst"));
1617 let mut writer =
1618 StreamingSstWriter::new(&sst_path, flags, entries.len() as u64, Compression::Lz4)?;
1619 for entry in entries {
1620 if let Err(err) = writer.add(entry) {
1621 writer.cancel();
1622 return Err(err);
1623 }
1624 }
1625 let (meta, _file) = writer.close()?;
1626 Ok(meta)
1627 }
1628
1629 fn assert_lookup(
1631 sst: &StaticSortedFile,
1632 entry: &TestEntry,
1633 kc: &BlockCache,
1634 vc: &BlockCache,
1635 ) -> Result<()> {
1636 let result = sst.lookup::<_, false>(entry.hash, &entry.key, kc, vc)?;
1637 match (&entry.value_kind, result) {
1638 (_, SstLookupResult::Found(values))
1639 if values.len() == 1 && matches!(values[0], LookupValue::Slice { .. }) =>
1640 {
1641 let LookupValue::Slice { value } = &values[0] else {
1642 unreachable!()
1643 };
1644 let expected = entry
1645 .expected_value()
1646 .expect("Got Slice but entry has no value");
1647 assert_eq!(
1648 value.as_ref(),
1649 expected,
1650 "value mismatch for key {:?}",
1651 std::str::from_utf8(&entry.key)
1652 );
1653 }
1654 (TestValueKind::Blob(expected_id), SstLookupResult::Found(values))
1655 if values.len() == 1 && matches!(values[0], LookupValue::Blob { .. }) =>
1656 {
1657 let LookupValue::Blob { sequence_number } = &values[0] else {
1658 unreachable!()
1659 };
1660 assert_eq!(*sequence_number, *expected_id);
1661 }
1662 (TestValueKind::KeyDeleted, SstLookupResult::Found(values))
1663 if values.len() == 1 && matches!(values[0], LookupValue::KeyDeleted) => {}
1664 (TestValueKind::KeyValueDeleted(expected), SstLookupResult::Found(values))
1665 if values.len() == 1
1666 && matches!(values[0], LookupValue::KeyValueDeleted { .. }) =>
1667 {
1668 let LookupValue::KeyValueDeleted { value } = &values[0] else {
1669 unreachable!()
1670 };
1671 assert_eq!(
1672 value.as_ref(),
1673 expected.as_slice(),
1674 "tombstone value mismatch for key {:?}",
1675 std::str::from_utf8(&entry.key)
1676 );
1677 }
1678 _ => {
1679 panic!(
1680 "Unexpected lookup result for key {:?}",
1681 std::str::from_utf8(&entry.key)
1682 );
1683 }
1684 }
1685 Ok(())
1686 }
1687
1688 #[test]
1689 fn single_inline_entry() -> Result<()> {
1690 let dir = tempfile::tempdir()?;
1691 let mut entries = vec![TestEntry::inline(b"key1", b"val1")];
1692 sort_entries(&mut entries);
1693
1694 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1695 assert_eq!(meta.entries, 1);
1696
1697 let sst = open_sst(dir.path(), 1, &meta)?;
1698 let kc = make_cache();
1699 let vc = make_cache();
1700 assert_lookup(&sst, &entries[0], &kc, &vc)?;
1701 Ok(())
1702 }
1703
1704 #[test]
1705 fn single_small_entry() -> Result<()> {
1706 let dir = tempfile::tempdir()?;
1707 let value = vec![0xAB; 100]; let mut entries = vec![TestEntry::small(b"skey", &value)];
1709 sort_entries(&mut entries);
1710
1711 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1712 assert_eq!(meta.entries, 1);
1713
1714 let sst = open_sst(dir.path(), 1, &meta)?;
1715 let kc = make_cache();
1716 let vc = make_cache();
1717 assert_lookup(&sst, &entries[0], &kc, &vc)?;
1718 Ok(())
1719 }
1720
1721 #[test]
1722 fn single_medium_entry() -> Result<()> {
1723 let dir = tempfile::tempdir()?;
1724 let value = vec![0xCD; 8192]; let mut entries = vec![TestEntry::medium(b"mkey", &value)];
1726 sort_entries(&mut entries);
1727
1728 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1729 assert_eq!(meta.entries, 1);
1730
1731 let sst = open_sst(dir.path(), 1, &meta)?;
1732 let kc = make_cache();
1733 let vc = make_cache();
1734 assert_lookup(&sst, &entries[0], &kc, &vc)?;
1735 Ok(())
1736 }
1737
1738 #[test]
1739 fn single_blob_entry() -> Result<()> {
1740 let dir = tempfile::tempdir()?;
1741 let mut entries = vec![TestEntry::blob(b"bkey", 42)];
1742 sort_entries(&mut entries);
1743
1744 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1745 assert_eq!(meta.entries, 1);
1746
1747 let sst = open_sst(dir.path(), 1, &meta)?;
1748 let kc = make_cache();
1749 let vc = make_cache();
1750 assert_lookup(&sst, &entries[0], &kc, &vc)?;
1751 Ok(())
1752 }
1753
1754 #[test]
1755 fn single_deleted_entry() -> Result<()> {
1756 let dir = tempfile::tempdir()?;
1757 let mut entries = vec![TestEntry::deleted(b"dkey")];
1758 sort_entries(&mut entries);
1759
1760 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1761 assert_eq!(meta.entries, 1);
1762
1763 let sst = open_sst(dir.path(), 1, &meta)?;
1764 let kc = make_cache();
1765 let vc = make_cache();
1766 assert_lookup(&sst, &entries[0], &kc, &vc)?;
1767 Ok(())
1768 }
1769
1770 #[test]
1771 fn many_small_values() -> Result<()> {
1772 let dir = tempfile::tempdir()?;
1773 let count = 200;
1776 let mut entries: Vec<TestEntry> = (0..count)
1777 .map(|i| {
1778 let key = format!("key-{i:04}");
1779 let value = vec![(i & 0xFF) as u8; 200];
1780 TestEntry::small(key.as_bytes(), &value)
1781 })
1782 .collect();
1783 sort_entries(&mut entries);
1784
1785 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1786 assert_eq!(meta.entries, count as u64);
1787
1788 let sst = open_sst(dir.path(), 1, &meta)?;
1789 let kc = make_cache();
1790 let vc = make_cache();
1791
1792 for entry in &entries {
1793 assert_lookup(&sst, entry, &kc, &vc)?;
1794 }
1795 Ok(())
1796 }
1797
1798 #[test]
1799 fn many_medium_values() -> Result<()> {
1800 let dir = tempfile::tempdir()?;
1801 let count = 50;
1802 let mut entries: Vec<TestEntry> = (0..count)
1803 .map(|i| {
1804 let key = format!("mkey-{i:04}");
1805 let value = vec![(i & 0xFF) as u8; 8192];
1806 TestEntry::medium(key.as_bytes(), &value)
1807 })
1808 .collect();
1809 sort_entries(&mut entries);
1810
1811 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1812 assert_eq!(meta.entries, count as u64);
1813
1814 let sst = open_sst(dir.path(), 1, &meta)?;
1815 let kc = make_cache();
1816 let vc = make_cache();
1817
1818 for entry in &entries {
1819 assert_lookup(&sst, entry, &kc, &vc)?;
1820 }
1821 Ok(())
1822 }
1823
1824 #[test]
1825 fn mixed_value_types() -> Result<()> {
1826 let dir = tempfile::tempdir()?;
1827 let mut entries = vec![
1828 TestEntry::inline(b"a-inline", b"tiny"),
1829 TestEntry::small(b"b-small", &[0x11; 200]),
1830 TestEntry::medium(b"c-medium", &[0x22; 8192]),
1831 TestEntry::blob(b"d-blob", 99),
1832 TestEntry::deleted(b"e-deleted"),
1833 TestEntry::small(b"f-small2", &[0x33; 300]),
1834 TestEntry::inline(b"g-inline2", b"mini"),
1835 TestEntry::medium(b"h-medium2", &[0x44; 16384]),
1836 ];
1837 sort_entries(&mut entries);
1838
1839 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
1840 assert_eq!(meta.entries, 8);
1841
1842 let sst = open_sst(dir.path(), 1, &meta)?;
1843 let kc = make_cache();
1844 let vc = make_cache();
1845
1846 for entry in &entries {
1847 assert_lookup(&sst, entry, &kc, &vc)?;
1848 }
1849 Ok(())
1850 }
1851
1852 #[test]
1853 #[cfg(debug_assertions)]
1854 fn cancel_after_failed_add_does_not_trigger_drop_assertion() {
1855 let dir = tempfile::tempdir().unwrap();
1856 let sst_path = dir.path().join("test.sst");
1857 let mut writer =
1858 StreamingSstWriter::new(&sst_path, MetaEntryFlags::default(), 1, Compression::Lz4)
1859 .unwrap();
1860
1861 drop(writer.file.take());
1864 writer.file = Some(BufWriter::with_capacity(0, File::open(&sst_path).unwrap()));
1865 let error = writer
1866 .add(TestEntry::medium(b"key", &[0; 8192]))
1867 .unwrap_err();
1868 assert!(format!("{error:#}").contains("Failed to write value block"));
1869
1870 let cancel_result =
1871 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| writer.cancel()));
1872 assert!(
1873 cancel_result.is_ok(),
1874 "cancelling a writer after an add error must not trigger the lifecycle assertion"
1875 );
1876 assert!(!sst_path.exists(), "cancel should remove the partial SST");
1877 }
1878
1879 #[test]
1880 fn failed_close_removes_partial_file() {
1881 let dir = tempfile::tempdir().unwrap();
1882 let sst_path = dir.path().join("test.sst");
1883 let mut writer =
1884 StreamingSstWriter::new(&sst_path, MetaEntryFlags::default(), 1, Compression::Lz4)
1885 .unwrap();
1886 writer.add(TestEntry::inline(b"key", b"value")).unwrap();
1887
1888 drop(writer.file.take());
1890 writer.file = Some(BufWriter::with_capacity(0, File::open(&sst_path).unwrap()));
1891 let error = writer.close().unwrap_err();
1892
1893 assert!(format!("{error:#}").contains("Failed to write key block"));
1894 assert!(
1895 !sst_path.exists(),
1896 "a failed close should remove the partial SST"
1897 );
1898 }
1899
1900 #[test]
1901 fn is_full_entry_count_limit() {
1902 let dir = tempfile::tempdir().unwrap();
1903 let sst_path = dir.path().join("test.sst");
1904 let mut writer =
1905 StreamingSstWriter::new(&sst_path, MetaEntryFlags::default(), 100, Compression::Lz4)
1906 .unwrap();
1907
1908 let max_entries = 50;
1909 for i in 0..max_entries {
1910 let key = format!("k{i:06}");
1911 let entry = TestEntry::inline(key.as_bytes(), &[0; 4]);
1912 writer.add(entry).unwrap();
1913 }
1914
1915 assert_eq!(writer.entry_count, max_entries as u64);
1916 assert!(
1917 writer.is_full(max_entries, usize::MAX),
1918 "Should be full when entry count reaches max_entries"
1919 );
1920 assert!(
1921 !writer.is_full(max_entries + 1, usize::MAX),
1922 "Should not be full when limit is higher"
1923 );
1924 writer.close().unwrap();
1925 }
1926
1927 #[test]
1928 fn is_full_data_size_limit() {
1929 let dir = tempfile::tempdir().unwrap();
1930 let sst_path = dir.path().join("test.sst");
1931 let mut writer =
1932 StreamingSstWriter::new(&sst_path, MetaEntryFlags::default(), 100, Compression::Lz4)
1933 .unwrap();
1934
1935 let value = vec![0u8; 1000];
1936 for i in 0..10 {
1937 let key = format!("k{i:06}");
1938 let entry = TestEntry::small(key.as_bytes(), &value);
1939 writer.add(entry).unwrap();
1940 }
1941
1942 let total = writer.total_key_size + writer.total_value_size;
1943 assert!(total > 10_000, "total data should exceed 10KB");
1944 assert!(writer.is_full(usize::MAX, total - 1));
1945 assert!(!writer.is_full(usize::MAX, total + 1));
1946 writer.close().unwrap();
1947 }
1948
1949 #[test]
1950 fn write_static_stored_file_matches_streaming() -> Result<()> {
1951 let dir = tempfile::tempdir()?;
1952
1953 let mut entries: Vec<TestEntry> = (0..100)
1954 .map(|i| {
1955 let key = format!("rkey-{i:04}");
1956 if i % 3 == 0 {
1957 TestEntry::inline(key.as_bytes(), &[(i & 0xFF) as u8; 4])
1958 } else if i % 3 == 1 {
1959 TestEntry::small(key.as_bytes(), &[(i & 0xFF) as u8; 200])
1960 } else {
1961 TestEntry::medium(key.as_bytes(), &[(i & 0xFF) as u8; 8192])
1962 }
1963 })
1964 .collect();
1965 sort_entries(&mut entries);
1966
1967 let batch_path = dir.path().join("00000001.sst");
1969 let (meta1, _) = write_static_stored_file(
1970 &entries,
1971 &batch_path,
1972 MetaEntryFlags::default(),
1973 Compression::Lz4,
1974 )?;
1975
1976 let streaming_path = dir.path().join("00000002.sst");
1978 let mut writer = StreamingSstWriter::new(
1979 &streaming_path,
1980 MetaEntryFlags::default(),
1981 entries.len() as u64,
1982 Compression::Lz4,
1983 )?;
1984 for entry in &entries {
1985 if let Err(err) = writer.add(entry) {
1986 writer.cancel();
1987 return Err(err);
1988 }
1989 }
1990 let (meta2, _) = writer.close()?;
1991
1992 assert_eq!(meta1.entries, meta2.entries);
1994 assert_eq!(meta1.min_hash, meta2.min_hash);
1995 assert_eq!(meta1.max_hash, meta2.max_hash);
1996 assert_eq!(meta1.block_count, meta2.block_count);
1997
1998 let sst1 = StaticSortedFile::open(
2000 dir.path(),
2001 StaticSortedFileMetaData {
2002 sequence_number: 1,
2003 block_count: meta1.block_count,
2004 },
2005 Compression::Lz4,
2006 crate::mmap_access_mode(),
2007 )?;
2008 let sst2 = StaticSortedFile::open(
2009 dir.path(),
2010 StaticSortedFileMetaData {
2011 sequence_number: 2,
2012 block_count: meta2.block_count,
2013 },
2014 Compression::Lz4,
2015 crate::mmap_access_mode(),
2016 )?;
2017 let kc = make_cache();
2018 let vc = make_cache();
2019
2020 for entry in &entries {
2021 let r1 = sst1.lookup::<_, false>(entry.hash, &entry.key, &kc, &vc)?;
2022 let r2 = sst2.lookup::<_, false>(entry.hash, &entry.key, &kc, &vc)?;
2023 match (&r1, &r2) {
2024 (SstLookupResult::Found(v1), SstLookupResult::Found(v2))
2025 if v1.len() == 1 && v2.len() == 1 =>
2026 {
2027 match (&v1[0], &v2[0]) {
2028 (
2029 LookupValue::Slice { value: val1 },
2030 LookupValue::Slice { value: val2 },
2031 ) => {
2032 assert_eq!(
2033 val1.as_ref(),
2034 val2.as_ref(),
2035 "Value mismatch for key {:?}",
2036 std::str::from_utf8(&entry.key)
2037 );
2038 }
2039 (LookupValue::KeyDeleted, LookupValue::KeyDeleted) => {}
2040 (
2041 LookupValue::Blob {
2042 sequence_number: s1,
2043 },
2044 LookupValue::Blob {
2045 sequence_number: s2,
2046 },
2047 ) => {
2048 assert_eq!(s1, s2);
2049 }
2050 _ => panic!(
2051 "Mismatched results for key {:?}",
2052 std::str::from_utf8(&entry.key)
2053 ),
2054 }
2055 }
2056 _ => panic!(
2057 "Mismatched results for key {:?}",
2058 std::str::from_utf8(&entry.key)
2059 ),
2060 }
2061 }
2062 Ok(())
2063 }
2064
2065 #[test]
2066 #[should_panic(expected = "StreamingSstWriter::close() called with no entries")]
2067 fn close_empty_writer_panics() {
2068 let dir = tempfile::tempdir().unwrap();
2069 let sst_path = dir.path().join("empty.sst");
2070 let writer = StreamingSstWriter::<TestEntry>::new(
2071 &sst_path,
2072 MetaEntryFlags::default(),
2073 0,
2074 Compression::Lz4,
2075 )
2076 .unwrap();
2077 writer.close().unwrap();
2078 }
2079
2080 #[test]
2081 fn key_block_boundary_at_max_entries() -> Result<()> {
2082 let dir = tempfile::tempdir()?;
2083 let count = MAX_KEY_BLOCK_ENTRIES + 1;
2084 let mut entries: Vec<TestEntry> = (0..count)
2085 .map(|i| {
2086 let key = format!("boundary-{i:06}");
2087 TestEntry::inline(key.as_bytes(), &[0u8; 4])
2088 })
2089 .collect();
2090 sort_entries(&mut entries);
2091
2092 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
2093 assert_eq!(meta.entries, count as u64);
2094 assert!(
2096 meta.block_count >= 3,
2097 "expected at least 2 key blocks + 1 index block"
2098 );
2099
2100 let sst = open_sst(dir.path(), 1, &meta)?;
2101 let kc = make_cache();
2102 let vc = make_cache();
2103 for entry in &entries {
2104 assert_lookup(&sst, entry, &kc, &vc)?;
2105 }
2106 Ok(())
2107 }
2108
2109 fn read_first_block(dir: &Path, seq: u32, block_count: u16) -> Result<Vec<u8>> {
2113 let data = fs_err::read(dir.join(format!("{seq:08}.sst")))?;
2114 let offsets_start = data.len() - block_count as usize * size_of::<u32>();
2115 let end = BE::read_u32(&data[offsets_start..]) as usize;
2116 let raw = &data[..end];
2117 let uncompressed_size = BE::read_u32(raw) as usize;
2120 let body = &raw[BLOCK_HEADER_SIZE..];
2121 Ok(if uncompressed_size == 0 {
2122 body.to_vec()
2123 } else {
2124 let mut out = vec![0u8; uncompressed_size];
2125 lz4_flex::block::decompress_into(body, &mut out)?;
2126 out
2127 })
2128 }
2129
2130 #[test]
2136 fn fixed_layout_survives_mixed_value_types_of_equal_size() -> Result<()> {
2137 let dir = tempfile::tempdir()?;
2138
2139 let mut entries: Vec<TestEntry> = (0..64u64)
2141 .map(|i| {
2142 let key = format!("k-{i:06}");
2143 if i % 4 == 0 {
2144 TestEntry::new(
2145 key.as_bytes(),
2146 TestValueKind::KeyValueDeleted(vec![0xAAu8; 4]),
2147 )
2148 } else {
2149 TestEntry::inline(key.as_bytes(), &[0xBBu8; 4])
2150 }
2151 })
2152 .collect();
2153 sort_entries(&mut entries);
2154
2155 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
2156 let block = read_first_block(dir.path(), 1, meta.block_count)?;
2157
2158 assert_eq!(
2159 KeyBlockLayout::from_block_type(block[0]).map(|(_, fixed)| fixed),
2160 Some(true),
2161 "mixed value types of equal size should stay in fixed layout, got block type {}",
2162 block[0]
2163 );
2164 assert_eq!(
2165 block[5], FIXED_KEY_BLOCK_MIXED_VALUE_TYPE,
2166 "block should be marked mixed-type"
2167 );
2168 assert_eq!(block[6], 4, "value size should be recorded in the header");
2169
2170 let sst = open_sst(dir.path(), 1, &meta)?;
2172 let kc = make_cache();
2173 let vc = make_cache();
2174 for entry in &entries {
2175 assert_lookup(&sst, entry, &kc, &vc)?;
2176 }
2177 Ok(())
2178 }
2179
2180 #[test]
2183 fn mixed_value_sizes_fall_back_to_variable_layout() -> Result<()> {
2184 let dir = tempfile::tempdir()?;
2185
2186 let mut entries: Vec<TestEntry> = (0..64u64)
2187 .map(|i| {
2188 let key = format!("k-{i:06}");
2189 let len = if i % 4 == 0 { 2 } else { 4 };
2190 TestEntry::inline(key.as_bytes(), &vec![0xCCu8; len])
2191 })
2192 .collect();
2193 sort_entries(&mut entries);
2194
2195 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
2196 let block = read_first_block(dir.path(), 1, meta.block_count)?;
2197 assert_eq!(
2198 KeyBlockLayout::from_block_type(block[0]).map(|(_, fixed)| fixed),
2199 Some(false),
2200 "differing value sizes should use variable layout, got block type {}",
2201 block[0]
2202 );
2203
2204 let sst = open_sst(dir.path(), 1, &meta)?;
2205 let kc = make_cache();
2206 let vc = make_cache();
2207 for entry in &entries {
2208 assert_lookup(&sst, entry, &kc, &vc)?;
2209 }
2210 Ok(())
2211 }
2212
2213 #[test]
2214 fn single_medium_raw_entry() -> Result<()> {
2215 let dir = tempfile::tempdir()?;
2216 let value = vec![0xBE; 8192];
2217 let mut entries = vec![TestEntry::medium_raw(b"rkey", &value)];
2218 sort_entries(&mut entries);
2219
2220 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default())?;
2221 assert_eq!(meta.entries, 1);
2222
2223 let sst = open_sst(dir.path(), 1, &meta)?;
2224 let kc = make_cache();
2225 let vc = make_cache();
2226 assert_lookup(&sst, &entries[0], &kc, &vc)?;
2227 Ok(())
2228 }
2229
2230 fn corrupt_sst_byte(dir: &Path, seq: u32, pos: u64) {
2232 use std::io::{Seek, SeekFrom, Write as _};
2233
2234 let sst_path = dir.join(format!("{seq:08}.sst"));
2235 let file_bytes = std::fs::read(&sst_path).unwrap();
2236 let original = file_bytes[pos as usize];
2237 let mut file = std::fs::OpenOptions::new()
2238 .write(true)
2239 .open(&sst_path)
2240 .unwrap();
2241 file.seek(SeekFrom::Start(pos)).unwrap();
2242 file.write_all(&[original ^ 0xFF]).unwrap();
2243 file.sync_all().unwrap();
2244 }
2245
2246 fn assert_corruption_detected(
2248 dir: &Path,
2249 seq: u32,
2250 meta: &StaticSortedFileBuilderMeta<'_>,
2251 entries: &[TestEntry],
2252 ) {
2253 let sst = open_sst(dir, seq, meta).unwrap();
2254 let kc = make_cache();
2255 let vc = make_cache();
2256 match sst.lookup::<_, false>(entries[0].hash, &entries[0].key, &kc, &vc) {
2257 Err(err) => {
2258 let msg = format!("{err}");
2259 assert!(
2260 msg.contains("corruption"),
2261 "Expected corruption error, got: {msg}"
2262 );
2263 }
2264 Ok(_) => panic!("Expected checksum error, but lookup succeeded"),
2265 }
2266 }
2267
2268 #[test]
2269 fn checksum_detects_corrupted_compressed_block() {
2270 let dir = tempfile::tempdir().unwrap();
2271 let value = vec![0xCD; 8192];
2273 let entries = vec![TestEntry::medium(b"mkey", &value)];
2274
2275 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default()).unwrap();
2276
2277 corrupt_sst_byte(dir.path(), 1, 4);
2280 assert_corruption_detected(dir.path(), 1, &meta, &entries);
2281 }
2282
2283 #[test]
2284 fn checksum_detects_corrupted_uncompressed_block() {
2285 let dir = tempfile::tempdir().unwrap();
2286 let entries = vec![TestEntry::inline(b"key1", b"val1")];
2288
2289 let meta = write_sst(dir.path(), 1, &entries, MetaEntryFlags::default()).unwrap();
2290
2291 corrupt_sst_byte(dir.path(), 1, BLOCK_HEADER_SIZE as u64 + 1);
2293 assert_corruption_detected(dir.path(), 1, &meta, &entries);
2294 }
2295
2296 #[test]
2297 fn be_key_order_matches_byte_order() {
2298 let keys4: Vec<[u8; 4]> = vec![
2299 [0, 0, 0, 0],
2300 [0, 0, 0, 1],
2301 [0, 0, 1, 0],
2302 [0x7f, 0xff, 0xff, 0xff],
2303 [0x80, 0, 0, 0],
2304 [0xff, 0xfe, 0, 0],
2305 [0xff, 0xff, 0xff, 0xff],
2306 ];
2307 for a in &keys4 {
2308 for b in &keys4 {
2309 assert_eq!(
2310 be_key_u32(a).cmp(&be_key_u32(b)),
2311 a[..].cmp(&b[..]),
2312 "u32 order disagrees with byte order for {a:?} vs {b:?}"
2313 );
2314 }
2315 }
2316 let keys8: Vec<[u8; 8]> = vec![
2317 [0; 8],
2318 [0, 0, 0, 0, 0, 0, 0, 1],
2319 [0x7f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff],
2320 [0x80, 0, 0, 0, 0, 0, 0, 0],
2321 [0xff; 8],
2322 ];
2323 for a in &keys8 {
2324 for b in &keys8 {
2325 assert_eq!(
2326 be_key_u64(a).cmp(&be_key_u64(b)),
2327 a[..].cmp(&b[..]),
2328 "u64 order disagrees with byte order for {a:?} vs {b:?}"
2329 );
2330 }
2331 }
2332 }
2333
2334 #[test]
2337 fn uniform_key_len_requires_equal_lengths() {
2338 let mut acc = KeyBlockAccumulator::new();
2339 assert_eq!(acc.flush_info().uniform_key_len(), None, "empty block");
2340
2341 let ty = EntryType(KEY_BLOCK_ENTRY_TYPE_INLINE_MIN);
2342 acc.add(8, 1, ty);
2343 acc.add(8, 2, ty);
2344 assert_eq!(acc.flush_info().uniform_key_len(), Some(8));
2345
2346 acc.add(4, 3, ty);
2347 assert_eq!(
2348 acc.flush_info().uniform_key_len(),
2349 None,
2350 "mixed lengths must not report a uniform length"
2351 );
2352
2353 acc.reset();
2354 acc.add(4, 4, ty);
2355 assert_eq!(
2356 acc.flush_info().uniform_key_len(),
2357 Some(4),
2358 "reset clears min"
2359 );
2360 }
2361}