diff --git a/doc/bin/relay/config.md b/doc/bin/relay/config.md index ab039495c64..44706065d36 100644 --- a/doc/bin/relay/config.md +++ b/doc/bin/relay/config.md @@ -61,8 +61,8 @@ idle_timeout = "10s" # Drop a connection after this long with no keep_alive = "3s" # Ping interval; "0s" disables it. iroh still pings at least every 5s. Default. gso = true # UDP segmentation offload. iroh cannot turn it off. mtu_discovery = false # Path MTU discovery. Default. -receive_window = 67108864 # Whole-connection window, in bytes. Default (64 MiB). -stream_receive_window = 8388608 # Per-stream window. Omit for the backend default. +receive_window = 134217728 # Whole-connection window, in bytes. Default (128 MiB). +stream_receive_window = 8388608 # Per-stream window. Omit to follow receive_window. send_window = 33554432 # Unacknowledged outgoing data. Omit for the backend default. qlog = "/var/log/moq/qlog" # Existing directory. Needs the `qlog` build feature. ``` @@ -72,11 +72,17 @@ sessions, and so its [cluster routes](/bin/relay/cluster#failure-detection). QUIC uses the smaller of the two endpoints' values, so this also bounds the other end. Keep `keep_alive` under a third of it. -Raise the receive windows when a fat, long path idles below the link rate. The -64 MiB `receive_window` default carries several Gbps at a 100 ms RTT and bounds -how much unread data one peer can make the relay buffer. Keep -`stream_receive_window` well under it so one slow group cannot starve the -connection. +Raise `receive_window` when a fat, long path idles below the link rate. The +128 MiB default carries several Gbps at a 100 ms RTT and bounds how much unread +data and incomplete messages one peer can make the relay buffer. A message holds +its share until it is complete, so the window must fit the largest one, just over +64 MiB, and the relay refuses to start with less. When incomplete messages fill +it, the relay cancels the least urgent stream and carries on. iroh returns credit +as it reads, so there the window bounds only unread data, and incomplete messages +are bounded by `max_streams` streams of up to the largest message each. + +Stream credit returns as data is read, so a `stream_receive_window` under +`receive_window` only adds round trips to a large frame. ## \[session] and \[peer] diff --git a/quest/m1/quic/README.md b/quest/m1/quic/README.md index 11c3c74df3c..0a513d7f7d9 100644 --- a/quest/m1/quic/README.md +++ b/quest/m1/quic/README.md @@ -63,7 +63,7 @@ cache limits remain separate; bitrate pacing composes with this memory cap. ## Required - [Regression net](/quest/m1/quic/coverage/README.md) - benchmarks, fuzzing, characterization tests, and a seeded simulation, so later QUIC changes see every behavior or cost change -- [Bound message assembly](/quest/m1/quic/receive-memory.md) - complete headers and frames return credit, and receiver priority cancellation restores progress +- [Prove receive memory](/quest/m1/quic/receive-memory-coverage.md) - a JavaScript sender exhausts both native receivers and they recover; nightly measures the cost - [A reset group keeps its header](/quest/m1/quic/group-reset.md) - a lite publisher resets a group with `RESET_STREAM_AT`, so its subscriber drops the tail grace diff --git a/quest/m1/quic/peer-limits.md b/quest/m1/quic/peer-limits.md index 7a97758dab9..a11e140826f 100644 --- a/quest/m1/quic/peer-limits.md +++ b/quest/m1/quic/peer-limits.md @@ -35,9 +35,10 @@ its own. - Refuse a `peer` value below the client default rather than silently shrinking. -The [receive-memory plan](/quest/m1/quic/receive-memory.md) requires client -and peer memory allowances to fit the largest supported encoded unit, -including framing. Preserve that validation and retained-credit accounting +Message assembly holds connection credit until each unit is complete, so +client and peer receive windows must fit the largest supported encoded unit, +`moq_net::transport::MIN_RECEIVE_WINDOW`; the relay refuses a smaller +`[quic]` window at config load. Preserve that validation and retained-credit accounting when applying runtime limits; wider stream counts must not create a second unbounded assembly pool. diff --git a/quest/m1/quic/receive-memory-coverage.md b/quest/m1/quic/receive-memory-coverage.md new file mode 100644 index 00000000000..784ba357ffd --- /dev/null +++ b/quest/m1/quic/receive-memory-coverage.md @@ -0,0 +1,39 @@ +# [M] Prove receive memory with JavaScript senders and benchmarks + +## Goal + +A JavaScript sender drives both native receivers into receive-credit +exhaustion and they recover, and nightly measures what bounded receive memory +costs. This is coverage for the message-assembly work in #5735, with no new +transport behavior. + +## Plan + +Split out of #5735, which lands native credit holding, pressure recovery, and +the minimum receive window with Rust-only end-to-end tests over tokio and +io_uring. + +A JavaScript sender over the stock receivers cannot stall them. The minimum +window fits the largest unit, and a strict-priority sender finishes one +stream before the next, so every frame completes. The test therefore needs a +receiver whose connection window is far below the minimum. Neither stock +config allows that, so the test assembles a WebTransport (HTTP/3) server +from each raw driver: tokio's `quic::driver` with `quic::web::Server`, and +an io_uring endpoint in web mode. Keep it test-only; do not add a public way +to configure a window below `MIN_RECEIVE_WINDOW`. + +The sender is a Bun script using `@moq/web-transport`, pinned to the test +server's certificate hash with WebSocket disabled. It mirrors the Rust tests +in `rs/moq-tokio/tests/integration/receive_credit.rs`. A low-priority track sends a frame +larger than the window. A high-priority frame then has to complete, and a +later low-priority group has to arrive after the low stream is stopped. Run +it from Rust the way `rs/moq-tokio/tests/integration/max_age.rs` drives its Bun client, +and wire it into the interop workflow, since normal CI has no Bun. + +The benchmark sweeps concurrent streams and frame sizes, and compares a +finite stream window with one that follows the connection window. Report +retained bytes (held plus unread credit), throughput, and frame completion +latency. Sweep streams and frame sizes on separate axes, so a cost that grows +with the stream count shows up as a slope. Wire it into `sh/gh/nightly.sh` +next to `benchmark_held_reads`. Unit tests stay on mocked time; the benchmark +is the only wall-clock part. diff --git a/quest/m1/quic/receive-memory.md b/quest/m1/quic/receive-memory.md deleted file mode 100644 index da985e4eb66..00000000000 --- a/quest/m1/quic/receive-memory.md +++ /dev/null @@ -1,110 +0,0 @@ -# [XL] Bound incomplete messages with QUIC connection credit - -## Goal - -Native MoQ receivers bound incomplete frames and protocol messages by the -connection memory allowance, returning connection credit only when each -complete unit is parsed and handed off or discarded. When concurrent partial -messages prevent progress, the receiver cancels streams by MoQ priority and -continues. This works with Rust and JavaScript senders through tokio and -io_uring; completed media caching remains separately governed. - -## Plan - -Decided in the 2026-10-10 interview: cover headers and frame/group data alike. -The frame preallocation budget only bounds speculative allocation, and the -cache has soft limits; neither bounds aggregate incomplete messages. Use -peer-visible QUIC flow control instead of a separate header quota error. -The native transport primitive is a separate prerequisite; do not activate -whole-message retention without its recovery policy in this quest. - -That primitive landed as a provisional seam to finalize here: both native -drivers' receive streams have a copying `poll_read_held` that yields a -`Credit` for exactly the bytes read. Dropping a `Credit` returns it, so a -reader buffer keeps one merged `Credit` and drops `split(n)` for each -complete unit; read-ahead stays in the remainder. io_uring's `poll_closed` -read-ahead keeps its credit until the application reads it. Both drivers' -`Connection::stats().credit` reports held and still-available credit. -`poll_read_held` stays separate until this quest folds it into `poll_read`. -moq-net's transport trait and the WebTransport wrappers do not carry it yet; -reshape any of it to fit the consumer. - -Hold connection credit through reader buffers, type/length prefixes, headers, -extensions, and frame assembly. Attribute read-ahead bytes to the next unit -rather than releasing them with the completed one. Release on complete-unit -handoff, not after application playback or cache eviction. Avoid retaining -an entire multi-frame group as one unit. Account for received-but-unread and -out-of-order transport bytes as well as parser-held bytes; they share the -same allowance. Speculative allocations and copies need separate accounting -or bounds: QUIC byte credit is not a claim about total process RSS. - -Keep the existing stream-window setting and honor explicit overrides. -Give the native path a very large default so it normally adds no round trips; -stream credit is replenished on reads, independently of message completion. -Do not remove or ignore the setting without benchmark evidence and a new -maintainer decision. The connection memory allowance is the shared cap. - -Decided 2026-10-16: once [Segmented frames](/quest/m2/unbounded-groups/segmented-frames.md) -lands, a frame payload releases credit per segment rather than per frame, -so frames of any size never need a matching allowance. Do not build on -whole-frame units beyond what that change replaces. - -Require that allowance to fit every supported encoded message, including -framing and extensions. Validate configuration and align native defaults -rather than discovering an impossible frame after receiving it. The current -32 MiB frame limit exceeds io_uring's 16 MiB connection default; a legacy -64 MiB announcement body also needs room for framing beyond 64 MiB. Derive -the minimum from the supported units, not just the payload limit. Lowering -a protocol/frame size limit is a separate maintainer decision, not an implicit -way around this requirement. Runtime memory-limit updates obey the same -minimum; already-advertised credit cannot be revoked. A rate limit can pace -a smaller amount of outstanding credit while the memory allowance remains -large enough for the complete unit. - -MoQ selects victims; the transport only reports pressure and sends -STOP_SENDING. Do not add sender-initiated shedding. This keeps recovery -available with browser senders and avoids treating intentional rate pacing -as receiver deadlock. The sender's required RESET_STREAM response is still -normal QUIC behavior, not a second recovery policy. - -Use the receiver's assembly and credit accounting to recover when partial -units exhaust progress capacity. Do not require peer DATA_BLOCKED, an -arbitrary timeout, loss, or a deadline. Demonstrate that useful work can -complete after cancellation, including when incoming streams keep arriving; -merely cycling through victims is not recovery. Distinguish bytes still able -to arrive under advertised credit from memory held by incomplete units. - -Within the existing broadcast fairness domains, cancel lower-priority media -and less useful groups first, respecting subscription priority and group -order. Prefer low-priority attributed media over unidentified headers so -headers have a chance to identify themselves; unidentified streams remain -eligible if they prevent progress. Control streams have the highest priority -but are eligible too. If canceling an essential control stream terminates the -session under its protocol, surface that outcome instead of leaving tasks -parked. There is no permanently protected class or separate control reserve. -Use deterministic tie-breaking without comparing unrelated broadcasts' raw -priority values. Do not invent another subscription-priority API here. - -Regression coverage must include many partial headers and frames filling the -shared allowance, a finite stream window smaller than a frame, the largest -supported encoded unit, read-ahead across unit boundaries, cancellation and -late reset/data races, unidentified and control-only pressure, and continued -progress after selecting victims. Exercise a JavaScript sender against both -native receive paths and run all-language interop. Measure retained bytes, -throughput, and completion latency across concurrent streams and frame sizes, -including finite versus very large stream windows. Wire checks and benchmarks -into CI; no wall-clock sleeps in unit tests. - -Public API: native receive-credit integration and validation of the existing -window configuration; no new header-specific quota. Wire: existing QUIC -flow-control and STOP_SENDING/RESET_STREAM behavior, no MoQ format change. -Browser/iroh receive-side guarantees and completed-cache limits are outside -this quest. Update existing configuration and transport docs inline; no -separate guide quest is needed. - -## Related - -- [Unbounded groups](/quest/m2/unbounded-groups/README.md) - removes the frame cap; frame payload credit becomes per segment -- [Scope track priority](/quest/m1/track-priority-scope.md) - broadcast fairness domains and subscription priority semantics -- [QUIC bitrate caps](/quest/m2/rate-quic.md) - pace credit grants inside the independently bounded memory allowance -- [Relay peers get wider limits](/quest/m1/quic/peer-limits.md) - client and peer windows must both fit supported units diff --git a/quest/m2/rate-quic.md b/quest/m2/rate-quic.md index 029b8414182..a67fa943c6c 100644 --- a/quest/m2/rate-quic.md +++ b/quest/m2/rate-quic.md @@ -18,8 +18,8 @@ backpressure the peer cannot ignore, rather than metering and closing. so a modified client cannot exceed the cap. An honest one goes flow-control limited (its congestion controller does not back off) and queues until [the grant](/quest/m2/rate-grant.md) clamps its encoder. - Keep rate pacing separate from the connection memory allowance in - [message assembly](/quest/m1/quic/receive-memory.md). The memory allowance + Keep rate pacing separate from the connection memory allowance that + message assembly holds credit against. The memory allowance must fit the largest supported encoded unit even when it takes many rate intervals to receive. Limit each grant by both rate allowance and available memory backing; do not wait for a whole-message credit release to grant @@ -90,4 +90,3 @@ that a one-second rate window can also be the complete assembly allowance. ## Required - [Peer limits](/quest/m1/quic/peer-limits.md) - adds the runtime `set_limits(Limits)` seam and shrink-as-debt behavior this extends -- [Bound message assembly](/quest/m1/quic/receive-memory.md) - separate retained memory from paced outstanding credit diff --git a/quest/m2/unbounded-groups/segmented-frames.md b/quest/m2/unbounded-groups/segmented-frames.md index 7ca5d55219b..bf45bdaeac1 100644 --- a/quest/m2/unbounded-groups/segmented-frames.md +++ b/quest/m2/unbounded-groups/segmented-frames.md @@ -22,8 +22,9 @@ See the [epic](/quest/m2/unbounded-groups/README.md) for the decisions. - Choose the segment size from benchmarks, and keep it a storage detail rather than a public knob. Coalescing the chunks the transport delivers is fine. -- Update [Receive memory](/quest/m1/quic/receive-memory.md) so a frame's - payload releases credit per segment. The frame header stays a whole unit. +- Change `moq-net`'s `coding::Reader`, which holds connection credit until a + whole frame is handed off (#5735), so a frame's payload releases credit per + segment. The frame header stays a whole unit. - Measure `group.rs`, `frame.rs` (`frame_recv`), and `memory.rs` before and after. Add a large-frame case swept over frame size. diff --git a/rs/moq-ffi/src/transport/adapter.rs b/rs/moq-ffi/src/transport/adapter.rs index 46a2c8ff250..72fe88c7ad6 100644 --- a/rs/moq-ffi/src/transport/adapter.rs +++ b/rs/moq-ffi/src/transport/adapter.rs @@ -107,6 +107,9 @@ impl + Clone + 'st fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll { backend::Session::poll_closed(&mut self.0, cx).map(Error) } + fn poll_credit_exhausted(&mut self, _: &mut Context<'_>) -> Poll<()> { + Poll::Pending + } fn stats(&self) -> impl moq::Stats { Stats(backend::Session::stats(&self.0)) } @@ -141,20 +144,39 @@ impl poll::SendStream for SendStream { } /// An incoming WebTransport stream adapted to moq-net. pub struct RecvStream(S); +// The WebTransport backend returns connection credit as it reads, so there is none to hold. impl poll::RecvStream for RecvStream { type Error = Error; - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { - backend::RecvStream::poll_read(&mut self.0, cx, dst).map_err(Error) + type Credit = moq::Unheld; + type Stopper = (); + fn stopper(&self) {} + fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + backend::RecvStream::poll_read(&mut self.0, cx, dst) + .map_err(Error) + .map_ok(|n| n.map(|n| moq::Unheld(n as u64))) } fn poll_read_buf( &mut self, cx: &mut Context<'_>, buf: &mut B, - ) -> Poll, Self::Error>> { - backend::RecvStream::poll_read_buf(&mut self.0, cx, buf).map_err(Error) + ) -> Poll, Self::Error>> { + backend::RecvStream::poll_read_buf(&mut self.0, cx, buf) + .map_err(Error) + .map_ok(|n| n.map(|n| moq::Unheld(n as u64))) } - fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Self::Error>> { - backend::RecvStream::poll_read_chunk(&mut self.0, cx, max).map_err(Error) + fn poll_read_chunk( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Self::Error>> { + backend::RecvStream::poll_read_chunk(&mut self.0, cx, max) + .map_err(Error) + .map_ok(|chunk| { + chunk.map(|b| { + let n = b.len() as u64; + (b, moq::Unheld(n)) + }) + }) } fn stop(&mut self, code: u32) { backend::RecvStream::stop(&mut self.0, code) @@ -309,7 +331,7 @@ mod tests { let Poll::Ready(Ok(Some(chunk))) = poll::RecvStream::poll_read_chunk(&mut recv, &mut cx, 2) else { panic!("read failed") }; - assert_eq!(chunk, b"hi"[..]); + assert_eq!(chunk.0, b"hi"[..]); poll::RecvStream::stop(&mut recv, 13); assert_eq!(state.get(), 13); let Poll::Ready(Err(error)) = poll::SendStream::poll_closed(&mut send, &mut cx) else { diff --git a/rs/moq-net/src/client.rs b/rs/moq-net/src/client.rs index 00e3d42a6d8..3b0d29c676d 100644 --- a/rs/moq-net/src/client.rs +++ b/rs/moq-net/src/client.rs @@ -657,6 +657,10 @@ mod tests { Poll::Ready(FakeError) } + fn poll_credit_exhausted(&mut self, _cx: &mut std::task::Context<'_>) -> std::task::Poll<()> { + std::task::Poll::Pending + } + fn stats(&self) -> impl crate::transport::Stats { FakeStats { send_rate: *self.state.send_rate.lock().unwrap(), @@ -717,8 +721,15 @@ mod tests { impl crate::transport::poll::RecvStream for FakeRecvStream { type Error = FakeError; + type Credit = crate::transport::Unheld; + type Stopper = (); + fn stopper(&self) {} - fn poll_read(&mut self, _cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + fn poll_read( + &mut self, + _cx: &mut Context<'_>, + dst: &mut [u8], + ) -> Poll, Self::Error>> { if self.data.is_empty() { return Poll::Ready(Ok(None)); } @@ -727,7 +738,7 @@ mod tests { for slot in dst.iter_mut().take(size) { *slot = self.data.pop_front().unwrap(); } - Poll::Ready(Ok(Some(size))) + Poll::Ready(Ok(Some(crate::transport::Unheld(size as u64)))) } fn stop(&mut self, _code: u32) {} @@ -1056,6 +1067,10 @@ mod tests { self.inner.poll_closed(cx) } + fn poll_credit_exhausted(&mut self, _cx: &mut std::task::Context<'_>) -> std::task::Poll<()> { + std::task::Poll::Pending + } + fn stats(&self) -> impl crate::transport::Stats { self.inner.stats() } @@ -1092,8 +1107,15 @@ mod tests { impl crate::transport::poll::RecvStream for LocalRecv { type Error = FakeError; + type Credit = crate::transport::Unheld; + type Stopper = (); + fn stopper(&self) {} - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + fn poll_read( + &mut self, + cx: &mut Context<'_>, + dst: &mut [u8], + ) -> Poll, Self::Error>> { self.inner.poll_read(cx, dst) } diff --git a/rs/moq-net/src/coding/reader.rs b/rs/moq-net/src/coding/reader.rs index 5c3bf08ec8a..9e2e6fdee82 100644 --- a/rs/moq-net/src/coding/reader.rs +++ b/rs/moq-net/src/coding/reader.rs @@ -1,12 +1,18 @@ use std::{ cmp, fmt::Debug, + ops::{Deref, DerefMut}, + sync::Arc, task::{Context, Poll, ready}, }; -use bytes::{Buf, Bytes, BytesMut}; +use bytes::{Bytes, BytesMut}; -use crate::{Error, StreamError, coding::*}; +use crate::{ + Error, StreamError, + coding::*, + pressure::{Erase, Hold, Holder, Lent, Locked, Pressure, Rank}, +}; /// A reader for decoding messages from a stream. /// @@ -14,12 +20,27 @@ use crate::{Error, StreamError, coding::*}; /// wrappers, so the reader can be driven from another poll function without /// pinning a future. Partial bytes accumulate in the internal buffer, so a /// `Pending` (or a cancelled wrapper) never desynchronizes the stream. +/// +/// A [tracked](Self::track) reader holds the transport's connection credit for its +/// buffered bytes, returning each unit's as it is decoded and handed off, so read-ahead +/// stays with the next unit. Its buffer sits beside that credit in its [`Hold`], so an +/// eviction frees the memory with the credit even while the reader is parked. A partial +/// frame's payload carries its own credit in a [`HeldFrame`], which returns it when the +/// frame goes. An untracked reader returns credit as soon as it reads. pub struct Reader { stream: S, + /// The buffer until tracked, then empty: tracking moves it into the hold. buffer: BytesMut, version: V, + /// While tracked, the buffer and its credit. + hold: Option>>, + /// The receive allowance filled and this stream was chosen to free it. + evicted: bool, } +/// A chunk of a unit, with the credit lent to it if tracked. +type LentChunk = (Bytes, Option>); + /// How much payload may pile up behind a parked consumer before /// [`Reader::poll_read_frame`] wakes it, even while the transport keeps handing over /// ready chunks. @@ -37,9 +58,45 @@ impl Reader { stream, buffer: Default::default(), version, + hold: None, + evicted: false, + } + } + + /// Hold connection credit until each unit completes, evictable by `pressure` at `rank`. + pub(crate) fn track(&mut self, pressure: &Pressure, rank: Rank) + where + Hold: Erase, + { + if self.hold.is_some() { + return self.rank(rank); + } + let code = self.version.encode_stream_code(&StreamError::Evicted); + let hold = pressure.track(rank, self.stream.stopper(), code); + hold.lock().adopt(std::mem::take(&mut self.buffer)); + self.hold = Some(hold); + } + + /// Re-rank a tracked stream once its header identifies it. + pub(crate) fn rank(&self, rank: Rank) { + if let Some(hold) = &self.hold { + hold.rank(rank); } } + /// Whether this stream was evicted, registering `cx` to hear when it is. The eviction + /// already stopped the stream and took its buffer and credit. + fn poll_evicted(&mut self, cx: &Context<'_>) -> bool { + self.evicted |= self.hold.as_ref().is_some_and(|hold| hold.poll_evicted(cx)); + self.evicted + } + + /// The buffered bytes, for tests. + #[cfg(test)] + fn buffered(&mut self) -> usize { + Buffer::of(&mut self.buffer, &self.hold).len() + } + /// Poll `decode` against the buffered bytes, reading more until it stops coming up /// short. `consume` drops what it read; otherwise the bytes stay for the next read. fn poll_with( @@ -52,22 +109,24 @@ impl Reader { V: Into
+ Copy, { loop { - let mut r = Decoder::new(&self.buffer, self.version.into()); - match decode(&mut r) { + let mut buffer = Buffer::of(&mut self.buffer, &self.hold); + let mut r = Decoder::new(&buffer, self.version.into()); + let decoded = decode(&mut r); + let used = buffer.len() - r.remaining(); + match decoded { Ok(value) => { if consume { - let used = self.buffer.len() - r.remaining(); - self.buffer.advance(used); + buffer.consume(used); } return Poll::Ready(Ok(value)); } - // Stream closed while we still need more data. - Err(DecodeError::Short) if !ready!(self.poll_read_more(cx))? => { - return Poll::Ready(Err(DecodeError::Short.into())); - } - Err(DecodeError::Short) => {} + Err(DecodeError::Short) => drop(buffer), Err(e) => return Poll::Ready(Err(e.into())), } + // Stream closed while we still need more data. + if !ready!(self.poll_read_more(cx))? { + return Poll::Ready(Err(DecodeError::Short.into())); + } } } @@ -190,22 +249,48 @@ impl Reader { std::future::poll_fn(|cx| self.poll_varint_peek(cx)).await } - /// Poll for the next chunk, draining the reader's internal buffer first. + /// Poll for the next chunk, draining the reader's internal buffer first. The chunk is + /// a complete unit: its credit returns as it is handed over. pub fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Error>> { - if !self.buffer.is_empty() { - let n = cmp::min(self.buffer.len(), max); - return Poll::Ready(Ok(Some(self.buffer.split_to(n).freeze()))); + let chunk = ready!(self.poll_unit_chunk(cx, max))?; + Poll::Ready(Ok(chunk.map(|(chunk, _credit)| chunk))) + } + + /// The next chunk of a unit, draining the buffer first, with the credit lent to it: + /// whoever keeps the bytes keeps the credit, so an eviction frees only the buffer. + fn poll_unit_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll>, Error>> { + let mut buffer = Buffer::of(&mut self.buffer, &self.hold); + if !buffer.is_empty() { + let n = cmp::min(buffer.len(), max); + return Poll::Ready(Ok(Some(buffer.lend_buffered(n)))); + } + if self.evicted || buffer.is_evicted() { + self.evicted = true; + return Poll::Ready(Err(Error::Evicted)); } - self.stream - .poll_read_chunk(cx, max) - .map_err(|err| self.version.transport_error(err)) + let read = match self.stream.poll_read_chunk(cx, max) { + Poll::Pending => { + drop(buffer); + return match self.poll_evicted(cx) { + true => Poll::Ready(Err(Error::Evicted)), + false => Poll::Pending, + }; + } + Poll::Ready(read) => read, + }; + let Some((chunk, credit)) = read.map_err(|err| self.version.transport_error(err))? else { + return Poll::Ready(Ok(None)); + }; + let credit = buffer.lend(credit); + Poll::Ready(Ok(Some((chunk, credit)))) } /// Fill a frame's payload from the stream, returning `Pending` once it would block. /// /// `Ready(Ok(()))` means the frame is full and the caller should commit it; an error /// means the caller should abort it. Either way the frame is untouched otherwise, so - /// the caller keeps ownership of the commit/abort decision. + /// the caller keeps ownership of the commit/abort decision. The payload's credit moves + /// into the frame, so it returns only as the frame goes, whichever way that is. /// /// Consumers are woken where the loop yields rather than once per chunk. The wire /// hands over everything it has already buffered in one turn, and each wake pays a @@ -218,49 +303,50 @@ impl Reader { pub(crate) fn poll_read_frame( &mut self, cx: &mut Context<'_>, - frame: &mut crate::frame::ProducerOwned, + frame: &mut HeldFrame, ) -> Poll> { let mut owed = 0; - let result = loop { - if frame.remaining() == 0 { - break Ok(()); + loop { + if frame.frame.remaining() == 0 { + return Poll::Ready(Ok(())); } - match self.poll_read_chunk(cx, frame.remaining()) { + match self.poll_unit_chunk(cx, frame.frame.remaining()) { Poll::Pending => { if owed > 0 { - frame.notify(); + frame.frame.notify(); } return Poll::Pending; } - Poll::Ready(Ok(Some(chunk))) if !chunk.is_empty() => { + Poll::Ready(Ok(Some((chunk, credit)))) if !chunk.is_empty() => { + frame.keep(credit); owed += chunk.len(); - if let Err(err) = frame.write(chunk) { - break Err(err); - } + frame.frame.write(chunk)?; if owed >= WAKE_BUDGET { - frame.notify(); + frame.frame.notify(); owed = 0; } } // A FIN mid-payload: the declared size never arrived. - Poll::Ready(Ok(_)) => break Err(Error::WrongSize), - Poll::Ready(Err(err)) => break Err(err), + Poll::Ready(Ok(_)) => return Poll::Ready(Err(Error::WrongSize)), + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), } - }; - - Poll::Ready(result) + } } /// Poll for exactly `size` bytes, accumulating partial reads in the buffer. pub fn poll_read_exact(&mut self, cx: &mut Context<'_>, size: usize) -> Poll> { - while self.buffer.len() < size { + loop { + let mut buffer = Buffer::of(&mut self.buffer, &self.hold); + if buffer.len() >= size { + return Poll::Ready(Ok(buffer.hand_off(size))); + } + drop(buffer); if !ready!(self.poll_read_more(cx))? { return Poll::Ready(Err(DecodeError::Short.into())); } } - Poll::Ready(Ok(self.buffer.split_to(size).freeze())) } /// Poll to discard `remaining` bytes, counting it down as they arrive. Nothing is @@ -268,8 +354,10 @@ impl Reader { /// memory. pub(crate) fn poll_skip(&mut self, cx: &mut Context<'_>, remaining: &mut usize) -> Poll> { loop { - let skipped = self.buffer.len().min(*remaining); - self.buffer.advance(skipped); + let mut buffer = Buffer::of(&mut self.buffer, &self.hold); + let skipped = buffer.len().min(*remaining); + buffer.consume(skipped); + drop(buffer); *remaining -= skipped; if *remaining == 0 { return Poll::Ready(Ok(())); @@ -296,7 +384,7 @@ impl Reader { /// Poll for whether data is available in the buffer or stream: `false` once it finishes. pub(crate) fn poll_has_more(&mut self, cx: &mut Context<'_>) -> Poll> { - if !self.buffer.is_empty() { + if !Buffer::of(&mut self.buffer, &self.hold).is_empty() { return Poll::Ready(Ok(true)); } @@ -306,19 +394,48 @@ impl Reader { /// Poll for more data from the stream. `true` if data was read, `false` if the /// stream finished. fn poll_read_more(&mut self, cx: &mut Context<'_>) -> Poll> { - match ready!(self.stream.poll_read_buf(cx, &mut self.buffer)) { - Ok(Some(_)) => Poll::Ready(Ok(true)), + let mut buffer = Buffer::of(&mut self.buffer, &self.hold); + if self.evicted || buffer.is_evicted() { + self.evicted = true; + return Poll::Ready(Err(Error::Evicted)); + } + let read = match self.stream.poll_read_buf(cx, &mut *buffer) { + Poll::Pending => { + drop(buffer); + return match self.poll_evicted(cx) { + true => Poll::Ready(Err(Error::Evicted)), + false => Poll::Pending, + }; + } + Poll::Ready(read) => read, + }; + match read { + Ok(Some(credit)) => { + buffer.take(credit); + Poll::Ready(Ok(true)) + } Ok(None) => Poll::Ready(Ok(false)), Err(e) => Poll::Ready(Err(self.version.transport_error(e))), } } - /// Abort the stream with the given error. + /// Abort the stream with the given error, discarding what it buffered and returning + /// the credit it holds. pub fn abort(&mut self, err: impl Into) { - // STOP_SENDING is a stream operation, so it carries a stream code from the - // negotiated protocol's registry. A session code, or the other protocol's, would be - // read against the wrong table and mean something else. - self.stream.stop(self.version.encode_stream_code(&err.into())); + // An evicted stream was stopped with its own code already. + if !self.hold.as_ref().is_some_and(|hold| hold.is_evicted()) { + // STOP_SENDING is a stream operation, so it carries a stream code from the + // negotiated protocol's registry. A session code, or the other protocol's, would + // be read against the wrong table and mean something else. + self.stream.stop(self.version.encode_stream_code(&err.into())); + } + // Nothing buffered is worth keeping, and a reader retained after its abort must not + // keep memory whose credit returned. A partial frame's credit is its own, returned + // as its owner aborts it. + self.buffer = BytesMut::new(); + if let Some(hold) = &self.hold { + hold.clear(); + } } /// Cast the reader to a different version, used during version negotiation. @@ -327,6 +444,128 @@ impl Reader { stream: self.stream, buffer: self.buffer, version, + hold: self.hold, + evicted: self.evicted, + } + } +} + +/// A reader's buffered bytes: its own while untracked, else locked in its hold beside +/// their credit, so no eviction runs between a check and a read. +enum Buffer<'a, S: crate::transport::poll::RecvStream> { + Own(&'a mut BytesMut), + Held(Locked<'a, S>), +} + +impl<'a, S: crate::transport::poll::RecvStream> Buffer<'a, S> { + fn of(own: &'a mut BytesMut, hold: &'a Option>>) -> Self { + match hold { + Some(hold) => Self::Held(hold.lock()), + None => Self::Own(own), + } + } + + /// Whether an eviction emptied the buffer, after which nothing may be read into it. + fn is_evicted(&self) -> bool { + matches!(self, Self::Held(held) if held.is_evicted()) + } + + /// Keep the credit for bytes just read into the buffer, if tracked; otherwise it + /// returns now. + fn take(&mut self, credit: S::Credit) { + if let Self::Held(held) = self { + held.take(credit); + } + } + + /// Discard the first `bytes`, a unit consumed whole, then return their credit. + fn consume(&mut self, bytes: usize) { + match self { + Self::Own(buffer) => bytes::Buf::advance(buffer, bytes), + Self::Held(held) => held.consume(bytes), + } + } + + /// Hand off the first `bytes` as a complete unit, whose credit returns as it leaves. + fn hand_off(&mut self, bytes: usize) -> Bytes { + match self { + Self::Own(buffer) => buffer.split_to(bytes).freeze(), + Self::Held(held) => held.hand_off(bytes), + } + } + + /// Move the first `bytes` into a partial unit, with their credit if tracked. + fn lend_buffered(&mut self, bytes: usize) -> LentChunk { + match self { + Self::Own(buffer) => (buffer.split_to(bytes).freeze(), None), + Self::Held(held) => held.lend_buffered(bytes), + } + } + + /// The credit for bytes read straight into a partial unit, if tracked; otherwise it + /// returns now. + fn lend(&self, credit: S::Credit) -> Option> { + match self { + Self::Own(_) => None, + Self::Held(held) => Some(held.lend(credit)), + } + } +} + +/// A frame filled from a stream by [`Reader::poll_read_frame`], with the connection credit +/// for the payload it holds so far. The credit is the frame's own, so it returns only as +/// the frame finishes, aborts or drops, on every exit alike. +pub(crate) struct HeldFrame { + /// Declared before `credit`: fields drop in order, so a dropped frame frees its payload + /// before the credit returns and lets the peer send more. + frame: crate::frame::ProducerOwned, + credit: Option>, +} + +impl HeldFrame { + pub fn new(frame: crate::frame::ProducerOwned) -> Self { + Self { frame, credit: None } + } + + /// Keep the credit for payload just written into the frame. + fn keep(&mut self, credit: Option>) { + match (&mut self.credit, credit) { + (Some(held), Some(credit)) => held.merge(credit), + (held, credit @ Some(_)) => *held = credit, + (_, None) => {} + } + } + + /// Publish the complete frame, returning its credit once it is handed off. + pub fn finish(self) -> crate::Result<()> { + // `_credit` lives to the end of the scope, past the frame `finish` consumes. + let Self { frame, credit: _credit } = self; + frame.finish() + } + + /// Fail the frame with `err`, freeing its payload before its credit returns. + pub fn abort(self, err: Error) -> crate::Result<()> { + let Self { frame, credit: _credit } = self; + frame.abort(err) + } +} + +impl Deref for Buffer<'_, S> { + type Target = BytesMut; + + fn deref(&self) -> &BytesMut { + match self { + Self::Own(buffer) => buffer, + Self::Held(held) => held.buffer(), + } + } +} + +impl DerefMut for Buffer<'_, S> { + fn deref_mut(&mut self) -> &mut BytesMut { + match self { + Self::Own(buffer) => buffer, + Self::Held(held) => held.buffer_mut(), } } } @@ -348,12 +587,15 @@ mod tests { impl crate::transport::poll::RecvStream for StopLog { type Error = crate::lite::test_transport::SinkError; + type Credit = crate::transport::Unheld; + type Stopper = (); + fn stopper(&self) {} fn poll_read( &mut self, _cx: &mut std::task::Context<'_>, _dst: &mut [u8], - ) -> std::task::Poll, Self::Error>> { + ) -> std::task::Poll, Self::Error>> { std::task::Poll::Ready(Ok(None)) } @@ -420,14 +662,9 @@ mod tests { } /// A group with one open frame, plus a consumer parked on its payload. - fn parked_frame( + fn parked_frame( size: usize, - ) -> ( - crate::frame::ProducerOwned, - crate::frame::Consumer, - kio::Waiter, - Arc, - ) { + ) -> (HeldFrame, crate::frame::Consumer, kio::Waiter, Arc) { let mut group = crate::group::Info { sequence: 0 }.produce(); let mut consumer = group.consume(); let frame = group @@ -444,7 +681,7 @@ mod tests { let wakes = Arc::new(CountWaker::default()); let waiter = kio::Waiter::new(wakes.clone().into()); - (frame, payload, waiter, wakes) + (HeldFrame::new(frame), payload, waiter, wakes) } /// Hands back one queued chunk per poll, then blocks: a stream whose socket has @@ -454,8 +691,15 @@ mod tests { impl crate::transport::poll::RecvStream for Chunks { type Error = crate::lite::test_transport::SinkError; + type Credit = crate::transport::Unheld; + type Stopper = (); + fn stopper(&self) {} - fn poll_read(&mut self, _cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + fn poll_read( + &mut self, + _cx: &mut Context<'_>, + dst: &mut [u8], + ) -> Poll, Self::Error>> { let Some(chunk) = self.0.pop_front() else { return Poll::Pending; }; @@ -467,7 +711,7 @@ mod tests { self.0.push_front(&chunk[n..]); } - Poll::Ready(Ok(Some(n))) + Poll::Ready(Ok(Some(crate::transport::Unheld(n as u64)))) } fn stop(&mut self, _code: u32) {} @@ -493,7 +737,7 @@ mod tests { let mut remaining = 5; assert!(reader.poll_skip(&mut cx, &mut remaining).is_pending()); assert_eq!(remaining, 1); - assert!(reader.buffer.is_empty(), "nothing skipped is kept"); + assert_eq!(reader.buffered(), 0, "nothing skipped is kept"); } fn refuses_frame + Debug, V: StreamCodes + Into + Copy>( @@ -519,7 +763,7 @@ mod tests { "{result:?}" ); } - assert_eq!(reader.buffer.len(), prefix.len()); + assert_eq!(reader.buffered(), prefix.len()); assert_eq!(reader.stream.0.len(), 1, "read past the complete frame"); } @@ -604,7 +848,7 @@ mod tests { reader.poll_decode::(&mut cx), Poll::Ready(Ok(_)) )); - assert!(reader.buffer.is_empty()); + assert_eq!(reader.buffered(), 0); } /// A consumer parked on the group cannot run until this task yields, so a burst of @@ -655,8 +899,15 @@ mod tests { impl crate::transport::poll::RecvStream for Burst { type Error = crate::lite::test_transport::SinkError; + type Credit = crate::transport::Unheld; + type Stopper = (); + fn stopper(&self) {} - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + fn poll_read( + &mut self, + cx: &mut Context<'_>, + dst: &mut [u8], + ) -> Poll, Self::Error>> { // Start each turn with none of the last turn's registrations left. self.waiter = kio::Waiter::new(self.waiter.waker().clone()); while self.payload.poll_read_chunk(&self.waiter).is_ready() {} @@ -696,4 +947,366 @@ mod tests { // Two budgets' worth, and the boundary is owed nothing by the time it arrives. assert_eq!(wakes.count(), 2, "the burst was withheld past its budget"); } + + mod credit { + use std::task::Waker; + + use super::*; + use crate::lite::Version; + use crate::transport::{ + memory, + poll::{SendStream as _, Session as _}, + }; + + fn noop() -> Context<'static> { + Context::from_waker(Waker::noop()) + } + + /// A registry of memory streams. + type Tracked = Pressure>; + + /// A memory stream whose writer stays open, read by a reader the returned pressure + /// tracks. `window` bounds what the reader's side accepts. + fn tracked( + window: Option, + ) -> ( + memory::SendStream, + Reader, + Tracked, + (memory::Session, memory::Session), + ) { + let (mut a, mut b) = memory::pair(memory::Config { + receive_window: window, + ..Default::default() + }); + let Poll::Ready(Ok(send)) = a.poll_open_uni(&mut noop()) else { + panic!("open") + }; + let Poll::Ready(Ok(recv)) = b.poll_accept_uni(&mut noop()) else { + panic!("accept") + }; + let pressure = Tracked::default(); + let mut reader = Reader::new(recv, Version::Lite05); + reader.track(&pressure, Rank::Unidentified); + (send, reader, pressure, (a, b)) + } + + fn varint(poll: Poll>) -> u64 { + match poll { + Poll::Ready(Ok(value)) => value, + other => panic!("{other:?}"), + } + } + + fn write(send: &mut memory::SendStream, bytes: &'static [u8]) -> Poll { + send.poll_write(&mut noop(), bytes).map(|res| res.expect("write")) + } + + /// An incomplete unit keeps its credit, and the bytes read past a unit stay held + /// for the next one. + #[test] + fn credit_returns_per_complete_unit() { + let (mut send, mut reader, pressure, _sessions) = tracked(None); + + // The first byte of a two-byte varint. + assert_eq!(write(&mut send, &[0x40]), Poll::Ready(1)); + assert!(reader.poll_varint(&mut noop()).is_pending()); + assert_eq!(pressure.held(), 1); + + // The rest of it, plus a second varint read ahead in the same chunk. + assert_eq!(write(&mut send, &[0x05, 0x07]), Poll::Ready(2)); + assert_eq!(varint(reader.poll_varint(&mut noop())), 5); + assert_eq!(pressure.held(), 1); + assert_eq!(varint(reader.poll_varint(&mut noop())), 7); + assert_eq!(pressure.held(), 0); + } + + /// A peek holds what it looked at until the unit is consumed. + #[test] + fn a_peek_keeps_the_credit() { + let (mut send, mut reader, pressure, _sessions) = tracked(None); + assert_eq!(write(&mut send, &[0x07, 0x00, 0x01]), Poll::Ready(3)); + assert_eq!(varint(reader.poll_varint_peek(&mut noop())), 7); + assert_eq!(pressure.held(), 3); + let bytes = reader.poll_read_exact(&mut noop(), 2); + assert!(matches!(bytes, Poll::Ready(Ok(bytes)) if bytes.as_ref() == [0x07, 0x00])); + assert_eq!(pressure.held(), 1); + } + + /// Bytes buffered before tracking began already returned their credit. + #[test] + fn tracking_starts_from_what_is_buffered() { + let (mut a, mut b) = memory::pair(memory::Config::default()); + let Poll::Ready(Ok(mut send)) = a.poll_open_uni(&mut noop()) else { + panic!("open") + }; + let Poll::Ready(Ok(recv)) = b.poll_accept_uni(&mut noop()) else { + panic!("accept") + }; + let mut reader = Reader::new(recv, Version::Lite05); + let _sessions = (a, b); + + assert_eq!(write(&mut send, &[0x07, 0x40]), Poll::Ready(2)); + assert_eq!(varint(reader.poll_varint(&mut noop())), 7); + let pressure = Tracked::default(); + reader.track(&pressure, Rank::Unidentified); + assert_eq!(pressure.held(), 0); + + assert_eq!(write(&mut send, &[0x09, 0x01]), Poll::Ready(2)); + assert_eq!(varint(reader.poll_varint(&mut noop())), 9); + assert_eq!(pressure.held(), 1); + } + + /// An untracked reader returns credit as it reads, so an incomplete unit never + /// stalls the peer. + #[test] + fn an_untracked_reader_holds_nothing() { + let (mut a, mut b) = memory::pair(memory::Config { + receive_window: Some(2), + ..Default::default() + }); + let Poll::Ready(Ok(mut send)) = a.poll_open_uni(&mut noop()) else { + panic!("open") + }; + let Poll::Ready(Ok(recv)) = b.poll_accept_uni(&mut noop()) else { + panic!("accept") + }; + let mut reader = Reader::new(recv, Version::Lite05); + let _sessions = (a, b); + + // An eight-byte varint, two bytes at a time. + assert_eq!(write(&mut send, &[0xc0, 0x00]), Poll::Ready(2)); + assert!(reader.poll_varint(&mut noop()).is_pending()); + assert_eq!(write(&mut send, &[0x00, 0x00]), Poll::Ready(2)); + } + + /// A tracked reader's incomplete unit spends the window until it completes. + #[test] + fn a_tracked_reader_holds_the_window() { + let (mut send, mut reader, _pressure, _sessions) = tracked(Some(2)); + assert_eq!(write(&mut send, &[0xc0, 0x00]), Poll::Ready(2)); + assert!(reader.poll_varint(&mut noop()).is_pending()); + assert!(write(&mut send, &[0x00]).is_pending()); + } + + /// Aborting lets go of the credit and of what was buffered. + #[test] + fn abort_returns_the_credit() { + let (mut send, mut reader, pressure, _sessions) = tracked(None); + assert_eq!(write(&mut send, &[0x07, 0x40]), Poll::Ready(2)); + assert!(reader.poll_has_more(&mut noop()).is_ready()); + assert_eq!(pressure.held(), 2); + reader.abort(&Error::Cancel); + assert_eq!(pressure.held(), 0); + assert_eq!(reader.buffered(), 0); + } + + /// Hands out `max`-byte reads of a looping buffer of units, so only the reader costs. + struct Endless { + data: Bytes, + pos: usize, + max: usize, + } + + impl crate::transport::poll::RecvStream for Endless { + type Error = crate::lite::test_transport::SinkError; + type Credit = crate::transport::Unheld; + type Stopper = (); + fn stopper(&self) {} + + fn poll_read( + &mut self, + _: &mut Context<'_>, + dst: &mut [u8], + ) -> Poll, Self::Error>> { + let n = dst.len().min(self.max).min(self.data.len() - self.pos); + dst[..n].copy_from_slice(&self.data[self.pos..self.pos + n]); + self.pos = (self.pos + n) % self.data.len(); + Poll::Ready(Ok(Some(crate::transport::Unheld(n as u64)))) + } + + fn stop(&mut self, _code: u32) {} + + fn poll_closed(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Pending + } + } + + /// Decode size-prefixed units, untracked and tracked, as messages and as frames: the + /// per-read and per-unit cost of holding credit in the reader. Run by nightly. + #[test] + #[ignore = "release benchmark, run by nightly"] + fn benchmark_tracked_reads() { + const UNIT: usize = 1000; + const UNITS: usize = 1 << 16; + const RUNS: usize = 5; + let mut unit = Vec::new(); + Encoder::new(&mut unit, Version::Lite05.into()) + .varint(UNIT as u64) + .unwrap(); + unit.resize(unit.len() + UNIT, 7); + let data = Bytes::from(unit.repeat(1024)); + + println!("max_read,untracked_ns_per_unit,tracked_ns_per_unit,untracked_ns_per_frame,tracked_ns_per_frame"); + for max in [1200usize, 16 * 1024, 64 * 1024] { + let mut ns = [f64::MAX; 4]; + for _ in 0..RUNS { + let modes = [(false, false), (true, false), (false, true), (true, true)]; + for (mode, (tracked, framed)) in modes.into_iter().enumerate() { + let stream = Endless { + data: data.clone(), + pos: 0, + max, + }; + let mut reader = Reader::new(stream, Version::Lite05); + let pressure = Pressure::>::default(); + if tracked { + reader.track(&pressure, Rank::Unidentified); + } + let mut group = crate::group::Info { sequence: 0 }.produce(); + let start = std::time::Instant::now(); + for i in 0..UNITS { + let size = varint(reader.poll_varint(&mut noop())); + if !framed { + let Poll::Ready(Ok(payload)) = reader.poll_read_exact(&mut noop(), size as usize) + else { + panic!("short unit"); + }; + std::hint::black_box(payload); + continue; + } + // A fresh group now and then, so the cache stays small. + if i % 64 == 0 { + group = crate::group::Info { sequence: i as u64 }.produce(); + } + let info = crate::frame::Info { + size, + timestamp: Some(crate::Timestamp::ZERO), + }; + let frame = group.create_frame_owned(info, &Default::default()).unwrap(); + let mut frame = HeldFrame::new(frame); + let Poll::Ready(Ok(())) = reader.poll_read_frame(&mut noop(), &mut frame) else { + panic!("short frame"); + }; + frame.finish().unwrap(); + } + ns[mode] = ns[mode].min(start.elapsed().as_nanos() as f64 / UNITS as f64); + } + } + println!("{max},{:.1},{:.1},{:.1},{:.1}", ns[0], ns[1], ns[2], ns[3]); + } + } + + /// An evicted reader stops its stream and fails its next read. + #[test] + fn eviction_stops_the_stream() { + let (mut send, mut reader, pressure, _sessions) = tracked(None); + assert_eq!(write(&mut send, &[0x40]), Poll::Ready(1)); + assert!(reader.poll_varint(&mut noop()).is_pending()); + + assert!(pressure.evict()); + assert!(matches!( + reader.poll_varint(&mut noop()), + Poll::Ready(Err(Error::Evicted)) + )); + assert_eq!(pressure.held(), 0); + assert!(matches!( + send.poll_write(&mut noop(), &[0x05]), + Poll::Ready(Err(memory::Error::Stopped(_))) + )); + } + + /// An evicted reader mid-frame keeps the credit for the payload already in its + /// partial frame, which is still allocated, until it gives the frame up. + #[test] + fn a_partial_frame_keeps_its_credit_through_eviction() { + let (mut send, mut reader, pressure, _sessions) = tracked(None); + let (mut frame, _payload, _waiter, _wakes) = parked_frame(9); + + // Six bytes read ahead into the buffer, then moved into the frame, and two more + // read straight into it. + assert_eq!(write(&mut send, b"abcdef"), Poll::Ready(6)); + assert!(reader.poll_has_more(&mut noop()).is_ready()); + assert!(reader.poll_read_frame(&mut noop(), &mut frame).is_pending()); + assert_eq!(write(&mut send, b"gh"), Poll::Ready(2)); + assert!(reader.poll_read_frame(&mut noop(), &mut frame).is_pending()); + assert_eq!(pressure.held(), 8); + + assert!(pressure.evict()); + assert_eq!( + pressure.held(), + 8, + "the partial frame's credit returned while it was allocated" + ); + assert!(matches!( + reader.poll_read_frame(&mut noop(), &mut frame), + Poll::Ready(Err(Error::Evicted)) + )); + reader.abort(&Error::Evicted); + assert_eq!(pressure.held(), 8, "the reader returned its frame's credit"); + frame.abort(Error::Evicted).unwrap(); + assert_eq!(pressure.held(), 0); + } + + /// A FIN mid-payload fails the frame, whose credit returns only once its owner drops + /// it: the payload is still allocated until then. + #[test] + fn a_truncated_frame_keeps_its_credit_until_dropped() { + let (mut send, mut reader, pressure, _sessions) = tracked(None); + let (mut frame, _payload, _waiter, _wakes) = parked_frame(9); + + assert_eq!(write(&mut send, b"abc"), Poll::Ready(3)); + assert!(reader.poll_read_frame(&mut noop(), &mut frame).is_pending()); + send.finish().expect("finish"); + assert!(matches!( + reader.poll_read_frame(&mut noop(), &mut frame), + Poll::Ready(Err(Error::WrongSize)) + )); + assert_eq!(pressure.held(), 3, "the credit returned before its frame dropped"); + drop(frame); + assert_eq!(pressure.held(), 0); + } + + /// Parked readers filled and evicted in turn never buffer more than the window. An + /// eviction that returned the credit but left a parked reader its read-ahead would + /// let the peer refill the window while every evicted reader kept its bytes. + #[test] + fn evicted_parked_readers_free_their_read_ahead() { + const WINDOW: usize = 16; + let (mut a, mut b) = memory::pair(memory::Config { + receive_window: Some(WINDOW), + ..Default::default() + }); + let pressure = Tracked::default(); + // One message, then the start of the next, which the reader keeps as read-ahead. + const MESSAGE: &[u8] = &[0x07, 0x40, 0x01, 0x40, 0x02, 0x40, 0x03, 0x40]; + let mut parked = Vec::new(); + for _ in 0..64 { + let Poll::Ready(Ok(mut send)) = a.poll_open_uni(&mut noop()) else { + panic!("open") + }; + let mut rest = MESSAGE; + while !rest.is_empty() { + match write(&mut send, rest) { + Poll::Ready(n) => rest = &rest[n..], + // Full of parked read-ahead: evict one, as the monitor would. + Poll::Pending => assert!(pressure.evict(), "a full window has a holder"), + } + } + let Poll::Ready(Ok(recv)) = b.poll_accept_uni(&mut noop()) else { + panic!("accept") + }; + let mut reader = Reader::new(recv, Version::Lite05); + reader.track(&pressure, Rank::Unidentified); + assert_eq!(varint(reader.poll_varint(&mut noop())), 7); + parked.push((send, reader)); + + let buffered: usize = parked.iter_mut().map(|(_, reader)| reader.buffered()).sum(); + assert!( + buffered <= WINDOW, + "{buffered} bytes buffered past a {WINDOW}-byte window" + ); + } + } + } } diff --git a/rs/moq-net/src/fuzz.rs b/rs/moq-net/src/fuzz.rs index 8b9ff9f52f9..89ad792a018 100644 --- a/rs/moq-net/src/fuzz.rs +++ b/rs/moq-net/src/fuzz.rs @@ -1080,6 +1080,9 @@ pub fn seeds() -> Vec { seeds } +/// A stream's reader and the frame it fills. +type Receiving = (coding::Reader, coding::HeldFrame); + /// Frames arriving on concurrent group streams, as the frame bench times them: one /// frame per stream, read a chunk per stream per poll turn so the streams interleave /// the way a session's do. @@ -1088,7 +1091,7 @@ pub fn seeds() -> Vec { /// every declared size when the first byte lands and `Some(0)` grows every buffer with /// the bytes received. pub struct FrameRecv { - streams: Vec, frame::ProducerOwned)>>, + streams: Vec>, groups: Vec, // Keeps the groups' track alive. _track: track::Producer, @@ -1127,7 +1130,10 @@ impl FrameRecv { chunk, ready: true, }; - Some((coding::Reader::new(stream, lite::Version::Lite05), frame)) + Some(( + coding::Reader::new(stream, lite::Version::Lite05), + coding::HeldFrame::new(frame), + )) }) .collect(); @@ -1169,17 +1175,20 @@ struct Chunks { impl crate::transport::poll::RecvStream for Chunks { type Error = NoError; + type Credit = crate::transport::Unheld; + type Stopper = (); + fn stopper(&self) {} fn poll_read( &mut self, cx: &mut std::task::Context<'_>, dst: &mut [u8], - ) -> std::task::Poll, NoError>> { + ) -> std::task::Poll, NoError>> { let max = dst.len(); self.poll_read_chunk(cx, max).map_ok(|chunk| { - chunk.map(|chunk| { + chunk.map(|(chunk, credit)| { dst[..chunk.len()].copy_from_slice(&chunk); - chunk.len() + credit }) }) } @@ -1189,7 +1198,7 @@ impl crate::transport::poll::RecvStream for Chunks { &mut self, _cx: &mut std::task::Context<'_>, max: usize, - ) -> std::task::Poll, NoError>> { + ) -> std::task::Poll, NoError>> { self.ready = !self.ready; if self.ready { return std::task::Poll::Pending; @@ -1198,7 +1207,7 @@ impl crate::transport::poll::RecvStream for Chunks { return std::task::Poll::Ready(Ok(None)); } let n = self.chunk.min(max).min(self.payload.len()); - std::task::Poll::Ready(Ok(Some(self.payload.split_to(n)))) + std::task::Poll::Ready(Ok(Some((self.payload.split_to(n), crate::transport::Unheld(n as u64))))) } fn stop(&mut self, _code: u32) {} diff --git a/rs/moq-net/src/ietf/adapter.rs b/rs/moq-net/src/ietf/adapter.rs index 7c7bb3959bb..9b41bd4badb 100644 --- a/rs/moq-net/src/ietf/adapter.rs +++ b/rs/moq-net/src/ietf/adapter.rs @@ -10,6 +10,7 @@ use crate::{ Error, PathOwned, coding::{Decode, Decoder, Encoder, Reader, Writer}, ietf::{self, RequestId}, + transport::{Credit, Unheld}, }; use super::{Control, Message, Version, control::RequestPermit}; @@ -159,12 +160,20 @@ impl VirtualRecvStream { } } +/// The routed control messages a virtual stream holds were complete when the real control +/// stream handed them off, so their connection credit is already returned. impl crate::transport::poll::RecvStream for VirtualRecvStream { type Error = crate::Error; + type Credit = Unheld; + type Stopper = VirtualStopper; - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + fn stopper(&self) -> VirtualStopper { + VirtualStopper(self.rx.clone()) + } + + fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { if dst.is_empty() { - return Poll::Ready(Ok(Some(0))); + return Poll::Ready(Ok(Some(Unheld(0)))); } std::task::ready!(self.poll_fill(cx)); @@ -175,16 +184,16 @@ impl crate::transport::poll::RecvStream for VirtualRecvStream { let n = dst.len().min(self.buffer.len()); dst[..n].copy_from_slice(&self.buffer[..n]); self.buffer.advance(n); - Poll::Ready(Ok(Some(n))) + Poll::Ready(Ok(Some(Unheld(n as u64)))) } fn poll_read_buf( &mut self, cx: &mut Context<'_>, buf: &mut B, - ) -> Poll, Self::Error>> { + ) -> Poll, Self::Error>> { if !buf.has_remaining_mut() { - return Poll::Ready(Ok(Some(0))); + return Poll::Ready(Ok(Some(Unheld(0)))); } std::task::ready!(self.poll_fill(cx)); @@ -194,12 +203,16 @@ impl crate::transport::poll::RecvStream for VirtualRecvStream { let n = buf.remaining_mut().min(self.buffer.len()); buf.put(self.buffer.split_to(n)); - Poll::Ready(Ok(Some(n))) + Poll::Ready(Ok(Some(Unheld(n as u64)))) } - fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Self::Error>> { + fn poll_read_chunk( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Self::Error>> { if max == 0 { - return Poll::Ready(Ok(Some(Bytes::new()))); + return Poll::Ready(Ok(Some((Bytes::new(), Unheld(0))))); } std::task::ready!(self.poll_fill(cx)); @@ -208,7 +221,7 @@ impl crate::transport::poll::RecvStream for VirtualRecvStream { } let n = max.min(self.buffer.len()); - Poll::Ready(Ok(Some(self.buffer.split_to(n)))) + Poll::Ready(Ok(Some((self.buffer.split_to(n), Unheld(n as u64))))) } fn stop(&mut self, _code: u32) { @@ -235,6 +248,15 @@ impl crate::transport::poll::RecvStream for VirtualRecvStream { } } +/// Stops a [`VirtualRecvStream`] without borrowing it, discarding what is routed to it. +pub struct VirtualStopper(Queue); + +impl crate::transport::Stopper for VirtualStopper { + fn stop(&self, _code: u32) { + self.0.close(); + } +} + impl Drop for VirtualRecvStream { fn drop(&mut self) { // The reader is gone: close the queue so routed follow-ups are discarded @@ -469,13 +491,73 @@ pub enum AdapterRecv { Virtual(VirtualRecvStream), } +/// The credit of an [`AdapterRecv`]: the real stream's, or none for a virtual one. +pub enum AdapterCredit { + Real(C), + Virtual(Unheld), +} + +impl Credit for AdapterCredit { + fn len(&self) -> u64 { + match self { + Self::Real(c) => c.len(), + Self::Virtual(c) => c.len(), + } + } + + fn split(&mut self, bytes: u64) -> Self { + match self { + Self::Real(c) => Self::Real(c.split(bytes)), + Self::Virtual(c) => Self::Virtual(c.split(bytes)), + } + } + + fn merge(&mut self, other: Self) { + match (self, other) { + (Self::Real(c), Self::Real(other)) => c.merge(other), + (Self::Virtual(c), Self::Virtual(other)) => c.merge(other), + _ => panic!("merged a real stream's credit with a virtual one's"), + } + } +} + +type RealCredit = <::RecvStream as crate::transport::poll::RecvStream>::Credit; +type RealStopper = + <::RecvStream as crate::transport::poll::RecvStream>::Stopper; + +/// Stops an [`AdapterRecv`] without borrowing it. +pub enum AdapterStopper { + Real(T), + Virtual(VirtualStopper), +} + +impl crate::transport::Stopper for AdapterStopper { + fn stop(&self, code: u32) { + match self { + Self::Real(s) => s.stop(code), + Self::Virtual(s) => s.stop(code), + } + } +} + impl crate::transport::poll::RecvStream for AdapterRecv { type Error = crate::Error; + type Credit = AdapterCredit>; + type Stopper = AdapterStopper>; + + fn stopper(&self) -> Self::Stopper { + match self { + Self::Real(s) => AdapterStopper::Real(s.stopper()), + Self::Virtual(s) => AdapterStopper::Virtual(s.stopper()), + } + } - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { match self { - Self::Real(s) => s.poll_read(cx, dst).map(|r| r.map_err(|_| crate::Error::Closed)), - Self::Virtual(s) => s.poll_read(cx, dst), + Self::Real(s) => s + .poll_read(cx, dst) + .map(|r| r.map(|c| c.map(AdapterCredit::Real)).map_err(|_| crate::Error::Closed)), + Self::Virtual(s) => s.poll_read(cx, dst).map_ok(|c| c.map(AdapterCredit::Virtual)), } } @@ -483,17 +565,28 @@ impl crate::transport::poll::RecvStream for &mut self, cx: &mut Context<'_>, buf: &mut B, - ) -> Poll, Self::Error>> { + ) -> Poll, Self::Error>> { match self { - Self::Real(s) => s.poll_read_buf(cx, buf).map(|r| r.map_err(|_| crate::Error::Closed)), - Self::Virtual(s) => s.poll_read_buf(cx, buf), + Self::Real(s) => s + .poll_read_buf(cx, buf) + .map(|r| r.map(|c| c.map(AdapterCredit::Real)).map_err(|_| crate::Error::Closed)), + Self::Virtual(s) => s.poll_read_buf(cx, buf).map_ok(|c| c.map(AdapterCredit::Virtual)), } } - fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Self::Error>> { + fn poll_read_chunk( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Self::Error>> { match self { - Self::Real(s) => s.poll_read_chunk(cx, max).map(|r| r.map_err(|_| crate::Error::Closed)), - Self::Virtual(s) => s.poll_read_chunk(cx, max), + Self::Real(s) => s.poll_read_chunk(cx, max).map(|r| { + r.map(|c| c.map(|(b, c)| (b, AdapterCredit::Real(c)))) + .map_err(|_| crate::Error::Closed) + }), + Self::Virtual(s) => s + .poll_read_chunk(cx, max) + .map_ok(|c| c.map(|(b, c)| (b, AdapterCredit::Virtual(c)))), } } @@ -1213,6 +1306,10 @@ impl crate::transport::poll::Session for Con self.inner.poll_closed(cx).map(|_| crate::Error::Closed) } + fn poll_credit_exhausted(&mut self, cx: &mut Context<'_>) -> Poll<()> { + self.inner.poll_credit_exhausted(cx) + } + fn stats(&self) -> impl crate::transport::Stats { self.inner.stats() } @@ -1392,14 +1489,14 @@ mod tests { .unwrap(); let (_, mut recv) = shared.incoming.pop().await.unwrap(); assert_eq!( - recv.read_chunk(usize::MAX).await.unwrap(), + recv.read_chunk(usize::MAX).await.unwrap().map(|(b, _)| b), Some(Bytes::from_static(b"request")) ); assert_eq!(routed, target_id, "{version}: update {update_id} routed to {routed}"); shared.push(routed, raw.clone()); assert_eq!( - recv.read_chunk(usize::MAX).await.unwrap(), + recv.read_chunk(usize::MAX).await.unwrap().map(|(b, _)| b), Some(raw), "{version}: update {update_id} did not reach {target_id}" ); @@ -1463,12 +1560,12 @@ mod tests { // Read initial data let mut buf = [0u8; 32]; - let n = stream.read(&mut buf).await.unwrap().unwrap(); + let n = stream.read(&mut buf).await.unwrap().unwrap().0 as usize; assert_eq!(&buf[..n], b"initial"); // Send follow-up tx.push(Bytes::from_static(b"followup")); - let n = stream.read(&mut buf).await.unwrap().unwrap(); + let n = stream.read(&mut buf).await.unwrap().unwrap().0 as usize; assert_eq!(&buf[..n], b"followup"); // Drop the writer → FIN @@ -1484,14 +1581,14 @@ mod tests { // Read small chunks let mut buf = [0u8; 5]; - let n = stream.read(&mut buf).await.unwrap().unwrap(); + let n = stream.read(&mut buf).await.unwrap().unwrap().0 as usize; assert_eq!(&buf[..n], b"hello"); - let n = stream.read(&mut buf).await.unwrap().unwrap(); + let n = stream.read(&mut buf).await.unwrap().unwrap().0 as usize; assert_eq!(&buf[..n], b" worl"); let mut buf = [0u8; 1]; - let n = stream.read(&mut buf).await.unwrap().unwrap(); + let n = stream.read(&mut buf).await.unwrap().unwrap().0 as usize; assert_eq!(&buf[..n], b"d"); } @@ -1589,7 +1686,7 @@ mod tests { // Drain the advertisement the stream opened with, so both halves are quiescent // and a later read only sees what the terminal message routes. assert_eq!( - recv.read_chunk(usize::MAX).await.unwrap(), + recv.read_chunk(usize::MAX).await.unwrap().map(|(b, _)| b), Some(encode_msg(&msg, version)) ); recv diff --git a/rs/moq-net/src/ietf/auth.rs b/rs/moq-net/src/ietf/auth.rs index 0cceeaed4bd..8a173239544 100644 --- a/rs/moq-net/src/ietf/auth.rs +++ b/rs/moq-net/src/ietf/auth.rs @@ -235,6 +235,7 @@ pub(super) fn peer_grant( /// says it negotiated the extension. Never returns: a session without it just fails /// the tokens as unsupported. #[cfg(feature = "async")] +#[allow(clippy::too_many_arguments)] pub(super) async fn run_present( runtime: crate::time::Deadlines, session: S, @@ -243,6 +244,7 @@ pub(super) async fn run_present( peer_setup: peer::PeerSetup, version: Version, going_away: crate::goaway::GoingAway, + pressure: crate::pressure::Pressure, ) -> Result<(), Error> where S: crate::transport::poll::Boxable, @@ -273,6 +275,7 @@ where handle: handle.clone(), version, id, + pressure: pressure.clone(), }; let going_away = going_away.clone(); let violation = violation.clone(); @@ -308,6 +311,7 @@ struct Present { handle: Handle, version: Version, id: u64, + pressure: crate::pressure::Pressure, } #[cfg(feature = "async")] @@ -323,6 +327,7 @@ impl Present { Ok(stream) => stream, Err(err) => return err, }; + stream.reader.track(&self.pressure, crate::pressure::Rank::Control); if let Err(err) = stream.writer.encode_message(&Auth { request_id, token }).await { return err; } diff --git a/rs/moq-net/src/ietf/publisher.rs b/rs/moq-net/src/ietf/publisher.rs index f86cfcb4e6e..25fd8d3890c 100644 --- a/rs/moq-net/src/ietf/publisher.rs +++ b/rs/moq-net/src/ietf/publisher.rs @@ -614,6 +614,8 @@ pub(super) struct Publisher { // Subscriptions the peer may hold at once (`session::Limits::subscriptions`). pub(super) subscriptions: crate::session::Slots, pub(super) fetches: crate::session::Slots, + /// Tracks each stream this half opens. + pub(super) pressure: crate::pressure::Pressure, // Ranks every subscription's group streams against each other for the transport. priority: PriorityQueue, } @@ -690,6 +692,7 @@ where owed: Default::default(), subscriptions: Default::default(), fetches: Default::default(), + pressure: Default::default(), priority: Default::default(), } } @@ -3323,7 +3326,12 @@ where /// peer is waiting on us to retire. Failing instead keeps the loop moving. async fn open_request(&self) -> Result>, Error> { let mut session = self.session.clone(); - let mut open = std::pin::pin!(Stream::open(&mut session, self.version)); + let pressure = self.pressure.clone(); + let mut open = std::pin::pin!(async { + let mut stream = Stream::open(&mut session, self.version).await?; + stream.reader.track(&pressure, crate::pressure::Rank::Control); + Ok(stream) + }); let mut timeout = crate::time::Deadline::after(&self.runtime, ADVERTISE_TIMEOUT); kio::wait(|waiter| { diff --git a/rs/moq-net/src/ietf/session.rs b/rs/moq-net/src/ietf/session.rs index c04e850beef..a77d587ef4b 100644 --- a/rs/moq-net/src/ietf/session.rs +++ b/rs/moq-net/src/ietf/session.rs @@ -3,6 +3,7 @@ use crate::{ Error, Hop, SessionError, StreamError, coding::{Decode, DecodeError, Encode, Reader, Stream, Writer}, ietf::{self, FetchHeader, RequestId}, + pressure::Rank, setup, util::{MaybeBoxedExt, MaybeSendBox, TaskSet, err_only}, }; @@ -178,9 +179,14 @@ where None => {} } + // Every stream the session reads holds its credit until a whole message arrives, and + // the monitor cancels one by rank whenever that fills the receive allowance. + let pressure = crate::pressure::Pressure::::default(); + let driver = async move { // Held for the life of the session. let _setup_token = setup_token; + let mut monitor = crate::pressure::Monitor::new(session.clone(), pressure.clone()); // Released on any exit, so nothing waits on a token the session will never answer. let _auth_close = AuthClose(auth.clone()); // Decided once, before any Auth request can be accepted: the app took the requests @@ -205,7 +211,7 @@ where let res = match version { Version::Draft14 | Version::Draft15 | Version::Draft16 => { - let Some(setup) = setup else { + let Some(mut setup) = setup else { let err = Error::ProtocolViolation; session.close( super::error::to_session_code(&SessionError::from(&err)), @@ -213,6 +219,7 @@ where ); return Err(err); }; + setup.reader.track(&pressure, Rank::Control); let control = Control::new(request_id_max, client).with_window(limits.requests(), client, version); let adapter = ControlStreamAdapter::new(session.clone(), control.clone(), version); @@ -228,6 +235,7 @@ where ) .with_auth(auth.clone()); let (tasks, mut task_set) = TaskSet::new(); + publisher.pressure = pressure.clone(); publisher.withdrawal = withdrawing.clone(); publisher.owed = serving.clone(); publisher.subscriptions = subscriptions; @@ -246,6 +254,7 @@ where tasks.clone(), goaway.going_away.clone(), ) + .with_pressure(pressure.clone()) .with_auth(auth.clone()) .with_solicit(extensions.solicit); subscriber.announces = announces; @@ -315,7 +324,7 @@ where let mut sub_ns = sub_ns.clone(); let sub_ns_adapter = sub_ns_adapter.clone(); prefixes.push(async move { - let stream = match version { + let mut stream = match version { Version::Draft16 => { let mut sub_ns_adapter = sub_ns_adapter; let (send, recv) = sub_ns_adapter.open_native_bi().await?; @@ -326,6 +335,7 @@ where } _ => Stream::open(&mut sub_ns_adapter.clone(), version).await?, }; + stream.reader.track(&sub_ns.pressure, Rank::Control); if let Err(err) = sub_ns.run_subscribe_namespace(stream, prefix).await { // The peer breaking the protocol is fatal, and the driver // below turns this into the session close the draft wants. @@ -345,6 +355,7 @@ where let res = kio::wait(|waiter| { use std::task::Poll; + let _ = monitor.poll(&mut waiter.context()); if let Poll::Ready(err) = waiter.poll_future(adapter_run.as_mut()) { return Poll::Ready(Err::<(), Error>(err)); } @@ -435,6 +446,7 @@ where ) .with_auth(auth.clone()); let (tasks, mut task_set) = TaskSet::new(); + publisher.pressure = pressure.clone(); publisher.withdrawal = withdrawing.clone(); publisher.owed = serving.clone(); publisher.subscriptions = subscriptions; @@ -453,6 +465,7 @@ where tasks, goaway.going_away.clone(), ) + .with_pressure(pressure.clone()) .with_auth(auth.clone()) .with_solicit(extensions.solicit); @@ -465,6 +478,7 @@ where peer_setup.clone(), version, goaway.going_away.clone(), + pressure.clone(), ); // Only a side that offered MoQ Auth answers the peer's Auth requests: one that // arrives otherwise falls through to the protocol violation in `run_dispatch`. @@ -483,9 +497,14 @@ where // arrives on the wire. let goaway_recv = { let goaway = goaway.clone(); + let pressure = pressure.clone(); async move { match peer_setup_stream { - Some(reader) => run_goaway(reader.with_version(version), version, client, goaway).await, + Some(reader) => { + let mut reader = reader.with_version(version); + reader.track(&pressure, Rank::Control); + run_goaway(reader, version, client, goaway).await + } None => std::future::pending().await, } } @@ -530,7 +549,8 @@ where let sub_ns_session = sub_ns_session.clone(); prefixes.push(async move { let mut sub_ns_session = sub_ns_session; - let stream = Stream::open(&mut sub_ns_session, version).await?; + let mut stream = Stream::open(&mut sub_ns_session, version).await?; + stream.reader.track(&sub_ns.pressure, Rank::Control); if let Err(err) = sub_ns.run_subscribe_namespace(stream, prefix).await { // The peer breaking the protocol is fatal, and the driver // below turns this into the session close the draft wants. @@ -550,6 +570,7 @@ where let res = kio::wait(|waiter| { use std::task::Poll; + let _ = monitor.poll(&mut waiter.context()); if let Poll::Ready(err) = waiter.poll_future(unis.as_mut()) { return Poll::Ready(Err::<(), Error>(err)); } @@ -942,7 +963,7 @@ where let mut early = early.into_iter(); loop { - let reader = match early.next() { + let mut reader = match early.next() { Some(reader) => reader, None => { let recv = tasks @@ -978,6 +999,7 @@ where Reader::new(recv, outer_version) } }; + reader.track(&subscriber.pressure, Rank::Unidentified); let uni = uni.clone(); let fatal = fatal.clone(); @@ -1063,6 +1085,7 @@ impl Uni { if seen_setup.swap(true, std::sync::atomic::Ordering::Relaxed) { return Err(Error::ProtocolViolation); } + reader.rank(Rank::Control); // The negotiation gates the announce and dispatch loops, so a SETUP we // cannot read must end the session rather than leave them parked on a @@ -1215,10 +1238,11 @@ where Stream::poll_accept(&mut session, version, &mut cx) }) .await; - let stream = match accepted { + let mut stream = match accepted { Ok(stream) => stream, Err(err) => break err, }; + stream.reader.track(&bidi.subscriber.pressure, Rank::Control); // QUIC may expose an earlier stream before its bytes arrive. Reading the // header here could block later streams whose unread data holds flow control. diff --git a/rs/moq-net/src/ietf/subscriber.rs b/rs/moq-net/src/ietf/subscriber.rs index d237a236743..90b50c78e5f 100644 --- a/rs/moq-net/src/ietf/subscriber.rs +++ b/rs/moq-net/src/ietf/subscriber.rs @@ -10,7 +10,9 @@ use crate::{ coding::{Decode, DecodeError, Decoder, Reader, Stream, Writer}, frame, group, ietf::{self, Control, FetchType, Filter, GroupOrder, RequestId}, - origin, track, + origin, + pressure::{Pressure, Rank}, + track, util::{MaybeBoxedExt, MaybeSendBox, TaskSet, Tasks}, }; @@ -500,8 +502,9 @@ struct TrackState { // The broadcast this track was subscribed from. With the track name it forms the full // track name, which is what decides whether a repeated alias is the fatal collision - // (one alias, two tracks) or the legal sharing of an alias across subscriptions. - broadcast: PathOwned, + // (one alias, two tracks) or the legal sharing of an alias across subscriptions. Shared + // with the rank of every stream it carries, so victim choice compares one pointer. + broadcast: std::sync::Arc, // Units for this track's object Timestamps, from the TIMESCALE Track Property in // SUBSCRIBE_OK. `None` until it arrives, and for a track that declares none: its @@ -524,6 +527,11 @@ struct TrackState { // The data streams read so far, which PUBLISH_DONE's Stream Count is checked against. tail: kio::Producer, + + // The priority and group order this subscription was sent with, which rank its data + // streams when held receive credit fills. + priority: u8, + order: track::Order, } impl TrackState { @@ -545,7 +553,7 @@ impl TrackState { pending: None, name, alias: None, - broadcast, + broadcast: std::sync::Arc::new(broadcast), timescale: None, fill, largest: None, @@ -553,6 +561,18 @@ impl TrackState { joining, resume: None, tail: Default::default(), + priority: 0, + order: Default::default(), + } + } + + /// The rank of this subscription's data stream for `group`, starting at `start`. + fn rank(&self, subscribe: RequestId, group: u64, start: u64) -> Rank { + Rank::Media { + broadcast: self.broadcast.clone(), + priority: crate::priority::Priority::new(self.priority, subscribe.0, group) + .with_order(self.order) + .with_start(start), } } } @@ -698,6 +718,8 @@ pub(super) struct Subscriber { frames: frame::Budget, // Namespaces the peer may have announced at once (`session::Limits::announces`). pub(super) announces: crate::session::Slots, + // Every stream this half reads, cancelled by rank when held receive credit fills. + pub(super) pressure: Pressure, } /// The prefixes to issue SUBSCRIBE_NAMESPACE for: `origin`'s permitted scope, @@ -782,10 +804,18 @@ where auth: crate::auth::Handle::new(false), frames: Default::default(), announces: Default::default(), + pressure: Default::default(), validity: Default::default(), } } + /// Open a request stream, its reads tracked as control. + async fn open_stream(&self) -> Result, Error> { + let mut stream = Stream::open(&mut self.session.clone(), self.version).await?; + stream.reader.track(&self.pressure, Rank::Control); + Ok(stream) + } + /// Bound what we subscribe to by the grant this session's tokens earn (MoQ Auth), /// and what the peer may publish to us by the session's limit, as either changes. pub fn with_auth(mut self, auth: crate::auth::Handle) -> Self { @@ -795,6 +825,13 @@ where self } + /// Track the streams this half reads with the session's `pressure`. Set before the + /// halves that clone this subscriber start. + pub fn with_pressure(mut self, pressure: Pressure) -> Self { + self.pressure = pressure; + self + } + /// Whether our SETUP declared MoQ Solicit. Without it an unsolicited PUBLISH_NAMESPACE /// is what we invited, even from a peer that implements the extension. pub fn with_solicit(mut self, solicit: bool) -> Self { @@ -2521,7 +2558,7 @@ where } }; - let mut stream = match Stream::open(&mut self.session.clone(), self.version).await { + let mut stream = match self.open_stream().await { Ok(s) => s, Err(err) => { tracing::debug!(%err, "failed to open subscribe stream"); @@ -2537,6 +2574,11 @@ where true => Fill::Requested, false => Fill::Done, }); + // The aggregate is read now: a subscriber can join while the request ID and stream + // were awaited, and nothing updates the priority or order after. + let aggregate = target.subscription(); + let priority = aggregate.as_ref().map(|s| s.priority).unwrap_or(0); + let order = aggregate.map(|s| s.order).unwrap_or_default(); { let mut state = self.state.lock(); state.subscribes.insert( @@ -2548,18 +2590,15 @@ where .filter(|start| start.frame != 0), // A resumed copy is the session's to abort from here on. producer: target.producer().cloned(), + priority, + order, ..TrackState::pending(track_name.to_owned(), broadcast_path.to_owned(), fill.clone(), joining) }, ); } - // Write Subscribe message. The aggregate is read now: a subscriber can join while - // the request ID and stream were awaited, and nothing updates the priority or - // order after. - let aggregate = target.subscription(); - let priority = aggregate.as_ref().map(|s| s.priority).unwrap_or(0); - let order = aggregate.map(|s| s.order).unwrap_or_default(); - // Without an Established subscription, a relay includes the request it holds. + // Write Subscribe message. Without an Established subscription, a relay includes the + // request it holds. let new_group = target .new_group() .filter(|_| super::subscribe::has_new_group_request(self.version)); @@ -3275,7 +3314,7 @@ where let work = async move { // No ID is spent while opening a stream, so dropping a blocked open // cannot leave a gap in the peer's sequential request IDs. - let mut stream = match Stream::open(&mut this.session.clone(), this.version).await { + let mut stream = match this.open_stream().await { Ok(stream) => stream, Err(err) => { tracing::debug!(%err, "failed to open joining FETCH stream"); @@ -3476,7 +3515,7 @@ where return None; } }; - let mut stream = match Stream::open(&mut self.session.clone(), self.version).await { + let mut stream = match self.open_stream().await { Ok(stream) => stream, Err(err) => { request.reject(err); @@ -3622,6 +3661,7 @@ where let (mut track, timescale, fill, resume, mut reading) = { let state = self.state.lock(); let track = state.subscribes.get(&request_id).ok_or(Error::NotFound)?; + stream.rank(track.rank(request_id, group.group_id, 0)); ( track.producer.clone().ok_or(Error::NotFound)?, track.timescale, @@ -4093,7 +4133,7 @@ where /// Pumps moq-transport subgroup objects from a reader into a group producer: /// the id delta, the extension headers (carrying the timestamp), the size, the /// status for empty objects, and the streamed payload. -struct GroupIngest { +struct GroupIngest { has_extensions: bool, timescale: Option, version: Version, @@ -4103,11 +4143,11 @@ struct GroupIngest { gap: Option, /// The object being read is one the group started past, so it is read and dropped. dropping: bool, - phase: IngestPhase, + phase: IngestPhase, budget: frame::Budget, } -enum IngestPhase { +enum IngestPhase { /// Reading the object id delta. Stream end here ends the group. Delta, /// Reading the extension block's size. @@ -4119,14 +4159,14 @@ enum IngestPhase { /// Reading the status of an empty object. Status { timestamp: Option }, /// Streaming the object payload. - Payload { frame: frame::ProducerOwned }, + Payload { frame: crate::coding::HeldFrame }, /// Discarding a dropped object's payload: the bytes of it still to read. Skip { size: usize }, /// An explicit end-of-group or end-of-track status arrived. Finished(Ended), } -impl GroupIngest { +impl GroupIngest { fn new( subscriber: &Subscriber, group: &ietf::GroupHeader, @@ -4177,6 +4217,9 @@ where let joined = state.fetches.get(&header.request_id).copied(); let subscribe_id = joined.unwrap_or(header.request_id); let track = state.subscribes.get(&subscribe_id).ok_or(Error::NotFound)?; + // The head of the group the subscription joined, as urgent as its live edge. + let group = track.largest.map_or(0, |largest| largest.group); + stream.rank(track.rank(subscribe_id, group, 0)); // A fill is one of the subscription's own data streams, so PUBLISH_DONE counts // it. A joining FETCH is a request of its own. let counted = joined @@ -4523,7 +4566,7 @@ where info: frame::Info { size, timestamp }, fields, }; - let mut frame = producer.create_frame_owned(header, &self.frames)?; + let mut frame = crate::coding::HeldFrame::new(producer.create_frame_owned(header, &self.frames)?); if let Err(err) = std::future::poll_fn(|cx| stream.poll_read_frame(cx, &mut frame)).await { let _ = frame.abort(err.clone()); return Err(err); @@ -4561,7 +4604,7 @@ where let slot = kio::Producer::new(GroupFetch::Pending); let registered = GroupFetchEntry::new(&self.state, fetch_id, slot.clone()); - let mut stream = match Stream::open(&mut self.session.clone(), self.version).await { + let mut stream = match self.open_stream().await { Ok(stream) => stream, Err(err) => return request.reject(err), }; @@ -5238,7 +5281,7 @@ fn object_time( } } -impl GroupIngest { +impl GroupIngest { fn header(&mut self, group: &group::Producer, info: frame::Info) -> Result { let id = self.prior_object.expect("payload follows object ID"); group.mark_gap(self.prior_object.expect("properties follow Object ID"), self.gap)?; @@ -5257,7 +5300,7 @@ impl GroupIngest { /// `Ready(Ok(_))` once the stream FINs on an object boundary, or an explicit /// end-of-group or end-of-track status arrives. The caller finishes or aborts the /// group; an object cut short mid-payload was already aborted here with the reason. - fn poll( + fn poll( &mut self, reader: &mut Reader, group: &mut group::Producer, @@ -5309,7 +5352,9 @@ impl GroupIngest { let timestamp = object_time(self.timescale, *timestamp)?; let header = self.header(group, frame::Info { size, timestamp })?; let frame = group.create_frame_owned(header, &self.budget)?; - self.phase = IngestPhase::Payload { frame }; + self.phase = IngestPhase::Payload { + frame: crate::coding::HeldFrame::new(frame), + }; } IngestPhase::Status { timestamp } => { let status = ready!(reader.poll_varint(&mut cx))?; @@ -11037,6 +11082,45 @@ mod stitch_tests { assert!(matches!(*h.fill.read(), Fill::Done), "the head was claimed"); } + /// A tail parked on its head, holding the read-ahead it peeked, lets go the moment the + /// pressure monitor chooses it: the eviction itself returns the credit, which the head may + /// need to arrive, and stops the stream, without the tail being polled. The head still + /// lands, and the tail finishes from what it had already buffered. + #[moq_clock::test] + async fn an_evicted_tail_waiting_on_its_fill_lets_go() { + use crate::coding::StreamCodes as _; + + let h = Harness::new( + Fill::Serving(Some(Timescale::MICRO)), + vec![ + tail_stream(SEQUENCE, 2, &[b"tail-2"]), + fill_stream(SEQUENCE, &[b"head-0", b"head-1"]), + ], + ); + let pressure = h.subscriber.pressure.clone(); + let mut tail = h.stream().await; + tail.track(&pressure, Rank::Unidentified); + let evicted = vec![VERSION.encode_stream_code(&crate::StreamError::Evicted)]; + + let mut serve_tail = h.subscriber.clone(); + let mut parked = std::pin::pin!(serve_tail.recv_group(&mut tail)); + let waited = moq_clock::timeout(std::time::Duration::from_secs(1), parked.as_mut()).await; + assert!(waited.is_err(), "the tail did not wait for its head: {waited:?}"); + assert!(pressure.held() > 0, "the parked tail holds its read-ahead"); + + assert!(pressure.evict()); + assert_eq!(pressure.held(), 0, "the eviction let go of the parked tail's credit"); + assert_eq!(h.session.log.stops(), evicted, "the eviction stopped the parked tail"); + + let mut fill = h.stream().await; + h.subscriber.clone().recv_fill(&mut fill).await.expect("the head lands"); + moq_clock::timeout(std::time::Duration::from_secs(1), parked) + .await + .expect("the tail never took its head") + .expect("the tail finishes from its buffer"); + assert_eq!(h.session.log.stops(), evicted, "an evicted stream is stopped once"); + } + /// A tail parked on its head stops holding the end open, so the finished fill does until /// the head is claimed: otherwise the subscription looks settled between the fill /// finishing and the tail waking, and ends before the tail is read. diff --git a/rs/moq-net/src/lib.rs b/rs/moq-net/src/lib.rs index c7999766c73..79fef0cc377 100644 --- a/rs/moq-net/src/lib.rs +++ b/rs/moq-net/src/lib.rs @@ -100,6 +100,8 @@ mod lite; mod model; pub mod path; #[cfg(feature = "async")] +mod pressure; +#[cfg(feature = "async")] mod priority; #[cfg(feature = "async")] mod recv; diff --git a/rs/moq-net/src/lite/publisher.rs b/rs/moq-net/src/lite/publisher.rs index 292d49d18c8..e974289af65 100644 --- a/rs/moq-net/src/lite/publisher.rs +++ b/rs/moq-net/src/lite/publisher.rs @@ -173,6 +173,8 @@ pub(super) struct Publisher { children: kio::Tasks>, /// Aborts the session on a publication outside our grant (dialing side only). enforce: Option, + /// Tracks each accepted control stream's reads. + pub(super) pressure: crate::pressure::Of, } impl Publisher { @@ -208,6 +210,7 @@ impl Publisher { runtime: config.runtime, accept, children: kio::Tasks::new(), + pressure: Default::default(), } } } @@ -229,7 +232,8 @@ where let mut cx = waiter.context(); loop { match Stream::poll_accept(&mut self.accept, self.shared.version, &mut cx) { - Poll::Ready(Ok(stream)) => { + Poll::Ready(Ok(mut stream)) => { + stream.reader.track(&self.pressure, crate::pressure::Rank::Control); self.children.push(Control { shared: self.shared.clone(), runtime: self.runtime.clone(), diff --git a/rs/moq-net/src/lite/session.rs b/rs/moq-net/src/lite/session.rs index 756871b1ac5..fbc5782fb34 100644 --- a/rs/moq-net/src/lite/session.rs +++ b/rs/moq-net/src/lite/session.rs @@ -3,6 +3,7 @@ use crate::{ Error, Hop, SessionError, bandwidth, coding::{Reader, Stream, Writer}, lite::SessionInfo, + pressure::{Monitor, Pressure, Rank}, }; use std::task::{Context, Poll, ready}; @@ -202,7 +203,7 @@ where runtime, client, session, - setup_stream, + mut setup_stream, publish, subscribe, peer_hop, @@ -293,8 +294,15 @@ where false => None, }; + // Every stream the session reads holds its credit until a whole message arrives, and + // the monitor cancels one by rank whenever that fills the receive allowance. + let pressure = Pressure::default(); + if let Some(stream) = &mut setup_stream { + stream.reader.track(&pressure, Rank::Control); + } + let estimate = kio::Producer::default(); - let publisher = Publisher::new(PublisherConfig { + let mut publisher = Publisher::new(PublisherConfig { estimate: estimate.weak(), runtime: runtime.clone(), session: session.clone(), @@ -325,6 +333,8 @@ where auth: auth.clone(), }); subscriber.announces = announces; + subscriber.pressure = pressure.clone(); + publisher.pressure = pressure.clone(); let driver = Driver { auth: Present { @@ -337,7 +347,9 @@ where violation: Default::default(), started: false, _setup: setup_token, + pressure: pressure.clone(), }, + monitor: Monitor::new(session.clone(), pressure), local_close: Default::default(), setup: version .has_setup_stream() @@ -363,6 +375,8 @@ where pub(crate) struct Driver { /// Presenting our tokens, one AUTH stream each. auth: Present, + /// Cancels a stream whenever held receive credit fills the allowance. + monitor: Monitor>, pub(crate) local_close: std::sync::Arc, /// Advertising our capabilities, or `None` once sent (or on a version with no /// Setup Stream). @@ -428,6 +442,9 @@ where fn poll_protocol(&mut self, waiter: &kio::Waiter) -> Poll> { let mut cx = waiter.context(); + // Never resolves; it only cancels streams, whose own handlers report it. + let _ = self.monitor.poll(&mut cx); + // Presenting tokens ends the session only when the peer breaks the protocol. if let Poll::Ready(err) = self.auth.poll(waiter) { return Poll::Ready(Err(err)); @@ -481,6 +498,8 @@ struct Present { started: bool, /// The connection's own credential, held for the life of the session. _setup: Option, + /// Tracks each AUTH stream's reads. + pressure: crate::pressure::Of, } impl Present { @@ -500,6 +519,7 @@ impl Present { handle: self.handle.clone(), going_away: self.going_away.clone(), violation: self.violation.clone(), + pressure: self.pressure.clone(), id, state: PresentState::Open { token }, }); @@ -522,6 +542,7 @@ struct PresentToken { handle: crate::auth::Handle, going_away: crate::goaway::GoingAway, violation: kio::Shared>, + pressure: crate::pressure::Of, id: u64, state: PresentState, } @@ -580,6 +601,7 @@ impl PresentToken { Ok(stream) => stream, Err(err) => return Poll::Ready(err), }; + stream.reader.track(&self.pressure, Rank::Control); let res = stream .writer .buffer(&super::ControlType::Auth) diff --git a/rs/moq-net/src/lite/subscriber.rs b/rs/moq-net/src/lite/subscriber.rs index 5f59f20715d..64a6c34d36d 100644 --- a/rs/moq-net/src/lite/subscriber.rs +++ b/rs/moq-net/src/lite/subscriber.rs @@ -10,6 +10,7 @@ use crate::{ AsPath, Error, Path, PathOwned, Timescale, Timestamp, bandwidth, coding::{Reader, Stream}, lite, + pressure::Rank, track::{Position, Subscription}, }; @@ -88,6 +89,8 @@ pub(super) struct Subscriber { frames: frame::Budget, /// Broadcasts the peer may have announced at once (`session::Limits::announces`). pub(super) announces: crate::session::Slots, + /// Every stream this half reads, cancelled by rank when held receive credit fills. + pub(super) pressure: crate::pressure::Of, } #[derive(Clone)] @@ -98,6 +101,27 @@ struct TrackEntry { timescale: Option, /// The groups received so far, so the subscription's end can wait for the ones owed. tail: kio::Producer, + /// Ranks the subscription's group streams when held receive credit fills. + rank: GroupRank, +} + +/// What ranks a subscription's group streams for receive-side recovery. +#[derive(Clone, Default)] +struct GroupRank { + broadcast: Arc, + priority: u8, + order: track::Order, +} + +impl GroupRank { + fn of(&self, subscribe: u64, group: &lite::Group) -> Rank { + Rank::Media { + broadcast: self.broadcast.clone(), + priority: crate::priority::Priority::new(self.priority, subscribe, group.sequence) + .with_order(self.order) + .with_start(group.frame_start), + } + } } impl Subscriber { @@ -125,6 +149,7 @@ impl Subscriber { auth: config.auth, frames: Default::default(), announces: Default::default(), + pressure: Default::default(), } } @@ -681,7 +706,8 @@ impl UniAccept { fn new(subscriber: Subscriber, early: Vec>) -> Self { let accept = subscriber.session.clone(); let mut children = kio::Tasks::new(); - for reader in early { + for mut reader in early { + reader.track(&subscriber.pressure, Rank::Unidentified); children.push(UniServe { subscriber: subscriber.clone(), state: UniState::Start { reader }, @@ -701,11 +727,11 @@ impl UniAccept { loop { match self.accept.poll_accept_uni(&mut cx) { Poll::Ready(Ok(stream)) => { + let mut reader = Reader::new(stream, self.subscriber.version); + reader.track(&self.subscriber.pressure, Rank::Unidentified); self.children.push(UniServe { subscriber: self.subscriber.clone(), - state: UniState::Start { - reader: Reader::new(stream, self.subscriber.version), - }, + state: UniState::Start { reader }, }); } Poll::Ready(Err(err)) => return Poll::Ready(Err(crate::lite::error::from_transport(err))), @@ -786,7 +812,10 @@ impl UniServe { }; self.state = match kind { lite::DataType::Group => UniState::Group(GroupRecv::new(self.subscriber.clone(), reader)), - lite::DataType::Setup => UniState::Setup { reader }, + lite::DataType::Setup => { + reader.rank(Rank::Control); + UniState::Setup { reader } + } }; } UniState::Setup { reader } => { @@ -829,13 +858,13 @@ impl UniServe { struct GroupRecv { subscriber: Subscriber, reader: Reader, - state: GroupRecvState, + state: GroupRecvState, } // A state machine's enum is its storage: one transient instance per stream, so the // big variant is the working state, not padding held in bulk. #[allow(clippy::large_enum_variant)] -enum GroupRecvState { +enum GroupRecvState { /// Reading the GROUP header. Header, /// Filling the group, bailing if the track or group dies first. @@ -843,7 +872,7 @@ enum GroupRecvState { /// Guarded: dropping this machine mid-group is a cancellation, not a clean end. group: crate::recv::Group, track: track::Producer, - ingest: FrameIngest, + ingest: FrameIngest, /// Whether readers see the group yet; see [`track::Producer::receive_group`]. shown: bool, _reading: Reading, @@ -870,6 +899,7 @@ impl GroupRecv { let (group, track, timescale, reading) = { let mut subs = self.subscriber.subscribes.lock(); let entry = subs.get_mut(&hdr.subscribe).ok_or(Error::Cancel)?; + self.reader.rank(entry.rank.of(hdr.subscribe, &hdr)); let group_info = group::Info { sequence: hdr.sequence }; // Stats (groups/frames/bytes) are counted in the model as the group @@ -997,20 +1027,20 @@ impl GroupRecv { } } -/// Pumps bare FRAME messages from a reader into a group producer: the wire +/// Pumps bare FRAME messages from a reader of `R` into a group producer: the wire /// format shared by GROUP streams and FETCH responses. -struct FrameIngest { +struct FrameIngest { /// `Some` decodes the lite-05 zigzag-delta timestamp prefix; `None` leaves frames /// untimed (pre-lite-05). timescale: Option, /// Previous frame's raw timestamp value (in `timescale` units), for the /// zigzag-delta decode. The first frame's delta is absolute (prev = 0). prev_ts: u64, - phase: IngestPhase, + phase: IngestPhase, budget: frame::Budget, } -enum IngestPhase { +enum IngestPhase { /// Reading the timestamp delta (skipped without a timescale). Stream end here /// means the group has no more frames. Timing, @@ -1018,10 +1048,10 @@ enum IngestPhase { /// where there is no timing prefix to act as the sentinel). Size { timestamp: Option }, /// Streaming the frame payload. - Payload { frame: frame::ProducerOwned }, + Payload { frame: crate::coding::HeldFrame }, } -impl FrameIngest { +impl FrameIngest { fn new(subscriber: &Subscriber, timescale: Option) -> Self { Self { timescale, @@ -1034,7 +1064,7 @@ impl FrameIngest { /// `Ready(Ok(()))` once the stream FINs on a frame boundary. The caller /// finishes or aborts the group; a frame cut short mid-payload was already /// aborted here with the reason. - fn poll( + fn poll( &mut self, reader: &mut Reader, group: &mut group::Producer, @@ -1077,7 +1107,9 @@ impl FrameIngest { }, &self.budget, )?; - self.phase = IngestPhase::Payload { frame }; + self.phase = IngestPhase::Payload { + frame: crate::coding::HeldFrame::new(frame), + }; } IngestPhase::Payload { frame } => { let failed = ready!(reader.poll_read_frame(&mut cx, frame)).err(); @@ -1254,6 +1286,7 @@ impl ProbeStream { return Poll::Ready(Ok(())); } let mut stream = ready!(Stream::poll_open(&mut self.session, self.subscriber.version, &mut cx))?; + stream.reader.track(&self.subscriber.pressure, Rank::Control); stream.writer.buffer(&lite::ControlType::Probe)?; self.state = ProbeState::Send { stream }; } @@ -1348,6 +1381,7 @@ impl AnnouncePrefix { self.subscriber.version, &mut cx ))?; + stream.reader.track(&self.subscriber.pressure, Rank::Control); stream.writer.buffer(&lite::ControlType::Announce)?; // Lite04/05: ask the peer to filter out announces that already passed @@ -2143,6 +2177,7 @@ mod tests { producer: track, timescale: Some(Timescale::MILLI), tail: Default::default(), + rank: Default::default(), }, ); let mut accept = UniAccept::new(subscriber, early); @@ -2199,6 +2234,7 @@ mod tests { producer: track.clone(), timescale: Some(Timescale::MILLI), tail: Default::default(), + rank: Default::default(), }, ); @@ -2267,6 +2303,7 @@ mod tests { producer: track.clone(), timescale: Some(Timescale::MILLI), tail: tail.clone(), + rank: Default::default(), }, ); @@ -2485,6 +2522,7 @@ mod tests { producer, timescale: Some(Timescale::MILLI), tail: Default::default(), + rank: Default::default(), }, ); @@ -4227,6 +4265,62 @@ mod tests { let _a = slots.acquire().expect("the failed stream kept a slot"); let _b = slots.acquire().expect("the failed stream kept a slot"); } + + /// A prefix stream waiting on the peer's SETUP for its link cost holds what it read + /// past ANNOUNCE_OK. The pressure monitor choosing it returns that credit and stops the + /// stream itself, without the prefix being polled. + #[moq_clock::test] + async fn an_evicted_prefix_waiting_on_its_link_cost_lets_go() { + use crate::coding::StreamCodes as _; + const VERSION: Version = Version::Lite06; + let mut script = Vec::new(); + lite::AnnounceOk { + origin: crate::Hop::new(9).unwrap(), + active: 1, + } + .encode(&mut crate::coding::Encoder::new(&mut script, VERSION.into()), VERSION) + .unwrap(); + lite::AnnounceBroadcast::Active { + epoch: None, + suffix: lite::PathRef::literal(Path::new("a")), + hops: lite::HopsRef::literal(crate::Hops::new()), + cost: crate::origin::Cost::default(), + } + .encode(&mut crate::coding::Encoder::new(&mut script, VERSION.into()), VERSION) + .unwrap(); + + let origin = origin::Config::new(crate::Hop::new(1).unwrap()).produce(); + let session = crate::lite::test_transport::ScriptedSession::new(script); + let log = session.log.clone(); + let subscriber = Subscriber::new(SubscriberConfig { + runtime: crate::time::sim(), + session, + origin, + recv_bandwidth: None, + version: VERSION, + // No configured cost, and a SETUP that never arrives. + peer_setup: crate::lite::PeerSetup::new(false, crate::transport::Binding::InBand), + cost: None, + peer_hop: None, + going_away: Default::default(), + auth: crate::auth::Handle::new(false), + }); + let pressure = subscriber.pressure.clone(); + let mut prefix = AnnouncePrefix::new(subscriber, Path::new("").to_owned()); + + let mut run = std::pin::pin!(kio::wait(|waiter| prefix.poll(waiter))); + let waited = moq_clock::timeout(std::time::Duration::from_secs(1), run.as_mut()).await; + assert!(waited.is_err(), "the prefix did not wait for its link cost: {waited:?}"); + assert!(pressure.held() > 0, "the parked prefix holds its read-ahead"); + + assert!(pressure.evict()); + assert_eq!(pressure.held(), 0, "the eviction let go of the parked prefix's credit"); + assert_eq!( + log.stops(), + vec![VERSION.encode_stream_code(&crate::StreamError::Evicted)], + "the eviction stopped the parked prefix" + ); + } } /// The four wire fields a subscription's half-open range encodes to. @@ -4314,6 +4408,8 @@ impl SubStream { } } +// One per subscribed track, and the big variant is the working state. +#[allow(clippy::large_enum_variant)] enum Sub { None, Active(SubStream), @@ -4747,6 +4843,11 @@ impl TrackServe { let start_moved = active.start != subscription.start; active.priority = subscription.priority; active.order = subscription.order; + // Streams already open keep the rank they were read at. + if let Some(entry) = self.subscriber.subscribes.lock().get_mut(&active.id) { + entry.rank.priority = subscription.priority; + entry.rank.order = subscription.order; + } active.max_delay = subscription.max_delay; if let Ok(mut tail) = active.tail.write() { tail.set_grace(tail::grace(subscription.max_delay)); @@ -4833,6 +4934,11 @@ impl TrackServe { producer: producer.clone(), timescale, tail: tail.clone(), + rank: GroupRank { + broadcast: Arc::new(self.path.clone()), + priority: subscription.priority, + order: subscription.order, + }, }, ); @@ -4975,6 +5081,7 @@ impl Establish { self.serve.subscriber.version, &mut cx ))?; + stream.reader.track(&self.serve.subscriber.pressure, Rank::Control); let bounds = WireBounds::new(self.subscription.start, self.subscription.end); let msg = lite::Subscribe { @@ -5273,6 +5380,7 @@ impl TrackInfoFetch { match &mut self.state { TrackInfoState::Open => { let mut stream = ready!(Stream::poll_open(&mut self.session, serve.subscriber.version, &mut cx))?; + stream.reader.track(&serve.subscriber.pressure, Rank::Control); stream.writer.buffer(&lite::ControlType::Track)?; stream.writer.buffer(&lite::Track { epoch: serve.epoch.clone(), @@ -5869,7 +5977,7 @@ enum FetchRunState { // picked the group up from the cache. joined: kio::Producer, // Boxed so the other states stay small; one allocation per fetch. - ingest: Box, + ingest: Box>, }, Done, } @@ -5992,6 +6100,14 @@ impl kio::Task for FetchServeRun { tracing::info!(broadcast = %self.serve.subscriber.log_path(&self.serve.path), track = %self.serve.name, group = self.group, "fetch started"); let request = request.take().expect("request pending"); + // The response is the group's frames, ranked like a group stream of the track. + stream.reader.track( + &self.serve.subscriber.pressure, + Rank::Media { + broadcast: Arc::new(self.serve.path.clone()), + priority: crate::priority::Priority::new(request.priority(), u64::MAX, self.group), + }, + ); // A peer that predates lite-06 addresses whole groups only, so ask for // the whole group and number the response from 0. The wider group still diff --git a/rs/moq-net/src/lite/test_transport.rs b/rs/moq-net/src/lite/test_transport.rs index 967f72e621a..ffb4c880312 100644 --- a/rs/moq-net/src/lite/test_transport.rs +++ b/rs/moq-net/src/lite/test_transport.rs @@ -315,8 +315,11 @@ pub struct PendingRecv; impl poll::RecvStream for PendingRecv { type Error = SinkError; + type Credit = crate::transport::Unheld; + type Stopper = (); + fn stopper(&self) {} - fn poll_read(&mut self, _cx: &mut Context<'_>, _dst: &mut [u8]) -> Poll, Self::Error>> { + fn poll_read(&mut self, _cx: &mut Context<'_>, _dst: &mut [u8]) -> Poll, Self::Error>> { Poll::Pending } @@ -362,8 +365,11 @@ pub struct DeadRecv(ResetError); impl poll::RecvStream for DeadRecv { type Error = ResetError; + type Credit = crate::transport::Unheld; + type Stopper = (); + fn stopper(&self) {} - fn poll_read(&mut self, _cx: &mut Context<'_>, _dst: &mut [u8]) -> Poll, Self::Error>> { + fn poll_read(&mut self, _cx: &mut Context<'_>, _dst: &mut [u8]) -> Poll, Self::Error>> { Poll::Ready(Err(self.0)) } @@ -491,6 +497,10 @@ impl poll::Session for DeadStreamSession { Poll::Pending } + fn poll_credit_exhausted(&mut self, _cx: &mut std::task::Context<'_>) -> std::task::Poll<()> { + std::task::Poll::Pending + } + fn stats(&self) -> impl crate::transport::Stats { SinkStats::default() } @@ -705,6 +715,10 @@ impl poll::Session for SinkSession { Poll::Pending } + fn poll_credit_exhausted(&mut self, _cx: &mut std::task::Context<'_>) -> std::task::Poll<()> { + std::task::Poll::Pending + } + fn stats(&self) -> impl crate::transport::Stats { *self.stats.lock().unwrap() } @@ -761,10 +775,26 @@ pub enum Close { Reset, } +/// Records a [`ScriptedRecv`]'s stops from outside it, as its own `stop` does. +#[derive(Clone)] +pub struct ScriptedStopper(Log); + +impl crate::transport::Stopper for ScriptedStopper { + fn stop(&self, code: u32) { + self.0.stops.lock().unwrap().push(code); + } +} + impl poll::RecvStream for ScriptedRecv { type Error = SinkError; + type Credit = crate::transport::Unheld; + type Stopper = ScriptedStopper; + + fn stopper(&self) -> ScriptedStopper { + ScriptedStopper(self.log.clone()) + } - fn poll_read(&mut self, _cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + fn poll_read(&mut self, _cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { let take = { let mut script = self.script.lock().unwrap(); if script.is_empty() { @@ -783,7 +813,7 @@ impl poll::RecvStream for ScriptedRecv { Some(Close::Reset) => Poll::Ready(Err(SinkError)), None => Poll::Pending, }, - take => Poll::Ready(Ok(Some(take))), + take => Poll::Ready(Ok(Some(crate::transport::Unheld(take as u64)))), } } @@ -1036,6 +1066,10 @@ impl poll::Session for ScriptedSession { Poll::Pending } + fn poll_credit_exhausted(&mut self, _cx: &mut std::task::Context<'_>) -> std::task::Poll<()> { + std::task::Poll::Pending + } + fn stats(&self) -> impl crate::transport::Stats { self.stats } diff --git a/rs/moq-net/src/pressure.rs b/rs/moq-net/src/pressure.rs new file mode 100644 index 00000000000..fdb6dc8e3b1 --- /dev/null +++ b/rs/moq-net/src/pressure.rs @@ -0,0 +1,921 @@ +//! Receive-side recovery for when held connection credit fills the allowance. +//! +//! A [`Reader`](crate::coding::Reader) tracked here keeps its buffered bytes and their +//! connection credit in a [`Hold`], and lends the credit for a partial unit's bytes to +//! whatever owns them as a [`Lent`], so incomplete messages are bounded by the +//! transport's receive allowance. When the transport reports +//! that the allowance is spent ([`poll_credit_exhausted`](crate::transport::poll::Session::poll_credit_exhausted)), +//! nothing can arrive to complete any unit, so the [`Monitor`] evicts the least useful +//! holder: the eviction itself drops the buffer and the credit and stops the stream with +//! STOP_SENDING, whatever the stream's handler is waiting on, and the handler's next read +//! fails with [`Error::Evicted`](crate::Error::Evicted). +//! +//! Victims, first to last: +//! - Media, from the broadcast holding the most credit, lowest priority first. Within a +//! broadcast the subscription's priority and group order rank the streams, as on the +//! send side; raw priorities of unrelated broadcasts are never compared. +//! - Unidentified streams, whose header has yet to name a subscription, so a header +//! gets the chance to identify itself before it is cancelled. +//! - Control streams. Cancelling one may end the session under its protocol, which the +//! stream's handler reports like any other error. +//! +//! Ties go to the larger holder, then to the newest stream. There is no protected class: +//! the allowance fits the largest supported unit, so cancelling every other holder +//! always leaves room for one to complete. + +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, MutexGuard, PoisonError, Weak, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, + task::{Context, Waker}, +}; + +use bytes::{Buf, Bytes, BytesMut}; + +use crate::{ + PathOwned, + priority::Priority, + transport::{Credit, MaybeSend, MaybeSync, Stopper, poll::RecvStream}, +}; + +/// How useful a tracked stream is, which orders the victims. +#[derive(Clone, Debug)] +pub(crate) enum Rank { + /// A stream whose header has not identified it yet. + Unidentified, + /// A media stream of one broadcast, ranked within it. + Media { + /// The broadcast, the domain inside which priorities compare. Shared, since every + /// group stream of a subscription names the same one. + broadcast: Arc, + /// The stream's rank within its broadcast. + priority: Priority, + }, + /// A control stream. + Control, +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} + +/// A tracked stream's buffer and credit, and what evicting it takes, shared by its reader +/// and the registry. +pub(crate) struct Hold { + /// Bytes of credit the slot covers, read by the registry without the lock. + held: AtomicU64, + /// Bytes of credit lent out with partial units, until each [`Lent`] drops. + lent: AtomicU64, + /// The monitor chose this stream: its buffer and credit are gone and the stream is + /// stopped. + evicted: AtomicBool, + /// The reader's task, woken on eviction. + waker: Mutex>, + rank: Mutex, + /// Registration order, the last tiebreak: the newest stream goes first. + seq: u64, + slot: Mutex>, + stopper: R::Stopper, + /// The STOP_SENDING code for an eviction, in the stream's negotiated registry. + code: u32, +} + +/// A tracked reader's bytes and their credit, under one lock. Credit travels with the +/// bytes it covers and returns only once they are freed: credit returned for bytes still +/// allocated would let a peer refill the window while readers keep the memory. +struct Slot { + /// Declared before `credit`: fields drop in order, so a dropped slot frees the bytes + /// before their credit returns. + buffer: BytesMut, + /// Leading bytes of `buffer` read before tracking began, which hold no credit. + free: u64, + /// The credit for `buffer` past `free`, which an eviction frees with it. + credit: Option, +} + +impl Slot { + fn held(&self) -> u64 { + self.credit.as_ref().map_or(0, C::len) + } + + /// Split off the credit for the first `bytes` of the buffer, which are leaving it. + fn split(&mut self, bytes: u64) -> Option { + let free = self.free.min(bytes); + self.free -= free; + split(&mut self.credit, bytes - free) + } + + /// Empty the buffer, handing back its bytes and their credit to drop in that order. + fn empty(&mut self) -> (BytesMut, Option) { + self.free = 0; + (std::mem::take(&mut self.buffer), self.credit.take()) + } +} + +/// Add `credit` to `into`. +fn merge(into: &mut Option, credit: C) { + match into { + Some(held) => held.merge(credit), + None => *into = Some(credit), + } +} + +/// Split the credit for `bytes` off `from`, if it holds any. +fn split(from: &mut Option, bytes: u64) -> Option { + from.as_mut().map(|credit| credit.split(bytes)) +} + +/// A [`Hold`]'s slot, locked by its reader. +pub(crate) struct Locked<'a, R: RecvStream> { + hold: &'a Arc>, + slot: MutexGuard<'a, Slot>, +} + +impl Locked<'_, R> { + /// The reader's buffered bytes. + pub fn buffer(&self) -> &BytesMut { + &self.slot.buffer + } + + /// The reader's buffered bytes, to read into or consume. + pub fn buffer_mut(&mut self) -> &mut BytesMut { + &mut self.slot.buffer + } + + /// Whether the monitor chose this stream. An eviction sets it under the lock, so it + /// holds for as long as this does: a `false` means bytes read now keep their credit + /// until an eviction frees both. + pub fn is_evicted(&self) -> bool { + self.hold.is_evicted() + } + + /// Take over bytes buffered before tracking began, whose credit already returned. + pub fn adopt(&mut self, buffer: BytesMut) { + self.slot.free = buffer.len() as u64; + self.slot.buffer = buffer; + } + + /// Keep the credit for bytes just read into the buffer. An eviction sets its flag under + /// this lock, so a reader that checked it first cannot be reading into an evicted slot. + pub fn take(&mut self, credit: R::Credit) { + debug_assert!(!self.hold.is_evicted(), "read into an evicted buffer"); + merge(&mut self.slot.credit, credit); + self.store(); + } + + /// Discard the first `bytes` of the buffer, a unit consumed whole, then return their + /// credit. + pub fn consume(&mut self, bytes: usize) { + self.slot.buffer.advance(bytes); + let credit = self.slot.split(bytes as u64); + self.store(); + drop(credit); + } + + /// Hand off the first `bytes` of the buffer as a complete unit, whose credit returns as + /// it leaves. + pub fn hand_off(&mut self, bytes: usize) -> Bytes { + let unit = self.slot.buffer.split_to(bytes).freeze(); + let credit = self.slot.split(bytes as u64); + self.store(); + drop(credit); + unit + } + + /// Move the first `bytes` of the buffer into a partial unit, lending their credit to + /// whatever owns them next. + pub fn lend_buffered(&mut self, bytes: usize) -> (Bytes, Option>) { + let unit = self.slot.buffer.split_to(bytes).freeze(); + let lent = self.slot.split(bytes as u64).map(|credit| self.hold.lend(credit)); + self.store(); + (unit, lent) + } + + /// Lend the credit for bytes read straight into a partial unit to whatever owns them. + pub fn lend(&self, credit: R::Credit) -> Lent { + self.hold.lend(credit) + } + + fn store(&self) { + self.hold.held.store(self.slot.held(), Ordering::Relaxed); + } +} + +/// Credit lent out of a [`Hold`] with the bytes of a partial unit: still counted toward +/// the hold, and returned when dropped. Whatever owns those bytes owns this beside them, so +/// the credit returns only once they are given up. +pub(crate) struct Lent { + credit: Option, + hold: Arc>, +} + +impl Lent { + /// Take over `other`'s credit, lent by the same hold for the same unit. + pub fn merge(&mut self, mut other: Self) { + if let Some(credit) = other.credit.take() { + merge(&mut self.credit, credit); + } + } +} + +impl Drop for Lent { + fn drop(&mut self) { + let Some(credit) = self.credit.take() else { + return; + }; + let len = credit.len(); + // Returned before it stops counting, so a fill reported in between still sees an + // eviction's release in flight rather than cancelling another stream. + drop(credit); + self.hold.lent.fetch_sub(len, Ordering::Relaxed); + } +} + +impl Hold { + /// Lock the buffer and credit, which an eviction drops together. + pub fn lock(self: &Arc) -> Locked<'_, R> { + Locked { + hold: self, + slot: lock(&self.slot), + } + } + + fn lend(self: &Arc, credit: R::Credit) -> Lent { + self.lent.fetch_add(credit.len(), Ordering::Relaxed); + Lent { + credit: Some(credit), + hold: self.clone(), + } + } + + /// Drop the buffer and return its credit: the reader gave up on the stream. Credit lent + /// to a partial unit returns when its owner gives that up. + pub fn clear(&self) { + let mut slot = lock(&self.slot); + let (buffer, credit) = slot.empty(); + self.held.store(0, Ordering::Relaxed); + drop(slot); + // Freed before the credit returns, as on eviction. + drop(buffer); + drop(credit); + } + + /// Re-rank the stream, once its header identifies it. + pub fn rank(&self, rank: Rank) { + *lock(&self.rank) = rank; + } + + /// Whether the monitor chose this stream, without registering to hear when it does. + pub fn is_evicted(&self) -> bool { + self.evicted.load(Ordering::Acquire) + } + + /// Whether the monitor chose this stream, registering `cx` to hear when it does. + pub fn poll_evicted(&self, cx: &Context<'_>) -> bool { + if self.is_evicted() { + return true; + } + let mut waker = lock(&self.waker); + match waker.as_mut() { + Some(waker) => waker.clone_from(cx.waker()), + None => *waker = Some(cx.waker().clone()), + } + drop(waker); + // Checked again: an eviction between the load and the store took no waker. + self.is_evicted() + } +} + +/// What the registry needs from a [`Hold`], whatever its stream type. +pub(crate) trait Holder { + /// Bytes of credit held. + fn held(&self) -> u64; + /// Whether it was evicted already. + fn is_evicted(&self) -> bool; + /// How useful the stream is. + fn rank(&self) -> Rank; + /// Registration order. + fn seq(&self) -> u64; + /// Drop the buffer and credit, stop the stream, and wake its reader to fail. + fn evict(&self); +} + +impl Holder for Hold { + fn held(&self) -> u64 { + self.held.load(Ordering::Relaxed) + self.lent.load(Ordering::Relaxed) + } + + fn is_evicted(&self) -> bool { + Hold::is_evicted(self) + } + + fn rank(&self) -> Rank { + lock(&self.rank).clone() + } + + fn seq(&self) -> u64 { + self.seq + } + + fn evict(&self) { + let mut slot = lock(&self.slot); + // Set under the lock, so a reader holding it reads and keeps that credit as one + // step, and one locking after sees the flag and reads nothing more in. + self.evicted.store(true, Ordering::Release); + let (buffer, credit) = slot.empty(); + // Credit lent to a partial unit stays counted until the owner of its bytes, woken + // below through the reader, gives the unit up. + self.held.store(0, Ordering::Relaxed); + drop(slot); + let freed = credit.as_ref().map_or(0, Credit::len); + tracing::debug!(rank = ?Holder::rank(self), freed, "receive credit exhausted, cancelling a stream"); + // Stopped first, so the window the credit frees does not invite more of this stream. + self.stopper.stop(self.code); + // The bytes go before their credit: returned credit lets the peer send again at + // once, and the memory it bounds must already be free. + drop(buffer); + drop(credit); + let waker = lock(&self.waker).take(); + if let Some(waker) = waker { + waker.wake(); + } + } +} + +/// The registry entry of a session whose tracked readers differ in stream type: the +/// moq-transport path, which reads both raw streams and its control adapter's. +#[cfg(not(target_family = "wasm"))] +pub(crate) type Erased = dyn Holder + Send + Sync; +#[cfg(target_family = "wasm")] +pub(crate) type Erased = dyn Holder; + +/// The registry of a session whose tracked readers all read its own streams: moq-lite's. +/// +/// An alias rather than the spelled-out type, so a `derive` on a struct holding it bounds +/// only `S`, not its stream type. +pub(crate) type Of = Pressure::RecvStream>>; + +/// Turns a reader's [`Hold`] into the registry's entry type. +pub(crate) trait Erase { + fn erase(self: Arc) -> Arc; +} + +impl Erase> for Hold { + fn erase(self: Arc) -> Arc { + self + } +} + +impl Erase for Hold +where + Self: MaybeSend + MaybeSync, +{ + fn erase(self: Arc) -> Arc { + self + } +} + +/// One session's tracked receivers, each a [`Hold`] seen as `H`. +pub(crate) struct Pressure { + state: kio::Lock>, +} + +impl Clone for Pressure { + fn clone(&self) -> Self { + Self { + state: self.state.clone(), + } + } +} + +impl Default for Pressure { + fn default() -> Self { + Self { + state: kio::Lock::new(State { + entries: Vec::new(), + next: 0, + prune: PRUNE, + }), + } + } +} + +/// Entries kept before the first sweep of dropped readers. +const PRUNE: usize = 64; + +struct State { + /// Weak, so a reader unregisters by being dropped, and swept on growth. + entries: Vec>, + next: u64, + /// Sweep dropped readers once `entries` reaches this. + prune: usize, +} + +/// A live holder, as of the choice. +struct Candidate { + hold: Arc, + rank: Rank, + held: u64, + seq: u64, +} + +impl Pressure { + /// Track a new receiver at `rank`, which `stopper` stops with `code` on eviction. + pub fn track(&self, rank: Rank, stopper: R::Stopper, code: u32) -> Arc> + where + Hold: Erase, + { + let mut state = self.state.lock(); + let hold = Arc::new(Hold { + held: AtomicU64::new(0), + lent: AtomicU64::new(0), + evicted: AtomicBool::new(false), + waker: Mutex::new(None), + rank: Mutex::new(rank), + seq: state.next, + slot: Mutex::new(Slot { + buffer: BytesMut::new(), + free: 0, + credit: None, + }), + stopper, + code, + }); + state.next += 1; + if state.entries.len() >= state.prune { + state.entries.retain(|entry| entry.strong_count() > 0); + state.prune = PRUNE.max(2 * state.entries.len()); + } + state.entries.push(Arc::downgrade(&hold.clone().erase())); + hold + } + + /// Evict the least useful holder of credit, returning whether there was one. + /// + /// An evicted stream still holding a partial unit is one: its reader, woken by the + /// eviction, gives the unit up and frees what this fill needs, so cancelling another + /// stream meanwhile would cost it for nothing. + pub fn evict(&self) -> bool { + let victim = { + let state = self.state.lock(); + let releasing = state + .entries + .iter() + .filter_map(Weak::upgrade) + .any(|hold| hold.is_evicted() && hold.held() > 0); + if releasing { + return true; + } + victim(&state.entries) + }; + // Evicted without the registry lock: dropping credit and stopping call into the + // transport, and readers track and drop meanwhile. + match victim { + Some(hold) => { + hold.evict(); + true + } + None => false, + } + } + + /// The bytes every tracked receiver holds, for tests. + #[cfg(test)] + pub fn held(&self) -> u64 { + let state = self.state.lock(); + state + .entries + .iter() + .filter_map(Weak::upgrade) + .map(|hold| hold.held()) + .sum() + } +} + +/// Pick the victim: see the module docs for the order. +fn victim(entries: &[Weak]) -> Option> { + let holding: Vec<_> = entries + .iter() + .filter_map(Weak::upgrade) + .filter_map(|hold| { + let held = hold.held(); + (held > 0 && !hold.is_evicted()).then(|| Candidate { + rank: hold.rank(), + seq: hold.seq(), + held, + hold, + }) + }) + .collect(); + + // Media: the broadcast holding the most, then its least useful stream. Keyed by path + // rather than by `Arc`, since separate subscriptions to one broadcast share no `Arc`. + // A run of streams sharing one `Arc` sums before it is hashed. + let mut domains: HashMap<&PathOwned, u64> = HashMap::new(); + let mut run: Option<(&Arc, u64)> = None; + for entry in &holding { + if let Rank::Media { broadcast, .. } = &entry.rank { + match &mut run { + Some((path, total)) if Arc::ptr_eq(path, broadcast) => *total += entry.held, + _ => { + if let Some((path, total)) = run.replace((broadcast, entry.held)) { + *domains.entry(path.as_ref()).or_default() += total; + } + } + } + } + } + if let Some((path, total)) = run { + *domains.entry(path.as_ref()).or_default() += total; + } + // The larger holder, then the lower path, so the choice never depends on order. + let domain = domains + .into_iter() + .max_by(|(a, a_held), (b, b_held)| a_held.cmp(b_held).then_with(|| b.as_str().cmp(a.as_str()))); + let chosen = match domain { + Some((domain, _)) => holding + .iter() + .filter_map(|entry| match &entry.rank { + Rank::Media { broadcast, priority } if broadcast.as_ref() == domain => Some((priority, entry)), + _ => None, + }) + // `Priority` sorts the most urgent first, so the greatest is the least useful. + .max_by(|(a, a_entry), (b, b_entry)| { + a.cmp(b) + .then_with(|| a_entry.held.cmp(&b_entry.held)) + .then_with(|| a_entry.seq.cmp(&b_entry.seq)) + }) + .map(|(_, entry)| entry), + None => { + let largest = |control: bool| { + holding + .iter() + .filter(|entry| matches!(entry.rank, Rank::Control) == control) + .max_by(|a, b| a.held.cmp(&b.held).then_with(|| a.seq.cmp(&b.seq))) + }; + largest(false).or_else(|| largest(true)) + } + }; + chosen.map(|entry| entry.hold.clone()) +} + +/// Cancels a stream each time the session's receive allowance fills. +pub(crate) struct Monitor { + session: S, + pressure: Pressure, +} + +impl Monitor { + /// Watch `session`, cancelling streams tracked by `pressure`. + pub fn new(session: S, pressure: Pressure) -> Self { + Self { session, pressure } + } + + /// Run until dropped. Each fill cancels one stream: the next fill can only come after + /// it let go of its credit and the peer spent what that freed. + pub fn poll(&mut self, cx: &mut Context<'_>) -> std::task::Poll<()> { + while self.session.poll_credit_exhausted(cx).is_ready() { + if !self.pressure.evict() { + // Untracked readers or transport read-ahead hold it all; whoever holds it + // lets go on their own, and a later fill is reported again. + tracing::warn!("receive credit exhausted with no tracked holder to cancel"); + } + } + std::task::Poll::Pending + } +} + +#[cfg(test)] +mod tests { + use std::task::{Poll, Waker}; + + use super::*; + use crate::{Path, track::Order, transport::Unheld}; + + /// Logs the codes its streams were stopped with. + #[derive(Clone, Default)] + struct Stops(Arc>>); + + impl Stopper for Stops { + fn stop(&self, code: u32) { + lock(&self.0).push(code); + } + } + + /// A stream that only exists to name its credit and stopper. + struct Probe; + + impl RecvStream for Probe { + type Error = crate::lite::test_transport::SinkError; + type Credit = Unheld; + type Stopper = Stops; + + fn poll_read(&mut self, _: &mut Context<'_>, _: &mut [u8]) -> Poll, Self::Error>> { + Poll::Pending + } + + fn stop(&mut self, _: u32) {} + + fn stopper(&self) -> Stops { + Stops::default() + } + + fn poll_closed(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Pending + } + } + + type Probes = Pressure>; + + /// Credit that records, as it returns, whether the bytes it covers were freed first: + /// it shares their allocation, so it is the last owner once the buffer is gone. + struct Watched { + len: u64, + bytes: Option, + freed_first: Arc, + } + + impl Credit for Watched { + fn len(&self) -> u64 { + self.len + } + + fn split(&mut self, bytes: u64) -> Self { + self.len -= bytes; + Self { + len: bytes, + bytes: None, + freed_first: self.freed_first.clone(), + } + } + + fn merge(&mut self, other: Self) { + self.len += other.len; + } + } + + impl Drop for Watched { + fn drop(&mut self) { + if let Some(bytes) = &self.bytes { + self.freed_first.store(bytes.is_unique(), Ordering::Relaxed); + } + } + } + + /// A stream whose credit is [`Watched`]. + struct Watch; + + impl RecvStream for Watch { + type Error = crate::lite::test_transport::SinkError; + type Credit = Watched; + type Stopper = Stops; + + fn poll_read(&mut self, _: &mut Context<'_>, _: &mut [u8]) -> Poll, Self::Error>> { + Poll::Pending + } + + fn stop(&mut self, _: u32) {} + + fn stopper(&self) -> Stops { + Stops::default() + } + + fn poll_closed(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Pending + } + } + + /// An evicted stream still holding a partial unit answers the next fill: its reader + /// frees that credit once it wakes, and another stream goes only if the window is still + /// full after. + #[test] + fn a_release_in_flight_spares_the_next_stream() { + let pressure = Probes::default(); + let partial = pressure.track::(Rank::Unidentified, Stops::default(), 0x35); + let lent = partial.lock().lend(Unheld(10)); + let other = hold(&pressure, Rank::Unidentified, 5); + + assert!(pressure.evict()); + assert!(Holder::is_evicted(&*partial)); + assert_eq!(pressure.held(), 15, "the partial unit stays counted"); + + assert!(pressure.evict()); + assert!( + !Holder::is_evicted(&*other), + "cancelled while the partial unit was letting go" + ); + + drop(lent); + assert!(pressure.evict()); + assert!(Holder::is_evicted(&*other)); + } + + /// Eviction and abort free the buffered bytes before their credit returns, so the + /// window the peer regains never covers memory still allocated. + #[test] + fn buffered_bytes_are_freed_before_their_credit_returns() { + let pressure = Pressure::>::default(); + for evict in [true, false] { + let hold = pressure.track::(Rank::Unidentified, Stops::default(), 0x35); + let freed_first = Arc::new(AtomicBool::new(false)); + { + let mut locked = hold.lock(); + locked.buffer_mut().extend_from_slice(&[7; 64]); + let bytes = locked.buffer_mut().split_off(32).freeze(); + locked.take(Watched { + len: 32, + bytes: Some(bytes), + freed_first: freed_first.clone(), + }); + } + match evict { + true => assert!(pressure.evict()), + false => hold.clear(), + } + assert!( + freed_first.load(Ordering::Relaxed), + "evict={evict}: the credit returned while its bytes were still allocated" + ); + assert!(hold.lock().buffer().is_empty()); + } + } + + fn media(broadcast: &str, track: u8, group: u64) -> Rank { + Rank::Media { + broadcast: Arc::new(Path::new(broadcast).to_owned()), + priority: Priority::new(track, 0, group), + } + } + + /// Track a receiver holding `held` bytes. + fn hold(pressure: &Probes, rank: Rank, held: u64) -> Arc> { + let hold = pressure.track::(rank, Stops::default(), 0x35); + hold.lock().take(Unheld(held)); + hold + } + + /// Evict once, returning every hold chosen so far. + fn evict(pressure: &Probes, holds: &[&Arc>]) -> Vec { + assert!(pressure.evict()); + holds + .iter() + .enumerate() + .filter(|(_, hold)| hold.is_evicted()) + .map(|(i, _)| i) + .collect() + } + + #[test] + fn media_goes_before_unidentified_before_control() { + let pressure = Probes::default(); + let control = hold(&pressure, Rank::Control, 100); + let unidentified = hold(&pressure, Rank::Unidentified, 50); + let media = hold(&pressure, media("room", 255, 0), 10); + let all = [&control, &unidentified, &media]; + + assert_eq!(evict(&pressure, &all), [2]); + assert_eq!(evict(&pressure, &all), [1, 2]); + assert_eq!(evict(&pressure, &all), [0, 1, 2]); + // Everything holding was chosen already. + assert!(!pressure.evict()); + } + + #[test] + fn the_least_urgent_stream_of_the_fullest_broadcast_goes() { + let pressure = Probes::default(); + // `small` holds less in total, so its less urgent track is not compared. + let small = hold(&pressure, media("small", 0, 0), 30); + let urgent = hold(&pressure, media("big", 200, 0), 20); + let lax = hold(&pressure, media("big", 100, 0), 20); + let all = [&small, &urgent, &lax]; + + assert_eq!(evict(&pressure, &all), [2]); + // `big` is left with 20, now less than `small`. + assert_eq!(evict(&pressure, &all), [0, 2]); + } + + #[test] + fn group_order_ranks_within_a_subscription() { + let pressure = Probes::default(); + let old = hold(&pressure, media("room", 1, 1), 10); + let new = hold(&pressure, media("room", 1, 2), 10); + // Newest first by default, so the older group is the less useful one. + assert_eq!(evict(&pressure, &[&old, &new]), [0]); + + let ascending = |group| Rank::Media { + broadcast: Arc::new(Path::new("room").to_owned()), + priority: Priority::new(1, 0, group).with_order(Order::Ascending), + }; + let pressure = Probes::default(); + let old = hold(&pressure, ascending(1), 10); + let new = hold(&pressure, ascending(2), 10); + assert_eq!(evict(&pressure, &[&old, &new]), [1]); + } + + #[test] + fn ties_go_to_the_larger_then_the_newer() { + let pressure = Probes::default(); + let first = hold(&pressure, Rank::Unidentified, 10); + let larger = hold(&pressure, Rank::Unidentified, 20); + let newest = hold(&pressure, Rank::Unidentified, 10); + let all = [&first, &larger, &newest]; + assert_eq!(evict(&pressure, &all), [1]); + assert_eq!(evict(&pressure, &all), [1, 2]); + } + + #[test] + fn empty_and_dropped_receivers_are_never_chosen() { + let pressure = Probes::default(); + let idle = hold(&pressure, media("room", 0, 0), 0); + drop(hold(&pressure, media("room", 0, 0), 10)); + assert!(!pressure.evict()); + assert!(!idle.is_evicted()); + } + + /// Dropped readers are swept as the registry grows, so it stays the size of what is live. + #[test] + fn dropped_receivers_are_swept() { + let pressure = Probes::default(); + let live = hold(&pressure, Rank::Control, 1); + for _ in 0..10 * PRUNE { + drop(hold(&pressure, Rank::Control, 1)); + } + assert!(pressure.state.lock().entries.len() <= PRUNE + 1); + assert_eq!(pressure.held(), 1); + drop(live); + } + + #[test] + fn rerank_once_identified() { + let pressure = Probes::default(); + let header = hold(&pressure, Rank::Unidentified, 10); + let control = hold(&pressure, Rank::Control, 10); + header.rank(media("room", 0, 0)); + assert_eq!(evict(&pressure, &[&header, &control]), [0]); + } + + /// The eviction itself returns the credit and stops the stream, so a reader parked + /// on something else holds nothing; it hears of it on its next poll. + #[test] + fn eviction_lets_go_without_the_reader() { + #[derive(Default)] + struct Count(std::sync::atomic::AtomicUsize); + impl std::task::Wake for Count { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + let pressure = Probes::default(); + let stops = Stops::default(); + let hold = pressure.track::(Rank::Unidentified, stops.clone(), 0x35); + hold.lock().take(Unheld(10)); + let count = Arc::new(Count::default()); + let waker = Waker::from(count.clone()); + assert!(!hold.poll_evicted(&Context::from_waker(&waker))); + + assert!(pressure.evict()); + assert_eq!(pressure.held(), 0); + assert_eq!(*lock(&stops.0), [0x35]); + assert_eq!(count.0.load(Ordering::Relaxed), 1); + assert!(hold.poll_evicted(&Context::from_waker(&waker))); + } + + /// Choose a victim among streams spread across broadcasts. The cost per stream must stay + /// flat along both axes. Run by nightly. + #[test] + #[ignore = "release benchmark, run by nightly"] + fn benchmark_victim() { + println!("streams,broadcasts,ns_per_choice,ns_per_stream"); + for streams in [64usize, 1024, 16384] { + for broadcasts in [1usize, 64, 1024, 16384] { + if broadcasts > streams { + continue; + } + let paths: Vec<_> = (0..broadcasts) + .map(|b| Arc::new(Path::new(&format!("broadcast/{b}")).to_owned())) + .collect(); + let pressure = Probes::default(); + let _holds: Vec<_> = (0..streams) + .map(|i| { + let rank = Rank::Media { + broadcast: paths[i % broadcasts].clone(), + priority: Priority::new((i % 256) as u8, 0, i as u64), + }; + hold(&pressure, rank, 1000 + i as u64) + }) + .collect(); + let rounds = (1 << 22) / streams; + let state = pressure.state.lock(); + let start = std::time::Instant::now(); + for _ in 0..rounds { + std::hint::black_box(victim(std::hint::black_box(&state.entries))); + } + let ns = start.elapsed().as_nanos() as f64 / rounds as f64; + println!("{streams},{broadcasts},{ns:.0},{:.1}", ns / streams as f64); + } + } + } +} diff --git a/rs/moq-net/src/server.rs b/rs/moq-net/src/server.rs index 50fe8bf01e1..84c6ea365f9 100644 --- a/rs/moq-net/src/server.rs +++ b/rs/moq-net/src/server.rs @@ -1102,6 +1102,10 @@ mod tests { fn poll_closed(&mut self, _cx: &mut std::task::Context<'_>) -> std::task::Poll { std::task::Poll::Pending } + fn poll_credit_exhausted(&mut self, _cx: &mut std::task::Context<'_>) -> std::task::Poll<()> { + std::task::Poll::Pending + } + fn stats(&self) -> impl crate::transport::Stats { crate::transport::StatsUnavailable } @@ -1139,11 +1143,14 @@ mod tests { } impl crate::transport::poll::RecvStream for FakeRecv { type Error = FakeError; + type Credit = crate::transport::Unheld; + type Stopper = (); + fn stopper(&self) {} fn poll_read( &mut self, _cx: &mut std::task::Context<'_>, dst: &mut [u8], - ) -> std::task::Poll, Self::Error>> { + ) -> std::task::Poll, Self::Error>> { if self.data.is_empty() { return std::task::Poll::Ready(Ok(None)); } @@ -1151,7 +1158,7 @@ mod tests { for slot in dst.iter_mut().take(size) { *slot = self.data.pop_front().unwrap(); } - std::task::Poll::Ready(Ok(Some(size))) + std::task::Poll::Ready(Ok(Some(crate::transport::Unheld(size as u64)))) } fn stop(&mut self, code: u32) { self.stops.lock().unwrap().push(code); diff --git a/rs/moq-net/src/transport.rs b/rs/moq-net/src/transport.rs index 8cf1ce6f551..f9b3e7380fd 100644 --- a/rs/moq-net/src/transport.rs +++ b/rs/moq-net/src/transport.rs @@ -97,6 +97,77 @@ pub trait Stats { pub struct StatsUnavailable; impl Stats for StatsUnavailable {} +/// The smallest connection receive window a transport that holds [`Credit`] may offer. +/// +/// A session holds the credit of each unit it is still assembling, so the window has to +/// fit the largest one with its framing: a legacy announcement set, past any frame payload. +pub const MIN_RECEIVE_WINDOW: u64 = { + // A type and a size prefix, each a varint of at most 8 bytes. + const PREFIXES: u64 = 16; + let message = as crate::lite::Message>::MAX_SIZE as u64 + PREFIXES; + let frame = crate::group::MAX_CACHE_BYTES; + if message > frame { message } else { frame } +}; + +/// Connection flow-control credit for bytes read from a [`poll::RecvStream`]. +/// +/// A transport that bounds receive memory keeps the connection window those bytes +/// occupy spent until their credit is dropped, so a receiver holds it until the unit +/// the bytes belong to is complete. Stream credit is unaffected and advances on reads. +/// Dropping a credit returns it, even after its stream is gone. +pub trait Credit: 'static { + /// The number of bytes this credit covers. + fn len(&self) -> u64; + + /// Whether this credit covers no bytes. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Split off the credit for the first `bytes`, keeping the rest. + /// + /// Panics if `bytes` exceeds [`len`](Self::len). + fn split(&mut self, bytes: u64) -> Self; + + /// Take over the bytes covered by `other`, read from the same session. + fn merge(&mut self, other: Self); +} + +/// The credit of a transport that returns connection credit as soon as bytes are read: +/// only a byte count. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct Unheld(pub u64); + +impl Credit for Unheld { + fn len(&self) -> u64 { + self.0 + } + + fn split(&mut self, bytes: u64) -> Self { + assert!(bytes <= self.0, "split {bytes} of {} bytes", self.0); + self.0 -= bytes; + Self(bytes) + } + + fn merge(&mut self, other: Self) { + self.0 += other.0; + } +} + +/// Stops a [`poll::RecvStream`] without borrowing it, as [`stop`](poll::RecvStream::stop) +/// would, for when the stream's reader is parked somewhere else. +/// +/// A no-op once the stream has ended. A transport whose reads return their credit (an +/// [`Unheld`] credit) never needs a stream stopped this way, so it may use `()`. +pub trait Stopper: 'static { + /// Send STOP_SENDING with a stream `code`, unless the stream already ended. + fn stop(&self, code: u32); +} + +impl Stopper for () { + fn stop(&self, _code: u32) {} +} + /// The typed transport condition when no application code was received. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] #[non_exhaustive] @@ -135,7 +206,7 @@ pub mod poll { pub type BiStreams = (::SendStream, ::RecvStream); use std::task::{Context, Poll, ready}; - use super::{Error, MaybeSend, MaybeSync, Stats}; + use super::{Credit, Error, MaybeSend, MaybeSync, Stats}; use bytes::BytesMut; use bytes::{Buf, BufMut, Bytes}; @@ -231,6 +302,17 @@ pub mod poll { /// Poll until the connection is closed by either side. fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll; + /// Poll until held [`Credit`] fills the receive allowance: the peer may send + /// nothing more and every byte it sent is held, so no stream progresses until + /// some credit is dropped. + /// + /// Resolves once each time it starts holding, then pends until it holds again + /// after some credit was dropped, so a receiver evicts once per stall. A transport whose reads + /// return their credit (an [`Unheld`](super::Unheld) credit) never fills it, so + /// it always pends. Required rather than defaulted: a transport that holds + /// credit and forgets to report this parks its receiver forever. + fn poll_credit_exhausted(&mut self, cx: &mut Context<'_>) -> Poll<()>; + /// Return connection-level statistics. /// /// Return [`super::StatsUnavailable`] if the transport does not track them. Required @@ -317,7 +399,7 @@ pub mod poll { pub trait Boxable: Session< SendStream: MaybeSend + SendStream, - RecvStream: MaybeSend, + RecvStream: MaybeSend + RecvStream, Error: MaybeSend, > + MaybeSend + MaybeSync @@ -327,7 +409,7 @@ pub mod poll { impl Boxable for S where S: Session< SendStream: MaybeSend + SendStream, - RecvStream: MaybeSend, + RecvStream: MaybeSend + RecvStream, Error: MaybeSend, > + MaybeSend + MaybeSync @@ -508,19 +590,32 @@ pub mod poll { } /// An incoming stream with reads, stop codes, and closure notification. + /// + /// Every read returns the [`Credit`] for exactly the bytes it read. A caller that + /// does not bound its receive memory drops it straight away. pub trait RecvStream: 'static { /// The error type for every operation on this stream. type Error: Error; + /// The connection credit each read returns. + type Credit: Credit; + + /// What [`stopper`](Self::stopper) returns. + type Stopper: super::Stopper; + /// Poll to read some data into the provided slice. /// - /// Returns the number of bytes read, or `None` once the peer has finished the - /// stream. An empty `dst` reads nothing and returns `Some(0)`: asking for no - /// bytes is not end of stream. - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>>; + /// Returns the credit for the bytes read ([`Credit::len`] counts them), or + /// `None` once the peer has finished the stream. An empty `dst` reads nothing + /// and returns an empty credit: asking for no bytes is not end of stream. + fn poll_read( + &mut self, + cx: &mut Context<'_>, + dst: &mut [u8], + ) -> Poll, Self::Error>>; - /// Poll to read some data into the provided buffer, advancing it by the number - /// of bytes read. + /// Poll to read some data into the provided buffer, advancing it by the bytes + /// read, which the returned credit covers. /// /// Override this to avoid a copy when the underlying transport already owns the /// bytes as a [`Bytes`], which can be handed to [`BufMut::put`] directly. @@ -528,7 +623,7 @@ pub mod poll { &mut self, cx: &mut Context<'_>, buf: &mut B, - ) -> Poll, Self::Error>> { + ) -> Poll, Self::Error>> { // Cap the slice: it is zeroed on every poll, Pending included, and a // read may be partial anyway. let len = buf.chunk_mut().len().min(64 * 1024); @@ -536,7 +631,7 @@ pub mod poll { // A destination with no room is not a closed stream. Collapsing the two // would turn "buffer full" into "stream ended", which reads as truncation. if len == 0 { - return Poll::Ready(Ok(Some(0))); + return self.poll_read(cx, &mut []); } // `poll_read` is safe and may inspect its input, so initialize spare @@ -547,26 +642,32 @@ pub mod poll { std::slice::from_raw_parts_mut(chunk.as_mut_ptr(), len) }; - let size = match ready!(self.poll_read(cx, dst))? { - Some(size) if size > 0 => size, - Some(_) => return Poll::Ready(Ok(Some(0))), - None => return Poll::Ready(Ok(None)), + let Some(credit) = ready!(self.poll_read(cx, dst))? else { + return Poll::Ready(Ok(None)); }; + let size = credit.len() as usize; assert!(size <= len, "transport read exceeded its destination"); unsafe { buf.advance_mut(size) }; - Poll::Ready(Ok(Some(size))) + Poll::Ready(Ok(Some(credit))) } - /// Poll for the next chunk of data, up to `max` bytes. + /// Poll for the next chunk of data, up to `max` bytes, with its credit. /// /// Override this when the transport can hand over a [`Bytes`] it already owns; /// the default allocates and copies. - fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Self::Error>> { + #[allow(clippy::type_complexity)] + fn poll_read_chunk( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Self::Error>> { // As in `poll_read_buf`: asking for nothing is not end of stream. if max == 0 { - return Poll::Ready(Ok(Some(Bytes::new()))); + return self + .poll_read(cx, &mut []) + .map_ok(|read| read.map(|credit| (Bytes::new(), credit))); } // Don't allocate too much. Override this to avoid the copy, or to use a @@ -574,16 +675,15 @@ pub mod poll { let capacity = max.min(8 * 1024); let mut buf = BytesMut::zeroed(capacity); - let size = match ready!(self.poll_read(cx, &mut buf))? { - Some(size) if size > 0 => size, - Some(_) => return Poll::Ready(Ok(Some(Bytes::new()))), - None => return Poll::Ready(Ok(None)), + let Some(credit) = ready!(self.poll_read(cx, &mut buf))? else { + return Poll::Ready(Ok(None)); }; + let size = credit.len() as usize; assert!(size <= capacity, "transport read exceeded its destination"); buf.truncate(size); - Poll::Ready(Ok(Some(buf.freeze()))) + Poll::Ready(Ok(Some((buf.freeze(), credit)))) } /// Send a `STOP_SENDING` QUIC code, informing the peer that no more data will be @@ -593,6 +693,9 @@ pub mod poll { /// Call this method manually if you want to specify a code yourself. fn stop(&mut self, code: u32); + /// A handle that stops this stream as [`stop`](Self::stop) does, without borrowing it. + fn stopper(&self) -> Self::Stopper; + /// Poll until the stream has been closed by either side. /// /// This includes: @@ -601,17 +704,17 @@ pub mod poll { /// - We received a FIN via [`SendStream::finish`] and read all data. fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll>; - /// Read some bytes into the slice, or `None` once the stream is finished. + /// Read some bytes into the slice, returning their credit, or `None` once finished. fn read<'a>(&'a mut self, dst: &'a mut [u8]) -> Read<'a, Self> { Read { stream: self, dst } } - /// Read some bytes into the buffer, advancing it, or `None` once finished. + /// Read some bytes into the buffer, advancing it and returning their credit, or `None`. fn read_buf<'a, B: BufMut>(&'a mut self, buf: &'a mut B) -> ReadBuf<'a, Self, B> { ReadBuf { stream: self, buf } } - /// Read the next chunk of data, up to `max` bytes, or `None` once finished. + /// Read the next chunk, up to `max` bytes, with its credit, or `None` once finished. fn read_chunk(&mut self, max: usize) -> ReadChunk<'_, Self> { ReadChunk { stream: self, max } } @@ -629,7 +732,7 @@ pub mod poll { } impl Future for Read<'_, S> { - type Output = Result, S::Error>; + type Output = Result, S::Error>; fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = &mut *self; @@ -644,7 +747,7 @@ pub mod poll { } impl Future for ReadBuf<'_, S, B> { - type Output = Result, S::Error>; + type Output = Result, S::Error>; fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = &mut *self; @@ -659,7 +762,7 @@ pub mod poll { } impl Future for ReadChunk<'_, S> { - type Output = Result, S::Error>; + type Output = Result, S::Error>; fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = &mut *self; @@ -670,18 +773,21 @@ pub mod poll { #[cfg(test)] mod tests { - use super::poll::RecvStream as _; + use super::{Unheld, poll::RecvStream as _}; use bytes::{Bytes, BytesMut}; use std::task::{Context, Poll}; struct Read; impl super::poll::RecvStream for Read { type Error = crate::Error; - fn poll_read(&mut self, _: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + type Credit = Unheld; + type Stopper = (); + fn stopper(&self) {} + fn poll_read(&mut self, _: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { // A safe transport may read its destination before overwriting it. assert!(dst.iter().all(|byte| *byte == 0)); let size = dst.len().min(2); dst[..size].copy_from_slice(&b"hi"[..size]); - Poll::Ready(Ok(Some(size))) + Poll::Ready(Ok(Some(Unheld(size as u64)))) } fn stop(&mut self, _: u32) {} fn poll_closed(&mut self, _: &mut Context<'_>) -> Poll> { @@ -695,21 +801,23 @@ mod tests { let mut buf = BytesMut::with_capacity(8); assert!(matches!( read.poll_read_buf(&mut cx, &mut buf), - Poll::Ready(Ok(Some(2))) + Poll::Ready(Ok(Some(Unheld(2)))) )); assert_eq!(buf, b"hi"[..]); - let Poll::Ready(Ok(Some(chunk))) = read.poll_read_chunk(&mut cx, 1) else { + let Poll::Ready(Ok(Some((chunk, credit)))) = read.poll_read_chunk(&mut cx, 1) else { panic!("missing chunk") }; assert_eq!(chunk, Bytes::from_static(b"h")); - let Poll::Ready(Ok(Some(chunk))) = read.poll_read_chunk(&mut cx, 0) else { + assert_eq!(credit, Unheld(1)); + let Poll::Ready(Ok(Some((chunk, credit)))) = read.poll_read_chunk(&mut cx, 0) else { panic!("missing empty chunk") }; assert!(chunk.is_empty()); + assert_eq!(credit, Unheld(0)); let mut empty: &mut [u8] = &mut []; assert!(matches!( read.poll_read_buf(&mut cx, &mut empty), - Poll::Ready(Ok(Some(0))) + Poll::Ready(Ok(Some(Unheld(0)))) )); } } diff --git a/rs/moq-net/src/transport/memory.rs b/rs/moq-net/src/transport/memory.rs index e1713331c47..3a492115c98 100644 --- a/rs/moq-net/src/transport/memory.rs +++ b/rs/moq-net/src/transport/memory.rs @@ -6,7 +6,13 @@ //! //! Each stream's channel is bounded by [`Config::stream_buffer`] queued bytes, and a write //! pends while it is full, so a slow reader throttles only its own stream. A write is -//! never split: one accepted while the channel has room may overshoot the bound. +//! never split for it: one accepted while the channel has room may overshoot the bound. +//! +//! [`Config::receive_window`] optionally bounds each side's incoming bytes across every +//! stream, like QUIC's connection flow control: queued bytes and bytes whose read +//! [`Credit`] is still held count against it, a write takes only what fits, and +//! [`poll_credit_exhausted`](poll::Session::poll_credit_exhausted) reports each time held +//! credit fills it. //! //! [`finish`](poll::SendStream::finish) closes the channel cleanly once its messages are //! read. A [`reset`](poll::SendStream::reset) or [`stop`](poll::RecvStream::stop) closes it @@ -42,6 +48,9 @@ pub struct Config { pub stream_buffer: usize, /// The application protocol both sides report, standing in for a negotiated ALPN. pub protocol: Option, + /// Incoming bytes, queued or held by an undropped [`Credit`], each side accepts across + /// all of its streams before the peer's writes pend. `None` is unbounded. + pub receive_window: Option, } impl Default for Config { @@ -51,15 +60,26 @@ impl Default for Config { Self { stream_buffer: 1024 * 1024, protocol: None, + receive_window: None, } } } /// Create two connected sessions: streams one opens, the other accepts. pub fn pair(config: Config) -> (Session, Session) { + let fills = [Arc::new(Fills::default()), Arc::new(Fills::default())]; + let window = |fills: &Arc| { + kio::Shared::new(Window { + limit: config.receive_window, + fills: fills.clone(), + ..Default::default() + }) + }; let conn = Arc::new(Conn { close: Default::default(), incoming: [Incoming::default(), Incoming::default()], + windows: [window(&fills[0]), window(&fills[1])], + fills, config, }); @@ -144,6 +164,129 @@ struct Conn { close: Close, /// What each side has yet to accept. incoming: [Incoming; 2], + /// Each side's receive allowance. + windows: [kio::Shared; 2], + /// Each side's fills of that allowance, not yet reported. + fills: [Arc; 2], +} + +/// One side's receive allowance, shared by every stream it reads. +#[derive(Default)] +struct Window { + /// The most bytes `queued + held` may reach, if bounded. + limit: Option, + /// Bytes written to this side's streams and not yet read. + queued: usize, + /// Bytes read whose credit is not dropped yet. + held: usize, + /// Held credit fills the limit. + exhausted: bool, + /// Where a new fill is reported. Its own lock and waker, so the session's monitor + /// is not woken by every read and write that changes the window. + fills: Arc, +} + +impl Window { + /// How many more bytes a writer may queue. + fn room(&self) -> usize { + match self.limit { + Some(limit) => limit.saturating_sub(self.queued + self.held), + None => usize::MAX, + } + } + + /// Latch a fill: everything the peer may send is held. + fn check(&mut self) { + let full = self.limit.is_some_and(|limit| self.held > 0 && self.held >= limit); + if full && !self.exhausted { + self.fills.report(); + } + self.exhausted = full; + } +} + +/// Fills of one side's receive allowance, waiting for +/// [`poll::Session::poll_credit_exhausted`]. +#[derive(Default)] +struct Fills { + unreported: std::sync::atomic::AtomicBool, + waker: std::sync::Mutex>, +} + +impl Fills { + fn report(&self) { + self.unreported.store(true, std::sync::atomic::Ordering::Release); + let waker = self.waker.lock().unwrap_or_else(|err| err.into_inner()).take(); + if let Some(waker) = waker { + waker.wake(); + } + } + + fn poll(&self, cx: &Context<'_>) -> Poll<()> { + use std::sync::atomic::Ordering; + if self.unreported.swap(false, Ordering::AcqRel) { + return Poll::Ready(()); + } + *self.waker.lock().unwrap_or_else(|err| err.into_inner()) = Some(cx.waker().clone()); + // Checked again: a fill between the swap and the store took no waker. + match self.unreported.swap(false, Ordering::AcqRel) { + true => Poll::Ready(()), + false => Poll::Pending, + } + } +} + +/// The receive allowance taken by bytes read from an in-memory stream, returned on drop. +pub struct Credit { + window: kio::Shared, + bytes: usize, +} + +impl Credit { + fn new(window: &kio::Shared, bytes: usize) -> Self { + Self { + window: window.clone(), + bytes, + } + } +} + +impl super::Credit for Credit { + fn len(&self) -> u64 { + self.bytes as u64 + } + + fn split(&mut self, bytes: u64) -> Self { + let bytes = bytes as usize; + assert!(bytes <= self.bytes, "split {bytes} of {} bytes", self.bytes); + self.bytes -= bytes; + Self::new(&self.window, bytes) + } + + fn merge(&mut self, mut other: Self) { + assert!( + self.window.same_channel(&other.window), + "merged credit from another session" + ); + self.bytes += std::mem::take(&mut other.bytes); + } +} + +impl Drop for Credit { + fn drop(&mut self) { + if self.bytes == 0 { + return; + } + let mut window = self.window.lock(); + window.held -= self.bytes; + window.check(); + } +} + +impl std::fmt::Debug for Credit { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Credit").field("len", &self.bytes).finish() + } } impl Conn { @@ -154,10 +297,13 @@ impl Conn { } } - fn stream(&self) -> (SendStream, RecvStream) { + /// A stream that side `reader` reads. + fn stream(&self, reader: usize) -> (SendStream, RecvStream) { let channel = kio::Shared::::default(); + let window = self.windows[reader].clone(); let send = SendStream { channel: channel.clone(), + window: window.clone(), close: self.close.clone(), buffer: self.config.stream_buffer, done: false, @@ -165,6 +311,7 @@ impl Conn { }; let recv = RecvStream { channel, + window, close: self.close.clone(), stopped: false, park: Default::default(), @@ -283,13 +430,13 @@ impl poll::Session for Session { } fn poll_open_uni(&mut self, _cx: &mut Context<'_>) -> Poll> { - let (send, recv) = self.conn().stream(); + let (send, recv) = self.conn().stream(1 - self.side.local); Poll::Ready(self.open(&self.side.outgoing().uni, recv).map(|()| send)) } fn poll_open_bi(&mut self, _cx: &mut Context<'_>) -> Poll> { - let (send, peer_recv) = self.conn().stream(); - let (peer_send, recv) = self.conn().stream(); + let (send, peer_recv) = self.conn().stream(1 - self.side.local); + let (peer_send, recv) = self.conn().stream(self.side.local); let opened = self.open(&self.side.outgoing().bi, (peer_send, peer_recv)); Poll::Ready(opened.map(|()| (send, recv))) } @@ -335,6 +482,18 @@ impl poll::Session for Session { poll_close(&self.side.conn.close, self.closed.hold(cx)) } + fn poll_credit_exhausted(&mut self, cx: &mut Context<'_>) -> Poll<()> { + let local = self.side.local; + // A fill whose credit returned before this poll relieves nothing, as in the QUIC + // drivers: report it only while the window is still exhausted. + loop { + ready!(self.side.conn.fills[local].poll(cx)); + if self.side.conn.windows[local].read().exhausted { + return Poll::Ready(()); + } + } + } + fn stats(&self) -> impl super::Stats { StatsUnavailable } @@ -361,13 +520,16 @@ enum End { } impl Channel { - /// Close with a reset or stop code and drop what is queued. An earlier end stands, - /// unless it was a finish the reader has not fully read. - fn abort(&mut self, end: End) { + /// Close with a reset or stop code and drop what is queued, returning its room to + /// `window`. An earlier end stands, unless it was a finish the reader has not fully read. + fn abort(&mut self, end: End, window: &kio::Shared) { if self.end.is_none() || !self.messages.is_empty() { self.end = Some(end); } self.messages.clear(); + if self.queued > 0 { + window.lock().queued -= self.queued; + } self.queued = 0; } @@ -385,6 +547,8 @@ impl Channel { /// The writing half of an in-memory stream. pub struct SendStream { channel: kio::Shared, + /// The reading side's receive allowance. + window: kio::Shared, close: Close, buffer: usize, /// This side finished or reset the stream. @@ -401,7 +565,8 @@ impl poll::SendStream for SendStream { self.poll_write_buf(cx, &mut buf) } - /// Queue all of `buf` as one message, waiting while the channel is full. + /// Queue as much of `buf` as the receive window allows as one message, waiting while + /// the channel or the window is full. fn poll_write_buf(&mut self, cx: &mut Context<'_>, buf: &mut B) -> Poll> { if self.done { return Poll::Ready(Err(Error::Closed)); @@ -425,9 +590,20 @@ impl poll::SendStream for SendStream { if let Some(End::Stop(code)) = channel.end { return Poll::Ready(Err(Error::Stopped(code))); } + // A reset stream's bytes are dropped, never read, so they take no room. Otherwise + // wait for some; the channel lock keeps the stream from ending meanwhile. + let mut window = match channel.end { + Some(_) => self.window.lock(), + None => ready!(self.window.poll(waiter, |window| match window.room() > 0 { + true => Poll::Ready(()), + false => Poll::Pending, + })), + }; // Take ownership rather than copy: free for a `Bytes` source. - let size = buf.remaining(); + let size = buf.remaining().min(window.room()); + window.queued += size; + drop(window); channel.messages.push_back(buf.copy_to_bytes(size)); channel.queued += size; @@ -465,7 +641,7 @@ impl poll::SendStream for SendStream { // close is checked under the channel lock so a concurrent close can't land between. let mut channel = self.channel.lock(); if closed(&self.close).is_none() { - channel.abort(End::Reset(code)); + channel.abort(End::Reset(code), &self.window); } } @@ -493,41 +669,63 @@ impl Drop for SendStream { /// The reading half of an in-memory stream. pub struct RecvStream { channel: kio::Shared, + /// This side's receive allowance. + window: kio::Shared, close: Close, /// This side stopped the stream. stopped: bool, park: kio::Park, } +/// Stops a [`RecvStream`] without borrowing it. +#[derive(Clone)] +pub struct Stopper { + channel: kio::Shared, + window: kio::Shared, + close: Close, +} + +impl super::Stopper for Stopper { + fn stop(&self, code: u32) { + // A closed session leaves its streams as they were; checked under the channel lock, + // as in `reset`. + let mut channel = self.channel.lock(); + if closed(&self.close).is_none() { + channel.abort(End::Stop(code), &self.window); + } + } +} + impl poll::RecvStream for RecvStream { type Error = Error; + type Credit = Credit; + type Stopper = Stopper; - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Error>> { - let Some(chunk) = ready!(self.poll_read_chunk(cx, dst.len()))? else { + fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Error>> { + let Some((chunk, credit)) = ready!(self.poll_read_chunk(cx, dst.len()))? else { return Poll::Ready(Ok(None)); }; dst[..chunk.len()].copy_from_slice(&chunk); - Poll::Ready(Ok(Some(chunk.len()))) + Poll::Ready(Ok(Some(credit))) } - fn poll_read_buf(&mut self, cx: &mut Context<'_>, buf: &mut B) -> Poll, Error>> { + fn poll_read_buf(&mut self, cx: &mut Context<'_>, buf: &mut B) -> Poll, Error>> { // The same cap the default uses, but without zeroing a destination first. let max = buf.remaining_mut().min(64 * 1024); - let Some(chunk) = ready!(self.poll_read_chunk(cx, max))? else { + let Some((chunk, credit)) = ready!(self.poll_read_chunk(cx, max))? else { return Poll::Ready(Ok(None)); }; - let size = chunk.len(); buf.put(chunk); - Poll::Ready(Ok(Some(size))) + Poll::Ready(Ok(Some(credit))) } /// Pop the next message, or split off its first `max` bytes. - fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Error>> { + fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Error>> { if self.stopped { return Poll::Ready(Err(Error::Closed)); } if max == 0 { - return Poll::Ready(Ok(Some(Bytes::new()))); + return Poll::Ready(Ok(Some((Bytes::new(), Credit::new(&self.window, 0))))); } let waiter = self.park.hold(cx); @@ -556,18 +754,28 @@ impl poll::RecvStream for RecvStream { }; channel.queued -= chunk.len(); - Poll::Ready(Ok(Some(chunk))) + // Read bytes move from queued to held, still spending the window. + let mut window = self.window.lock(); + window.queued -= chunk.len(); + window.held += chunk.len(); + window.check(); + drop(window); + + let credit = Credit::new(&self.window, chunk.len()); + Poll::Ready(Ok(Some((chunk, credit)))) } fn stop(&mut self, code: u32) { - if std::mem::replace(&mut self.stopped, true) { - return; + if !std::mem::replace(&mut self.stopped, true) { + super::Stopper::stop(&self.stopper(), code); } - // A closed session leaves its streams as they were; checked under the channel lock, - // as in `reset`. - let mut channel = self.channel.lock(); - if closed(&self.close).is_none() { - channel.abort(End::Stop(code)); + } + + fn stopper(&self) -> Stopper { + Stopper { + channel: self.channel.clone(), + window: self.window.clone(), + close: self.close.clone(), } } @@ -607,6 +815,7 @@ mod tests { }; use super::*; + use crate::transport::Credit as _; use poll::{RecvStream as _, SendStream as _, Session as _}; /// Counts its wakes. @@ -647,8 +856,10 @@ mod tests { send.poll_write_buf(&mut noop(), &mut chunk.clone()) } + /// Read a chunk, returning its credit straight away. fn read(recv: &mut RecvStream, max: usize) -> Poll, Error>> { recv.poll_read_chunk(&mut noop(), max) + .map_ok(|chunk| chunk.map(|(chunk, _)| chunk)) } fn text(chunk: Poll, Error>>) -> Bytes { @@ -888,4 +1099,105 @@ mod tests { } assert!(b.poll_recv_datagram(&mut noop()).is_pending()); } + + fn window(receive_window: usize) -> Config { + Config { + receive_window: Some(receive_window), + ..Default::default() + } + } + + /// Queued bytes and held credit both spend the window; dropping the credit returns it. + #[test] + fn the_window_bounds_queued_and_held_bytes() { + let (mut a, mut b) = pair(window(10)); + let (mut send, mut recv) = uni(&mut a, &mut b); + let (mut other, _other_recv) = uni(&mut a, &mut b); + + // A write takes only what fits, and the window is shared by every stream. + assert_eq!(ready(write(&mut send, &Bytes::from(vec![1u8; 16]))).unwrap(), 10); + assert!(write(&mut other, &Bytes::from_static(b"x")).is_pending()); + + // Reading moves bytes from queued to held, which still spends the window. + let (chunk, credit) = ready(recv.poll_read_chunk(&mut noop(), 4)).unwrap().unwrap(); + assert_eq!(chunk.len(), 4); + assert!(write(&mut other, &Bytes::from_static(b"x")).is_pending()); + + drop(credit); + assert_eq!(ready(write(&mut other, &Bytes::from(vec![2u8; 8]))).unwrap(), 4); + } + + /// A fill is reported once, and again only after some credit returned and the window + /// filled again. + #[test] + fn exhaustion_is_reported_once_per_fill() { + let (mut a, mut b) = pair(window(8)); + let (mut send, mut recv) = uni(&mut a, &mut b); + + assert_eq!(ready(write(&mut send, &Bytes::from(vec![1u8; 8]))).unwrap(), 8); + // Unread bytes are still on their way to the application, so nothing is stuck. + assert!(b.poll_credit_exhausted(&mut noop()).is_pending()); + + let (_, mut held) = ready(recv.poll_read_chunk(&mut noop(), 8)).unwrap().unwrap(); + ready(b.poll_credit_exhausted(&mut noop())); + assert!(b.poll_credit_exhausted(&mut noop()).is_pending()); + + // Returning some credit clears the fill; spending it again is a new one. + drop(held.split(2)); + assert!(b.poll_credit_exhausted(&mut noop()).is_pending()); + assert_eq!(ready(write(&mut send, &Bytes::from(vec![1u8; 8]))).unwrap(), 2); + let (_, more) = ready(recv.poll_read_chunk(&mut noop(), 8)).unwrap().unwrap(); + ready(b.poll_credit_exhausted(&mut noop())); + + // The other direction has its own window. + assert!(a.poll_credit_exhausted(&mut noop()).is_pending()); + drop((held, more)); + } + + /// A fill whose credit returns before the monitor polls is not reported, so nothing is + /// cancelled once progress is possible again. + #[test] + fn a_fill_released_before_the_poll_is_not_reported() { + let (mut a, mut b) = pair(window(8)); + let (mut send, mut recv) = uni(&mut a, &mut b); + + assert_eq!(ready(write(&mut send, &Bytes::from(vec![1u8; 8]))).unwrap(), 8); + let (_, mut held) = ready(recv.poll_read_chunk(&mut noop(), 8)).unwrap().unwrap(); + drop(held.split(4)); + assert!(b.poll_credit_exhausted(&mut noop()).is_pending()); + + // Filling the window again is still reported. + assert_eq!(ready(write(&mut send, &Bytes::from(vec![1u8; 8]))).unwrap(), 4); + let (_, more) = ready(recv.poll_read_chunk(&mut noop(), 8)).unwrap().unwrap(); + ready(b.poll_credit_exhausted(&mut noop())); + drop((held, more)); + } + + /// The peer waiting on a fill wakes when it happens. + #[test] + fn exhaustion_wakes_the_monitor() { + let (mut a, mut b) = pair(window(4)); + let (mut send, mut recv) = uni(&mut a, &mut b); + + let count = Arc::new(Count::default()); + let waker = Waker::from(count.clone()); + assert!(b.poll_credit_exhausted(&mut Context::from_waker(&waker)).is_pending()); + + ready(write(&mut send, &Bytes::from(vec![1u8; 4]))).unwrap(); + let _held = ready(recv.poll_read_chunk(&mut noop(), 4)).unwrap().unwrap(); + assert!(count.0.load(Ordering::Relaxed) > 0); + } + + /// Stopping a stream returns the window its queued bytes took. + #[test] + fn stop_returns_queued_room() { + let (mut a, mut b) = pair(window(8)); + let (mut send, mut recv) = uni(&mut a, &mut b); + let (mut other, _other_recv) = uni(&mut a, &mut b); + + assert_eq!(ready(write(&mut send, &Bytes::from(vec![1u8; 8]))).unwrap(), 8); + assert!(write(&mut other, &Bytes::from_static(b"x")).is_pending()); + recv.stop(1); + assert_eq!(ready(write(&mut other, &Bytes::from(vec![2u8; 8]))).unwrap(), 8); + } } diff --git a/rs/moq-net/tests/integration/main.rs b/rs/moq-net/tests/integration/main.rs index d155bb218f5..bf11c99ba25 100644 --- a/rs/moq-net/tests/integration/main.rs +++ b/rs/moq-net/tests/integration/main.rs @@ -33,6 +33,7 @@ mod mesh_withdraw; mod model; mod multi_cdn; mod quiet_catalog; +mod receive_credit; mod rejoin; mod replacement; mod request_caps; diff --git a/rs/moq-net/tests/integration/receive_credit.rs b/rs/moq-net/tests/integration/receive_credit.rs new file mode 100644 index 00000000000..c784e1ac310 --- /dev/null +++ b/rs/moq-net/tests/integration/receive_credit.rs @@ -0,0 +1,258 @@ +#![cfg(feature = "async")] + +//! Incomplete messages hold the receiver's connection credit, and when they fill the +//! allowance the receiver cancels the least useful stream and keeps going. +//! +//! Deterministic: paused time over the mock transport, whose finite receive allowance +//! stands in for QUIC connection flow control. Without recovery each of these stalls: +//! the sender has spent every byte of credit on messages that can never complete. + +use crate::support; + +use std::time::Duration; + +use bytes::Bytes; +use moq_net::transport::poll::{SendStream as _, Session as _}; +use moq_net::{Hop, Timestamp, Version, frame, origin, track::Subscription}; +use support::harness::{MockConnectOptions, MockPair, connect_mock, subscribed}; + +const TIMEOUT: Duration = Duration::from_secs(10); + +/// Each side's receive allowance. +const WINDOW: usize = 64 * 1024; + +const VERSIONS: &[&str] = &[ + "moq-lite-03", + "moq-lite-06", + "moq-transport-14", + "moq-transport-16", + "moq-transport-17", + "moq-transport-22", +]; + +fn timed() -> moq_net::track::Info { + moq_net::track::Info::default().with_timescale(moq_net::Timescale::MILLI) +} + +fn produce_origin(hop: u64) -> origin::Producer { + let (producer, driver) = origin::Producer::new(origin::Config::new(Hop::new(hop).unwrap())); + support::harness::spawn(driver); + producer +} + +/// A client publishing broadcast `room` to a server that subscribes to it, over links +/// that each accept [`WINDOW`] bytes of credit. +struct Setup { + pair: MockPair, + broadcast: moq_net::broadcast::Producer, + remote: moq_net::broadcast::Consumer, +} + +async fn setup(version: Version) -> Setup { + let publisher = produce_origin(1); + let broadcast = publisher.create_broadcast("room").unwrap(); + broadcast.announce(Default::default()).unwrap(); + + let subscriber = produce_origin(2); + let mut options = MockConnectOptions::new(version); + options.client_publish = Some(publisher.consume()); + options.server_subscribe = Some(subscriber.clone()); + options.transport.receive_window = Some(WINDOW); + let pair = connect_mock(options).await; + + let consumer = subscriber.consume(); + moq_clock::timeout(TIMEOUT, consumer.routed("room")) + .await + .expect("announce timeout") + .expect("routed"); + let remote = moq_clock::timeout(TIMEOUT, consumer.request_broadcast("room", None)) + .await + .expect("resolve timeout") + .expect("broadcast resolves"); + + Setup { + pair, + broadcast, + remote, + } +} + +impl Setup { + /// Subscribe to a new track at `priority`, once the publisher sees the demand. + async fn subscribe(&self, name: &str, priority: u8) -> (moq_net::track::Producer, moq_net::track::Subscriber) { + let track = self.broadcast.create_track(name, timed()).unwrap(); + let sub = moq_clock::timeout( + TIMEOUT, + self.remote + .track(name) + .unwrap() + .subscribe(Subscription::default().with_priority(priority)), + ) + .await + .expect("subscribe timeout") + .expect("subscribe"); + subscribed(&track).await; + (track, sub) + } +} + +/// The next frame's payload, which must arrive in time. +async fn next_frame(sub: &mut moq_net::track::Subscriber) -> moq_net::Result { + let mut group = moq_clock::timeout(TIMEOUT, sub.recv_group()) + .await + .expect("group timed out")? + .expect("track ended"); + let frame = moq_clock::timeout(TIMEOUT, group.read_frame()) + .await + .expect("frame timed out")? + .expect("group ended"); + Ok(frame.payload) +} + +/// Two partial frames that together exceed the allowance: the lower-priority one is +/// cancelled, and the higher-priority one completes on the credit it returns. +#[moq_clock::test] +async fn the_less_urgent_partial_frame_gives_way() { + for version in VERSIONS { + let setup = setup(version.parse().unwrap()).await; + let (low, mut low_sub) = setup.subscribe("low", 1).await; + let (high, mut high_sub) = setup.subscribe("high", 200).await; + + // Three quarters of the allowance, held by a frame that never finishes. + let mut low_group = low.append_group().unwrap(); + let mut low_frame = low_group + .create_frame(frame::Info { + size: WINDOW as u64, + timestamp: Some(Timestamp::ZERO), + }) + .unwrap(); + low_frame.write(vec![1u8; WINDOW * 3 / 4]).unwrap(); + + // Half the allowance, which only fits once the low frame lets go. + let payload = Bytes::from(vec![2u8; WINDOW / 2]); + let mut high_group = high.append_group().unwrap(); + high_group.write_frame(Timestamp::ZERO, payload.clone()).unwrap(); + high_group.finish().unwrap(); + + let got = next_frame(&mut high_sub).await.expect("high frame"); + assert_eq!(got, payload, "{version}"); + + // The session survives, and the low track keeps delivering what fits. Its + // cancelled group never completes, so the next frame the subscriber reads is this. + drop(low_frame); + drop(low_group); + let mut group = low.append_group().unwrap(); + group.write_frame(Timestamp::ZERO, &b"after"[..]).unwrap(); + group.finish().unwrap(); + let got = next_frame(&mut low_sub).await.expect("low frame"); + assert_eq!(got, &b"after"[..], "{version}"); + drop(setup.pair); + } +} + +/// Write `bytes` on a new uni stream from the client, in the background: it pends once +/// the allowance runs out. +fn raw_uni(pair: &MockPair, bytes: Vec) { + let mut session = pair.client_transport.clone(); + moq_clock::sim::spawn(async move { + let Ok(mut send) = session.open_uni().await else { return }; + let _ = send.write_chunk(Bytes::from(bytes)).await; + // Hold the stream open, so only the receiver can end it. + std::future::pending::<()>().await; + }); +} + +/// The prefix of a lite GROUP stream whose header claims `size` bytes and carries `body` +/// of them, so it never identifies itself. +fn partial_group_header(size: u32, body: usize) -> Vec { + assert!(body < size as usize); + let mut bytes = vec![0x00]; // GROUP + bytes.extend_from_slice(&(0x8000_0000 | size).to_be_bytes()); // A 4-byte varint size. + bytes.resize(bytes.len() + body, 0); + bytes +} + +/// Headers that never complete fill the allowance: one is cancelled, and a subscription +/// that arrives afterwards still delivers. +#[moq_clock::test] +async fn unidentified_headers_are_cancelled() { + let setup = setup("moq-lite-06".parse().unwrap()).await; + + // Together more than the allowance, and each short of its declared size. + raw_uni(&setup.pair, partial_group_header(u16::MAX as u32, 40_000)); + raw_uni(&setup.pair, partial_group_header(u16::MAX as u32, 40_000)); + + let (track, mut sub) = setup.subscribe("video", 1).await; + let mut group = track.append_group().unwrap(); + group.write_frame(Timestamp::ZERO, vec![3u8; 4096]).unwrap(); + group.finish().unwrap(); + assert_eq!(next_frame(&mut sub).await.expect("frame").len(), 4096); +} + +/// Control requests that never complete fill the allowance: one is cancelled with its +/// stream, and the session carries on. +#[moq_clock::test] +async fn partial_control_requests_are_cancelled() { + let setup = setup("moq-lite-06".parse().unwrap()).await; + + for _ in 0..2 { + let mut session = setup.pair.client_transport.clone(); + moq_clock::sim::spawn(async move { + let Ok((mut send, _recv)) = session.open_bi().await else { + return; + }; + // ANNOUNCE_PLEASE, with a body that claims more than arrives. + let mut bytes = vec![0x01]; + bytes.extend_from_slice(&(0x8000_0000u32 | u16::MAX as u32).to_be_bytes()); + bytes.resize(bytes.len() + 40_000, 0); + let _ = send.write_chunk(Bytes::from(bytes)).await; + // Hold the stream open, so only the receiver can end it. + std::future::pending::<()>().await; + }); + } + + let (track, mut sub) = setup.subscribe("video", 1).await; + let mut group = track.append_group().unwrap(); + group.write_frame(Timestamp::ZERO, vec![3u8; 4096]).unwrap(); + group.finish().unwrap(); + assert_eq!(next_frame(&mut sub).await.expect("frame").len(), 4096); +} + +/// A frame as large as the allowance allows completes while partial frames of a less +/// urgent track keep arriving, each cancelled as it gets in the way. +#[moq_clock::test] +async fn the_largest_unit_completes_under_pressure() { + for version in VERSIONS { + let setup = setup(version.parse().unwrap()).await; + let (low, _low_sub) = setup.subscribe("low", 1).await; + let (high, mut high_sub) = setup.subscribe("high", 200).await; + + // Several groups' worth of frames that never finish, opened before and after. + let mut partial = Vec::new(); + let mut open_partial = || { + // Leaked, so the frame can borrow its group for the rest of the test. + let group = Box::leak(Box::new(low.append_group().unwrap())); + let mut frame = group + .create_frame(frame::Info { + size: WINDOW as u64, + timestamp: Some(Timestamp::ZERO), + }) + .unwrap(); + frame.write(vec![1u8; WINDOW / 4]).unwrap(); + partial.push(frame); + }; + open_partial(); + open_partial(); + + // Room for the stream's headers and nothing else. + let payload = Bytes::from(vec![2u8; WINDOW - 64]); + let mut group = high.append_group().unwrap(); + group.write_frame(Timestamp::ZERO, payload.clone()).unwrap(); + group.finish().unwrap(); + open_partial(); + + let got = next_frame(&mut high_sub).await.expect("high frame"); + assert_eq!(got.len(), payload.len(), "{version}"); + drop(setup.pair); + } +} diff --git a/rs/moq-net/tests/integration/subscribe_tracks.rs b/rs/moq-net/tests/integration/subscribe_tracks.rs index 64e0a247fbc..c6026f54380 100644 --- a/rs/moq-net/tests/integration/subscribe_tracks.rs +++ b/rs/moq-net/tests/integration/subscribe_tracks.rs @@ -116,7 +116,7 @@ async fn subscribe_tracks(pair: &MockPair) -> Vec { .expect("reply timed out") .expect("read") { - reply.extend_from_slice(&buf[..n]); + reply.extend_from_slice(&buf[..moq_net::transport::Credit::len(&n) as usize]); } reply } diff --git a/rs/moq-net/tests/integration/support/mock.rs b/rs/moq-net/tests/integration/support/mock.rs index 2301a706073..2e53c03c882 100644 --- a/rs/moq-net/tests/integration/support/mock.rs +++ b/rs/moq-net/tests/integration/support/mock.rs @@ -23,7 +23,7 @@ use std::{ use bytes::{Buf, BufMut, Bytes}; use moq_net::{ time::Instant, - transport::{memory, poll}, + transport::{Credit as _, memory, poll}, }; use super::harness::now; @@ -283,20 +283,22 @@ impl Drop for MockSendStream { // ── RecvStream ────────────────────────────────────────────────────── /// What the transport handed over: a chunk, the FIN, or an error. -type Event = Result, memory::Error>; +type Event = Result, memory::Error>; /// A receive stream that delivers each event once its arrival has passed. pub struct MockRecvStream { /// Taken once this side stops reading. inner: Option, + /// What [`MockStopper`] stops, kept because `inner` goes once this side stops. + stopper: memory::Stopper, line: Arc, conn: Arc, /// The next event, read from the transport but still crossing the link. flight: Option<(Instant, Event)>, /// The arrival a wake is already scheduled for, so each flight schedules one. landing: Option, - /// The rest of a chunk a smaller read split. - buf: Bytes, + /// The rest of a chunk a smaller read split, with its credit. + buf: Option<(Bytes, memory::Credit)>, /// Data chunks handed to the reader so far. chunks: usize, park: kio::Park, @@ -305,12 +307,13 @@ pub struct MockRecvStream { impl MockRecvStream { fn new(inner: memory::RecvStream, line: Arc, conn: Arc) -> Self { Self { + stopper: poll::RecvStream::stopper(&inner), inner: Some(inner), line, conn, flight: None, landing: None, - buf: Bytes::new(), + buf: None, chunks: 0, park: Default::default(), } @@ -381,43 +384,79 @@ impl MockRecvStream { } } +/// Stops a [`MockRecvStream`] without borrowing it, telling the writer once the latency +/// has passed. +#[derive(Clone)] +pub struct MockStopper { + inner: memory::Stopper, + conn: Arc, +} + +impl moq_net::transport::Stopper for MockStopper { + fn stop(&self, code: u32) { + let latency = self.conn.latency(); + if latency.is_zero() { + return moq_net::transport::Stopper::stop(&self.inner, code); + } + let inner = self.inner.clone(); + drop(moq_clock::sim::spawn(async move { + moq_clock::sleep(latency).await; + moq_net::transport::Stopper::stop(&inner, code); + })); + } +} + impl poll::RecvStream for MockRecvStream { type Error = memory::Error; + type Credit = memory::Credit; + type Stopper = MockStopper; + + fn stopper(&self) -> MockStopper { + MockStopper { + inner: self.stopper.clone(), + conn: self.conn.clone(), + } + } - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { - let Some(chunk) = ready!(self.poll_read_chunk(cx, dst.len()))? else { + fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + let Some((chunk, credit)) = ready!(self.poll_read_chunk(cx, dst.len()))? else { return Poll::Ready(Ok(None)); }; dst[..chunk.len()].copy_from_slice(&chunk); - Poll::Ready(Ok(Some(chunk.len()))) + Poll::Ready(Ok(Some(credit))) } fn poll_read_buf( &mut self, cx: &mut Context<'_>, buf: &mut B, - ) -> Poll, Self::Error>> { + ) -> Poll, Self::Error>> { let max = buf.remaining_mut().min(64 * 1024); - let Some(chunk) = ready!(self.poll_read_chunk(cx, max))? else { + let Some((chunk, credit)) = ready!(self.poll_read_chunk(cx, max))? else { return Poll::Ready(Ok(None)); }; - let size = chunk.len(); buf.put(chunk); - Poll::Ready(Ok(Some(size))) + Poll::Ready(Ok(Some(credit))) } - fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Self::Error>> { - if max == 0 { - return Poll::Ready(Ok(Some(Bytes::new()))); - } - if self.buf.is_empty() { + fn poll_read_chunk( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Self::Error>> { + if self.buf.is_none() { match ready!(self.poll_event(cx))? { - Some(chunk) => self.buf = chunk, + Some(chunk) => self.buf = Some(chunk), None => return Poll::Ready(Ok(None)), } } - let size = max.min(self.buf.len()); - Poll::Ready(Ok(Some(self.buf.split_to(size)))) + let (buf, credit) = self.buf.as_mut().expect("filled above"); + let size = max.min(buf.len()); + let chunk = (buf.split_to(size), credit.split(size as u64)); + if buf.is_empty() { + self.buf = None; + } + Poll::Ready(Ok(Some(chunk))) } fn stop(&mut self, code: u32) { @@ -636,6 +675,10 @@ impl poll::Session for MockSession { self.inner.poll_closed(cx) } + fn poll_credit_exhausted(&mut self, cx: &mut Context<'_>) -> Poll<()> { + self.inner.poll_credit_exhausted(cx) + } + fn stats(&self) -> impl moq_net::transport::Stats { assert!( !self.side.panic_on_stats.load(Ordering::Relaxed), diff --git a/rs/moq-quic/src/connection/mod.rs b/rs/moq-quic/src/connection/mod.rs index b2f5fdb6049..5a2cb1e6b69 100644 --- a/rs/moq-quic/src/connection/mod.rs +++ b/rs/moq-quic/src/connection/mod.rs @@ -1394,6 +1394,15 @@ impl Connection { transmit } + /// Whether held credit came to fill the receive window since the last call. + /// + /// True when [`CreditStats::exhausted()`] holds now and some read, release, or discarded + /// stream since the last call left it so: nothing more can arrive until held credit is + /// released. Check after holding, releasing, stopping, or handling incoming data. + pub fn poll_credit_exhausted(&mut self) -> bool { + mem::take(&mut self.streams.stalled) && self.streams.credit_stats().exhausted() + } + /// Ping the remote endpoint /// /// Causes an ACK-eliciting packet to be transmitted. diff --git a/rs/moq-quic/src/connection/stats.rs b/rs/moq-quic/src/connection/stats.rs index 1a6dab4ba3b..bf3f424661f 100644 --- a/rs/moq-quic/src/connection/stats.rs +++ b/rs/moq-quic/src/connection/stats.rs @@ -192,4 +192,15 @@ pub struct CreditStats { /// Bytes the peer may still send before it needs more credit, including credit not yet /// advertised pub available: u64, + /// Bytes received, or owed for a gap below a received offset, that the application has not + /// read yet + pub unread: u64, +} + +impl CreditStats { + /// Whether held credit is all the peer may send: nothing more can arrive and nothing + /// received is left to read, so only releasing [`held`](Self::held) credit makes progress. + pub fn exhausted(&self) -> bool { + self.held > 0 && self.available == 0 && self.unread == 0 + } } diff --git a/rs/moq-quic/src/connection/streams/state.rs b/rs/moq-quic/src/connection/streams/state.rs index 6c1f57c8d96..000a0ae483e 100644 --- a/rs/moq-quic/src/connection/streams/state.rs +++ b/rs/moq-quic/src/connection/streams/state.rs @@ -134,6 +134,10 @@ pub(crate) struct StreamsState { /// Bytes read with [`Chunks::hold`](super::Chunks::hold) whose connection credit the /// application has not released yet pub(super) held: u64, + /// Connection credit returned for received bytes so far: read, released, or discarded + credited: u64, + /// Held credit was left filling the window since the application last asked + pub(crate) stalled: bool, /// Which connection [`Credit`]s handed out here belong to pub(super) tag: ConnectionTag, /// Total quantity of outgoing data retained in send buffers, including acknowledged data @@ -202,6 +206,8 @@ impl StreamsState { data_sent: 0, data_recvd: 0, held: 0, + credited: 0, + stalled: false, tag: ConnectionTag::next(), buffered_data: 0, send_window, @@ -1060,6 +1066,14 @@ impl StreamsState { /// suppress sending further updates until the window increases significantly /// again. pub(super) fn add_read_credits(&mut self, credits: u64) -> ShouldTransmit { + self.credited += credits; + debug_assert!( + self.credited + self.held <= self.data_recvd, + "credited {} and held {} past the {} bytes received", + self.credited, + self.held, + self.data_recvd + ); if credits > self.receive_window_shrink_debt { let net_credits = credits - self.receive_window_shrink_debt; self.local_max_data = self.local_max_data.saturating_add(net_credits); @@ -1068,6 +1082,10 @@ impl StreamsState { self.receive_window_shrink_debt -= credits; } + // Every change to held or returned credit ends here, a hold included (its `finalize` + // returns nothing), so this sees each moment that leaves the window stalled. + self.stalled |= self.held > 0 && self.credit_stats().exhausted(); + ShouldTransmit(self.max_data_due()) } @@ -1107,6 +1125,7 @@ impl StreamsState { CreditStats { held: self.held, available: self.local_max_data.saturating_sub(self.data_recvd), + unread: self.data_recvd - self.credited - self.held, } } @@ -2424,6 +2443,8 @@ mod tests { // credit, expecting all of them added to local_max_data let credits = 1024u64; + // As if that many bytes arrived and were read: credit never outruns what was received. + server.data_recvd += credits; let should_transmit = server.add_read_credits(credits); assert_eq!(server.receive_window_shrink_debt, 0); assert_eq!(server.local_max_data, prev_local_max_data + credits); @@ -2449,6 +2470,7 @@ mod tests { let credits = 1024u64; for _ in 0..2 { let expected_receive_window_shrink_debt = server.receive_window_shrink_debt - credits; + server.data_recvd += credits; let should_transmit = server.add_read_credits(credits); assert_eq!( server.receive_window_shrink_debt, @@ -2463,6 +2485,7 @@ mod tests { let prev_local_max_data = server.local_max_data; let expected_local_max_data = server.local_max_data + (credits - server.receive_window_shrink_debt); + server.data_recvd += credits; let _should_transmit = server.add_read_credits(credits); assert_eq!(server.receive_window_shrink_debt, 0); assert_eq!(server.local_max_data, expected_local_max_data); @@ -2471,6 +2494,7 @@ mod tests { // credit again, all should be added to local_max_data let credits = 1024 * 512; let expected_local_max_data = server.local_max_data + credits; + server.data_recvd += credits; let should_transmit = server.add_read_credits(credits); assert_eq!(server.receive_window_shrink_debt, 0); assert_eq!(server.local_max_data, expected_local_max_data); @@ -3124,6 +3148,8 @@ mod tests { Stop, Reset, StopThenReset, + ResetThenStop, + ReliableReset, } fn data(client: &mut StreamsState, id: StreamId, offset: u64, len: usize, fin: bool) { @@ -3170,7 +3196,14 @@ mod tests { /// it takes and whether the held bytes are released before or after its state is reclaimed. #[test] fn held_credit_returns_once() { - for end in [End::Fin, End::Stop, End::Reset, End::StopThenReset] { + for end in [ + End::Fin, + End::Stop, + End::Reset, + End::StopThenReset, + End::ResetThenStop, + End::ReliableReset, + ] { let mut client = make(Side::Client); let initial = client.local_max_data; let id = StreamId::new(Side::Server, Dir::Uni, 0); @@ -3236,6 +3269,34 @@ mod tests { reset(&mut client, id, 128); 128 } + End::ResetThenStop => { + // The reset returns what it abandons past its Reliable Size; the stop adds + // only the gap below it, which the reset still owed. + let _ = client + .received_reset(reset_at(id, 0, 128, Some(48))) + .unwrap(); + let mut pending = Retransmits::default(); + RecvStream { + id, + state: &mut client, + pending: &mut pending, + } + .stop(0u32.into()) + .unwrap(); + 128 + } + End::ReliableReset => { + // Data past the Reliable Size is dropped on arrival, and reading the owed + // bytes reaches the reset. + let _ = client + .received_reset(reset_at(id, 0, 128, Some(48))) + .unwrap(); + data(&mut client, id, 32, 32, false); + let owed = read_held(&mut client, id); + assert_eq!(owed.len(), 16, "{end:?}"); + let _ = client.release(owed); + 128 + } }; assert!( !client.recv.contains_key(&id), @@ -3265,6 +3326,101 @@ mod tests { assert_eq!(client.local_max_data, initial + 16); } + /// Exhausted only once every byte the peer may send arrived and is held: unread bytes, gaps + /// still owed, and unadvertised credit all mean something can still move. + #[test] + fn exhausted_when_every_byte_is_held() { + let mut client = StreamsState::new( + Side::Client, + 128u32.into(), + 128u32.into(), + 1024 * 1024, + 64u32.into(), + (1024 * 1024u32).into(), + ); + let id = StreamId::new(Side::Server, Dir::Uni, 0); + assert!(!client.credit_stats().exhausted(), "nothing held"); + + // A readable prefix, then a gap below the last byte: the gap is still owed. + data(&mut client, id, 0, 32, false); + data(&mut client, id, 48, 16, false); + let prefix = read_held(&mut client, id); + let stats = client.credit_stats(); + assert_eq!((stats.held, stats.available, stats.unread), (32, 0, 32)); + assert!(!stats.exhausted(), "the gap can still arrive"); + + // The gap fills, so everything is readable, but not read yet. + data(&mut client, id, 32, 16, false); + assert!( + !client.credit_stats().exhausted(), + "unread bytes can still be read" + ); + + let rest = read_held(&mut client, id); + let stats = client.credit_stats(); + assert_eq!((stats.held, stats.available, stats.unread), (64, 0, 0)); + assert!(stats.exhausted()); + + // Releasing any of it lets the peer send again. + let _ = client.release(prefix); + assert!(!client.credit_stats().exhausted()); + let _ = client.release(rest); + assert_eq!(client.credit_stats().unread, 0); + } + + /// A stall reports once, and again whenever credit moves without relieving it. + #[test] + fn stall_reports_once_per_change() { + let mut client = StreamsState::new( + Side::Client, + 128u32.into(), + 128u32.into(), + 1024 * 1024, + 64u32.into(), + (1024 * 1024u32).into(), + ); + let take = |state: &mut StreamsState| { + mem::take(&mut state.stalled) && state.credit_stats().exhausted() + }; + let id = StreamId::new(Side::Server, Dir::Uni, 0); + data(&mut client, id, 0, 64, false); + let mut held = read_held(&mut client, id); + assert!(take(&mut client)); + assert!(!take(&mut client), "reported already"); + + // Shrinking the window absorbs a release, so the peer still may send nothing. + client.set_receive_window(32u32.into()); + let _ = client.release(held.split(16)); + assert!(take(&mut client), "still stalled after a release"); + + // A release past the debt relieves it. + let _ = client.release(held.split(32)); + assert!(!take(&mut client)); + let _ = client.release(held); + } + + /// Bytes read without holding return their credit, so they never count as held or unread. + #[test] + fn plain_reads_are_neither_held_nor_unread() { + let mut client = make(Side::Client); + let id = StreamId::new(Side::Server, Dir::Uni, 0); + data(&mut client, id, 0, 32, false); + assert_eq!(client.credit_stats().unread, 32); + + let mut pending = Retransmits::default(); + let mut recv = RecvStream { + id, + state: &mut client, + pending: &mut pending, + }; + let mut chunks = recv.read(true).unwrap(); + while let Ok(Some(_)) = chunks.next(usize::MAX) {} + let _ = chunks.finalize(); + + let stats = client.credit_stats(); + assert_eq!((stats.held, stats.unread), (0, 0)); + } + /// A group token indexes its own connection's schedule, so another connection refuses it /// rather than reach whichever of its own groups shares the index. #[test] diff --git a/rs/moq-relay/src/config.rs b/rs/moq-relay/src/config.rs index f9fea858c21..80d5f945d7e 100644 --- a/rs/moq-relay/src/config.rs +++ b/rs/moq-relay/src/config.rs @@ -926,7 +926,7 @@ congestion_control = "delay" let toml = r#" [quic] -receive_window = 67108864 +receive_window = 134217728 stream_receive_window = 8388608 send_window = 33554432 "#; @@ -938,7 +938,7 @@ send_window = 33554432 let args = vec![std::ffi::OsString::from("moq-relay"), std::ffi::OsString::from(&path)]; let config = Config::parse_and_merge(args).expect("config load"); - assert_eq!(config.quic.receive_window, Some(67108864)); + assert_eq!(config.quic.receive_window, Some(134217728)); assert_eq!(config.quic.stream_receive_window, Some(8388608)); assert_eq!( config.quic.send_window, diff --git a/rs/moq-relay/src/uring.rs b/rs/moq-relay/src/uring.rs index bcca3f32614..2d8cf48bdb4 100644 --- a/rs/moq-relay/src/uring.rs +++ b/rs/moq-relay/src/uring.rs @@ -410,6 +410,15 @@ fn transport(quic: &moq_tokio::quic::Resolved) -> anyhow::Result= moq_net::transport::MIN_RECEIVE_WINDOW, + "quic.receive_window {window} cannot fit the largest message, {} bytes", + moq_net::transport::MIN_RECEIVE_WINDOW + ); + } let mut transport = moq_uring::quic::Transport::default(); #[cfg(feature = "qlog")] @@ -863,16 +872,33 @@ mod tests { #[allow(clippy::field_reassign_with_default)] fn windows_are_applied() { let mut quic = moq_tokio::quic::Config::default(); - quic.receive_window = Some(64 << 20); + quic.receive_window = Some(128 << 20); quic.stream_receive_window = Some(8 << 20); quic.send_window = Some(32 << 20); let transport = transport(&quic.resolve()).expect("windows must be accepted"); - assert_eq!(transport.receive_window, Some(64 << 20)); + assert_eq!(transport.receive_window, Some(128 << 20)); assert_eq!(transport.stream_receive_window, Some(8 << 20)); assert_eq!(transport.send_window, Some(32 << 20)); } + /// Incomplete messages hold connection credit, so the window has to fit the largest. + #[test] + #[allow(clippy::field_reassign_with_default)] + fn a_receive_window_below_the_largest_message_is_refused() { + let min = moq_net::transport::MIN_RECEIVE_WINDOW; + let mut quic = moq_tokio::quic::Config::default(); + quic.receive_window = Some(min - 1); + let err = transport(&quic.resolve()).expect_err("below the largest message"); + assert!(err.to_string().contains("largest message"), "{err}"); + + quic.receive_window = Some(min); + assert_eq!( + transport(&quic.resolve()).expect("the minimum").receive_window, + Some(min) + ); + } + /// Leaving the windows unset is the ordinary case. The worker then keeps /// its own credits rather than inheriting another backend's default. #[test] diff --git a/rs/moq-relay/tests/drills.rs b/rs/moq-relay/tests/drills.rs index bce5ce6144f..3e09794d4b0 100644 --- a/rs/moq-relay/tests/drills.rs +++ b/rs/moq-relay/tests/drills.rs @@ -794,7 +794,7 @@ const SETTLE: Duration = Duration::from_secs(10); /// The cross-relay drill's link: the drills' impairment behind a bottleneck a /// burst overruns, so its tail waits in the queue or is dropped from it. It is -/// slower than [`credit`]'s receive window lets a sender go. The first burst +/// slower than [`credit`]'s stream window lets a sender go. The first burst /// is collected at ingress after setup, so scheduling cannot drain the bucket /// between its datagrams; its surviving bytes exceed the burst and 50ms queue /// allowance before they are accounted at one instant. @@ -811,18 +811,16 @@ fn bursty() -> moq_shaper::Profile { /// QUIC credit tighter than a burst on the impaired lane: fewer concurrent /// streams than a burst has groups (each group is a stream, each FETCH another), -/// and receive windows of a few groups, so senders wait on MAX_STREAMS and -/// MAX_DATA mid-burst, over a path that loses the frames granting them. +/// and stream windows of a few groups, so senders wait on MAX_STREAMS and +/// MAX_STREAM_DATA mid-burst, over a path that loses the frames granting them. +/// The connection window keeps its default, since it has to fit the largest +/// message a session assembles. fn credit(lane: Lane, mut quic: moq_tokio::quic::Config) -> moq_tokio::quic::Config { if let Lane::Impaired = lane { quic.max_streams = Some(16); // The one-shot ingress burst must fit inside transport credit before - // anything is released. Keep two bucket allowances, rounded to groups, - // while still forcing a full application burst to wait for more credit. - let receive_window = ingress_bytes(); - assert!(receive_window < BURST_GROUPS * PAYLOAD as u64); - quic.receive_window = Some(receive_window); - quic.stream_receive_window = Some(receive_window); + // anything is released. Keep two bucket allowances, rounded to groups. + quic.stream_receive_window = Some(ingress_bytes()); } quic } diff --git a/rs/moq-tokio/src/error.rs b/rs/moq-tokio/src/error.rs index 7cca9e4b60a..d9606ae871c 100644 --- a/rs/moq-tokio/src/error.rs +++ b/rs/moq-tokio/src/error.rs @@ -105,6 +105,11 @@ pub enum Error { #[error("quic {0} must be non-zero and fit a QUIC varint")] WindowRange(&'static str), + /// The QUIC connection receive window can't fit the largest message a session + /// assembles. The payload is the minimum, in bytes. + #[error("quic receive_window must be at least {0} bytes to fit the largest message")] + ReceiveWindowTooSmall(u64), + /// Every backend we tried gave up without reporting why. #[error("failed to connect to server")] ConnectFailed, @@ -251,6 +256,7 @@ impl Error { | Self::Deprecated(_) | Self::IdleTimeoutRange | Self::WindowRange(_) + | Self::ReceiveWindowTooSmall(_) | Self::QlogUnsupported | Self::Connect(_) => true, #[cfg(feature = "quic")] diff --git a/rs/moq-tokio/src/iroh.rs b/rs/moq-tokio/src/iroh.rs index 058f33ca0c9..6e563179cd1 100644 --- a/rs/moq-tokio/src/iroh.rs +++ b/rs/moq-tokio/src/iroh.rs @@ -288,14 +288,14 @@ fn transport_config(quic: &crate::quic::Resolved) -> iroh::endpoint::QuicTranspo if !quic.mtu_discovery { transport = transport.mtu_discovery_config(None); } - // iroh's connection window is unlimited, so apply ours when unset. + // iroh's connection window is unlimited, so apply ours when unset, and the stream + // window follows it as on the native backend. let window = quic.receive_window.unwrap_or(crate::quic::DEFAULT_RECEIVE_WINDOW); + let stream = quic.stream_receive_window.unwrap_or(window); let window = iroh::endpoint::VarInt::from_u64(window).unwrap_or(iroh::endpoint::VarInt::MAX); + let stream = iroh::endpoint::VarInt::from_u64(stream).unwrap_or(iroh::endpoint::VarInt::MAX); transport = transport.receive_window(window); - if let Some(window) = quic.stream_receive_window { - let window = iroh::endpoint::VarInt::from_u64(window).unwrap_or(iroh::endpoint::VarInt::MAX); - transport = transport.stream_receive_window(window); - } + transport = transport.stream_receive_window(stream); if let Some(window) = quic.send_window { transport = transport.send_window(window); } @@ -500,7 +500,6 @@ mod tests { async fn bind_rejects_invalid_windows_before_reading_secret() { let dir = tempfile::tempdir().unwrap(); for (name, value) in [ - ("receive_window", 0), ("stream_receive_window", 0), ("send_window", 0), ("receive_window", 1 << 62), diff --git a/rs/moq-tokio/src/quic.rs b/rs/moq-tokio/src/quic.rs index f1e9d9fb739..223fc64cafb 100644 --- a/rs/moq-tokio/src/quic.rs +++ b/rs/moq-tokio/src/quic.rs @@ -107,11 +107,15 @@ pub const DEFAULT_MAX_STREAMS: u64 = 10_000; /// Default connection-wide receive window, used when [`Config::receive_window`] is unset. /// -/// Bounds how much unread data one peer can make us buffer across all of its streams, -/// which the backends otherwise leave unlimited. Sized so a relay-to-relay session, which -/// shares one connection, can still sustain several Gbps at intercontinental RTTs. +/// Bounds how much unread data and incomplete messages one peer can make us buffer across +/// all of its streams, which the backends otherwise leave unlimited. On iroh, which returns +/// credit as it reads, it bounds only the unread data. Sized to fit the +/// largest message with room to spare, so a relay-to-relay session, which shares one +/// connection, can still sustain several Gbps at intercontinental RTTs. #[cfg(any(feature = "quic", feature = "iroh"))] -pub(crate) const DEFAULT_RECEIVE_WINDOW: u64 = 64 << 20; +pub(crate) const DEFAULT_RECEIVE_WINDOW: u64 = 128 << 20; +#[cfg(any(feature = "quic", feature = "iroh"))] +const _: () = assert!(DEFAULT_RECEIVE_WINDOW >= moq_net::transport::MIN_RECEIVE_WINDOW); /// Default idle timeout before an inactive connection is dropped. /// @@ -231,7 +235,12 @@ pub struct Config { pub congestion_control: Option, /// Connection-wide receive window, in bytes: how much data the peer may have - /// in flight across every stream at once. Defaults to 64 MiB. + /// in flight or in incomplete messages across every stream at once. Defaults to + /// 128 MiB, and must fit the largest message a session assembles, just over 64 MiB. + /// + /// iroh returns credit as it reads, so there this bounds only unread data, and + /// incomplete messages are bounded by [`max_streams`](Self::max_streams) streams + /// of up to the largest message each. /// /// Raise it when a fat, long path is idling below the link rate: the window /// has to cover the bandwidth-delay product or the peer stalls waiting for @@ -245,10 +254,13 @@ pub struct Config { )] pub receive_window: Option, - /// Per-stream receive window, in bytes. Unset leaves the backend default. + /// Per-stream receive window, in bytes. Unset follows + /// [`receive_window`](Self::receive_window). /// - /// Keep it below [`receive_window`](Self::receive_window) so one slow group - /// can't starve the rest of the connection. + /// Stream credit returns as bytes are read, whether or not their message is + /// complete, so a smaller window only adds round trips to a large frame. The + /// connection window bounds what incomplete messages hold, and when they fill + /// it the least urgent stream is cancelled. #[serde(skip_serializing_if = "Option::is_none")] #[usage( name = "quic-stream-receive-window", @@ -607,9 +619,17 @@ const MAX_VARINT: u64 = (1 << 62) - 1; /// Zero is rejected everywhere: a window of zero credits nothing, so it would wedge /// the connection rather than tighten it. The receive windows additionally have to /// fit a varint, since they are sent as transport parameters; the send window is -/// local bookkeeping and only has to be non-zero. +/// local bookkeeping and only has to be non-zero. The connection receive window has +/// to fit the largest message, which holds its credit until it is complete. #[cfg(feature = "_transport")] fn validate_windows(quic: &Config) -> crate::Result<()> { + if let Some(window) = quic.receive_window + && window < moq_net::transport::MIN_RECEIVE_WINDOW + { + return Err(crate::Error::ReceiveWindowTooSmall( + moq_net::transport::MIN_RECEIVE_WINDOW, + )); + } for (name, window, wire) in [ ("receive_window", quic.receive_window, true), ("stream_receive_window", quic.stream_receive_window, true), @@ -655,9 +675,9 @@ pub struct Resolved { pub mtu_discovery: bool, /// Congestion control override, or `None` for the default. pub congestion_control: Option, - /// Connection-wide receive window in bytes, or `None` for the 64 MiB default. + /// Connection-wide receive window in bytes, or `None` for the 128 MiB default. pub receive_window: Option, - /// Per-stream receive window in bytes, or `None` for the backend default. + /// Per-stream receive window in bytes, or `None` to follow the connection window. pub stream_receive_window: Option, /// Cap on unacknowledged outgoing data in bytes, or `None` for the backend default. pub send_window: Option, @@ -922,13 +942,6 @@ mod tests { #[test] fn zero_windows_are_rejected() { for (name, config) in [ - ( - "receive_window", - Config { - receive_window: Some(0), - ..Default::default() - }, - ), ( "stream_receive_window", Config { @@ -981,6 +994,29 @@ mod tests { assert!(at_limit.validate().is_ok()); } + /// Incomplete messages hold connection credit, so a connection window that can't fit + /// the largest one could stall it forever. + #[cfg(feature = "_transport")] + #[test] + fn the_receive_window_fits_the_largest_message() { + let min = moq_net::transport::MIN_RECEIVE_WINDOW; + for window in [0, min - 1] { + let config = Config { + receive_window: Some(window), + ..Default::default() + }; + assert!(matches!( + config.validate(), + Err(crate::Error::ReceiveWindowTooSmall(got)) if got == min + )); + } + let config = Config { + receive_window: Some(min), + ..Default::default() + }; + assert!(config.validate().is_ok()); + } + #[test] fn window_flags_parse() { let quic = parse(&[ diff --git a/rs/moq-tokio/src/quic/backend.rs b/rs/moq-tokio/src/quic/backend.rs index 1dbb82af289..442d56385ad 100644 --- a/rs/moq-tokio/src/quic/backend.rs +++ b/rs/moq-tokio/src/quic/backend.rs @@ -74,15 +74,15 @@ fn apply_transport(transport: &mut driver::TransportConfig, quic: &Resolved) { } /// Apply the flow-control windows. The connection window defaults to ours, since the -/// backend's is unlimited; the others stay at the backend default when unset. +/// backend's is unlimited, and the stream window follows it; the send window stays at +/// the backend default when unset. fn apply_windows(transport: &mut driver::TransportConfig, quic: &Resolved) { // Saturating rather than erroring: `Config::validate` already rejects a window // past the varint, so this only bites a `Resolved` built without it. let window = quic.receive_window.unwrap_or(crate::quic::DEFAULT_RECEIVE_WINDOW); + let stream = quic.stream_receive_window.unwrap_or(window); transport.receive_window(driver::VarInt::from_u64(window).unwrap_or(driver::VarInt::MAX)); - if let Some(window) = quic.stream_receive_window { - transport.stream_receive_window(driver::VarInt::from_u64(window).unwrap_or(driver::VarInt::MAX)); - } + transport.stream_receive_window(driver::VarInt::from_u64(stream).unwrap_or(driver::VarInt::MAX)); if let Some(window) = quic.send_window { transport.send_window(window); } @@ -1064,7 +1064,8 @@ mod tests { /// moq-quic exposes no getters for the flow-control windows, but its `Debug` prints /// them, which is enough to prove each one reached the transport config, that an /// unset connection window gets our finite default instead of moq-quic's unlimited one, - /// and that the other unset knobs leave moq-quic's own defaults in place. + /// that an unset stream window follows it, and that an unset send window leaves + /// moq-quic's own default in place. #[test] fn apply_windows_writes_each_field() { let quic = crate::quic::Config { @@ -1089,7 +1090,9 @@ mod tests { let mut untouched = driver::TransportConfig::default(); apply_windows(&mut untouched, &crate::quic::Config::default().resolve()); let mut expected = driver::TransportConfig::default(); - expected.receive_window(driver::VarInt::from_u64(crate::quic::DEFAULT_RECEIVE_WINDOW).unwrap()); + let window = driver::VarInt::from_u64(crate::quic::DEFAULT_RECEIVE_WINDOW).unwrap(); + expected.receive_window(window); + expected.stream_receive_window(window); assert_eq!(format!("{:?}", untouched), format!("{:?}", expected)); assert_ne!( format!("{:?}", untouched), diff --git a/rs/moq-tokio/src/quic/driver/connection.rs b/rs/moq-tokio/src/quic/driver/connection.rs index 8de19c76139..6d12988a581 100644 --- a/rs/moq-tokio/src/quic/driver/connection.rs +++ b/rs/moq-tokio/src/quic/driver/connection.rs @@ -2,7 +2,7 @@ use std::{ any::Any, fmt, future::Future, - io, + io, mem, net::{IpAddr, SocketAddr}, pin::Pin, sync::{ @@ -647,6 +647,25 @@ impl Connection { self.0.state.lock("stats").inner.stats() } + /// Poll until held [`Credit`](crate::quic::driver::Credit) fills the receive window: the peer + /// may send nothing more and every byte it sent is held, so no stream makes progress until + /// some credit is dropped. + /// + /// Resolves once per stall, and again only after more credit is held or credit drops without + /// relieving it. Only the most recent task to poll is woken. + pub fn poll_credit_exhausted(&self, cx: &mut Context<'_>) -> Poll<()> { + let mut conn = self.0.state.lock("poll_credit_exhausted"); + // Credit may have dropped since the stall was noted. + if mem::take(&mut conn.credit_exhausted) && conn.inner.stats().credit.exhausted() { + return Poll::Ready(()); + } + match conn.credit_waker.as_mut() { + Some(waker) => waker.clone_from(cx.waker()), + None => conn.credit_waker = Some(cx.waker().clone()), + } + Poll::Pending + } + /// Current state of the congestion control algorithm, for debugging purposes pub fn congestion_state(&self) -> Box { self.0 @@ -1059,6 +1078,11 @@ pub(crate) struct State { pub(crate) blocked_writers: FxHashMap, pub(crate) blocked_readers: FxHashMap, pub(crate) stopped: FxHashMap>, + /// Held credit filled the receive window since [`Connection::poll_credit_exhausted`] last + /// resolved + credit_exhausted: bool, + /// The task watching for held credit to fill the receive window + credit_waker: Option, /// Always set to Some before the connection becomes drained pub(crate) error: Option, sender: Pin>, @@ -1094,6 +1118,8 @@ impl State { blocked_writers: FxHashMap::default(), blocked_readers: FxHashMap::default(), stopped: FxHashMap::default(), + credit_exhausted: false, + credit_waker: None, error: None, sender, send_buffer: Vec::new(), @@ -1262,6 +1288,9 @@ impl State { _ => {} } } + // Data for a stopped stream or a reset returns its credit, possibly without relieving + // a stall. + self.check_credit(); } fn drive_timer(&mut self, cx: &mut Context<'_>) -> bool { @@ -1303,6 +1332,18 @@ impl State { true } + /// Note when held credit comes to fill the receive window, waking the task watching for it. + /// + /// Call after anything that holds, returns, or discards connection credit. + pub(crate) fn check_credit(&mut self) { + if self.inner.poll_credit_exhausted() { + self.credit_exhausted = true; + if let Some(waker) = self.credit_waker.take() { + waker.wake(); + } + } + } + /// Wake up a blocked `Driver` task to process I/O pub(crate) fn wake(&mut self) { if let Some(x) = self.driver.take() { diff --git a/rs/moq-tokio/src/quic/driver/mod.rs b/rs/moq-tokio/src/quic/driver/mod.rs index a20c51650ff..cff743d15e5 100644 --- a/rs/moq-tokio/src/quic/driver/mod.rs +++ b/rs/moq-tokio/src/quic/driver/mod.rs @@ -74,7 +74,7 @@ pub use crate::quic::driver::connection::{ pub use crate::quic::driver::endpoint::{Accept, Endpoint, EndpointStats}; pub use crate::quic::driver::incoming::{Incoming, IncomingFuture, RetryError}; pub use crate::quic::driver::recv_stream::{ - Credit, ReadError, ReadExactError, ReadToEndError, RecvStream, ResetError, + Credit, ReadError, ReadExactError, ReadToEndError, RecvStream, ResetError, Stopper, }; #[cfg(test)] pub(crate) use crate::quic::driver::runtime::wrap_udp_socket; diff --git a/rs/moq-tokio/src/quic/driver/recv_stream.rs b/rs/moq-tokio/src/quic/driver/recv_stream.rs index bde840c96c6..f9d89289ade 100644 --- a/rs/moq-tokio/src/quic/driver/recv_stream.rs +++ b/rs/moq-tokio/src/quic/driver/recv_stream.rs @@ -75,15 +75,14 @@ impl RecvStream { /// Read data contiguously from the stream. /// - /// Yields the number of bytes read into `buf` on success, or `None` if the stream was finished. + /// Yields the [`Credit`] for exactly the bytes read into `buf`, which [`Credit::len()`] + /// counts, or `None` if the stream was finished. Stream credit advances on the read, while + /// the connection window these bytes occupy stays spent until the credit is dropped, so drop + /// it once the bytes are consumed. /// /// This operation is cancel-safe. - pub async fn read(&mut self, buf: &mut [u8]) -> Result, ReadError> { - Read { - stream: self, - buf: ReadBuf::new(buf), - } - .await + pub async fn read(&mut self, buf: &mut [u8]) -> Result, ReadError> { + poll_fn(|cx| self.poll_read(cx, buf)).await } /// Read an exact number of bytes contiguously from the stream. @@ -99,27 +98,16 @@ impl RecvStream { .await } - /// Read data contiguously from the stream, keeping its connection flow control credit. + /// Attempts to read from the stream into the provided buffer /// - /// Like [`read()`], but yields the [`Credit`] for exactly the bytes read into `buf`, or `None` - /// if the stream was finished. Stream credit still advances, so a stream window smaller than - /// an application message keeps the message flowing, while the connection window these bytes - /// occupy stays spent until the credit is dropped. + /// See [`read()`] for what the yielded credit means. /// - /// This operation is cancel-safe. + /// If no data is available for reading, this returns `Poll::Pending` and arranges for the + /// current task (via `cx.waker()`) to be notified when the stream becomes readable or is + /// closed. /// /// [`read()`]: RecvStream::read - pub async fn read_held(&mut self, buf: &mut [u8]) -> Result, ReadError> { - poll_fn(|cx| self.poll_read_held(cx, buf)).await - } - - /// Attempts to read from the stream into the provided buffer, keeping the bytes' connection - /// flow control credit. - /// - /// See [`read_held()`] for details. - /// - /// [`read_held()`]: RecvStream::read_held - pub fn poll_read_held( + pub fn poll_read( &mut self, cx: &mut Context<'_>, buf: &mut [u8], @@ -143,33 +131,12 @@ impl RecvStream { }) } - /// Attempts to read from the stream into the provided buffer - /// - /// On success, returns `Poll::Ready(Ok(num_bytes_read))` and places data into `buf`. If this - /// returns zero bytes read (and `buf` has a non-zero length), that indicates that the remote - /// side has [`finish`]ed the stream and the local side has already read all bytes. - /// - /// If no data is available for reading, this returns `Poll::Pending` and arranges for the - /// current task (via `cx.waker()`) to be notified when the stream becomes readable or is - /// closed. - /// - /// [`finish`]: crate::quic::driver::SendStream::finish - pub fn poll_read( - &mut self, - cx: &mut Context<'_>, - buf: &mut [u8], - ) -> Poll> { - let mut buf = ReadBuf::new(buf); - ready!(self.poll_read_buf(cx, &mut buf))?; - Poll::Ready(Ok(buf.filled().len())) - } - /// Attempts to read from the stream into the provided buffer, which may be uninitialized /// /// On success, returns `Poll::Ready(Ok(()))` and places data into the unfilled portion of /// `buf`. If this does not write any bytes to `buf` (and `buf.remaining()` is non-zero), that /// indicates that the remote side has [`finish`]ed the stream and the local side has already - /// read all bytes. + /// read all bytes. The bytes' connection credit returns as they are read. /// /// If no data is available for reading, this returns `Poll::Pending` and arranges for the /// current task (via `cx.waker()`) to be notified when the stream becomes readable or is @@ -192,30 +159,40 @@ impl RecvStream { /// Read the next segment of data /// /// Yields `None` if the stream was finished. Otherwise, yields a segment of data and its - /// offset in the stream. If `ordered` is `true`, the chunk's offset will be immediately after - /// the last data yielded by `read()` or `read_chunk()`. If `ordered` is `false`, segments may - /// be received in any order, and the `Chunk`'s `offset` field can be used to determine - /// ordering in the caller. Unordered reads are less prone to head-of-line blocking within a - /// stream, but require the application to manage reassembling the original data. + /// offset in the stream, with the [`Credit`] for its bytes (see [`read()`]). If `ordered` is + /// `true`, the chunk's offset will be immediately after the last data yielded by `read()` or + /// `read_chunk()`. If `ordered` is `false`, segments may be received in any order, and the + /// `Chunk`'s `offset` field can be used to determine ordering in the caller. Unordered reads + /// are less prone to head-of-line blocking within a stream, but require the application to + /// manage reassembling the original data. /// /// Slightly more efficient than `read` due to not copying. Chunk boundaries do not correspond /// to peer writes, and hence cannot be used as framing. /// /// This operation is cancel-safe. + /// + /// [`read()`]: RecvStream::read pub async fn read_chunk( &mut self, max_length: usize, ordered: bool, - ) -> Result, ReadError> { - ReadChunk { - stream: self, - max_length, - ordered, - } + ) -> Result, ReadError> { + poll_fn(|cx| { + self.poll_read_generic(cx, ordered, |chunks| match chunks.next(max_length) { + Ok(Some(chunk)) => ReadStatus::Readable((chunk, chunks.hold())), + res => (None, res.err()).into(), + }) + .map_ok(|read| { + read.map(|(chunk, inner)| { + let conn = self.conn.0.clone(); + (chunk, Credit { conn, inner }) + }) + }) + }) .await } - /// Attempts to read a chunk from the stream. + /// Attempts to read a chunk from the stream, returning its credit as it reads. /// /// On success, returns `Poll::Ready(Ok(Some(chunk)))`. If `Poll::Ready(Ok(None))` /// is returned, it implies that EOF has been reached. @@ -223,7 +200,7 @@ impl RecvStream { /// If no data is available for reading, the method returns `Poll::Pending` /// and arranges for the current task (via cx.waker()) to receive a notification /// when the stream becomes readable or is closed. - fn poll_read_chunk( + pub(super) fn poll_read_chunk( &mut self, cx: &mut Context<'_>, max_length: usize, @@ -308,6 +285,7 @@ impl RecvStream { return Ok(()); } conn.inner.recv_stream(self.stream).stop(error_code)?; + conn.check_credit(); conn.wake(); self.all_data_read = true; // Clean up shared state that might be left over from a cancalled read @@ -316,6 +294,15 @@ impl RecvStream { Ok(()) } + /// A handle that stops this stream as [`stop`](Self::stop) does, without borrowing it. + pub fn stopper(&self) -> Stopper { + Stopper { + conn: self.conn.0.clone(), + stream: self.stream, + is_0rtt: self.is_0rtt, + } + } + /// Check if this stream predates completion of the handshake on an incoming connection. /// /// True only if the stream was accepted before the handshake completed, which is only possible @@ -416,6 +403,7 @@ impl RecvStream { if chunks.finalize().should_transmit() { conn.wake(); } + conn.check_credit(); status } }; @@ -481,7 +469,8 @@ fn fill(chunks: &mut Chunks<'_>, buf: &mut ReadBuf<'_>) -> ReadStatus<()> { } } -/// Connection flow control credit for bytes read with [`RecvStream::read_held`]. +/// Connection flow control credit for bytes read with [`RecvStream::read`] or +/// [`RecvStream::read_chunk`]. /// /// The peer cannot reuse the connection window these bytes occupy until this is dropped, which /// returns it even after the stream or connection is gone. @@ -538,6 +527,7 @@ impl Drop for Credit { { conn.wake(); } + conn.check_credit(); } } @@ -645,6 +635,45 @@ impl tokio::io::AsyncRead for RecvStream { } } +/// Stops a [`RecvStream`] without borrowing it. +/// +/// It does not keep the connection open: only [`RecvStream`] and the other handles count. +#[derive(Clone)] +pub struct Stopper { + conn: Arc, + stream: StreamId, + is_0rtt: bool, +} + +impl Stopper { + /// Stop accepting data, as [`RecvStream::stop`] does. + /// + /// The stream's own handle learns of it on its next read, which yields `ClosedStream`; + /// a read parked on the stream is woken for it. + pub fn stop(&self, error_code: VarInt) -> Result<(), ClosedStream> { + let mut conn = self.conn.state.lock("Stopper::stop"); + if self.is_0rtt && conn.check_0rtt().is_err() { + return Ok(()); + } + conn.inner.recv_stream(self.stream).stop(error_code)?; + conn.check_credit(); + // A stop produces no readable event, so the parked read would never hear of it. + if let Some(waker) = conn.blocked_readers.remove(&self.stream) { + waker.wake(); + } + conn.wake(); + Ok(()) + } +} + +impl fmt::Debug for Stopper { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Stopper") + .field("stream", &self.stream) + .finish() + } +} + impl Drop for RecvStream { fn drop(&mut self) { if self.all_data_read { @@ -764,27 +793,6 @@ impl From for io::Error { } } -/// Future produced by [`RecvStream::read()`]. -/// -/// [`RecvStream::read()`]: crate::quic::driver::RecvStream::read -struct Read<'a> { - stream: &'a mut RecvStream, - buf: ReadBuf<'a>, -} - -impl Future for Read<'_> { - type Output = Result, ReadError>; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.get_mut(); - ready!(this.stream.poll_read_buf(cx, &mut this.buf))?; - match this.buf.filled().len() { - 0 if this.buf.capacity() != 0 => Poll::Ready(Ok(None)), - n => Poll::Ready(Ok(Some(n))), - } - } -} - /// Future produced by [`RecvStream::read_exact()`]. /// /// [`RecvStream::read_exact()`]: crate::quic::driver::RecvStream::read_exact @@ -822,23 +830,6 @@ pub enum ReadExactError { ReadError(#[from] ReadError), } -/// Future produced by [`RecvStream::read_chunk()`]. -/// -/// [`RecvStream::read_chunk()`]: crate::quic::driver::RecvStream::read_chunk -struct ReadChunk<'a> { - stream: &'a mut RecvStream, - max_length: usize, - ordered: bool, -} - -impl Future for ReadChunk<'_> { - type Output = Result, ReadError>; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let (max_length, ordered) = (self.max_length, self.ordered); - self.stream.poll_read_chunk(cx, max_length, ordered) - } -} - /// Future produced by [`RecvStream::read_chunks()`]. /// /// [`RecvStream::read_chunks()`]: crate::quic::driver::RecvStream::read_chunks diff --git a/rs/moq-tokio/src/quic/driver/tests.rs b/rs/moq-tokio/src/quic/driver/tests.rs index d7846d30e9f..409908d2f55 100755 --- a/rs/moq-tokio/src/quic/driver/tests.rs +++ b/rs/moq-tokio/src/quic/driver/tests.rs @@ -1770,6 +1770,48 @@ async fn stream_drop_removes_blocked_reader() { } } +/// A read parked on a stream wakes when a `Stopper` stops it, and fails as closed. +#[tokio::test] +async fn stopper_wakes_a_parked_read() { + let _guard = subscribe(); + let endpoint_factory = EndpointFactory::new(); + let server = endpoint_factory.endpoint(); + let server_address = server.local_addr().unwrap(); + let client = endpoint_factory.endpoint(); + + let server_task = tokio::spawn(async move { + let conn = server.accept().await.unwrap().await.unwrap(); + let mut stream = conn.accept_uni().await.unwrap(); + let mut buf = [0u8; 5]; + stream.read_exact(&mut buf).await.unwrap(); + + let stopper = stream.stopper(); + let (waker, wake_counter) = new_count_waker(); + let mut cx = Context::from_waker(&waker); + let mut buf = [0u8; 64]; + let read_fut = stream.read(&mut buf); + tokio::pin!(read_fut); + assert!(matches!(read_fut.as_mut().poll(&mut cx), Poll::Pending)); + + stopper.stop(0u32.into()).unwrap(); + assert_eq!(wake_counter.wakes(), 1, "the parked read was not woken"); + assert!(matches!( + read_fut.as_mut().poll(&mut cx), + Poll::Ready(Err(crate::quic::driver::ReadError::ClosedStream)) + )); + conn.close(0u32.into(), b"done"); + }); + + let conn = client + .connect(server_address, "localhost") + .unwrap() + .await + .unwrap(); + let mut stream = conn.open_uni().await.unwrap(); + stream.write_all(b"hello").await.unwrap(); + server_task.await.unwrap(); +} + /// Test that dropping a `RecvStream` after cancelling a read and then /// explicitly `stop`ing it doesn't panic. #[tokio::test] @@ -1902,12 +1944,129 @@ async fn held_reads_spend_the_connection_window() { writer.await.unwrap(); } +/// Held chunks filling the window resolve `poll_credit_exhausted` once, and plain reads after the +/// credit drops never do. +#[tokio::test] +async fn filling_the_window_with_held_credit_is_reported_once() { + const WINDOW: usize = 32 * 1024; + let _guard = subscribe(); + let mut transport = TransportConfig::default(); + transport.receive_window((WINDOW as u32).into()); + let endpoint = endpoint_with_config(transport); + + let (client, server) = tokio::try_join!( + endpoint + .connect(endpoint.local_addr().unwrap(), "localhost") + .unwrap(), + async { endpoint.accept().await.unwrap().await } + ) + .unwrap(); + let mut send = client.open_uni().await.unwrap(); + let writer = tokio::spawn(async move { + // The receiver walks away before this finishes. + let _ = send.write_all(&[7; 4 * WINDOW]).await; + }); + + let (waker, counter) = new_count_waker(); + let mut cx = Context::from_waker(&waker); + assert!(server.poll_credit_exhausted(&mut cx).is_pending()); + + let mut recv = server.accept_uni().await.unwrap(); + let mut held = recv.read_chunk(usize::MAX, true).await.unwrap().unwrap().1; + while (held.len() as usize) < WINDOW { + held.merge(recv.read_chunk(usize::MAX, true).await.unwrap().unwrap().1); + } + assert_eq!(counter.wakes(), 1); + assert!(server.poll_credit_exhausted(&mut cx).is_ready()); + assert!(server.poll_credit_exhausted(&mut cx).is_pending()); + + drop(held); + let mut buf = vec![0; WINDOW]; + recv.read_exact(&mut buf).await.unwrap(); + assert!(server.poll_credit_exhausted(&mut cx).is_pending()); + assert_eq!(counter.wakes(), 1); + drop(recv); + writer.await.unwrap(); +} + +/// Stream 256 MiB over loopback, reading chunks whose credit returns inside the read, or drops +/// straight after it as `read_chunk` does: the per-read cost of handing credit to the caller. +/// Run by nightly. +#[tokio::test(flavor = "multi_thread")] +#[ignore = "release benchmark, run by nightly"] +async fn benchmark_held_reads() { + const BYTES: usize = 256 << 20; + const RUNS: usize = 3; + let _guard = subscribe(); + println!("max_chunk,reads,plain_ns_per_read,held_ns_per_read"); + for max in [1200usize, 16 * 1024, 64 * 1024] { + let mut ns = [f64::MAX; 2]; + let mut reads = 0; + // Alternate the modes so drift in the host hits both alike; keep each best run. + for _ in 0..RUNS { + for (mode, held) in [false, true].into_iter().enumerate() { + let endpoint = endpoint(); + let (client, server) = tokio::try_join!( + endpoint + .connect(endpoint.local_addr().unwrap(), "localhost") + .unwrap(), + async { endpoint.accept().await.unwrap().await } + ) + .unwrap(); + let mut send = client.open_uni().await.unwrap(); + let writer = tokio::spawn(async move { + let block = vec![7u8; 1 << 20]; + for _ in 0..BYTES / block.len() { + send.write_all(&block).await.unwrap(); + } + send.finish().unwrap(); + // Keep the connection up until the reader is done. + let _ = send.stopped().await; + }); + + let mut recv = server.accept_uni().await.unwrap(); + let (mut read, mut count) = (0, 0); + let start = std::time::Instant::now(); + loop { + let len = if held { + match recv.read_chunk(max, true).await.unwrap() { + Some((chunk, credit)) => { + drop(credit); + chunk.bytes.len() + } + None => break, + } + } else { + // `read_to_end`'s path, which returns the credit before handing over. + match std::future::poll_fn(|cx| recv.poll_read_chunk(cx, max, true)) + .await + .unwrap() + { + Some(chunk) => chunk.bytes.len(), + None => break, + } + }; + read += len; + count += 1; + } + let elapsed = start.elapsed().as_nanos() as f64 / count as f64; + assert_eq!(read, BYTES); + ns[mode] = ns[mode].min(elapsed); + reads = count; + drop(recv); + writer.await.unwrap(); + } + } + println!("{max},{reads},{:.1},{:.1}", ns[0], ns[1]); + } +} + /// Fill `buf` with held reads, returning the credit for all of it. async fn read_held_exact(recv: &mut RecvStream, buf: &mut [u8]) -> Credit { - let mut held = recv.read_held(buf).await.unwrap().unwrap(); + let mut held = recv.read(buf).await.unwrap().unwrap(); while (held.len() as usize) < buf.len() { let read = held.len() as usize; - held.merge(recv.read_held(&mut buf[read..]).await.unwrap().unwrap()); + held.merge(recv.read(&mut buf[read..]).await.unwrap().unwrap()); } held } diff --git a/rs/moq-tokio/src/quic/web/recv.rs b/rs/moq-tokio/src/quic/web/recv.rs index 1e46bac784b..27412b41646 100644 --- a/rs/moq-tokio/src/quic/web/recv.rs +++ b/rs/moq-tokio/src/quic/web/recv.rs @@ -16,6 +16,20 @@ pub struct RecvStream { close: Arc, } +/// Stops a [`RecvStream`] without borrowing it. +#[derive(Clone)] +pub struct Stopper { + inner: crate::quic::driver::Stopper, + close: Arc, +} + +impl Stopper { + /// Stop accepting data, as [`RecvStream::stop`] does. + pub fn stop(&self, code: u32) -> Result<(), crate::quic::driver::ClosedStream> { + self.inner.stop(self.close.encode_stream_code(code)) + } +} + impl RecvStream { pub(crate) fn new(stream: crate::quic::driver::RecvStream, close: Arc) -> Self { Self { inner: stream, close } @@ -43,10 +57,19 @@ impl RecvStream { self.inner.stop(code) } + /// A handle that stops this stream as [`stop`](Self::stop) does, without borrowing it. + pub fn stopper(&self) -> Stopper { + Stopper { + inner: self.inner.stopper(), + close: self.close.clone(), + } + } + // Unfortunately, we have to wrap ReadError for a bunch of functions. - /// Read some data into the buffer and return the amount read. See [`crate::quic::driver::RecvStream::read`]. - pub async fn read(&mut self, buf: &mut [u8]) -> Result, ReadError> { + /// Read some data into the buffer, returning the credit for the bytes read. See + /// [`crate::quic::driver::RecvStream::read`]. + pub async fn read(&mut self, buf: &mut [u8]) -> Result, ReadError> { self.inner.read(buf).await.map_err(|e| self.map_error(e)) } @@ -58,12 +81,16 @@ impl RecvStream { }) } - /// Read a chunk of data from the stream. See [`crate::quic::driver::RecvStream::read_chunk`]. - pub async fn read_chunk(&mut self, max_length: usize) -> Result, ReadError> { + /// Read a chunk of data from the stream, with the credit for its bytes. See + /// [`crate::quic::driver::RecvStream::read_chunk`]. + pub async fn read_chunk( + &mut self, + max_length: usize, + ) -> Result, ReadError> { self.inner .read_chunk(max_length, true) .await - .map(|chunk| chunk.map(|chunk| chunk.bytes)) + .map(|read| read.map(|(chunk, credit)| (chunk.bytes, credit))) .map_err(|e| self.map_error(e)) } @@ -116,12 +143,13 @@ impl web_transport_trait::RecvStream for RecvStream { Self::stop(self, code).ok(); } + // This interface has no credit, so it returns each read's as it reads. async fn read(&mut self, dst: &mut [u8]) -> Result, Self::Error> { - self.read(dst).await + Ok(self.read(dst).await?.map(|credit| credit.len() as usize)) } async fn read_chunk(&mut self, max: usize) -> Result, Self::Error> { - self.read_chunk(max).await + Ok(self.read_chunk(max).await?.map(|(bytes, _)| bytes)) } async fn closed(&mut self) -> Result<(), Self::Error> { diff --git a/rs/moq-tokio/src/server.rs b/rs/moq-tokio/src/server.rs index d9c33cd4646..4cac0d912f1 100644 --- a/rs/moq-tokio/src/server.rs +++ b/rs/moq-tokio/src/server.rs @@ -966,7 +966,7 @@ async fn accept_request( deadline: Option, ) -> crate::Result<(PendingRequest>, Option)> where - S: crate::transport::Grouped, + S: crate::transport::HeldSession + crate::transport::Grouped, crate::transport::Session: moq_net::transport::poll::Boxable, as moq_net::transport::poll::Session>::SendStream: moq_net::transport::MaybeSync, as moq_net::transport::poll::Session>::RecvStream: moq_net::transport::MaybeSync, @@ -2213,6 +2213,9 @@ mod tests { } } + #[cfg(feature = "websocket")] + crate::transport::held::unheld!(Scripted, ScriptedRecv); + #[cfg(feature = "websocket")] impl web_transport_trait::Session for Scripted { type SendStream = ScriptedSend; diff --git a/rs/moq-tokio/src/transport.rs b/rs/moq-tokio/src/transport.rs index e908eb41d18..cfe3617435d 100644 --- a/rs/moq-tokio/src/transport.rs +++ b/rs/moq-tokio/src/transport.rs @@ -14,13 +14,68 @@ use std::task::{Context, Poll, ready}; use bytes::Bytes; use futures::FutureExt; +use moq_net::transport::Credit as _; use web_transport_trait::poll as wt_poll; mod grouped; +pub(crate) mod held; mod owned; pub(crate) use grouped::{Grouped, GroupedSend}; pub use owned::Error; +/// A backend receive stream that can keep the connection credit of what it reads. +/// +/// A backend whose reads return their credit straight away reads plain chunks and counts them +/// as [`Unheld`](moq_net::transport::Unheld). +pub trait HeldRecvStream: web_transport_trait::RecvStream { + /// The connection credit a read keeps until it is dropped. + type Credit: moq_net::transport::Credit + Send + Sync; + + /// What [`stopper`](Self::stopper) returns. + type Stopper: moq_net::transport::Stopper + Clone + Send + Sync; + + /// Read the next chunk, up to `max` bytes, with its credit, or `None` once finished. + fn read_chunk_held( + &mut self, + max: usize, + ) -> impl Future, Self::Error>> + Send; + + /// A handle that stops this stream without borrowing it. + fn stopper(&self) -> Self::Stopper; +} + +/// A backend session that reports when held credit fills its receive window. +pub trait HeldSession: web_transport_trait::Session { + /// See [`poll_credit_exhausted`](moq_net::transport::poll::Session::poll_credit_exhausted). + fn poll_credit_exhausted(&self, cx: &mut Context<'_>) -> Poll<()>; +} + +/// The connection credit for bytes read through a [`RecvStream`], none for an empty read. +#[derive(Debug)] +pub struct Credit(Option); + +impl moq_net::transport::Credit for Credit { + fn len(&self) -> u64 { + self.0.as_ref().map_or(0, |credit| credit.len()) + } + + fn split(&mut self, bytes: u64) -> Self { + match &mut self.0 { + _ if bytes == 0 => Self(None), + Some(credit) => Self(Some(credit.split(bytes))), + None => panic!("split {bytes} of 0 bytes"), + } + } + + fn merge(&mut self, other: Self) { + match (&mut self.0, other.0) { + (_, None) => {} + (Some(credit), Some(other)) => credit.merge(other), + (empty, other) => *empty = other, + } + } +} + /// The network transport carrying a MoQ session, on either side of it. #[non_exhaustive] #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] @@ -159,7 +214,7 @@ impl Clone for Session { impl wt_poll::Session for Session where - S: Grouped, + S: Grouped, S::SendStream: 'static, S::RecvStream: 'static, { @@ -622,16 +677,17 @@ impl Drop for SendStream { /// The receive half produced by [`Session`]: an async-interface stream adapted to /// the poll interface. See [`Session`] for the deferred `stop` behavior. -pub struct RecvStream { +pub struct RecvStream { // See [`SendStream::state`] for why this Mutex exists; it is never locked. state: std::sync::Mutex>>, - /// Bytes already read from the transport but not yet handed to the caller. - buffer: Bytes, + /// A chunk already read from the transport but not yet handed to the caller, never empty, + /// with the credit it keeps. + buffer: Option<(Bytes, S::Credit)>, /// Read-ahead pulled in by the closed-watch while payload sat undelivered, so /// a FIN right behind buffered data still resolves the watch. Drained by the /// read methods before the transport is polled again; bounded by - /// [`READ_AHEAD_CAP`]. - queued: std::collections::VecDeque, + /// [`READ_AHEAD_CAP`]. It keeps its credit until the caller reads it. + queued: std::collections::VecDeque<(Bytes, S::Credit)>, /// Total bytes across `queued`, so the cap check is O(1). queued_len: usize, /// The peer finished the stream. @@ -642,6 +698,8 @@ pub struct RecvStream { /// waiting behind blocked I/O. Fired by [`stop`](wt_poll::RecvStream::stop) /// and by [`Drop`]. interrupt: Option>, + /// Taken up front, since an in-flight read owns the stream. + stopper: S::Stopper, } /// How much the closed-watch may read ahead of the caller before it parks and @@ -649,19 +707,23 @@ pub struct RecvStream { /// stream nobody is reading. const READ_AHEAD_CAP: usize = 64 * 1024; -enum RecvState { +/// A chunk read from a backend with its credit, or `None` once finished. +type Held = Option<(Bytes, ::Credit)>; + +enum RecvState { Idle(S), /// A read in flight; a `None` result means it was interrupted by a stop /// (see [`RecvStream::interrupt`]). - Reading(#[allow(clippy::type_complexity)] OpBox<(S, Option, S::Error>>)>), + Reading(#[allow(clippy::type_complexity)] OpBox<(S, Option, S::Error>>)>), } -impl RecvStream { +impl RecvStream { /// Wrap an async-interface receive stream. pub fn new(stream: S) -> Self { Self { + stopper: stream.stopper(), state: std::sync::Mutex::new(Some(RecvState::Idle(stream))), - buffer: Bytes::new(), + buffer: None, queued: std::collections::VecDeque::new(), queued_len: 0, fin: false, @@ -670,18 +732,49 @@ impl RecvStream { } } - /// Pop read-ahead into `self.buffer` if the head is empty. - fn unqueue(&mut self) { - if self.buffer.is_empty() - && let Some(next) = self.queued.pop_front() - { - self.queued_len -= next.len(); - self.buffer = next; + /// Take up to `max` buffered bytes with their credit, read-ahead first. + fn take(&mut self, max: usize) -> Option<(Bytes, S::Credit)> { + if self.buffer.is_none() { + let next = self.queued.pop_front()?; + self.queued_len -= next.0.len(); + self.buffer = Some(next); + } + let (bytes, credit) = self.buffer.as_mut().expect("just filled"); + if bytes.len() <= max { + return self.buffer.take(); + } + Some((bytes.split_to(max), credit.split(max as u64))) + } + + /// Read up to `max` bytes with their credit, or `None` once the peer finished the stream. + #[allow(clippy::type_complexity)] + pub(crate) fn poll_read_held( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll)>, S::Error>> { + if max == 0 { + return Poll::Ready(Ok(Some((Bytes::new(), Credit(None))))); + } + + loop { + if let Some((bytes, credit)) = self.take(max) { + return Poll::Ready(Ok(Some((bytes, Credit(Some(credit)))))); + } + if self.fin { + return Poll::Ready(Ok(None)); + } + // An empty chunk carries nothing worth buffering. + if let Some(chunk) = ready!(self.poll_fill(cx, max))? + && !chunk.0.is_empty() + { + self.buffer = Some(chunk); + } } } /// Read the next chunk from the transport, setting `self.fin` on the FIN. - fn poll_fill(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, S::Error>> { + fn poll_fill(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, S::Error>> { loop { match self.state.get_mut().unwrap().take().expect("in-flight") { RecvState::Idle(mut stream) => { @@ -692,7 +785,7 @@ impl RecvStream { self.interrupt = Some(tx); let fut = async move { let res = { - let mut op = std::pin::pin!(stream.read_chunk(max).fuse()); + let mut op = std::pin::pin!(stream.read_chunk_held(max).fuse()); let mut rx = rx.fuse(); loop { futures::select_biased! { @@ -740,48 +833,20 @@ impl RecvStream { } } -impl wt_poll::RecvStream for RecvStream { +// This view returns connection credit as it reads; moq-net reads through the held one. +impl wt_poll::RecvStream for RecvStream { type Error = S::Error; fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { - if dst.is_empty() { - return Poll::Ready(Ok(Some(0))); - } - - loop { - self.unqueue(); - if !self.buffer.is_empty() { - let n = dst.len().min(self.buffer.len()); - dst[..n].copy_from_slice(&self.buffer.split_to(n)); - return Poll::Ready(Ok(Some(n))); - } - if self.fin { - return Poll::Ready(Ok(None)); - } - if let Some(bytes) = ready!(self.poll_fill(cx, dst.len()))? { - self.buffer = bytes; - } - } + let read = ready!(self.poll_read_held(cx, dst.len()))?; + Poll::Ready(Ok(read.map(|(bytes, _)| { + dst[..bytes.len()].copy_from_slice(&bytes); + bytes.len() + }))) } fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Self::Error>> { - if max == 0 { - return Poll::Ready(Ok(Some(Bytes::new()))); - } - - loop { - self.unqueue(); - if !self.buffer.is_empty() { - let n = max.min(self.buffer.len()); - return Poll::Ready(Ok(Some(self.buffer.split_to(n)))); - } - if self.fin { - return Poll::Ready(Ok(None)); - } - if let Some(bytes) = ready!(self.poll_fill(cx, max))? { - self.buffer = bytes; - } - } + self.poll_read_held(cx, max).map_ok(|read| read.map(|(bytes, _)| bytes)) } fn stop(&mut self, code: u32) { @@ -809,15 +874,15 @@ impl wt_poll::RecvStream for RecvS if self.fin { return Poll::Ready(Ok(())); } - if self.buffer.len() + self.queued_len >= READ_AHEAD_CAP { + if self.buffer.as_ref().map_or(0, |(bytes, _)| bytes.len()) + self.queued_len >= READ_AHEAD_CAP { // The caller's own reads are what drain the backlog, and their // re-poll of this watch is what resumes it. return Poll::Pending; } match ready!(self.poll_fill(cx, 8 * 1024)) { - Ok(Some(bytes)) => { + Ok(Some((bytes, credit))) => { self.queued_len += bytes.len(); - self.queued.push_back(bytes); + self.queued.push_back((bytes, credit)); } Ok(None) => {} // A read error is also a closed stream. @@ -827,7 +892,7 @@ impl wt_poll::RecvStream for RecvS } } -impl Drop for RecvStream { +impl Drop for RecvStream { fn drop(&mut self) { match self.state.get_mut().unwrap().take() { // Apply a deferred stop if the stream is idle. @@ -1422,6 +1487,7 @@ mod tests { #[derive(Clone)] struct FakeSession; + super::held::unheld!(FakeSession, FakeRecv); impl Grouped for FakeSession {} impl web_transport_trait::Session for FakeSession { diff --git a/rs/moq-tokio/src/transport/held.rs b/rs/moq-tokio/src/transport/held.rs new file mode 100644 index 00000000000..e3085dd408c --- /dev/null +++ b/rs/moq-tokio/src/transport/held.rs @@ -0,0 +1,97 @@ +//! How each backend reads with connection credit. + +/// Read a plain chunk, whose connection credit the backend returned as it read. +#[cfg(any(test, feature = "tcp", feature = "websocket", feature = "iroh"))] +pub(crate) async fn read_unheld( + stream: &mut S, + max: usize, +) -> Result, S::Error> { + let chunk = stream.read_chunk(max).await?; + Ok(chunk.map(|bytes| { + let len = bytes.len() as u64; + (bytes, moq_net::transport::Unheld(len)) + })) +} + +/// Give a backend that returns connection credit as it reads, so never fills its window with +/// held credit, the held interface. +#[cfg(any(test, feature = "tcp", feature = "websocket", feature = "iroh"))] +macro_rules! unheld { + ($session:ty, $recv:ty) => { + impl $crate::transport::HeldRecvStream for $recv { + type Credit = moq_net::transport::Unheld; + type Stopper = (); + + fn stopper(&self) {} + + fn read_chunk_held( + &mut self, + max: usize, + ) -> impl Future, Self::Error>> + Send { + $crate::transport::held::read_unheld(self, max) + } + } + + impl $crate::transport::HeldSession for $session { + fn poll_credit_exhausted(&self, _: &mut std::task::Context<'_>) -> std::task::Poll<()> { + std::task::Poll::Pending + } + } + }; +} +#[cfg(test)] +pub(crate) use unheld; + +#[cfg(any(feature = "tcp", feature = "websocket"))] +unheld!(qmux::Session, qmux::RecvStream); +// iroh's upstream quinn has no way to hold credit, so its incomplete messages are bounded +// only by its stream limit, each up to the largest message. +#[cfg(feature = "iroh")] +unheld!(web_transport_iroh::Session, web_transport_iroh::RecvStream); + +#[cfg(feature = "quic")] +impl moq_net::transport::Credit for crate::quic::driver::Credit { + fn len(&self) -> u64 { + self.len() + } + + fn split(&mut self, bytes: u64) -> Self { + self.split(bytes) + } + + fn merge(&mut self, other: Self) { + self.merge(other) + } +} + +#[cfg(feature = "quic")] +impl moq_net::transport::Stopper for crate::quic::web::Stopper { + fn stop(&self, code: u32) { + // Err means the stream is already gone, which is what we wanted. + let _ = crate::quic::web::Stopper::stop(self, code); + } +} + +#[cfg(feature = "quic")] +impl super::HeldRecvStream for crate::quic::web::RecvStream { + type Credit = crate::quic::driver::Credit; + type Stopper = crate::quic::web::Stopper; + + fn stopper(&self) -> Self::Stopper { + crate::quic::web::RecvStream::stopper(self) + } + + fn read_chunk_held( + &mut self, + max: usize, + ) -> impl Future, Self::Error>> + Send { + crate::quic::web::RecvStream::read_chunk(self, max) + } +} + +#[cfg(feature = "quic")] +impl super::HeldSession for crate::quic::web::Session { + fn poll_credit_exhausted(&self, cx: &mut std::task::Context<'_>) -> std::task::Poll<()> { + crate::quic::driver::Connection::poll_credit_exhausted(self, cx) + } +} diff --git a/rs/moq-tokio/src/transport/owned.rs b/rs/moq-tokio/src/transport/owned.rs index 433022db697..ed52d9d3b6f 100644 --- a/rs/moq-tokio/src/transport/owned.rs +++ b/rs/moq-tokio/src/transport/owned.rs @@ -1,6 +1,6 @@ use bytes::{Buf, BufMut, Bytes}; use moq_net::transport::{self as moq, poll}; -use std::task::{Context, Poll}; +use std::task::{Context, Poll, ready}; use web_transport_trait::{self as wt, poll as backend}; /// A backend failure preserving the session and stream code registries. @@ -47,7 +47,9 @@ impl moq::Stats for Stats { self.0.estimated_send_rate() } } -impl> poll::Session for super::Session { +impl + super::Grouped> poll::Session + for super::Session +{ type SendStream = super::SendStream; type RecvStream = super::RecvStream; type Error = Error; @@ -92,6 +94,9 @@ impl> poll::Session fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll { backend::Session::poll_closed(self, cx).map(Error) } + fn poll_credit_exhausted(&mut self, cx: &mut Context<'_>) -> Poll<()> { + super::HeldSession::poll_credit_exhausted(&self.session, cx) + } fn stats(&self) -> impl moq::Stats { Stats(backend::Session::stats(self)) } @@ -121,20 +126,39 @@ impl poll::SendStream for super::SendStream backend::SendStream::poll_closed(self, cx).map_err(Error) } } -impl poll::RecvStream for super::RecvStream { +impl poll::RecvStream for super::RecvStream { type Error = Error; - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { - backend::RecvStream::poll_read(self, cx, dst).map_err(Error) + type Credit = super::Credit; + type Stopper = S::Stopper; + fn stopper(&self) -> S::Stopper { + self.stopper.clone() + } + fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + let read = ready!(self.poll_read_held(cx, dst.len())).map_err(Error)?; + Poll::Ready(Ok(read.map(|(bytes, credit)| { + dst[..bytes.len()].copy_from_slice(&bytes); + credit + }))) } fn poll_read_buf( &mut self, cx: &mut Context<'_>, buf: &mut B, - ) -> Poll, Self::Error>> { - backend::RecvStream::poll_read_buf(self, cx, buf).map_err(Error) - } - fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Self::Error>> { - backend::RecvStream::poll_read_chunk(self, cx, max).map_err(Error) + ) -> Poll, Self::Error>> { + // Hand over the chunk the backend owns, without zeroing a destination first. + let max = buf.remaining_mut().min(64 * 1024); + let read = ready!(self.poll_read_held(cx, max)).map_err(Error)?; + Poll::Ready(Ok(read.map(|(bytes, credit)| { + buf.put(bytes); + credit + }))) + } + fn poll_read_chunk( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Self::Error>> { + self.poll_read_held(cx, max).map_err(Error) } fn stop(&mut self, code: u32) { backend::RecvStream::stop(self, code) diff --git a/rs/moq-tokio/tests/integration/backend.rs b/rs/moq-tokio/tests/integration/backend.rs index e08488a9572..3fc4a193d38 100644 --- a/rs/moq-tokio/tests/integration/backend.rs +++ b/rs/moq-tokio/tests/integration/backend.rs @@ -1328,16 +1328,17 @@ async fn connect_once( // ── flow control ──────────────────────────────────────────────────── -/// Run a connect with every flow-control window pinned well under the payload, so -/// the frame only arrives if the receiver keeps issuing credit as it drains. +/// Run a connect with the stream and send windows pinned well under the payload, so +/// the frame only arrives if the receiver keeps issuing stream credit as it drains. +/// The connection window sits at the smallest that fits the largest message, which +/// holds its credit until it is complete. /// /// The size is the point: a frame that fits inside one window would pass whether or /// not the setting ever reached the backend. -/// #[cfg(feature = "quic")] async fn window_test(scheme: &str) { let mut quic = moq_tokio::quic::Config::default(); - quic.receive_window = Some(64 * 1024); + quic.receive_window = Some(moq_net::transport::MIN_RECEIVE_WINDOW); quic.stream_receive_window = Some(16 * 1024); quic.send_window = Some(32 * 1024); diff --git a/rs/moq-tokio/tests/integration/main.rs b/rs/moq-tokio/tests/integration/main.rs index 79552d9449e..46e1819dbfa 100644 --- a/rs/moq-tokio/tests/integration/main.rs +++ b/rs/moq-tokio/tests/integration/main.rs @@ -8,6 +8,7 @@ mod backend; mod broadcast; mod close_code; mod max_age; +mod receive_credit; mod reconnect; mod stream_code; mod subscription_end_integrity; diff --git a/rs/moq-tokio/tests/integration/receive_credit.rs b/rs/moq-tokio/tests/integration/receive_credit.rs new file mode 100644 index 00000000000..d9c146d81c8 --- /dev/null +++ b/rs/moq-tokio/tests/integration/receive_credit.rs @@ -0,0 +1,246 @@ +//! Over native QUIC, incomplete frames hold the receiver's connection credit, and when +//! they fill its window the receiver stops the least urgent stream and keeps going. +//! +//! The server publishes with the stock config. The client dials over the raw driver so +//! its connection window can be far smaller than the stack allows, which keeps the +//! payloads small. Without recovery the high-priority frame never arrives: the window +//! is spent on a low-priority frame that can never complete. + +#![cfg(feature = "quic")] + +use std::{sync::Arc, time::Duration}; + +use bytes::Bytes; +use moq_tokio::moq_net::{self, Timestamp, frame, track::Subscription}; +use moq_tokio::quic::driver; + +const TIMEOUT: Duration = Duration::from_secs(10); + +/// The client's connection receive window. +const WINDOW: usize = 256 * 1024; + +fn timed() -> moq_net::track::Info { + moq_net::track::Info::default().with_timescale(moq_net::Timescale::MILLI) +} + +/// A server publishing broadcast `room`, and a client subscribed to its origin over a +/// connection whose receive window is [`WINDOW`]. +struct Setup { + broadcast: moq_net::broadcast::Producer, + remote: moq_net::broadcast::Consumer, + _client: moq_net::Session, + _endpoint: driver::Endpoint, +} + +async fn setup(version: &str) -> Setup { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + let version: moq_net::Version = version.parse().expect("version"); + + let publisher = moq_tokio::origin::spawn(); + let broadcast = publisher.create_broadcast("room").expect("broadcast"); + broadcast.announce(Default::default()).expect("announce"); + + let mut config = moq_tokio::listen::Config::default(); + config.bind = Some("127.0.0.1:0".parse().unwrap()); + config.tls.generate = vec!["localhost".into()]; + config.version = vec![version]; + let mut server = config + .init(Default::default()) + .expect("init server") + .listen() + .await + .expect("listen"); + let port = server.local_addr().expect("addr").port(); + tokio::spawn(async move { + let request = server.accept().await.expect("accept"); + let (session, _driver) = request.with_publisher(&publisher).ok().await.expect("session"); + let _ = session.closed().await; + }); + + let mut tls = rustls::ClientConfig::builder() + .dangerous() + .with_custom_certificate_verifier(Arc::new(AnyCert)) + .with_no_client_auth(); + tls.alpn_protocols = vec![version.alpn().as_bytes().to_vec()]; + let quic = driver::crypto::rustls::QuicClientConfig::try_from(tls).expect("quic config"); + let mut transport = driver::TransportConfig::default(); + transport.receive_window((WINDOW as u32).into()); + let mut client = driver::ClientConfig::new(Arc::new(quic)); + client.transport_config(Arc::new(transport)); + + let endpoint = driver::Endpoint::client("127.0.0.1:0".parse().unwrap()).expect("endpoint"); + let conn = endpoint + .connect_with(client, ([127, 0, 0, 1], port).into(), "localhost") + .expect("connect") + .await + .expect("handshake"); + + let subscriber = moq_tokio::origin::spawn(); + let session = moq_tokio::transport::Session::new( + moq_tokio::quic::web::Session::raw(conn), + moq_net::transport::Binding::InBand, + ); + let (client, driver) = moq_net::Client::new() + .with_versions(version.into()) + .with_subscriber(subscriber.clone()) + .connect(moq_net::time::now(), session) + .await + .expect("client handshake"); + tokio::spawn(moq_net::time::run(driver)); + + let consumer = subscriber.consume(); + tokio::time::timeout(TIMEOUT, consumer.routed("room")) + .await + .expect("announce timeout") + .expect("routed"); + let remote = tokio::time::timeout(TIMEOUT, consumer.request_broadcast("room", None)) + .await + .expect("resolve timeout") + .expect("broadcast resolves"); + + Setup { + broadcast, + remote, + _client: client, + _endpoint: endpoint, + } +} + +impl Setup { + /// Subscribe to a new track at `priority`. + async fn subscribe(&self, name: &str, priority: u8) -> (moq_net::track::Producer, moq_net::track::Subscriber) { + let track = self.broadcast.create_track(name, timed()).unwrap(); + let sub = tokio::time::timeout( + TIMEOUT, + self.remote + .track(name) + .unwrap() + .subscribe(Subscription::default().with_priority(priority)), + ) + .await + .expect("subscribe timeout") + .expect("subscribe"); + (track, sub) + } +} + +/// The next group, which must arrive in time. +async fn next_group(sub: &mut moq_net::track::Subscriber) -> moq_net::group::Consumer { + tokio::time::timeout(TIMEOUT, sub.recv_group()) + .await + .expect("group timed out") + .expect("group") + .expect("track ended") +} + +/// The next frame's payload, which must arrive in time. +async fn next_frame(sub: &mut moq_net::track::Subscriber) -> Bytes { + let mut group = next_group(sub).await; + tokio::time::timeout(TIMEOUT, group.read_frame()) + .await + .expect("frame timed out") + .expect("frame") + .expect("group ended") + .payload +} + +/// Two partial frames that together exceed the window: the lower-priority one is +/// stopped, and the higher-priority one completes on the credit it returns. +async fn the_less_urgent_partial_frame_gives_way(version: &str) { + let setup = setup(version).await; + let (low, mut low_sub) = setup.subscribe("low", 1).await; + let (high, mut high_sub) = setup.subscribe("high", 200).await; + + // Three quarters of the window, held by a frame that never finishes. + let mut low_group = low.append_group().unwrap(); + let mut low_frame = low_group + .create_frame(frame::Info { + size: WINDOW as u64, + timestamp: Some(Timestamp::ZERO), + }) + .unwrap(); + low_frame.write(vec![1u8; WINDOW * 3 / 4]).unwrap(); + + // Wait until the receiver holds all of it, so nothing else got there first. + let mut partial = next_group(&mut low_sub).await; + let mut partial = tokio::time::timeout(TIMEOUT, partial.next_frame()) + .await + .expect("frame header timed out") + .expect("frame") + .expect("group ended"); + let mut seen = 0; + while seen < WINDOW * 3 / 4 { + let chunk = tokio::time::timeout(TIMEOUT, partial.read_chunk()) + .await + .expect("partial payload timed out") + .expect("chunk") + .expect("frame ended"); + seen += chunk.len(); + } + + // Half the window, which only fits once the low frame lets go. + let payload = Bytes::from(vec![2u8; WINDOW / 2]); + let mut high_group = high.append_group().unwrap(); + high_group.write_frame(Timestamp::ZERO, payload.clone()).unwrap(); + high_group.finish().unwrap(); + assert_eq!(next_frame(&mut high_sub).await, payload, "{version}"); + + // The session survives, and the low track keeps delivering what fits. + drop(low_frame); + drop(low_group); + let mut group = low.append_group().unwrap(); + group.write_frame(Timestamp::ZERO, &b"after"[..]).unwrap(); + group.finish().unwrap(); + assert_eq!(next_frame(&mut low_sub).await, &b"after"[..], "{version}"); +} + +#[tokio::test] +async fn lite_partial_frame_gives_way() { + the_less_urgent_partial_frame_gives_way("moq-lite-06").await; +} + +#[tokio::test] +async fn ietf_partial_frame_gives_way() { + the_less_urgent_partial_frame_gives_way("moq-transport-16").await; +} + +/// Accepts any server certificate: the test dials its own self-signed server. +#[derive(Debug)] +struct AnyCert; + +impl rustls::client::danger::ServerCertVerifier for AnyCert { + fn verify_server_cert( + &self, + _: &rustls::pki_types::CertificateDer<'_>, + _: &[rustls::pki_types::CertificateDer<'_>], + _: &rustls::pki_types::ServerName<'_>, + _: &[u8], + _: rustls::pki_types::UnixTime, + ) -> Result { + Ok(rustls::client::danger::ServerCertVerified::assertion()) + } + + fn verify_tls12_signature( + &self, + _: &[u8], + _: &rustls::pki_types::CertificateDer<'_>, + _: &rustls::DigitallySignedStruct, + ) -> Result { + Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) + } + + fn verify_tls13_signature( + &self, + _: &[u8], + _: &rustls::pki_types::CertificateDer<'_>, + _: &rustls::DigitallySignedStruct, + ) -> Result { + Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) + } + + fn supported_verify_schemes(&self) -> Vec { + rustls::crypto::aws_lc_rs::default_provider() + .signature_verification_algorithms + .supported_schemes() + } +} diff --git a/rs/moq-uring/src/quic/driver/connection.rs b/rs/moq-uring/src/quic/driver/connection.rs index da60852433b..01c3a0894f8 100644 --- a/rs/moq-uring/src/quic/driver/connection.rs +++ b/rs/moq-uring/src/quic/driver/connection.rs @@ -8,7 +8,7 @@ use std::time::Instant; use bytes::Bytes; use moq_quic::{ConnectionHandle, Dir, StreamId, VarInt}; -use rustc_hash::FxHashMap; +use rustc_hash::{FxHashMap, FxHashSet}; use super::super::{Error, SEGMENT}; use super::endpoint; @@ -58,6 +58,9 @@ pub(crate) struct State { /// driver's per-event path, and moq-quic assigns the ids. readable: FxHashMap, writable: FxHashMap, + /// Receive streams a [`Stopper`](super::Stopper) stopped, so their handle's next read + /// fails rather than reading the discarded tail as the stream's end. + stopped: FxHashSet, /// Send streams waiting for their end: a FIN the peer acknowledged, or a /// `STOP_SENDING`. finishing: FxHashMap, @@ -73,6 +76,11 @@ pub(crate) struct State { datagram_recv_waiters: kio::WaiterList, + /// Held credit filled the receive window since + /// [`Connection::poll_credit_exhausted`] last resolved. + credit_exhausted: bool, + credit_waiters: kio::WaiterList, + /// The terminal error, set exactly once; everything fails with it after. closed: Option, closed_waiters: kio::WaiterList, @@ -99,9 +107,12 @@ impl State { open_waiters: kio::WaiterList::new(), readable: FxHashMap::default(), writable: FxHashMap::default(), + stopped: FxHashSet::default(), finishing: FxHashMap::default(), sends: FxHashMap::default(), datagram_recv_waiters: kio::WaiterList::new(), + credit_exhausted: false, + credit_waiters: kio::WaiterList::new(), closed: None, closed_waiters: kio::WaiterList::new(), dead: false, @@ -121,6 +132,7 @@ impl State { /// be parked. fn forget_recv(&mut self, id: StreamId) { self.readable.remove(&id); + self.stopped.remove(&id); } /// Terminate with `err` (the first one wins) and wake absolutely everyone, @@ -155,6 +167,18 @@ impl Inner { self.state.borrow().closed.clone() } + /// Note when held credit comes to fill the receive window, waking whoever + /// watches for it. + /// + /// Call after anything that holds, returns, or discards connection credit. + pub(crate) fn check_credit(&self) { + if self.conn.borrow_mut().poll_credit_exhausted() { + let mut state = self.state.borrow_mut(); + state.credit_exhausted = true; + state.credit_waiters.wake(); + } + } + /// Wake the driver so it flushes what a handle just queued. pub(crate) fn kick(&self) { self.state.borrow_mut().driver.wake(); @@ -236,6 +260,21 @@ impl Inner { self.state.borrow_mut().forget_recv(id); } + /// Record that a [`Stopper`](super::Stopper) stopped receive stream `id`, waking a read + /// parked on it: a stop produces no readable event. + pub(crate) fn stop_recv(&self, id: StreamId) { + let mut state = self.state.borrow_mut(); + state.stopped.insert(id); + if let Some(mut waiters) = state.readable.remove(&id) { + waiters.wake(); + } + } + + /// Whether a [`Stopper`](super::Stopper) stopped receive stream `id`. + pub(crate) fn recv_stopped(&self, id: StreamId) -> bool { + self.state.borrow().stopped.contains(&id) + } + /// Wake anyone parked on stream `id` being readable. pub(crate) fn wake_readable(&self, id: StreamId) { let mut state = self.state.borrow_mut(); @@ -376,6 +415,23 @@ impl Connection { pub fn stats(&self) -> moq_quic::ConnectionStats { self.shared.conn.borrow().stats() } + + /// Poll until held [`Credit`](super::Credit) fills the receive window: the + /// peer may send nothing more and every byte it sent is held, so no stream + /// makes progress until some credit is dropped. + /// + /// Resolves once per stall, and again only after more credit is held or + /// credit drops without relieving it. + pub fn poll_credit_exhausted(&mut self, cx: &mut Context<'_>) -> Poll<()> { + let waiter = self.park.hold(cx); + let mut state = self.shared.state.borrow_mut(); + // Credit may have dropped since the stall was noted. + if std::mem::take(&mut state.credit_exhausted) && self.shared.conn.borrow().stats().credit.exhausted() { + return Poll::Ready(()); + } + waiter.register(&mut state.credit_waiters); + Poll::Pending + } } impl Clone for Connection { @@ -736,6 +792,9 @@ impl Driver { loop { self.endpoint_events(); self.sweep(); + // Data for a stopped stream or a reset returns its credit, possibly + // without relieving a stall. + self.shared.check_credit(); if self.shared.state.borrow().dead { return Poll::Ready(()); diff --git a/rs/moq-uring/src/quic/driver/mod.rs b/rs/moq-uring/src/quic/driver/mod.rs index cd6d5a77ae0..fc5ca26dd18 100644 --- a/rs/moq-uring/src/quic/driver/mod.rs +++ b/rs/moq-uring/src/quic/driver/mod.rs @@ -17,7 +17,7 @@ mod stream; pub use connection::Connection; pub use endpoint::Endpoint; pub use group::SendGroup; -pub use stream::{Credit, RecvStream, SendStream}; +pub use stream::{Credit, RecvStream, SendStream, Stopper}; pub(crate) use connection::{End, Shared}; @@ -32,10 +32,12 @@ use super::qlog; use super::{Congestion, Error, Identity, SEGMENT, Transport, client, server}; use crate::udp; -/// Per-stream flow control credit. -const STREAM_WINDOW: u32 = 4 * 1024 * 1024; -/// Per-connection flow control credit. -const CONNECTION_WINDOW: u32 = 16 * 1024 * 1024; +/// Per-connection flow control credit, which also bounds incomplete messages, so it has to fit +/// the largest one. Each stream may use all of it. +const CONNECTION_WINDOW: u32 = 128 * 1024 * 1024; +const _: () = assert!(CONNECTION_WINDOW as u64 >= moq_net::transport::MIN_RECEIVE_WINDOW); +/// Unacknowledged outgoing data to buffer, whatever the peer allows. +const SEND_WINDOW: u32 = 16 * 1024 * 1024; /// How many datagrams to buffer in each direction. const DATAGRAM_WINDOW: usize = 64 * SEGMENT; @@ -263,17 +265,17 @@ fn transport_config(config: &Transport) -> Result varint_window("stream receive window", window)?, - None => STREAM_WINDOW.into(), - }; + // A value that cannot go on the wire is an error: clamping it would ignore + // what was set. let connection_window = match config.receive_window { Some(window) => varint_window("receive window", window)?, None => CONNECTION_WINDOW.into(), }; - let send_window = config.send_window.unwrap_or(u64::from(CONNECTION_WINDOW)); + let stream_window = match config.stream_receive_window { + Some(window) => varint_window("stream receive window", window)?, + None => connection_window, + }; + let send_window = config.send_window.unwrap_or(u64::from(SEND_WINDOW)); transport.stream_receive_window(stream_window); transport.receive_window(connection_window); transport.send_window(send_window); @@ -387,7 +389,7 @@ impl rustls::client::danger::ServerCertVerifier for NoVerify { #[cfg(test)] mod tests { - use super::{CONNECTION_WINDOW, STREAM_WINDOW, Transport, transport_config}; + use super::{CONNECTION_WINDOW, SEND_WINDOW, Transport, transport_config}; /// moq-quic keeps the windows private, but its `Debug` prints them. The leading /// space keeps `receive_window` from matching `stream_receive_window`, and @@ -402,23 +404,23 @@ mod tests { assert!(has_window( &transport, "stream_receive_window", - u64::from(STREAM_WINDOW) + u64::from(CONNECTION_WINDOW) )); assert!(has_window(&transport, "receive_window", u64::from(CONNECTION_WINDOW))); - assert!(has_window(&transport, "send_window", u64::from(CONNECTION_WINDOW))); + assert!(has_window(&transport, "send_window", u64::from(SEND_WINDOW))); } #[test] fn configured_windows_replace_the_defaults() { let config = Transport { - receive_window: Some(64 << 20), + receive_window: Some(256 << 20), stream_receive_window: Some(8 << 20), send_window: Some(32 << 20), ..Default::default() }; let transport = transport_config(&config).expect("windows"); - assert!(has_window(&transport, "receive_window", 64 << 20)); + assert!(has_window(&transport, "receive_window", 256 << 20)); assert!(has_window(&transport, "stream_receive_window", 8 << 20)); assert!(has_window(&transport, "send_window", 32 << 20)); assert!(!has_window(&transport, "receive_window", u64::from(CONNECTION_WINDOW))); @@ -441,6 +443,16 @@ mod tests { assert!(has_window(&transport, "send_window", u64::MAX)); } + #[test] + fn the_stream_window_follows_the_receive_window() { + let config = Transport { + receive_window: Some(256 << 20), + ..Default::default() + }; + let transport = transport_config(&config).expect("windows"); + assert!(has_window(&transport, "stream_receive_window", 256 << 20)); + } + #[test] fn a_receive_window_past_a_varint_is_an_error() { let over = Transport { diff --git a/rs/moq-uring/src/quic/driver/stream.rs b/rs/moq-uring/src/quic/driver/stream.rs index 36157e1cee6..43291be9d3c 100644 --- a/rs/moq-uring/src/quic/driver/stream.rs +++ b/rs/moq-uring/src/quic/driver/stream.rs @@ -269,40 +269,35 @@ impl RecvStream { } } - /// Read into `dst`, keeping the connection flow control credit for exactly - /// the bytes read, or `None` once the peer has finished the stream. + /// Read up to `max` bytes, keeping the connection flow control credit for + /// exactly the bytes read, or `None` once the peer has finished the stream. /// /// Stream credit still advances, so a stream window smaller than an /// application message keeps the message flowing, while the connection /// window these bytes occupy stays spent until the credit is dropped. - pub fn poll_read_held(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Error>> { + pub fn poll_read_held(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Error>> { let mut held = moq_quic::Credit::default(); - self.poll_read_into(cx, dst, Some(&mut held)).map_ok(|read| { - read.map(|_| Credit { + let chunk = ready!(self.poll_chunk(cx, max, Some(&mut held)))?; + Poll::Ready(Ok(chunk.map(|bytes| { + let credit = Credit { shared: self.shared.clone(), inner: held, - }) - }) + }; + (bytes, credit) + }))) } - /// [`poll_read`](web_transport_trait::poll::RecvStream::poll_read), with - /// the credit for the bytes read moved into `held` instead of returned to - /// the peer. - fn poll_read_into( - &mut self, - cx: &mut Context<'_>, - dst: &mut [u8], - held: Option<&mut moq_quic::Credit>, - ) -> Poll, Error>> { - let Some(chunk) = ready!(self.poll_chunk(cx, dst.len(), held))? else { + /// [`poll_read`](web_transport_trait::poll::RecvStream::poll_read), which + /// returns the credit for the bytes read to the peer. + fn poll_read_plain(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Error>> { + let Some(chunk) = ready!(self.poll_chunk(cx, dst.len(), None))? else { return Poll::Ready(Ok(None)); }; dst[..chunk.len()].copy_from_slice(&chunk); Poll::Ready(Ok(Some(chunk.len()))) } - /// [`poll_read_chunk`](web_transport_trait::poll::RecvStream::poll_read_chunk), - /// with the credit for the bytes read moved into `held` instead of + /// Up to `max` bytes, with the credit for them moved into `held` or else /// returned to the peer. fn poll_chunk( &mut self, @@ -323,6 +318,11 @@ impl RecvStream { match Self::read(&self.shared, self.id, max, held) { // moq-quic hands out at most `max`. Read::Chunk(bytes) => Poll::Ready(Ok(Some(bytes))), + // What a stop discarded is not the stream's end. + Read::Finished if self.shared.recv_stopped(self.id) => { + self.stopped = true; + Poll::Ready(Err(Error::Quic("stream stopped".to_string()))) + } Read::Finished => { self.finished = true; Poll::Ready(Ok(None)) @@ -344,6 +344,14 @@ impl RecvStream { self.finished || self.stopped } + /// A handle that stops this stream without borrowing it. + pub fn stopper(&self) -> Stopper { + Stopper { + shared: self.shared.clone(), + id: self.id, + } + } + /// [`stop`](web_transport_trait::poll::RecvStream::stop) with a /// full-width code, for the WebTransport HTTP/3 error mapping. pub(crate) fn stop_code(&mut self, code: u64) { @@ -360,6 +368,7 @@ impl RecvStream { .recv_stream(self.id) .stop(VarInt::from_u64(code).unwrap_or(VarInt::MAX)); self.stopped = true; + self.shared.check_credit(); self.shared.kick(); } @@ -393,6 +402,7 @@ impl RecvStream { if transmit { shared.kick(); } + shared.check_credit(); read } @@ -419,7 +429,7 @@ impl web_transport_trait::poll::RecvStream for RecvStream { type Error = Error; fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { - self.poll_read_into(cx, dst, None) + self.poll_read_plain(cx, dst) } fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Self::Error>> { @@ -455,6 +465,11 @@ impl web_transport_trait::poll::RecvStream for RecvStream { } match Self::read(&self.shared, self.id, READ_CHUNK, Some(&mut self.backlog_credit)) { Read::Chunk(bytes) => self.backlog.extend_from_slice(&bytes), + // What a stop discarded is not the stream's end, so the next read fails. + Read::Finished if self.shared.recv_stopped(self.id) => { + self.stopped = true; + return Poll::Ready(Ok(())); + } Read::Finished => { self.finished = true; return Poll::Ready(Ok(())); @@ -472,6 +487,41 @@ impl web_transport_trait::poll::RecvStream for RecvStream { } } +/// Stops a [`RecvStream`] without borrowing it. +/// +/// The stream's own handle learns of it on its next read, which fails as stopped, and a +/// read parked on the stream is woken for it. What its closed-watch read ahead stays with +/// that handle until it stops or drops. +#[derive(Clone)] +pub struct Stopper { + shared: Shared, + id: StreamId, +} + +impl Stopper { + /// Stop the stream with `code`, as [`stop`](web_transport_trait::poll::RecvStream::stop) + /// does. + pub fn stop(&self, code: u32) { + self.stop_code(code.into()); + } + + /// [`RecvStream::stop_code`] from outside the stream. + pub(crate) fn stop_code(&self, code: u64) { + let stopped = self + .shared + .conn + .borrow_mut() + .recv_stream(self.id) + .stop(VarInt::from_u64(code).unwrap_or(VarInt::MAX)); + // Err means the stream is already gone, which is what we wanted. + if stopped.is_ok() { + self.shared.stop_recv(self.id); + } + self.shared.check_credit(); + self.shared.kick(); + } +} + impl Drop for RecvStream { fn drop(&mut self) { release(&self.shared, std::mem::take(&mut self.backlog_credit)); @@ -502,6 +552,7 @@ fn release(shared: &Shared, credit: moq_quic::Credit) { if shared.conn.borrow_mut().release(credit).should_transmit() { shared.kick(); } + shared.check_credit(); } /// Connection flow control credit for bytes read with diff --git a/rs/moq-uring/src/quic/mod.rs b/rs/moq-uring/src/quic/mod.rs index 70e2b099644..16efae87315 100644 --- a/rs/moq-uring/src/quic/mod.rs +++ b/rs/moq-uring/src/quic/mod.rs @@ -30,7 +30,7 @@ pub mod web; mod driver; -pub use driver::{Connection, Credit, RecvStream, SendGroup, SendStream}; +pub use driver::{Connection, Credit, RecvStream, SendGroup, SendStream, Stopper}; pub use endpoint::Endpoint; /// The QUIC payload size every full datagram in a GSO train uses, and the @@ -115,14 +115,17 @@ pub struct Transport { /// may have open at once. MoQ opens a stream per group, so busy endpoints /// want this high. pub max_streams: u64, - /// How much unread data one peer may have in flight across every stream, - /// in bytes. `None` leaves the 16 MiB default. + /// How much unread data and incomplete messages one peer may have in flight + /// across every stream, in bytes. `None` leaves the 128 MiB default. A moq + /// session needs it to fit the largest message it assembles, + /// [`MIN_RECEIVE_WINDOW`](moq_net::transport::MIN_RECEIVE_WINDOW), or a + /// partial one can stall the connection. pub receive_window: Option, /// How much unread data one peer may have in flight on a single stream, - /// in bytes. `None` leaves the 4 MiB default. + /// in bytes. `None` follows [`receive_window`](Self::receive_window). /// - /// Keep it under [`receive_window`](Self::receive_window) so one stream - /// cannot fill the connection. + /// Stream credit returns as bytes are read, whether or not their message is + /// complete, so a smaller window only adds round trips to a large frame. pub stream_receive_window: Option, /// How much unacknowledged outgoing data to buffer, in bytes, whatever /// the peer allows. `None` leaves the 16 MiB default. diff --git a/rs/moq-uring/src/quic/web.rs b/rs/moq-uring/src/quic/web.rs index 911d940887c..31bbafaf935 100644 --- a/rs/moq-uring/src/quic/web.rs +++ b/rs/moq-uring/src/quic/web.rs @@ -545,6 +545,11 @@ impl Session { } } + /// See [`Connection::poll_credit_exhausted`]. + pub fn poll_credit_exhausted(&mut self, cx: &mut Context<'_>) -> Poll<()> { + self.conn.poll_credit_exhausted(cx) + } + /// Rewrite error codes out of the HTTP/3 mapping in web mode. fn map_err(&self, err: Error) -> Error { match self.web { @@ -939,13 +944,51 @@ pub struct RecvStream { web: bool, } +/// Stops a [`RecvStream`] without borrowing it, with the same code mapping as its `stop`. +#[derive(Clone)] +pub struct Stopper { + inner: super::Stopper, + web: bool, +} + +impl Stopper { + /// Stop the stream with an application `code`. + pub fn stop(&self, code: u32) { + match self.web { + true => self.inner.stop_code(proto::error_to_http3(code)), + false => self.inner.stop_code(u64::from(code)), + } + } +} + impl RecvStream { + /// A handle that stops this stream without borrowing it. + pub fn stopper(&self) -> Stopper { + Stopper { + inner: self.inner.stopper(), + web: self.web, + } + } + fn map(&self, err: Error) -> Error { match self.web { true => unmap_err(err), false => unmap_legacy(err), } } + + /// Read up to `max` bytes with their connection credit. See + /// [`super::RecvStream::poll_read_held`]. + pub fn poll_read_held( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Error>> { + match self.inner.poll_read_held(cx, max) { + Poll::Ready(Err(err)) => Poll::Ready(Err(self.map(err))), + other => other, + } + } } impl web_transport_trait::poll::RecvStream for RecvStream { diff --git a/rs/moq-uring/src/transport.rs b/rs/moq-uring/src/transport.rs index b77abb15754..259bda7649e 100644 --- a/rs/moq-uring/src/transport.rs +++ b/rs/moq-uring/src/transport.rs @@ -1,6 +1,8 @@ //! Adapt backend poll sessions to moq-net without changing thread affinity. mod adapter; -pub use adapter::{Error, RecvStream, SendStream, Session}; +#[cfg(feature = "quic")] +mod held; +pub use adapter::{Error, HeldRecvStream, HeldSession, RecvStream, SendStream, Session}; #[cfg(feature = "quic")] use adapter::{Grouped, GroupedSend}; diff --git a/rs/moq-uring/src/transport/adapter.rs b/rs/moq-uring/src/transport/adapter.rs index 0f4ff4dc4ce..dfc776ae66f 100644 --- a/rs/moq-uring/src/transport/adapter.rs +++ b/rs/moq-uring/src/transport/adapter.rs @@ -1,6 +1,6 @@ use bytes::{Buf, BufMut, Bytes}; use moq_net::transport::{self as moq, poll}; -use std::task::{Context, Poll}; +use std::task::{Context, Poll, ready}; use web_transport_trait::{self as wt, poll as backend}; /// A backend failure preserving the session and stream code registries. @@ -47,6 +47,32 @@ impl moq::Stats for Stats { self.0.estimated_send_rate() } } +/// A backend receive stream that can keep the connection credit of what it reads. +pub trait HeldRecvStream: backend::RecvStream { + /// The connection credit a read keeps until it is dropped. + type Credit: moq::Credit; + + /// What [`stopper`](Self::stopper) returns. + type Stopper: moq::Stopper; + + /// A handle that stops this stream without borrowing it. + fn stopper(&self) -> Self::Stopper; + + /// Read up to `max` bytes with their credit, or `None` once finished. + #[allow(clippy::type_complexity)] + fn poll_read_held( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Self::Error>>; +} + +/// A backend session that reports when held credit fills its receive window. +pub trait HeldSession: backend::Session { + /// See [`poll_credit_exhausted`](poll::Session::poll_credit_exhausted). + fn poll_credit_exhausted(&mut self, cx: &mut Context<'_>) -> Poll<()>; +} + /// A backend session whose send streams may join send groups, as native QUIC's do. /// /// `pub` only because `pub(crate)` is refused (E0446): the adapter's `type Group` projects @@ -105,7 +131,9 @@ impl Session { Self(session, binding) } } -impl + Clone + 'static> poll::Session for Session { +impl + Grouped + Clone + 'static> poll::Session + for Session +{ type SendStream = SendStream; type RecvStream = RecvStream; type Error = Error; @@ -158,6 +186,9 @@ impl + Clone + 'static> pol fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll { backend::Session::poll_closed(&mut self.0, cx).map(Error) } + fn poll_credit_exhausted(&mut self, cx: &mut Context<'_>) -> Poll<()> { + HeldSession::poll_credit_exhausted(&mut self.0, cx) + } fn stats(&self) -> impl moq::Stats { Stats(backend::Session::stats(&self.0)) } @@ -191,20 +222,39 @@ impl poll::SendStream for SendStream { } /// An incoming WebTransport stream adapted to moq-net. pub struct RecvStream(S); -impl poll::RecvStream for RecvStream { +impl poll::RecvStream for RecvStream { type Error = Error; - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { - backend::RecvStream::poll_read(&mut self.0, cx, dst).map_err(Error) + type Credit = S::Credit; + type Stopper = S::Stopper; + fn stopper(&self) -> S::Stopper { + self.0.stopper() + } + fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + let read = ready!(self.0.poll_read_held(cx, dst.len())).map_err(Error)?; + Poll::Ready(Ok(read.map(|(bytes, credit)| { + dst[..bytes.len()].copy_from_slice(&bytes); + credit + }))) } fn poll_read_buf( &mut self, cx: &mut Context<'_>, buf: &mut B, - ) -> Poll, Self::Error>> { - backend::RecvStream::poll_read_buf(&mut self.0, cx, buf).map_err(Error) - } - fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Self::Error>> { - backend::RecvStream::poll_read_chunk(&mut self.0, cx, max).map_err(Error) + ) -> Poll, Self::Error>> { + // Hand over the chunk the backend owns, without zeroing a destination first. + let max = buf.remaining_mut().min(64 * 1024); + let read = ready!(self.0.poll_read_held(cx, max)).map_err(Error)?; + Poll::Ready(Ok(read.map(|(bytes, credit)| { + buf.put(bytes); + credit + }))) + } + fn poll_read_chunk( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Self::Error>> { + self.0.poll_read_held(cx, max).map_err(Error) } fn stop(&mut self, code: u32) { backend::RecvStream::stop(&mut self.0, code) @@ -279,6 +329,28 @@ mod tests { Poll::Ready(Err(Failure)) } } + impl HeldRecvStream for Stream { + type Credit = moq::Unheld; + type Stopper = (); + fn stopper(&self) {} + fn poll_read_held( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Failure>> { + backend::RecvStream::poll_read_chunk(self, cx, max).map_ok(|chunk| { + chunk.map(|bytes| { + let len = bytes.len() as u64; + (bytes, moq::Unheld(len)) + }) + }) + } + } + impl HeldSession for Backend { + fn poll_credit_exhausted(&mut self, _: &mut Context<'_>) -> Poll<()> { + Poll::Pending + } + } impl GroupedSend for Stream { type Group = poll::Ungrouped; fn set_group(&mut self, group: &Self::Group) { @@ -378,7 +450,7 @@ mod tests { let Poll::Ready(Ok(Some(chunk))) = poll::RecvStream::poll_read_chunk(&mut recv, &mut cx, 2) else { panic!("read failed") }; - assert_eq!(chunk, b"hi"[..]); + assert_eq!(chunk.0, b"hi"[..]); poll::RecvStream::stop(&mut recv, 13); assert_eq!(state.get(), 13); let Poll::Ready(Err(error)) = poll::SendStream::poll_closed(&mut send, &mut cx) else { diff --git a/rs/moq-uring/src/transport/held.rs b/rs/moq-uring/src/transport/held.rs new file mode 100644 index 00000000000..2de4bcd52f4 --- /dev/null +++ b/rs/moq-uring/src/transport/held.rs @@ -0,0 +1,50 @@ +//! How the io_uring QUIC backend reads with connection credit. + +use std::task::{Context, Poll}; + +use bytes::Bytes; + +use crate::quic::{self as driver, web}; + +impl moq_net::transport::Credit for driver::Credit { + fn len(&self) -> u64 { + self.len() + } + + fn split(&mut self, bytes: u64) -> Self { + self.split(bytes) + } + + fn merge(&mut self, other: Self) { + self.merge(other) + } +} + +impl moq_net::transport::Stopper for web::Stopper { + fn stop(&self, code: u32) { + web::Stopper::stop(self, code) + } +} + +impl super::HeldRecvStream for web::RecvStream { + type Credit = driver::Credit; + type Stopper = web::Stopper; + + fn stopper(&self) -> web::Stopper { + web::RecvStream::stopper(self) + } + + fn poll_read_held( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Self::Error>> { + web::RecvStream::poll_read_held(self, cx, max) + } +} + +impl super::HeldSession for web::Session { + fn poll_credit_exhausted(&mut self, cx: &mut Context<'_>) -> Poll<()> { + web::Session::poll_credit_exhausted(self, cx) + } +} diff --git a/rs/moq-uring/tests/integration/credit.rs b/rs/moq-uring/tests/integration/credit.rs index 49cca1c4a90..4ff0167e326 100644 --- a/rs/moq-uring/tests/integration/credit.rs +++ b/rs/moq-uring/tests/integration/credit.rs @@ -1,7 +1,8 @@ //! Credit through the io_uring driver: held reads keep the connection window //! spent until their credit drops, bytes `poll_closed` reads ahead keep theirs -//! until the application reads them, and a chunked write takes no more from -//! its buffer than the peer's window admits. +//! until the application reads them, the connection reports each time held +//! credit fills the window, and a chunked write takes no more from its buffer +//! than the peer's window admits. //! //! Unsupported local kernels skip loudly; MOQ_STRICT makes CI refuse them. @@ -51,18 +52,18 @@ async fn within(handle: &moq_uring::Handle, what: &str, future: impl Future quic::Credit { - let mut buf = vec![0u8; len]; let mut held: Option = None; while held.as_ref().map_or(0, quic::Credit::len) < len as u64 { let filled = held.as_ref().map_or(0, |held| held.len() as usize); - let credit = within( + let (bytes, credit) = within( handle, "stream data", - std::future::poll_fn(|cx| recv.poll_read_held(cx, &mut buf[filled..])), + std::future::poll_fn(|cx| recv.poll_read_held(cx, len - filled)), ) .await .expect("read") .expect("not finished"); + assert_eq!(bytes.len() as u64, credit.len(), "credit covers exactly the bytes read"); match &mut held { Some(held) => held.merge(credit), None => held = Some(credit), @@ -148,14 +149,22 @@ fn held_reads_spend_the_connection_window() { .await .expect("accept_uni"); + let mut cx = Context::from_waker(Waker::noop()); + let mut exhausted = |conn: &mut quic::Connection| conn.poll_credit_exhausted(&mut cx).is_ready(); + assert!(!exhausted(&mut accepted)); + let mut held = read_held(&handle, &mut recv, WINDOW).await; let stats = accepted.stats().credit; assert_eq!((stats.held, stats.available), (WINDOW as u64, 0)); + assert!(exhausted(&mut accepted), "held credit filled the window"); + assert!(!exhausted(&mut accepted), "reported once per stall"); // Far below the usual update threshold, and still enough to unblock. drop(held.split(100)); + assert!(!exhausted(&mut accepted), "a release relieves the stall"); held.merge(read_held(&handle, &mut recv, 100).await); assert_eq!(accepted.stats().credit.held, WINDOW as u64); + assert!(exhausted(&mut accepted), "holding the room again stalls again"); // Read-ahead keeps the credit the application has not taken yet. drop(held); @@ -187,6 +196,94 @@ fn held_reads_spend_the_connection_window() { .expect("worker"); } +/// A read parked on a stream wakes when its `Stopper` stops it, and fails rather than +/// reading what the stop discarded as the stream's end, even after the closed-watch saw +/// the stop first. +#[test] +fn a_stopper_wakes_and_fails_a_parked_read() { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + #[derive(Default)] + struct Wakes(AtomicUsize); + + impl std::task::Wake for Wakes { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + let Some(mut worker) = worker() else { return }; + let handle = worker.handle(); + let certs = support::certs().expect("certificates"); + + let socket = |handle: &moq_uring::Handle| { + handle + .udp(UdpSocket::bind("127.0.0.1:0").expect("bind"), udp::Config::default()) + .expect("socket") + }; + let server_config = quic::server::Config::new( + quic::Identity::open(&certs.cert, &certs.key) + .expect("identity") + .server(&[ALPN], quic::server::ClientAuth::None) + .expect("tls"), + ); + let server = quic::Endpoint::new( + socket(&handle), + quic::endpoint::Config::default().with_server(server_config), + ) + .expect("server endpoint"); + let client = quic::Endpoint::new(socket(&handle), quic::endpoint::Config::default()).expect("dial-only endpoint"); + let mut dial = quic::client::Config::new(server.local_addr(), "localhost"); + dial.alpn = vec![ALPN.to_string()]; + dial.verify = false; + + worker + .block_on(async move { + let mut dialed = client.connect(&dial).await.expect("dial"); + let mut accepted = server.accept().await.expect("accept"); + + // Held open, so the stream never finishes on its own. + let mut send = std::future::poll_fn(|cx| dialed.poll_open_uni(cx)) + .await + .expect("open_uni"); + let written = within( + &handle, + "the write", + std::future::poll_fn(|cx| send.poll_write(cx, b"hello")), + ) + .await + .expect("write"); + let mut recv = within( + &handle, + "the stream", + std::future::poll_fn(|cx| accepted.poll_accept_uni(cx)), + ) + .await + .expect("accept_uni"); + read(&handle, &mut recv, written).await; + + let stopper = recv.stopper(); + let wakes = Arc::new(Wakes::default()); + let waker = Waker::from(wakes.clone()); + let mut cx = Context::from_waker(&waker); + let mut buf = [0u8; 64]; + assert!(recv.poll_read(&mut cx, &mut buf).is_pending()); + + stopper.stop(0); + assert_eq!(wakes.0.load(Ordering::Relaxed), 1, "the parked read was not woken"); + assert!(recv.poll_closed(&mut cx).is_ready(), "the closed-watch missed the stop"); + assert!( + matches!(recv.poll_read(&mut cx, &mut buf), Poll::Ready(Err(_))), + "a stopped stream read as finished" + ); + drop(send); + }) + .expect("worker"); +} + #[test] fn chunked_writes_stop_at_the_window() { const STREAM_WINDOW: usize = 4096; diff --git a/rs/moq-uring/tests/integration/main.rs b/rs/moq-uring/tests/integration/main.rs index 39715fff97e..83833a76650 100644 --- a/rs/moq-uring/tests/integration/main.rs +++ b/rs/moq-uring/tests/integration/main.rs @@ -15,6 +15,7 @@ mod ecn; mod endpoint; mod identity; mod qlog; +mod receive_credit; mod session; mod support; mod teardown; diff --git a/rs/moq-uring/tests/integration/receive_credit.rs b/rs/moq-uring/tests/integration/receive_credit.rs new file mode 100644 index 00000000000..085ed546aff --- /dev/null +++ b/rs/moq-uring/tests/integration/receive_credit.rs @@ -0,0 +1,223 @@ +//! Over the worker's QUIC path, incomplete frames hold the subscriber's connection +//! credit, and when they fill its window the subscriber stops the least urgent stream +//! and keeps going. +//! +//! The subscriber's connection window is far smaller than a moq session would allow, +//! which keeps the payloads small. Without recovery the high-priority frame never +//! arrives: the window is spent on a low-priority frame that can never complete. +//! +//! Unsupported local kernels skip loudly; MOQ_STRICT makes CI refuse them. + +#![cfg(all(target_os = "linux", feature = "quic"))] + +use crate::support; + +use std::net::UdpSocket; +use std::task::Poll; +use std::time::Duration; + +use bytes::Bytes; +use moq_net::{Timestamp, frame, origin, track::Subscription}; +use moq_uring::{Config, Error, Worker, quic, udp}; + +const ALPN: &str = "moq-lite-06"; +/// The subscriber's connection receive window. +const WINDOW: usize = 256 * 1024; + +fn worker() -> Option { + match Worker::new(Config::default()) { + Ok(worker) => Some(worker), + Err(Error::Unsupported(reason)) if std::env::var_os("MOQ_STRICT").is_none() => { + eprintln!("skipping io_uring receive credit test: {reason}"); + None + } + Err(err) => panic!("worker setup failed: {err}"), + } +} + +fn timed() -> moq_net::track::Info { + moq_net::track::Info::default().with_timescale(moq_net::Timescale::MILLI) +} + +/// Await `future`, failing loudly rather than hanging if it never resolves. +async fn within(handle: &moq_uring::Handle, what: &str, future: impl Future) -> T { + let mut deadline = moq_uring::Timer::after(handle, Duration::from_secs(10)); + let mut future = std::pin::pin!(future); + kio::wait(|waiter| { + let mut cx = waiter.context(); + if let Poll::Ready(value) = future.as_mut().poll(&mut cx) { + return Poll::Ready(Some(value)); + } + deadline.poll(waiter).map(|()| None) + }) + .await + .unwrap_or_else(|| panic!("timed out waiting for {what}")) +} + +/// The next frame's payload. +async fn next_frame(handle: &moq_uring::Handle, sub: &mut moq_net::track::Subscriber) -> Bytes { + let mut group = within(handle, "a group", sub.recv_group()) + .await + .expect("group") + .expect("track ended"); + within(handle, "a frame", group.read_frame()) + .await + .expect("frame") + .expect("group ended") + .payload +} + +/// Two partial frames that together exceed the window: the lower-priority one is +/// stopped, and the higher-priority one completes on the credit it returns. +#[test] +fn partial_frame_gives_way() { + let Some(mut worker) = worker() else { return }; + let handle = worker.handle(); + + // The model and its origins, driven on a tokio thread as in `session.rs`. + let (pub_origin, pub_driver) = origin::Producer::new(origin::Config::default()); + let (sub_origin, sub_driver) = origin::Producer::new(origin::Config::default()); + let origins = std::thread::spawn(move || { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .expect("tokio runtime"); + rt.block_on(async move { + tokio::join!(moq_net::time::run(pub_driver), moq_net::time::run(sub_driver)); + }); + }); + + let broadcast = pub_origin.create_broadcast("room").expect("create broadcast"); + broadcast.announce(Default::default()).expect("announce"); + let low = broadcast.create_track("low", timed()).expect("low track"); + let high = broadcast.create_track("high", timed()).expect("high track"); + + let certs = support::certs().expect("certificates"); + let server_config = quic::server::Config::new( + quic::Identity::open(&certs.cert, &certs.key) + .expect("identity") + .server(&[ALPN], quic::server::ClientAuth::None) + .expect("tls"), + ); + let socket = |handle: &moq_uring::Handle| { + handle + .udp(UdpSocket::bind("127.0.0.1:0").expect("bind"), udp::Config::default()) + .expect("socket") + }; + let server_sock = socket(&handle); + let mut dial = quic::client::Config::new(server_sock.local_addr().expect("server addr"), "localhost"); + dial.alpn = vec![ALPN.to_string()]; + dial.verify = false; + dial.transport.receive_window = Some(WINDOW as u64); + let client_sock = socket(&handle); + + let server_handle = handle.clone(); + handle.spawn(async move { + let conn = quic::server::accept(server_sock, &server_config) + .await + .expect("quic accept"); + let (session, driver) = moq_net::Server::new() + .with_publisher(&pub_origin) + .accept_lite( + moq_net::time::now(), + moq_uring::transport::Session::new(quic::web::Session::raw(conn), moq_net::transport::Binding::InBand), + ) + .await + .expect("accept_lite"); + let _ = server_handle.run(driver).await; + session.closed().await; + }); + + let sub = sub_origin.clone(); + worker + .block_on(async move { + let conn = quic::client::connect(client_sock, &dial).await.expect("quic connect"); + let (session, driver) = moq_net::Client::new() + .with_subscriber(sub.clone()) + .connect_lite( + moq_net::time::now(), + moq_uring::transport::Session::new( + quic::web::Session::raw(conn), + moq_net::transport::Binding::InBand, + ), + ) + .await + .expect("connect_lite"); + let task_handle = handle.clone(); + handle.spawn(async move { + let _ = task_handle.run(driver).await; + }); + + let consumer = sub.consume(); + within(&handle, "the route", consumer.routed("room")) + .await + .expect("routed"); + let remote = within(&handle, "the broadcast", consumer.request_broadcast("room", None)) + .await + .expect("broadcast resolves"); + let subscribe = |name: &str, priority: u8| { + let track = remote.track(name).expect("track"); + async move { + track + .subscribe(Subscription::default().with_priority(priority)) + .await + .expect("subscribe") + } + }; + let mut low_sub = within(&handle, "the low subscription", subscribe("low", 1)).await; + let mut high_sub = within(&handle, "the high subscription", subscribe("high", 200)).await; + + // Three quarters of the window, held by a frame that never finishes. + let mut low_group = low.append_group().expect("low group"); + let mut low_frame = low_group + .create_frame(frame::Info { + size: WINDOW as u64, + timestamp: Some(Timestamp::ZERO), + }) + .expect("low frame"); + low_frame.write(vec![1u8; WINDOW * 3 / 4]).expect("partial payload"); + + // Wait until the subscriber holds all of it, so nothing else got there first. + let mut partial = within(&handle, "the low group", low_sub.recv_group()) + .await + .expect("group") + .expect("track ended"); + let mut partial = within(&handle, "the low frame header", partial.next_frame()) + .await + .expect("frame") + .expect("group ended"); + let mut seen = 0; + while seen < WINDOW * 3 / 4 { + seen += within(&handle, "the partial payload", partial.read_chunk()) + .await + .expect("chunk") + .expect("frame ended") + .len(); + } + + // Half the window, which only fits once the low frame lets go. + let payload = Bytes::from(vec![2u8; WINDOW / 2]); + let mut high_group = high.append_group().expect("high group"); + high_group + .write_frame(Timestamp::ZERO, payload.clone()) + .expect("high frame"); + high_group.finish().expect("finish high"); + assert_eq!(next_frame(&handle, &mut high_sub).await, payload); + + // The session survives, and the low track keeps delivering what fits. + drop(low_frame); + drop(low_group); + let mut group = low.append_group().expect("after group"); + group.write_frame(Timestamp::ZERO, &b"after"[..]).expect("after frame"); + group.finish().expect("finish after"); + assert_eq!(next_frame(&handle, &mut low_sub).await, &b"after"[..]); + + session.abort(moq_net::Error::Cancel); + }) + .expect("worker"); + + drop(worker); + drop(broadcast); + drop(sub_origin); + origins.join().expect("origin drivers"); +} diff --git a/rs/moq-uring/tests/integration/web.rs b/rs/moq-uring/tests/integration/web.rs index b7b6635b315..057b06ecc2f 100644 --- a/rs/moq-uring/tests/integration/web.rs +++ b/rs/moq-uring/tests/integration/web.rs @@ -82,8 +82,8 @@ type ClientFuture = std::pin::Pin + Send>>; async fn read_all(recv: &mut moq_tokio::quic::web::RecvStream) -> Vec { let mut out = Vec::new(); let mut buf = [0u8; 4096]; - while let Some(n) = recv.read(&mut buf).await.expect("read") { - out.extend_from_slice(&buf[..n]); + while let Some(credit) = recv.read(&mut buf).await.expect("read") { + out.extend_from_slice(&buf[..credit.len() as usize]); } out } @@ -495,7 +495,7 @@ fn a_dropped_stream_carries_a_webtransport_code() { // The server writes this and then drops the stream unfinished. let mut recv = session.accept_uni().await.expect("accept_uni"); let mut buf = [0u8; 4096]; - let n = recv.read(&mut buf).await.expect("read").expect("payload"); + let n = recv.read(&mut buf).await.expect("read").expect("payload").len() as usize; assert_eq!(&buf[..n], PAYLOAD); // Tell it we have the payload, so the reset below cannot race it. diff --git a/rs/moq-wasm/src/transport/adapter.rs b/rs/moq-wasm/src/transport/adapter.rs index 46a2c8ff250..72fe88c7ad6 100644 --- a/rs/moq-wasm/src/transport/adapter.rs +++ b/rs/moq-wasm/src/transport/adapter.rs @@ -107,6 +107,9 @@ impl + Clone + 'st fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll { backend::Session::poll_closed(&mut self.0, cx).map(Error) } + fn poll_credit_exhausted(&mut self, _: &mut Context<'_>) -> Poll<()> { + Poll::Pending + } fn stats(&self) -> impl moq::Stats { Stats(backend::Session::stats(&self.0)) } @@ -141,20 +144,39 @@ impl poll::SendStream for SendStream { } /// An incoming WebTransport stream adapted to moq-net. pub struct RecvStream(S); +// The WebTransport backend returns connection credit as it reads, so there is none to hold. impl poll::RecvStream for RecvStream { type Error = Error; - fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { - backend::RecvStream::poll_read(&mut self.0, cx, dst).map_err(Error) + type Credit = moq::Unheld; + type Stopper = (); + fn stopper(&self) {} + fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll, Self::Error>> { + backend::RecvStream::poll_read(&mut self.0, cx, dst) + .map_err(Error) + .map_ok(|n| n.map(|n| moq::Unheld(n as u64))) } fn poll_read_buf( &mut self, cx: &mut Context<'_>, buf: &mut B, - ) -> Poll, Self::Error>> { - backend::RecvStream::poll_read_buf(&mut self.0, cx, buf).map_err(Error) + ) -> Poll, Self::Error>> { + backend::RecvStream::poll_read_buf(&mut self.0, cx, buf) + .map_err(Error) + .map_ok(|n| n.map(|n| moq::Unheld(n as u64))) } - fn poll_read_chunk(&mut self, cx: &mut Context<'_>, max: usize) -> Poll, Self::Error>> { - backend::RecvStream::poll_read_chunk(&mut self.0, cx, max).map_err(Error) + fn poll_read_chunk( + &mut self, + cx: &mut Context<'_>, + max: usize, + ) -> Poll, Self::Error>> { + backend::RecvStream::poll_read_chunk(&mut self.0, cx, max) + .map_err(Error) + .map_ok(|chunk| { + chunk.map(|b| { + let n = b.len() as u64; + (b, moq::Unheld(n)) + }) + }) } fn stop(&mut self, code: u32) { backend::RecvStream::stop(&mut self.0, code) @@ -309,7 +331,7 @@ mod tests { let Poll::Ready(Ok(Some(chunk))) = poll::RecvStream::poll_read_chunk(&mut recv, &mut cx, 2) else { panic!("read failed") }; - assert_eq!(chunk, b"hi"[..]); + assert_eq!(chunk.0, b"hi"[..]); poll::RecvStream::stop(&mut recv, 13); assert_eq!(state.get(), 13); let Poll::Ready(Err(error)) = poll::SendStream::poll_closed(&mut send, &mut cx) else { diff --git a/sh/gh/nightly.sh b/sh/gh/nightly.sh index 2e42f51e400..99fb385439f 100755 --- a/sh/gh/nightly.sh +++ b/sh/gh/nightly.sh @@ -25,6 +25,8 @@ case "${1:-}" in --run-ignored only --success-output immediate run just rs bench-feed-fetch run just rs test -p moq-quic --release -E 'test(benchmark_receive_credit)' --run-ignored only --success-output immediate + run just rs test -p moq-tokio --release -E 'test(benchmark_held_reads)' --run-ignored only --success-output immediate + run just rs test -p moq-net --release -E 'test(benchmark_victim) | test(benchmark_tracked_reads)' --run-ignored only --success-output immediate run just rs test -p moq-relay -E 'test(benchmark_stun_sources)' --run-ignored only --success-output immediate run just rs test -p moq-archive -p moq-mux -p moq-hls --no-fail-fast \ -E 'test(benchmark_timeline_fetch_fanout) | test(benchmark_damage_pid_fanout) | test(benchmark_pes_staging) | test(benchmark_reference_handoff_fanout)' \