nodedb_lite/sync/outbound/
columnar.rs1use std::collections::HashMap;
22use std::sync::Arc;
23use std::sync::atomic::{AtomicU64, Ordering};
24
25use nodedb_types::Namespace;
26use nodedb_types::value::Value;
27use tokio::sync::Mutex;
28
29use super::durable_queue::DurableOutboundQueue;
30use crate::error::LiteError;
31use crate::storage::engine::StorageEngine;
32
33#[derive(
35 Debug,
36 Clone,
37 serde::Serialize,
38 serde::Deserialize,
39 zerompk::ToMessagePack,
40 zerompk::FromMessagePack,
41)]
42pub struct PendingColumnarBatch {
43 pub batch_id: u64,
45 pub collection: String,
47 pub rows: Vec<Vec<Value>>,
49 pub schema_bytes: Vec<u8>,
51 #[serde(default)]
55 pub seq: u64,
56}
57
58pub struct ColumnarOutbound<S: StorageEngine> {
65 queue: DurableOutboundQueue<S>,
66 ids: AtomicU64,
67 in_flight: Mutex<HashMap<u64, Vec<u8>>>,
70}
71
72impl<S: StorageEngine> ColumnarOutbound<S> {
73 pub async fn open(storage: Arc<S>) -> Result<Self, LiteError> {
75 Self::open_with_cap(storage, DurableOutboundQueue::<S>::DEFAULT_CAP).await
76 }
77
78 pub async fn open_with_cap(storage: Arc<S>, cap: usize) -> Result<Self, LiteError> {
80 let queue =
81 DurableOutboundQueue::open_with_cap(storage, Namespace::ColumnarPending, cap).await?;
82 Ok(Self {
83 queue,
84 ids: AtomicU64::new(1),
85 in_flight: Mutex::new(HashMap::new()),
86 })
87 }
88
89 pub async fn enqueue_row(
97 &self,
98 collection: &str,
99 row: Vec<Value>,
100 schema_bytes: Vec<u8>,
101 ) -> Result<(), LiteError> {
102 let batch_id = self.ids.fetch_add(1, Ordering::Relaxed);
103 let batch = PendingColumnarBatch {
104 batch_id,
105 collection: collection.to_string(),
106 rows: vec![row],
107 schema_bytes,
108 seq: 0,
109 };
110 let payload = zerompk::to_msgpack_vec(&batch).map_err(|e| LiteError::Serialization {
111 detail: format!("columnar outbound encode: {e}"),
112 })?;
113 self.queue.enqueue(&payload).await
114 }
115
116 pub async fn drain_batch(
125 &self,
126 limit: usize,
127 ) -> Result<Vec<(Vec<u8>, PendingColumnarBatch)>, LiteError> {
128 let in_flight = self.in_flight.lock().await;
129 let pairs = self.queue.drain_batch(limit).await?;
130 let mut out = Vec::with_capacity(pairs.len());
131 for (key, payload) in pairs {
132 if in_flight.values().any(|k| k == &key) {
133 continue;
134 }
135 let batch: PendingColumnarBatch =
136 zerompk::from_msgpack(&payload).map_err(|e| LiteError::Serialization {
137 detail: format!("columnar outbound decode: {e}"),
138 })?;
139 out.push((key, batch));
140 }
141 Ok(out)
142 }
143
144 pub async fn mark_in_flight(&self, batch_id: u64, durable_key: Vec<u8>) {
148 self.in_flight.lock().await.insert(batch_id, durable_key);
149 }
150
151 pub async fn ack_in_flight(&self, batch_id: u64) -> Option<Vec<u8>> {
156 self.in_flight.lock().await.remove(&batch_id)
157 }
158
159 pub async fn clear_in_flight(&self) {
164 self.in_flight.lock().await.clear();
165 }
166
167 pub async fn ack_keys(&self, keys: &[Vec<u8>]) -> Result<(), LiteError> {
169 self.queue.ack_keys(keys).await
170 }
171
172 pub async fn update_entry(
174 &self,
175 key: &[u8],
176 batch: &PendingColumnarBatch,
177 ) -> Result<(), LiteError> {
178 let payload = zerompk::to_msgpack_vec(batch).map_err(|e| LiteError::Serialization {
179 detail: format!("columnar outbound update encode: {e}"),
180 })?;
181 self.queue.update_entry(key, &payload).await
182 }
183
184 pub async fn len(&self) -> Result<u64, LiteError> {
186 self.queue.len().await
187 }
188
189 pub async fn is_empty(&self) -> Result<bool, LiteError> {
191 self.queue.is_empty().await
192 }
193}
194
195#[cfg(test)]
196mod tests {
197 use super::*;
198 use crate::storage::pagedb_storage::PagedbStorageMem;
199
200 async fn make_queue() -> ColumnarOutbound<PagedbStorageMem> {
201 let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
202 ColumnarOutbound::open(storage).await.unwrap()
203 }
204
205 #[tokio::test]
206 async fn enqueue_and_drain() {
207 let q = make_queue().await;
208 q.enqueue_row("metrics", vec![Value::Integer(1)], Vec::new())
209 .await
210 .unwrap();
211 q.enqueue_row("metrics", vec![Value::Integer(2)], Vec::new())
212 .await
213 .unwrap();
214
215 let pairs = q.drain_batch(usize::MAX).await.unwrap();
216 assert_eq!(pairs.len(), 2);
217 assert_eq!(pairs[0].1.collection, "metrics");
218 assert_eq!(pairs[0].1.rows[0][0], Value::Integer(1));
219 assert_eq!(pairs[1].1.rows[0][0], Value::Integer(2));
220 }
221
222 #[tokio::test]
223 async fn ack_keys_removes_entries() {
224 let q = make_queue().await;
225 q.enqueue_row("m", vec![Value::Integer(1)], Vec::new())
226 .await
227 .unwrap();
228 q.enqueue_row("m", vec![Value::Integer(2)], Vec::new())
229 .await
230 .unwrap();
231
232 let pairs = q.drain_batch(1).await.unwrap();
233 assert_eq!(pairs.len(), 1);
234 let keys: Vec<Vec<u8>> = pairs.into_iter().map(|(k, _)| k).collect();
235 q.ack_keys(&keys).await.unwrap();
236
237 assert_eq!(q.len().await.unwrap(), 1);
238 }
239
240 #[tokio::test]
241 async fn cap_returns_backpressure() {
242 let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
243 let q = ColumnarOutbound::open_with_cap(storage, 2).await.unwrap();
244 q.enqueue_row("m", vec![Value::Integer(1)], Vec::new())
245 .await
246 .unwrap();
247 q.enqueue_row("m", vec![Value::Integer(2)], Vec::new())
248 .await
249 .unwrap();
250 let err = q
251 .enqueue_row("m", vec![Value::Integer(3)], Vec::new())
252 .await
253 .unwrap_err();
254 assert!(matches!(err, LiteError::Backpressure { .. }));
255 }
256}