Skip to main content

nodedb_lite/sync/
shapes.rs

1//! Shape subscription management on the edge side.
2//!
3//! Tracks which shapes the client is subscribed to, handles initial
4//! snapshot loading and incremental delta application.
5
6use std::collections::HashMap;
7
8use nodedb_types::sync::shape::ShapeDefinition;
9
10/// Manages shape subscriptions for a Lite client.
11///
12/// On reconnect, the client re-sends all active subscriptions in the
13/// handshake. The ShapeManager tracks what's active and the LSN
14/// watermark for each shape (so the Origin knows what deltas to send).
15pub struct ShapeManager {
16    /// Active subscriptions: shape_id → (definition, last_seen_lsn).
17    subscriptions: HashMap<String, ShapeSubscription>,
18}
19
20/// A single active shape subscription.
21#[derive(Debug, Clone)]
22pub struct ShapeSubscription {
23    /// The shape definition (what data this subscription covers).
24    pub definition: ShapeDefinition,
25    /// LSN of the last delta received for this shape.
26    /// Used to resume sync after reconnect.
27    pub last_lsn: u64,
28    /// Whether the initial snapshot has been loaded.
29    pub snapshot_loaded: bool,
30}
31
32impl ShapeManager {
33    pub fn new() -> Self {
34        Self {
35            subscriptions: HashMap::new(),
36        }
37    }
38
39    /// Subscribe to a shape. Returns the shape_id.
40    pub fn subscribe(&mut self, definition: ShapeDefinition) -> String {
41        let shape_id = definition.shape_id.clone();
42        self.subscriptions.insert(
43            shape_id.clone(),
44            ShapeSubscription {
45                definition,
46                last_lsn: 0,
47                snapshot_loaded: false,
48            },
49        );
50        shape_id
51    }
52
53    /// Unsubscribe from a shape.
54    pub fn unsubscribe(&mut self, shape_id: &str) -> bool {
55        self.subscriptions.remove(shape_id).is_some()
56    }
57
58    /// Mark a shape's snapshot as loaded.
59    pub fn mark_snapshot_loaded(&mut self, shape_id: &str, lsn: u64) {
60        if let Some(sub) = self.subscriptions.get_mut(shape_id) {
61            sub.snapshot_loaded = true;
62            sub.last_lsn = lsn;
63        }
64    }
65
66    /// Update the last-seen LSN for a shape (after receiving a delta).
67    pub fn advance_lsn(&mut self, shape_id: &str, lsn: u64) {
68        if let Some(sub) = self.subscriptions.get_mut(shape_id)
69            && lsn > sub.last_lsn
70        {
71            sub.last_lsn = lsn;
72        }
73    }
74
75    /// Get a subscription by shape_id.
76    pub fn get(&self, shape_id: &str) -> Option<&ShapeSubscription> {
77        self.subscriptions.get(shape_id)
78    }
79
80    /// All active shape IDs (for handshake).
81    pub fn active_shape_ids(&self) -> Vec<String> {
82        self.subscriptions.keys().cloned().collect()
83    }
84
85    /// All active shape definitions (for re-subscription on reconnect).
86    pub fn active_definitions(&self) -> Vec<&ShapeDefinition> {
87        self.subscriptions.values().map(|s| &s.definition).collect()
88    }
89
90    /// Number of active subscriptions.
91    pub fn count(&self) -> usize {
92        self.subscriptions.len()
93    }
94
95    /// Check if a specific shape is subscribed.
96    pub fn is_subscribed(&self, shape_id: &str) -> bool {
97        self.subscriptions.contains_key(shape_id)
98    }
99}
100
101impl Default for ShapeManager {
102    fn default() -> Self {
103        Self::new()
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use nodedb_types::sync::shape::ShapeType;
111
112    fn make_shape(id: &str, collection: &str) -> ShapeDefinition {
113        ShapeDefinition {
114            shape_id: id.into(),
115            tenant_id: 1,
116            shape_type: ShapeType::Document {
117                collection: collection.into(),
118                predicate: Vec::new(),
119            },
120            description: format!("all {collection}"),
121            field_filter: vec![],
122        }
123    }
124
125    #[test]
126    fn subscribe_and_query() {
127        let mut mgr = ShapeManager::new();
128        let id = mgr.subscribe(make_shape("s1", "orders"));
129
130        assert_eq!(id, "s1");
131        assert!(mgr.is_subscribed("s1"));
132        assert_eq!(mgr.count(), 1);
133
134        let sub = mgr.get("s1").unwrap();
135        assert!(!sub.snapshot_loaded);
136        assert_eq!(sub.last_lsn, 0);
137    }
138
139    #[test]
140    fn unsubscribe() {
141        let mut mgr = ShapeManager::new();
142        mgr.subscribe(make_shape("s1", "orders"));
143        assert!(mgr.unsubscribe("s1"));
144        assert!(!mgr.is_subscribed("s1"));
145        assert_eq!(mgr.count(), 0);
146        assert!(!mgr.unsubscribe("s1")); // Already removed.
147    }
148
149    #[test]
150    fn snapshot_and_lsn_tracking() {
151        let mut mgr = ShapeManager::new();
152        mgr.subscribe(make_shape("s1", "orders"));
153
154        mgr.mark_snapshot_loaded("s1", 100);
155        let sub = mgr.get("s1").unwrap();
156        assert!(sub.snapshot_loaded);
157        assert_eq!(sub.last_lsn, 100);
158
159        mgr.advance_lsn("s1", 150);
160        assert_eq!(mgr.get("s1").unwrap().last_lsn, 150);
161
162        // Should not decrease.
163        mgr.advance_lsn("s1", 50);
164        assert_eq!(mgr.get("s1").unwrap().last_lsn, 150);
165    }
166
167    #[test]
168    fn active_ids_for_handshake() {
169        let mut mgr = ShapeManager::new();
170        mgr.subscribe(make_shape("s1", "orders"));
171        mgr.subscribe(make_shape("s2", "users"));
172
173        let mut ids = mgr.active_shape_ids();
174        ids.sort();
175        assert_eq!(ids, vec!["s1", "s2"]);
176    }
177
178    #[test]
179    fn active_definitions() {
180        let mut mgr = ShapeManager::new();
181        mgr.subscribe(make_shape("s1", "orders"));
182        mgr.subscribe(make_shape("s2", "users"));
183
184        let defs = mgr.active_definitions();
185        assert_eq!(defs.len(), 2);
186    }
187}