1use std::collections::hash_map::Entry;
10use std::collections::{BTreeMap, HashMap};
11use std::task::Poll;
12
13use moq_net::kio::{self, Pending, Waiter};
14use moq_net::stats::{Presence, Role, Tier, Traffic};
15use moq_net::track::Subscribing;
16use moq_net::{PathOwned, origin};
17
18use crate::{Result, SessionsFrame, TrafficFrame, parse_node_path, sessions_track, traffic_track};
19
20#[derive(Debug, Clone)]
28#[non_exhaustive]
29pub struct Config {
30 pub prefix: PathOwned,
33 pub depth: usize,
37 pub compression: bool,
41}
42
43impl Config {
44 pub fn new() -> Self {
47 Self::default()
48 }
49
50 pub fn with_prefix(mut self, prefix: impl Into<PathOwned>) -> Self {
52 self.prefix = prefix.into();
53 self
54 }
55
56 pub fn with_depth(mut self, depth: usize) -> Self {
58 self.depth = depth;
59 self
60 }
61
62 pub fn with_compression(mut self, compression: bool) -> Self {
64 self.compression = compression;
65 self
66 }
67}
68
69impl Default for Config {
70 fn default() -> Self {
71 Self {
72 prefix: PathOwned::from(".stats"),
73 depth: 0,
74 compression: false,
75 }
76 }
77}
78
79pub struct Consumer {
92 origin: origin::Consumer,
93 config: Config,
94}
95
96impl Consumer {
97 pub fn new(origin: origin::Consumer, config: Config) -> Self {
100 Self { origin, config }
101 }
102
103 pub fn traffic(&self, tier: &Tier, role: Role) -> TrafficConsumer {
107 let name = traffic_track(tier, role, self.config.compression);
108 TrafficConsumer {
109 inner: Merged::new(self.origin.clone(), &self.config, name),
110 }
111 }
112
113 pub fn sessions(&self, tier: &Tier) -> SessionsConsumer {
115 let name = sessions_track(tier, self.config.compression);
116 SessionsConsumer {
117 inner: Merged::new(self.origin.clone(), &self.config, name),
118 }
119 }
120}
121
122pub struct TrafficConsumer {
126 inner: Merged<Traffic>,
127}
128
129impl TrafficConsumer {
130 pub async fn next(&mut self) -> Result<Option<TrafficFrame>> {
133 kio::wait(|waiter| self.inner.poll_next(waiter)).await
134 }
135}
136
137pub struct SessionsConsumer {
140 inner: Merged<Presence>,
141}
142
143impl SessionsConsumer {
144 pub async fn next(&mut self) -> Result<Option<SessionsFrame>> {
146 kio::wait(|waiter| self.inner.poll_next(waiter)).await
147 }
148}
149
150trait Mergeable: serde::de::DeserializeOwned + Default + Copy + 'static {
152 fn merge(acc: &mut Self, other: Self);
154
155 const STICKY: bool;
160
161 fn retire(&mut self) -> bool;
165}
166
167impl Mergeable for Traffic {
168 const STICKY: bool = true;
169
170 fn merge(acc: &mut Self, other: Self) {
171 acc.add(other);
172 }
173
174 fn retire(&mut self) -> bool {
178 let changed = self.announces_ended < self.announces_started
179 || self.broadcasts_ended < self.broadcasts_started
180 || self.subscriptions_ended < self.subscriptions_started;
181 self.announces_ended = self.announces_ended.max(self.announces_started);
182 self.broadcasts_ended = self.broadcasts_ended.max(self.broadcasts_started);
183 self.subscriptions_ended = self.subscriptions_ended.max(self.subscriptions_started);
184 changed
185 }
186}
187
188impl Mergeable for Presence {
189 const STICKY: bool = false;
190
191 fn merge(acc: &mut Self, other: Self) {
192 acc.add(other);
193 }
194
195 fn retire(&mut self) -> bool {
197 false
198 }
199}
200
201enum Reader<V: Mergeable> {
203 Resolving {
209 pending: Pending<origin::Requesting>,
210 queued: bool,
211 },
212 Subscribing(Pending<Subscribing>),
214 Active(Box<moq_json::snapshot::Consumer<BTreeMap<String, V>>>),
217 Ended,
221}
222
223struct Node<V: Mergeable> {
226 reader: Reader<V>,
227 path: PathOwned,
230 last: Option<BTreeMap<String, V>>,
231}
232
233impl<V: Mergeable> Node<V> {
234 fn depart(&mut self) -> bool {
240 self.reader = Reader::Ended;
241 if !V::STICKY {
242 return self.last.take().is_some();
243 }
244 let mut changed = false;
245 if let Some(last) = &mut self.last {
246 for value in last.values_mut() {
247 changed |= value.retire();
248 }
249 }
250 changed
251 }
252}
253
254struct Merged<V: Mergeable> {
256 origin: origin::Consumer,
258 announce: moq_net::announce::Consumer,
259 prefix: PathOwned,
260 depth: usize,
261 name: String,
263 config: moq_json::snapshot::consumer::Config,
264 nodes: HashMap<PathOwned, Node<V>>,
266}
267
268impl<V: Mergeable> Merged<V> {
269 fn new(origin: origin::Consumer, config: &Config, name: String) -> Self {
270 Self {
271 announce: origin.announced(),
272 origin,
273 prefix: config.prefix.clone(),
274 depth: config.depth,
275 name,
276 config: {
277 let mut json = moq_json::snapshot::consumer::Config::default();
278 if config.compression {
279 json.compression = moq_json::Compression::Deflate;
280 }
281 json
282 },
283 nodes: HashMap::new(),
284 }
285 }
286
287 fn poll_next(&mut self, waiter: &Waiter) -> Poll<Result<Option<BTreeMap<String, V>>>> {
292 let mut changed = false;
293
294 loop {
297 match self.announce.poll_next(waiter) {
298 Poll::Ready(Some(update)) => changed |= self.apply_announce(update),
299 Poll::Ready(None) => return Poll::Ready(Ok(None)),
300 Poll::Pending => break,
301 }
302 }
303
304 let config = &self.config;
306 let name = self.name.as_str();
307 let origin = &self.origin;
308 for node in self.nodes.values_mut() {
309 changed |= advance(node, origin, config, name, waiter);
310 }
311
312 if changed {
313 Poll::Ready(Ok(Some(self.merged())))
314 } else {
315 Poll::Pending
316 }
317 }
318
319 fn apply_announce(&mut self, update: moq_net::announce::Update) -> bool {
323 let path = update.prefix;
324 let absolute = self.announce.absolute(&path).to_owned();
325
326 if parse_node_path(&self.prefix, self.depth, &absolute).is_none() {
330 return false;
331 }
332
333 if update.kind.is_active() {
334 match self.nodes.entry(absolute) {
342 Entry::Occupied(mut entry) => {
343 let node = entry.get_mut();
344 if matches!(node.reader, Reader::Ended) {
345 node.reader = resolve(&self.origin, &node.path);
346 }
347 false
348 }
349 Entry::Vacant(entry) => {
350 entry.insert(Node {
351 reader: resolve(&self.origin, &path),
352 path,
353 last: None,
354 });
355 false
356 }
357 }
358 } else if V::STICKY {
359 match self.nodes.get_mut(&absolute) {
362 Some(node) => node.depart(),
363 None => false,
364 }
365 } else {
366 self.nodes.remove(&absolute).is_some_and(|old| old.last.is_some())
369 }
370 }
371
372 fn merged(&self) -> BTreeMap<String, V> {
374 let mut acc: BTreeMap<String, V> = BTreeMap::new();
375 for node in self.nodes.values() {
376 if let Some(last) = &node.last {
377 for (key, value) in last {
378 V::merge(acc.entry(key.clone()).or_default(), *value);
379 }
380 }
381 }
382 acc
383 }
384}
385
386fn advance<V: Mergeable>(
389 node: &mut Node<V>,
390 origin: &origin::Consumer,
391 config: &moq_json::snapshot::consumer::Config,
392 name: &str,
393 waiter: &Waiter,
394) -> bool {
395 let mut changed = false;
396 let mut rearmed = false;
399 loop {
400 match &mut node.reader {
401 Reader::Resolving { pending, queued } => match pending.poll_ok(waiter) {
402 Poll::Ready(Ok(broadcast)) => match broadcast.track(name) {
403 Ok(track) => node.reader = Reader::Subscribing(track.subscribe(None)),
404 Err(err) => {
405 tracing::debug!(?err, name, "stats: node missing track");
406 return changed | node.depart();
407 }
408 },
409 Poll::Ready(Err(moq_net::Error::Unroutable)) if *queued => {
416 node.reader = resolve(origin, &node.path);
417 }
418 Poll::Ready(Err(err)) => {
419 tracing::debug!(?err, name, "stats: node broadcast unresolvable");
420 return changed | node.depart();
421 }
422 Poll::Pending => return changed,
423 },
424 Reader::Subscribing(pending) => match pending.poll_ok(waiter) {
425 Poll::Ready(Ok(subscriber)) => {
426 node.reader =
427 Reader::Active(Box::new(moq_json::snapshot::Consumer::new(subscriber, config.clone())));
428 }
429 Poll::Ready(Err(err)) => {
430 tracing::debug!(?err, name, "stats: node subscribe failed");
431 return changed | node.depart();
432 }
433 Poll::Pending => return changed,
434 },
435 Reader::Active(reader) => match reader.poll_next(waiter) {
436 Poll::Ready(Ok(Some(frame))) => {
437 node.last = Some(frame);
438 changed = true;
439 }
440 Poll::Ready(result @ (Ok(None) | Err(_))) => {
447 if let Err(err) = result {
448 tracing::debug!(?err, name, "stats: node read error");
451 }
452 changed |= node.depart();
457 if rearmed {
458 return changed;
459 }
460 rearmed = true;
461 node.reader = resolve(origin, &node.path);
462 }
463 Poll::Pending => return changed,
464 },
465 Reader::Ended => return changed,
466 }
467 }
468}
469
470fn resolve<V: Mergeable>(origin: &origin::Consumer, path: &PathOwned) -> Reader<V> {
472 let pending = origin.request_broadcast(path);
473 let queued = pending.is_queued();
474 Reader::Resolving { pending, queued }
475}
476
477#[cfg(test)]
478mod tests {
479 fn produce_origin() -> moq_net::origin::Producer {
481 let (producer, driver) = moq_net::origin::Producer::new(moq_net::origin::Config::default());
482 if tokio::runtime::Handle::try_current().is_ok() {
483 tokio::spawn(moq_net::time::run(driver));
484 } else {
485 std::mem::forget(driver);
488 }
489 producer
490 }
491
492 use std::time::Duration;
493
494 use moq_net::{PathOwned, Timestamp, announce, broadcast, origin, track};
495
496 use crate::{Producer, produce};
497
498 use super::*;
499
500 fn node_producer(origin: &origin::Producer, node: &str) -> Producer {
504 Producer::new(
505 produce::Config::new()
506 .with_origin(origin.clone())
507 .with_node(PathOwned::from(node.to_string()))
508 .with_depth(1),
509 )
510 }
511
512 #[allow(dead_code)]
515 struct Feed {
516 announced: announce::Consumer,
517 source: broadcast::Producer,
518 consumer: broadcast::Consumer,
519 sub: track::Subscriber,
520 ctx: moq_net::stats::Session,
521 }
522
523 async fn feed(producer: &Producer, tier: Tier, root: &str, path: &str, bytes: usize) -> Feed {
527 let ctx = producer.registry().tier(tier).session(root);
528 let feed_origin = produce_origin();
529 let egress = feed_origin.consume().with_stats(ctx.clone());
530
531 let mut announced = egress.announced();
532 let source = feed_origin.create_broadcast(path).expect("create_broadcast");
533 source.announce(origin::Route::default()).expect("announce");
534 let track = source.create_track("video", None).expect("create_track");
535
536 let update = announced.next().await.expect("announce");
537 assert!(update.kind.is_active());
538 let consumer = egress.request_broadcast(path).await.expect("resolve");
539 let mut sub = consumer.track("video").unwrap().subscribe(None).await.unwrap();
540
541 let mut group = track.append_group().unwrap();
542 group.write_frame(Timestamp::ZERO, vec![0u8; bytes]).unwrap();
543 group.finish().unwrap();
544 let mut group = sub.recv_group().await.unwrap().unwrap();
545 while group.read_frame().await.unwrap().is_some() {}
546
547 Feed {
548 announced,
549 source,
550 consumer,
551 sub,
552 ctx,
553 }
554 }
555
556 async fn drive_tick() {
558 tokio::time::advance(Duration::from_millis(1100)).await;
559 for _ in 0..8 {
560 tokio::task::yield_now().await;
561 }
562 }
563
564 async fn read_until_bytes(consumer: &mut TrafficConsumer, path: &str, want: u64) -> TrafficFrame {
567 loop {
568 let frame = consumer.next().await.expect("read").expect("frame");
569 if frame.get(path).map(|t| t.bytes).unwrap_or(0) >= want {
570 return frame;
571 }
572 }
573 }
574
575 async fn read_monotonic_until(consumer: &mut TrafficConsumer, path: &str, min: u64, want: u64) -> TrafficFrame {
579 loop {
580 let frame = consumer.next().await.expect("read").expect("frame");
581 let bytes = frame.get(path).map(|t| t.bytes).unwrap_or(0);
582 assert!(bytes >= min, "traffic regressed below {min}: {bytes}");
583 if bytes >= want {
584 return frame;
585 }
586 }
587 }
588
589 #[allow(dead_code)]
594 struct NodeBroadcast {
595 source: broadcast::Producer,
596 traffic: moq_json::snapshot::Producer<TrafficFrame>,
597 track: track::Producer,
598 frame: TrafficFrame,
599 }
600
601 impl NodeBroadcast {
602 fn new(origin: &origin::Producer, group: &str, node: &str) -> Self {
603 let path = format!(".stats/{group}/node/{node}");
604 let source = origin.create_broadcast(path.as_str()).expect("create broadcast");
605 source.announce(origin::Route::default()).expect("announce");
606 let name = traffic_track(&Tier::default(), Role::Publisher, false);
607 let track = source.create_track(name, None).expect("create track");
608 let config = moq_json::snapshot::Config::default().with_delta_ratio(0);
609 Self {
610 traffic: moq_json::snapshot::Producer::new(track.clone(), config),
611 track,
612 source,
613 frame: TrafficFrame::new(),
614 }
615 }
616
617 fn publish(&mut self, path: &str, bytes: u64) {
620 let entry = self.frame.entry(path.to_string()).or_default();
621 entry.bytes += bytes;
622 self.traffic.update(&self.frame).expect("publish");
623 }
624
625 fn fail_traffic(&mut self) {
628 let mut group = self.track.append_group().expect("append group");
629 group
630 .write_frame(Timestamp::ZERO, b"not json".to_vec())
631 .expect("write frame");
632 group.finish().expect("finish group");
633 }
634 }
635
636 #[tokio::test(start_paused = true)]
637 async fn merges_traffic_across_nodes() {
638 let origin = produce_origin();
641 let node_a = node_producer(&origin, "a");
642 let node_b = node_producer(&origin, "b");
643
644 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
645 let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
646 drive_tick().await;
647
648 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
649 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
650
651 let frame = read_until_bytes(&mut traffic, "acme/room", 140).await;
652 let snap = frame.get("acme/room").expect("entry");
653 assert_eq!(snap.bytes, 140, "bytes sum across both nodes");
654 assert_eq!(snap.subscriptions_started, 2, "one subscription per node");
655 assert_eq!(snap.broadcasts_started, 2, "one viewer per node");
656 }
657
658 #[tokio::test(start_paused = true)]
659 async fn node_drop_keeps_the_traffic_total() {
660 let origin = produce_origin();
664 let node_a = node_producer(&origin, "a");
665 let node_b = node_producer(&origin, "b");
666
667 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
668 let fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
669 drive_tick().await;
670
671 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
672 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
673 read_until_bytes(&mut traffic, "acme/room", 140).await;
674
675 drop(fb);
677 drop(node_b);
678 drive_tick().await;
679
680 let _fa2 = feed(&node_a, Tier::default(), "acme", "acme/other", 10).await;
683 drive_tick().await;
684
685 let frame = read_until_bytes(&mut traffic, "acme/other", 10).await;
686 assert_eq!(
687 frame.get("acme/room").map(|t| t.bytes),
688 Some(140),
689 "the departed node's contribution stays in the total",
690 );
691 }
692
693 #[tokio::test(start_paused = true)]
694 async fn reannounce_with_higher_counter_stays_monotonic() {
695 let origin = produce_origin();
699 let node_a = node_producer(&origin, "a");
700 let node_b = node_producer(&origin, "b");
701
702 let fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
703 let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
704 drive_tick().await;
705
706 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
707 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
708 read_until_bytes(&mut traffic, "acme/room", 140).await;
709
710 drop(fa);
712 drop(node_a);
713 drive_tick().await;
714
715 let node_a = node_producer(&origin, "a");
716 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 120).await;
717 drive_tick().await;
718
719 let frame = read_monotonic_until(&mut traffic, "acme/room", 140, 160).await;
722 assert_eq!(frame.get("acme/room").expect("entry").bytes, 160);
723 }
724
725 #[tokio::test(start_paused = true)]
726 async fn restarted_node_regresses_the_total() {
727 let origin = produce_origin();
731 let node_a = node_producer(&origin, "a");
732 let node_b = node_producer(&origin, "b");
733
734 let fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
735 let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
736 drive_tick().await;
737
738 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
739 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
740 read_until_bytes(&mut traffic, "acme/room", 140).await;
741
742 drop(fa);
744 drop(node_a);
745 drive_tick().await;
746
747 let node_a = node_producer(&origin, "a");
748 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 30).await;
749 drive_tick().await;
750
751 loop {
755 let frame = traffic.next().await.expect("read").expect("frame");
756 if frame.get("acme/room").map(|t| t.bytes) == Some(70) {
757 break;
758 }
759 }
760 }
761
762 #[tokio::test(start_paused = true)]
763 async fn reader_failure_keeps_the_traffic_total() {
764 let origin = produce_origin();
767 let mut node_a = NodeBroadcast::new(&origin, "acme", "a");
768 let mut node_b = NodeBroadcast::new(&origin, "acme", "b");
769 node_a.publish("acme/room", 100);
770 node_b.publish("acme/room", 40);
771
772 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
773 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
774 read_until_bytes(&mut traffic, "acme/room", 140).await;
775
776 node_a.fail_traffic();
778
779 node_b.publish("acme/other", 10);
781
782 let frame = read_until_bytes(&mut traffic, "acme/other", 10).await;
783 assert_eq!(
784 frame.get("acme/room").map(|t| t.bytes),
785 Some(140),
786 "the failed node's contribution stays in the total",
787 );
788 }
789
790 #[tokio::test(start_paused = true)]
791 async fn unannounce_retires_live_gauges() {
792 let origin = produce_origin();
796 let mut node_a = NodeBroadcast::new(&origin, "acme", "a");
797
798 let mut published = Traffic::default();
799 published.announces_started = 2;
800 published.announces_ended = 1;
801 published.broadcasts_started = 3;
802 published.broadcasts_ended = 1;
803 published.subscriptions_started = 4;
804 published.subscriptions_ended = 1;
805 published.bytes = 100;
806 node_a.frame.insert("acme/room".to_string(), published);
807 node_a.traffic.update(&node_a.frame).expect("publish");
808
809 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
810 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
811 let frame = read_until_bytes(&mut traffic, "acme/room", 100).await;
812 let snap = frame.get("acme/room").expect("entry");
813 assert!(snap.is_announced());
814 assert_eq!(snap.active_broadcasts(), 2);
815 assert_eq!(snap.active_subscriptions(), 3);
816
817 drop(node_a);
819
820 let frame = traffic.next().await.expect("read").expect("frame");
822 let snap = frame.get("acme/room").expect("entry");
823 assert_eq!(snap.bytes, 100, "cumulative totals stay");
824 assert!(!snap.is_announced(), "no phantom announcement");
825 assert_eq!(snap.active_broadcasts(), 0, "no phantom broadcasts");
826 assert_eq!(snap.active_subscriptions(), 0, "no phantom subscriptions");
827 }
828
829 #[tokio::test(start_paused = true)]
830 async fn merges_sessions_across_nodes() {
831 let origin = produce_origin();
833 let node_a = node_producer(&origin, "a");
834 let node_b = node_producer(&origin, "b");
835
836 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 8).await;
839 let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 8).await;
840 let _sa = node_a.registry().tier(Tier::default()).session("acme");
841 let _sb = node_b.registry().tier(Tier::default()).session("acme");
842 drive_tick().await;
843
844 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
845 let mut sessions = agg.sessions(&Tier::default());
846
847 loop {
848 let frame = sessions.next().await.expect("read").expect("frame");
849 if frame.get("acme").map(|p| p.active()) >= Some(4) {
852 break;
853 }
854 }
855 }
856
857 #[tokio::test(start_paused = true)]
858 async fn unannounce_drops_presence_immediately() {
859 let origin = produce_origin();
862 let node_a = node_producer(&origin, "a");
863 let node_b = node_producer(&origin, "b");
864
865 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 8).await;
866 let fb = feed(&node_b, Tier::default(), "acme", "acme/room", 8).await;
867 let _sa = node_a.registry().tier(Tier::default()).session("acme");
868 let sb = node_b.registry().tier(Tier::default()).session("acme");
869 drive_tick().await;
870
871 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
872 let mut sessions = agg.sessions(&Tier::default());
873
874 loop {
875 let frame = sessions.next().await.expect("read").expect("frame");
876 if frame.get("acme").map(|p| p.active()) >= Some(4) {
879 break;
880 }
881 }
882
883 drop(fb);
885 drop(sb);
886 drop(node_b);
887 drive_tick().await;
888
889 let frame = sessions.next().await.expect("read").expect("frame");
890 assert_eq!(
891 frame.get("acme").map(|p| p.active()),
892 Some(2),
893 "presence drops immediately"
894 );
895 }
896}