reifydb_cdc/storage/
memory.rs1use std::{
5 collections::{BTreeMap, Bound},
6 sync::{
7 Arc,
8 atomic::{AtomicU64, Ordering},
9 },
10};
11
12use reifydb_core::{
13 common::CommitVersion,
14 interface::cdc::{Cdc, CdcBatch},
15};
16use reifydb_runtime::sync::rwlock::RwLock;
17
18use super::{
19 CdcStorage, CdcStorageResult, DropBeforeResult, aggregate_evictions, normalize_range_inclusive,
20 total_evicted_count,
21};
22
23#[derive(Clone)]
24pub struct MemoryCdcStorage {
25 inner: Arc<RwLock<BTreeMap<CommitVersion, Cdc>>>,
26 truncated_before: Arc<AtomicU64>,
27}
28
29impl MemoryCdcStorage {
30 pub fn new() -> Self {
31 Self {
32 inner: Arc::new(RwLock::new(BTreeMap::new())),
33 truncated_before: Arc::new(AtomicU64::new(0)),
34 }
35 }
36
37 pub fn with_entries(entries: impl IntoIterator<Item = Cdc>) -> Self {
38 let map: BTreeMap<CommitVersion, Cdc> = entries.into_iter().map(|cdc| (cdc.version, cdc)).collect();
39 Self {
40 inner: Arc::new(RwLock::new(map)),
41 truncated_before: Arc::new(AtomicU64::new(0)),
42 }
43 }
44
45 pub fn len(&self) -> usize {
46 self.inner.read().len()
47 }
48
49 pub fn is_empty(&self) -> bool {
50 self.inner.read().is_empty()
51 }
52
53 pub fn clear(&self) {
54 self.inner.write().clear();
55 }
56}
57
58impl Default for MemoryCdcStorage {
59 fn default() -> Self {
60 Self::new()
61 }
62}
63
64impl CdcStorage for MemoryCdcStorage {
65 fn write(&self, cdc: &Cdc) -> CdcStorageResult<()> {
66 self.inner.write().insert(cdc.version, cdc.clone());
67 Ok(())
68 }
69
70 fn read(&self, version: CommitVersion) -> CdcStorageResult<Option<Cdc>> {
71 Ok(self.inner.read().get(&version).cloned())
72 }
73
74 fn read_range(
75 &self,
76 start: Bound<CommitVersion>,
77 end: Bound<CommitVersion>,
78 batch_size: u64,
79 ) -> CdcStorageResult<CdcBatch> {
80 let Some((lo_inc, hi_inc)) = normalize_range_inclusive(start, end) else {
81 return Ok(CdcBatch {
82 items: Vec::new(),
83 has_more: false,
84 });
85 };
86 let guard = self.inner.read();
87 let (items, has_more) = collect_range_into(&guard, lo_inc, hi_inc, batch_size as usize);
88 Ok(CdcBatch {
89 items,
90 has_more,
91 })
92 }
93
94 fn count(&self, version: CommitVersion) -> CdcStorageResult<usize> {
95 Ok(self.inner.read().get(&version).map(|cdc| cdc.system_changes.len()).unwrap_or(0))
96 }
97
98 fn min_version(&self) -> CdcStorageResult<Option<CommitVersion>> {
99 Ok(self.inner.read().keys().next().copied())
100 }
101
102 fn max_version(&self) -> CdcStorageResult<Option<CommitVersion>> {
103 Ok(self.inner.read().keys().next_back().copied())
104 }
105
106 fn drop_before(&self, version: CommitVersion, limit: usize) -> CdcStorageResult<DropBeforeResult> {
107 let mut guard = self.inner.write();
108 let keys_to_remove: Vec<_> = guard.range(..version).take(limit).map(|(k, _)| *k).collect();
109 let more_remaining = keys_to_remove.len() == limit && guard.range(..version).nth(limit).is_some();
110 let entries = aggregate_evictions(
111 keys_to_remove.iter().filter_map(|k| guard.get(k)).flat_map(|cdc| cdc.system_changes.iter()),
112 );
113 let count = total_evicted_count(&entries);
114 for key in &keys_to_remove {
115 guard.remove(key);
116 }
117 if let Some(max_deleted) = keys_to_remove.last() {
118 self.truncated_before.fetch_max(max_deleted.0.saturating_add(1), Ordering::Release);
119 }
120 Ok(DropBeforeResult {
121 count,
122 entries,
123 more_remaining,
124 })
125 }
126
127 fn truncated_before(&self) -> CdcStorageResult<CommitVersion> {
128 Ok(CommitVersion(self.truncated_before.load(Ordering::Acquire)))
129 }
130}
131
132#[inline]
133fn collect_range_into(
134 guard: &BTreeMap<CommitVersion, Cdc>,
135 lo_inc: CommitVersion,
136 hi_inc: CommitVersion,
137 batch_size: usize,
138) -> (Vec<Cdc>, bool) {
139 let mut items: Vec<Cdc> = Vec::with_capacity(batch_size.min(64));
140 for (count, (_, cdc)) in guard.range(lo_inc..=hi_inc).enumerate() {
141 if count >= batch_size {
142 return (items, true);
143 }
144 items.push(cdc.clone());
145 }
146 (items, false)
147}