1use crate::{
7 db::{
8 data::{CanonicalRow, RawDataStoreKey, RawRow},
9 direction::Direction,
10 ordered_overlay::{OrderedOverlayEntry, ordered_overlay_entries},
11 positioned_overlay::{
12 JournalOverlayPosition, PositionedOverlayMetadata, PositionedOverlayRetirement,
13 },
14 },
15 types::EntityTag,
16};
17use ic_memory::{DefaultMemoryImpl, RuntimeMemory};
18use ic_storage::adapter::StableBTreeMap;
19#[cfg(test)]
20use std::cell::Cell;
21use std::collections::{BTreeMap as HeapBTreeMap, BTreeSet};
22use std::convert::Infallible;
23use std::ops::{Bound, RangeBounds};
24
25#[cfg(test)]
26thread_local! {
27 static DATA_STORE_GET_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
28}
29
30#[cfg(test)]
31fn record_data_store_get_call() {
32 DATA_STORE_GET_CALL_COUNT.with(|count| {
33 count.set(count.get().saturating_add(1));
34 });
35}
36
37pub struct DataStore {
47 backend: DataStoreBackend,
48 generation: u64,
49 entity_cardinality: EntityCardinality,
50}
51
52enum DataStoreBackend {
53 Heap(HeapBTreeMap<RawDataStoreKey, RawRow>),
54 Journaled {
55 canonical: StableBTreeMap<RawDataStoreKey, RawRow, DefaultMemoryImpl>,
56 live: HeapBTreeMap<RawDataStoreKey, RawRow>,
57 tombstones: BTreeSet<RawDataStoreKey>,
58 positions: PositionedOverlayMetadata<RawDataStoreKey>,
59 entity_cardinality_delta: EntityCardinalityDelta,
60 },
61}
62
63#[cfg(feature = "migration")]
65pub(in crate::db) struct PreparedDataPositionPublication {
66 keys: Vec<RawDataStoreKey>,
67 position: JournalOverlayPosition,
68}
69
70pub(in crate::db) struct PreparedDataPositionRetirement {
72 entries: Vec<(RawDataStoreKey, PositionedOverlayRetirement)>,
73}
74
75pub(in crate::db) enum StoredRowRead<'a> {
81 Missing,
82 Borrowed(&'a RawRow),
83 Owned(RawRow),
84}
85
86impl StoredRowRead<'_> {
87 #[must_use]
89 pub(in crate::db) const fn as_row(&self) -> Option<&RawRow> {
90 match self {
91 Self::Missing => None,
92 Self::Borrowed(row) => Some(row),
93 Self::Owned(row) => Some(row),
94 }
95 }
96
97 #[must_use]
99 fn into_owned(self) -> Option<RawRow> {
100 match self {
101 Self::Missing => None,
102 Self::Borrowed(row) => Some(row.clone()),
103 Self::Owned(row) => Some(row),
104 }
105 }
106}
107
108#[derive(Clone, Copy, Debug, Eq, PartialEq)]
110pub(in crate::db) enum StoreVisit {
111 Continue,
112 Stop,
113}
114
115impl StoreVisit {
116 const fn should_stop(self) -> bool {
117 matches!(self, Self::Stop)
118 }
119}
120
121impl DataStore {
122 #[must_use]
124 pub const fn init_heap() -> Self {
125 Self {
126 backend: DataStoreBackend::Heap(HeapBTreeMap::new()),
127 generation: 0,
128 entity_cardinality: EntityCardinality::empty(),
129 }
130 }
131
132 #[must_use]
138 pub fn init_journaled(memory: RuntimeMemory<DefaultMemoryImpl>) -> Self {
139 let canonical = StableBTreeMap::init(memory);
140 let entity_cardinality = if canonical.is_empty() {
141 EntityCardinality::empty()
142 } else {
143 EntityCardinality::unavailable()
144 };
145 Self {
146 backend: DataStoreBackend::Journaled {
147 canonical,
148 live: HeapBTreeMap::new(),
149 tombstones: BTreeSet::new(),
150 positions: PositionedOverlayMetadata::new(),
151 entity_cardinality_delta: EntityCardinalityDelta::empty(),
152 },
153 generation: 0,
154 entity_cardinality,
158 }
159 }
160
161 pub(in crate::db) fn insert(
163 &mut self,
164 key: RawDataStoreKey,
165 row: CanonicalRow,
166 ) -> Option<RawRow> {
167 let row = row.into_raw_row();
168 let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
169 self.get(&key)
170 } else {
171 None
172 };
173 let cardinality_key = key.clone();
174 let previous = match &mut self.backend {
175 DataStoreBackend::Heap(map) => map.insert(key, row),
176 DataStoreBackend::Journaled {
177 live, tombstones, ..
178 } => {
179 tombstones.remove(&key);
180 live.insert(key, row);
181 previous_journaled
182 }
183 };
184 self.entity_cardinality
185 .apply_insert(&cardinality_key, previous.is_some());
186 self.apply_entity_overlay_delta(&cardinality_key, previous.is_some(), true);
187 self.bump_generation();
188 previous
189 }
190
191 #[cfg(test)]
193 pub(in crate::db) fn insert_raw_for_test(
194 &mut self,
195 key: RawDataStoreKey,
196 row: RawRow,
197 ) -> Option<RawRow> {
198 let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
199 self.get(&key)
200 } else {
201 None
202 };
203 let cardinality_key = key.clone();
204 let previous = match &mut self.backend {
205 DataStoreBackend::Heap(map) => map.insert(key, row),
206 DataStoreBackend::Journaled {
207 live, tombstones, ..
208 } => {
209 tombstones.remove(&key);
210 live.insert(key, row);
211 previous_journaled
212 }
213 };
214 self.entity_cardinality
215 .apply_insert(&cardinality_key, previous.is_some());
216 self.apply_entity_overlay_delta(&cardinality_key, previous.is_some(), true);
217 self.bump_generation();
218 previous
219 }
220
221 pub(in crate::db) fn remove(&mut self, key: &RawDataStoreKey) -> Option<RawRow> {
223 let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
224 self.get(key)
225 } else {
226 None
227 };
228 let previous = match &mut self.backend {
229 DataStoreBackend::Heap(map) => map.remove(key),
230 DataStoreBackend::Journaled {
231 live, tombstones, ..
232 } => {
233 live.remove(key);
234 tombstones.insert(key.clone());
235 previous_journaled
236 }
237 };
238 self.entity_cardinality
239 .apply_remove(key, previous.is_some());
240 self.apply_entity_overlay_delta(key, previous.is_some(), false);
241 self.bump_generation();
242 previous
243 }
244
245 pub(in crate::db) fn reset_journaled_live_projection(
248 &mut self,
249 ) -> Result<(), crate::error::InternalError> {
250 let DataStoreBackend::Journaled {
251 canonical,
252 live,
253 tombstones,
254 positions,
255 entity_cardinality_delta,
256 } = &mut self.backend
257 else {
258 return Err(crate::error::InternalError::store_invariant());
259 };
260
261 live.clear();
262 tombstones.clear();
263 positions.clear();
264 *entity_cardinality_delta = EntityCardinalityDelta::empty();
265 self.entity_cardinality = if canonical.is_empty() {
266 EntityCardinality::empty()
267 } else {
268 EntityCardinality::unavailable()
269 };
270 self.bump_generation();
271
272 Ok(())
273 }
274
275 #[cfg(any(test, feature = "migration"))]
277 pub(in crate::db) fn apply_recovered_journal_put(
278 &mut self,
279 key: RawDataStoreKey,
280 row: RawRow,
281 ) -> Result<Option<RawRow>, crate::error::InternalError> {
282 let DataStoreBackend::Journaled {
283 canonical,
284 live,
285 tombstones,
286 ..
287 } = &mut self.backend
288 else {
289 return Err(crate::error::InternalError::store_invariant());
290 };
291
292 let previous = if tombstones.contains(&key) {
293 None
294 } else {
295 live.get(&key).cloned().or_else(|| canonical.get(&key))
296 };
297 tombstones.remove(&key);
298 let cardinality_key = key.clone();
299 live.insert(key, row);
300 self.entity_cardinality
301 .apply_insert(&cardinality_key, previous.is_some());
302 self.apply_entity_overlay_delta(&cardinality_key, previous.is_some(), true);
303 self.bump_generation();
304
305 Ok(previous)
306 }
307
308 pub(in crate::db) fn publish_preflighted_journal_entry(
310 &mut self,
311 key: RawDataStoreKey,
312 row: Option<RawRow>,
313 position: JournalOverlayPosition,
314 ) -> Result<(), crate::error::InternalError> {
315 let previous_present = self.contains(&key);
317 let DataStoreBackend::Journaled {
318 live,
319 tombstones,
320 positions,
321 entity_cardinality_delta,
322 ..
323 } = &mut self.backend
324 else {
325 return Err(crate::error::InternalError::store_invariant());
326 };
327 let next_present = row.is_some();
328 if let Some(row) = row {
329 tombstones.remove(&key);
330 live.insert(key.clone(), row);
331 self.entity_cardinality.apply_insert(&key, previous_present);
332 } else {
333 live.remove(&key);
334 tombstones.insert(key.clone());
335 self.entity_cardinality.apply_remove(&key, previous_present);
336 }
337 entity_cardinality_delta.apply_presence_transition(&key, previous_present, next_present);
338 positions.publish_preflighted(key, position);
339 self.bump_generation();
340
341 Ok(())
342 }
343
344 #[cfg(test)]
346 pub(in crate::db) fn publish_positioned_journal_entry(
347 &mut self,
348 key: RawDataStoreKey,
349 row: Option<RawRow>,
350 position: JournalOverlayPosition,
351 ) -> Result<(), crate::error::InternalError> {
352 self.preflight_positioned_journal_entry(&key, position)?;
353 self.publish_preflighted_journal_entry(key, row, position)
354 }
355
356 pub(in crate::db) fn preflight_positioned_journal_entry(
358 &self,
359 key: &RawDataStoreKey,
360 position: JournalOverlayPosition,
361 ) -> Result<(), crate::error::InternalError> {
362 let DataStoreBackend::Journaled { positions, .. } = &self.backend else {
363 return Err(crate::error::InternalError::store_invariant());
364 };
365 positions.preflight_publish(key, position)
366 }
367
368 #[cfg(feature = "migration")]
370 pub(in crate::db) fn prepare_position_publication(
371 &self,
372 keys: impl IntoIterator<Item = RawDataStoreKey>,
373 position: JournalOverlayPosition,
374 ) -> Result<PreparedDataPositionPublication, crate::error::InternalError> {
375 let DataStoreBackend::Journaled { positions, .. } = &self.backend else {
376 return Err(crate::error::InternalError::store_invariant());
377 };
378 let keys = keys.into_iter().collect::<BTreeSet<_>>();
379 for key in &keys {
380 positions.preflight_publish(key, position)?;
381 }
382 Ok(PreparedDataPositionPublication {
383 keys: keys.into_iter().collect(),
384 position,
385 })
386 }
387
388 #[cfg(feature = "migration")]
390 pub(in crate::db) fn publish_prepared_positions(
391 &mut self,
392 prepared: PreparedDataPositionPublication,
393 ) {
394 let DataStoreBackend::Journaled { positions, .. } = &mut self.backend else {
395 debug_assert!(false, "preflighted row positions require a journaled store");
396 return;
397 };
398 for key in prepared.keys {
399 positions.publish_preflighted(key, prepared.position);
400 }
401 }
402
403 pub(in crate::db) fn prepare_position_retirement(
405 &self,
406 keys: impl IntoIterator<Item = RawDataStoreKey>,
407 position: JournalOverlayPosition,
408 ) -> Result<PreparedDataPositionRetirement, crate::error::InternalError> {
409 let DataStoreBackend::Journaled { positions, .. } = &self.backend else {
410 return Err(crate::error::InternalError::store_invariant());
411 };
412 let entries = keys
413 .into_iter()
414 .collect::<BTreeSet<_>>()
415 .into_iter()
416 .map(|key| {
417 positions
418 .preflight_retirement(&key, position)
419 .map(|retirement| (key, retirement))
420 })
421 .collect::<Result<Vec<_>, _>>()?;
422 Ok(PreparedDataPositionRetirement { entries })
423 }
424
425 pub(in crate::db) fn apply_prepared_position_retirement(
427 &mut self,
428 prepared: PreparedDataPositionRetirement,
429 ) {
430 let DataStoreBackend::Journaled {
431 live,
432 tombstones,
433 positions,
434 ..
435 } = &mut self.backend
436 else {
437 debug_assert!(
438 false,
439 "preflighted row retirement requires a journaled store"
440 );
441 return;
442 };
443 for (key, retirement) in prepared.entries {
444 if retirement == PositionedOverlayRetirement::Exact {
445 live.remove(&key);
446 tombstones.remove(&key);
447 positions.retire_preflighted(&key, retirement);
448 }
449 }
450 }
451
452 pub(in crate::db) fn fold_recovered_journal_put(
454 &mut self,
455 key: RawDataStoreKey,
456 row: RawRow,
457 ) -> Result<Option<RawRow>, crate::error::InternalError> {
458 let DataStoreBackend::Journaled {
459 canonical,
460 live,
461 tombstones,
462 ..
463 } = &mut self.backend
464 else {
465 return Err(crate::error::InternalError::store_invariant());
466 };
467
468 let visible = !live.contains_key(&key) && !tombstones.contains(&key);
469 let cardinality_key = key.clone();
470 let previous = canonical.insert(key, row);
471 if visible {
472 self.entity_cardinality
473 .apply_insert(&cardinality_key, previous.is_some());
474 } else if previous.is_none() {
475 self.apply_entity_overlay_delta(&cardinality_key, true, false);
476 }
477 self.bump_generation();
478
479 Ok(previous)
480 }
481
482 pub(in crate::db) fn fold_recovered_journal_delete(
484 &mut self,
485 key: &RawDataStoreKey,
486 ) -> Result<Option<RawRow>, crate::error::InternalError> {
487 let DataStoreBackend::Journaled {
488 canonical,
489 live,
490 tombstones,
491 ..
492 } = &mut self.backend
493 else {
494 return Err(crate::error::InternalError::store_invariant());
495 };
496
497 let visible = !live.contains_key(key) && !tombstones.contains(key);
498 let previous = canonical.remove(key);
499 if visible {
500 self.entity_cardinality
501 .apply_remove(key, previous.is_some());
502 } else if previous.is_some() {
503 self.apply_entity_overlay_delta(key, false, true);
504 }
505 self.bump_generation();
506
507 Ok(previous)
508 }
509
510 pub(in crate::db) fn preflight_fold_recovered_journal(
512 &self,
513 ) -> Result<(), crate::error::InternalError> {
514 match self.backend {
515 DataStoreBackend::Journaled { .. } => Ok(()),
516 DataStoreBackend::Heap(_) => Err(crate::error::InternalError::store_invariant()),
517 }
518 }
519
520 pub(in crate::db) fn get(&self, key: &RawDataStoreKey) -> Option<RawRow> {
522 self.read(key).into_owned()
523 }
524
525 pub(in crate::db) fn get_canonical(&self, key: &RawDataStoreKey) -> Option<RawRow> {
531 match &self.backend {
532 DataStoreBackend::Heap(map) => map.get(key).cloned(),
533 DataStoreBackend::Journaled { canonical, .. } => canonical.get(key),
534 }
535 }
536
537 pub(in crate::db) fn read<'a>(&'a self, key: &RawDataStoreKey) -> StoredRowRead<'a> {
539 #[cfg(test)]
540 record_data_store_get_call();
541
542 match &self.backend {
543 DataStoreBackend::Heap(map) => map
544 .get(key)
545 .map_or(StoredRowRead::Missing, StoredRowRead::Borrowed),
546 DataStoreBackend::Journaled {
547 canonical,
548 live,
549 tombstones,
550 ..
551 } => {
552 if tombstones.contains(key) {
553 StoredRowRead::Missing
554 } else if let Some(row) = live.get(key) {
555 StoredRowRead::Borrowed(row)
556 } else {
557 canonical
558 .get(key)
559 .map_or(StoredRowRead::Missing, StoredRowRead::Owned)
560 }
561 }
562 }
563 }
564
565 #[must_use]
567 pub(in crate::db) fn contains(&self, key: &RawDataStoreKey) -> bool {
568 match &self.backend {
569 DataStoreBackend::Heap(map) => map.contains_key(key),
570 DataStoreBackend::Journaled {
571 canonical,
572 live,
573 tombstones,
574 ..
575 } => {
576 !tombstones.contains(key) && (live.contains_key(key) || canonical.contains_key(key))
577 }
578 }
579 }
580
581 #[must_use]
586 pub(in crate::db) fn is_empty(&self) -> bool {
587 let mut empty = true;
588 let _: Result<(), Infallible> = self.visit_key_range(.., |_key| {
589 empty = false;
590 Ok(StoreVisit::Stop)
591 });
592 empty
593 }
594
595 #[cfg(test)]
597 #[must_use]
598 pub(in crate::db) fn len(&self) -> u64 {
599 match &self.backend {
600 DataStoreBackend::Heap(map) => u64::try_from(map.len()).unwrap_or(u64::MAX),
601 DataStoreBackend::Journaled { .. } => {
602 let mut count = 0_u64;
603 let _: Result<(), Infallible> = self.visit_entries(|_key, _row| {
604 count = count.saturating_add(1);
605 Ok(StoreVisit::Continue)
606 });
607 count
608 }
609 }
610 }
611
612 #[must_use]
614 pub(in crate::db) const fn generation(&self) -> u64 {
615 self.generation
616 }
617
618 #[must_use]
620 pub(in crate::db) fn exact_entity_count(&self, entity: EntityTag) -> Option<u64> {
621 self.entity_cardinality.exact_count(entity)
622 }
623
624 #[must_use]
626 pub(in crate::db) fn exact_entity_cardinality_delta(&self, entity: EntityTag) -> Option<i64> {
627 match &self.backend {
628 DataStoreBackend::Heap(_) => Some(0),
629 DataStoreBackend::Journaled {
630 entity_cardinality_delta,
631 ..
632 } => entity_cardinality_delta.exact_delta(entity),
633 }
634 }
635
636 pub(in crate::db) fn visit_entries<E>(
638 &self,
639 mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
640 ) -> Result<(), E> {
641 match &self.backend {
642 DataStoreBackend::Heap(map) => {
643 for (key, row) in map {
644 if visitor(key, row)?.should_stop() {
645 break;
646 }
647 }
648 }
649 DataStoreBackend::Journaled { .. } => Self::visit_journaled_entries_in_bounds(
650 &self.backend,
651 (Bound::Unbounded, Bound::Unbounded),
652 &mut visitor,
653 )?,
654 }
655
656 Ok(())
657 }
658
659 pub(in crate::db) fn visit_canonical_entries_after(
664 &self,
665 checkpoint: Option<&RawDataStoreKey>,
666 mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<bool, crate::error::InternalError>,
667 ) -> Result<(), crate::error::InternalError> {
668 let lower = checkpoint.map_or(Bound::Unbounded, Bound::Excluded);
669 match &self.backend {
670 DataStoreBackend::Heap(map) => {
671 for (key, row) in map.range((lower, Bound::Unbounded)) {
672 if visitor(key, row)? {
673 break;
674 }
675 }
676 }
677 DataStoreBackend::Journaled { canonical, .. } => {
678 for entry in canonical.range((lower, Bound::Unbounded)) {
679 if visitor(entry.key(), &entry.value())? {
680 break;
681 }
682 }
683 }
684 }
685 Ok(())
686 }
687
688 pub(in crate::db) fn canonical_is_empty(&self) -> Result<bool, crate::error::InternalError> {
693 match &self.backend {
694 DataStoreBackend::Journaled { canonical, .. } => Ok(canonical.is_empty()),
695 DataStoreBackend::Heap(_) => Err(crate::error::InternalError::store_invariant()),
696 }
697 }
698
699 pub(in crate::db) fn visit_range<E>(
701 &self,
702 key_range: impl RangeBounds<RawDataStoreKey>,
703 mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
704 ) -> Result<(), E> {
705 let bounds = (key_range.start_bound(), key_range.end_bound());
706 match &self.backend {
707 DataStoreBackend::Heap(map) => {
708 for (key, row) in map.range(bounds) {
709 if visitor(key, row)?.should_stop() {
710 break;
711 }
712 }
713 }
714 DataStoreBackend::Journaled { .. } => {
715 Self::visit_journaled_entries_in_bounds(&self.backend, bounds, &mut visitor)?;
716 }
717 }
718
719 Ok(())
720 }
721
722 pub(in crate::db) fn try_visit_range_with_row_preflight<E>(
729 &self,
730 key_range: impl RangeBounds<RawDataStoreKey>,
731 mut preflight: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
732 mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
733 ) -> Result<Option<bool>, E> {
734 let bounds = (key_range.start_bound(), key_range.end_bound());
735 let mut stopped = false;
736 match &self.backend {
737 DataStoreBackend::Heap(map) => {
738 for (key, row) in map.range(bounds) {
739 if preflight(key)?.should_stop() {
740 stopped = true;
741 break;
742 }
743 if visitor(key, row)?.should_stop() {
744 stopped = true;
745 break;
746 }
747 }
748 }
749 DataStoreBackend::Journaled {
750 canonical,
751 live,
752 tombstones,
753 ..
754 } if canonical.is_empty() => {
755 for (key, row) in live.range(bounds) {
756 if tombstones.contains(key) {
757 continue;
758 }
759 if preflight(key)?.should_stop() {
760 stopped = true;
761 break;
762 }
763 if visitor(key, row)?.should_stop() {
764 stopped = true;
765 break;
766 }
767 }
768 }
769 DataStoreBackend::Journaled {
770 canonical,
771 live,
772 tombstones,
773 ..
774 } if live.is_empty() && tombstones.is_empty() => {
775 for entry in canonical.range(bounds) {
776 if preflight(entry.key())?.should_stop() {
777 stopped = true;
778 break;
779 }
780 if visitor(entry.key(), &entry.value())?.should_stop() {
781 stopped = true;
782 break;
783 }
784 }
785 }
786 DataStoreBackend::Journaled { .. } => return Ok(None),
787 }
788
789 Ok(Some(!stopped))
790 }
791
792 pub(in crate::db) fn visit_key_range<E>(
799 &self,
800 key_range: impl RangeBounds<RawDataStoreKey>,
801 visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
802 ) -> Result<(), E> {
803 self.visit_keys_in_bounds(
804 (key_range.start_bound(), key_range.end_bound()),
805 false,
806 visitor,
807 )
808 }
809
810 pub(in crate::db) fn visit_key_range_rev<E>(
812 &self,
813 key_range: impl RangeBounds<RawDataStoreKey>,
814 visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
815 ) -> Result<(), E> {
816 self.visit_keys_in_bounds(
817 (key_range.start_bound(), key_range.end_bound()),
818 true,
819 visitor,
820 )
821 }
822
823 #[cfg(test)]
825 pub(in crate::db) fn memory_bytes(&self) -> u64 {
826 let mut bytes = 0u64;
828 let _: Result<(), Infallible> = self.visit_entries(|key, row| {
829 bytes = bytes.saturating_add(key.as_bytes().len() as u64 + row.len() as u64);
830 Ok(StoreVisit::Continue)
831 });
832 bytes
833 }
834
835 const fn bump_generation(&mut self) {
836 self.generation = self.generation.saturating_add(1);
837 }
838
839 #[cfg(test)]
840 fn rebuild_entity_cardinality_from_entries(&mut self) {
841 let mut cardinality = EntityCardinality::empty();
842 let _: Result<(), Infallible> = self.visit_entries(|key, _row| {
843 cardinality.apply_present_key(key);
844 Ok(StoreVisit::Continue)
845 });
846 self.entity_cardinality = cardinality;
847 }
848
849 fn apply_entity_overlay_delta(
850 &mut self,
851 key: &RawDataStoreKey,
852 previous_present: bool,
853 next_present: bool,
854 ) {
855 let DataStoreBackend::Journaled {
856 entity_cardinality_delta,
857 ..
858 } = &mut self.backend
859 else {
860 return;
861 };
862 entity_cardinality_delta.apply_presence_transition(key, previous_present, next_present);
863 }
864
865 #[cfg(test)]
867 pub(in crate::db) fn current_get_call_count() -> u64 {
868 DATA_STORE_GET_CALL_COUNT.with(Cell::get)
869 }
870
871 fn visit_keys_in_bounds<E>(
872 &self,
873 bounds: (Bound<&RawDataStoreKey>, Bound<&RawDataStoreKey>),
874 reverse: bool,
875 mut visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
876 ) -> Result<(), E> {
877 match &self.backend {
878 DataStoreBackend::Heap(map) => {
879 if reverse {
880 for (key, _row) in map.range(bounds).rev() {
881 if visitor(key)?.should_stop() {
882 break;
883 }
884 }
885 } else {
886 for (key, _row) in map.range(bounds) {
887 if visitor(key)?.should_stop() {
888 break;
889 }
890 }
891 }
892 }
893 DataStoreBackend::Journaled { .. } => {
894 Self::visit_journaled_keys_in_bounds(&self.backend, bounds, reverse, &mut visitor)?;
895 }
896 }
897
898 Ok(())
899 }
900
901 fn visit_journaled_keys_in_bounds<E>(
904 backend: &DataStoreBackend,
905 bounds: (Bound<&RawDataStoreKey>, Bound<&RawDataStoreKey>),
906 reverse: bool,
907 visitor: &mut dyn FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
908 ) -> Result<(), E> {
909 let DataStoreBackend::Journaled {
910 canonical,
911 live,
912 tombstones,
913 ..
914 } = backend
915 else {
916 return Ok(());
917 };
918
919 if canonical.is_empty() {
920 if reverse {
921 for (key, _row) in live.range(bounds).rev() {
922 if visitor(key)?.should_stop() {
923 return Ok(());
924 }
925 }
926 } else {
927 for (key, _row) in live.range(bounds) {
928 if visitor(key)?.should_stop() {
929 return Ok(());
930 }
931 }
932 }
933 return Ok(());
934 }
935
936 if live.is_empty() && tombstones.is_empty() {
937 if reverse {
938 for entry in canonical.range(bounds).rev() {
939 if visitor(entry.key())?.should_stop() {
940 return Ok(());
941 }
942 }
943 } else {
944 for entry in canonical.range(bounds) {
945 if visitor(entry.key())?.should_stop() {
946 return Ok(());
947 }
948 }
949 }
950 return Ok(());
951 }
952
953 let direction = if reverse {
954 Direction::Desc
955 } else {
956 Direction::Asc
957 };
958 match direction {
959 Direction::Asc => {
960 for entry in ordered_overlay_entries(
961 canonical.range(bounds),
962 live.range(bounds),
963 direction,
964 |entry| entry.key(),
965 |entry| entry.0,
966 tombstones,
967 ) {
968 let visit = match entry {
969 OrderedOverlayEntry::Canonical(canonical_entry) => {
970 visitor(canonical_entry.key())?
971 }
972 OrderedOverlayEntry::Live((key, _row)) => visitor(key)?,
973 };
974 if visit.should_stop() {
975 return Ok(());
976 }
977 }
978 }
979 Direction::Desc => {
980 for entry in ordered_overlay_entries(
981 canonical.range(bounds).rev(),
982 live.range(bounds).rev(),
983 direction,
984 |entry| entry.key(),
985 |entry| entry.0,
986 tombstones,
987 ) {
988 let visit = match entry {
989 OrderedOverlayEntry::Canonical(canonical_entry) => {
990 visitor(canonical_entry.key())?
991 }
992 OrderedOverlayEntry::Live((key, _row)) => visitor(key)?,
993 };
994 if visit.should_stop() {
995 return Ok(());
996 }
997 }
998 }
999 }
1000
1001 Ok(())
1002 }
1003
1004 fn visit_journaled_entries_in_bounds<E>(
1007 backend: &DataStoreBackend,
1008 bounds: (Bound<&RawDataStoreKey>, Bound<&RawDataStoreKey>),
1009 visitor: &mut dyn FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
1010 ) -> Result<(), E> {
1011 let DataStoreBackend::Journaled {
1012 canonical,
1013 live,
1014 tombstones,
1015 ..
1016 } = backend
1017 else {
1018 return Ok(());
1019 };
1020
1021 if canonical.is_empty() {
1022 for (key, row) in live.range(bounds) {
1023 if visitor(key, row)?.should_stop() {
1024 return Ok(());
1025 }
1026 }
1027 return Ok(());
1028 }
1029
1030 if live.is_empty() && tombstones.is_empty() {
1031 for entry in canonical.range(bounds) {
1032 if visitor(entry.key(), &entry.value())?.should_stop() {
1033 return Ok(());
1034 }
1035 }
1036 return Ok(());
1037 }
1038
1039 for entry in ordered_overlay_entries(
1040 canonical.range(bounds),
1041 live.range(bounds),
1042 Direction::Asc,
1043 |entry| entry.key(),
1044 |entry| entry.0,
1045 tombstones,
1046 ) {
1047 let visit = match entry {
1048 OrderedOverlayEntry::Canonical(canonical_entry) => {
1049 visitor(canonical_entry.key(), &canonical_entry.value())?
1050 }
1051 OrderedOverlayEntry::Live((key, row)) => visitor(key, row)?,
1052 };
1053 if visit.should_stop() {
1054 return Ok(());
1055 }
1056 }
1057
1058 Ok(())
1059 }
1060}
1061
1062#[derive(Clone, Debug)]
1063struct EntityCardinality {
1064 counts: HeapBTreeMap<EntityTag, u64>,
1065 decodable: bool,
1066}
1067
1068#[derive(Clone, Debug)]
1069struct EntityCardinalityDelta {
1070 counts: HeapBTreeMap<EntityTag, i64>,
1071 decodable: bool,
1072}
1073
1074impl EntityCardinalityDelta {
1075 const fn empty() -> Self {
1076 Self {
1077 counts: HeapBTreeMap::new(),
1078 decodable: true,
1079 }
1080 }
1081
1082 fn exact_delta(&self, entity: EntityTag) -> Option<i64> {
1083 self.decodable
1084 .then(|| self.counts.get(&entity).copied().unwrap_or(0))
1085 }
1086
1087 fn apply_presence_transition(
1088 &mut self,
1089 key: &RawDataStoreKey,
1090 previous_present: bool,
1091 next_present: bool,
1092 ) {
1093 if !self.decodable || previous_present == next_present {
1094 return;
1095 }
1096 let Some(entity) = key.entity_tag_prefix() else {
1097 self.counts.clear();
1098 self.decodable = false;
1099 return;
1100 };
1101 let delta = if next_present { 1_i64 } else { -1_i64 };
1102 let count = self.counts.entry(entity).or_insert(0);
1103 let Some(next) = count.checked_add(delta) else {
1104 self.counts.clear();
1105 self.decodable = false;
1106 return;
1107 };
1108 *count = next;
1109 if *count == 0 {
1110 self.counts.remove(&entity);
1111 }
1112 }
1113}
1114
1115impl EntityCardinality {
1116 const fn empty() -> Self {
1117 Self {
1118 counts: HeapBTreeMap::new(),
1119 decodable: true,
1120 }
1121 }
1122
1123 const fn unavailable() -> Self {
1124 Self {
1125 counts: HeapBTreeMap::new(),
1126 decodable: false,
1127 }
1128 }
1129
1130 fn exact_count(&self, entity: EntityTag) -> Option<u64> {
1131 self.decodable
1132 .then(|| self.counts.get(&entity).copied().unwrap_or(0))
1133 }
1134
1135 fn apply_insert(&mut self, key: &RawDataStoreKey, previous_present: bool) {
1136 if previous_present {
1137 return;
1138 }
1139 self.apply_present_key(key);
1140 }
1141
1142 fn apply_remove(&mut self, key: &RawDataStoreKey, previous_present: bool) {
1143 if !previous_present {
1144 return;
1145 }
1146 self.apply_removed_key(key);
1147 }
1148
1149 fn apply_present_key(&mut self, key: &RawDataStoreKey) {
1150 if !self.decodable {
1151 return;
1152 }
1153 let Some(entity) = key.entity_tag_prefix() else {
1154 self.invalidate();
1155 return;
1156 };
1157
1158 let count = self.counts.entry(entity).or_insert(0);
1159 *count = count.saturating_add(1);
1160 }
1161
1162 fn apply_removed_key(&mut self, key: &RawDataStoreKey) {
1163 if !self.decodable {
1164 return;
1165 }
1166 let Some(entity) = key.entity_tag_prefix() else {
1167 self.invalidate();
1168 return;
1169 };
1170
1171 if let Some(count) = self.counts.get_mut(&entity) {
1172 *count = count.saturating_sub(1);
1173 if *count == 0 {
1174 self.counts.remove(&entity);
1175 }
1176 }
1177 }
1178
1179 fn invalidate(&mut self) {
1180 self.counts.clear();
1181 self.decodable = false;
1182 }
1183}
1184
1185#[cfg(test)]
1186mod tests;