use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use crate::Result;
use zenoh::Session;
use crate::bus::monitor::SampleView;
use crate::report::SeedCoverage;
#[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)]
pub enum SeedItem {
Sample(SampleView),
Dropped(u64),
SeedComplete(SeedCoverage),
}
const SEED_CAPACITY: usize = 1024;
#[derive(Clone)]
struct SeedSender {
tx: tokio::sync::mpsc::Sender<SeedItem>,
dropped: Arc<std::sync::atomic::AtomicU64>,
}
impl SeedSender {
fn send_sample(&self, view: SampleView) {
use tokio::sync::mpsc::error::TrySendError;
match self.tx.try_send(SeedItem::Sample(view)) {
Ok(()) => {}
Err(TrySendError::Full(_)) => {
self.dropped
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
Err(TrySendError::Closed(_)) => {}
}
}
async fn send_boundary(&self, coverage: SeedCoverage) {
let _ = self.tx.send(SeedItem::SeedComplete(coverage)).await;
}
}
struct SeedReceiver {
rx: tokio::sync::mpsc::Receiver<SeedItem>,
dropped: Arc<std::sync::atomic::AtomicU64>,
}
impl SeedReceiver {
async fn recv(&mut self) -> Option<SeedItem> {
let missed = self.dropped.swap(0, std::sync::atomic::Ordering::Relaxed);
if missed > 0 {
return Some(SeedItem::Dropped(missed));
}
self.rx.recv().await
}
}
fn seed_channel(capacity: usize) -> (SeedSender, SeedReceiver) {
let (tx, rx) = tokio::sync::mpsc::channel::<SeedItem>(capacity);
let dropped = Arc::new(std::sync::atomic::AtomicU64::new(0));
(
SeedSender {
tx,
dropped: Arc::clone(&dropped),
},
SeedReceiver { rx, dropped },
)
}
pub struct SeededSubscriber {
rx: SeedReceiver,
_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::of(sample)
}
pub(crate) async fn seed_get(
session: &Session,
selector: &str,
timeout: Duration,
merge: &Merge,
mut deliver: impl FnMut(SampleView),
) -> usize {
let mut n = 0usize;
let opts = crate::bus::query::GetOpts::new(timeout).accept_any();
if let Ok(replies) = crate::bus::query::disciplined_get(session, selector, &opts).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) = seed_channel(SEED_CAPACITY);
let merge = Arc::new(Merge::new());
let subscriber = crate::bus::teardown::declared(
"seeded subscribe",
selector,
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) {
tx.send_sample(view);
}
}
}),
)
.await?;
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| {
tx.send_sample(view);
})
.await,
)
} else {
None
}
};
let storage = async {
if policy.storage {
Some(
seed_get(&session, &selector, policy.timeout, &merge, |view| {
tx.send_sample(view);
})
.await,
)
} else {
None
}
};
let (history_replies, storage_replies) = tokio::join!(history, storage);
tx.send_boundary(SeedCoverage {
history_replies,
storage_replies,
superseded: merge.superseded(),
})
.await;
})
};
Ok(SeededSubscriber {
rx,
_subscriber: subscriber,
task,
})
}
#[cfg(test)]
mod tests {
use super::*;
fn view(key: &str) -> SampleView {
SampleView {
key: key.to_string(),
payload: zenoh::bytes::ZBytes::from(vec![0u8; 1]),
encoding: "zenoh/bytes".to_string(),
kind: zenoh::sample::SampleKind::Put,
timestamp: None,
stamped_by: None,
attachment: None,
priority: zenoh::qos::Priority::DEFAULT,
congestion_control: zenoh::qos::CongestionControl::DEFAULT,
reliability: zenoh::qos::Reliability::DEFAULT,
express: false,
source: None,
received: std::time::Instant::now(),
}
}
#[tokio::test]
async fn a_slow_seed_consumer_is_told_what_it_missed() {
let (tx, mut rx) = seed_channel(4);
for i in 0..10 {
tx.send_sample(view(&format!("k/{i}")));
}
let Some(SeedItem::Dropped(n)) = rx.recv().await else {
panic!("expected the dropped count first");
};
assert_eq!(n, 6, "every refusal is counted, exactly once");
for i in 0..4 {
let Some(SeedItem::Sample(v)) = rx.recv().await else {
panic!("expected the retained samples");
};
assert_eq!(v.key, format!("k/{i}"), "the retained head is in order");
}
tx.send_sample(view("k/late"));
let Some(SeedItem::Sample(v)) = rx.recv().await else {
panic!("the stream resumes");
};
assert_eq!(v.key, "k/late");
for i in 0..4 {
tx.send_sample(view(&format!("b/{i}")));
}
let boundary = tokio::spawn(async move {
tx.send_boundary(SeedCoverage {
history_replies: Some(0),
storage_replies: Some(0),
superseded: 0,
})
.await;
});
let mut seen_boundary = false;
while let Some(item) = rx.recv().await {
if let SeedItem::SeedComplete(c) = item {
assert_eq!(c.superseded, 0);
seen_boundary = true;
break;
}
}
assert!(seen_boundary, "the boundary is never among the dropped");
boundary.await.expect("boundary task");
}
}