Skip to main content

moq_stats/
aggregate.rs

1//! The aggregating half: fold a group's per-node stats broadcasts into one view.
2//!
3//! A single-broadcast [`Consumer`](crate::Consumer) reads one
4//! `<prefix>/<group>/node/<node>` broadcast. This reader watches an origin's
5//! announce stream for *every* node broadcast in a group and folds their
6//! cumulative counters into one merged frame per `(tier, role)`, so a downstream
7//! sees a project's whole live traffic as if it came from a single node.
8
9use 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/// Configuration for an [`Consumer`]. Construct with [`Config::new`] and chain
21/// the `with_*` setters.
22///
23/// The `prefix` and `depth` must match the producing side's
24/// [`produce::Config`](crate::produce::Config): they are how announced paths are
25/// recognized as node broadcasts and filtered from sibling categories under the
26/// same prefix.
27#[derive(Debug, Clone)]
28#[non_exhaustive]
29pub struct Config {
30	/// Top-level path stats are published under (default `.stats`). Must match
31	/// the producer's prefix.
32	pub prefix: PathOwned,
33	/// The producer's grouping depth (default `0`). Announced paths whose group
34	/// is deeper than this are not recognized as node broadcasts and are
35	/// skipped. Must match the producer's depth.
36	pub depth: usize,
37	/// Read the compressed `.json.z` tracks instead of the plain `.json` ones.
38	/// Same data for a fraction of the bytes, but requires a producer that
39	/// publishes them. Defaults to `false`.
40	pub compression: bool,
41}
42
43impl Config {
44	/// A config with default settings: the `.stats` prefix, depth `0`, and the
45	/// plain `.json` tracks.
46	pub fn new() -> Self {
47		Self::default()
48	}
49
50	/// Override the top-level prefix (default `.stats`). Must match the producer.
51	pub fn with_prefix(mut self, prefix: impl Into<PathOwned>) -> Self {
52		self.prefix = prefix.into();
53		self
54	}
55
56	/// Override the grouping depth (default `0`). Must match the producer.
57	pub fn with_depth(mut self, depth: usize) -> Self {
58		self.depth = depth;
59		self
60	}
61
62	/// Read the compressed `.json.z` tracks instead of the plain `.json` ones.
63	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
79/// Folds a group's per-node stats broadcasts into one merged view.
80///
81/// Scope an [`origin::Consumer`] to a single group (e.g. `.stats/<pid>`) and
82/// hand it here; each [`Self::traffic`] / [`Self::sessions`] call opens its own
83/// announce cursor and subscribes to that track on every node broadcast in the
84/// group, summing the cumulative counters per key. Traffic is sticky: a node
85/// dropping out (its broadcast unannounces or its reader ends) keeps its last
86/// contribution, so a relay that returns with its boot-lifetime counters
87/// intact never looks like new traffic. Only a genuine per-node counter
88/// regression (a restarted relay) regresses the merged counter, the same reset
89/// contract a single node's own restart follows. Presence is not sticky: a
90/// departed node stops counting sessions immediately.
91pub struct Consumer {
92	origin: origin::Consumer,
93	config: Config,
94}
95
96impl Consumer {
97	/// Wrap an origin consumer, ideally already scoped to one group. `config`'s
98	/// `prefix` and `depth` must match the producing side.
99	pub fn new(origin: origin::Consumer, config: Config) -> Self {
100		Self { origin, config }
101	}
102
103	/// A merged reader over the traffic track for `(tier, role)`, folding every
104	/// node broadcast in the group. Nodes are subscribed lazily as they announce,
105	/// so this returns without a handshake.
106	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	/// A merged reader over the sessions track for `tier`; see [`Self::traffic`].
114	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
122/// A merged reader over one traffic track across every node in the group. Yields
123/// the latest merged [`TrafficFrame`]; a slow reader collapses intermediate
124/// frames, which is safe because the counters are cumulative.
125pub struct TrafficConsumer {
126	inner: Merged<Traffic>,
127}
128
129impl TrafficConsumer {
130	/// The next merged frame, or `None` once the announce stream ends (the
131	/// source origin went away).
132	pub async fn next(&mut self) -> Result<Option<TrafficFrame>> {
133		kio::wait(|waiter| self.inner.poll_next(waiter)).await
134	}
135}
136
137/// A merged reader over one sessions track across every node in the group; see
138/// [`TrafficConsumer`].
139pub struct SessionsConsumer {
140	inner: Merged<Presence>,
141}
142
143impl SessionsConsumer {
144	/// The next merged frame, or `None` once the announce stream ends.
145	pub async fn next(&mut self) -> Result<Option<SessionsFrame>> {
146		kio::wait(|waiter| self.inner.poll_next(waiter)).await
147	}
148}
149
150/// A per-key counter that folds across nodes: the two wire counter types.
151trait Mergeable: serde::de::DeserializeOwned + Default + Copy + 'static {
152	/// Fold `other` into `acc`.
153	fn merge(acc: &mut Self, other: Self);
154
155	/// Whether a node's last contribution survives its departure. Cumulative
156	/// counters do: a relay that leaves and returns with its boot-lifetime
157	/// counters intact must not look like new traffic. Gauges don't: a
158	/// departed node stops counting the moment it's gone.
159	const STICKY: bool;
160
161	/// Retire any outstanding live state in a kept contribution, returning
162	/// whether it changed. Cumulative totals are untouched; this is how a
163	/// departed node stops reporting open counters without losing its history.
164	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	/// Close every open counter pair: a departed relay can no longer carry its
175	/// broadcasts or subscriptions, so the merged view must not keep counting
176	/// them as live. The cumulative counters, bytes included, stay.
177	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	/// Presence is not sticky, so it never reaches here; its entry is dropped.
196	fn retire(&mut self) -> bool {
197		false
198	}
199}
200
201/// One node's subscription to the merged track.
202enum Reader<V: Mergeable> {
203	/// Resolving the announced path into a broadcast. `queued` records whether
204	/// the request was handed to a serving route (fixed at request time): a
205	/// queued request that fails `Unroutable` was killed by its serving route
206	/// retracting (the table has already changed), while an unqueued
207	/// `Unroutable` means nothing serves the path at all.
208	Resolving {
209		pending: Pending<origin::Requesting>,
210		queued: bool,
211	},
212	/// Awaiting the subscription handshake.
213	Subscribing(Pending<Subscribing>),
214	/// Reading frames. Boxed: the snapshot consumer dwarfs the other variants,
215	/// and one lives per node in a map.
216	Active(Box<moq_json::snapshot::Consumer<BTreeMap<String, V>>>),
217	/// The subscription failed or the track ended; the node no longer reads. It
218	/// lingers until it unannounces or reannounces, still contributing its last
219	/// frame when [`Mergeable::STICKY`].
220	Ended,
221}
222
223/// One node's reader plus the last frame it produced (the value folded into the
224/// merged view).
225struct Node<V: Mergeable> {
226	reader: Reader<V>,
227	/// The announced path, relative to the announce cursor: what a re-resolve
228	/// after the reader ends requests again.
229	path: PathOwned,
230	last: Option<BTreeMap<String, V>>,
231}
232
233impl<V: Mergeable> Node<V> {
234	/// The node's broadcast went away (unannounced, replaced by one without
235	/// this track, or its subscription ended): stop reading, and keep the last
236	/// frame only when sticky. A kept frame retires its live counters, so a
237	/// departed node stops reporting open state while its totals stay. Returns
238	/// whether the merged view changed.
239	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
254/// Watches a group's node announces and folds one track across all of them.
255struct Merged<V: Mergeable> {
256	/// Resolves announced node paths into broadcasts.
257	origin: origin::Consumer,
258	announce: moq_net::announce::Consumer,
259	prefix: PathOwned,
260	depth: usize,
261	/// Track name subscribed on each node broadcast.
262	name: String,
263	config: moq_json::snapshot::consumer::Config,
264	/// One entry per live node broadcast, keyed by absolute announced path.
265	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			// Stats live under a `.`-named prefix, which discovery hides by default.
272			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	/// Poll for the next merged frame. Returns `Ready(Some(_))` whenever the
289	/// merged view changed (a node produced a frame, or a non-sticky
290	/// contribution left), `Ready(None)` once the announce stream closes, else
291	/// `Pending`.
292	fn poll_next(&mut self, waiter: &Waiter) -> Poll<Result<Option<BTreeMap<String, V>>>> {
293		let mut changed = false;
294
295		// Drain announce membership updates: add/replace on announce, drop on
296		// unannounce. A closed stream ends the merged view.
297		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		// Advance each node's reader, collapsing any backlog to its latest frame.
306		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	/// Apply one announce update to the node set. Returns whether the merged view
321	/// changed (only a non-sticky contribution leaving does; a sticky one is
322	/// kept).
323	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		// Only fold node-category routes; skip sibling categories a producer
328		// may publish under the same prefix. A route names a prefix, and node
329		// broadcasts are announced at their exact path by convention.
330		if parse_node_path(&self.prefix, self.depth, &absolute).is_none() {
331			return false;
332		}
333
334		if update.kind.is_active() {
335			// A route update on a node already tracked (a reprice, or a takeover
336			// with different metadata) keeps the live reader: existing
337			// subscriptions survive a takeover, and the reader re-resolves through
338			// the current best route the moment it actually ends. Only an ended
339			// node re-arms here, since a fresh route is new evidence that a
340			// re-resolve could succeed. A sticky contribution carries across the
341			// re-arm, holding the total until a fresh frame replaces it.
342			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			// Unannounce: keep the cumulative totals, retire the live gauges, and
361			// stop reading. The entry stays so a reannounce re-arms it.
362			match self.nodes.get_mut(&absolute) {
363				Some(node) => node.depart(),
364				None => false,
365			}
366		} else {
367			// A gauge drops its contribution, and its entry, so a departed node
368			// stops counting and its path is not retained.
369			self.nodes.remove(&absolute).is_some_and(|old| old.last.is_some())
370		}
371	}
372
373	/// Sum every node's last frame, per key.
374	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
387/// Drive one node's reader as far as it goes, updating its `last` frame. Returns
388/// whether that node's contribution to the merged view changed.
389fn 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	// At most one re-resolve per call: a subscription that terminates
398	// synchronously twice in a row is done, not failing over.
399	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				// A queued request killed by its route retracting: an identical
411				// standby swaps in without any announce update, so re-resolve
412				// through the already-updated table. Each retry consumed a real
413				// retraction, so this cannot spin; an instant Unroutable instead
414				// means nothing serves the path (the retraction that empties the
415				// table also unannounces this node).
416				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				// The subscription ended: the serving session died (a failover to
442				// an identical route delivers no announce update) or the track
443				// finished. A non-sticky gauge drops its last frame so a stale
444				// value stops pinning the sum, while a sticky counter keeps it.
445				// Then re-resolve through the current best route; an
446				// authoritative refusal on the way ends the node instead.
447				Poll::Ready(result @ (Ok(None) | Err(_))) => {
448					if let Err(err) = result {
449						// One bad node must not tear down the whole merged view;
450						// re-resolve just this node and keep folding the rest.
451						tracing::debug!(?err, name, "stats: node read error");
452					}
453					// `depart` drops the dead reader before the re-request: our
454					// own handle is what keeps a dying served broadcast cached,
455					// and releasing it first lets the request materialize a
456					// fresh one.
457					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
471/// Start resolving `path` (relative to the announce cursor) into a broadcast.
472fn 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	/// Build an origin producer, spawning its driver on the ambient runtime.
481	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			// A sync test: nothing polls the driver, and dropping it would tear
487			// the origin down, so leak it and rely on the synchronous half.
488			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	/// A stats producer publishing one node's broadcasts on `origin`, grouped at
502	/// depth 1 (so feeding a broadcast under `<group>/...` announces
503	/// `.stats/<group>/node/<node>`).
504	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	/// Kept-alive handles from [`feed`]: dropping them closes the subscription and
514	/// announce, bumping the `_closed` counters.
515	#[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	/// Record `bytes` of egress traffic on `path` in `producer`'s registry under
525	/// `tier`/`root`, by driving a throwaway tagged broadcast. The broadcast lives
526	/// on its own origin so it never lands on the stats origin.
527	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	/// Advance past one publish interval so every producer task drains and writes.
558	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	/// Read merged traffic frames until `path`'s byte count reaches `want` (each
566	/// node folds in independently, so a partial frame can arrive first).
567	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	/// Read merged traffic frames until `path`'s byte count reaches `want`,
577	/// asserting it never regresses below `min` along the way: a departed node's
578	/// kept contribution must hold the total through every intermediate frame.
579	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	/// A hand-published node broadcast at `.stats/<group>/node/<node>` with a
591	/// plain default-tier traffic track, so a test controls the exact frames and
592	/// can fail one node's reader alone (the registry-driven producer only
593	/// publishes whole broadcasts). Dropping it unannounces the node.
594	#[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		/// Add `bytes` to `path`'s cumulative counter and publish the node's
619		/// whole snapshot, like a real node's registry drain would.
620		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		/// Fail the node's reader: append a frame the snapshot decoder can't
627		/// parse, so the subscription errors while the broadcast stays announced.
628		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		// Two nodes each serve the same broadcast; the merged view sums their
640		// cumulative counters per path.
641		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		// Dropping a node unannounces its broadcast, but traffic is sticky: its
662		// last contribution stays in the total, so a relay that returns with its
663		// boot-lifetime counters intact never looks like new traffic.
664		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 node B entirely: its publish task ends and finishes the broadcast.
677		drop(fb);
678		drop(node_b);
679		drive_tick().await;
680
681		// Node A serves more traffic on another path; the merged frame still
682		// carries node B's kept contribution for acme/room.
683		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		// A node's stats session reconnects with its cumulative counters intact
697		// and still growing: the total holds through the swap, then advances,
698		// never dipping.
699		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		// Node A's broadcast goes away and comes back with a higher counter.
712		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		// The kept contribution holds the total at 140 until the new frame
721		// replaces it, landing on 120 + 40.
722		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		// A node that returns with a fresh counter (it restarted) replaces its
729		// contribution wholesale: the total regresses once, the existing
730		// fresh-segment contract.
731		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		// Node A restarts: its broadcast returns with a lower counter.
744		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		// The restarted frame replaces A's contribution, so the total drops to
753		// 30 + 40, a genuine per-node regression downstream treats as a fresh
754		// segment. An earlier frame may retire A's live gauges first.
755		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		// A node's reader failing while its broadcast is still announced keeps
766		// its last contribution in the total.
767		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's subscription errors under its still-announced broadcast.
778		node_a.fail_traffic();
779
780		// Node B reports more traffic; node A's kept 100 stays in the total.
781		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		// A departed node's totals stay in the merged view, but its live gauges
794		// retire: the aggregate must not show phantom viewers or broadcasts for
795		// a node that is gone.
796		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		// The node departs with those sessions still open.
819		drop(node_a);
820
821		// The totals stay; the live gauges retire.
822		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		// Session presence sums per auth root across nodes.
833		let origin = produce_origin();
834		let node_a = node_producer(&origin, "a");
835		let node_b = node_producer(&origin, "b");
836
837		// Each node needs a live broadcast to announce; the sessions ride the same
838		// group.
839		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			// Each feed opens one session ("acme" root) plus the explicit ones,
851			// summed across both nodes.
852			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		// Unlike traffic, presence is a gauge: a node's unannounce must stop
861		// counting its sessions immediately, not pin a stale gauge.
862		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			// Each feed opens one session ("acme" root) plus the explicit ones,
878			// summed across both nodes.
879			if frame.get("acme").map(|p| p.active()) >= Some(4) {
880				break;
881			}
882		}
883
884		// Node B departs: its sessions leave the gauge at once.
885		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}