use std::collections::hash_map::Entry;
use std::collections::{BTreeMap, HashMap};
use std::task::Poll;
use moq_net::kio::{self, Pending, Waiter};
use moq_net::stats::{Presence, Role, Tier, Traffic};
use moq_net::track::Subscribing;
use moq_net::{PathOwned, origin};
use crate::{Result, SessionsFrame, TrafficFrame, parse_node_path, sessions_track, traffic_track};
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct Config {
pub prefix: PathOwned,
pub depth: usize,
pub compression: bool,
}
impl Config {
pub fn new() -> Self {
Self::default()
}
pub fn with_prefix(mut self, prefix: impl Into<PathOwned>) -> Self {
self.prefix = prefix.into();
self
}
pub fn with_depth(mut self, depth: usize) -> Self {
self.depth = depth;
self
}
pub fn with_compression(mut self, compression: bool) -> Self {
self.compression = compression;
self
}
}
impl Default for Config {
fn default() -> Self {
Self {
prefix: PathOwned::from(".stats"),
depth: 0,
compression: false,
}
}
}
pub struct Consumer {
origin: origin::Consumer,
config: Config,
}
impl Consumer {
pub fn new(origin: origin::Consumer, config: Config) -> Self {
Self { origin, config }
}
pub fn traffic(&self, tier: &Tier, role: Role) -> TrafficConsumer {
let name = traffic_track(tier, role, self.config.compression);
TrafficConsumer {
inner: Merged::new(self.origin.clone(), &self.config, name),
}
}
pub fn sessions(&self, tier: &Tier) -> SessionsConsumer {
let name = sessions_track(tier, self.config.compression);
SessionsConsumer {
inner: Merged::new(self.origin.clone(), &self.config, name),
}
}
}
pub struct TrafficConsumer {
inner: Merged<Traffic>,
}
impl TrafficConsumer {
pub async fn next(&mut self) -> Result<Option<TrafficFrame>> {
kio::wait(|waiter| self.inner.poll_next(waiter)).await
}
}
pub struct SessionsConsumer {
inner: Merged<Presence>,
}
impl SessionsConsumer {
pub async fn next(&mut self) -> Result<Option<SessionsFrame>> {
kio::wait(|waiter| self.inner.poll_next(waiter)).await
}
}
trait Mergeable: serde::de::DeserializeOwned + Default + Copy + 'static {
fn merge(acc: &mut Self, other: Self);
const STICKY: bool;
fn retire(&mut self) -> bool;
}
impl Mergeable for Traffic {
const STICKY: bool = true;
fn merge(acc: &mut Self, other: Self) {
acc.add(other);
}
fn retire(&mut self) -> bool {
let changed = self.announces_ended < self.announces_started
|| self.broadcasts_ended < self.broadcasts_started
|| self.subscriptions_ended < self.subscriptions_started;
self.announces_ended = self.announces_ended.max(self.announces_started);
self.broadcasts_ended = self.broadcasts_ended.max(self.broadcasts_started);
self.subscriptions_ended = self.subscriptions_ended.max(self.subscriptions_started);
changed
}
}
impl Mergeable for Presence {
const STICKY: bool = false;
fn merge(acc: &mut Self, other: Self) {
acc.add(other);
}
fn retire(&mut self) -> bool {
false
}
}
enum Reader<V: Mergeable> {
Resolving {
pending: Pending<origin::Requesting>,
queued: bool,
},
Subscribing(Pending<Subscribing>),
Active(Box<moq_json::snapshot::Consumer<BTreeMap<String, V>>>),
Ended,
}
struct Node<V: Mergeable> {
reader: Reader<V>,
path: PathOwned,
last: Option<BTreeMap<String, V>>,
}
impl<V: Mergeable> Node<V> {
fn depart(&mut self) -> bool {
self.reader = Reader::Ended;
if !V::STICKY {
return self.last.take().is_some();
}
let mut changed = false;
if let Some(last) = &mut self.last {
for value in last.values_mut() {
changed |= value.retire();
}
}
changed
}
}
struct Merged<V: Mergeable> {
origin: origin::Consumer,
announce: moq_net::announce::Consumer,
prefix: PathOwned,
depth: usize,
name: String,
config: moq_json::snapshot::consumer::Config,
nodes: HashMap<PathOwned, Node<V>>,
}
impl<V: Mergeable> Merged<V> {
fn new(origin: origin::Consumer, config: &Config, name: String) -> Self {
Self {
announce: origin.clone().with_hidden(true).announced(),
origin,
prefix: config.prefix.clone(),
depth: config.depth,
name,
config: {
let mut json = moq_json::snapshot::consumer::Config::default();
if config.compression {
json.compression = moq_json::Compression::Deflate;
}
json
},
nodes: HashMap::new(),
}
}
fn poll_next(&mut self, waiter: &Waiter) -> Poll<Result<Option<BTreeMap<String, V>>>> {
let mut changed = false;
loop {
match self.announce.poll_next(waiter) {
Poll::Ready(Some(update)) => changed |= self.apply_announce(update),
Poll::Ready(None) => return Poll::Ready(Ok(None)),
Poll::Pending => break,
}
}
let config = &self.config;
let name = self.name.as_str();
let origin = &self.origin;
for node in self.nodes.values_mut() {
changed |= advance(node, origin, config, name, waiter);
}
if changed {
Poll::Ready(Ok(Some(self.merged())))
} else {
Poll::Pending
}
}
fn apply_announce(&mut self, update: moq_net::announce::Update) -> bool {
let path = update.prefix;
let absolute = self.announce.absolute(&path).to_owned();
if parse_node_path(&self.prefix, self.depth, &absolute).is_none() {
return false;
}
if update.kind.is_active() {
match self.nodes.entry(absolute) {
Entry::Occupied(mut entry) => {
let node = entry.get_mut();
if matches!(node.reader, Reader::Ended) {
node.reader = resolve(&self.origin, &node.path);
}
false
}
Entry::Vacant(entry) => {
entry.insert(Node {
reader: resolve(&self.origin, &path),
path,
last: None,
});
false
}
}
} else if V::STICKY {
match self.nodes.get_mut(&absolute) {
Some(node) => node.depart(),
None => false,
}
} else {
self.nodes.remove(&absolute).is_some_and(|old| old.last.is_some())
}
}
fn merged(&self) -> BTreeMap<String, V> {
let mut acc: BTreeMap<String, V> = BTreeMap::new();
for node in self.nodes.values() {
if let Some(last) = &node.last {
for (key, value) in last {
V::merge(acc.entry(key.clone()).or_default(), *value);
}
}
}
acc
}
}
fn advance<V: Mergeable>(
node: &mut Node<V>,
origin: &origin::Consumer,
config: &moq_json::snapshot::consumer::Config,
name: &str,
waiter: &Waiter,
) -> bool {
let mut changed = false;
let mut rearmed = false;
loop {
match &mut node.reader {
Reader::Resolving { pending, queued } => match pending.poll_ok(waiter) {
Poll::Ready(Ok(broadcast)) => match broadcast.track(name) {
Ok(track) => node.reader = Reader::Subscribing(track.subscribe(None)),
Err(err) => {
tracing::debug!(?err, name, "stats: node missing track");
return changed | node.depart();
}
},
Poll::Ready(Err(moq_net::Error::Unroutable)) if *queued => {
node.reader = resolve(origin, &node.path);
}
Poll::Ready(Err(err)) => {
tracing::debug!(?err, name, "stats: node broadcast unresolvable");
return changed | node.depart();
}
Poll::Pending => return changed,
},
Reader::Subscribing(pending) => match pending.poll_ok(waiter) {
Poll::Ready(Ok(subscriber)) => {
node.reader =
Reader::Active(Box::new(moq_json::snapshot::Consumer::new(subscriber, config.clone())));
}
Poll::Ready(Err(err)) => {
tracing::debug!(?err, name, "stats: node subscribe failed");
return changed | node.depart();
}
Poll::Pending => return changed,
},
Reader::Active(reader) => match reader.poll_next(waiter) {
Poll::Ready(Ok(Some(frame))) => {
node.last = Some(frame);
changed = true;
}
Poll::Ready(result @ (Ok(None) | Err(_))) => {
if let Err(err) = result {
tracing::debug!(?err, name, "stats: node read error");
}
changed |= node.depart();
if rearmed {
return changed;
}
rearmed = true;
node.reader = resolve(origin, &node.path);
}
Poll::Pending => return changed,
},
Reader::Ended => return changed,
}
}
}
fn resolve<V: Mergeable>(origin: &origin::Consumer, path: &PathOwned) -> Reader<V> {
let pending = origin.request_broadcast(path);
let queued = pending.is_queued();
Reader::Resolving { pending, queued }
}
#[cfg(test)]
mod tests {
fn produce_origin() -> moq_net::origin::Producer {
let (producer, driver) = moq_net::origin::Producer::new(moq_net::origin::Config::default());
if tokio::runtime::Handle::try_current().is_ok() {
tokio::spawn(moq_net::time::run(driver));
} else {
std::mem::forget(driver);
}
producer
}
use std::time::Duration;
use moq_net::{PathOwned, Timestamp, announce, broadcast, origin, track};
use crate::{Producer, produce};
use super::*;
fn node_producer(origin: &origin::Producer, node: &str) -> Producer {
Producer::new(
produce::Config::new()
.with_origin(origin.clone())
.with_node(PathOwned::from(node.to_string()))
.with_depth(1),
)
}
#[allow(dead_code)]
struct Feed {
announced: announce::Consumer,
source: broadcast::Producer,
consumer: broadcast::Consumer,
sub: track::Subscriber,
ctx: moq_net::stats::Session,
}
async fn feed(producer: &Producer, tier: Tier, root: &str, path: &str, bytes: usize) -> Feed {
let ctx = producer.registry().tier(tier).session(root);
let feed_origin = produce_origin();
let egress = feed_origin.consume().with_stats(ctx.clone());
let mut announced = egress.announced();
let source = feed_origin.create_broadcast(path).expect("create_broadcast");
source.announce(origin::Route::default()).expect("announce");
let track = source.create_track("video", None).expect("create_track");
let update = announced.next().await.expect("announce");
assert!(update.kind.is_active());
let consumer = egress.request_broadcast(path).await.expect("resolve");
let mut sub = consumer.track("video").unwrap().subscribe(None).await.unwrap();
let mut group = track.append_group().unwrap();
group.write_frame(Timestamp::ZERO, vec![0u8; bytes]).unwrap();
group.finish().unwrap();
let mut group = sub.recv_group().await.unwrap().unwrap();
while group.read_frame().await.unwrap().is_some() {}
Feed {
announced,
source,
consumer,
sub,
ctx,
}
}
async fn drive_tick() {
tokio::time::advance(Duration::from_millis(1100)).await;
for _ in 0..8 {
tokio::task::yield_now().await;
}
}
async fn read_until_bytes(consumer: &mut TrafficConsumer, path: &str, want: u64) -> TrafficFrame {
loop {
let frame = consumer.next().await.expect("read").expect("frame");
if frame.get(path).map(|t| t.bytes).unwrap_or(0) >= want {
return frame;
}
}
}
async fn read_monotonic_until(consumer: &mut TrafficConsumer, path: &str, min: u64, want: u64) -> TrafficFrame {
loop {
let frame = consumer.next().await.expect("read").expect("frame");
let bytes = frame.get(path).map(|t| t.bytes).unwrap_or(0);
assert!(bytes >= min, "traffic regressed below {min}: {bytes}");
if bytes >= want {
return frame;
}
}
}
#[allow(dead_code)]
struct NodeBroadcast {
source: broadcast::Producer,
traffic: moq_json::snapshot::Producer<TrafficFrame>,
track: track::Producer,
frame: TrafficFrame,
}
impl NodeBroadcast {
fn new(origin: &origin::Producer, group: &str, node: &str) -> Self {
let path = format!(".stats/{group}/node/{node}");
let source = origin.create_broadcast(path.as_str()).expect("create broadcast");
source.announce(origin::Route::default()).expect("announce");
let name = traffic_track(&Tier::default(), Role::Publisher, false);
let track = source.create_track(name, None).expect("create track");
let config = moq_json::snapshot::Config::default().with_delta_ratio(0);
Self {
traffic: moq_json::snapshot::Producer::new(track.clone(), config),
track,
source,
frame: TrafficFrame::new(),
}
}
fn publish(&mut self, path: &str, bytes: u64) {
let entry = self.frame.entry(path.to_string()).or_default();
entry.bytes += bytes;
self.traffic.update(&self.frame).expect("publish");
}
fn fail_traffic(&mut self) {
let mut group = self.track.append_group().expect("append group");
group
.write_frame(Timestamp::ZERO, b"not json".to_vec())
.expect("write frame");
group.finish().expect("finish group");
}
}
#[tokio::test(start_paused = true)]
async fn merges_traffic_across_nodes() {
let origin = produce_origin();
let node_a = node_producer(&origin, "a");
let node_b = node_producer(&origin, "b");
let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
drive_tick().await;
let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
let frame = read_until_bytes(&mut traffic, "acme/room", 140).await;
let snap = frame.get("acme/room").expect("entry");
assert_eq!(snap.bytes, 140, "bytes sum across both nodes");
assert_eq!(snap.subscriptions_started, 2, "one subscription per node");
assert_eq!(snap.broadcasts_started, 2, "one viewer per node");
}
#[tokio::test(start_paused = true)]
async fn node_drop_keeps_the_traffic_total() {
let origin = produce_origin();
let node_a = node_producer(&origin, "a");
let node_b = node_producer(&origin, "b");
let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
let fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
drive_tick().await;
let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
read_until_bytes(&mut traffic, "acme/room", 140).await;
drop(fb);
drop(node_b);
drive_tick().await;
let _fa2 = feed(&node_a, Tier::default(), "acme", "acme/other", 10).await;
drive_tick().await;
let frame = read_until_bytes(&mut traffic, "acme/other", 10).await;
assert_eq!(
frame.get("acme/room").map(|t| t.bytes),
Some(140),
"the departed node's contribution stays in the total",
);
}
#[tokio::test(start_paused = true)]
async fn reannounce_with_higher_counter_stays_monotonic() {
let origin = produce_origin();
let node_a = node_producer(&origin, "a");
let node_b = node_producer(&origin, "b");
let fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
drive_tick().await;
let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
read_until_bytes(&mut traffic, "acme/room", 140).await;
drop(fa);
drop(node_a);
drive_tick().await;
let node_a = node_producer(&origin, "a");
let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 120).await;
drive_tick().await;
let frame = read_monotonic_until(&mut traffic, "acme/room", 140, 160).await;
assert_eq!(frame.get("acme/room").expect("entry").bytes, 160);
}
#[tokio::test(start_paused = true)]
async fn restarted_node_regresses_the_total() {
let origin = produce_origin();
let node_a = node_producer(&origin, "a");
let node_b = node_producer(&origin, "b");
let fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
drive_tick().await;
let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
read_until_bytes(&mut traffic, "acme/room", 140).await;
drop(fa);
drop(node_a);
drive_tick().await;
let node_a = node_producer(&origin, "a");
let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 30).await;
drive_tick().await;
loop {
let frame = traffic.next().await.expect("read").expect("frame");
if frame.get("acme/room").map(|t| t.bytes) == Some(70) {
break;
}
}
}
#[tokio::test(start_paused = true)]
async fn reader_failure_keeps_the_traffic_total() {
let origin = produce_origin();
let mut node_a = NodeBroadcast::new(&origin, "acme", "a");
let mut node_b = NodeBroadcast::new(&origin, "acme", "b");
node_a.publish("acme/room", 100);
node_b.publish("acme/room", 40);
let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
read_until_bytes(&mut traffic, "acme/room", 140).await;
node_a.fail_traffic();
node_b.publish("acme/other", 10);
let frame = read_until_bytes(&mut traffic, "acme/other", 10).await;
assert_eq!(
frame.get("acme/room").map(|t| t.bytes),
Some(140),
"the failed node's contribution stays in the total",
);
}
#[tokio::test(start_paused = true)]
async fn unannounce_retires_live_gauges() {
let origin = produce_origin();
let mut node_a = NodeBroadcast::new(&origin, "acme", "a");
let mut published = Traffic::default();
published.announces_started = 2;
published.announces_ended = 1;
published.broadcasts_started = 3;
published.broadcasts_ended = 1;
published.subscriptions_started = 4;
published.subscriptions_ended = 1;
published.bytes = 100;
node_a.frame.insert("acme/room".to_string(), published);
node_a.traffic.update(&node_a.frame).expect("publish");
let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
let frame = read_until_bytes(&mut traffic, "acme/room", 100).await;
let snap = frame.get("acme/room").expect("entry");
assert!(snap.is_announced());
assert_eq!(snap.active_broadcasts(), 2);
assert_eq!(snap.active_subscriptions(), 3);
drop(node_a);
let frame = traffic.next().await.expect("read").expect("frame");
let snap = frame.get("acme/room").expect("entry");
assert_eq!(snap.bytes, 100, "cumulative totals stay");
assert!(!snap.is_announced(), "no phantom announcement");
assert_eq!(snap.active_broadcasts(), 0, "no phantom broadcasts");
assert_eq!(snap.active_subscriptions(), 0, "no phantom subscriptions");
}
#[tokio::test(start_paused = true)]
async fn merges_sessions_across_nodes() {
let origin = produce_origin();
let node_a = node_producer(&origin, "a");
let node_b = node_producer(&origin, "b");
let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 8).await;
let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 8).await;
let _sa = node_a.registry().tier(Tier::default()).session("acme");
let _sb = node_b.registry().tier(Tier::default()).session("acme");
drive_tick().await;
let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
let mut sessions = agg.sessions(&Tier::default());
loop {
let frame = sessions.next().await.expect("read").expect("frame");
if frame.get("acme").map(|p| p.active()) >= Some(4) {
break;
}
}
}
#[tokio::test(start_paused = true)]
async fn unannounce_drops_presence_immediately() {
let origin = produce_origin();
let node_a = node_producer(&origin, "a");
let node_b = node_producer(&origin, "b");
let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 8).await;
let fb = feed(&node_b, Tier::default(), "acme", "acme/room", 8).await;
let _sa = node_a.registry().tier(Tier::default()).session("acme");
let sb = node_b.registry().tier(Tier::default()).session("acme");
drive_tick().await;
let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
let mut sessions = agg.sessions(&Tier::default());
loop {
let frame = sessions.next().await.expect("read").expect("frame");
if frame.get("acme").map(|p| p.active()) >= Some(4) {
break;
}
}
drop(fb);
drop(sb);
drop(node_b);
drive_tick().await;
let frame = sessions.next().await.expect("read").expect("frame");
assert_eq!(
frame.get("acme").map(|p| p.active()),
Some(2),
"presence drops immediately"
);
}
}