feat(relay): expose wire and control byte counters

doc/bin/relay/http.md

@@ -42,10 +42,23 @@ listen = "127.0.0.1:9101"
 
 ### GET /metrics
 
-Prometheus text: bytes, frames, groups, announces, subscriptions, viewers, requests,
-and sessions split by `tier` and `role`, plus accept-loop counters per TCP listener. Traffic
-counters accumulate for the node's lifetime and leave out the stats feed's own
-traffic. Host CPU and memory belong to a node exporter.
+Prometheus text: payload bytes, frames, groups, announces, subscriptions, and
+viewers split by `tier` and `role`; requests and sessions by `tier`; and
+accept-loop counters per TCP listener. Payload counters accumulate for the
+node's lifetime and leave out the stats feed's own traffic. Host CPU and memory
+belong to a node exporter.
+
+`moq_relay_wire_tx_bytes_total` and `moq_relay_wire_rx_bytes_total` sum the
+per-session [transport counters](/concept/stats#sessions) by `tier`, including
+internal sessions and the stats feed. Compare them with payload to spot
+transport overhead and idle traffic. Transports that do not report them
+contribute zero.
+
+`moq_relay_control_tx_bytes_total` and `moq_relay_control_rx_bytes_total`
+count control-message bytes by `tier` and `class` (`session`, `announce`, `track_info`, `subscribe`, `fetch`,
+`probe`, `auth`, `goaway`, `other`) to spot control-heavy sessions. They
+include message size prefixes but leave out media payload and transport
+overhead such as QUIC framing, ACKs, and retransmissions.
 
 A few worth alerting on:
 

rs/moq-net/src/client.rs

@@ -226,6 +226,7 @@ impl Client {
 
 		let start = lite::start(lite::Config {
 			runtime: runtime.clone(),
+			stats: self.stats.clone(),
 			client: true,
 			limits: self.limits,
 			session: session.clone(),
@@ -314,6 +315,7 @@ impl Client {
 				let auth = crate::auth::Handle::ietf(self.extensions.auth);
 				let (protocol, goaway) = ietf::start(ietf::Config {
 					runtime: runtime.clone(),
+					stats: self.stats.clone(),
 					limits: self.limits,
 					session: session.clone(),
 					setup: None,
@@ -398,7 +400,7 @@ impl Client {
 			Some(p) => return Err(Error::UnknownAlpn(p.to_string())),
 		};
 
-		let mut stream = Stream::open(&mut session, encoding).await?;
+		let mut stream = Stream::open(&mut session, encoding).await?.with_stats(&self.stats);
 
 		// The encoding is always an IETF version for SETUP negotiation.
 		let ietf_encoding = ietf::Version::try_from(encoding).map_err(|_| Error::Version)?;
@@ -446,6 +448,7 @@ impl Client {
 				let stream = stream.with_version(v);
 				let start = lite::start(lite::Config {
 					runtime: runtime.clone(),
+					stats: self.stats.clone(),
 					client: true,
 					limits: self.limits,
 					session: session.clone(),
@@ -492,6 +495,7 @@ impl Client {
 				// Draft 14-16: the path rode in the bidi SETUP above, not the uni one.
 				let (protocol, goaway) = ietf::start(ietf::Config {
 					runtime: runtime.clone(),
+					stats: self.stats.clone(),
 					limits: self.limits,
 					session: session.clone(),
 					setup: Some(stream),

rs/moq-net/src/coding/decode.rs

@@ -141,6 +141,11 @@ pub struct Decoder<'a> {
 	len: usize,
 	/// See [`Self::need`].
 	need: usize,
+	/// The control message read, counted once the read is consumed.
+	pub(crate) control: Option<crate::stats::Control>,
+	/// Leading bytes left out of the `control` count: a type selector the sender writes
+	/// apart from the message, and so never counts.
+	pub(crate) uncounted: usize,
 }
 
 impl<'a> Decoder<'a> {
@@ -152,6 +157,8 @@ impl<'a> Decoder<'a> {
 			form,
 			len: buf.len(),
 			need: 0,
+			control: None,
+			uncounted: 0,
 		}
 	}
 
@@ -176,6 +183,12 @@ impl<'a> Decoder<'a> {
 		self.form
 	}
 
+	/// How many bytes have been read.
+	#[cfg_attr(not(feature = "async"), allow(dead_code))]
+	pub(crate) fn position(&self) -> usize {
+		self.len - self.buf.len()
+	}
+
 	/// The number of unread bytes.
 	#[cfg_attr(not(feature = "async"), allow(dead_code))]
 	pub fn remaining(&self) -> usize {

rs/moq-net/src/coding/encode.rs

@@ -59,12 +59,17 @@ pub trait Encode<V> {
 pub struct Encoder<'a> {
 	buf: &'a mut Vec<u8>,
 	form: Form,
+	pub(crate) control: Option<crate::stats::Control>,
 }
 
 impl<'a> Encoder<'a> {
 	/// Append to `buf`, with varints in the given form.
 	pub fn new(buf: &'a mut Vec<u8>, form: Form) -> Self {
-		Self { buf, form }
+		Self {
+			buf,
+			form,
+			control: None,
+		}
 	}
 
 	/// The varint form this encoder writes.

rs/moq-net/src/coding/reader.rs

@@ -26,6 +26,7 @@ pub struct Reader<S: RecvStream, V> {
 	want: usize,
 	/// A contiguous copy of a message that spans transport chunks, empty between decodes.
 	scratch: Vec<u8>,
+	stats: crate::stats::Session,
 }
 
 /// How much payload may pile up behind a parked consumer before
@@ -49,9 +50,20 @@ impl<S: RecvStream, V: StreamCodes> Reader<S, V> {
 			version,
 			want: 0,
 			scratch: Vec::new(),
+			stats: Default::default(),
 		}
 	}
 
+	pub(crate) fn with_stats(mut self, stats: &crate::stats::Session) -> Self {
+		self.stats = stats.clone();
+		self
+	}
+
+	/// Where control messages read here are counted.
+	pub(crate) fn stats(&self) -> &crate::stats::Session {
+		&self.stats
+	}
+
 	/// Poll the stream for a lease of at least `want` bytes.
 	fn poll_lease(&mut self, cx: &mut Context<'_>, want: usize) -> Poll<Result<S::Lease<'_>, Error>> {
 		let version = &self.version;
@@ -68,6 +80,19 @@ impl<S: RecvStream, V: StreamCodes> Reader<S, V> {
 		consume: bool,
 		decode: impl Fn(&mut Decoder<'_>) -> Result<T, DecodeError>,
 	) -> Poll<Result<T, Error>>
+	where
+		V: Into<Form> + Copy,
+	{
+		self.poll_with_sized(cx, consume, decode).map_ok(|(value, _)| value)
+	}
+
+	/// [`Self::poll_with`], also returning how many bytes the value took.
+	fn poll_with_sized<T>(
+		&mut self,
+		cx: &mut Context<'_>,
+		consume: bool,
+		decode: impl Fn(&mut Decoder<'_>) -> Result<T, DecodeError>,
+	) -> Poll<Result<(T, usize), Error>>
 	where
 		V: Into<Form> + Copy,
 	{
@@ -101,15 +126,21 @@ impl<S: RecvStream, V: StreamCodes> Reader<S, V> {
 				let need = match decode(&mut r) {
 					Ok(value) => {
 						let used = buf.len() - r.remaining();
+						let control = r.control.map(|class| (class, used - r.uncounted));
 						self.want = 0;
 						if self.scratch.capacity() > SCRATCH_KEEP {
 							self.scratch = Vec::new();
 						}
 						if consume {
+							// Counted only once the bytes are released, so a peek or a retry
+							// after a short read never counts the same message twice.
+							if let Some((class, size)) = control {
+								self.stats.control(class, size, true);
+							}
 							lease.advance(used);
 							lease.release();
 						}
-						return Poll::Ready(Ok(value));
+						return Poll::Ready(Ok((value, used)));
 					}
 					// A short read inside a complete body is malformed, but reports what it
 					// needed from the body, so always ask for at least one byte more.
@@ -170,6 +201,27 @@ impl<S: RecvStream, V: StreamCodes> Reader<S, V> {
 		std::future::poll_fn(|cx| self.poll_decode(cx)).await
 	}
 
+	/// [`Self::poll_decode`], also returning the bytes it consumed, for a handshake to
+	/// count once the session's stats are known.
+	pub(crate) fn poll_decode_sized<T: Decode<V> + Debug>(
+		&mut self,
+		cx: &mut Context<'_>,
+	) -> Poll<Result<(T, usize), Error>>
+	where
+		V: Into<Form> + Copy,
+	{
+		let version = self.version;
+		self.poll_with_sized(cx, true, |r| T::decode(r, version))
+	}
+
+	/// [`Self::decode`], also returning the bytes it consumed.
+	pub(crate) async fn decode_sized<T: Decode<V> + Debug>(&mut self) -> Result<(T, usize), Error>
+	where
+		V: Into<Form> + Copy,
+	{
+		std::future::poll_fn(|cx| self.poll_decode_sized(cx)).await
+	}
+
 	/// Poll for the next message unless the stream is closed cleanly first.
 	pub fn poll_decode_maybe<T: Decode<V> + Debug>(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<T>, Error>>
 	where
@@ -405,6 +457,7 @@ impl<S: RecvStream, V: StreamCodes> Reader<S, V> {
 			version,
 			want: self.want,
 			scratch: self.scratch,
+			stats: self.stats,
 		}
 	}
 }
@@ -441,6 +494,75 @@ mod tests {
 		Reader::new(Buffered::new(StopLog::default()), version)
 	}
 
+	/// A control message counts once, when it is consumed: never on a peek, a retry after
+	/// a short read, or a failed decode, and in full when it spans chunks.
+	#[test]
+	fn control_counts_consumed_messages_once() {
+		use crate::{
+			lite,
+			stats::{Bytes, Control, Registry, Tier},
+		};
+		let registry = Registry::new(Default::default());
+		let session = registry.tier(Tier::default()).session("");
+		let version = lite::Version::Lite07;
+		let msg = lite::Probe {
+			bitrate: Some(123),
+			rtt: Some(4),
+		};
+		let bytes: &'static [u8] = msg.encode_bytes(version).unwrap().to_vec().leak();
+		let size = bytes.len() as u64;
+		let (head, tail) = bytes.split_at(2);
+		let counted = || registry.snapshot().control();
+		let probes = |n: u64| vec![(Tier::default(), Control::Probe, Bytes { tx: 0, rx: size * n })];
+
+		let mut reader = queue(&[bytes, head], version).with_stats(&session);
+		let mut cx = Context::from_waker(std::task::Waker::noop());
+		assert!(matches!(
+			reader.poll_decode_peek_maybe::<lite::Probe>(&mut cx),
+			Poll::Ready(Ok(Some(_)))
+		));
+		assert!(counted().is_empty());
+		assert!(matches!(reader.poll_decode::<lite::Probe>(&mut cx), Poll::Ready(Ok(_))));
+		assert_eq!(counted(), probes(1));
+
+		// The second probe arrives in two chunks, so the reader parks and retries.
+		assert!(reader.poll_decode::<lite::Probe>(&mut cx).is_pending());
+		assert!(reader.poll_decode::<lite::Probe>(&mut cx).is_pending());
+		assert_eq!(counted(), probes(1));
+		reader.stream.get_mut().chunks.push_back(tail);
+		assert!(matches!(reader.poll_decode::<lite::Probe>(&mut cx), Poll::Ready(Ok(_))));
+		assert_eq!(counted(), probes(2));
+
+		// A wrapped message counts its type too, as the sender's one encode does.
+		let reply = lite::AuthReply::Error(lite::AuthError {
+			code: 1,
+			reason: "no".into(),
+		});
+		let reply: &'static [u8] = reply.encode_bytes(version).unwrap().to_vec().leak();
+		reader.stream.get_mut().chunks.push_back(reply);
+		assert!(matches!(
+			reader.poll_decode::<lite::AuthReply>(&mut cx),
+			Poll::Ready(Ok(_))
+		));
+		let auth = (
+			Tier::default(),
+			Control::Auth,
+			Bytes {
+				tx: 0,
+				rx: reply.len() as u64,
+			},
+		);
+		assert_eq!(counted(), [vec![auth.clone()], probes(2)].concat());
+
+		// A body cut short inside its declared size is malformed, and never counted.
+		reader.stream.get_mut().chunks.push_back(&[1, 0x80]);
+		assert!(matches!(
+			reader.poll_decode::<lite::Probe>(&mut cx),
+			Poll::Ready(Err(_))
+		));
+		assert_eq!(counted(), [vec![auth], probes(2)].concat());
+	}
+
 	/// STOP_SENDING refuses a stream, so it must carry a stream code. Reusing the session
 	/// table here would have the peer decode it against the wrong registry: `Cancel` would
 	/// arrive as INTERNAL_ERROR, and `Unauthorized` as DELIVERY_TIMEOUT.

rs/moq-net/src/coding/stream.rs

@@ -77,6 +77,13 @@ impl<S: crate::transport::poll::Session, V: StreamCodes> Stream<S, V> {
 		Self { writer, reader }
 	}
 
+	pub(crate) fn with_stats(self, stats: &crate::stats::Session) -> Self {
+		Self {
+			writer: self.writer.with_stats(stats),
+			reader: self.reader.with_stats(stats),
+		}
+	}
+
 	/// Cast the stream to a different version, used during version negotiation.
 	pub fn with_version<V2: Clone>(self, version: V2) -> Stream<S, V2> {
 		Stream {