Skip to main content

nodedb_lite/sync/outbound/
vector.rs

1//! Vector insert / delete outbound queue for Lite sync.
2//!
3//! When `NodeDbLite::vector_insert_impl` or `vector_delete_impl` is called,
4//! the operation is durably enqueued here. The sync transport drains this queue
5//! on every tick and sends `VectorInsert` (0xA2) / `VectorDelete` (0xA4) wire
6//! frames to Origin.
7//!
8//! # Durability
9//!
10//! Inserts are backed by [`DurableOutboundQueue`] in
11//! [`Namespace::VectorInsertPending`]; deletes use
12//! [`Namespace::VectorDeletePending`]. Each entry is persisted before the
13//! enqueue call returns. Entries survive process restarts; un-acked entries are
14//! re-drained on reconnect (at-least-once delivery).
15//!
16//! # Backpressure
17//!
18//! `enqueue_insert` / `enqueue_delete` return [`LiteError::Backpressure`] when
19//! their respective queue is at cap, propagating to the write caller so the
20//! device can pause until the sync transport drains the backlog.
21//!
22//! # Design: two queues, two namespaces
23//!
24//! Inserts and deletes are kept in separate queues so the push transport can
25//! drain and ack them independently. A shared queue with a 1-byte op tag would
26//! require the push path to branch on type after decoding; two queues keep each
27//! drain path homogeneous and mirror the columnar/timeseries pattern exactly.
28
29use std::collections::HashMap;
30use std::sync::Arc;
31use std::sync::atomic::{AtomicU64, Ordering};
32
33use nodedb_types::Namespace;
34use tokio::sync::Mutex;
35
36use super::durable_queue::DurableOutboundQueue;
37use crate::error::LiteError;
38use crate::storage::engine::StorageEngine;
39
40/// A single pending vector insert awaiting sync to Origin.
41#[derive(
42    Debug,
43    Clone,
44    serde::Serialize,
45    serde::Deserialize,
46    zerompk::ToMessagePack,
47    zerompk::FromMessagePack,
48)]
49pub struct PendingVectorInsert {
50    /// Monotonic batch ID for ACK correlation.
51    pub batch_id: u64,
52    /// Collection name.
53    pub collection: String,
54    /// Document ID.
55    pub id: String,
56    /// Raw embedding values.
57    pub vector: Vec<f32>,
58    /// Embedding dimension.
59    pub dim: usize,
60    /// Named vector field; empty string = default (no named field).
61    pub field_name: String,
62    /// Stable idempotent-producer seq for this entry. 0 = not yet assigned;
63    /// assigned at first drain and persisted so re-sends after reconnect reuse
64    /// the same seq (Origin dedups instead of double-applying).
65    #[serde(default)]
66    pub seq: u64,
67}
68
69/// A single pending vector delete awaiting sync to Origin.
70#[derive(
71    Debug,
72    Clone,
73    serde::Serialize,
74    serde::Deserialize,
75    zerompk::ToMessagePack,
76    zerompk::FromMessagePack,
77)]
78pub struct PendingVectorDelete {
79    /// Monotonic batch ID for ACK correlation.
80    pub batch_id: u64,
81    /// Collection name.
82    pub collection: String,
83    /// Document ID.
84    pub id: String,
85    /// Named vector field; empty string = default (no named field).
86    pub field_name: String,
87    /// Stable idempotent-producer seq for this entry. 0 = not yet assigned;
88    /// assigned at first drain and persisted so re-sends after reconnect reuse
89    /// the same seq (Origin dedups instead of double-applying).
90    #[serde(default)]
91    pub seq: u64,
92}
93
94/// Durable outbound queue for vector insert and delete sync.
95///
96/// Held as `Arc<VectorOutbound<S>>` by `NodeDbLite` and shared with the sync
97/// transport. The inner storage is accessed only for `enqueue_insert` /
98/// `enqueue_delete` (from the write path) and `drain_inserts` /
99/// `drain_deletes` / `ack_*` (from the async sync transport task).
100pub struct VectorOutbound<S: StorageEngine> {
101    inserts: DurableOutboundQueue<S>,
102    deletes: DurableOutboundQueue<S>,
103    ids: AtomicU64,
104    /// batch_id → durable_key for in-flight insert entries.
105    in_flight_inserts: Mutex<HashMap<u64, Vec<u8>>>,
106    /// batch_id → durable_key for in-flight delete entries.
107    in_flight_deletes: Mutex<HashMap<u64, Vec<u8>>>,
108}
109
110impl<S: StorageEngine> VectorOutbound<S> {
111    /// Open durable queues backed by [`Namespace::VectorInsertPending`] and
112    /// [`Namespace::VectorDeletePending`].
113    pub async fn open(storage: Arc<S>) -> Result<Self, LiteError> {
114        Self::open_with_cap(storage, DurableOutboundQueue::<S>::DEFAULT_CAP).await
115    }
116
117    /// Open with a custom cap (useful for tests with small limits).
118    pub async fn open_with_cap(storage: Arc<S>, cap: usize) -> Result<Self, LiteError> {
119        let inserts = DurableOutboundQueue::open_with_cap(
120            Arc::clone(&storage),
121            Namespace::VectorInsertPending,
122            cap,
123        )
124        .await?;
125        let deletes = DurableOutboundQueue::open_with_cap(
126            Arc::clone(&storage),
127            Namespace::VectorDeletePending,
128            cap,
129        )
130        .await?;
131        Ok(Self {
132            inserts,
133            deletes,
134            ids: AtomicU64::new(1),
135            in_flight_inserts: Mutex::new(HashMap::new()),
136            in_flight_deletes: Mutex::new(HashMap::new()),
137        })
138    }
139
140    /// Durably enqueue a vector insert.
141    ///
142    /// Returns [`LiteError::Backpressure`] when the insert queue is at cap.
143    pub async fn enqueue_insert(
144        &self,
145        collection: &str,
146        id: &str,
147        vector: Vec<f32>,
148        dim: usize,
149        field_name: &str,
150    ) -> Result<(), LiteError> {
151        let batch_id = self.ids.fetch_add(1, Ordering::Relaxed);
152        let entry = PendingVectorInsert {
153            batch_id,
154            collection: collection.to_string(),
155            id: id.to_string(),
156            vector,
157            dim,
158            field_name: field_name.to_string(),
159            seq: 0,
160        };
161        let payload = zerompk::to_msgpack_vec(&entry).map_err(|e| LiteError::Serialization {
162            detail: format!("vector insert outbound encode: {e}"),
163        })?;
164        self.inserts.enqueue(&payload).await
165    }
166
167    /// Durably enqueue a vector delete.
168    ///
169    /// Returns [`LiteError::Backpressure`] when the delete queue is at cap.
170    pub async fn enqueue_delete(
171        &self,
172        collection: &str,
173        id: &str,
174        field_name: &str,
175    ) -> Result<(), LiteError> {
176        let batch_id = self.ids.fetch_add(1, Ordering::Relaxed);
177        let entry = PendingVectorDelete {
178            batch_id,
179            collection: collection.to_string(),
180            id: id.to_string(),
181            field_name: field_name.to_string(),
182            seq: 0,
183        };
184        let payload = zerompk::to_msgpack_vec(&entry).map_err(|e| LiteError::Serialization {
185            detail: format!("vector delete outbound encode: {e}"),
186        })?;
187        self.deletes.enqueue(&payload).await
188    }
189
190    /// Drain up to `limit` pending inserts in FIFO order, skipping in-flight entries.
191    ///
192    /// Returns `(durable_key, insert)` pairs. On send success, call
193    /// [`mark_insert_in_flight`]. The durable entry is deleted only on Origin
194    /// ack via [`ack_insert_in_flight`].
195    ///
196    /// Does **not** remove entries from storage.
197    pub async fn drain_inserts(
198        &self,
199        limit: usize,
200    ) -> Result<Vec<(Vec<u8>, PendingVectorInsert)>, LiteError> {
201        let in_flight = self.in_flight_inserts.lock().await;
202        let pairs = self.inserts.drain_batch(limit).await?;
203        let mut out = Vec::with_capacity(pairs.len());
204        for (key, payload) in pairs {
205            if in_flight.values().any(|k| k == &key) {
206                continue;
207            }
208            let entry: PendingVectorInsert =
209                zerompk::from_msgpack(&payload).map_err(|e| LiteError::Serialization {
210                    detail: format!("vector insert outbound decode: {e}"),
211                })?;
212            out.push((key, entry));
213        }
214        Ok(out)
215    }
216
217    /// Drain up to `limit` pending deletes in FIFO order, skipping in-flight entries.
218    ///
219    /// Returns `(durable_key, delete)` pairs. On send success, call
220    /// [`mark_delete_in_flight`]. The durable entry is deleted only on Origin
221    /// ack via [`ack_delete_in_flight`].
222    ///
223    /// Does **not** remove entries from storage.
224    pub async fn drain_deletes(
225        &self,
226        limit: usize,
227    ) -> Result<Vec<(Vec<u8>, PendingVectorDelete)>, LiteError> {
228        let in_flight = self.in_flight_deletes.lock().await;
229        let pairs = self.deletes.drain_batch(limit).await?;
230        let mut out = Vec::with_capacity(pairs.len());
231        for (key, payload) in pairs {
232            if in_flight.values().any(|k| k == &key) {
233                continue;
234            }
235            let entry: PendingVectorDelete =
236                zerompk::from_msgpack(&payload).map_err(|e| LiteError::Serialization {
237                    detail: format!("vector delete outbound decode: {e}"),
238                })?;
239            out.push((key, entry));
240        }
241        Ok(out)
242    }
243
244    /// Record that an insert has been sent to Origin and is awaiting its ack.
245    pub async fn mark_insert_in_flight(&self, batch_id: u64, durable_key: Vec<u8>) {
246        self.in_flight_inserts
247            .lock()
248            .await
249            .insert(batch_id, durable_key);
250    }
251
252    /// Remove the in-flight record for an insert and return its durable key.
253    pub async fn ack_insert_in_flight(&self, batch_id: u64) -> Option<Vec<u8>> {
254        self.in_flight_inserts.lock().await.remove(&batch_id)
255    }
256
257    /// Record that a delete has been sent to Origin and is awaiting its ack.
258    pub async fn mark_delete_in_flight(&self, batch_id: u64, durable_key: Vec<u8>) {
259        self.in_flight_deletes
260            .lock()
261            .await
262            .insert(batch_id, durable_key);
263    }
264
265    /// Remove the in-flight record for a delete and return its durable key.
266    pub async fn ack_delete_in_flight(&self, batch_id: u64) -> Option<Vec<u8>> {
267        self.in_flight_deletes.lock().await.remove(&batch_id)
268    }
269
270    /// Clear all in-flight records on reconnect so entries are re-drained.
271    pub async fn clear_in_flight(&self) {
272        self.in_flight_inserts.lock().await.clear();
273        self.in_flight_deletes.lock().await.clear();
274    }
275
276    /// Delete the durable insert entries identified by `keys` (Origin ack path).
277    pub async fn ack_insert_keys(&self, keys: &[Vec<u8>]) -> Result<(), LiteError> {
278        self.inserts.ack_keys(keys).await
279    }
280
281    /// Delete the durable delete entries identified by `keys` (Origin ack path).
282    pub async fn ack_delete_keys(&self, keys: &[Vec<u8>]) -> Result<(), LiteError> {
283        self.deletes.ack_keys(keys).await
284    }
285
286    /// Update the durable insert payload for `key` with the new seq.
287    pub async fn update_insert_entry(
288        &self,
289        key: &[u8],
290        insert: &PendingVectorInsert,
291    ) -> Result<(), LiteError> {
292        let payload = zerompk::to_msgpack_vec(insert).map_err(|e| LiteError::Serialization {
293            detail: format!("vector insert outbound update encode: {e}"),
294        })?;
295        self.inserts.update_entry(key, &payload).await
296    }
297
298    /// Update the durable delete payload for `key` with the new seq.
299    pub async fn update_delete_entry(
300        &self,
301        key: &[u8],
302        delete: &PendingVectorDelete,
303    ) -> Result<(), LiteError> {
304        let payload = zerompk::to_msgpack_vec(delete).map_err(|e| LiteError::Serialization {
305            detail: format!("vector delete outbound update encode: {e}"),
306        })?;
307        self.deletes.update_entry(key, &payload).await
308    }
309
310    /// Number of pending insert entries in durable storage.
311    pub async fn pending_insert_count(&self) -> Result<u64, LiteError> {
312        self.inserts.len().await
313    }
314
315    /// Number of pending delete entries in durable storage.
316    pub async fn pending_delete_count(&self) -> Result<u64, LiteError> {
317        self.deletes.len().await
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324    use crate::storage::pagedb_storage::PagedbStorageMem;
325
326    async fn make_queue() -> VectorOutbound<PagedbStorageMem> {
327        let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
328        VectorOutbound::open(storage).await.unwrap()
329    }
330
331    #[tokio::test]
332    async fn enqueue_and_drain_inserts() {
333        let q = make_queue().await;
334        q.enqueue_insert("vecs", "v1", vec![1.0, 0.0], 2, "")
335            .await
336            .unwrap();
337        q.enqueue_insert("vecs", "v2", vec![0.0, 1.0], 2, "")
338            .await
339            .unwrap();
340
341        let pairs = q.drain_inserts(usize::MAX).await.unwrap();
342        assert_eq!(pairs.len(), 2);
343        assert_eq!(pairs[0].1.id, "v1");
344        assert_eq!(pairs[1].1.id, "v2");
345    }
346
347    #[tokio::test]
348    async fn enqueue_and_drain_deletes() {
349        let q = make_queue().await;
350        q.enqueue_delete("vecs", "v1", "").await.unwrap();
351
352        let pairs = q.drain_deletes(usize::MAX).await.unwrap();
353        assert_eq!(pairs.len(), 1);
354        assert_eq!(pairs[0].1.id, "v1");
355    }
356
357    #[tokio::test]
358    async fn ack_insert_keys_removes_entries() {
359        let q = make_queue().await;
360        q.enqueue_insert("vecs", "v1", vec![1.0], 1, "")
361            .await
362            .unwrap();
363        q.enqueue_insert("vecs", "v2", vec![2.0], 1, "")
364            .await
365            .unwrap();
366
367        let pairs = q.drain_inserts(1).await.unwrap();
368        assert_eq!(pairs.len(), 1);
369        let keys: Vec<Vec<u8>> = pairs.into_iter().map(|(k, _)| k).collect();
370        q.ack_insert_keys(&keys).await.unwrap();
371
372        assert_eq!(q.pending_insert_count().await.unwrap(), 1);
373    }
374
375    #[tokio::test]
376    async fn ack_delete_keys_removes_entries() {
377        let q = make_queue().await;
378        q.enqueue_delete("vecs", "v1", "").await.unwrap();
379        q.enqueue_delete("vecs", "v2", "").await.unwrap();
380
381        let pairs = q.drain_deletes(1).await.unwrap();
382        assert_eq!(pairs.len(), 1);
383        let keys: Vec<Vec<u8>> = pairs.into_iter().map(|(k, _)| k).collect();
384        q.ack_delete_keys(&keys).await.unwrap();
385
386        assert_eq!(q.pending_delete_count().await.unwrap(), 1);
387    }
388
389    #[tokio::test]
390    async fn insert_cap_returns_backpressure() {
391        let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
392        let q = VectorOutbound::open_with_cap(storage, 2).await.unwrap();
393        q.enqueue_insert("vecs", "a", vec![1.0], 1, "")
394            .await
395            .unwrap();
396        q.enqueue_insert("vecs", "b", vec![2.0], 1, "")
397            .await
398            .unwrap();
399        let err = q
400            .enqueue_insert("vecs", "c", vec![3.0], 1, "")
401            .await
402            .unwrap_err();
403        assert!(matches!(err, LiteError::Backpressure { .. }));
404    }
405
406    #[tokio::test]
407    async fn delete_cap_returns_backpressure() {
408        let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
409        let q = VectorOutbound::open_with_cap(storage, 2).await.unwrap();
410        q.enqueue_delete("vecs", "a", "").await.unwrap();
411        q.enqueue_delete("vecs", "b", "").await.unwrap();
412        let err = q.enqueue_delete("vecs", "c", "").await.unwrap_err();
413        assert!(matches!(err, LiteError::Backpressure { .. }));
414    }
415
416    #[tokio::test]
417    async fn batch_ids_are_unique() {
418        let q = make_queue().await;
419        q.enqueue_insert("vecs", "a", vec![1.0], 1, "")
420            .await
421            .unwrap();
422        q.enqueue_delete("vecs", "b", "").await.unwrap();
423        q.enqueue_insert("vecs", "c", vec![2.0], 1, "")
424            .await
425            .unwrap();
426
427        let inserts = q.drain_inserts(usize::MAX).await.unwrap();
428        let deletes = q.drain_deletes(usize::MAX).await.unwrap();
429
430        let mut all_ids: Vec<u64> = inserts.iter().map(|(_, e)| e.batch_id).collect();
431        all_ids.extend(deletes.iter().map(|(_, e)| e.batch_id));
432        let mut sorted = all_ids.clone();
433        sorted.sort_unstable();
434        sorted.dedup();
435        assert_eq!(sorted.len(), all_ids.len(), "batch_ids must be unique");
436        assert!(all_ids.iter().all(|&id| id > 0));
437    }
438
439    #[tokio::test]
440    async fn survives_reload() {
441        let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
442        {
443            let q = VectorOutbound::open(Arc::clone(&storage)).await.unwrap();
444            q.enqueue_insert("vecs", "v1", vec![1.0], 1, "")
445                .await
446                .unwrap();
447        }
448        let q = VectorOutbound::open(Arc::clone(&storage)).await.unwrap();
449        assert_eq!(q.pending_insert_count().await.unwrap(), 1);
450        // new entry must not overwrite the existing key
451        q.enqueue_insert("vecs", "v2", vec![2.0], 1, "")
452            .await
453            .unwrap();
454        assert_eq!(q.pending_insert_count().await.unwrap(), 2);
455    }
456}