use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use anyhow::{Result, anyhow};
use zenoh::Session;
use crate::sub::SampleView;
#[derive(Debug, Clone, Copy)]
pub struct SeedPolicy {
pub history: bool,
pub storage: bool,
pub timeout: Duration,
}
impl Default for SeedPolicy {
fn default() -> Self {
SeedPolicy {
history: true,
storage: true,
timeout: Duration::from_secs(3),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
pub struct SeedCoverage {
pub history_replies: Option<usize>,
pub storage_replies: Option<usize>,
pub superseded: u64,
}
#[derive(Debug, Clone)]
pub enum SeedItem {
Sample(SampleView),
SeedComplete(SeedCoverage),
}
pub struct SeededSubscriber {
rx: tokio::sync::mpsc::UnboundedReceiver<SeedItem>,
_subscriber: zenoh::pubsub::Subscriber<()>,
task: tokio::task::JoinHandle<()>,
}
impl std::fmt::Debug for SeededSubscriber {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SeededSubscriber").finish_non_exhaustive()
}
}
impl Drop for SeededSubscriber {
fn drop(&mut self) {
self.task.abort();
}
}
impl SeededSubscriber {
pub async fn recv(&mut self) -> Option<SeedItem> {
self.rx.recv().await
}
}
pub(crate) struct Merge {
latest: Mutex<HashMap<String, Option<zenoh::time::Timestamp>>>,
superseded: std::sync::atomic::AtomicU64,
}
impl Merge {
pub(crate) fn new() -> Merge {
Merge {
latest: Mutex::new(HashMap::new()),
superseded: std::sync::atomic::AtomicU64::new(0),
}
}
pub(crate) fn superseded(&self) -> u64 {
self.superseded.load(std::sync::atomic::Ordering::Relaxed)
}
pub(crate) fn admit(&self, view: &SampleView) -> bool {
let mut latest = self.latest.lock().expect("merge lock");
let entry = latest.entry(view.key.clone()).or_insert(None);
let admit = match (&entry, &view.timestamp) {
(None, _) => true,
(Some(_), None) => false, (Some(prev), Some(ts)) => ts > prev,
};
if admit {
if view.timestamp.is_some() || entry.is_none() {
*entry = view.timestamp;
}
} else {
self.superseded
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
admit
}
}
pub(crate) fn view_of(sample: &zenoh::sample::Sample) -> SampleView {
SampleView {
key: sample.key_expr().as_str().to_string(),
payload: sample.payload().clone(),
encoding: sample.encoding().to_string(),
kind: sample.kind(),
timestamp: sample.timestamp().copied(),
attachment: sample.attachment().cloned(),
priority: sample.priority(),
congestion_control: sample.congestion_control(),
reliability: sample.reliability(),
express: sample.express(),
source: sample.source_info().map(|si| crate::sub::SampleSource {
zid: si.source_id().zid(),
eid: si.source_id().eid(),
sn: si.source_sn(),
}),
received: std::time::Instant::now(),
}
}
pub(crate) async fn seed_get(
session: &Session,
selector: &str,
timeout: Duration,
merge: &Merge,
mut deliver: impl FnMut(SampleView),
) -> usize {
let mut n = 0usize;
if let Ok(replies) = session
.get(selector)
.target(zenoh::query::QueryTarget::All)
.consolidation(zenoh::query::ConsolidationMode::None)
.accept_replies(zenoh::query::ReplyKeyExpr::Any)
.timeout(timeout)
.await
{
while let Ok(reply) = replies.recv_async().await {
let Ok(sample) = reply.result() else { continue };
n += 1;
let view = view_of(sample);
if merge.admit(&view) {
deliver(view);
}
}
}
n
}
pub(crate) fn cache_selector(selector: &str) -> String {
format!("{selector}/@adv/**")
}
pub async fn seed_subscribe(
session: &Session,
selector: &str,
policy: SeedPolicy,
) -> Result<SeededSubscriber> {
let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<SeedItem>();
let merge = Arc::new(Merge::new());
let subscriber = session
.declare_subscriber(selector.to_string())
.callback({
let tx = tx.clone();
let merge = Arc::clone(&merge);
move |sample| {
let view = view_of(&sample);
if merge.admit(&view) {
let _ = tx.send(SeedItem::Sample(view));
}
}
})
.await
.map_err(|e| anyhow!("seeded subscribe {selector}: {e}"))?;
let task = {
let session = session.clone();
let selector = selector.to_string();
let merge = Arc::clone(&merge);
tokio::spawn(async move {
let history = async {
if policy.history {
let sel = cache_selector(&selector);
Some(
seed_get(&session, &sel, policy.timeout, &merge, |view| {
let _ = tx.send(SeedItem::Sample(view));
})
.await,
)
} else {
None
}
};
let storage = async {
if policy.storage {
Some(
seed_get(&session, &selector, policy.timeout, &merge, |view| {
let _ = tx.send(SeedItem::Sample(view));
})
.await,
)
} else {
None
}
};
let (history_replies, storage_replies) = tokio::join!(history, storage);
let _ = tx.send(SeedItem::SeedComplete(SeedCoverage {
history_replies,
storage_replies,
superseded: merge.superseded(),
}));
})
};
Ok(SeededSubscriber {
rx,
_subscriber: subscriber,
task,
})
}