Skip to main content

reifydb_cdc/storage/
memory.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use 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}