1use 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#[derive(
42 Debug,
43 Clone,
44 serde::Serialize,
45 serde::Deserialize,
46 zerompk::ToMessagePack,
47 zerompk::FromMessagePack,
48)]
49pub struct PendingVectorInsert {
50 pub batch_id: u64,
52 pub collection: String,
54 pub id: String,
56 pub vector: Vec<f32>,
58 pub dim: usize,
60 pub field_name: String,
62 #[serde(default)]
66 pub seq: u64,
67}
68
69#[derive(
71 Debug,
72 Clone,
73 serde::Serialize,
74 serde::Deserialize,
75 zerompk::ToMessagePack,
76 zerompk::FromMessagePack,
77)]
78pub struct PendingVectorDelete {
79 pub batch_id: u64,
81 pub collection: String,
83 pub id: String,
85 pub field_name: String,
87 #[serde(default)]
91 pub seq: u64,
92}
93
94pub struct VectorOutbound<S: StorageEngine> {
101 inserts: DurableOutboundQueue<S>,
102 deletes: DurableOutboundQueue<S>,
103 ids: AtomicU64,
104 in_flight_inserts: Mutex<HashMap<u64, Vec<u8>>>,
106 in_flight_deletes: Mutex<HashMap<u64, Vec<u8>>>,
108}
109
110impl<S: StorageEngine> VectorOutbound<S> {
111 pub async fn open(storage: Arc<S>) -> Result<Self, LiteError> {
114 Self::open_with_cap(storage, DurableOutboundQueue::<S>::DEFAULT_CAP).await
115 }
116
117 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 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 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 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 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 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 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 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 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 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 pub async fn ack_insert_keys(&self, keys: &[Vec<u8>]) -> Result<(), LiteError> {
278 self.inserts.ack_keys(keys).await
279 }
280
281 pub async fn ack_delete_keys(&self, keys: &[Vec<u8>]) -> Result<(), LiteError> {
283 self.deletes.ack_keys(keys).await
284 }
285
286 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 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 pub async fn pending_insert_count(&self) -> Result<u64, LiteError> {
312 self.inserts.len().await
313 }
314
315 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 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}