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.clone().with_hidden(true).announced(),
273 origin,
274 prefix: config.prefix.clone(),
275 depth: config.depth,
276 name,
277 config: {
278 let mut json = moq_json::snapshot::consumer::Config::default();
279 if config.compression {
280 json.compression = moq_json::Compression::Deflate;
281 }
282 json
283 },
284 nodes: HashMap::new(),
285 }
286 }
287
288 fn poll_next(&mut self, waiter: &Waiter) -> Poll<Result<Option<BTreeMap<String, V>>>> {
293 let mut changed = false;
294
295 loop {
298 match self.announce.poll_next(waiter) {
299 Poll::Ready(Some(update)) => changed |= self.apply_announce(update),
300 Poll::Ready(None) => return Poll::Ready(Ok(None)),
301 Poll::Pending => break,
302 }
303 }
304
305 let config = &self.config;
307 let name = self.name.as_str();
308 let origin = &self.origin;
309 for node in self.nodes.values_mut() {
310 changed |= advance(node, origin, config, name, waiter);
311 }
312
313 if changed {
314 Poll::Ready(Ok(Some(self.merged())))
315 } else {
316 Poll::Pending
317 }
318 }
319
320 fn apply_announce(&mut self, update: moq_net::announce::Update) -> bool {
324 let path = update.prefix;
325 let absolute = self.announce.absolute(&path).to_owned();
326
327 if parse_node_path(&self.prefix, self.depth, &absolute).is_none() {
331 return false;
332 }
333
334 if update.kind.is_active() {
335 match self.nodes.entry(absolute) {
343 Entry::Occupied(mut entry) => {
344 let node = entry.get_mut();
345 if matches!(node.reader, Reader::Ended) {
346 node.reader = resolve(&self.origin, &node.path);
347 }
348 false
349 }
350 Entry::Vacant(entry) => {
351 entry.insert(Node {
352 reader: resolve(&self.origin, &path),
353 path,
354 last: None,
355 });
356 false
357 }
358 }
359 } else if V::STICKY {
360 match self.nodes.get_mut(&absolute) {
363 Some(node) => node.depart(),
364 None => false,
365 }
366 } else {
367 self.nodes.remove(&absolute).is_some_and(|old| old.last.is_some())
370 }
371 }
372
373 fn merged(&self) -> BTreeMap<String, V> {
375 let mut acc: BTreeMap<String, V> = BTreeMap::new();
376 for node in self.nodes.values() {
377 if let Some(last) = &node.last {
378 for (key, value) in last {
379 V::merge(acc.entry(key.clone()).or_default(), *value);
380 }
381 }
382 }
383 acc
384 }
385}
386
387fn advance<V: Mergeable>(
390 node: &mut Node<V>,
391 origin: &origin::Consumer,
392 config: &moq_json::snapshot::consumer::Config,
393 name: &str,
394 waiter: &Waiter,
395) -> bool {
396 let mut changed = false;
397 let mut rearmed = false;
400 loop {
401 match &mut node.reader {
402 Reader::Resolving { pending, queued } => match pending.poll_ok(waiter) {
403 Poll::Ready(Ok(broadcast)) => match broadcast.track(name) {
404 Ok(track) => node.reader = Reader::Subscribing(track.subscribe(None)),
405 Err(err) => {
406 tracing::debug!(?err, name, "stats: node missing track");
407 return changed | node.depart();
408 }
409 },
410 Poll::Ready(Err(moq_net::Error::Unroutable)) if *queued => {
417 node.reader = resolve(origin, &node.path);
418 }
419 Poll::Ready(Err(err)) => {
420 tracing::debug!(?err, name, "stats: node broadcast unresolvable");
421 return changed | node.depart();
422 }
423 Poll::Pending => return changed,
424 },
425 Reader::Subscribing(pending) => match pending.poll_ok(waiter) {
426 Poll::Ready(Ok(subscriber)) => {
427 node.reader =
428 Reader::Active(Box::new(moq_json::snapshot::Consumer::new(subscriber, config.clone())));
429 }
430 Poll::Ready(Err(err)) => {
431 tracing::debug!(?err, name, "stats: node subscribe failed");
432 return changed | node.depart();
433 }
434 Poll::Pending => return changed,
435 },
436 Reader::Active(reader) => match reader.poll_next(waiter) {
437 Poll::Ready(Ok(Some(frame))) => {
438 node.last = Some(frame);
439 changed = true;
440 }
441 Poll::Ready(result @ (Ok(None) | Err(_))) => {
448 if let Err(err) = result {
449 tracing::debug!(?err, name, "stats: node read error");
452 }
453 changed |= node.depart();
458 if rearmed {
459 return changed;
460 }
461 rearmed = true;
462 node.reader = resolve(origin, &node.path);
463 }
464 Poll::Pending => return changed,
465 },
466 Reader::Ended => return changed,
467 }
468 }
469}
470
471fn resolve<V: Mergeable>(origin: &origin::Consumer, path: &PathOwned) -> Reader<V> {
473 let pending = origin.request_broadcast(path);
474 let queued = pending.is_queued();
475 Reader::Resolving { pending, queued }
476}
477
478#[cfg(test)]
479mod tests {
480 fn produce_origin() -> moq_net::origin::Producer {
482 let (producer, driver) = moq_net::origin::Producer::new(moq_net::origin::Config::default());
483 if tokio::runtime::Handle::try_current().is_ok() {
484 tokio::spawn(moq_net::time::run(driver));
485 } else {
486 std::mem::forget(driver);
489 }
490 producer
491 }
492
493 use std::time::Duration;
494
495 use moq_net::{PathOwned, Timestamp, announce, broadcast, origin, track};
496
497 use crate::{Producer, produce};
498
499 use super::*;
500
501 fn node_producer(origin: &origin::Producer, node: &str) -> Producer {
505 Producer::new(
506 produce::Config::new()
507 .with_origin(origin.clone())
508 .with_node(PathOwned::from(node.to_string()))
509 .with_depth(1),
510 )
511 }
512
513 #[allow(dead_code)]
516 struct Feed {
517 announced: announce::Consumer,
518 source: broadcast::Producer,
519 consumer: broadcast::Consumer,
520 sub: track::Subscriber,
521 ctx: moq_net::stats::Session,
522 }
523
524 async fn feed(producer: &Producer, tier: Tier, root: &str, path: &str, bytes: usize) -> Feed {
528 let ctx = producer.registry().tier(tier).session(root);
529 let feed_origin = produce_origin();
530 let egress = feed_origin.consume().with_stats(ctx.clone());
531
532 let mut announced = egress.announced();
533 let source = feed_origin.create_broadcast(path).expect("create_broadcast");
534 source.announce(origin::Route::default()).expect("announce");
535 let track = source.create_track("video", None).expect("create_track");
536
537 let update = announced.next().await.expect("announce");
538 assert!(update.kind.is_active());
539 let consumer = egress.request_broadcast(path).await.expect("resolve");
540 let mut sub = consumer.track("video").unwrap().subscribe(None).await.unwrap();
541
542 let mut group = track.append_group().unwrap();
543 group.write_frame(Timestamp::ZERO, vec![0u8; bytes]).unwrap();
544 group.finish().unwrap();
545 let mut group = sub.recv_group().await.unwrap().unwrap();
546 while group.read_frame().await.unwrap().is_some() {}
547
548 Feed {
549 announced,
550 source,
551 consumer,
552 sub,
553 ctx,
554 }
555 }
556
557 async fn drive_tick() {
559 tokio::time::advance(Duration::from_millis(1100)).await;
560 for _ in 0..8 {
561 tokio::task::yield_now().await;
562 }
563 }
564
565 async fn read_until_bytes(consumer: &mut TrafficConsumer, path: &str, want: u64) -> TrafficFrame {
568 loop {
569 let frame = consumer.next().await.expect("read").expect("frame");
570 if frame.get(path).map(|t| t.bytes).unwrap_or(0) >= want {
571 return frame;
572 }
573 }
574 }
575
576 async fn read_monotonic_until(consumer: &mut TrafficConsumer, path: &str, min: u64, want: u64) -> TrafficFrame {
580 loop {
581 let frame = consumer.next().await.expect("read").expect("frame");
582 let bytes = frame.get(path).map(|t| t.bytes).unwrap_or(0);
583 assert!(bytes >= min, "traffic regressed below {min}: {bytes}");
584 if bytes >= want {
585 return frame;
586 }
587 }
588 }
589
590 #[allow(dead_code)]
595 struct NodeBroadcast {
596 source: broadcast::Producer,
597 traffic: moq_json::snapshot::Producer<TrafficFrame>,
598 track: track::Producer,
599 frame: TrafficFrame,
600 }
601
602 impl NodeBroadcast {
603 fn new(origin: &origin::Producer, group: &str, node: &str) -> Self {
604 let path = format!(".stats/{group}/node/{node}");
605 let source = origin.create_broadcast(path.as_str()).expect("create broadcast");
606 source.announce(origin::Route::default()).expect("announce");
607 let name = traffic_track(&Tier::default(), Role::Publisher, false);
608 let track = source.create_track(name, None).expect("create track");
609 let config = moq_json::snapshot::Config::default().with_delta_ratio(0);
610 Self {
611 traffic: moq_json::snapshot::Producer::new(track.clone(), config),
612 track,
613 source,
614 frame: TrafficFrame::new(),
615 }
616 }
617
618 fn publish(&mut self, path: &str, bytes: u64) {
621 let entry = self.frame.entry(path.to_string()).or_default();
622 entry.bytes += bytes;
623 self.traffic.update(&self.frame).expect("publish");
624 }
625
626 fn fail_traffic(&mut self) {
629 let mut group = self.track.append_group().expect("append group");
630 group
631 .write_frame(Timestamp::ZERO, b"not json".to_vec())
632 .expect("write frame");
633 group.finish().expect("finish group");
634 }
635 }
636
637 #[tokio::test(start_paused = true)]
638 async fn merges_traffic_across_nodes() {
639 let origin = produce_origin();
642 let node_a = node_producer(&origin, "a");
643 let node_b = node_producer(&origin, "b");
644
645 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
646 let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
647 drive_tick().await;
648
649 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
650 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
651
652 let frame = read_until_bytes(&mut traffic, "acme/room", 140).await;
653 let snap = frame.get("acme/room").expect("entry");
654 assert_eq!(snap.bytes, 140, "bytes sum across both nodes");
655 assert_eq!(snap.subscriptions_started, 2, "one subscription per node");
656 assert_eq!(snap.broadcasts_started, 2, "one viewer per node");
657 }
658
659 #[tokio::test(start_paused = true)]
660 async fn node_drop_keeps_the_traffic_total() {
661 let origin = produce_origin();
665 let node_a = node_producer(&origin, "a");
666 let node_b = node_producer(&origin, "b");
667
668 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
669 let fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
670 drive_tick().await;
671
672 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
673 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
674 read_until_bytes(&mut traffic, "acme/room", 140).await;
675
676 drop(fb);
678 drop(node_b);
679 drive_tick().await;
680
681 let _fa2 = feed(&node_a, Tier::default(), "acme", "acme/other", 10).await;
684 drive_tick().await;
685
686 let frame = read_until_bytes(&mut traffic, "acme/other", 10).await;
687 assert_eq!(
688 frame.get("acme/room").map(|t| t.bytes),
689 Some(140),
690 "the departed node's contribution stays in the total",
691 );
692 }
693
694 #[tokio::test(start_paused = true)]
695 async fn reannounce_with_higher_counter_stays_monotonic() {
696 let origin = produce_origin();
700 let node_a = node_producer(&origin, "a");
701 let node_b = node_producer(&origin, "b");
702
703 let fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
704 let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
705 drive_tick().await;
706
707 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
708 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
709 read_until_bytes(&mut traffic, "acme/room", 140).await;
710
711 drop(fa);
713 drop(node_a);
714 drive_tick().await;
715
716 let node_a = node_producer(&origin, "a");
717 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 120).await;
718 drive_tick().await;
719
720 let frame = read_monotonic_until(&mut traffic, "acme/room", 140, 160).await;
723 assert_eq!(frame.get("acme/room").expect("entry").bytes, 160);
724 }
725
726 #[tokio::test(start_paused = true)]
727 async fn restarted_node_regresses_the_total() {
728 let origin = produce_origin();
732 let node_a = node_producer(&origin, "a");
733 let node_b = node_producer(&origin, "b");
734
735 let fa = feed(&node_a, Tier::default(), "acme", "acme/room", 100).await;
736 let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 40).await;
737 drive_tick().await;
738
739 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
740 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
741 read_until_bytes(&mut traffic, "acme/room", 140).await;
742
743 drop(fa);
745 drop(node_a);
746 drive_tick().await;
747
748 let node_a = node_producer(&origin, "a");
749 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 30).await;
750 drive_tick().await;
751
752 loop {
756 let frame = traffic.next().await.expect("read").expect("frame");
757 if frame.get("acme/room").map(|t| t.bytes) == Some(70) {
758 break;
759 }
760 }
761 }
762
763 #[tokio::test(start_paused = true)]
764 async fn reader_failure_keeps_the_traffic_total() {
765 let origin = produce_origin();
768 let mut node_a = NodeBroadcast::new(&origin, "acme", "a");
769 let mut node_b = NodeBroadcast::new(&origin, "acme", "b");
770 node_a.publish("acme/room", 100);
771 node_b.publish("acme/room", 40);
772
773 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
774 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
775 read_until_bytes(&mut traffic, "acme/room", 140).await;
776
777 node_a.fail_traffic();
779
780 node_b.publish("acme/other", 10);
782
783 let frame = read_until_bytes(&mut traffic, "acme/other", 10).await;
784 assert_eq!(
785 frame.get("acme/room").map(|t| t.bytes),
786 Some(140),
787 "the failed node's contribution stays in the total",
788 );
789 }
790
791 #[tokio::test(start_paused = true)]
792 async fn unannounce_retires_live_gauges() {
793 let origin = produce_origin();
797 let mut node_a = NodeBroadcast::new(&origin, "acme", "a");
798
799 let mut published = Traffic::default();
800 published.announces_started = 2;
801 published.announces_ended = 1;
802 published.broadcasts_started = 3;
803 published.broadcasts_ended = 1;
804 published.subscriptions_started = 4;
805 published.subscriptions_ended = 1;
806 published.bytes = 100;
807 node_a.frame.insert("acme/room".to_string(), published);
808 node_a.traffic.update(&node_a.frame).expect("publish");
809
810 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
811 let mut traffic = agg.traffic(&Tier::default(), Role::Publisher);
812 let frame = read_until_bytes(&mut traffic, "acme/room", 100).await;
813 let snap = frame.get("acme/room").expect("entry");
814 assert!(snap.is_announced());
815 assert_eq!(snap.active_broadcasts(), 2);
816 assert_eq!(snap.active_subscriptions(), 3);
817
818 drop(node_a);
820
821 let frame = traffic.next().await.expect("read").expect("frame");
823 let snap = frame.get("acme/room").expect("entry");
824 assert_eq!(snap.bytes, 100, "cumulative totals stay");
825 assert!(!snap.is_announced(), "no phantom announcement");
826 assert_eq!(snap.active_broadcasts(), 0, "no phantom broadcasts");
827 assert_eq!(snap.active_subscriptions(), 0, "no phantom subscriptions");
828 }
829
830 #[tokio::test(start_paused = true)]
831 async fn merges_sessions_across_nodes() {
832 let origin = produce_origin();
834 let node_a = node_producer(&origin, "a");
835 let node_b = node_producer(&origin, "b");
836
837 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 8).await;
840 let _fb = feed(&node_b, Tier::default(), "acme", "acme/room", 8).await;
841 let _sa = node_a.registry().tier(Tier::default()).session("acme");
842 let _sb = node_b.registry().tier(Tier::default()).session("acme");
843 drive_tick().await;
844
845 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
846 let mut sessions = agg.sessions(&Tier::default());
847
848 loop {
849 let frame = sessions.next().await.expect("read").expect("frame");
850 if frame.get("acme").map(|p| p.active()) >= Some(4) {
853 break;
854 }
855 }
856 }
857
858 #[tokio::test(start_paused = true)]
859 async fn unannounce_drops_presence_immediately() {
860 let origin = produce_origin();
863 let node_a = node_producer(&origin, "a");
864 let node_b = node_producer(&origin, "b");
865
866 let _fa = feed(&node_a, Tier::default(), "acme", "acme/room", 8).await;
867 let fb = feed(&node_b, Tier::default(), "acme", "acme/room", 8).await;
868 let _sa = node_a.registry().tier(Tier::default()).session("acme");
869 let sb = node_b.registry().tier(Tier::default()).session("acme");
870 drive_tick().await;
871
872 let agg = Consumer::new(origin.consume(), Config::new().with_depth(1));
873 let mut sessions = agg.sessions(&Tier::default());
874
875 loop {
876 let frame = sessions.next().await.expect("read").expect("frame");
877 if frame.get("acme").map(|p| p.active()) >= Some(4) {
880 break;
881 }
882 }
883
884 drop(fb);
886 drop(sb);
887 drop(node_b);
888 drive_tick().await;
889
890 let frame = sessions.next().await.expect("read").expect("frame");
891 assert_eq!(
892 frame.get("acme").map(|p| p.active()),
893 Some(2),
894 "presence drops immediately"
895 );
896 }
897}