use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use nodedb_types::Namespace;
use nodedb_types::value::Value;
use tokio::sync::Mutex;
use super::durable_queue::DurableOutboundQueue;
use crate::error::LiteError;
use crate::storage::engine::StorageEngine;
#[derive(
Debug,
Clone,
serde::Serialize,
serde::Deserialize,
zerompk::ToMessagePack,
zerompk::FromMessagePack,
)]
pub struct PendingTimeseriesBatch {
pub batch_id: u64,
pub collection: String,
pub column_names: Vec<String>,
pub rows: Vec<Vec<Value>>,
#[serde(default)]
pub seq: u64,
}
pub struct TimeseriesOutbound<S: StorageEngine> {
queue: DurableOutboundQueue<S>,
ids: AtomicU64,
in_flight: Mutex<HashMap<u64, InFlightBatch>>,
}
#[derive(Debug, Clone)]
struct InFlightBatch {
batch_id: u64,
durable_key: Vec<u8>,
}
impl<S: StorageEngine> TimeseriesOutbound<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 queue =
DurableOutboundQueue::open_with_cap(storage, Namespace::TimeseriesPending, cap).await?;
Ok(Self {
queue,
ids: AtomicU64::new(1),
in_flight: Mutex::new(HashMap::new()),
})
}
pub async fn enqueue_row(
&self,
collection: &str,
column_names: Vec<String>,
row: Vec<Value>,
) -> Result<(), LiteError> {
let batch_id = self.ids.fetch_add(1, Ordering::Relaxed);
let batch = PendingTimeseriesBatch {
batch_id,
collection: collection.to_string(),
column_names,
rows: vec![row],
seq: 0,
};
let payload = zerompk::to_msgpack_vec(&batch).map_err(|e| LiteError::Serialization {
detail: format!("timeseries outbound encode: {e}"),
})?;
self.queue.enqueue(&payload).await
}
pub async fn drain_batch(
&self,
limit: usize,
) -> Result<Vec<(Vec<u8>, PendingTimeseriesBatch)>, LiteError> {
let in_flight = self.in_flight.lock().await;
let pairs = self.queue.drain_batch(limit).await?;
let mut out = Vec::with_capacity(pairs.len());
for (key, payload) in pairs {
if in_flight.values().any(|f| f.durable_key == key) {
continue;
}
let batch: PendingTimeseriesBatch =
zerompk::from_msgpack(&payload).map_err(|e| LiteError::Serialization {
detail: format!("timeseries outbound decode: {e}"),
})?;
out.push((key, batch));
}
Ok(out)
}
pub async fn mark_in_flight_by_seq(
&self,
stream_seq: u64,
batch_id: u64,
durable_key: Vec<u8>,
) {
self.in_flight.lock().await.insert(
stream_seq,
InFlightBatch {
batch_id,
durable_key,
},
);
}
pub async fn ack_in_flight_through_seq(&self, applied_seq: u64) -> Vec<Vec<u8>> {
let mut guard = self.in_flight.lock().await;
let to_ack: Vec<u64> = guard
.keys()
.copied()
.filter(|&seq| seq <= applied_seq)
.collect();
to_ack
.into_iter()
.filter_map(|seq| guard.remove(&seq))
.map(|f| f.durable_key)
.collect()
}
pub async fn ack_in_flight_by_batch_id(&self, batch_id: u64) -> Option<Vec<u8>> {
let mut guard = self.in_flight.lock().await;
let seq = guard
.iter()
.find(|(_, f)| f.batch_id == batch_id)
.map(|(seq, _)| *seq)?;
guard.remove(&seq).map(|f| f.durable_key)
}
pub async fn clear_in_flight(&self) {
self.in_flight.lock().await.clear();
}
pub async fn ack_keys(&self, keys: &[Vec<u8>]) -> Result<(), LiteError> {
self.queue.ack_keys(keys).await
}
pub async fn update_entry(
&self,
key: &[u8],
batch: &PendingTimeseriesBatch,
) -> Result<(), LiteError> {
let payload = zerompk::to_msgpack_vec(batch).map_err(|e| LiteError::Serialization {
detail: format!("timeseries outbound update encode: {e}"),
})?;
self.queue.update_entry(key, &payload).await
}
pub async fn len(&self) -> Result<u64, LiteError> {
self.queue.len().await
}
pub async fn is_empty(&self) -> Result<bool, LiteError> {
self.queue.is_empty().await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::pagedb_storage::PagedbStorageMem;
async fn make_queue() -> TimeseriesOutbound<PagedbStorageMem> {
let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
TimeseriesOutbound::open(storage).await.unwrap()
}
#[tokio::test]
async fn enqueue_and_drain() {
let q = make_queue().await;
q.enqueue_row(
"metrics",
vec!["time".into(), "value".into()],
vec![Value::Integer(1000), Value::Float(1.0)],
)
.await
.unwrap();
q.enqueue_row(
"metrics",
vec!["time".into(), "value".into()],
vec![Value::Integer(2000), Value::Float(2.0)],
)
.await
.unwrap();
let pairs = q.drain_batch(usize::MAX).await.unwrap();
assert_eq!(pairs.len(), 2);
assert_eq!(pairs[0].1.collection, "metrics");
assert_eq!(pairs[1].1.rows[0][0], Value::Integer(2000));
}
#[tokio::test]
async fn ack_keys_removes_entries() {
let q = make_queue().await;
q.enqueue_row(
"m",
vec!["time".into(), "value".into()],
vec![Value::Integer(1000), Value::Float(1.0)],
)
.await
.unwrap();
q.enqueue_row(
"m",
vec!["time".into(), "value".into()],
vec![Value::Integer(2000), Value::Float(2.0)],
)
.await
.unwrap();
let pairs = q.drain_batch(1).await.unwrap();
let keys: Vec<Vec<u8>> = pairs.into_iter().map(|(k, _)| k).collect();
q.ack_keys(&keys).await.unwrap();
assert_eq!(q.len().await.unwrap(), 1);
}
#[tokio::test]
async fn a_rejected_batch_is_retirable_even_though_the_frontier_never_reached_it() {
let q = make_queue().await;
q.mark_in_flight_by_seq(5, 100, b"key-5".to_vec()).await;
q.mark_in_flight_by_seq(6, 101, b"key-6".to_vec()).await;
assert!(
q.ack_in_flight_through_seq(4).await.is_empty(),
"the cumulative sweep must not retire a batch the frontier never reached"
);
assert_eq!(
q.ack_in_flight_by_batch_id(101).await,
Some(b"key-6".to_vec()),
"the rejected batch must be retirable by the id the ack echoed"
);
assert_eq!(
q.ack_in_flight_by_batch_id(101).await,
None,
"a duplicate rejection ack must be a no-op, not a second retirement"
);
assert_eq!(
q.ack_in_flight_through_seq(5).await,
vec![b"key-5".to_vec()],
"the untouched batch must still retire normally once the frontier advances"
);
}
#[tokio::test]
async fn cap_returns_backpressure() {
let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
let q = TimeseriesOutbound::open_with_cap(storage, 2).await.unwrap();
q.enqueue_row("m", vec!["t".into()], vec![Value::Integer(1)])
.await
.unwrap();
q.enqueue_row("m", vec!["t".into()], vec![Value::Integer(2)])
.await
.unwrap();
let err = q
.enqueue_row("m", vec!["t".into()], vec![Value::Integer(3)])
.await
.unwrap_err();
assert!(matches!(err, LiteError::Backpressure { .. }));
}
}