use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use nodedb_types::Namespace;
use tokio::sync::Mutex;
use super::durable_queue::DurableOutboundQueue;
use super::queue::PendingQueue;
use crate::error::LiteError;
use crate::storage::engine::StorageEngine;
const STAGING_CAP: usize = 4_096;
#[derive(
Debug,
Clone,
serde::Serialize,
serde::Deserialize,
zerompk::ToMessagePack,
zerompk::FromMessagePack,
)]
pub struct PendingSpatialInsert {
pub batch_id: u64,
pub collection: String,
pub field: String,
pub doc_id: String,
pub geometry_bytes: Vec<u8>,
#[serde(default)]
pub seq: u64,
}
#[derive(
Debug,
Clone,
serde::Serialize,
serde::Deserialize,
zerompk::ToMessagePack,
zerompk::FromMessagePack,
)]
pub struct PendingSpatialDelete {
pub batch_id: u64,
pub collection: String,
pub field: String,
pub doc_id: String,
#[serde(default)]
pub seq: u64,
}
pub struct SpatialOutbound<S: StorageEngine> {
staging_inserts: PendingQueue<PendingSpatialInsert>,
staging_deletes: PendingQueue<PendingSpatialDelete>,
durable_inserts: DurableOutboundQueue<S>,
durable_deletes: DurableOutboundQueue<S>,
ids: AtomicU64,
in_flight_inserts: Mutex<HashMap<u64, Vec<u8>>>,
in_flight_deletes_map: Mutex<HashMap<u64, Vec<u8>>>,
}
impl<S: StorageEngine> SpatialOutbound<S> {
pub async fn open(storage: Arc<S>) -> Result<Self, LiteError> {
Self::open_with_cap(storage, DurableOutboundQueue::<S>::DEFAULT_CAP).await
}
pub async fn open_with_cap(storage: Arc<S>, cap: usize) -> Result<Self, LiteError> {
let durable_inserts = DurableOutboundQueue::open_with_cap(
Arc::clone(&storage),
Namespace::SpatialInsertPending,
cap,
)
.await?;
let durable_deletes = DurableOutboundQueue::open_with_cap(
Arc::clone(&storage),
Namespace::SpatialDeletePending,
cap,
)
.await?;
Ok(Self {
staging_inserts: PendingQueue::new(),
staging_deletes: PendingQueue::new(),
durable_inserts,
durable_deletes,
ids: AtomicU64::new(1),
in_flight_inserts: Mutex::new(HashMap::new()),
in_flight_deletes_map: Mutex::new(HashMap::new()),
})
}
pub fn stage_insert(
&self,
collection: &str,
field: &str,
doc_id: &str,
geometry: &nodedb_types::geometry::Geometry,
) {
if self.staging_inserts.len() >= STAGING_CAP {
tracing::warn!(
collection,
field,
doc_id,
"spatial_outbound: staging buffer full; dropping insert sync entry \
(will re-sync on next catch-up)"
);
return;
}
let geometry_bytes = match zerompk::to_msgpack_vec(geometry) {
Ok(b) => b,
Err(e) => {
tracing::warn!(
collection,
field,
doc_id,
error = %e,
"spatial_outbound: failed to serialise geometry; skipping sync enqueue"
);
return;
}
};
let batch_id = self.ids.fetch_add(1, Ordering::Relaxed);
self.staging_inserts.push(PendingSpatialInsert {
batch_id,
collection: collection.to_string(),
field: field.to_string(),
doc_id: doc_id.to_string(),
geometry_bytes,
seq: 0,
});
}
pub fn stage_delete(&self, collection: &str, field: &str, doc_id: &str) {
if self.staging_deletes.len() >= STAGING_CAP {
tracing::warn!(
collection,
field,
doc_id,
"spatial_outbound: staging buffer full; dropping delete sync entry \
(will re-sync on next catch-up)"
);
return;
}
let batch_id = self.ids.fetch_add(1, Ordering::Relaxed);
self.staging_deletes.push(PendingSpatialDelete {
batch_id,
collection: collection.to_string(),
field: field.to_string(),
doc_id: doc_id.to_string(),
seq: 0,
});
}
pub async fn flush_staging(&self) -> Result<(), LiteError> {
let staged_inserts = self.staging_inserts.drain();
let mut failed_inserts: Vec<PendingSpatialInsert> = Vec::new();
let mut backpressure: Option<LiteError> = None;
for entry in staged_inserts {
if backpressure.is_some() {
failed_inserts.push(entry);
continue;
}
let payload =
zerompk::to_msgpack_vec(&entry).map_err(|e| LiteError::Serialization {
detail: format!("spatial insert outbound encode: {e}"),
})?;
match self.durable_inserts.enqueue(&payload).await {
Ok(()) => {}
Err(e @ LiteError::Backpressure { .. }) => {
failed_inserts.push(entry);
backpressure = Some(e);
}
Err(e) => return Err(e),
}
}
for entry in failed_inserts.into_iter().rev() {
self.staging_inserts.requeue(entry);
}
let staged_deletes = self.staging_deletes.drain();
let mut failed_deletes: Vec<PendingSpatialDelete> = Vec::new();
for entry in staged_deletes {
if backpressure.is_some() {
failed_deletes.push(entry);
continue;
}
let payload =
zerompk::to_msgpack_vec(&entry).map_err(|e| LiteError::Serialization {
detail: format!("spatial delete outbound encode: {e}"),
})?;
match self.durable_deletes.enqueue(&payload).await {
Ok(()) => {}
Err(e @ LiteError::Backpressure { .. }) => {
failed_deletes.push(entry);
backpressure = Some(e);
}
Err(e) => return Err(e),
}
}
for entry in failed_deletes.into_iter().rev() {
self.staging_deletes.requeue(entry);
}
match backpressure {
Some(e) => Err(e),
None => Ok(()),
}
}
pub async fn drain_inserts(
&self,
limit: usize,
) -> Result<Vec<(Vec<u8>, PendingSpatialInsert)>, LiteError> {
if let Err(e) = self.flush_staging().await {
tracing::debug!(error = %e, "spatial_outbound: staging spill backpressured");
}
let in_flight = self.in_flight_inserts.lock().await;
let pairs = self.durable_inserts.drain_batch(limit).await?;
let mut out = Vec::with_capacity(pairs.len());
for (key, payload) in pairs {
if in_flight.values().any(|k| k == &key) {
continue;
}
let entry: PendingSpatialInsert =
zerompk::from_msgpack(&payload).map_err(|e| LiteError::Serialization {
detail: format!("spatial insert outbound decode: {e}"),
})?;
out.push((key, entry));
}
Ok(out)
}
pub async fn drain_deletes(
&self,
limit: usize,
) -> Result<Vec<(Vec<u8>, PendingSpatialDelete)>, LiteError> {
if let Err(e) = self.flush_staging().await {
tracing::debug!(error = %e, "spatial_outbound: staging spill backpressured");
}
let in_flight = self.in_flight_deletes_map.lock().await;
let pairs = self.durable_deletes.drain_batch(limit).await?;
let mut out = Vec::with_capacity(pairs.len());
for (key, payload) in pairs {
if in_flight.values().any(|k| k == &key) {
continue;
}
let entry: PendingSpatialDelete =
zerompk::from_msgpack(&payload).map_err(|e| LiteError::Serialization {
detail: format!("spatial delete outbound decode: {e}"),
})?;
out.push((key, entry));
}
Ok(out)
}
pub async fn mark_insert_in_flight(&self, batch_id: u64, durable_key: Vec<u8>) {
self.in_flight_inserts
.lock()
.await
.insert(batch_id, durable_key);
}
pub async fn ack_insert_in_flight(&self, batch_id: u64) -> Option<Vec<u8>> {
self.in_flight_inserts.lock().await.remove(&batch_id)
}
pub async fn mark_delete_in_flight(&self, batch_id: u64, durable_key: Vec<u8>) {
self.in_flight_deletes_map
.lock()
.await
.insert(batch_id, durable_key);
}
pub async fn ack_delete_in_flight(&self, batch_id: u64) -> Option<Vec<u8>> {
self.in_flight_deletes_map.lock().await.remove(&batch_id)
}
pub async fn clear_in_flight(&self) {
self.in_flight_inserts.lock().await.clear();
self.in_flight_deletes_map.lock().await.clear();
}
pub async fn ack_insert_keys(&self, keys: &[Vec<u8>]) -> Result<(), LiteError> {
self.durable_inserts.ack_keys(keys).await
}
pub async fn ack_delete_keys(&self, keys: &[Vec<u8>]) -> Result<(), LiteError> {
self.durable_deletes.ack_keys(keys).await
}
pub async fn update_insert_entry(
&self,
key: &[u8],
insert: &PendingSpatialInsert,
) -> Result<(), LiteError> {
let payload = zerompk::to_msgpack_vec(insert).map_err(|e| LiteError::Serialization {
detail: format!("spatial insert outbound update encode: {e}"),
})?;
self.durable_inserts.update_entry(key, &payload).await
}
pub async fn update_delete_entry(
&self,
key: &[u8],
delete: &PendingSpatialDelete,
) -> Result<(), LiteError> {
let payload = zerompk::to_msgpack_vec(delete).map_err(|e| LiteError::Serialization {
detail: format!("spatial delete outbound update encode: {e}"),
})?;
self.durable_deletes.update_entry(key, &payload).await
}
pub async fn pending_insert_count(&self) -> Result<u64, LiteError> {
self.durable_inserts.len().await
}
pub async fn pending_delete_count(&self) -> Result<u64, LiteError> {
self.durable_deletes.len().await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::pagedb_storage::PagedbStorageMem;
fn point_geometry() -> nodedb_types::geometry::Geometry {
nodedb_types::geometry::Geometry::point(1.0, 2.0)
}
async fn make_queue() -> SpatialOutbound<PagedbStorageMem> {
let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
SpatialOutbound::open(storage).await.unwrap()
}
#[tokio::test]
async fn stage_and_flush_inserts() {
let q = make_queue().await;
q.stage_insert("places", "loc", "doc1", &point_geometry());
q.stage_insert("places", "loc", "doc2", &point_geometry());
assert_eq!(q.pending_insert_count().await.unwrap(), 0);
q.flush_staging().await.unwrap();
let pairs = q.drain_inserts(usize::MAX).await.unwrap();
assert_eq!(pairs.len(), 2);
assert_eq!(pairs[0].1.doc_id, "doc1");
assert_eq!(pairs[1].1.doc_id, "doc2");
}
#[tokio::test]
async fn stage_and_flush_deletes() {
let q = make_queue().await;
q.stage_delete("places", "loc", "doc1");
q.flush_staging().await.unwrap();
let pairs = q.drain_deletes(usize::MAX).await.unwrap();
assert_eq!(pairs.len(), 1);
assert_eq!(pairs[0].1.doc_id, "doc1");
}
#[tokio::test]
async fn ack_insert_keys_removes_entries() {
let q = make_queue().await;
q.stage_insert("places", "loc", "doc1", &point_geometry());
q.stage_insert("places", "loc", "doc2", &point_geometry());
q.flush_staging().await.unwrap();
let pairs = q.drain_inserts(1).await.unwrap();
assert_eq!(pairs.len(), 1);
let keys: Vec<Vec<u8>> = pairs.into_iter().map(|(k, _)| k).collect();
q.ack_insert_keys(&keys).await.unwrap();
assert_eq!(q.pending_insert_count().await.unwrap(), 1);
}
#[tokio::test]
async fn ack_delete_keys_removes_entries() {
let q = make_queue().await;
q.stage_delete("places", "loc", "doc1");
q.stage_delete("places", "loc", "doc2");
q.flush_staging().await.unwrap();
let pairs = q.drain_deletes(1).await.unwrap();
assert_eq!(pairs.len(), 1);
let keys: Vec<Vec<u8>> = pairs.into_iter().map(|(k, _)| k).collect();
q.ack_delete_keys(&keys).await.unwrap();
assert_eq!(q.pending_delete_count().await.unwrap(), 1);
}
#[tokio::test]
async fn geometry_bytes_round_trip() {
let geom = point_geometry();
let q = make_queue().await;
q.stage_insert("places", "loc", "doc1", &geom);
q.flush_staging().await.unwrap();
let pairs = q.drain_inserts(usize::MAX).await.unwrap();
assert!(!pairs[0].1.geometry_bytes.is_empty());
let restored: nodedb_types::geometry::Geometry =
zerompk::from_msgpack(&pairs[0].1.geometry_bytes)
.expect("geometry round-trip deserialise");
assert_eq!(geom, restored);
}
#[tokio::test]
async fn durable_cap_returns_backpressure() {
let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
let q = SpatialOutbound::open_with_cap(storage, 2).await.unwrap();
q.stage_insert("places", "loc", "a", &point_geometry());
q.stage_insert("places", "loc", "b", &point_geometry());
q.stage_insert("places", "loc", "c", &point_geometry());
let err = q.flush_staging().await.unwrap_err();
assert!(matches!(err, LiteError::Backpressure { .. }));
assert_eq!(q.staging_inserts.len(), 1);
}
#[tokio::test]
async fn survives_reload() {
let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
{
let q = SpatialOutbound::open(Arc::clone(&storage)).await.unwrap();
q.stage_insert("places", "loc", "v1", &point_geometry());
q.flush_staging().await.unwrap();
}
let q = SpatialOutbound::open(Arc::clone(&storage)).await.unwrap();
assert_eq!(q.pending_insert_count().await.unwrap(), 1);
q.stage_insert("places", "loc", "v2", &point_geometry());
q.flush_staging().await.unwrap();
assert_eq!(q.pending_insert_count().await.unwrap(), 2);
}
}