Skip to main content

nodedb_lite/sync/
stream_seq.rs

1//! Per-stream monotonic sequence frontier — durable state for outbound frame
2//! stamping and inbound ack tracking.
3//!
4//! # Persistence
5//!
6//! Each per-stream frontier is persisted under `Namespace::Meta` with key
7//! `"sync.stream_seq:{stream_id:016x}"`. The 16-byte value encodes
8//! `last_assigned || last_acked` as two big-endian u64s.
9//!
10//! [`StreamSeqTracker::load`] restores all entries on startup so the first
11//! outbound frame after a restart continues from where it left off (the
12//! persist-before-send invariant is maintained across process restarts).
13
14use std::collections::HashMap;
15use std::sync::{Arc, Mutex};
16
17use nodedb_types::Namespace;
18
19use crate::error::LiteError;
20use crate::storage::engine::StorageEngine;
21
22/// Storage key prefix for per-stream sequence frontier entries.
23const STREAM_SEQ_PREFIX: &str = "sync.stream_seq:";
24
25/// Storage key for a given stream_id.
26fn stream_seq_key(stream_id: u64) -> Vec<u8> {
27    format!("{STREAM_SEQ_PREFIX}{stream_id:016x}").into_bytes()
28}
29
30/// In-memory frontier for a single stream.
31#[derive(Debug, Clone, Copy, Default)]
32struct Frontier {
33    /// Highest sequence number assigned (and persisted) for outbound frames.
34    last_assigned: u64,
35    /// Highest sequence number acknowledged by Origin.
36    last_acked: u64,
37}
38
39/// Tracks the per-stream monotonic sequence frontier for outbound frame
40/// stamping and ack recording.
41///
42/// Thread-safe via an internal [`Mutex`]. Persist-before-send is enforced
43/// in [`next_seq`]: storage is flushed before the sequence number is returned
44/// to the caller.
45pub struct StreamSeqTracker<S: StorageEngine> {
46    storage: Arc<S>,
47    state: Mutex<HashMap<u64, Frontier>>,
48}
49
50impl<S: StorageEngine> StreamSeqTracker<S> {
51    /// Load all persisted stream sequence entries from `Namespace::Meta`.
52    ///
53    /// Scans keys starting with `b"sync.stream_seq:"` and parses each
54    /// 16-byte value as `last_assigned || last_acked` (two big-endian u64s).
55    /// Returns an empty tracker if no entries are stored yet.
56    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    /// Assign the next sequence number for `stream_id`.
99    ///
100    /// Computes `next = last_assigned + 1`, persists the updated frontier to
101    /// `Namespace::Meta` BEFORE returning (persist-before-send invariant), then
102    /// returns `next`.
103    ///
104    /// The in-memory lock is dropped before the storage await to avoid holding
105    /// it across async I/O.
106    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    /// Record that Origin has applied `applied_seq` for `stream_id`.
124    ///
125    /// Advances `last_acked` (and `last_assigned` if needed to stay consistent),
126    /// then persists. No-op if `applied_seq <= current last_acked`.
127    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    /// Return the highest acknowledged sequence for `stream_id`, or 0 if unknown.
150    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    /// Return the highest assigned sequence for `stream_id`, or 0 if unknown.
159    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
168/// Encode a `Frontier` as 16 bytes: `last_assigned || last_acked` big-endian.
169fn 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// ─── Tests ────────────────────────────────────────────────────────────────────
177
178#[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            // Must resume from 2, not restart from 0.
236            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(); // assigns 1
260            tracker.next_seq(5).await.unwrap(); // assigns 2
261            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        // Recording same or lower is a no-op.
290        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        // Stream 20 starts independently at 1.
301        assert_eq!(tracker.next_seq(20).await.unwrap(), 1);
302        assert_eq!(tracker.next_seq(20).await.unwrap(), 2);
303        // Stream 10 continues from where it left off.
304        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        // Stream 20's ack is unaffected.
309        assert_eq!(tracker.last_acked(20).await, 0);
310    }
311}