1use std::collections::{HashMap, HashSet};
4use std::sync::{Arc, Weak};
5use std::time::Duration;
6
7use std::task::Poll;
8
9use moq_net::stats::{Presence, Registry, Role, Tier, Traffic};
10use moq_net::{Path, PathOwned, broadcast, kio, origin, track};
11use serde::Serialize;
12use web_async::spawn;
13
14use crate::{COMPRESSED_SUFFIX, SessionsFrame, TrafficFrame, sessions_track, traffic_track};
15
16#[derive(Clone)]
25#[non_exhaustive]
26pub struct ProducerConfig {
27 pub origin: Option<origin::Producer>,
30 pub prefix: PathOwned,
35 pub node: Option<PathOwned>,
41 pub interval: Duration,
43 pub depth: usize,
51}
52
53impl ProducerConfig {
54 pub fn new() -> Self {
58 Self {
59 origin: None,
60 prefix: PathOwned::from(".stats"),
61 node: None,
62 interval: Duration::from_secs(1),
63 depth: 0,
64 }
65 }
66
67 pub fn with_origin(mut self, origin: impl Into<Option<origin::Producer>>) -> Self {
70 self.origin = origin.into();
71 self
72 }
73
74 pub fn with_prefix(mut self, prefix: impl Into<PathOwned>) -> Self {
76 self.prefix = prefix.into();
77 self
78 }
79
80 pub fn with_interval(mut self, interval: Duration) -> Self {
82 self.interval = interval;
83 self
84 }
85
86 pub fn with_node(mut self, node: impl Into<Option<PathOwned>>) -> Self {
88 self.node = node.into();
89 self
90 }
91
92 pub fn with_depth(mut self, depth: usize) -> Self {
94 self.depth = depth;
95 self
96 }
97}
98
99impl Default for ProducerConfig {
100 fn default() -> Self {
101 Self::new()
102 }
103}
104
105const MAX_REQUESTED_TRACKS: usize = 64;
115
116const MAX_PARKED_REQUESTS: usize = 256;
123
124struct Keepalive;
127
128#[derive(Clone)]
136pub struct Producer {
137 registry: Registry,
138 _keepalive: Option<Arc<Keepalive>>,
141}
142
143impl Producer {
144 pub fn new(config: ProducerConfig) -> Self {
152 let ProducerConfig {
153 origin,
154 prefix,
155 node,
156 interval,
157 depth,
158 } = config;
159 let node = node.filter(|p| !p.is_empty());
164
165 let Some(origin) = origin else {
166 return Self {
167 registry: Registry::disabled(),
168 _keepalive: None,
169 };
170 };
171
172 let registry = Registry::new(moq_net::stats::Config::new().with_exclude(prefix.clone()));
173 let keepalive = Arc::new(Keepalive);
174 let task = Task {
175 registry: registry.clone(),
176 origin,
177 prefix,
178 node,
179 depth,
180 interval,
181 };
182 spawn(task.run(Arc::downgrade(&keepalive)));
183
184 Self {
185 registry,
186 _keepalive: Some(keepalive),
187 }
188 }
189
190 pub fn registry(&self) -> &Registry {
194 &self.registry
195 }
196}
197
198struct Task {
200 registry: Registry,
201 origin: origin::Producer,
202 prefix: PathOwned,
203 node: Option<PathOwned>,
204 depth: usize,
205 interval: Duration,
206}
207
208impl Task {
209 async fn run(self, weak: Weak<Keepalive>) {
212 let node = self.node.as_ref().map(moq_net::Path::as_str);
213 let mut groups: HashMap<PathOwned, GroupPublisher> = HashMap::new();
214
215 if self.depth == 0 {
216 let Some(group) = GroupPublisher::create(&self.origin, &self.prefix, &Path::empty(), node) else {
217 return;
218 };
219 groups.insert(Path::empty().to_owned(), group);
220 }
221
222 let mut ticker = web_async::time::interval(self.interval);
223 ticker.set_missed_tick_behavior(web_async::time::MissedTickBehavior::Delay);
224
225 loop {
226 ticker.tick().await;
227
228 if weak.upgrade().is_none() {
229 for (_, publisher) in groups.drain() {
230 publisher.finish();
231 }
232 return;
233 }
234
235 let report = self.registry.report();
238
239 let mut entries_by_group: HashMap<PathOwned, Vec<&moq_net::stats::TrafficEntry>> = HashMap::new();
240 for entry in &report.traffic {
241 entries_by_group
242 .entry(group_key(entry.path.as_str(), self.depth))
243 .or_default()
244 .push(entry);
245 }
246
247 let mut sessions_by_group: HashMap<PathOwned, Vec<&moq_net::stats::SessionEntry>> = HashMap::new();
248 for entry in &report.sessions {
249 sessions_by_group
250 .entry(group_key(entry.root.as_str(), self.depth))
251 .or_default()
252 .push(entry);
253 }
254
255 let mut active: HashSet<PathOwned> = HashSet::new();
256 active.extend(entries_by_group.keys().cloned());
257 active.extend(sessions_by_group.keys().cloned());
258 if self.depth == 0 {
259 active.insert(Path::empty().to_owned());
260 }
261
262 for group in &active {
263 if !groups.contains_key(group) {
264 let Some(publisher) = GroupPublisher::create(&self.origin, &self.prefix, group, node) else {
265 continue;
266 };
267 groups.insert(group.clone(), publisher);
268 }
269 let publisher = groups.get_mut(group).expect("just inserted");
270
271 let mut frames: HashMap<String, TrafficFrame> = HashMap::new();
272 if let Some(group_entries) = entries_by_group.get(group) {
273 for entry in group_entries {
274 let slots = publisher
275 .local
276 .entry(entry.path.clone())
277 .or_default()
278 .entry(entry.tier.clone())
279 .or_default();
280 process_slot(entry.publisher, &mut slots.publisher, |snap| {
281 frames
282 .entry(traffic_track(&entry.tier, Role::Publisher, false))
283 .or_default()
284 .insert(entry.path.as_str().to_string(), snap);
285 });
286 process_slot(entry.subscriber, &mut slots.subscriber, |snap| {
287 frames
288 .entry(traffic_track(&entry.tier, Role::Subscriber, false))
289 .or_default()
290 .insert(entry.path.as_str().to_string(), snap);
291 });
292 }
293 }
294
295 let mut session_frames: HashMap<String, SessionsFrame> = HashMap::new();
296 if let Some(group_sessions) = sessions_by_group.get(group) {
297 for entry in group_sessions {
298 let state = publisher
299 .session_local
300 .entry(entry.tier.clone())
301 .or_default()
302 .entry(entry.root.clone())
303 .or_default();
304 process_session_slot(entry.presence, state, |snap| {
305 session_frames
306 .entry(sessions_track(&entry.tier, false))
307 .or_default()
308 .insert(entry.root.as_str().to_string(), snap);
309 });
310 }
311 }
312
313 for name in frames.keys().chain(session_frames.keys()) {
317 publisher.requested.remove(name);
318 }
319
320 publisher.traffic.flush(&mut publisher.broadcast, &frames);
321 publisher.sessions.flush(&mut publisher.broadcast, &session_frames);
322 }
323
324 for publisher in groups.values_mut() {
330 publisher.serve_requests();
331 }
332
333 let reported: HashSet<(&PathOwned, &Tier)> =
336 report.traffic.iter().map(|entry| (&entry.path, &entry.tier)).collect();
337 let reported_sessions: HashSet<(&Tier, &PathOwned)> =
338 report.sessions.iter().map(|entry| (&entry.tier, &entry.root)).collect();
339 for publisher in groups.values_mut() {
340 publisher.local.retain(|path, tiers| {
341 tiers.retain(|tier, _| reported.contains(&(path, tier)));
342 !tiers.is_empty()
343 });
344 publisher.session_local.retain(|tier, roots| {
345 roots.retain(|root, _| reported_sessions.contains(&(tier, root)));
346 !roots.is_empty()
347 });
348 }
349
350 let evicted: Vec<PathOwned> = groups
354 .keys()
355 .filter(|group| !active.contains(*group))
356 .cloned()
357 .collect();
358 for group in evicted {
359 if let Some(publisher) = groups.remove(&group) {
360 publisher.finish();
361 }
362 }
363 }
364 }
365}
366
367struct TrackPair<T> {
372 plain: moq_json::snapshot::Producer<T>,
373 compressed: moq_json::snapshot::Producer<T>,
374}
375
376impl<T: Serialize> TrackPair<T> {
377 fn create(broadcast: &mut broadcast::Producer, name: &str) -> Result<Self, moq_net::Error> {
378 let plain_track = broadcast.create_track(name, None)?;
379 let compressed_track = broadcast.create_track(format!("{name}{COMPRESSED_SUFFIX}").as_str(), None)?;
380 Ok(Self::from_tracks(plain_track, compressed_track))
381 }
382
383 fn adopt(broadcast: &mut broadcast::Producer, name: &str, pending: PendingPair) -> Result<Self, moq_net::Error> {
389 let PendingPair { plain, compressed } = pending;
390 let plain_track = match plain {
391 Some(request) => request.accept(None),
392 None => broadcast.create_track(name, None)?,
393 };
394 let compressed_track = match compressed {
395 Some(request) => request.accept(None),
396 None => broadcast.create_track(format!("{name}{COMPRESSED_SUFFIX}").as_str(), None)?,
397 };
398 Ok(Self::from_tracks(plain_track, compressed_track))
399 }
400
401 fn from_tracks(plain_track: track::Producer, compressed_track: track::Producer) -> Self {
402 let plain_config = moq_json::snapshot::ProducerConfig::default().with_delta_ratio(0);
403 let compressed_config = moq_json::snapshot::ProducerConfig::default().with_compression(true);
404
405 Self {
406 plain: moq_json::snapshot::Producer::new(plain_track, plain_config),
407 compressed: moq_json::snapshot::Producer::new(compressed_track, compressed_config),
408 }
409 }
410
411 fn is_used(&self) -> bool {
413 self.plain.is_used() || self.compressed.is_used()
414 }
415
416 fn update(&mut self, name: &str, frame: &T) {
418 if let Err(err) = self.plain.update(frame) {
419 tracing::debug!(?err, name, "stats: failed to write frame");
420 }
421 if let Err(err) = self.compressed.update(frame) {
422 tracing::debug!(?err, name, "stats: failed to write compressed frame");
423 }
424 }
425
426 fn finish(&mut self) {
430 let _ = self.plain.finish();
431 let _ = self.compressed.finish();
432 }
433}
434
435#[derive(Default)]
438struct PendingPair {
439 plain: Option<track::Request>,
440 compressed: Option<track::Request>,
441}
442
443impl PendingPair {
444 fn reject(self, err: moq_net::Error) {
445 if let Some(request) = self.plain {
446 request.reject(err.clone());
447 }
448 if let Some(request) = self.compressed {
449 request.reject(err);
450 }
451 }
452
453 fn is_used(&self, waiter: &kio::Waiter) -> bool {
458 self.plain
459 .iter()
460 .chain(self.compressed.iter())
461 .any(|request| request.poll_unused(waiter).is_pending())
462 }
463}
464
465struct TrackFamily<T> {
468 tracks: HashMap<String, TrackPair<T>>,
469 parked: HashMap<String, PendingPair>,
473}
474
475impl<T: Serialize + Default> TrackFamily<T> {
476 fn new() -> Self {
477 Self {
478 tracks: HashMap::new(),
479 parked: HashMap::new(),
480 }
481 }
482
483 fn flush(&mut self, broadcast: &mut broadcast::Producer, frames: &HashMap<String, T>) {
492 for name in frames.keys() {
493 if !self.tracks.contains_key(name) {
494 let result = match self.parked.remove(name) {
495 Some(pending) => TrackPair::adopt(broadcast, name, pending),
496 None => TrackPair::create(broadcast, name),
497 };
498 match result {
499 Ok(pair) => {
500 self.tracks.insert(name.clone(), pair);
501 }
502 Err(err) => tracing::warn!(?err, name, "stats: failed to create track"),
503 }
504 }
505 }
506
507 let empty = T::default();
508 for (name, pair) in self.tracks.iter_mut() {
509 pair.update(name, frames.get(name).unwrap_or(&empty));
510 }
511 }
512
513 fn reclaim(&mut self, requested: &mut HashSet<String>) {
519 self.tracks.retain(|name, pair| {
520 if !requested.contains(name) || pair.is_used() {
521 return true;
522 }
523 requested.remove(name);
524 pair.finish();
525 false
526 });
527 }
528
529 fn park(&mut self, plain: String, compressed: bool, request: track::Request, full: bool) {
532 match self.parked.get_mut(&plain) {
533 Some(pending) => {
534 let slot = match compressed {
535 true => &mut pending.compressed,
536 false => &mut pending.plain,
537 };
538 if slot.is_none() {
543 *slot = Some(request);
544 }
545 }
546 None if full => request.reject(moq_net::Error::NotFound),
547 None => {
548 let mut pending = PendingPair::default();
549 match compressed {
550 true => pending.compressed = Some(request),
551 false => pending.plain = Some(request),
552 }
553 self.parked.insert(plain, pending);
554 }
555 }
556 }
557
558 fn adopt_parked(&mut self, broadcast: &mut broadcast::Producer, requested: &mut HashSet<String>) {
563 let noop = kio::Waiter::noop();
564 let mut parked = std::mem::take(&mut self.parked);
565 parked.retain(|plain, pending| {
566 if !pending.is_used(&noop) {
567 return false;
568 }
569 if requested.len() >= MAX_REQUESTED_TRACKS {
570 return true;
571 }
572 self.adopt_pair(broadcast, requested, plain.clone(), std::mem::take(pending));
573 false
574 });
575 self.parked = parked;
576 }
577
578 fn adopt_pair(
582 &mut self,
583 broadcast: &mut broadcast::Producer,
584 requested: &mut HashSet<String>,
585 plain: String,
586 pending: PendingPair,
587 ) {
588 if self.tracks.contains_key(&plain) {
593 pending.reject(moq_net::Error::NotFound);
594 return;
595 }
596 match TrackPair::adopt(broadcast, &plain, pending) {
597 Ok(mut pair) => {
598 pair.update(&plain, &T::default());
600 self.tracks.insert(plain.clone(), pair);
601 requested.insert(plain);
602 }
603 Err(err) => tracing::warn!(?err, name = %plain, "stats: failed to adopt requested track"),
604 }
605 }
606
607 fn finish(&mut self) {
609 for pair in self.tracks.values_mut() {
610 pair.finish();
611 }
612 }
613}
614
615struct GroupPublisher {
617 broadcast: broadcast::Producer,
618 dynamic: broadcast::Dynamic,
622 requested: HashSet<String>,
627 traffic: TrackFamily<TrafficFrame>,
628 sessions: TrackFamily<SessionsFrame>,
629 local: HashMap<PathOwned, HashMap<Tier, SideSlots>>,
630 session_local: HashMap<Tier, HashMap<PathOwned, SessionSlotState>>,
631}
632
633impl GroupPublisher {
634 fn create(origin: &origin::Producer, prefix: &Path, group: &Path, node: Option<&str>) -> Option<Self> {
635 let advertised = advertised_path(prefix, group, node);
636 let mut broadcast = match origin.create_broadcast(&advertised, broadcast::Route::new().with_announce(true)) {
637 Ok(broadcast) => broadcast,
638 Err(err) => {
639 tracing::warn!(advertised = %advertised, ?err, "stats: origin rejected stats broadcast");
640 return None;
641 }
642 };
643 tracing::debug!(advertised = %advertised, "stats: publishing broadcast");
644
645 let mut traffic = TrackFamily::new();
646 let mut sessions = TrackFamily::new();
647
648 let tier = Tier::default();
650 for role in [Role::Publisher, Role::Subscriber] {
651 let name = traffic_track(&tier, role, false);
652 match TrackPair::create(&mut broadcast, &name) {
653 Ok(pair) => {
654 traffic.tracks.insert(name, pair);
655 }
656 Err(err) => {
657 tracing::warn!(?err, name, "stats: failed to create track");
658 return None;
659 }
660 }
661 }
662 let name = sessions_track(&tier, false);
663 match TrackPair::create(&mut broadcast, &name) {
664 Ok(pair) => {
665 sessions.tracks.insert(name, pair);
666 }
667 Err(err) => {
668 tracing::warn!(?err, name, "stats: failed to create track");
669 return None;
670 }
671 }
672
673 let dynamic = broadcast.dynamic();
674
675 Some(Self {
676 broadcast,
677 dynamic,
678 requested: HashSet::new(),
679 traffic,
680 sessions,
681 local: HashMap::new(),
682 session_local: HashMap::new(),
683 })
684 }
685
686 fn serve_requests(&mut self) {
698 self.traffic.reclaim(&mut self.requested);
701 self.sessions.reclaim(&mut self.requested);
702
703 let noop = kio::Waiter::noop();
707 while let Poll::Ready(Ok(request)) = self.dynamic.poll_requested_track(&noop) {
708 let Some(shape) = requested_track_shape(request.name()) else {
709 request.reject(moq_net::Error::NotFound);
710 continue;
711 };
712 let full = self.traffic.parked.len() + self.sessions.parked.len() >= MAX_PARKED_REQUESTS;
713 match shape.sessions {
714 true => self.sessions.park(shape.plain, shape.compressed, request, full),
715 false => self.traffic.park(shape.plain, shape.compressed, request, full),
716 }
717 }
718
719 self.traffic.adopt_parked(&mut self.broadcast, &mut self.requested);
720 self.sessions.adopt_parked(&mut self.broadcast, &mut self.requested);
721 }
722
723 fn finish(mut self) {
726 self.traffic.finish();
727 self.sessions.finish();
728 self.broadcast.finish();
729 }
730}
731
732struct RequestedShape {
734 plain: String,
736 compressed: bool,
738 sessions: bool,
740}
741
742fn requested_track_shape(name: &str) -> Option<RequestedShape> {
746 let (base, compressed) = match name.strip_suffix(COMPRESSED_SUFFIX) {
747 Some(base) => (base, true),
748 None => (name, false),
749 };
750 let (tier, kind) = match base.rsplit_once('/') {
751 Some((tier, kind)) => (Some(tier), kind),
752 None => (None, base),
753 };
754 let sessions = match kind {
755 "publisher.json" | "subscriber.json" => false,
756 "sessions.json" => true,
757 _ => return None,
758 };
759 if let Some(tier) = tier
762 && (tier.is_empty() || tier.starts_with('/') || tier.ends_with('/') || tier.contains("//"))
763 {
764 return None;
765 }
766 Some(RequestedShape {
767 plain: base.to_string(),
768 compressed,
769 sessions,
770 })
771}
772
773#[derive(Default)]
776struct SlotState {
777 prev_emitted: Option<Traffic>,
780}
781
782#[derive(Default)]
784struct SideSlots {
785 publisher: SlotState,
786 subscriber: SlotState,
787}
788
789#[derive(Default)]
791struct SessionSlotState {
792 prev_emitted: Option<Presence>,
793}
794
795fn process_slot(snap: Traffic, slot_state: &mut SlotState, emit: impl FnOnce(Traffic)) {
799 let live = !snap.is_idle();
806
807 let prev_snap = slot_state.prev_emitted.unwrap_or_default();
818 let changed = snap != prev_snap;
819 if changed {
820 slot_state.prev_emitted = Some(snap);
821 }
822 if live || changed {
823 emit(snap);
824 }
825}
826
827fn process_session_slot(snap: Presence, slot_state: &mut SessionSlotState, emit: impl FnOnce(Presence)) {
830 let live = snap.active() > 0;
831 let prev_snap = slot_state.prev_emitted.unwrap_or_default();
832 let changed = snap != prev_snap;
833 if changed {
834 slot_state.prev_emitted = Some(snap);
835 }
836 if live || changed {
837 emit(snap);
838 }
839}
840
841fn group_key(path: &str, depth: usize) -> PathOwned {
842 if depth == 0 {
843 return Path::empty().to_owned();
844 }
845
846 let mut seen = 0;
847 let mut end = path.len();
848 for (i, b) in path.bytes().enumerate() {
849 if b == b'/' {
850 seen += 1;
851 if seen == depth {
852 end = i;
853 break;
854 }
855 }
856 }
857 Path::new(&path[..end]).to_owned()
858}
859
860fn advertised_path(prefix: &Path, group: &Path, node: Option<&str>) -> PathOwned {
861 let mut out = prefix.as_str().to_string();
865 if !group.is_empty() {
866 out.push('/');
867 out.push_str(group.as_str());
868 }
869 out.push_str("/node");
870 if let Some(node) = node {
871 out.push('/');
872 out.push_str(node);
873 }
874 PathOwned::from(out)
875}
876
877#[cfg(test)]
878mod tests {
879 use std::collections::BTreeMap;
880
881 use moq_net::stats::{Registry, Tier};
882 use moq_net::{Origin, Timestamp, announce, broadcast, track};
883
884 use super::*;
885
886 fn test_producer(node: Option<&str>) -> (Producer, origin::Producer) {
887 let origin = Origin::random().produce();
888 let producer = Producer::new(
889 ProducerConfig::new()
890 .with_origin(origin.clone())
891 .with_node(node.map(|s| PathOwned::from(s.to_string()))),
892 );
893 (producer, origin)
894 }
895
896 #[allow(dead_code)]
899 struct Feed {
900 announced: announce::Consumer,
901 source: broadcast::Producer,
902 consumer: broadcast::Consumer,
903 sub: Option<track::Subscriber>,
904 }
905
906 async fn feed(
913 registry: &Registry,
914 tier: Tier,
915 path: &str,
916 subscribe: bool,
917 frames: usize,
918 frame_size: usize,
919 ) -> Feed {
920 let ctx = registry.tier(tier).session("feed");
921 let origin = Origin::random().produce();
922 let egress = origin.consume().with_stats(ctx);
924
925 let mut announced = egress.announced();
926 let mut source = origin
927 .create_broadcast(path, broadcast::Route::announced())
928 .expect("create_broadcast");
929 let mut producer = source.create_track("video", None).expect("create_track");
930
931 tokio::time::sleep(Duration::from_millis(1)).await;
934 tokio::time::sleep(Duration::from_millis(1)).await;
935
936 let announce::Update { broadcast, .. } = announced.next().await.expect("announce");
937 let consumer = broadcast.expect("active");
938
939 let sub = if subscribe {
940 let mut sub = consumer
941 .track("video")
942 .expect("track")
943 .subscribe(None)
944 .await
945 .expect("subscribe");
946
947 if frames > 0 {
948 let mut group = producer.append_group().expect("group");
949 for _ in 0..frames {
950 group
951 .write_frame(Timestamp::ZERO, vec![0u8; frame_size])
952 .expect("write");
953 }
954 group.finish().expect("finish");
955
956 let mut group = sub.recv_group().await.expect("recv").expect("group");
957 while group.read_frame().await.expect("read").is_some() {}
958 }
959 Some(sub)
960 } else {
961 None
962 };
963
964 Feed {
965 announced,
966 source,
967 consumer,
968 sub,
969 }
970 }
971
972 async fn announced(origin: &origin::Producer) -> (String, moq_net::broadcast::Consumer) {
974 let mut consumer = origin.consume().announced();
975 tokio::time::advance(Duration::from_millis(1)).await;
976 let announce::Update { path, broadcast } = consumer.next().await.expect("expected announce");
977 (path.as_str().to_string(), broadcast.expect("active"))
978 }
979
980 async fn drive_tick() {
982 tokio::time::advance(Duration::from_millis(1100)).await;
983 for _ in 0..4 {
986 tokio::task::yield_now().await;
987 }
988 }
989
990 async fn read_frame(broadcast: &moq_net::broadcast::Consumer, name: &str) -> BTreeMap<String, Traffic> {
993 let mut track = subscribe(broadcast, name).await;
994 let frame = track.read_frame().await.expect("ok").expect("frame");
995 serde_json::from_slice(&frame.payload).expect("json parse")
996 }
997
998 async fn read_last_frame(broadcast: &moq_net::broadcast::Consumer, name: &str) -> BTreeMap<String, Traffic> {
1002 use futures::FutureExt;
1003 let mut track = subscribe(broadcast, name).await;
1004 let mut last = track.read_frame().await.expect("ok").expect("frame");
1005 while let Some(Ok(Some(frame))) = track.read_frame().now_or_never() {
1006 last = frame;
1007 }
1008 serde_json::from_slice(&last.payload).expect("json parse")
1009 }
1010
1011 async fn read_session_frame(broadcast: &moq_net::broadcast::Consumer, name: &str) -> BTreeMap<String, Presence> {
1012 let mut track = subscribe(broadcast, name).await;
1013 let frame = track.read_frame().await.expect("ok").expect("frame");
1014 serde_json::from_slice(&frame.payload).expect("json parse")
1015 }
1016
1017 async fn subscribe(broadcast: &moq_net::broadcast::Consumer, name: &str) -> track::Subscriber {
1018 broadcast
1019 .track(name)
1020 .expect("track")
1021 .subscribe(None)
1022 .await
1023 .expect("subscribe")
1024 }
1025
1026 #[tokio::test(start_paused = true)]
1030 async fn new_normalizes_and_drops_empty_node() {
1031 let (_producer, origin) = test_producer(Some("/sjc//1/"));
1032 assert_eq!(announced(&origin).await.0, ".stats/node/sjc/1");
1033
1034 let (_producer, origin) = test_producer(Some("///"));
1035 assert_eq!(announced(&origin).await.0, ".stats/node");
1036 }
1037
1038 #[tokio::test(start_paused = true)]
1039 async fn single_broadcast_path_announced() {
1040 let (producer, origin) = test_producer(Some("sjc/1"));
1043
1044 let _f1 = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 8).await;
1045 let _f2 = feed(producer.registry(), Tier::default(), "baz/qux", true, 1, 8).await;
1046
1047 assert_eq!(announced(&origin).await.0, ".stats/node/sjc/1");
1048 }
1049
1050 #[tokio::test(start_paused = true)]
1051 async fn task_announces_without_node_suffix() {
1052 let (producer, origin) = test_producer(None);
1053 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 8).await;
1054 assert_eq!(announced(&origin).await.0, ".stats/node");
1055 }
1056
1057 #[tokio::test(start_paused = true)]
1058 async fn frame_emits_expected_counters() {
1059 let (producer, origin) = test_producer(Some("sjc"));
1060 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 42).await;
1062
1063 drive_tick().await;
1064
1065 let (_, broadcast) = announced(&origin).await;
1066 let frame = read_last_frame(&broadcast, "publisher.json").await;
1067 let snap = frame.get("foo/bar").expect("foo/bar entry");
1068 assert_eq!(snap.announced, 1, "egress announce stream bumps announced");
1069 assert_eq!(snap.broadcasts, 1, "one session subscribed");
1070 assert_eq!(snap.subscriptions, 1);
1071 assert_eq!(snap.bytes, 42);
1072 assert_eq!(snap.frames, 1);
1073 }
1074
1075 #[tokio::test(start_paused = true)]
1076 async fn announced_bytes_surfaces_in_frame() {
1077 let (producer, origin) = test_producer(Some("sjc"));
1078 let _f = feed(producer.registry(), Tier::default(), "foo/bar", false, 0, 0).await;
1080
1081 drive_tick().await;
1082
1083 let (_, broadcast) = announced(&origin).await;
1084 let frame = read_last_frame(&broadcast, "publisher.json").await;
1085 let snap = frame.get("foo/bar").expect("foo/bar entry");
1086 assert_eq!(snap.announced, 1);
1087 assert_eq!(
1088 snap.announced_bytes,
1089 "foo/bar".len() as u64,
1090 "name length recorded on announce"
1091 );
1092 }
1093
1094 #[tokio::test(start_paused = true)]
1095 async fn announced_decouples_from_broadcasts() {
1096 let (producer, origin) = test_producer(Some("sjc"));
1099 let _f = feed(producer.registry(), Tier::default(), "foo/bar", false, 0, 0).await;
1100
1101 drive_tick().await;
1102
1103 let (_, broadcast) = announced(&origin).await;
1104 let frame = read_last_frame(&broadcast, "publisher.json").await;
1105 let snap = frame.get("foo/bar").expect("foo/bar entry");
1106 assert_eq!(snap.announced, 1);
1107 assert_eq!(snap.broadcasts, 0, "no subscription, no broadcasts sentinel");
1108 assert_eq!(snap.subscriptions, 0);
1109 }
1110
1111 #[tokio::test(start_paused = true)]
1112 async fn short_lived_sub_is_surfaced() {
1113 let (producer, origin) = test_producer(Some("sjc"));
1119 {
1120 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 123).await;
1123 }
1124
1125 drive_tick().await;
1126
1127 let (_, broadcast) = announced(&origin).await;
1128 let frame = read_last_frame(&broadcast, "publisher.json").await;
1129 let snap = frame.get("foo/bar").expect("foo/bar entry");
1130 assert_eq!(snap.subscriptions, 1);
1132 assert_eq!(snap.subscriptions_closed, 1);
1133 assert_eq!(snap.broadcasts, 1, "one session subscribed");
1134 assert_eq!(snap.broadcasts_closed, 1);
1135 assert_eq!(snap.bytes, 123);
1136 assert_eq!(snap.frames, 1);
1137 }
1138
1139 #[tokio::test(start_paused = true)]
1140 async fn session_track_surfaces_by_root() {
1141 let (producer, origin) = test_producer(Some("sjc"));
1142 let _a = producer.registry().tier(Tier::default()).session("acme");
1143 let _b = producer.registry().tier(Tier::default()).session("acme");
1144 let _c = producer.registry().tier(Tier::new("region/sjc")).session("peer");
1145
1146 drive_tick().await;
1147
1148 let (_, broadcast) = announced(&origin).await;
1149 let frame = read_session_frame(&broadcast, "sessions.json").await;
1150 let snap = frame.get("acme").expect("root entry");
1151 assert_eq!(snap.sessions, 2);
1152 assert_eq!(snap.sessions_closed, 0);
1153 assert!(
1154 !frame.contains_key("peer"),
1155 "regional session must not appear on the default track"
1156 );
1157
1158 let snap = *read_session_frame(&broadcast, "region/sjc/sessions.json")
1159 .await
1160 .get("peer")
1161 .expect("regional entry");
1162 assert_eq!(snap.sessions, 1);
1163 }
1164
1165 #[tokio::test(start_paused = true)]
1166 async fn unused_slots_dont_surface() {
1167 let (producer, origin) = test_producer(Some("sjc"));
1171 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 8).await;
1174
1175 drive_tick().await;
1176 drive_tick().await;
1177
1178 let (_, broadcast) = announced(&origin).await;
1179
1180 assert!(
1182 read_last_frame(&broadcast, "publisher.json")
1183 .await
1184 .contains_key("foo/bar"),
1185 "publisher.json must include the active foo/bar entry"
1186 );
1187
1188 let frame = read_frame(&broadcast, "subscriber.json").await;
1191 assert!(frame.is_empty(), "subscriber.json must be empty, got {frame:?}");
1192
1193 for name in ["publisher.json.z", "subscriber.json.z", "sessions.json.z"] {
1195 assert!(broadcast.track(name).is_ok(), "{name} must exist");
1196 }
1197
1198 let subscribing = broadcast
1203 .track("region/sjc/publisher.json")
1204 .expect("logical track")
1205 .subscribe(None);
1206 drive_tick().await;
1207 let mut sub = subscribing.await.expect("an idle tier's track is held open");
1208 let frame = sub.read_frame().await.expect("ok").expect("frame");
1209 let parsed: BTreeMap<String, Traffic> = serde_json::from_slice(&frame.payload).expect("json");
1210 assert!(parsed.is_empty(), "an idle tier serves zeros, got {parsed:?}");
1211 }
1212
1213 #[test]
1214 fn advertised_path_with_and_without_node() {
1215 let prefix = Path::new(".stats");
1216 let empty = Path::empty();
1217 assert_eq!(
1218 advertised_path(&prefix, &empty, Some("sjc")).as_str(),
1219 ".stats/node/sjc"
1220 );
1221 assert_eq!(
1222 advertised_path(&prefix, &empty, Some("sjc/1")).as_str(),
1223 ".stats/node/sjc/1"
1224 );
1225 assert_eq!(advertised_path(&prefix, &empty, None).as_str(), ".stats/node");
1226 assert_eq!(
1227 advertised_path(&prefix, &Path::new("acme"), Some("sjc")).as_str(),
1228 ".stats/acme/node/sjc"
1229 );
1230
1231 let prefix = Path::new("metrics");
1232 assert_eq!(
1233 advertised_path(&prefix, &Path::new("demo/room"), Some("lon")).as_str(),
1234 "metrics/demo/room/node/lon"
1235 );
1236 }
1237
1238 #[test]
1239 fn group_key_uses_leading_segments() {
1240 assert_eq!(group_key("acme/room/cam", 0), Path::empty().to_owned());
1241 assert_eq!(group_key("acme/room/cam", 1), Path::new("acme").to_owned());
1242 assert_eq!(group_key("acme/room/cam", 2), Path::new("acme/room").to_owned());
1243 assert_eq!(group_key("acme/room", 3), Path::new("acme/room").to_owned());
1244 }
1245
1246 #[test]
1247 fn requested_track_shape_classifies() {
1248 let shape = requested_track_shape("rtmp/publisher.json").expect("valid");
1249 assert_eq!(shape.plain, "rtmp/publisher.json");
1250 assert!(!shape.compressed);
1251 assert!(!shape.sessions);
1252
1253 let shape = requested_track_shape("region/sjc/subscriber.json.z").expect("valid");
1254 assert_eq!(shape.plain, "region/sjc/subscriber.json");
1255 assert!(shape.compressed);
1256 assert!(!shape.sessions);
1257
1258 let shape = requested_track_shape("sessions.json").expect("default tier");
1259 assert_eq!(shape.plain, "sessions.json");
1260 assert!(shape.sessions);
1261
1262 assert!(requested_track_shape("bogus.json").is_none());
1263 assert!(requested_track_shape("xpublisher.json").is_none());
1264 assert!(requested_track_shape("/publisher.json").is_none());
1265 assert!(requested_track_shape("rtmp//publisher.json").is_none());
1266 assert!(requested_track_shape("rtmp/publisher.json.z.z").is_none());
1267 }
1268
1269 #[tokio::test(start_paused = true)]
1273 async fn idle_tier_track_resolves_with_zeros() {
1274 let (producer, origin) = test_producer(Some("sjc"));
1275 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 42).await;
1277 drive_tick().await;
1278 let (_, broadcast) = announced(&origin).await;
1279
1280 let subscribing = broadcast.track("rtmp/publisher.json").expect("track").subscribe(None);
1282 drive_tick().await;
1283 let mut sub = subscribing.await.expect("held open, not rejected");
1284 let frame = sub.read_frame().await.expect("ok").expect("frame");
1285 let parsed: BTreeMap<String, Traffic> = serde_json::from_slice(&frame.payload).expect("json");
1286 assert!(parsed.is_empty(), "an idle tier serves zeros");
1287
1288 let _rtmp = feed(producer.registry(), Tier::new("rtmp"), "foo/live", true, 1, 7).await;
1290 drive_tick().await;
1291 let frame = sub.read_frame().await.expect("ok").expect("frame");
1292 let parsed: BTreeMap<String, Traffic> = serde_json::from_slice(&frame.payload).expect("json");
1293 assert_eq!(parsed.get("foo/live").expect("entry").bytes, 7);
1294 }
1295
1296 #[tokio::test(start_paused = true)]
1299 async fn compressed_tier_request_creates_the_pair() {
1300 let (producer, origin) = test_producer(Some("sjc"));
1301 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 42).await;
1302 drive_tick().await;
1303 let (_, broadcast) = announced(&origin).await;
1304
1305 let subscribing = broadcast.track("srt/subscriber.json.z").expect("track").subscribe(None);
1306 drive_tick().await;
1307 subscribing.await.expect("compressed flavor held open");
1308
1309 subscribe(&broadcast, "srt/subscriber.json").await;
1311 }
1312
1313 #[tokio::test(start_paused = true)]
1315 async fn idle_tier_sessions_track_resolves_with_zeros() {
1316 let (producer, origin) = test_producer(Some("sjc"));
1317 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 42).await;
1318 drive_tick().await;
1319 let (_, broadcast) = announced(&origin).await;
1320
1321 let subscribing = broadcast.track("webrtc/sessions.json").expect("track").subscribe(None);
1322 drive_tick().await;
1323 let mut sub = subscribing.await.expect("held open, not rejected");
1324 let frame = sub.read_frame().await.expect("ok").expect("frame");
1325 let parsed: BTreeMap<String, Presence> = serde_json::from_slice(&frame.payload).expect("json");
1326 assert!(parsed.is_empty());
1327 }
1328
1329 #[tokio::test(start_paused = true)]
1331 async fn malformed_track_name_rejected() {
1332 let (producer, origin) = test_producer(Some("sjc"));
1333 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 42).await;
1334 drive_tick().await;
1335 let (_, broadcast) = announced(&origin).await;
1336
1337 let subscribing = broadcast.track("bogus.json").expect("track").subscribe(None);
1338 drive_tick().await;
1339 assert!(subscribing.await.is_err(), "a non-stats name is rejected");
1340 }
1341
1342 #[tokio::test(start_paused = true)]
1347 async fn request_racing_first_traffic_is_fulfilled() {
1348 let (producer, origin) = test_producer(Some("sjc"));
1349 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 42).await;
1350 drive_tick().await;
1351 let (_, broadcast) = announced(&origin).await;
1352
1353 let subscribing = broadcast.track("rtmp/publisher.json").expect("track").subscribe(None);
1356 assert!(subscribing.poll_ok(&moq_net::kio::Waiter::noop()).is_pending());
1357 for _ in 0..8 {
1358 tokio::task::yield_now().await;
1359 }
1360
1361 let _rtmp = feed(producer.registry(), Tier::new("rtmp"), "foo/live", true, 1, 7).await;
1363 drive_tick().await;
1364
1365 let mut sub = subscribing.await.expect("fulfilled by the tick's own creation");
1367 let frame = sub.read_frame().await.expect("ok").expect("frame");
1368 let parsed: BTreeMap<String, Traffic> = serde_json::from_slice(&frame.payload).expect("json");
1369 assert_eq!(parsed.get("foo/live").expect("entry").bytes, 7);
1370 }
1371
1372 #[tokio::test(start_paused = true)]
1377 async fn requested_quota_recovers_after_disconnect() {
1378 let (producer, origin) = test_producer(Some("sjc"));
1379 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 42).await;
1380 drive_tick().await;
1381 let (_, broadcast) = announced(&origin).await;
1382
1383 let mut held = Vec::new();
1385 for i in 0..MAX_REQUESTED_TRACKS {
1386 let name = format!("junk{i}/publisher.json");
1387 let subscribing = broadcast.track(&name).expect("track").subscribe(None);
1388 drive_tick().await;
1389 held.push(subscribing.await.expect("within the cap"));
1390 }
1391
1392 let subscribing = broadcast.track("real/publisher.json").expect("track").subscribe(None);
1394 assert!(subscribing.poll_ok(&moq_net::kio::Waiter::noop()).is_pending());
1395 drive_tick().await;
1396 assert!(
1397 subscribing.poll_ok(&moq_net::kio::Waiter::noop()).is_pending(),
1398 "an over-quota request parks instead of being rejected"
1399 );
1400
1401 drop(held);
1407 for _ in 0..4 {
1408 tokio::task::yield_now().await;
1409 }
1410 tokio::time::advance(Duration::from_secs(31)).await;
1411 for _ in 0..3 {
1412 drive_tick().await;
1413 }
1414 subscribing
1415 .await
1416 .expect("the parked request is adopted once the quota frees");
1417 }
1418
1419 #[tokio::test(start_paused = true)]
1424 async fn parked_request_is_adopted_by_first_traffic() {
1425 let (producer, origin) = test_producer(Some("sjc"));
1426 let _f = feed(producer.registry(), Tier::default(), "foo/bar", true, 1, 42).await;
1427 drive_tick().await;
1428 let (_, broadcast) = announced(&origin).await;
1429
1430 let mut held = Vec::new();
1432 for i in 0..MAX_REQUESTED_TRACKS {
1433 let name = format!("junk{i}/publisher.json");
1434 let subscribing = broadcast.track(&name).expect("track").subscribe(None);
1435 drive_tick().await;
1436 held.push(subscribing.await.expect("within the cap"));
1437 }
1438 let subscribing = broadcast.track("rt/publisher.json").expect("track").subscribe(None);
1439 assert!(subscribing.poll_ok(&moq_net::kio::Waiter::noop()).is_pending());
1440 drive_tick().await;
1441 assert!(
1442 subscribing.poll_ok(&moq_net::kio::Waiter::noop()).is_pending(),
1443 "parked"
1444 );
1445
1446 let _rt = feed(producer.registry(), Tier::new("rt"), "foo/live", true, 1, 9).await;
1448 drive_tick().await;
1449 let mut sub = subscribing.await.expect("adopted by the flush");
1450 let frame = sub.read_frame().await.expect("ok").expect("frame");
1451 let parsed: BTreeMap<String, Traffic> = serde_json::from_slice(&frame.payload).expect("json");
1452 assert_eq!(parsed.get("foo/live").expect("entry").bytes, 9);
1453 }
1454}