use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use anyhow::{Result, anyhow};
use arc_swap::ArcSwap;
use tokio::sync::broadcast;
use zenoh::Session;
use zenoh::sample::SampleKind;
use crate::stats::StatsTable;
use crate::tree::KeyTreeSnapshot;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SampleSource {
pub zid: zenoh::config::ZenohId,
pub eid: u32,
pub sn: u32,
}
#[derive(Debug, Clone)]
pub struct SampleView {
pub key: String,
pub payload: zenoh::bytes::ZBytes,
pub encoding: String,
pub kind: SampleKind,
pub timestamp: Option<zenoh::time::Timestamp>,
pub attachment: Option<zenoh::bytes::ZBytes>,
pub priority: zenoh::qos::Priority,
pub congestion_control: zenoh::qos::CongestionControl,
pub reliability: zenoh::qos::Reliability,
pub express: bool,
pub source: Option<SampleSource>,
pub received: Instant,
}
impl SampleView {
pub fn qos_matches(&self, profile: zenkey::qos::QosProfile) -> bool {
self.priority == profile.priority()
&& self.congestion_control == profile.congestion_control()
&& self.reliability == profile.reliability()
&& self.express == profile.express()
}
}
#[derive(Debug, Clone)]
pub enum FleetEvent {
Sample(Arc<SampleView>),
NodeUp(String),
NodeDown(String),
StatsTick,
WatchChanged,
WatchSeeded {
id: WatchId,
coverage: crate::seed::SeedCoverage,
},
}
#[derive(Debug, Clone)]
pub struct MonitorSpec {
pub selectors: Vec<String>,
pub liveliness: Vec<String>,
pub stats_tick: Duration,
pub capacity: usize,
pub max_keys: usize,
}
impl Default for MonitorSpec {
fn default() -> Self {
MonitorSpec {
selectors: Vec::new(),
liveliness: Vec::new(),
stats_tick: Duration::from_millis(250),
capacity: 1024,
max_keys: crate::stats::DEFAULT_MAX_KEYS,
}
}
}
pub struct MonitorCore {
tx: broadcast::Sender<FleetEvent>,
stats: Mutex<StatsTable>,
tree: ArcSwap<KeyTreeSnapshot>,
dropped: AtomicU64,
}
impl MonitorCore {
pub fn new(capacity: usize) -> Arc<MonitorCore> {
MonitorCore::bounded(capacity, crate::stats::DEFAULT_MAX_KEYS)
}
pub fn bounded(capacity: usize, max_keys: usize) -> Arc<MonitorCore> {
let (tx, _) = broadcast::channel(capacity.max(2));
Arc::new(MonitorCore {
tx,
stats: Mutex::new(StatsTable::with_capacity(max_keys)),
tree: ArcSwap::from_pointee(KeyTreeSnapshot::default()),
dropped: AtomicU64::new(0),
})
}
pub fn ingest(&self, view: SampleView, sn: Option<u32>) {
{
let latency_us = view.timestamp.map(|t| {
let published = t.get_time().to_system_time();
match std::time::SystemTime::now().duration_since(published) {
Ok(d) => i64::try_from(d.as_micros()).unwrap_or(i64::MAX),
Err(e) => -i64::try_from(e.duration().as_micros()).unwrap_or(i64::MAX),
}
});
let mut stats = self.stats.lock().expect("stats lock");
stats.record(
&view.key,
view.payload.len(),
sn,
Instant::now(),
latency_us,
);
}
let _ = self.tx.send(FleetEvent::Sample(Arc::new(view)));
}
pub fn node_event(&self, key: String, up: bool) {
let _ = self.tx.send(if up {
FleetEvent::NodeUp(key)
} else {
FleetEvent::NodeDown(key)
});
}
pub fn tick(&self) {
let snapshot = {
let stats = self.stats.lock().expect("stats lock");
KeyTreeSnapshot::build(&stats)
};
self.tree.store(Arc::new(snapshot));
let _ = self.tx.send(FleetEvent::StatsTick);
}
pub fn tree(&self) -> Arc<KeyTreeSnapshot> {
self.tree.load_full()
}
pub fn with_stats<R>(&self, f: impl FnOnce(&StatsTable) -> R) -> R {
f(&self.stats.lock().expect("stats lock"))
}
pub fn with_stats_mut<R>(&self, f: impl FnOnce(&mut StatsTable) -> R) -> R {
f(&mut self.stats.lock().expect("stats lock"))
}
pub fn keys_unwatched(&self) -> u64 {
self.with_stats(|s| s.unwatched())
}
pub fn dropped(&self) -> u64 {
self.dropped.load(Ordering::Relaxed)
}
pub fn keys_evicted(&self) -> u64 {
self.with_stats(|s| s.evicted())
}
pub fn events(self: &Arc<Self>) -> EventStream {
EventStream {
rx: self.tx.subscribe(),
core: Arc::clone(self),
}
}
}
pub struct EventStream {
rx: broadcast::Receiver<FleetEvent>,
core: Arc<MonitorCore>,
}
#[derive(Debug, Clone)]
pub enum StreamItem {
Event(FleetEvent),
Dropped(u64),
}
impl EventStream {
pub async fn recv(&mut self) -> Option<StreamItem> {
match self.rx.recv().await {
Ok(ev) => Some(StreamItem::Event(ev)),
Err(broadcast::error::RecvError::Lagged(n)) => {
self.core.dropped.fetch_add(n, Ordering::Relaxed);
Some(StreamItem::Dropped(n))
}
Err(broadcast::error::RecvError::Closed) => None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct WatchId(u64);
struct WatchEntry {
selector: String,
subscriber: zenoh::pubsub::Subscriber<()>,
seed_task: Option<tokio::task::JoinHandle<()>>,
}
pub struct Monitor {
core: Arc<MonitorCore>,
session: Session,
watches: tokio::sync::Mutex<std::collections::HashMap<WatchId, WatchEntry>>,
next_watch: AtomicU64,
tasks: Vec<tokio::task::JoinHandle<()>>,
}
impl std::fmt::Debug for Monitor {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Monitor").finish_non_exhaustive()
}
}
impl Monitor {
pub async fn start(session: &Session, spec: MonitorSpec) -> Result<Monitor> {
let core = MonitorCore::bounded(spec.capacity, spec.max_keys);
let mut tasks = Vec::new();
for liveliness_sel in &spec.liveliness {
let subscriber = session
.liveliness()
.declare_subscriber(liveliness_sel)
.history(true)
.await
.map_err(|e| anyhow!("liveliness subscribe {liveliness_sel}: {e}"))?;
let core = Arc::clone(&core);
tasks.push(tokio::spawn(async move {
while let Ok(sample) = subscriber.recv_async().await {
let key = sample.key_expr().as_str().to_string();
core.node_event(key, sample.kind() == SampleKind::Put);
}
}));
}
{
let core = Arc::clone(&core);
let period = spec.stats_tick;
tasks.push(tokio::spawn(async move {
let mut interval = tokio::time::interval(period);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
interval.tick().await;
core.tick();
}
}));
}
let monitor = Monitor {
core,
session: session.clone(),
watches: tokio::sync::Mutex::new(std::collections::HashMap::new()),
next_watch: AtomicU64::new(0),
tasks,
};
for selector in &spec.selectors {
monitor.watch(selector).await?;
}
Ok(monitor)
}
pub async fn watch(&self, selector: &str) -> Result<WatchId> {
let core = Arc::clone(&self.core);
let subscriber = self
.session
.declare_subscriber(selector)
.callback(move |sample| {
let source = sample.source_info().map(|si| SampleSource {
zid: si.source_id().zid(),
eid: si.source_id().eid(),
sn: si.source_sn(),
});
core.ingest(
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,
received: Instant::now(),
},
source.map(|s| s.sn),
);
})
.await
.map_err(|e| anyhow!("subscribe {selector}: {e}"))?;
let id = WatchId(self.next_watch.fetch_add(1, Ordering::Relaxed));
self.watches.lock().await.insert(
id,
WatchEntry {
selector: selector.to_string(),
subscriber,
seed_task: None,
},
);
let _ = self.core.tx.send(FleetEvent::WatchChanged);
Ok(id)
}
pub async fn watch_seeded(
&self,
selector: &str,
policy: crate::seed::SeedPolicy,
) -> Result<WatchId> {
use crate::seed::{Merge, SeedCoverage, cache_selector, seed_get, view_of};
let gate: Arc<arc_swap::ArcSwapOption<Merge>> =
Arc::new(arc_swap::ArcSwapOption::from_pointee(Merge::new()));
let core = Arc::clone(&self.core);
let cb_gate = Arc::clone(&gate);
let subscriber = self
.session
.declare_subscriber(selector)
.callback(move |sample| {
let sn = sample.source_info().map(|si| si.source_sn());
let view = view_of(&sample);
if let Some(merge) = cb_gate.load_full()
&& !merge.admit(&view)
{
return;
}
core.ingest(view, sn);
})
.await
.map_err(|e| anyhow!("seeded subscribe {selector}: {e}"))?;
let id = WatchId(self.next_watch.fetch_add(1, Ordering::Relaxed));
let seed_task = {
let session = self.session.clone();
let core = Arc::clone(&self.core);
let selector = selector.to_string();
tokio::spawn(async move {
let merge = gate
.load_full()
.expect("gate holds the merge while seeding");
let history = async {
if policy.history {
let sel = cache_selector(&selector);
Some(
seed_get(&session, &sel, policy.timeout, &merge, |view| {
core.ingest(view, None);
})
.await,
)
} else {
None
}
};
let storage = async {
if policy.storage {
Some(
seed_get(&session, &selector, policy.timeout, &merge, |view| {
core.ingest(view, None);
})
.await,
)
} else {
None
}
};
let (history_replies, storage_replies) = tokio::join!(history, storage);
let coverage = SeedCoverage {
history_replies,
storage_replies,
superseded: merge.superseded(),
};
gate.store(None);
core.tick();
let _ = core.tx.send(FleetEvent::WatchSeeded { id, coverage });
})
};
self.watches.lock().await.insert(
id,
WatchEntry {
selector: selector.to_string(),
subscriber,
seed_task: Some(seed_task),
},
);
let _ = self.core.tx.send(FleetEvent::WatchChanged);
Ok(id)
}
pub async fn unwatch(&self, id: WatchId) -> Result<()> {
let mut entry = {
let mut watches = self.watches.lock().await;
watches
.remove(&id)
.ok_or_else(|| anyhow!("unknown watch id {id:?}"))?
};
if let Some(task) = entry.seed_task.take() {
task.abort();
}
entry
.subscriber
.undeclare()
.await
.map_err(|e| anyhow!("undeclare {}: {e}", entry.selector))?;
let kept: Vec<String> = {
let watches = self.watches.lock().await;
watches.values().map(|w| w.selector.clone()).collect()
};
self.core.with_stats_mut(|stats| {
stats.retire_unwatched(&entry.selector, &kept);
});
self.core.tick();
let _ = self.core.tx.send(FleetEvent::WatchChanged);
Ok(())
}
pub async fn watched(&self) -> Vec<(WatchId, String)> {
let watches = self.watches.lock().await;
let mut v: Vec<(WatchId, String)> = watches
.iter()
.map(|(id, w)| (*id, w.selector.clone()))
.collect();
v.sort();
v
}
pub fn core(&self) -> &Arc<MonitorCore> {
&self.core
}
pub fn events(&self) -> EventStream {
self.core.events()
}
pub fn tree(&self) -> Arc<KeyTreeSnapshot> {
self.core.tree()
}
pub fn stop(self) {
drop(self);
}
}
impl Drop for Monitor {
fn drop(&mut self) {
for t in &self.tasks {
t.abort();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn view(key: &str, len: usize) -> SampleView {
SampleView {
key: key.to_string(),
payload: zenoh::bytes::ZBytes::from(vec![0u8; len]),
encoding: "zenoh/bytes".to_string(),
kind: SampleKind::Put,
timestamp: None,
attachment: None,
priority: zenoh::qos::Priority::DEFAULT,
congestion_control: zenoh::qos::CongestionControl::DEFAULT,
reliability: zenoh::qos::Reliability::DEFAULT,
express: false,
source: None,
received: Instant::now(),
}
}
#[tokio::test]
async fn events_flow_and_snapshots_rebuild_on_tick() {
let core = MonitorCore::new(8);
let mut events = core.events();
core.ingest(view("zs/v1/h-a/telemetry/x/m", 4), None);
core.tick();
let Some(StreamItem::Event(FleetEvent::Sample(s))) = events.recv().await else {
panic!("expected sample");
};
assert_eq!(s.key, "zs/v1/h-a/telemetry/x/m");
assert_eq!(s.payload.len(), 4);
let Some(StreamItem::Event(FleetEvent::StatsTick)) = events.recv().await else {
panic!("expected tick");
};
let snap = core.tree();
assert_eq!(snap.keys, 1);
assert_eq!(snap.root.subtree_count, 1);
}
#[tokio::test]
async fn overflow_surfaces_as_dropped_counts() {
let core = MonitorCore::new(2);
let mut slow = core.events();
for i in 0..10 {
core.ingest(view(&format!("zs/v1/h-a/telemetry/x/m{i}"), 1), None);
}
let Some(StreamItem::Dropped(n)) = slow.recv().await else {
panic!("expected a dropped count first");
};
assert!(n >= 8, "missed at least 8, reported {n}");
assert_eq!(core.dropped(), n);
let Some(StreamItem::Event(FleetEvent::Sample(_))) = slow.recv().await else {
panic!("expected a sample after the gap report");
};
}
}