1use crate::{
7 db::{
8 data::{CanonicalRow, RawDataStoreKey, RawRow},
9 direction::Direction,
10 ordered_overlay::{OrderedOverlayEntry, OrderedOverlayVisit, visit_ordered_overlay},
11 },
12 types::EntityTag,
13};
14use ic_stable_structures::{
15 BTreeMap as StableBTreeMap, DefaultMemoryImpl, memory_manager::VirtualMemory,
16};
17#[cfg(all(feature = "sql", feature = "diagnostics"))]
18use std::cell::Cell;
19use std::collections::{BTreeMap as HeapBTreeMap, BTreeSet};
20use std::convert::Infallible;
21use std::ops::{Bound, RangeBounds};
22
23#[cfg(all(feature = "sql", feature = "diagnostics"))]
24thread_local! {
25 static DATA_STORE_GET_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
26}
27
28#[cfg(all(feature = "sql", feature = "diagnostics"))]
29fn record_data_store_get_call() {
30 DATA_STORE_GET_CALL_COUNT.with(|count| {
31 count.set(count.get().saturating_add(1));
32 });
33}
34
35pub struct DataStore {
45 backend: DataStoreBackend,
46 generation: u64,
47 entity_cardinality: EntityCardinality,
48}
49
50enum DataStoreBackend {
51 Heap(HeapBTreeMap<RawDataStoreKey, RawRow>),
52 Journaled {
53 canonical: StableBTreeMap<RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>>,
54 live: HeapBTreeMap<RawDataStoreKey, RawRow>,
55 tombstones: BTreeSet<RawDataStoreKey>,
56 },
57}
58
59#[derive(Clone, Copy, Debug, Eq, PartialEq)]
61pub(in crate::db) enum StoreVisit {
62 Continue,
63 Stop,
64}
65
66impl StoreVisit {
67 const fn should_stop(self) -> bool {
68 matches!(self, Self::Stop)
69 }
70}
71
72impl DataStore {
73 #[must_use]
75 pub const fn init_heap() -> Self {
76 Self {
77 backend: DataStoreBackend::Heap(HeapBTreeMap::new()),
78 generation: 0,
79 entity_cardinality: EntityCardinality::empty(),
80 }
81 }
82
83 #[must_use]
89 pub fn init_journaled(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
90 let mut store = Self {
91 backend: DataStoreBackend::Journaled {
92 canonical: StableBTreeMap::init(memory),
93 live: HeapBTreeMap::new(),
94 tombstones: BTreeSet::new(),
95 },
96 generation: 0,
97 entity_cardinality: EntityCardinality::empty(),
98 };
99 store.rebuild_entity_cardinality_from_entries();
100 store
101 }
102
103 pub(in crate::db) fn insert(
105 &mut self,
106 key: RawDataStoreKey,
107 row: CanonicalRow,
108 ) -> Option<RawRow> {
109 let row = row.into_raw_row();
110 let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
111 self.get(&key)
112 } else {
113 None
114 };
115 let cardinality_key = key.clone();
116 let previous = match &mut self.backend {
117 DataStoreBackend::Heap(map) => map.insert(key, row),
118 DataStoreBackend::Journaled {
119 live, tombstones, ..
120 } => {
121 tombstones.remove(&key);
122 live.insert(key, row);
123 previous_journaled
124 }
125 };
126 self.entity_cardinality
127 .apply_insert(&cardinality_key, previous.as_ref());
128 self.bump_generation();
129 previous
130 }
131
132 #[cfg(test)]
134 pub(in crate::db) fn insert_raw_for_test(
135 &mut self,
136 key: RawDataStoreKey,
137 row: RawRow,
138 ) -> Option<RawRow> {
139 let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
140 self.get(&key)
141 } else {
142 None
143 };
144 let cardinality_key = key.clone();
145 let previous = match &mut self.backend {
146 DataStoreBackend::Heap(map) => map.insert(key, row),
147 DataStoreBackend::Journaled {
148 live, tombstones, ..
149 } => {
150 tombstones.remove(&key);
151 live.insert(key, row);
152 previous_journaled
153 }
154 };
155 self.entity_cardinality
156 .apply_insert(&cardinality_key, previous.as_ref());
157 self.bump_generation();
158 previous
159 }
160
161 pub(in crate::db) fn remove(&mut self, key: &RawDataStoreKey) -> Option<RawRow> {
163 let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
164 self.get(key)
165 } else {
166 None
167 };
168 let previous = match &mut self.backend {
169 DataStoreBackend::Heap(map) => map.remove(key),
170 DataStoreBackend::Journaled {
171 live, tombstones, ..
172 } => {
173 live.remove(key);
174 tombstones.insert(key.clone());
175 previous_journaled
176 }
177 };
178 self.entity_cardinality.apply_remove(key, previous.as_ref());
179 self.bump_generation();
180 previous
181 }
182
183 pub(in crate::db) fn reset_journaled_live_projection(
186 &mut self,
187 ) -> Result<(), crate::error::InternalError> {
188 let DataStoreBackend::Journaled {
189 live, tombstones, ..
190 } = &mut self.backend
191 else {
192 return Err(crate::error::InternalError::store_invariant());
193 };
194
195 live.clear();
196 tombstones.clear();
197 self.rebuild_entity_cardinality_from_entries();
198 self.bump_generation();
199
200 Ok(())
201 }
202
203 pub(in crate::db) fn apply_recovered_journal_put(
205 &mut self,
206 key: RawDataStoreKey,
207 row: RawRow,
208 ) -> Result<Option<RawRow>, crate::error::InternalError> {
209 let DataStoreBackend::Journaled {
210 canonical,
211 live,
212 tombstones,
213 } = &mut self.backend
214 else {
215 return Err(crate::error::InternalError::store_invariant());
216 };
217
218 let previous = if tombstones.contains(&key) {
219 None
220 } else {
221 live.get(&key).cloned().or_else(|| canonical.get(&key))
222 };
223 tombstones.remove(&key);
224 let cardinality_key = key.clone();
225 live.insert(key, row);
226 self.entity_cardinality
227 .apply_insert(&cardinality_key, previous.as_ref());
228 self.bump_generation();
229
230 Ok(previous)
231 }
232
233 pub(in crate::db) fn apply_recovered_journal_delete(
235 &mut self,
236 key: &RawDataStoreKey,
237 ) -> Result<Option<RawRow>, crate::error::InternalError> {
238 let DataStoreBackend::Journaled {
239 canonical,
240 live,
241 tombstones,
242 } = &mut self.backend
243 else {
244 return Err(crate::error::InternalError::store_invariant());
245 };
246
247 let previous = if tombstones.contains(key) {
248 None
249 } else {
250 live.get(key).cloned().or_else(|| canonical.get(key))
251 };
252 live.remove(key);
253 tombstones.insert(key.clone());
254 self.entity_cardinality.apply_remove(key, previous.as_ref());
255 self.bump_generation();
256
257 Ok(previous)
258 }
259
260 pub(in crate::db) fn fold_recovered_journal_put(
262 &mut self,
263 key: RawDataStoreKey,
264 row: RawRow,
265 ) -> Result<Option<RawRow>, crate::error::InternalError> {
266 let DataStoreBackend::Journaled {
267 canonical,
268 live,
269 tombstones,
270 } = &mut self.backend
271 else {
272 return Err(crate::error::InternalError::store_invariant());
273 };
274
275 let visible = !live.contains_key(&key) && !tombstones.contains(&key);
276 let cardinality_key = key.clone();
277 let previous = canonical.insert(key, row);
278 if visible {
279 self.entity_cardinality
280 .apply_insert(&cardinality_key, previous.as_ref());
281 }
282 self.bump_generation();
283
284 Ok(previous)
285 }
286
287 pub(in crate::db) fn fold_recovered_journal_delete(
289 &mut self,
290 key: &RawDataStoreKey,
291 ) -> Result<Option<RawRow>, crate::error::InternalError> {
292 let DataStoreBackend::Journaled {
293 canonical,
294 live,
295 tombstones,
296 } = &mut self.backend
297 else {
298 return Err(crate::error::InternalError::store_invariant());
299 };
300
301 let visible = !live.contains_key(key) && !tombstones.contains(key);
302 let previous = canonical.remove(key);
303 if visible {
304 self.entity_cardinality.apply_remove(key, previous.as_ref());
305 }
306 self.bump_generation();
307
308 Ok(previous)
309 }
310
311 pub(in crate::db) fn get(&self, key: &RawDataStoreKey) -> Option<RawRow> {
313 #[cfg(all(feature = "sql", feature = "diagnostics"))]
314 record_data_store_get_call();
315
316 match &self.backend {
317 DataStoreBackend::Heap(map) => map.get(key).cloned(),
318 DataStoreBackend::Journaled { .. } => Self::journaled_get_raw(&self.backend, key),
319 }
320 }
321
322 #[must_use]
324 pub(in crate::db) fn contains(&self, key: &RawDataStoreKey) -> bool {
325 match &self.backend {
326 DataStoreBackend::Heap(map) => map.contains_key(key),
327 DataStoreBackend::Journaled { .. } => {
328 Self::journaled_get_raw(&self.backend, key).is_some()
329 }
330 }
331 }
332
333 #[must_use]
335 pub(in crate::db) fn len(&self) -> u64 {
336 match &self.backend {
337 DataStoreBackend::Heap(map) => u64::try_from(map.len()).unwrap_or(u64::MAX),
338 DataStoreBackend::Journaled { .. } => {
339 let mut count = 0_u64;
340 let _: Result<(), Infallible> = self.visit_entries(|_key, _row| {
341 count = count.saturating_add(1);
342 Ok(StoreVisit::Continue)
343 });
344 count
345 }
346 }
347 }
348
349 #[must_use]
351 pub(in crate::db) const fn generation(&self) -> u64 {
352 self.generation
353 }
354
355 #[must_use]
357 pub(in crate::db) fn exact_entity_count(&self, entity: EntityTag) -> Option<u64> {
358 self.entity_cardinality.exact_count(entity)
359 }
360
361 pub(in crate::db) fn visit_entries<E>(
363 &self,
364 mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
365 ) -> Result<(), E> {
366 match &self.backend {
367 DataStoreBackend::Heap(map) => {
368 for (key, row) in map {
369 if visitor(key, row)?.should_stop() {
370 break;
371 }
372 }
373 }
374 DataStoreBackend::Journaled {
375 canonical: _,
376 live: _,
377 tombstones: _,
378 } => Self::visit_journaled_entries_in_bounds(
379 &self.backend,
380 (Bound::Unbounded, Bound::Unbounded),
381 visitor,
382 )?,
383 }
384
385 Ok(())
386 }
387
388 pub(in crate::db) fn visit_range<E>(
390 &self,
391 key_range: impl RangeBounds<RawDataStoreKey>,
392 mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
393 ) -> Result<(), E> {
394 let bounds = Self::owned_range_bounds(&key_range);
395 match &self.backend {
396 DataStoreBackend::Heap(map) => {
397 for (key, row) in map.range((bounds.0.clone(), bounds.1)) {
398 if visitor(key, row)?.should_stop() {
399 break;
400 }
401 }
402 }
403 DataStoreBackend::Journaled {
404 canonical: _,
405 live: _,
406 tombstones: _,
407 } => Self::visit_journaled_entries_in_bounds(&self.backend, bounds, visitor)?,
408 }
409
410 Ok(())
411 }
412
413 pub(in crate::db) fn visit_key_range<E>(
420 &self,
421 key_range: impl RangeBounds<RawDataStoreKey>,
422 visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
423 ) -> Result<(), E> {
424 self.visit_keys_in_bounds(Self::owned_range_bounds(&key_range), false, visitor)
425 }
426
427 pub(in crate::db) fn visit_key_range_rev<E>(
429 &self,
430 key_range: impl RangeBounds<RawDataStoreKey>,
431 visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
432 ) -> Result<(), E> {
433 self.visit_keys_in_bounds(Self::owned_range_bounds(&key_range), true, visitor)
434 }
435
436 pub(in crate::db) fn memory_bytes(&self) -> u64 {
438 let mut bytes = 0u64;
440 let _: Result<(), Infallible> = self.visit_entries(|key, row| {
441 bytes = bytes.saturating_add(key.as_bytes().len() as u64 + row.len() as u64);
442 Ok(StoreVisit::Continue)
443 });
444 bytes
445 }
446
447 const fn bump_generation(&mut self) {
448 self.generation = self.generation.saturating_add(1);
449 }
450
451 fn rebuild_entity_cardinality_from_entries(&mut self) {
452 let mut cardinality = EntityCardinality::empty();
453 let _: Result<(), Infallible> = self.visit_entries(|key, _row| {
454 cardinality.apply_present_key(key);
455 Ok(StoreVisit::Continue)
456 });
457 self.entity_cardinality = cardinality;
458 }
459
460 #[cfg(all(feature = "sql", feature = "diagnostics"))]
462 pub(in crate::db) fn current_get_call_count() -> u64 {
463 DATA_STORE_GET_CALL_COUNT.with(Cell::get)
464 }
465
466 fn journaled_get_raw(backend: &DataStoreBackend, key: &RawDataStoreKey) -> Option<RawRow> {
467 let DataStoreBackend::Journaled {
468 canonical,
469 live,
470 tombstones,
471 } = backend
472 else {
473 return None;
474 };
475
476 if tombstones.contains(key) {
477 return None;
478 }
479 live.get(key).cloned().or_else(|| canonical.get(key))
480 }
481
482 fn owned_range_bounds(
483 key_range: &impl RangeBounds<RawDataStoreKey>,
484 ) -> (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>) {
485 let lower = match key_range.start_bound() {
486 Bound::Included(key) => Bound::Included(key.clone()),
487 Bound::Excluded(key) => Bound::Excluded(key.clone()),
488 Bound::Unbounded => Bound::Unbounded,
489 };
490 let upper = match key_range.end_bound() {
491 Bound::Included(key) => Bound::Included(key.clone()),
492 Bound::Excluded(key) => Bound::Excluded(key.clone()),
493 Bound::Unbounded => Bound::Unbounded,
494 };
495
496 (lower, upper)
497 }
498
499 fn visit_keys_in_bounds<E>(
500 &self,
501 bounds: (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>),
502 reverse: bool,
503 mut visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
504 ) -> Result<(), E> {
505 match &self.backend {
506 DataStoreBackend::Heap(map) => {
507 if reverse {
508 for (key, _row) in map.range(bounds).rev() {
509 if visitor(key)?.should_stop() {
510 break;
511 }
512 }
513 } else {
514 for (key, _row) in map.range(bounds) {
515 if visitor(key)?.should_stop() {
516 break;
517 }
518 }
519 }
520 }
521 DataStoreBackend::Journaled { .. } => {
522 Self::visit_journaled_keys_in_bounds(&self.backend, bounds, reverse, visitor)?;
523 }
524 }
525
526 Ok(())
527 }
528
529 fn visit_journaled_keys_in_bounds<E>(
530 backend: &DataStoreBackend,
531 bounds: (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>),
532 reverse: bool,
533 mut visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
534 ) -> Result<(), E> {
535 let DataStoreBackend::Journaled {
536 canonical,
537 live,
538 tombstones,
539 } = backend
540 else {
541 return Ok(());
542 };
543
544 if canonical.is_empty() {
545 if reverse {
546 for (key, _row) in live.range(bounds).rev() {
547 if visitor(key)?.should_stop() {
548 return Ok(());
549 }
550 }
551 } else {
552 for (key, _row) in live.range(bounds) {
553 if visitor(key)?.should_stop() {
554 return Ok(());
555 }
556 }
557 }
558 return Ok(());
559 }
560
561 if live.is_empty() && tombstones.is_empty() {
562 if reverse {
563 for entry in canonical.range(bounds).rev() {
564 if visitor(entry.key())?.should_stop() {
565 return Ok(());
566 }
567 }
568 } else {
569 for entry in canonical.range(bounds) {
570 if visitor(entry.key())?.should_stop() {
571 return Ok(());
572 }
573 }
574 }
575 return Ok(());
576 }
577
578 let direction = if reverse {
579 Direction::Desc
580 } else {
581 Direction::Asc
582 };
583 match direction {
584 Direction::Asc => visit_ordered_overlay(
585 canonical.range((bounds.0.clone(), bounds.1.clone())),
586 live.range((bounds.0, bounds.1)),
587 direction,
588 |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
589 |canonical_entry| !tombstones.contains(canonical_entry.key()),
590 |live_entry| !tombstones.contains(live_entry.0),
591 |entry| {
592 let visit = match entry {
593 OrderedOverlayEntry::Canonical(canonical_entry) => {
594 visitor(canonical_entry.key())?
595 }
596 OrderedOverlayEntry::Live((key, _row)) => visitor(key)?,
597 };
598 Ok(if visit.should_stop() {
599 OrderedOverlayVisit::Stop
600 } else {
601 OrderedOverlayVisit::Continue
602 })
603 },
604 ),
605 Direction::Desc => visit_ordered_overlay(
606 canonical.range((bounds.0.clone(), bounds.1.clone())).rev(),
607 live.range((bounds.0, bounds.1)).rev(),
608 direction,
609 |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
610 |canonical_entry| !tombstones.contains(canonical_entry.key()),
611 |live_entry| !tombstones.contains(live_entry.0),
612 |entry| {
613 let visit = match entry {
614 OrderedOverlayEntry::Canonical(canonical_entry) => {
615 visitor(canonical_entry.key())?
616 }
617 OrderedOverlayEntry::Live((key, _row)) => visitor(key)?,
618 };
619 Ok(if visit.should_stop() {
620 OrderedOverlayVisit::Stop
621 } else {
622 OrderedOverlayVisit::Continue
623 })
624 },
625 ),
626 }
627 }
628
629 fn visit_journaled_entries_in_bounds<E>(
630 backend: &DataStoreBackend,
631 bounds: (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>),
632 mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
633 ) -> Result<(), E> {
634 let DataStoreBackend::Journaled {
635 canonical,
636 live,
637 tombstones,
638 } = backend
639 else {
640 return Ok(());
641 };
642
643 if canonical.is_empty() {
644 for (key, row) in live.range(bounds) {
645 if visitor(key, row)?.should_stop() {
646 return Ok(());
647 }
648 }
649 return Ok(());
650 }
651
652 if live.is_empty() && tombstones.is_empty() {
653 for entry in canonical.range(bounds) {
654 if visitor(entry.key(), &entry.value())?.should_stop() {
655 return Ok(());
656 }
657 }
658 return Ok(());
659 }
660
661 visit_ordered_overlay(
662 canonical.range((bounds.0.clone(), bounds.1.clone())),
663 live.range((bounds.0, bounds.1)),
664 Direction::Asc,
665 |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
666 |canonical_entry| !tombstones.contains(canonical_entry.key()),
667 |live_entry| !tombstones.contains(live_entry.0),
668 |entry| {
669 let visit = match entry {
670 OrderedOverlayEntry::Canonical(canonical_entry) => {
671 visitor(canonical_entry.key(), &canonical_entry.value())?
672 }
673 OrderedOverlayEntry::Live((key, row)) => visitor(key, row)?,
674 };
675 Ok(if visit.should_stop() {
676 OrderedOverlayVisit::Stop
677 } else {
678 OrderedOverlayVisit::Continue
679 })
680 },
681 )
682 }
683}
684
685#[derive(Clone, Debug)]
686struct EntityCardinality {
687 counts: HeapBTreeMap<EntityTag, u64>,
688 decodable: bool,
689}
690
691impl EntityCardinality {
692 const fn empty() -> Self {
693 Self {
694 counts: HeapBTreeMap::new(),
695 decodable: true,
696 }
697 }
698
699 fn exact_count(&self, entity: EntityTag) -> Option<u64> {
700 self.decodable
701 .then(|| self.counts.get(&entity).copied().unwrap_or(0))
702 }
703
704 fn apply_insert(&mut self, key: &RawDataStoreKey, previous: Option<&RawRow>) {
705 if previous.is_some() {
706 return;
707 }
708 self.apply_present_key(key);
709 }
710
711 fn apply_remove(&mut self, key: &RawDataStoreKey, previous: Option<&RawRow>) {
712 if previous.is_none() {
713 return;
714 }
715 self.apply_removed_key(key);
716 }
717
718 fn apply_present_key(&mut self, key: &RawDataStoreKey) {
719 if !self.decodable {
720 return;
721 }
722 let Some(entity) = key.entity_tag_prefix() else {
723 self.invalidate();
724 return;
725 };
726
727 let count = self.counts.entry(entity).or_insert(0);
728 *count = count.saturating_add(1);
729 }
730
731 fn apply_removed_key(&mut self, key: &RawDataStoreKey) {
732 if !self.decodable {
733 return;
734 }
735 let Some(entity) = key.entity_tag_prefix() else {
736 self.invalidate();
737 return;
738 };
739
740 if let Some(count) = self.counts.get_mut(&entity) {
741 *count = count.saturating_sub(1);
742 if *count == 0 {
743 self.counts.remove(&entity);
744 }
745 }
746 }
747
748 fn invalidate(&mut self) {
749 self.counts.clear();
750 self.decodable = false;
751 }
752}
753
754#[cfg(test)]
755mod tests;