nodedb_lite/sync/
stream_seq.rs1use std::collections::HashMap;
15use std::sync::{Arc, Mutex};
16
17use nodedb_types::Namespace;
18
19use crate::error::LiteError;
20use crate::storage::engine::StorageEngine;
21
22const STREAM_SEQ_PREFIX: &str = "sync.stream_seq:";
24
25fn stream_seq_key(stream_id: u64) -> Vec<u8> {
27 format!("{STREAM_SEQ_PREFIX}{stream_id:016x}").into_bytes()
28}
29
30#[derive(Debug, Clone, Copy, Default)]
32struct Frontier {
33 last_assigned: u64,
35 last_acked: u64,
37}
38
39pub struct StreamSeqTracker<S: StorageEngine> {
46 storage: Arc<S>,
47 state: Mutex<HashMap<u64, Frontier>>,
48}
49
50impl<S: StorageEngine> StreamSeqTracker<S> {
51 pub async fn load(storage: Arc<S>) -> Result<Self, LiteError> {
57 let prefix = STREAM_SEQ_PREFIX.as_bytes();
58 let pairs = storage
59 .scan_range(Namespace::Meta, prefix, usize::MAX)
60 .await?;
61
62 let mut state = HashMap::new();
63 for (key, value) in pairs {
64 if !key.starts_with(prefix) {
65 break;
66 }
67 let id_bytes = &key[prefix.len()..];
68 let id_str = std::str::from_utf8(id_bytes).map_err(|e| LiteError::Storage {
69 detail: format!("stream_seq_tracker: non-UTF8 stream_id in storage: {e}"),
70 })?;
71 let stream_id = u64::from_str_radix(id_str, 16).map_err(|e| LiteError::Storage {
72 detail: format!("stream_seq_tracker: invalid hex stream_id '{id_str}': {e}"),
73 })?;
74
75 let arr: [u8; 16] = value.try_into().map_err(|v: Vec<u8>| LiteError::Storage {
76 detail: format!(
77 "stream_seq_tracker: frontier wrong length ({}) for stream {stream_id:016x}",
78 v.len()
79 ),
80 })?;
81 let last_assigned = u64::from_be_bytes(arr[..8].try_into().unwrap_or([0; 8]));
82 let last_acked = u64::from_be_bytes(arr[8..].try_into().unwrap_or([0; 8]));
83 state.insert(
84 stream_id,
85 Frontier {
86 last_assigned,
87 last_acked,
88 },
89 );
90 }
91
92 Ok(Self {
93 storage,
94 state: Mutex::new(state),
95 })
96 }
97
98 pub async fn next_seq(&self, stream_id: u64) -> Result<u64, LiteError> {
107 let (next, encoded) = {
108 let mut state = self.state.lock().map_err(|_| LiteError::LockPoisoned)?;
109 let entry = state.entry(stream_id).or_default();
110 let next = entry.last_assigned + 1;
111 entry.last_assigned = next;
112 let encoded = encode_frontier(*entry);
113 (next, encoded)
114 };
115
116 self.storage
117 .put(Namespace::Meta, &stream_seq_key(stream_id), &encoded)
118 .await?;
119
120 Ok(next)
121 }
122
123 pub async fn record_ack(&self, stream_id: u64, applied_seq: u64) -> Result<(), LiteError> {
128 let (updated, encoded) = {
129 let mut state = self.state.lock().map_err(|_| LiteError::LockPoisoned)?;
130 let entry = state.entry(stream_id).or_default();
131 if applied_seq <= entry.last_acked {
132 return Ok(());
133 }
134 entry.last_acked = applied_seq;
135 entry.last_assigned = entry.last_assigned.max(applied_seq);
136 let encoded = encode_frontier(*entry);
137 (true, encoded)
138 };
139
140 if updated {
141 self.storage
142 .put(Namespace::Meta, &stream_seq_key(stream_id), &encoded)
143 .await?;
144 }
145
146 Ok(())
147 }
148
149 pub async fn last_acked(&self, stream_id: u64) -> u64 {
151 self.state
152 .lock()
153 .ok()
154 .and_then(|s| s.get(&stream_id).map(|f| f.last_acked))
155 .unwrap_or(0)
156 }
157
158 pub async fn last_assigned(&self, stream_id: u64) -> u64 {
160 self.state
161 .lock()
162 .ok()
163 .and_then(|s| s.get(&stream_id).map(|f| f.last_assigned))
164 .unwrap_or(0)
165 }
166}
167
168fn encode_frontier(f: Frontier) -> Vec<u8> {
170 let mut buf = Vec::with_capacity(16);
171 buf.extend_from_slice(&f.last_assigned.to_be_bytes());
172 buf.extend_from_slice(&f.last_acked.to_be_bytes());
173 buf
174}
175
176#[cfg(test)]
179mod tests {
180 use super::*;
181 use crate::storage::pagedb_storage::{PagedbStorageDefault, PagedbStorageMem};
182
183 async fn make_tracker() -> StreamSeqTracker<PagedbStorageMem> {
184 let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
185 StreamSeqTracker::load(storage).await.unwrap()
186 }
187
188 #[tokio::test]
189 async fn next_seq_starts_at_one() {
190 let tracker = make_tracker().await;
191 let seq = tracker.next_seq(1).await.unwrap();
192 assert_eq!(seq, 1);
193 }
194
195 #[tokio::test]
196 async fn next_seq_is_monotonic() {
197 let tracker = make_tracker().await;
198 let s1 = tracker.next_seq(42).await.unwrap();
199 let s2 = tracker.next_seq(42).await.unwrap();
200 let s3 = tracker.next_seq(42).await.unwrap();
201 assert_eq!(s1, 1);
202 assert_eq!(s2, 2);
203 assert_eq!(s3, 3);
204 }
205
206 #[tokio::test]
207 async fn next_seq_survives_reload() {
208 let dir = tempfile::tempdir().unwrap();
209 let path = dir.path().join("stream_seq_test.pagedb");
210
211 {
212 let storage = Arc::new(
213 PagedbStorageDefault::open(
214 &path,
215 crate::storage::encryption::Encryption::Plaintext,
216 )
217 .await
218 .unwrap(),
219 );
220 let tracker = StreamSeqTracker::load(Arc::clone(&storage)).await.unwrap();
221 assert_eq!(tracker.next_seq(7).await.unwrap(), 1);
222 assert_eq!(tracker.next_seq(7).await.unwrap(), 2);
223 }
224
225 {
226 let storage = Arc::new(
227 PagedbStorageDefault::open(
228 &path,
229 crate::storage::encryption::Encryption::Plaintext,
230 )
231 .await
232 .unwrap(),
233 );
234 let tracker = StreamSeqTracker::load(storage).await.unwrap();
235 assert_eq!(
237 tracker.next_seq(7).await.unwrap(),
238 3,
239 "seq must resume from durable last_assigned on reload"
240 );
241 }
242 }
243
244 #[tokio::test]
245 async fn record_ack_advances_and_survives_reload() {
246 let dir = tempfile::tempdir().unwrap();
247 let path = dir.path().join("stream_seq_ack_test.pagedb");
248
249 {
250 let storage = Arc::new(
251 PagedbStorageDefault::open(
252 &path,
253 crate::storage::encryption::Encryption::Plaintext,
254 )
255 .await
256 .unwrap(),
257 );
258 let tracker = StreamSeqTracker::load(Arc::clone(&storage)).await.unwrap();
259 tracker.next_seq(5).await.unwrap(); tracker.next_seq(5).await.unwrap(); tracker.record_ack(5, 2).await.unwrap();
262 assert_eq!(tracker.last_acked(5).await, 2);
263 }
264
265 {
266 let storage = Arc::new(
267 PagedbStorageDefault::open(
268 &path,
269 crate::storage::encryption::Encryption::Plaintext,
270 )
271 .await
272 .unwrap(),
273 );
274 let tracker = StreamSeqTracker::load(storage).await.unwrap();
275 assert_eq!(
276 tracker.last_acked(5).await,
277 2,
278 "last_acked must survive storage restart"
279 );
280 }
281 }
282
283 #[tokio::test]
284 async fn record_ack_is_noop_if_not_advancing() {
285 let tracker = make_tracker().await;
286 tracker.next_seq(3).await.unwrap();
287 tracker.record_ack(3, 1).await.unwrap();
288 assert_eq!(tracker.last_acked(3).await, 1);
289 tracker.record_ack(3, 1).await.unwrap();
291 tracker.record_ack(3, 0).await.unwrap();
292 assert_eq!(tracker.last_acked(3).await, 1);
293 }
294
295 #[tokio::test]
296 async fn two_stream_ids_are_independent() {
297 let tracker = make_tracker().await;
298 assert_eq!(tracker.next_seq(10).await.unwrap(), 1);
299 assert_eq!(tracker.next_seq(10).await.unwrap(), 2);
300 assert_eq!(tracker.next_seq(20).await.unwrap(), 1);
302 assert_eq!(tracker.next_seq(20).await.unwrap(), 2);
303 assert_eq!(tracker.next_seq(10).await.unwrap(), 3);
305
306 tracker.record_ack(10, 2).await.unwrap();
307 assert_eq!(tracker.last_acked(10).await, 2);
308 assert_eq!(tracker.last_acked(20).await, 0);
310 }
311}