1use super::{Config, Error, Identifier};
2use crate::{
3 journal::segmented::oversized::{
4 Config as OversizedConfig, Oversized, Record as OversizedRecord,
5 },
6 Context,
7};
8use commonware_codec::{CodecShared, FixedArray, FixedSize, Read, ReadExt, Write as CodecWrite};
9use commonware_cryptography::{crc32, Crc32, Hasher};
10use commonware_runtime::{
11 buffer,
12 iobuf::EncodeExt,
13 telemetry::metrics::{Counter, MetricsExt as _},
14 Blob, Buf, BufMut, BufferPooler, IoBuf,
15};
16use commonware_utils::{Array, Span};
17use futures::future::try_join;
18use std::{cmp::Ordering, collections::BTreeSet, num::NonZeroUsize, ops::Deref};
19use tracing::debug;
20
21const RESIZE_THRESHOLD: u64 = 50;
24
25#[derive(Hash, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, FixedArray)]
30#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
31#[repr(transparent)]
32pub struct Cursor([u8; u64::SIZE + u64::SIZE + u32::SIZE]);
33
34impl Cursor {
35 fn new(section: u64, offset: u64, size: u32) -> Self {
37 let mut buf = [0u8; u64::SIZE + u64::SIZE + u32::SIZE];
38 buf[..u64::SIZE].copy_from_slice(§ion.to_be_bytes());
39 buf[u64::SIZE..u64::SIZE + u64::SIZE].copy_from_slice(&offset.to_be_bytes());
40 buf[u64::SIZE + u64::SIZE..].copy_from_slice(&size.to_be_bytes());
41 Self(buf)
42 }
43
44 fn section(&self) -> u64 {
46 u64::from_be_bytes(self.0[..u64::SIZE].try_into().unwrap())
47 }
48
49 fn offset(&self) -> u64 {
51 u64::from_be_bytes(self.0[u64::SIZE..u64::SIZE + u64::SIZE].try_into().unwrap())
52 }
53
54 fn size(&self) -> u32 {
56 u32::from_be_bytes(self.0[u64::SIZE + u64::SIZE..].try_into().unwrap())
57 }
58}
59
60impl Read for Cursor {
61 type Cfg = ();
62
63 fn read_cfg(buf: &mut impl Buf, _: &Self::Cfg) -> Result<Self, commonware_codec::Error> {
64 <[u8; u64::SIZE + u64::SIZE + u32::SIZE]>::read(buf).map(Self)
65 }
66}
67
68impl CodecWrite for Cursor {
69 fn write(&self, buf: &mut impl BufMut) {
70 self.0.write(buf);
71 }
72}
73
74impl FixedSize for Cursor {
75 const SIZE: usize = u64::SIZE + u64::SIZE + u32::SIZE;
76}
77
78impl Span for Cursor {}
79
80impl Array for Cursor {}
81
82impl Deref for Cursor {
83 type Target = [u8];
84 fn deref(&self) -> &Self::Target {
85 &self.0
86 }
87}
88
89impl AsRef<[u8]> for Cursor {
90 fn as_ref(&self) -> &[u8] {
91 &self.0
92 }
93}
94
95impl std::fmt::Debug for Cursor {
96 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97 write!(
98 f,
99 "Cursor(section={}, offset={}, size={})",
100 self.section(),
101 self.offset(),
102 self.size()
103 )
104 }
105}
106
107impl std::fmt::Display for Cursor {
108 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109 write!(
110 f,
111 "Cursor(section={}, offset={}, size={})",
112 self.section(),
113 self.offset(),
114 self.size()
115 )
116 }
117}
118
119#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Copy)]
124#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
125pub struct Checkpoint {
126 epoch: u64,
128 section: u64,
130 oversized_size: u64,
132 table_size: u32,
134}
135
136impl Checkpoint {
137 const fn init(table_size: u32) -> Self {
139 Self {
140 table_size,
141 epoch: 0,
142 section: 0,
143 oversized_size: 0,
144 }
145 }
146
147 const fn is_empty(&self) -> bool {
149 self.epoch == 0 && self.section == 0 && self.oversized_size == 0 && self.table_size == 0
150 }
151}
152
153impl Read for Checkpoint {
154 type Cfg = ();
155 fn read_cfg(buf: &mut impl Buf, _: &()) -> Result<Self, commonware_codec::Error> {
156 let epoch = u64::read(buf)?;
157 let section = u64::read(buf)?;
158 let oversized_size = u64::read(buf)?;
159 let table_size = u32::read(buf)?;
160 Ok(Self {
161 epoch,
162 section,
163 oversized_size,
164 table_size,
165 })
166 }
167}
168
169impl CodecWrite for Checkpoint {
170 fn write(&self, buf: &mut impl BufMut) {
171 self.epoch.write(buf);
172 self.section.write(buf);
173 self.oversized_size.write(buf);
174 self.table_size.write(buf);
175 }
176}
177
178impl FixedSize for Checkpoint {
179 const SIZE: usize = u64::SIZE + u64::SIZE + u64::SIZE + u32::SIZE;
180}
181
182const TABLE_BLOB_NAME: &[u8] = b"table";
184
185#[derive(Debug, Clone, PartialEq)]
187#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
188struct Entry {
189 epoch: u64,
191 section: u64,
193 position: u64,
195 added: u8,
197 crc: u32,
199}
200
201impl Entry {
202 const FULL_SIZE: usize = Self::SIZE * 2;
204
205 fn compute_crc(epoch: u64, section: u64, position: u64, added: u8) -> u32 {
207 let mut hasher = Crc32::new();
208 hasher.update(&epoch.to_be_bytes());
209 hasher.update(§ion.to_be_bytes());
210 hasher.update(&position.to_be_bytes());
211 hasher.update(&added.to_be_bytes());
212 hasher.finalize().as_u32()
213 }
214
215 fn new(epoch: u64, section: u64, position: u64, added: u8) -> Self {
217 Self {
218 epoch,
219 section,
220 position,
221 added,
222 crc: Self::compute_crc(epoch, section, position, added),
223 }
224 }
225
226 const fn new_empty() -> Self {
228 Self {
229 epoch: 0,
230 section: 0,
231 position: 0,
232 added: 0,
233 crc: 0,
234 }
235 }
236
237 const fn is_empty(&self) -> bool {
239 self.epoch == 0
240 && self.section == 0
241 && self.position == 0
242 && self.added == 0
243 && self.crc == 0
244 }
245
246 fn is_valid(&self) -> bool {
250 Self::compute_crc(self.epoch, self.section, self.position, self.added) == self.crc
251 }
252}
253
254impl FixedSize for Entry {
255 const SIZE: usize = u64::SIZE + u64::SIZE + u64::SIZE + u8::SIZE + crc32::Digest::SIZE;
256}
257
258impl CodecWrite for Entry {
259 fn write(&self, buf: &mut impl BufMut) {
260 self.epoch.write(buf);
261 self.section.write(buf);
262 self.position.write(buf);
263 self.added.write(buf);
264 self.crc.write(buf);
265 }
266}
267
268impl Read for Entry {
269 type Cfg = ();
270 fn read_cfg(buf: &mut impl Buf, _: &Self::Cfg) -> Result<Self, commonware_codec::Error> {
271 let epoch = u64::read(buf)?;
272 let section = u64::read(buf)?;
273 let position = u64::read(buf)?;
274 let added = u8::read(buf)?;
275 let crc = u32::read(buf)?;
276
277 Ok(Self {
278 epoch,
279 section,
280 position,
281 added,
282 crc,
283 })
284 }
285}
286
287const NO_NEXT_SECTION: u64 = u64::MAX;
289const NO_NEXT_POSITION: u64 = u64::MAX;
290
291#[derive(Debug, Clone, PartialEq)]
299struct Record<K: Array> {
300 key: K,
302 next_section: u64,
305 next_position: u64,
306 value_offset: u64,
308 value_size: u32,
310}
311
312impl<K: Array> Record<K> {
313 fn new(key: K, next: Option<(u64, u64)>, value_offset: u64, value_size: u32) -> Self {
315 let (next_section, next_position) = next.unwrap_or((NO_NEXT_SECTION, NO_NEXT_POSITION));
316 Self {
317 key,
318 next_section,
319 next_position,
320 value_offset,
321 value_size,
322 }
323 }
324
325 const fn next(&self) -> Option<(u64, u64)> {
327 if self.next_section == NO_NEXT_SECTION && self.next_position == NO_NEXT_POSITION {
328 None
329 } else {
330 Some((self.next_section, self.next_position))
331 }
332 }
333}
334
335impl<K: Array> CodecWrite for Record<K> {
336 fn write(&self, buf: &mut impl BufMut) {
337 self.key.write(buf);
338 self.next_section.write(buf);
339 self.next_position.write(buf);
340 self.value_offset.write(buf);
341 self.value_size.write(buf);
342 }
343}
344
345impl<K: Array> Read for Record<K> {
346 type Cfg = ();
347 fn read_cfg(buf: &mut impl Buf, _: &Self::Cfg) -> Result<Self, commonware_codec::Error> {
348 let key = K::read(buf)?;
349 let next_section = u64::read(buf)?;
350 let next_position = u64::read(buf)?;
351 let value_offset = u64::read(buf)?;
352 let value_size = u32::read(buf)?;
353
354 Ok(Self {
355 key,
356 next_section,
357 next_position,
358 value_offset,
359 value_size,
360 })
361 }
362}
363
364impl<K: Array> FixedSize for Record<K> {
365 const SIZE: usize = K::SIZE + u64::SIZE + u64::SIZE + u64::SIZE + u32::SIZE;
367}
368
369impl<K: Array> OversizedRecord for Record<K> {
370 fn value_location(&self) -> (u64, u32) {
371 (self.value_offset, self.value_size)
372 }
373
374 fn with_location(mut self, offset: u64, size: u32) -> Self {
375 self.value_offset = offset;
376 self.value_size = size;
377 self
378 }
379}
380
381#[cfg(feature = "arbitrary")]
382impl<K: Array> arbitrary::Arbitrary<'_> for Record<K>
383where
384 K: for<'a> arbitrary::Arbitrary<'a>,
385{
386 fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
387 Ok(Self {
388 key: K::arbitrary(u)?,
389 next_section: u64::arbitrary(u)?,
390 next_position: u64::arbitrary(u)?,
391 value_offset: u64::arbitrary(u)?,
392 value_size: u32::arbitrary(u)?,
393 })
394 }
395}
396
397pub struct Freezer<E: BufferPooler + Context, K: Array, V: CodecShared> {
399 context: E,
401
402 table_partition: String,
404 table_size: u32,
405 table_resize_threshold: u64,
406 table_resize_frequency: u8,
407 table_resize_chunk_size: u32,
408
409 table: E::Blob,
411
412 oversized: Oversized<E, Record<K>, V>,
414
415 blob_target_size: u64,
417
418 current_section: u64,
420 next_epoch: u64,
421
422 modified_sections: BTreeSet<u64>,
424 resizable: u32,
425 resize_progress: Option<u32>,
426
427 puts: Counter,
429 gets: Counter,
430 has: Counter,
431 unnecessary_reads: Counter,
432 unnecessary_writes: Counter,
433 resizes: Counter,
434}
435
436impl<E: BufferPooler + Context, K: Array, V: CodecShared> Freezer<E, K, V> {
437 #[inline]
439 const fn table_offset(table_index: u32) -> u64 {
440 table_index as u64 * Entry::FULL_SIZE as u64
441 }
442
443 fn parse_entries(mut buf: impl Buf) -> Result<(Entry, Entry), Error> {
445 let entry1 = Entry::read(&mut buf)?;
446 let entry2 = Entry::read(&mut buf)?;
447 Ok((entry1, entry2))
448 }
449
450 async fn read_table(blob: &E::Blob, table_index: u32) -> Result<(Entry, Entry), Error> {
452 let offset = Self::table_offset(table_index);
453 let read_buf = blob.read_at(offset, Entry::FULL_SIZE).await?;
454
455 Self::parse_entries(read_buf)
456 }
457
458 async fn recover_entry(
460 blob: &E::Blob,
461 entry: &mut Entry,
462 entry_offset: u64,
463 max_valid_epoch: Option<u64>,
464 max_epoch: &mut u64,
465 max_section: &mut u64,
466 ) -> Result<bool, Error> {
467 if entry.is_empty() {
468 return Ok(false);
469 }
470
471 if !entry.is_valid()
472 || (max_valid_epoch.is_some() && entry.epoch > max_valid_epoch.unwrap())
473 {
474 debug!(
475 valid_epoch = max_valid_epoch,
476 entry_epoch = entry.epoch,
477 "found invalid table entry"
478 );
479 *entry = Entry::new_empty();
480 let zero_buf = IoBuf::from(&[0u8; Entry::SIZE]);
481 blob.write_at(entry_offset, zero_buf).await?;
482 Ok(true)
483 } else if max_valid_epoch.is_none() && entry.epoch > *max_epoch {
484 *max_epoch = entry.epoch;
486 *max_section = entry.section;
487 Ok(false)
488 } else {
489 Ok(false)
490 }
491 }
492
493 async fn recover_table(
501 pooler: &impl BufferPooler,
502 blob: &E::Blob,
503 table_size: u32,
504 table_resize_frequency: u8,
505 max_valid_epoch: Option<u64>,
506 table_replay_buffer: NonZeroUsize,
507 ) -> Result<(bool, u64, u64, u32), Error> {
508 let blob_size = Self::table_offset(table_size);
510 let mut reader =
511 buffer::Read::from_pooler(pooler, blob.clone(), blob_size, table_replay_buffer);
512
513 let mut modified = false;
515 let mut max_epoch = 0u64;
516 let mut max_section = 0u64;
517 let mut resizable = 0u32;
518 for table_index in 0..table_size {
519 let offset = Self::table_offset(table_index);
520
521 let entry_buf = reader.read(Entry::FULL_SIZE).await?;
523 let (mut entry1, mut entry2) = Self::parse_entries(entry_buf)?;
524
525 let entry1_cleared = Self::recover_entry(
527 blob,
528 &mut entry1,
529 offset,
530 max_valid_epoch,
531 &mut max_epoch,
532 &mut max_section,
533 )
534 .await?;
535 let entry2_cleared = Self::recover_entry(
536 blob,
537 &mut entry2,
538 offset + Entry::SIZE as u64,
539 max_valid_epoch,
540 &mut max_epoch,
541 &mut max_section,
542 )
543 .await?;
544 modified |= entry1_cleared || entry2_cleared;
545
546 if let Some((_, _, added)) = Self::read_latest_entry(&entry1, &entry2) {
548 if added >= table_resize_frequency {
549 resizable += 1;
550 }
551 }
552 }
553
554 Ok((modified, max_epoch, max_section, resizable))
555 }
556
557 const fn compute_write_offset(entry1: &Entry, entry2: &Entry, epoch: u64) -> u64 {
559 if !entry1.is_empty() && entry1.epoch == epoch {
561 return 0;
562 }
563 if !entry2.is_empty() && entry2.epoch == epoch {
564 return Entry::SIZE as u64;
565 }
566
567 match (entry1.is_empty(), entry2.is_empty()) {
569 (true, _) => 0, (_, true) => Entry::SIZE as u64, (false, false) => {
572 if entry1.epoch < entry2.epoch {
573 0
574 } else {
575 Entry::SIZE as u64
576 }
577 }
578 }
579 }
580
581 fn read_latest_entry(entry1: &Entry, entry2: &Entry) -> Option<(u64, u64, u8)> {
583 match (
584 !entry1.is_empty() && entry1.is_valid(),
585 !entry2.is_empty() && entry2.is_valid(),
586 ) {
587 (true, true) => match entry1.epoch.cmp(&entry2.epoch) {
588 Ordering::Greater => Some((entry1.section, entry1.position, entry1.added)),
589 Ordering::Less => Some((entry2.section, entry2.position, entry2.added)),
590 Ordering::Equal => {
591 unreachable!("two valid entries with the same epoch")
592 }
593 },
594 (true, false) => Some((entry1.section, entry1.position, entry1.added)),
595 (false, true) => Some((entry2.section, entry2.position, entry2.added)),
596 (false, false) => None,
597 }
598 }
599
600 async fn update_head(
602 pooler: &impl BufferPooler,
603 table: &E::Blob,
604 table_index: u32,
605 entry1: &Entry,
606 entry2: &Entry,
607 update: Entry,
608 ) -> Result<(), Error> {
609 let table_offset = Self::table_offset(table_index);
611
612 let start = Self::compute_write_offset(entry1, entry2, update.epoch);
614
615 table
617 .write_at(
618 table_offset + start,
619 update.encode_with_pool_mut(pooler.storage_buffer_pool()),
620 )
621 .await
622 .map_err(Error::Runtime)
623 }
624
625 async fn init_table(blob: &E::Blob, table_size: u32) -> Result<(), Error> {
627 let table_len = Self::table_offset(table_size);
628 blob.resize(table_len).await?;
629 blob.sync().await?;
630 Ok(())
631 }
632
633 pub async fn init(
637 context: E,
638 config: Config<V::Cfg>,
639 checkpoint: Option<Checkpoint>,
640 ) -> Result<Self, Error> {
641 assert!(
643 config.table_initial_size > 0 && config.table_initial_size.is_power_of_two(),
644 "table_initial_size must be a power of 2"
645 );
646
647 let reset = checkpoint.is_none_or(|checkpoint| checkpoint.is_empty());
649 if reset {
650 for partition in [
651 &config.key_partition,
652 &config.value_partition,
653 &config.table_partition,
654 ] {
655 match context.remove(partition, None).await {
656 Ok(()) | Err(commonware_runtime::Error::PartitionMissing(_)) => {}
657 Err(err) => return Err(Error::Runtime(err)),
658 }
659 }
660 }
661
662 let oversized_cfg = OversizedConfig {
664 index_partition: config.key_partition.clone(),
665 value_partition: config.value_partition.clone(),
666 index_page_cache: config.key_page_cache.clone(),
667 index_write_buffer: config.key_write_buffer,
668 value_write_buffer: config.value_write_buffer,
669 compression: config.value_compression,
670 codec_config: config.codec_config,
671 };
672 let mut oversized: Oversized<E, Record<K>, V> =
673 Oversized::init(context.child("oversized"), oversized_cfg).await?;
674
675 let (table, table_len) = context
677 .open(&config.table_partition, TABLE_BLOB_NAME)
678 .await?;
679
680 let (checkpoint, resizable) = match checkpoint {
682 Some(checkpoint) if !checkpoint.is_empty() => {
684 if table_len == 0 {
686 return Err(Error::CheckpointMismatch);
687 }
688 assert!(
689 checkpoint.table_size > 0 && checkpoint.table_size.is_power_of_two(),
690 "table_size must be a power of 2"
691 );
692
693 oversized
695 .rewind(checkpoint.section, checkpoint.oversized_size)
696 .await?;
697
698 oversized.sync(checkpoint.section).await?;
700
701 let expected_table_len = Self::table_offset(checkpoint.table_size);
703 let mut modified = if table_len != expected_table_len {
704 table.resize(expected_table_len).await?;
705 true
706 } else {
707 false
708 };
709
710 let (table_modified, _, _, resizable) = Self::recover_table(
712 &context,
713 &table,
714 checkpoint.table_size,
715 config.table_resize_frequency,
716 Some(checkpoint.epoch),
717 config.table_replay_buffer,
718 )
719 .await?;
720 if table_modified {
721 modified = true;
722 }
723
724 if modified {
726 table.sync().await?;
727 }
728
729 (checkpoint, resizable)
730 }
731
732 _ => {
734 Self::init_table(&table, config.table_initial_size).await?;
735 (Checkpoint::init(config.table_initial_size), 0)
736 }
737 };
738
739 let puts = context.counter("puts", "number of put operations");
741 let gets = context.counter("gets", "number of get operations");
742 let has = context.counter("has", "number of has operations");
743 let unnecessary_reads = context.counter(
744 "unnecessary_reads",
745 "number of unnecessary reads performed during key lookups",
746 );
747 let unnecessary_writes = context.counter(
748 "unnecessary_writes",
749 "number of unnecessary writes performed during resize",
750 );
751 let resizes = context.counter("resizes", "number of table resizing operations");
752
753 Ok(Self {
754 context,
755 table_partition: config.table_partition,
756 table_size: checkpoint.table_size,
757 table_resize_threshold: checkpoint.table_size as u64 * RESIZE_THRESHOLD / 100,
758 table_resize_frequency: config.table_resize_frequency,
759 table_resize_chunk_size: config.table_resize_chunk_size,
760 table,
761 oversized,
762 blob_target_size: config.value_target_size,
763 current_section: checkpoint.section,
764 next_epoch: checkpoint.epoch.checked_add(1).expect("epoch overflow"),
765 modified_sections: BTreeSet::new(),
766 resizable,
767 resize_progress: None,
768 puts,
769 gets,
770 has,
771 unnecessary_reads,
772 unnecessary_writes,
773 resizes,
774 })
775 }
776
777 fn table_index(&self, key: &K) -> u32 {
789 let hash = Crc32::checksum(key.as_ref());
790 hash & (self.table_size - 1)
791 }
792
793 const fn should_resize(&self) -> bool {
795 self.resizable as u64 >= self.table_resize_threshold
796 }
797
798 async fn update_section(&mut self) -> Result<(), Error> {
800 let value_size = self.oversized.value_size(self.current_section).await?;
802
803 if value_size >= self.blob_target_size {
805 self.current_section += 1;
806 debug!(
807 size = value_size,
808 section = self.current_section,
809 "updated section"
810 );
811 }
812
813 Ok(())
814 }
815
816 pub async fn put(&mut self, key: K, value: V) -> Result<Cursor, Error> {
819 self.puts.inc();
820
821 self.update_section().await?;
823
824 let table_index = self.table_index(&key);
826 let (entry1, entry2) = Self::read_table(&self.table, table_index).await?;
827 let head = Self::read_latest_entry(&entry1, &entry2);
828
829 let key_entry = Record::new(
831 key,
832 head.map(|(section, position, _)| (section, position)),
833 0,
834 0,
835 );
836
837 let (position, value_offset, value_size) = self
839 .oversized
840 .append(self.current_section, key_entry, &value)
841 .await?;
842
843 let mut added = head.map(|(_, _, added)| added).unwrap_or(0);
847 added = added.saturating_add(1);
848
849 if added == self.table_resize_frequency {
851 self.resizable += 1;
852 }
853
854 self.modified_sections.insert(self.current_section);
856 let new_entry = Entry::new(self.next_epoch, self.current_section, position, added);
857 Self::update_head(
858 &self.context,
859 &self.table,
860 table_index,
861 &entry1,
862 &entry2,
863 new_entry,
864 )
865 .await?;
866
867 if let Some(resize_progress) = self.resize_progress {
869 if table_index < resize_progress {
870 self.unnecessary_writes.inc();
871
872 if added == self.table_resize_frequency {
874 self.resizable += 1;
875 }
876
877 let new_table_index = self.table_size + table_index;
881 let new_entry = Entry::new(self.next_epoch, self.current_section, position, added);
882 Self::update_head(
883 &self.context,
884 &self.table,
885 new_table_index,
886 &entry1,
887 &entry2,
888 new_entry,
889 )
890 .await?;
891 }
892 }
893
894 Ok(Cursor::new(self.current_section, value_offset, value_size))
895 }
896
897 async fn get_cursor(&self, cursor: Cursor) -> Result<V, Error> {
899 let value = self
900 .oversized
901 .get_value(cursor.section(), cursor.offset(), cursor.size())
902 .await?;
903
904 Ok(value)
905 }
906
907 async fn find_key(&self, key: &K) -> Result<Option<(u64, Record<K>)>, Error> {
911 let table_index = self.table_index(key);
913 let (entry1, entry2) = Self::read_table(&self.table, table_index).await?;
914 let Some((mut section, mut position, _)) = Self::read_latest_entry(&entry1, &entry2) else {
915 return Ok(None);
916 };
917
918 loop {
920 let key_entry = self.oversized.get(section, position).await?;
922
923 if key_entry.key.as_ref() == key.as_ref() {
925 return Ok(Some((section, key_entry)));
926 }
927
928 self.unnecessary_reads.inc();
930
931 let Some(next) = key_entry.next() else {
933 break; };
935 section = next.0;
936 position = next.1;
937 }
938
939 Ok(None)
940 }
941
942 async fn get_key(&self, key: &K) -> Result<Option<V>, Error> {
944 self.gets.inc();
945
946 let Some((section, key_entry)) = self.find_key(key).await? else {
947 return Ok(None);
948 };
949 let value = self
950 .oversized
951 .get_value(section, key_entry.value_offset, key_entry.value_size)
952 .await?;
953 Ok(Some(value))
954 }
955
956 pub async fn get<'a>(&'a self, identifier: Identifier<'a, K>) -> Result<Option<V>, Error> {
961 match identifier {
962 Identifier::Cursor(cursor) => self.get_cursor(cursor).await.map(Some),
963 Identifier::Key(key) => self.get_key(key).await,
964 }
965 }
966
967 pub async fn has(&self, key: &K) -> Result<bool, Error> {
972 self.has.inc();
973
974 Ok(self.find_key(key).await?.is_some())
975 }
976
977 async fn start_resize(&mut self) -> Result<(), Error> {
979 self.resizes.inc();
980
981 let old_size = self.table_size;
983 let Some(new_size) = old_size.checked_mul(2) else {
984 return Ok(());
985 };
986 self.table.resize(Self::table_offset(new_size)).await?;
987
988 self.resize_progress = Some(0);
990 debug!(old = old_size, new = new_size, "table resize started");
991
992 Ok(())
993 }
994
995 fn rewrite_entries(buf: &mut impl BufMut, entry1: &Entry, entry2: &Entry, new_entry: &Entry) {
997 if Self::compute_write_offset(entry1, entry2, new_entry.epoch) == 0 {
998 new_entry.write(buf);
999 entry2.write(buf);
1000 } else {
1001 entry1.write(buf);
1002 new_entry.write(buf);
1003 }
1004 }
1005
1006 async fn advance_resize(&mut self) -> Result<(), Error> {
1011 let current_index = self.resize_progress.unwrap();
1013 let old_size = self.table_size;
1014 let chunk_end = (current_index + self.table_resize_chunk_size).min(old_size);
1015 let chunk_size = chunk_end - current_index;
1016
1017 let chunk_bytes = chunk_size as usize * Entry::FULL_SIZE;
1019 let read_offset = Self::table_offset(current_index);
1020 let mut read_buf = self.table.read_at(read_offset, chunk_bytes).await?;
1021
1022 let mut writes = self.context.storage_buffer_pool().alloc(chunk_bytes);
1024 for _ in 0..chunk_size {
1025 let (entry1, entry2) = Self::parse_entries(&mut read_buf)?;
1027
1028 let head = Self::read_latest_entry(&entry1, &entry2);
1030
1031 let reset_entry = match head {
1033 Some((section, position, added)) => {
1034 if added >= self.table_resize_frequency {
1036 self.resizable -= 1;
1037 }
1038 Entry::new(self.next_epoch, section, position, 0)
1039 }
1040 None => Entry::new_empty(),
1041 };
1042
1043 Self::rewrite_entries(&mut writes, &entry1, &entry2, &reset_entry);
1045 }
1046
1047 let writes = writes.freeze();
1049 let old_write = self.table.write_at(read_offset, writes.clone());
1050 let new_offset = (old_size as usize * Entry::FULL_SIZE) as u64 + read_offset;
1051 let new_write = self.table.write_at(new_offset, writes);
1052 try_join(old_write, new_write).await?;
1053
1054 if chunk_end >= old_size {
1056 self.table_size = old_size * 2;
1058 self.table_resize_threshold = self.table_size as u64 * RESIZE_THRESHOLD / 100;
1059 self.resize_progress = None;
1060 debug!(
1061 old = old_size,
1062 new = self.table_size,
1063 "table resize completed"
1064 );
1065 } else {
1066 self.resize_progress = Some(chunk_end);
1068 debug!(current = current_index, chunk_end, "table resize progress");
1069 }
1070
1071 Ok(())
1072 }
1073
1074 pub async fn sync(&mut self) -> Result<Checkpoint, Error> {
1084 self.oversized.sync(&self.modified_sections).await?;
1086 self.modified_sections.clear();
1087
1088 if self.should_resize() && self.resize_progress.is_none() {
1090 self.start_resize().await?;
1091 }
1092
1093 if self.resize_progress.is_some() {
1095 self.advance_resize().await?;
1096 }
1097
1098 self.table.sync().await?;
1100 let stored_epoch = self.next_epoch;
1101 self.next_epoch = self.next_epoch.checked_add(1).expect("epoch overflow");
1102
1103 let oversized_size = self.oversized.size(self.current_section)?;
1105
1106 Ok(Checkpoint {
1107 epoch: stored_epoch,
1108 section: self.current_section,
1109 oversized_size,
1110 table_size: self.table_size,
1111 })
1112 }
1113
1114 pub async fn close(mut self) -> Result<Checkpoint, Error> {
1116 while self.resize_progress.is_some() {
1118 self.advance_resize().await?;
1119 }
1120
1121 let checkpoint = self.sync().await?;
1123
1124 Ok(checkpoint)
1125 }
1126
1127 pub async fn destroy(self) -> Result<(), Error> {
1129 self.oversized.destroy().await?;
1131
1132 drop(self.table);
1134 self.context
1135 .remove(&self.table_partition, Some(TABLE_BLOB_NAME))
1136 .await?;
1137 self.context.remove(&self.table_partition, None).await?;
1138
1139 Ok(())
1140 }
1141
1142 #[cfg(test)]
1146 pub const fn resizing(&self) -> Option<u32> {
1147 self.resize_progress
1148 }
1149
1150 #[cfg(test)]
1152 pub const fn resizable(&self) -> u32 {
1153 self.resizable
1154 }
1155}
1156
1157#[cfg(all(test, feature = "arbitrary"))]
1158mod conformance {
1159 use super::*;
1160 use commonware_codec::conformance::CodecConformance;
1161 use commonware_utils::sequence::U64;
1162
1163 commonware_conformance::conformance_tests! {
1164 CodecConformance<Cursor>,
1165 CodecConformance<Checkpoint>,
1166 CodecConformance<Entry>,
1167 CodecConformance<Record<U64>>
1168 }
1169}
1170
1171#[cfg(test)]
1172mod tests {
1173 use super::*;
1174 use commonware_codec::DecodeExt;
1175 use commonware_macros::test_traced;
1176 use commonware_runtime::{
1177 buffer::paged::CacheRef, deterministic, deterministic::Context, Runner, Storage,
1178 Supervisor as _,
1179 };
1180 use commonware_utils::{
1181 sequence::{FixedBytes, U64},
1182 NZUsize, NZU16,
1183 };
1184
1185 fn test_key(key: &str) -> FixedBytes<64> {
1186 let mut buf = [0u8; 64];
1187 let key = key.as_bytes();
1188 assert!(key.len() <= buf.len());
1189 buf[..key.len()].copy_from_slice(key);
1190 FixedBytes::decode(buf.as_ref()).unwrap()
1191 }
1192
1193 fn test_key_at_index(table_size: u32, table_index: u32) -> FixedBytes<64> {
1194 assert!(table_size.is_power_of_two());
1195 assert!(table_index < table_size);
1196
1197 for value in 0u64.. {
1198 let mut buf = [0u8; 64];
1199 let bytes = value.to_be_bytes();
1200 buf[..bytes.len()].copy_from_slice(&bytes);
1201 let key = FixedBytes::new(buf);
1202 if Crc32::checksum(key.as_ref()) & (table_size - 1) == table_index {
1203 return key;
1204 }
1205 }
1206
1207 unreachable!("u64 key space exhausted");
1208 }
1209
1210 type TestFreezer = Freezer<Context, U64, u64>;
1211
1212 fn is_send<T: Send>(_: T) {}
1213
1214 #[allow(dead_code)]
1215 fn assert_freezer_futures_are_send(freezer: &mut TestFreezer, key: U64) {
1216 is_send(freezer.get(Identifier::Key(&key)));
1217 is_send(freezer.put(key, 0u64));
1218 }
1219
1220 #[allow(dead_code)]
1221 fn assert_freezer_destroy_is_send(freezer: TestFreezer) {
1222 is_send(freezer.destroy());
1223 }
1224
1225 #[test_traced]
1226 fn issue_2966_regression() {
1227 let executor = deterministic::Runner::default();
1228 executor.start(|context| async move {
1229 let cfg = super::super::Config {
1230 key_partition: "test-key-index".into(),
1231 key_write_buffer: NZUsize!(1024),
1232 key_page_cache: CacheRef::from_pooler(&context, NZU16!(1024), NZUsize!(10)),
1233 value_partition: "test-value-journal".into(),
1234 value_compression: None,
1235 value_write_buffer: NZUsize!(1024),
1236 value_target_size: 10 * 1024 * 1024,
1237 table_partition: "test-table".into(),
1238 table_initial_size: 4,
1240 table_resize_frequency: 1,
1241 table_resize_chunk_size: 4,
1242 table_replay_buffer: NZUsize!(64 * 1024),
1243 codec_config: (),
1244 };
1245 let mut freezer =
1246 Freezer::<_, FixedBytes<64>, i32>::init(context.child("first"), cfg.clone(), None)
1247 .await
1248 .unwrap();
1249
1250 freezer.put(test_key("key0"), 0).await.unwrap();
1253 freezer.put(test_key("key2"), 1).await.unwrap();
1254 freezer.close().await.unwrap();
1255
1256 let (blob, size) = context.open(&cfg.table_partition, b"table").await.unwrap();
1257 let table_data = blob.read_at(0, size as usize).await.unwrap().coalesce();
1258
1259 let num_entries = size as usize / Entry::FULL_SIZE;
1261 assert_eq!(num_entries, 8);
1262
1263 let mut both_empty_count = 0;
1266 for entry_idx in 0..num_entries {
1267 let offset = entry_idx * Entry::FULL_SIZE;
1268 let buf = &table_data.as_ref()[offset..offset + Entry::FULL_SIZE];
1269 let (slot0, slot1) =
1270 Freezer::<Context, FixedBytes<64>, i32>::parse_entries(buf).unwrap();
1271 if slot0.is_empty() && slot1.is_empty() {
1272 both_empty_count += 1;
1273 }
1274 }
1275 assert_eq!(both_empty_count, 4);
1277 });
1278 }
1279
1280 #[test_traced]
1281 fn issue_2955_regression() {
1282 let executor = deterministic::Runner::default();
1283 executor.start(|context| async move {
1284 let cfg = super::super::Config {
1285 key_partition: "test-key-index".into(),
1286 key_write_buffer: NZUsize!(1024),
1287 key_page_cache: CacheRef::from_pooler(&context, NZU16!(1024), NZUsize!(10)),
1288 value_partition: "test-value-journal".into(),
1289 value_compression: None,
1290 value_write_buffer: NZUsize!(1024),
1291 value_target_size: 10 * 1024 * 1024,
1292 table_partition: "test-table".into(),
1293 table_initial_size: 4,
1294 table_resize_frequency: 1,
1295 table_resize_chunk_size: 4,
1296 table_replay_buffer: NZUsize!(64 * 1024),
1297 codec_config: (),
1298 };
1299
1300 let checkpoint = {
1302 let mut freezer = Freezer::<_, FixedBytes<64>, i32>::init(
1303 context.child("first"),
1304 cfg.clone(),
1305 None,
1306 )
1307 .await
1308 .unwrap();
1309 freezer.put(test_key("key0"), 42).await.unwrap();
1310 freezer.sync().await.unwrap();
1311 freezer.close().await.unwrap()
1312 };
1313
1314 {
1316 let (blob, _) = context.open(&cfg.table_partition, b"table").await.unwrap();
1317 let entry_data = blob.read_at(0, Entry::FULL_SIZE).await.unwrap();
1318 let mut corrupted = entry_data.coalesce();
1319 corrupted.as_mut()[Entry::SIZE - 4] ^= 0xFF;
1321 corrupted.as_mut()[Entry::FULL_SIZE - 4] ^= 0xFF;
1323 blob.write_at_sync(0, corrupted).await.unwrap();
1324 }
1325
1326 let freezer = Freezer::<_, FixedBytes<64>, i32>::init(
1331 context.child("second"),
1332 cfg.clone(),
1333 Some(checkpoint),
1334 )
1335 .await
1336 .unwrap();
1337 drop(freezer);
1338 });
1339 }
1340
1341 #[test_traced]
1342 fn no_checkpoint_deletes_partial_sync_resize() {
1343 let executor = deterministic::Runner::default();
1344 executor.start(|context| async move {
1345 let cfg = super::super::Config {
1346 key_partition: "test-key-index".into(),
1347 key_write_buffer: NZUsize!(1024),
1348 key_page_cache: CacheRef::from_pooler(&context, NZU16!(1024), NZUsize!(10)),
1349 value_partition: "test-value-journal".into(),
1350 value_compression: None,
1351 value_write_buffer: NZUsize!(1024),
1352 value_target_size: 10 * 1024 * 1024,
1353 table_partition: "test-table".into(),
1354 table_initial_size: 2,
1355 table_resize_frequency: 1,
1356 table_resize_chunk_size: 1,
1357 table_replay_buffer: NZUsize!(64 * 1024),
1358 codec_config: (),
1359 };
1360 let key = test_key_at_index(4, 3);
1361
1362 {
1363 let mut freezer = Freezer::<_, FixedBytes<64>, i32>::init(
1364 context.child("first"),
1365 cfg.clone(),
1366 None,
1367 )
1368 .await
1369 .unwrap();
1370 freezer.put(key.clone(), 42).await.unwrap();
1371 freezer.sync().await.unwrap();
1372
1373 assert_eq!(freezer.resizing(), Some(1));
1374 assert_eq!(freezer.get(Identifier::Key(&key)).await.unwrap(), Some(42));
1375 }
1376
1377 let freezer =
1378 Freezer::<_, FixedBytes<64>, i32>::init(context.child("second"), cfg.clone(), None)
1379 .await
1380 .unwrap();
1381 assert_eq!(freezer.table_size, 2);
1382 assert_eq!(freezer.resizing(), None);
1383 assert_eq!(freezer.get(Identifier::Key(&key)).await.unwrap(), None);
1384 });
1385 }
1386
1387 #[test_traced]
1388 fn empty_checkpoint_deletes_existing_data() {
1389 let executor = deterministic::Runner::default();
1390 executor.start(|context| async move {
1391 let cfg = super::super::Config {
1392 key_partition: "test-key-index".into(),
1393 key_write_buffer: NZUsize!(1024),
1394 key_page_cache: CacheRef::from_pooler(&context, NZU16!(1024), NZUsize!(10)),
1395 value_partition: "test-value-journal".into(),
1396 value_compression: None,
1397 value_write_buffer: NZUsize!(1024),
1398 value_target_size: 10 * 1024 * 1024,
1399 table_partition: "test-table".into(),
1400 table_initial_size: 2,
1401 table_resize_frequency: 1,
1402 table_resize_chunk_size: 1,
1403 table_replay_buffer: NZUsize!(64 * 1024),
1404 codec_config: (),
1405 };
1406 let key = test_key_at_index(4, 3);
1407
1408 {
1409 let mut freezer = Freezer::<_, FixedBytes<64>, i32>::init(
1410 context.child("first"),
1411 cfg.clone(),
1412 None,
1413 )
1414 .await
1415 .unwrap();
1416 freezer.put(key.clone(), 42).await.unwrap();
1417 freezer.sync().await.unwrap();
1418 assert_eq!(freezer.get(Identifier::Key(&key)).await.unwrap(), Some(42));
1419 }
1420
1421 let checkpoint = Checkpoint {
1422 epoch: 0,
1423 section: 0,
1424 oversized_size: 0,
1425 table_size: 0,
1426 };
1427 let freezer = Freezer::<_, FixedBytes<64>, i32>::init(
1428 context.child("second"),
1429 cfg.clone(),
1430 Some(checkpoint),
1431 )
1432 .await
1433 .unwrap();
1434 assert_eq!(freezer.table_size, 2);
1435 assert_eq!(freezer.get(Identifier::Key(&key)).await.unwrap(), None);
1436 });
1437 }
1438
1439 #[test_traced]
1440 fn no_checkpoint_deletes_close_started_partial_resize() {
1441 let executor = deterministic::Runner::default();
1442 executor.start(|context| async move {
1443 let cfg = super::super::Config {
1444 key_partition: "test-key-index".into(),
1445 key_write_buffer: NZUsize!(1024),
1446 key_page_cache: CacheRef::from_pooler(&context, NZU16!(1024), NZUsize!(10)),
1447 value_partition: "test-value-journal".into(),
1448 value_compression: None,
1449 value_write_buffer: NZUsize!(1024),
1450 value_target_size: 10 * 1024 * 1024,
1451 table_partition: "test-table".into(),
1452 table_initial_size: 2,
1453 table_resize_frequency: 1,
1454 table_resize_chunk_size: 1,
1455 table_replay_buffer: NZUsize!(64 * 1024),
1456 codec_config: (),
1457 };
1458 let key = test_key_at_index(4, 3);
1459
1460 {
1461 let mut freezer = Freezer::<_, FixedBytes<64>, i32>::init(
1462 context.child("first"),
1463 cfg.clone(),
1464 None,
1465 )
1466 .await
1467 .unwrap();
1468 freezer.put(key.clone(), 42).await.unwrap();
1469 let checkpoint = freezer.close().await.unwrap();
1470 assert_eq!(checkpoint.table_size, 2);
1471 }
1472
1473 let freezer =
1474 Freezer::<_, FixedBytes<64>, i32>::init(context.child("second"), cfg.clone(), None)
1475 .await
1476 .unwrap();
1477 assert_eq!(freezer.table_size, 2);
1478 assert_eq!(freezer.resizing(), None);
1479 assert_eq!(freezer.get(Identifier::Key(&key)).await.unwrap(), None);
1480 });
1481 }
1482
1483 #[test_traced]
1484 fn no_checkpoint_deletes_completed_resize() {
1485 let executor = deterministic::Runner::default();
1486 executor.start(|context| async move {
1487 let cfg = super::super::Config {
1488 key_partition: "test-key-index".into(),
1489 key_write_buffer: NZUsize!(1024),
1490 key_page_cache: CacheRef::from_pooler(&context, NZU16!(1024), NZUsize!(10)),
1491 value_partition: "test-value-journal".into(),
1492 value_compression: None,
1493 value_write_buffer: NZUsize!(1024),
1494 value_target_size: 10 * 1024 * 1024,
1495 table_partition: "test-table".into(),
1496 table_initial_size: 2,
1497 table_resize_frequency: 1,
1498 table_resize_chunk_size: 2,
1499 table_replay_buffer: NZUsize!(64 * 1024),
1500 codec_config: (),
1501 };
1502 let key = test_key_at_index(4, 3);
1503
1504 {
1505 let mut freezer = Freezer::<_, FixedBytes<64>, i32>::init(
1506 context.child("first"),
1507 cfg.clone(),
1508 None,
1509 )
1510 .await
1511 .unwrap();
1512 freezer.put(key.clone(), 42).await.unwrap();
1513 let checkpoint = freezer.sync().await.unwrap();
1514
1515 assert_eq!(checkpoint.table_size, 4);
1516 assert_eq!(freezer.resizing(), None);
1517 assert_eq!(freezer.get(Identifier::Key(&key)).await.unwrap(), Some(42));
1518 }
1519
1520 let freezer =
1521 Freezer::<_, FixedBytes<64>, i32>::init(context.child("second"), cfg.clone(), None)
1522 .await
1523 .unwrap();
1524 assert_eq!(freezer.table_size, 2);
1525 assert_eq!(freezer.get(Identifier::Key(&key)).await.unwrap(), None);
1526 });
1527 }
1528
1529 #[test_traced]
1530 fn checkpoint_rewinds_completed_resize() {
1531 let executor = deterministic::Runner::default();
1532 executor.start(|context| async move {
1533 let cfg = super::super::Config {
1534 key_partition: "test-key-index".into(),
1535 key_write_buffer: NZUsize!(1024),
1536 key_page_cache: CacheRef::from_pooler(&context, NZU16!(1024), NZUsize!(10)),
1537 value_partition: "test-value-journal".into(),
1538 value_compression: None,
1539 value_write_buffer: NZUsize!(1024),
1540 value_target_size: 10 * 1024 * 1024,
1541 table_partition: "test-table".into(),
1542 table_initial_size: 2,
1543 table_resize_frequency: 1,
1544 table_resize_chunk_size: 2,
1545 table_replay_buffer: NZUsize!(64 * 1024),
1546 codec_config: (),
1547 };
1548 let key = test_key_at_index(4, 3);
1549
1550 let stale_checkpoint = {
1551 let mut freezer = Freezer::<_, FixedBytes<64>, i32>::init(
1552 context.child("first"),
1553 cfg.clone(),
1554 None,
1555 )
1556 .await
1557 .unwrap();
1558 let stale_checkpoint = freezer.sync().await.unwrap();
1559 assert_eq!(stale_checkpoint.table_size, 2);
1560
1561 freezer.put(key.clone(), 42).await.unwrap();
1562 let checkpoint = freezer.sync().await.unwrap();
1563 assert_eq!(checkpoint.table_size, 4);
1564 assert_eq!(freezer.resizing(), None);
1565 assert_eq!(freezer.get(Identifier::Key(&key)).await.unwrap(), Some(42));
1566
1567 stale_checkpoint
1568 };
1569
1570 let freezer = Freezer::<_, FixedBytes<64>, i32>::init(
1571 context.child("second"),
1572 cfg.clone(),
1573 Some(stale_checkpoint),
1574 )
1575 .await
1576 .unwrap();
1577 assert_eq!(freezer.table_size, 2);
1578 assert_eq!(freezer.get(Identifier::Key(&key)).await.unwrap(), None);
1579 });
1580 }
1581
1582 #[test_traced]
1583 fn non_empty_checkpoint_against_empty_table_errors() {
1584 let executor = deterministic::Runner::default();
1585 executor.start(|context| async move {
1586 let cfg = super::super::Config {
1587 key_partition: "test-key-index".into(),
1588 key_write_buffer: NZUsize!(1024),
1589 key_page_cache: CacheRef::from_pooler(&context, NZU16!(1024), NZUsize!(10)),
1590 value_partition: "test-value-journal".into(),
1591 value_compression: None,
1592 value_write_buffer: NZUsize!(1024),
1593 value_target_size: 10 * 1024 * 1024,
1594 table_partition: "test-table".into(),
1595 table_initial_size: 2,
1596 table_resize_frequency: 1,
1597 table_resize_chunk_size: 1,
1598 table_replay_buffer: NZUsize!(64 * 1024),
1599 codec_config: (),
1600 };
1601
1602 let checkpoint = Checkpoint {
1603 epoch: 1,
1604 section: 0,
1605 oversized_size: 0,
1606 table_size: 2,
1607 };
1608 let result = Freezer::<_, FixedBytes<64>, i32>::init(
1609 context.child("storage"),
1610 cfg.clone(),
1611 Some(checkpoint),
1612 )
1613 .await;
1614 assert!(matches!(result, Err(Error::CheckpointMismatch)));
1615 });
1616 }
1617}