diff --git a/Cargo.lock b/Cargo.lock index c4085febfb5..128570fd142 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2945,6 +2945,7 @@ dependencies = [ "anyhow", "bytes", "derive_more 2.1.1", + "hang", "hex", "kio", "moq-json", @@ -4549,6 +4550,7 @@ dependencies = [ "data-encoding", "hex", "kio", + "moq-crypto", "moq-net", "rind", "rind-rustls", @@ -4636,6 +4638,7 @@ dependencies = [ "kio", "m3u8-rs", "moq-archive", + "moq-hls", "moq-json", "moq-mux", "moq-net", @@ -4680,6 +4683,7 @@ dependencies = [ "futures", "json-patch", "kio", + "moq-json", "moq-net", "proptest", "serde", @@ -4730,6 +4734,7 @@ dependencies = [ "moq-json", "moq-loc", "moq-msf", + "moq-mux", "moq-net", "mp4-atom", "mpeg2ts", @@ -4957,6 +4962,7 @@ dependencies = [ "moq-auth 0.2.1", "moq-json", "moq-net", + "moq-room", "moq-tokio", "thiserror 2.0.21", "tokio", @@ -5107,6 +5113,7 @@ dependencies = [ "kio", "moq-json", "moq-net", + "moq-stats", "serde", "serde_json", "serde_with", diff --git a/doc/lib/rs/hang.md b/doc/lib/rs/hang.md index 28e3e127089..1662ba86bf0 100644 --- a/doc/lib/rs/hang.md +++ b/doc/lib/rs/hang.md @@ -20,6 +20,9 @@ carries a timestamp with each payload. cargo add hang ``` +The `poll_*` methods work by default. Enable the `async` Cargo feature to await +`Catalog::subscribe` and the catalog consumer's `next()` instead. + hang is the format, not the pipeline. [`moq-mux`](/lib/rs/moq-mux) produces the catalog for you and decodes any container in group order (`moq_mux::container::Consumer`). Its `Delay { min, max }` says how far behind diff --git a/doc/lib/rs/moq-json.md b/doc/lib/rs/moq-json.md index 84e0611a307..f678c6fd095 100644 --- a/doc/lib/rs/moq-json.md +++ b/doc/lib/rs/moq-json.md @@ -13,6 +13,9 @@ stream are the catalog modes on the [hang](/concept/hang#data-tracks) page. Window is a third framing for a raw track: both ends opt into it, and a generic catalog reader will not pick it up. +The `poll_*` methods work by default. Enable the `async` Cargo feature to await +`next()`, `latest()`, or `refresh()` instead. + - **snapshot**: lossy latest value, with RFC 7396 merge-patch deltas. An unchanged update publishes nothing. Several owners can edit their own keys in place instead of clobbering one document. Deltas ride one group until a joiner wants a fresh one: set the track's `refresh` ratio (8 is typical) and drive `Producer::refresh`, which restates the value in a new group. Without a ratio a group rolls only at its caps. - **stream**: lossless append-log in a single group. A reader that falls behind fails the read rather than resuming mid-log. - **window**: a bounded run of records a reader can join at any point. A new group restates what it keeps explicitly, so a reader that was keeping up is not handed a record twice. Its ops roll on the track's `refresh` ratio the same way. diff --git a/doc/lib/rs/moq-mux.md b/doc/lib/rs/moq-mux.md index 7e1da45c713..eef5d2f641d 100644 --- a/doc/lib/rs/moq-mux.md +++ b/doc/lib/rs/moq-mux.md @@ -61,5 +61,9 @@ only when they share that clock. cargo add moq-mux ``` +The `poll_*` methods work by default. Enable the `async` Cargo feature to await +the exporters' `next()`, a container's `read()`, and the async constructors such +as `catalog::Consumer::new` and `Source::catalog` instead. + API: [docs.rs/moq-mux](https://docs.rs/moq-mux). The commands that exercise it are [`moq-cli`](/bin/cli). diff --git a/doc/lib/rs/moq-net.md b/doc/lib/rs/moq-net.md index e96543e9395..634ebe17553 100644 --- a/doc/lib/rs/moq-net.md +++ b/doc/lib/rs/moq-net.md @@ -25,14 +25,15 @@ What you use it for, beyond what the concept page already describes: - **Authored tracks default to a zero `max_age`.** `track::Info::default()` sets `max_age` to `Some(Duration::ZERO)`, the live edge only; a publisher that wants readers to reach history sets a longer bound, or `with_max_age(None)` for no publisher limit. Media tracks from `hang::container::track_info` declare `hang::container::MAX_AGE` (30 seconds). `max_age` is how long a superseded group stays fresh, never the newest group; it frees no bytes, which the cache pool's expiry reclaims. Received metadata keeps absence as no limit, and an origin's `cache_duration` can only cap it lower. - **End content by finishing tracks.** A broadcast has no finished state. When the input ends, finish each track; the broadcast stays announced and keeps serving them until you close it or drop its last producer. How long to keep it up for late readers is yours to decide: `moq_mux::Linger` holds a hang broadcast for a fixed window. - **Demand says whether anyone is watching.** `demand()` on a track, group, or broadcast says whether a subscriber is attached. First-party capture does not wait for it: video encodes while enabled and audio keeps raw PCM to encode when someone listens, so a late subscription's `delay.max` or a FETCH finds the recent past. Demand gates bandwidth reservations, transcode rungs, and on-demand video where a publisher opts in, which saves the encode but leaves nothing from before the first subscriber. Holding a broadcast or track consumer is not demand: a track consumer requests nothing until it queries, subscribes, or fetches, and those count until dropped. A shared fetch stays up until its last reader leaves. Demand rolls up: a frame being read keeps its group wanted, a group a subscription or fetch handed out keeps its track wanted until read to its end, and any wanted track keeps its broadcast wanted. -- **Use the model without a runtime.** Disable default features on `moq-net` to use its `poll_*` APIs and origin driver with caller-supplied time. The default `async` feature adds async helpers, `time::run`, and both wire sessions. +- **Use the model without a runtime.** By default `moq-net` is its `poll_*` APIs and origin driver with caller-supplied time. The `async` feature adds their async wrappers, `time::run`, and both wire sessions. - **You drive the session, or `moq-tokio` does.** `connect` and `accept` return a session plus a driver that never reads the clock itself. `moq_net::time::run` polls it on tokio. `moq-tokio` hides that. A custom transport implements `moq_net::transport::poll`, and `moq_net::transport::memory::pair` links two sessions in one process with no sockets, for tests or in-process links. - **A session bounds what its peer can hold.** Each request holds a stream, so the transport's stream credit bounds subscriptions and FETCHes. `session::Limits` (via `Client::with_limits` and `Server::with_limits`) bounds the announcements per session. The default suits a relay mesh, so lower it for untrusted peers. Past it, the session closes with `TOO_MANY_REQUESTS`. Peer-declared lengths are capped before they are buffered. - **Authorization is in band.** On moq-lite 07 (`moq-lite-07-wip`, opt-in) and on moq-transport draft-18+ with the [MoQ Auth extension](/draft/moq-auth), `session.auth()` presents tokens without reconnecting and watches the grant the peer sent back. `add(token)` resolves with the token and how the peer took it: `Answer::Granted` on its AUTH stream, or `Answer::Requests` on moq-transport without AUTH, where only a request that names it carries it. A handle passed to `Client::with_auth` outlives each session, so its tokens, and the routes naming them, carry across your own reconnects. A server verifies tokens itself by taking `handshake.requests()` before `ok()`. `authorize(&grant)` narrows or widens a live session on any version: what falls outside resets with `Unauthorized` and the session stays up. A client that publishes outside its grant closes the session, naming the path. - **A graceful close waits.** `session.close().await` withdraws announcements and drains for up to one second. `moq-lite-05` and newer wait for the subscriber to receive finished tails; older Lite versions and IETF wait only for transport ACKs. `abort` ends immediately. ```bash -cargo add moq-net moq-tokio +cargo add moq-net --features async +cargo add moq-tokio ``` See the [Rust quick start](/lib/rs/#quick-start) and diff --git a/quest/m1/README.md b/quest/m1/README.md index 19d7bfe6fa3..e5a8e4b22e6 100644 --- a/quest/m1/README.md +++ b/quest/m1/README.md @@ -176,7 +176,6 @@ their end report, so an embedder's drain alert pages on real cuts. - [Catalog contributions become includes on the root's timeline](/quest/m1/catalog-include.md) - rename the unreleased `contributions` map before the cut and state that included broadcasts share the root's timeline - [Unbounded groups](/quest/m1/unbounded-groups/README.md) - groups and frames lose their size caps; a group keeps a 32 MB tail, and readers behind it replay through paced flights or refresh - [Session policy lives in moq-net](/quest/m1/session-policy/README.md) - accept and connect deadlines, admission, window widening and reconnect live once in moq-net; moq-tokio and moq-relay's io_uring path call them -- [Poll-first async](/quest/m1/poll-first/README.md) - library crates compile no coroutines by default: async wrappers are generated from `poll_*` and gated behind an `async` feature - [moq auth serve scopes each client certificate by its chain](/quest/m1/auth-cert-rules.md) - per-certificate and per-intermediate rules in a reloaded TOML config, keyed by the chain hashes the relay sends - [A replacement that fails before delivering never ends a live track](/quest/m1/route-swap-keeps-incumbent.md) - a route change whose replacement aborts before its first group leaves readers on the source that was serving - [A Rust relay never replays from group 0 to a lite-03/04 subscriber](/quest/m1/relay-inflight-replay.md) - A Rust relay or publisher serving a lite-03 or lite-04 SUBSCRIBE from a track copy that is still resolving starts at the newest group once the copy resolves diff --git a/quest/m1/poll-first/README.md b/quest/m1/poll-first/README.md deleted file mode 100644 index 738e8a26822..00000000000 --- a/quest/m1/poll-first/README.md +++ /dev/null @@ -1,77 +0,0 @@ -# Poll-first async - -## Goal - -moq's library crates write asynchronous logic as kio `poll_*` functions and -compile no coroutines by default. Their `async fn` wrappers are generated and -behind an `async` cargo feature that binaries, tests, and tokio glue enable. A -design report says how far to take it: whether a kio proc macro should write -`poll_*` bodies from straight-line code, and whether moq-net's protocol core -should become sans-IO. Everything stays interoperable with the Rust ecosystem at -the edges, through std's `Waker` and `Context`. - -## Plan - -Decided with the maintainer on 2026-10-25, from a survey of async Rust's problems -and their alternatives: - -- **Problems:** - - function coloring; - - `Pin` and `Unpin`; - - cancel safety and no async drop; - - `Send` inference and the `Send` bound problem; - - `'static` spawn; - - runtime coupling; - - future size (a 32 MiB test stack in moq-relay); - - debugging; - - compile cost: rmeta encodes every `async fn` body's coroutine MIR, so any - body edit changes the crate's metadata and rebuilds its dependents. -- **Upstream won't fix them soon:** - - return type notation stalled in Dec 2025; - - pin ergonomics was discontinued in Aug 2026; - - async drop is nightly-only, and the 2026 roadmap keeps destructors sync; - - scoped spawn waits on guaranteed destructors (2026-27). -- **kio already avoids:** `Pin` in user code, hidden cancellation points (progress - lives in named fields), leaked wakers, runtime coupling in the library, and - per-future allocation in `Tasks`. -- **What kio keeps:** coloring (`poll_x` plus `async fn x`), `'static` spawn, - hand-written state machines, and 184 non-test `async fn` - wrappers in moq-net that still compile into coroutines. -- **Rejected:** stackful coroutines (green threads). They break wasm, rust2's MIR - TypeScript translation, and thread-local soundness. -- **Zig's model** (I/O passed as an explicit `std.Io`, like an allocator) removes - coloring only because its evented backend parks whole stacks. kio's `&Waiter` - parameter is the stackless half of that idea. Sans-IO is the other: the protocol - never does I/O and each runtime drives it. -- **Sans-IO compared** in [moq-dev/moq#7058](https://github.com/moq-dev/moq/pull/7058), - on the lite subscriber's upstream SUBSCRIBE (an unmerged prototype, since deleted). - moq-net already takes I/O through its transport traits and runs on any executor, - so sans-IO only buys a core without waiters. That core translated through rust2-ts - with 4 refusals, but the serve loop driving it still had 919, so translating a - session still waits on the model. Rules became plain `#[test]`s, while - interleaving bugs stay in driver tests. Compile cost and the benches were flat. - The cost is an effect enum per model touch (+206 lines on this path). - **Decided (maintainer, 2026-10-28):** moq-net's protocol core stays kio poll and - does not move toward sans-IO. Pull pure wire rules into plain functions where - that turns scripted async tests into plain ones, as ordinary refactors. -- **Gate the wrappers regardless** of the build measurement, because poll-first is - the style moq wants. moq-dev/rust2's - [semantic API hash quest](https://github.com/moq-dev/rust2/blob/main/quest/m3/api-cutoff.md) - replays moq-net's history with and without the wrappers, to report the build - effect. -- **Adopt `#[kio::machine]` for straight-line code only**, decided with the - maintainer on 2026-10-28 from the spike in - [moq-dev/moq#7057](https://github.com/moq-dev/moq/pull/7057). It lands with its - first users; loop support waits until conversions show it's needed. The spike - converted TrackInfoFetch (108 to 78 lines, no `mem::replace` transitions). Its - type-check and borrow-check cost matches hand-written code; it adds about 1.5 ms - of expansion per machine plus a one-time 1.3 s build of `kio-macros`, below what - async costs. The cost is a dialect: stored locals need type annotations and are - `&mut` inside a wait's poll, and a `!Send` error names the type and the wait but - not the local: the generated state keeps Rust's automatic `Send`, with no - `unsafe impl` (decided 2026-10-28). Its users are moq-net's 23 straight-line multi-await `async fn` bodies and 12 straight-line - hand-written machines; 59 bodies await in loops or selects. - -## Required - -- [Generated async wrappers behind a feature](/quest/m1/poll-first/async-feature.md) - each library's `async fn` wrappers are generated from its `poll_*` signatures and compile only with its `async` feature diff --git a/quest/m1/poll-first/async-feature.md b/quest/m1/poll-first/async-feature.md deleted file mode 100644 index fd71c7603ad..00000000000 --- a/quest/m1/poll-first/async-feature.md +++ /dev/null @@ -1,68 +0,0 @@ -# [L] Generated async wrappers behind a feature - -## Goal - -Every `async fn x(..) { kio::wait(|w| self.poll_x(w, ..)).await }` wrapper in -moq's library crates (kio, moq-net, and their poll-based siblings) is generated -from its `poll_*` method, not written by hand, and compiles only with the crate's -`async` cargo feature. With the feature off, the crate compiles no coroutines -(`async fn` bodies). moq's binaries, tests, tokio glue, and the bindings enable it, -and the bindings and JS keep their APIs. - -## Plan - -Decided with the maintainer on 2026-10-25: poll is canonical, async is the -ergonomic layer, and the gating proceeds regardless of the build measurement. - -Decided with the maintainer on 2026-11-01, from the -[generator spike](https://github.com/moq-dev/moq/blob/91b1ec62e1e35295ef74d0f7cee1afc49739bf00/bench/async-wrappers/RESULTS.md): -the generator is a method attribute in kio-macros, `#[kio::wrapper]`. The poll -stays ordinary Rust, the wrapper's name, signature, and docs are derived from it, -and a standalone kio build with `async` pays the measured proc-macro cost (about -3.5 seconds clean). - -Done in [moq-dev/moq#7506](https://github.com/moq-dev/moq/pull/7506): -`#[kio::wrapper]` with its contract in kio-macros' behavior and UI tests; kio's -default-off `async` feature, under which kio compiles no coroutines; moq-net's -`async` enabling `kio/async`; moq-net's four `broadcast` wrappers generated; and -`just rs coroutines`, in the nightly feature matrix, which scans the MIR of kio's -and moq-net's poll-only builds. Callers gate the attribute with -`#[cfg_attr(feature = "async", kio::wrapper)]`, so a poll-only build skips the -proc macro, and the scan catches a wrapper left ungated. - -What remains: - -- **moq-net's remaining wrappers:** about 47 hand-written - `kio::wait(|waiter| self.poll_x(waiter, ..)).await` wrappers become the attribute - on their polls, each poll's docs rewritten to read for both forms. A few stay - hand-written: those that take a closure by value and poll it by `&mut`, and the - `#[cfg(test)]` helpers the poll-only tests use. -- **moq-net's poll-only coroutines:** `auth::until` (behind `Grant::guard`), and - `model::origin`'s `run_front` and `serve_front`, still compile without `async`. - Gate `until` with its async callers, or make it a poll, and turn the front - driver into a hand-written poll machine (its body is one large `kio::wait`). - Then drop their `--allow`s from `just rs coroutines`. -- **Logic in `async fn` bodies:** an `async fn` that does more than call one - `poll_*` becomes a `poll_*` first, or moves out of the library to the edges. List - each one in the PR. moq-net's straight-line ones became `#[kio::machine]`s in - [moq-dev/moq#7284](https://github.com/moq-dev/moq/pull/7284), except those - that await a looped body with no `poll_*` yet, such as `Server::accept`; the - bodies with a loop or a branch around a wait (about 59 in moq-net) are - converted by hand, since the macro stays straight-line only (decided - 2026-10-28). -- **Sibling libraries:** the poll-based libraries with hand-written wrappers - (moq-mux, moq-hls, moq-auth, moq-stats, moq-room, moq-json, moq-crypto, - qmux, and others) gain an `async` feature, generate their wrappers, and join - `just rs coroutines`. -- **Default off:** moq-net's `async` feature is still on by default. Turn it off, - and enable it at the edges: binaries, tests, tokio glue, and the bindings. -- **Tests:** `#[kio::test]` and `#[tokio::test]` bodies stay async, so test - targets enable the feature. -- **Cargo:** feature unification means one dependent enabling `async` enables it - for the whole build. Make sure library-only builds, such as `cargo check -p - moq-net` and the rust2 bench's library-only build, leave it off. - -Public API: every library crate gains an `async` feature, and its async wrappers -move behind it. Add an Impact migration note for the published crates whenever -a consumer must enable `async`, including with `default-features = false`, as -`rs/kio/MIGRATE.md` does for kio. diff --git a/quest/m1/session-policy/accept.md b/quest/m1/session-policy/accept.md index 44e7d35c7a2..c32b085d28b 100644 --- a/quest/m1/session-policy/accept.md +++ b/quest/m1/session-policy/accept.md @@ -38,7 +38,3 @@ see it for the decisions and research. Public API: breaking in moq-net (`Server`, `Handshake`, the transport session trait), with MIGRATE notes. Wire: none. - -## Required - -- [Generated async wrappers](/quest/m1/poll-first/async-feature.md) - gives the handshake a poll form these machines wait on diff --git a/rs/hang/Cargo.toml b/rs/hang/Cargo.toml index cba18cab5eb..219b305193e 100644 --- a/rs/hang/Cargo.toml +++ b/rs/hang/Cargo.toml @@ -12,6 +12,14 @@ rust-version.workspace = true keywords = ["quic", "http3", "webtransport", "media", "live"] categories = ["multimedia", "network-programming", "web-programming"] +[package.metadata.docs.rs] +features = ["async"] + +[features] +default = [] +# Generated async wrappers for the poll methods, and `Catalog::subscribe`. Enable where you `.await` them. +async = ["kio/async"] + [dependencies] bytes = { workspace = true } hex = { workspace = true } @@ -30,7 +38,9 @@ features = ["from", "display", "debug"] [dev-dependencies] anyhow = { workspace = true } -moq-mux = { path = "../moq-mux" } +# Compile the async API used by the tests. +hang = { path = ".", features = ["async"] } +moq-mux = { path = "../moq-mux", features = ["async"] } moq-tokio = { path = "../moq-tokio", default-features = false, features = ["quic", "websocket"] } tokio = { workspace = true, features = ["macros", "fs"] } tracing = { workspace = true } diff --git a/rs/hang/MIGRATE.md b/rs/hang/MIGRATE.md index 2000b63d707..2a99b33c714 100644 --- a/rs/hang/MIGRATE.md +++ b/rs/hang/MIGRATE.md @@ -4,6 +4,14 @@ How to update your code for each breaking change, newest first. ## Unreleased +### The async API is opt-in + +Enable the `async` Cargo feature to keep awaiting `catalog::Consumer::next` and `Catalog::subscribe`: +`hang = { version = "...", features = ["async"] }`. This also applies with `default-features = false`. +The feature is off by default; `catalog::Consumer::poll_next` stays available without it. +The generated `next` no longer requires `Catalog: Unpin`. +The catalog format and JS and binding APIs are unchanged. + ### The catalog track advertises a zero max age ([#7729](https://github.com/moq-dev/moq/pull/7729)) `Catalog::default_track_info` takes moq-net's new zero default, so the catalog track advertises the live edge only where it advertised no limit. A late joiner still reads the newest snapshot, which never goes stale. diff --git a/rs/hang/src/catalog/consumer.rs b/rs/hang/src/catalog/consumer.rs index eebd5b48726..20ac6b6768f 100644 --- a/rs/hang/src/catalog/consumer.rs +++ b/rs/hang/src/catalog/consumer.rs @@ -31,7 +31,7 @@ impl Consumer { } } - /// Poll for the next catalog update. + /// Get the next catalog update, or `None` when the track ends. /// /// Fails with [`Error::EscapingBroadcast`](crate::Error::EscapingBroadcast) if any rendition's /// `broadcast` reference, contribution, or echo prefix walks above the root. Such a reference names no broadcast, and the @@ -40,6 +40,7 @@ impl Consumer { /// publisher that emits one has a bug (or is hostile), and silently serving the rest hides /// that from the operator while the missing rendition surfaces much later as a track that /// never fills. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>>> { let value = ready!(self.inner.poll_latest(waiter)) .map_err(|err| match err { @@ -55,14 +56,6 @@ impl Consumer { } Poll::Ready(Ok(catalog)) } - - /// Wait for the next catalog update, or return `None` when the track ends. - pub async fn next(&mut self) -> Result>> - where - Catalog: Unpin, - { - kio::wait(|waiter| self.poll_next(waiter)).await - } } /// Reject the catalog if a track reference, contribution, or echo prefix walks above the root. diff --git a/rs/hang/src/catalog/root.rs b/rs/hang/src/catalog/root.rs index 73b752b1819..495fda5738b 100644 --- a/rs/hang/src/catalog/root.rs +++ b/rs/hang/src/catalog/root.rs @@ -344,6 +344,7 @@ impl FromStr for Catalog { impl Catalog { /// Subscribe to catalog updates on a broadcast. + #[cfg(feature = "async")] pub async fn subscribe(broadcast: &moq_net::broadcast::Consumer) -> Result> { Ok(super::Consumer::new(broadcast.track(Catalog::DEFAULT_NAME)?)) } diff --git a/rs/justfile b/rs/justfile index f9246effc3d..552fe6ff9ea 100644 --- a/rs/justfile +++ b/rs/justfile @@ -264,11 +264,25 @@ net-features *args: bash sh/rs/net-features.sh "$@" # Prove the poll-only library builds compile no coroutines and document cleanly. -# moq-net allows the async fns still waiting on a poll conversion. coroutines: bash sh/rs/coroutines.sh kio -- --no-default-features --features macros,runtime,manual,serde RUSTDOCFLAGS="-D warnings" bash sh/local.sh cargo doc --locked -p kio --no-deps --no-default-features --features macros,runtime,manual,serde - bash sh/rs/coroutines.sh moq-net --allow until --allow run_front --allow serve_front -- --no-default-features + bash sh/rs/coroutines.sh moq-net -- --no-default-features + bash sh/rs/coroutines.sh moq-json -- --no-default-features + RUSTDOCFLAGS="-D warnings" bash sh/local.sh cargo doc --locked -p moq-json --no-deps --no-default-features + bash sh/rs/coroutines.sh hang -- --no-default-features + RUSTDOCFLAGS="-D warnings" bash sh/local.sh cargo doc --locked -p hang --no-deps --no-default-features + bash sh/rs/coroutines.sh moq-mux -- --no-default-features + RUSTDOCFLAGS="-D warnings" bash sh/local.sh cargo doc --locked -p moq-mux --no-deps --no-default-features + bash sh/rs/coroutines.sh moq-hls -- --no-default-features + RUSTDOCFLAGS="-D warnings" bash sh/local.sh cargo doc --locked -p moq-hls --no-deps --no-default-features + bash sh/rs/coroutines.sh moq-room -- --no-default-features + RUSTDOCFLAGS="-D warnings" bash sh/local.sh cargo doc --locked -p moq-room --no-deps --no-default-features + bash sh/rs/coroutines.sh moq-stats -- --no-default-features + RUSTDOCFLAGS="-D warnings" bash sh/local.sh cargo doc --locked -p moq-stats --no-deps --no-default-features + bash sh/rs/coroutines.sh moq-crypto -- --no-default-features --features sign,e2ee,provider,serde + RUSTDOCFLAGS="-D warnings" bash sh/local.sh cargo doc --locked -p moq-crypto --no-deps --no-default-features --features sign,e2ee,provider,serde + bash sh/rs/coroutines.sh rs/moq-auth/Cargo.toml -- --no-default-features # Prove moq-crypto's key-only build, which moq-auth depends on, leaves out the transport and aws-lc-rs, then lint and test it alone; a workspace build unifies the default features back in. crypto-key: diff --git a/rs/kio/src/clock/timer.rs b/rs/kio/src/clock/timer.rs index c296390792f..f9cf5b4b94f 100644 --- a/rs/kio/src/clock/timer.rs +++ b/rs/kio/src/clock/timer.rs @@ -20,6 +20,10 @@ pub struct Sleep { parked: Option, } +#[allow( + clippy::large_enum_variant, + reason = "one per sleep, and the size gap varies with the clock features enabled; an allow, since some builds balance it" +)] enum Timer { Queue(crate::clock::deadline::Timer), Host(crate::clock::runtime::Timer), diff --git a/rs/moq-archive/Cargo.toml b/rs/moq-archive/Cargo.toml index 3ac1a856cfe..adc0dfeee9c 100644 --- a/rs/moq-archive/Cargo.toml +++ b/rs/moq-archive/Cargo.toml @@ -16,7 +16,7 @@ categories = ["multimedia", "filesystem"] bytes = { workspace = true } futures = { workspace = true } kio = { workspace = true, features = ["runtime"] } -moq-net = { workspace = true } +moq-net = { workspace = true, features = ["async"] } object_store = { workspace = true } percent-encoding = { workspace = true } quick_cache = "0.7" diff --git a/rs/moq-audio/Cargo.toml b/rs/moq-audio/Cargo.toml index bf6392e84c0..9ec534937ba 100644 --- a/rs/moq-audio/Cargo.toml +++ b/rs/moq-audio/Cargo.toml @@ -77,8 +77,8 @@ bytes = { workspace = true } # without either pulls no new dependency. cpal = { version = "0.18", optional = true } hang = { workspace = true } -kio = { workspace = true, features = ["runtime"] } -moq-mux = { workspace = true } +kio = { workspace = true, features = ["async", "runtime"] } +moq-mux = { workspace = true, features = ["async"] } moq-net = { workspace = true } rand = { workspace = true } # Fixed-capacity SPSC rings keep buffer handoff lock-free on the capture @@ -139,7 +139,7 @@ objc2-screen-capture-kit = { workspace = true, features = ["SCStream", "SCSharea [dev-dependencies] criterion = { package = "bench-harness", path = "../bench-harness" } # An HLS-only viewer starting an on-demand capture. -moq-hls = { workspace = true } +moq-hls = { workspace = true, features = ["async"] } # Reads the jitter conformance corpus at doc/concept/audio-jitter/. serde = { workspace = true } serde_json = { workspace = true } diff --git a/rs/moq-audio/src/bus.rs b/rs/moq-audio/src/bus.rs index bd56f3cb513..4ad5fbe436a 100644 --- a/rs/moq-audio/src/bus.rs +++ b/rs/moq-audio/src/bus.rs @@ -702,6 +702,7 @@ impl Reader { /// The next block, or `None` once the mix has ended and every block kept /// for this reader has been read. + #[kio::wrapper] pub fn poll_read(&mut self, waiter: &kio::Waiter) -> Poll> { let next = self.next; let result = ready!(self.history.poll(waiter, |history| { @@ -719,12 +720,6 @@ impl Reader { Err(_) => Poll::Ready(None), } } - - /// The next block, or `None` once the mix has ended and every block kept - /// for this reader has been read. - pub async fn read(&mut self) -> Option { - kio::wait(|waiter| self.poll_read(waiter)).await - } } /// One source's PCM on its way into a [`Bus`]. diff --git a/rs/moq-audio/src/capture/fanout.rs b/rs/moq-audio/src/capture/fanout.rs index 1cef6cd191f..ffc32982be5 100644 --- a/rs/moq-audio/src/capture/fanout.rs +++ b/rs/moq-audio/src/capture/fanout.rs @@ -145,7 +145,8 @@ impl Consumer { self.device.as_ref() } - /// Read the next buffer, or the shared error when capture stops or fails. + /// The next buffer, or the shared error when capture stops or fails. + #[kio::wrapper] pub fn poll_read(&mut self, waiter: &kio::Waiter) -> Poll> { let result = ready!(self.state.poll(waiter, |state| { if let Some(error) = &state.failure { @@ -168,11 +169,6 @@ impl Consumer { Err(_) => Poll::Ready(Err(Error::Capture("audio capture stopped".into()))), } } - - /// Read the next buffer, or the shared error when capture stops or fails. - pub async fn read(&mut self) -> Result { - kio::wait(|waiter| self.poll_read(waiter)).await - } } #[cfg(test)] diff --git a/rs/moq-auth/Cargo.toml b/rs/moq-auth/Cargo.toml index 10feaf13e35..deda27f2d56 100644 --- a/rs/moq-auth/Cargo.toml +++ b/rs/moq-auth/Cargo.toml @@ -11,17 +11,29 @@ rust-version.workspace = true [features] default = ["client"] +# Generated async wrappers for the lease poll methods, and `lease::Consumer::reported`. +# Enable where you `.await` them; `client` and `serve` are async and enable it. +async = ["kio/async"] # The HTTP client that drives a lease against an auth server. -client = ["dep:reqwest", "dep:rind-rustls", "dep:rustls", "dep:rustls-platform-verifier", "dep:tokio", "kio/runtime", "dep:tracing", "dep:rand"] +client = [ + "async", + "dep:reqwest", + "dep:rind-rustls", + "dep:rustls", + "dep:rustls-platform-verifier", + "dep:tokio", + "kio/runtime", + "dep:tracing", + "dep:rand", +] # The reference auth server behind `moq auth serve`: the policy a relay held, on axum. -serve = ["dep:axum", "dep:tokio", "dep:tracing", "tokio/net"] +serve = ["async", "dep:axum", "dep:tokio", "dep:tracing", "tokio/net"] [dependencies] axum = { workspace = true, optional = true } base64 = { workspace = true } jsonwebtoken = { workspace = true } -# `async` for `Shared::wait`, until moq-auth grows an `async` feature of its own. -kio = { workspace = true, features = ["async", "serde"] } +kio = { workspace = true, features = ["serde"] } moq-crypto = { workspace = true, features = ["serde"] } moq-pattern = { workspace = true, features = ["serde"] } pkcs1 = { workspace = true, features = ["std"] } diff --git a/rs/moq-auth/MIGRATE.md b/rs/moq-auth/MIGRATE.md index ab52472ff57..96bfbac39fb 100644 --- a/rs/moq-auth/MIGRATE.md +++ b/rs/moq-auth/MIGRATE.md @@ -4,6 +4,16 @@ How to update your code for each breaking change, newest first. ## Unreleased +### Lease async wrappers are opt-in + +The lease's `changed`, `closed`, and `revalidate_requested` wrappers and +`lease::Consumer::reported` need the `async` Cargo feature, which `client` (a default) and +`serve` enable. A dependency with `default-features = false` that awaits them adds it: +`moq-auth = { version = "...", default-features = false, features = ["async"] }`. Without it, +drive the `poll_*` methods. +`lease::Producer::poll_revalidate` is `poll_revalidate_requested`, after the wrapper it backs. +The JSON contract and tokens are unchanged. + ### Mesh endpoints use the signing public key Replace `moq_auth::mesh::EndpointId` with `moq_crypto::sign::PublicKey` and diff --git a/rs/moq-auth/src/client.rs b/rs/moq-auth/src/client.rs index 72cf6fe93cb..276213f3ef9 100644 --- a/rs/moq-auth/src/client.rs +++ b/rs/moq-auth/src/client.rs @@ -6,7 +6,7 @@ use std::time::Duration; use kio::clock::Instant; use url::Url; -use crate::lease::{self, Due, Reason}; +use crate::lease::{self, Due, DueTimers, Reason}; use crate::{Bytes, Error, Event, Grant, Request}; /// Every request is bounded so a hung server refuses rather than parks the session. @@ -218,16 +218,14 @@ impl Context { } } - /// Resolve after this credential's end was reported, or its driver stopped. - pub(crate) async fn reported(&self) { - drop( - self.lifecycle() - .wait(|lifecycle| match lifecycle.reported { - true => std::task::Poll::Ready(()), - false => std::task::Poll::Pending, - }) - .await, - ); + /// Ready once this credential's end was reported, or its driver stopped. + pub(crate) fn poll_reported(&self, waiter: &kio::Waiter) -> std::task::Poll<()> { + self.lifecycle() + .poll(waiter, |lifecycle| match lifecycle.reported { + true => std::task::Poll::Ready(()), + false => std::task::Poll::Pending, + }) + .map(drop) } /// Resolve with the connection's end reason, for an attached credential only. @@ -435,6 +433,7 @@ impl Driver { // Keep an in-flight request alive while expiry or a push is observed. let mut inflight: Option> + Send>>> = None; let mut pending = false; + let mut timers = DueTimers::default(); let ended = loop { let reply = async { @@ -448,7 +447,7 @@ impl Driver { ended = producer.closed() => break ended, // An attached credential ends with its connection. reason = self.context.connection_ended() => break producer.finish(reason, Bytes::default()), - due = producer.due() => match due { + due = kio::wait(|waiter| producer.poll_due(waiter, &mut timers)) => match due { Due::Expired => break producer.finish(Reason::Expired, Bytes::default()), Due::Revalidate if inflight.is_some() => pending = true, Due::Revalidate => inflight = Some(self.post_revalidate()), diff --git a/rs/moq-auth/src/lease.rs b/rs/moq-auth/src/lease.rs index 7725e43b3da..a2b1af9b160 100644 --- a/rs/moq-auth/src/lease.rs +++ b/rs/moq-auth/src/lease.rs @@ -7,8 +7,6 @@ //! by any in-process logic when the embedder decides itself. No trait, no callbacks. use std::sync::Arc; -#[cfg(feature = "client")] -use std::sync::Mutex; use std::task::Poll; #[cfg(feature = "client")] use std::time::Duration; @@ -122,7 +120,39 @@ struct Clock { expires: Option, cadence: Option, failures: u32, - revision: u64, +} + +/// The timers a [`Producer::poll_due`] waits on, kept by its caller between polls. +#[cfg(feature = "client")] +#[derive(Default)] +pub(crate) struct DueTimers { + revalidate: Option, + expire: Option, +} + +#[cfg(feature = "client")] +impl DueTimers { + /// Ready once the clock reaches `at`; never for no instant. Re-arms when `at` moves. + fn poll_at( + timer: &mut Option, + waiter: &kio::Waiter, + at: Option, + ) -> Poll<()> { + let Some(at) = at else { + *timer = None; + return Poll::Pending; + }; + let sleep = match timer { + Some(sleep) => { + if sleep.deadline() != at { + sleep.reset(at); + } + sleep + } + None => timer.insert(kio::clock::sleep_until(at)), + }; + waiter.poll_future(std::pin::Pin::new(sleep)) + } } #[cfg(feature = "client")] @@ -136,7 +166,6 @@ impl Clock { expires: grant.deadline(), cadence: grant.revalidate, failures: 0, - revision: 0, } } } @@ -147,17 +176,16 @@ impl Clock { #[derive(Debug)] pub struct Producer { state: kio::Shared, + // Its changes wake a parked `poll_due`, which re-arms on the new deadlines. #[cfg(feature = "client")] - clock: Mutex, - #[cfg(feature = "client")] - clock_changed: tokio::sync::Notify, + clock: kio::Shared, } impl Producer { /// Start a lease on `grant`, returning both handles. pub fn new(grant: Grant) -> (Self, Consumer) { #[cfg(feature = "client")] - let clock = Mutex::new(Clock::new(&grant)); + let clock = kio::Shared::new(Clock::new(&grant)); let state = kio::Shared::new(State { grant, epoch: 0, @@ -169,8 +197,6 @@ impl Producer { state: state.clone(), #[cfg(feature = "client")] clock, - #[cfg(feature = "client")] - clock_changed: tokio::sync::Notify::new(), }, Consumer::new(state), ) @@ -184,63 +210,37 @@ impl Producer { } #[cfg(feature = "client")] { - let mut clock = self.clock.lock().expect("lease clock"); - let revision = clock.revision.wrapping_add(1); - *clock = Clock::new(&grant); - clock.revision = revision; + *self.clock.lock() = Clock::new(&grant); } state.grant = grant; state.epoch += 1; - drop(state); - #[cfg(feature = "client")] - self.clock_changed.notify_waiters(); } - /// Wait until the grant needs a re-check or expires. A session may request an - /// earlier re-check through [`Consumer::revalidate`]. + /// What the grant's clock requires next: a re-check, or its expiry. A session may + /// request an earlier re-check through [`Consumer::revalidate`]. #[cfg(feature = "client")] - pub(crate) async fn due(&self) -> Due { - loop { - let changed = self.clock_changed.notified(); - tokio::pin!(changed); - changed.as_mut().enable(); - let (next, expires, revision) = { - let clock = self.clock.lock().expect("lease clock"); - (clock.next, clock.expires, clock.revision) - }; - let revalidate = async { - match next { - Some(at) => kio::clock::sleep_until(at).await, - None => std::future::pending().await, - } - }; - let expire = async { - match expires { - Some(at) => kio::clock::sleep_until(at).await, - None => std::future::pending().await, - } - }; - tokio::select! { - biased; - () = expire => { - if self.clock.lock().expect("lease clock").revision == revision { - return Due::Expired; - } - }, - () = changed => continue, - () = revalidate => { - let mut clock = self.clock.lock().expect("lease clock"); - if clock.revision == revision { - clock.next = None; - return Due::Revalidate; - } - }, - () = self.revalidate_requested() => { - self.clock.lock().expect("lease clock").next = None; - return Due::Revalidate; - }, + #[cfg_attr(test, kio::wrapper)] + pub(crate) fn poll_due(&self, waiter: &kio::Waiter, timers: &mut DueTimers) -> Poll { + // Judged under the clock's lock, so a change after it parks this poll on the new deadlines. + let mut due = None; + let clock = self.clock.poll(waiter, |clock| { + if DueTimers::poll_at(&mut timers.expire, waiter, clock.expires).is_ready() { + due = Some(Due::Expired); + } else if DueTimers::poll_at(&mut timers.revalidate, waiter, clock.next).is_ready() { + due = Some(Due::Revalidate); + } + if due.is_some() { Poll::Ready(()) } else { Poll::Pending } + }); + if let Poll::Ready(mut clock) = clock { + let due = due.expect("ready with a due"); + if due == Due::Revalidate { + clock.next = None; } + return Poll::Ready(due); } + std::task::ready!(self.poll_revalidate_requested(waiter)); + self.clock.lock().next = None; + Poll::Ready(Due::Revalidate) } /// Keep the current grant after a failed re-check and schedule a jittered retry. @@ -248,26 +248,19 @@ impl Producer { pub(crate) fn failed(&self) -> Duration { use rand::RngExt; const BACKOFF_MAX: Duration = Duration::from_secs(60); - let mut clock = self.clock.lock().expect("lease clock"); + let mut clock = self.clock.lock(); clock.failures = clock.failures.saturating_add(1); let base = Duration::from_secs(1) * 2u32.saturating_pow(clock.failures.saturating_sub(1).min(16)); let base = base.min(clock.cadence.unwrap_or(BACKOFF_MAX)).min(BACKOFF_MAX); let delay = base.mul_f64(kio::clock::rng().random_range(0.75..=1.25)); clock.next = Some(kio::clock::now() + delay); - clock.revision = clock.revision.wrapping_add(1); - drop(clock); - self.clock_changed.notify_waiters(); delay } /// Consume an already-asked re-check after an in-flight response. #[cfg(feature = "client")] pub(crate) fn immediate(&self) { - let mut clock = self.clock.lock().expect("lease clock"); - clock.next = None; - clock.revision = clock.revision.wrapping_add(1); - drop(clock); - self.clock_changed.notify_waiters(); + self.clock.lock().next = None; } /// End the lease with `reason`, consuming the handle, and return the reason @@ -276,31 +269,23 @@ impl Producer { self.finish(reason, Bytes::default()).0 } - /// Poll for the lease ending, from either side: why it ended, and the totals - /// the session reported. A producer-side revoke, or a drop, reports zero bytes. + /// The lease's end, from either side: why it ended, and the totals the session reported. + /// A producer-side revoke, or a drop, reports zero bytes. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll<(Reason, Bytes)> { let state = std::task::ready!(self.state.poll(waiter, |state| ready_if(state.closed.is_some()))); Poll::Ready(state.closed.clone().expect("waited for a close")) } - /// Wait for the lease to end, from either side. - pub async fn closed(&self) -> (Reason, Bytes) { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - - /// Poll until at least one re-check has been asked since the last observation. + /// Resolves once at least one re-check has been asked since the last observation. /// A burst of asks resolves once. - pub fn poll_revalidate(&self, waiter: &kio::Waiter) -> Poll<()> { + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_revalidate_requested(&self, waiter: &kio::Waiter) -> Poll<()> { let mut state = std::task::ready!(self.state.poll(waiter, |state| ready_if(state.revalidate > 0))); state.revalidate = 0; Poll::Ready(()) } - /// Wait until a re-check has been asked since the last observation. - pub async fn revalidate_requested(&self) { - kio::wait(|waiter| self.poll_revalidate(waiter)).await - } - pub(crate) fn finish(&self, reason: Reason, bytes: Bytes) -> (Reason, Bytes) { self.state.lock().closed.get_or_insert((reason, bytes)).clone() } @@ -367,7 +352,8 @@ impl Consumer { self.hold.state.read().grant.clone() } - /// Poll for the next update: the new grant, or the reason the lease ended. + /// The next update: the new grant, or the reason the lease ended. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_changed(&mut self, waiter: &kio::Waiter) -> Poll> { let seen = self.seen; let state = std::task::ready!( @@ -382,22 +368,13 @@ impl Consumer { Poll::Ready(Ok(state.grant.clone())) } - /// Wait for the next update: the new grant, or the reason the lease ended. - pub async fn changed(&mut self) -> Result { - kio::wait(|waiter| self.poll_changed(waiter)).await - } - - /// Poll for the lease ending. + /// The reason the lease ended, once it does. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll { let state = std::task::ready!(self.hold.state.poll(waiter, |state| ready_if(state.closed.is_some()))); Poll::Ready(state.closed.as_ref().expect("waited for a close").0.clone()) } - /// Wait for the lease to end. - pub async fn closed(&self) -> Reason { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - /// Ask the producer to re-check now. A no-op on a [`fixed`](Self::fixed) /// lease, which has nobody to ask. pub fn revalidate(&self) { @@ -408,17 +385,15 @@ impl Consumer { /// `end` POST returns or fails, so its request timeout bounds the wait, and as soon /// as the lease ends when nobody needs telling. Take it before [`close`](Self::close) /// to learn when that close was reported. + #[cfg(feature = "async")] pub fn reported(&self) -> impl Future + Send + 'static { - let state = self.hold.state.clone(); - #[cfg(feature = "client")] - let context = self.hold.context.clone(); - async move { - drop(state.wait(|state| ready_if(state.closed.is_some())).await); + // Holds the state, not the lease, so dropping every consumer still ends it. + let reported = Reported { + state: self.hold.state.clone(), #[cfg(feature = "client")] - if let Some(context) = context { - context.reported().await; - } - } + context: self.hold.context.clone(), + }; + kio::wait(move |waiter| reported.poll(waiter)) } /// End the lease with the session's own close classification and the totals it @@ -447,6 +422,28 @@ impl Drop for Held { } } +/// What [`Consumer::reported`] waits on. +#[cfg(feature = "async")] +struct Reported { + state: kio::Shared, + #[cfg(feature = "client")] + context: Option, +} + +#[cfg(feature = "async")] +impl Reported { + fn poll(&self, waiter: &kio::Waiter) -> Poll<()> { + drop(std::task::ready!( + self.state.poll(waiter, |state| ready_if(state.closed.is_some())) + )); + #[cfg(feature = "client")] + if let Some(context) = &self.context { + std::task::ready!(context.poll_reported(waiter)); + } + Poll::Ready(()) + } +} + fn ready_if(condition: bool) -> Poll<()> { if condition { Poll::Ready(()) } else { Poll::Pending } } @@ -595,21 +592,33 @@ mod tests { grant.revalidate = Some(Duration::from_secs(10)); grant.expires = Some(kio::clock::wall() + Duration::from_secs(60)); let (producer, consumer) = Producer::new(grant.clone()); - assert!(kio::clock::timeout(Duration::ZERO, producer.due()).await.is_err()); + assert!( + kio::clock::timeout(Duration::ZERO, producer.due(&mut DueTimers::default())) + .await + .is_err() + ); tokio::time::advance(Duration::from_secs(10)).await; - assert_eq!(producer.due().await, Due::Revalidate); - assert!(kio::clock::timeout(Duration::ZERO, producer.due()).await.is_err()); + assert_eq!(producer.due(&mut DueTimers::default()).await, Due::Revalidate); + assert!( + kio::clock::timeout(Duration::ZERO, producer.due(&mut DueTimers::default())) + .await + .is_err() + ); let delay = producer.failed(); assert!(delay >= Duration::from_millis(750) && delay <= Duration::from_millis(1250)); tokio::time::advance(Duration::from_secs(2)).await; - assert_eq!(producer.due().await, Due::Revalidate); + assert_eq!(producer.due(&mut DueTimers::default()).await, Due::Revalidate); grant.revalidate = Some(Duration::from_secs(5)); producer.update(grant); consumer.revalidate(); - assert_eq!(producer.due().await, Due::Revalidate); - assert!(kio::clock::timeout(Duration::ZERO, producer.due()).await.is_err()); + assert_eq!(producer.due(&mut DueTimers::default()).await, Due::Revalidate); + assert!( + kio::clock::timeout(Duration::ZERO, producer.due(&mut DueTimers::default())) + .await + .is_err() + ); } #[cfg(feature = "client")] @@ -623,7 +632,7 @@ mod tests { let producer = std::sync::Arc::new(producer); let waiter = tokio::spawn({ let producer = producer.clone(); - async move { producer.due().await } + async move { producer.due(&mut DueTimers::default()).await } }); tokio::task::yield_now().await; grant.revalidate = Some(Duration::from_secs(5)); @@ -657,7 +666,11 @@ mod tests { grant.revalidate = Some(Duration::MAX); let (producer, _consumer) = Producer::new(grant); tokio::time::advance(Duration::from_secs(3600)).await; - assert!(kio::clock::timeout(Duration::ZERO, producer.due()).await.is_err()); + assert!( + kio::clock::timeout(Duration::ZERO, producer.due(&mut DueTimers::default())) + .await + .is_err() + ); } #[cfg(feature = "client")] @@ -667,7 +680,7 @@ mod tests { let mut grant = grant("a/**"); grant.expires = Some(kio::clock::wall() - Duration::from_secs(4)); let (producer, _consumer) = Producer::new(grant); - let task = tokio::spawn(async move { producer.due().await }); + let task = tokio::spawn(async move { producer.due(&mut DueTimers::default()).await }); tokio::task::yield_now().await; tokio::time::advance(Duration::from_secs(2)).await; assert_eq!(task.await.unwrap(), Due::Expired); diff --git a/rs/moq-auth/src/lib.rs b/rs/moq-auth/src/lib.rs index a348829f9fa..8a8d4ba3676 100644 --- a/rs/moq-auth/src/lib.rs +++ b/rs/moq-auth/src/lib.rs @@ -18,6 +18,10 @@ //! Grants and claims name paths with [`Pattern`]s from `moq-pattern`, re-exported here: //! `foo` is one broadcast, `foo/**` is a subtree, `**` is everything. Each role may also //! carry a bitrate cap, a [`Rate`]. +//! +//! The lease `poll_*` methods are always available. Enable the `async` feature for their +//! generated async wrappers (`changed`, `closed`, and `revalidate_requested`) and +//! `lease::Consumer::reported`; `client` and `serve` enable it. #![warn(missing_docs)] diff --git a/rs/moq-boy/Cargo.toml b/rs/moq-boy/Cargo.toml index f714c7e23de..060242210db 100644 --- a/rs/moq-boy/Cargo.toml +++ b/rs/moq-boy/Cargo.toml @@ -35,7 +35,7 @@ bytes = { workspace = true } hang = { workspace = true } kio = { workspace = true, features = ["runtime"] } moq-audio = { workspace = true } -moq-json = { workspace = true } +moq-json = { workspace = true, features = ["async"] } moq-mux = { workspace = true } moq-net = { workspace = true } moq-tokio = { workspace = true, default-features = false, features = ["net"] } diff --git a/rs/moq-c/Cargo.toml b/rs/moq-c/Cargo.toml index 2481bb6cb7a..ab9f4a20424 100644 --- a/rs/moq-c/Cargo.toml +++ b/rs/moq-c/Cargo.toml @@ -25,11 +25,11 @@ openh264 = ["moq-video/openh264"] anyhow = { workspace = true, features = ["backtrace"] } bytes = { workspace = true } hang = { workspace = true } -kio = { workspace = true } +kio = { workspace = true, features = ["async"] } # Codecs only: OBS and the other C consumers do their own device I/O. moq-audio = { workspace = true, features = ["aac"] } -moq-json = { workspace = true } -moq-mux = { workspace = true } +moq-json = { workspace = true, features = ["async"] } +moq-mux = { workspace = true, features = ["async"] } moq-net = { workspace = true } moq-tokio = { workspace = true, features = ["quic", "websocket", "tcp", "uds"] } # Select the NVIDIA and VAAPI hardware codecs explicitly because workspace diff --git a/rs/moq-c/src/ffi.rs b/rs/moq-c/src/ffi.rs index ca0a3d78553..b085bdf2141 100644 --- a/rs/moq-c/src/ffi.rs +++ b/rs/moq-c/src/ffi.rs @@ -132,7 +132,9 @@ impl Background { .unwrap_or(Ok(())) } - fn poll_finished(&self, waiter: &kio::Waiter) -> std::task::Poll> { + /// The background operation's result, once it finishes. + #[kio::wrapper] + pub fn poll_finished(&self, waiter: &kio::Waiter) -> std::task::Poll> { let Some(result) = &self.result else { return std::task::Poll::Ready(Ok(())); }; @@ -145,11 +147,6 @@ impl Background { } } - /// Wait until the background operation finishes and return its result. - pub async fn finished(&self) -> Result<(), Error> { - kio::wait(|waiter| self.poll_finished(waiter)).await - } - /// Fail an operation when its owned background task fails. pub async fn watch(&self, future: impl std::future::Future>) -> Result { self.check()?; diff --git a/rs/moq-cli/Cargo.toml b/rs/moq-cli/Cargo.toml index 5855b6af1dd..de56dc9b04d 100644 --- a/rs/moq-cli/Cargo.toml +++ b/rs/moq-cli/Cargo.toml @@ -104,7 +104,7 @@ axum-server = { workspace = true, features = ["tls-rustls-no-provider"] } base64 = { workspace = true } bytes = { workspace = true } futures = { workspace = true } -hang = { workspace = true } +hang = { workspace = true, features = ["async"] } humantime = { workspace = true } kio = { workspace = true, features = ["runtime"] } moq-archive = { workspace = true } @@ -113,8 +113,8 @@ moq-auth = { workspace = true, features = ["serve"] } # `server` enables the HTTP egress server for `moq export hls`; the importer is always available. moq-hls = { workspace = true, features = ["server"] } # Writes and reads the stats and feedback snapshots. -moq-json = { workspace = true } -moq-mux = { workspace = true } +moq-json = { workspace = true, features = ["async"] } +moq-mux = { workspace = true, features = ["async"] } # Cluster, LAN mesh, and inbound `/.cluster` auth. default-features off so the # CLI's own feature set (quic, cluster-lan) is what selects them. moq-relay = { workspace = true } diff --git a/rs/moq-crypto/Cargo.toml b/rs/moq-crypto/Cargo.toml index ad079403164..3a31818d768 100644 --- a/rs/moq-crypto/Cargo.toml +++ b/rs/moq-crypto/Cargo.toml @@ -28,6 +28,8 @@ e2ee = ["dep:base64", "dep:bytes", "dep:kio", "dep:moq-net", "dep:rind", "dep:uu provider = ["dep:rind", "dep:rind-rustls", "dep:rustls"] # `sign::PublicKey` serializes as its canonical z-base-32 text. serde = ["dep:serde"] +# Generated async wrappers for the `sign` and `e2ee` poll methods. Enable where you `.await` them. +async = ["kio?/async"] [dependencies] base64 = { workspace = true, optional = true } @@ -47,6 +49,9 @@ zeroize = { workspace = true, optional = true } criterion = { package = "bench-harness", path = "../bench-harness" } hex = { workspace = true } kio = { workspace = true, features = ["runtime"] } +# Compile the generated wrappers used by the tests and benches, keeping the features they pick. +moq-crypto = { path = ".", default-features = false, features = ["async"] } +moq-net = { path = "../moq-net", features = ["async"] } serde_json = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync", "test-util", "time"] } diff --git a/rs/moq-crypto/MIGRATE.md b/rs/moq-crypto/MIGRATE.md index 765efc831b3..d16eb693e71 100644 --- a/rs/moq-crypto/MIGRATE.md +++ b/rs/moq-crypto/MIGRATE.md @@ -4,6 +4,14 @@ How to update your code for each breaking change, newest first. ## Unreleased +### Async wrappers are opt-in + +Enable the `async` Cargo feature to keep awaiting `read_frame`, `recv_group`, `recv`, +`requested_track`, `requested`, and `closed`: +`moq-crypto = { version = "...", features = ["async"] }`. This also applies with +`default-features = false`. The feature is off by default; the matching `poll_*` methods stay +available without it. Signatures, ciphertext, and track names are unchanged. + ### Datagrams are their own track kind ([#7752](https://github.com/moq-dev/moq/pull/7752)) `e2ee::track::Producer::{append_datagram, insert_datagram}` and `e2ee::track::Consumer::{poll_recv_datagram, recv_datagram}` are removed. Create a datagram track with `e2ee::broadcast::Producer::create_datagrams(name, info)` and read it with `e2ee::broadcast::Consumer::datagrams(name)`, which return `datagram::Producer` (`append`, `insert`) and `datagram::Consumer` (`recv`, yielding the same `datagram::Event`). Its sequence is its own, not the stream track's group counter, and its key domain stays separate. diff --git a/rs/moq-crypto/src/e2ee/broadcast.rs b/rs/moq-crypto/src/e2ee/broadcast.rs index 106f1ea80c8..574ecb32de1 100644 --- a/rs/moq-crypto/src/e2ee/broadcast.rs +++ b/rs/moq-crypto/src/e2ee/broadcast.rs @@ -124,11 +124,12 @@ impl Requests { self.inner.info() } - /// Poll for the next consumer-requested track, without blocking. + /// The next track a consumer requests, as a [`track::Request`] to serve. /// /// # Errors /// /// [`moq_net::Error::Closed`] once the broadcast has ended. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_requested_track(&mut self, waiter: &kio::Waiter) -> Poll> { loop { let request = ready!(self.inner.poll_requested_track(waiter))?; @@ -142,15 +143,6 @@ impl Requests { } } - /// Block until a consumer requests a track, returning a [`track::Request`] to serve. - /// - /// # Errors - /// - /// Same as [`Self::poll_requested_track`]. - pub async fn requested_track(&mut self) -> Result { - kio::wait(|waiter| self.poll_requested_track(waiter)).await - } - /// Create a handler for requests to datagram tracks nobody created, named by their semantic /// names. See [`moq_net::broadcast::Requests::datagrams`]. pub fn datagrams(&self) -> DatagramRequests { @@ -178,12 +170,8 @@ impl Requests { }) } - /// Block until the broadcast ends. - pub async fn closed(&self) { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - - /// Poll until the broadcast ends. + /// Resolves once the broadcast ends. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll<()> { self.inner.poll_closed(waiter) } @@ -215,11 +203,12 @@ pub struct DatagramRequests { } impl DatagramRequests { - /// Poll for the next datagram track a subscriber asks for, to accept or reject. + /// The next datagram track a subscriber asks for, as a [`datagram::Request`] to accept or reject. /// /// # Errors /// /// [`moq_net::Error::Closed`] once the broadcast has ended. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_requested(&mut self, waiter: &kio::Waiter) -> Poll> { loop { let request = ready!(self.inner.poll_requested(waiter))?; @@ -232,15 +221,6 @@ impl DatagramRequests { } } } - - /// Block until a subscriber asks for a datagram track, returning a [`datagram::Request`] to serve. - /// - /// # Errors - /// - /// Same as [`Self::poll_requested`]. - pub async fn requested(&mut self) -> Result { - kio::wait(|waiter| self.poll_requested(waiter)).await - } } /// A handler's answer to a protected broadcast's listings, taking semantic names and listing diff --git a/rs/moq-crypto/src/e2ee/datagram.rs b/rs/moq-crypto/src/e2ee/datagram.rs index faf1c5ef69f..e62d6d69ad5 100644 --- a/rs/moq-crypto/src/e2ee/datagram.rs +++ b/rs/moq-crypto/src/e2ee/datagram.rs @@ -241,11 +241,12 @@ impl Consumer { self.inner.set_priority(priority); } - /// Poll for the next datagram event, or `None` once the track finished. + /// The next datagram event, or `None` once the track finished. /// /// # Errors /// /// [`Error::Exhausted`] once the datagram key is spent, or a net error. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_recv(&mut self, waiter: &kio::Waiter) -> Poll>> { let Some(datagram) = ready!(self.inner.poll_recv(waiter)?) else { return Poll::Ready(Ok(None)); @@ -275,15 +276,6 @@ impl Consumer { }; Poll::Ready(Ok(Some(event))) } - - /// Receive the next datagram event. - /// - /// # Errors - /// - /// Same as [`Self::poll_recv`]. - pub async fn recv(&mut self) -> Result> { - kio::wait(|waiter| self.poll_recv(waiter)).await - } } impl fmt::Debug for Consumer { diff --git a/rs/moq-crypto/src/e2ee/group.rs b/rs/moq-crypto/src/e2ee/group.rs index f0b94d601ec..b3f4680b3d0 100644 --- a/rs/moq-crypto/src/e2ee/group.rs +++ b/rs/moq-crypto/src/e2ee/group.rs @@ -158,12 +158,13 @@ impl Consumer { self.sequence } - /// Read the next decrypted frame, without blocking. + /// The next decrypted frame, or `None` once the group finished. /// /// # Errors /// /// [`Error::Authentication`] ends this track. [`Error::Oversize`] or /// [`Error::Exhausted`] as for open, and net errors from the underlying group. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_read_frame(&mut self, waiter: &kio::Waiter) -> Poll>> { if self.terminal.poll_failed(waiter) { self.release(); @@ -197,15 +198,6 @@ impl Consumer { } } } - - /// Read the next decrypted frame. - /// - /// # Errors - /// - /// Same as [`Self::poll_read_frame`]. - pub async fn read_frame(&mut self) -> Result> { - kio::wait(|waiter| self.poll_read_frame(waiter)).await - } } impl fmt::Debug for Consumer { diff --git a/rs/moq-crypto/src/e2ee/track.rs b/rs/moq-crypto/src/e2ee/track.rs index d5da6ead31f..63d0cc4d369 100644 --- a/rs/moq-crypto/src/e2ee/track.rs +++ b/rs/moq-crypto/src/e2ee/track.rs @@ -226,11 +226,12 @@ impl Consumer { &self.name } - /// Poll for the next protected group in arrival order. + /// The next protected group in arrival order, or `None` once the track finished. /// /// # Errors /// /// [`Error::Authentication`](crate::e2ee::Error::Authentication) once a grouped frame has failed to open, or a net error. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_recv_group(&mut self, waiter: &kio::Waiter) -> Poll>> { let terminal = Arc::clone(&self.terminal); let key = Arc::clone(&self.key); @@ -245,15 +246,6 @@ impl Consumer { )))) }) } - - /// Receive the next protected group in arrival order. - /// - /// # Errors - /// - /// Same as [`Self::poll_recv_group`]. - pub async fn recv_group(&mut self) -> Result> { - kio::wait(|waiter| self.poll_recv_group(waiter)).await - } } // Groups share the terminal, so without this a held group would keep the diff --git a/rs/moq-crypto/src/lib.rs b/rs/moq-crypto/src/lib.rs index 0070d40b27d..605bd3b4285 100644 --- a/rs/moq-crypto/src/lib.rs +++ b/rs/moq-crypto/src/lib.rs @@ -11,6 +11,9 @@ //! gates the provider and needs no `moq-net`. Without any of them the crate holds only //! [`sign::PublicKey`], the endpoint ID, and its canonical z-base-32 text, with no `moq-net` //! or crypto backend; `serde` adds its string serialization. +//! +//! The `poll_*` methods are always available. Enable the `async` feature for their generated +//! async wrappers (such as `read_frame` and `recv_group`); it is off by default. #![warn(missing_docs)] diff --git a/rs/moq-crypto/src/sign/group.rs b/rs/moq-crypto/src/sign/group.rs index 855aa2587ce..07e3d4ac6c2 100644 --- a/rs/moq-crypto/src/sign/group.rs +++ b/rs/moq-crypto/src/sign/group.rs @@ -122,11 +122,12 @@ impl Consumer { self.sequence } - /// Read the next verified frame, without blocking. + /// The next verified frame, or `None` once the group finished. /// /// # Errors /// /// [`Error::Signature`] once a frame has failed to verify, or a net error from the group. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_read_frame(&mut self, waiter: &kio::Waiter) -> Poll>> { let Some(inner) = self.inner.as_mut() else { return Poll::Ready(Err(Error::Signature)); @@ -151,15 +152,6 @@ impl Consumer { } } } - - /// Read the next verified frame. - /// - /// # Errors - /// - /// Same as [`Self::poll_read_frame`]. - pub async fn read_frame(&mut self) -> Result> { - kio::wait(|waiter| self.poll_read_frame(waiter)).await - } } impl fmt::Debug for Consumer { diff --git a/rs/moq-crypto/src/sign/track.rs b/rs/moq-crypto/src/sign/track.rs index 91cd5b577b6..77a41d869cc 100644 --- a/rs/moq-crypto/src/sign/track.rs +++ b/rs/moq-crypto/src/sign/track.rs @@ -172,26 +172,18 @@ impl Consumer { self.inner.name() } - /// Poll for the next group in arrival order. + /// The next group in arrival order, or `None` once the track finished. /// /// # Errors /// /// A net error if the track fails. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_recv_group(&mut self, waiter: &kio::Waiter) -> Poll>> { let Some(inner) = ready!(self.inner.poll_recv_group(waiter)?) else { return Poll::Ready(Ok(None)); }; Poll::Ready(Ok(Some(group::Consumer::new(inner, self.key, Arc::clone(&self.scope))))) } - - /// Receive the next group in arrival order. - /// - /// # Errors - /// - /// Same as [`Self::poll_recv_group`]. - pub async fn recv_group(&mut self) -> Result> { - kio::wait(|waiter| self.poll_recv_group(waiter)).await - } } impl fmt::Debug for Consumer { diff --git a/rs/moq-ffi/Cargo.toml b/rs/moq-ffi/Cargo.toml index f2e5f404bde..7bb722bca7b 100644 --- a/rs/moq-ffi/Cargo.toml +++ b/rs/moq-ffi/Cargo.toml @@ -42,8 +42,8 @@ hang = { workspace = true } kio = { workspace = true, features = ["runtime"] } # Codecs only: the language wrappers do their own device I/O. moq-audio = { workspace = true, optional = true, features = ["aac"] } -moq-json = { workspace = true } -moq-mux = { workspace = true } +moq-json = { workspace = true, features = ["async"] } +moq-mux = { workspace = true, features = ["async"] } moq-net = { workspace = true } moq-tokio = { workspace = true, default-features = true } # `nvidia` and `vaapi` back on, matching rs/moq-c: the workspace turns diff --git a/rs/moq-gst/Cargo.toml b/rs/moq-gst/Cargo.toml index b37a2eac1b9..051fb298fb3 100644 --- a/rs/moq-gst/Cargo.toml +++ b/rs/moq-gst/Cargo.toml @@ -24,7 +24,7 @@ gst = { package = "gstreamer", version = "0.25" } gst-audio = { package = "gstreamer-audio", version = "0.25" } hang = { workspace = true } kio = { workspace = true, features = ["runtime"] } -moq-mux = { workspace = true } +moq-mux = { workspace = true, features = ["async"] } moq-net = { workspace = true } moq-tokio = { workspace = true, default-features = true } tokio = { workspace = true, features = ["full"] } diff --git a/rs/moq-hls/Cargo.toml b/rs/moq-hls/Cargo.toml index bd0242b988f..61f2365a1da 100644 --- a/rs/moq-hls/Cargo.toml +++ b/rs/moq-hls/Cargo.toml @@ -14,7 +14,9 @@ categories = ["multimedia", "network-programming", "web-programming"] [features] default = ["server"] -server = ["dep:axum", "dep:http-body"] +# The import and export gateways, which drive HTTP fetches and background watchers as async code. +async = ["kio/async", "moq-mux/async"] +server = ["async", "dep:axum", "dep:http-body"] [dependencies] # Only needed by the HTTP export server (gated by `server`). @@ -45,6 +47,8 @@ async-trait = { workspace = true } criterion = { package = "bench-harness", path = "../bench-harness" } futures = { workspace = true } moq-archive = { path = "../moq-archive" } +# Compile the gateways for the tests, with or without the `server` feature. +moq-hls = { path = ".", default-features = false, features = ["async"] } moq-json = { path = "../moq-json" } moq-tokio = { path = "../moq-tokio", default-features = false, features = ["tcp", "quic"] } mp4-atom = { workspace = true } diff --git a/rs/moq-hls/MIGRATE.md b/rs/moq-hls/MIGRATE.md index a7dc51e8f70..3b2877431ee 100644 --- a/rs/moq-hls/MIGRATE.md +++ b/rs/moq-hls/MIGRATE.md @@ -4,6 +4,18 @@ How to update your code for each breaking change, newest first. ## Unreleased +### The gateways need the `async` feature + +The `import` and `export` modules compile only with the new `async` Cargo feature, which the default +`server` feature enables. With `default-features = false`, enable it to keep them: +`moq-hls = { version = "...", default-features = false, features = ["async"] }`. +Nothing else changes for a build with `server`. + +The export waits have polls for a caller that drives them itself: `Broadcaster::poll_closed` +and `poll_ready`, `Rendition::poll_playable`, and `poll_next` on the `renditions::Consumer` and +`segments::Consumer` cursors. A `segments::Consumer::next` dropped mid-fetch leaves that fetch to the +next call, which resumes it rather than starting over. + ### HTTP HLS lists contributions without starting their media `Server` follows the root catalog's contributions and lists their renditions from their diff --git a/rs/moq-hls/src/error.rs b/rs/moq-hls/src/error.rs index 80c16faf38d..ed216837467 100644 --- a/rs/moq-hls/src/error.rs +++ b/rs/moq-hls/src/error.rs @@ -185,6 +185,7 @@ pub enum Error { impl Error { /// Carries rind's message rather than its type, so rind stays out of this crate's API. + #[cfg(feature = "async")] pub(crate) fn unsupported_cpu(err: rind::cpu::MissingFeature) -> Self { Self::UnsupportedCpu(err.to_string()) } @@ -207,6 +208,7 @@ impl Error { /// another attempt: request timeout, rate limit, and the gateway/overload statuses. Every other /// status, `404` and `403` included, is settled, as is a playlist that had already ended when /// the import joined it. A failure with no answer at all is left to the backoff budget. + #[cfg(feature = "async")] pub(crate) fn settled(&self) -> bool { match self { Self::PlaylistEnded => true, @@ -221,6 +223,7 @@ impl Error { /// /// Every variant naming a resource stores it through here rather than redacting at display, so /// its `Debug` print is as safe as its `Display`. +#[cfg(feature = "async")] pub(crate) fn redact(url: &url::Url) -> url::Url { let mut url = url.clone(); // Fails only for a URL that cannot carry userinfo at all, such as `file:`. diff --git a/rs/moq-hls/src/export/contributions.rs b/rs/moq-hls/src/export/contributions.rs index e08b341c928..a0c35d66f1f 100644 --- a/rs/moq-hls/src/export/contributions.rs +++ b/rs/moq-hls/src/export/contributions.rs @@ -2,9 +2,7 @@ use std::collections::BTreeMap; use std::sync::Arc; -use std::task::Poll; -#[cfg(feature = "server")] -use std::task::ready; +use std::task::{Poll, ready}; use std::time::Duration; use hang::catalog::{Contribution, Media}; @@ -97,11 +95,11 @@ impl Producer { } let current = kio::Producer::new(Current::Pending); let owner = match source.follow(Some(&reference.broadcast), reference.instance.as_ref()) { - Ok(instances) => Some(self.context.spawner.spawn(follow_instances( + Ok(instances) => Some(self.context.spawner.spawn(FollowInstances { instances, - current.clone(), - self.context.clone(), - ))), + current: current.clone(), + context: self.context.clone(), + })), Err(err) => { *current.write().unwrap_or_else(|_| unreachable!("new entry is open")) = Current::Failed(err.into()); @@ -204,17 +202,8 @@ impl Producer { } #[cfg(feature = "server")] - pub async fn rendition( - &self, - at: &super::route::Contribution, - kind: Kind, - name: &str, - ) -> crate::Result>> { - kio::wait(|waiter| self.poll_rendition(waiter, at, kind, name)).await - } - - #[cfg(feature = "server")] - fn poll_rendition( + #[kio::wrapper] + pub fn poll_rendition( &self, waiter: &kio::Waiter, at: &super::route::Contribution, @@ -242,55 +231,69 @@ impl Producer { } } -async fn follow_instances(mut instances: follow::Instances, current: kio::Producer, context: Context) { - loop { - let next = instances.next().await; - let next = match next { - Ok(Some(broadcast)) => { - let id = broadcast - .info() - .instance - .clone() - .unwrap_or_else(moq_net::Instance::random) - .to_string(); - let source = moq_mux::Source::new(context.origin.clone(), broadcast); - let renditions = renditions::Producer::lazy(context.window, context.spawner.clone()); - let owner = context.spawner.lead(super::watch_catalog( - source, - renditions.clone(), - super::Role::Contribution(context.root.consume()), - )); - Current::Online(Arc::new(Instance { - id, - renditions, - _owner: owner, - })) - } - Ok(None) => Current::Offline, - // An authorization refusal is final: the contribution stays out and its URLs answer 403. - Err(err @ moq_mux::Error::Moq(moq_net::Error::Unauthorized)) => { - tracing::debug!(%err, "contribution refused"); - let Ok(mut state) = current.write() else { return }; - *state = Current::Failed(err.into()); - return; - } - Err(moq_mux::Error::Moq(moq_net::Error::Closed)) => { - let Ok(mut state) = current.write() else { return }; - *state = Current::Offline; - return; - } - // Not routed yet, or refused otherwise: wait for the path's next announcement, which - // the follower hands out, however long it takes. Until then nothing is listed. - Err(err) => { - tracing::debug!(%err, "contribution not routed; waiting for its next announcement"); - match &*current.read() { - Current::Pending => Current::Offline, - _ => continue, +/// Follows a contribution's instances into its [`Current`], watching each one's catalog. +struct FollowInstances { + instances: follow::Instances, + current: kio::Producer, + context: Context, +} + +impl kio::Task for FollowInstances { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + loop { + let next = match ready!(self.instances.poll_next(waiter)) { + Ok(Some(broadcast)) => { + let id = broadcast + .info() + .instance + .clone() + .unwrap_or_else(moq_net::Instance::random) + .to_string(); + let source = moq_mux::Source::new(self.context.origin.clone(), broadcast); + let renditions = renditions::Producer::lazy(self.context.window, self.context.spawner.clone()); + let owner = self.context.spawner.lead(super::WatchCatalog::new( + source, + renditions.clone(), + super::Role::Contribution(self.context.root.consume()), + )); + Current::Online(Arc::new(Instance { + id, + renditions, + _owner: owner, + })) } - } - }; - let Ok(mut state) = current.write() else { return }; - *state = next; + Ok(None) => Current::Offline, + // An authorization refusal is final: the contribution stays out and its URLs answer 403. + Err(err @ moq_mux::Error::Moq(moq_net::Error::Unauthorized)) => { + tracing::debug!(%err, "contribution refused"); + if let Ok(mut state) = self.current.write() { + *state = Current::Failed(err.into()); + } + return Poll::Ready(()); + } + Err(moq_mux::Error::Moq(moq_net::Error::Closed)) => { + if let Ok(mut state) = self.current.write() { + *state = Current::Offline; + } + return Poll::Ready(()); + } + // Not routed yet, or refused otherwise: wait for the path's next announcement, which + // the follower hands out, however long it takes. Until then nothing is listed. + Err(err) => { + tracing::debug!(%err, "contribution not routed; waiting for its next announcement"); + match &*self.current.read() { + Current::Pending => Current::Offline, + _ => continue, + } + } + }; + let Ok(mut state) = self.current.write() else { + return Poll::Ready(()); + }; + *state = next; + } } } diff --git a/rs/moq-hls/src/export/driver.rs b/rs/moq-hls/src/export/driver.rs index 262786421fb..abbbdf20005 100644 --- a/rs/moq-hls/src/export/driver.rs +++ b/rs/moq-hls/src/export/driver.rs @@ -1,7 +1,5 @@ //! The export's background work, as `kio` tasks one [`Driver`] polls. -use std::future::Future; -use std::pin::Pin; use std::sync::Arc; use std::task::Poll; @@ -44,7 +42,7 @@ impl Driver { /// Run `task` ahead of every other task (see [`leads`](Self::leads)), until it finishes or /// the returned [`Owner`] drops. - pub(crate) fn lead(&mut self, task: impl Future + Send + 'static) -> Owner { + pub(crate) fn lead(&mut self, task: impl kio::Task + Send + 'static) -> Owner { let (job, owner) = Job::new(task); self.leads.push(job); owner @@ -104,17 +102,17 @@ impl Spawner { } /// Run `task` on the driver until it finishes or the returned [`Owner`] drops. - pub fn spawn(&self, task: impl Future + Send + 'static) -> Owner { + pub fn spawn(&self, task: impl kio::Task + Send + 'static) -> Owner { self.start(Lane::Other, task) } /// [`spawn`](Self::spawn) `task` among the driver's leads, polled before every other task, /// as a catalog watcher must be. - pub fn lead(&self, task: impl Future + Send + 'static) -> Owner { + pub fn lead(&self, task: impl kio::Task + Send + 'static) -> Owner { self.start(Lane::Lead, task) } - fn start(&self, lane: Lane, task: impl Future + Send + 'static) -> Owner { + fn start(&self, lane: Lane, task: impl kio::Task + Send + 'static) -> Owner { let (job, owner) = Job::new(task); // Refused once the driver is gone, which drops the task unrun. let _ = self.inbox.0.try_push((lane, job)); @@ -128,15 +126,15 @@ pub(crate) struct Owner(#[expect(dead_code, reason = "held only to close on drop /// A task and the [`Owner`] that ends it. struct Job { owner: kio::Consumer<()>, - task: Pin + Send>>, + task: Box + Send>, } impl Job { - fn new(task: impl Future + Send + 'static) -> (Self, Owner) { + fn new(task: impl kio::Task + Send + 'static) -> (Self, Owner) { let owner = kio::Producer::new(()); let job = Self { owner: owner.consume(), - task: Box::pin(task), + task: Box::new(task), }; (job, Owner(owner)) } @@ -149,7 +147,7 @@ impl kio::Task for Job { if self.owner.poll_closed(waiter).is_ready() { return Poll::Ready(()); } - waiter.poll_future(self.task.as_mut()) + self.task.poll(waiter) } } @@ -165,9 +163,9 @@ mod tests { let order = Arc::new(std::sync::Mutex::new(Vec::new())); let record = |name: &'static str| { let order = order.clone(); - async move { + move |_: &kio::Waiter| { order.lock().unwrap().push(name); - std::future::pending::<()>().await; + Poll::Pending } }; let _other = spawner.spawn(record("other")); @@ -182,8 +180,8 @@ mod tests { fn tasks_end_with_their_owner() { let (spawner, mut driver) = Spawner::new(); let waiter = kio::Waiter::noop(); - let forever = spawner.spawn(std::future::pending()); - let lead = driver.lead(std::future::pending()); + let forever = spawner.spawn(|_: &kio::Waiter| Poll::Pending); + let lead = driver.lead(|_: &kio::Waiter| Poll::Pending); drop(spawner); assert!(driver.poll(&waiter).is_pending(), "the owned tasks still run"); drop(lead); @@ -199,7 +197,10 @@ mod tests { let held = Arc::new(()); let _owner = spawner.spawn({ let held = held.clone(); - async move { drop(held) } + move |_: &kio::Waiter| { + let _ = &held; + Poll::Ready(()) + } }); drop(driver); assert_eq!(Arc::strong_count(&held), 1, "the queued task was dropped"); diff --git a/rs/moq-hls/src/export/mod.rs b/rs/moq-hls/src/export/mod.rs index 39291ffeea7..de216f0fc45 100644 --- a/rs/moq-hls/src/export/mod.rs +++ b/rs/moq-hls/src/export/mod.rs @@ -45,9 +45,10 @@ pub mod renditions; pub mod segments; use std::sync::{Arc, Mutex}; +use std::task::{Poll, ready}; use std::time::Duration; -use moq_mux::catalog::{self, CatalogFormat, Stream}; +use moq_mux::catalog::{self, Stream}; use rand::RngExt; pub use driver::Driver; @@ -102,13 +103,6 @@ impl Default for Config { /// last clone drops. #[derive(Clone)] pub struct Broadcaster { - #[cfg_attr( - not(feature = "server"), - allow( - dead_code, - reason = "read only by the serve path's `is_closed`/`closed` (server-gated); a recording consumer observes close through its `renditions()` cursor ending instead" - ) - )] broadcast: moq_net::broadcast::Consumer, /// The stream served, which every URL the master playlist and manifest list starts with: /// its instance, or an id minted for a broadcast that came through no route. @@ -170,7 +164,7 @@ impl Broadcaster { let renditions = renditions::Producer::new(config.window, spawner); // The watcher owns its own producer clone; dropping the `Broadcaster` ends it so the // standing catalog subscription stops when nobody's serving from this broadcaster. - let watcher = driver.lead(watch_catalog( + let watcher = driver.lead(WatchCatalog::new( source, renditions.clone(), Role::Root(contributions.clone().map(Box::new)), @@ -228,14 +222,9 @@ impl Broadcaster { self.broadcast.is_closed() } - /// Resolve once the source broadcast closes, so a pool can evict a dead broadcaster. - pub async fn closed(&self) { - self.broadcast.closed().await; - } - - /// [`closed`](Self::closed), polled: ready once the source broadcast closes. - #[cfg(feature = "server")] - pub(crate) fn poll_closed(&self, waiter: &kio::Waiter) -> std::task::Poll<()> { + /// Ready once the source broadcast closes, so a pool can evict a dead broadcaster. + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll<()> { self.broadcast.poll_closed(waiter) } @@ -255,32 +244,31 @@ impl Broadcaster { self.renditions.subscribe() } - /// Resolve once at least one rendition has a playable media playlist or a settled failure. + /// Ready once at least one rendition has a playable media playlist or a settled failure. /// How long to wait is the caller's policy: wrap this in a timeout as needed. - pub async fn ready(&self) { - self.renditions.ready().await; + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_ready(&self, waiter: &kio::Waiter) -> Poll<()> { + self.renditions.poll_ready(waiter) } - /// Resolve once the master playlist lists at least one rendition. + /// Ready once the master playlist lists at least one rendition. #[cfg_attr( not(feature = "server"), allow(dead_code, reason = "only the `server` feature uses it") )] - pub(crate) async fn advertised(&self) { - kio::wait(|waiter| { - let own = self.renditions.poll_advertised(waiter); - let contribution = self - .contributions - .as_ref() - .is_some_and(|contributions| contributions.poll_advertised(waiter).is_ready()); - // Contributions alone answer only while no own rendition could still be listed. - if own.is_ready() || (!self.renditions.listable() && contribution) { - std::task::Poll::Ready(()) - } else { - std::task::Poll::Pending - } - }) - .await; + #[cfg_attr(any(test, feature = "server"), kio::wrapper)] + pub(crate) fn poll_advertised(&self, waiter: &kio::Waiter) -> Poll<()> { + let own = self.renditions.poll_advertised(waiter); + let contribution = self + .contributions + .as_ref() + .is_some_and(|contributions| contributions.poll_advertised(waiter).is_ready()); + // Contributions alone answer only while no own rendition could still be listed. + if own.is_ready() || (!self.renditions.listable() && contribution) { + Poll::Ready(()) + } else { + Poll::Pending + } } /// Version every segment URL with `generation`, a label for the publisher run the @@ -359,17 +347,20 @@ impl Broadcaster { (video, audio) } + /// The contribution rendition `at` names, once the catalog is read. #[cfg(feature = "server")] - pub(crate) async fn contribution( + #[kio::wrapper] + pub(crate) fn poll_contribution( &self, + waiter: &kio::Waiter, at: &Contribution, kind: Kind, name: &str, - ) -> crate::Result>> { - self.renditions.catalogued().await; + ) -> Poll>>> { + ready!(self.renditions.poll_catalogued(waiter)); match &self.contributions { - Some(contributions) => contributions.rendition(at, kind, name).await, - None => Ok(None), + Some(contributions) => contributions.poll_rendition(waiter, at, kind, name), + None => Poll::Ready(Ok(None)), } } @@ -475,8 +466,9 @@ impl Broadcaster { not(feature = "server"), allow(dead_code, reason = "only the `server` feature uses it") )] - pub(crate) async fn video_init(&self) { - self.renditions.video_init().await; + #[cfg_attr(any(test, feature = "server"), kio::wrapper)] + pub(crate) fn poll_video_init(&self, waiter: &kio::Waiter) -> Poll<()> { + self.renditions.poll_video_init(waiter) } /// Resolve once the broadcast has something the DASH manifest lists (its first complete @@ -489,7 +481,7 @@ impl Broadcaster { )] pub(crate) async fn represented(&self) { if let Some(rendition) = self.renditions.snapshot().into_iter().next() { - rendition.represented().await; + kio::wait(|waiter| rendition.poll_represented(waiter)).await; } } @@ -528,72 +520,166 @@ enum Role { Contribution(kio::Consumer), } -async fn watch_catalog(source: moq_mux::Source, renditions: renditions::Producer, role: Role) { - // The source's catalog broadcast, so the catalog, the media subscriptions, and the closed - // check below all refer to the one broadcast it was built from. - let broadcast = &source.broadcast(); - let mut delay = CATALOG_RETRY_MIN; - let mut consumer = loop { - match catalog::Consumer::<()>::new(broadcast, CatalogFormat::Hang).await { - Ok(consumer) => break consumer, - Err(err) => { - tracing::warn!(%err, "failed to subscribe to broadcast catalog, retrying"); - let wait = delay.mul_f64(0.5 + kio::clock::rng().random::() / 2.0); - delay = (delay * 2).min(CATALOG_RETRY_MAX); - tokio::select! { - _ = kio::clock::sleep(wait) => {} - _ = kio::wait(|waiter| broadcast.poll_closed(waiter)) => { - renditions.close(); - return; - } - } - } +/// Follows a broadcast's catalog into its rendition set, as a driver task. +struct WatchCatalog { + source: moq_mux::Source, + renditions: renditions::Producer, + role: Role, + /// The source's catalog broadcast, taken on the first poll, so the catalog, the media + /// subscriptions, and the closed check all refer to the one broadcast it was built from. + broadcast: Option, + /// The next retry's backoff, before jitter. + delay: Duration, + state: Watch, +} + +#[allow( + clippy::large_enum_variant, + reason = "one per catalog watcher, and the reading variant is the working state" +)] +enum Watch { + Start, + /// Waiting for the catalog subscription, or the failure it started with. The consumer + /// subscribes itself; its track's info resolving rides on that subscription, so a refused + /// catalog fails here. + Subscribing(moq_mux::Result<(catalog::hang::Consumer<()>, moq_net::track::Resolving)>), + /// Waiting out a failed subscription's backoff, unless the broadcast closes first. + Backoff(kio::clock::Sleep), + Reading { + consumer: catalog::Consumer<()>, + /// A contribution's latest catalog, re-synced when the root's media changes. + latest: Option, + }, +} + +impl WatchCatalog { + fn new(source: moq_mux::Source, renditions: renditions::Producer, role: Role) -> Self { + Self { + source, + renditions, + role, + broadcast: None, + delay: CATALOG_RETRY_MIN, + state: Watch::Start, } - }; - - // A contribution's latest catalog, re-synced when the root's media changes. - let mut latest: Option = None; - let result = loop { - let next = kio::wait(|waiter| { - if let std::task::Poll::Ready(next) = consumer.poll_next(waiter) { - return std::task::Poll::Ready(Some(next)); - } - if let (Role::Contribution(root), Some(catalog)) = (&role, &latest) - && let std::task::Poll::Ready(Ok(())) = root.poll(waiter, |media| match **media != catalog.media { - true => std::task::Poll::Ready(()), - false => std::task::Poll::Pending, - }) { - return std::task::Poll::Ready(None); - } - std::task::Poll::Pending - }) - .await; - let mut catalog = match next { - Some(Ok(Some(catalog))) => catalog, - Some(Ok(None)) => break Ok(()), - Some(Err(err)) => { - tracing::warn!(%err, "broadcast catalog stream ended with error"); - break Err(err.into()); - } - None => latest.take().expect("a root media change follows a catalog"), + } + + fn broadcast(&self) -> &moq_net::broadcast::Consumer { + self.broadcast.as_ref().expect("taken on the first poll") + } + + /// Subscribe to the broadcast's hang catalog track. + fn subscribe(&self) -> Watch { + let track = self.broadcast().track(hang::Catalog::DEFAULT_NAME).map(|track| { + let info = track.info().into_inner(); + (catalog::hang::Consumer::new(track), info) + }); + Watch::Subscribing(track.map_err(Into::into)) + } + + /// The next catalog to sync, `Ok(None)` once the stream ends cleanly. + fn poll_catalog(&mut self, waiter: &kio::Waiter) -> Poll>> { + let Watch::Reading { consumer, latest } = &mut self.state else { + unreachable!("reading the catalog"); }; - match &role { - Role::Root(Some(contributions)) => contributions.sync(&source, &catalog), + if let Poll::Ready(next) = consumer.poll_next(waiter) { + return Poll::Ready(next.map_err(|err| { + tracing::warn!(%err, "broadcast catalog stream ended with error"); + err.into() + })); + } + if let (Role::Contribution(root), Some(catalog)) = (&self.role, &latest) + && let Poll::Ready(Ok(())) = root.poll(waiter, |media| match **media != catalog.media { + true => Poll::Ready(()), + false => Poll::Pending, + }) { + return Poll::Ready(Ok(latest.take())); + } + Poll::Pending + } + + /// Reconcile the renditions with `catalog`. + fn sync(&mut self, mut catalog: catalog::hang::Catalog) { + match &self.role { + Role::Root(Some(contributions)) => contributions.sync(&self.source, &catalog), Role::Root(None) => {} Role::Contribution(root) => catalog.media = *root.read(), } - renditions.sync(&source, &catalog); - if let Role::Contribution(_) = role { - latest = Some(catalog); + self.renditions.sync(&self.source, &catalog); + if let (Role::Contribution(_), Watch::Reading { latest, .. }) = (&self.role, &mut self.state) { + *latest = Some(catalog); + } + } + + /// Record the catalog's end in the renditions. + fn finish(&self, result: crate::Result<()>) { + let failed = result.is_err(); + self.renditions.finish(result); + // The root catalog lists the contributions: once it fails, stop following and advertising + // them too. A clean end freezes the final catalog, which keeps its entries, as the own + // renditions keep their final set. + if failed && let Role::Root(Some(contributions)) = &self.role { + contributions.clear(); + } + } +} + +impl kio::Task for WatchCatalog { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + loop { + match &mut self.state { + Watch::Start => { + self.broadcast = Some(self.source.broadcast()); + self.state = self.subscribe(); + } + Watch::Subscribing(subscribing) => { + let subscribed = match subscribing { + Ok((_, info)) => ready!(info.poll_ok(waiter)).map(drop).map_err(moq_mux::Error::from), + Err(err) => Err(err.clone()), + }; + let watch = std::mem::replace(&mut self.state, Watch::Start); + self.state = match subscribed { + Ok(()) => { + let Watch::Subscribing(Ok((consumer, _))) = watch else { + unreachable!("subscribed above"); + }; + Watch::Reading { + consumer: consumer.into(), + latest: None, + } + } + Err(err) => { + tracing::warn!(%err, "failed to subscribe to broadcast catalog, retrying"); + let wait = self.delay.mul_f64(0.5 + kio::clock::rng().random::() / 2.0); + self.delay = (self.delay * 2).min(CATALOG_RETRY_MAX); + Watch::Backoff(kio::clock::sleep(wait)) + } + }; + } + Watch::Backoff(sleep) => { + let broadcast = self.broadcast.as_ref().expect("taken on the first poll"); + if broadcast.poll_closed(waiter).is_ready() { + self.renditions.close(); + return Poll::Ready(()); + } + ready!(waiter.poll_future(std::pin::Pin::new(sleep))); + self.state = self.subscribe(); + } + Watch::Reading { .. } => match ready!(self.poll_catalog(waiter)) { + Ok(Some(catalog)) => self.sync(catalog), + Ok(None) => { + self.finish(Ok(())); + return Poll::Ready(()); + } + Err(err) => { + self.finish(Err(err)); + return Poll::Ready(()); + } + }, + } } - }; - let failed = result.is_err(); - renditions.finish(result); - // The root catalog lists the contributions: once it fails, stop following and advertising - // them too. A clean end freezes the final catalog, which keeps its entries, as the own - // renditions keep their final set. - if failed && let Role::Root(Some(contributions)) = &role { - contributions.clear(); } } @@ -2694,7 +2780,7 @@ mod tests { drop((catalog, media, broadcast)); } - /// A one-video-rendition catalog snapshot reconciled by hand, which is what `watch_catalog` + /// A one-video-rendition catalog snapshot reconciled by hand, which is what `WatchCatalog` /// does with each snapshot: the rendition follows its media on its own. fn export( upstream: &moq_mux::Source, @@ -2891,7 +2977,7 @@ mod tests { let upstream = crate::resolved(&origin, "live").await; // The publisher goes away and comes back before the catalog snapshot is reconciled. That - // gap is wide in practice: `watch_catalog` retries a not-yet-written catalog track with + // gap is wide in practice: `WatchCatalog` retries a not-yet-written catalog track with // backoff, so seconds can pass between resolving the broadcast and syncing its first // snapshot. The old broadcast ends first: a path still routed is one broadcast, so a // publisher back before then would resume it instead. @@ -3470,7 +3556,7 @@ mod tests { } } - /// Renditions over hand-published media groups, reconciled the way `watch_catalog` does, on a + /// Renditions over hand-published media groups, reconciled the way `WatchCatalog` does, on a /// two-second join interval. struct Groups { broadcast: moq_net::broadcast::Producer, diff --git a/rs/moq-hls/src/export/rendition.rs b/rs/moq-hls/src/export/rendition.rs index 49e6686c3bf..66b60292dea 100644 --- a/rs/moq-hls/src/export/rendition.rs +++ b/rs/moq-hls/src/export/rendition.rs @@ -4,7 +4,7 @@ use kio::clock::Wall; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, RwLock, Weak}; -use std::task::Poll; +use std::task::{Poll, ready}; use std::time::Duration; use bytes::Bytes; @@ -217,20 +217,21 @@ impl Media { /// never looks the path up ad hoc: a same-path republish installs a new broadcast whose /// group numbering restarts, and those bytes must not be served under rows produced for /// the publisher it replaced. - async fn track( + fn poll_track_or_refusal( &self, binding: &moq_mux::Binding, - ) -> Result> { - let track = kio::wait(|waiter| self.poll_track(binding, waiter)).await; + waiter: &kio::Waiter, + ) -> Poll>> { + let track = ready!(self.poll_track(binding, waiter)); // The bind can be refused while this read waits. Preserve that exact instance's // authorization error; other failed or disconnected references stay missing. if track.is_none() && let Poll::Ready(Err(err @ moq_mux::Error::Moq(moq_net::Error::Unauthorized))) = binding.poll_broadcast(&kio::Waiter::noop()) { - return Err(err.into()); + return Poll::Ready(Err(err.into())); } - Ok(track) + Poll::Ready(Ok(track)) } fn poll_track( @@ -405,7 +406,7 @@ pub struct Rendition { /// Follows the media track's groups into [`live`](Self::live); ends when the rendition drops. watcher: Mutex, /// Set once the catalog drops or replaces the rendition (see [`retire`](Self::retire)). - retiring: tokio::sync::watch::Sender, + retiring: kio::Producer, /// Which segments this instance lists. claim: Mutex, /// The retired instance this one replaced under the same playlist URI, whose segments its playlist @@ -421,9 +422,9 @@ pub struct Rendition { /// generation drops the init, since a restarted inline-codec publisher may carry new /// parameter sets under the same catalog. Shared state so a render can wait for the init. run: kio::Producer, - /// Serializes init builds, so concurrent first requests fetch the keyframe group once. Shared - /// so a background build (see [`start_init`](Self::start_init)) holds it for its whole run. - building: Arc>, + /// Serializes init builds, so concurrent first requests fetch the keyframe group once. A + /// background build (see [`start_init`](Self::start_init)) holds it for its whole run. + building: Building, /// The background init build; ends when the rendition drops. build: Mutex>, } @@ -437,6 +438,45 @@ struct Watcher { stopped: bool, } +/// Admits one init build at a time; the others wait for it. +#[derive(Clone, Default)] +struct Building(kio::Shared); + +/// A held [`Building`], released on drop. +struct Built(kio::Shared); + +impl Building { + /// Hold it now, unless a build already does. + fn try_hold(&self) -> Option { + let mut held = self.0.lock(); + if *held { + return None; + } + *held = true; + Some(Built(self.0.clone())) + } + + /// Hold it once no build does. + fn poll_hold(&self, waiter: &kio::Waiter) -> Poll { + let mut held = ready!(self.0.poll(waiter, |held| match **held { + true => Poll::Pending, + false => Poll::Ready(()), + })); + *held = true; + Poll::Ready(Built(self.0.clone())) + } + + fn is_held(&self) -> bool { + *self.0.read() + } +} + +impl Drop for Built { + fn drop(&mut self) { + *self.0.lock() = false; + } +} + /// Keeps lazy media demanded for the whole response, including its first encode. #[cfg(feature = "server")] pub(crate) struct Demand(Option>); @@ -460,40 +500,61 @@ impl Drop for Demand { let duration = target .saturating_mul(3) .max(rendition.fanout.window().saturating_add(target)); - let weak = Arc::downgrade(&rendition); - watcher.idle = Some(rendition.fanout.spawn(async move { - kio::clock::sleep(duration).await; - if let Some(rendition) = weak.upgrade() { - let mut claim = rendition.claim.lock().expect("claim lock poisoned"); - let mut watcher = rendition.watcher.lock().expect("watcher lock poisoned"); - if watcher.requests == 0 { - // A finished window keeps listing what it listed, inherited predecessor rows - // included, and keeps its track pin: reads of its retained groups go through - // it, since a closed sibling refuses any later lookup. Its finished watcher - // stays too, so no request follows the ended track again. - if rendition.live.finished() { - return; - } - rendition.media.handle.lock().expect("media lock poisoned").track.take(); - watcher.media.take(); - // A resumed watcher never rewrites already-addressed segments or lists a - // partial segment it had open when the last player left. - claim.start = claim.next; - claim.pending = false; - if claim.retired { - // A replacement waits for this window to close before reading its own - // track. Stopping the watcher must finish that handover too. - rendition.live.cut(); - } else { - rendition.live.clear(); - rendition.previous.lock().expect("previous lock poisoned").take(); - } - } - } + watcher.idle = Some(rendition.fanout.spawn(Idle { + rendition: Arc::downgrade(&rendition), + duration, + sleep: None, })); } } +/// Releases a lazy rendition's media once it went `duration` without a request. +#[cfg(feature = "server")] +struct Idle { + rendition: Weak, + duration: Duration, + /// Started on the first poll. + sleep: Option, +} + +#[cfg(feature = "server")] +impl kio::Task for Idle { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + let duration = self.duration; + let sleep = self.sleep.get_or_insert_with(|| kio::clock::sleep(duration)); + std::task::ready!(waiter.poll_future(std::pin::Pin::new(sleep))); + let Some(rendition) = self.rendition.upgrade() else { + return Poll::Ready(()); + }; + let mut claim = rendition.claim.lock().expect("claim lock poisoned"); + let mut watcher = rendition.watcher.lock().expect("watcher lock poisoned"); + // A finished window keeps listing what it listed, inherited predecessor rows included, + // and keeps its track pin: reads of its retained groups go through it, since a closed + // sibling refuses any later lookup. Its finished watcher stays too, so no request follows + // the ended track again. + if watcher.requests != 0 || rendition.live.finished() { + return Poll::Ready(()); + } + rendition.media.handle.lock().expect("media lock poisoned").track.take(); + watcher.media.take(); + // A resumed watcher never rewrites already-addressed segments or lists a partial segment + // it had open when the last player left. + claim.start = claim.next; + claim.pending = false; + if claim.retired { + // A replacement waits for this window to close before reading its own track. + // Stopping the watcher must finish that handover too. + rendition.live.cut(); + } else { + rendition.live.clear(); + rendition.previous.lock().expect("previous lock poisoned").take(); + } + Poll::Ready(()) + } +} + impl Rendition { /// Whether `config` describes the same media this rendition is already serving, i.e. whether /// it decodes and muxes identically from the same track. Fields the publisher revises @@ -623,12 +684,12 @@ impl Rendition { live: Arc::new(segments::Producer::new(0)), fanout, watcher: Mutex::default(), - retiring: tokio::sync::watch::Sender::new(false), + retiring: kio::Producer::new(false), claim: Mutex::default(), previous: Mutex::new(None), media: Arc::new(media), run: kio::Producer::default(), - building: Arc::new(tokio::sync::Mutex::new(())), + building: Building::default(), build: Mutex::new(None), }) } @@ -666,21 +727,19 @@ impl Rendition { live: Arc::new(segments::Producer::new(0)), fanout, watcher: Mutex::default(), - retiring: tokio::sync::watch::Sender::new(false), + retiring: kio::Producer::new(false), claim: Mutex::default(), previous: Mutex::new(None), media: Arc::new(media), run: kio::Producer::default(), - building: Arc::new(tokio::sync::Mutex::new(())), + building: Building::default(), build: Mutex::new(None), }) } /// Start following this rendition's media track onto the grid. pub(crate) fn watch(self: &Arc) { - let watcher = self - .fanout - .spawn(watch_media(Arc::downgrade(self), self.retiring.subscribe())); + let watcher = self.fanout.spawn(WatchMedia::new(self)); self.watcher.lock().expect("watcher lock poisoned").media = Some(watcher); } @@ -698,10 +757,7 @@ impl Rendition { watcher.idle.take(); watcher.requests += 1; if watcher.media.is_none() && !watcher.stopped { - watcher.media = Some( - self.fanout - .spawn(watch_media(Arc::downgrade(self), self.retiring.subscribe())), - ); + watcher.media = Some(self.fanout.spawn(WatchMedia::new(self))); } Demand(Some(self.clone())) } @@ -759,7 +815,10 @@ impl Rendition { } (claim.next, claim.pending) }; - self.retiring.send_replace(true); + *self + .retiring + .write() + .unwrap_or_else(|_| unreachable!("the rendition holds its producer")) = true; // Once retired, one that neither lists nor waits to list a segment never will, so the // replacement keeps only what this one was keeping: churn can't stack empty instances. let lists = pending || !self.live.window().segments.is_empty(); @@ -1004,16 +1063,12 @@ impl Rendition { self.live.close(); } - /// Resolve once the playlist is renderable or failed, i.e. once - /// [`media_playlist`](Self::media_playlist) would return `Some` or an error. Bounding the wait - /// is the caller's policy (the serve path wraps this in its own timeout). - pub async fn playable(&self) { - kio::wait(|waiter| self.poll_playable(waiter)).await; - } - - /// Poll until this rendition is playable (it lists a segment, or its media track ended), its - /// media is denied, or its window failed. - pub(crate) fn poll_playable(&self, waiter: &kio::Waiter) -> Poll<()> { + /// Ready once the playlist is renderable or failed, i.e. once + /// [`media_playlist`](Self::media_playlist) would return `Some` or an error: it lists a + /// segment, or its media track ended, its media is denied, or its window failed. Bounding the + /// wait is the caller's policy (the serve path wraps this in its own timeout). + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_playable(&self, waiter: &kio::Waiter) -> Poll<()> { self.poll_media(waiter); if self.media.refusal().is_some() { return Poll::Ready(()); @@ -1030,15 +1085,12 @@ impl Rendition { Poll::Pending } - /// Resolve once the DASH [representation](Self::representation) lists a segment or this - /// rendition fails, as [`playable`](Self::playable) does for the media playlist. - pub(crate) async fn represented(&self) { - kio::wait(|waiter| self.poll_represented(waiter)).await; - } - - /// Poll until this instance lists a segment of its own (or its media track ended), its media is - /// denied, or its window failed. Unlike [`poll_playable`](Self::poll_playable), a - /// predecessor's segments don't count: the DASH manifest lists none of them. + /// Ready once the DASH [representation](Self::representation) lists a segment or this + /// rendition fails, as [`poll_playable`](Self::poll_playable) does for the media playlist: + /// this instance lists a segment of its own (or its media track ended), its media is denied, + /// or its window failed. A predecessor's segments don't count: the DASH manifest lists none of + /// them. + #[cfg_attr(test, kio::wrapper)] pub(crate) fn poll_represented(&self, waiter: &kio::Waiter) -> Poll<()> { self.poll_media(waiter); if self.media.refusal().is_some() { @@ -1108,35 +1160,8 @@ impl Rendition { /// /// A failed rendition or catalog fails it as in [`media_playlist`](Self::media_playlist). pub async fn playlist(&self, query: Option<&str>) -> Result> { - // A denied reference is refused before anything else, so no other error masks it. - self.authorize()?; - self.playable().await; - // A failed media track may never supply the keyframe an inline init waits for. - self.check()?; - // Build the inits the listed segments map. An inline-codec init needs a keyframe group - // fetched first. Building it caches, so the player's follow-up GET of the init is free. - let own = self.live.window().segments; - let mut inherited = self.inherited(&own); - inherited.dedup_by(|(a, _), (b, _)| Arc::ptr_eq(a, b)); - for (rendition, _) in &inherited { - if rendition.load_init().await?.is_none() { - return Ok(None); - } - } - // This instance's own init is mapped only once it lists a segment (see `media_playlist`). - let mapped = !own.is_empty() || inherited.is_empty(); - if mapped && self.load_init().await?.is_none() { - return Ok(None); - } - let rendered = self.media_playlist(query)?; - if rendered.is_some() || mapped { - return Ok(rendered); - } - // It listed its first segment meanwhile. - if self.load_init().await?.is_none() { - return Ok(None); - } - self.media_playlist(query) + let mut playlist = Playlist::Start; + kio::wait(|waiter| playlist.poll(self, waiter, query)).await } /// Snapshot the media playlist from the current window. @@ -1344,7 +1369,7 @@ impl Rendition { /// Fetch and transmux the segment of the track tagged `tag` at `at`: its number, or /// its start in milliseconds (DASH `$Time$`). - async fn segment_at(&self, tag: &str, at: segments::At) -> Result> { + fn fetch_at(&self, tag: &str, at: segments::At) -> Result> { self.authorize()?; let sequence = match at { segments::At::Sequence(sequence) => sequence, @@ -1353,7 +1378,7 @@ impl Rendition { None => return Ok(None), }, }; - self.numbered(tag, sequence).await + self.fetch_numbered(tag, sequence) } /// The wall-clock time of a media timestamp, when the catalog advertises its broadcast @@ -1416,13 +1441,16 @@ impl Rendition { return Ok(Some(init.bytes.clone())); } } - let init = self.load_init().await?; + let mut load = LoadInit::default(); + let init = kio::wait(|waiter| load.poll(self, waiter)).await?; Ok(init.filter(|init| init.hash == hash).map(|init| init.bytes.clone())) } /// The current CMAF init segment, whatever its hash, built on first request and cached. - pub(crate) async fn latest_init(&self) -> Result> { - Ok(self.load_init().await?.map(|init| init.bytes.clone())) + /// `load` holds the progress across polls, starting from [`LoadInit::default`]. + pub(crate) fn poll_latest_init(&self, waiter: &kio::Waiter, load: &mut LoadInit) -> Poll>> { + let init = ready!(load.poll(self, waiter))?; + Poll::Ready(Ok(init.map(|init| init.bytes.clone()))) } /// `run`'s init if it is already built, or can be without media: an out-of-band codec @@ -1431,7 +1459,7 @@ impl Rendition { if let Some(init) = &run.init { return Some(init.clone()); } - // A muxer error surfaces from `bootstrap`, which builds from media. + // A muxer error surfaces from a `LoadInit`, which builds from media. let bytes = self.muxer().ok()?.init().ok()??; Some(self.cache_init(run.epoch, bytes)) } @@ -1461,7 +1489,7 @@ impl Rendition { return; } // Taken before spawning, so renders never queue tasks behind a build stuck on an open group. - let Ok(building) = self.building.clone().try_lock_owned() else { + let Some(building) = self.building.try_hold() else { return; }; let read = match self.muxer().map(|muxer| self.init_read(muxer)) { @@ -1472,20 +1500,14 @@ impl Rendition { return; } }; - let media = self.media.clone(); - let rendition = Arc::downgrade(self); - let name = self.name.clone(); - let build = self.fanout.spawn(async move { - let _building = building; - match read.read(&media, &binding).await { - Ok(Some(bytes)) => { - if let Some(rendition) = rendition.upgrade() { - rendition.cache_init(run.epoch, bytes); - } - } - Ok(None) => {} - Err(err) => tracing::warn!(rendition = %name, %err, "failed to build init segment"), - } + let build = self.fanout.spawn(InitBuild { + _building: building, + read, + media: self.media.clone(), + binding, + rendition: Arc::downgrade(self), + name: self.name.clone(), + epoch: run.epoch, }); *self.build.lock().expect("build lock poisoned") = Some(build); } @@ -1493,46 +1515,13 @@ impl Rendition { /// Whether the init is missing while a build of it is running, so a later render may list /// this rendition. pub(crate) fn init_pending(&self) -> bool { - self.run().init.is_none() && self.building.try_lock().is_err() - } - - async fn load_init(&self) -> Result>> { - let binding = self.authorize()?; - // A transmuxed segment can cache the init while another bootstrap holds `building`, - // blocked on an unfinished group. - if let Some(init) = self.run().init { - return Ok(Some(init)); - } - let _building = self.building.lock().await; - self.bootstrap(&binding).await - } - - /// Build and cache the init, holding `building`. - async fn bootstrap(&self, binding: &moq_mux::Binding) -> Result>> { - let run = self.run(); - if let Some(init) = run.init { - return Ok(Some(init)); - } - - // An out-of-band codec builds its init straight from the catalog; an inline one returns - // `None` until a keyframe group is read. - let muxer = self.muxer()?; - if let Some(bytes) = muxer.init()? { - return Ok(Some(self.cache_init(run.epoch, bytes))); - } - let Some(read) = self.init_read(muxer) else { - return Ok(None); - }; - let Some(bytes) = read.read(&self.media, binding).await? else { - return Ok(None); - }; - Ok(Some(self.cache_init(run.epoch, bytes))) + self.run().init.is_none() && self.building.is_held() } /// The listed groups an inline-codec init is read from, or `None` while none is listed. fn init_read(&self, muxer: Muxer) -> Option { let groups = self.live.init_groups(); - (!groups.is_empty()).then_some(InitRead { muxer, groups }) + (!groups.is_empty()).then(|| InitRead::new(muxer, groups)) } /// Fetch and transmux the segment a [`Resource::Segment`](super::Resource::Segment) URL names. @@ -1550,111 +1539,425 @@ impl Rendition { pub async fn segment(&self, id: &segments::Id) -> Result> { // A predecessor's segments still listed are muxed under its config. A config that returns // gets its tag back, but the instances' segment numbers never overlap. - if let Some(rendition) = self + let owner = self .lineage() .into_iter() - .find(|rendition| *rendition.tag == *id.tag && rendition.live.holds(&id.tag, id.at)) - { - return rendition.own_segment(id).await; - } - self.own_segment(id).await + .find(|rendition| *rendition.tag == *id.tag && rendition.live.holds(&id.tag, id.at)); + let owner = owner.as_deref().unwrap_or(self); + let Some(mut fetch) = owner.fetch_own(id)? else { + return Ok(None); + }; + let segment = kio::wait(|waiter| fetch.poll(owner, waiter)).await?; + // A generation change while fetching may have swapped the rows under the lookup. + Ok(segment.filter(|_| owner.generation().as_deref() == id.generation.as_deref())) } - /// Fetch and transmux the segment `id` names among this instance's own rows; see + /// The fetch of the segment `id` names among this instance's own rows; see /// [`segment`](Self::segment). - async fn own_segment(&self, id: &segments::Id) -> Result> { + fn fetch_own(&self, id: &segments::Id) -> Result> { self.authorize()?; - let current = || self.generation().as_deref() == id.generation.as_deref(); - if !current() { + if self.generation().as_deref() != id.generation.as_deref() { return Ok(None); } - let segment = self.segment_at(&id.tag, id.at).await?; - // A generation change while fetching may have swapped the rows under the lookup. - Ok(segment.filter(|_| current())) + self.fetch_at(&id.tag, id.at) } - /// Fetch and transmux the segment numbered `segment` on the track tagged `tag`, - /// whatever the run; see [`segment`](Self::segment). - pub(crate) async fn numbered(&self, tag: &str, segment: u64) -> Result> { + /// The fetch of the segment numbered `segment` on the track tagged `tag`, whatever the run; + /// see [`segment`](Self::segment). + fn fetch_numbered(&self, tag: &str, segment: u64) -> Result> { self.authorize()?; - let Some(row) = self.live.row(tag, segment) else { - return Ok(None); - }; - self.fetch(&row).await + Ok(self.live.row(tag, segment).map(Fetch::new)) } +} - /// Fetch and transmux `row`'s groups on this rendition; see [`segment`](Self::segment). - pub(crate) async fn fetch(&self, row: &segments::Row) -> Result> { - let binding = self.authorize()?; - // Read the run before resolving the row (see `restart`), so an init decoded from its - // groups is never cached for a run that started while they were being fetched. - let run = self.run(); - if row.stop { - return Ok(None); - } - let groups = row.groups.clone(); +/// Fetches and transmuxes a row's groups on a rendition; see [`Rendition::segment`]. +pub(crate) struct Fetch { + row: segments::Row, + state: Fetching, +} - // A segment must be served whole. If its group span exceeds the safety cap (a publisher - // grouping far finer than its interval), 404 rather than silently truncating it to a - // segment shorter than its advertised duration. The live window lists such a segment as a - // gap instead, so only a row from elsewhere reaches this. - if groups.end.saturating_sub(groups.start) > self.max_groups() { - tracing::warn!( - ?groups, - "segment spans more groups than one request fetches; refusing to truncate" - ); - return Ok(None); +#[allow( + clippy::large_enum_variant, + reason = "one per segment request, and the reading variant is the working state" +)] +enum Fetching { + Start, + Track { + binding: Arc, + epoch: u64, + }, + Groups { + _binding: Arc, + epoch: u64, + broadcast: moq_net::broadcast::Consumer, + track: moq_net::track::Consumer, + muxer: Muxer, + /// Every group's frames, accumulated into ONE fragment, so duration inference sees each + /// frame's true successor. A per-group fragment would mis-time the trailing sample of + /// every group (its real successor lives in the next group), audible on audio segments + /// that span many groups. + frames: Vec, + /// The group being read, and the ones after it. + groups: std::ops::Range, + group: Option, + }, + Done, +} + +impl Fetch { + pub(crate) fn new(row: segments::Row) -> Self { + Self { + row, + state: Fetching::Start, } + } - let Some((broadcast, track)) = self.media.track(&binding).await? else { - return Ok(None); - }; + /// The row fetched. + pub(crate) fn row(&self) -> &segments::Row { + &self.row + } - let mut muxer = self.muxer()?; - // Accumulate every group's frames into ONE fragment, so duration inference sees each - // frame's true successor. A per-group fragment would mis-time the trailing sample of every - // group (its real successor lives in the next group), audible on audio segments that span - // many groups. - let mut frames = Vec::new(); - for sequence in groups { - let miss = match fetch(&track, sequence, &broadcast).await? { - Ok(mut group) => match read_group(&mut muxer, &mut group, &broadcast).await? { - Ok(mut group_frames) => { - frames.append(&mut group_frames); - continue; + /// The row's CMAF fragment, or `None` when it is a stop or its media is gone. + pub(crate) fn poll(&mut self, rendition: &Rendition, waiter: &kio::Waiter) -> Poll>> { + let result = ready!(self.poll_state(rendition, waiter)); + self.state = Fetching::Done; + Poll::Ready(result) + } + + fn poll_state(&mut self, rendition: &Rendition, waiter: &kio::Waiter) -> Poll>> { + let row = &self.row; + loop { + match &mut self.state { + Fetching::Start => { + let binding = rendition.authorize()?; + // Read the run before resolving the row (see `restart`), so an init decoded from + // its groups is never cached for a run that started while they were being fetched. + let epoch = rendition.run().epoch; + if row.stop { + return Poll::Ready(Ok(None)); } - Err(err) => err, - }, - Err(err) => err, - }; - tracing::warn!( - segment = row.segment, - group = sequence, - err = %miss, - "abandoning an advertised segment: a group is unavailable" - ); - return Ok(None); + // A segment must be served whole. If its group span exceeds the safety cap (a + // publisher grouping far finer than its interval), 404 rather than silently + // truncating it to a segment shorter than its advertised duration. The live window + // lists such a segment as a gap instead, so only a row from elsewhere reaches this. + if row.groups.end.saturating_sub(row.groups.start) > rendition.max_groups() { + tracing::warn!( + groups = ?row.groups, + "segment spans more groups than one request fetches; refusing to truncate" + ); + return Poll::Ready(Ok(None)); + } + self.state = Fetching::Track { binding, epoch }; + } + Fetching::Track { binding, epoch } => { + let Some((broadcast, track)) = ready!(rendition.media.poll_track_or_refusal(binding, waiter))? + else { + return Poll::Ready(Ok(None)); + }; + self.state = Fetching::Groups { + _binding: binding.clone(), + epoch: *epoch, + broadcast, + track, + muxer: rendition.muxer()?, + frames: Vec::new(), + groups: row.groups.clone(), + group: None, + }; + } + Fetching::Groups { + broadcast, + track, + muxer, + frames, + groups, + group, + .. + } => { + let read = match group { + Some(read) => read, + None => match groups.next() { + Some(sequence) => group.insert(GroupRead::new(track, sequence)), + None => break, + }, + }; + let sequence = read.sequence; + let miss = match ready!(read.poll(waiter, muxer, broadcast))? { + Ok(mut read) => { + frames.append(&mut read); + *group = None; + continue; + } + Err(miss) => miss, + }; + tracing::warn!( + segment = row.segment, + group = sequence, + err = %miss, + "abandoning an advertised segment: a group is unavailable" + ); + return Poll::Ready(Ok(None)); + } + Fetching::Done => panic!("polled a finished segment fetch"), + } } + let Fetching::Groups { + epoch, muxer, frames, .. + } = &mut self.state + else { + unreachable!("read every group above"); + }; if frames.is_empty() { - return Ok(None); + return Poll::Ready(Ok(None)); } // The last presented sample has no successor to time it, so the fragment ends where the // next group starts, unless its codec or a duration marker times it. let end = row .content .and_then(|content| moq_net::Timestamp::from_micros(content.as_micros() as u64).ok()); - let fragment = muxer.fragment(row.segment as u32, &frames, end)?; + let fragment = muxer.fragment(row.segment as u32, frames, end)?; // Inline codec metadata was already decoded from these completed groups. // Reuse it instead of fetching a newer, possibly unfinished group. if let Some(bytes) = muxer.init()? { - self.cache_init(run.epoch, bytes); + rendition.cache_init(*epoch, bytes); + } + Poll::Ready(Ok(Some(fragment))) + } +} + +/// Loads the current run's init, building and caching it from media if needed. +#[derive(Default)] +pub(crate) struct LoadInit(Loading); + +#[derive(Default)] +#[allow( + clippy::large_enum_variant, + reason = "one per init load, and the reading variant is the working state" +)] +enum Loading { + #[default] + Start, + /// Waiting for any other build to finish. + Holding(Arc), + /// Reading media for the init, holding the build. + Reading { + binding: Arc, + _building: Built, + epoch: u64, + read: InitRead, + }, + Done, +} + +impl LoadInit { + fn poll(&mut self, rendition: &Rendition, waiter: &kio::Waiter) -> Poll>>> { + let result = ready!(self.0.poll(rendition, waiter)); + // Releases the build as the load returns. + self.0 = Loading::Done; + Poll::Ready(result) + } +} + +impl Loading { + fn poll(&mut self, rendition: &Rendition, waiter: &kio::Waiter) -> Poll>>> { + loop { + match self { + Self::Start => { + let binding = rendition.authorize()?; + // A transmuxed segment can cache the init while another bootstrap holds + // `building`, blocked on an unfinished group. + if let Some(init) = rendition.run().init { + return Poll::Ready(Ok(Some(init))); + } + *self = Self::Holding(binding); + } + Self::Holding(binding) => { + let building = ready!(rendition.building.poll_hold(waiter)); + let run = rendition.run(); + if let Some(init) = run.init { + return Poll::Ready(Ok(Some(init))); + } + // An out-of-band codec builds its init straight from the catalog; an inline one + // returns `None` until a keyframe group is read. + let muxer = rendition.muxer()?; + if let Some(bytes) = muxer.init()? { + return Poll::Ready(Ok(Some(rendition.cache_init(run.epoch, bytes)))); + } + let Some(read) = rendition.init_read(muxer) else { + return Poll::Ready(Ok(None)); + }; + *self = Self::Reading { + binding: binding.clone(), + _building: building, + epoch: run.epoch, + read, + }; + } + Self::Reading { + binding, epoch, read, .. + } => { + let Some(bytes) = ready!(read.poll(waiter, &rendition.media, binding))? else { + return Poll::Ready(Ok(None)); + }; + return Poll::Ready(Ok(Some(rendition.cache_init(*epoch, bytes)))); + } + Self::Done => panic!("polled a finished init load"), + } + } + } +} + +/// Renders a media playlist once it is servable; see [`Rendition::playlist`]. +enum Playlist { + Start, + Playable, + /// Building the init of each inherited instance in turn, then this instance's own. + Inits { + inherited: Vec>, + /// This instance's own init is mapped only once it lists a segment (see `media_playlist`). + mapped: bool, + load: LoadInit, + }, + /// Building this instance's own init, which it mapped since the render was checked. + Again(LoadInit), + Done, +} + +impl Playlist { + fn poll( + &mut self, + rendition: &Rendition, + waiter: &kio::Waiter, + query: Option<&str>, + ) -> Poll>> { + let result = ready!(self.poll_state(rendition, waiter, query)); + *self = Self::Done; + Poll::Ready(result) + } + + fn poll_state( + &mut self, + rendition: &Rendition, + waiter: &kio::Waiter, + query: Option<&str>, + ) -> Poll>> { + loop { + match self { + Self::Start => { + // A denied reference is refused before anything else, so no other error masks it. + rendition.authorize()?; + *self = Self::Playable; + } + Self::Playable => { + ready!(rendition.poll_playable(waiter)); + // A failed media track may never supply the keyframe an inline init waits for. + rendition.check()?; + // Build the inits the listed segments map. An inline-codec init needs a keyframe + // group fetched first. Building it caches, so the player's follow-up GET of the + // init is free. + let own = rendition.live.window().segments; + let mut inherited = rendition.inherited(&own); + inherited.dedup_by(|(a, _), (b, _)| Arc::ptr_eq(a, b)); + let mapped = !own.is_empty() || inherited.is_empty(); + // Loaded in reverse, popping the oldest first. + let inherited = inherited.into_iter().rev().map(|(rendition, _)| rendition).collect(); + *self = Self::Inits { + inherited, + mapped, + load: LoadInit::default(), + }; + } + Self::Inits { + inherited, + mapped, + load, + } => { + let loading = match inherited.last() { + Some(previous) => &**previous, + None if *mapped => rendition, + None => { + let rendered = rendition.media_playlist(query)?; + if rendered.is_some() { + return Poll::Ready(Ok(rendered)); + } + // It listed its first segment meanwhile. + *self = Self::Again(LoadInit::default()); + continue; + } + }; + if ready!(load.poll(loading, waiter))?.is_none() { + return Poll::Ready(Ok(None)); + } + *load = LoadInit::default(); + if inherited.pop().is_none() { + // This instance's own init, which it maps. + return Poll::Ready(rendition.media_playlist(query)); + } + } + Self::Again(load) => { + if ready!(load.poll(rendition, waiter))?.is_none() { + return Poll::Ready(Ok(None)); + } + return Poll::Ready(rendition.media_playlist(query)); + } + Self::Done => panic!("polled a finished playlist render"), + } } - Ok(Some(fragment)) + } +} + +/// Builds an init in the background, holding the build until it ends; see +/// [`Rendition::start_init`]. +struct InitBuild { + _building: Built, + read: InitRead, + media: Arc, + binding: Arc, + rendition: Weak, + name: String, + epoch: u64, +} + +impl kio::Task for InitBuild { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + match ready!(self.read.poll(waiter, &self.media, &self.binding)) { + Ok(Some(bytes)) => { + if let Some(rendition) = self.rendition.upgrade() { + rendition.cache_init(self.epoch, bytes); + } + } + Ok(None) => {} + Err(err) => tracing::warn!(rendition = %self.name, %err, "failed to build init segment"), + } + Poll::Ready(()) } } #[cfg(test)] impl Rendition { + /// The current CMAF init segment, whatever its hash, built on first request and cached. + pub(crate) async fn latest_init(&self) -> Result> { + let mut load = LoadInit::default(); + kio::wait(|waiter| self.poll_latest_init(waiter, &mut load)).await + } + + /// Fetch and transmux the segment numbered `segment` on the track tagged `tag`, whatever the + /// run. + pub(crate) async fn numbered(&self, tag: &str, segment: u64) -> Result> { + let Some(mut fetch) = self.fetch_numbered(tag, segment)? else { + return Ok(None); + }; + kio::wait(|waiter| fetch.poll(self, waiter)).await + } + + /// Fetch and transmux `row`'s groups. + pub(crate) async fn fetch(&self, row: &segments::Row) -> Result> { + let mut fetch = Fetch::new(row.clone()); + kio::wait(|waiter| fetch.poll(self, waiter)).await + } + /// Fetch the listed segment numbered `segment`. pub(crate) async fn listed_segment(&self, segment: u64) -> Result> { self.numbered(&self.tag.clone(), segment).await @@ -1662,7 +1965,10 @@ impl Rendition { /// Fetch the listed segment starting at `time` milliseconds. pub(crate) async fn listed_segment_at(&self, time: u64) -> Result> { - self.segment_at(&self.tag.clone(), segments::At::Pts(time)).await + let Some(mut fetch) = self.fetch_at(&self.tag.clone(), segments::At::Pts(time))? else { + return Ok(None); + }; + kio::wait(|waiter| fetch.poll(self, waiter)).await } } @@ -1678,267 +1984,431 @@ fn broken(segments: &[u64], inherited: usize) -> usize { .unwrap_or(0) } -/// Resolve once `rendition` is retired with no segment left waiting to list, or is gone. +/// Ready once `rendition` is retired with no segment left waiting to list, or is gone. /// -/// The claim is read on every change, so a retire that leaves nothing pending wakes it. -async fn retired(retiring: &mut tokio::sync::watch::Receiver, rendition: &Weak) { - let done = |retiring: &bool| { - *retiring - && rendition - .upgrade() - .is_none_or(|rendition| !rendition.claim.lock().expect("claim lock poisoned").pending) +/// The claim is read on every poll, so a retire that leaves nothing pending wakes it. +fn poll_retired(retiring: &kio::Consumer, rendition: &Weak, waiter: &kio::Waiter) -> Poll<()> { + // Upgraded outside the channel's lock, so the last strong reference never drops under it. + let Some(rendition) = rendition.upgrade() else { + return Poll::Ready(()); }; - // A dropped sender means the rendition is gone, which ends the watch all the same. - let _ = retiring.wait_for(done).await; + let done = |retiring: &bool| *retiring && !rendition.claim.lock().expect("claim lock poisoned").pending; + // A closed channel means the rendition is gone, which ends the watch all the same. + retiring + .poll(waiter, |retiring| match done(retiring) { + true => Poll::Ready(()), + false => Poll::Pending, + }) + .map(|_| ()) } -/// Follow a rendition's media track into its window, numbering its groups on the grid, until the +/// Follows a rendition's media track into its window, numbering its groups on the grid, until the /// track ends or fails. /// /// A segment is listed once a later one starts, which bounds its groups and its `EXTINF`, or once /// the track ends. A watcher joining mid-segment can't know the groups before its first, so it /// lists nothing before the first segment it sees start (see [`grid::Grid::starts`]): another edge /// lists that segment whole, and the same URL must name the same bytes everywhere. -async fn watch_media(rendition: Weak, mut retiring: tokio::sync::watch::Receiver) { - let Some((track, interval, segment, container, mut probe)) = rendition.upgrade().map(|r| { - ( - r.track().to_string(), - r.interval, - r.segment, - r.container.clone(), - r.muxer().ok(), - ) - }) else { - return; - }; - let result: Result<()> = async { - // A replacement lists nothing until its predecessor has listed all it will, so their - // segments never interleave and its own start a new discontinuity: they decode under - // another init. It then starts where the predecessor finally stopped, and on the same - // track past the last group the predecessor read, which another config may have muxed. - let previous = rendition.upgrade().and_then(|rendition| rendition.previous()); - let mut broken = previous.is_some(); - let mut after: Option = None; - if let Some(previous) = previous { - tokio::select! { - _ = previous.live.closed() => {} - _ = retired(&mut retiring, &rendition) => return Ok(()), +struct WatchMedia { + rendition: Weak, + retiring: kio::Consumer, + /// Taken from the rendition on the first poll. + follow: Option, + /// The rendition's media and the binding it is read through, held while the watch runs. + held: Option<(Arc, Arc)>, + groups: Option, + step: Step, +} + +/// Where a [`WatchMedia`] is. +#[allow( + clippy::large_enum_variant, + reason = "one per watched rendition, and the subscribing variant is short-lived" +)] +enum Step { + Start, + /// Waiting for the predecessor to list all it will. + Previous(Arc), + Bind, + /// Waiting for the bound broadcast to serve the track. + Track, + Subscribing { + _track: moq_net::track::Consumer, + subscribing: moq_net::track::Subscribing, + }, + /// Waiting for the next group. + Next, + /// Reading a group's first frame, which places it. + First(moq_net::group::Consumer), + /// Reading the last group to its end once the track ended, for where its content ends. + End(moq_net::group::Consumer, Option), +} + +/// A group read, and its first pts unless it is a marker. +type Read = (u64, Option); + +/// What a [`WatchMedia`] numbers its groups with, and how far it got. +struct Follow { + track: String, + interval: Duration, + container: Arc, + /// Reads each group's picture size. + probe: Option, + /// The next segment listed starts a new discontinuity. + broken: bool, + /// The last group a predecessor on the same track read: only later ones are this instance's. + after: Option, + grid: grid::Grid, + /// The newest group, its first pts, and what its first frame reached, read again for its + /// content's end if the track ends after it, and the gap from the group before it. + latest: Option<(moq_net::group::Consumer, moq_net::Timestamp, Reach)>, + spacing: Option, + /// Nothing is listed before the first segment this watcher saw start. + whole: bool, +} + +impl Follow { + fn new(rendition: &Rendition) -> Self { + Self { + track: rendition.track().to_string(), + interval: rendition.interval, + container: rendition.container.clone(), + probe: rendition.muxer().ok(), + broken: false, + after: None, + grid: grid::Grid::new(rendition.segment, segments::MAX_SEGMENTS), + latest: None, + spacing: None, + whole: false, + } + } + + /// Place a group by its first frame and list what that places; `Some` once the watch is done. + fn first( + &mut self, + rendition: &Weak, + group: moq_net::group::Consumer, + first: Option, + ) -> Option> { + // A group's place is its first frame's pts: the transport's stamp, or the container's own + // where the session carries none. + let Some(first) = first.and_then(|frame| first_frame(&self.container, frame, self.probe.as_mut())) else { + tracing::warn!(track = %self.track, group = group.sequence, "skipping a group without a timed first frame"); + return None; + }; + let (slots, read, size) = match first { + // A discontinuity marker carries no media: the content before it ends there, and the + // grid numbers what resumes by its own timestamps, as every edge does. + First::Marker(at) => { + self.latest = None; + (self.grid.stop(at), (group.sequence, None), None) } - let (next, boundary) = { - let claim = previous.claim.lock().expect("claim lock poisoned"); - (claim.next, claim.boundary) - }; - // Another track, or the same name in another broadcast, numbers its groups afresh. - let continues = rendition - .upgrade() - .is_some_and(|rendition| rendition.continues(&previous)); - after = boundary.filter(|_| continues); - if let Some(rendition) = rendition.upgrade() { - let mut claim = rendition.claim.lock().expect("claim lock poisoned"); - claim.start = claim.start.max(next); - claim.next = claim.next.max(next); + First::Media(pts, reach, declared) => { + let sequence = group.sequence; + if let Some((_, previous, _)) = &self.latest { + self.spacing = pts.checked_sub(*previous).ok().map(Duration::from).or(self.spacing); + } + self.latest = Some((group, pts, reach)); + (self.grid.group(sequence, pts), (sequence, Some(pts)), declared) } - } - let (media, binding) = match rendition.upgrade() { - Some(rendition) => (rendition.media.clone(), rendition.authorize()?), - None => return Ok(()), }; - let subscribe = async { - let Some((_, track)) = media.track(&binding).await? else { - return Ok(None); - }; - let subscription = moq_net::track::Subscription::default() - .with_delay(moq_net::track::Delay::default().with_max(FEED_MAX_DELAY)); - Result::Ok(Some(track.subscribe(subscription).await?.ordered())) - }; - let mut groups = tokio::select! { - groups = subscribe => match groups? { - Some(groups) => groups, - // The bound broadcast never serves the track: nothing to list. - None => return Ok(()), - }, - _ = retired(&mut retiring, &rendition) => return Ok(()), + self.list(rendition, slots, Some(read), size, false) + } + + /// List what `slots` place, after reading the group `read` whose keyframe declares `size`; + /// `Some` once the watch is done. + fn list( + &mut self, + rendition: &Weak, + slots: Vec, + read: Option, + size: Option<(u32, u32)>, + ended: bool, + ) -> Option> { + let Self { + track, + grid, + broken, + whole, + after, + .. + } = self; + let after = *after; + let Some(rendition) = rendition.upgrade() else { + return Some(Ok(())); }; + // Held while listing, so a retire sees exactly what this instance lists. + let mut claim = rendition.claim.lock().expect("claim lock poisoned"); + // Retired, a group past the last it had read is the replacement's: this instance lists + // the segments and gaps before the first holding one, and its own part of that one, then + // stops there. + let boundary = claim.boundary; + let theirs = |groups: &std::ops::Range| boundary.is_none_or(|boundary| groups.end > boundary + 1); + if claim.retired + && let Some((sequence, Some(pts))) = read + && theirs(&(sequence..sequence + 1)) + { + claim.theirs_at = claim.theirs_at.or(Some(pts)); + } + let theirs_at = claim.theirs_at; + for slot in slots { + let Some(segment) = slot.segment() else { + tracing::warn!(%track, "media timestamps went back; resetting the playlist window"); + rendition.live.clear(); + *broken = true; + *whole = false; + continue; + }; + if segment < claim.start { + continue; + } + let (slot, last) = match slot { + grid::Slot::Media { + segment, + groups, + pts, + start, + .. + } if claim.retired && theirs(&groups) => match (boundary, theirs_at) { + // The replacement can't list this segment's number again, so this instance + // lists the groups it had read, ending where the replacement's media starts. + (Some(boundary), Some(at)) if groups.start <= boundary => { + let slot = grid::Slot::Media { + segment, + groups: groups.start..boundary + 1, + pts, + start, + end: at, + content: Some(at), + }; + (slot, true) + } + _ => return Some(Ok(())), + }, + slot => (slot, false), + }; + if !*whole { + // The first segment listed holds only this instance's groups, and starts whole. + let own = match &slot { + grid::Slot::Media { groups, .. } => after.is_none_or(|after| groups.start > after), + _ => false, + }; + if !own || !grid.starts(&slot) { + if last { + return Some(Ok(())); + } + continue; + } + *whole = true; + } + claim.next = segment + 1; + rendition.push(row(&rendition, slot, *broken)); + *broken = false; + if last { + return Some(Ok(())); + } + } + if claim.retired + && grid + .open_groups() + .is_some_and(|groups| theirs(&(groups.start..groups.start + 1))) + { + return Some(Ok(())); + } + if let Some((sequence, _)) = read { + claim.seen = Some(sequence); + } + if let (Some((group, Some(pts))), Some(size)) = (read, size) + && claim.size != Some(size) + { + claim.size = Some(size); + claim.resized = Some(Resize { group, pts, size }); + } + claim.unplaced = grid.unplaced().map(|first| (first, grid.open())); + if !claim.retired + && let Some(open) = grid.open() + { + claim.next = claim.next.max(open + 1); + } + claim.pending = grid.waiting(); + if ended { + return Some(Ok(())); + } + None + } +} + +impl WatchMedia { + fn new(rendition: &Arc) -> Self { + Self { + rendition: Arc::downgrade(rendition), + retiring: rendition.retiring.consume(), + follow: None, + held: None, + groups: None, + step: Step::Start, + } + } - let mut grid = grid::Grid::new(segment, segments::MAX_SEGMENTS); - // The newest group, its first pts, and what its first frame reached, read again for its - // content's end if the track ends after it, and the gap from the group before it. - let mut latest: Option<(moq_net::group::Consumer, moq_net::Timestamp, Reach)> = None; - let mut spacing: Option = None; - // Nothing is listed before the first segment this watcher saw start. - let mut whole = false; + fn poll_watch(&mut self, waiter: &kio::Waiter) -> Poll> { + let follow = self.follow.as_mut().expect("taken on the first poll"); loop { - let next = tokio::select! { - next = groups.next_group() => next?, - _ = retired(&mut retiring, &rendition) => return Ok(()), - }; - let ended = next.is_none(); - // The group read and, unless it is a marker, its first pts. - let mut read: Option<(u64, Option)> = None; - // The picture size the group's keyframe declares. - let mut size = None; - let slots = match next { - Some(mut group) => { - // A group's place is its first frame's pts: the transport's stamp, or the - // container's own where the session carries none. - let first = tokio::select! { - first = group.read_frame() => first?, - _ = retired(&mut retiring, &rendition) => return Ok(()), + let retired = || poll_retired(&self.retiring, &self.rendition, waiter).is_ready(); + match &mut self.step { + Step::Start => { + // A replacement lists nothing until its predecessor has listed all it will, so + // their segments never interleave and its own start a new discontinuity: they + // decode under another init. It then starts where the predecessor finally stopped, + // and on the same track past the last group the predecessor read, which another + // config may have muxed. + let previous = self.rendition.upgrade().and_then(|rendition| rendition.previous()); + follow.broken = previous.is_some(); + self.step = match previous { + Some(previous) => Step::Previous(previous), + None => Step::Bind, }; - let Some(first) = first.and_then(|frame| first_frame(&container, frame, probe.as_mut())) else { - tracing::warn!(%track, group = group.sequence, "skipping a group without a timed first frame"); - continue; + } + Step::Previous(previous) => { + if retired() { + return Poll::Ready(Ok(())); + } + ready!(previous.live.poll_closed(waiter)); + let (next, boundary) = { + let claim = previous.claim.lock().expect("claim lock poisoned"); + (claim.next, claim.boundary) }; - match first { - // A discontinuity marker carries no media: the content before it ends there, - // and the grid numbers what resumes by its own timestamps, as every edge does. - First::Marker(at) => { - read = Some((group.sequence, None)); - latest = None; - grid.stop(at) - } - First::Media(pts, reach, declared) => { - read = Some((group.sequence, Some(pts))); - size = declared; - if let Some((_, previous, _)) = &latest { - spacing = pts.checked_sub(*previous).ok().map(Duration::from).or(spacing); - } - latest = Some((group, pts, reach)); - grid.group(read.expect("set above").0, pts) - } + // Another track, or the same name in another broadcast, numbers its groups afresh. + let continues = self + .rendition + .upgrade() + .is_some_and(|rendition| rendition.continues(previous)); + follow.after = boundary.filter(|_| continues); + if let Some(rendition) = self.rendition.upgrade() { + let mut claim = rendition.claim.lock().expect("claim lock poisoned"); + claim.start = claim.start.max(next); + claim.next = claim.next.max(next); } + self.step = Step::Bind; } - None => { - let (reach, end) = match latest.take() { - Some((group, _, reach)) => content_end(group, &container, reach, spacing, interval).await?, - None => (None, None), + Step::Bind => { + let Some(rendition) = self.rendition.upgrade() else { + return Poll::Ready(Ok(())); }; - grid.finish(reach, end) + self.held = Some((rendition.media.clone(), rendition.authorize()?)); + self.step = Step::Track; } - }; - let Some(rendition) = rendition.upgrade() else { - return Ok(()); - }; - // Held while listing, so a retire sees exactly what this instance lists. - let mut claim = rendition.claim.lock().expect("claim lock poisoned"); - // Retired, a group past the last it had read is the replacement's: this instance lists - // the segments and gaps before the first holding one, and its own part of that one, then - // stops there. - let boundary = claim.boundary; - let theirs = |groups: &std::ops::Range| boundary.is_none_or(|boundary| groups.end > boundary + 1); - if claim.retired - && let Some((sequence, Some(pts))) = read - && theirs(&(sequence..sequence + 1)) - { - claim.theirs_at = claim.theirs_at.or(Some(pts)); - } - let theirs_at = claim.theirs_at; - for slot in slots { - let Some(segment) = slot.segment() else { - tracing::warn!(%track, "media timestamps went back; resetting the playlist window"); - rendition.live.clear(); - broken = true; - whole = false; - continue; - }; - if segment < claim.start { - continue; + Step::Track => { + if retired() { + return Poll::Ready(Ok(())); + } + let (media, binding) = self.held.as_ref().expect("bound above"); + // The bound broadcast never serves the track: nothing to list. + let Some((_, track)) = ready!(media.poll_track_or_refusal(binding, waiter))? else { + return Poll::Ready(Ok(())); + }; + let subscription = moq_net::track::Subscription::default() + .with_delay(moq_net::track::Delay::default().with_max(FEED_MAX_DELAY)); + let subscribing = track.subscribe(subscription).into_inner(); + self.step = Step::Subscribing { + _track: track, + subscribing, + }; } - let (slot, last) = match slot { - grid::Slot::Media { - segment, - groups, - pts, - start, - .. - } if claim.retired && theirs(&groups) => match (boundary, theirs_at) { - // The replacement can't list this segment's number again, so this instance - // lists the groups it had read, ending where the replacement's media starts. - (Some(boundary), Some(at)) if groups.start <= boundary => { - let slot = grid::Slot::Media { - segment, - groups: groups.start..boundary + 1, - pts, - start, - end: at, - content: Some(at), - }; - (slot, true) - } - _ => return Ok(()), - }, - slot => (slot, false), - }; - if !whole { - // The first segment listed holds only this instance's groups, and starts whole. - let own = match &slot { - grid::Slot::Media { groups, .. } => after.is_none_or(|after| groups.start > after), - _ => false, + Step::Subscribing { subscribing, .. } => { + if retired() { + return Poll::Ready(Ok(())); + } + let subscriber = ready!(subscribing.poll_ok(waiter))?; + self.groups = Some(subscriber.ordered()); + self.step = Step::Next; + } + Step::Next => { + if retired() { + return Poll::Ready(Ok(())); + } + let groups = self.groups.as_mut().expect("subscribed above"); + match ready!(groups.poll_next_group(waiter))? { + Some(group) => self.step = Step::First(group), + None => match follow.latest.take() { + Some((group, _, reach)) => self.step = Step::End(group, Some(reach)), + None => { + let slots = follow.grid.finish(None, None); + let done = follow.list(&self.rendition, slots, None, None, true); + return Poll::Ready(done.expect("an ended track ends the watch")); + } + }, + } + } + Step::First(group) => { + if retired() { + return Poll::Ready(Ok(())); + } + let first = ready!(group.poll_read_frame(waiter))?; + let Step::First(group) = std::mem::replace(&mut self.step, Step::Next) else { + unreachable!("reading the first frame"); }; - if !own || !grid.starts(&slot) { - if last { - return Ok(()); - } - continue; + if let Some(done) = follow.first(&self.rendition, group, first) { + return Poll::Ready(done); } - whole = true; } - claim.next = segment + 1; - rendition.push(row(&rendition, slot, broken)); - broken = false; - if last { - return Ok(()); + Step::End(group, reach) => { + // Read once the track has ended so every frame is in. Fails when the group does: + // its last segment would otherwise be listed complete but truncated. + let reach_mut = reach.as_mut().expect("read to its end once"); + while let Some(frame) = ready!(group.poll_read_frame(waiter))? { + // The container says what a frame holds: a marker's payload is empty only once decoded. + if let Ok(decoded) = moq_mux::container::Container::decode(&*follow.container, frame) { + reach_mut.add(decoded); + } + } + let (reach, end) = reach + .take() + .expect("read to its end once") + .finish(follow.spacing, follow.interval); + let slots = follow.grid.finish(reach, end); + let done = follow.list(&self.rendition, slots, None, None, true); + return Poll::Ready(done.expect("an ended track ends the watch")); } } - if claim.retired - && grid - .open_groups() - .is_some_and(|groups| theirs(&(groups.start..groups.start + 1))) - { - return Ok(()); - } - if let Some((sequence, _)) = read { - claim.seen = Some(sequence); - } - if let (Some((group, Some(pts))), Some(size)) = (read, size) - && claim.size != Some(size) - { - claim.size = Some(size); - claim.resized = Some(Resize { group, pts, size }); - } - claim.unplaced = grid.unplaced().map(|first| (first, grid.open())); - if !claim.retired - && let Some(open) = grid.open() - { - claim.next = claim.next.max(open + 1); - } - claim.pending = grid.waiting(); - if ended { - return Ok(()); + } + } +} + +impl kio::Task for WatchMedia { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + if self.follow.is_none() { + let Some(follow) = self.rendition.upgrade().map(|rendition| Follow::new(&rendition)) else { + return Poll::Ready(()); + }; + self.follow = Some(follow); + } + let result = ready!(self.poll_watch(waiter)); + let Some(rendition) = self.rendition.upgrade() else { + return Poll::Ready(()); + }; + let track = self.follow.take().expect("taken on the first poll").track; + match result { + Ok(()) => rendition.end(), + // A bound sibling that ended takes its track down with it: that ends the rendition, + // whose retained rows drain, rather than failing it. + Err(_) if !rendition.media.admits(&rendition.live) => rendition.end(), + // Never retried (the origin already rides out transient source failures), and not + // papered over with gaps: the rendition's playlist answers the error. + Err(err) => { + tracing::error!(%track, %err, "media track failed; failing the rendition's playlist"); + rendition.fail(crate::Error::MediaFailed { + track, + source: Box::new(err), + }); } } + rendition.close(); + Poll::Ready(()) } - .await; - let Some(rendition) = rendition.upgrade() else { - return; - }; - match result { - Ok(()) => rendition.end(), - // A bound sibling that ended takes its track down with it: that ends the rendition, whose - // retained rows drain, rather than failing it. - Err(_) if !rendition.media.admits(&rendition.live) => rendition.end(), - // Never retried (the origin already rides out transient source failures), and not - // papered over with gaps: the rendition's playlist answers the error. - Err(err) => { - tracing::error!(%track, %err, "media track failed; failing the rendition's playlist"); - rendition.fail(crate::Error::MediaFailed { - track, - source: Box::new(err), - }); - } - } - rendition.close(); } + /// What a group's frames reach: the timestamps of its samples, where one ends explicitly, and a /// discontinuity marker, read as they are decoded. #[derive(Default)] @@ -2031,27 +2501,6 @@ fn first_frame( Some(First::Media(pts, reach, size)) } -/// What a group's content reaches, read once the track has ended so every frame is in, starting -/// from what its first frame reached (see [`Reach::finish`]). -/// -/// Fails when the group does: its last segment would otherwise be listed complete but truncated. -async fn content_end( - mut group: moq_net::group::Consumer, - container: &moq_mux::catalog::hang::Container, - mut reach: Reach, - spacing: Option, - interval: std::time::Duration, -) -> Result<(Option, Option)> { - use moq_mux::container::Container as _; - while let Some(frame) = group.read_frame().await? { - // The container says what a frame holds: a marker's payload is empty only once decoded. - if let Ok(decoded) = container.decode(frame) { - reach.add(decoded); - } - } - Ok(reach.finish(spacing, interval)) -} - /// The row listing `slot`, starting a new discontinuity when `broken`. fn row(rendition: &Rendition, slot: grid::Slot, broken: bool) -> segments::Row { let (segment, groups, pts, end, content, stop) = match slot { @@ -2115,62 +2564,103 @@ struct InitRead { muxer: Muxer, /// Groups to try in order (see [`segments::Producer::init_groups`]): a complete group first, /// since reading an open one waits for the publisher to finish it. - groups: Vec, + groups: std::vec::IntoIter, + track: Option<(moq_net::broadcast::Consumer, moq_net::track::Consumer)>, + /// The group being read. + group: Option, } impl InitRead { + fn new(muxer: Muxer, groups: Vec) -> Self { + Self { + muxer, + groups: groups.into_iter(), + track: None, + group: None, + } + } + /// Read the groups until one yields the init. - async fn read(mut self, media: &Media, binding: &moq_mux::Binding) -> Result> { - let Some((broadcast, track)) = media.track(binding).await? else { - return Ok(None); + fn poll(&mut self, waiter: &kio::Waiter, media: &Media, binding: &moq_mux::Binding) -> Poll>> { + let (broadcast, track) = match &self.track { + Some(track) => track, + None => match ready!(media.poll_track_or_refusal(binding, waiter))? { + Some(track) => self.track.insert(track), + None => return Poll::Ready(Ok(None)), + }, }; - for sequence in self.groups { - // A group gone from the cache, before or mid-read, leaves the next candidate to try. - let Ok(mut group) = fetch(&track, sequence, &broadcast).await? else { - continue; + loop { + let read = match &mut self.group { + Some(read) => read, + None => match self.groups.next() { + Some(sequence) => self.group.insert(GroupRead::new(track, sequence)), + None => return Poll::Ready(Ok(None)), + }, }; - if read_group(&mut self.muxer, &mut group, &broadcast).await?.is_err() { + let read = ready!(read.poll(waiter, &mut self.muxer, broadcast))?; + self.group = None; + // A group gone from the cache, before or mid-read, leaves the next candidate to try. + if read.is_err() { continue; } if let Some(bytes) = self.muxer.init()? { - return Ok(Some(bytes)); + return Poll::Ready(Ok(Some(bytes))); } } - Ok(None) } } -/// Fetch one group, mapping "no longer (or not yet) servable" to the error that says so -/// rather than a bare `None`: a caller abandoning a segment can only explain itself with it. -async fn fetch( - track: &moq_net::track::Consumer, +/// Fetches one group, then reads its frames through a muxer. +struct GroupRead { sequence: u64, - broadcast: &moq_net::broadcast::Consumer, -) -> Result> { - match track.fetch_group(sequence, None).await { - Ok(group) => Ok(Ok(group)), - Err(err) if is_unservable(&err, broadcast) => Ok(Err(err)), - Err(err) => Err(err.into()), - } + state: GroupReading, +} + +enum GroupReading { + Fetching(moq_net::track::Fetching), + Reading(moq_net::group::Consumer, Vec), } -/// Read a fetched group's frames, mapping a mid-read cache eviction (the group aged out of the -/// relay while we were reading it) to `None` rather than an error. Without this the eviction would -/// surface as a 500; it's really the same "group gone" 404 as a fetch-time miss. -async fn read_group( - muxer: &mut Muxer, - group: &mut moq_net::group::Consumer, - broadcast: &moq_net::broadcast::Consumer, -) -> Result, moq_net::Error>> { - match muxer.read(group).await { - Ok(frames) => Ok(Ok(frames)), - Err(err) => { - let miss = net_errors(&err).find(|err| is_unservable(err, broadcast)).cloned(); - match miss { - Some(miss) => Ok(Err(miss)), - None => Err(err.into()), +impl GroupRead { + fn new(track: &moq_net::track::Consumer, sequence: u64) -> Self { + Self { + sequence, + state: GroupReading::Fetching(track.fetch_group(sequence, None).into_inner()), + } + } + + /// The group's frames, or the error saying it is no longer (or not yet) servable rather than a + /// bare `None`: a caller abandoning a segment can only explain itself with it. + /// + /// A mid-read cache eviction (the group aged out of the relay while we were reading it) is + /// such a miss too. Without this the eviction would surface as a 500; it's really the same + /// "group gone" 404 as a fetch-time miss. + fn poll( + &mut self, + waiter: &kio::Waiter, + muxer: &mut Muxer, + broadcast: &moq_net::broadcast::Consumer, + ) -> Poll, moq_net::Error>>> { + if let GroupReading::Fetching(fetching) = &mut self.state { + match ready!(kio::Task::poll(fetching, waiter)) { + Ok(group) => self.state = GroupReading::Reading(group, Vec::new()), + Err(err) if is_unservable(&err, broadcast) => return Poll::Ready(Ok(Err(err))), + Err(err) => return Poll::Ready(Err(err.into())), } } + let GroupReading::Reading(group, frames) = &mut self.state else { + unreachable!("fetched above"); + }; + Poll::Ready(match ready!(muxer.poll_read(waiter, group, frames)) { + Ok(frames) => Ok(Ok(frames)), + Err(err) => { + let miss = net_errors(&err).find(|err| is_unservable(err, broadcast)).cloned(); + match miss { + Some(miss) => Ok(Err(miss)), + None => Err(err.into()), + } + } + }) } } diff --git a/rs/moq-hls/src/export/renditions.rs b/rs/moq-hls/src/export/renditions.rs index 84a293945cf..1bb35005031 100644 --- a/rs/moq-hls/src/export/renditions.rs +++ b/rs/moq-hls/src/export/renditions.rs @@ -135,22 +135,70 @@ struct State { /// Where a rendition's media track is: its broadcast, instance, and track name. type Held = (Option, Option, String); -/// Hold `track` on `binding` subscribed until dropped, as demand for its media. +/// Holds `track` on `binding` subscribed until dropped, as demand for its media. /// /// A publisher that measures its join interval does so from media, and an on-demand one produces /// media only for a subscriber: without this, an export that subscribes only once the interval /// is declared would never see one. -async fn hold(binding: moq_mux::Binding, track: String) { - let Ok(broadcast) = binding.broadcast().await else { - return; - }; - let Ok(track) = broadcast.track(&track) else { - return; - }; - let Ok(_subscriber) = track.subscribe(moq_net::track::Subscription::default()).await else { - return; - }; - std::future::pending::<()>().await +struct Hold { + binding: moq_mux::Binding, + track: String, + state: Holding, +} + +#[allow( + clippy::large_enum_variant, + reason = "one per track a catalog lists before its join interval, held in place until dropped" +)] +enum Holding { + Binding, + Subscribing { + _broadcast: moq_net::broadcast::Consumer, + _track: moq_net::track::Consumer, + subscribing: moq_net::track::Subscribing, + }, + Held { + _broadcast: moq_net::broadcast::Consumer, + _track: moq_net::track::Consumer, + _subscriber: moq_net::track::Subscriber, + }, +} + +impl kio::Task for Hold { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + if let Holding::Binding = self.state { + let Ok(broadcast) = std::task::ready!(self.binding.poll_broadcast(waiter)) else { + return Poll::Ready(()); + }; + let Ok(track) = broadcast.track(&self.track) else { + return Poll::Ready(()); + }; + let subscribing = track.subscribe(moq_net::track::Subscription::default()).into_inner(); + self.state = Holding::Subscribing { + _broadcast: broadcast, + _track: track, + subscribing, + }; + } + if let Holding::Subscribing { subscribing, .. } = &self.state { + let Ok(subscriber) = std::task::ready!(subscribing.poll_ok(waiter)) else { + return Poll::Ready(()); + }; + let Holding::Subscribing { _broadcast, _track, .. } = std::mem::replace(&mut self.state, Holding::Binding) + else { + unreachable!("subscribing above"); + }; + self.state = Holding::Held { + _broadcast, + _track, + _subscriber: subscriber, + }; + } + // Held until the owner drops it. + Poll::Pending + } } /// The shared state handle each rendition's record watcher reaches. @@ -362,7 +410,7 @@ impl Producer { impl Fanout { /// Run `task` on the broadcaster's driver until it finishes or the returned [`Owner`] drops. - pub fn spawn(&self, task: impl Future + Send + 'static) -> Owner { + pub fn spawn(&self, task: impl kio::Task + Send + 'static) -> Owner { self.spawner.spawn(task) } @@ -449,11 +497,10 @@ impl Fanout { } impl Producer { - /// Resolve once at least one rendition has a playable media playlist. Bounding how long to - /// wait is the caller's policy, so wrap this in a timeout rather than passing one in. - pub async fn ready(&self) { - let _ = kio::wait(|waiter| { - self.state.poll_ref(waiter, |current| { + /// Ready once at least one rendition has a playable media playlist, or the set closed. + pub fn poll_ready(&self, waiter: &kio::Waiter) -> Poll<()> { + self.state + .poll_ref(waiter, |current| { if current .renditions .values() @@ -464,8 +511,7 @@ impl Producer { Poll::Pending } }) - }) - .await; + .map(|_| ()) } /// Resolve once at least one enabled rendition is [advertised](Rendition::is_advertised), so a @@ -488,21 +534,21 @@ impl Producer { .map(|_| ()) } + /// Ready once a catalog was read, or the set closed. #[cfg(feature = "server")] - pub(crate) async fn catalogued(&self) { - let _ = kio::wait(|waiter| { - self.state.poll_ref(waiter, |current| { + pub(crate) fn poll_catalogued(&self, waiter: &kio::Waiter) -> Poll<()> { + self.state + .poll_ref(waiter, |current| { if current.media.is_some() { Poll::Ready(()) } else { Poll::Pending } }) - }) - .await; + .map(|_| ()) } - /// Start every enabled video rendition's init build, then resolve once one of them has its + /// Start every enabled video rendition's init build, then be ready once one of them has its /// init and lists a segment of its own, so a manifest lists it, or none is left /// that listings do not refuse: a denied or failed one can never be listed. A build that has /// nothing to read from yet starts on the wake that indexes its media. @@ -510,8 +556,9 @@ impl Producer { not(feature = "server"), allow(dead_code, reason = "only the `server` feature uses it") )] - pub(crate) async fn video_init(&self) { - kio::wait(|waiter| { + #[cfg_attr(test, kio::wrapper)] + pub(crate) fn poll_video_init(&self, waiter: &kio::Waiter) -> Poll<()> { + { let poll = self.state.poll_ref(waiter, |current| { let mut video = current .renditions @@ -540,8 +587,7 @@ impl Producer { // and the predicate registered `waiter` on each of them. Poll::Ready(Err(_)) | Poll::Pending => Poll::Pending, } - }) - .await + } } /// Enroll freshly created renditions with what their siblings share, and start following @@ -655,7 +701,11 @@ impl Producer { let Ok(binding) = source.bind(broadcast.as_ref(), instance.as_ref()) else { continue; }; - let owner = self.fanout.spawn(hold(binding, track.clone())); + let owner = self.fanout.spawn(Hold { + binding, + track: track.clone(), + state: Holding::Binding, + }); waiting.insert(key, (held, owner)); } } @@ -792,9 +842,10 @@ pub struct Consumer { impl Consumer { /// The next rendition change; `None` only after the catalog finishes cleanly. - pub async fn next(&mut self) -> crate::Result> { - let Some(event) = kio::wait(|waiter| self.poll_next(waiter)).await? else { - return Ok(None); + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { + let Some(event) = std::task::ready!(self.poll_event(waiter))? else { + return Poll::Ready(Ok(None)); }; match &event { Event::Added(rendition) => { @@ -805,10 +856,10 @@ impl Consumer { self.seen.remove(&(*kind, name.clone())); } } - Ok(Some(event)) + Poll::Ready(Ok(Some(event))) } - fn poll_next(&self, waiter: &kio::Waiter) -> Poll>> { + fn poll_event(&self, waiter: &kio::Waiter) -> Poll>> { // A failed catalog cannot wait for the change stream to close. if let Poll::Ready(Ok(err)) = self .feed diff --git a/rs/moq-hls/src/export/segments.rs b/rs/moq-hls/src/export/segments.rs index 1a426db9cc2..795b4e8f1e2 100644 --- a/rs/moq-hls/src/export/segments.rs +++ b/rs/moq-hls/src/export/segments.rs @@ -24,6 +24,7 @@ use std::time::Duration; use bytes::Bytes; use super::Rendition; +use super::rendition::{Fetch, LoadInit}; use crate::Result; /// A segment as its URL names it: `seg/[{generation}.]{tag}.{sequence}.m4s`, or `t{pts}` in @@ -405,9 +406,9 @@ impl Producer { } } - /// Resolve once the window is closed: its media watcher lists no more rows. - pub async fn closed(&self) { - self.state.closed().await; + /// Ready once the window is closed: its media watcher lists no more rows. + pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll<()> { + self.state.poll_closed(waiter) } /// Whether the window lists segment `at` of the record track tagged `tag`. @@ -442,6 +443,7 @@ impl Producer { rendition, position: (state.resets, state.produced - state.rows.len() as u64), feed: None, + fetch: None, } } } @@ -482,6 +484,8 @@ pub struct Consumer { /// The reset epoch and next append position; errors leave it unchanged. position: (u64, u64), feed: Option, + /// The next row's media being fetched, kept across polls. + fetch: Option, } /// The live subscription a [`Consumer`] holds on its rendition's media track. @@ -504,7 +508,8 @@ impl Consumer { /// The rendition's CMAF init segment, built once and cached; `None` until it can be built /// (an inline-parameter-set codec needs the first segment first). pub async fn init(&self) -> Result> { - self.rendition.latest_init().await + let mut load = LoadInit::default(); + kio::wait(|waiter| self.rendition.poll_latest_init(waiter, &mut load)).await } /// The next segment, with its media; `None` only once the rendition ends cleanly, so a @@ -524,27 +529,36 @@ impl Consumer { /// silently omitting segments. Fetch/transmux errors leave the cursor in place so the caller /// can retry. A bound instance's retained media drains after End or Restart; a replacement's /// media is never read under the old rows. - pub async fn next(&mut self) -> Result> { - let Some(row) = kio::wait(|waiter| self.poll_next(waiter)).await? else { - return Ok(None); - }; - let media = if row.stop { - None - } else { - let media = self.rendition.fetch(&row).await?; - Some(media.ok_or(crate::Error::SegmentUnavailable { segment: row.segment })?) + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { + let (row, media) = loop { + let Some(fetch) = &mut self.fetch else { + let Some(row) = ready!(self.poll_row(waiter))? else { + return Poll::Ready(Ok(None)); + }; + if row.stop { + break (row, None); + } + self.fetch = Some(Fetch::new(row)); + continue; + }; + let media = ready!(fetch.poll(&self.rendition, waiter)); + // A failure leaves the position for a retry to fetch the row afresh. + let row = self.fetch.take().expect("fetching above").row().clone(); + let media = media?.ok_or(crate::Error::SegmentUnavailable { segment: row.segment })?; + break (row, Some(media)); }; self.position.1 += 1; - Ok(Some(Segment { + Poll::Ready(Ok(Some(Segment { segment: row.segment, media, duration: row.duration, program_date_time: self.rendition.wall_clock(row.pts), discontinuity: row.discontinuity, - })) + }))) } - fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { + fn poll_row(&mut self, waiter: &kio::Waiter) -> Poll>> { self.poll_feed(waiter); // A refused binding never held media, so there is nothing to drain. `poll_feed` registered // the wake for a refusal that arrives while the window waits for rows. diff --git a/rs/moq-hls/src/lib.rs b/rs/moq-hls/src/lib.rs index f3091832e62..96189b3ed9b 100644 --- a/rs/moq-hls/src/lib.rs +++ b/rs/moq-hls/src/lib.rs @@ -20,11 +20,23 @@ //! //! Its cryptography is [rind](https://docs.rs/rind), so every constructor that first needs it //! refuses a CPU missing a feature rind requires with [`Error::UnsupportedCpu`]. +//! +//! The gateway drives its HTTP fetches and background watchers as async code, so [`import`] and +//! [`export`] need the `async` feature. The default `server` feature enables it. #![warn(missing_docs)] +#![cfg_attr( + not(feature = "async"), + allow( + rustdoc::broken_intra_doc_links, + reason = "the docs link the gateways, which a poll-only build leaves out" + ) +)] mod error; +#[cfg(feature = "async")] pub mod export; +#[cfg(feature = "async")] pub mod import; #[cfg(feature = "server")] pub mod server; diff --git a/rs/moq-hls/src/server/mod.rs b/rs/moq-hls/src/server/mod.rs index 8b64de537d4..98399973d82 100644 --- a/rs/moq-hls/src/server/mod.rs +++ b/rs/moq-hls/src/server/mod.rs @@ -265,13 +265,14 @@ impl Server { broadcasts.insert(name.to_string(), followed.clone()); } - let task = tokio::spawn(run_follow( - Arc::downgrade(&self.inner), - name.to_string(), - Arc::downgrade(&followed), + let mut follow = Follow { + inner: Arc::downgrade(&self.inner), + name: name.to_string(), + followed: Arc::downgrade(&followed), instances, - broadcast, - )); + served: Some(broadcast), + }; + let task = tokio::spawn(kio::wait(move |waiter| follow.poll(waiter))); *followed.task.lock().unwrap() = Some(task); Some(followed) } @@ -291,59 +292,68 @@ impl Inner { } } -/// Follow `name`'s instances into `followed` until nothing serves the path, then forget it, so the -/// next request follows the path afresh. Holds both weakly, so it never keeps the server's state -/// alive; dropping `followed` aborts it. `served` is the instance `followed` starts on, which the -/// follower hands out first unless another took over since, or folds away if it closed first. -async fn run_follow( +/// Follows `name`'s instances into `followed` until nothing serves the path, then forgets it, so +/// the next request follows the path afresh. Holds both weakly, so it never keeps the server's +/// state alive; dropping `followed` aborts it. `served` is the instance `followed` starts on, which +/// the follower hands out first unless another took over since, or folds away if it closed first. +struct Follow { inner: Weak, name: String, followed: Weak, - mut instances: moq_net::origin::Following, - served: moq_net::broadcast::Consumer, -) { - let mut served = Some(served); - let (inner, followed) = loop { - let next = kio::wait(|waiter| { - let Some(followed) = followed.upgrade() else { - return Poll::Ready(None); - }; - followed.forget_closed(waiter); - if let Poll::Ready(next) = instances.poll_available(waiter) { - return Poll::Ready(Some(next)); - } - // The follower folds away an instance that started and ended before its first poll, - // so the served instance closing before it is handed out means nothing serves the path. - match &served { - Some(served) => served.poll_closed(waiter).map(|()| Some(Ok(None))), - None => Poll::Pending, - } - }) - .await; - let (Some(inner), Some(followed)) = (inner.upgrade(), followed.upgrade()) else { - return; + instances: moq_net::origin::Following, + served: Option, +} + +impl Follow { + /// The next instance event, or `None` once `followed` is gone. + fn poll_next( + &mut self, + waiter: &kio::Waiter, + ) -> Poll, moq_net::Error>>> { + let Some(followed) = self.followed.upgrade() else { + return Poll::Ready(None); }; - match next { - Some(Ok(Some(broadcast))) => { - if !served.take().is_some_and(|served| served.is_clone(&broadcast)) { - followed.start(inner.serve(broadcast)); + followed.forget_closed(waiter); + if let Poll::Ready(next) = self.instances.poll_available(waiter) { + return Poll::Ready(Some(next)); + } + // The follower folds away an instance that started and ended before its first poll, + // so the served instance closing before it is handed out means nothing serves the path. + match &self.served { + Some(served) => served.poll_closed(waiter).map(|()| Some(Ok(None))), + None => Poll::Pending, + } + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + let (inner, followed) = loop { + let next = std::task::ready!(self.poll_next(waiter)); + let (Some(inner), Some(followed)) = (self.inner.upgrade(), self.followed.upgrade()) else { + return Poll::Ready(()); + }; + match next { + Some(Ok(Some(broadcast))) => { + if !self.served.take().is_some_and(|served| served.is_clone(&broadcast)) { + followed.start(inner.serve(broadcast)); + } } + Some(Ok(None)) => break (inner, followed), + Some(Err(err)) => { + tracing::warn!(%err, name = %self.name, "stopped following broadcast"); + break (inner, followed); + } + None => return Poll::Ready(()), } - Some(Ok(None)) => break (inner, followed), - Some(Err(err)) => { - tracing::warn!(%err, %name, "stopped following broadcast"); - break (inner, followed); - } - None => return, + }; + + let mut broadcasts = inner.broadcasts.lock().unwrap(); + if broadcasts + .get(&self.name) + .is_some_and(|current| Arc::ptr_eq(current, &followed)) + { + broadcasts.remove(&self.name); } - }; - - let mut broadcasts = inner.broadcasts.lock().unwrap(); - if broadcasts - .get(&name) - .is_some_and(|current| Arc::ptr_eq(current, &followed)) - { - broadcasts.remove(&name); + Poll::Ready(()) } } diff --git a/rs/moq-json/Cargo.toml b/rs/moq-json/Cargo.toml index e62e4bc8cb4..cb0a9ce1465 100644 --- a/rs/moq-json/Cargo.toml +++ b/rs/moq-json/Cargo.toml @@ -12,6 +12,14 @@ rust-version.workspace = true keywords = ["quic", "http3", "webtransport", "json", "live"] categories = ["multimedia", "network-programming", "web-programming"] +[package.metadata.docs.rs] +features = ["async"] + +[features] +default = [] +# Generated async wrappers for the poll methods. Enable where you `.await` them. +async = ["kio/async"] + [dependencies] bytes = { workspace = true } kio = { workspace = true } @@ -26,6 +34,9 @@ tracing = { workspace = true } criterion = { package = "bench-harness", path = "../bench-harness" } futures = { workspace = true } json-patch = "4" +# Compile the generated wrappers used by the tests. +moq-json = { path = ".", features = ["async"] } +moq-net = { path = "../moq-net", features = ["async"] } proptest = { workspace = true, features = ["handle-panics"] } workspace-hack = { path = "../workspace-hack" } diff --git a/rs/moq-json/MIGRATE.md b/rs/moq-json/MIGRATE.md index 9a3c7527ec2..6e96c0c3eed 100644 --- a/rs/moq-json/MIGRATE.md +++ b/rs/moq-json/MIGRATE.md @@ -4,6 +4,15 @@ How to update your code for each breaking change, newest first. ## Unreleased +### Async wrappers are opt-in + +Enable the `async` Cargo feature to keep awaiting consumer `next`/`latest` and producer +`refresh`: `moq-json = { version = "...", features = ["async"] }`. This also applies +with `default-features = false`. The feature is off by default; `poll_next`, +`poll_latest`, and `poll_refresh` stay available without it. +The generated consumer wrappers no longer require `T: Unpin`. +The JSON framing and JS and binding APIs are unchanged. + ### Snapshot consumers subscribe themselves at a zero delay ([#7736](https://github.com/moq-dev/moq/pull/7736)) `snapshot::Consumer::new` takes the track, a `moq_net::track::Consumer`, instead of a `track::Subscriber`, and subscribes to it itself at a zero `delay.max` and default preferences otherwise; the new `Config::priority` sets the subscription's priority. diff --git a/rs/moq-json/src/lib.rs b/rs/moq-json/src/lib.rs index 64dcbd4de26..587a64d613e 100644 --- a/rs/moq-json/src/lib.rs +++ b/rs/moq-json/src/lib.rs @@ -16,6 +16,9 @@ //! (and back), with the encoder saying where the group boundaries fall. Reach for the codec layer //! when something else already owns the track, such as a `moq_mux::container::Producer` also //! managing a timeline and a catalog estimate. +//! +//! The `poll_*` methods are always available. Enable the `async` feature for their generated +//! async wrappers (`next`, `latest`, and `refresh`); it is off by default. #![warn(missing_docs)] diff --git a/rs/moq-json/src/snapshot/consumer.rs b/rs/moq-json/src/snapshot/consumer.rs index 909ba7e9408..bcea338bfbf 100644 --- a/rs/moq-json/src/snapshot/consumer.rs +++ b/rs/moq-json/src/snapshot/consumer.rs @@ -73,33 +73,19 @@ impl Consumer { } /// Get the next state and its exact frame timestamp, or `None` once the track ends. - pub async fn next(&mut self) -> Result>> - where - T: Unpin, - { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for each state in order, skipping every group a newer one has superseded. /// - /// Finish the current group's buffered frames before taking the next. Lost groups - /// resynchronize from a later snapshot. + /// Skips every group a newer one has superseded. Finish the current group's buffered frames + /// before taking the next. Lost groups resynchronize from a later snapshot. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>>> { self.poll_read(waiter, false) } - /// Skip to the newest buffered state and its timestamp, or wait for a new one. - pub async fn latest(&mut self) -> Result>> - where - T: Unpin, - { - kio::wait(|waiter| self.poll_latest(waiter)).await - } - - /// Poll for the newest state, discarding superseded groups and intermediate yields. + /// Get the newest state and its timestamp, discarding superseded groups and intermediate yields. /// /// Frames inside the newest group are still decoded in order. What this skips is gone; /// a subsequent `next` resumes after it. An untimed frame keeps `at: None`. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_latest(&mut self, waiter: &kio::Waiter) -> Poll>>> { self.poll_read(waiter, true) } diff --git a/rs/moq-json/src/snapshot/mod.rs b/rs/moq-json/src/snapshot/mod.rs index b0943dbebb0..9a6f522378a 100644 --- a/rs/moq-json/src/snapshot/mod.rs +++ b/rs/moq-json/src/snapshot/mod.rs @@ -20,7 +20,7 @@ //! [`refresh`](moq_net::track::Info::refresh) ratio asks for one when a subscriber joins while the //! open group's deltas outweigh that ratio times its snapshot. The joiner still reads the cached //! group meanwhile: the ratio limits how often groups roll, not what a joiner downloads. -//! [`Producer::refresh`] answers by restating the value as a snapshot in +//! [`Producer::poll_refresh`] answers by restating the value as a snapshot in //! a new group; [`Encoder::snapshot`] does the same for a caller that owns the track. Without a ratio, //! or behind a relay that cannot forward the ask, a group runs to the encoder's frame and byte caps. //! diff --git a/rs/moq-json/src/snapshot/producer.rs b/rs/moq-json/src/snapshot/producer.rs index 66ebd029bdc..eef6c72fbc2 100644 --- a/rs/moq-json/src/snapshot/producer.rs +++ b/rs/moq-json/src/snapshot/producer.rs @@ -30,7 +30,7 @@ fn take(inner: &Mutex>) -> MutexGuard<'_, Inner> { /// encoder emits a snapshot. When something else already owns the track, use the [`Encoder`] /// directly. /// -/// A group runs until a joiner wants a fresh one, which [`refresh`](Self::refresh) answers, or until +/// A group runs until a joiner wants a fresh one, which [`poll_refresh`](Self::poll_refresh) answers, or until /// it reaches the encoder's caps. Set the track's [`refresh`](moq_net::track::Info::refresh) ratio /// when creating it, or joins never ask. /// @@ -97,11 +97,7 @@ impl Producer { /// /// The producer has no runtime of its own, so the owner drives this: loop on it, or select on it /// beside the updates. Errors once the track is closed. - pub async fn refresh(&self) -> Result<()> { - kio::wait(|waiter| self.poll_refresh(waiter)).await - } - - /// Poll for [`refresh`](Self::refresh), for an owner driving the producer synchronously. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_refresh(&self, waiter: &kio::Waiter) -> std::task::Poll> { take(&self.inner).poll_refresh(waiter) } @@ -192,7 +188,7 @@ impl Producer { /// Finish the open group, so the deltas already written stop being provisional. /// - /// No replacement group opens until the next [`update`](Self::update) or [`refresh`](Self::refresh), which emits a full + /// No replacement group opens until the next [`update`](Self::update) or [`poll_refresh`](Self::poll_refresh), which emits a full /// snapshot as its first frame even when the value is unchanged. A consumer joining at that group /// therefore reads the whole value without the deltas that preceded it. /// diff --git a/rs/moq-json/src/stream/consumer.rs b/rs/moq-json/src/stream/consumer.rs index 4fac76ae20d..ec9976be16e 100644 --- a/rs/moq-json/src/stream/consumer.rs +++ b/rs/moq-json/src/stream/consumer.rs @@ -48,14 +48,7 @@ impl Consumer { } /// Get the next record and its exact frame timestamp, or `None` once the track ends. - pub async fn next(&mut self) -> Result>> - where - T: Unpin, - { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next record, without blocking. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>>> { loop { if self.rolled { diff --git a/rs/moq-json/src/window/consumer.rs b/rs/moq-json/src/window/consumer.rs index e5283158def..3104a7c65b8 100644 --- a/rs/moq-json/src/window/consumer.rs +++ b/rs/moq-json/src/window/consumer.rs @@ -47,14 +47,7 @@ impl Consumer { } /// Get the next event and its frame's timestamp, or `None` once the track ends. - pub async fn next(&mut self) -> Result>>> - where - T: Unpin, - { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next event, without blocking. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>>>> { loop { // Drain what the last frame produced before reading more, so every queued event is that diff --git a/rs/moq-json/src/window/mod.rs b/rs/moq-json/src/window/mod.rs index 71a4b0897c3..4aed9546cfe 100644 --- a/rs/moq-json/src/window/mod.rs +++ b/rs/moq-json/src/window/mod.rs @@ -28,10 +28,10 @@ //! The publisher rolls a group when a joiner wants a fresh one, exactly as //! [`snapshot`](crate::snapshot) does: on a track created with a //! [`refresh`](moq_net::track::Info::refresh) ratio, a subscriber joining once the open group's ops -//! outweigh that ratio times its header asks for one, and [`Producer::refresh`] restates the window -//! in it. Otherwise a group runs to the frame and byte caps. A [`Consumer`] never surfaces a roll: -//! a header restating records a reader already has yields nothing, so however often the publisher -//! rolls, the reader sees one continuous stream of [`Event`]s. +//! outweigh that ratio times its header asks for one, and [`Producer::poll_refresh`] restates the +//! window in it. Otherwise a group runs to the frame and byte caps. A [`Consumer`] never surfaces a +//! roll: a header restating records a reader already has yields nothing, so however often the +//! publisher rolls, the reader sees one continuous stream of [`Event`]s. //! [`ProducerConfig::checkpoint_records`] bounds the suffix repeated on each roll for a long-lived //! window. //! diff --git a/rs/moq-json/src/window/producer.rs b/rs/moq-json/src/window/producer.rs index f155afd4286..10fdd074267 100644 --- a/rs/moq-json/src/window/producer.rs +++ b/rs/moq-json/src/window/producer.rs @@ -16,7 +16,7 @@ use crate::refresh::Refresh; /// An [`Encoder`] that owns its track: it writes each encoded frame and rolls a group whenever the /// encoder emits a header. When something else already owns the track, use the [`Encoder`] directly. /// -/// A group runs until a joiner wants a fresh one, which [`refresh`](Self::refresh) answers, or until +/// A group runs until a joiner wants a fresh one, which [`poll_refresh`](Self::poll_refresh) answers, or until /// it reaches the encoder's caps. Set the track's [`refresh`](moq_net::track::Info::refresh) ratio /// when creating it, or joins never ask. /// @@ -97,11 +97,7 @@ impl Producer { /// /// The producer has no runtime of its own, so the owner drives this: loop on it, or select on it /// beside the edits. Errors once the track is finished or closed. - pub async fn refresh(&self) -> Result<()> { - kio::wait(|waiter| self.poll_refresh(waiter)).await - } - - /// Poll for [`refresh`](Self::refresh), for an owner driving the producer synchronously. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_refresh(&self, waiter: &kio::Waiter) -> std::task::Poll> { self.inner.lock().unwrap().poll_refresh(waiter) } diff --git a/rs/moq-mux/Cargo.toml b/rs/moq-mux/Cargo.toml index a8966e275ec..c76e1758c23 100644 --- a/rs/moq-mux/Cargo.toml +++ b/rs/moq-mux/Cargo.toml @@ -12,7 +12,14 @@ rust-version.workspace = true keywords = ["quic", "http3", "webtransport", "media", "live"] categories = ["multimedia", "network-programming", "web-programming"] +[package.metadata.docs.rs] +features = ["async"] + [features] +default = [] +# Generated async wrappers for the poll methods, and the async constructors and helpers with no +# poll of their own. Enable where you `.await` them. +async = ["kio/async", "moq-net/async"] # Exposes the media parsers' fuzz target bodies and seeds to the `fuzz/` harness through the # hidden `fuzz` module. Off by default and hidden from the docs. fuzz = [] @@ -51,7 +58,11 @@ webm-iterable = "0.7" [dev-dependencies] criterion = { package = "bench-harness", path = "../bench-harness" } futures = { workspace = true } +# The tests await hang's catalog subscription. +hang = { path = "../hang", features = ["async"] } kio = { workspace = true, features = ["sim"] } +# Compile the async API used by the tests, benches, and examples. +moq-mux = { path = ".", features = ["async"] } proptest = { workspace = true } tokio = { workspace = true, features = ["test-util"] } tracing-test = { workspace = true } diff --git a/rs/moq-mux/MIGRATE.md b/rs/moq-mux/MIGRATE.md index 2817ab45268..383df17a179 100644 --- a/rs/moq-mux/MIGRATE.md +++ b/rs/moq-mux/MIGRATE.md @@ -4,6 +4,24 @@ How to adopt new APIs and update for breaking changes, newest first. ## Unreleased +### The async API is opt-in + +Enable the `async` Cargo feature to keep awaiting moq-mux: +`moq-mux = { version = "...", features = ["async"] }`. This also applies with `default-features = false`. +The feature is off by default, and every `poll_*` method stays available without it. +It covers the `next`, `read`, `latest`, `refresh`, `ready`, `wait`, and `broadcast` methods beside each poll: +the catalog consumers and `catalog::Stream::next`, the container `Consumer` and `GroupConsumer`, +every exporter (`next`, and fMP4's `next_chunk`), `ts::Follower`, `ts::Release`, `follow::Instances`, +`json::Consumer`, `json::Snapshot`, `rate::Follower`, `timeline::Resume`, `Binding`, and `Linger`. +It also covers the async functions with no poll of their own: `catalog::Consumer::new`, +`Source::{catalog, resolve, subscribe_track}`, `ts::Release::new`, and the JSON catalog entry's `subscribe`. +`fmp4::Muxer::read` is async too; without the feature, drive `fmp4::Muxer::poll_read`, passing the same +empty `Vec` to every poll of one group. + +`timeline::Resume::ready` borrows the resume instead of consuming it, and the hang catalog consumer's +`next` no longer requires `Catalog: Unpin`. +moq-mux no longer enables `moq-json/async`: enable it yourself to await moq-json's readers directly. + ### `follow::Instances` folds into `Following` ([#7678](https://github.com/moq-dev/moq/pull/7678)) `follow::Instances::new` is removed: call `origin.follow(path, instance)` and read `Following::available()`, which returns `Some(broadcast)` where `Event::Start` did and `None` where `Event::End` did. diff --git a/rs/moq-mux/src/catalog/consumer.rs b/rs/moq-mux/src/catalog/consumer.rs index e31c4288384..94bc8578a42 100644 --- a/rs/moq-mux/src/catalog/consumer.rs +++ b/rs/moq-mux/src/catalog/consumer.rs @@ -4,7 +4,7 @@ //! advertises and yields [`Catalog`](super::hang::Catalog) snapshots so callers //! and exporters only deal with one shape. -use std::task::Poll; +use std::task::{Poll, ready}; use super::hang::{Catalog, CatalogExt}; use super::{CatalogFormat, Stream}; @@ -48,23 +48,59 @@ impl Consumer { } /// Subscribe to the catalog track advertised by `format`. + #[cfg(feature = "async")] pub async fn new(broadcast: &moq_net::broadcast::Consumer, format: CatalogFormat) -> Result { - Ok(Self { - inner: match format { - CatalogFormat::Hang => { - let track = broadcast.track(hang::Catalog::DEFAULT_NAME)?; - let consumer = super::hang::Consumer::new(track.clone()); - // The consumer subscribes itself; resolving the info rides on that subscription, - // so a refused catalog fails here as it does for MSF. - track.info().await?; - State::Hang(consumer) - } - CatalogFormat::Msf => { - let track = broadcast.track(moq_msf::DEFAULT_NAME)?.subscribe(None).await?; - State::Msf(super::msf::Consumer::new(track)) - } + let mut subscribing = Subscribing::new(broadcast, format); + kio::wait(|waiter| subscribing.poll_ok(waiter)).await + } +} + +/// A catalog subscription awaiting its track: [`Consumer::new`] for a caller that polls. +pub(crate) struct Subscribing { + pending: Pending, +} + +enum Pending { + /// The hang consumer subscribes itself; its track's info resolving rides on that + /// subscription, so a refused catalog fails here as it does for MSF. + Hang { + consumer: Option>>, + info: moq_net::track::Resolving, + }, + Msf(moq_net::track::Subscribing), + Failed(crate::Error), +} + +impl Subscribing { + /// Subscribe to the catalog track advertised by `format`. + pub(crate) fn new(broadcast: &moq_net::broadcast::Consumer, format: CatalogFormat) -> Self { + let pending = match format { + CatalogFormat::Hang => match broadcast.track(hang::Catalog::DEFAULT_NAME) { + Ok(track) => Pending::Hang { + consumer: Some(Box::new(super::hang::Consumer::new(track.clone()))), + info: track.info().into_inner(), + }, + Err(err) => Pending::Failed(err.into()), + }, + CatalogFormat::Msf => match broadcast.track(moq_msf::DEFAULT_NAME) { + Ok(track) => Pending::Msf(track.subscribe(None).into_inner()), + Err(err) => Pending::Failed(err.into()), }, - }) + }; + Self { pending } + } + + /// The catalog consumer, once the track subscription is confirmed. + pub(crate) fn poll_ok(&mut self, waiter: &kio::Waiter) -> Poll, crate::Error>> { + let inner = match &mut self.pending { + Pending::Hang { consumer, info } => { + ready!(info.poll_ok(waiter))?; + State::Hang(*consumer.take().expect("resolves once")) + } + Pending::Msf(subscribing) => State::Msf(super::msf::Consumer::new(ready!(subscribing.poll_ok(waiter))?)), + Pending::Failed(err) => return Poll::Ready(Err(err.clone())), + }; + Poll::Ready(Ok(Consumer { inner })) } } diff --git a/rs/moq-mux/src/catalog/hang/consumer.rs b/rs/moq-mux/src/catalog/hang/consumer.rs index b1a5233c109..13efa1bec8e 100644 --- a/rs/moq-mux/src/catalog/hang/consumer.rs +++ b/rs/moq-mux/src/catalog/hang/consumer.rs @@ -29,12 +29,13 @@ impl Consumer { } } - /// Poll for the next catalog update. + /// Get the next catalog update, or `None` once the track ends. /// /// Fails with [`Error::EscapingBroadcast`](crate::Error::EscapingBroadcast) if any rendition's /// `broadcast` reference or echo prefix walks above the root, rejecting the whole catalog /// (see [`hang::catalog::Consumer::poll_next`]). A transport failure is /// [`Error::Moq`](crate::Error::Moq). + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>>> { self.inner.poll_next(waiter).map_err(|err| match err { // Lifted to this crate's variants, which `Source` reports for the same faults. @@ -43,17 +44,6 @@ impl Consumer { err => crate::Error::Hang(err), }) } - - /// Get the next catalog update. - /// - /// This method waits for the next catalog publication and returns the - /// catalog data. If there are no more updates, `None` is returned. - pub async fn next(&mut self) -> Result>> - where - Catalog: Unpin, - { - kio::wait(|waiter| self.poll_next(waiter)).await - } } impl From for Consumer { diff --git a/rs/moq-mux/src/catalog/mod.rs b/rs/moq-mux/src/catalog/mod.rs index 6b3673ce9cb..96f55ab8261 100644 --- a/rs/moq-mux/src/catalog/mod.rs +++ b/rs/moq-mux/src/catalog/mod.rs @@ -38,6 +38,7 @@ pub(crate) mod tracks; pub(crate) use claim::Claim; pub use consumer::Consumer; +pub(crate) use consumer::Subscribing; pub use data::IntoRendition; pub(crate) use data::Listing; pub use entry::Entry; diff --git a/rs/moq-mux/src/catalog/msf/consumer.rs b/rs/moq-mux/src/catalog/msf/consumer.rs index 6092a2e570a..7f786fbdea1 100644 --- a/rs/moq-mux/src/catalog/msf/consumer.rs +++ b/rs/moq-mux/src/catalog/msf/consumer.rs @@ -34,7 +34,9 @@ impl Consumer { } } - /// Poll for the next catalog update. + /// Get the next catalog update, converted to a [`Catalog`](crate::catalog::hang::Catalog), + /// or `None` once the track ends with no further updates. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>>> { // Drain pending groups, keeping only the newest. Remember whether the track is done // so we can distinguish "more groups may arrive" from "no more groups, ever". @@ -72,15 +74,6 @@ impl Consumer { Poll::Pending } } - - /// Get the next catalog update. - /// - /// Waits for the next MSF catalog publication and returns it converted to a - /// [`Catalog`](crate::catalog::hang::Catalog). Returns `None` when the track has ended - /// with no further updates. - pub async fn next(&mut self) -> Result>> { - kio::wait(|waiter| self.poll_next(waiter)).await - } } impl From for Consumer { diff --git a/rs/moq-mux/src/catalog/producer.rs b/rs/moq-mux/src/catalog/producer.rs index 8a289303fd0..4f25c3fedaa 100644 --- a/rs/moq-mux/src/catalog/producer.rs +++ b/rs/moq-mux/src/catalog/producer.rs @@ -1001,7 +1001,7 @@ impl Producer { /// configures the encoder and is not written into the catalog entry. With /// [`deltas`](crate::json::Options::deltas), create `track` with a /// [`refresh`](moq_net::track::Info::refresh) ratio and drive - /// [`Snapshot::refresh`](crate::json::Snapshot::refresh), so joiners get a fresh group. + /// [`Snapshot::poll_refresh`](crate::json::Snapshot::poll_refresh), so joiners get a fresh group. /// /// Errors if the entry's section already carries that name, for example an entry seeded /// through [`Config::with_catalog`] or one pointing at a sibling broadcast, if the entry diff --git a/rs/moq-mux/src/catalog/stream.rs b/rs/moq-mux/src/catalog/stream.rs index e070ff28ac0..826b8340825 100644 --- a/rs/moq-mux/src/catalog/stream.rs +++ b/rs/moq-mux/src/catalog/stream.rs @@ -30,6 +30,7 @@ pub trait Stream: Send + 'static { fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>>>; /// Wait for the next snapshot. + #[cfg(feature = "async")] fn next(&mut self) -> impl std::future::Future>>> + Send where Self: Sized, diff --git a/rs/moq-mux/src/codec/h264/export.rs b/rs/moq-mux/src/codec/h264/export.rs index 7beebe83b98..e39d43bad10 100644 --- a/rs/moq-mux/src/codec/h264/export.rs +++ b/rs/moq-mux/src/codec/h264/export.rs @@ -45,12 +45,8 @@ impl Export { self } - /// Wait for the next Annex B access unit, or `None` once the export ends. - pub async fn next(&mut self) -> crate::Result> { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next Annex B access unit, or `None` once the export ends. + /// Get the next Annex B access unit, or `None` once the export ends. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { self.inner.poll_next(waiter) } diff --git a/rs/moq-mux/src/codec/h265/export.rs b/rs/moq-mux/src/codec/h265/export.rs index bd2e4bc1438..b4278ef48e3 100644 --- a/rs/moq-mux/src/codec/h265/export.rs +++ b/rs/moq-mux/src/codec/h265/export.rs @@ -40,12 +40,8 @@ impl Export { self } - /// Wait for the next Annex B access unit, or `None` once the export ends. - pub async fn next(&mut self) -> crate::Result> { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next Annex B access unit, or `None` once the export ends. + /// Get the next Annex B access unit, or `None` once the export ends. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { self.inner.poll_next(waiter) } diff --git a/rs/moq-mux/src/container/consumer.rs b/rs/moq-mux/src/container/consumer.rs index bcdc7893b04..fba63982554 100644 --- a/rs/moq-mux/src/container/consumer.rs +++ b/rs/moq-mux/src/container/consumer.rs @@ -55,7 +55,7 @@ pub(crate) enum Item { /// `Consumer` wraps a [`moq_net::track::Subscriber`] and a [`Container`] /// format implementation, typically /// [`catalog::hang::Container`](crate::catalog::hang::Container). Yields -/// decoded frames via [`read`](Self::read). +/// decoded frames via [`poll_read`](Self::poll_read). /// /// ## Ordering & age skipping /// @@ -98,7 +98,7 @@ pub(crate) enum Item { /// /// ## Gaps and ends /// -/// [`read`](Self::read) yields [`Read::Gap`] wherever media after an item already read +/// [`poll_read`](Self::poll_read) yields [`Read::Gap`] wherever media after an item already read /// will never arrive: a latency skip, an evicted current group, or a sequence hole whose /// boundary is not contiguous within 1 ms. Whatever is lost before the first item is no /// gap, since nothing was read before it. A marker group (no decodable frames: one empty @@ -420,8 +420,18 @@ impl Consumer { /// /// Returns `None` when the track has ended, and after an error. See [`Read`] for the items between frames; /// an undeclared forward timestamp jump is not a gap, since nothing says media is missing. - pub async fn read(&mut self) -> Result, F::Error> { - kio::wait(|waiter| self.poll_read(waiter)).await + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_read(&mut self, waiter: &kio::Waiter) -> Poll, F::Error>> { + loop { + let read = match ready!(self.poll_item(waiter))? { + Some(Item::Frame { group, frame, .. }) => Read::Frame { group, frame }, + Some(Item::Gap { timestamp, duration }) => Read::Gap { timestamp, duration }, + Some(Item::End { group, end }) => Read::End { group, end }, + Some(Item::FrameEnd(_) | Item::GroupEnd | Item::Cut { .. }) => continue, + None => return Poll::Ready(Ok(None)), + }; + return Poll::Ready(Ok(Some(read))); + } } /// The next frame, passing over gaps and ends. @@ -436,23 +446,6 @@ impl Consumer { } } - /// Poll-based implementation of the read loop. - /// - /// Uses a single waiter that gets registered on all relevant kio channels, - /// avoiding the need for `tokio::select!` or `FuturesUnordered`. - pub fn poll_read(&mut self, waiter: &kio::Waiter) -> Poll, F::Error>> { - loop { - let read = match ready!(self.poll_item(waiter))? { - Some(Item::Frame { group, frame, .. }) => Read::Frame { group, frame }, - Some(Item::Gap { timestamp, duration }) => Read::Gap { timestamp, duration }, - Some(Item::End { group, end }) => Read::End { group, end }, - Some(Item::FrameEnd(_) | Item::GroupEnd | Item::Cut { .. }) => continue, - None => return Poll::Ready(Ok(None)), - }; - return Poll::Ready(Ok(Some(read))); - } - } - /// Read items while preserving boundaries needed to finish exporter fragments. pub(crate) fn poll_item(&mut self, waiter: &kio::Waiter) -> Poll, F::Error>> { if let Some(item) = self.next.take() { diff --git a/rs/moq-mux/src/container/flv/export.rs b/rs/moq-mux/src/container/flv/export.rs index 93e1d71e23b..4f1e9dd1662 100644 --- a/rs/moq-mux/src/container/flv/export.rs +++ b/rs/moq-mux/src/container/flv/export.rs @@ -128,7 +128,7 @@ impl Flavor { /// Produce an FLV byte stream from a broadcast, following a catalog stream. /// -/// Use [`next`](Self::next) to pull byte chunks. The first chunk is the FLV file +/// Use [`poll_next`](Self::poll_next) to pull byte chunks. The first chunk is the FLV file /// header followed by the AVC/AAC sequence headers; each subsequent chunk is the /// tag for one media frame. Returns `None` when the broadcast ends. Narrow the /// renditions first with [`Stream::select`](crate::catalog::Stream::select). @@ -380,12 +380,8 @@ impl Export { self } - /// Get the next byte chunk. - pub async fn next(&mut self) -> anyhow::Result> { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next byte chunk. + /// Get the next FLV byte chunk, or `None` once the export ends. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { // 1. Drain catalog updates to discover the track layout. while let Some(catalog) = self.catalog.as_mut() { diff --git a/rs/moq-mux/src/container/fmp4/export.rs b/rs/moq-mux/src/container/fmp4/export.rs index b52bcaa6660..d95921a2e3b 100644 --- a/rs/moq-mux/src/container/fmp4/export.rs +++ b/rs/moq-mux/src/container/fmp4/export.rs @@ -32,10 +32,10 @@ const INIT_QUEUE: Duration = Duration::from_secs(30); /// (un)subscribes per-rendition tracks as the catalog changes, decodes both Legacy and /// CMAF tracks via a per-track source, and re-encodes everything as a merged init /// segment + moof+mdat fragments, each track in presentation-timestamp order (see -/// [`next`](Self::next) for how far that holds across tracks). This is what an fMP4 +/// [`poll_next`](Self::poll_next) for how far that holds across tracks). This is what an fMP4 /// player (e.g. ffplay, MSE) expects. /// -/// Use [`next`](Self::next) to pull byte chunks: the first call returns the merged +/// Use [`poll_next`](Self::poll_next) to pull byte chunks: the first call returns the merged /// init segment (ftyp + multi-track moov), subsequent calls return moof+mdat /// fragments. By default a fragment is a group, emitted when the group closes. A video /// keyframe also rolls the fragment within a group; audio follows publisher cuts. [`with_fragment_duration`](Self::with_fragment_duration) caps the @@ -55,7 +55,7 @@ const INIT_QUEUE: Duration = Duration::from_secs(30); /// the new track's first sync sample past what was read. A new rendition, a changed sample /// entry, or a return that replays media already written fails the export. /// -/// [`next_chunk`](Self::next_chunk) returns the same bytes as a [`Chunk`], which +/// [`poll_next_chunk`](Self::poll_next_chunk) returns the same bytes as a [`Chunk`], which /// separates the init segment from a [`Fragment`] carrying whether it begins at a /// sync sample and how long it lasts. A segmenting consumer (e.g. an HLS/LL-HLS /// packager) needs that to map fragments onto segments and parts; narrow the catalog @@ -612,24 +612,16 @@ impl Export { /// needs: a demuxer indexes by track and reads `tfdt`. Narrow the catalog to a /// single rendition with [`Stream::select`](crate::catalog::Stream::select) where /// the order matters. - pub async fn next(&mut self) -> Result> { - Ok(self.next_chunk().await?.map(Chunk::into_data)) - } - - /// Poll-based variant of [`Self::next`]. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { Poll::Ready(Ok( std::task::ready!(self.poll_next_chunk(waiter)?).map(Chunk::into_data) )) } - /// Like [`next`](Self::next) but returns a [`Chunk`], which tells the init segment + /// Like [`poll_next`](Self::poll_next) but gets a [`Chunk`], which tells the init segment /// apart from a [`Fragment`] and carries that fragment's segmenting metadata. - pub async fn next_chunk(&mut self) -> Result> { - kio::wait(|waiter| self.poll_next_chunk(waiter)).await - } - - /// Poll-based variant of [`Self::next_chunk`]. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next_chunk(&mut self, waiter: &kio::Waiter) -> Poll>> { // Appending a frame to a track's buffer restarts the search for work by // going around this loop again, never by re-entering the function: a diff --git a/rs/moq-mux/src/container/fmp4/muxer.rs b/rs/moq-mux/src/container/fmp4/muxer.rs index 6bdc004b08d..5a8da38e1f4 100644 --- a/rs/moq-mux/src/container/fmp4/muxer.rs +++ b/rs/moq-mux/src/container/fmp4/muxer.rs @@ -1,5 +1,6 @@ //! One-shot CMAF muxing for individually fetched groups. +use std::task::{Poll, ready}; use std::time::Duration; use bytes::Bytes; @@ -177,11 +178,23 @@ impl Muxer { /// block until the publisher closes it). Parameter-set frames are absorbed into the codec /// config record; the group's first emitted frame is marked a keyframe (a group opens on /// one by convention). + #[cfg(feature = "async")] pub async fn read(&mut self, group: &mut moq_net::group::Consumer) -> crate::Result> { + let mut out = Vec::new(); + kio::wait(|waiter| self.poll_read(waiter, group, &mut out)).await + } + + /// [`read`](Self::read), polled: ready with the group's frames once it ends. Pass the same + /// empty `out` to every poll of one group, which holds the frames decoded so far. + pub fn poll_read( + &mut self, + waiter: &kio::Waiter, + group: &mut moq_net::group::Consumer, + out: &mut Vec, + ) -> Poll>> { use crate::container::Container as _; - let mut out: Vec = Vec::new(); - while let Some(wire) = group.read_frame().await? { + while let Some(wire) = ready!(group.poll_read_frame(waiter))? { for frame in self.container.decode(wire)? { if let Some(bound) = self.container.end(&frame) { if matches!(self.kind, Kind::Video(_)) @@ -211,10 +224,11 @@ impl Muxer { } } } + let mut out = std::mem::take(out); if let Some(first) = out.first_mut() { first.keyframe = true; } - Ok(out) + Poll::Ready(Ok(out)) } /// The picture size `keyframe` decodes to: from its in-band parameter sets (Annex-B, or diff --git a/rs/moq-mux/src/container/group.rs b/rs/moq-mux/src/container/group.rs index 55d9d8a2025..f2957c9dc5c 100644 --- a/rs/moq-mux/src/container/group.rs +++ b/rs/moq-mux/src/container/group.rs @@ -41,11 +41,7 @@ impl GroupConsumer { } /// Read the next frame, or `None` once the group ends. - pub async fn read(&mut self) -> Result, F::Error> { - kio::wait(|waiter| self.poll_read(waiter)).await - } - - /// Poll for the next frame, without blocking. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_read(&mut self, waiter: &kio::Waiter) -> Poll, F::Error>> { // Peek past an untimed video frame for a marker or successor that times it, but only // at what has already landed: a live group returns the frame now rather than waiting. diff --git a/rs/moq-mux/src/container/mkv/export.rs b/rs/moq-mux/src/container/mkv/export.rs index 2e45f293fcd..b3bc30b0972 100644 --- a/rs/moq-mux/src/container/mkv/export.rs +++ b/rs/moq-mux/src/container/mkv/export.rs @@ -26,7 +26,7 @@ const TIMESTAMP_SCALE_NS: u64 = 1_000_000; /// re-encodes everything as EBML + Segment + Tracks + Cluster/SimpleBlock tags ready /// for any Matroska-aware consumer (ffplay, libwebm, browser MSE for WebM). /// -/// Use [`next`](Self::next) to pull byte chunks. The first chunk is the file +/// Use [`poll_next`](Self::poll_next) to pull byte chunks. The first chunk is the file /// header (EBML + unknown-size Segment + Info + Tracks); subsequent chunks are /// complete Cluster elements. By default a Cluster contains one GOP (rolled /// over on each video keyframe, or on i16 timestamp overflow); @@ -374,12 +374,8 @@ impl Export { self } - /// Get the next byte chunk. - pub async fn next(&mut self) -> Result> { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next MKV byte chunk, or `None` once the export ends. + /// Get the next MKV byte chunk, or `None` once the export ends. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { // 1. Drain catalog updates. while let Some(catalog) = self.catalog.as_mut() { diff --git a/rs/moq-mux/src/container/mp4.rs b/rs/moq-mux/src/container/mp4.rs index f6fe25118bb..f4453203a44 100644 --- a/rs/moq-mux/src/container/mp4.rs +++ b/rs/moq-mux/src/container/mp4.rs @@ -36,7 +36,7 @@ const NON_SYNC: u32 = 0x0001_0000; /// Subscribe to a moq broadcast and record it as a crash-safe regular MP4. /// -/// [`next`](Self::next) returns the bytes to append to the file, the same fMP4 an +/// [`poll_next`](Self::poll_next) returns the bytes to append to the file, the same fMP4 an /// [`fmp4::Export`] writes plus the reserved box and track descriptions, so the file plays if /// the recording never finishes. [`finish`](Self::finish) then returns the trailing `moov` and /// the `mdat` header that make it a regular MP4. Writing those needs a file that can be @@ -60,7 +60,7 @@ pub struct Export { #[derive(Clone, Debug)] pub struct Finish { /// Where [`moov`](Self::moov) is written, from the start of the file, which is cut there - /// first: the end of everything [`Export::next`] returned, or of a repaired file's last + /// first: the end of everything [`Export::poll_next`] returned, or of a repaired file's last /// complete fragment. pub moov_offset: u64, @@ -101,11 +101,7 @@ impl Export { /// The next bytes to append to the file, or `None` once the broadcast ends or the export /// [stops](Self::stop). - pub async fn next(&mut self) -> Result> { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll-based variant of [`Self::next`]. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { let Some(chunk) = std::task::ready!(self.fmp4.poll_next_chunk(waiter)?) else { return Poll::Ready(Ok(None)); @@ -118,7 +114,7 @@ impl Export { Poll::Ready(Ok(Some(data))) } - /// Stop recording before the broadcast ends. [`next`](Self::next) still returns the media + /// Stop recording before the broadcast ends. [`poll_next`](Self::poll_next) still returns the media /// already read, the fragment each track was buffering included, then `None`. pub fn stop(&mut self) { self.fmp4.stop(); diff --git a/rs/moq-mux/src/container/ts/export.rs b/rs/moq-mux/src/container/ts/export.rs index 75806816de2..67065dab5be 100644 --- a/rs/moq-mux/src/container/ts/export.rs +++ b/rs/moq-mux/src/container/ts/export.rs @@ -101,7 +101,7 @@ fn sanitize_mux_rate(rate: u64) -> Option { /// Subscribe to a broadcast and produce an MPEG-TS byte stream. /// -/// Use [`next`](Self::next) to pull one [`Frame`] per PCR grid slot: its `payload` +/// Use [`poll_next`](Self::poll_next) to pull one [`Frame`] per PCR grid slot: its `payload` /// is the TS packets belonging to that slot, stamped at the slot's boundary. The /// leading PAT/PMT rides on the first frame (so it inherits a real timestamp), and /// is re-emitted at video keyframes and periodically for mid-stream tune-in. SI and @@ -1715,11 +1715,7 @@ where /// re-emitted at video keyframes and periodically for mid-stream tune-in. /// Returns `None` when the broadcast ends. `duration` is always `None`: the /// muxer has no use for it. - pub async fn next(&mut self) -> crate::Result> { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next MPEG-TS packet frame, or `None` once the export ends. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { if self.pending.is_some() { self.stop_old(waiter)?; diff --git a/rs/moq-mux/src/container/ts/follower.rs b/rs/moq-mux/src/container/ts/follower.rs index 4ffe82c9cd2..3ceb2887dfe 100644 --- a/rs/moq-mux/src/container/ts/follower.rs +++ b/rs/moq-mux/src/container/ts/follower.rs @@ -1,6 +1,5 @@ //! [`Follower`]: an [`Export`] driven by the publisher instances serving its path, as a player is. -use std::future::Future; use std::pin::Pin; use std::task::Poll; use std::time::Duration; @@ -10,8 +9,6 @@ use kio::clock::Instant; use super::{Export, catalog}; use crate::container::Frame; -type Subscribing = Pin>> + Send>>; - /// What the instances say serves the exported path. enum Serving { /// The instance being exported, with no change since the export carried on into it. @@ -27,7 +24,7 @@ enum Serving { /// Subscribing to an instance's catalog, to carry the export on into it. struct Subscription { target: moq_net::broadcast::Consumer, - catalog: Subscribing, + catalog: crate::catalog::Subscribing, } /// Why the export awaits the first catalog snapshot of the instance it carried on into. @@ -148,11 +145,7 @@ impl Follower { } /// Get the next muxed frame, or `None` once the broadcast ends with nothing to follow. - pub async fn next(&mut self) -> crate::Result> { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next muxed frame, registering `waiter` when there is none yet. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { let next = self.poll_state(waiter); if matches!(next, Poll::Ready(Ok(None) | Err(_))) { @@ -245,7 +238,7 @@ impl Follower { self.state = State::Settling { end, deadline }; continue; } - match waiter.poll_future(subscription.catalog.as_mut()) { + match subscription.catalog.poll_ok(waiter) { Poll::Ready(Ok(catalog)) => { self.carry_on(&subscription.target, catalog)?; self.state = State::Running { @@ -302,7 +295,7 @@ impl Follower { *switching = Some(self.subscribe(target.clone())); } let subscription = switching.as_mut().expect("just set"); - match waiter.poll_future(subscription.catalog.as_mut()) { + match subscription.catalog.poll_ok(waiter) { Poll::Ready(Ok(catalog)) => { let subscription = switching.take().expect("just polled"); self.carry_on(&subscription.target, catalog)?; @@ -371,10 +364,9 @@ impl Follower { /// Begin subscribing to `target`'s catalog. fn subscribe(&self, target: moq_net::broadcast::Consumer) -> Subscription { - let (broadcast, format) = (target.clone(), self.export.catalog_format()); Subscription { + catalog: crate::catalog::Subscribing::new(&target, self.export.catalog_format()), target, - catalog: Box::pin(async move { crate::catalog::Consumer::new(&broadcast, format).await }), } } diff --git a/rs/moq-mux/src/container/ts/release.rs b/rs/moq-mux/src/container/ts/release.rs index 6c319ec00ce..97b8f43d667 100644 --- a/rs/moq-mux/src/container/ts/release.rs +++ b/rs/moq-mux/src/container/ts/release.rs @@ -5,7 +5,9 @@ use std::time::Duration; use kio::clock::Instant; -use crate::catalog::{CatalogFormat, Stream}; +#[cfg(feature = "async")] +use crate::catalog::CatalogFormat; +use crate::catalog::Stream; use crate::container::consumer::Event; use crate::container::{ExportSource, Frame}; use crate::jitter::{self, Arrival, Push}; @@ -42,13 +44,15 @@ impl Release { /// Subscribe to the catalog of `source`, waiting for the `m2ts` section that names its track. /// /// Fails for [`CatalogFormat::Msf`], which has no `m2ts` section. + #[cfg(feature = "async")] pub async fn new(source: crate::Source, catalog_format: CatalogFormat) -> crate::Result { if catalog_format == CatalogFormat::Msf { return Err( anyhow::anyhow!("the MSF catalog has no `m2ts` passthrough track; read the hang catalog").into(), ); } - let catalog = source.catalog(catalog_format).await?; + let mut subscribing = crate::catalog::Subscribing::new(&source.broadcast(), catalog_format); + let catalog = kio::wait(|waiter| subscribing.poll_ok(waiter)).await?; let mut release = Self { source, catalog: Some(catalog), @@ -123,11 +127,7 @@ impl Release { /// The next object, at its time: whole TS packets, verbatim, stamped with the PCR time of /// its first byte. Returns `None` once the track, or a catalog that never named one, ends. - pub async fn next(&mut self) -> crate::Result> { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next original MPEG-TS frame at its release time, or `None` at end. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { while let Some(catalog) = self.catalog.as_mut() { match catalog.poll_next(waiter)? { diff --git a/rs/moq-mux/src/follow.rs b/rs/moq-mux/src/follow.rs index 675cc5e18c7..76294b1183f 100644 --- a/rs/moq-mux/src/follow.rs +++ b/rs/moq-mux/src/follow.rs @@ -45,13 +45,14 @@ impl Instances { Self(Feed::Fixed(Some(broadcast))) } - /// Poll for the next change, registering `waiter` when there is none yet. + /// The next change. /// /// Returns each instance as it comes online or takes over, and `None` once the path goes /// offline or, with an instance named, another one takes it over; see /// [`Following::available`](moq_net::origin::Following::available). Fails once the origin /// closes, and with any resolution failure, such as a handler refusing the path; the call /// after a refusal waits for the serving route to change. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { let (following, instance, live) = match &mut self.0 { Feed::Following { @@ -84,11 +85,6 @@ impl Instances { } } } - - /// The next change; see [`Self::poll_next`]. - pub async fn next(&mut self) -> crate::Result> { - kio::wait(|waiter| self.poll_next(waiter)).await - } } #[cfg(test)] diff --git a/rs/moq-mux/src/json.rs b/rs/moq-mux/src/json.rs index fe4711723ab..b32a1a34229 100644 --- a/rs/moq-mux/src/json.rs +++ b/rs/moq-mux/src/json.rs @@ -78,11 +78,15 @@ //! # } //! ``` +use std::task::{Poll, ready}; + use moq_net::Timed; use serde::Serialize; use serde::de::DeserializeOwned; -use hang::catalog::{JsonConfig, Mode, Reference}; +#[cfg(feature = "async")] +use hang::catalog::Reference; +use hang::catalog::{JsonConfig, Mode}; use crate::catalog::hang::CatalogExt; use crate::catalog::{IntoRendition, Listing, RenditionConfig}; @@ -201,14 +205,10 @@ impl Snapshot { } /// Wait for a joiner to want a fresh group, then restate the value in one; see - /// [`moq_json::snapshot::Producer::refresh`]. + /// [`moq_json::snapshot::Producer::poll_refresh`]. /// /// A restatement is not a new value, so it measures neither `jitter` nor `delay`. - pub async fn refresh(&self) -> crate::Result<()> { - Ok(self.inner.refresh().await?) - } - - /// Poll for [`refresh`](Self::refresh), for an owner driving the producer synchronously. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_refresh(&self, waiter: &kio::Waiter) -> std::task::Poll> { self.inner.poll_refresh(waiter).map_err(Into::into) } @@ -368,20 +368,10 @@ impl Consumer { /// /// Errors if the entry declares a mode this build doesn't implement, which is a track a consumer /// must skip rather than guess at. + #[cfg(feature = "async")] pub async fn from_track(track: moq_net::track::Consumer, config: &JsonConfig) -> crate::Result { - let inner = match &config.mode { - Mode::Snapshot => Inner::Snapshot(Box::new(moq_json::snapshot::Consumer::new(track, Default::default()))), - Mode::Stream => Inner::Stream(Box::new(moq_json::stream::Consumer::new( - track.subscribe(None).await?, - Default::default(), - ))), - other => return Err(crate::Error::UnsupportedMode(other.to_string())), - }; - - Ok(Self { - inner, - mode: config.mode.clone(), - }) + let mut opening = Opening::new(track, config)?; + kio::wait(|waiter| opening.poll(waiter)).await } /// The mode the publisher chose. @@ -394,11 +384,7 @@ impl Consumer { } /// Get the next value, or `None` once the track ends. - pub async fn next(&mut self) -> crate::Result>> { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next value, without blocking. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> std::task::Poll>>> { let result = match &mut self.inner { Inner::Snapshot(consumer) => consumer.poll_next(waiter), @@ -411,12 +397,9 @@ impl Consumer { } } - /// Skip to the newest snapshot state; stream mode returns `UnsupportedMode`. - pub async fn latest(&mut self) -> crate::Result>> { - kio::wait(|waiter| self.poll_latest(waiter)).await - } - - /// Poll for the newest snapshot state without discarding records from a lossless stream. + /// Skip to the newest snapshot state; stream mode returns `UnsupportedMode` rather than + /// discard records from a lossless stream. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_latest(&mut self, waiter: &kio::Waiter) -> std::task::Poll>>> { match &mut self.inner { Inner::Snapshot(consumer) => consumer.poll_latest(waiter).map(|result| result.map_err(Into::into)), @@ -432,12 +415,86 @@ impl crate::catalog::Entry<'_, JsonConfig> { /// field names a sibling broadcast is followed rather than silently read from the catalog's own. /// /// Errors if the entry declares a mode this build doesn't implement. + #[cfg(feature = "async")] pub async fn subscribe(&self, source: &crate::Source) -> crate::Result> { - let broadcast = source - .bind(self.config().broadcast.as_ref(), self.config().instance.as_ref())? - .broadcast() - .await?; - Consumer::from_track(broadcast.track(self.config().track(self.name()))?, self.config()).await + let binding = source.bind(self.config().broadcast.as_ref(), self.config().instance.as_ref())?; + let mut subscribe = EntrySubscribe::Binding { + binding, + name: self.config().track(self.name()).to_string(), + config: self.config().clone(), + }; + kio::wait(|waiter| subscribe.poll(waiter)).await + } +} + +/// A [`Consumer`] waiting for its subscription: a stream's here, a snapshot's on its first read. +struct Opening { + mode: Mode, + state: OpeningState, +} + +enum OpeningState { + Snapshot(Option>>), + Stream(Box), +} + +#[cfg_attr(not(feature = "async"), allow(dead_code, reason = "async-only"))] +impl Opening { + /// Errors if the entry declares a mode this build doesn't implement. + fn new(track: moq_net::track::Consumer, config: &JsonConfig) -> crate::Result { + let state = match &config.mode { + Mode::Snapshot => OpeningState::Snapshot(Some(Box::new(moq_json::snapshot::Consumer::new( + track, + Default::default(), + )))), + Mode::Stream => OpeningState::Stream(Box::new(track.subscribe(None).into_inner())), + other => return Err(crate::Error::UnsupportedMode(other.to_string())), + }; + Ok(Self { + mode: config.mode.clone(), + state, + }) + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll>> { + let inner = match &mut self.state { + OpeningState::Snapshot(consumer) => Inner::Snapshot(consumer.take().expect("resolves once")), + OpeningState::Stream(subscribing) => { + let track = ready!(subscribing.poll_ok(waiter))?; + Inner::Stream(Box::new(moq_json::stream::Consumer::new(track, Default::default()))) + } + }; + Poll::Ready(Ok(Consumer { + inner, + mode: self.mode.clone(), + })) + } +} + +/// [`Entry::subscribe`](crate::catalog::Entry::subscribe): the entry's broadcast, then its track. +#[cfg_attr(not(feature = "async"), allow(dead_code, reason = "async-only"))] +#[expect(clippy::large_enum_variant, reason = "a machine's enum is its storage")] +enum EntrySubscribe { + Binding { + binding: crate::Binding, + name: String, + config: JsonConfig, + }, + Opening(Opening), +} + +#[cfg_attr(not(feature = "async"), allow(dead_code, reason = "async-only"))] +impl EntrySubscribe { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll>> { + if let Self::Binding { binding, name, config } = self { + let broadcast = ready!(binding.poll_broadcast(waiter))?; + // The binding drops as soon as it answers, so its request ends with it. + *self = Self::Opening(Opening::new(broadcast.track(name)?, config)?); + } + let Self::Opening(opening) = self else { + unreachable!("bound above"); + }; + opening.poll(waiter) } } diff --git a/rs/moq-mux/src/lib.rs b/rs/moq-mux/src/lib.rs index 657ed82189e..9247b609a88 100644 --- a/rs/moq-mux/src/lib.rs +++ b/rs/moq-mux/src/lib.rs @@ -1,4 +1,11 @@ #![warn(missing_docs)] +#![cfg_attr( + not(feature = "async"), + allow( + rustdoc::broken_intra_doc_links, + reason = "the docs link the async-only functions, which a poll-only build leaves out" + ) +)] //! Media muxers and demuxers for MoQ. //! @@ -28,6 +35,9 @@ //! serve it, for a reader that restarts on each one. //! - [`Linger`] keeps a cleanly ended broadcast up for a fixed time, so a //! subscriber that arrives after the end still gets its tail. +//! +//! The `poll_*` methods are always available. Enable the `async` feature for their generated +//! async wrappers and the async constructors with no poll of their own; it is off by default. #[cfg(any(test, feature = "fuzz"))] mod allocation; diff --git a/rs/moq-mux/src/linger.rs b/rs/moq-mux/src/linger.rs index ce8ee43b173..b13cba9821d 100644 --- a/rs/moq-mux/src/linger.rs +++ b/rs/moq-mux/src/linger.rs @@ -61,6 +61,9 @@ impl Linger { } /// Wait for the deadline on the ambient clock, then close the broadcasts. + /// + /// Hand-written rather than generated, since it consumes the handle to close them. + #[cfg(feature = "async")] pub async fn wait(mut self) { kio::wait(|waiter| self.poll_wait(waiter)).await } diff --git a/rs/moq-mux/src/rate.rs b/rs/moq-mux/src/rate.rs index ea5abe49a6f..c61d8250dec 100644 --- a/rs/moq-mux/src/rate.rs +++ b/rs/moq-mux/src/rate.rs @@ -463,8 +463,8 @@ impl Follower { } } - /// Poll for the next target the sender should produce at, or the error that - /// closed the grant. + /// The next target the sender should produce at, or the error that closed the grant. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll> { loop { let estimate = if std::mem::take(&mut self.refeed) { @@ -484,12 +484,7 @@ impl Follower { } } - /// The next target the sender should produce at, or the error that closed the grant. - pub async fn next(&mut self) -> Result { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Take back the target [`next`](Self::next) last returned, because the sender + /// Take back the target [`poll_next`](Self::poll_next) last returned, because the sender /// failed to apply it. /// /// The previous target is restored and asked for again no sooner than one ramp diff --git a/rs/moq-mux/src/source.rs b/rs/moq-mux/src/source.rs index c90d9955712..2eaf862d879 100644 --- a/rs/moq-mux/src/source.rs +++ b/rs/moq-mux/src/source.rs @@ -191,11 +191,13 @@ impl Source { /// The stream rejects a catalog carrying a `broadcast` reference that walks above the root /// (see [`Error::EscapingBroadcast`](crate::Error::EscapingBroadcast)), so every snapshot a /// consumer sees is one whose references all name a broadcast it may reach. + #[cfg(feature = "async")] pub async fn catalog( &self, format: crate::catalog::CatalogFormat, ) -> crate::Result> { - crate::catalog::Consumer::new(&self.broadcast(), format).await + let mut subscribing = crate::catalog::Subscribing::new(&self.broadcast(), format); + kio::wait(|waiter| subscribing.poll_ok(waiter)).await } /// The catalog broadcast, as a binding. @@ -310,12 +312,14 @@ impl Source { /// A reference that escapes above the origin root is /// [`Error::EscapingBroadcast`](crate::Error::EscapingBroadcast): the rendition names /// no broadcast, so there is nothing to resolve. + #[cfg(feature = "async")] pub async fn resolve( &self, rel: Option<&moq_net::path::Relative>, instance: Option<&moq_net::Instance>, ) -> crate::Result { - self.bind(rel, instance)?.broadcast().await + let binding = self.bind(rel, instance)?; + kio::wait(|waiter| binding.poll_broadcast(waiter)).await } /// Start one broadcast request now and retain its result for repeated reads. @@ -354,14 +358,29 @@ impl Source { /// This is the async counterpart to the poll-driven container exporters: consumers /// that wrap a raw [`moq_net::track::Subscriber`] themselves (e.g. the WebRTC egress) /// use it to honor cross-broadcast renditions without reimplementing the path math. + #[cfg(feature = "async")] pub async fn subscribe_track( &self, rel: Option<&moq_net::path::Relative>, instance: Option<&moq_net::Instance>, name: &str, ) -> crate::Result { - let broadcast = self.bind(rel, instance)?.broadcast().await?; - Ok(broadcast.track(name)?.subscribe(None).await?) + let mut subscribe = self.subscribe(rel, instance, name)?; + kio::wait(|waiter| subscribe.poll(waiter)).await + } + + /// Start [`subscribe_track`](Self::subscribe_track) for a caller that polls. + #[cfg_attr(not(feature = "async"), allow(dead_code, reason = "async-only"))] + pub(crate) fn subscribe( + &self, + rel: Option<&moq_net::path::Relative>, + instance: Option<&moq_net::Instance>, + name: &str, + ) -> crate::Result { + Ok(Subscribe::Binding { + binding: self.bind(rel, instance)?, + name: name.to_string(), + }) } /// Follow the broadcast a rendition's reference names across its announcements. @@ -461,9 +480,10 @@ impl Binding { self.sibling } - /// Poll for the bound broadcast without blocking. + /// The bound broadcast, once the origin answers. /// - /// Concurrent polls and cancelled waits retain the same request and result. + /// Concurrent reads and cancelled waits retain the same request and result. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_broadcast(&self, waiter: &kio::Waiter) -> Poll> { let broadcast = match &self.bound { Bound::Ready(broadcast) => broadcast.clone(), @@ -487,12 +507,56 @@ impl Binding { }; Poll::Ready(check_instance(self.instance.as_ref(), &broadcast).map(|()| broadcast)) } +} - /// The bound broadcast, awaiting the origin's answer when it hasn't arrived yet. - /// - /// Concurrent reads and cancelled waits retain the same request and result. - pub async fn broadcast(&self) -> crate::Result { - kio::wait(|waiter| self.poll_broadcast(waiter)).await +/// A track subscription waiting for its broadcast, then for SUBSCRIBE_OK. +#[cfg_attr(not(feature = "async"), allow(dead_code, reason = "async-only"))] +pub(crate) enum Subscribe { + Binding { + binding: Binding, + name: String, + }, + // The broadcast and track handles stay live until the subscription answers. + Subscribing { + _broadcast: moq_net::broadcast::Consumer, + _track: moq_net::track::Consumer, + subscribing: moq_net::track::Subscribing, + }, + Done, +} + +#[cfg_attr(not(feature = "async"), allow(dead_code, reason = "async-only"))] +impl Subscribe { + /// The track subscriber, once its broadcast resolves and the subscription is confirmed. + pub(crate) fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + if let Self::Binding { binding, name } = self { + let broadcast = ready!(binding.poll_broadcast(waiter)); + let subscribing = broadcast.and_then(|broadcast| { + let track = broadcast.track(name)?; + let subscribing = track.subscribe(None).into_inner(); + Ok((broadcast, track, subscribing)) + }); + // The binding drops as soon as it answers, so its request ends with it. + match subscribing { + Ok((broadcast, track, subscribing)) => { + *self = Self::Subscribing { + _broadcast: broadcast, + _track: track, + subscribing, + } + } + Err(err) => { + *self = Self::Done; + return Poll::Ready(Err(err)); + } + } + } + let Self::Subscribing { subscribing, .. } = self else { + panic!("polled a finished track subscription"); + }; + let result = ready!(subscribing.poll_ok(waiter)).map_err(Into::into); + *self = Self::Done; + Poll::Ready(result) } } diff --git a/rs/moq-mux/src/timeline.rs b/rs/moq-mux/src/timeline.rs index 97e97edcefa..0e4ff14a20d 100644 --- a/rs/moq-mux/src/timeline.rs +++ b/rs/moq-mux/src/timeline.rs @@ -491,7 +491,8 @@ pub struct Resume<'a> { } impl Resume<'_> { - /// Resolve the replacement's mapping once all declared old media has drained. + /// The replacement's mapping, once all declared old media has drained. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_ready(&mut self, waiter: &kio::Waiter) -> Poll> { let _ = self.track.presentation.changed.poll(waiter, |_| Poll::Pending); let now = kio::clock::now(); @@ -549,11 +550,6 @@ impl Resume<'_> { } Poll::Ready(Ok(mapping)) } - - /// Wait for the old media to drain and return the replacement's mapping. - pub async fn ready(mut self) -> Result { - kio::wait(|waiter| self.poll_ready(waiter)).await - } } /// Whether a track feeds the picture or the speaker. diff --git a/rs/moq-net-tests/Cargo.toml b/rs/moq-net-tests/Cargo.toml index cbc6aa233cd..1cead4b99ae 100644 --- a/rs/moq-net-tests/Cargo.toml +++ b/rs/moq-net-tests/Cargo.toml @@ -15,7 +15,7 @@ ignored = ["tokio"] [dev-dependencies] futures = { workspace = true } kio = { path = "../kio", features = ["sim"] } -moq-json = { path = "../moq-json" } +moq-json = { path = "../moq-json", features = ["async"] } moq-net = { path = "../moq-net", features = ["fuzz"] } serde_json = { workspace = true } tokio = { workspace = true, features = ["rt"] } diff --git a/rs/moq-net/Cargo.toml b/rs/moq-net/Cargo.toml index b34d923d7a8..ae98fd93c11 100644 --- a/rs/moq-net/Cargo.toml +++ b/rs/moq-net/Cargo.toml @@ -12,9 +12,14 @@ rust-version.workspace = true keywords = ["quic", "http3", "webtransport", "media", "live"] categories = ["multimedia", "network-programming", "web-programming"] +[package.metadata.docs.rs] +features = ["async"] + [features] -default = ["async"] -# Runtime-backed helpers and Lite/IETF sessions; explicit-time model drivers stay available. +default = [] +# The Lite/IETF sessions, runtime-backed helpers, and the generated async wrappers for the +# poll methods. Off by default: enable it where you connect, accept, or `.await`. The poll +# methods and explicit-time model drivers stay available without it. async = ["dep:futures", "kio/async", "kio/runtime"] # Exposes the wire codecs to the `fuzz/` harness and the announce, codec, path, varint, frame, and canonical storage benches, diff --git a/rs/moq-net/MIGRATE.md b/rs/moq-net/MIGRATE.md index 9ee0dae2aa3..0ea701763f4 100644 --- a/rs/moq-net/MIGRATE.md +++ b/rs/moq-net/MIGRATE.md @@ -4,6 +4,15 @@ How to update your code for each breaking change, newest first. ## Unreleased +### The `async` feature is opt-in ([#7607](https://github.com/moq-dev/moq/pull/7607)) + +`async` is no longer a default feature. Enable it to keep the sessions (`Client`, `Server`, `Session`), `time::run`, the announcement `Stream` adapter, and every async wrapper such as `next_group`, `read_frame`, or `closed`: +`moq-net = { version = "...", features = ["async"] }`. +Each wrapper is now generated from its `poll_*` method, which stays available without the feature, and its docs live on that poll. +`frame::Consumer::read_all`, `track::Demand::closed`, and `track::Ordered::finished` gain theirs as the new `poll_read_all`, `poll_closed`, and `poll_finished`. +`moq-tokio`'s `net` feature enables it for you. +`group::Producer::ready` now needs the feature too. + ### Admission completion records its own time `server::Activity::drain()` begins closing admission once its active set empties. diff --git a/rs/moq-net/README.md b/rs/moq-net/README.md index 6d81ccf8935..be0a1f66305 100644 --- a/rs/moq-net/README.md +++ b/rs/moq-net/README.md @@ -15,11 +15,11 @@ Live media is built on top of this layer using something like [hang](https://git - **Groups**: A sequential collection of frames, usually starting with a keyframe. - **Frame**: A timed chunk of data. -The default `async` feature includes Lite/IETF sessions, async convenience methods, -and the announcement `Stream` adapter. -Use `default-features = false` for the model's `poll_*` APIs and origin driver, -polled with explicit time without a native runtime. This excludes sessions and -`time::run`; the origin's core serving logic remains available. +The `async` feature, off by default, includes Lite/IETF sessions, the async wrappers +for the `poll_*` methods, and the announcement `Stream` adapter. +Without it, use the model's `poll_*` APIs and origin driver, polled with explicit +time without a native runtime. This excludes sessions and `time::run`; the origin's +core serving logic remains available. ## Examples diff --git a/rs/moq-net/src/auth.rs b/rs/moq-net/src/auth.rs index 4c5744cdf03..a79093e520a 100644 --- a/rs/moq-net/src/auth.rs +++ b/rs/moq-net/src/auth.rs @@ -1431,17 +1431,19 @@ impl Credential { /// handle's end. A session ending is not one: the next session presents it again. #[cfg(feature = "async")] pub async fn closed(&self) -> Error { + kio::wait(|waiter| self.poll_closed(waiter)).await + } + + /// See [`Self::closed`]. + fn poll_closed(&self, waiter: &kio::Waiter) -> Poll { let selector = self.presented.selector(); - kio::wait(|waiter| { - let state = ready_or!(self.presented.state.poll(waiter, |state| { - match selector.ended(state) { - Some(_) => Poll::Ready(()), - None => Poll::Pending, - } - })); - Poll::Ready(selector.ended(&state).unwrap_or(Error::Cancel)) - }) - .await + let state = ready_or!(self.presented.state.poll(waiter, |state| { + match selector.ended(state) { + Some(_) => Poll::Ready(()), + None => Poll::Pending, + } + })); + Poll::Ready(selector.ended(&state).unwrap_or(Error::Cancel)) } /// Poll for the first verdict on a handle's token: a grant on any session, or the @@ -1554,11 +1556,14 @@ impl Watch { self.selector.current(&self.state.read()).cloned() } - /// Poll for the grant to differ from the last one this watch returned (or saw - /// at creation). + /// The grant once it differs from the last one this watch returned (or saw at + /// creation). + /// + /// # Errors /// - /// `Err` once nothing will change it again: the handle ended, or for a - /// token's watch, the token ended for good. + /// Once nothing will change it again: the handle ended, or for a token's watch, the + /// token ended for good. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_changed(&mut self, waiter: &kio::Waiter) -> Poll>> { let last = &self.last; let selector = self.selector; @@ -1575,16 +1580,6 @@ impl Watch { } Poll::Ready(Err(selector.ended(&state).unwrap_or(Error::Cancel))) } - - /// Wait for the grant to change, returning the new one. - /// - /// # Errors - /// - /// See [`poll_changed`](Self::poll_changed). - #[cfg(feature = "async")] - pub async fn changed(&mut self) -> Result> { - kio::wait(|waiter| self.poll_changed(waiter)).await - } } /// The tokens the peer presents, for an application that answers them itself. @@ -1595,16 +1590,11 @@ pub struct Requests { } impl Requests { - /// Poll for the next token, or `None` once the handle ends. + /// The next token the peer presents, or `None` once the handle ends. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll> { self.queue.poll_next(waiter) } - - /// Wait for the next token the peer presents, or `None` once the handle ends. - #[cfg(feature = "async")] - pub async fn next(&mut self) -> Option { - kio::wait(|waiter| self.poll_next(waiter)).await - } } impl Drop for Requests { @@ -1844,7 +1834,7 @@ impl Request { /// close the session. #[cfg(feature = "async")] pub async fn closed(&self) -> Error { - closed(self.issue()).await + kio::wait(|waiter| poll_closed(self.issue(), waiter)).await } /// Grant the token. The grant holds until the returned [`Issued`] is revoked @@ -2241,18 +2231,6 @@ impl RequestGrant { latch(&held.ended, ended.clone()); Poll::Ready(ended) } - - /// Run `body`, everything a request does once admitted, until it finishes or `grant` - /// (when its token admitted it) ends. See [`until`]. - pub(crate) fn guard( - grant: Option<&Self>, - body: F, - ) -> impl std::future::Future> + Send { - crate::util::sendable(until(body, move |waiter| match grant { - Some(grant) => grant.poll_ended(waiter), - None => Poll::Pending, - })) - } } /// Record `refusal` as why a request grant ended, unless an earlier end already is. @@ -2264,29 +2242,6 @@ fn latch(ended: &kio::Producer>, refusal: Refusal) { } } -/// Run `body` until it finishes, or until `ended` resolves first, which drops `body` wherever -/// it waits. -/// -/// An admitted request runs everything it does under this, its setup included, so authority -/// ending at any await ends the request there, rather than only where a loop remembers to -/// check. `ended` is polled first, so a request never acts on authority that already ended. -/// A dropped `body` releases what it holds; anything it was writing stays buffered whole in -/// the stream's [`Writer`](crate::coding::Writer) as long as each message is buffered before -/// its first await, so the caller can still end the request on the wire. -pub(crate) async fn until( - body: F, - mut ended: impl FnMut(&kio::Waiter) -> Poll, -) -> std::result::Result { - let mut body = std::pin::pin!(body); - kio::wait(|waiter| { - if let Poll::Ready(refusal) = ended(waiter) { - return Poll::Ready(Err(refusal)); - } - waiter.poll_future(body.as_mut()).map(Ok) - }) - .await -} - /// Whether `grant` covers a `kind` request at `path` and has not already expired: admitting /// an expired one until its deadline fires would let others act on it meanwhile. fn live(clock: &crate::time::Deadlines, kind: RequestKind, path: &str, grant: &Grant) -> bool { @@ -2324,7 +2279,7 @@ impl Issued { /// to end ([`Error::Cancel`]), or its stream or the session to end. #[cfg(feature = "async")] pub async fn closed(&self) -> Error { - closed(&self.issue).await + kio::wait(|waiter| poll_closed(&self.issue, waiter)).await } } @@ -2334,19 +2289,16 @@ impl Drop for Issued { } } -#[cfg(feature = "async")] -async fn closed(issue: &kio::Shared) -> Error { - kio::wait(|waiter| { - let issue = ready_or!(issue.poll(waiter, |issue| { - if issue.peer.is_some() { - Poll::Ready(()) - } else { - Poll::Pending - } - })); - Poll::Ready(issue.peer.clone().expect("checked above")) - }) - .await +/// Resolve with why the peer ended the token behind `issue`. +fn poll_closed(issue: &kio::Shared, waiter: &kio::Waiter) -> Poll { + let issue = ready_or!(issue.poll(waiter, |issue| { + if issue.peer.is_some() { + Poll::Ready(()) + } else { + Poll::Pending + } + })); + Poll::Ready(issue.peer.clone().expect("checked above")) } /// Which way media flows, from this side's view. @@ -2593,22 +2545,14 @@ impl Gate { Poll::Ready(()) } - /// Resolve one operation under this gate without dropping the handler's refusal path. - #[cfg(feature = "async")] - pub(crate) async fn guard(&mut self, operation: impl std::future::Future>) -> Result { - let mut operation = std::pin::pin!(operation); - kio::wait(|waiter| { - std::task::ready!(self.poll_ready(waiter)); - // Ahead of the gate it also closes, so a lapsed token is refused as expired. - if let Poll::Ready(refusal) = self.poll_token_ended(waiter) { - return Poll::Ready(Err(refusal.error)); - } - if self.poll_denied(waiter).is_ready() { - return Poll::Ready(Err(Error::Unauthorized)); - } - waiter.poll_future(operation.as_mut()) - }) - .await + /// Resolve with why an admitted request must stop: its token's refusal once it ended, or + /// [`Error::Unauthorized`] once the session no longer allows the path. + pub(crate) fn poll_refused(&mut self, waiter: &kio::Waiter) -> Poll { + // Ahead of the gate it also closes, so a lapsed token is refused as expired. + if let Poll::Ready(refusal) = self.poll_token_ended(waiter) { + return Poll::Ready(refusal.error); + } + self.poll_denied(waiter).map(|()| Error::Unauthorized) } /// Resolve once the session no longer allows the path. A union that is still diff --git a/rs/moq-net/src/auth/tokens.rs b/rs/moq-net/src/auth/tokens.rs index 6d587d42ee3..4bf94415663 100644 --- a/rs/moq-net/src/auth/tokens.rs +++ b/rs/moq-net/src/auth/tokens.rs @@ -128,7 +128,19 @@ impl Tokens { state.hold(); } - /// Poll for the peer's answer to the tokens set last; see [`answered`](Self::answered). + /// The peer's first verdict on the tokens the last [`set`](Self::set) newly presented, + /// or on the initial tokens before any set: `Ok` once each is granted, in band or as the + /// connection credential a later dial carried. + /// + /// A token the set kept, the connection credential included, counts only once the peer + /// refuses it. Nothing is re-checked after the answer: a later revocation, or a new + /// session that refuses a token an earlier one granted, does not change it. + /// + /// # Errors + /// + /// The peer's refusal of a token as [`Error::Session`], [`Error::Unsupported`] when the + /// session cannot carry an in-band token, or the connection's end, such as a refused dial. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_answered(&self, waiter: &kio::Waiter) -> Poll> { let answer = |state: &State| { state @@ -147,23 +159,6 @@ impl Tokens { Poll::Ready(answer(&state).expect("checked answered")) } - /// Wait for the peer's first verdict on the tokens the last [`set`](Self::set) newly - /// presented, or on the initial tokens before any set: `Ok` once each is granted, in band - /// or as the connection credential a later dial carried. - /// - /// A token the set kept, the connection credential included, counts only once the peer - /// refuses it. Nothing is re-checked after the answer: a later revocation, or a new - /// session that refuses a token an earlier one granted, does not change it. - /// - /// # Errors - /// - /// The peer's refusal of a token as [`Error::Session`], [`Error::Unsupported`] when the - /// session cannot carry an in-band token, or the connection's end, such as a refused dial. - #[cfg(feature = "async")] - pub async fn answered(&self) -> Result<()> { - kio::wait(|waiter| self.poll_answered(waiter)).await - } - /// Decide what a dial to `url` offering `versions` presents, and hold every other token in /// band on the session it builds. Call it before each dial. pub fn dial(&self, url: &Url, versions: &crate::Versions) -> Dial { diff --git a/rs/moq-net/src/client.rs b/rs/moq-net/src/client.rs index 5047e1d1afb..bbc75ffcf02 100644 --- a/rs/moq-net/src/client.rs +++ b/rs/moq-net/src/client.rs @@ -1,3 +1,5 @@ +use std::task::{Poll, ready}; + use crate::driver::Transport; use crate::origin; use crate::time::{Deadlines, Instant}; @@ -385,8 +387,21 @@ impl Client { self.connect_boxed(now, Transport::new(session)) } - async fn connect_boxed(&self, now: Instant, session: Transport) -> Result<(Session, crate::Driver), Error> { - let runtime = Deadlines::new(now); + fn connect_boxed( + &self, + now: Instant, + session: Transport, + ) -> impl Future> { + let mut connect = Connect::Start { + runtime: Deadlines::new(now), + session, + }; + kio::wait(move |waiter| connect.poll(self, waiter)) + } + + /// Start the session the transport's ALPN picks, or begin negotiating moq-lite on a bidi + /// SETUP when it picks none. + fn start(&self, runtime: Deadlines, session: Transport) -> Result { let peer_hop = self.assigned_hop(); let (publish, subscribe) = self.origins(peer_hop); @@ -412,8 +427,8 @@ impl Client { limits: self.limits, session: session.clone(), client: true, - publish: publish.clone(), - subscribe: subscribe.clone(), + publish, + subscribe, peer_hop: Some(peer_hop), cost: self.cost, version: draft, @@ -430,7 +445,7 @@ impl Client { })?; tracing::debug!(version = ?v, "connected"); - return Ok(Session::new( + return Ok(Connect::Started(Session::new( runtime, session, v, @@ -447,7 +462,7 @@ impl Client { viewers: self.viewers(), peer_implementation, }, - )); + ))); } Some(alpn @ (ALPN_LITE_05 | ALPN_LITE_06 | ALPN_LITE_07_WIP)) => { let version = match alpn { @@ -456,57 +471,60 @@ impl Client { _ => lite::Version::Lite05, }; self.versions.select(Version::Lite(version)).ok_or(Error::Version)?; - return self.start_lite(runtime, session, version, peer_hop); + return self + .start_lite(runtime, session, version, peer_hop) + .map(Connect::Started); } Some(ALPN_LITE_04) => { self.versions .select(Version::Lite(lite::Version::Lite04)) .ok_or(Error::Version)?; - return self.start_lite(runtime, session, lite::Version::Lite04, peer_hop); + return self + .start_lite(runtime, session, lite::Version::Lite04, peer_hop) + .map(Connect::Started); } Some(ALPN_LITE_03) => { self.versions .select(Version::Lite(lite::Version::Lite03)) .ok_or(Error::Version)?; - return self.start_lite(runtime, session, lite::Version::Lite03, peer_hop); + return self + .start_lite(runtime, session, lite::Version::Lite03, peer_hop) + .map(Connect::Started); } // No ALPN, or `moql`: negotiate moq-lite in a bidi SETUP. Some(ALPN_LITE) | None => self.versions.filter(&NEGOTIATED.into()).ok_or(Error::Version)?, Some(p) => return Err(Error::UnknownAlpn(p.to_string())), }; - // The SETUP keeps lite-01's framing whichever version it offers. - let mut stream = Stream::open(&session, Version::Lite(lite::Version::Lite01)) - .await? - .with_stats(&self.stats); - - let client = setup::Client { - versions: supported.clone().into(), - parameters: lite::Parameters::default().encode_bytes(lite::Version::Lite01)?, - }; - - stream.writer.encode(&client).await?; - - let server: setup::Server = stream.reader.decode().await?; - - let version = supported - .iter() - .find(|v| coding::Version::from(**v) == server.version) - .copied() - .ok_or(Error::Version)?; - let Version::Lite(v) = version else { - return Err(Error::Version); - }; - - // The server ends the stream after its SETUP. Reading to the end first keeps the drop - // from stopping a reply the server is still waiting to have acknowledged. - if version.hands_off() { - kio::wait(|waiter| stream.reader.poll_closed(waiter)).await?; - drop(stream); - return self.start_lite(runtime, session, v, peer_hop); - } + Ok(Connect::Open { + negotiation: Negotiation { + runtime, + session, + peer_hop, + publish, + subscribe, + supported, + }, + open: Default::default(), + }) + } - let stream = stream.with_version(v); + /// Start a lite-01/02 session on the negotiation stream the server's SETUP picked it on. + fn start_negotiated( + &self, + negotiation: Negotiation, + stream: Stream, + version: lite::Version, + ) -> Result<(Session, crate::Driver), Error> { + let Negotiation { + runtime, + session, + peer_hop, + publish, + subscribe, + .. + } = negotiation; + let stream = stream.with_version(version); // One registry, which the subscriber sends and the session hands out. let viewers = self.viewers(); let start = lite::start(lite::Config { @@ -517,21 +535,21 @@ impl Client { viewers: viewers.clone(), session: session.clone(), setup_stream: Some(stream), - publish: publish.clone(), - subscribe: subscribe.clone(), + publish, + subscribe, peer_hop: Some(peer_hop), - version: v, + version, // Lite 01/02 continue on the bidi SETUP stream and have no Setup Stream. our_setup: lite::Setup::default(), peer_setup: None, // Lite 01/02 carry no AUTH. - auth: self.link(v.has_auth(), false)?, + auth: self.link(version.has_auth(), false)?, })?; Ok(Session::new( runtime, session, - version, + Version::Lite(version), start.recv_bandwidth, crate::driver::Protocol::Lite(Box::new(start.driver)), crate::session::Handles { @@ -549,6 +567,139 @@ impl Client { } } +/// What a client without a versioned ALPN negotiates moq-lite with, held until it starts. +struct Negotiation { + runtime: Deadlines, + session: Transport, + peer_hop: crate::Hop, + publish: Option, + subscribe: Option, + /// The versions offered in the client SETUP. + supported: Versions, +} + +/// The client's side of a handshake: [`Client::connect`] as a poll. Dropped while pending, it +/// releases what its step holds, resetting a negotiation stream still open. +enum Connect { + /// Not polled yet: the transport's ALPN picks the path. + Start { runtime: Deadlines, session: Transport }, + /// Started without waiting on the peer. + Started((Session, crate::Driver)), + /// Waiting its turn for credit to open the bidi stream the SETUP is negotiated on. + Open { + negotiation: Negotiation, + open: crate::transport::open::Open, + }, + /// Flushing the client SETUP, already buffered, then reading the server's. + Exchange { + negotiation: Negotiation, + stream: Stream, + }, + /// Reading the stream to its end once the server's SETUP hands the session off to a + /// versioned ALPN path, which keeps the drop from stopping a reply the server is still + /// waiting to have acknowledged. + HandOff { + negotiation: Negotiation, + stream: Stream, + version: lite::Version, + }, + /// Resolved. + Done, +} + +impl Connect { + /// The started session, or the error the handshake failed with. + fn poll(&mut self, client: &Client, waiter: &kio::Waiter) -> Poll> { + let res = ready!(self.poll_step(client, waiter)); + // Release what the handshake held, even on a failure. + *self = Self::Done; + Poll::Ready(res) + } + + fn poll_step(&mut self, client: &Client, waiter: &kio::Waiter) -> Poll> { + loop { + match self { + Self::Start { .. } => { + let Self::Start { runtime, session } = std::mem::replace(self, Self::Done) else { + unreachable!("matched above"); + }; + *self = client.start(runtime, session)?; + } + Self::Started(_) => { + let Self::Started(started) = std::mem::replace(self, Self::Done) else { + unreachable!("matched above"); + }; + return Poll::Ready(Ok(started)); + } + Self::Open { negotiation, open } => { + // The SETUP keeps lite-01's framing whichever version it offers. + let stream = ready!(Stream::poll_open( + &negotiation.session, + open, + Version::Lite(lite::Version::Lite01), + waiter + ))?; + let mut stream = stream.with_stats(&client.stats); + let setup = setup::Client { + versions: negotiation.supported.clone().into(), + parameters: lite::Parameters::default().encode_bytes(lite::Version::Lite01)?, + }; + stream.writer.buffer(&setup)?; + let Self::Open { negotiation, .. } = std::mem::replace(self, Self::Done) else { + unreachable!("matched above"); + }; + *self = Self::Exchange { negotiation, stream }; + } + Self::Exchange { negotiation, stream } => { + ready!(stream.writer.poll_flush(waiter))?; + let server: setup::Server = ready!(stream.reader.poll_decode(waiter))?; + let version = negotiation + .supported + .iter() + .find(|v| coding::Version::from(**v) == server.version) + .copied() + .ok_or(Error::Version)?; + let Version::Lite(version) = version else { + return Poll::Ready(Err(Error::Version)); + }; + let Self::Exchange { negotiation, stream } = std::mem::replace(self, Self::Done) else { + unreachable!("matched above"); + }; + if !Version::Lite(version).hands_off() { + return Poll::Ready(client.start_negotiated(negotiation, stream, version)); + } + *self = Self::HandOff { + negotiation, + stream, + version, + }; + } + Self::HandOff { stream, .. } => { + // The server ends the stream after its SETUP. + ready!(stream.reader.poll_closed(waiter))?; + let Self::HandOff { + negotiation, + stream, + version, + } = std::mem::replace(self, Self::Done) + else { + unreachable!("matched above"); + }; + drop(stream); + let Negotiation { + runtime, + session, + peer_hop, + .. + } = negotiation; + return Poll::Ready(client.start_lite(runtime, session, version, peer_hop)); + } + Self::Done => panic!("handshake polled after it resolved"), + } + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/rs/moq-net/src/coding/reader.rs b/rs/moq-net/src/coding/reader.rs index f14253f0e04..c4e343f17df 100644 --- a/rs/moq-net/src/coding/reader.rs +++ b/rs/moq-net/src/coding/reader.rs @@ -231,7 +231,8 @@ impl Reader { self.poll_with(waiter, consume, decode).map_ok(Some) } - /// Poll for the next message on the stream. + /// The next message on the stream. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_decode Decode<'a, V> + Debug>(&mut self, waiter: &kio::Waiter) -> Poll> where V: Into
+ Copy, @@ -266,16 +267,9 @@ impl Reader { self.poll_with_maybe(waiter, true, decode) } - /// Decode the next message from the stream. - pub async fn decode Decode<'a, V> + Debug>(&mut self) -> Result - where - V: Into + Copy, - { - kio::wait(|waiter| self.poll_decode(waiter)).await - } - /// [`Self::poll_decode`], also returning the bytes it consumed, for a handshake to /// count once the session's stats are known. + #[cfg_attr(feature = "async", kio::wrapper)] pub(crate) fn poll_decode_sized Decode<'a, V> + Debug>( &mut self, waiter: &kio::Waiter, @@ -287,15 +281,8 @@ impl Reader { self.poll_with_sized(waiter, true, |r| T::decode(r, version)) } - /// [`Self::decode`], also returning the bytes it consumed. - pub(crate) async fn decode_sized Decode<'a, V> + Debug>(&mut self) -> Result<(T, usize), Error> - where - V: Into + Copy, - { - kio::wait(|waiter| self.poll_decode_sized(waiter)).await - } - - /// Poll for the next message unless the stream is closed cleanly first. + /// The next message, or `None` if the stream closes cleanly first. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_decode_maybe Decode<'a, V> + Debug>( &mut self, waiter: &kio::Waiter, @@ -307,15 +294,8 @@ impl Reader { self.poll_with_maybe(waiter, true, |r| T::decode(r, version)) } - /// Decode the next message unless the stream is closed. - pub async fn decode_maybe Decode<'a, V> + Debug>(&mut self) -> Result, Error> - where - V: Into + Copy, - { - kio::wait(|waiter| self.poll_decode_maybe(waiter)).await - } - - /// Poll for the next message without consuming it unless the stream closes cleanly first. + /// The next message without consuming it, or `None` if the stream closes cleanly first. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_decode_peek_maybe Decode<'a, V> + Debug>( &mut self, waiter: &kio::Waiter, @@ -327,15 +307,8 @@ impl Reader { self.poll_with_maybe(waiter, false, |r| T::decode(r, version)) } - /// Peek the next message unless the stream is closed. - pub async fn decode_peek_maybe Decode<'a, V> + Debug>(&mut self) -> Result, Error> - where - V: Into + Copy, - { - kio::wait(|waiter| self.poll_decode_peek_maybe(waiter)).await - } - - /// Poll for the next varint on the stream. + /// The next varint on the stream. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_varint(&mut self, waiter: &kio::Waiter) -> Poll> where V: Into + Copy, @@ -343,15 +316,8 @@ impl Reader { self.poll_with(waiter, true, |r| r.varint()) } - /// Read the next varint from the stream. - pub async fn varint(&mut self) -> Result - where - V: Into + Copy, - { - kio::wait(|waiter| self.poll_varint(waiter)).await - } - - /// Poll for the next varint unless the stream is closed cleanly first. + /// The next varint, or `None` if the stream closes cleanly first. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_varint_maybe(&mut self, waiter: &kio::Waiter) -> Poll, Error>> where V: Into + Copy, @@ -359,15 +325,8 @@ impl Reader { self.poll_with_maybe(waiter, true, |r| r.varint()) } - /// Read the next varint unless the stream is closed. - pub async fn varint_maybe(&mut self) -> Result, Error> - where - V: Into + Copy, - { - kio::wait(|waiter| self.poll_varint_maybe(waiter)).await - } - - /// Poll for the next varint without consuming it. + /// The next varint without consuming it. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_varint_peek(&mut self, waiter: &kio::Waiter) -> Poll> where V: Into + Copy, @@ -375,14 +334,6 @@ impl Reader { self.poll_with(waiter, false, |r| r.varint()) } - /// Read the next varint without consuming it. - pub async fn varint_peek(&mut self) -> Result - where - V: Into + Copy, - { - kio::wait(|waiter| self.poll_varint_peek(waiter)).await - } - /// Poll for the next chunk of up to `max` bytes, or `None` once the stream finished. /// Its credit returns as it is read. pub fn poll_read_chunk(&mut self, waiter: &kio::Waiter, max: usize) -> Poll, Error>> { diff --git a/rs/moq-net/src/coding/stream.rs b/rs/moq-net/src/coding/stream.rs index 73a8b4c9538..318cf9f26da 100644 --- a/rs/moq-net/src/coding/stream.rs +++ b/rs/moq-net/src/coding/stream.rs @@ -49,8 +49,9 @@ impl Stream { Poll::Ready(Ok(Self::build(send, recv, version))) } - /// Open a control stream with the given version, waiting its turn for stream credit. - pub async fn open(session: &crate::driver::Transport, version: V) -> Result { + /// [`Self::poll_open`] as a future, for tests. + #[cfg(test)] + pub(crate) async fn open(session: &crate::driver::Transport, version: V) -> Result { let mut open = crate::transport::open::Open::default(); kio::wait(|waiter| Self::poll_open(session, &mut open, version.clone(), waiter)).await } @@ -66,13 +67,13 @@ impl Stream { Poll::Ready(Ok(Self::accepted(send, recv, version))) } - /// Accept a new stream with the given version. - pub async fn accept(session: &mut S, version: V) -> Result + /// [`Self::poll_accept`] as a future, for tests. + #[cfg(test)] + pub(crate) async fn accept(session: &mut S, version: V) -> Result where V: Clone, { - let (send, recv) = session.accept_bi().await.map_err(|err| version.transport_error(err))?; - Ok(Self::accepted(send, recv, version)) + kio::wait(|waiter| Self::poll_accept(session, version.clone(), waiter)).await } /// The peer's stream, answered at the control send order an opened one starts at. diff --git a/rs/moq-net/src/coding/writer.rs b/rs/moq-net/src/coding/writer.rs index d822e8030bb..01bb04c8da2 100644 --- a/rs/moq-net/src/coding/writer.rs +++ b/rs/moq-net/src/coding/writer.rs @@ -107,6 +107,7 @@ impl Writer { /// cancelled `encode` must not be retried with the same message, or two /// copies go on the wire; resume with [`Self::poll_flush`] (or any later /// write) instead. + #[cfg(test)] pub async fn encode + Debug>(&mut self, msg: &T) -> Result<(), Error> where V: Into + Copy, @@ -116,6 +117,7 @@ impl Writer { } /// Encode a varint to the stream, with the same cancellation rules as [`Self::encode`]. + #[cfg(test)] pub async fn varint(&mut self, value: u64) -> Result<(), Error> where V: Into + Copy, @@ -135,9 +137,10 @@ impl Writer { .map_err(|err| self.version.transport_error(err)) } - /// Poll until the entire `Buf` has been written to the stream. + /// Write the entire `Buf` to the stream. /// /// NOTE: This can avoid performing a copy when using `Bytes`. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_write_all(&mut self, waiter: &kio::Waiter, buf: &mut Buf) -> Poll> { while buf.has_remaining() { ready!(self.poll_write(waiter, buf))?; @@ -145,13 +148,6 @@ impl Writer { Poll::Ready(Ok(())) } - /// Write the entire `Buf` to the stream. - /// - /// NOTE: This can avoid performing a copy when using `Bytes`. - pub async fn write_all(&mut self, buf: &mut Buf) -> Result<(), Error> { - kio::wait(|waiter| self.poll_write_all(waiter, buf)).await - } - /// Whether every buffered byte has hit the stream, so a FIN would not cut a message short. pub fn is_flushed(&self) -> bool { self.buffer.is_empty() @@ -201,21 +197,35 @@ impl Writer { /// A message still partly buffered resets the stream with [`Error::Cancel`] instead: a FIN /// would hand the peer a truncated message, and flushing the rest could wait forever on /// credit the peer withholds. - pub async fn close(mut self) -> Result<(), Error> { + #[cfg(test)] + pub async fn close(self) -> Result<(), Error> + where + V: Clone, + { + let mut closing = self.closing()?; + kio::wait(|waiter| closing.poll(waiter)).await + } + + /// Finish the stream, returning the wait for the peer's acknowledgement: [`Self::close`] + /// for a caller that polls. + pub fn closing(mut self) -> Result, Error> + where + V: Clone, + { if !self.is_flushed() { self.abort(&Error::Cancel); return Err(Error::Cancel); } - let Some(mut stream) = self.stream.take() else { - return Ok(()); - }; - - stream.finish().map_err(|err| self.version.transport_error(err))?; - let version = &self.version; - kio::wait(|waiter| stream.poll_closed(waiter).map_err(|err| version.transport_error(err))).await + let version = self.version.clone(); + let mut stream = self.stream.take(); + if let Some(stream) = &mut stream { + stream.finish().map_err(|err| version.transport_error(err))?; + } + Ok(Closing { stream, version }) } - /// Poll until the stream is closed, or the [Self::finish] is acknowledged by the peer. + /// Resolves once the stream is closed, or the [Self::finish] is acknowledged by the peer. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_closed(&mut self, waiter: &kio::Waiter) -> Poll> { let version = &self.version; self.stream @@ -234,11 +244,6 @@ impl Writer { Poll::Ready(res) } - /// Wait for the stream to be closed, or the [Self::finish] to be acknowledged by the peer. - pub async fn closed(&mut self) -> Result<(), Error> { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - /// Set the priority within the stream's send group, or session-wide when ungrouped. /// /// Streams with HIGHER values are transmitted first. This is the transport trait's @@ -263,6 +268,26 @@ impl Writer { } } +/// A [`Writer`] whose stream has sent its FIN, waiting for the peer to acknowledge it. +/// +/// Dropping it never resets the stream, so what was written still arrives. +pub struct Closing { + stream: Option, + version: V, +} + +impl Closing { + /// Resolves once the peer acknowledges everything written, or the stream dies. + pub fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + let Some(stream) = &self.stream else { + return Poll::Ready(Ok(())); + }; + stream + .poll_closed(waiter) + .map_err(|err| self.version.transport_error(err)) + } +} + impl Writer { /// Encode `msg` into the write buffer as `[id][size][body]`, to be sent by /// [`Self::poll_flush`]. The type and body go in together or not at all, so a failed encode @@ -282,6 +307,7 @@ impl Writer { /// Encode an IETF `Message` to the stream as `[type_id][size][body]`, with the same /// all-or-nothing buffering as [`Self::buffer_message`] and cancellation rules as [`Self::encode`]. + #[cfg(test)] pub async fn encode_message<'m, T: ietf::Message<'m>>(&mut self, msg: &T) -> Result<(), Error> { self.buffer_message(T::ID, msg)?; kio::wait(|waiter| self.poll_flush(waiter)).await diff --git a/rs/moq-net/src/driver.rs b/rs/moq-net/src/driver.rs index 7548022b985..7fc72dfaf09 100644 --- a/rs/moq-net/src/driver.rs +++ b/rs/moq-net/src/driver.rs @@ -91,7 +91,7 @@ impl Protocol { fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { match self { Self::Lite(driver) => driver.poll(waiter), - Self::Ietf(driver) => waiter.poll_future(std::pin::Pin::new(driver)), + Self::Ietf(driver) => driver.poll(waiter), } } diff --git a/rs/moq-net/src/ietf/auth.rs b/rs/moq-net/src/ietf/auth.rs index 0a3627f5425..4bb98854770 100644 --- a/rs/moq-net/src/ietf/auth.rs +++ b/rs/moq-net/src/ietf/auth.rs @@ -11,7 +11,7 @@ use std::time::Duration; use bytes::Bytes; -use crate::auth::{Grant, Issue, Link, Reply, Request, Serving}; +use crate::auth::{Grant, Link, Reply, Request, Serving}; use crate::coding::{Decode, DecodeError, Decoder, Encode, EncodeError, Encoder}; use crate::{Error, Patterns, SessionError}; @@ -184,13 +184,13 @@ fn from_code(code: u64) -> Error { /// Refuse an AUTH past our GOAWAY cutoff with AUTH_ERROR, the only answer its presenter reads. #[cfg(feature = "async")] -pub(super) async fn refuse_going_away(mut stream: Stream) -> Result<(), Error> { +pub(super) fn going_away(mut stream: Stream) -> super::request_stream::Reject { let reply = AuthError { code: GOING_AWAY, reason: "going away".to_string(), }; - stream.writer.buffer_message(AuthError::ID, &reply)?; - super::request_stream::deliver(&mut stream.writer).await + let buffered = stream.writer.buffer_message(AuthError::ID, &reply); + super::request_stream::Reject::buffered(stream, buffered) } /// What the peer's connection credential earns by default: publishing what our @@ -206,208 +206,316 @@ pub(super) fn peer_grant( } } -/// Present every token this side adds, one AUTH request each, once the peer's SETUP -/// says it negotiated the extension. Never returns: a session without it just fails -/// the tokens as unsupported. +/// What every one of our Auth requests is sent with. #[cfg(feature = "async")] -pub(super) async fn run_present( +#[derive(Clone)] +struct Presenter { + stats: crate::stats::Session, runtime: crate::time::Deadlines, session: Transport, control: Control, handle: Link, - peer_setup: peer::PeerSetup, version: Version, - stats: crate::stats::Session, -) -> Result<(), Error> { - if !peer_setup.get().await.auth { - handle.unsupported(); - return std::future::pending().await; +} + +/// Present every token this side adds, one AUTH request each as a child task, once the +/// peer's SETUP says it negotiated the extension. Never finishes: a session without it +/// just fails the tokens as unsupported. +#[cfg(feature = "async")] +pub(super) struct Presenting { + presenter: Presenter, + peer_setup: peer::PeerSetup, + /// Whether the peer's SETUP negotiated the extension, once read. + negotiated: Option, + tasks: crate::util::TaskSet, + /// A protocol violation on any token's stream, which ends the session. + violation: kio::Shared>, + /// The session ended, so no token opens another stream. + ended: bool, +} + +#[cfg(feature = "async")] +impl Presenting { + pub(super) fn new( + runtime: crate::time::Deadlines, + session: Transport, + control: Control, + handle: Link, + peer_setup: peer::PeerSetup, + version: Version, + stats: crate::stats::Session, + ) -> Self { + Self { + presenter: Presenter { + stats, + runtime, + session, + control, + handle, + version, + }, + peer_setup, + negotiated: None, + tasks: crate::util::TaskSet::owned(), + violation: Default::default(), + ended: false, + } } - let mut tasks = crate::util::TaskSet::owned(); - let violation: kio::Shared> = Default::default(); - while let Some((id, token)) = tasks - .drive(|waiter| { - if let Poll::Ready(mut err) = violation.poll(waiter, |err| match err.is_some() { + /// Ready only with the protocol violation a token's stream hit. + pub(super) fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + let negotiated = match self.negotiated { + Some(negotiated) => negotiated, + None => { + let negotiated = std::task::ready!(self.peer_setup.poll_get(waiter)).auth; + if !negotiated { + self.presenter.handle.unsupported(); + } + *self.negotiated.insert(negotiated) + } + }; + if !negotiated || self.ended { + return Poll::Pending; + } + + loop { + if let Poll::Ready(mut err) = self.violation.poll(waiter, |err| match err.is_some() { true => Poll::Ready(()), false => Poll::Pending, }) { return Poll::Ready(Err(err.take().unwrap())); } - handle.poll_opening(waiter).map(Ok) - }) - .await? - { - let present = Present { - stats: stats.clone(), - runtime: runtime.clone(), - session: session.clone(), - control: control.clone(), - handle: handle.clone(), - version, - id, - }; - let violation = violation.clone(); - tasks.push(async move { - let mut token = Ok(token); - let mut answered = false; - loop { - // After a GOAWAY, sent or received, the peer must not see new requests. - let err = match token { - Ok(token) if present.handle.active() => present.run(token, &mut answered).await, - Ok(_) => Error::GoingAway, - Err(err) => err, - }; - let err = match err { - // The peer negotiated the extension, but a reset before any reply still tells - // us nothing about the token: it earns nothing here, as on moq-lite. - Error::Stream(_) if !answered => Error::Unsupported, - err => crate::auth::protocol_error(err), - }; - // AUTH_ERROR carries no code that may pass, so unlike moq-lite only a lost - // request is presented again. - if let Some(delay) = present.handle.retry(present.id, &err) { - tracing::debug!(%err, ?delay, "auth token presenting again"); - token = present.reopen(delay).await; - continue; + match self.presenter.handle.poll_opening(waiter) { + Poll::Ready(Some((id, token))) => { + let mut present = Present { + presenter: self.presenter.clone(), + violation: self.violation.clone(), + id, + answered: false, + state: PresentState::Start(Ok(token)), + }; + self.tasks.push(move |waiter: &kio::Waiter| present.poll(waiter)); } - if matches!(err, Error::ProtocolViolation) { - violation.lock().get_or_insert(err.clone()); + // The children are left where they wait, no longer polled. + Poll::Ready(None) => { + self.ended = true; + return Poll::Pending; } - - match &err { - Error::Cancel - | Error::Unsupported - | Error::Transport(_) - | Error::IdleTimeout - | Error::Session(_) => { - tracing::debug!(%err, "auth token ended") - } - err => tracing::warn!(%err, "auth token ended"), + Poll::Pending => { + // The children never end the session; `Ready` just means they're drained + // until the next token. + let _ = self.tasks.poll(waiter); + return Poll::Pending; } - present.handle.ended(present.id, err); - return; } - }); + } } - std::future::pending().await } -/// One token's Auth request. +/// One token's Auth requests: presented, tracked until either side ends it, and +/// presented again while [`Link::retry`] asks. #[cfg(feature = "async")] struct Present { - stats: crate::stats::Session, - runtime: crate::time::Deadlines, - session: Transport, - control: Control, - handle: Link, - version: Version, + presenter: Presenter, + violation: kio::Shared>, id: u64, + /// The peer replied to the token, on this request or an earlier one, so a request it + /// closes or resets before replying is lost rather than unsupported. + answered: bool, + state: PresentState, } #[cfg(feature = "async")] -impl Present { - /// Wait out a retry's `delay`, then take the token to present again. - async fn reopen(&self, delay: Duration) -> Result { - let mut deadline = crate::time::Deadline::after(&self.runtime, delay); - kio::wait(|waiter| { - if self.handle.poll_withdrawn(self.id, waiter).is_ready() { - return Poll::Ready(Err(Error::Cancel)); - } - std::task::ready!(deadline.poll(waiter)); - Poll::Ready(self.handle.reopen(self.id)) - }) - .await - } - - /// Send the token, then track its grant until either side ends it, resolving with why. - /// - /// `answered` records that the peer replied to the token, on this request or an - /// earlier one, so a request it closes or resets before replying is lost rather than - /// unsupported. - async fn run(&self, token: Bytes, answered: &mut bool) -> Error { - let mut stream = match Stream::open(&self.session, self.version).await { - Ok(stream) => stream, - Err(err) => return err, - }; - stream = stream.with_stats(&self.stats); - let request_id = self.control.next_request_id(); - if let Err(err) = stream.writer.encode_message(&Auth { request_id, token }).await { - return err; - } +enum PresentState { + /// The token to present next, or why it cannot be. + Start(Result), + /// Waiting for a stream to send it on. + Opening { + open: crate::transport::open::Open, + token: Bytes, + }, + /// Flushing the AUTH. + Sending(Stream), + /// Reading the peer's replies. + Reading { + stream: Stream, + messages: super::request_stream::Messages, + budget: kio::coop::Budget, + }, + /// Waiting out a retry's delay. + Reopening(crate::time::Deadline), + Done, +} - let mut messages = super::request_stream::Messages::new(self.session.clone(), self.control.client()); - let mut budget = kio::coop::Budget::new(32); +#[cfg(feature = "async")] +impl Present { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { loop { - enum Next { - Withdrawn, - Reply(Result, Error>), - } - - let next = { - let mut read = std::pin::pin!(messages.next(&mut stream.reader)); - kio::wait(|waiter| { - if self.handle.poll_withdrawn(self.id, waiter).is_ready() { - return Poll::Ready(Next::Withdrawn); + let err = match &mut self.state { + PresentState::Start(token) => match std::mem::replace(token, Err(Error::Cancel)) { + // After a GOAWAY, sent or received, the peer must not see new requests. + Ok(token) if self.presenter.handle.active() => { + self.state = PresentState::Opening { + open: Default::default(), + token, + }; + continue; } - waiter.poll_future(read.as_mut()).map(Next::Reply) - }) - .await + Ok(_) => Error::GoingAway, + Err(err) => err, + }, + PresentState::Reopening(deadline) => { + let token = match self.presenter.handle.poll_withdrawn(self.id, waiter) { + Poll::Ready(()) => Err(Error::Cancel), + Poll::Pending => { + std::task::ready!(deadline.poll(waiter)); + self.presenter.handle.reopen(self.id) + } + }; + self.state = PresentState::Start(token); + continue; + } + PresentState::Done => return Poll::Ready(()), + PresentState::Opening { .. } | PresentState::Sending(_) | PresentState::Reading { .. } => { + let err = std::task::ready!(self.poll_request(waiter)); + // Drops the stream, whose reset is what tells the peer. + self.state = PresentState::Done; + err + } }; - let (id, body) = match next { - // Resetting our side is what tells the peer. - Next::Withdrawn => return Error::Cancel, - Next::Reply(Ok(Some(super::request_stream::Framed(id, body)))) => (id, body), - // The peer ended the grant without revoking it. - Next::Reply(Ok(None)) if *answered => return Error::Cancel, - // A peer that closes before answering could not tell us anything. - Next::Reply(Ok(None)) => return Error::Unsupported, - Next::Reply(Err(err)) => return err, + let err = match err { + // The peer negotiated the extension, but a reset before any reply still tells + // us nothing about the token: it earns nothing here, as on moq-lite. + Error::Stream(_) if !self.answered => Error::Unsupported, + err => crate::auth::protocol_error(err), }; + let handle = &self.presenter.handle; + // AUTH_ERROR carries no code that may pass, so unlike moq-lite only a lost + // request is presented again. + if let Some(delay) = handle.retry(self.id, &err) { + tracing::debug!(%err, ?delay, "auth token presenting again"); + self.state = PresentState::Reopening(crate::time::Deadline::after(&self.presenter.runtime, delay)); + continue; + } + if matches!(err, Error::ProtocolViolation) { + self.violation.lock().get_or_insert(err.clone()); + } + + match &err { + Error::Cancel | Error::Unsupported | Error::Transport(_) | Error::IdleTimeout | Error::Session(_) => { + tracing::debug!(%err, "auth token ended") + } + err => tracing::warn!(%err, "auth token ended"), + } + handle.ended(self.id, err); + return Poll::Ready(()); + } + } - match id { - AuthOk::ID => { - let ok: AuthOk = match decode_reply(&body, self.version) { - Ok(ok) => ok, - Err(err) => return err, + /// Send the token, then track its grant until either side ends it, ready with why. + fn poll_request(&mut self, waiter: &kio::Waiter) -> Poll { + let presenter = &self.presenter; + loop { + match &mut self.state { + PresentState::Opening { open, token } => { + let stream = + match std::task::ready!(Stream::poll_open(&presenter.session, open, presenter.version, waiter)) + { + Ok(stream) => stream, + Err(err) => return Poll::Ready(err), + }; + let mut stream = stream.with_stats(&presenter.stats); + let auth = Auth { + request_id: presenter.control.next_request_id(), + token: std::mem::take(token), }; - let now = self.runtime.now(); - // The expiry is the peer's number: one past the local clock's range is - // malformed, not a reason to panic. - let expires = match ok.expires.map(|expires| now.checked_add(expires)) { - Some(None) => return Error::ProtocolViolation, - expires => expires.flatten(), + if let Err(err) = stream.writer.buffer_message(Auth::ID, &auth) { + return Poll::Ready(err); + } + self.state = PresentState::Sending(stream); + } + PresentState::Sending(stream) => { + if let Err(err) = std::task::ready!(stream.writer.poll_flush(waiter)) { + return Poll::Ready(err); + } + let PresentState::Sending(stream) = std::mem::replace(&mut self.state, PresentState::Done) else { + unreachable!() + }; + self.state = PresentState::Reading { + stream, + messages: super::request_stream::Messages::new( + presenter.session.clone(), + presenter.control.client(), + ), + // Each grant recomputes the session's permits, so a peer streaming them + // must not hold the session's other tasks. + budget: kio::coop::Budget::new(32), }; - *answered = true; - self.handle.granted( - self.id, - Grant { - publish: ok.publish, - subscribe: ok.subscribe, - expires, - }, - ); - // Each grant recomputes the session's permits, so a peer streaming them - // must not hold the session's other tasks. - kio::wait(|waiter| budget.poll_yield(waiter)).await; } - AuthError::ID => { - let refused: AuthError = match decode_reply(&body, self.version) { - Ok(refused) => refused, - Err(err) => return err, + PresentState::Reading { + stream, + messages, + budget, + } => loop { + std::task::ready!(budget.poll_yield(waiter)); + if presenter.handle.poll_withdrawn(self.id, waiter).is_ready() { + // Resetting our side is what tells the peer. + return Poll::Ready(Error::Cancel); + } + let (id, body) = match std::task::ready!(messages.poll_next(&mut stream.reader, waiter)) { + Ok(Some(super::request_stream::Framed(id, body))) => (id, body), + // The peer ended the grant without revoking it. + Ok(None) if self.answered => return Poll::Ready(Error::Cancel), + // A peer that closes before answering could not tell us anything. + Ok(None) => return Poll::Ready(Error::Unsupported), + Err(err) => return Poll::Ready(err), }; - let err = from_code(refused.code); - tracing::warn!(%err, code = refused.code, reason = %refused.reason, "auth token refused"); - // A grant the acceptor could not tell, or never judged, is unknown, not - // empty, so only a real refusal settles the union. - if !matches!(err, Error::Unsupported | Error::GoingAway) { - self.handle.refused(self.id); + + match id { + AuthOk::ID => { + let ok: AuthOk = match decode_reply(&body, presenter.version) { + Ok(ok) => ok, + Err(err) => return Poll::Ready(err), + }; + let now = presenter.runtime.now(); + // The expiry is the peer's number: one past the local clock's range is + // malformed, not a reason to panic. + let expires = match ok.expires.map(|expires| now.checked_add(expires)) { + Some(None) => return Poll::Ready(Error::ProtocolViolation), + expires => expires.flatten(), + }; + self.answered = true; + presenter.handle.granted( + self.id, + Grant { + publish: ok.publish, + subscribe: ok.subscribe, + expires, + }, + ); + } + AuthError::ID => { + let refused: AuthError = match decode_reply(&body, presenter.version) { + Ok(refused) => refused, + Err(err) => return Poll::Ready(err), + }; + let err = from_code(refused.code); + tracing::warn!(%err, code = refused.code, reason = %refused.reason, "auth token refused"); + // A grant the acceptor could not tell, or never judged, is unknown, not + // empty, so only a real refusal settles the union. + if !matches!(err, Error::Unsupported | Error::GoingAway) { + presenter.handle.refused(self.id); + } + self.answered = true; + return Poll::Ready(err); + } + _ => return Poll::Ready(Error::UnexpectedMessage), } - *answered = true; - return err; + }, + PresentState::Start(_) | PresentState::Reopening(_) | PresentState::Done => { + unreachable!("not a request") } - _ => return Error::UnexpectedMessage, } } } @@ -434,21 +542,16 @@ pub(super) struct Serve { } impl Serve { - /// Answer one Auth request. The stream lives as long as the token. + /// Start answering one Auth request. The stream lives as long as the token. #[cfg(feature = "async")] - pub(super) async fn run( - self, - mut stream: Stream, - msg: Auth, - version: Version, - ) -> Result<(), Error> { - let mut goaways = super::request_stream::GoAways::new(msg.request_id.0 % 2 == 1); + pub(super) fn answer(self, mut stream: Stream, msg: Auth, version: Version) -> Answering { + let goaways = super::request_stream::GoAways::new(msg.request_id.0 % 2 == 1); let serving = Serving::new(self.handle.clone()); - let issue = &serving.issue; let mut default = None; - match self.handle.acceptor() { + let state = match self.handle.acceptor() { Some(requests) => { - requests.push(Request::new(msg.token, issue.clone())); + requests.push(Request::new(msg.token, serving.issue.clone())); + AnsweringState::Serving } // Only the connection's own credential has a default answer; a token needs // someone to verify it. @@ -457,146 +560,227 @@ impl Serve { code: NOT_SUPPORTED, reason: "tokens are not verified in band".to_string(), }; - if stream.writer.encode_message(&reply).await.is_ok() { - let _ = stream.writer.close().await; + match stream.writer.buffer_message(AuthError::ID, &reply) { + Ok(()) => AnsweringState::Unverified, + Err(_) => AnsweringState::Done, } - return Ok(()); } // Re-sent whenever a new limit changes it. - None => default = Some(crate::auth::DefaultGrant::new(self.handle.clone(), self.peer_grant)), - } - - let err = match serve_issue( - &self.runtime, - issue, - default.as_mut(), - &mut stream, - version, - &mut goaways, - ) - .await - { - Ok(()) => Error::Cancel, - Err(err) => { - match &err { - Error::Cancel - | Error::Unsupported - | Error::Stream(_) - | Error::Session(_) - | Error::Transport(_) - | Error::IdleTimeout => { - tracing::debug!(%err, "auth request ended") - } - err => tracing::warn!(%err, "auth request error"), - } - err + None => { + default = Some(crate::auth::DefaultGrant::new(self.handle.clone(), self.peer_grant)); + AnsweringState::Serving } }; - let err = crate::auth::protocol_error(err); - serving.end(err.clone()); - if matches!(err, Error::ProtocolViolation) { - return Err(err); + Answering { + runtime: self.runtime, + serving, + default, + stream: Some(stream), + version, + goaways, + state, } - Ok(()) } } -/// Anything but a GOAWAY after the presenter's Auth: refused once its type is read, -/// rather than waiting on a body that cannot make it legal. +/// One Auth request being answered: [`Serve::answer`]. #[cfg(feature = "async")] -#[derive(Debug)] -struct Unexpected; +pub(super) struct Answering { + runtime: crate::time::Deadlines, + serving: Serving, + /// Feeds the default acceptor's grant into the outbox as it changes. + default: Option, + /// Taken only to close it after refusing an unverified token. + stream: Option>, + version: Version, + goaways: super::request_stream::GoAways, + state: AnsweringState, +} #[cfg(feature = "async")] -impl<'a> crate::coding::Decode<'a, Version> for Unexpected { - fn decode(r: &mut crate::coding::Decoder<'_>, _: Version) -> Result { - Err(crate::coding::DecodeError::InvalidMessage(r.varint()?)) - } +enum AnsweringState { + /// Flushing the refusal of a token nobody verifies, then closing. + Unverified, + /// Waiting for the refusal's acknowledgement. + UnverifiedClosing { + closing: crate::coding::Closing, + _reader: crate::coding::Reader, + }, + /// Writing the acceptor's replies until either side ends the token. + Serving, + /// Flushing a reply; `last` ends the replies once it is out. + Flushing { + last: bool, + }, + /// Waiting for the acknowledgement of our FIN. + Closing, + Done, } -/// Write the acceptor's replies until either side ends the token. `default` feeds -/// the default acceptor's grant into the outbox as it changes. #[cfg(feature = "async")] -async fn serve_issue( - runtime: &crate::time::Deadlines, - issue: &kio::Shared, - mut default: Option<&mut crate::auth::DefaultGrant>, - stream: &mut Stream, - version: Version, - goaways: &mut super::request_stream::GoAways, -) -> Result<(), Error> { - enum Next { - Withdrawn(Result<(), Error>), - Reply(Option), - } - - loop { - let next = kio::wait(|waiter| { - if let Poll::Ready(res) = goaways - .poll_next::(&mut stream.reader, waiter) - .map(|res| res.map(|_| ())) - { - return Poll::Ready(Next::Withdrawn(res)); - } - if let Some(default) = default.as_deref_mut() - && let Poll::Ready(grant) = default.poll(waiter) - { - issue.lock().outbox.push_back(Reply::Grant(grant)); - } - match issue.poll(waiter, |issue| match issue.outbox.is_empty() && !issue.done { - true => Poll::Pending, - false => Poll::Ready(()), - }) { - Poll::Ready(mut issue) => Poll::Ready(Next::Reply(issue.outbox.pop_front())), - Poll::Pending => Poll::Pending, - } - }) - .await; - - match next { - // The presenter withdrew the token, by FIN or by a cancelling reset. That is - // why it ended, whether or not our own FIN reaches a presenter that left. - Next::Withdrawn(Ok(()) | Err(Error::Stream(crate::StreamError::Cancel))) => { - issue.lock().peer.get_or_insert(Error::Cancel); - break; +impl Answering { + /// Ready once the token's stream ends, failing only on a protocol violation. + pub(super) fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + loop { + match &mut self.state { + AnsweringState::Unverified => { + let stream = self.stream.as_mut().expect("unverified holds the stream"); + let flushed = std::task::ready!(stream.writer.poll_flush(waiter)); + self.state = AnsweringState::Done; + if flushed.is_ok() { + let Stream { writer, reader } = self.stream.take().expect("unverified holds the stream"); + if let Ok(closing) = writer.closing() { + self.state = AnsweringState::UnverifiedClosing { + closing, + _reader: reader, + }; + } + } + } + AnsweringState::UnverifiedClosing { closing, .. } => { + let _ = std::task::ready!(closing.poll(waiter)); + self.state = AnsweringState::Done; + } + AnsweringState::Done => return Poll::Ready(Ok(())), + AnsweringState::Serving | AnsweringState::Flushing { .. } | AnsweringState::Closing => { + let err = match std::task::ready!(self.poll_issue(waiter)) { + Ok(()) => Error::Cancel, + Err(err) => { + match &err { + Error::Cancel + | Error::Unsupported + | Error::Stream(_) + | Error::Session(_) + | Error::Transport(_) + | Error::IdleTimeout => { + tracing::debug!(%err, "auth request ended") + } + err => tracing::warn!(%err, "auth request error"), + } + err + } + }; + self.state = AnsweringState::Done; + let err = crate::auth::protocol_error(err); + self.serving.end(err.clone()); + if matches!(err, Error::ProtocolViolation) { + return Poll::Ready(Err(err)); + } + return Poll::Ready(Ok(())); + } } - Next::Withdrawn(Err(err)) => return Err(err), - Next::Reply(Some(Reply::Grant(grant))) => { - let now = runtime.now(); - let ok = AuthOk { - publish: grant.publish, - subscribe: grant.subscribe, - expires: grant.expires.map(|at| at.saturating_duration_since(now)), - }; - // Sized as the full message before anything is written, so a grant over the u16 - // message size never leaves half a message on the wire. Never trim it: refuse, - // which revokes whatever this stream granted before. - if let Err(err) = ok.encode(&mut Encoder::new(&mut Vec::new(), version.into()), version) { - tracing::debug!(%err, "auth grant cannot be encoded as one AUTH_OK; refusing the token"); - issue.lock().done = true; + } + } + + /// Write the acceptor's replies until either side ends the token. + fn poll_issue(&mut self, waiter: &kio::Waiter) -> Poll> { + let issue = &self.serving.issue; + let stream = self.stream.as_mut().expect("serving holds the stream"); + loop { + let reply = match &mut self.state { + AnsweringState::Flushing { last } => { + std::task::ready!(stream.writer.poll_flush(waiter))?; + match *last { + true => { + stream.writer.finish()?; + self.state = AnsweringState::Closing; + } + false => self.state = AnsweringState::Serving, + } + continue; + } + AnsweringState::Closing => return stream.writer.poll_closed(waiter), + AnsweringState::Serving => { + if let Poll::Ready(res) = self + .goaways + .poll_next::(&mut stream.reader, waiter) + .map(|res| res.map(|_| ())) + { + match res { + // The presenter withdrew the token, by FIN or by a cancelling reset. + // That is why it ended, whether or not our own FIN reaches a + // presenter that left. + Ok(()) | Err(Error::Stream(crate::StreamError::Cancel)) => { + issue.lock().peer.get_or_insert(Error::Cancel); + stream.writer.finish()?; + self.state = AnsweringState::Closing; + continue; + } + Err(err) => return Poll::Ready(Err(err)), + } + } + if let Some(default) = self.default.as_mut() + && let Poll::Ready(grant) = default.poll(waiter) + { + issue.lock().outbox.push_back(Reply::Grant(grant)); + } + let mut issue = std::task::ready!(issue.poll(waiter, |issue| { + match issue.outbox.is_empty() && !issue.done { + true => Poll::Pending, + false => Poll::Ready(()), + } + })); + issue.outbox.pop_front() + } + AnsweringState::Unverified | AnsweringState::UnverifiedClosing { .. } | AnsweringState::Done => { + unreachable!("not serving") + } + }; + + match reply { + Some(Reply::Grant(grant)) => { + let now = self.runtime.now(); + let ok = AuthOk { + publish: grant.publish, + subscribe: grant.subscribe, + expires: grant.expires.map(|at| at.saturating_duration_since(now)), + }; + // Sized as the full message before anything is written, so a grant over the u16 + // message size never leaves half a message on the wire. Never trim it: refuse, + // which revokes whatever this stream granted before. + if let Err(err) = ok.encode(&mut Encoder::new(&mut Vec::new(), self.version.into()), self.version) { + tracing::debug!(%err, "auth grant cannot be encoded as one AUTH_OK; refusing the token"); + issue.lock().done = true; + let refused = AuthError { + code: NOT_SUPPORTED, + reason: "grant cannot be encoded as one AUTH_OK".to_string(), + }; + stream.writer.buffer_message(AuthError::ID, &refused)?; + self.state = AnsweringState::Flushing { last: true }; + continue; + } + stream.writer.buffer_message(AuthOk::ID, &ok)?; + self.state = AnsweringState::Flushing { last: false }; + } + Some(Reply::Refuse { code, reason }) => { let refused = AuthError { - code: NOT_SUPPORTED, - reason: "grant cannot be encoded as one AUTH_OK".to_string(), + code: to_code(code), + reason, }; - stream.writer.encode_message(&refused).await?; - break; + stream.writer.buffer_message(AuthError::ID, &refused)?; + self.state = AnsweringState::Flushing { last: false }; + } + // The app is done with the grant, or refused the token: close our side. + None => { + stream.writer.finish()?; + self.state = AnsweringState::Closing; } - stream.writer.encode_message(&ok).await?; - } - Next::Reply(Some(Reply::Refuse { code, reason })) => { - let refused = AuthError { - code: to_code(code), - reason, - }; - stream.writer.encode_message(&refused).await?; } - // The app is done with the grant, or refused the token: close our side. - Next::Reply(None) => break, } } +} + +/// Anything but a GOAWAY after the presenter's Auth: refused once its type is read, +/// rather than waiting on a body that cannot make it legal. +#[cfg(feature = "async")] +#[derive(Debug)] +struct Unexpected; - stream.writer.finish()?; - stream.writer.closed().await +#[cfg(feature = "async")] +impl<'a> crate::coding::Decode<'a, Version> for Unexpected { + fn decode(r: &mut crate::coding::Decoder<'_>, _: Version) -> Result { + Err(crate::coding::DecodeError::InvalidMessage(r.varint()?)) + } } #[cfg(test)] @@ -611,20 +795,26 @@ mod tests { for client in [false, true] { let first = if client { 0 } else { 1 }; let control = Control::new(client); - let present = Present { - runtime: crate::time::sim(), - stats: Default::default(), - session: crate::driver::Transport::new( - crate::lite::test_transport::ScriptedSession::new(Vec::new()).without_credit(), - ), - control: control.clone(), - handle: Link::new(false), - version, + let mut present = Present { + presenter: Presenter { + runtime: crate::time::sim(), + stats: Default::default(), + session: crate::driver::Transport::new( + crate::lite::test_transport::ScriptedSession::new(Vec::new()).without_credit(), + ), + control: control.clone(), + handle: Link::new(false), + version, + }, + violation: Default::default(), id: 0, + answered: false, + state: PresentState::Opening { + open: Default::default(), + token: Bytes::new(), + }, }; - let mut answered = false; - let mut run = std::pin::pin!(present.run(Bytes::new(), &mut answered)); - let polled = kio::wait(|waiter| Poll::Ready(waiter.poll_future(run.as_mut()).is_pending())).await; + let polled = kio::wait(|waiter| Poll::Ready(present.poll_request(waiter).is_pending())).await; assert!(polled, "the request waits for credit"); assert_eq!(control.next_request_id(), RequestId(first)); } diff --git a/rs/moq-net/src/ietf/peer.rs b/rs/moq-net/src/ietf/peer.rs index bc518ddc8d5..011ea9ec76c 100644 --- a/rs/moq-net/src/ietf/peer.rs +++ b/rs/moq-net/src/ietf/peer.rs @@ -95,12 +95,12 @@ impl PeerSetup { } } - /// Poll for the peer's SETUP. + /// The peer's SETUP. /// - /// The peer MUST send exactly one, so this is ready once that stream is read. Pending - /// forever if it never does; the caller is a session task, cancelled when the driver - /// drops. + /// The peer MUST send exactly one, so this resolves once that stream is read, and never + /// if the peer sends none; the caller is a session task, cancelled when the driver drops. #[cfg(feature = "async")] + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_get(&self, waiter: &kio::Waiter) -> std::task::Poll { let slot = std::task::ready!(self.slot.poll(waiter, |peer| match peer.is_some() { true => std::task::Poll::Ready(()), @@ -108,12 +108,6 @@ impl PeerSetup { })); std::task::Poll::Ready((*slot).expect("polled for Some")) } - - /// Await the peer's SETUP; see [`Self::poll_get`]. - #[cfg(feature = "async")] - pub async fn get(&self) -> Peer { - kio::wait(|waiter| self.poll_get(waiter)).await - } } #[cfg(all(test, feature = "async"))] diff --git a/rs/moq-net/src/ietf/publisher.rs b/rs/moq-net/src/ietf/publisher.rs index 10d1d790f23..f49f04b8416 100644 --- a/rs/moq-net/src/ietf/publisher.rs +++ b/rs/moq-net/src/ietf/publisher.rs @@ -22,7 +22,6 @@ use crate::{ ietf::{self, Control, EndLocation, FetchHeader, FetchType, Filter, Location, RequestId}, priority::{Priority, PriorityHandle, PriorityQueue}, track::{Order, Subscription}, - util::BoxFuture, }; use super::{Message, Version, cluster, error::request, peer, request_update}; @@ -147,108 +146,140 @@ struct PreparedFetch { last: FetchedGroup, } -/// Resolve a one-group FETCH and read its objects before promising them in FETCH_OK. +/// Resolves a one-group FETCH and reads its objects before promising them in FETCH_OK. /// A whole-group request waits for an open group to finish. -async fn prepare_fetch( - track: &track::Consumer, - mut range: FetchRange, +struct PrepareFetch { + range: FetchRange, largest_filter: Option, joined: bool, priority: i8, compression: bool, -) -> Result { - let failed = |err: Error| { - let reason = err.to_string(); - (err, reason) - }; - // Asking for the group needs nothing from the track's info, only judging and answering - // it: a relay's FETCH goes upstream alongside the request for the info, driven while - // that resolves. - let mut ahead = largest_filter.is_none().then(|| { - Ahead::Fetching(track.fetch_group( - range.start.group, - group::Fetch { - priority, - ..Default::default() - }, - )) - }); - let info = track.info(); - kio::wait(|waiter| { - if let Some(ahead) = &mut ahead { - ahead.poll(waiter); + /// The whole group, when the request names it. Asking for the group needs nothing from + /// the track's info, only judging and answering it: a relay's FETCH goes upstream + /// alongside the request for the info, driven while that resolves. + ahead: Option, + info: track::Resolving, + state: PrepareFetchState, +} + +#[kio::machine(self.state: PrepareFetchState)] +impl PrepareFetch { + fn new( + track: &track::Consumer, + range: FetchRange, + largest_filter: Option, + joined: bool, + priority: i8, + compression: bool, + ) -> Self { + let ahead = largest_filter.is_none().then(|| { + Ahead::Fetching(track.fetch_group( + range.start.group, + group::Fetch { + priority, + ..Default::default() + }, + )) + }); + Self { + range, + largest_filter, + joined, + priority, + compression, + ahead, + info: track.info().into_inner(), + state: PrepareFetchState::Start, } - info.poll_ok(waiter) - }) - .await - .map_err(failed)?; - // The cache's newest object is Largest Object while a live feed writes the track, or once - // the track ended at its newest group. A fetch-only copy retains its publisher's report - // independently. A joining FETCH ends where its subscription starts instead. - let finished = track.final_sequence(); - let known = - track.is_live() || finished.is_some_and(|end| Some(end) == track.latest_sequence().map(|latest| latest + 1)); - let largest = match !joined && (known || track.reported_object().is_some()) { - true => live_edge(track).largest, - false => None, - }; - if let Some(filter) = largest_filter { - let Some(largest) = largest else { - return Err(( - Error::Unsupported, - "FETCH relative to Largest Object needs a known publisher edge".into(), - )); + } + + fn poll(&mut self, waiter: &kio::Waiter, track: &track::Consumer) -> Poll> { + let resolved = wait!(Info => { + if let Some(ahead) = &mut self.ahead { + ahead.poll(waiter); + } + self.info.poll_ok(waiter) + }); + resolved.map_err(with_reason)?; + let range: FetchRange = self.judge(track)?; + let mut read: ReadFetch = ReadFetch::new( + track, + range.start.group, + range.start.object, + range.until(range.start.group), + self.priority, + self.compression, + self.ahead.take(), + ); + let last = wait!(Read => read.poll(waiter, track)).map_err(with_reason)?; + Poll::Ready(Ok(PreparedFetch { range, last })) + } + + /// Judge the range against the track, now that its info resolved. + fn judge(&mut self, track: &track::Consumer) -> Result { + let range = &mut self.range; + // The cache's newest object is Largest Object while a live feed writes the track, or once + // the track ended at its newest group. A fetch-only copy retains its publisher's report + // independently. A joining FETCH ends where its subscription starts instead. + let finished = track.final_sequence(); + let known = track.is_live() + || finished.is_some_and(|end| Some(end) == track.latest_sequence().map(|latest| latest + 1)); + let largest = match !self.joined && (known || track.reported_object().is_some()) { + true => live_edge(track).largest, + false => None, }; - range.start = match filter { - Filter::Unfiltered => Location { group: 0, object: 0 }, - Filter::NextObject => Location { + if let Some(filter) = self.largest_filter { + let Some(largest) = largest else { + return Err(( + Error::Unsupported, + "FETCH relative to Largest Object needs a known publisher edge".into(), + )); + }; + range.start = match filter { + Filter::Unfiltered => Location { group: 0, object: 0 }, + Filter::NextObject => Location { + group: largest.group, + object: largest.object.saturating_add(1), + }, + Filter::Relative(groups) => Location { + group: largest.group.saturating_add(1).saturating_sub(groups), + object: 0, + }, + Filter::Absolute { start, .. } => start, + }; + range.end = Location { group: largest.group, - object: largest.object.saturating_add(1), - }, - Filter::Relative(groups) => Location { - group: largest.group.saturating_add(1).saturating_sub(groups), object: 0, - }, - Filter::Absolute { start, .. } => start, - }; - range.end = Location { - group: largest.group, - object: 0, - }; - } - if range.start.group != range.end.group { - return Err((Error::Unsupported, "FETCH spans more than one group".into())); - } - - if let Some(largest) = largest { - // A start past Largest Object in a group the cache holds is refused INVALID_RANGE - // (draft-22 section 9.11), rather than waiting on the live group. A group past every - // one held does not exist yet. - if (range.start.group, range.start.object) > (largest.group, largest.object) { - let held = range - .start - .group - .checked_add(1) - .is_some_and(|end| !track.cached_groups(range.start.group..end).is_empty()); - return Err(match held { - true => (Error::InvalidRange, "start past Largest Object".into()), - false => failed(Error::NotFound), - }); + }; + } + if range.start.group != range.end.group { + return Err((Error::Unsupported, "FETCH spans more than one group".into())); + } + + if let Some(largest) = largest { + // A start past Largest Object in a group the cache holds is refused INVALID_RANGE + // (draft-22 section 9.11), rather than waiting on the live group. A group past every + // one held does not exist yet. + if (range.start.group, range.start.object) > (largest.group, largest.object) { + let held = range + .start + .group + .checked_add(1) + .is_some_and(|end| !track.cached_groups(range.start.group..end).is_empty()); + return Err(match held { + true => (Error::InvalidRange, "start past Largest Object".into()), + false => with_reason(Error::NotFound), + }); + } } + Ok(*range) } +} - let last = read_fetch( - track, - range.start.group, - range.start.object, - range.until(range.start.group), - priority, - compression, - ahead, - ) - .await - .map_err(failed)?; - Ok(PreparedFetch { range, last }) +/// A refusal's error with its reason phrase. +fn with_reason(err: Error) -> (Error, String) { + let reason = err.to_string(); + (err, reason) } /// A whole-group fetch made before the track's info resolved. @@ -265,13 +296,6 @@ impl Ahead { *self = Self::Fetched(result); } } - - async fn fetched(self) -> Result { - match self { - Self::Fetching(fetching) => fetching.await, - Self::Fetched(result) => result, - } - } } /// The group a FETCH answers with, read out of the cache so a later eviction cannot @@ -339,105 +363,223 @@ impl FetchedGroup { } } -/// Read group `sequence` of `track` for a FETCH, from object `skip` up to the exclusive +/// Reads group `sequence` of a track for a FETCH, from object `skip` up to the exclusive /// `until`, or through the end of the group without one. /// /// This is a [`track::Consumer::fetch_group`], so a relay fetches a miss upstream. The /// track's info resolves first, as it would for a SUBSCRIBE; `ahead` is the whole group /// already asked for, used for a miss. -async fn read_fetch( - track: &track::Consumer, +struct ReadFetch { sequence: u64, skip: u64, until: Option, priority: i8, compression: bool, + /// Held until the read ends, even when the read asks for its own part of the group. ahead: Option, -) -> Result { - let info = track.info().await?; - let cached = track.peek_group(sequence); - let wire = cached.as_ref().and_then(|group| group.wire()); - if wire.as_ref().is_some_and(|wire| !wire.fetchable) { - return Err(Error::Unsupported); - } - // Relayed Object IDs need not match frame ordinals, so a relayed group is read from - // its first frame and filtered by ID. Its marked gaps cover the IDs between where it - // starts and its first object; a start below that misses to the upstream. A miss may - // be relayed too, which only the upstream knows, so it fetches the whole group. - let frame_start = match (&wire, &cached) { - (Some(_), Some(group)) => skip.min(group.first_frame()), - (None, Some(_)) => skip, - (_, None) => 0, - }; - let mut group = match ahead { - Some(ahead) if frame_start == 0 => ahead.fetched().await?, - _ => { - track - .fetch_group( - sequence, - group::Fetch { - priority, - frame_start, - ..Default::default() - }, - ) - .await? + state: ReadFetchState, +} + +enum ReadFetchState { + /// Resolving the track's info. + Info(track::Resolving), + /// Waiting for the group. + Fetch { + info: track::Info, + fetching: kio::Pending, + }, + /// Reading the group's objects in the range. + Read { + info: track::Info, + group: group::Consumer, + /// The Object ID of the first frame read. + first: u64, + frames: Vec, + /// A long cached group is read a slice per poll, so the session's other tasks still run. + budget: kio::coop::Budget, + }, + Done, +} + +impl ReadFetch { + fn new( + track: &track::Consumer, + sequence: u64, + skip: u64, + until: Option, + priority: i8, + compression: bool, + ahead: Option, + ) -> Self { + Self { + sequence, + skip, + until, + priority, + compression, + ahead, + state: ReadFetchState::Info(track.info().into_inner()), } - }; - if compression { - group = group.compressed(); - } - let mut first = skip; - let mut frames = Vec::new(); - // Reading an object past the range proves the range missed the group's end. - let mut overshot = false; - // A long cached group is read a slice per poll, so the session's other tasks still run. - let mut budget = kio::coop::Budget::new(32); - loop { - kio::wait(|waiter| budget.poll_yield(waiter)).await; + } + + fn poll(&mut self, waiter: &kio::Waiter, track: &track::Consumer) -> Poll> { + let res = ready!(self.poll_read(waiter, track)); + self.state = ReadFetchState::Done; + Poll::Ready(res) + } + + fn poll_read(&mut self, waiter: &kio::Waiter, track: &track::Consumer) -> Poll> { + loop { + match &mut self.state { + ReadFetchState::Info(resolving) => { + let info = ready!(resolving.poll_ok(waiter))?; + let cached = track.peek_group(self.sequence); + let wire = cached.as_ref().and_then(|group| group.wire()); + if wire.as_ref().is_some_and(|wire| !wire.fetchable) { + return Poll::Ready(Err(Error::Unsupported)); + } + // Relayed Object IDs need not match frame ordinals, so a relayed group is read + // from its first frame and filtered by ID. Its marked gaps cover the IDs between + // where it starts and its first object; a start below that misses to the + // upstream. A miss may be relayed too, which only the upstream knows, so it + // fetches the whole group. + let frame_start = match (&wire, &cached) { + (Some(_), Some(group)) => self.skip.min(group.first_frame()), + (None, Some(_)) => self.skip, + (_, None) => 0, + }; + let fetching = match self.ahead.take_if(|_| frame_start == 0) { + Some(Ahead::Fetched(result)) => { + self.state = self.reading(info, result?); + continue; + } + Some(Ahead::Fetching(fetching)) => fetching, + None => track.fetch_group( + self.sequence, + group::Fetch { + priority: self.priority, + frame_start, + ..Default::default() + }, + ), + }; + self.state = ReadFetchState::Fetch { info, fetching }; + } + ReadFetchState::Fetch { info, fetching } => { + let group = ready!(kio::Task::poll(&mut **fetching, waiter))?; + let info = info.clone(); + self.state = self.reading(info, group); + } + ReadFetchState::Read { + group, + first, + frames, + budget, + .. + } => { + // Reading an object past the range proves the range missed the group's end. + let mut overshot = false; + loop { + ready!(budget.poll_yield(waiter)); + if group.wire().is_some_and(|wire| !wire.fetchable) { + return Poll::Ready(Err(Error::Unsupported)); + } + let ordinal = group.index(); + let Some(frame) = ready!(group.poll_read_object(waiter))? else { + break; + }; + let id = frame.fields.as_ref().map_or(ordinal, |fields| fields.id); + if self.until.is_some_and(|until| id >= until) { + overshot = true; + break; + } + if id < self.skip { + continue; + } + if frames.is_empty() { + *first = id; + } + frames.push(frame); + if self.until.is_some_and(|until| id.saturating_add(1) >= until) { + break; + } + } + let ReadFetchState::Read { + info, + group, + first, + frames, + .. + } = std::mem::replace(&mut self.state, ReadFetchState::Done) + else { + unreachable!("matched above"); + }; + return Poll::Ready(self.read(info, group, first, frames, overshot)); + } + ReadFetchState::Done => panic!("polled after it resolved"), + } + } + } + + /// Start reading `group`, the one the FETCH answers with. + fn reading(&self, info: track::Info, mut group: group::Consumer) -> ReadFetchState { + if self.compression { + group = group.compressed(); + } + ReadFetchState::Read { + info, + group, + first: self.skip, + frames: Vec::new(), + budget: kio::coop::Budget::new(32), + } + } + + /// What the read found, once it reached the end of its range or of the group. + fn read( + &self, + info: track::Info, + mut group: group::Consumer, + first: u64, + frames: Vec, + overshot: bool, + ) -> Result { + // A bounded range can reach the group's end without reading its terminal None, and a + // start past the end is already there. + let complete = !overshot + && matches!(group.finished().now_or_never(), Some(Ok(_))) + && group.index() >= group.frame_count() as u64; + let end_object = group + .end_object() + .filter(|end| complete && end.id >= self.skip && self.until.is_none_or(|until| end.id < until)); + if group.wire().is_some_and(|wire| !wire.fetchable) { return Err(Error::Unsupported); } - let ordinal = group.index(); - let Some(frame) = kio::wait(|waiter| group.poll_read_object(waiter)).await? else { - break; - }; - let id = frame.fields.as_ref().map_or(ordinal, |fields| fields.id); - if until.is_some_and(|until| id >= until) { - overshot = true; - break; - } - if id < skip { - continue; - } - if frames.is_empty() { - first = id; - } - frames.push(frame); - if until.is_some_and(|until| id.saturating_add(1) >= until) { - break; - } - } - // A bounded range can reach the group's end without reading its terminal None, and a - // start past the end is already there. - let complete = !overshot - && matches!(group.finished().now_or_never(), Some(Ok(_))) - && group.index() >= group.frame_count() as u64; - let end_object = group - .end_object() - .filter(|end| complete && end.id >= skip && until.is_none_or(|until| end.id < until)); - - if group.wire().is_some_and(|wire| !wire.fetchable) { - return Err(Error::Unsupported); - } - Ok(FetchedGroup { - sequence, - first, - frames, - end_object, - complete, - info, - }) + Ok(FetchedGroup { + sequence: self.sequence, + first, + frames, + end_object, + complete, + info, + }) + } +} + +/// Await [`ReadFetch`]. +#[cfg(test)] +async fn read_fetch( + track: &track::Consumer, + sequence: u64, + skip: u64, + until: Option, + priority: i8, + compression: bool, + ahead: Option, +) -> Result { + let mut read = ReadFetch::new(track, sequence, skip, until, priority, compression, ahead); + kio::wait(|waiter| read.poll(waiter, track)).await } /// A broadcast whose route table is watched for changes in what we advertise: the @@ -722,6 +864,7 @@ enum UpdateRead { const TOKEN_ON_SESSION_REQUEST: &str = "AUTHORIZATION TOKEN on a request admitted without one"; /// A subscription's update state, which survives its pauses. +#[derive(Default)] struct SubscriptionUpdates { /// The grant the subscription's own request token earned, if one admitted it, which its /// renewals replace in place. A request admitted without one never names one later. @@ -967,1165 +1110,791 @@ impl Drop for Join { } } -impl Publisher { - pub fn new( - runtime: crate::time::Deadlines, - session: Transport, - origin: origin::Consumer, - control: Control, - peer_hop: Option, - peer_setup: peer::PeerSetup, - version: Version, - ) -> Self { - Self { - withdrawal: Default::default(), - runtime, - session, - self_origin: origin.hop(), - // Serving the peer is not a local reader asking to discover anything. - origin: origin.serving(), - control, - peer_hop, - peer_setup, - joins: Default::default(), - version, - auth: crate::auth::Link::new(false), - owed: Default::default(), - subscriptions: Default::default(), - priority: Default::default(), - } - } +/// Serves a SUBSCRIBE on its bidi stream, from the checks on its first poll to the last +/// message the request carries. +struct SubscribeServe { + request: SubscribeRequest, + state: SubscribeState, +} - /// Mark a paused subscription for joining FETCHes: draft-18 section 5.1 refuses one - /// unless the subscription forwards. - fn pause_join(&self, join: Option<&Join>>) { - if let Some(join) = join { - join.joins.lock().insert(join.request_id, Some(Joined::Empty)); +/// What a SUBSCRIBE holds whichever step it is at. +struct SubscribeRequest { + publisher: Publisher, + /// Taken to end the request with PUBLISH_DONE, which closes it. + stream: Option>, + /// Its serving parameters follow the updates applied while paused. + msg: ietf::Subscribe<'static>, + /// Taken by a [`Retain`] that keeps the join. + gate: Option, + /// Registered at dispatch, before either request task can be polled. + join: Option>>, + /// Held for the life of the subscription, and by every token verification the + /// application takes for it, initial or renewal, even past a cancel. The request keeps + /// its stream, and so its stream credit or request window slot, until all of them let + /// go, so a peer cancelling and resubscribing cannot leave more verifications in flight + /// than its credit admits. + credit: kio::Producer<()>, + /// The subscription's slot, held with the credit once the request passes its checks. + slot: Option, + request_finished: bool, + goaways: super::request_stream::Messages, + /// The grant the request's own token earned, if one admitted it. + grant: Option, + updates: SubscriptionUpdates, + /// The group stream count once SUBSCRIBE_OK is buffered, which decides how a request + /// whose authority ends is told. + acked: Option>, + /// Set once the track ran to its end, which is all that leaves a join worth keeping. + track_ended: bool, + /// Set once the request token lapses, just before the subscription's streams drop. + lapse: Lapse, +} + +enum SubscribeState { + /// Not polled yet: the request's checks run on its first poll. + Start, + /// Refused: delivering the REQUEST_ERROR already buffered. + Refuse(super::request_stream::Deliver), + /// Verifying the request's AUTHORIZATION TOKEN. + Verify(crate::auth::RequestVerdict), + /// Withdrawn before an answer: waiting for its verifications to let go of the credit. + Release, + /// Admitted, serving under its authority. + Serve(Box), + /// A refused renewal ended the request: flushing its REQUEST_ERROR before PUBLISH_DONE. + RenewalRefused { + error: Error, + streams: u64, + retain: Option, + }, + /// Ending the subscription with PUBLISH_DONE. + PublishDone(WritePublishDone), + Done, +} + +impl SubscribeServe { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + let request = &mut self.request; + loop { + match &mut self.state { + SubscribeState::Start => self.state = request.start()?, + SubscribeState::Refuse(deliver) => { + let stream = request.stream.as_mut().expect("open until PUBLISH_DONE"); + return deliver.poll(waiter, &mut stream.writer); + } + SubscribeState::Verify(verdict) => { + let stream = request.stream.as_mut().expect("open until PUBLISH_DONE"); + let fin_cancels = super::request_stream::fin_cancels(request.publisher.version); + // The subscriber cancelling or the session closing withdraws the token. + let verified = ready!(request.publisher.poll_until_withdrawn( + stream, + fin_cancels, + waiter, + |waiter| verdict.poll_grant(waiter) + )); + let SubscribeState::Verify(verdict) = std::mem::replace(&mut self.state, SubscribeState::Release) + else { + unreachable!("matched above"); + }; + let Some(verified) = verified else { + drop((verdict, request.slot.take())); + continue; + }; + let admitted = crate::auth::RequestGrant::new( + &request.publisher.runtime, + verdict, + verified, + request.msg.track_namespace.clone(), + crate::auth::RequestKind::Subscribe, + request.slot.clone(), + ); + self.state = match admitted { + Ok(grant) => { + request.gate = request.gate.take().map(crate::auth::Gate::named); + request.grant = Some(grant); + request.admitted() + } + Err(refusal) => request.refuse(&refusal.error, &refusal.reason)?, + }; + } + SubscribeState::Release => { + ready!(request.publisher.poll_released(waiter, &request.credit)); + self.state = SubscribeState::Done; + return Poll::Ready(Ok(())); + } + SubscribeState::Serve(serving) => { + // Polled first, so the request never acts on authority that already ended. + let subscribed = match request.poll_ended(waiter) { + Poll::Ready(refusal) => Err(refusal), + Poll::Pending => Ok(ready!(serving.poll(waiter, request))), + }; + // Dropping the serve releases what it holds upstream; only then is the peer told. + self.state = SubscribeState::Done; + self.state = request.served(subscribed)?; + } + SubscribeState::RenewalRefused { .. } => { + let stream = request.stream.as_mut().expect("open until PUBLISH_DONE"); + let flushed = ready!(stream.writer.poll_flush(waiter)); + let SubscribeState::RenewalRefused { error, streams, retain } = + std::mem::replace(&mut self.state, SubscribeState::Done) + else { + unreachable!("matched above"); + }; + let served = match flushed { + Ok(()) => Served::UpdateFailed(error), + Err(err) => Served::Track(Err(err)), + }; + self.state = request.publish_done(served, streams, retain); + } + SubscribeState::PublishDone(done) => { + let res = ready!(done.poll(waiter, &request.publisher, &request.credit)); + self.state = SubscribeState::Done; + return Poll::Ready(res); + } + SubscribeState::Done => panic!("polled after it resolved"), + } } } +} - /// Bound what we publish by the grant this session's tokens earn (MoQ Auth). - pub fn with_auth(mut self, auth: crate::auth::Link) -> Self { - self.auth = auth; - self +impl SubscribeRequest { + /// Check the request before admitting it, refusing what it cannot have. + fn start(&mut self) -> Result { + let publisher = &self.publisher; + let msg = &self.msg; + if let Some(error) = super::namespace::refusal(&msg.track_namespace, publisher.version) { + let stream = self.stream.as_mut().expect("open until PUBLISH_DONE"); + stream.writer.buffer_message(ietf::RequestError::ID, &error)?; + return Ok(SubscribeState::Refuse(super::request_stream::Deliver::new())); + } + // Logged beside the root rather than joined, so a namespace too long under it is + // refused by the lookup, on this request alone. + tracing::debug!(id = %msg.request_id, broadcast = %self.broadcast(), track = %msg.track_name, "subscribe started"); + + // Serve only what our grant lets us publish (MoQ Auth), and stop once it no + // longer does. Checked before resolving, so a denied request never reaches the + // origin. Our own authorization of the peer (the limit) admits a request without a + // token; an AUTHORIZATION TOKEN on the request goes to the application to verify for + // this request alone (MoQ request-token), even where the limit covers it. What the + // peer granted us bounds what we send either way, and never stands in for our + // authorization of it. + let gate = self.gate.as_ref().expect("held until the subscription ends"); + let token = match &msg.authorization_token { + Some(token) + if publisher + .auth + .within_grant(crate::auth::Direction::Publish, msg.track_namespace.as_str()) + && gate.token_decides() => + { + Some(token.0.clone()) + } + _ => None, + }; + if token.is_none() && !gate.allows() { + return self.refuse(&Error::Unauthorized, "not granted"); + } + // Legal requests we can't honor are refused one at a time, never by closing the + // session. Serving a Range Filter unfiltered would deliver objects the subscriber + // excluded. + if msg.range_filters { + let stream = self.stream.as_mut().expect("open until PUBLISH_DONE"); + stream.writer.buffer_message( + ietf::RequestError::ID, + &ietf::RequestError { + error_code: request::INVALID_FILTER, + reason_phrase: "Range Filters exceed MAX_FILTER_RANGES=0".into(), + retry_interval: 0, + }, + )?; + return Ok(SubscribeState::Refuse(super::request_stream::Deliver::new())); + } + let slot: crate::auth::Hold = Arc::new((publisher.subscriptions.hold(), self.credit.consume())); + self.slot = Some(slot.clone()); + // A token-authorized subscription is held to its own grant, which ends it when the + // grant lapses, is revoked, or stops covering it. Every subscription is gated so it + // ends once what the peer granted us stops covering it, and a limit-authorized one + // once the limit does too. + Ok(match token { + Some(token) => SubscribeState::Verify(gate.verify(token, Some(slot))), + None => self.admitted(), + }) } - /// What the peer declared in its SETUP. - /// - /// Blocks until the peer's SETUP arrives, because the extension changes the NAMESPACE - /// encoding: nothing can be advertised until we know whether the peer speaks it. - fn poll_peer(&self, waiter: &kio::Waiter) -> Poll { - self.peer_setup.poll_get(waiter).map(|peer| peer.cluster) + /// The broadcast the request names, for its logs. + fn broadcast(&self) -> crate::path::Rooted<'_> { + crate::path::Rooted(self.publisher.origin.root(), &self.msg.track_namespace) } - /// Await [`Self::poll_peer`]. - async fn peer(&self) -> cluster::Peer { - kio::wait(|waiter| self.poll_peer(waiter)).await + /// Serve the admitted request. + fn admitted(&mut self) -> SubscribeState { + self.updates = SubscriptionUpdates { + grant: self.grant.clone(), + renewal: Default::default(), + }; + SubscribeState::Serve(Box::new(Serving::new(&self.publisher, &self.msg))) } - /// Whether the peer requires advertisements to be solicited, from the same SETUP. - /// - /// Blocks on it for the same reason [`Self::peer`] does: this decides whether the - /// first advertisement is sent unasked, so it cannot be guessed and corrected later. - async fn requires_solicitation(&self) -> bool { - self.peer_setup.get().await.solicit.unwrap_or(false) + /// Refuse the request with `err`, delivering its REQUEST_ERROR next. + fn refuse(&mut self, err: &Error, reason: &str) -> Result { + let stream = self.stream.as_mut().expect("open until PUBLISH_DONE"); + let deliver = self.publisher.reject_subscribe(&mut stream.writer, err, reason)?; + Ok(SubscribeState::Refuse(deliver)) } - /// The origin to serve this peer's subscriptions from: sources whose hop chain flows - /// through the peer are excluded, so a subscription is never handed data that already - /// flowed through the subscriber. - /// - /// The same exclusion the announce path applies (see [`Self::select`]), which is what - /// keeps advertised paths truthful and prevents subscription cycles of any length. - pub(super) fn poll_serving_origin(&self, waiter: &kio::Waiter) -> Poll { - self.poll_peer(waiter).map(|peer| self.excluding(&peer)) + /// Ready once the subscription's authority ends: its request token's grant, or the gate + /// on what the peer granted us and on the limit. + fn poll_ended(&mut self, waiter: &kio::Waiter) -> Poll { + if let Some(grant) = &self.grant + && let Poll::Ready(refusal) = grant.poll_ended(waiter) + { + // PUBLISH_DONE carries a fixed reason per status, so only the code reaches + // the peer once SUBSCRIBE_OK went out. + tracing::info!(broadcast = %self.broadcast(), track = %self.msg.track_name, err = %refusal.error, reason = %refusal.reason, "request token grant ended"); + if super::error::token_expired(&refusal.error) { + self.lapse.set(); + } + return Poll::Ready(refusal); + } + let gate = self.gate.as_mut().expect("held until the subscription ends"); + if gate.poll_denied(waiter).is_ready() { + tracing::info!(broadcast = %self.broadcast(), track = %self.msg.track_name, "subscription no longer authorized"); + return Poll::Ready(Error::Unauthorized.into()); + } + Poll::Pending } - /// Resolve a request's broadcast while the peer still wants the answer, a FIN withdrawing it - /// where `fin_cancels`. - async fn resolve( - &self, - stream: &mut Stream, - fin_cancels: bool, - path: &crate::Path, - ) -> Result, Error> { - let mut resolve = Resolve::new(self.session.clone(), path.clone(), fin_cancels); - kio::wait(|waiter| resolve.poll(waiter, self, stream)).await - } + /// How the request ends once serving stopped, or why its authority ended first. + fn served( + &mut self, + subscribed: std::result::Result, crate::auth::Refusal>, + ) -> Result { + // No update follows the subscription's end, so its pending renewals no longer matter. + drop(std::mem::take(&mut self.updates)); + // A track that ended cleanly keeps what a joining FETCH needs (see [`Retain`]); any + // other end drops the join, refusing a joining FETCH that has not answered yet. + let retain = match (&subscribed, self.join.take()) { + (Ok(Ok(Subscribed::Served(Served::Track(Ok(())), _))), Some(join)) + if self.track_ended && join.answers() => + { + Some(Retain { + _join: join, + grant: self.grant.take(), + gate: self.gate.take().expect("held until the subscription ends"), + _slot: self.slot.take().expect("held since admission"), + }) + } + _ => { + // The subscription's demand is gone; only its credit may still wait on a + // verification. + drop((self.grant.take(), self.slot.take())); + None + } + }; - /// Our origin handle with [`Self::exclude`] applied, the view both the data plane - /// and the announce loops read this peer's routes through. - fn excluding(&self, peer: &cluster::Peer) -> origin::Consumer { - self.origin.clone().excluding(self.exclude(peer)) + Ok(match subscribed { + Ok(Ok(Subscribed::Served(Served::RenewalRefused(refusal), streams))) => { + let stream = self.stream.as_mut().expect("open until PUBLISH_DONE"); + match self + .publisher + .buffer_subscribe_error(&mut stream.writer, &refusal.error, &refusal.reason) + { + Ok(()) => SubscribeState::RenewalRefused { + error: refusal.error, + streams, + retain, + }, + Err(err) => self.publish_done(Served::Track(Err(err)), streams, retain), + } + } + Ok(Ok(Subscribed::Served(served, streams))) => self.publish_done(served, streams, retain), + Ok(Ok(Subscribed::Refused(reason, err))) => self.refuse(&err, &reason)?, + // Nobody is left to answer. + Ok(Ok(Subscribed::Withdrawn)) => SubscribeState::Release, + Ok(Err(err)) => return Err(err), + Err(refusal) => match &self.acked { + Some(opened) => { + let streams = opened.load(Ordering::Relaxed); + self.publish_done(Served::Track(Err(refusal.error)), streams, retain) + } + None => self.refuse(&refusal.error, &refusal.reason)?, + }, + }) } - /// `origin` as advertised to this peer: never its own routes back (split horizon), - /// and under MoQ Cluster only the routes with room for our hop under the origin's - /// cap, so selection falls back to a shorter route instead of withdrawing. - fn advertising(&self, origin: origin::Consumer, peer: &cluster::Peer) -> origin::Consumer { - let origin = origin.excluding(self.exclude(peer)); - match peer.negotiated() { - true => origin.forwardable(), - false => origin, - } + /// End the subscription with the PUBLISH_DONE `served` reports, after `streams` group + /// streams. + fn publish_done(&mut self, served: Served, streams: u64, retain: Option) -> SubscribeState { + let stream = self.stream.take().expect("open until PUBLISH_DONE"); + SubscribeState::PublishDone(WritePublishDone::new(stream, served, streams, retain)) } - /// The Hop ID whose paths must not be advertised (or served) back to this peer. - /// - /// A peer that declared an identity supplies its own; otherwise fall back to the one - /// we assigned it (a fresh per-session id, or one pinned with `Client::with_peer_hop` - /// or `Request::with_peer_hop`), since moq-transport carries no identity - /// of its own. A peer that declared the reserved 0 declared no identity, so it takes - /// the fallback like any other anonymous peer. - fn exclude(&self, peer: &cluster::Peer) -> crate::Hop { - peer.identity().or(self.peer_hop).unwrap_or(crate::Hop::UNKNOWN) + /// Mark the subscription for joining FETCHes from the boundary it serves from: the + /// prefix ending right after `largest`. + fn save_join( + &self, + filter: Filter, + largest: Option, + cache: &track::Consumer, + ranking: &Ranking>, + ) { + let Some(join) = &self.join else { + return; + }; + let joined = match (filter, largest) { + (Filter::NextObject, Some(largest)) => Joined::Group { + end: Location { + group: largest.group, + object: largest.object + 1, + }, + cache: cache.clone(), + namespace: self.msg.track_namespace.clone(), + grant: self.grant.clone(), + ranking: ranking.clone(), + }, + (Filter::NextObject, None) => Joined::Empty, + _ => Joined::Unsupported, + }; + join.joins.lock().insert(self.msg.request_id, Some(joined)); } +} - /// Pick what to advertise to this peer for one route. +/// An admitted SUBSCRIBE's serve, which ends however its authority does: dropping it +/// releases everything it holds upstream. +struct Serving { + /// The filter the SUBSCRIBE asked for, which every saved join boundary follows. /// - /// The origin's announce cursor already filters routes through the peer - /// (control-plane split horizon); this only shapes what the wire can carry. - fn select(&self, route: &crate::origin::Route, peer: &cluster::Peer) -> Advert { - // A route that already passed through us is a reflection. The origin - // filters these on receive, so this is defensive. - if self.self_origin != crate::Hop::UNKNOWN && route.hops.contains(&self.self_origin) { - return Advert::None; - } + /// Draft-22 replaced joining FETCH with subscription fills. Older drafts save the same + /// boundary used by the subscription so the two streams never overlap, again on every + /// resume. + filter: Filter, + broadcast: Option, + /// The subscription being served, `None` while paused. + track: Option, + step: ServingStep, +} - if !peer.negotiated() { - return Advert::Some(ietf::AdvertParams { cluster: None }); - } +#[allow( + clippy::large_enum_variant, + reason = "a state machine's enum is its storage: one per subscription, boxed in its request; an allow, since the size gap varies by target" +)] +enum ServingStep { + /// Resolving the broadcast. + Resolve(Resolve), + /// Subscribing to the track. + Subscribe { + track: track::Consumer, + subscribe: SubscribeTrack, + }, + /// Waiting on the peer's SETUP for whether SUBSCRIBE_OK's properties may compress. + Properties { + /// A fill reads the group cache through its own consumer, independent of the + /// subscription's cursor. + cache: track::Consumer, + edge: LiveEdge, + }, + /// Flushing SUBSCRIBE_OK, then starting the fill it promised, if any. + SubscribeOk(Option), + /// Forwarding. A resumed subscription holds the track it resubscribed to. + Forward { + run: RunSubscription, + _track: Option, + }, + /// Flushing the REQUEST_OK answering the update that paused the subscription. + Pausing, + /// Paused, answering updates until one resumes it. + Paused(WaitResume), + /// Subscribing again for the update resuming it. + Resume { + track: track::Consumer, + subscribe: SubscribeTrack, + }, + /// Flushing the error that fails the resuming update. + ResumeFailed(Error), + /// Flushing the REQUEST_OK answering the resuming update. + Resumed(track::Consumer), + /// Sending the END_OF_TRACK a subscription that ran to the track's end still owes: + /// PUBLISH_DONE carries no end location, so the marker is what tells the subscriber + /// where the track ended. One cancelled while the marker waits for stream credit + /// abandons it. + EndOfTrack { + served: Served, + marker: EndOfTrack, + }, + Done, +} - let cost = route.cost.value(); - // Our own Hop ID is always the last entry, so the peer reconstructs the full - // path. A chain with no room left is a loop in all but name. - match cluster::Advert::forward(&route.hops, cost, self.self_origin) { - Ok(advert) => Advert::Some(ietf::AdvertParams { cluster: Some(advert) }), - Err(_) => Advert::None, +impl Serving { + fn new(publisher: &Publisher, msg: &ietf::Subscribe<'_>) -> Self { + let fin_cancels = super::request_stream::fin_cancels(publisher.version); + Self { + filter: msg.filter, + broadcast: None, + track: None, + // We just received a subscribe for this exact namespace, so the peer must have + // already seen the announcement. The immediate lookup resolves it at once, or + // falls back to the route covering it (an `origin::Requests`), if any. + step: ServingStep::Resolve(Resolve::new( + publisher.session.clone(), + msg.track_namespace.clone(), + fin_cancels, + )), } } - /// Handle an incoming bidi stream dispatched by the session. - pub fn handle_stream( - &self, - id: u64, - body: ietf::Body, - stream: Stream, - ) -> Result, Error> { - let this = self.clone(); - let mut data = body.decoder(this.version); - // Count at dispatch so close cannot pass a request waiting for its first poll. - let serve = matches!(id, ietf::Subscribe::ID | ietf::Fetch::ID).then(|| { - self.owed.fetch_add(1, Ordering::Relaxed); - Serve { - owed: self.owed.clone(), - _working: self.session.work().start(), - } - }); - let task = match id { - ietf::Subscribe::ID => { - let msg = match ietf::Subscribe::decode_msg(&mut data, this.version) { - Err(err) if super::token::refuses_request(&err) => { - // Junk after the fields is fatal, which a token refusal must not hide. - if !data.is_empty() { - return Err(Error::WrongSize); + /// Stats (subscriptions, viewer refcount, groups/frames/bytes) are counted in the model, + /// through the tagged `origin::Consumer` the broadcast resolves from. + fn poll(&mut self, waiter: &kio::Waiter, request: &mut SubscribeRequest) -> Poll> { + let refused = |err: Error, rendezvous, resolved| { + Subscribed::Refused(err.to_string(), rendezvous_error(err, rendezvous, resolved)) + }; + let version = request.publisher.version; + loop { + let stream = request.stream.as_mut().expect("open until PUBLISH_DONE"); + match &mut self.step { + ServingStep::Resolve(resolve) => { + let broadcast = match ready!(resolve.poll(waiter, &request.publisher, stream)) { + Ok(Some(broadcast)) => broadcast, + // The reason is the origin's, not ours: a dynamic router refusing on + // authorization or a drain says so, and only an unroutable path means the + // broadcast is not here. + Err(err) => return Poll::Ready(Ok(refused(err, request.msg.rendezvous, false))), + Ok(None) => return Poll::Ready(Ok(Subscribed::Withdrawn)), + }; + let track = match broadcast.track(&request.msg.track_name) { + Ok(track) => track, + Err(err) => return Poll::Ready(Ok(refused(err, request.msg.rendezvous, true))), + }; + self.broadcast = Some(broadcast); + self.step = ServingStep::Subscribe { + track, + subscribe: SubscribeTrack::new(), + }; + } + ServingStep::Subscribe { track, subscribe } => { + let subscribed = + ready!(subscribe.poll(waiter, &request.publisher, stream, track, &request.msg, &request.lapse)); + let (serve, edge) = match subscribed { + Some(Ok(subscribed)) => subscribed, + Some(Err(err)) => return Poll::Ready(Ok(refused(err, request.msg.rendezvous, true))), + None => return Poll::Ready(Ok(Subscribed::Withdrawn)), + }; + let ServingStep::Subscribe { track: cache, .. } = + std::mem::replace(&mut self.step, ServingStep::Done) + else { + unreachable!("matched above"); + }; + self.track = Some(serve); + self.step = ServingStep::Properties { cache, edge }; + } + ServingStep::Properties { edge, .. } => { + let flate = ready!(request.publisher.peer_setup.poll_get(waiter)).flate; + let edge = *edge; + let ServingStep::Properties { cache, .. } = std::mem::replace(&mut self.step, ServingStep::Done) + else { + unreachable!("matched above"); + }; + let serve = self.track.as_mut().expect("subscribed above"); + let properties = + match track_properties(serve.track.info(), request.msg.properties_wanted, version, flate) { + Ok(properties) => properties, + Err(err) => return Poll::Ready(Ok(Subscribed::Refused(err.to_string(), err))), + }; + request.save_join(self.filter, edge.largest, &cache, &serve.ranking); + + // The fill (when one was requested) runs alongside on its own fetch stream; + // its failures reset that stream and never touch the subscription. Every + // subscription a resume makes counts toward the same PUBLISH_DONE. + let fill = request + .msg + .fill + .filter(|_| request.msg.forward && Filter::is_draft20(version)) + .map(|fill| Fill { + run: RunFill::new( + fill_range(fill, request.msg.filter, edge.largest), + cache, + // The fill stamps its objects as the subscription does. + serve.timescale, + serve.ranking.request_id, + ), + ranking: serve.ranking.clone(), + opened: serve.opened.clone(), + }); + + // Buffered whole before its first wait, so however the subscription is + // dropped from here, the request can still end on the wire. + request.acked = Some(serve.opened.clone()); + let stream = request.stream.as_mut().expect("open until PUBLISH_DONE"); + stream.writer.buffer_message( + ietf::SubscribeOk::ID, + &ietf::SubscribeOk { + track_alias: request.msg.request_id.0, + // The subscription floor and FETCH/FILL cap use this same snapshot. + largest: edge.largest, + properties, + }, + )?; + self.step = ServingStep::SubscribeOk(fill); + } + ServingStep::SubscribeOk(fill) => { + ready!(stream.writer.poll_flush(waiter))?; + let fill = fill.take(); + match request.msg.forward { + true => { + self.step = ServingStep::Forward { + run: RunSubscription::new(fill), + _track: None, + } + } + // Unstarted, the fill would still hold the track it serves from. + false => { + drop(fill); + if let Some(subscribed) = self.forwarded(request, Some(Served::Paused(false)))? { + return Poll::Ready(Ok(subscribed)); + } } - return Ok(this.reject_token(stream, request::Kind::Subscribe, err, serve)); } - msg => msg?, - }; - if !data.is_empty() { - return Err(Error::WrongSize); } - tracing::debug!(message = ?msg, "received subscribe"); - let gate = crate::auth::Gate::new( - this.auth.clone(), - msg.track_namespace.clone(), - crate::auth::Direction::Publish, - ); - let task = this.serve_subscribe_stream(stream, msg.into_owned(), gate); - async move { - let _serve = serve; - if let Err(err) = task.await { - tracing::debug!(%err, "subscribe stream error"); + ServingStep::Forward { run, .. } => { + let serve = self.track.as_mut().expect("forwarding"); + let served = ready!(run.poll( + waiter, + &request.publisher, + stream, + serve, + &mut request.request_finished, + &mut request.goaways, + &mut request.updates, + )); + self.step = ServingStep::Done; + if let Some(subscribed) = self.forwarded(request, served)? { + return Poll::Ready(Ok(subscribed)); } } - .boxed() - } - ietf::Fetch::ID => { - let msg = match ietf::Fetch::decode_msg(&mut data, this.version) { - Err(err) if super::token::refuses_request(&err) => { - // Junk after the fields is fatal, which a token refusal must not hide. - if !data.is_empty() { - return Err(Error::WrongSize); - } - return Ok(this.reject_token(stream, request::Kind::Fetch, err, serve)); - } - msg => msg?, - }; - if !data.is_empty() { - return Err(Error::WrongSize); + ServingStep::Pausing => { + ready!(stream.writer.poll_flush(waiter))?; + self.step = ServingStep::Paused(WaitResume::new()); } - tracing::debug!(message = ?msg, "received fetch"); - let gate = this.fetch_gate(&msg); - // Owned, since the task outlives the body it was read from. - let msg = msg.into_owned(); - async move { - let _serve = serve; - if let Err(err) = this.serve_fetch_stream(stream, msg, gate).await { - tracing::debug!(%err, "fetch stream error"); + ServingStep::Paused(paused) => { + let resumed = ready!(paused.poll( + waiter, + &request.publisher, + stream, + &mut request.request_finished, + &mut request.goaways, + &mut request.msg, + &mut request.updates, + )); + if let Err(served) = resumed { + let served = served.unwrap_or(Served::Track(Ok(()))); + return Poll::Ready(Ok(Subscribed::Served(served, request.opened()))); } + let broadcast = self.broadcast.as_ref().expect("resolved before serving"); + self.step = match broadcast.track(&request.msg.track_name) { + Ok(track) => ServingStep::Resume { + track, + subscribe: SubscribeTrack::new(), + }, + Err(err) => Self::resume_failed(request, err)?, + }; } - .boxed() - } - ietf::SubscribeNamespace::ID => { - let msg = match ietf::SubscribeNamespace::decode_msg(&mut data, this.version) { - Err(err) if super::token::refuses_request(&err) => { - // Junk after the fields is fatal, which a token refusal must not hide. - if !data.is_empty() { - return Err(Error::WrongSize); + ServingStep::Resume { track, subscribe } => { + let subscribed = + ready!(subscribe.poll(waiter, &request.publisher, stream, track, &request.msg, &request.lapse)); + let ServingStep::Resume { track, .. } = std::mem::replace(&mut self.step, ServingStep::Done) else { + unreachable!("matched above"); + }; + self.step = match subscribed { + // Withdrawn mid-setup, it ends as it would have while paused. + None => { + return Poll::Ready(Ok(Subscribed::Served(Served::Track(Ok(())), request.opened()))); } - let kind = request::Kind::SubscribeNamespace; - return Ok(this.reject_token(stream, kind, err, serve)); - } - decoded => decoded?, - }; - if !data.is_empty() { - return Err(Error::WrongSize); + // The resuming update is owed an answer, so it fails. + Some(Err(err)) => Self::resume_failed(request, err)?, + Some(Ok((mut serve, edge))) => { + serve.opened = request.acked.clone().expect("acked before serving"); + request.save_join(self.filter, edge.largest, &track, &serve.ranking); + let stream = request.stream.as_mut().expect("open until PUBLISH_DONE"); + Publisher::buffer_request_ok(&mut stream.writer, edge.largest)?; + self.track = Some(serve); + ServingStep::Resumed(track) + } + }; } - tracing::debug!(message = ?msg, "received subscribe_namespace"); - let gate = crate::auth::Gate::new( - this.auth.clone(), - msg.namespace.clone(), - crate::auth::Direction::Publish, - ); - async move { - let session = this.session.clone(); - match this.serve_subscribe_namespace_stream(stream, msg, gate).await { - // Such as a malformed answer to an advertisement this request solicited. - Err(err) if super::subscriber::is_protocol_violation(&err) => session.close( - super::error::to_session_code(&crate::SessionError::from(&err)), - &err.to_string(), - ), - Err(err) => tracing::debug!(%err, "subscribe_namespace stream error"), - Ok(()) => {} - } + ServingStep::ResumeFailed(_) => { + ready!(stream.writer.poll_flush(waiter))?; + let ServingStep::ResumeFailed(err) = std::mem::replace(&mut self.step, ServingStep::Done) else { + unreachable!("matched above"); + }; + return Poll::Ready(Ok(Subscribed::Served(Served::UpdateFailed(err), request.opened()))); } - .boxed() - } - ietf::SubscribeTracks::ID => { - let msg = match ietf::SubscribeTracks::decode_msg(&mut data, this.version) { - Err(err) if super::token::refuses_request(&err) => { - // Junk after the fields is fatal, which a token refusal must not hide. - if !data.is_empty() { - return Err(Error::WrongSize); + ServingStep::Resumed(_) => { + ready!(stream.writer.poll_flush(waiter))?; + let ServingStep::Resumed(track) = std::mem::replace(&mut self.step, ServingStep::Done) else { + unreachable!("matched above"); + }; + self.step = ServingStep::Forward { + run: RunSubscription::new(None), + _track: Some(track), + }; + } + ServingStep::EndOfTrack { marker, .. } => { + let serve = self.track.as_mut().expect("served to its end"); + let written = match marker.poll(waiter, serve) { + Poll::Ready(res) => res, + Poll::Pending => { + let fin_cancels = super::request_stream::fin_cancels(version); + let cancelled = super::request_stream::poll_cancel( + stream, + &mut request.request_finished, + &mut request.goaways, + fin_cancels, + waiter, + ) + .is_ready() || request.publisher.session.poll_closed(waiter).is_ready(); + if !cancelled { + return Poll::Pending; + } + Err(Error::Cancel) } - let kind = request::Kind::SubscribeNamespace; - return Ok(this.reject_token(stream, kind, err, serve)); + }; + // A failure only costs the subscriber the early boundary. + if let Err(err) = written { + tracing::debug!(%err, id = %serve.request_id, "end of track failed"); } - msg => msg?, - }; - if !data.is_empty() { - return Err(Error::WrongSize); + let ServingStep::EndOfTrack { served, .. } = std::mem::replace(&mut self.step, ServingStep::Done) + else { + unreachable!("matched above"); + }; + return Poll::Ready(Ok(self.ended(served))); } - tracing::debug!(message = ?msg, "received subscribe_tracks"); - // We list no tracks by PUBLISH, so each listing is refused on its own and the - // session stays up. Decoded first, so a refusal still admits its Request ID. - let error = - super::namespace::refusal(&msg.prefix, this.version).unwrap_or_else(|| ietf::RequestError { - error_code: request::NOT_SUPPORTED, - reason_phrase: "SUBSCRIBE_TRACKS is not supported".into(), - retry_interval: 0, - }); - async move { - if let Err(err) = super::request_stream::reject(stream, error).await { - tracing::debug!(%err, "subscribe_tracks refusal failed"); - } - } - .boxed() - } - ietf::TrackStatus::ID => { - let msg = match ietf::TrackStatus::decode_msg(&mut data, this.version) { - Err(err) if super::token::refuses_request(&err) => { - // Junk after the fields is fatal, which a token refusal must not hide. - if !data.is_empty() { - return Err(Error::WrongSize); - } - return Ok(this.reject_token(stream, request::Kind::TrackStatus, err, serve)); - } - msg => msg?, - }; - if !data.is_empty() { - return Err(Error::WrongSize); - } - tracing::debug!(message = ?msg, "received track_status"); - let gate = crate::auth::Gate::new( - this.auth.clone(), - msg.track_namespace.clone(), - crate::auth::Direction::Publish, - ); - // Owned, since the task outlives the body it was read from. - let msg = msg.into_owned(); - async move { - if let Err(err) = this.serve_track_status_stream(stream, msg, gate).await { - tracing::debug!(%err, "track_status stream error"); - } - } - .boxed() - } - _ => { - tracing::warn!(id, "unexpected bidi stream type for publisher"); - return Err(Error::UnexpectedStream); - } - }; - Ok(task) - } - - /// Refuse a request whose AUTHORIZATION TOKEN repeats (section 8.9), which unlike any - /// other decode error leaves the session open. - fn reject_token( - &self, - stream: Stream, - kind: request::Kind, - err: DecodeError, - serve: Option, - ) -> BoxFuture<'static, ()> { - let version = self.version; - async move { - let _serve = serve; - if let Err(err) = super::token::reject(stream, kind, err, version).await { - tracing::debug!(%err, "token refusal failed"); + ServingStep::Done => panic!("polled after it resolved"), } } - .boxed() } - /// Handle a SUBSCRIBE on its bidi stream. - #[cfg(test)] - fn run_subscribe_stream( - self, - stream: Stream, - msg: ietf::Subscribe<'static>, - ) -> impl std::future::Future> { - let gate = crate::auth::Gate::new( - self.auth.clone(), - msg.track_namespace.clone(), - crate::auth::Direction::Publish, - ); - self.serve_subscribe_stream(stream, msg, gate) - } - - fn serve_subscribe_stream( - self, - mut stream: Stream, - mut msg: ietf::Subscribe<'static>, - gate: crate::auth::Gate, - ) -> impl std::future::Future> { - // Register during dispatch, before either request task can be polled. - let join = (!Filter::is_draft20(self.version)).then(|| self.register_join(msg.request_id)); - async move { - let _join = join; - if let Some(error) = super::namespace::refusal(&msg.track_namespace, self.version) { - return super::request_stream::reject(stream, error).await; - } - let request_id = msg.request_id; - let track_name = msg.track_name.clone(); - // Logged beside the root rather than joined, so a namespace too long under it is - // refused by the lookup, on this request alone. - let (root, namespace) = (self.origin.root().clone(), msg.track_namespace.clone()); - let broadcast = || crate::path::Rooted(&root, &namespace); - tracing::debug!(id = %request_id, broadcast = %broadcast(), track = %track_name, "subscribe started"); - - // Serve only what our grant lets us publish (MoQ Auth), and stop once it no - // longer does. Checked before resolving, so a denied request never reaches the - // origin. Our own authorization of the peer (the limit) admits a request without a - // token; an AUTHORIZATION TOKEN on the request goes to the application to verify for - // this request alone (MoQ request-token), even where the limit covers it. What the - // peer granted us bounds what we send either way, and never stands in for our - // authorization of it. - let token = match &msg.authorization_token { - Some(token) - if self - .auth - .within_grant(crate::auth::Direction::Publish, msg.track_namespace.as_str()) - && gate.token_decides() => - { - Some(token.0.clone()) - } - _ => None, - }; - if token.is_none() && !gate.allows() { - let err = Error::Unauthorized; - return self.reject_subscribe(stream, &err, "not granted").await; - } - // Legal requests we can't honor are refused one at a time, never by closing the - // session. Serving a Range Filter unfiltered would deliver objects the subscriber - // excluded. - if msg.range_filters { - return super::request_stream::reject( - stream, - ietf::RequestError { - error_code: request::INVALID_FILTER, - reason_phrase: "Range Filters exceed MAX_FILTER_RANGES=0".into(), - retry_interval: 0, - }, - ) - .await; - } - // Held for the life of the subscription, and by every token verification the - // application takes for it, initial or renewal, even past a cancel. The request keeps - // its stream, and so its stream credit or request window slot, until all of them let - // go, so a peer cancelling and resubscribing cannot leave more verifications in flight - // than its credit admits. - let credit = kio::Producer::<()>::default(); - let slot: crate::auth::Hold = std::sync::Arc::new((self.subscriptions.hold(), credit.consume())); - - let mut request_finished = false; - let mut goaways = super::request_stream::Messages::new(self.session.clone(), self.control.client()); - // A token-authorized subscription is held to its own grant, which ends it when the - // grant lapses, is revoked, or stops covering it. Every subscription is gated so it - // ends once what the peer granted us stops covering it, and a limit-authorized one - // once the limit does too. - let (mut gate, request_grant) = match token { - Some(token) => { - let verdict = gate.verify(token, Some(slot.clone())); - // The subscriber cancelling or the session closing withdraws the token. - let verified = self - .until_withdrawn( - &mut stream, - super::request_stream::fin_cancels(self.version), - |waiter| verdict.poll_grant(waiter), - ) - .await; - let Some(verified) = verified else { - drop((verdict, slot)); - self.released(&credit).await; - return Ok(()); - }; - let admitted = crate::auth::RequestGrant::new( - &self.runtime, - verdict, - verified, - msg.track_namespace.clone(), - crate::auth::RequestKind::Subscribe, - Some(slot.clone()), - ); - match admitted { - Ok(grant) => (gate.named(), Some(grant)), - Err(refusal) => { - return self.reject_subscribe(stream, &refusal.error, &refusal.reason).await; - } - } + /// Go on from forwarding, which ended as `served` says (`None` when cancelled): pause, or + /// end the subscription. + fn forwarded( + &mut self, + request: &mut SubscribeRequest, + served: Option, + ) -> Result, Error> { + // A pause drops the subscription, and its demand with it; a resume subscribes again + // as a new SUBSCRIBE would, so nothing carries across but the request. + if let Some(Served::Paused(answer)) = served { + let serve = self.track.take().expect("forwarding"); + // What an update changed while forwarding outlives the pause. + let applied = serve.track.subscription(); + let msg = &mut request.msg; + msg.subscriber_priority = super::priority::to_wire(applied.priority); + msg.filter = serve.filter; + msg.max_age = applied.delay.max.and_then(|delay| delay.as_millis().try_into().ok()); + let mut serve = serve; + serve.close_groups(); + request.publisher.pause_join(request.join.as_ref()); + drop(serve); + self.step = match answer { + true => { + let stream = request.stream.as_mut().expect("open until PUBLISH_DONE"); + Publisher::buffer_request_ok(&mut stream.writer, None)?; + ServingStep::Pausing } - None => (gate, None), - }; - let mut updates = SubscriptionUpdates { - grant: request_grant.clone(), - renewal: Default::default(), + false => ServingStep::Paused(WaitResume::new()), }; - let _slot = slot; - - // Everything the subscription does once admitted, its setup included, runs under its - // authority: a grant that ends or a gate that closes at any await drops it right there, - // releasing what it holds upstream, and only then is the peer told. `acked` holds the - // group stream count once SUBSCRIBE_OK is buffered, which decides how. - let mut acked: Option> = None; - // Set once the track ran to its end, which is all that leaves a join worth keeping. - let mut track_ended = false; - // Set once the request token lapses, just before the subscription's streams drop. - let lapse = Lapse::default(); - let rendezvous = msg.rendezvous; - let subscribed = { - let serve = async { - // Stats (subscriptions, viewer refcount, groups/frames/bytes) are counted in - // the model, through the tagged `origin::Consumer` the broadcast resolves from. - - let refused = |err: Error, resolved| { - Ok(Subscribed::Refused( - err.to_string(), - rendezvous_error(err, rendezvous, resolved), - )) - }; - // We just received a subscribe for this exact namespace, so the peer must have - // already seen the announcement. The immediate lookup resolves it at once, or - // falls back to the route covering it (an `origin::Requests`), if any. - let broadcast = match self - .resolve( - &mut stream, - super::request_stream::fin_cancels(self.version), - &msg.track_namespace, - ) - .await - { - Ok(Some(broadcast)) => broadcast, - // The reason is the origin's, not ours: a dynamic router refusing on - // authorization or a drain says so, and only an unroutable path means the - // broadcast is not here. - Err(err) => return refused(err, false), - Ok(None) => return Ok(Subscribed::Withdrawn), - }; - let track = match broadcast.track(&msg.track_name) { - Ok(track) => track, - Err(err) => return refused(err, true), - }; - let (mut track_serve, edge) = match self.subscribe_track(&mut stream, &track, &msg, &lapse).await { - Some(Ok(subscribed)) => subscribed, - Some(Err(err)) => return refused(err, true), - None => return Ok(Subscribed::Withdrawn), - }; - let cache = track; - // The fill stamps its objects as the subscription does. - let timescale = track_serve.timescale; - let properties = match track_properties( - track_serve.track.info(), - msg.properties_wanted, - self.version, - self.peer_setup.get().await.flate, - ) { - Ok(properties) => properties, - Err(err) => return Ok(Subscribed::Refused(err.to_string(), err)), - }; - // Draft-22 replaced joining FETCH with subscription fills. Older drafts save - // the same boundary used by the subscription so the two streams never overlap, - // again on every resume. - let filter = msg.filter; - let namespace = msg.track_namespace.clone(); - let grant = request_grant.clone(); - let save_join = |largest: Option, - cache: &crate::track::Consumer, - ranking: &Ranking>| { - if let Some(join) = &_join { - let joined = match (filter, largest) { - (Filter::NextObject, Some(largest)) => Joined::Group { - end: Location { - group: largest.group, - object: largest.object + 1, - }, - cache: cache.clone(), - namespace: namespace.clone(), - grant: grant.clone(), - ranking: ranking.clone(), - }, - (Filter::NextObject, None) => Joined::Empty, - _ => Joined::Unsupported, - }; - join.joins.lock().insert(request_id, Some(joined)); - } - }; - save_join(edge.largest, &cache, &track_serve.ranking); - - // A fill reads the group cache through its own consumer, independent of the - // subscription's cursor. - let fill = msg - .fill - .filter(|_| msg.forward && Filter::is_draft20(self.version)) - .map(|fill| (fill_range(fill, msg.filter, edge.largest), cache, timescale)); - - // Buffered whole before its first await, so however the subscription is - // dropped from here, the request can still end on the wire. - acked = Some(track_serve.opened.clone()); - stream - .writer - .encode_message(&ietf::SubscribeOk { - track_alias: request_id.0, - // The subscription floor and FETCH/FILL cap use this same snapshot. - largest: edge.largest, - properties, - }) - .await?; - - // The fill (when one was requested) runs alongside on its own fetch stream; - // its failures reset that stream and never touch the subscription. Every - // subscription a resume makes counts toward the same PUBLISH_DONE. - let opened = track_serve.opened.clone(); - let ranking = track_serve.ranking.clone(); - let fill = async { - if let Some((fill, cache, timescale)) = fill { - self.run_fill(fill, &ranking, cache, timescale, &opened, ranking.request_id) - .await; - } - }; - let mut served = match msg.forward { - true => { - self.run_subscription( - &mut stream, - &mut track_serve, - &mut request_finished, - &mut goaways, - fill, - &mut updates, - ) - .await - } - false => { - // Unpolled, it would still hold the track it serves from. - drop(fill); - Some(Served::Paused(false)) - } - }; - // A pause drops the subscription, and its demand with it; a resume subscribes - // again as a new SUBSCRIBE would, so nothing carries across but the request. - while let Some(Served::Paused(answer)) = served { - // What an update changed while forwarding outlives the pause. - let applied = track_serve.track.subscription(); - msg.subscriber_priority = super::priority::to_wire(applied.priority); - msg.filter = track_serve.filter; - msg.max_age = applied.delay.max.and_then(|delay| delay.as_millis().try_into().ok()); - track_serve.close_groups(); - self.pause_join(_join.as_ref()); - drop(track_serve); - if answer { - self.write_request_ok(&mut stream.writer, None).await?; - } - match self - .wait_resume(&mut stream, &mut request_finished, &mut goaways, &mut msg, &mut updates) - .await - { - Ok(()) => {} - Err(served) => { - let served = served.unwrap_or(Served::Track(Ok(()))); - return Ok(Subscribed::Served(served, opened.load(Ordering::Relaxed))); - } - } - let resumed = match broadcast.track(&msg.track_name) { - Ok(track) => self - .subscribe_track(&mut stream, &track, &msg, &lapse) - .await - .map(|subscribed| subscribed.map(|(serve, edge)| (serve, edge, track))), - Err(err) => Some(Err(err)), - }; - // Withdrawn mid-setup, it ends as it would have while paused. - let Some(resumed) = resumed else { - return Ok(Subscribed::Served( - Served::Track(Ok(())), - opened.load(Ordering::Relaxed), - )); - }; - let (edge, track); - (track_serve, edge, track) = match resumed { - Ok(resumed) => resumed, - // The resuming update is owed an answer, so it fails. - Err(err) => { - self.write_subscribe_error(&mut stream.writer, &err, "resume failed") - .await?; - return Ok(Subscribed::Served( - Served::UpdateFailed(err), - opened.load(Ordering::Relaxed), - )); - } - }; - track_serve.opened = opened.clone(); - save_join(edge.largest, &track, &track_serve.ranking); - self.write_request_ok(&mut stream.writer, edge.largest).await?; - served = self - .run_subscription( - &mut stream, - &mut track_serve, - &mut request_finished, - &mut goaways, - async {}, - &mut updates, - ) - .await; - } - - let completed = served.is_some(); - let served = served.unwrap_or(Served::Track(Ok(()))); - - if completed && matches!(served, Served::Track(Ok(()))) { - track_ended = true; - self.write_end_of_track(&mut stream, &mut track_serve, &mut request_finished, &mut goaways) - .await; - } + return Ok(None); + } - // PUBLISH_DONE must follow the close of every data stream. A subscription that - // completed has drained its groups; one that ended early resets the rest here. - // The fill stream closed when `run_subscription` dropped it. - track_serve.close_groups(); - Ok::<_, Error>(Subscribed::Served(served, track_serve.opened())) + let completed = served.is_some(); + let served = served.unwrap_or(Served::Track(Ok(()))); + if completed && matches!(served, Served::Track(Ok(()))) { + request.track_ended = true; + if let Some(end) = self.track.as_ref().expect("forwarding").end() { + self.step = ServingStep::EndOfTrack { + served, + marker: EndOfTrack::new(end), }; - crate::auth::until(serve, |waiter| { - if let Some(grant) = &request_grant - && let Poll::Ready(refusal) = grant.poll_ended(waiter) - { - // PUBLISH_DONE carries a fixed reason per status, so only the code reaches - // the peer once SUBSCRIBE_OK went out. - tracing::info!(broadcast = %broadcast(), track = %track_name, err = %refusal.error, reason = %refusal.reason, "request token grant ended"); - if super::error::token_expired(&refusal.error) { - lapse.set(); - } - return Poll::Ready(refusal); - } - if gate.poll_denied(waiter).is_ready() { - tracing::info!(broadcast = %broadcast(), track = %track_name, "subscription no longer authorized"); - return Poll::Ready(Error::Unauthorized.into()); - } - Poll::Pending - }) - .await - }; - // No update follows the subscription's end, so its pending renewals no longer matter. - drop(updates); - // A track that ended cleanly keeps what a joining FETCH needs (see [`Retain`]); any - // other end drops the join, refusing a joining FETCH that has not answered yet. - let retain = match (&subscribed, _join) { - (Ok(Ok(Subscribed::Served(Served::Track(Ok(())), _))), Some(join)) if track_ended && join.answers() => { - Some(Retain { - _join: join, - grant: request_grant, - gate, - _slot, - }) - } - _ => { - // The subscription's demand is gone; only its credit may still wait on a - // verification. - drop((request_grant, _slot)); - None - } - }; - - let (served, streams) = match subscribed { - Ok(Ok(Subscribed::Served(Served::RenewalRefused(refusal), streams))) => { - let res = self - .write_subscribe_error(&mut stream.writer, &refusal.error, &refusal.reason) - .await; - match res { - Ok(()) => (Served::UpdateFailed(refusal.error), streams), - Err(err) => (Served::Track(Err(err)), streams), - } - } - Ok(Ok(Subscribed::Served(served, streams))) => (served, streams), - Ok(Ok(Subscribed::Refused(reason, err))) => { - return self.reject_subscribe(stream, &err, &reason).await; - } - // Nobody is left to answer. - Ok(Ok(Subscribed::Withdrawn)) => { - self.released(&credit).await; - return Ok(()); - } - Ok(Err(err)) => return Err(err), - Err(refusal) => match acked { - Some(opened) => (Served::Track(Err(refusal.error)), opened.load(Ordering::Relaxed)), - None => { - return self.reject_subscribe(stream, &refusal.error, &refusal.reason).await; - } - }, - }; - self.write_publish_done(stream, served, streams, Some(&credit), retain) - .await + return Ok(None); + } } + Ok(Some(self.ended(served))) } - /// Ready once the requester withdraws the request on `stream` or `session` closes, without - /// reading anything it sent meanwhile. - fn poll_withdrawn( - &self, - session: &mut Transport, - stream: &mut Stream, - fin_cancels: bool, - waiter: &kio::Waiter, - ) -> Poll<()> { - if session.poll_closed(waiter).is_ready() { - return Poll::Ready(()); - } - super::request_stream::poll_withdrawn(stream, fin_cancels, waiter) + /// Fail the update resuming the subscription with `err`, flushing its answer next. + fn resume_failed(request: &mut SubscribeRequest, err: Error) -> Result { + let stream = request.stream.as_mut().expect("open until PUBLISH_DONE"); + request + .publisher + .buffer_subscribe_error(&mut stream.writer, &err, "resume failed")?; + Ok(ServingStep::ResumeFailed(err)) } - /// Verify an incoming request token while keeping its stream credit until verification ends. - async fn admit_request( - &self, - stream: &mut Stream, - gate: &mut crate::auth::Gate, - token: Option, - fin_cancels: bool, - ) -> std::result::Result { - let Some(token) = token.filter(|_| gate.token_decides()) else { - return Ok(true); - }; - let credit = kio::Producer::<()>::default(); - let hold: crate::auth::Hold = Arc::new(credit.consume()); - let verdict = gate.verify(token.0, Some(hold)); - let answer = self - .until_withdrawn(stream, fin_cancels, |waiter| verdict.poll_grant(waiter)) - .await; - let Some(answer) = answer else { - drop(verdict); - self.released(&credit).await; - return Ok(false); - }; - gate.admit(&self.runtime, (verdict, answer))?; - Ok(true) + /// The subscription ended as `served` says. + fn ended(&mut self, served: Served) -> Subscribed { + // PUBLISH_DONE must follow the close of every data stream. A subscription that + // completed has drained its groups; one that ended early resets the rest here. The + // fill stream closed when its [`RunSubscription`] dropped it. + let serve = self.track.as_mut().expect("forwarding"); + serve.close_groups(); + Subscribed::Served(served, serve.opened()) } +} - /// Poll `step` of a request's setup until ready, or `None` once [`Self::poll_withdrawn`] is. - async fn until_withdrawn( - &self, - stream: &mut Stream, - fin_cancels: bool, - mut step: impl FnMut(&kio::Waiter) -> Poll, - ) -> Option { - let mut session = self.session.clone(); - kio::wait(|waiter| { - if self - .poll_withdrawn(&mut session, stream, fin_cancels, waiter) - .is_ready() - { - return Poll::Ready(None); - } - step(waiter).map(Some) - }) - .await +impl SubscribeRequest { + /// The group streams the subscription opened, counted since SUBSCRIBE_OK. + fn opened(&self) -> u64 { + self.acked + .as_ref() + .expect("acked before serving") + .load(Ordering::Relaxed) } +} - /// Wait until no token verification holds `credit`'s request any more, or the session closes. - async fn released(&self, credit: &kio::Producer<()>) { - let session = self.session.clone(); - kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() || credit.poll_wanted_changed(waiter, true).is_ready() { - return Poll::Ready(()); - } - Poll::Pending - }) - .await +/// Subscribes to a track as a SUBSCRIBE asks, resolving its filter at the live edge, as a +/// SUBSCRIBE does and each resume of one does again. `None` once the subscriber withdraws +/// the request, since the producer may never answer. +struct SubscribeTrack { + state: SubscribeTrackState, +} + +#[kio::machine(self.state: SubscribeTrackState)] +impl SubscribeTrack { + fn new() -> Self { + Self { + state: SubscribeTrackState::Start, + } } - /// Send the END_OF_TRACK a subscription that ran to the track's end still owes: - /// PUBLISH_DONE carries no end location, so the marker is what tells the subscriber - /// where the track ended. One cancelled while the marker waits for stream credit abandons - /// it. - async fn write_end_of_track( - &self, + fn poll( + &mut self, + waiter: &kio::Waiter, + publisher: &Publisher, stream: &mut Stream, - track_serve: &mut TrackServe, - request_finished: &mut bool, - goaways: &mut super::request_stream::Messages, - ) { - let Some(end) = track_serve.end() else { - return; + track: &track::Consumer, + msg: &ietf::Subscribe<'_>, + lapse: &Lapse, + ) -> Poll>> { + let setup = wait!(Setup => publisher.poll_until_withdrawn( + stream, + super::request_stream::fin_cancels(publisher.version), + waiter, + |waiter| publisher.peer_setup.poll_get(waiter), + )); + let flate: bool = match setup { + Some(declared) => declared.flate, + None => return Poll::Ready(None), }; - let request_id = track_serve.request_id; - let mut marker = EndOfTrack::new(end); - let closed_session = self.session.clone(); - let written = kio::wait(|waiter| { - if let Poll::Ready(res) = marker.poll(waiter, track_serve) { - return Poll::Ready(res); - } - if super::request_stream::poll_cancel( - stream, - request_finished, - goaways, - super::request_stream::fin_cancels(self.version), - waiter, - ) - .is_ready() || closed_session.poll_closed(waiter).is_ready() - { - return Poll::Ready(Err(Error::Cancel)); - } - Poll::Pending - }) - .await; - // A failure only costs the subscriber the early boundary. - if let Err(err) = written { - tracing::debug!(%err, id = %request_id, "end of track failed"); - } - } - - /// End a subscription on the wire with the PUBLISH_DONE `served` reports, after `streams` - /// group streams, returning how serving ended. The stream closes only once no verification - /// holds `credit`'s request, if it has one, and once `retain` lets go of its join. - async fn write_publish_done( - &self, - mut stream: Stream, - served: Served, - streams: u64, - credit: Option<&kio::Producer<()>>, - retain: Option, - ) -> Result<(), Error> { - let (status, reason) = served.status(); - let res = match served { - Served::Track(res) => res, - Served::UpdateFailed(err) => Err(err), - Served::RenewalRefused(refusal) => Err(refusal.error), - Served::Paused(_) => Ok(()), - }; - - let mut retain = retain; - { - let done = ietf::PublishDone { - status_code: status.code(), - stream_count: streams, - reason_phrase: reason.into(), - }; - // A peer holding back flow control parks PUBLISH_DONE, and the kept join must still - // go once the request closes or its authorization lapses meanwhile. - let mut encode = std::pin::pin!(stream.writer.encode_message(&done)); - let reader = &mut stream.reader; - let _ = kio::wait(|waiter| { - if let Some(kept) = &mut retain - && (reader.poll_discard(waiter).is_ready() || kept.poll_lapsed(waiter).is_ready()) - { - retain = None; - } - waiter.poll_future(encode.as_mut()) - }) - .await; - } - - if let Some(retain) = retain { - self.retain_join(&mut stream, retain).await; - } - - // The peer learns the subscription ended at once, but the stream holds its credit until - // the application lets go of every token it took for it. - if let Some(credit) = credit { - self.released(credit).await; - } - // PUBLISH_DONE is the last thing on this stream, so it needs the acknowledgement too. - let _ = stream.writer.close().await; - - res - } - - /// Keep a finished subscription's join until the subscriber closes the request, the - /// session closes, or what authorized the subscription stops covering it. - async fn retain_join(&self, stream: &mut Stream, mut retain: Retain) { - let session = self.session.clone(); - kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() || stream.writer.poll_closed(waiter).is_ready() { - return Poll::Ready(()); - } - // An update that crossed PUBLISH_DONE is skipped, since nothing answers it now; the - // subscriber's FIN or a reset ends the request. - if stream.reader.poll_discard(waiter).is_ready() { - return Poll::Ready(()); - } - retain.poll_lapsed(waiter) - }) - .await; - } - - async fn run_subscription( - &self, - stream: &mut Stream, - serve: &mut TrackServe, - finished: &mut bool, - goaways: &mut super::request_stream::Messages, - fill: impl std::future::Future, - updates: &mut SubscriptionUpdates, - ) -> Option { - let session = self.session.clone(); - let mut fill = std::pin::pin!(fill); - let mut fill_done = false; - let mut served = None; - let mut fills = crate::util::TaskSet::owned(); - let mut fills_done = true; - let mut waiting: Option = None; - let mut renewal_largest = None; - - // Renewal state carried across turns of the loop. After SUBSCRIBE_OK the peer may send a - // REQUEST_UPDATE refreshing the request's token. A token-bearing update's verify is raced - // against serving (never a bare await), so media keeps flowing while the acceptor is - // slow; the caller's guard still ends the request when the old grant does. While a verdict is - // pending the loop keeps reading so a cancellation still ends the request, stashing every - // further update, token or not, to handle once the verdict resolves: answers - // carry no Request ID, so they must go out in the order the updates arrived. The stash is - // a FIFO bounded by MAX_REQUEST_UPDATES (counting the one being verified). - let SubscriptionUpdates { grant, renewal } = updates; - let mut pausing = false; - - // What one turn of the loop resolved to. - enum Turn { - // The track and any fill finished, or the stream failed. - Served(Result<(), Error>), - // Reading the requester's next update failed. - Malformed(Error), - // The peer cancelled, or the session ended. - Cancelled, - // A REQUEST_UPDATE, already decoded. - Update(super::request_stream::Update), - // A pending renewal's verdict resolved, for the update with this request id. - Renewal(request_update::Verdict), - } - - // Each turn may apply a REQUEST_UPDATE, so a peer streaming them must not hold the - // session's other tasks. - let mut budget = kio::coop::Budget::new(32); - loop { - kio::wait(|waiter| budget.poll_yield(waiter)).await; - let mut queued = waiting.is_none().then(|| renewal.next()).flatten(); - - let turn = kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() { - return Poll::Ready(Turn::Cancelled); - } - if stream.writer.poll_closed(waiter).is_ready() { - return Poll::Ready(Turn::Cancelled); - } - if let Poll::Ready(renewed) = renewal.poll(waiter) { - return Poll::Ready(Turn::Renewal(renewed)); - } - if !fill_done && waiter.poll_future(fill.as_mut()).is_ready() { - fill_done = true; - } - if !fills_done && fills.poll(waiter).is_ready() { - fills_done = true; - } - if fill_done - && fills_done && serve.tail.is_none() - && serve.opening.is_none() - && let Some(update) = waiting.take() - { - return Poll::Ready(Turn::Update(update)); - } - if let Some(update) = queued.take() { - return Poll::Ready(Turn::Update(update)); - } - if waiting.is_none() { - match self.poll_update(stream, finished, goaways, waiter) { - Poll::Ready(UpdateRead::Update(update)) => return Poll::Ready(Turn::Update(update)), - Poll::Ready(UpdateRead::Cancelled) => return Poll::Ready(Turn::Cancelled), - Poll::Ready(UpdateRead::Malformed(err)) => return Poll::Ready(Turn::Malformed(err)), - Poll::Ready(UpdateRead::Failed(err)) => return Poll::Ready(Turn::Served(Err(err))), - Poll::Pending => {} - } - } - - if served.is_none() - && let Poll::Ready(done) = serve.poll(waiter) - { - served = Some(done); - } - if fill_done - && fills_done && serve.tail.is_none() - && serve.opening.is_none() - && let Some(update) = waiting.take() - { - return Poll::Ready(Turn::Update(update)); - } - match (&served, fill_done && fills_done && waiting.is_none()) { - (Some(Ok(())), true) => Poll::Ready(Turn::Served(Ok(()))), - (Some(Err(err)), true) => Poll::Ready(Turn::Served(Err(err.clone()))), - _ => Poll::Pending, - } - }) - .await; - - match turn { - // A malformed update, such as one whose AUTHORIZATION TOKEN names an alias, - // closes the session rather than only ending this subscription, as the same - // token on the request itself would. A source track's decode failure is not - // this requester's fault, so it reaches Served and ends only this subscription. - Turn::Malformed(err) => { - if super::subscriber::is_protocol_violation(&err) { - self.session.clone().close( - super::error::to_session_code(&crate::SessionError::from(&err)), - &err.to_string(), - ); - } - return Some(Served::Track(Err(err))); - } - Turn::Served(res) => return Some(Served::Track(res)), - Turn::Cancelled => return None, - Turn::Renewal(request_update::Verdict { verdict, grant: res }) => { - // Only a token-authorized request starts a renewal. - let rg = grant.as_ref().expect("a renewal without a request grant"); - // On accept the old grant is dropped (ending the old token) and the deadline - // re-armed at the new expiry. - let Err(refusal) = rg.renew(verdict, res) else { - if pausing { - return Some(Served::Paused(true)); - } - if let Err(err) = self.write_request_ok(&mut stream.writer, renewal_largest.take()).await { - return Some(Served::Track(Err(err))); - } - continue; - }; - // Per draft-18 section 10.9.1, a refused or uncovered - // renewal ends only this subscription: REQUEST_ERROR, then PUBLISH_DONE - // UPDATE_FAILED. The session stays up and the old grant does not survive, - // so a FETCH joined to it stops at once rather than finishing under it. - rg.end(refusal.clone()); - return Some(Served::RenewalRefused(refusal)); - } - // An update behind a pending verdict waits its turn, so its answer follows the - // verdict's. - Turn::Update(update) if renewal.pending() => match renewal.queue(update, self.version) { - Ok(()) => {} - Err(request_update::Overflow::Session) => { - return Some(Served::Track(Err(request_update::too_many(&mut self.session.clone())))); - } - // Below draft 22 no limit is negotiated, so this is a local memory guard: end - // this request, never the session. - Err(request_update::Overflow::Request) => { - return Some(Served::Track(Err(Error::ProtocolViolation))); - } - }, - Turn::Update(update) - if update.fill.is_some() - && (!fill_done || !fills_done || serve.tail.is_some() || serve.opening.is_some()) => - { - waiting = Some(update); - } - Turn::Update(update) => { - pausing = update.forward == Some(false); - match self - .apply_subscription_update(stream, serve, grant, renewal, update) - .await - { - Err(served) => return Some(served), - Ok(Some(fill)) => { - if renewal.pending() { - renewal_largest = fill.largest; - } - let publisher = self.clone(); - fills.push(async move { - publisher - .run_fill( - fill.fill, - &fill.ranking, - fill.track, - fill.timescale, - &fill.opened, - fill.request_id, - ) - .await; - }); - fills_done = false; - } - Ok(None) => {} - } - } - } - } - } - - /// Subscribe to `track` as `msg` asks, resolving its filter at the live edge, as a - /// SUBSCRIBE does and each resume of one does again. `None` once the subscriber withdraws - /// the request on `stream`, since the producer may never answer. - async fn subscribe_track( - &self, - stream: &mut Stream, - track: &crate::track::Consumer, - msg: &ietf::Subscribe<'_>, - lapse: &Lapse, - ) -> Option> { - let declared = self - .until_withdrawn(stream, super::request_stream::fin_cancels(self.version), |waiter| { - self.peer_setup.poll_get(waiter) - }) - .await?; - // A live SUBSCRIBE's NEW_GROUP_REQUEST asks nothing more: subscribing is itself a join, - // which the track's refresh ratio judges. - let mut subscription = serving_subscription(msg.subscriber_priority).with_order(msg.group_order.to_order()); - if let Some(max_age) = msg.max_age { - subscription.delay.max = Some(Duration::from_millis(max_age)); + // A live SUBSCRIBE's NEW_GROUP_REQUEST asks nothing more: subscribing is itself a join, + // which the track's refresh ratio judges. + let mut subscription: Subscription = + serving_subscription(msg.subscriber_priority).with_order(msg.group_order.to_order()); + if let Some(max_age) = msg.max_age { + subscription.delay.max = Some(Duration::from_millis(max_age)); } // A range that resolves without the live edge has a fixed start, which is no join: // bound the subscription before subscribing so the ratio never judges it. From the // start of its group, so a relay's copy holds the group's head for its other readers, // as a live join would; the resolved range trims below. - let fixed = subscribe_range(msg.filter, LiveEdge::default(), self.version); + let fixed: ServeRange = subscribe_range(msg.filter, LiveEdge::default(), publisher.version); if let Some(start) = fixed.start { ServeRange { start: Some(Location { object: 0, ..start }), @@ -2138,35 +1907,41 @@ impl Publisher { // becomes readable once the subscription's demand attaches a route, so the edge snapshot // has to come after. The resolved range is applied to the preference right below, before // anything is served. - let subscribing = track.subscribe(subscription.clone()); - let mut subscribed = match self - .until_withdrawn(stream, super::request_stream::fin_cancels(self.version), |waiter| { - subscribing.poll_ok(waiter) - }) - .await? - { + let subscribing: track::Subscribing = track.subscribe(subscription.clone()).into_inner(); + let answer = wait!(Subscribe => publisher.poll_until_withdrawn( + stream, + super::request_stream::fin_cancels(publisher.version), + waiter, + |waiter| subscribing.poll_ok(waiter), + )); + let mut subscribed: track::Subscriber = match answer { // A peer reads the end of the subscription as the track's. - Ok(subscribed) => subscribed.parked(), - Err(err) => return Some(Err(err)), + Some(Ok(found)) => found.parked(), + Some(Err(err)) => return Poll::Ready(Some(Err(err))), + None => return Poll::Ready(None), }; - drop(subscribing); // A relay's copy that went idle cannot say where its live edge is until its route // answers again; answering from its cache would advertise a stale one. - self.until_withdrawn(stream, super::request_stream::fin_cancels(self.version), |waiter| { - subscribed.poll_live(waiter) - }) - .await?; + let live = wait!(Live => publisher.poll_until_withdrawn( + stream, + super::request_stream::fin_cancels(publisher.version), + waiter, + |waiter| subscribed.poll_live(waiter), + )); + if live.is_none() { + return Poll::Ready(None); + } // The filter and any fill are relative to the live edge, so snapshot it once: the fill // ends exactly where a Next Object subscription begins, which is what lets the draft's // current-group join (Next Object plus a StartGroup=1 fill) cover the group with no gap // and no overlap. let edge = live_edge(track); - let mut range = subscribe_range(msg.filter, edge, self.version); + let mut range = subscribe_range(msg.filter, edge, publisher.version); let fill = msg .fill - .filter(|_| msg.forward && Filter::is_draft20(self.version)) + .filter(|_| msg.forward && Filter::is_draft20(publisher.version)) .map(|fill| fill_range(fill, msg.filter, edge.largest)); if let Some(next) = range.fill_tail(fill, edge) { range.start = Some(next); @@ -2179,17 +1954,16 @@ impl Publisher { // Serve the track while reading updates; only abrupt closure cancels on draft-22. A // fill or joining FETCH ranks in the same subscription as the group it carries. - let serve = TrackServe::new( - self.session.clone(), + let mut serve = TrackServe::new( + publisher.session.clone(), subscribed, msg.request_id, - self.version, + publisher.version, range, timescale, ) - .compression(declared.flate) - .ranked(self.priority.clone()); - let mut serve = serve; + .compression(flate) + .ranked(publisher.priority.clone()); serve.filter = msg.filter; serve.ranking.lapse = lapse.clone(); // A fixed start is no join, so its NEW_GROUP_REQUEST asks as an update's would. @@ -2198,113 +1972,443 @@ impl Publisher { { serve.request_new_group(value); } - Some(Ok((serve, edge))) + Poll::Ready(Some(Ok((serve, edge)))) } +} - /// Hold a paused SUBSCRIBE, answering its updates, until one resumes it, or return how the - /// request ended (`None` when cancelled). - /// - /// A paused subscription holds nothing upstream. Its priority, max age and authority can - /// change while it waits; serving parameters are applied on resume. - async fn wait_resume( - &self, +/// Serves a subscription's track while answering its updates, until the track and its fills +/// finish, the request ends, or an update pauses it. `None` once the subscriber cancels or +/// the session closes. +/// +/// After SUBSCRIBE_OK the peer may send a REQUEST_UPDATE refreshing the request's token. A +/// token-bearing update's verify is raced against serving (never a bare wait), so media keeps +/// flowing while the acceptor is slow; the caller's guard still ends the request when the old +/// grant does. While a verdict is pending the loop keeps reading so a cancellation still ends +/// the request, stashing every further update, token or not, to handle once the verdict +/// resolves: answers carry no Request ID, so they must go out in the order the updates +/// arrived. The stash is a FIFO bounded by MAX_REQUEST_UPDATES (counting the one being +/// verified). +struct RunSubscription { + /// The fill the SUBSCRIBE asked for, until it finishes. + fill: Option, + /// The fills updates asked for. + fills: crate::util::TaskSet, + fills_done: bool, + /// How serving the track ended, once it did. + served: Option>, + /// An update asking for a fill, held until the fills and the group being opened finish. + waiting: Option, + /// The Largest Location the REQUEST_OK for a renewal being verified reports. + renewal_largest: Option, + /// The last update applied pauses the subscription once its renewal is accepted. + pausing: bool, + /// Each turn may apply a REQUEST_UPDATE, so a peer streaming them must not hold the + /// session's other tasks. + budget: kio::coop::Budget, + state: RunState, +} + +enum RunState { + /// Starting a turn. + Next, + /// Waiting for what the turn resolves to, with the update a verdict held back, if any. + Turn(Option), + /// Flushing an update's answer, then starting its fill, or ending as given. + Answer(std::result::Result, Served>), +} + +impl RunSubscription { + fn new(fill: Option) -> Self { + Self { + fill, + fills: crate::util::TaskSet::owned(), + fills_done: true, + served: None, + waiting: None, + renewal_largest: None, + pausing: false, + budget: kio::coop::Budget::new(32), + state: RunState::Next, + } + } + + #[expect( + clippy::too_many_arguments, + reason = "the request state it shares with the serve that owns it" + )] + fn poll( + &mut self, + waiter: &kio::Waiter, + publisher: &Publisher, stream: &mut Stream, + serve: &mut TrackServe, finished: &mut bool, goaways: &mut super::request_stream::Messages, - msg: &mut ietf::Subscribe<'_>, updates: &mut SubscriptionUpdates, - ) -> Result<(), Option> { - let session = self.session.clone(); - let SubscriptionUpdates { grant, renewal } = updates; - let mut resuming = false; + ) -> Poll> { + // What one turn of the loop resolved to. enum Turn { - Read(UpdateRead), + // The track and any fill finished, or the stream failed. + Served(Result<(), Error>), + // Reading the requester's next update failed. + Malformed(Error), + // The peer cancelled, or the session ended. + Cancelled, + // A REQUEST_UPDATE, already decoded. + Update(super::request_stream::Update), + // A pending renewal's verdict resolved, for the update with this request id. Renewal(request_update::Verdict), } - // Each update is answered, so a peer streaming them must not hold the session's - // other tasks. - let mut budget = kio::coop::Budget::new(32); + + let SubscriptionUpdates { grant, renewal } = updates; loop { - kio::wait(|waiter| budget.poll_yield(waiter)).await; - let mut queued = renewal.next(); - let turn = kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() { - return Poll::Ready(Turn::Read(UpdateRead::Cancelled)); - } - if stream.writer.poll_closed(waiter).is_ready() { - return Poll::Ready(Turn::Read(UpdateRead::Cancelled)); - } - if let Poll::Ready(verdict) = renewal.poll(waiter) { - return Poll::Ready(Turn::Renewal(verdict)); - } - if let Some(update) = queued.take() { - return Poll::Ready(Turn::Read(UpdateRead::Update(update))); - } - self.poll_update(stream, finished, goaways, waiter).map(Turn::Read) - }) - .await; - let read = match turn { - Turn::Read(read) => read, - Turn::Renewal(request_update::Verdict { verdict, grant: res }) => { - let rg = grant.as_ref().expect("a renewal without a request grant"); - if let Err(refusal) = rg.renew(verdict, res) { - // As while serving: the grant ends, and its REQUEST_ERROR goes out once - // the guard is gone. - rg.end(refusal.clone()); - return Err(Some(Served::RenewalRefused(refusal))); - } - if resuming { - return Ok(()); - } - self.write_request_ok(&mut stream.writer, None) - .await - .map_err(|err| Some(Served::Track(Err(err))))?; - continue; + match &mut self.state { + RunState::Next => { + ready!(self.budget.poll_yield(waiter)); + let queued = self.waiting.is_none().then(|| renewal.next()).flatten(); + self.state = RunState::Turn(queued); } - }; - let update = match read { - UpdateRead::Update(update) => update, - UpdateRead::Cancelled => return Err(None), - UpdateRead::Malformed(err) => { - if super::subscriber::is_protocol_violation(&err) { - self.session.clone().close( - super::error::to_session_code(&crate::SessionError::from(&err)), - &err.to_string(), - ); + RunState::Turn(queued) => { + let turn = 'turn: { + if publisher.session.poll_closed(waiter).is_ready() { + break 'turn Turn::Cancelled; + } + if stream.writer.poll_closed(waiter).is_ready() { + break 'turn Turn::Cancelled; + } + if let Poll::Ready(renewed) = renewal.poll(waiter) { + break 'turn Turn::Renewal(renewed); + } + if let Some(fill) = &mut self.fill + && fill.poll(waiter, publisher).is_ready() + { + self.fill = None; + } + if !self.fills_done && self.fills.poll(waiter).is_ready() { + self.fills_done = true; + } + if self.fill.is_none() + && self.fills_done && serve.tail.is_none() + && serve.opening.is_none() + && let Some(update) = self.waiting.take() + { + break 'turn Turn::Update(update); + } + if let Some(update) = queued.take() { + break 'turn Turn::Update(update); + } + if self.waiting.is_none() { + match publisher.poll_update(stream, finished, goaways, waiter) { + Poll::Ready(UpdateRead::Update(update)) => break 'turn Turn::Update(update), + Poll::Ready(UpdateRead::Cancelled) => break 'turn Turn::Cancelled, + Poll::Ready(UpdateRead::Malformed(err)) => break 'turn Turn::Malformed(err), + Poll::Ready(UpdateRead::Failed(err)) => break 'turn Turn::Served(Err(err)), + Poll::Pending => {} + } + } + + if self.served.is_none() + && let Poll::Ready(done) = serve.poll(waiter) + { + self.served = Some(done); + } + if self.fill.is_none() + && self.fills_done && serve.tail.is_none() + && serve.opening.is_none() + && let Some(update) = self.waiting.take() + { + break 'turn Turn::Update(update); + } + match ( + &self.served, + self.fill.is_none() && self.fills_done && self.waiting.is_none(), + ) { + (Some(Ok(())), true) => Turn::Served(Ok(())), + (Some(Err(err)), true) => Turn::Served(Err(err.clone())), + _ => return Poll::Pending, + } + }; + self.state = RunState::Next; + + match turn { + // A malformed update, such as one whose AUTHORIZATION TOKEN names an alias, + // closes the session rather than only ending this subscription, as the same + // token on the request itself would. A source track's decode failure is not + // this requester's fault, so it reaches Served and ends only this subscription. + Turn::Malformed(err) => { + if super::subscriber::is_protocol_violation(&err) { + publisher.session.close( + super::error::to_session_code(&crate::SessionError::from(&err)), + &err.to_string(), + ); + } + return Poll::Ready(Some(Served::Track(Err(err)))); + } + Turn::Served(res) => return Poll::Ready(Some(Served::Track(res))), + Turn::Cancelled => return Poll::Ready(None), + Turn::Renewal(request_update::Verdict { verdict, grant: res }) => { + // Only a token-authorized request starts a renewal. + let rg = grant.as_ref().expect("a renewal without a request grant"); + // On accept the old grant is dropped (ending the old token) and the + // deadline re-armed at the new expiry. + let Err(refusal) = rg.renew(verdict, res) else { + if self.pausing { + return Poll::Ready(Some(Served::Paused(true))); + } + let largest = self.renewal_largest.take(); + if let Err(err) = Publisher::buffer_request_ok(&mut stream.writer, largest) { + return Poll::Ready(Some(Served::Track(Err(err)))); + } + self.state = RunState::Answer(Ok(None)); + continue; + }; + // Per draft-18 section 10.9.1, a refused or uncovered renewal ends only + // this subscription: REQUEST_ERROR, then PUBLISH_DONE UPDATE_FAILED. The + // session stays up and the old grant does not survive, so a FETCH joined + // to it stops at once rather than finishing under it. + rg.end(refusal.clone()); + return Poll::Ready(Some(Served::RenewalRefused(refusal))); + } + // An update behind a pending verdict waits its turn, so its answer follows + // the verdict's. + Turn::Update(update) if renewal.pending() => match renewal.queue(update, publisher.version) { + Ok(()) => {} + Err(request_update::Overflow::Session) => { + let err = request_update::too_many(&mut publisher.session.clone()); + return Poll::Ready(Some(Served::Track(Err(err)))); + } + // Below draft 22 no limit is negotiated, so this is a local memory guard: + // end this request, never the session. + Err(request_update::Overflow::Request) => { + return Poll::Ready(Some(Served::Track(Err(Error::ProtocolViolation)))); + } + }, + Turn::Update(update) + if update.fill.is_some() + && (self.fill.is_some() + || !self.fills_done || serve.tail.is_some() + || serve.opening.is_some()) => + { + self.waiting = Some(update); + } + Turn::Update(update) => { + self.pausing = update.forward == Some(false); + let applied = match publisher.apply_subscription_update( + &mut stream.writer, + serve, + grant, + renewal, + update, + ) { + Ok(applied) => applied, + Err(err) => return Poll::Ready(Some(Served::Track(Err(err)))), + }; + match (applied.answered, applied.then) { + (true, then) => self.state = RunState::Answer(then), + (false, Err(served)) => return Poll::Ready(Some(served)), + (false, Ok(fill)) => self.start_fill(publisher, renewal, fill), + } + } } - return Err(Some(Served::Track(Err(err)))); } - UpdateRead::Failed(err) => return Err(Some(Served::Track(Err(err)))), - }; - if renewal.pending() { - if let Err(overflow) = renewal.queue(update, self.version) { - let err = match overflow { - request_update::Overflow::Session => request_update::too_many(&mut self.session.clone()), - request_update::Overflow::Request => Error::ProtocolViolation, + RunState::Answer(_) => { + let flushed = ready!(stream.writer.poll_flush(waiter)); + let RunState::Answer(then) = std::mem::replace(&mut self.state, RunState::Next) else { + unreachable!("matched above"); }; - return Err(Some(Served::Track(Err(err)))); + if let Err(err) = flushed { + return Poll::Ready(Some(Served::Track(Err(err)))); + } + match then { + Err(served) => return Poll::Ready(Some(served)), + Ok(fill) => self.start_fill(publisher, renewal, fill), + } } - continue; } - let paused = "REQUEST_UPDATE parameters not supported while paused"; - let refused = match update.refused { - Some(err) => Some((err, paused)), - None if update.filter.is_some() || update.fill.is_some() || update.new_group.is_some() => { - Some((Error::Unsupported, paused)) - } - None if update.authorization_token.is_some() - && grant.is_none() - && self.auth.token_decides(crate::auth::Direction::Publish) => - { + } + } + + /// Start the fill an update asked for, if any, alongside the others. + fn start_fill( + &mut self, + publisher: &Publisher, + renewal: &request_update::Renewal, + fill: Option, + ) { + let Some(fill) = fill else { + return; + }; + if renewal.pending() { + self.renewal_largest = fill.largest; + } + let publisher = publisher.clone(); + let mut fill = Fill { + run: RunFill::new(fill.fill, fill.track, fill.timescale, fill.request_id), + ranking: fill.ranking, + opened: fill.opened, + }; + self.fills + .push(move |waiter: &kio::Waiter| fill.poll(waiter, &publisher)); + self.fills_done = false; + } +} + +/// Holds a paused SUBSCRIBE, answering its updates, until one resumes it, or returns how the +/// request ended (`None` when cancelled). +/// +/// A paused subscription holds nothing upstream. Its priority, max age and authority can +/// change while it waits; serving parameters are applied on resume. +struct WaitResume { + /// The last update applied resumes the subscription once its renewal is accepted. + resuming: bool, + /// Each update is answered, so a peer streaming them must not hold the session's other + /// tasks. + budget: kio::coop::Budget, + state: PausedState, +} + +enum PausedState { + /// Starting a turn. + Next, + /// Waiting for what the turn resolves to, with the update a verdict held back, if any. + Turn(Option), + /// Flushing an update's answer, then waiting on, or ending as given. + Answer(Option), +} + +impl WaitResume { + fn new() -> Self { + Self { + resuming: false, + budget: kio::coop::Budget::new(32), + state: PausedState::Next, + } + } + + #[expect( + clippy::too_many_arguments, + reason = "the request state it shares with the serve that owns it" + )] + fn poll( + &mut self, + waiter: &kio::Waiter, + publisher: &Publisher, + stream: &mut Stream, + finished: &mut bool, + goaways: &mut super::request_stream::Messages, + msg: &mut ietf::Subscribe<'_>, + updates: &mut SubscriptionUpdates, + ) -> Poll>> { + enum Turn { + Read(UpdateRead), + Renewal(request_update::Verdict), + } + + let failed = |err| Poll::Ready(Err(Some(Served::Track(Err(err))))); + let SubscriptionUpdates { grant, renewal } = updates; + loop { + let queued = match &mut self.state { + PausedState::Next => { + ready!(self.budget.poll_yield(waiter)); + self.state = PausedState::Turn(renewal.next()); + continue; + } + PausedState::Answer(_) => { + let flushed = ready!(stream.writer.poll_flush(waiter)); + let PausedState::Answer(then) = std::mem::replace(&mut self.state, PausedState::Next) else { + unreachable!("matched above"); + }; + if let Err(err) = flushed { + return failed(err); + } + if let Some(served) = then { + return Poll::Ready(Err(Some(served))); + } + continue; + } + PausedState::Turn(queued) => queued, + }; + let turn = 'turn: { + if publisher.session.poll_closed(waiter).is_ready() { + break 'turn Turn::Read(UpdateRead::Cancelled); + } + if stream.writer.poll_closed(waiter).is_ready() { + break 'turn Turn::Read(UpdateRead::Cancelled); + } + if let Poll::Ready(verdict) = renewal.poll(waiter) { + break 'turn Turn::Renewal(verdict); + } + if let Some(update) = queued.take() { + break 'turn Turn::Read(UpdateRead::Update(update)); + } + Turn::Read(ready!(publisher.poll_update(stream, finished, goaways, waiter))) + }; + self.state = PausedState::Next; + + let read = match turn { + Turn::Read(read) => read, + Turn::Renewal(request_update::Verdict { verdict, grant: res }) => { + let rg = grant.as_ref().expect("a renewal without a request grant"); + if let Err(refusal) = rg.renew(verdict, res) { + // As while serving: the grant ends, and its REQUEST_ERROR goes out once + // the guard is gone. + rg.end(refusal.clone()); + return Poll::Ready(Err(Some(Served::RenewalRefused(refusal)))); + } + if self.resuming { + return Poll::Ready(Ok(())); + } + if let Err(err) = Publisher::buffer_request_ok(&mut stream.writer, None) { + return failed(err); + } + self.state = PausedState::Answer(None); + continue; + } + }; + let update = match read { + UpdateRead::Update(update) => update, + UpdateRead::Cancelled => return Poll::Ready(Err(None)), + UpdateRead::Malformed(err) => { + if super::subscriber::is_protocol_violation(&err) { + publisher.session.close( + super::error::to_session_code(&crate::SessionError::from(&err)), + &err.to_string(), + ); + } + return failed(err); + } + UpdateRead::Failed(err) => return failed(err), + }; + if renewal.pending() { + if let Err(overflow) = renewal.queue(update, publisher.version) { + let err = match overflow { + request_update::Overflow::Session => request_update::too_many(&mut publisher.session.clone()), + request_update::Overflow::Request => Error::ProtocolViolation, + }; + return failed(err); + } + continue; + } + let paused = "REQUEST_UPDATE parameters not supported while paused"; + let refused = match update.refused { + Some(err) => Some((err, paused)), + None if update.filter.is_some() || update.fill.is_some() || update.new_group.is_some() => { + Some((Error::Unsupported, paused)) + } + None if update.authorization_token.is_some() + && grant.is_none() + && publisher.auth.token_decides(crate::auth::Direction::Publish) => + { Some((Error::Unsupported, TOKEN_ON_SESSION_REQUEST)) } None => None, }; if let Some((err, reason)) = refused { - self.write_subscribe_error(&mut stream.writer, &err, reason) - .await - .map_err(|err| Some(Served::Track(Err(err))))?; - return Err(Some(Served::UpdateFailed(Error::Unsupported))); + if let Err(err) = publisher.buffer_subscribe_error(&mut stream.writer, &err, reason) { + return failed(err); + } + self.state = PausedState::Answer(Some(Served::UpdateFailed(Error::Unsupported))); + continue; } if let Some(priority) = update.priority { msg.subscriber_priority = priority; @@ -2312,451 +2416,954 @@ impl Publisher { if let Some(max_age) = update.max_age { msg.max_age = Some(max_age); } - resuming = update.forward == Some(true); + self.resuming = update.forward == Some(true); if let (Some(token), Some(rg)) = ( update .authorization_token - .filter(|_| self.auth.token_decides(crate::auth::Direction::Publish)), + .filter(|_| publisher.auth.token_decides(crate::auth::Direction::Publish)), grant.as_ref(), ) { - renewal.verify(rg.verify_renewal(&self.auth, token.0)); + renewal.verify(rg.verify_renewal(&publisher.auth, token.0)); continue; } - if resuming { - return Ok(()); + if self.resuming { + return Poll::Ready(Ok(())); } - self.write_request_ok(&mut stream.writer, None) - .await - .map_err(|err| Some(Served::Track(Err(err))))?; + if let Err(err) = Publisher::buffer_request_ok(&mut stream.writer, None) { + return failed(err); + } + self.state = PausedState::Answer(None); } } +} - /// Read the requester's next REQUEST_UPDATE on its request stream, admitting its - /// Request ID. Pending once the requester finished sending where a FIN does not cancel. - fn poll_update( - &self, - stream: &mut Stream, - finished: &mut bool, - goaways: &mut super::request_stream::Messages, +/// An update [`Publisher::apply_subscription_update`] applied. +struct Applied { + /// Its answer is buffered, to flush before going on. + answered: bool, + /// How the request goes on: a fill to start, or how it ends. + then: std::result::Result, Served>, +} + +/// Ends a subscription on the wire with the PUBLISH_DONE its end reports, resolving to how +/// serving ended. The stream closes only once no verification holds the request's credit, +/// and once a [`Retain`] lets go of its join. +struct WritePublishDone { + /// How serving ended. + res: Option>, + /// Let go once the request closes or its authorization lapses, even while a peer holding + /// back flow control parks PUBLISH_DONE. + retain: Option, + state: PublishDoneState, +} + +enum PublishDoneState { + /// Flushing PUBLISH_DONE, unless it failed to buffer. + Flush { + stream: Stream, + buffered: bool, + }, + /// Waiting for the kept join and the request's credit to let go. + Release(Stream), + /// PUBLISH_DONE is the last thing on this stream, so it needs the acknowledgement too. + /// `None` when the stream failed to finish. + Close { + _reader: crate::coding::Reader, + closing: Option>, + }, + Done, +} + +impl WritePublishDone { + fn new(mut stream: Stream, served: Served, streams: u64, retain: Option) -> Self { + let (status, reason) = served.status(); + let res = match served { + Served::Track(res) => res, + Served::UpdateFailed(err) => Err(err), + Served::RenewalRefused(refusal) => Err(refusal.error), + Served::Paused(_) => Ok(()), + }; + let done = ietf::PublishDone { + status_code: status.code(), + stream_count: streams, + reason_phrase: reason.into(), + }; + let buffered = stream.writer.buffer_message(ietf::PublishDone::ID, &done).is_ok(); + Self { + res: Some(res), + retain, + state: PublishDoneState::Flush { stream, buffered }, + } + } + + fn poll( + &mut self, waiter: &kio::Waiter, - ) -> Poll { - use super::request_stream::{Next, Update}; - while !*finished { - match ready!(stream.reader.poll_decode_maybe::>(waiter)) { - Ok(Some(Next::Message(mut update))) => { - if let Err(err) = self.control.accept(update.request_id) { - if matches!(err, Error::GoingAway) { - update.refused = Some(err); - return Poll::Ready(UpdateRead::Update(update)); - } - self.session.clone().close( - super::error::to_session_code(&crate::SessionError::from(&err)), - "invalid request ID", - ); - return Poll::Ready(UpdateRead::Failed(err)); + publisher: &Publisher, + credit: &kio::Producer<()>, + ) -> Poll> { + loop { + match &mut self.state { + PublishDoneState::Flush { stream, buffered } => { + if let Some(kept) = &mut self.retain + && (stream.reader.poll_discard(waiter).is_ready() || kept.poll_lapsed(waiter).is_ready()) + { + self.retain = None; } - return Poll::Ready(UpdateRead::Update(update)); + if *buffered { + let _ = ready!(stream.writer.poll_flush(waiter)); + } + let PublishDoneState::Flush { stream, .. } = + std::mem::replace(&mut self.state, PublishDoneState::Done) + else { + unreachable!("matched above"); + }; + self.state = PublishDoneState::Release(stream); } - Ok(Some(Next::GoAway { redirect })) => { - if let Err(err) = goaways.receive(redirect) { - return Poll::Ready(UpdateRead::Malformed(err)); + PublishDoneState::Release(stream) => { + if let Some(retain) = &mut self.retain { + ready!(publisher.poll_retain_join(waiter, stream, retain)); + self.retain = None; + } + // The peer learns the subscription ended at once, but the stream holds its + // credit until the application lets go of every token it took for it. + ready!(publisher.poll_released(waiter, credit)); + let PublishDoneState::Release(stream) = std::mem::replace(&mut self.state, PublishDoneState::Done) + else { + unreachable!("matched above"); + }; + let Stream { writer, reader } = stream; + self.state = PublishDoneState::Close { + _reader: reader, + closing: writer.closing().ok(), + }; + } + PublishDoneState::Close { closing, .. } => { + if let Some(closing) = closing { + let _ = ready!(closing.poll(waiter)); } + self.state = PublishDoneState::Done; + return Poll::Ready(self.res.take().expect("resolves once")); } - Ok(None) if !super::request_stream::fin_cancels(self.version) => *finished = true, - Ok(None) => return Poll::Ready(UpdateRead::Cancelled), - Err(err) => return Poll::Ready(UpdateRead::Malformed(err)), + PublishDoneState::Done => panic!("polled after it resolved"), } } - Poll::Pending } +} - /// Apply one REQUEST_UPDATE on an in-flight subscription, with no verdict pending: refuse - /// unsupported parameters, apply the serving parameters, then verify its token or answer - /// at once. Each answer restores the peer's MAX_REQUEST_UPDATES credit. - /// - /// An error is how the request ends. - async fn apply_subscription_update( - &self, - stream: &mut Stream, - serve: &mut TrackServe, - request_grant: &Option, - renewal: &mut request_update::Renewal, - update: super::request_stream::Update, - ) -> Result, Served> { - let failed = |err| Served::Track(Err(err)); - // A token renews a token-authorized request. Where tokens decide, one arriving on a - // request admitted without one is refused rather than taken unread, so every token this - // side admits is verified: a request is admitted by the session or by a token, and that - // never changes. - let refused = match update.refused.clone() { - Some(Error::Decode(err)) => Some((super::token::reason(&err), Error::Decode(err))), - Some(Error::GoingAway) => Some(("going away", Error::GoingAway)), - Some(err) => Some(("REQUEST_UPDATE parameters not supported", err)), - // A SUBSCRIBE serves the range it asked for. A subscription we initiated with - // PUBLISH decodes those parameters itself. - None if update.fill.is_some() && !Filter::is_draft20(self.version) => { - Some(("REQUEST_UPDATE parameters not supported", Error::Unsupported)) - } - None if update.authorization_token.is_some() - && request_grant.is_none() - && self.auth.token_decides(crate::auth::Direction::Publish) => - { - Some((TOKEN_ON_SESSION_REQUEST, Error::Unsupported)) - } - // The draft forbids it on a track that did not advertise DYNAMIC_GROUPS. - None if update.new_group.is_some() && !serve.dynamic => Some(( - "NEW_GROUP_REQUEST on a track without DYNAMIC_GROUPS", - Error::Unsupported, - )), - // A paused subscription opens no groups, so one is asked for only as it resumes. - None if update.new_group.is_some() && update.forward == Some(false) => { - Some(("NEW_GROUP_REQUEST while pausing", Error::Unsupported)) +/// Serves the largest group's head on its own draft-22 fetch stream. +/// +/// A fill is a promise once requested. An empty range opens no stream, but a range we +/// cannot serve still opens one and resets it right after the FETCH_HEADER, the +/// draft's fill-failure signal. Nothing here touches the subscription either way. +/// +/// Counts its stream into `opened` once it opens, like a group stream, so +/// PUBLISH_DONE's Stream Count includes it however the subscription ends. It ranks in +/// its subscription's `ranking` as the group it fills. +struct RunFill { + fill: FillServe, + track: track::Consumer, + timescale: Option, + request_id: RequestId, + open: crate::transport::open::Open, + state: RunFillState, +} + +/// The fetch stream a fill writes, ranked as the group it fills. +type FillStream = Writer, Version>; + +#[kio::machine(self.state: RunFillState)] +impl RunFill { + fn new(fill: FillServe, track: track::Consumer, timescale: Option, request_id: RequestId) -> Self { + Self { + fill, + track, + timescale, + request_id, + open: Default::default(), + state: RunFillState::Start, + } + } + + fn poll( + &mut self, + waiter: &kio::Waiter, + publisher: &Publisher, + ranking: &Ranking>, + opened: &AtomicU64, + ) -> Poll<()> { + let priority: i8 = *ranking.priority.read(); + let mut rank: Option>> = match self.fill { + FillServe::Empty => return Poll::Ready(()), + FillServe::Group { sequence, skip, .. } => Some(ranking.rank(Location { + group: sequence, + object: skip, + })), + // Only the refusal goes out on its stream, so it competes with nothing. + FillServe::Unsupported => None, + }; + let opening = wait!(Open => publisher.session.poll_open_uni( + &mut self.open, + waiter, + rank.as_mut().map(Rank::options).unwrap_or_default(), + )); + let stream = match opening { + Ok(stream) => stream, + Err(err) => { + tracing::debug!(err = %Error::from_transport(err), fill = %self.request_id, "fill stream failed to open"); + return Poll::Ready(()); } - None => None, }; - if let Some((reason, err)) = refused { - self.write_subscribe_error(&mut stream.writer, &err, reason) - .await - .map_err(failed)?; - // The update failed, which ends the subscription. - return Err(Served::UpdateFailed(Error::Unsupported)); + opened.fetch_add(1, Ordering::Relaxed); + let mut writer: FillStream = Writer::new(Ranked::new(stream, rank).lapsing(&ranking.lapse), publisher.version); + + if let Err(err) = writer.buffer_varint(FetchHeader::TYPE) { + return self.fail(writer, err); } - let fill = if update.filter.is_some() || update.fill.is_some() { - serve.update_range(&update).map_err(failed)? - } else { - None - }; - let largest = fill.as_ref().and_then(|fill| fill.largest); - if let Some(priority) = update.priority { - serve - .set_priority(super::priority::from_wire(priority)) - .map_err(failed)?; + let typed = wait!(Type => writer.poll_flush(waiter)); + if let Err(err) = typed.and_then(|()| { + writer.buffer(&FetchHeader { + request_id: self.request_id, + }) + }) { + return self.fail(writer, err); } - if let Some(max_age) = update.max_age { - let mut subscription = serve.track.subscription(); - subscription.delay.max = Some(Duration::from_millis(max_age)); - serve.track.update(subscription).map_err(failed)?; - } - if let Some(value) = update.new_group { - serve.request_new_group(value); + let headed = wait!(Header => writer.poll_flush(waiter)); + if let Err(err) = headed { + return self.fail(writer, err); } - if let (Some(token), Some(grant)) = ( - update - .authorization_token - .filter(|_| self.auth.token_decides(crate::auth::Direction::Publish)), - request_grant, - ) { - renewal.verify(grant.verify_renewal(&self.auth, token.0)); - return Ok(fill); + + // Resolved already: the subscription this fills described the track. + let resolving: track::Resolving = self.track.info().into_inner(); + let resolved = wait!(Info => resolving.poll_ok(waiter)); + let publisher_priority: u8 = match resolved { + Ok(info) => super::priority::to_wire(info.priority), + Err(err) => return self.fail(writer, err), + }; + let FillServe::Group { sequence, skip, until } = self.fill else { + return self.fail(writer, Error::Unsupported); + }; + let fetching: kio::Pending = self.track.fetch_group( + sequence, + group::Fetch { + priority, + ..Default::default() + }, + ); + let filled: u64 = sequence; + let slice: GroupSlice = GroupSlice { skip, until }; + let fetched = wait!(Fetch => kio::Task::poll(&mut **fetching, waiter)); + let group: group::Consumer = match fetched { + Ok(group) => group, + Err(err) => return self.fail(writer, err), + }; + let setup = wait!(Setup => publisher.peer_setup.poll_get(waiter)); + let compressed = if setup.flate { group.compressed() } else { group }; + let (timescale, version) = (self.timescale, publisher.version); + let started = FetchGroupWrite::new(compressed, filled, slice, timescale, version, publisher_priority); + let mut write: FetchGroupWrite = match started { + Ok(write) => write, + Err(err) => return self.fail(writer, err), + }; + let done = wait!(Write => write.poll(waiter, writer)); + if let Err(err) = done { + return self.fail(writer, err); + } + // Keeps the stream, and its rank, until the acknowledgement: unlike [`Writer::close`], + // cancelling resets it as a cancelled group stream's is. + if let Err(err) = writer.finish() { + tracing::debug!(%err, fill = %self.request_id, "fill stream close failed"); + return Poll::Ready(()); } - // The caller stops forwarding, then answers. - if update.forward == Some(false) { - return Err(Served::Paused(true)); + let closed = wait!(Settle => writer.poll_close(waiter)); + match closed { + Ok(()) => tracing::debug!(fill = %self.request_id, "fill complete"), + Err(err) => tracing::debug!(%err, fill = %self.request_id, "fill stream close failed"), } - self.write_request_ok(&mut stream.writer, largest) - .await - .map_err(failed)?; - Ok(fill) + Poll::Ready(()) } - /// Answer a REQUEST_UPDATE with REQUEST_OK, carrying the `largest` Location a resumed - /// subscription now starts from. - async fn write_request_ok( - &self, - writer: &mut Writer, - largest: Option, - ) -> Result<(), Error> { - writer.encode_message(&ietf::RequestOk { largest, active: None }).await + /// Reset the fill's `stream` with `err`, the draft's fill-failure signal. + fn fail(&self, stream: FillStream, err: Error) -> Poll<()> { + tracing::debug!(%err, fill = %self.request_id, "fill failed, resetting its stream"); + stream.abort(&err); + Poll::Ready(()) } +} - /// Reject a SUBSCRIBE, ending the request stream. - /// - /// Takes the whole stream because delivering the error is the other half of the job: - /// [`Writer`] resets on drop, and a reset discards data the peer has not acknowledged, so - /// returning here without its acknowledgement leaves the subscriber waiting on a request - /// we already refused. - async fn reject_subscribe( - &self, - mut stream: Stream, - err: &Error, - reason: &str, - ) -> Result<(), Error> { - let kind = request::Kind::Subscribe; - super::request_stream::buffer_error(&mut stream.writer, kind, err, reason, self.version)?; - super::request_stream::deliver(&mut stream.writer).await - } +/// A [`RunFill`] with the subscription state it ranks and counts in, so it runs on its own. +struct Fill { + run: RunFill, + ranking: Ranking>, + opened: Arc, +} - /// Write a subscribe error on the bidi stream writer. - async fn write_subscribe_error( - &self, - writer: &mut Writer, - err: &Error, - reason: &str, - ) -> Result<(), Error> { - super::request_stream::write_error(writer, request::Kind::Subscribe, err, reason, self.version).await +impl Fill { + fn poll(&mut self, waiter: &kio::Waiter, publisher: &Publisher) -> Poll<()> { + self.run.poll(waiter, publisher, &self.ranking, &self.opened) } +} - /// Serve the largest group's head on its own draft-22 fetch stream. - /// - /// A fill is a promise once requested. An empty range opens no stream, but a range we - /// cannot serve still opens one and resets it right after the FETCH_HEADER, the - /// draft's fill-failure signal. Nothing here touches the subscription either way. - /// - /// Counts its stream into `opened` once it opens, like a group stream, so - /// PUBLISH_DONE's Stream Count includes it however the subscription ends. It ranks in - /// its subscription's `ranking` as the group it fills. - async fn run_fill( - &self, - fill: FillServe, - ranking: &Ranking>, - track: track::Consumer, - timescale: Option, - opened: &AtomicU64, - request_id: RequestId, - ) { - let priority = *ranking.priority.read(); - let rank = match fill { - FillServe::Empty => return, - FillServe::Group { sequence, skip, .. } => Some(ranking.rank(Location { - group: sequence, - object: skip, - })), - // Only the refusal goes out on its stream, so it competes with nothing. - FillServe::Unsupported => None, - }; - - let mut rank = rank; - let options = rank.as_mut().map(Rank::options).unwrap_or_default(); - let session = self.session.clone(); - let stream = match session.open_uni(options).await { - Ok(stream) => stream, - Err(err) => { - tracing::debug!(err = %Error::from_transport(err), fill = %request_id, "fill stream failed to open"); - return; - } +/// Deliver the refusal `buffered` on `stream`, all a refused request carries, logging a +/// failure as `failed`. A counted request stays in `serve` until it ends. +fn deliver_refusal( + mut stream: Stream, + buffered: Result<(), Error>, + mut serve: Option, + failed: &'static str, +) -> crate::util::BoxTask { + let mut buffered = buffered.err(); + let mut deliver = super::request_stream::Deliver::new(); + Box::new(move |waiter: &kio::Waiter| { + let res = match buffered.take() { + Some(err) => Err(err), + None => ready!(deliver.poll(waiter, &mut stream.writer)), }; - opened.fetch_add(1, Ordering::Relaxed); - let mut stream = Writer::new(Ranked::new(stream, rank).lapsing(&ranking.lapse), self.version); - - let res = async { - stream.varint(FetchHeader::TYPE).await?; - stream.encode(&FetchHeader { request_id }).await?; - // Resolved already: the subscription this fills described the track. - let publisher_priority = super::priority::to_wire(track.info().await?.priority); + if let Err(err) = res { + tracing::debug!(%err, "{failed}"); + } + drop(serve.take()); + Poll::Ready(()) + }) +} - let FillServe::Group { sequence, skip, until } = fill else { - return Err(Error::Unsupported); - }; - let group = track - .fetch_group( - sequence, - group::Fetch { - priority, - ..Default::default() - }, - ) - .await?; - let group = if self.peer_setup.get().await.flate { - group.compressed() - } else { - group - }; - Self::write_fetch_group( - &mut stream, - group, - sequence, - GroupSlice { skip, until }, - timescale, - self.version, - publisher_priority, - ) - .await?; - Ok(()) +impl Publisher { + pub fn new( + runtime: crate::time::Deadlines, + session: Transport, + origin: origin::Consumer, + control: Control, + peer_hop: Option, + peer_setup: peer::PeerSetup, + version: Version, + ) -> Self { + Self { + withdrawal: Default::default(), + runtime, + session, + self_origin: origin.hop(), + // Serving the peer is not a local reader asking to discover anything. + origin: origin.serving(), + control, + peer_hop, + peer_setup, + joins: Default::default(), + version, + auth: crate::auth::Link::new(false), + owed: Default::default(), + subscriptions: Default::default(), + priority: Default::default(), } - .await; + } - match res { - Ok(()) => { - if let Err(err) = settle(&mut stream).await { - tracing::debug!(%err, fill = %request_id, "fill stream close failed"); - } else { - tracing::debug!(fill = %request_id, "fill complete"); - } - } - Err(err) => { - tracing::debug!(%err, fill = %request_id, "fill failed, resetting its stream"); - stream.abort(&err); - } + /// Mark a paused subscription for joining FETCHes: draft-18 section 5.1 refuses one + /// unless the subscription forwards. + fn pause_join(&self, join: Option<&Join>>) { + if let Some(join) = join { + join.joins.lock().insert(join.request_id, Some(Joined::Empty)); } } - /// Write one group's frames in the negotiated draft's FETCH object layout. + /// Bound what we publish by the grant this session's tokens earn (MoQ Auth). + pub fn with_auth(mut self, auth: crate::auth::Link) -> Self { + self.auth = auth; + self + } + + /// What the peer declared in its SETUP. /// - /// The first object carries its absolute Group and Object IDs plus the priority; - /// every later one inherits them and increments the Object ID, so only the - /// properties (the timestamp) and the payload go on the wire. A fetch object has no - /// status field. - async fn write_fetch_group( - stream: &mut Writer, Version>, - mut group: group::Consumer, - sequence: u64, - slice: GroupSlice, - timescale: Option, - version: Version, - priority: u8, - ) -> Result<(), Error> { - let GroupSlice { skip, until } = slice; - let mut index: u64 = 0; - let mut prior = None; + /// Resolves only once the peer's SETUP arrives, because the extension changes the + /// NAMESPACE encoding: nothing can be advertised until we know whether the peer speaks it. + #[cfg_attr(feature = "async", kio::wrapper)] + fn poll_peer(&self, waiter: &kio::Waiter) -> Poll { + self.peer_setup.poll_get(waiter).map(|peer| peer.cluster) + } - let mut buf = arrayvec::ArrayVec::::new(); - if group.wire().is_some_and(|wire| !wire.fetchable) { - return Err(Error::Unsupported); - } - // A long cached group is written a slice per poll, so the session's other tasks still run. - let mut budget = kio::coop::Budget::new(32); - 'serve: loop { - kio::wait(|waiter| budget.poll_yield(waiter)).await; - // The cap is the Largest Object snapshot: the group may keep growing, but - // everything past the snapshot belongs to the subscription, not the fill. - if until.is_some_and(|until| index >= until) { - break; - } + /// Whether the peer requires advertisements to be solicited, from the same SETUP. + /// + /// Blocks on it for the same reason [`Self::peer`] does: this decides whether the + /// first advertisement is sent unasked, so it cannot be guessed and corrected later. + fn poll_requires_solicitation(&self, waiter: &kio::Waiter) -> Poll { + self.peer_setup + .poll_get(waiter) + .map(|peer| peer.solicit.unwrap_or(false)) + } - let step = { - let mut closed = std::pin::pin!(stream.closed()); - kio::wait(|waiter| { - if waiter.poll_future(closed.as_mut()).is_ready() { - return Poll::Ready(Err(Error::Cancel)); - } - match group.poll_read_objects(waiter, &mut buf) { - Poll::Pending => match group.poll_next_frame(waiter) { - Poll::Ready(Ok(Some(frame))) => Poll::Ready(Ok(FillStep::Partial(frame))), - Poll::Ready(Ok(None)) => Poll::Ready(Ok(FillStep::Done)), - Poll::Ready(Err(err)) => Poll::Ready(Err(err)), - Poll::Pending => Poll::Pending, - }, - Poll::Ready(Ok(0)) => Poll::Ready(Ok(FillStep::Done)), - Poll::Ready(Ok(_)) => Poll::Ready(Ok(FillStep::Batch)), - Poll::Ready(Err(err)) => Poll::Ready(Err(err)), - } - }) - .await - }; + /// The origin to serve this peer's subscriptions from: sources whose hop chain flows + /// through the peer are excluded, so a subscription is never handed data that already + /// flowed through the subscriber. + /// + /// The same exclusion the announce path applies (see [`Self::select`]), which is what + /// keeps advertised paths truthful and prevents subscription cycles of any length. + pub(super) fn poll_serving_origin(&self, waiter: &kio::Waiter) -> Poll { + self.poll_peer(waiter).map(|peer| self.excluding(&peer)) + } - match step? { - FillStep::Batch => { - for frame in &buf { - if until.is_some_and(|until| index >= until) { - break 'serve; - } - let frame = frame.clone(); - index = frame.fields.as_ref().map_or(index, |fields| fields.id); - if until.is_some_and(|until| index >= until) { - break 'serve; - } - if index >= skip { - Self::write_fetch_object( - stream, - FetchObjectInfo { - sequence, - object: index, - prior: fetch_prior(prior, index)?, - priority, - timestamp: frame.timestamp, - inflated: frame.inflated, - timescale, - version, - fields: frame.fields.as_deref(), - source: group.ietf(), - }, - ) - .await?; - Self::write_fetch_payload(stream, frame.frame.payload).await?; - prior = Some(index); - } - index += 1; - group.keep_alive(); - } - } - FillStep::Partial(mut frame) => { - index = frame.fields.as_ref().map_or(index, |fields| fields.id); - if until.is_some_and(|until| index >= until) { - break; - } - if index < skip { - // A skipped frame still has to be drained to advance the cursor. - loop { - let chunk = { - let mut closed = std::pin::pin!(stream.closed()); - kio::wait(|waiter| { - if waiter.poll_future(closed.as_mut()).is_ready() { - return Poll::Ready(Err(Error::Cancel)); - } - frame.poll_read_chunk(waiter) - }) - .await - }; - if chunk?.is_none() { - break; - } - } - index += 1; - continue; - } + /// Resolve a request's broadcast while the peer still wants the answer, a FIN withdrawing it + /// where `fin_cancels`. + #[cfg(test)] + async fn resolve( + &self, + stream: &mut Stream, + fin_cancels: bool, + path: &crate::Path, + ) -> Result, Error> { + let mut resolve = Resolve::new(self.session.clone(), path.clone(), fin_cancels); + kio::wait(|waiter| resolve.poll(waiter, self, stream)).await + } - Self::write_fetch_object( - stream, - FetchObjectInfo { - sequence, - object: index, - prior: fetch_prior(prior, index)?, - priority, - timestamp: frame.timestamp, - inflated: frame.inflated, - timescale, - version, - fields: frame.fields.as_deref(), - source: group.ietf(), - }, - ) - .await?; - prior = Some(index); - index += 1; + /// Our origin handle with [`Self::exclude`] applied, the view both the data plane + /// and the announce loops read this peer's routes through. + fn excluding(&self, peer: &cluster::Peer) -> origin::Consumer { + self.origin.clone().excluding(self.exclude(peer)) + } - stream.varint(frame.size).await?; - loop { - let chunk = { - let mut closed = std::pin::pin!(stream.closed()); - kio::wait(|waiter| { - if waiter.poll_future(closed.as_mut()).is_ready() { - return Poll::Ready(Err(Error::Cancel)); - } - frame.poll_read_chunk(waiter) - }) - .await - }; + /// `origin` as advertised to this peer: never its own routes back (split horizon), + /// and under MoQ Cluster only the routes with room for our hop under the origin's + /// cap, so selection falls back to a shorter route instead of withdrawing. + fn advertising(&self, origin: origin::Consumer, peer: &cluster::Peer) -> origin::Consumer { + let origin = origin.excluding(self.exclude(peer)); + match peer.negotiated() { + true => origin.forwardable(), + false => origin, + } + } - match chunk? { - Some(mut chunk) => stream.write_all(&mut chunk).await?, - None => break, - } - } - } - FillStep::Done => break, - } + /// The Hop ID whose paths must not be advertised (or served) back to this peer. + /// + /// A peer that declared an identity supplies its own; otherwise fall back to the one + /// we assigned it (a fresh per-session id, or one pinned with `Client::with_peer_hop` + /// or `Request::with_peer_hop`), since moq-transport carries no identity + /// of its own. A peer that declared the reserved 0 declared no identity, so it takes + /// the fallback like any other anonymous peer. + fn exclude(&self, peer: &cluster::Peer) -> crate::Hop { + peer.identity().or(self.peer_hop).unwrap_or(crate::Hop::UNKNOWN) + } + + /// Pick what to advertise to this peer for one route. + /// + /// The origin's announce cursor already filters routes through the peer + /// (control-plane split horizon); this only shapes what the wire can carry. + fn select(&self, route: &crate::origin::Route, peer: &cluster::Peer) -> Advert { + // A route that already passed through us is a reflection. The origin + // filters these on receive, so this is defensive. + if self.self_origin != crate::Hop::UNKNOWN && route.hops.contains(&self.self_origin) { + return Advert::None; } - if until.is_some_and(|until| index < until) { - return Err(Error::NotFound); + if !peer.negotiated() { + return Advert::Some(ietf::AdvertParams { cluster: None }); } - Ok(()) - } - /// Write one fetch object's header and timestamp properties in the negotiated layout. - async fn write_fetch_object( - stream: &mut Writer, Version>, - header: FetchObjectInfo<'_>, + let cost = route.cost.value(); + // Our own Hop ID is always the last entry, so the peer reconstructs the full + // path. A chain with no room left is a loop in all but name. + match cluster::Advert::forward(&route.hops, cost, self.self_origin) { + Ok(advert) => Advert::Some(ietf::AdvertParams { cluster: Some(advert) }), + Err(_) => Advert::None, + } + } + + /// The task serving an incoming bidi stream dispatched by the session. + pub(super) fn handle_stream( + &self, + id: u64, + body: ietf::Body, + stream: Stream, + ) -> Result { + let this = self.clone(); + let mut data = body.decoder(this.version); + // Count at dispatch so close cannot pass a request waiting for its first poll. + let serve = matches!(id, ietf::Subscribe::ID | ietf::Fetch::ID).then(|| { + self.owed.fetch_add(1, Ordering::Relaxed); + Serve { + owed: self.owed.clone(), + _working: self.session.work().start(), + } + }); + let task: crate::util::BoxTask = match id { + ietf::Subscribe::ID => { + let msg = match ietf::Subscribe::decode_msg(&mut data, this.version) { + Err(err) if super::token::refuses_request(&err) => { + // Junk after the fields is fatal, which a token refusal must not hide. + if !data.is_empty() { + return Err(Error::WrongSize); + } + return Ok(this.reject_token(stream, request::Kind::Subscribe, err, serve)); + } + msg => msg?, + }; + if !data.is_empty() { + return Err(Error::WrongSize); + } + tracing::debug!(message = ?msg, "received subscribe"); + let gate = crate::auth::Gate::new( + this.auth.clone(), + msg.track_namespace.clone(), + crate::auth::Direction::Publish, + ); + let mut task = this.serve_subscribe_stream(stream, msg.into_owned(), gate); + let mut serve = serve; + Box::new(move |waiter: &kio::Waiter| { + if let Err(err) = ready!(task.poll(waiter)) { + tracing::debug!(%err, "subscribe stream error"); + } + drop(serve.take()); + Poll::Ready(()) + }) + } + ietf::Fetch::ID => { + let msg = match ietf::Fetch::decode_msg(&mut data, this.version) { + Err(err) if super::token::refuses_request(&err) => { + // Junk after the fields is fatal, which a token refusal must not hide. + if !data.is_empty() { + return Err(Error::WrongSize); + } + return Ok(this.reject_token(stream, request::Kind::Fetch, err, serve)); + } + msg => msg?, + }; + if !data.is_empty() { + return Err(Error::WrongSize); + } + tracing::debug!(message = ?msg, "received fetch"); + let gate = this.fetch_gate(&msg); + // Owned, since the task outlives the body it was read from. + let msg = msg.into_owned(); + let mut task = FetchServe::new(this, stream, msg, gate); + let mut serve = serve; + Box::new(move |waiter: &kio::Waiter| { + if let Err(err) = ready!(task.poll(waiter)) { + tracing::debug!(%err, "fetch stream error"); + } + drop(serve.take()); + Poll::Ready(()) + }) + } + ietf::SubscribeNamespace::ID => { + let msg = match ietf::SubscribeNamespace::decode_msg(&mut data, this.version) { + Err(err) if super::token::refuses_request(&err) => { + // Junk after the fields is fatal, which a token refusal must not hide. + if !data.is_empty() { + return Err(Error::WrongSize); + } + let kind = request::Kind::SubscribeNamespace; + return Ok(this.reject_token(stream, kind, err, serve)); + } + decoded => decoded?, + }; + if !data.is_empty() { + return Err(Error::WrongSize); + } + tracing::debug!(message = ?msg, "received subscribe_namespace"); + let gate = crate::auth::Gate::new( + this.auth.clone(), + msg.namespace.clone(), + crate::auth::Direction::Publish, + ); + let session = this.session.clone(); + let mut task = NamespaceServe::new(this, stream, msg, gate); + Box::new(move |waiter: &kio::Waiter| { + match ready!(task.poll(waiter)) { + // Such as a malformed answer to an advertisement this request solicited. + Err(err) if super::subscriber::is_protocol_violation(&err) => session.close( + super::error::to_session_code(&crate::SessionError::from(&err)), + &err.to_string(), + ), + Err(err) => tracing::debug!(%err, "subscribe_namespace stream error"), + Ok(()) => {} + } + Poll::Ready(()) + }) + } + ietf::SubscribeTracks::ID => { + let msg = match ietf::SubscribeTracks::decode_msg(&mut data, this.version) { + Err(err) if super::token::refuses_request(&err) => { + // Junk after the fields is fatal, which a token refusal must not hide. + if !data.is_empty() { + return Err(Error::WrongSize); + } + let kind = request::Kind::SubscribeNamespace; + return Ok(this.reject_token(stream, kind, err, serve)); + } + msg => msg?, + }; + if !data.is_empty() { + return Err(Error::WrongSize); + } + tracing::debug!(message = ?msg, "received subscribe_tracks"); + // We list no tracks by PUBLISH, so each listing is refused on its own and the + // session stays up. Decoded first, so a refusal still admits its Request ID. + let error = + super::namespace::refusal(&msg.prefix, this.version).unwrap_or_else(|| ietf::RequestError { + error_code: request::NOT_SUPPORTED, + reason_phrase: "SUBSCRIBE_TRACKS is not supported".into(), + retry_interval: 0, + }); + let mut stream = stream; + let buffered = stream.writer.buffer_message(ietf::RequestError::ID, &error); + deliver_refusal(stream, buffered, serve, "subscribe_tracks refusal failed") + } + ietf::TrackStatus::ID => { + let msg = match ietf::TrackStatus::decode_msg(&mut data, this.version) { + Err(err) if super::token::refuses_request(&err) => { + // Junk after the fields is fatal, which a token refusal must not hide. + if !data.is_empty() { + return Err(Error::WrongSize); + } + return Ok(this.reject_token(stream, request::Kind::TrackStatus, err, serve)); + } + msg => msg?, + }; + if !data.is_empty() { + return Err(Error::WrongSize); + } + tracing::debug!(message = ?msg, "received track_status"); + let gate = crate::auth::Gate::new( + this.auth.clone(), + msg.track_namespace.clone(), + crate::auth::Direction::Publish, + ); + // Owned, since the task outlives the body it was read from. + let msg = msg.into_owned(); + let mut task = TrackStatusServe::new(this, stream, msg, gate); + Box::new(move |waiter: &kio::Waiter| { + if let Err(err) = ready!(task.poll(waiter)) { + tracing::debug!(%err, "track_status stream error"); + } + Poll::Ready(()) + }) + } + _ => { + tracing::warn!(id, "unexpected bidi stream type for publisher"); + return Err(Error::UnexpectedStream); + } + }; + Ok(task) + } + + /// Refuse a request whose AUTHORIZATION TOKEN repeats (section 8.9), which unlike any + /// other decode error leaves the session open. + fn reject_token( + &self, + mut stream: Stream, + kind: request::Kind, + err: DecodeError, + serve: Option, + ) -> crate::util::BoxTask { + let reason = super::token::reason(&err); + let buffered = + super::request_stream::buffer_error(&mut stream.writer, kind, &Error::Decode(err), reason, self.version); + deliver_refusal(stream, buffered, serve, "token refusal failed") + } + + /// Handle a SUBSCRIBE on its bidi stream. + #[cfg(test)] + fn run_subscribe_stream( + self, + stream: Stream, + msg: ietf::Subscribe<'static>, + ) -> impl std::future::Future> { + let gate = crate::auth::Gate::new( + self.auth.clone(), + msg.track_namespace.clone(), + crate::auth::Direction::Publish, + ); + let mut serve = self.serve_subscribe_stream(stream, msg, gate); + kio::wait(move |waiter| serve.poll(waiter)) + } + + /// Serve a SUBSCRIBE on its bidi stream. + fn serve_subscribe_stream( + self, + stream: Stream, + msg: ietf::Subscribe<'static>, + gate: crate::auth::Gate, + ) -> SubscribeServe { + // Register during dispatch, before either request task can be polled. + let join = (!Filter::is_draft20(self.version)).then(|| self.register_join(msg.request_id)); + let goaways = super::request_stream::Messages::new(self.session.clone(), self.control.client()); + SubscribeServe { + request: SubscribeRequest { + publisher: self, + stream: Some(stream), + msg, + gate: Some(gate), + join, + credit: Default::default(), + slot: None, + request_finished: false, + goaways, + grant: None, + updates: Default::default(), + acked: None, + track_ended: false, + lapse: Lapse::default(), + }, + state: SubscribeState::Start, + } + } + + /// Ready once the requester withdraws the request on `stream` or `session` closes, without + /// reading anything it sent meanwhile. + fn poll_withdrawn( + &self, + session: &mut Transport, + stream: &mut Stream, + fin_cancels: bool, + waiter: &kio::Waiter, + ) -> Poll<()> { + if session.poll_closed(waiter).is_ready() { + return Poll::Ready(()); + } + super::request_stream::poll_withdrawn(stream, fin_cancels, waiter) + } + + /// Poll `step` of a request's setup, or `None` once [`Self::poll_withdrawn`] is ready. + fn poll_until_withdrawn( + &self, + stream: &mut Stream, + fin_cancels: bool, + waiter: &kio::Waiter, + step: impl FnOnce(&kio::Waiter) -> Poll, + ) -> Poll> { + if self.session.poll_closed(waiter).is_ready() + || super::request_stream::poll_withdrawn(stream, fin_cancels, waiter).is_ready() + { + return Poll::Ready(None); + } + step(waiter).map(Some) + } + + /// Ready once no token verification holds `credit`'s request any more, or the session + /// closes. + #[cfg_attr(feature = "async", kio::wrapper)] + fn poll_released(&self, waiter: &kio::Waiter, credit: &kio::Producer<()>) -> Poll<()> { + if self.session.poll_closed(waiter).is_ready() || credit.poll_wanted_changed(waiter, true).is_ready() { + return Poll::Ready(()); + } + Poll::Pending + } + + /// Ready once the subscriber closes a finished subscription's request, the session closes, + /// or what authorized the subscription stops covering it, ending the join `retain` keeps. + fn poll_retain_join( + &self, + waiter: &kio::Waiter, + stream: &mut Stream, + retain: &mut Retain, + ) -> Poll<()> { + if self.session.poll_closed(waiter).is_ready() || stream.writer.poll_closed(waiter).is_ready() { + return Poll::Ready(()); + } + // An update that crossed PUBLISH_DONE is skipped, since nothing answers it now; the + // subscriber's FIN or a reset ends the request. + if stream.reader.poll_discard(waiter).is_ready() { + return Poll::Ready(()); + } + retain.poll_lapsed(waiter) + } + + /// Await [`RunSubscription`] with no fill. + #[cfg(test)] + async fn run_subscription( + &self, + stream: &mut Stream, + serve: &mut TrackServe, + finished: &mut bool, + goaways: &mut super::request_stream::Messages, + updates: &mut SubscriptionUpdates, + ) -> Option { + let mut run = RunSubscription::new(None); + kio::wait(|waiter| run.poll(waiter, self, stream, serve, finished, goaways, updates)).await + } + + /// Read the requester's next REQUEST_UPDATE on its request stream, admitting its + /// Request ID. Pending once the requester finished sending where a FIN does not cancel. + fn poll_update( + &self, + stream: &mut Stream, + finished: &mut bool, + goaways: &mut super::request_stream::Messages, + waiter: &kio::Waiter, + ) -> Poll { + use super::request_stream::{Next, Update}; + while !*finished { + match ready!(stream.reader.poll_decode_maybe::>(waiter)) { + Ok(Some(Next::Message(mut update))) => { + if let Err(err) = self.control.accept(update.request_id) { + if matches!(err, Error::GoingAway) { + update.refused = Some(err); + return Poll::Ready(UpdateRead::Update(update)); + } + self.session.clone().close( + super::error::to_session_code(&crate::SessionError::from(&err)), + "invalid request ID", + ); + return Poll::Ready(UpdateRead::Failed(err)); + } + return Poll::Ready(UpdateRead::Update(update)); + } + Ok(Some(Next::GoAway { redirect })) => { + if let Err(err) = goaways.receive(redirect) { + return Poll::Ready(UpdateRead::Malformed(err)); + } + } + Ok(None) if !super::request_stream::fin_cancels(self.version) => *finished = true, + Ok(None) => return Poll::Ready(UpdateRead::Cancelled), + Err(err) => return Poll::Ready(UpdateRead::Malformed(err)), + } + } + Poll::Pending + } + + /// Apply one REQUEST_UPDATE on an in-flight subscription, with no verdict pending: refuse + /// unsupported parameters, apply the serving parameters, then verify its token or buffer + /// its answer at once. Each answer restores the peer's MAX_REQUEST_UPDATES credit. + /// + /// An error, or an [`Applied`] that goes on with a [`Served`], is how the request ends. + fn apply_subscription_update( + &self, + writer: &mut Writer, + serve: &mut TrackServe, + request_grant: &Option, + renewal: &mut request_update::Renewal, + update: super::request_stream::Update, + ) -> Result { + // A token renews a token-authorized request. Where tokens decide, one arriving on a + // request admitted without one is refused rather than taken unread, so every token this + // side admits is verified: a request is admitted by the session or by a token, and that + // never changes. + let refused = match update.refused.clone() { + Some(Error::Decode(err)) => Some((super::token::reason(&err), Error::Decode(err))), + Some(Error::GoingAway) => Some(("going away", Error::GoingAway)), + Some(err) => Some(("REQUEST_UPDATE parameters not supported", err)), + // A SUBSCRIBE serves the range it asked for. A subscription we initiated with + // PUBLISH decodes those parameters itself. + None if update.fill.is_some() && !Filter::is_draft20(self.version) => { + Some(("REQUEST_UPDATE parameters not supported", Error::Unsupported)) + } + None if update.authorization_token.is_some() + && request_grant.is_none() + && self.auth.token_decides(crate::auth::Direction::Publish) => + { + Some((TOKEN_ON_SESSION_REQUEST, Error::Unsupported)) + } + // The draft forbids it on a track that did not advertise DYNAMIC_GROUPS. + None if update.new_group.is_some() && !serve.dynamic => Some(( + "NEW_GROUP_REQUEST on a track without DYNAMIC_GROUPS", + Error::Unsupported, + )), + // A paused subscription opens no groups, so one is asked for only as it resumes. + None if update.new_group.is_some() && update.forward == Some(false) => { + Some(("NEW_GROUP_REQUEST while pausing", Error::Unsupported)) + } + None => None, + }; + if let Some((reason, err)) = refused { + self.buffer_subscribe_error(writer, &err, reason)?; + // The update failed, which ends the subscription. + return Ok(Applied { + answered: true, + then: Err(Served::UpdateFailed(Error::Unsupported)), + }); + } + let fill = if update.filter.is_some() || update.fill.is_some() { + serve.update_range(&update)? + } else { + None + }; + let largest = fill.as_ref().and_then(|fill| fill.largest); + if let Some(priority) = update.priority { + serve.set_priority(super::priority::from_wire(priority))?; + } + if let Some(max_age) = update.max_age { + let mut subscription = serve.track.subscription(); + subscription.delay.max = Some(Duration::from_millis(max_age)); + serve.track.update(subscription)?; + } + if let Some(value) = update.new_group { + serve.request_new_group(value); + } + if let (Some(token), Some(grant)) = ( + update + .authorization_token + .filter(|_| self.auth.token_decides(crate::auth::Direction::Publish)), + request_grant, + ) { + renewal.verify(grant.verify_renewal(&self.auth, token.0)); + return Ok(Applied { + answered: false, + then: Ok(fill), + }); + } + // The caller stops forwarding, then answers. + if update.forward == Some(false) { + return Ok(Applied { + answered: false, + then: Err(Served::Paused(true)), + }); + } + Self::buffer_request_ok(writer, largest)?; + Ok(Applied { + answered: true, + then: Ok(fill), + }) + } + + /// Buffer the REQUEST_OK answering a REQUEST_UPDATE, carrying the `largest` Location a + /// resumed subscription now starts from. + fn buffer_request_ok( + writer: &mut Writer, + largest: Option, + ) -> Result<(), Error> { + writer.buffer_message(ietf::RequestOk::ID, &ietf::RequestOk { largest, active: None }) + } + + /// Refuse a SUBSCRIBE, returning the delivery that ends its request stream. + /// + /// Delivering the error is the other half of the job: [`Writer`] resets on drop, and a + /// reset discards data the peer has not acknowledged, so ending without its + /// acknowledgement leaves the subscriber waiting on a request we already refused. + fn reject_subscribe( + &self, + writer: &mut Writer, + err: &Error, + reason: &str, + ) -> Result { + self.buffer_subscribe_error(writer, err, reason)?; + Ok(super::request_stream::Deliver::new()) + } + + /// Buffer a subscribe error on the bidi stream writer. + fn buffer_subscribe_error( + &self, + writer: &mut Writer, + err: &Error, + reason: &str, + ) -> Result<(), Error> { + super::request_stream::buffer_error(writer, request::Kind::Subscribe, err, reason, self.version) + } + + /// Await [`RunFill`]. + #[cfg(test)] + async fn run_fill( + &self, + fill: FillServe, + ranking: &Ranking>, + track: track::Consumer, + timescale: Option, + opened: &AtomicU64, + request_id: RequestId, + ) { + let mut run = RunFill::new(fill, track, timescale, request_id); + kio::wait(|waiter| run.poll(waiter, self, ranking, opened)).await + } + + /// Buffer one fetch object's header and timestamp properties in the negotiated layout. + fn buffer_fetch_object( + stream: &mut Writer, Version>, + header: FetchObjectInfo<'_>, ) -> Result<(), Error> { let FetchObjectInfo { sequence, @@ -2812,2022 +3419,2959 @@ impl Publisher { }, }; - stream.encode(&header).await?; + stream.buffer(&header) + } - Ok(()) + /// Register a pending subscription until the returned guard drops. + fn register_join(&self, request_id: RequestId) -> Join> { + self.joins.lock().insert(request_id, None); + Join { + joins: self.joins.clone(), + request_id, + } } - /// Write one group's objects to a fetch stream, which carries no other group. - async fn write_fetched_group( - writer: &mut Writer, Version>, - group: FetchedGroup, - version: Version, - timescale: Option, + /// Answer a FETCH: a standalone or filtered range of the named track, or a joining + /// FETCH's groups up to where its subscription starts. + /// + /// Its highest group is read before replying: FETCH_OK names where the response ends, + /// which a range running past the track only learns by reading it, and a refusal can + /// still replace it until then. + fn fetch_gate(&self, msg: &ietf::Fetch<'_>) -> FetchGate { + match &msg.fetch_type { + FetchType::Standalone { namespace, .. } | FetchType::Filtered { namespace, .. } => FetchGate::Standalone( + crate::auth::Gate::new(self.auth.clone(), namespace.clone(), crate::auth::Direction::Publish), + ), + FetchType::RelativeJoining { + subscriber_request_id, .. + } + | FetchType::AbsoluteJoining { + subscriber_request_id, .. + } => FetchGate::Joining { + id: *subscriber_request_id, + auth: self.auth.clone(), + }, + } + } + + #[cfg(test)] + async fn run_fetch_stream( + self, + stream: Stream, + msg: ietf::Fetch<'static>, ) -> Result<(), Error> { - let source = group.info.ietf(); - let priority = super::priority::to_wire(group.info.priority); - let mut prior = None; - let mut budget = kio::coop::Budget::new(32); - for (ordinal, frame) in (group.first..).zip(group.frames) { - kio::wait(|waiter| budget.poll_yield(waiter)).await; - let object = frame.fields.as_ref().map_or(ordinal, |fields| fields.id); - let position = fetch_prior(prior, object)?; - prior = Some(object); - Self::write_fetch_object( - writer, - FetchObjectInfo { - sequence: group.sequence, - object, - prior: position, - priority, - timestamp: frame.timestamp, - inflated: frame.inflated, - timescale, - version, - fields: frame.fields.as_deref(), - source, - }, - ) - .await?; - Self::write_fetch_payload(writer, frame.frame.payload).await?; + let gate = self.fetch_gate(&msg); + self.serve_fetch_stream(stream, msg, gate).await + } + + /// Await [`FetchServe`]. + #[cfg(test)] + async fn serve_fetch_stream( + self, + stream: Stream, + msg: ietf::Fetch<'static>, + admission: FetchGate, + ) -> Result<(), Error> { + let mut serve = FetchServe::new(self, stream, msg, admission); + kio::wait(|waiter| serve.poll(waiter)).await + } + + /// Answer a TRACK_STATUS with what a SUBSCRIBE_OK for the track would carry, without + /// subscribing. + /// + /// The track resolves as it would for a SUBSCRIBE, so a relay asks its upstream for a + /// track it does not hold yet and the answer waits on that. The Largest Location is + /// whatever the cache holds now: nothing waits for a fresher one. + #[cfg(test)] + fn run_track_status_stream( + self, + stream: Stream, + msg: ietf::TrackStatus<'static>, + ) -> impl std::future::Future> { + let gate = crate::auth::Gate::new( + self.auth.clone(), + msg.track_namespace.clone(), + crate::auth::Direction::Publish, + ); + self.serve_track_status_stream(stream, msg, gate) + } + + /// Await [`TrackStatusServe`]. + #[cfg(test)] + async fn serve_track_status_stream( + self, + stream: Stream, + msg: ietf::TrackStatus<'static>, + gate: crate::auth::Gate, + ) -> Result<(), Error> { + let mut serve = TrackStatusServe::new(self, stream, msg, gate); + kio::wait(|waiter| serve.poll(waiter)).await + } + + /// Bring the peer's view of one namespace in line with the current selection. + /// + /// A loop writing inline (answering a SUBSCRIBE_NAMESPACE) re-sends + /// NAMESPACE on that stream, which the receiver treats as a replacement, and + /// retracts with NAMESPACE_DONE. Otherwise each advertisement rides its own + /// PUBLISH_NAMESPACE request: an update is a REQUEST_UPDATE **on the stream that + /// already carries it**, since a second stream would leave two claiming one + /// namespace, and a withdrawal cancels the request. + /// + /// Inline messages are only buffered, for the loop to flush before its next event. An + /// advertisement that needs a stream of its own is returned, for the loop to open one + /// and finish it with [`Self::advertised`]. + fn sync_namespace( + &self, + ns: &mut Namespaces, + suffix: &crate::Path, + path: &crate::Path, + ) -> Result, Error> { + let Some((advert, auth)) = self.offer(ns, suffix, path) else { + return Ok(None); + }; + let negotiated = ns.negotiated(); + let Some(watch) = ns.watched.get(suffix) else { + return Ok(None); + }; + let refused = watch.refused; + let wanted = advert.wanted(); + let held = watch.sent.wanted(); + let request = ns.requests.get(suffix); + let held_auth = request.as_ref().and_then(|request| request.auth.clone()); + let unchanged = advert == watch.sent && auth == held_auth; + // moq-transport renames a request's auth with REQUEST_UPDATE, but cannot name one on + // a request that had none, nor drop one: those withdraw the request and make it + // again. + let restart = held_auth != auth && (held_auth.is_none() || auth.is_none()); + + // Keep one opening advertisement per namespace. Its answer reconciles the latest + // route, so route churn cannot accumulate replacement streams behind a stalled write. + if wanted && request.is_some_and(|request| !request.accepted) { + return Ok(None); } - Ok(()) + + // An update already waiting or in flight carries whatever is current once it goes + // out, and its answer brings the loop back here for anything after that. + if wanted + && !restart + && let Some(request) = ns.requests.get_mut(suffix) + && !unchanged + { + if !request.queued && request.pending.is_none() { + request.queued = true; + ns.updates.push_back(suffix.clone()); + } + return Ok(None); + } + + if unchanged { + // Nothing to send. A namespace that is no longer advertisable is no longer + // pending either, and leaving that set would keep the retry timer armed + // forever for a wire message that can never happen. + if !wanted { + ns.watched.get_mut(suffix).expect("watched").deferred = false; + } + return Ok(None); + } + + // A fresh offer waits for what the refusal asked for, whatever brought us back. + // Only withdrawing and re-announcing clears it, since that builds a fresh entry. + if wanted && !held && !refused.offerable(self.runtime.now()) { + return Ok(None); + } + + let absolute = self.origin.root().rejoin(path); + // Only a PUBLISH_NAMESPACE request or an update to one can be refused; everything + // else below either rides a stream the peer already accepted or says nothing at all. + let refused = watch.refused; + let advertise = |advert: Advert| Advertise { + suffix: suffix.clone(), + path: path.clone(), + advert, + auth, + wanted, + }; + let sent = match &mut ns.target { + Target::Requests => { + match (advert.wanted(), ns.requests.get_mut(suffix)) { + (false, _) => { + if held { + tracing::debug!(broadcast = %absolute, "namespace_done"); + } + self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone())?; + } + (true, Some(_)) => { + tracing::debug!(broadcast = %absolute, "publish_namespace restart"); + self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone())?; + return Ok(Some(advertise(advert))); + } + (true, None) => { + tracing::debug!(broadcast = %absolute, "publish_namespace"); + return Ok(Some(advertise(advert))); + } + } + Self::requested(ns, suffix, advert) + } + Target::Inline(stream) => { + match (advert.wanted(), held) { + (true, _) => { + tracing::debug!(broadcast = %absolute, "namespace"); + stream.writer.buffer_message( + ietf::Namespace::ID, + &ietf::Namespace { + suffix: suffix.clone(), + params: negotiated.extended().then(|| advert.params()), + }, + )?; + } + (false, true) => { + tracing::debug!(broadcast = %absolute, "namespace_done"); + stream + .writer + .buffer_message(ietf::NamespaceDone::ID, &ietf::NamespaceDone { suffix: suffix.clone() })?; + } + // Never advertised and still not advertisable: nothing to say. + (false, false) => {} + } + advert + } + }; + Self::record(ns, suffix, wanted, refused, sent); + Ok(None) } - /// Write a whole fetch object's length and payload, after its header. - async fn write_fetch_payload( - stream: &mut Writer, Version>, - payload: crate::frame::Payload, + /// Finish [`Self::sync_namespace`]'s advertisement once its stream opened, or gave up. + fn advertised( + &self, + ns: &mut Namespaces, + advertise: Advertise, + stream: Option>, ) -> Result<(), Error> { - let mut write = FetchPayloadWrite::new(payload); - kio::wait(|waiter| write.poll(waiter, stream)).await + let Advertise { + suffix, + path, + advert, + auth, + wanted, + } = advertise; + let refused = self.advertise_namespace(ns, &path, suffix.clone(), advert.params(), auth, stream)?; + let sent = Self::requested(ns, &suffix, advert); + Self::record(ns, &suffix, wanted, refused, sent); + Ok(()) } - /// Register a pending subscription until the returned guard drops. - fn register_join(&self, request_id: RequestId) -> Join> { - self.joins.lock().insert(request_id, None); - Join { - joins: self.joins.clone(), - request_id, + /// What the peer holds once a PUBLISH_NAMESPACE request was made, withdrawn, or refused. + fn requested(ns: &mut Namespaces, suffix: &crate::Path, advert: Advert) -> Advert { + // Drive the first write too, and keep reading once idle for GOAWAY. + if ns.requests.contains_key(suffix) { + ns.woken.suffixes.lock().unwrap().push(suffix.clone()); } - } - - /// Answer a FETCH: a standalone or filtered range of the named track, or a joining - /// FETCH's groups up to where its subscription starts. - /// - /// Its highest group is read before replying: FETCH_OK names where the response ends, - /// which a range running past the track only learns by reading it, and a refusal can - /// still replace it until then. - fn fetch_gate(&self, msg: &ietf::Fetch<'_>) -> FetchGate { - match &msg.fetch_type { - FetchType::Standalone { namespace, .. } | FetchType::Filtered { namespace, .. } => FetchGate::Standalone( - crate::auth::Gate::new(self.auth.clone(), namespace.clone(), crate::auth::Direction::Publish), - ), - FetchType::RelativeJoining { - subscriber_request_id, .. - } - | FetchType::AbsoluteJoining { - subscriber_request_id, .. - } => FetchGate::Joining { - id: *subscriber_request_id, - auth: self.auth.clone(), - }, + // The peer can reject a fresh PUBLISH_NAMESPACE or an update, which leaves + // no request behind. Record each buffered advertisement now, so withdrawal + // cancels it even before its first message leaves. + match ns.requests.contains_key(suffix) { + true => advert, + false => Advert::None, } } - #[cfg(test)] - async fn run_fetch_stream( - self, - stream: Stream, - msg: ietf::Fetch<'static>, - ) -> Result<(), Error> { - let gate = self.fetch_gate(&msg); - self.serve_fetch_stream(stream, msg, gate).await + /// Record what was `sent` for one namespace the peer should hold if `wanted`. + fn record(ns: &mut Namespaces, suffix: &crate::Path, wanted: bool, refused: Refused, sent: Advert) { + // A peer that asked not to be offered this again outranks the retry timer; + // anything else it should hold and does not comes back on one. + let watch = ns.watched.get_mut(suffix).expect("watched"); + watch.refused = refused; + watch.deferred = wanted && !sent.wanted() && refused.pending(); + watch.sent = sent; } - async fn serve_fetch_stream( - mut self, - mut stream: Stream, - msg: ietf::Fetch<'static>, - admission: FetchGate, - ) -> Result<(), Error> { - let namespace = match &msg.fetch_type { - FetchType::Standalone { namespace, .. } | FetchType::Filtered { namespace, .. } => Some(namespace), - _ => None, + /// Buffer a PUBLISH_NAMESPACE on `stream`, the one [`Opening`] got for it, and register + /// that with the announce loop, so a withdrawal never waits for send credit or the + /// peer's answer. + fn advertise_namespace( + &self, + ns: &mut Namespaces, + path: &crate::Path, + suffix: crate::Path, + params: ietf::AdvertParams, + auth: Option, + stream: Option>, + ) -> Result { + let Some(mut request) = stream else { + tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), path), "no stream for the advertisement"); + return Ok(Refused::No); }; - if let Some(error) = namespace.and_then(|namespace| super::namespace::refusal(namespace, self.version)) { - return super::request_stream::reject(stream, error).await; + let request_id = self.control.next_request_id(); + if let Err(err) = request.writer.buffer_message( + ietf::PublishNamespace::ID, + &ietf::PublishNamespace { + request_id, + track_namespace: path.clone(), + params, + authorization_token: auth.as_ref().and_then(|auth| self.request_token(auth)), + }, + ) { + tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), path), %err, "advertisement failed"); + return Ok(ns.watched.get_mut(&suffix).expect("watched").lapsed(self.runtime.now())); } - // What the requester sent on this stream so far, kept across every phase that watches - // it for a cancel, so a second GOAWAY is caught wherever it arrives. - let mut finished = false; - let mut messages = super::request_stream::Messages::new(self.session.clone(), self.control.client()); - let priority = super::priority::from_wire(msg.subscriber_priority); + ns.requests.insert( + suffix.clone(), + NamespaceRequest { + messages: super::request_stream::Messages::new(self.session.clone(), self.control.client()), + answer: None, + path: path.clone(), + stream: request, + auth, + accepted: false, + queued: false, + pending: Some(NamespacePending::Writing { _auth: None }), + waiter: None, + }, + ); + Ok(Refused::No) + } - // Serving a Range Filter unfiltered would deliver objects the subscriber excluded. - if msg.range_filters { - return super::request_stream::reject( - stream, - ietf::RequestError { - error_code: request::INVALID_FILTER, - reason_phrase: "Range Filters exceed MAX_FILTER_RANGES=0".into(), - retry_interval: 0, - }, - ) - .await; - } + /// What the peer should hold for one namespace: the advertisement, and the auth its + /// PUBLISH_NAMESPACE request carries. `None` when the namespace is not watched. + /// + /// Nothing our grant does not cover reaches the wire, and a shrink withdraws it. Only a + /// request carries an auth, and only one this session's handle issued: any other route + /// is advertised under the session's scope. + fn offer( + &self, + ns: &Namespaces, + suffix: &crate::Path, + path: &crate::Path, + ) -> Option<(Advert, Option)> { + let watch = ns.watched.get(suffix)?; + let advert = match ns.permitted(path) { + true => self.select(&watch.route, &ns.peer), + false => Advert::None, + }; + let auth = match (&ns.target, advert.wanted()) { + (Target::Requests, true) => watch + .route + .auth + .as_ref() + .filter(|auth| self.auth.token(auth).is_some()) + .cloned(), + _ => None, + }; + Some((advert, auth)) + } - // FILL_TIMEOUT is accepted and ignored: a one-group read ends with the group. + /// The `AUTHORIZATION TOKEN` a request names `auth` with, by value. + fn request_token(&self, auth: &crate::auth::Credential) -> Option { + let token = self.auth.token(auth)?; + Some(super::token::RequestToken(crate::setup::Token { + kind: crate::setup::Token::OUT_OF_BAND, + value: token.to_vec(), + })) + } - // Draft-22's LOCATION_FILTER is answered as the standalone range it spells. Every other - // filter ends at the Largest Object, which resolves once the track does. - let mut largest_filter = None; - let fetch_type = match msg.fetch_type { - FetchType::Filtered { - namespace, - track, - filter: Filter::Absolute { start, end: Some(end) }, - } => { - // The filter's End Object is inclusive, where a standalone one is the last - // object plus one. Without one, both include the whole End Group. - let Some(object) = end.object.map_or(Some(0), |object| object.checked_add(1)) else { - return self - .reject_fetch(stream, &Error::InvalidRange, "End Object overflows") - .await; - }; - FetchType::Standalone { - namespace, - track, - start, - end: Location { - group: end.group, - object, - }, - } + /// Drop queued updates that no longer have anything to send, so the next one to take + /// a request ID is one that will use it: drafts before 17 require the next ID. + fn prune_updates(&self, ns: &mut Namespaces, prefix: &crate::Path) { + while let Some(suffix) = ns.updates.front() { + let wants = ns.requests.get(suffix).is_some_and(|request| { + request.queued + && request.pending.is_none() + && ns.watched.get(suffix).is_some_and(|watch| { + self.offer(ns, suffix, &prefix.rejoin(suffix)) + .is_some_and(|(advert, auth)| { + advert.wanted() && (advert != watch.sent || auth != request.auth) + }) + }) + }); + if wants { + return; } - FetchType::Filtered { - namespace, - track, - filter, - } => { - largest_filter = Some(filter); - // Placeholders until the filter resolves against the Largest Object. - let whole = Location { group: 0, object: 0 }; - FetchType::Standalone { - namespace, - track, - start: whole, - end: whole, - } + let suffix = ns.updates.pop_front().expect("checked non-empty"); + if let Some(request) = ns.requests.get_mut(&suffix) { + request.queued = false; } - other => other, - }; - - // Every FETCH is held to the grant while its groups load and its response is written. - // A standalone one is also checked here, before it reaches the origin; a joining one - // was checked when its subscription was. - // - // A joining one is refused if its subscription ends before FETCH_OK. After that the - // head is already held and the requester can still cancel, so it finishes on its own - // authority: the gate, and the token grant if one authorized the subscription, which - // ends it on the token's expiry, revocation, or a refused renewal. - let mut lapses_with = None; - let mut token_grant = None; - let (track, start, end, joined, ranking, mut gate) = match fetch_type { - FetchType::Standalone { - namespace, - track, - start, - end, - } => { - let FetchGate::Standalone(mut gate) = admission else { - return Err(Error::Unauthorized); - }; - match self - .admit_request( - &mut stream, - &mut gate, - msg.authorization_token.clone(), - super::request_stream::fin_cancels(self.version), - ) - .await - { - Ok(true) => {} - Ok(false) => return Ok(()), - Err(refusal) => return self.reject_fetch(stream, &refusal.error, &refusal.reason).await, - } - if !gate.allows() { - return self.reject_fetch(stream, &Error::Unauthorized, "not granted").await; - } - - // End Object 0 includes the entire End Group, even at the largest Group ID. - if largest_filter.is_none() - && (start.group > end.group - || (start.group == end.group && end.object != 0 && start.object >= end.object)) - { - return self.reject_fetch(stream, &Error::InvalidRange, "empty range").await; - } + } + } - if largest_filter.is_none() && start.group != end.group { - return self - .reject_fetch(stream, &Error::Unsupported, "FETCH spans more than one group") - .await; - } + /// Send the update waiting at the front of the queue with `request_id`: REQUEST_UPDATE + /// on the request that carries the namespace, with only what changed since the peer's + /// copy. Its answer arrives through [`Self::poll_answer`] while the loop carries on. + /// + /// One update is outstanding per request, which satisfies any MAX_REQUEST_UPDATES the + /// peer set without reading it. A different original publisher is an update like any + /// other: the receiver drains what it already serves from the old one and never splices + /// the two. + fn send_update(&self, ns: &mut Namespaces, prefix: &crate::Path, request_id: RequestId) -> Result<(), Error> { + let suffix = ns.updates.pop_front().expect("allocated for a queued update"); + let Some((advert, auth)) = self.offer(ns, &suffix, &prefix.rejoin(&suffix)) else { + return Ok(()); + }; + let (Some(request), Some(watch)) = (ns.requests.get_mut(&suffix), ns.watched.get_mut(&suffix)) else { + return Ok(()); + }; + request.queued = false; + let mut update = match (advert.params().cluster, watch.sent.params().cluster) { + (Some(next), Some(held)) => ietf::PublishNamespaceUpdate::between(request_id, &held, &next), + _ => ietf::PublishNamespaceUpdate { + request_id, + hops: None, + cost: None, + authorization_token: None, + }, + }; + if auth != request.auth { + update.authorization_token = auth.as_ref().and_then(|auth| self.request_token(auth)); + } + tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), &request.path), "publish_namespace update"); + let written = request + .stream + .writer + .buffer_message(ietf::PublishNamespaceUpdate::ID, &update); + if let Err(err) = written { + // An unencodable update ends only this advertisement; the retry re-offers it + // fresh, carrying the complete parameters. + tracing::debug!(%err, "publish_namespace update not sent"); + self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone())?; + if let Some(watch) = ns.watched.get_mut(&suffix) { + watch.refused = Refused::No; + watch.deferred = Refused::No.pending(); + watch.sent = Advert::None; + } + return Ok(()); + } - // The peer must have seen the announcement to name this namespace, so this - // resolves like a SUBSCRIBE does. - let broadcast = match gate - .guard(self.resolve( - &mut stream, - super::request_stream::fin_cancels(self.version), - &namespace, - )) - .await - { - Ok(Some(broadcast)) => broadcast, - Ok(None) => return Ok(()), - Err(err) => return self.reject_fetch(stream, &err, &err.to_string()).await, - }; - let track = match broadcast.track(&track) { - Ok(track) => track, - Err(err) => return self.reject_fetch(stream, &err, &err.to_string()).await, - }; + // Keep one buffered update until the peer answers it. + request.pending = Some(NamespacePending::Writing { + _auth: std::mem::replace(&mut request.auth, auth), + }); + watch.sent = advert; + // Marked woken, so the loop's next pass reads the answer it may already have. + ns.woken.suffixes.lock().unwrap().push(suffix); + Ok(()) + } - // Its own entry in the queue, or its own send group, so it competes with group - // streams by its priority alone, its request id breaking ties in the queue. - // Nothing updates a FETCH's priority, and it ranks before the track's info - // resolves, so the publisher's priority is the unset default. - let ranking = Ranking { - queue: self.priority.clone(), - group: self.session.clone().send_group(crate::priority::group(priority, 0)), - request_id: msg.request_id, - priority: kio::Producer::new(priority).consume(), - publisher: 0, - order: Order::Descending, - lapse: Lapse::default(), - }; - (track, start, end, false, ranking, gate) - } - FetchType::RelativeJoining { - subscriber_request_id, .. + /// Flush writes and read answers, polling only the requests that woke. + fn poll_answer( + requests: &mut HashMap, + inflight: &mut VecDeque<(crate::time::Instant, crate::Path)>, + woken: &Arc, + now: crate::time::Instant, + waiter: &kio::Waiter, + ) -> Poll<(crate::Path, Result)> { + // Registered before the take, so a wake landing between them is not lost. + *woken.owner.lock().unwrap() = Some(waiter.waker().clone()); + let mut suffixes = std::mem::take(&mut *woken.suffixes.lock().unwrap()); + while let Some(suffix) = suffixes.pop() { + let Some(request) = requests.get_mut(&suffix) else { + continue; + }; + // Kept on the request, so its registrations outlive this pass. + let request_waiter = request + .waiter + .get_or_insert_with(|| kio::Waiter::new(woken.waker(suffix.clone()))); + let written = request.stream.writer.poll_flush(request_waiter); + if matches!(&written, Poll::Ready(Ok(()))) + && matches!(request.pending, Some(NamespacePending::Writing { .. })) + { + let deadline = now + ADVERTISE_TIMEOUT; + request.pending = Some(NamespacePending::Answer(deadline)); + inflight.push_back((deadline, suffix.clone())); } - | FetchType::AbsoluteJoining { - subscriber_request_id, .. - } => { - let FetchGate::Joining { id, auth } = admission else { - return Err(Error::Unauthorized); - }; - debug_assert_eq!(id, subscriber_request_id); - let (end, cache, namespace, ranking, grant) = match self - .joined(&mut stream, &mut finished, &mut messages, subscriber_request_id) - .await? - { - Ok(joined) => joined, - Err((err, reason)) => return self.reject_fetch(stream, &err, reason).await, - }; - // The cache outlives the subscription, so the fetch holds its own gate rather than - // trusting the cache. One the subscription's request token authorized is held to - // the peer's grant alone, as the subscription is. - let gate = crate::auth::Gate::new(auth, namespace, crate::auth::Direction::Publish); - let gate = if grant.is_some() { gate.named() } else { gate }; - lapses_with = Some(subscriber_request_id); - token_grant = grant; - let start = match fetch_type { - FetchType::RelativeJoining { group_offset: 0, .. } => end.group, - FetchType::RelativeJoining { .. } => { - return self - .reject_fetch( - stream, - &Error::Unsupported, - "joining FETCH reaches before the joined group", - ) - .await; + let answer = match written { + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Ready(Ok(())) if request.answer.is_some() => Poll::Ready(Ok(request.answer.take())), + _ => loop { + let answer = request.messages.poll_next(&mut request.stream.reader, request_waiter); + match answer { + Poll::Ready(Ok(Some(_))) if request.answer.is_some() => { + break Poll::Ready(Err(request.messages.failed(Error::UnexpectedMessage))); + } + Poll::Ready(Ok(Some(framed))) if written.is_pending() && framed.0 == ietf::RequestOk::ID => { + // An early OK must not admit another update or replacement behind + // this blocked write. Keep one, while still reading peer terminals. + request.answer = Some(framed); + } + answer => break answer, } - FetchType::AbsoluteJoining { group_id, .. } if group_id <= end.group => group_id, - _ => { - return self - .reject_fetch(stream, &Error::InvalidRange, "joining group past the subscription") - .await; + }, + }; + if let Poll::Ready(answer) = answer { + // The rest stay woken for the next pass. + woken.suffixes.lock().unwrap().append(&mut suffixes); + let answer = answer.and_then(|answer| { + if answer.is_some() && request.pending.is_none() { + return Err(request.messages.failed(Error::UnexpectedMessage)); } - }; - ( - cache, - Location { - group: start, - object: 0, - }, - end, - true, - ranking, - gate, - ) + Ok(answer.map_or(Answer::Ended, Answer::Message)) + }); + woken.suffixes.lock().unwrap().push(suffix.clone()); + return Poll::Ready((suffix, answer)); } - // Rewritten as standalone above. - FetchType::Filtered { .. } => { - return self.reject_fetch(stream, &Error::Unsupported, "not supported").await; + } + while let Some((deadline, suffix)) = inflight.front() { + let current = matches!( + requests.get(suffix).and_then(|request| request.pending.as_ref()), + Some(NamespacePending::Answer(at)) if at == deadline + ); + if current && now < *deadline { + break; } - }; + let (_, suffix) = inflight.pop_front().expect("just peeked"); + if current { + return Poll::Ready((suffix, Ok(Answer::Late))); + } + } + Poll::Pending + } - let joins = self.joins.clone(); - let unsubscribed = |waiter: &kio::Waiter| { - lapses_with.is_some_and(|id| { - joins - .poll(waiter, |joins| match joins.contains_key(&id) { - true => Poll::Pending, - false => Poll::Ready(()), - }) - .is_ready() - }) - }; - // A standalone FETCH's own token, which its gate holds, or a joining one's subscription's. - let token_grant = token_grant.or_else(|| gate.token_guard()); - // Set once that token lapses, so the streams this FETCH answers on say so. - let lapse = Lapse::default(); - // Ready with why the token's grant ended, ahead of the gate it also closes, so a lapse - // is refused as one. - let token_ended = |waiter: &kio::Waiter| match &token_grant { - Some(grant) => grant.poll_ended(waiter).map(|refusal| { - if super::error::token_expired(&refusal.error) { - lapse.set(); - } - refusal - }), - None => Poll::Pending, + /// Apply an advertisement or update's answer: an OK keeps the namespace, and a + /// REQUEST_ERROR closed the stream and withdrew the advertisement (moq-transport Section + /// 9.5.1), so the request is dropped and the retry re-offers it fresh. An unanswered + /// update is dropped the same way, since a peer that ignored it cannot be assumed to hold + /// either version, and so is one the peer cancelled instead of answering. An idle + /// advertisement the peer ends waits out its own growing backoff before coming back. + /// + /// Returns the advertisement an OK leaves to make, as [`Self::sync_namespace`] does. + fn answered( + &self, + ns: &mut Namespaces, + prefix: &crate::Path, + suffix: crate::Path, + answer: Result, + ) -> Result, Error> { + let Some(request) = ns.requests.get_mut(&suffix) else { + return Ok(None); }; + let initial = !request.accepted; + let idle = request.pending.take().is_none(); + let absolute = self.origin.root().rejoin(&request.path); - // Read the group before FETCH_OK promises it, while a refusal is still possible. - let compression = self.peer_setup.get().await.flate; - let prepared = { - let mut prepare = std::pin::pin!(prepare_fetch( - &track, - FetchRange { start, end }, - largest_filter, - joined, - priority, - compression, - )); - kio::wait(|waiter| { - if super::request_stream::poll_cancel( - &mut stream, - &mut finished, - &mut messages, - super::request_stream::fin_cancels(self.version), - waiter, - ) - .is_ready() - { - return Poll::Ready(None); - } - if let Poll::Ready(refusal) = token_ended(waiter) { - return Poll::Ready(Some(Err((refusal.error, refusal.reason)))); - } - if gate.poll_denied(waiter).is_ready() || unsubscribed(waiter) { - let err = Error::Unauthorized; - let reason = err.to_string(); - return Poll::Ready(Some(Err((err, reason)))); - } - waiter.poll_future(prepare.as_mut()).map(Some) - }) - .await - }; - let PreparedFetch { range, last: group } = match prepared { - Some(Ok(prepared)) => prepared, - Some(Err((err, reason))) => return self.reject_fetch(stream, &err, &reason).await, - None => return Ok(()), - }; - let end = range.end; - let group = group.without_status(); - // Only the TIMESCALE in FETCH_OK opts the answer into timestamps, so a requester that - // opted out of the properties gets an untimed track. - let timescale = group.info.timescale.filter(|_| msg.properties_wanted); - let frames = group - .frames - .iter() - .filter_map(|frame| Some((frame.fields.as_deref()?, frame.timestamp))); - for (fields, frame) in frames.chain(group.end_object.as_deref().map(|end| (end, None))) { - if let Err(err) = - super::object::write(&fields.properties, group.info.ietf(), frame, self.version, timescale) - { - return self.reject_fetch(stream, &err, &err.to_string()).await; + // A reset stream ends only this request, like a FIN; a malformed answer is the peer's + // violation and still ends the loop. + let answer = match answer { + Err(err) if !super::subscriber::is_protocol_violation(&err) => { + tracing::debug!(broadcast = %absolute, %err, "advertisement stream ended"); + Answer::Ended } - } - let properties = match track_properties(&group.info, msg.properties_wanted, self.version, compression) { - Ok(properties) => properties, - Err(err) => return self.reject_fetch(stream, &err, &err.to_string()).await, + answer => answer?, }; - let draft20 = Filter::is_draft20(self.version); - // Only a finished track's final object ends it, never a live track's Largest Object. - let end_of_track = !joined && group.complete && track.final_sequence() == group.sequence.checked_add(1); - - let end_location = if joined { - // The subscription starts at the saved Largest Object, so the prefix of that - // group has to be here in full or the two leave a gap. - if group.sequence != end.group || group.end() != end.object { - return self - .reject_fetch(stream, &Error::NotFound, "joining prefix unavailable") - .await; + let refused = match answer { + Answer::Late => { + tracing::debug!(broadcast = %absolute, "no answer to the advertisement"); + self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone())?; + Refused::No } - end - } else if draft20 { - // Draft-22's End Location is inclusive. An empty answer still covers its start. - let object = match range.until(group.sequence) { - // The track's last object as this draft writes it, without a status object. - Some(until) if end_of_track => (until - 1).min(group.end().saturating_sub(1)), - Some(until) => until - 1, - None => group.end().saturating_sub(1).max(range.skip(group.sequence)), - }; - Location { - group: group.sequence, - object, + Answer::Ended if idle || initial => { + tracing::debug!(broadcast = %absolute, "advertisement ended"); + self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone())?; + let now = self.runtime.now(); + ns.watched + .get_mut(&suffix) + .map_or(Refused::No, |watch| watch.lapsed(now)) } - } else if end_of_track && end.object > 0 { - // A bounded response cannot promise objects past the track's end. - Location { - group: group.sequence, - object: group.end(), + Answer::Ended => { + tracing::debug!(broadcast = %absolute, "update cancelled"); + // The peer closed its side; closing ours completes the withdrawal. + self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone())?; + Refused::No + } + Answer::Message(super::request_stream::Framed(type_id, body)) => { + match type_id { + ietf::RequestOk::ID => { + let msg: ietf::RequestOk = body + .message(self.version) + .map_err(|err| request.messages.failed(err.into()))?; + // ACTIVE_COUNT only answers SUBSCRIBE_NAMESPACE (MoQ Active Count). + if msg.active.is_some() { + self.session.clone().close( + super::error::to_session_code(&crate::SessionError::ProtocolViolation), + "ACTIVE_COUNT on a PUBLISH_NAMESPACE answer", + ); + return Err(Error::ProtocolViolation); + } + tracing::debug!(message = ?msg, "publish_namespace ok"); + request.accepted = true; + // Whatever changed while it was in flight goes out next. + return self.sync_namespace(ns, &suffix, &prefix.rejoin(&suffix)); + } + ietf::RequestError::ID => { + let msg: ietf::RequestError = body + .message(self.version) + .map_err(|err| request.messages.failed(err.into()))?; + tracing::warn!(message = ?msg, "publish_namespace error"); + // The peer closed its side; closing ours completes the withdrawal. + self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone())?; + self.refusal(msg.retry_interval) + } + _ => return Err(request.messages.failed(Error::UnexpectedMessage)), + } } - } else { - end }; - // The response, raced against the grant: a standalone FETCH whose grant narrows while it - // waits on stream credit stops there instead of sending what it no longer may. - let ok = async { - // FETCH_OK on every draft, never REQUEST_OK: section 5.2 allows exactly one FETCH_OK or - // REQUEST_ERROR in answer to a FETCH, and REQUEST_OK's own definition lists the other - // requests it answers without ever naming this one. - stream - .writer - .encode_message(&ietf::FetchOk { - end_of_track, - end_location, - properties, - }) - .await - }; - let res = { - let mut ok = std::pin::pin!(ok); - kio::wait(|waiter| { - if let Poll::Ready(refusal) = token_ended(waiter) { - return Poll::Ready(Err(refusal.error)); - } - if gate.poll_denied(waiter).is_ready() || unsubscribed(waiter) { - return Poll::Ready(Err(Error::Unauthorized)); - } - waiter.poll_future(ok.as_mut()) - }) - .await - }; - if let Err(err) = res { - lapse.abort(stream.writer, &err); - return Err(err); + if let Some(watch) = ns.watched.get_mut(&suffix) { + watch.refused = refused; + watch.deferred = refused.pending(); + watch.sent = Advert::None; } + Ok(None) + } - // Cancellation releases the buffered answer even while its stream is flow controlled. - // A joining one no longer ends with its subscription. - let respond = async { - let mut rank = ranking.rank(range.start); - let uni = self - .session - .open_uni(rank.options()) - .await - .map_err(Error::from_transport)?; - let mut writer = Writer::new(Ranked::new(uni, Some(rank)).lapsing(&lapse), self.version); - writer.varint(FetchHeader::TYPE).await?; - writer - .encode(&FetchHeader { - request_id: msg.request_id, - }) - .await?; - Self::write_fetched_group(&mut writer, group, self.version, timescale).await?; - settle(&mut writer).await?; - Ok::<(), Error>(()) - }; - let res = { - let mut respond = std::pin::pin!(respond); - kio::wait(|waiter| { - if super::request_stream::poll_cancel( - &mut stream, - &mut finished, - &mut messages, - super::request_stream::fin_cancels(self.version), - waiter, - ) - .is_ready() - { - return Poll::Ready(None); - } - if let Poll::Ready(refusal) = token_ended(waiter) { - return Poll::Ready(Some(Err(refusal.error))); - } - if gate.poll_denied(waiter).is_ready() { - return Poll::Ready(Some(Err(Error::Unauthorized))); - } - waiter.poll_future(respond.as_mut()).map(Some) - }) - .await - }; - // Cancelled: the fetch stream, if open, reset as it dropped with the response. - let Some(res) = res else { - return Ok(()); - }; - if let Err(err) = res { - // The fetch stream, if open, reset as it dropped with the response. Once FETCH_OK - // is out, no REQUEST_ERROR can follow it, so the reset is the whole refusal. - lapse.abort(stream.writer, &err); - return Err(err); + /// How to read a refusal's retry interval, in milliseconds. + /// + /// 0 is the peer asking not to be offered this again, which is what keeps a permanent + /// refusal (unauthorized, uninterested) from becoming a request every few seconds for + /// the life of the session. + fn refusal(&self, retry_interval: u64) -> Refused { + match retry_interval { + 0 => Refused::Never, + ms => Refused::Until(self.runtime.now() + Duration::from_millis(ms)), } + } - // FETCH_OK is the last thing this stream has to say, and [`Writer`] resets on drop: - // without the finish the answer is discarded before the peer ever reads it, exactly - // as it would be for a refusal. The peer dropping the stream first is a normal end. - let _ = stream.writer.close().await; + pub(super) fn stats(&self) -> &crate::stats::Session { + self.origin.stats() + } + /// Withdraw an advertised namespace: NAMESPACE_DONE inline, buffered for the loop to + /// flush, or end the request that carried it. + fn withdraw_namespace( + &self, + target: &mut Target, + requests: &mut HashMap, + suffix: crate::Path, + ) -> Result<(), Error> { + match target { + Target::Requests => { + if let Some(request) = requests.remove(&suffix) { + end_request(request.stream); + } + } + Target::Inline(stream) => { + stream + .writer + .buffer_message(ietf::NamespaceDone::ID, &ietf::NamespaceDone { suffix })?; + } + } Ok(()) } - /// Resolve the subscription a joining FETCH names to its saved start, or the refusal - /// to answer the FETCH with. - async fn joined( - &mut self, - stream: &mut Stream, - finished: &mut bool, - messages: &mut super::request_stream::Messages, - subscribe_id: RequestId, - ) -> Result< - Result< - ( - Location, - track::Consumer, - crate::Path, - Ranking>, - Option, - ), - (Error, &'static str), - >, - Error, - > { - // Request streams can arrive out of order. Wait on registration, while bounding - // the lifetime of a request whose subscription never arrives or resolves. - let joined = { - let mut pending = false; - let mut deadline = crate::time::Deadline::after(&self.runtime, Duration::from_secs(10)); - kio::wait(|waiter| { - // The request reader is what the subscriber FINs or resets. - if super::request_stream::poll_cancel( - stream, - finished, - messages, - super::request_stream::fin_cancels(self.version), - waiter, - ) - .is_ready() - { - return Poll::Ready(Err(Error::Cancel)); + /// Close out every open PUBLISH_NAMESPACE request. A no-op for a loop whose entries + /// ride the SUBSCRIBE_NAMESPACE stream itself, which retracts them by ending. + fn withdraw_requests(&self, target: &mut Target, requests: &mut HashMap) { + let suffixes: Vec = requests.keys().cloned().collect(); + for suffix in suffixes { + let _ = self.withdraw_namespace(target, requests, suffix); + } + } + + /// [`PublishNamespaces`] as a future, for tests. + #[cfg(test)] + pub async fn run_publish_namespaces(self) -> Result<(), Error> { + let mut run = PublishNamespaces::new(self); + kio::wait(|waiter| run.poll(waiter)).await + } + + /// Handle a SUBSCRIBE_NAMESPACE on its bidi stream. + /// + /// All the announce state is local to this task (mirroring `lite::Publisher`'s + /// announce handling): whatever this subscription advertised is withdrawn + /// when its stream ends. Every match is a NAMESPACE on this stream, whatever the + /// peer's SETUP said. + #[cfg(test)] + async fn run_subscribe_namespace_stream( + self, + stream: Stream, + msg: ietf::SubscribeNamespace, + ) -> Result<(), Error> { + let gate = crate::auth::Gate::new( + self.auth.clone(), + msg.namespace.clone(), + crate::auth::Direction::Publish, + ); + self.serve_subscribe_namespace_stream(stream, msg, gate).await + } + + /// Await [`NamespaceServe`]. + #[cfg(test)] + async fn serve_subscribe_namespace_stream( + self, + stream: Stream, + msg: ietf::SubscribeNamespace, + gate: crate::auth::Gate, + ) -> Result<(), Error> { + let mut serve = NamespaceServe::new(self, stream, msg, gate); + kio::wait(|waiter| serve.poll(waiter)).await + } + + /// Our publish permit as it stands now, with the auth epoch it was read at. + fn permit_now(&self, gate: Option<&crate::auth::Gate>) -> (crate::auth::Permit, u64) { + let mut epoch = 0; + let permit = match gate { + Some(gate) => gate.poll_permit(&mut epoch, &kio::Waiter::noop()), + None => self + .auth + .poll_permit(crate::auth::Direction::Publish, &mut epoch, &kio::Waiter::noop()), + }; + match permit { + Poll::Ready(permit) => (permit, epoch), + Poll::Pending => (Default::default(), 0), + } + } + + /// The routes an announce cursor holds right now, without waiting for more. + /// + /// A route announced and retracted within the snapshot is left out, and a repeat + /// keeps its latest metadata, so each path appears once. + fn snapshot(announced: &mut crate::origin::Routes) -> Vec { + let mut initial = std::collections::BTreeMap::new(); + while let Some(event) = announced.try_next() { + match event { + crate::origin::RouteEvent::Start(update) + | crate::origin::RouteEvent::Update(update) + | crate::origin::RouteEvent::Restart(update) + | crate::origin::RouteEvent::Takeover { announce: update, .. } => { + initial.insert(update.prefix.clone(), update); } - let joins = self.joins.poll(waiter, |joins| match joins.get(&subscribe_id) { - Some(Some(_)) => Poll::Ready(()), - Some(None) => { - pending = true; - Poll::Pending - } - None if pending => Poll::Ready(()), - None => Poll::Pending, - }); - if let Poll::Ready(joins) = joins { - return Poll::Ready(Ok(joins.get(&subscribe_id).cloned().flatten())); + crate::origin::RouteEvent::End(update) => { + initial.remove(&update.prefix); } - if deadline.poll(waiter).is_ready() { - return Poll::Ready(if pending { Err(Error::Timeout) } else { Ok(None) }); + } + } + initial.into_values().collect() + } + + /// Drop an unanswered verification, ready once the acceptor lets go of the request's + /// stream credit. Pending, it leaves the credit in `ns` for the next call. + fn poll_abandon_verification(&self, ns: &mut Namespaces, waiter: &kio::Waiter) -> Poll<()> { + if let Some(Verifying { verdict, credit, .. }) = ns.verifying.take() { + drop(verdict); + ns.abandoned = Some(credit); + } + if let Some(credit) = &ns.abandoned { + ready!(self.poll_released(waiter, credit)); + ns.abandoned = None; + } + Poll::Ready(()) + } + + /// Answer an admitted namespace update: move to the prefix it names, if that is new, + /// then acknowledge it, buffered for the loop to flush. Returns whether the prefix moved. + fn rescope( + &self, + ns: &mut Namespaces, + prefix: &mut crate::Path, + announced: &mut crate::origin::Routes, + next: Option, + ) -> Result { + let moved = match next.filter(|next| next != prefix) { + Some(next) => { + // DONE still names the old prefix. After OK every suffix is relative to the + // new prefix, so replay the new cursor only after acknowledging. + let held: Vec<_> = ns + .watched + .iter() + .filter(|(_, watch)| watch.sent.wanted()) + .map(|(suffix, _)| suffix.clone()) + .collect(); + for suffix in held { + self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix)?; } - Poll::Pending - }) - .await - }; - let refusal = match joined { - Err(Error::Timeout) => (Error::Timeout, "subscription not ready"), - Err(err) => return Err(err), - Ok(Some(Joined::Group { - end, - cache, - namespace, - ranking, - grant, - })) => return Ok(Ok((end, cache, namespace, ranking, grant))), - Ok(None) => ( - match self.version { - Version::Draft18 => Error::InvalidJoiningRequestId, - _ => Error::NotFound, - }, - "no such subscription", - ), - Ok(Some(Joined::Unsupported)) => (Error::Unsupported, "joining filter not supported"), - Ok(Some(Joined::Empty)) => (Error::InvalidRange, "no objects at subscription start"), + ns.watched.clear(); + *prefix = next; + let scope = crate::Pattern::subtree(prefix.as_str()) + .map(|p| p.rebase(prefix.as_str())) + .unwrap_or_default(); + let origin = self + .origin + .scope(&*prefix, &scope) + .unwrap_or_else(|_| self.origin.empty()); + *announced = self.advertising(origin, &ns.peer).discovery(ns.hidden).routes(); + true + } + None => false, }; - Ok(Err(refusal)) + let stream = ns.target.stream().expect("update has a stream"); + stream.writer.buffer_message( + ietf::RequestOk::ID, + &ietf::RequestOk { + largest: None, + active: None, + }, + )?; + Ok(moved) + } + + /// Reconcile what the peer holds with one route (un)announce. Returns the advertisement + /// left to make, as [`Self::sync_namespace`] does. + fn apply_update( + &self, + ns: &mut Namespaces, + prefix: &crate::Path, + update: crate::announce::Announce, + active: bool, + ) -> Result, Error> { + let suffix = update.prefix; + let path = prefix.rejoin(&suffix); + + if active { + // A repeat for a live suffix is a metadata update: keep the + // peer's refusal state and re-run the selection. + match ns.watched.get_mut(&suffix) { + Some(watch) => watch.route = update.route, + None => { + ns.watched.insert(suffix.clone(), Watched::new(update.route)); + } + } + self.sync_namespace(ns, &suffix, &path) + } else { + // Only close out namespaces the peer actually saw. + let held = ns.watched.remove(&suffix).is_some_and(|watch| watch.sent.wanted()); + if held { + tracing::debug!(route = %crate::path::Rooted(self.origin.root(), &path), "namespace_done"); + self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix)?; + } + Ok(None) + } } +} - /// Answer a TRACK_STATUS with what a SUBSCRIBE_OK for the track would carry, without - /// subscribing. - /// - /// The track resolves as it would for a SUBSCRIBE, so a relay asks its upstream for a - /// track it does not hold yet and the answer waits on that. The Largest Location is - /// whatever the cache holds now: nothing waits for a fresher one. - #[cfg(test)] - fn run_track_status_stream( - self, - stream: Stream, - msg: ietf::TrackStatus<'static>, - ) -> impl std::future::Future> { - let gate = crate::auth::Gate::new( - self.auth.clone(), - msg.track_namespace.clone(), - crate::auth::Direction::Publish, - ); - self.serve_track_status_stream(stream, msg, gate) +/// The track's end, shared by a subscription's group writers so the one serving the end's +/// group can carry the END_OF_TRACK inline. +#[derive(Default)] +struct Terminal { + end: std::sync::OnceLock, + /// A group stream carrying the marker was acknowledged, so no standalone one is needed. + sent: AtomicBool, +} + +/// Whether the request token that admitted a request lapsed, shared by the streams the +/// request answers on. +/// +/// A grant's end drops the request's serving from outside, so its streams learn why only +/// here: each one torn down after a lapse resets with EXPIRED_AUTH_TOKEN instead of its usual +/// code, telling the requester to renew its token rather than give up. A revoked or narrowed +/// grant leaves it unset, as does every other end, and so does an error that only says a +/// token expired elsewhere, such as an upstream's. +#[derive(Clone, Default)] +struct Lapse(Arc); + +impl Lapse { + /// Record the lapse, before the subscription's streams are dropped. + fn set(&self) { + self.0.store(true, Ordering::Relaxed); } - async fn serve_track_status_stream( - self, - mut stream: Stream, - msg: ietf::TrackStatus<'static>, - mut gate: crate::auth::Gate, - ) -> Result<(), Error> { - if let Some(error) = super::namespace::refusal(&msg.track_namespace, self.version) { - return super::request_stream::reject(stream, error).await; - } - match self - .admit_request(&mut stream, &mut gate, msg.authorization_token.clone(), false) - .await - { - Ok(true) => {} - Ok(false) => return Ok(()), - Err(refusal) => return self.reject_track_status(stream, &refusal.error).await, - } - if !gate.allows() { - return self.reject_track_status(stream, &Error::Unauthorized).await; + fn lapsed(&self) -> bool { + self.0.load(Ordering::Relaxed) + } + + /// Cancel a request whose answer is already out, with EXPIRED_AUTH_TOKEN once its token + /// lapsed, as no [`StreamError`](crate::StreamError) says, and otherwise as `err` does. + fn abort(&self, writer: Writer, err: &Error) { + match self.lapsed() { + true => writer.reset(super::error::EXPIRED_AUTH_TOKEN), + false => writer.abort(err), } - // Nothing can follow a TRACK_STATUS, so a FIN only ends the requester's input; a reset or - // STOP_SENDING withdraws it. - let fin_cancels = false; - let broadcast = match gate - .guard(self.resolve(&mut stream, fin_cancels, &msg.track_namespace)) - .await - { - Ok(Some(broadcast)) => broadcast, - Ok(None) => return Ok(()), - Err(err) => return self.reject_track_status(stream, &err).await, - }; + } +} - let track = match broadcast.track(&msg.track_name) { - Ok(track) => track, - Err(err) => return self.reject_track_status(stream, &err).await, - }; +struct SubscriptionFill { + fill: FillServe, + largest: Option, + request_id: RequestId, + track: track::Consumer, + ranking: Ranking>, + timescale: Option, + opened: Arc, +} - // The requester abandoning the request is owed nothing. - let info = { - let query = track.info(); - let mut finished = false; - let mut messages = super::request_stream::Messages::new(self.session.clone(), self.control.client()); - kio::wait(|waiter| { - if super::request_stream::poll_cancel(&mut stream, &mut finished, &mut messages, fin_cancels, waiter) - .is_ready() - { - return Poll::Ready(None); - } - if gate.poll_denied(waiter).is_ready() { - return Poll::Ready(Some(Err(Error::Unauthorized))); - } - query.poll_ok(waiter).map(Some) - }) - .await - }; - let negotiated = self.peer_setup.get().await.flate; - let properties = - match info.map(|info| track_properties(&info?, msg.properties_wanted, self.version, negotiated)) { - Some(Ok(properties)) => properties, - Some(Err(err)) => return self.reject_track_status(stream, &err).await, - None => return Ok(()), - }; +/// The current group's tail, even when its head was filled rather than arrived live. +struct FillTail { + start: Location, + fetch: track::Fetching, +} - // The answer stays gated until the transport takes it: one parked on stream credit - // when the grant narrows is reset rather than delivered. - let respond = async { - let ok = ietf::TrackStatusOk { - largest: live_edge(&track).largest, - properties, - }; - stream.writer.buffer_message(ietf::TrackStatusOk::ID, &ok)?; - kio::wait(|waiter| stream.writer.poll_flush(waiter)).await?; - Ok::<(), Error>(()) - }; - let res = { - let mut respond = std::pin::pin!(respond); - kio::wait(|waiter| { - if gate.poll_denied(waiter).is_ready() { - return Poll::Ready(Err(Error::Unauthorized)); - } - waiter.poll_future(respond.as_mut()) - }) - .await - }; - if let Err(err) = res { - stream.writer.abort(&err); - return Err(err); - } +/// The stream serving a subscription's highest group, as it was opened. +struct LiveGroup { + sequence: u64, + slice: GroupSlice, + /// Set once a narrowing gave the stream up before its FIN. + retired: Arc, +} - // The answer is all this stream carries. See [`Self::reject_subscribe`] for why the - // close is not optional. - let _ = stream.writer.close().await; - Ok(()) - } +/// Resolve a request's broadcast while the peer still wants the answer, `None` once the +/// requester withdraws or the session closes. +struct Resolve { + session: Transport, + fin_cancels: bool, + path: crate::Path, + state: ResolveState, +} - async fn reject_track_status(&self, mut stream: Stream, err: &Error) -> Result<(), Error> { - let (kind, reason) = (request::Kind::TrackStatus, err.to_string()); - super::request_stream::buffer_error(&mut stream.writer, kind, err, &reason, self.version)?; - super::request_stream::deliver(&mut stream.writer).await +#[kio::machine(self.state: ResolveState)] +impl Resolve { + fn new(session: Transport, path: crate::Path, fin_cancels: bool) -> Self { + Self { + session, + path, + fin_cancels, + state: ResolveState::Start, + } } - /// Reject a FETCH, ending the request stream. See [`Self::reject_subscribe`] for why the - /// close is not optional. - async fn reject_fetch( - &self, - mut stream: Stream, - err: &Error, - reason: &str, - ) -> Result<(), Error> { - let kind = request::Kind::Fetch; - super::request_stream::buffer_error(&mut stream.writer, kind, err, reason, self.version)?; - super::request_stream::deliver(&mut stream.writer).await + fn poll( + &mut self, + waiter: &kio::Waiter, + publisher: &Publisher, + stream: &mut Stream, + ) -> Poll, Error>> { + let origin = wait!(Origin => match publisher.poll_withdrawn(&mut self.session, stream, self.fin_cancels, waiter) { + Poll::Ready(()) => Poll::Ready(None), + Poll::Pending => publisher.poll_serving_origin(waiter).map(Some), + }); + let Some(origin) = origin else { + return Poll::Ready(Ok(None)); + }; + // The immediate lookup: answering a peer never starts discovery of its own. + let requesting: kio::Pending = origin.lookup(self.path.clone(), None); + let found = wait!(Lookup => match publisher.poll_withdrawn(&mut self.session, stream, self.fin_cancels, waiter) { + Poll::Ready(()) => Poll::Ready(None), + Poll::Pending => requesting.poll_ok(waiter).map(Some), + }); + Poll::Ready(found.transpose()) } +} - /// Bring the peer's view of one namespace in line with the current selection. - /// - /// A loop writing inline (answering a SUBSCRIBE_NAMESPACE) re-sends - /// NAMESPACE on that stream, which the receiver treats as a replacement, and - /// retracts with NAMESPACE_DONE. Otherwise each advertisement rides its own - /// PUBLISH_NAMESPACE request: an update is a REQUEST_UPDATE **on the stream that - /// already carries it**, since a second stream would leave two claiming one - /// namespace, and a withdrawal cancels the request. - async fn sync_namespace(&self, ns: &mut Namespaces, suffix: &crate::Path, path: &crate::Path) -> Result<(), Error> { - let Some((advert, auth)) = self.offer(ns, suffix, path) else { - return Ok(()); - }; - let negotiated = ns.negotiated(); - let Some(watch) = ns.watched.get(suffix) else { - return Ok(()); - }; - let refused = watch.refused; - let wanted = advert.wanted(); - let held = watch.sent.wanted(); - let request = ns.requests.get(suffix); - let held_auth = request.as_ref().and_then(|request| request.auth.clone()); - let unchanged = advert == watch.sent && auth == held_auth; - // moq-transport renames a request's auth with REQUEST_UPDATE, but cannot name one on - // a request that had none, nor drop one: those withdraw the request and make it - // again. - let restart = held_auth != auth && (held_auth.is_none() || auth.is_none()); +/// Verify an incoming request token while keeping its stream credit until verification ends. +struct Admit { + session: Transport, + fin_cancels: bool, + /// Held by the verification, so the request keeps its stream credit until it ends. + credit: kio::Producer<()>, + state: AdmitState, +} - // Keep one opening advertisement per namespace. Its answer reconciles the latest - // route, so route churn cannot accumulate replacement streams behind a stalled write. - if wanted && request.is_some_and(|request| !request.accepted) { - return Ok(()); - } +enum AdmitState { + /// No token to verify, or its verdict admitted. + Admitted, + Verifying(crate::auth::RequestVerdict), + /// Withdrawn while verifying: the stream is owed until the acceptor lets go of it. + Releasing, +} - // An update already waiting or in flight carries whatever is current once it goes - // out, and its answer brings the loop back here for anything after that. - if wanted - && !restart - && let Some(request) = ns.requests.get_mut(suffix) - && !unchanged - { - if !request.queued && request.pending.is_none() { - request.queued = true; - ns.updates.push_back(suffix.clone()); +impl Admit { + fn new( + publisher: &Publisher, + gate: &crate::auth::Gate, + token: Option, + fin_cancels: bool, + ) -> Self { + let credit = kio::Producer::<()>::default(); + let state = match token.filter(|_| gate.token_decides()) { + Some(token) => { + let hold: crate::auth::Hold = Arc::new(credit.consume()); + AdmitState::Verifying(gate.verify(token.0, Some(hold))) } - return Ok(()); + None => AdmitState::Admitted, + }; + Self { + session: publisher.session.clone(), + fin_cancels, + credit, + state, } + } - if unchanged { - // Nothing to send. A namespace that is no longer advertisable is no longer - // pending either, and leaving that set would keep the retry timer armed - // forever for a wire message that can never happen. - if !wanted { - ns.watched.get_mut(suffix).expect("watched").deferred = false; + /// Ready with whether to serve the request, false once the requester withdrew it, or + /// with the refusal to answer it with. + fn poll( + &mut self, + waiter: &kio::Waiter, + publisher: &Publisher, + stream: &mut Stream, + gate: &mut crate::auth::Gate, + ) -> Poll> { + if let AdmitState::Verifying(verdict) = &self.state { + if publisher + .poll_withdrawn(&mut self.session, stream, self.fin_cancels, waiter) + .is_ready() + { + self.state = AdmitState::Releasing; + } else { + let answer = ready!(verdict.poll_grant(waiter)); + let AdmitState::Verifying(verdict) = std::mem::replace(&mut self.state, AdmitState::Admitted) else { + unreachable!("verifying"); + }; + gate.admit(&publisher.runtime, (verdict, answer))?; } - return Ok(()); } + match self.state { + AdmitState::Admitted => Poll::Ready(Ok(true)), + AdmitState::Verifying(_) => unreachable!("answered above"), + AdmitState::Releasing => { + ready!(publisher.poll_released(waiter, &self.credit)); + Poll::Ready(Ok(false)) + } + } + } +} + +/// Pending until the gate settles the tokens that might admit the request, then why it no +/// longer allows it, if it doesn't. +fn poll_gate_refusal(gate: &mut crate::auth::Gate, waiter: &kio::Waiter) -> Poll> { + ready!(gate.poll_ready(waiter)); + Poll::Ready(match gate.poll_refused(waiter) { + Poll::Ready(err) => Some(err), + Poll::Pending => None, + }) +} + +/// A request stream's answer, finished and waiting for the peer to acknowledge it, while the +/// requester's side stays open. +struct Answered { + _reader: crate::coding::Reader, + closing: crate::coding::Closing, +} + +impl Answered { + /// Finish `stream`'s answer, or `None` once that failed: a peer dropping the stream first + /// is a normal end. + fn new(stream: Stream) -> Option { + let Stream { reader, writer } = stream; + let closing = writer.closing().ok()?; + Some(Self { + _reader: reader, + closing, + }) + } +} + +/// [`PrepareFetch`] owning the track it reads until it is ready. +struct FetchPreparing { + track: track::Consumer, + prepare: PrepareFetch, +} + +/// The track [`FetchPreparing`] read, with what it found or the refusal and its reason. +type Prepared = (track::Consumer, Result); + +impl FetchPreparing { + fn new( + track: track::Consumer, + range: FetchRange, + largest_filter: Option, + joined: bool, + priority: i8, + compression: bool, + ) -> Self { + let prepare = PrepareFetch::new(&track, range, largest_filter, joined, priority, compression); + Self { track, prepare } + } - // A fresh offer waits for what the refusal asked for, whatever brought us back. - // Only withdrawing and re-announcing clears it, since that builds a fresh entry. - if wanted && !held && !refused.offerable(self.runtime.now()) { - return Ok(()); + fn poll(&mut self, waiter: &kio::Waiter) -> Poll { + let prepared = ready!(self.prepare.poll(waiter, &self.track)); + Poll::Ready((self.track.clone(), prepared)) + } +} + +/// What a joining FETCH inherits from the subscription it names. +type JoinedFetch = ( + Location, + track::Consumer, + crate::Path, + Ranking>, + Option, +); + +/// A joining FETCH's subscription, or the refusal to answer the FETCH with. +type JoinAnswer = Result; + +/// Resolve the subscription a joining FETCH names to its saved start, or the refusal to +/// answer the FETCH with. +/// +/// Request streams can arrive out of order, so this waits on the subscription's +/// registration, while bounding the lifetime of a request whose subscription never arrives +/// or resolves. +struct Joining { + /// The subscription registered and is still resolving. + pending: bool, + deadline: crate::time::Deadline, +} + +impl Joining { + fn new(runtime: &crate::time::Deadlines) -> Self { + Self { + pending: false, + deadline: crate::time::Deadline::after(runtime, Duration::from_secs(10)), } + } - let absolute = self.origin.root().rejoin(path); - // Only a PUBLISH_NAMESPACE request or an update to one can be refused; everything - // else below either rides a stream the peer already accepted or says nothing at all. - let mut refused = watch.refused; - let sent = match &mut ns.target { - Target::Requests => { - match (advert.wanted(), ns.requests.get_mut(suffix)) { - (false, _) => { - if held { - tracing::debug!(broadcast = %absolute, "namespace_done"); - } - self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone()) - .await?; - } - (true, Some(_)) => { - tracing::debug!(broadcast = %absolute, "publish_namespace restart"); - self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone()) - .await?; - refused = self - .advertise_namespace(ns, path, suffix.clone(), advert.params(), auth) - .await?; - } - (true, None) => { - tracing::debug!(broadcast = %absolute, "publish_namespace"); - refused = self - .advertise_namespace(ns, path, suffix.clone(), advert.params(), auth) - .await?; - } - } - // Drive the first write too, and keep reading once idle for GOAWAY. - if ns.requests.contains_key(suffix) { - ns.woken.suffixes.lock().unwrap().push(suffix.clone()); - } - // The peer can reject a fresh PUBLISH_NAMESPACE or an update, which leaves - // no request behind. Record each buffered advertisement now, so withdrawal - // cancels it even before its first message leaves. - match ns.requests.contains_key(suffix) { - true => advert, - false => Advert::None, - } + fn poll( + &mut self, + waiter: &kio::Waiter, + publisher: &Publisher, + stream: &mut Stream, + finished: &mut bool, + messages: &mut super::request_stream::Messages, + subscribe_id: RequestId, + ) -> Poll> { + let joined = 'joined: { + // The request reader is what the subscriber FINs or resets. + if super::request_stream::poll_cancel( + stream, + finished, + messages, + super::request_stream::fin_cancels(publisher.version), + waiter, + ) + .is_ready() + { + break 'joined Err(Error::Cancel); } - Target::Inline(stream) => { - match (advert.wanted(), held) { - (true, _) => { - tracing::debug!(broadcast = %absolute, "namespace"); - stream - .writer - .encode_message(&ietf::Namespace { - suffix: suffix.clone(), - params: negotiated.extended().then(|| advert.params()), - }) - .await?; - } - (false, true) => { - tracing::debug!(broadcast = %absolute, "namespace_done"); - stream - .writer - .encode_message(&ietf::NamespaceDone { suffix: suffix.clone() }) - .await?; - } - // Never advertised and still not advertisable: nothing to say. - (false, false) => {} + let pending = &mut self.pending; + let joins = publisher.joins.poll(waiter, |joins| match joins.get(&subscribe_id) { + Some(Some(_)) => Poll::Ready(()), + Some(None) => { + *pending = true; + Poll::Pending } - advert + None if *pending => Poll::Ready(()), + None => Poll::Pending, + }); + if let Poll::Ready(joins) = joins { + break 'joined Ok(joins.get(&subscribe_id).cloned().flatten()); + } + if self.deadline.poll(waiter).is_ready() { + break 'joined if self.pending { Err(Error::Timeout) } else { Ok(None) }; } + return Poll::Pending; }; - - // A peer that asked not to be offered this again outranks the retry timer; - // anything else it should hold and does not comes back on one. - let watch = ns.watched.get_mut(suffix).expect("watched"); - watch.refused = refused; - watch.deferred = wanted && !sent.wanted() && refused.pending(); - watch.sent = sent; - Ok(()) + let refusal = match joined { + Err(Error::Timeout) => (Error::Timeout, "subscription not ready"), + Err(err) => return Poll::Ready(Err(err)), + Ok(Some(Joined::Group { + end, + cache, + namespace, + ranking, + grant, + })) => return Poll::Ready(Ok(Ok((end, cache, namespace, ranking, grant)))), + Ok(None) => ( + match publisher.version { + Version::Draft18 => Error::InvalidJoiningRequestId, + _ => Error::NotFound, + }, + "no such subscription", + ), + Ok(Some(Joined::Unsupported)) => (Error::Unsupported, "joining filter not supported"), + Ok(Some(Joined::Empty)) => (Error::InvalidRange, "no objects at subscription start"), + }; + Poll::Ready(Ok(Err(refusal))) } +} - /// Buffer a PUBLISH_NAMESPACE and register its stream with the announce loop, so a - /// withdrawal never waits for send credit or the peer's answer. - async fn advertise_namespace( - &self, - ns: &mut Namespaces, - path: &crate::Path, - suffix: crate::Path, - params: ietf::AdvertParams, - auth: Option, - ) -> Result { - let Some(mut request) = self.open_request().await? else { - tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), path), "no stream for the advertisement"); - return Ok(Refused::No); - }; - let request_id = self.control.next_request_id(); - if let Err(err) = request.writer.buffer_message( - ietf::PublishNamespace::ID, - &ietf::PublishNamespace { - request_id, - track_namespace: path.clone(), - params, - authorization_token: auth.as_ref().and_then(|auth| self.request_token(auth)), - }, - ) { - tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), path), %err, "advertisement failed"); - return Ok(ns.watched.get_mut(&suffix).expect("watched").lapsed(self.runtime.now())); +/// A standalone FETCH's range, until its track resolves. +struct FetchRequest { + namespace: crate::Path, + track: std::borrow::Cow<'static, str>, + start: Location, + end: Location, + /// Draft-22's filter, resolved against the Largest Object once the track is read. + largest_filter: Option, +} + +/// What an admitted FETCH stays held to while its groups load and its response is written. +/// +/// A joining one is refused if its subscription ends before FETCH_OK. After that the head is +/// already held and the requester can still cancel, so it finishes on its own authority: the +/// gate, and the token grant if one authorized the subscription, which ends it on the token's +/// expiry, revocation, or a refused renewal. +struct Admitted { + gate: crate::auth::Gate, + /// A standalone FETCH's own token, which its gate holds, or a joining one's subscription's. + token: Option, + /// Set once that token lapses, so the streams this FETCH answers on say so. + lapse: Lapse, + /// The subscription a joining FETCH is refused without, until FETCH_OK. + subscription: Option, + ranking: Ranking>, +} + +impl Admitted { + fn new( + gate: crate::auth::Gate, + token: Option, + subscription: Option, + ranking: Ranking>, + ) -> Self { + let token = token.or_else(|| gate.token_guard()); + Self { + gate, + token, + lapse: Lapse::default(), + subscription, + ranking, } - ns.requests.insert( - suffix.clone(), - NamespaceRequest { - messages: super::request_stream::Messages::new(self.session.clone(), self.control.client()), - answer: None, - path: path.clone(), - stream: request, - auth, - accepted: false, - queued: false, - pending: Some(NamespacePending::Writing { _auth: None }), - waiter: None, - }, - ); - Ok(Refused::No) } - /// What the peer should hold for one namespace: the advertisement, and the auth its - /// PUBLISH_NAMESPACE request carries. `None` when the namespace is not watched. - /// - /// Nothing our grant does not cover reaches the wire, and a shrink withdraws it. Only a - /// request carries an auth, and only one this session's handle issued: any other route - /// is advertised under the session's scope. - fn offer( - &self, - ns: &Namespaces, - suffix: &crate::Path, - path: &crate::Path, - ) -> Option<(Advert, Option)> { - let watch = ns.watched.get(suffix)?; - let advert = match ns.permitted(path) { - true => self.select(&watch.route, &ns.peer), - false => Advert::None, - }; - let auth = match (&ns.target, advert.wanted()) { - (Target::Requests, true) => watch - .route - .auth - .as_ref() - .filter(|auth| self.auth.token(auth).is_some()) - .cloned(), - _ => None, - }; - Some((advert, auth)) + /// Ready with why the token's grant ended, ahead of the gate it also closes, so a lapse + /// is refused as one. + fn poll_token_ended(&self, waiter: &kio::Waiter) -> Poll { + match &self.token { + Some(grant) => grant.poll_ended(waiter).map(|refusal| { + if super::error::token_expired(&refusal.error) { + self.lapse.set(); + } + refusal + }), + None => Poll::Pending, + } } - /// The `AUTHORIZATION TOKEN` a request names `auth` with, by value. - fn request_token(&self, auth: &crate::auth::Credential) -> Option { - let token = self.auth.token(auth)?; - Some(super::token::RequestToken(crate::setup::Token { - kind: crate::setup::Token::OUT_OF_BAND, - value: token.to_vec(), - })) + /// Whether a joining FETCH's subscription ended. + fn unsubscribed(&self, publisher: &Publisher, waiter: &kio::Waiter) -> bool { + self.subscription.is_some_and(|id| { + publisher + .joins + .poll(waiter, |joins| match joins.contains_key(&id) { + true => Poll::Pending, + false => Poll::Ready(()), + }) + .is_ready() + }) } +} - /// Drop queued updates that no longer have anything to send, so the next one to take - /// a request ID is one that will use it: drafts before 17 require the next ID. - fn prune_updates(&self, ns: &mut Namespaces, prefix: &crate::Path) { - while let Some(suffix) = ns.updates.front() { - let wants = ns.requests.get(suffix).is_some_and(|request| { - request.queued - && request.pending.is_none() - && ns.watched.get(suffix).is_some_and(|watch| { - self.offer(ns, suffix, &prefix.rejoin(suffix)) - .is_some_and(|(advert, auth)| { - advert.wanted() && (advert != watch.sent || auth != request.auth) - }) - }) - }); - if wants { - return; - } - let suffix = ns.updates.pop_front().expect("checked non-empty"); - if let Some(request) = ns.requests.get_mut(&suffix) { - request.queued = false; - } +/// Answer a FETCH: a standalone or filtered range of the named track, or a joining FETCH's +/// groups up to where its subscription starts. +/// +/// Its highest group is read before replying: FETCH_OK names where the response ends, which +/// a range running past the track only learns by reading it, and a refusal can still replace +/// it until then. +struct FetchServe { + publisher: Publisher, + /// The request stream, until its answer is finished or reset. + stream: Option>, + request_id: RequestId, + priority: i8, + properties_wanted: bool, + /// What the requester sent on this stream so far, kept across every phase that watches + /// it for a cancel, so a second GOAWAY is caught wherever it arrives. + finished: bool, + messages: super::request_stream::Messages, + /// The track once resolved, held until the request ends. + track: Option, + state: FetchState, +} + +#[expect(clippy::large_enum_variant, reason = "a machine's enum is its storage")] +enum FetchState { + Start { + msg: ietf::Fetch<'static>, + admission: FetchGate, + }, + /// Verifying a standalone FETCH's own token. + Admit { + admit: Admit, + gate: crate::auth::Gate, + request: FetchRequest, + }, + /// Resolving a standalone FETCH's broadcast under its gate. + Resolve { + resolve: Resolve, + gate: crate::auth::Gate, + request: FetchRequest, + }, + /// Waiting on the subscription a joining FETCH names. + Join { + joining: Joining, + fetch_type: FetchType<'static>, + auth: crate::auth::Link, + id: RequestId, + }, + /// Waiting on the peer's SETUP, which says whether it reads compressed objects. + Setup { + admitted: Admitted, + track: track::Consumer, + range: FetchRange, + largest_filter: Option, + joined: bool, + }, + /// Reading the group before FETCH_OK promises it, while a refusal is still possible. + Prepare { + admitted: Admitted, + prepare: FetchPreparing, + compression: bool, + joined: bool, + }, + /// FETCH_OK, raced against the grant: a standalone FETCH whose grant narrows while it + /// waits on stream credit stops there instead of sending what it no longer may. + Answer { + admitted: Admitted, + /// Taken once buffered. + ok: Option, + /// Where the response starts, which ranks its stream. + start: Location, + write: FetchedGroupWrite, + }, + /// The response on its own stream. Cancellation releases the buffered answer even while + /// its stream is flow controlled, and a joining one no longer ends with its subscription. + Respond { + admitted: Admitted, + respond: FetchRespond, + }, + /// FETCH_OK is the last thing this stream has to say, and [`Writer`] resets on drop: + /// without the finish the answer is discarded before the peer ever reads it, exactly as + /// it would be for a refusal. + Close { + _admitted: Admitted, + answered: Answered, + }, + Reject(super::request_stream::Deliver), + Done, +} + +impl FetchServe { + fn new( + publisher: Publisher, + stream: Stream, + msg: ietf::Fetch<'static>, + admission: FetchGate, + ) -> Self { + let messages = super::request_stream::Messages::new(publisher.session.clone(), publisher.control.client()); + Self { + request_id: msg.request_id, + priority: super::priority::from_wire(msg.subscriber_priority), + properties_wanted: msg.properties_wanted, + publisher, + stream: Some(stream), + finished: false, + messages, + track: None, + state: FetchState::Start { msg, admission }, } } - /// Send the update waiting at the front of the queue with `request_id`: REQUEST_UPDATE - /// on the request that carries the namespace, with only what changed since the peer's - /// copy. Its answer arrives through [`Self::poll_answer`] while the loop carries on. - /// - /// One update is outstanding per request, which satisfies any MAX_REQUEST_UPDATES the - /// peer set without reading it. A different original publisher is an update like any - /// other: the receiver drains what it already serves from the old one and never splices - /// the two. - async fn send_update(&self, ns: &mut Namespaces, prefix: &crate::Path, request_id: RequestId) -> Result<(), Error> { - let suffix = ns.updates.pop_front().expect("allocated for a queued update"); - let Some((advert, auth)) = self.offer(ns, &suffix, &prefix.rejoin(&suffix)) else { - return Ok(()); - }; - let (Some(request), Some(watch)) = (ns.requests.get_mut(&suffix), ns.watched.get_mut(&suffix)) else { - return Ok(()); - }; - request.queued = false; - let mut update = match (advert.params().cluster, watch.sent.params().cluster) { - (Some(next), Some(held)) => ietf::PublishNamespaceUpdate::between(request_id, &held, &next), - _ => ietf::PublishNamespaceUpdate { - request_id, - hops: None, - cost: None, - authorization_token: None, - }, + fn fin_cancels(&self) -> bool { + super::request_stream::fin_cancels(self.publisher.version) + } + + /// Refuse the request with `error`, delivered on the following polls. + fn refuse(&mut self, error: ietf::RequestError<'_>) -> Result<(), Error> { + let stream = self.stream.as_mut().expect("refused before answering"); + stream.writer.buffer_message(ietf::RequestError::ID, &error)?; + self.state = FetchState::Reject(super::request_stream::Deliver::new()); + Ok(()) + } + + /// Refuse the FETCH with `err`. See [`Publisher::reject_subscribe`] for why the close is + /// not optional. + fn reject(&mut self, err: &Error, reason: &str) -> Result<(), Error> { + let stream = self.stream.as_mut().expect("refused before answering"); + let kind = request::Kind::Fetch; + super::request_stream::buffer_error(&mut stream.writer, kind, err, reason, self.publisher.version)?; + self.state = FetchState::Reject(super::request_stream::Deliver::new()); + Ok(()) + } + + /// Reset the request stream once its answer is out: no REQUEST_ERROR can follow FETCH_OK. + fn abort(&mut self, admitted: Admitted, err: Error) -> Poll> { + let stream = self.stream.take().expect("aborted once"); + admitted.lapse.abort(stream.writer, &err); + Poll::Ready(Err(err)) + } + + /// Everything up to the first wait: refusals the request alone decides, then admission. + fn start(&mut self, msg: ietf::Fetch<'static>, admission: FetchGate) -> Result<(), Error> { + let version = self.publisher.version; + let namespace = match &msg.fetch_type { + FetchType::Standalone { namespace, .. } | FetchType::Filtered { namespace, .. } => Some(namespace), + _ => None, }; - if auth != request.auth { - update.authorization_token = auth.as_ref().and_then(|auth| self.request_token(auth)); + if let Some(error) = namespace.and_then(|namespace| super::namespace::refusal(namespace, version)) { + return self.refuse(error); } - tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), &request.path), "publish_namespace update"); - let written = request - .stream - .writer - .buffer_message(ietf::PublishNamespaceUpdate::ID, &update); - if let Err(err) = written { - // An unencodable update ends only this advertisement; the retry re-offers it - // fresh, carrying the complete parameters. - tracing::debug!(%err, "publish_namespace update not sent"); - self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone()) - .await?; - if let Some(watch) = ns.watched.get_mut(&suffix) { - watch.refused = Refused::No; - watch.deferred = Refused::No.pending(); - watch.sent = Advert::None; - } - return Ok(()); + + // Serving a Range Filter unfiltered would deliver objects the subscriber excluded. + if msg.range_filters { + return self.refuse(ietf::RequestError { + error_code: request::INVALID_FILTER, + reason_phrase: "Range Filters exceed MAX_FILTER_RANGES=0".into(), + retry_interval: 0, + }); } - // Keep one buffered update until the peer answers it. - request.pending = Some(NamespacePending::Writing { - _auth: std::mem::replace(&mut request.auth, auth), - }); - watch.sent = advert; - // Marked woken, so the loop's next pass reads the answer it may already have. - ns.woken.suffixes.lock().unwrap().push(suffix); - Ok(()) - } + // FILL_TIMEOUT is accepted and ignored: a one-group read ends with the group. - /// Flush writes and read answers, polling only the requests that woke. - fn poll_answer( - requests: &mut HashMap, - inflight: &mut VecDeque<(crate::time::Instant, crate::Path)>, - woken: &Arc, - now: crate::time::Instant, - waiter: &kio::Waiter, - ) -> Poll<(crate::Path, Result)> { - // Registered before the take, so a wake landing between them is not lost. - *woken.owner.lock().unwrap() = Some(waiter.waker().clone()); - let mut suffixes = std::mem::take(&mut *woken.suffixes.lock().unwrap()); - while let Some(suffix) = suffixes.pop() { - let Some(request) = requests.get_mut(&suffix) else { - continue; - }; - // Kept on the request, so its registrations outlive this pass. - let request_waiter = request - .waiter - .get_or_insert_with(|| kio::Waiter::new(woken.waker(suffix.clone()))); - let written = request.stream.writer.poll_flush(request_waiter); - if matches!(&written, Poll::Ready(Ok(()))) - && matches!(request.pending, Some(NamespacePending::Writing { .. })) - { - let deadline = now + ADVERTISE_TIMEOUT; - request.pending = Some(NamespacePending::Answer(deadline)); - inflight.push_back((deadline, suffix.clone())); + // Draft-22's LOCATION_FILTER is answered as the standalone range it spells. Every other + // filter ends at the Largest Object, which resolves once the track does. + let mut largest_filter = None; + let fetch_type = match msg.fetch_type { + FetchType::Filtered { + namespace, + track, + filter: Filter::Absolute { start, end: Some(end) }, + } => { + // The filter's End Object is inclusive, where a standalone one is the last + // object plus one. Without one, both include the whole End Group. + let Some(object) = end.object.map_or(Some(0), |object| object.checked_add(1)) else { + return self.reject(&Error::InvalidRange, "End Object overflows"); + }; + FetchType::Standalone { + namespace, + track, + start, + end: Location { + group: end.group, + object, + }, + } } - let answer = match written { - Poll::Ready(Err(err)) => Poll::Ready(Err(err)), - Poll::Ready(Ok(())) if request.answer.is_some() => Poll::Ready(Ok(request.answer.take())), - _ => loop { - let answer = request.messages.poll_next(&mut request.stream.reader, request_waiter); - match answer { - Poll::Ready(Ok(Some(_))) if request.answer.is_some() => { - break Poll::Ready(Err(request.messages.failed(Error::UnexpectedMessage))); - } - Poll::Ready(Ok(Some(framed))) if written.is_pending() && framed.0 == ietf::RequestOk::ID => { - // An early OK must not admit another update or replacement behind - // this blocked write. Keep one, while still reading peer terminals. - request.answer = Some(framed); - } - answer => break answer, - } - }, - }; - if let Poll::Ready(answer) = answer { - // The rest stay woken for the next pass. - woken.suffixes.lock().unwrap().append(&mut suffixes); - let answer = answer.and_then(|answer| { - if answer.is_some() && request.pending.is_none() { - return Err(request.messages.failed(Error::UnexpectedMessage)); - } - Ok(answer.map_or(Answer::Ended, Answer::Message)) - }); - woken.suffixes.lock().unwrap().push(suffix.clone()); - return Poll::Ready((suffix, answer)); + FetchType::Filtered { + namespace, + track, + filter, + } => { + largest_filter = Some(filter); + // Placeholders until the filter resolves against the Largest Object. + let whole = Location { group: 0, object: 0 }; + FetchType::Standalone { + namespace, + track, + start: whole, + end: whole, + } } - } - while let Some((deadline, suffix)) = inflight.front() { - let current = matches!( - requests.get(suffix).and_then(|request| request.pending.as_ref()), - Some(NamespacePending::Answer(at)) if at == deadline - ); - if current && now < *deadline { - break; + other => other, + }; + + // Every FETCH is held to the grant while its groups load and its response is written. + // A standalone one is also checked here, before it reaches the origin; a joining one + // was checked when its subscription was. + self.state = match fetch_type { + FetchType::Standalone { + namespace, + track, + start, + end, + } => { + let FetchGate::Standalone(gate) = admission else { + return Err(Error::Unauthorized); + }; + let admit = Admit::new(&self.publisher, &gate, msg.authorization_token, self.fin_cancels()); + FetchState::Admit { + admit, + gate, + request: FetchRequest { + namespace, + track, + start, + end, + largest_filter, + }, + } } - let (_, suffix) = inflight.pop_front().expect("just peeked"); - if current { - return Poll::Ready((suffix, Ok(Answer::Late))); + FetchType::RelativeJoining { + subscriber_request_id, .. + } + | FetchType::AbsoluteJoining { + subscriber_request_id, .. + } => { + let FetchGate::Joining { id, auth } = admission else { + return Err(Error::Unauthorized); + }; + debug_assert_eq!(id, subscriber_request_id); + FetchState::Join { + joining: Joining::new(&self.publisher.runtime), + fetch_type, + auth, + id: subscriber_request_id, + } } + // Rewritten as standalone above. + FetchType::Filtered { .. } => return self.reject(&Error::Unsupported, "not supported"), + }; + Ok(()) + } + + /// A standalone FETCH its token admitted: checked against the grant and its range, then + /// resolved like a SUBSCRIBE, since the peer must have seen the announcement to name it. + fn admitted(&mut self, gate: crate::auth::Gate, request: FetchRequest) -> Result<(), Error> { + if !gate.allows() { + return self.reject(&Error::Unauthorized, "not granted"); } - Poll::Pending + let FetchRequest { + start, + end, + largest_filter, + .. + } = request; + + // End Object 0 includes the entire End Group, even at the largest Group ID. + if largest_filter.is_none() + && (start.group > end.group || (start.group == end.group && end.object != 0 && start.object >= end.object)) + { + return self.reject(&Error::InvalidRange, "empty range"); + } + + if largest_filter.is_none() && start.group != end.group { + return self.reject(&Error::Unsupported, "FETCH spans more than one group"); + } + + let resolve = Resolve::new( + self.publisher.session.clone(), + request.namespace.clone(), + self.fin_cancels(), + ); + self.state = FetchState::Resolve { resolve, gate, request }; + Ok(()) } - /// Apply an advertisement or update's answer: an OK keeps the namespace, and a - /// REQUEST_ERROR closed the stream and withdrew the advertisement (moq-transport Section - /// 9.5.1), so the request is dropped and the retry re-offers it fresh. An unanswered - /// update is dropped the same way, since a peer that ignored it cannot be assumed to hold - /// either version, and so is one the peer cancelled instead of answering. An idle - /// advertisement the peer ends waits out its own growing backoff before coming back. - async fn answered( - &self, - ns: &mut Namespaces, - prefix: &crate::Path, - suffix: crate::Path, - answer: Result, + /// Judge the group read for the FETCH, then promise it in FETCH_OK or refuse it. + fn answer( + &mut self, + admitted: Admitted, + prepared: PreparedFetch, + compression: bool, + joined: bool, ) -> Result<(), Error> { - let Some(request) = ns.requests.get_mut(&suffix) else { - return Ok(()); - }; - let initial = !request.accepted; - let idle = request.pending.take().is_none(); - let absolute = self.origin.root().rejoin(&request.path); - - // A reset stream ends only this request, like a FIN; a malformed answer is the peer's - // violation and still ends the loop. - let answer = match answer { - Err(err) if !super::subscriber::is_protocol_violation(&err) => { - tracing::debug!(broadcast = %absolute, %err, "advertisement stream ended"); - Answer::Ended + let version = self.publisher.version; + let PreparedFetch { range, last: group } = prepared; + let end = range.end; + let group = group.without_status(); + // Only the TIMESCALE in FETCH_OK opts the answer into timestamps, so a requester that + // opted out of the properties gets an untimed track. + let timescale = group.info.timescale.filter(|_| self.properties_wanted); + let frames = group + .frames + .iter() + .filter_map(|frame| Some((frame.fields.as_deref()?, frame.timestamp))); + for (fields, frame) in frames.chain(group.end_object.as_deref().map(|end| (end, None))) { + if let Err(err) = super::object::write(&fields.properties, group.info.ietf(), frame, version, timescale) { + return self.reject(&err, &err.to_string()); } - answer => answer?, + } + let properties = match track_properties(&group.info, self.properties_wanted, version, compression) { + Ok(properties) => properties, + Err(err) => return self.reject(&err, &err.to_string()), }; - let refused = match answer { - Answer::Late => { - tracing::debug!(broadcast = %absolute, "no answer to the advertisement"); - self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone()) - .await?; - Refused::No + let draft20 = Filter::is_draft20(version); + // Only a finished track's final object ends it, never a live track's Largest Object. + let track = self.track.as_ref().expect("read before answering"); + let end_of_track = !joined && group.complete && track.final_sequence() == group.sequence.checked_add(1); + + let end_location = if joined { + // The subscription starts at the saved Largest Object, so the prefix of that + // group has to be here in full or the two leave a gap. + if group.sequence != end.group || group.end() != end.object { + return self.reject(&Error::NotFound, "joining prefix unavailable"); } - Answer::Ended if idle || initial => { - tracing::debug!(broadcast = %absolute, "advertisement ended"); - self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone()) - .await?; - let now = self.runtime.now(); - ns.watched - .get_mut(&suffix) - .map_or(Refused::No, |watch| watch.lapsed(now)) + end + } else if draft20 { + // Draft-22's End Location is inclusive. An empty answer still covers its start. + let object = match range.until(group.sequence) { + // The track's last object as this draft writes it, without a status object. + Some(until) if end_of_track => (until - 1).min(group.end().saturating_sub(1)), + Some(until) => until - 1, + None => group.end().saturating_sub(1).max(range.skip(group.sequence)), + }; + Location { + group: group.sequence, + object, } - Answer::Ended => { - tracing::debug!(broadcast = %absolute, "update cancelled"); - // The peer closed its side; closing ours completes the withdrawal. - self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone()) - .await?; - Refused::No + } else if end_of_track && end.object > 0 { + // A bounded response cannot promise objects past the track's end. + Location { + group: group.sequence, + object: group.end(), } - Answer::Message(super::request_stream::Framed(type_id, body)) => { - match type_id { - ietf::RequestOk::ID => { - let msg: ietf::RequestOk = body - .message(self.version) - .map_err(|err| request.messages.failed(err.into()))?; - // ACTIVE_COUNT only answers SUBSCRIBE_NAMESPACE (MoQ Active Count). - if msg.active.is_some() { - self.session.clone().close( - super::error::to_session_code(&crate::SessionError::ProtocolViolation), - "ACTIVE_COUNT on a PUBLISH_NAMESPACE answer", - ); - return Err(Error::ProtocolViolation); + } else { + end + }; + + // FETCH_OK on every draft, never REQUEST_OK: section 5.2 allows exactly one FETCH_OK or + // REQUEST_ERROR in answer to a FETCH, and REQUEST_OK's own definition lists the other + // requests it answers without ever naming this one. + self.state = FetchState::Answer { + admitted, + ok: Some(ietf::FetchOk { + end_of_track, + end_location, + properties, + }), + start: range.start, + write: FetchedGroupWrite::new(group, version, timescale), + }; + Ok(()) + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + loop { + if let FetchState::Close { answered, .. } = &mut self.state { + // The peer dropping the stream first is a normal end. + let _ = ready!(answered.closing.poll(waiter)); + self.state = FetchState::Done; + return Poll::Ready(Ok(())); + } + let fin_cancels = self.fin_cancels(); + let stream = self.stream.as_mut().expect("held until answered"); + match &mut self.state { + FetchState::Start { .. } => { + let FetchState::Start { msg, admission } = std::mem::replace(&mut self.state, FetchState::Done) + else { + unreachable!("starting"); + }; + self.start(msg, admission)?; + } + FetchState::Admit { admit, gate, .. } => { + let admitted = ready!(admit.poll(waiter, &self.publisher, stream, gate)); + let FetchState::Admit { gate, request, .. } = std::mem::replace(&mut self.state, FetchState::Done) + else { + unreachable!("admitting"); + }; + match admitted { + Ok(true) => self.admitted(gate, request)?, + Ok(false) => return Poll::Ready(Ok(())), + Err(refusal) => self.reject(&refusal.error, &refusal.reason)?, + } + } + FetchState::Resolve { resolve, gate, .. } => { + let resolved = match ready!(poll_gate_refusal(gate, waiter)) { + Some(err) => Err(err), + None => ready!(resolve.poll(waiter, &self.publisher, stream)), + }; + let FetchState::Resolve { gate, request, .. } = + std::mem::replace(&mut self.state, FetchState::Done) + else { + unreachable!("resolving"); + }; + let broadcast = match resolved { + Ok(Some(broadcast)) => broadcast, + Ok(None) => return Poll::Ready(Ok(())), + Err(err) => { + self.reject(&err, &err.to_string())?; + continue; } - tracing::debug!(message = ?msg, "publish_namespace ok"); - request.accepted = true; - // Whatever changed while it was in flight goes out next. - return self.sync_namespace(ns, &suffix, &prefix.rejoin(&suffix)).await; + }; + let track = match broadcast.track(&request.track) { + Ok(track) => track, + Err(err) => { + self.reject(&err, &err.to_string())?; + continue; + } + }; + + // Its own entry in the queue, or its own send group, so it competes with group + // streams by its priority alone, its request id breaking ties in the queue. + // Nothing updates a FETCH's priority, and it ranks before the track's info + // resolves, so the publisher's priority is the unset default. + let ranking = Ranking { + queue: self.publisher.priority.clone(), + group: self + .publisher + .session + .clone() + .send_group(crate::priority::group(self.priority, 0)), + request_id: self.request_id, + priority: kio::Producer::new(self.priority).consume(), + publisher: 0, + order: Order::Descending, + lapse: Lapse::default(), + }; + let range = FetchRange { + start: request.start, + end: request.end, + }; + self.state = FetchState::Setup { + admitted: Admitted::new(gate, None, None, ranking), + track, + range, + largest_filter: request.largest_filter, + joined: false, + }; + } + FetchState::Join { joining, id, .. } => { + let joined = ready!(joining.poll( + waiter, + &self.publisher, + stream, + &mut self.finished, + &mut self.messages, + *id + ))?; + let FetchState::Join { + fetch_type, auth, id, .. + } = std::mem::replace(&mut self.state, FetchState::Done) + else { + unreachable!("joining"); + }; + let (end, cache, namespace, ranking, grant) = match joined { + Ok(joined) => joined, + Err((err, reason)) => { + self.reject(&err, reason)?; + continue; + } + }; + // The cache outlives the subscription, so the fetch holds its own gate rather than + // trusting the cache. One the subscription's request token authorized is held to + // the peer's grant alone, as the subscription is. + let gate = crate::auth::Gate::new(auth, namespace, crate::auth::Direction::Publish); + let gate = if grant.is_some() { gate.named() } else { gate }; + let start = match fetch_type { + FetchType::RelativeJoining { group_offset: 0, .. } => end.group, + FetchType::RelativeJoining { .. } => { + self.reject(&Error::Unsupported, "joining FETCH reaches before the joined group")?; + continue; + } + FetchType::AbsoluteJoining { group_id, .. } if group_id <= end.group => group_id, + _ => { + self.reject(&Error::InvalidRange, "joining group past the subscription")?; + continue; + } + }; + let range = FetchRange { + start: Location { + group: start, + object: 0, + }, + end, + }; + self.state = FetchState::Setup { + admitted: Admitted::new(gate, grant, Some(id), ranking), + track: cache, + range, + largest_filter: None, + joined: true, + }; + } + FetchState::Setup { .. } => { + let compression = ready!(self.publisher.peer_setup.poll_get(waiter)).flate; + let FetchState::Setup { + admitted, + track, + range, + largest_filter, + joined, + } = std::mem::replace(&mut self.state, FetchState::Done) + else { + unreachable!("waiting on SETUP"); + }; + let prepare = FetchPreparing::new(track, range, largest_filter, joined, self.priority, compression); + self.state = FetchState::Prepare { + admitted, + prepare, + compression, + joined, + }; + } + FetchState::Prepare { admitted, prepare, .. } => { + if super::request_stream::poll_cancel( + stream, + &mut self.finished, + &mut self.messages, + fin_cancels, + waiter, + ) + .is_ready() + { + return Poll::Ready(Ok(())); } - ietf::RequestError::ID => { - let msg: ietf::RequestError = body - .message(self.version) - .map_err(|err| request.messages.failed(err.into()))?; - tracing::warn!(message = ?msg, "publish_namespace error"); - // The peer closed its side; closing ours completes the withdrawal. - self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix.clone()) - .await?; - self.refusal(msg.retry_interval) + let prepared = 'prepared: { + if let Poll::Ready(refusal) = admitted.poll_token_ended(waiter) { + break 'prepared Err((refusal.error, refusal.reason)); + } + if admitted.gate.poll_denied(waiter).is_ready() + || admitted.unsubscribed(&self.publisher, waiter) + { + let err = Error::Unauthorized; + let reason = err.to_string(); + break 'prepared Err((err, reason)); + } + Ok(ready!(prepare.poll(waiter))) + }; + let FetchState::Prepare { + admitted, + compression, + joined, + .. + } = std::mem::replace(&mut self.state, FetchState::Done) + else { + unreachable!("preparing"); + }; + let prepared = match prepared { + Ok((track, prepared)) => { + self.track = Some(track); + prepared + } + Err(refusal) => Err(refusal), + }; + match prepared { + Ok(prepared) => self.answer(admitted, prepared, compression, joined)?, + Err((err, reason)) => self.reject(&err, &reason)?, } - _ => return Err(request.messages.failed(Error::UnexpectedMessage)), } + FetchState::Answer { admitted, ok, .. } => { + let res = 'ok: { + if let Poll::Ready(refusal) = admitted.poll_token_ended(waiter) { + break 'ok Err(refusal.error); + } + if admitted.gate.poll_denied(waiter).is_ready() + || admitted.unsubscribed(&self.publisher, waiter) + { + break 'ok Err(Error::Unauthorized); + } + if let Some(ok) = ok.take() + && let Err(err) = stream.writer.buffer_message(ietf::FetchOk::ID, &ok) + { + break 'ok Err(err); + } + ready!(stream.writer.poll_flush(waiter)) + }; + let FetchState::Answer { + admitted, start, write, .. + } = std::mem::replace(&mut self.state, FetchState::Done) + else { + unreachable!("answering"); + }; + if let Err(err) = res { + return self.abort(admitted, err); + } + let respond = FetchRespond::new(&self.publisher, &admitted, start, self.request_id, write); + self.state = FetchState::Respond { admitted, respond }; + } + FetchState::Respond { admitted, respond } => { + if super::request_stream::poll_cancel( + stream, + &mut self.finished, + &mut self.messages, + fin_cancels, + waiter, + ) + .is_ready() + { + // The fetch stream, if open, resets as it drops with the response. + self.state = FetchState::Done; + return Poll::Ready(Ok(())); + } + let res = 'respond: { + if let Poll::Ready(refusal) = admitted.poll_token_ended(waiter) { + break 'respond Err(refusal.error); + } + if admitted.gate.poll_denied(waiter).is_ready() { + break 'respond Err(Error::Unauthorized); + } + ready!(respond.poll(waiter)) + }; + let FetchState::Respond { admitted, .. } = std::mem::replace(&mut self.state, FetchState::Done) + else { + unreachable!("responding"); + }; + if let Err(err) = res { + // The fetch stream, if open, reset as it dropped with the response. Once + // FETCH_OK is out, no REQUEST_ERROR can follow it, so the reset is the + // whole refusal. + return self.abort(admitted, err); + } + let stream = self.stream.take().expect("answered once"); + let Some(answered) = Answered::new(stream) else { + return Poll::Ready(Ok(())); + }; + self.state = FetchState::Close { + _admitted: admitted, + answered, + }; + } + FetchState::Close { .. } => unreachable!("closed above"), + FetchState::Reject(deliver) => { + let res = ready!(deliver.poll(waiter, &mut stream.writer)); + self.state = FetchState::Done; + return Poll::Ready(res); + } + FetchState::Done => panic!("polled after it finished"), } - }; - - if let Some(watch) = ns.watched.get_mut(&suffix) { - watch.refused = refused; - watch.deferred = refused.pending(); - watch.sent = Advert::None; } - Ok(()) } +} - /// How to read a refusal's retry interval, in milliseconds. - /// - /// 0 is the peer asking not to be offered this again, which is what keeps a permanent - /// refusal (unauthorized, uninterested) from becoming a request every few seconds for - /// the life of the session. - fn refusal(&self, retry_interval: u64) -> Refused { - match retry_interval { - 0 => Refused::Never, - ms => Refused::Until(self.runtime.now() + Duration::from_millis(ms)), +/// A FETCH's response on its own stream, ranked as the group it carries, held until the +/// peer acknowledges it. +struct FetchRespond { + session: Transport, + version: Version, + request_id: RequestId, + lapse: Lapse, + /// Taken by the stream once it opens. + rank: Option>>, + open: crate::transport::open::Open, + write: FetchedGroupWrite, + state: FetchRespondState, +} + +#[kio::machine(self.state: FetchRespondState)] +impl FetchRespond { + fn new( + publisher: &Publisher, + admitted: &Admitted, + start: Location, + request_id: RequestId, + write: FetchedGroupWrite, + ) -> Self { + Self { + session: publisher.session.clone(), + version: publisher.version, + request_id, + lapse: admitted.lapse.clone(), + rank: Some(admitted.ranking.rank(start)), + open: Default::default(), + write, + state: FetchRespondState::Start, } } - pub(super) fn stats(&self) -> &crate::stats::Session { - self.origin.stats() + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + let uni = wait!(Open => self.session.poll_open_uni( + &mut self.open, + waiter, + self.rank.as_mut().expect("ranked until opened").options(), + )) + .map_err(Error::from_transport)?; + let mut writer: Writer, Version> = + Writer::new(Ranked::new(uni, self.rank.take()).lapsing(&self.lapse), self.version); + writer.buffer_varint(FetchHeader::TYPE)?; + writer.buffer(&FetchHeader { + request_id: self.request_id, + })?; + wait!(Header => writer.poll_flush(waiter))?; + wait!(Body => self.write.poll(waiter, writer))?; + // Settle: the writer keeps the stream, and its rank, until the acknowledgement. + writer.finish()?; + Poll::Ready(wait!(Settle => writer.poll_close(waiter))) } +} - /// Open a stream for one advertisement, or `None` if it did not get the peer's stream - /// credit and take its header in time. - /// - /// The announce loop is single-threaded over origin updates, so an open that parks - /// forever parks everything, including the unannounces that release the streams the - /// peer is waiting on us to retire. Giving up instead keeps the loop moving. - /// - /// An advertisement that gave up comes back on the retry timer: it is our own state the - /// peer should hold, and credit returning raises no signal the loop could wait on. - async fn open_request(&self) -> Result>, Error> { - let mut open = std::pin::pin!(Stream::open(&self.session, self.version)); - let mut timeout = crate::time::Deadline::after(&self.runtime, ADVERTISE_TIMEOUT); - - kio::wait(|waiter| { - if let Poll::Ready(res) = waiter.poll_future(open.as_mut()) { - return Poll::Ready(match res { - Ok(stream) => Ok(Some(stream.with_stats(self.origin.stats()))), - Err(err) => Err(err), - }); - } - if timeout.poll(waiter).is_ready() { - return Poll::Ready(Ok(None)); - } - Poll::Pending - }) - .await +/// Answer a TRACK_STATUS: see [`Publisher::serve_track_status_stream`]. +struct TrackStatusServe { + publisher: Publisher, + /// The request stream, until its answer is finished or reset. + stream: Option>, + namespace: crate::Path, + track_name: std::borrow::Cow<'static, str>, + properties_wanted: bool, + gate: crate::auth::Gate, + /// The track once resolved, held until the request ends. + track: Option, + state: TrackStatusState, +} + +enum TrackStatusState { + Start(Option), + Admit(Admit), + Resolve(Resolve), + /// The requester abandoning the request is owed nothing. + Info { + query: kio::Pending, + finished: bool, + messages: super::request_stream::Messages, + }, + /// Waiting on the peer's SETUP, which says whether it reads compressed tracks. + Setup(Result), + /// The answer stays gated until the transport takes it: one parked on stream credit when + /// the grant narrows is reset rather than delivered. + Respond(Option), + /// The answer is all this stream carries. See [`Publisher::reject_subscribe`] for why the + /// close is not optional. + Close(Answered), + Reject(super::request_stream::Deliver), + Done, +} + +impl TrackStatusServe { + fn new( + publisher: Publisher, + stream: Stream, + msg: ietf::TrackStatus<'static>, + gate: crate::auth::Gate, + ) -> Self { + Self { + publisher, + stream: Some(stream), + namespace: msg.track_namespace, + track_name: msg.track_name, + properties_wanted: msg.properties_wanted, + gate, + track: None, + state: TrackStatusState::Start(msg.authorization_token), + } } - /// Withdraw an advertised namespace: NAMESPACE_DONE inline, or end the request that - /// carried it. - async fn withdraw_namespace( - &self, - target: &mut Target, - requests: &mut HashMap, - suffix: crate::Path, - ) -> Result<(), Error> { - match target { - Target::Requests => { - if let Some(request) = requests.remove(&suffix) { - end_request(request.stream); + /// Refuse the TRACK_STATUS with `err`, delivered on the following polls. + fn reject(&mut self, err: &Error) -> Result<(), Error> { + let stream = self.stream.as_mut().expect("refused before answering"); + let (kind, reason) = (request::Kind::TrackStatus, err.to_string()); + super::request_stream::buffer_error(&mut stream.writer, kind, err, &reason, self.publisher.version)?; + self.state = TrackStatusState::Reject(super::request_stream::Deliver::new()); + Ok(()) + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + // Nothing can follow a TRACK_STATUS, so a FIN only ends the requester's input; a reset or + // STOP_SENDING withdraws it. + let fin_cancels = false; + loop { + if let TrackStatusState::Close(answered) = &mut self.state { + let _ = ready!(answered.closing.poll(waiter)); + self.state = TrackStatusState::Done; + return Poll::Ready(Ok(())); + } + let stream = self.stream.as_mut().expect("held until answered"); + match &mut self.state { + TrackStatusState::Start(token) => { + if let Some(error) = super::namespace::refusal(&self.namespace, self.publisher.version) { + stream.writer.buffer_message(ietf::RequestError::ID, &error)?; + self.state = TrackStatusState::Reject(super::request_stream::Deliver::new()); + continue; + } + let admit = Admit::new(&self.publisher, &self.gate, token.take(), fin_cancels); + self.state = TrackStatusState::Admit(admit); + } + TrackStatusState::Admit(admit) => { + match ready!(admit.poll(waiter, &self.publisher, stream, &mut self.gate)) { + Ok(true) => {} + Ok(false) => return Poll::Ready(Ok(())), + Err(refusal) => { + self.reject(&refusal.error)?; + continue; + } + } + if !self.gate.allows() { + self.reject(&Error::Unauthorized)?; + continue; + } + let resolve = Resolve::new(self.publisher.session.clone(), self.namespace.clone(), fin_cancels); + self.state = TrackStatusState::Resolve(resolve); + } + TrackStatusState::Resolve(resolve) => { + let resolved = match ready!(poll_gate_refusal(&mut self.gate, waiter)) { + Some(err) => Err(err), + None => ready!(resolve.poll(waiter, &self.publisher, stream)), + }; + let broadcast = match resolved { + Ok(Some(broadcast)) => broadcast, + Ok(None) => return Poll::Ready(Ok(())), + Err(err) => { + self.reject(&err)?; + continue; + } + }; + let track = match broadcast.track(&self.track_name) { + Ok(track) => track, + Err(err) => { + self.reject(&err)?; + continue; + } + }; + let query = track.info(); + self.track = Some(track); + self.state = TrackStatusState::Info { + query, + finished: false, + messages: super::request_stream::Messages::new( + self.publisher.session.clone(), + self.publisher.control.client(), + ), + }; + } + TrackStatusState::Info { + query, + finished, + messages, + } => { + if super::request_stream::poll_cancel(stream, finished, messages, fin_cancels, waiter).is_ready() { + return Poll::Ready(Ok(())); + } + if self.gate.poll_denied(waiter).is_ready() { + self.reject(&Error::Unauthorized)?; + continue; + } + let info = ready!(query.poll_ok(waiter)); + self.state = TrackStatusState::Setup(info); + } + TrackStatusState::Setup(_) => { + let negotiated = ready!(self.publisher.peer_setup.poll_get(waiter)).flate; + let TrackStatusState::Setup(info) = std::mem::replace(&mut self.state, TrackStatusState::Done) + else { + unreachable!("waiting on SETUP"); + }; + let properties = info.and_then(|info| { + track_properties(&info, self.properties_wanted, self.publisher.version, negotiated) + }); + match properties { + Ok(properties) => self.state = TrackStatusState::Respond(Some(properties)), + Err(err) => self.reject(&err)?, + } + } + TrackStatusState::Respond(properties) => { + let res = 'respond: { + if self.gate.poll_denied(waiter).is_ready() { + break 'respond Err(Error::Unauthorized); + } + if let Some(properties) = properties.take() { + let track = self.track.as_ref().expect("resolved before answering"); + let ok = ietf::TrackStatusOk { + largest: live_edge(track).largest, + properties, + }; + if let Err(err) = stream.writer.buffer_message(ietf::TrackStatusOk::ID, &ok) { + break 'respond Err(err); + } + } + ready!(stream.writer.poll_flush(waiter)) + }; + let stream = self.stream.take().expect("answered once"); + if let Err(err) = res { + stream.writer.abort(&err); + return Poll::Ready(Err(err)); + } + let Some(answered) = Answered::new(stream) else { + return Poll::Ready(Ok(())); + }; + self.state = TrackStatusState::Close(answered); } + TrackStatusState::Close(_) => unreachable!("closed above"), + TrackStatusState::Reject(deliver) => { + let res = ready!(deliver.poll(waiter, &mut stream.writer)); + self.state = TrackStatusState::Done; + return Poll::Ready(res); + } + TrackStatusState::Done => panic!("polled after it finished"), } - Target::Inline(stream) => { - stream - .writer - .encode_message(&ietf::NamespaceDone { suffix: suffix.clone() }) - .await?; - } - } - Ok(()) - } - - /// Close out every open PUBLISH_NAMESPACE request. A no-op for a loop whose entries - /// ride the SUBSCRIBE_NAMESPACE stream itself, which retracts them by ending. - async fn withdraw_requests(&self, target: &mut Target, requests: &mut HashMap) { - let suffixes: Vec = requests.keys().cloned().collect(); - for suffix in suffixes { - let _ = self.withdraw_namespace(target, requests, suffix).await; } } +} - /// Advertise every namespace we can, without waiting to be asked. - /// - /// moq-transport itself says nothing about which of the two discovery messages a peer - /// expects, and the peers that never send SUBSCRIBE_NAMESPACE are exactly the ones - /// expecting a publisher to announce itself, so the default has to be to announce. A - /// peer that would rather ask says so with the MoQ Solicit extension - /// ([`solicit`](super::solicit)), - /// and then this loop does nothing. - /// [`Self::serve_subscribe_namespace_stream`] also carries every - /// match, so a peer that did not ask hears each namespace both ways. A NAMESPACE is - /// discovery only, not a second route. - pub async fn run_publish_namespaces(self) -> Result<(), Error> { - if self.requires_solicitation().await { - return Ok(()); - } - - // The cluster extension changes what an advertisement carries, so nothing can be - // sent until the peer's SETUP says whether it speaks it. - let peer = self.peer().await; - - // Split horizon, as the solicited loop applies it: never advertise a route back - // to the peer it came from. - let origin = self - .advertising(self.origin.clone(), &peer) - .discovery(!self.peer_setup.get().await.hidden || self.origin.includes_hidden()); +/// An advertisement [`Publisher::sync_namespace`] decided on, waiting for a stream of its own. +struct Advertise { + suffix: crate::Path, + path: crate::Path, + advert: Advert, + auth: Option, + /// Whether the peer should hold the namespace. + wanted: bool, +} - let ns = Namespaces::new(peer, Target::Requests); - self.run_namespaces(origin.routes(), crate::Path::empty(), ns, Vec::new()) - .await - } +/// Open a stream for one advertisement, or `None` if it did not get the peer's stream credit +/// and take its header in time. +/// +/// The announce loop is single-threaded over origin updates, so an open that parks forever +/// parks everything, including the unannounces that release the streams the peer is waiting +/// on us to retire. Giving up instead keeps the loop moving. +/// +/// An advertisement that gave up comes back on the retry timer: it is our own state the peer +/// should hold, and credit returning raises no signal the loop could wait on. +struct Opening { + advertise: Advertise, + open: crate::transport::open::Open, + timeout: crate::time::Deadline, +} - /// Handle a SUBSCRIBE_NAMESPACE on its bidi stream. - /// - /// All the announce state is local to this task (mirroring `lite::Publisher`'s - /// announce handling): whatever this subscription advertised is withdrawn - /// when its stream ends. Every match is a NAMESPACE on this stream, whatever the - /// peer's SETUP said. - #[cfg(test)] - async fn run_subscribe_namespace_stream( - self, - stream: Stream, - msg: ietf::SubscribeNamespace, - ) -> Result<(), Error> { - let gate = crate::auth::Gate::new( - self.auth.clone(), - msg.namespace.clone(), - crate::auth::Direction::Publish, - ); - self.serve_subscribe_namespace_stream(stream, msg, gate).await +impl Opening { + fn new(publisher: &Publisher, advertise: Advertise) -> Self { + Self { + advertise, + open: Default::default(), + timeout: crate::time::Deadline::after(&publisher.runtime, ADVERTISE_TIMEOUT), + } } - async fn serve_subscribe_namespace_stream( - self, - mut stream: Stream, - msg: ietf::SubscribeNamespace, - mut gate: crate::auth::Gate, - ) -> Result<(), Error> { - if let Some(error) = super::namespace::refusal(&msg.namespace, self.version) { - return super::request_stream::reject(stream, error).await; + fn poll( + &mut self, + waiter: &kio::Waiter, + publisher: &Publisher, + ) -> Poll>, Error>> { + if let Poll::Ready(res) = Stream::poll_open(&publisher.session, &mut self.open, publisher.version, waiter) { + return Poll::Ready(res.map(|stream| Some(stream.with_stats(publisher.origin.stats())))); } - gate = gate.discovery(); - match self - .admit_request( - &mut stream, - &mut gate, - msg.authorization_token.clone(), - super::request_stream::fin_cancels(self.version), - ) - .await - { - Ok(true) => {} - Ok(false) => return Ok(()), - Err(refusal) => { - super::request_stream::write_error( - &mut stream.writer, - request::Kind::SubscribeNamespace, - &refusal.error, - &refusal.reason, - self.version, - ) - .await?; - let _ = stream.writer.close().await; - return Ok(()); - } - } - let prefix = msg.namespace.clone(); - - tracing::debug!(prefix = %crate::path::Rooted(self.origin.root(), &prefix), "subscribe_namespace stream"); - - // A prefix outside our scope (empty origin, or a token that doesn't grant it) - // just means we have nothing to announce; respond with an empty set rather than - // erroring, which would look fatal to the peer. - // The wire prefix decodes as a literal path; convert it explicitly to its - // subtree grant, refusing anything that cannot be a subtree. - // The cursor is rooted at the prefix so every update arrives named as its - // wire suffix, a route covering the prefix included: that one presents at - // the root, as the empty suffix. - let scope = crate::Pattern::subtree(prefix.as_str()) - .map(|subtree| subtree.rebase(prefix.as_str())) - .unwrap_or_default(); - let origin = self - .origin - .scope(&prefix, &scope) - .unwrap_or_else(|_| self.origin.empty()); - - // The extension changes what an advertisement carries, so nothing can be - // sent until the peer's SETUP says whether it speaks it. The same SETUP says - // whether the OK counts what it sends first (MoQ Active Count). - let peer = self.peer().await; - let declared = self.peer_setup.get().await; - // Register the split-horizon peer on the announce cursor too. The origin - // model uses this exposure to park a reflected copy before it can replace - // the source we are currently advertising to that peer. - let origin = self.advertising(origin, &peer); - - // Peers that declared MoQ Hidden filter namespaces unless they opted in. A publish - // origin that already opted in (the caller's choice for this peer) keeps them. - let origin = origin.discovery(!declared.hidden || msg.hidden || self.origin.includes_hidden()); - - // A NAMESPACE is discovery, not a route, so a peer that also hears the unsolicited - // PUBLISH_NAMESPACE still gets the match here. - let mut announced = origin.routes(); - - // MoQ Active Count: take what is advertised now, so the OK can say how many - // NAMESPACE messages carry it. They go out first, ahead of any change, so with - // MoQ Auth the count waits for our grant and leaves out what it does not cover. - let (initial, active) = match declared.active_count { - true => { - if declared.auth { - kio::wait(|waiter| gate.poll_setup_answered(waiter)).await; - } - let (permit, _) = self.permit_now(Some(&gate)); - let initial = Self::snapshot(&mut announced); - let count = initial - .iter() - .filter(|update| permit.matches(prefix.rejoin(&update.prefix).as_str())) - .filter(|update| self.select(&update.route, &peer).wanted()) - .count(); - (initial, Some(count as u64)) - } - false => (Vec::new(), None), - }; + if self.timeout.poll(waiter).is_ready() { + return Poll::Ready(Ok(None)); + } + Poll::Pending + } +} - // Send the acknowledgement only while the token is live, even behind stream credit. - crate::auth::until( - stream.writer.encode_message(&ietf::RequestOk { largest: None, active }), - |waiter| gate.poll_token_ended(waiter), - ) - .await - .map_err(|refusal| refusal.error)??; +/// A step of an announce loop's event left for the loop to finish, in order, before it reads +/// the next event. +enum NamespaceWork { + /// Reconcile one route (un)announce, through [`Publisher::apply_update`]. + Apply(crate::announce::Announce, bool), + /// Bring one namespace in line, through [`Publisher::sync_namespace`]. + Sync(crate::Path), +} - let target = Target::Inline(stream); +/// Forward origin (un)announces to the peer until the loop ends, leaving a token still being +/// verified in the [`Namespaces`] it runs over: [`NamespaceLoop`]'s body. +struct ForwardNamespaces { + announced: crate::origin::Routes, + prefix: crate::Path, + /// The routes counted before the loop started, applied ahead of any change. + initial: Vec, + finished: bool, + goaways: super::request_stream::Messages, + /// Registered on the first poll, as the loop starts. + withdrawing: Option, + /// When to re-offer whatever the peer should hold and doesn't, and how long to wait the + /// next time that fails. Jittered so a relay's namespaces don't all come back on the same + /// tick. + retry: crate::time::Deadline, + retry_at: Option, + retry_delay: Duration, + /// When the earliest update in flight gives up on its answer. + answer_by: crate::time::Deadline, + /// A peer streaming REQUEST_UPDATEs must not hold the session's other tasks. + budget: kio::coop::Budget, + /// Whether this turn already yielded to the budget and armed the retry. + yielded: bool, + session: Transport, + /// What the last event left to do. + work: VecDeque, + /// The advertisement waiting on its stream, which holds up the work behind it. + opening: Option, + state: ForwardState, +} - let mut ns = Namespaces::new(peer, target); - ns.gate = Some(gate); - ns.hidden = !declared.hidden || msg.hidden || self.origin.includes_hidden(); - self.run_namespaces(announced, prefix, ns, initial).await - } +enum ForwardState { + Start, + /// Waiting on the peer's SETUP, and with MoQ Auth on the answer to the credential we + /// presented at setup, so the first advertisement, the initial set included, is already + /// checked against our grant. + Setup, + /// Streaming updates (origin route (un)announces), until the peer closes its side. + Run, + /// The origin is gone and everything is withdrawn: finish the stream and wait for + /// delivery, once a verification still holding the request's credit lets go of it. + Drain, + /// Ending the request with what is buffered on its stream, waiting for the peer to + /// acknowledge it. + Finish { + finished: bool, + }, +} - /// Our publish permit as it stands now, with the auth epoch it was read at. - fn permit_now(&self, gate: Option<&crate::auth::Gate>) -> (crate::auth::Permit, u64) { - let mut epoch = 0; - let permit = match gate { - Some(gate) => gate.poll_permit(&mut epoch, &kio::Waiter::noop()), - None => self - .auth - .poll_permit(crate::auth::Direction::Publish, &mut epoch, &kio::Waiter::noop()), - }; - match permit { - Poll::Ready(permit) => (permit, epoch), - Poll::Pending => (Default::default(), 0), +impl ForwardNamespaces { + fn new( + publisher: &Publisher, + announced: crate::origin::Routes, + prefix: crate::Path, + initial: Vec, + ) -> Self { + Self { + announced, + prefix, + initial, + finished: false, + goaways: super::request_stream::Messages::new(publisher.session.clone(), publisher.control.client()), + withdrawing: None, + retry: crate::time::Deadline::new(&publisher.runtime), + retry_at: None, + retry_delay: RETRY_BASE, + answer_by: crate::time::Deadline::new(&publisher.runtime), + budget: kio::coop::Budget::new(32), + yielded: false, + session: publisher.session.clone(), + work: VecDeque::new(), + opening: None, + state: ForwardState::Start, } } - /// The routes an announce cursor holds right now, without waiting for more. - /// - /// A route announced and retracted within the snapshot is left out, and a repeat - /// keeps its latest metadata, so each path appears once. - fn snapshot(announced: &mut crate::origin::Routes) -> Vec { - let mut initial = std::collections::BTreeMap::new(); - while let Some(event) = announced.try_next() { - match event { - crate::origin::RouteEvent::Start(update) - | crate::origin::RouteEvent::Update(update) - | crate::origin::RouteEvent::Restart(update) - | crate::origin::RouteEvent::Takeover { announce: update, .. } => { - initial.insert(update.prefix.clone(), update); + fn poll(&mut self, waiter: &kio::Waiter, publisher: &Publisher, ns: &mut Namespaces) -> Poll> { + loop { + match &mut self.state { + ForwardState::Start => { + self.withdrawing = Some(publisher.withdrawal.register()); + if publisher.withdrawal.poll(&kio::Waiter::noop()).is_ready() { + return Poll::Ready(Ok(())); + } + self.state = ForwardState::Setup; } - crate::origin::RouteEvent::End(update) => { - initial.remove(&update.prefix); + ForwardState::Setup => { + if ready!(publisher.peer_setup.poll_get(waiter)).auth { + ready!(publisher.auth.poll_setup_answered(waiter)); + } + (ns.permit, ns.epoch) = publisher.permit_now(ns.gate.as_ref()); + self.work = std::mem::take(&mut self.initial) + .into_iter() + .map(|update| NamespaceWork::Apply(update, true)) + .collect(); + self.state = ForwardState::Run; + } + ForwardState::Run => { + if let Some(res) = ready!(self.poll_run(waiter, publisher, ns))? { + return Poll::Ready(self.end(publisher, ns, res)); + } + } + ForwardState::Drain => { + ready!(publisher.poll_abandon_verification(ns, waiter)); + if ns.target.stream().is_none() { + return Poll::Ready(Ok(())); + } + self.state = ForwardState::Finish { finished: false }; + } + ForwardState::Finish { finished } => { + let stream = ns.target.stream().expect("finishing a stream"); + ready!(stream.writer.poll_flush(waiter))?; + if !*finished { + // Dropping a finished writer resets it, which could discard what it carries. + stream.writer.finish()?; + *finished = true; + } + let res = ready!(stream.writer.poll_closed(waiter)); + return Poll::Ready(self.end(publisher, ns, res)); } } } - initial.into_values().collect() } - /// Forward origin (un)announces to the peer until the loop ends. - /// - /// Shared by both announce paths: they differ in where the advertisements go - /// ([`Target`]) and where the origin is rooted (`prefix`, empty when nothing asked - /// for a subset). The token that admitted the request, if one did, guards the loop. - async fn run_namespaces( - &self, - announced: crate::origin::Routes, - prefix: crate::Path, - mut ns: Namespaces, - initial: Vec, - ) -> Result<(), Error> { - let token = ns.gate.as_ref().and_then(|gate| gate.token_guard()); - let forward = self.forward_namespaces(announced, prefix, &mut ns, initial); - let res = crate::auth::RequestGrant::guard(token.as_ref(), forward) - .await - .unwrap_or_else(|refusal| Err(refusal.error)); - // Every exit, a failed write's or an ended token's included, keeps the request's - // stream, and so its credit, until the acceptor lets go of a token it is still - // verifying. - self.abandon_verification(&mut ns).await; + /// This loop's advertisements die with it. + fn end(&mut self, publisher: &Publisher, ns: &mut Namespaces, res: Result<(), Error>) -> Result<(), Error> { + publisher.withdraw_requests(&mut ns.target, &mut ns.requests); res } - /// [`Self::run_namespaces`]'s loop, leaving a token still being verified in `ns`. - async fn forward_namespaces( - &self, - mut announced: crate::origin::Routes, - mut prefix: crate::Path, - ns: &mut Namespaces, - initial: Vec, - ) -> Result<(), Error> { - let mut finished = false; - let mut goaways = super::request_stream::Messages::new(self.session.clone(), self.control.client()); - let _withdrawing = self.withdrawal.register(); - if self.withdrawal.poll(&kio::Waiter::noop()).is_ready() { - return Ok(()); + /// Make the advertisement `sync` left, if any, before anything after it. + fn advertise(&mut self, publisher: &Publisher, advertise: Option) { + if let Some(advertise) = advertise { + self.opening = Some(Opening::new(publisher, advertise)); } - // With MoQ Auth, wait for the answer to the credential we presented at setup, so - // the first advertisement, the initial set included, is already checked against - // our grant. - if self.peer_setup.get().await.auth { - kio::wait(|waiter| self.auth.poll_setup_answered(waiter)).await; - } - (ns.permit, ns.epoch) = self.permit_now(ns.gate.as_ref()); - for update in initial { - self.apply_update(ns, &prefix, update, true).await?; - } - - // When to re-offer whatever the peer should hold and doesn't, and how long to wait - // the next time that fails. Jittered so a relay's namespaces don't all come back on - // the same tick. - let mut retry = crate::time::Deadline::new(&self.runtime); - let mut retry_at: Option = None; - // When the earliest update in flight gives up on its answer. - let mut answer_by = crate::time::Deadline::new(&self.runtime); - let mut retry_delay = RETRY_BASE; - - // Stream updates (origin route (un)announces), bailing if the peer closes - // its side first. A peer streaming REQUEST_UPDATEs must not hold the session's other - // tasks. - let mut budget = kio::coop::Budget::new(32); - let mut session = self.session.clone(); - let res = loop { - kio::wait(|waiter| budget.poll_yield(waiter)).await; - match ns.watched.values().any(|watch| watch.deferred) { - // Arm on the edge, so a turn that changes nothing else doesn't push the - // deadline out forever. - true => retry_at = retry_at.or_else(|| Some(self.runtime.now() + jitter(retry_delay))), - false => { - retry_at = None; - retry_delay = RETRY_BASE; - } + } + + /// Run the loop until an event ends it, ready with how; a state change to finish the + /// request is ready with `None`. + fn poll_run( + &mut self, + waiter: &kio::Waiter, + publisher: &Publisher, + ns: &mut Namespaces, + ) -> Poll>, Error>> { + loop { + if let Some(opening) = &mut self.opening { + let stream = ready!(opening.poll(waiter, publisher))?; + let Opening { advertise, .. } = self.opening.take().expect("opening"); + publisher.advertised(ns, advertise, stream)?; + continue; } - retry.set(retry_at); - self.prune_updates(ns, &prefix); - - let event = { - let Namespaces { - target, - requests, - inflight, - woken, - updates, - epoch, - gate, - verifying, - .. - } = &mut *ns; - kio::wait(|waiter| { - if self.withdrawal.poll(waiter).is_ready() { - return Poll::Ready(NamespaceEvent::Update(None)); - } - if let Some(gate) = gate - && let Poll::Ready(refusal) = gate.poll_token_ended(waiter) - { - return Poll::Ready(NamespaceEvent::Denied(refusal)); - } - // A grant change applies before the next update, so a namespace it no - // longer covers is withdrawn rather than re-sent. - let permit = match gate { - Some(gate) => gate.poll_permit(epoch, waiter), - None => self.auth.poll_permit(crate::auth::Direction::Publish, epoch, waiter), - }; - if let Poll::Ready(permit) = permit { - return Poll::Ready(NamespaceEvent::Regrant(permit)); - } - match verifying { - // Nothing more is read until the token's verdict, so the updates behind it - // wait in the stream, under its flow control, and are answered in order. - Some(verifying) => { - if let Some(stream) = target.stream() - && self - .poll_withdrawn( - &mut session, - stream, - super::request_stream::fin_cancels(self.version), - waiter, - ) - .is_ready() - { - return Poll::Ready(NamespaceEvent::Closed(Ok(()))); - } - if let Poll::Ready(answer) = verifying.verdict.poll_grant(waiter) { - return Poll::Ready(NamespaceEvent::Verified(answer)); - } - } - None => { - if let Poll::Ready(event) = - target.poll_event(&mut finished, &mut goaways, self.version, waiter) - { - return Poll::Ready(event); - } - } - } - // Advertisements and updates flush and read on their own streams, so none - // holds up the withdrawals behind it. - let now = self.runtime.now(); - if let Poll::Ready((suffix, answer)) = Self::poll_answer(requests, inflight, woken, now, waiter) { - return Poll::Ready(NamespaceEvent::Answered(suffix, answer)); + if let Some(work) = self.work.pop_front() { + let advertise = match work { + NamespaceWork::Apply(update, active) => publisher.apply_update(ns, &self.prefix, update, active)?, + NamespaceWork::Sync(suffix) => { + let path = self.prefix.rejoin(&suffix); + publisher.sync_namespace(ns, &suffix, &path)? } - // The front is now the earliest update still waiting. - answer_by.set(inflight.front().map(|(deadline, _)| *deadline)); - let _ = answer_by.poll(waiter); - if !updates.is_empty() { - return Poll::Ready(NamespaceEvent::Renew(self.control.next_request_id())); - } - if let Poll::Ready(next) = announced.poll_next(waiter) { - return Poll::Ready(NamespaceEvent::Update(next.map(|event| match event { - crate::origin::RouteEvent::Start(update) | crate::origin::RouteEvent::Update(update) => { - (update, true, false) - } - crate::origin::RouteEvent::Restart(update) => (update, true, true), - // moq-transport has no takeover: another instance, in place or not. - crate::origin::RouteEvent::Takeover { announce, predecessor } => { - let restart = predecessor == announce.prefix; - (announce, true, restart) - } - crate::origin::RouteEvent::End(update) => (update, false, false), - }))); + }; + self.advertise(publisher, advertise); + continue; + } + // What the last event wrote inline goes out before anything else happens. + if let Some(stream) = ns.target.stream() { + ready!(stream.writer.poll_flush(waiter))?; + } + + if !self.yielded { + ready!(self.budget.poll_yield(waiter)); + self.yielded = true; + match ns.watched.values().any(|watch| watch.deferred) { + // Arm on the edge, so a turn that changes nothing else doesn't push the + // deadline out forever. + true => { + self.retry_at = self + .retry_at + .or_else(|| Some(publisher.runtime.now() + jitter(self.retry_delay))) } - if retry.poll(waiter).is_ready() { - return Poll::Ready(NamespaceEvent::Retry); + false => { + self.retry_at = None; + self.retry_delay = RETRY_BASE; } - Poll::Pending - }) - .await - }; + } + self.retry.set(self.retry_at); + publisher.prune_updates(ns, &self.prefix); + } - match event { - // Ended as the request's guard would have, had it seen the token end first: the - // stream is established, so no REQUEST_ERROR may follow its REQUEST_OK. - NamespaceEvent::Denied(refusal) => return Err(refusal.error), - - // Only a solicited loop has a stream to close, and its request task closes the - // session over a protocol violation. - NamespaceEvent::Closed(res) => break res, - NamespaceEvent::Request(mut update) => { - if let Err(err) = self.control.accept(update.request_id) { - if matches!(err, Error::GoingAway) { - update.refused = Some(err); - } else { - self.session.clone().close( - super::error::to_session_code(&crate::SessionError::from(&err)), - "invalid request ID", - ); - break Err(err); - } - } - let refusal = update - .refused - .map(|error| ietf::RequestError { - error_code: request::to_code(&error, request::Kind::SubscribeNamespace, self.version), - reason_phrase: "namespace update refused".into(), - retry_interval: 0, - }) - .or_else(|| { - update - .prefix - .as_ref() - .and_then(|prefix| super::namespace::refusal(prefix, self.version)) - }); - if let Some(error) = refusal { - let stream = ns.target.stream().expect("update has a stream"); - stream.writer.encode_message(&error).await?; - // Dropping a finished writer resets it, which could discard the refusal. - stream.writer.finish()?; - break stream.writer.closed().await; - } - let gate = ns.gate.as_ref().expect("update has a gate"); - match update.authorization_token.filter(|_| gate.token_decides()) { - // A request the session admitted is never moved to a token's authority. - Some(_) if gate.token_guard().is_none() => { - let stream = ns.target.stream().expect("update has a stream"); - super::request_stream::write_error( - &mut stream.writer, - request::Kind::SubscribeNamespace, - &Error::Unsupported, - TOKEN_ON_SESSION_REQUEST, - self.version, - ) - .await?; - stream.writer.finish()?; - break stream.writer.closed().await; - } - // A renewal, answered once the acceptor decides; the loop runs on meanwhile. - Some(token) => { - let credit = kio::Producer::<()>::default(); - let hold: crate::auth::Hold = Arc::new(credit.consume()); - ns.verifying = Some(Verifying { - verdict: gate.verify(token.0, Some(hold)), - credit, - prefix: update.prefix, - }); - } - None => { - if self.rescope(ns, &mut prefix, &mut announced, update.prefix).await? { - retry_at = None; - retry_delay = RETRY_BASE; - } - } + let event = ready!(self.poll_event(waiter, publisher, ns)); + self.yielded = false; + if let Some(end) = self.handle(publisher, ns, event)? { + return Poll::Ready(Ok(Some(end))); + } + if !matches!(self.state, ForwardState::Run) { + return Poll::Ready(Ok(None)); + } + } + } + + fn poll_event(&mut self, waiter: &kio::Waiter, publisher: &Publisher, ns: &mut Namespaces) -> Poll { + let Namespaces { + target, + requests, + inflight, + woken, + updates, + epoch, + gate, + verifying, + .. + } = ns; + if publisher.withdrawal.poll(waiter).is_ready() { + return Poll::Ready(NamespaceEvent::Update(None)); + } + if let Some(gate) = gate + && let Poll::Ready(refusal) = gate.poll_token_ended(waiter) + { + return Poll::Ready(NamespaceEvent::Denied(refusal)); + } + // A grant change applies before the next update, so a namespace it no longer covers is + // withdrawn rather than re-sent. + let permit = match gate { + Some(gate) => gate.poll_permit(epoch, waiter), + None => publisher + .auth + .poll_permit(crate::auth::Direction::Publish, epoch, waiter), + }; + if let Poll::Ready(permit) = permit { + return Poll::Ready(NamespaceEvent::Regrant(permit)); + } + match verifying { + // Nothing more is read until the token's verdict, so the updates behind it wait in + // the stream, under its flow control, and are answered in order. + Some(verifying) => { + if let Some(stream) = target.stream() + && publisher + .poll_withdrawn( + &mut self.session, + stream, + super::request_stream::fin_cancels(publisher.version), + waiter, + ) + .is_ready() + { + return Poll::Ready(NamespaceEvent::Closed(Ok(()))); + } + if let Poll::Ready(answer) = verifying.verdict.poll_grant(waiter) { + return Poll::Ready(NamespaceEvent::Verified(answer)); + } + } + None => { + if let Poll::Ready(event) = + target.poll_event(&mut self.finished, &mut self.goaways, publisher.version, waiter) + { + return Poll::Ready(event); + } + } + } + // Advertisements and updates flush and read on their own streams, so none holds up the + // withdrawals behind it. + let now = publisher.runtime.now(); + if let Poll::Ready((suffix, answer)) = Publisher::poll_answer(requests, inflight, woken, now, waiter) { + return Poll::Ready(NamespaceEvent::Answered(suffix, answer)); + } + // The front is now the earliest update still waiting. + self.answer_by.set(inflight.front().map(|(deadline, _)| *deadline)); + let _ = self.answer_by.poll(waiter); + if !updates.is_empty() { + return Poll::Ready(NamespaceEvent::Renew(publisher.control.next_request_id())); + } + if let Poll::Ready(next) = self.announced.poll_next(waiter) { + return Poll::Ready(NamespaceEvent::Update(next.map(|event| match event { + crate::origin::RouteEvent::Start(update) | crate::origin::RouteEvent::Update(update) => { + (update, true, false) + } + crate::origin::RouteEvent::Restart(update) => (update, true, true), + // moq-transport has no takeover: another instance, in place or not. + crate::origin::RouteEvent::Takeover { announce, predecessor } => { + let restart = predecessor == announce.prefix; + (announce, true, restart) + } + crate::origin::RouteEvent::End(update) => (update, false, false), + }))); + } + if self.retry.poll(waiter).is_ready() { + return Poll::Ready(NamespaceEvent::Retry); + } + Poll::Pending + } + + /// Apply one event: `Some` with how the loop ends, or `None` to carry on, perhaps in a + /// state that ends it. + fn handle( + &mut self, + publisher: &Publisher, + ns: &mut Namespaces, + event: NamespaceEvent, + ) -> Result>, Error> { + let version = publisher.version; + match event { + // Ended as the request's guard would have, had it seen the token end first: the + // stream is established, so no REQUEST_ERROR may follow its REQUEST_OK. + NamespaceEvent::Denied(refusal) => return Err(refusal.error), + + // Only a solicited loop has a stream to close, and its request task closes the + // session over a protocol violation. + NamespaceEvent::Closed(res) => return Ok(Some(res)), + NamespaceEvent::Request(mut update) => { + if let Err(err) = publisher.control.accept(update.request_id) { + if matches!(err, Error::GoingAway) { + update.refused = Some(err); + } else { + publisher.session.clone().close( + super::error::to_session_code(&crate::SessionError::from(&err)), + "invalid request ID", + ); + return Ok(Some(Err(err))); } } - NamespaceEvent::Verified(answer) => { - let Verifying { - verdict, prefix: next, .. - } = ns.verifying.take().expect("a token is verifying"); - let gate = ns.gate.as_mut().expect("update has a gate"); - if let Err(refusal) = gate.admit(&self.runtime, (verdict, answer)) { + let refusal = update + .refused + .map(|error| ietf::RequestError { + error_code: request::to_code(&error, request::Kind::SubscribeNamespace, version), + reason_phrase: "namespace update refused".into(), + retry_interval: 0, + }) + .or_else(|| { + update + .prefix + .as_ref() + .and_then(|prefix| super::namespace::refusal(prefix, version)) + }); + if let Some(error) = refusal { + let stream = ns.target.stream().expect("update has a stream"); + stream.writer.buffer_message(ietf::RequestError::ID, &error)?; + self.state = ForwardState::Finish { finished: false }; + return Ok(None); + } + let gate = ns.gate.as_ref().expect("update has a gate"); + match update.authorization_token.filter(|_| gate.token_decides()) { + // A request the session admitted is never moved to a token's authority. + Some(_) if gate.token_guard().is_none() => { let stream = ns.target.stream().expect("update has a stream"); - super::request_stream::write_error( + super::request_stream::buffer_error( &mut stream.writer, request::Kind::SubscribeNamespace, - &refusal.error, - &refusal.reason, - self.version, - ) - .await?; - stream.writer.finish()?; - break stream.writer.closed().await; - } - ns.epoch = 0; - if self.rescope(ns, &mut prefix, &mut announced, next).await? { - retry_at = None; - retry_delay = RETRY_BASE; + &Error::Unsupported, + TOKEN_ON_SESSION_REQUEST, + version, + )?; + self.state = ForwardState::Finish { finished: false }; } - } - NamespaceEvent::Retry => { - retry_at = None; - retry_delay = (retry_delay * 2).min(RETRY_MAX); - - // A minimum wait the peer named is enforced by `sync_namespace`, which - // every path goes through, so a namespace still inside one simply - // makes no offer this turn. The next sweep is at most RETRY_MAX away. - let deferred: Vec = ns - .watched - .iter() - .filter(|(_, watch)| watch.deferred) - .map(|(suffix, _)| suffix.clone()) - .collect(); - - for suffix in deferred { - let path = prefix.rejoin(&suffix); - self.sync_namespace(ns, &suffix, &path).await?; + // A renewal, answered once the acceptor decides; the loop runs on meanwhile. + Some(token) => { + let credit = kio::Producer::<()>::default(); + let hold: crate::auth::Hold = Arc::new(credit.consume()); + ns.verifying = Some(Verifying { + verdict: gate.verify(token.0, Some(hold)), + credit, + prefix: update.prefix, + }); } - } - NamespaceEvent::Regrant(permit) => { - ns.permit = permit; - let suffixes: Vec = ns.watched.keys().cloned().collect(); - for suffix in suffixes { - let path = prefix.rejoin(&suffix); - self.sync_namespace(ns, &suffix, &path).await?; + None => { + if publisher.rescope(ns, &mut self.prefix, &mut self.announced, update.prefix)? { + self.retry_at = None; + self.retry_delay = RETRY_BASE; + } } } - NamespaceEvent::Renew(request_id) => self.send_update(ns, &prefix, request_id).await?, - NamespaceEvent::Answered(suffix, answer) => self.answered(ns, &prefix, suffix, answer).await?, - NamespaceEvent::Update(None) => { - // The origin is gone: withdraw everything, then finish the - // stream and wait for delivery. - self.withdraw_requests(&mut ns.target, &mut ns.requests).await; - self.abandon_verification(ns).await; - let Some(stream) = ns.target.stream() else { - return Ok(()); - }; - stream.writer.finish()?; - return stream.writer.closed().await; + } + NamespaceEvent::Verified(answer) => { + let Verifying { + verdict, prefix: next, .. + } = ns.verifying.take().expect("a token is verifying"); + let gate = ns.gate.as_mut().expect("update has a gate"); + if let Err(refusal) = gate.admit(&publisher.runtime, (verdict, answer)) { + let stream = ns.target.stream().expect("update has a stream"); + super::request_stream::buffer_error( + &mut stream.writer, + request::Kind::SubscribeNamespace, + &refusal.error, + &refusal.reason, + version, + )?; + self.state = ForwardState::Finish { finished: false }; + return Ok(None); } - NamespaceEvent::Update(Some((update, active, restart))) => { - // moq-transport has no restart: the peer sees the namespace end and start - // again, and resubscribes. - if restart { - self.apply_update(ns, &prefix, update.clone(), false).await?; - } - self.apply_update(ns, &prefix, update, active).await?; + ns.epoch = 0; + if publisher.rescope(ns, &mut self.prefix, &mut self.announced, next)? { + self.retry_at = None; + self.retry_delay = RETRY_BASE; } } - }; - - // This loop's advertisements die with it. - self.withdraw_requests(&mut ns.target, &mut ns.requests).await; + NamespaceEvent::Retry => { + self.retry_at = None; + self.retry_delay = (self.retry_delay * 2).min(RETRY_MAX); - res - } - - /// Drop an unanswered verification, keeping the request's stream credit until the - /// acceptor lets go of it. Cancelled, it leaves the credit in `ns` for the next call. - async fn abandon_verification(&self, ns: &mut Namespaces) { - if let Some(Verifying { verdict, credit, .. }) = ns.verifying.take() { - drop(verdict); - ns.abandoned = Some(credit); - } - if let Some(credit) = &ns.abandoned { - self.released(credit).await; - ns.abandoned = None; - } - } - - /// Answer an admitted namespace update: move to the prefix it names, if that is new, - /// then acknowledge it. Returns whether the prefix moved. - async fn rescope( - &self, - ns: &mut Namespaces, - prefix: &mut crate::Path, - announced: &mut crate::origin::Routes, - next: Option, - ) -> Result { - let moved = match next.filter(|next| next != prefix) { - Some(next) => { - // DONE still names the old prefix. After OK every suffix is relative to the - // new prefix, so replay the new cursor only after acknowledging. - let held: Vec<_> = ns - .watched - .iter() - .filter(|(_, watch)| watch.sent.wanted()) - .map(|(suffix, _)| suffix.clone()) - .collect(); - for suffix in held { - self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix) - .await?; - } - ns.watched.clear(); - *prefix = next; - let scope = crate::Pattern::subtree(prefix.as_str()) - .map(|p| p.rebase(prefix.as_str())) - .unwrap_or_default(); - let origin = self - .origin - .scope(&*prefix, &scope) - .unwrap_or_else(|_| self.origin.empty()); - *announced = self.advertising(origin, &ns.peer).discovery(ns.hidden).routes(); - true + // A minimum wait the peer named is enforced by `sync_namespace`, which every + // path goes through, so a namespace still inside one simply makes no offer this + // turn. The next sweep is at most RETRY_MAX away. + self.work.extend( + ns.watched + .iter() + .filter(|(_, watch)| watch.deferred) + .map(|(suffix, _)| NamespaceWork::Sync(suffix.clone())), + ); } - None => false, - }; - let stream = ns.target.stream().expect("update has a stream"); - stream - .writer - .encode_message(&ietf::RequestOk { - largest: None, - active: None, - }) - .await?; - Ok(moved) - } - - /// Reconcile what the peer holds with one route (un)announce. - async fn apply_update( - &self, - ns: &mut Namespaces, - prefix: &crate::Path, - update: crate::announce::Announce, - active: bool, - ) -> Result<(), Error> { - let suffix = update.prefix; - let path = prefix.rejoin(&suffix); - - if active { - // A repeat for a live suffix is a metadata update: keep the - // peer's refusal state and re-run the selection. - match ns.watched.get_mut(&suffix) { - Some(watch) => watch.route = update.route, - None => { - ns.watched.insert(suffix.clone(), Watched::new(update.route)); + NamespaceEvent::Regrant(permit) => { + ns.permit = permit; + self.work + .extend(ns.watched.keys().map(|suffix| NamespaceWork::Sync(suffix.clone()))); + } + NamespaceEvent::Renew(request_id) => publisher.send_update(ns, &self.prefix, request_id)?, + NamespaceEvent::Answered(suffix, answer) => { + let advertise = publisher.answered(ns, &self.prefix, suffix, answer)?; + self.advertise(publisher, advertise); + } + NamespaceEvent::Update(None) => { + // The origin is gone: withdraw everything, then finish the stream and wait for + // delivery. + publisher.withdraw_requests(&mut ns.target, &mut ns.requests); + self.state = ForwardState::Drain; + } + NamespaceEvent::Update(Some((update, active, restart))) => { + // moq-transport has no restart: the peer sees the namespace end and start + // again, and resubscribes. + if restart { + self.work.push_back(NamespaceWork::Apply(update.clone(), false)); } + self.work.push_back(NamespaceWork::Apply(update, active)); } - self.sync_namespace(ns, &suffix, &path).await - } else { - // Only close out namespaces the peer actually saw. - let held = ns.watched.remove(&suffix).is_some_and(|watch| watch.sent.wanted()); - if held { - tracing::debug!(route = %crate::path::Rooted(self.origin.root(), &path), "namespace_done"); - self.withdraw_namespace(&mut ns.target, &mut ns.requests, suffix) - .await?; - } - Ok(()) } + Ok(None) } } -/// The track's end, shared by a subscription's group writers so the one serving the end's -/// group can carry the END_OF_TRACK inline. -#[derive(Default)] -struct Terminal { - end: std::sync::OnceLock, - /// A group stream carrying the marker was acknowledged, so no standalone one is needed. - sent: AtomicBool, -} - -/// Whether the request token that admitted a request lapsed, shared by the streams the -/// request answers on. +/// Forward origin (un)announces to the peer until the loop ends. /// -/// A grant's end drops the request's serving from outside, so its streams learn why only -/// here: each one torn down after a lapse resets with EXPIRED_AUTH_TOKEN instead of its usual -/// code, telling the requester to renew its token rather than give up. A revoked or narrowed -/// grant leaves it unset, as does every other end, and so does an error that only says a -/// token expired elsewhere, such as an upstream's. -#[derive(Clone, Default)] -struct Lapse(Arc); - -impl Lapse { - /// Record the lapse, before the subscription's streams are dropped. - fn set(&self) { - self.0.store(true, Ordering::Relaxed); - } +/// Shared by both announce paths: they differ in where the advertisements go ([`Target`]) and +/// where the origin is rooted (`prefix`, empty when nothing asked for a subset). The token +/// that admitted the request, if one did, guards the loop. +struct NamespaceLoop { + ns: Namespaces, + token: Option, + /// `None` once the loop ended. + forward: Option, + /// How the loop ended, until any verification lets go of the request. + ended: Option>, +} - fn lapsed(&self) -> bool { - self.0.load(Ordering::Relaxed) +impl NamespaceLoop { + fn new( + publisher: &Publisher, + announced: crate::origin::Routes, + prefix: crate::Path, + ns: Namespaces, + initial: Vec, + ) -> Self { + Self { + token: ns.gate.as_ref().and_then(|gate| gate.token_guard()), + ns, + forward: Some(ForwardNamespaces::new(publisher, announced, prefix, initial)), + ended: None, + } } - /// Cancel a request whose answer is already out, with EXPIRED_AUTH_TOKEN once its token - /// lapsed, as no [`StreamError`](crate::StreamError) says, and otherwise as `err` does. - fn abort(&self, writer: Writer, err: &Error) { - match self.lapsed() { - true => writer.reset(super::error::EXPIRED_AUTH_TOKEN), - false => writer.abort(err), + fn poll(&mut self, waiter: &kio::Waiter, publisher: &Publisher) -> Poll> { + if let Some(forward) = &mut self.forward { + let ended = match &self.token { + Some(token) => token.poll_ended(waiter).map(|refusal| Err(refusal.error)), + None => Poll::Pending, + }; + let res = match ended { + Poll::Ready(res) => res, + Poll::Pending => ready!(forward.poll(waiter, publisher, &mut self.ns)), + }; + self.forward = None; + self.ended = Some(res); } + // Every exit, a failed write's or an ended token's included, keeps the request's + // stream, and so its credit, until the acceptor lets go of a token it is still + // verifying. + ready!(publisher.poll_abandon_verification(&mut self.ns, waiter)); + Poll::Ready(self.ended.take().expect("polled after it finished")) } } -struct SubscriptionFill { - fill: FillServe, - largest: Option, - request_id: RequestId, - track: track::Consumer, - ranking: Ranking>, - timescale: Option, - opened: Arc, +/// Advertise every namespace we can, without waiting to be asked. +/// +/// moq-transport itself says nothing about which of the two discovery messages a peer +/// expects, and the peers that never send SUBSCRIBE_NAMESPACE are exactly the ones +/// expecting a publisher to announce itself, so the default has to be to announce. A +/// peer that would rather ask says so with the MoQ Solicit extension +/// ([`solicit`](super::solicit)), +/// and then this loop does nothing. +/// [`Publisher::serve_subscribe_namespace_stream`] also carries every +/// match, so a peer that did not ask hears each namespace both ways. A NAMESPACE is +/// discovery only, not a second route. +pub(super) struct PublishNamespaces { + publisher: Publisher, + /// Started once the peer's SETUP says to advertise unasked. + run: Option, } -/// The current group's tail, even when its head was filled rather than arrived live. -struct FillTail { - start: Location, - fetch: track::Fetching, +impl PublishNamespaces { + pub(super) fn new(publisher: Publisher) -> Self { + Self { publisher, run: None } + } + + pub(super) fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + if self.run.is_none() { + if ready!(self.publisher.poll_requires_solicitation(waiter)) { + return Poll::Ready(Ok(())); + } + + // The cluster extension changes what an advertisement carries, so nothing can be + // sent until the peer's SETUP says whether it speaks it. + let peer = ready!(self.publisher.poll_peer(waiter)); + let hidden = ready!(self.publisher.peer_setup.poll_get(waiter)).hidden; + + // Split horizon, as the solicited loop applies it: never advertise a route back + // to the peer it came from. + let origin = self + .publisher + .advertising(self.publisher.origin.clone(), &peer) + .discovery(!hidden || self.publisher.origin.includes_hidden()); + + let ns = Namespaces::new(peer, Target::Requests); + let run = NamespaceLoop::new(&self.publisher, origin.routes(), crate::Path::empty(), ns, Vec::new()); + self.run = Some(run); + } + self.run.as_mut().expect("started").poll(waiter, &self.publisher) + } } -/// The stream serving a subscription's highest group, as it was opened. -struct LiveGroup { - sequence: u64, - slice: GroupSlice, - /// Set once a narrowing gave the stream up before its FIN. - retired: Arc, +/// Handle a SUBSCRIBE_NAMESPACE on its bidi stream. +/// +/// All the announce state is local to this task (mirroring `lite::Publisher`'s announce +/// handling): whatever this subscription advertised is withdrawn when its stream ends. Every +/// match is a NAMESPACE on this stream, whatever the peer's SETUP said. +struct NamespaceServe { + publisher: Publisher, + /// The request stream, until the loop's [`Target`] carries it. + stream: Option>, + prefix: crate::Path, + /// Whether the requester asked for hidden namespaces too. + wants_hidden: bool, + /// The request's gate, until the loop holds it. + gate: Option, + state: NamespaceServeState, } -/// Resolve a request's broadcast while the peer still wants the answer, `None` once the -/// requester withdraws or the session closes. -struct Resolve { - session: Transport, - fin_cancels: bool, - path: crate::Path, - state: ResolveState, +enum NamespaceServeState { + Start(Option), + Admit(Admit), + /// The extension changes what an advertisement carries, so nothing can be sent until the + /// peer's SETUP says whether it speaks it. The same SETUP says whether the OK counts what + /// it sends first (MoQ Active Count). + Setup, + /// With MoQ Auth, the count waits for our grant and leaves out what it does not cover. + Count { + peer: cluster::Peer, + hidden: bool, + announced: crate::origin::Routes, + }, + /// The acknowledgement, sent only while the token is live, even behind stream credit. + Ok { + peer: cluster::Peer, + hidden: bool, + announced: crate::origin::Routes, + initial: Vec, + /// Taken once buffered. + ok: Option, + }, + Run(Box), + Reject(super::request_stream::Deliver), + Done, } -#[kio::machine(self.state: ResolveState)] -impl Resolve { - fn new(session: Transport, path: crate::Path, fin_cancels: bool) -> Self { +impl NamespaceServe { + fn new( + publisher: Publisher, + stream: Stream, + msg: ietf::SubscribeNamespace, + gate: crate::auth::Gate, + ) -> Self { Self { - session, - path, - fin_cancels, - state: ResolveState::Start, + publisher, + stream: Some(stream), + prefix: msg.namespace, + wants_hidden: msg.hidden, + gate: Some(gate), + state: NamespaceServeState::Start(msg.authorization_token), } } - fn poll( - &mut self, - waiter: &kio::Waiter, - publisher: &Publisher, - stream: &mut Stream, - ) -> Poll, Error>> { - let origin = wait!(Origin => match publisher.poll_withdrawn(&mut self.session, stream, self.fin_cancels, waiter) { - Poll::Ready(()) => Poll::Ready(None), - Poll::Pending => publisher.poll_serving_origin(waiter).map(Some), - }); - let Some(origin) = origin else { - return Poll::Ready(Ok(None)); + /// MoQ Active Count: take what is advertised now, so the OK can say how many NAMESPACE + /// messages carry it. They go out first, ahead of any change. + fn count(&mut self, peer: cluster::Peer, hidden: bool, mut announced: crate::origin::Routes) { + let gate = self.gate.as_ref().expect("held until the loop starts"); + let (permit, _) = self.publisher.permit_now(Some(gate)); + let initial = Publisher::snapshot(&mut announced); + let count = initial + .iter() + .filter(|update| permit.matches(self.prefix.rejoin(&update.prefix).as_str())) + .filter(|update| self.publisher.select(&update.route, &peer).wanted()) + .count(); + self.state = NamespaceServeState::Ok { + peer, + hidden, + announced, + initial, + ok: Some(ietf::RequestOk { + largest: None, + active: Some(count as u64), + }), }; - // The immediate lookup: answering a peer never starts discovery of its own. - let requesting: kio::Pending = origin.lookup(self.path.clone(), None); - let found = wait!(Lookup => match publisher.poll_withdrawn(&mut self.session, stream, self.fin_cancels, waiter) { - Poll::Ready(()) => Poll::Ready(None), - Poll::Pending => requesting.poll_ok(waiter).map(Some), - }); - Poll::Ready(found.transpose()) + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + loop { + if let NamespaceServeState::Run(run) = &mut self.state { + return run.poll(waiter, &self.publisher); + } + let stream = self.stream.as_mut().expect("held until the loop starts"); + let gate = self.gate.as_mut().expect("held until the loop starts"); + match &mut self.state { + NamespaceServeState::Start(token) => { + if let Some(error) = super::namespace::refusal(&self.prefix, self.publisher.version) { + stream.writer.buffer_message(ietf::RequestError::ID, &error)?; + self.state = NamespaceServeState::Reject(super::request_stream::Deliver::new()); + continue; + } + let discovery = self.gate.take().expect("held until the loop starts").discovery(); + let fin_cancels = super::request_stream::fin_cancels(self.publisher.version); + let admit = Admit::new(&self.publisher, &discovery, token.take(), fin_cancels); + self.gate = Some(discovery); + self.state = NamespaceServeState::Admit(admit); + } + NamespaceServeState::Admit(admit) => { + match ready!(admit.poll(waiter, &self.publisher, stream, gate)) { + Ok(true) => {} + Ok(false) => return Poll::Ready(Ok(())), + Err(refusal) => { + super::request_stream::buffer_error( + &mut stream.writer, + request::Kind::SubscribeNamespace, + &refusal.error, + &refusal.reason, + self.publisher.version, + )?; + self.state = NamespaceServeState::Reject(super::request_stream::Deliver::new()); + continue; + } + } + tracing::debug!(prefix = %crate::path::Rooted(self.publisher.origin.root(), &self.prefix), "subscribe_namespace stream"); + self.state = NamespaceServeState::Setup; + } + NamespaceServeState::Setup => { + let peer = ready!(self.publisher.poll_peer(waiter)); + let declared = ready!(self.publisher.peer_setup.poll_get(waiter)); + let publisher = &self.publisher; + let prefix = &self.prefix; + + // A prefix outside our scope (empty origin, or a token that doesn't grant it) + // just means we have nothing to announce; respond with an empty set rather + // than erroring, which would look fatal to the peer. + // The wire prefix decodes as a literal path; convert it explicitly to its + // subtree grant, refusing anything that cannot be a subtree. + // The cursor is rooted at the prefix so every update arrives named as its + // wire suffix, a route covering the prefix included: that one presents at + // the root, as the empty suffix. + let scope = crate::Pattern::subtree(prefix.as_str()) + .map(|subtree| subtree.rebase(prefix.as_str())) + .unwrap_or_default(); + let origin = publisher + .origin + .scope(prefix, &scope) + .unwrap_or_else(|_| publisher.origin.empty()); + + // Register the split-horizon peer on the announce cursor too. The origin + // model uses this exposure to park a reflected copy before it can replace + // the source we are currently advertising to that peer. + let origin = publisher.advertising(origin, &peer); + + // Peers that declared MoQ Hidden filter namespaces unless they opted in. A + // publish origin that already opted in (the caller's choice for this peer) + // keeps them. + let hidden = !declared.hidden || self.wants_hidden || publisher.origin.includes_hidden(); + let origin = origin.discovery(hidden); + + // A NAMESPACE is discovery, not a route, so a peer that also hears the + // unsolicited PUBLISH_NAMESPACE still gets the match here. + let announced = origin.routes(); + + match declared.active_count { + true if declared.auth => { + self.state = NamespaceServeState::Count { + peer, + hidden, + announced, + } + } + true => self.count(peer, hidden, announced), + false => { + self.state = NamespaceServeState::Ok { + peer, + hidden, + announced, + initial: Vec::new(), + ok: Some(ietf::RequestOk { + largest: None, + active: None, + }), + } + } + } + } + NamespaceServeState::Count { .. } => { + ready!(gate.poll_setup_answered(waiter)); + let NamespaceServeState::Count { + peer, + hidden, + announced, + } = std::mem::replace(&mut self.state, NamespaceServeState::Done) + else { + unreachable!("counting"); + }; + self.count(peer, hidden, announced); + } + NamespaceServeState::Ok { ok, .. } => { + if let Poll::Ready(refusal) = gate.poll_token_ended(waiter) { + return Poll::Ready(Err(refusal.error)); + } + if let Some(ok) = ok.take() { + stream.writer.buffer_message(ietf::RequestOk::ID, &ok)?; + } + ready!(stream.writer.poll_flush(waiter))?; + let NamespaceServeState::Ok { + peer, + hidden, + announced, + initial, + .. + } = std::mem::replace(&mut self.state, NamespaceServeState::Done) + else { + unreachable!("acknowledging"); + }; + let target = Target::Inline(self.stream.take().expect("acknowledged once")); + let mut ns = Namespaces::new(peer, target); + ns.gate = self.gate.take(); + ns.hidden = hidden; + let prefix = self.prefix.clone(); + let run = NamespaceLoop::new(&self.publisher, announced, prefix, ns, initial); + self.state = NamespaceServeState::Run(Box::new(run)); + } + NamespaceServeState::Run(_) => unreachable!("run above"), + NamespaceServeState::Reject(deliver) => { + let res = ready!(deliver.poll(waiter, &mut stream.writer)); + self.state = NamespaceServeState::Done; + return Poll::Ready(res); + } + NamespaceServeState::Done => panic!("polled after it finished"), + } + } } } @@ -5382,6 +6926,296 @@ impl crate::transport::poll::SendStream f } } +/// Write one group's frames in the negotiated draft's FETCH object layout. +/// +/// The first object carries its absolute Group and Object IDs plus the priority; +/// every later one inherits them and increments the Object ID, so only the +/// properties (the timestamp) and the payload go on the wire. A fetch object has no +/// status field. +struct FetchGroupWrite { + group: group::Consumer, + sequence: u64, + slice: GroupSlice, + timescale: Option, + version: Version, + priority: u8, + /// The Object ID of the frame being read or written. + index: u64, + /// The Object ID written last. + prior: Option, + buf: arrayvec::ArrayVec, + /// A long cached group is written a slice per poll, so the session's other tasks still run. + budget: kio::coop::Budget, + state: FetchGroupState, +} + +enum FetchGroupState { + /// Yield to the budget before reading on. + Yield, + /// Read whole cached objects, or else the next frame as it arrives. + Read, + /// Write the batch in `buf` from `next`, one object's payload at a time. + Batch { + next: usize, + payload: Option, + }, + /// Drain a frame before the slice, which still advances the cursor. + Skip(frame::Consumer), + /// Stream a frame whose header and size are buffered, chunk by chunk. + Partial { + frame: frame::Consumer, + chunk: Option, + }, +} + +impl FetchGroupWrite { + fn new( + group: group::Consumer, + sequence: u64, + slice: GroupSlice, + timescale: Option, + version: Version, + priority: u8, + ) -> Result { + if group.wire().is_some_and(|wire| !wire.fetchable) { + return Err(Error::Unsupported); + } + Ok(Self { + group, + sequence, + index: 0, + slice, + timescale, + version, + priority, + prior: None, + buf: arrayvec::ArrayVec::new(), + budget: kio::coop::Budget::new(32), + state: FetchGroupState::Yield, + }) + } + + /// The cap is the Largest Object snapshot: the group may keep growing, but everything + /// past the snapshot belongs to the subscription, not the fill. + fn capped(slice: GroupSlice, index: u64) -> bool { + slice.until.is_some_and(|until| index >= until) + } + + /// The result once writing stops: a slice that ends before its cap lost objects. + fn finish(&self) -> Result<(), Error> { + match self.slice.until.is_some_and(|until| self.index < until) { + true => Err(Error::NotFound), + false => Ok(()), + } + } + + fn object<'a>(&self, frame: &'a frame::Object) -> Result, Error> { + Ok(FetchObjectInfo { + sequence: self.sequence, + object: self.index, + prior: fetch_prior(self.prior, self.index)?, + priority: self.priority, + timestamp: frame.timestamp, + inflated: frame.inflated, + timescale: self.timescale, + version: self.version, + fields: frame.fields.as_deref(), + source: self.group.ietf(), + }) + } + + fn poll( + &mut self, + waiter: &kio::Waiter, + stream: &mut Writer, Version>, + ) -> Poll> { + loop { + match &mut self.state { + FetchGroupState::Yield => { + ready!(self.budget.poll_yield(waiter)); + if Self::capped(self.slice, self.index) { + return Poll::Ready(self.finish()); + } + self.state = FetchGroupState::Read; + } + FetchGroupState::Read => { + if stream.poll_closed(waiter).is_ready() { + return Poll::Ready(Err(Error::Cancel)); + } + let step = match self.group.poll_read_objects(waiter, &mut self.buf) { + Poll::Pending => match ready!(self.group.poll_next_frame(waiter))? { + Some(frame) => FillStep::Partial(frame), + None => FillStep::Done, + }, + Poll::Ready(Ok(0)) => FillStep::Done, + Poll::Ready(Ok(_)) => FillStep::Batch, + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + }; + match step { + FillStep::Batch => self.state = FetchGroupState::Batch { next: 0, payload: None }, + FillStep::Partial(frame) => { + self.index = frame.fields.as_ref().map_or(self.index, |fields| fields.id); + if Self::capped(self.slice, self.index) { + return Poll::Ready(self.finish()); + } + if self.index < self.slice.skip { + self.state = FetchGroupState::Skip(frame); + continue; + } + let header = FetchObjectInfo { + sequence: self.sequence, + object: self.index, + prior: fetch_prior(self.prior, self.index)?, + priority: self.priority, + timestamp: frame.timestamp, + inflated: frame.inflated, + timescale: self.timescale, + version: self.version, + fields: frame.fields.as_deref(), + source: self.group.ietf(), + }; + Publisher::buffer_fetch_object(stream, header)?; + stream.buffer_varint(frame.size)?; + self.prior = Some(self.index); + self.index += 1; + self.state = FetchGroupState::Partial { frame, chunk: None }; + } + FillStep::Done => return Poll::Ready(self.finish()), + } + } + FetchGroupState::Batch { next, payload } => { + if let Some(write) = payload { + ready!(write.poll(waiter, stream))?; + *payload = None; + self.index += 1; + self.group.keep_alive(); + continue; + } + let Some(frame) = self.buf.get(*next) else { + self.state = FetchGroupState::Yield; + continue; + }; + *next += 1; + if Self::capped(self.slice, self.index) { + return Poll::Ready(self.finish()); + } + let frame = frame.clone(); + self.index = frame.fields.as_ref().map_or(self.index, |fields| fields.id); + if Self::capped(self.slice, self.index) { + return Poll::Ready(self.finish()); + } + if self.index < self.slice.skip { + self.index += 1; + self.group.keep_alive(); + continue; + } + let header = self.object(&frame)?; + Publisher::buffer_fetch_object(stream, header)?; + self.prior = Some(self.index); + let FetchGroupState::Batch { payload, .. } = &mut self.state else { + unreachable!("writing a batch"); + }; + *payload = Some(FetchPayloadWrite::new(frame.frame.payload)); + } + FetchGroupState::Skip(frame) => { + if stream.poll_closed(waiter).is_ready() { + return Poll::Ready(Err(Error::Cancel)); + } + if ready!(frame.poll_read_chunk(waiter))?.is_none() { + self.index += 1; + self.state = FetchGroupState::Yield; + } + } + FetchGroupState::Partial { frame, chunk } => { + ready!(stream.poll_flush(waiter))?; + if let Some(bytes) = chunk { + ready!(stream.poll_write_all(waiter, bytes))?; + *chunk = None; + } + if stream.poll_closed(waiter).is_ready() { + return Poll::Ready(Err(Error::Cancel)); + } + match ready!(frame.poll_read_chunk(waiter))? { + Some(bytes) => *chunk = Some(bytes), + None => self.state = FetchGroupState::Yield, + } + } + } + } + } +} + +/// Write one group's objects to a fetch stream, which carries no other group. +struct FetchedGroupWrite { + frames: std::vec::IntoIter, + /// The Object ID of the next frame without its own. + ordinal: u64, + sequence: u64, + source: Option, + priority: u8, + version: Version, + timescale: Option, + prior: Option, + budget: kio::coop::Budget, + payload: Option, +} + +impl FetchedGroupWrite { + fn new(group: FetchedGroup, version: Version, timescale: Option) -> Self { + Self { + source: group.info.ietf(), + priority: super::priority::to_wire(group.info.priority), + frames: group.frames.into_iter(), + ordinal: group.first, + sequence: group.sequence, + version, + timescale, + prior: None, + budget: kio::coop::Budget::new(32), + payload: None, + } + } + + fn poll( + &mut self, + waiter: &kio::Waiter, + writer: &mut Writer, Version>, + ) -> Poll> { + loop { + if let Some(write) = &mut self.payload { + ready!(write.poll(waiter, writer))?; + self.payload = None; + } + if self.frames.as_slice().is_empty() { + return Poll::Ready(Ok(())); + } + ready!(self.budget.poll_yield(waiter)); + let frame = self.frames.next().expect("checked non-empty"); + let ordinal = self.ordinal; + self.ordinal += 1; + let object = frame.fields.as_ref().map_or(ordinal, |fields| fields.id); + let position = fetch_prior(self.prior, object)?; + self.prior = Some(object); + Publisher::buffer_fetch_object( + writer, + FetchObjectInfo { + sequence: self.sequence, + object, + prior: position, + priority: self.priority, + timestamp: frame.timestamp, + inflated: frame.inflated, + timescale: self.timescale, + version: self.version, + fields: frame.fields.as_deref(), + source: self.source, + }, + )?; + self.payload = Some(FetchPayloadWrite::new(frame.frame.payload)); + } + } +} + /// Write a whole fetch object's length and payload, after its header. struct FetchPayloadWrite { payload: crate::frame::Payload, @@ -5459,16 +7293,6 @@ impl EndOfTrack { } } -/// Finish a ranked stream and wait for the peer to acknowledge it, holding its rank until then. -/// -/// Unlike [`Writer::close`], the writer keeps the stream while it waits, so cancelling resets -/// the stream as a cancelled group stream's is, rather than leaving it sending once its rank -/// is released. Once acknowledged, the stream is released and dropping the writer resets nothing. -async fn settle(writer: &mut Writer, Version>) -> Result<(), Error> { - writer.finish()?; - kio::wait(|waiter| writer.poll_close(waiter)).await -} - /// A group stream's place in the session's send order: its subscription's priority, then /// its sequence in the subscription's GROUP_ORDER. enum Rank { @@ -7215,14 +9039,7 @@ mod serve_tests { let mut updates = subscription_updates(); let served = h .publisher - .run_subscription( - &mut stream, - &mut serving, - &mut false, - &mut goaways, - async {}, - &mut updates, - ) + .run_subscription(&mut stream, &mut serving, &mut false, &mut goaways, &mut updates) .await; assert!(matches!(served, Some(Served::UpdateFailed(_))), "{served:?}"); assert!(h.log.closes().is_empty()); @@ -7257,14 +9074,9 @@ mod serve_tests { super::super::request_stream::Messages::new(crate::driver::Transport::new(h.session.clone()), false); { let mut updates = subscription_updates(); - let run = h.publisher.run_subscription( - &mut stream, - &mut serving, - &mut finished, - &mut goaways, - async {}, - &mut updates, - ); + let run = + h.publisher + .run_subscription(&mut stream, &mut serving, &mut finished, &mut goaways, &mut updates); let mut run = std::pin::pin!(run); assert!(futures::poll!(run.as_mut()).is_pending()); h.track.finish().unwrap(); @@ -7386,7 +9198,6 @@ mod serve_tests { &mut serving, &mut finished, &mut goaways, - async {}, &mut updates, )); assert_eq!( @@ -7427,14 +9238,9 @@ mod serve_tests { ); let mut finished = false; let mut updates = subscription_updates(); - let run = h.publisher.run_subscription( - &mut stream, - &mut serving, - &mut finished, - &mut goaways, - async {}, - &mut updates, - ); + let run = + h.publisher + .run_subscription(&mut stream, &mut serving, &mut finished, &mut goaways, &mut updates); assert!(futures::poll!(std::pin::pin!(run)).is_ready(), "{version}"); } } @@ -7483,7 +9289,7 @@ mod serve_tests { let stream = Stream::open(&crate::driver::Transport::new(h.session.clone()), version) .await .unwrap(); - let task = h.publisher.handle_stream(id, ietf::Body(body.clone()), stream).unwrap(); + let task = kio::wait(h.publisher.handle_stream(id, ietf::Body(body.clone()), stream).unwrap()); assert_eq!(h.publisher.owed.load(Ordering::Relaxed), 1, "count before polling"); drop(task); assert_eq!(h.publisher.owed.load(Ordering::Relaxed), 0, "release a cancelled task"); @@ -7502,7 +9308,7 @@ mod serve_tests { let stream = Stream::open(&crate::driver::Transport::new(h.session.clone()), version) .await .unwrap(); - h.publisher.handle_stream(id, ietf::Body(body), stream).unwrap().await; + kio::wait(h.publisher.handle_stream(id, ietf::Body(body), stream).unwrap()).await; assert_eq!(h.publisher.owed.load(Ordering::Relaxed), 0, "release a completed fetch"); } } @@ -7536,7 +9342,6 @@ mod serve_tests { &mut serving, &mut finished, &mut goaways, - async {}, &mut updates, )); assert!(futures::poll!(run.as_mut()).is_pending()); @@ -7578,7 +9383,6 @@ mod serve_tests { &mut serving, &mut finished, &mut goaways, - async {}, &mut updates, )); @@ -7935,14 +9739,7 @@ mod serve_tests { let mut updates = subscription_updates(); let served = h .publisher - .run_subscription( - &mut stream, - &mut serving, - &mut finished, - &mut goaways, - async {}, - &mut updates, - ) + .run_subscription(&mut stream, &mut serving, &mut finished, &mut goaways, &mut updates) .await .expect("the track's error ends the subscription"); assert!(matches!(served, Served::Track(Err(Error::Unsupported))), "{served:?}"); @@ -8348,7 +10145,6 @@ mod serve_tests { &mut serving, &mut finished, &mut goaways, - async {}, &mut updates )); assert!(futures::poll!(run.as_mut()).is_pending(), "{version}: update ended it"); @@ -8385,14 +10181,7 @@ mod serve_tests { let mut updates = subscription_updates(); let served = h .publisher - .run_subscription( - &mut stream, - &mut serving, - &mut finished, - &mut goaways, - async {}, - &mut updates, - ) + .run_subscription(&mut stream, &mut serving, &mut finished, &mut goaways, &mut updates) .await; assert!( matches!(served, Some(Served::UpdateFailed(Error::Unsupported))), @@ -8456,14 +10245,7 @@ mod serve_tests { let mut updates = subscription_updates(); let served = h .publisher - .run_subscription( - &mut stream, - &mut serving, - &mut finished, - &mut goaways, - async {}, - &mut updates, - ) + .run_subscription(&mut stream, &mut serving, &mut finished, &mut goaways, &mut updates) .await; assert!( matches!(served, Some(Served::Track(Err(_)) | Served::UpdateFailed(_))), @@ -8567,7 +10349,6 @@ mod serve_tests { &mut serving, &mut finished, &mut goaways, - async {}, &mut updates )); let mut served = None; @@ -8615,14 +10396,7 @@ mod serve_tests { let mut updates = subscription_updates(); let result = h .publisher - .run_subscription( - &mut stream, - &mut serving, - &mut false, - &mut goaways, - async {}, - &mut updates, - ) + .run_subscription(&mut stream, &mut serving, &mut false, &mut goaways, &mut updates) .await; assert!( matches!(result, Some(Served::Track(Err(Error::Decode(_))))), @@ -8662,14 +10436,7 @@ mod serve_tests { let mut updates = subscription_updates(); let result = h .publisher - .run_subscription( - &mut stream, - &mut serving, - &mut false, - &mut goaways, - async {}, - &mut updates, - ) + .run_subscription(&mut stream, &mut serving, &mut false, &mut goaways, &mut updates) .await; assert!( matches!(result, Some(Served::Track(Err(Error::Decode(_))))), @@ -9343,11 +11110,12 @@ mod serve_tests { .unwrap(); let mark = h.log.trail().len(); let body = track_status("room", true, version); - let task = h - .publisher - .clone() - .handle_stream(ietf::TrackStatus::ID, ietf::Body(bytes::Bytes::from(body)), stream) - .unwrap(); + let task = kio::wait( + h.publisher + .clone() + .handle_stream(ietf::TrackStatus::ID, ietf::Body(bytes::Bytes::from(body)), stream) + .unwrap(), + ); kio::clock::timeout(std::time::Duration::from_secs(1), task) .await .unwrap_or_else(|_| panic!("{version}: never answered")); @@ -9390,11 +11158,13 @@ mod serve_tests { .unwrap(); let mark = h.log.writes.lock().unwrap().len(); let body = track_status("room", wanted, version); - h.publisher - .clone() - .handle_stream(ietf::TrackStatus::ID, ietf::Body(bytes::Bytes::from(body)), stream) - .unwrap() - .await; + kio::wait( + h.publisher + .clone() + .handle_stream(ietf::TrackStatus::ID, ietf::Body(bytes::Bytes::from(body)), stream) + .unwrap(), + ) + .await; assert!(h.log.resets().is_empty(), "{version}: the answer was reset"); assert!(h.track.subscription().is_none(), "{version}: TRACK_STATUS subscribed"); @@ -9469,12 +11239,12 @@ mod serve_tests { .unwrap(); let mark = h.log.writes.lock().unwrap().len(); let body = track_status("live/cam", true, version); - let mut answer = std::pin::pin!( + let mut answer = std::pin::pin!(kio::wait( h.publisher .clone() .handle_stream(ietf::TrackStatus::ID, ietf::Body(bytes::Bytes::from(body)), stream,) .unwrap() - ); + )); assert!(futures::poll!(answer.as_mut()).is_pending(), "{version}"); // Parked resolving the broadcast when the grant narrows. @@ -9597,11 +11367,13 @@ mod serve_tests { .await .unwrap(); let mark = h.log.writes.lock().unwrap().len(); - h.publisher - .clone() - .handle_stream(ietf::TrackStatus::ID, ietf::Body(bytes::Bytes::from(body)), stream) - .unwrap() - .await; + kio::wait( + h.publisher + .clone() + .handle_stream(ietf::TrackStatus::ID, ietf::Body(bytes::Bytes::from(body)), stream) + .unwrap(), + ) + .await; refusal_code(&h.log.writes.lock().unwrap()[mark..], version) } @@ -9670,11 +11442,13 @@ mod serve_tests { .await .unwrap(); let mark = h.log.writes.lock().unwrap().len(); - h.publisher - .clone() - .handle_stream(id, ietf::Body(bytes::Bytes::from(body)), stream) - .unwrap_or_else(|e| panic!("{version}: the request closed the session: {e}")) - .await; + kio::wait( + h.publisher + .clone() + .handle_stream(id, ietf::Body(bytes::Bytes::from(body)), stream) + .unwrap_or_else(|e| panic!("{version}: the request closed the session: {e}")), + ) + .await; assert!(h.log.resets().is_empty(), "{version}: refusal was reset"); let wire = h.log.writes.lock().unwrap()[mark..].to_vec(); @@ -9783,12 +11557,12 @@ mod serve_tests { .unwrap(); let mark = h.log.writes.lock().unwrap().len(); let body = track_status(path, true, version); - let mut serving = std::pin::pin!( + let mut serving = std::pin::pin!(kio::wait( h.publisher .clone() .handle_stream(ietf::TrackStatus::ID, ietf::Body(body.into()), stream) .unwrap() - ); + )); let done = futures::poll!(serving.as_mut()); assert!( claim.requested_broadcast().now_or_never().is_none(), @@ -9905,15 +11679,16 @@ mod serve_tests { .await .unwrap(); let mark = h.log.writes.lock().unwrap().len(); - let task = h - .publisher - .clone() - .handle_stream( - ietf::SubscribeNamespace::ID, - ietf::Body(bytes::Bytes::from(body)), - stream, - ) - .unwrap_or_else(|e| panic!("{version}: the request closed the session: {e}")); + let task = kio::wait( + h.publisher + .clone() + .handle_stream( + ietf::SubscribeNamespace::ID, + ietf::Body(bytes::Bytes::from(body)), + stream, + ) + .unwrap_or_else(|e| panic!("{version}: the request closed the session: {e}")), + ); let mut task = std::pin::pin!(task); let mut done = false; for _ in 0..100 { @@ -10822,10 +12597,11 @@ mod serve_tests { subscribe(Filter::NextObject, None) .encode_msg(&mut Encoder::new(&mut data, version.into()), version) .unwrap(); - let mut serving = h - .publisher - .handle_stream(ietf::Subscribe::ID, ietf::Body(bytes::Bytes::from(data)), stream) - .unwrap(); + let mut serving = kio::wait( + h.publisher + .handle_stream(ietf::Subscribe::ID, ietf::Body(bytes::Bytes::from(data)), stream) + .unwrap(), + ); let mut fetching = Box::pin(joining_fetch(&h, 0)); assert!( futures::poll!(fetching.as_mut()).is_pending(), @@ -10859,10 +12635,11 @@ mod serve_tests { subscribe(Filter::NextObject, None) .encode_msg(&mut Encoder::new(&mut data, version.into()), version) .unwrap(); - let mut serving = h - .publisher - .handle_stream(ietf::Subscribe::ID, ietf::Body(bytes::Bytes::from(data)), stream) - .unwrap(); + let mut serving = kio::wait( + h.publisher + .handle_stream(ietf::Subscribe::ID, ietf::Body(bytes::Bytes::from(data)), stream) + .unwrap(), + ); match futures::future::select(&mut fetching, &mut serving).await { futures::future::Either::Left((response, _)) => { response.unwrap(); @@ -10884,10 +12661,11 @@ mod serve_tests { subscribe(Filter::NextObject, None) .encode_msg(&mut Encoder::new(&mut data, version.into()), version) .unwrap(); - let serving = h - .publisher - .handle_stream(ietf::Subscribe::ID, ietf::Body(bytes::Bytes::from(data)), stream) - .unwrap(); + let serving = kio::wait( + h.publisher + .handle_stream(ietf::Subscribe::ID, ietf::Body(bytes::Bytes::from(data)), stream) + .unwrap(), + ); let mut fetching = Box::pin(joining_fetch(&h, 0)); assert!(futures::poll!(fetching.as_mut()).is_pending()); drop(serving); @@ -11944,11 +13722,12 @@ mod serve_tests { .await .unwrap(); let mark = h.log.writes.lock().unwrap().len(); - let serving = h - .publisher - .clone() - .handle_stream(ietf::Fetch::ID, ietf::Body(body.into()), stream) - .unwrap(); + let serving = kio::wait( + h.publisher + .clone() + .handle_stream(ietf::Fetch::ID, ietf::Body(body.into()), stream) + .unwrap(), + ); let ((), issued) = futures::join!(serving, async { requests.next().await.unwrap().accept(crate::auth::Grant::all()) }); @@ -12167,15 +13946,17 @@ mod serve_tests { let stream = Stream::open(&crate::driver::Transport::new(h.session.clone()), version) .await .unwrap(); - h.publisher - .clone() - .handle_stream( - ietf::TrackStatus::ID, - ietf::Body(track_status("room", true, version).into()), - stream, - ) - .unwrap() - .await; + kio::wait( + h.publisher + .clone() + .handle_stream( + ietf::TrackStatus::ID, + ietf::Body(track_status("room", true, version).into()), + stream, + ) + .unwrap(), + ) + .await; let wire = h.log.writes.lock().unwrap().clone(); let mut buf = Decoder::new(&wire, version.into()); assert_eq!(buf.varint().unwrap(), ietf::TrackStatusOk::ID); @@ -14178,7 +15959,7 @@ mod serve_tests { } /// The REQUEST_ERROR a refused renewal writes before the subscription ends, matching - /// `write_subscribe_error(.., &Error::Unauthorized, reason)`. + /// `buffer_subscribe_error(.., &Error::Unauthorized, reason)`. async fn request_error_bytes(version: Version, reason: &str) -> Vec { let log = crate::lite::test_transport::Log::default(); let mut writer = crate::coding::Writer::new( @@ -16644,7 +18425,7 @@ mod tests { let mut ns = Namespaces::new(cluster::Peer::default(), Target::Requests); ns.watched.insert(suffix.clone(), watch); - publisher.sync_namespace(&mut ns, &suffix, &suffix).await.unwrap(); + publisher.sync_namespace(&mut ns, &suffix, &suffix).unwrap(); assert!( !ns.watched[&suffix].deferred, diff --git a/rs/moq-net/src/ietf/request_stream.rs b/rs/moq-net/src/ietf/request_stream.rs index 53c7038e106..c6bc9c98950 100644 --- a/rs/moq-net/src/ietf/request_stream.rs +++ b/rs/moq-net/src/ietf/request_stream.rs @@ -170,7 +170,8 @@ impl Messages { self.goaways.receive(redirect).map_err(|err| self.failed(err)) } - /// Read complete frames, consuming at most one GOAWAY without ending the request. + /// The next complete frame, consuming at most one GOAWAY without ending the request. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next( &mut self, reader: &mut crate::coding::Reader, @@ -179,20 +180,17 @@ impl Messages { self.goaways.poll_next(reader, waiter).map_err(|err| self.failed(err)) } - pub async fn next( + /// The next complete frame, where the request stream ending first breaks the protocol. + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_message( &mut self, reader: &mut crate::coding::Reader, - ) -> Result, Error> { - kio::wait(|waiter| self.poll_next(reader, waiter)).await - } - - pub async fn message( - &mut self, - reader: &mut crate::coding::Reader, - ) -> Result { - self.next(reader) - .await? - .ok_or_else(|| self.failed(Error::ProtocolViolation)) + waiter: &kio::Waiter, + ) -> Poll> { + match ready!(self.poll_next(reader, waiter))? { + Some(framed) => Poll::Ready(Ok(framed)), + None => Poll::Ready(Err(self.failed(Error::ProtocolViolation))), + } } pub fn poll_end( @@ -207,13 +205,82 @@ impl Messages { } } -/// Refuse a request with its complete REQUEST_ERROR, then finish the response. -pub(super) async fn reject( - mut stream: Stream, - error: super::RequestError<'_>, -) -> Result<(), Error> { - stream.writer.buffer_message(super::RequestError::ID, &error)?; - deliver(&mut stream.writer).await +/// End a request stream with the refusal buffered on it: [`Deliver`] owning the stream. +pub(super) struct Reject { + stream: Stream, + /// Why the refusal could not be buffered, which the first poll reports. + failed: Option, + deliver: Deliver, +} + +impl Reject { + /// Refuse a request with its complete REQUEST_ERROR. + pub(super) fn new(mut stream: Stream, error: super::RequestError<'_>) -> Self { + let buffered = stream.writer.buffer_message(super::RequestError::ID, &error); + Self::buffered(stream, buffered) + } + + /// Deliver the refusal `buffered` put on `stream`, or fail with why it could not. + pub(super) fn buffered(stream: Stream, buffered: Result<(), Error>) -> Self { + Self { + stream, + failed: buffered.err(), + deliver: Deliver::new(), + } + } + + /// Ready once the refusal is delivered; see [`Deliver`]. + pub(super) fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + if let Some(err) = self.failed.take() { + return Poll::Ready(Err(err)); + } + self.deliver.poll(waiter, &mut self.stream.writer) + } +} + +/// Abandon a request we sent: STOP_SENDING, then our FIN and its acknowledgement. +/// +/// There is no message for it, only the stream itself: a FIN is explicitly not a +/// cancellation (draft-22 section 6.4.2.2), so section 6.4.2.3's pair applies, an endpoint +/// that has already FINed its sending direction cancels with STOP_SENDING on the receiving +/// one. Finishing alone would leave the writer's Drop free to RESET_STREAM, discarding what +/// was already written, so the writer is closed instead; a request the peer's send window +/// held back resets. +pub(super) struct Cancel { + /// Held until the close settles, as the request stream it belongs to. + _reader: crate::coding::Reader, + closing: Option>, +} + +impl Cancel { + pub(super) fn new(stream: Stream) -> Self { + let Stream { writer, mut reader } = stream; + // STOP_SENDING needs no acknowledgement, so it goes first and the wait covers only + // what we still have to deliver. + reader.abort(&Error::Cancel); + let closing = match writer.closing() { + Ok(closing) => Some(closing), + Err(err) => { + tracing::debug!(%err, "failed to close the subscribe stream"); + None + } + }; + Self { + _reader: reader, + closing, + } + } + + /// Ready once the peer acknowledges our FIN, or the stream dies. + pub(super) fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + if let Some(closing) = &mut self.closing { + if let Err(err) = ready!(closing.poll(waiter)) { + tracing::debug!(%err, "failed to close the subscribe stream"); + } + self.closing = None; + } + Poll::Ready(()) + } } /// Deliver the last message buffered on a response: flush it, then finish the stream and @@ -248,12 +315,6 @@ impl Deliver { } } -/// Await [`Deliver`]. -pub(super) async fn deliver(writer: &mut Writer) -> Result<(), Error> { - let mut deliver = Deliver::new(); - kio::wait(|waiter| deliver.poll(waiter, writer)).await -} - /// Buffer the REQUEST_ERROR refusing a `kind` request with `err`. pub(super) fn buffer_error( writer: &mut Writer, @@ -272,18 +333,6 @@ pub(super) fn buffer_error( ) } -/// Write the REQUEST_ERROR refusing a `kind` request with `err`. -pub(super) async fn write_error( - writer: &mut Writer, - kind: super::error::request::Kind, - err: &Error, - reason: &str, - version: Version, -) -> Result<(), Error> { - buffer_error(writer, kind, err, reason, version)?; - kio::wait(|waiter| writer.poll_flush(waiter)).await -} - /// A framed subscription update, preserving omitted preferences. #[derive(Debug)] pub(crate) struct Update { diff --git a/rs/moq-net/src/ietf/session.rs b/rs/moq-net/src/ietf/session.rs index 9f0b120df75..ec105d59d5e 100644 --- a/rs/moq-net/src/ietf/session.rs +++ b/rs/moq-net/src/ietf/session.rs @@ -6,31 +6,18 @@ use crate::{ coding::{Decode, DecodeError, Encode, Reader, Stream, Writer}, ietf::{self, FetchHeader, RequestId}, setup, - util::{BoxFuture, TaskSet, err_only}, + util::TaskSet, }; -use futures::FutureExt as _; +use std::task::{Poll, ready}; use super::{ - Control, Message, Publisher, Subscriber, Version, active_count, auth, cluster, hidden, max_age, peer, - request_update, solicit, subscriber::is_protocol_violation, + Control, Message, Namespaces, PublishNamespaces, Publisher, RecvFill, RecvGroup, Subscriber, Version, active_count, + auth, cluster, hidden, max_age, peer, request_stream::Reject, request_update, solicit, + subscriber::is_protocol_violation, }; -/// Keep the session's SUBSCRIBE_NAMESPACE requests open as `discovery` decides, each on -/// a stream from `open`; with no subscribe half, ask nothing. Only an error ends it. -async fn run_namespaces( - subscriber: Subscriber, - discovery: Option, - open: F, -) -> Result<(), Error> -where - F: Fn() -> O, - O: std::future::Future, Error>>, -{ - match discovery { - Some(discovery) => subscriber.run_namespaces(discovery, open).await, - None => std::future::pending().await, - } -} +type RecvStream = crate::transport::boxed::RecvStream; +type SendStream = crate::transport::boxed::SendStream; /// Everything one moq-transport session needs to start. pub struct Config { @@ -82,9 +69,9 @@ pub struct Config { pub token: Option, /// The peer's SETUP stream, when it was already read before [`start`] (a server - /// that gated on the client's path via [`accept_setup`]). It becomes the + /// that gated on the client's path via [`SetupAccept`]). It becomes the /// GOAWAY channel; `None` lets the uni loop read the SETUP itself. - pub peer_setup_stream: Option>, + pub peer_setup_stream: Option>, /// What that pre-read SETUP declared, so the session does not have to parse it /// twice. `None` when [`Self::peer_setup_stream`] is. @@ -98,7 +85,7 @@ pub struct Config { /// Uni streams that arrived before that pre-read SETUP (see [`PeerSetup::early`]), /// classified by the session before it accepts any more. - pub early_unis: Vec>, + pub early_unis: Vec>, /// The extensions our SETUP offers. MoQ Auth is offered only while [`Self::auth`] /// supports it too. @@ -116,7 +103,10 @@ pub(crate) struct Driver { pub(crate) withdrawal: crate::session::Withdrawal, // Dispatched SUBSCRIBE and FETCH serves still owing the peer data. owed: std::sync::Arc, - future: BoxFuture<'static, Result<(), Error>>, + machine: Machine, + /// Keeps the test executor's registrations alive between polls of the future form. + #[cfg(test)] + park: Box, } impl Driver { @@ -124,12 +114,48 @@ impl Driver { pub(crate) fn drained(&self) -> bool { self.withdrawal.drained() && self.owed.load(std::sync::atomic::Ordering::Relaxed) == 0 } + + /// Run the session until it ends, with the error that ended it. + pub(crate) fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + self.machine.poll(waiter) + } } +#[cfg(test)] impl std::future::Future for Driver { type Output = Result<(), Error>; - fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll { - self.future.as_mut().poll(cx) + fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { + let Self { machine, park, .. } = &mut *self; + machine.poll(park.hold(cx)) + } +} + +/// Where the session is in its life. +enum Machine { + /// Not polled yet. The session is built on the first poll, since until then the app + /// may still take the peer's Auth requests. + Start(Box), + Running(Box), + Done, +} + +impl Machine { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + if matches!(self, Self::Start(_)) { + let Self::Start(launch) = std::mem::replace(self, Self::Done) else { + unreachable!("matched above"); + }; + *self = Self::Running(Box::new(launch.run())); + } + let Self::Running(running) = self else { + panic!("session driver polled after it finished"); + }; + let res = ready!(running.poll(waiter)); + let Self::Running(running) = std::mem::replace(self, Self::Done) else { + unreachable!("matched above"); + }; + running.finish(&res); + Poll::Ready(res) } } @@ -137,7 +163,7 @@ pub fn start(config: Config) -> Result<(Driver, crate::goaway::Handle, Option Result<(Driver, crate::goaway::Handle, Option Result<(Driver, crate::goaway::Handle, Option, + auth: crate::auth::Link, + runtime: crate::time::Deadlines, + stats: crate::stats::Session, + session: Transport, + limits: crate::session::Limits, + client: bool, + publish: Option, + subscribe: Option, + peer_hop: Option, + cost: Option, + version: Version, + setup: SetupParams, + peer_setup_stream: Option>, + early_unis: Vec>, + goaway: crate::goaway::Protocol, + peer_grant: crate::auth::Grant, + discovery: Option, + withdrawal: crate::session::Withdrawal, + owed: std::sync::Arc, + peer_setup: peer::PeerSetup, + setup_read: bool, +} + +impl Launch { + fn run(self) -> Running { + let Self { + setup_token, + auth, + runtime, + stats, + session, + limits, + client, + publish, + subscribe, + peer_hop, + cost, + version, + setup, + peer_setup_stream, + early_unis, + goaway, + peer_grant, + discovery, + withdrawal, + owed, + peer_setup, + setup_read, + } = self; + // Released on any exit, so nothing waits on a token the session will never answer. - let _auth_close = AuthClose(auth.clone()); + let auth_close = AuthClose(auth.clone()); // Decided once, before any Auth request can be accepted: the app took the requests // before running the driver, or the session answers itself. let _ = auth.acceptor(); @@ -234,223 +351,279 @@ pub fn start(config: Config) -> Result<(Driver, crate::goaway::Handle, Option enforce_grant(auth, publisher, session).await, - false => std::future::pending().await, - } - } - }; + let mut publisher = Publisher::new( + runtime.clone(), + session.clone(), + publish, + control.clone(), + peer_hop, + peer_setup.clone(), + version, + ) + .with_auth(auth.clone()); + let (tasks, task_set) = TaskSet::new(); + publisher.withdrawal = withdrawal; + publisher.owed = owed; + publisher.subscriptions = subscriptions; + + // Only the dialing side fails loud on a publication outside its grant: a + // server's publish origin is everything the peer may read, not what it + // intends to push. + let enforce = client.then(|| EnforceGrant::new(auth.clone(), publisher.clone(), session.clone())); + + let mut subscriber = Subscriber::new( + runtime.clone(), + session.clone(), + subscribe, + control.clone(), + peer_hop, + peer_setup.clone(), + self_origin, + cost, + version, + tasks, + goaway.going_away.clone(), + ) + .with_auth(auth.clone()) + .with_solicit(setup.extensions.solicit); + + // Our tokens, one Auth request each, once the peer's SETUP negotiates it. + let present = auth::Presenting::new( + runtime.clone(), + session.clone(), + control.clone(), + auth.clone(), + peer_setup.clone(), + version, + stats.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 `Dispatch`. + let serve = auth.supported().then(|| auth::Serve { + runtime: runtime.clone(), + handle: auth.clone(), + peer_grant, + }); + subscriber.announces = announces; - let mut subscriber = Subscriber::new( - runtime.clone(), - session.clone(), - subscribe, - control.clone(), - peer_hop, - peer_setup.clone(), - self_origin, - cost, + // When the peer's SETUP was pre-read (a gated server accept), monitor GOAWAY on + // that stream here; otherwise `Unis` does it when the SETUP arrives on the wire. + let goaway_recv = peer_setup_stream.map(|reader| { + ReadGoaways::new( + reader.with_stats(&stats).with_version(version), version, - tasks, - goaway.going_away.clone(), + client, + goaway.clone(), ) - .with_auth(auth.clone()) - .with_solicit(extensions.solicit); + }); - // Our tokens, one Auth request each, once the peer's SETUP negotiates it. - let present = auth::run_present( - runtime.clone(), - session.clone(), - control.clone(), - auth.clone(), - peer_setup.clone(), + let unis = Unis::new( + session.clone(), + subscriber.clone(), + UniSetup { + peer: Some(peer_setup.clone()), + read: setup_read, + early: early_unis, + client, version, - stats.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`. - let serve = auth.supported().then(|| auth::Serve { - runtime: runtime.clone(), - handle: auth.clone(), - peer_grant, - }); - subscriber.announces = announces; - - // When the peer's SETUP was pre-read (a gated server accept), monitor - // GOAWAY on that stream here; otherwise `run_unis` does it when the SETUP - // arrives on the wire. - let goaway_recv = { - let goaway = goaway.clone(); - let stats = stats.clone(); - async move { - match peer_setup_stream { - Some(reader) => { - run_goaway(reader.with_stats(&stats).with_version(version), version, client, goaway).await - } - None => std::future::pending().await, - } - } - }; + goaway, + }, + ); + let publish_namespaces = PublishNamespaces::new(publisher.clone()); + let dispatch = Dispatch::new( + session.clone(), + publisher, + subscriber.clone(), + peer_setup, + serve, + version, + ); + // A transport without datagrams never delivers one. + let datagrams = (session.max_datagram_size() != 0).then(|| Datagrams { + session: session.clone(), + subscriber: subscriber.clone(), + budget: kio::coop::Budget::new(32), + charged: false, + }); + // Keep the session's SUBSCRIBE_NAMESPACE requests open as discovery decides, each on + // its own stream; with no subscribe half, ask nothing. + let namespaces = discovery.map(|discovery| Namespaces::new(subscriber.clone(), discovery)); + + Running { + race: Race { + namespaces, + publish_namespaces: Some(publish_namespaces), + setup: Some(send_setup), + enforce, + present: Some(present), + goaway: goaway_recv, + datagrams, + dispatch: Some(dispatch), + unis, + subscriber, + tasks: task_set, + session, + }, + ending: None, + auth, + auth_close, + setup_token, + } + } +} - // Every half only ends the session on error (err_only parks on clean - // completion); `setup` never resolves (it holds the stream open) and the - // task set draining is the one clean exit. - let mut unis = Unis::new( - session.clone(), - subscriber.clone(), - UniSetup { - peer: Some(peer_setup.clone()), - read: setup_read, - early: early_unis, - client, - version, - goaway, - }, - ); - let mut dispatch = std::pin::pin!(err_only(run_dispatch( - session.clone(), - publisher.clone(), - subscriber.clone(), - peer_setup.clone(), - serve, - version - ))); - let mut datagrams = std::pin::pin!(err_only(run_datagrams(session.clone(), subscriber.clone()))); - let mut goaway_recv = std::pin::pin!(err_only(goaway_recv)); - let mut present = std::pin::pin!(err_only(present)); - let mut enforce = std::pin::pin!(err_only(enforce)); - let mut setup = std::pin::pin!(setup); - // Unsolicited PUBLISH_NAMESPACE unless the peer requires solicitation; - // see `Publisher::run_publish_namespaces`. - let mut pub_ns_run = std::pin::pin!(err_only(publisher.clone().run_publish_namespaces())); - let sub_ns_run = run_namespaces(subscriber.clone(), discovery, { - let session = session.clone(); - move || { - let session = session.clone(); - async move { Stream::open(&session, version).await } - } - }); - let mut sub_ns_run = std::pin::pin!(err_only(sub_ns_run)); +/// A running session: its halves race until one fails or the task set drains. +struct Running { + race: Race, + /// How the race ended, held while what already arrived drains. + ending: Option>, + auth: crate::auth::Link, + auth_close: AuthClose, + /// Held for the life of the session. + setup_token: Option, +} - let res = kio::wait(|waiter| { - use std::task::Poll; - if let Poll::Ready(res) = unis.poll(waiter) { - return Poll::Ready(res); - } - if let Poll::Ready(err) = waiter.poll_future(dispatch.as_mut()) { - return Poll::Ready(Err(err)); - } - if let Poll::Ready(err) = waiter.poll_future(datagrams.as_mut()) { - return Poll::Ready(Err(err)); - } - if let Poll::Ready(err) = waiter.poll_future(goaway_recv.as_mut()) { - return Poll::Ready(Err(err)); - } - // Malformed replies are session-fatal, unlike token refusals. - if let Poll::Ready(err) = waiter.poll_future(present.as_mut()) { - return Poll::Ready(Err(err)); - } - if let Poll::Ready(err) = waiter.poll_future(enforce.as_mut()) { - return Poll::Ready(Err(err)); - } - if let Poll::Ready(err) = waiter.poll_future(setup.as_mut()) { - return Poll::Ready(Err(err)); - } - if task_set.poll(waiter).is_ready() { - return Poll::Ready(Ok(())); - } - if let Poll::Ready(err) = waiter.poll_future(sub_ns_run.as_mut()) { - return Poll::Ready(Err(err)); - } - if let Poll::Ready(err) = waiter.poll_future(pub_ns_run.as_mut()) { - return Poll::Ready(Err(err)); - } - Poll::Pending - }) - .await; - close(&mut session, &res); - // Before this arm's Auth serve tasks drop, so each settles with the - // session's error rather than a bare cancel. - auth.close(match &res { +impl Running { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + if self.ending.is_none() { + let res = ready!(self.race.poll(waiter)); + close(&mut self.race.session, &res); + // Before the Auth serve tasks drop, so each settles with the session's error + // rather than a bare cancel. + self.auth.close(match &res { Ok(()) => Error::Cancel, Err(err) => err.clone(), }); - // Whichever side closed it, read what already arrived. - kio::wait(|waiter| unis.poll_drain(waiter)).await; - // Closed cleanly by dropping or draining: readers see what the peer's clean - // close would have shown them. - if closing.is_local() { - subscriber.abort(&Error::Session(SessionError::Cancel)); - } else if let Err(err) = &res { - // Every track this session was receiving ends with its error. - subscriber.abort(err); - } - res - }; + self.ending = Some(res); + } + // Whichever side closed it, read what already arrived. + ready!(self.race.unis.poll_drain(waiter)); + let res = self.ending.take().expect("draining an ended session"); + // Closed cleanly by dropping or draining: readers see what the peer's clean + // close would have shown them. + if self.race.session.closing().is_local() { + self.race.subscriber.abort(&Error::Session(SessionError::Cancel)); + } else if let Err(err) = &res { + // Every track this session was receiving ends with its error. + self.race.subscriber.abort(err); + } + Poll::Ready(res) + } - auth.close(match &res { + /// Drop the halves, then release the auth handle and our connection credential. + fn finish(self: Box, res: &Result<(), Error>) { + let Self { + race, + auth, + auth_close, + setup_token, + .. + } = *self; + drop(race); + auth.close(match res { Ok(()) => Error::Cancel, Err(err) => err.clone(), }); + drop(auth_close); + drop(setup_token); + } +} - res +/// The halves of a running session. Every half but the task set ends the session only by +/// failing: one that finishes cleanly drops out, and the task set draining is the one clean +/// exit. Fields drop in declaration order: the halves in the reverse of their poll order, +/// then the uni streams, the subscriber and the task set. +struct Race { + /// Our SUBSCRIBE_NAMESPACE requests; `None` without a subscribe half. + namespaces: Option, + /// Unsolicited PUBLISH_NAMESPACE unless the peer requires solicitation. + publish_namespaces: Option, + /// Never finishes while the session lives: it holds our SETUP stream open. + setup: Option, + /// `None` on a server, which publishes whatever the peer may read. + enforce: Option, + present: Option, + /// GOAWAY on a pre-read SETUP stream; `None` when `unis` reads the SETUP. + goaway: Option>, + datagrams: Option, + dispatch: Option, + unis: Unis, + subscriber: Subscriber, + tasks: TaskSet, + session: Transport, +} + +impl Race { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + if let Poll::Ready(res) = self.unis.poll(waiter) { + return Poll::Ready(res); + } + if let Poll::Ready(err) = poll_failed(&mut self.dispatch, |dispatch| dispatch.poll(waiter).map(Err)) { + return Poll::Ready(Err(err)); + } + if let Poll::Ready(err) = poll_failed(&mut self.datagrams, |datagrams| datagrams.poll(waiter)) { + return Poll::Ready(Err(err)); + } + if let Poll::Ready(err) = poll_failed(&mut self.goaway, |goaway| { + goaway.poll(waiter).map(|err| Err(control_stream_error(err))) + }) { + return Poll::Ready(Err(err)); + } + // Malformed replies are session-fatal, unlike token refusals. + if let Poll::Ready(err) = poll_failed(&mut self.present, |present| present.poll(waiter)) { + return Poll::Ready(Err(err)); + } + if let Poll::Ready(err) = poll_failed(&mut self.enforce, |enforce| enforce.poll(waiter)) { + return Poll::Ready(Err(err)); + } + if let Poll::Ready(err) = poll_failed(&mut self.setup, |setup| setup.poll(waiter)) { + return Poll::Ready(Err(err)); + } + if self.tasks.poll(waiter).is_ready() { + return Poll::Ready(Ok(())); + } + if let Poll::Ready(err) = poll_failed(&mut self.namespaces, |namespaces| namespaces.poll(waiter).map(Err)) { + return Poll::Ready(Err(err)); + } + if let Poll::Ready(err) = poll_failed(&mut self.publish_namespaces, |publish| publish.poll(waiter)) { + return Poll::Ready(Err(err)); + } + Poll::Pending } - .boxed(); +} - Ok(( - Driver { - withdrawal, - owed, - future: driver, - }, - goaway_handle, - credential, - )) +/// Poll a half of the race, `Ready` only with its failure. A half that finishes, either +/// way, drops right away and is never polled again. +fn poll_failed(half: &mut Option, poll: impl FnOnce(&mut T) -> Poll>) -> Poll { + let Some(running) = half else { + return Poll::Pending; + }; + let res = ready!(poll(running)); + *half = None; + match res { + Ok(()) => Poll::Pending, + Err(err) => Poll::Ready(err), + } } /// Close the session over the driver's result, while its streams are still open. @@ -481,7 +654,7 @@ fn close(session: &mut Transport, res: &Result<(), Error>) { /// What a peer's SETUP told us, beyond the stream it arrived on. pub struct PeerSetup { /// The SETUP stream, which becomes the GOAWAY channel. - pub stream: Reader, + pub stream: Reader, /// The SETUP's encoded size, counted once the session's stats are known. pub size: usize, @@ -492,7 +665,7 @@ pub struct PeerSetup { /// Uni streams that arrived before the SETUP, type peeked but unread. The drafts say /// to buffer early data until the control streams arrive, so the session classifies /// these once it starts. QUIC stream credit bounds how many can pile up. - pub early: Vec>, + pub early: Vec>, } /// The Hop ID this session declares and detects loops against. @@ -507,55 +680,107 @@ fn self_origin(publish: Option<&origin::Consumer>, subscribe: Option<&origin::Pr .unwrap_or_else(Hop::random) } -/// Server: read the peer's SETUP off its uni stream before starting the -/// session, returning that stream plus what it declared. +/// Server: reads the peer's SETUP off its uni stream before starting the +/// session, resolving to that stream plus what it declared. /// -/// Blocks on the peer's Setup Stream. Any other uni stream racing ahead of it (padding, +/// Waits for the peer's Setup Stream. Any other uni stream racing ahead of it (padding, /// or group data for a subscription the peer already holds) is held in /// [`PeerSetup::early`] for the session to classify. Pass the returned reader to /// [`start`] as its `peer_setup_stream` so GOAWAY monitoring continues without /// re-reading it. -pub async fn accept_setup(session: &mut Transport, version: Version) -> Result { - let outer_version = crate::Version::Ietf(version); - let mut early = Vec::new(); - +pub struct SetupAccept { + version: Version, // Each stream racing ahead of SETUP is held, so a peer opening them at once must not // hold the task for its whole MAX_STREAMS credit. - let mut budget = kio::coop::Budget::new(32); - loop { - kio::wait(|waiter| budget.poll_yield(waiter)).await; - let recv = session.accept_uni().await.map_err(Error::from_transport)?; - let mut reader: Reader = Reader::new(recv, outer_version); - - let kind = match reader.varint_peek().await { - Ok(kind) => kind, - Err(err) if died_before_header(&err) => { - tracing::debug!(%err, "dropping uni stream that died before its type"); - continue; - } - Err(err) => return Err(err), - }; - if kind != setup::SETUP { - early.push(reader); - continue; - } + budget: kio::coop::Budget, + early: Vec>, + step: SetupStep, +} + +/// Where [`SetupAccept`] waits. Streams are taken one at a time: the next is accepted only +/// once the last one's type has arrived. +enum SetupStep { + /// Spending a pass of the budget before accepting another stream. + Yield, + /// Waiting for the peer's next uni stream. + Accept, + /// Waiting for this stream's type. + Peek(Reader), + /// Waiting for the whole SETUP on its stream. + Decode(Reader), +} - reader.set_priority(crate::priority::RECV_CONTROL); - let (setup, size) = reader.decode_sized::().await?; - let options = decode_peer_setup(&setup.parameters, version, false)?; - if session.address() == crate::transport::Address::Uri && options.path.is_some() { - return Err(Error::ProtocolViolation); +impl SetupAccept { + /// Start waiting for a `version` peer's SETUP. + pub fn new(version: Version) -> Self { + Self { + version, + budget: kio::coop::Budget::new(32), + early: Vec::new(), + step: SetupStep::Yield, } + } - return Ok(PeerSetup { - stream: reader, - size, - options, - early, - }); + /// The peer's SETUP, with the streams that raced ahead of it. + pub fn poll(&mut self, session: &Transport, waiter: &kio::Waiter) -> Poll> { + loop { + match &mut self.step { + SetupStep::Yield => { + ready!(self.budget.poll_yield(waiter)); + self.step = SetupStep::Accept; + } + SetupStep::Accept => { + let recv = ready!(session.poll_accept_uni(waiter)).map_err(Error::from_transport)?; + self.step = SetupStep::Peek(Reader::new(recv, crate::Version::Ietf(self.version))); + } + SetupStep::Peek(reader) => { + let kind = match ready!(reader.poll_varint_peek(waiter)) { + Ok(kind) => kind, + Err(err) if died_before_header(&err) => { + tracing::debug!(%err, "dropping uni stream that died before its type"); + self.step = SetupStep::Yield; + continue; + } + Err(err) => return Poll::Ready(Err(err)), + }; + let SetupStep::Peek(mut reader) = std::mem::replace(&mut self.step, SetupStep::Yield) else { + unreachable!("matched above"); + }; + if kind != setup::SETUP { + self.early.push(reader); + continue; + } + reader.set_priority(crate::priority::RECV_CONTROL); + self.step = SetupStep::Decode(reader); + } + SetupStep::Decode(reader) => { + let (setup, size) = ready!(reader.poll_decode_sized::(waiter))?; + let options = decode_peer_setup(&setup.parameters, self.version, false)?; + if session.address() == crate::transport::Address::Uri && options.path.is_some() { + return Poll::Ready(Err(Error::ProtocolViolation)); + } + let SetupStep::Decode(stream) = std::mem::replace(&mut self.step, SetupStep::Yield) else { + unreachable!("matched above"); + }; + return Poll::Ready(Ok(PeerSetup { + stream, + size, + options, + early: std::mem::take(&mut self.early), + })); + } + } + } } } +/// [`SetupAccept`] as a future, for tests. +#[cfg(test)] +pub(crate) async fn accept_setup(session: &Transport, version: Version) -> Result { + let mut accept = SetupAccept::new(version); + kio::wait(|waiter| accept.poll(session, waiter)).await +} + /// What a peer's SETUP parameters say, parsed by [`decode_peer_setup`] for every caller. pub struct SetupOptions { /// The request path the peer advertised, for URL-less transports. @@ -610,125 +835,204 @@ pub(crate) fn decode_peer_setup(parameters: &[u8], version: Version, client: boo }) } -/// Send our SETUP on a uni stream and keep it alive: this stream is also our GOAWAY -/// channel, so a fired drain trigger encodes the GOAWAY here. -/// -/// `path` is the request path we advertise (clients on URL-less transports); a -/// server passes `None`, as it does for `token`, the client's connection credential. -/// `self_origin` and `cost` are the MoQ Cluster options, which -/// declare our identity and (client-only) what this link costs to cross. `extensions` -/// names the other extensions we offer. -#[expect( - clippy::too_many_arguments, - reason = "private plumbing; a single-use argument struct would only rename the parameters" -)] -async fn run_setup( - session: Transport, - version: Version, +/// What our SETUP declares beyond the MoQ Cluster options. +struct SetupParams { + /// The request path we advertise (clients on URL-less transports); a server passes `None`. path: Option, authority: Option, + /// The client's connection credential; a server passes `None`. token: Option, implementation: Option, + /// The other extensions we offer. + extensions: crate::setup::Extensions, +} + +/// Sends our SETUP on a uni stream and keeps it alive: this stream is also our GOAWAY +/// channel, so a fired drain trigger encodes the GOAWAY here. +/// +/// `self_origin` and `cost` are the MoQ Cluster options, which declare our identity and +/// (client-only) what this link costs to cross. +struct SendSetup { + session: Transport, + version: Version, self_origin: Hop, cost: Option, - extensions: crate::setup::Extensions, goaway: crate::goaway::Protocol, control: Control, stats: crate::stats::Session, -) -> Result<(), Error> { - let outer_version = crate::Version::Ietf(version); - - let send = session - .open_uni(Default::default()) - .await - .map_err(Error::from_transport)?; - let mut writer: Writer = - Writer::new(send, outer_version).with_stats(&stats); - - let mut parameters = ietf::Parameters::default(); - if let Some(implementation) = implementation { - parameters.set_bytes(ietf::ParameterBytes::Implementation, implementation.into_bytes()); - } - if let Some(path) = path { - parameters.set_bytes(ietf::ParameterBytes::Path, path.into_bytes()); - } - if let Some(authority) = authority { - parameters.set_bytes(ietf::ParameterBytes::Authority, authority.into_bytes()); - } - if let Some(token) = &token { - super::token::into_setup(&mut parameters, token, version)?; - } - cluster::peer_into_setup(&mut parameters, self_origin, cost); - if extensions.solicit { - solicit::into_setup(&mut parameters, version); - } - hidden::into_setup(&mut parameters, version); - max_age::into_setup(&mut parameters, version); - if extensions.auth { - auth::into_setup(&mut parameters); - } - active_count::into_setup(&mut parameters); - request_update::into_setup(&mut parameters, version); - super::flate::into_setup(&mut parameters); - let parameters = parameters.encode_bytes(version)?; - - writer - .encode(&setup::Setup { parameters }) - .await - .map_err(control_stream_error)?; - - // Hold the writer alive until the session closes, sending a GOAWAY if the - // drain trigger fires meanwhile. The trigger resolves `None` when the session - // drops without draining; keep holding either way, since a FIN mid-session is - // a protocol violation the peer closes over. - let mut untriggered = false; - let payload = kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() { - return std::task::Poll::Ready(Ok(None)); - } - if writer.poll_closed(waiter).is_ready() { - return std::task::Poll::Ready(Err(Error::ProtocolViolation)); - } - if !untriggered { - match goaway.poll_triggered(waiter) { - std::task::Poll::Ready(Some(payload)) => return std::task::Poll::Ready(Ok(Some(payload))), - std::task::Poll::Ready(None) => untriggered = true, - std::task::Poll::Pending => {} - } + step: SetupSend, +} + +/// Where [`SendSetup`] waits. +enum SetupSend { + /// Waiting for credit to open the stream. + Open(crate::transport::open::Open, SetupParams), + /// Flushing the SETUP. + Send(Writer), + /// Holding the stream until the session closes or the drain trigger fires. + Hold { + writer: Writer, + /// The trigger resolved without a GOAWAY: the session dropped without draining. + untriggered: bool, + }, + /// Flushing our GOAWAY. + GoAway(Writer), + /// Holding the stream after our GOAWAY until the session closes. + Linger(Writer), + /// Between steps. + Done, +} + +impl SendSetup { + #[expect( + clippy::too_many_arguments, + reason = "private plumbing; a single-use argument struct would only rename the parameters" + )] + fn new( + session: Transport, + version: Version, + params: SetupParams, + self_origin: Hop, + cost: Option, + goaway: crate::goaway::Protocol, + control: Control, + stats: crate::stats::Session, + ) -> Self { + Self { + session, + version, + self_origin, + cost, + goaway, + control, + stats, + step: SetupSend::Open(Default::default(), params), } - std::task::Poll::Pending - }) - .await?; + } - if let Some(payload) = payload { - let timeout_ms = payload.timeout.map(|d| d.as_millis() as u64).unwrap_or(0); - let msg = ietf::GoAway { - new_session_uri: std::borrow::Cow::Borrowed(payload.uri.as_ref()), - timeout: timeout_ms, - request_id: (version == Version::Draft18).then(|| control.goaway()).transpose()?, - }; + /// Encode what our SETUP declares. + fn parameters(&self, params: SetupParams) -> Result { + let version = self.version; + let mut parameters = ietf::Parameters::default(); + if let Some(implementation) = params.implementation { + parameters.set_bytes(ietf::ParameterBytes::Implementation, implementation.into_bytes()); + } + if let Some(path) = params.path { + parameters.set_bytes(ietf::ParameterBytes::Path, path.into_bytes()); + } + if let Some(authority) = params.authority { + parameters.set_bytes(ietf::ParameterBytes::Authority, authority.into_bytes()); + } + if let Some(token) = ¶ms.token { + super::token::into_setup(&mut parameters, token, version)?; + } + cluster::peer_into_setup(&mut parameters, self.self_origin, self.cost); + if params.extensions.solicit { + solicit::into_setup(&mut parameters, version); + } + hidden::into_setup(&mut parameters, version); + max_age::into_setup(&mut parameters, version); + if params.extensions.auth { + auth::into_setup(&mut parameters); + } + active_count::into_setup(&mut parameters); + request_update::into_setup(&mut parameters, version); + super::flate::into_setup(&mut parameters); + Ok(parameters.encode_bytes(version)?) + } - // Frame as [type_id varint][size u16][body], the same shape as the - // control-stream messages this channel otherwise carries. - let mut writer = writer.with_version(version); - writer.encode_message(&msg).await.map_err(control_stream_error)?; - let mut closed = std::pin::pin!(session.closed()); - kio::wait(|waiter| { - if waiter.poll_future(closed.as_mut()).is_ready() { - return std::task::Poll::Ready(Ok(())); - } - if writer.poll_closed(waiter).is_ready() { - return std::task::Poll::Ready(Err(Error::ProtocolViolation)); + /// `Ready` once the session closed with the stream held, or with why it broke. + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + loop { + match &mut self.step { + SetupSend::Open(open, _) => { + let send = ready!(self.session.poll_open_uni(open, waiter, Default::default())) + .map_err(Error::from_transport)?; + let SetupSend::Open(_, params) = std::mem::replace(&mut self.step, SetupSend::Done) else { + unreachable!("matched above"); + }; + let mut writer: Writer = + Writer::new(send, crate::Version::Ietf(self.version)).with_stats(&self.stats); + let parameters = self.parameters(params)?; + writer + .buffer(&setup::Setup { parameters }) + .map_err(control_stream_error)?; + self.step = SetupSend::Send(writer); + } + SetupSend::Send(writer) => { + ready!(writer.poll_flush(waiter)).map_err(control_stream_error)?; + let SetupSend::Send(writer) = std::mem::replace(&mut self.step, SetupSend::Done) else { + unreachable!("matched above"); + }; + self.step = SetupSend::Hold { + writer, + untriggered: false, + }; + } + // Hold the writer alive until the session closes, sending a GOAWAY if the + // drain trigger fires meanwhile. The trigger resolves `None` when the session + // drops without draining; keep holding either way, since a FIN mid-session is + // a protocol violation the peer closes over. + SetupSend::Hold { writer, untriggered } => { + if self.session.poll_closed(waiter).is_ready() { + writer.finish().ok(); + return Poll::Ready(Ok(())); + } + if writer.poll_closed(waiter).is_ready() { + return Poll::Ready(Err(Error::ProtocolViolation)); + } + if *untriggered { + return Poll::Pending; + } + let payload = match ready!(self.goaway.poll_triggered(waiter)) { + Some(payload) => payload, + None => { + *untriggered = true; + return Poll::Pending; + } + }; + let SetupSend::Hold { writer, .. } = std::mem::replace(&mut self.step, SetupSend::Done) else { + unreachable!("matched above"); + }; + + let timeout_ms = payload.timeout.map(|d| d.as_millis() as u64).unwrap_or(0); + let msg = ietf::GoAway { + new_session_uri: std::borrow::Cow::Borrowed(payload.uri.as_ref()), + timeout: timeout_ms, + request_id: (self.version == Version::Draft18) + .then(|| self.control.goaway()) + .transpose()?, + }; + + // Frame as [type_id varint][size u16][body], the same shape as the + // control-stream messages this channel otherwise carries. + let mut writer = writer.with_version(self.version); + writer + .buffer_message(ietf::GoAway::ID, &msg) + .map_err(control_stream_error)?; + self.step = SetupSend::GoAway(writer); + } + SetupSend::GoAway(writer) => { + ready!(writer.poll_flush(waiter)).map_err(control_stream_error)?; + let SetupSend::GoAway(writer) = std::mem::replace(&mut self.step, SetupSend::Done) else { + unreachable!("matched above"); + }; + self.step = SetupSend::Linger(writer); + } + SetupSend::Linger(writer) => { + if self.session.poll_closed(waiter).is_ready() { + writer.finish().ok(); + return Poll::Ready(Ok(())); + } + if writer.poll_closed(waiter).is_ready() { + return Poll::Ready(Err(Error::ProtocolViolation)); + } + return Poll::Pending; + } + SetupSend::Done => unreachable!("polled after it finished"), } - std::task::Poll::Pending - }) - .await?; - writer.finish().ok(); - } else { - writer.finish().ok(); + } } - - Ok(()) } /// The PADDING stream type (draft-18+): bytes a peer sends to probe for bandwidth. @@ -773,7 +1077,7 @@ struct UniSetup { peer: Option, read: bool, /// Streams accepted before the session started, handled before accepting more. - early: Vec>, + early: Vec>, /// Our role, which decides what the peer's SETUP and GOAWAY may carry. client: bool, version: Version, @@ -790,7 +1094,7 @@ struct UniSetup { struct Unis { session: Transport, uni: Uni, - early: std::collections::VecDeque>, + early: std::collections::VecDeque>, tasks: TaskSet, fatal: kio::Shared>, ended: Option, @@ -824,51 +1128,46 @@ impl Unis { let reader = match self.early.pop_front() { Some(reader) => reader, None => match self.session.poll_accept_uni(waiter) { - std::task::Poll::Ready(Ok(recv)) => Reader::new(recv, crate::Version::Ietf(self.uni.version)), - std::task::Poll::Ready(Err(err)) => { + Poll::Ready(Ok(recv)) => Reader::new(recv, crate::Version::Ietf(self.uni.version)), + Poll::Ready(Err(err)) => { self.ended = Some(Error::from_transport(err)); break; } - std::task::Poll::Pending => break, + Poll::Pending => break, }, }; - let uni = self.uni.clone(); - let fatal = self.fatal.clone(); - self.tasks.push(async move { - if let Err(err) = uni.serve(reader).await { - fatal.lock().get_or_insert(err); - } - }); + self.tasks + .push(UniServe::new(self.uni.clone(), reader, self.fatal.clone())); } } - fn poll(&mut self, waiter: &kio::Waiter) -> std::task::Poll> { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { let _ = self.tasks.poll(waiter); if let Some(err) = self.fatal.lock().take() { - return std::task::Poll::Ready(Err(err)); + return Poll::Ready(Err(err)); } self.accept(waiter); let _ = self.tasks.poll(waiter); if let Some(err) = self.fatal.lock().take() { - return std::task::Poll::Ready(Err(err)); + return Poll::Ready(Err(err)); } match &self.ended { - Some(err) => std::task::Poll::Ready(Err(err.clone())), - None => std::task::Poll::Pending, + Some(err) => Poll::Ready(Err(err.clone())), + None => Poll::Pending, } } - fn poll_drain(&mut self, waiter: &kio::Waiter) -> std::task::Poll<()> { + fn poll_drain(&mut self, waiter: &kio::Waiter) -> Poll<()> { self.accept(waiter); let drained = self.tasks.poll(waiter).is_ready(); if self.ended.is_none() { - return std::task::Poll::Pending; + return Poll::Pending; } // Every accepted child has first consumed its buffered header. A fill // still waiting to be promised or opened can no longer arrive. self.uni.subscriber.settle_control(); if drained { - std::task::Poll::Ready(()) + Poll::Ready(()) } else { self.tasks.poll(waiter) } @@ -887,92 +1186,155 @@ struct Uni { version: Version, } -impl Uni { - /// Read the stream's type and serve it, failing only when the stream breaks the session. - async fn serve(self, mut reader: Reader) -> Result<(), Error> { - reader = reader.with_stats(self.subscriber.stats()); - let Self { - mut subscriber, - peer_setup, - seen_setup, - goaway, - client, - version, - } = self; +/// Reads one incoming uni stream's type and serves it, reporting to `fatal` only when the +/// stream breaks the session. +struct UniServe { + uni: Uni, + step: UniStep, + fatal: kio::Shared>, +} - // A stream that dies before its type varint is that stream's failure, not the - // session's. RESET_STREAM is how a peer drops a group, and QUIC does not order - // the reset behind the data, so one can beat the first byte even of a stream - // the peer wrote to. Failing here would tear down the whole session over a - // single stream the peer had already given up on. Only death is tolerated: - // bytes that arrive and do not parse stay session-fatal. - // - // A transport error counts as death too. A reset whose code the transport cannot - // place in the stream registry surfaces as one: over raw QUIC a moq-transport - // peer resets with its own codes (moxygen's CANCELLED is 0x1), which the - // WebTransport code mapping rejects. If the connection itself died, the next - // accept reports it. - let kind: u64 = match kio::wait(|waiter| reader.poll_varint_peek(waiter)).await { - Ok(kind) => kind, - Err(err) if died_before_header(&err) => { - tracing::debug!(%err, "dropping uni stream that died before its type"); - return Ok(()); - } - Err(err) => return Err(err), - }; +/// Where [`UniServe`] waits. +enum UniStep { + /// Waiting for the stream's type. + Peek(Reader), + /// Waiting for the peer's whole SETUP. + Setup(Reader), + /// Monitoring the peer's SETUP stream for GOAWAY. + GoAway(ReadGoaways), + Group(Reader, RecvGroup), + Fill(Reader, RecvFill), + Padding(Reader), + /// Between steps. + Done, +} - let Some(ty) = UniType::classify(kind) else { - tracing::warn!(kind, "unknown uni stream type"); - return Err(Error::UnexpectedStream); - }; +impl UniServe { + fn new(uni: Uni, reader: Reader, fatal: kio::Shared>) -> Self { + let reader = reader.with_stats(uni.subscriber.stats()); + Self { + uni, + step: UniStep::Peek(reader), + fatal, + } + } - match ty { - // SETUP then becomes the GOAWAY channel. It is read in the background like any - // other stream; the one thing that does need it (the MoQ Cluster negotiation) - // waits on `peer_setup` instead, so a slow SETUP delays announcements rather - // than the whole session. - UniType::Setup => { - // Exactly one SETUP per endpoint. A second would let a peer restate its - // declared identity mid-session, silently re-attributing every route - // already built from the first. - if seen_setup.swap(true, std::sync::atomic::Ordering::Relaxed) { - return Err(Error::ProtocolViolation); - } + /// `Ready` once the stream is served, failing only when it breaks the session. + fn poll_serve(&mut self, waiter: &kio::Waiter) -> Poll> { + let version = self.uni.version; + loop { + match &mut self.step { + UniStep::Peek(reader) => { + // A stream that dies before its type varint is that stream's failure, not + // the session's. RESET_STREAM is how a peer drops a group, and QUIC does + // not order the reset behind the data, so one can beat the first byte even + // of a stream the peer wrote to. Failing here would tear down the whole + // session over a single stream the peer had already given up on. Only + // death is tolerated: bytes that arrive and do not parse stay + // session-fatal. + // + // A transport error counts as death too. A reset whose code the transport + // cannot place in the stream registry surfaces as one: over raw QUIC a + // moq-transport peer resets with its own codes (moxygen's CANCELLED is + // 0x1), which the WebTransport code mapping rejects. If the connection + // itself died, the next accept reports it. + let kind: u64 = match ready!(reader.poll_varint_peek(waiter)) { + Ok(kind) => kind, + Err(err) if died_before_header(&err) => { + tracing::debug!(%err, "dropping uni stream that died before its type"); + return Poll::Ready(Ok(())); + } + Err(err) => return Poll::Ready(Err(err)), + }; + + let Some(ty) = UniType::classify(kind) else { + tracing::warn!(kind, "unknown uni stream type"); + return Poll::Ready(Err(Error::UnexpectedStream)); + }; + + // Exactly one SETUP per endpoint. A second would let a peer restate its + // declared identity mid-session, silently re-attributing every route + // already built from the first. + if ty == UniType::Setup && self.uni.seen_setup.swap(true, std::sync::atomic::Ordering::Relaxed) { + return Poll::Ready(Err(Error::ProtocolViolation)); + } - // 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 - // slot nothing will ever fill. - reader.set_priority(crate::priority::RECV_CONTROL); - let msg = reader.decode::().await?; - if let Some(peer_setup) = peer_setup { - let options = decode_peer_setup(&msg.parameters, version, client)?; - peer_setup.set(options.declared); - peer_setup.set_implementation(options.implementation); + let UniStep::Peek(mut reader) = std::mem::replace(&mut self.step, UniStep::Done) else { + unreachable!("matched above"); + }; + self.step = match ty { + // SETUP then becomes the GOAWAY channel. It is read in the background + // like any other stream; the one thing that does need it (the MoQ + // Cluster negotiation) waits on `peer_setup` instead, so a slow SETUP + // delays announcements rather than the whole session. + // + // 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 + // slot nothing will ever fill. + UniType::Setup => { + reader.set_priority(crate::priority::RECV_CONTROL); + UniStep::Setup(reader) + } + UniType::Subgroup => UniStep::Group(reader.with_version(version), RecvGroup::new()), + // A fill fetch stream carries the head of the group a draft-22 + // subscription joined part way through. One answering no fill of ours + // is refused inside. + UniType::Fetch => UniStep::Fill(reader.with_version(version), RecvFill::new()), + UniType::Padding => UniStep::Padding(reader), + }; } + UniStep::Setup(reader) => { + let msg = ready!(reader.poll_decode::(waiter))?; + if let Some(peer_setup) = &self.uni.peer_setup { + let options = decode_peer_setup(&msg.parameters, version, self.uni.client)?; + peer_setup.set(options.declared); + peer_setup.set_implementation(options.implementation); + } - // Monitor for GOAWAY after setup completes. - run_goaway(reader.with_version(version), version, client, goaway).await?; - } - UniType::Subgroup => { - let mut reader = reader.with_version(version); - let res = subscriber.recv_group(&mut reader).await; - stop_on_error(&mut reader, res)?; - } - // A fill fetch stream carries the head of the group a draft-22 subscription - // joined part way through. One answering no fill of ours is refused inside. - UniType::Fetch => { - let mut reader = reader.with_version(version); - let res = subscriber.recv_fill(&mut reader).await; - stop_on_error(&mut reader, res)?; - } - // The receiver MUST discard padding. One that dies is not the session's failure. - UniType::Padding => { - if let Err(err) = kio::wait(|waiter| reader.poll_discard(waiter)).await { - tracing::debug!(%err, "padding stream ended"); + let UniStep::Setup(reader) = std::mem::replace(&mut self.step, UniStep::Done) else { + unreachable!("matched above"); + }; + // Monitor for GOAWAY after setup completes. + self.step = UniStep::GoAway(ReadGoaways::new( + reader.with_version(version), + version, + self.uni.client, + self.uni.goaway.clone(), + )); + } + UniStep::GoAway(goaways) => { + let err = ready!(goaways.poll(waiter)); + return Poll::Ready(Err(control_stream_error(err))); + } + UniStep::Group(reader, recv) => { + let res = ready!(recv.poll(waiter, &self.uni.subscriber, reader)); + return Poll::Ready(stop_on_error(reader, res)); + } + UniStep::Fill(reader, recv) => { + let res = ready!(recv.poll(waiter, &self.uni.subscriber, reader)); + return Poll::Ready(stop_on_error(reader, res)); + } + // The receiver MUST discard padding. One that dies is not the session's failure. + UniStep::Padding(reader) => { + if let Err(err) = ready!(reader.poll_discard(waiter)) { + tracing::debug!(%err, "padding stream ended"); + } + return Poll::Ready(Ok(())); } + UniStep::Done => unreachable!("polled after it finished"), } } - Ok(()) + } +} + +impl kio::Task for UniServe { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + if let Err(err) = ready!(self.poll_serve(waiter)) { + self.fatal.lock().get_or_insert(err); + } + Poll::Ready(()) } } @@ -992,18 +1354,27 @@ fn died_before_header(err: &Error) -> bool { /// Receive QUIC datagrams, each an OBJECT_DATAGRAM for one of our subscriptions. /// -/// A transport without datagrams never delivers one, so this parks. A transport failure -/// or a malformed datagram ends the session. -async fn run_datagrams(session: Transport, subscriber: Subscriber) -> Result<(), Error> { - if session.max_datagram_size() == 0 { - return Ok(()); - } - // A peer keeping datagrams ready must not hold the session's other tasks. - let mut budget = kio::coop::Budget::new(32); - loop { - kio::wait(|waiter| budget.poll_yield(waiter)).await; - let payload = session.recv_datagram().await.map_err(Error::from_transport)?; - subscriber.recv_datagram(payload)?; +/// A transport failure or a malformed datagram ends the session. +struct Datagrams { + session: Transport, + subscriber: Subscriber, + /// A peer keeping datagrams ready must not hold the session's other tasks. + budget: kio::coop::Budget, + /// This pass of the budget is spent, waiting for the next datagram. + charged: bool, +} + +impl Datagrams { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + loop { + if !self.charged { + ready!(self.budget.poll_yield(waiter)); + self.charged = true; + } + let payload = ready!(self.session.poll_recv_datagram(waiter)).map_err(Error::from_transport)?; + self.charged = false; + self.subscriber.recv_datagram(payload)?; + } } } @@ -1030,98 +1401,136 @@ pub(super) fn stop_on_error( Ok(()) } -/// Accept incoming bidi streams, leaving each request to read its own header. -async fn run_dispatch( +/// Accept incoming bidi streams, leaving each request to read its own header. Only an +/// error ends it. +struct Dispatch { session: Transport, - publisher: Publisher, - subscriber: Subscriber, - peer_setup: peer::PeerSetup, - // Answers the peer's Auth requests, on the versions that can negotiate them. - serve: Option, version: Version, -) -> Result<(), Error> { - use std::task::Poll; + peer_setup: peer::PeerSetup, + /// What `bidi` is built from once the peer's SETUP arrives. + start: Option<(Publisher, Subscriber, Option)>, + /// What each request is served with, once the peer's SETUP says how they decode. + bidi: Option, + tasks: TaskSet, + fatal: kio::Shared>, + /// MAX_STREAMS credit alone lets the peer keep thousands of streams ready, and the + /// children are polled only when accepting pends. A yield pends it. + budget: kio::coop::Budget, + /// One pass per accepted stream, not per wake. + charge: bool, + /// Each AUTH serve task settles `Issued::closed` from whatever the auth handle + /// holds when it drops, and the first reason wins. So the error that ends the + /// session has to reach the handle while `tasks` is still alive, or every grant + /// reports a bare cancel instead. `None` when the peer never negotiated AUTH, + /// which leaves no serve task to settle. + handle: Option, +} - // PUBLISH_NAMESPACE decodes differently once the MoQ Cluster extension is - // negotiated, so dispatch waits for SETUP, which `run_unis` reads independently. - let peer = subscriber.peer().await; +impl Dispatch { + fn new( + session: Transport, + publisher: Publisher, + subscriber: Subscriber, + peer_setup: peer::PeerSetup, + // Answers the peer's Auth requests, on the versions that can negotiate them. + serve: Option, + version: Version, + ) -> Self { + Self { + session, + version, + peer_setup, + start: Some((publisher, subscriber, serve)), + bidi: None, + tasks: TaskSet::owned(), + fatal: Default::default(), + budget: kio::coop::Budget::new(32), + charge: true, + handle: None, + } + } - // An AUTH without MoQ Auth negotiated, offered by both sides, is an unknown request, - // which falls through to the protocol violation in `Bidi::serve`. - let declared_setup = peer_setup.get().await; - let serve = match declared_setup.auth { - true => serve, - false => None, - }; + /// Build what each request is served with, once the peer's SETUP arrives. + fn poll_setup(&mut self, waiter: &kio::Waiter) -> Poll<()> { + let (_, subscriber, _) = self.start.as_ref().expect("set up once"); - // From the same slot, so this costs nothing extra: it decides whether an unsolicited - // advertisement is the peer ignoring our own SETUP (MoQ Solicit). - let declared = subscriber.solicit().await; - - // Each AUTH serve task settles `Issued::closed` from whatever the auth handle - // holds when it drops, and the first reason wins. So the error that ends the - // session has to reach the handle while `tasks` is still alive, or every grant - // reports a bare cancel instead. `None` when the peer never negotiated AUTH, - // which leaves no serve task to settle. - let handle = serve.as_ref().map(|serve| serve.handle.clone()); - - let bidi = Bidi { - publisher, - subscriber, - peer, - declared, - serve, - version, - }; - let mut tasks = TaskSet::owned(); - let fatal = kio::Shared::>::default(); - // MAX_STREAMS credit alone lets the peer keep thousands of streams ready, and the - // drive below polls the children only when accepting pends. A yield pends it. - let mut budget = kio::coop::Budget::new(32); - - let err = loop { - // One pass per accepted stream, not per wake of the drive. - let mut charge = true; - let accepted = tasks - .drive(|waiter| { - let broken = fatal.poll(waiter, |fatal| match fatal.is_some() { + // PUBLISH_NAMESPACE decodes differently once the MoQ Cluster extension is + // negotiated, so dispatch waits for SETUP, which `Unis` reads independently. + let peer = ready!(subscriber.poll_peer(waiter)); + + // An AUTH without MoQ Auth negotiated, offered by both sides, is an unknown request, + // which falls through to the protocol violation in `BidiServe`. + let negotiated = ready!(self.peer_setup.poll_get(waiter)).auth; + + // From the same slot, so this costs nothing extra: it decides whether an unsolicited + // advertisement is the peer ignoring our own SETUP (MoQ Solicit). + let declared = ready!(subscriber.poll_solicit(waiter)); + + let (publisher, subscriber, serve) = self.start.take().expect("set up once"); + let serve = match negotiated { + true => serve, + false => None, + }; + self.handle = serve.as_ref().map(|serve| serve.handle.clone()); + self.bidi = Some(Bidi { + publisher, + subscriber, + peer, + declared, + serve, + version: self.version, + }); + Poll::Ready(()) + } + + /// `Ready` only with the error that ends the session. + fn poll(&mut self, waiter: &kio::Waiter) -> Poll { + if self.bidi.is_none() { + ready!(self.poll_setup(waiter)); + } + let bidi = self.bidi.as_ref().expect("set up above"); + loop { + let accepted = 'accept: { + let broken = self.fatal.poll(waiter, |fatal| match fatal.is_some() { true => Poll::Ready(()), false => Poll::Pending, }); if let Poll::Ready(mut fatal) = broken && let Some(err) = fatal.take() { - return Poll::Ready(Err(err)); + break 'accept Poll::Ready(Err(err)); } - if std::mem::take(&mut charge) { - std::task::ready!(budget.poll_yield(waiter)); + if std::mem::take(&mut self.charge) && self.budget.poll_yield(waiter).is_pending() { + break 'accept Poll::Pending; + } + Stream::poll_accept(&self.session, self.version, waiter) + }; + match accepted { + // QUIC may expose an earlier stream before its bytes arrive. Reading the + // header here could block later streams whose unread data holds flow control. + Poll::Ready(Ok(stream)) => { + self.charge = true; + self.tasks + .push(BidiServe::new(bidi.clone(), stream, self.fatal.clone())); + } + Poll::Ready(Err(err)) => { + // `tasks` still holds the AUTH serve tasks, so the session's error lands on + // the handle before they drop. The driver closes it again with the same + // error, which then does nothing. + if let Some(handle) = &self.handle { + handle.close(err.clone()); + } + return Poll::Ready(err); + } + Poll::Pending => { + // The children never end the loop; a `Ready` here just means they're + // drained until the next stream. + let _ = self.tasks.poll(waiter); + return Poll::Pending; } - Stream::poll_accept(&session, version, waiter) - }) - .await; - let stream = match accepted { - Ok(stream) => stream, - Err(err) => break err, - }; - - // QUIC may expose an earlier stream before its bytes arrive. Reading the - // header here could block later streams whose unread data holds flow control. - let bidi = bidi.clone(); - let fatal = fatal.clone(); - tasks.push(async move { - if let Err(err) = bidi.serve(stream).await { - fatal.lock().get_or_insert(err); } - }); - }; - - // `tasks` still holds the AUTH serve tasks, so the session's error lands on the - // handle before they drop. The driver closes it again with the same error, which - // then does nothing. - if let Some(handle) = &handle { - handle.close(err.clone()); + } } - Err(err) } /// What serving one incoming bidi request needs from its session. @@ -1135,38 +1544,88 @@ struct Bidi { version: Version, } -impl Bidi { - /// Read and serve a request, failing only when it breaks the session. - async fn serve(self, mut stream: Stream) -> Result<(), Error> { - stream = stream.with_stats(self.publisher.stats()); - let Self { +/// Reads one request's header and serves it, reporting to `fatal` only when the request +/// breaks the session. +struct BidiServe { + bidi: Bidi, + step: BidiStep, + fatal: kio::Shared>, +} + +/// Where [`BidiServe`] waits. +enum BidiStep { + /// Reading the request's type, then its body. + Header { + stream: Stream, + id: Option, + }, + /// Refusing a request past our GOAWAY cutoff. + Refuse(Reject), + /// The publisher or subscriber serving it. + Serve(crate::util::BoxTask), + /// Answering an Auth request. + Auth(Box), + /// Between steps. + Done, +} + +impl BidiServe { + fn new(bidi: Bidi, stream: Stream, fatal: kio::Shared>) -> Self { + let stream = stream.with_stats(bidi.publisher.stats()); + Self { + bidi, + step: BidiStep::Header { stream, id: None }, + fatal, + } + } + + /// `Ready` once the request is served, failing only when it breaks the session. + fn poll_serve(&mut self, waiter: &kio::Waiter) -> Poll> { + loop { + match &mut self.step { + BidiStep::Header { stream, id } => { + // Abandoning a request before its header is not a session failure. Bytes + // that arrive and do not parse, and dispatch errors, remain session-fatal. + let (id, data) = match ready!(poll_header(stream, id, waiter)) { + Ok(header) => header, + Err(err) if died_before_header(&err) => { + tracing::debug!(%err, "dropping bidi stream that died before its header"); + return Poll::Ready(Ok(())); + } + Err(err) => return Poll::Ready(Err(err)), + }; + let BidiStep::Header { stream, .. } = std::mem::replace(&mut self.step, BidiStep::Done) else { + unreachable!("matched above"); + }; + self.step = self.dispatch(id, data, stream)?; + } + BidiStep::Refuse(reject) => return reject.poll(waiter), + BidiStep::Serve(task) => { + ready!(task(waiter)); + return Poll::Ready(Ok(())); + } + BidiStep::Auth(answering) => return answering.poll(waiter), + BidiStep::Done => unreachable!("polled after it finished"), + } + } + } + + /// Hand a request whose header arrived to whatever serves it. + fn dispatch( + &mut self, + id: u64, + data: ietf::Body, + mut stream: Stream, + ) -> Result { + let Bidi { publisher, - mut subscriber, + subscriber, peer, declared, serve, version, - } = self; - let header = async { - let id = stream.reader.varint().await?; - // The type arrived, so a body cut short by FIN is malformed, not the stream dying. - let body = match stream.reader.decode::().await { - Err(Error::Decode(DecodeError::Short)) => return Err(Error::WrongSize), - body => body?, - }; - Ok::<_, Error>((id, body)) - } - .await; - // Abandoning a request before its header is not a session failure. Bytes - // that arrive and do not parse, and dispatch errors, remain session-fatal. - let (id, data) = match header { - Ok(header) => header, - Err(err) if died_before_header(&err) => { - tracing::debug!(%err, "dropping bidi stream that died before its header"); - return Ok(()); - } - Err(err) => return Err(err), - }; + } = &mut self.bidi; + let version = *version; stream.reader = stream.reader.with_stats(publisher.stats()); publisher .stats() @@ -1184,49 +1643,75 @@ impl Bidi { ) { let request_id = RequestId::decode(&mut data.decoder(version), version)?; match publisher.control.accept(request_id) { - Err(Error::GoingAway) => return refuse_going_away(stream, id, version).await, + Err(Error::GoingAway) => return Ok(BidiStep::Refuse(refuse_going_away(stream, id, version))), result => { result?; } } } - match id { + Ok(match id { // Publisher handles: Subscribe, Fetch, SubscribeNamespace, SubscribeTracks, TrackStatus ietf::Subscribe::ID | ietf::Fetch::ID | ietf::SubscribeNamespace::ID | ietf::SubscribeTracks::ID - | ietf::TrackStatus::ID => { - publisher.handle_stream(id, data, stream)?.await; - } + | ietf::TrackStatus::ID => BidiStep::Serve(publisher.handle_stream(id, data, stream)?), // Subscriber handles: Publish, PublishNamespace ietf::Publish::ID | ietf::PublishNamespace::ID => { - subscriber.handle_stream(id, data, stream, peer, declared)?.await; + BidiStep::Serve(subscriber.handle_stream(id, data, stream, *peer, *declared)?) } - auth::Auth::ID if let Some(serve) = serve => { + auth::Auth::ID if let Some(serve) = serve.take() => { let mut data = data.decoder(version); let msg = auth::Auth::decode_msg(&mut data, version)?; if !data.is_empty() { return Err(Error::WrongSize); } - serve.run(stream, msg, version).await?; + BidiStep::Auth(Box::new(serve.answer(stream, msg, version))) } _ => { tracing::warn!(id, "unexpected bidi stream type"); return Err(Error::UnexpectedStream); } + }) + } +} + +impl kio::Task for BidiServe { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + if let Err(err) = ready!(self.poll_serve(waiter)) { + self.fatal.lock().get_or_insert(err); } - Ok(()) + Poll::Ready(()) } } +/// A request's type, then its body. `id` holds the type once it arrived. +fn poll_header( + stream: &mut Stream, + id: &mut Option, + waiter: &kio::Waiter, +) -> Poll> { + let id = match *id { + Some(id) => id, + None => *id.insert(ready!(stream.reader.poll_varint(waiter))?), + }; + // The type arrived, so a body cut short by FIN is malformed, not the stream dying. + let body = match ready!(stream.reader.poll_decode::(waiter)) { + Err(Error::Decode(DecodeError::Short)) => return Poll::Ready(Err(Error::WrongSize)), + body => body?, + }; + Poll::Ready(Ok((id, body))) +} + /// Refuse a request past our GOAWAY cutoff with the answer its kind reads: an AUTH /// presenter reads only AUTH_OK and AUTH_ERROR. -async fn refuse_going_away(stream: Stream, id: u64, version: Version) -> Result<(), Error> { +fn refuse_going_away(stream: Stream, id: u64, version: Version) -> Reject { if id == auth::Auth::ID { - return auth::refuse_going_away(stream).await; + return auth::going_away(stream); } - super::request_stream::reject( + Reject::new( stream, ietf::RequestError { error_code: super::error::request::to_code( @@ -1238,20 +1723,6 @@ async fn refuse_going_away(stream: Stream, id: u64, version: retry_interval: 0, }, ) - .await -} - -/// Monitor the peer's SETUP stream for GOAWAY until the session closes. -/// A FIN or reset of this control stream is a protocol violation. -async fn run_goaway( - reader: Reader, - version: Version, - client: bool, - goaway: crate::goaway::Protocol, -) -> Result<(), Error> { - read_goaways(reader, version, client, goaway) - .await - .map_err(control_stream_error) } /// A control stream cannot end independently of its session, so its own failure is the @@ -1264,60 +1735,90 @@ pub(super) fn control_stream_error(err: Error) -> Error { } } -/// Read every message after SETUP on the peer's SETUP stream, which may only be GOAWAY. +/// Reads every message after SETUP on the peer's SETUP stream, which may only be GOAWAY. +/// A FIN or reset of this control stream is a protocol violation, so only an error ends it; +/// [`control_stream_error`] says which one the session closes with. /// /// Reading continues past the first GOAWAY until the peer FINs or the session closes. /// Dropping the reader would STOP_SENDING the peer's SETUP uni stream, which draft-22 /// section 6.4.1 forbids closing at the transport layer mid-session, so a strict peer would /// tear the session down as a PROTOCOL_VIOLATION in the middle of the drain we are honoring. -async fn read_goaways( - mut reader: Reader, +struct ReadGoaways { + reader: Reader, version: Version, client: bool, goaway: crate::goaway::Protocol, -) -> Result<(), Error> { - while let Some(id) = reader.varint_maybe().await? { - // The type arrived, so a body cut short by FIN is malformed, not the stream dying. - let body = match reader.decode::().await { - Err(Error::Decode(DecodeError::Short)) => return Err(Error::WrongSize), - body => body?, - }; - if id != ietf::GoAway::ID { - return Err(Error::UnexpectedMessage); + /// The type of the message whose body is arriving. + id: Option, +} + +impl ReadGoaways { + fn new(reader: Reader, version: Version, client: bool, goaway: crate::goaway::Protocol) -> Self { + Self { + reader, + version, + client, + goaway, + id: None, } + } - let msg: ietf::GoAway = body.message(version)?; - reader - .stats() - .control(crate::stats::Message::Goaway, body.0.len() + 2, true); - tracing::info!(message = ?msg, "received GOAWAY"); - - // Draft-18 names the first of our own Request IDs it did not process: even for a - // client, odd for a server. The field is required here, on the control stream. - match msg.request_id { - None if version == Version::Draft18 => return Err(Error::ProtocolViolation), - Some(request_id) if request_id.0 % 2 != u64::from(!client) => { - return Err(DecodeError::Session(SessionError::InvalidRequestId).into()); + /// `Ready` only with the error that ends it. + fn poll(&mut self, waiter: &kio::Waiter) -> Poll { + self.poll_read(waiter).map(|Err(err)| err) + } + + fn poll_read(&mut self, waiter: &kio::Waiter) -> Poll> { + let version = self.version; + loop { + let id = match self.id { + Some(id) => id, + None => match ready!(self.reader.poll_varint_maybe(waiter))? { + Some(id) => *self.id.insert(id), + None => return Poll::Ready(Err(Error::ProtocolViolation)), + }, + }; + // The type arrived, so a body cut short by FIN is malformed, not the stream dying. + let body = match ready!(self.reader.poll_decode::(waiter)) { + Err(Error::Decode(DecodeError::Short)) => return Poll::Ready(Err(Error::WrongSize)), + body => body?, + }; + self.id = None; + if id != ietf::GoAway::ID { + return Poll::Ready(Err(Error::UnexpectedMessage)); } - _ => {} - } - // Nothing may follow a GOAWAY on this stream, so bytes already here are a violation. - // Caught before the GOAWAY surfaces: the application acting on it can close the - // session first, without our code. - let waiter = &kio::Waiter::noop(); - if let std::task::Poll::Ready(Ok(true)) = reader.poll_has_more(waiter) { - return Err(Error::ProtocolViolation); - } + let msg: ietf::GoAway = body.message(version)?; + self.reader + .stats() + .control(crate::stats::Message::Goaway, body.0.len() + 2, true); + tracing::info!(message = ?msg, "received GOAWAY"); + + // Draft-18 names the first of our own Request IDs it did not process: even for a + // client, odd for a server. The field is required here, on the control stream. + match msg.request_id { + None if version == Version::Draft18 => return Poll::Ready(Err(Error::ProtocolViolation)), + Some(request_id) if request_id.0 % 2 != u64::from(!self.client) => { + return Poll::Ready(Err(DecodeError::Session(SessionError::InvalidRequestId).into())); + } + _ => {} + } - let timeout = (msg.timeout > 0).then(|| std::time::Duration::from_millis(msg.timeout)); - // A second GOAWAY is a protocol violation the draft requires we close over. - goaway.record(crate::goaway::Goaway { - uri: msg.new_session_uri.into_owned(), - timeout, - })?; + // Nothing may follow a GOAWAY on this stream, so bytes already here are a violation. + // Caught before the GOAWAY surfaces: the application acting on it can close the + // session first, without our code. + if let Poll::Ready(Ok(true)) = self.reader.poll_has_more(&kio::Waiter::noop()) { + return Poll::Ready(Err(Error::ProtocolViolation)); + } + + let timeout = (msg.timeout > 0).then(|| std::time::Duration::from_millis(msg.timeout)); + // A second GOAWAY is a protocol violation the draft requires we close over. + self.goaway.record(crate::goaway::Goaway { + uri: msg.new_session_uri.into_owned(), + timeout, + })?; + } } - Err(Error::ProtocolViolation) } /// Closes the auth handle when the session driver ends without finishing, releasing @@ -1335,12 +1836,6 @@ impl Drop for AuthClose { /// /// Judges the publisher's serving origin rather than the raw one: an origin that also /// consumes holds what the peer sent us, which is not ours to publish back. -async fn enforce_grant(auth: crate::auth::Link, publisher: Publisher, session: Transport) -> Result<(), Error> { - let mut enforce = EnforceGrant::new(auth, publisher, session); - kio::wait(|waiter| enforce.poll(waiter)).await -} - -/// See [`enforce_grant`]. struct EnforceGrant { auth: crate::auth::Link, publisher: Publisher, @@ -2804,16 +3299,14 @@ mod tests { tasks, Default::default(), ); - let mut dispatch = std::pin::pin!(run_dispatch( - transport, publisher, subscriber, peer_setup, None, VERSION - )); + let mut dispatch = Dispatch::new(transport, publisher, subscriber, peer_setup, None, VERSION); // Held open and empty, so each accepted stream waits on its header. let mut opened = Vec::new(); for _ in 0..1024 { opened.push(peer.open_bi(Default::default()).await.unwrap()); } - assert!(kio::Waiter::noop().poll_future(dispatch.as_mut()).is_pending()); + assert!(dispatch.poll(&kio::Waiter::noop()).is_pending()); assert!( local.poll_accept_bi(&kio::Waiter::noop()).is_ready(), "one poll accepted every stream" @@ -3011,7 +3504,7 @@ mod tests { subgroup_header(VERSION, 7, 0).await, setup, ]); - let accepted = accept_setup(&mut crate::driver::Transport::new(session.clone()), VERSION) + let accepted = accept_setup(&crate::driver::Transport::new(session.clone()), VERSION) .await .expect("accept the SETUP"); assert_eq!(accepted.early.len(), 2, "both early streams are held"); @@ -3402,14 +3895,13 @@ mod tests { let session = DeadStreamSession::unis(1); let recv = session.accept_uni().await.unwrap(); let (_handle, protocol) = crate::goaway::Handle::new(false); - let err = run_goaway( + let mut goaways = ReadGoaways::new( Reader::new(crate::transport::boxed::RecvStream::new(recv), version), version, true, protocol, - ) - .await - .unwrap_err(); + ); + let err = control_stream_error(kio::wait(|waiter| goaways.poll(waiter)).await); assert_eq!(SessionError::from(&err), SessionError::ProtocolViolation, "{version:?}"); } } @@ -3425,14 +3917,14 @@ mod tests { let stream = kio::wait(|waiter| session.poll_accept_uni(waiter)).await.unwrap(); let (handle, protocol) = crate::goaway::Handle::new(false); - let err = read_goaways( + let mut goaways = ReadGoaways::new( Reader::new(crate::transport::boxed::RecvStream::new(stream), version), version, true, protocol, - ) - .await - .expect_err("a second GOAWAY is fatal"); + ); + // A second GOAWAY is fatal. + let err = kio::wait(|waiter| goaways.poll(waiter)).await; assert_eq!(SessionError::from(&err), SessionError::ProtocolViolation, "{version:?}"); assert!( !matches!( @@ -3527,22 +4019,24 @@ mod tests { control.accept(RequestId(first)).unwrap(); control.accept(RequestId(first + 4)).unwrap(); let (handle, goaway_protocol) = crate::goaway::Handle::new(!client); - let mut setup = std::pin::pin!(run_setup( + let mut setup = SendSetup::new( crate::driver::Transport::new(session.clone()), VERSION, - None, - None, - None, - None, + SetupParams { + path: None, + authority: None, + token: None, + implementation: None, + extensions: Default::default(), + }, Hop::new(1).unwrap(), None, - Default::default(), goaway_protocol, control.clone(), Default::default(), - )); + ); handle.producer().send(crate::goaway::Goaway::new()).unwrap(); - assert!(futures::poll!(setup.as_mut()).is_pending()); + assert!(setup.poll(&kio::Waiter::noop()).is_pending()); let sent = goaway(VERSION, Some(first + 6)); assert_eq!(occurrences(&log, &sent), 1, "client={client}"); control.accept(RequestId(first + 2)).unwrap(); @@ -3615,7 +4109,8 @@ mod tests { let stream = Stream::open(&crate::driver::Transport::new(session.clone()), VERSION) .await .unwrap(); - refuse_going_away(stream, id, VERSION).await.unwrap(); + let mut refusal = refuse_going_away(stream, id, VERSION); + kio::wait(|waiter| refusal.poll(waiter)).await.unwrap(); assert_eq!(*log.writes.lock().unwrap(), answer, "{id:#x}"); } } diff --git a/rs/moq-net/src/ietf/subscriber.rs b/rs/moq-net/src/ietf/subscriber.rs index 12b19e6d989..e5553fb2adb 100644 --- a/rs/moq-net/src/ietf/subscriber.rs +++ b/rs/moq-net/src/ietf/subscriber.rs @@ -1,9 +1,7 @@ use crate::driver::Transport; use crate::transport::poll::Session as _; -use futures::FutureExt as _; use std::{ collections::{BTreeMap, HashMap, hash_map::Entry}, - sync::atomic::{AtomicU64, Ordering}, task::{Poll, ready}, time::Duration, }; @@ -14,7 +12,7 @@ use crate::{ frame, group, ietf::{self, Control, FetchType, Filter, GroupOrder, RequestId}, origin, track, - util::{BoxFuture, TaskSet, Tasks}, + util::{TaskSet, Tasks}, }; use super::{Message, Version, cluster, error::request, group::ObjectExtensionsLength, peer, request_update}; @@ -24,6 +22,7 @@ use crate::util::Later; use kio::Lock; mod namespaces; +pub(super) use namespaces::Namespaces; mod tracks; const TRACK_ALIAS_TIMEOUT: Duration = Duration::from_secs(1); @@ -479,17 +478,17 @@ struct RangeFetch { } /// The groups one fetch stream writes into its fill, a group at a time. -struct RangeWriter<'a> { +struct RangeWriter { fetch: RangeFetch, - slot: &'a kio::Producer, + slot: kio::Producer, /// The group being written, with its producer unless its objects are read past. current: Option<(u64, Option)>, /// The first group not opened yet. next: u64, } -impl<'a> RangeWriter<'a> { - fn new(fetch: RangeFetch, slot: &'a kio::Producer) -> Self { +impl RangeWriter { + fn new(fetch: RangeFetch, slot: kio::Producer) -> Self { Self { next: fetch.first, fetch, @@ -575,18 +574,19 @@ impl<'a> RangeWriter<'a> { /// Finish the current group: its objects end where the next group's begin. The slot /// stays locked meanwhile, so an abandoning request never aborts a group as it finishes. fn finish(&mut self) -> Result<(), Error> { - let mut state = self.slot.write().ok(); + let Self { slot, current, .. } = self; + let mut state = slot.write().ok(); if let Some(state) = state.as_mut() && matches!(**state, GroupFetch::Receiving(_)) { **state = GroupFetch::Receiving(None); } - self.close_group() + Self::close_group(current) } /// Finish the current group, with the slot already settled. - fn close_group(&mut self) -> Result<(), Error> { - match self.current.take() { + fn close_group(current: &mut Option<(u64, Option)>) -> Result<(), Error> { + match current.take() { Some((_, Some(producer))) => producer.finish(), _ => Ok(()), } @@ -595,11 +595,11 @@ impl<'a> RangeWriter<'a> { /// Settle the run once its stream ends: the groups after the last one written do not /// exist, except that the last group bounded by an object may still hold later objects. /// A failed stream aborts the group it was writing and refuses the rest for now. - fn settle(mut self, res: Result<(), Error>) -> Result<(), Error> { + fn settle(&mut self, res: Result<(), Error>) -> Result<(), Error> { let fill = self.fetch.fill.clone(); let end = self.fetch.last + 1; let res = match res { - Ok(()) => self.close_group(), + Ok(()) => Self::close_group(&mut self.current), Err(err) => { if let Some((_, Some(producer))) = self.current.take() { let _ = producer.abort(err.clone()); @@ -968,7 +968,7 @@ struct Held { /// A joining request may be abandoned before it takes an ID. Once committed, the session /// owns ending it on the wire. struct JoiningTask { - work: Option>>, + work: Option, state: kio::Shared, tasks: Tasks, } @@ -985,7 +985,7 @@ impl JoiningTask { let Some(work) = self.work.as_mut() else { return Poll::Ready(Ok(())); }; - let res = ready!(waiter.poll_future(work.as_mut())); + let res = ready!(work.poll(waiter)); self.work = None; Poll::Ready(res) } @@ -998,13 +998,12 @@ impl Drop for JoiningTask { state.cancelled = true; state.committed }; - if committed && let Some(work) = self.work.take() { + if committed && let Some(mut work) = self.work.take() { // The session owns cancelling the request, including after the subscription // ends: closing it once its header is out, or resetting it while the peer's // send window still holds the header back. - self.tasks.push(async move { - let _ = work.await; - }); + self.tasks + .push(move |waiter: &kio::Waiter| work.poll(waiter).map(|_| ())); } } } @@ -1070,28 +1069,30 @@ struct NamespaceRequest { live: HashMap, } -/// The initial set of a SUBSCRIBE_NAMESPACE still to read, under MoQ Active Count. -struct Initial { - /// NAMESPACE messages still owed before the first change. - initial: u64, - /// Released once `initial` reaches zero. - snapshot: Option, -} - /// Resolve the subscription a data stream belongs to. /// /// SUBSCRIBE_OK can be reordered behind the stream it describes, so an alias we have not /// seen is worth waiting on briefly (draft-22 section 9.9). Three outcomes: /// the subscription, [`Error::Cancel`] for an alias we retired, and [`Error::NotFound`] /// once the wait expires without any binding at all. -async fn resolve_track_alias( - runtime: &crate::time::Deadlines, +struct ResolveAlias { aliases: kio::Consumer, alias: u64, -) -> Result { - let mut timeout = crate::time::Deadline::after(runtime, TRACK_ALIAS_TIMEOUT); - kio::wait(|waiter| { - let resolved = aliases.poll(waiter, |aliases| match aliases.map.get(&alias) { + timeout: crate::time::Deadline, +} + +impl ResolveAlias { + fn new(runtime: &crate::time::Deadlines, aliases: kio::Consumer, alias: u64) -> Self { + Self { + aliases, + alias, + timeout: crate::time::Deadline::after(runtime, TRACK_ALIAS_TIMEOUT), + } + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + let alias = self.alias; + let resolved = self.aliases.poll(waiter, |aliases| match aliases.map.get(&alias) { Some(Alias::Active(request_id)) => Poll::Ready(Ok(*request_id)), // A subscription we already cancelled, whose publisher has not caught up with // our STOP_SENDING. Discard the group now rather than waiting out the timeout @@ -1105,12 +1106,22 @@ async fn resolve_track_alias( if let Poll::Ready(result) = resolved { return Poll::Ready(result.unwrap_or(Err(Error::Cancel))); } - if timeout.poll(waiter).is_ready() { + if self.timeout.poll(waiter).is_ready() { return Poll::Ready(Err(Error::NotFound)); } Poll::Pending - }) - .await + } +} + +/// Await [`ResolveAlias`]. +#[cfg(test)] +async fn resolve_track_alias( + runtime: &crate::time::Deadlines, + aliases: kio::Consumer, + alias: u64, +) -> Result { + let mut resolve = ResolveAlias::new(runtime, aliases, alias); + kio::wait(|waiter| resolve.poll(waiter)).await } impl Subscriber { @@ -1158,7 +1169,9 @@ impl Subscriber { pub fn with_auth(mut self, auth: crate::auth::Link) -> Self { self.auth = auth; let this = self.clone(); - self.tasks.push(async move { this.run_limit().await }); + let mut epoch = 0; + self.tasks + .push(move |waiter: &kio::Waiter| this.poll_limit(&mut epoch, waiter)); self } @@ -1172,14 +1185,9 @@ impl Subscriber { /// Follow the session's limit, attaching every withheld namespace a new limit /// covers. Each route's own serve task holds itself back when a limit no longer /// covers it. - async fn run_limit(&self) { - let mut epoch = 0; + fn poll_limit(&self, epoch: &mut u64, waiter: &kio::Waiter) -> Poll<()> { loop { - let permit = kio::wait(|waiter| { - self.auth - .poll_permit(crate::auth::Direction::Subscribe, &mut epoch, waiter) - }) - .await; + let permit = ready!(self.auth.poll_permit(crate::auth::Direction::Subscribe, epoch, waiter)); let mut state = self.state.lock(); let mut attached = Vec::new(); for (path, entry) in state.broadcasts.iter_mut() { @@ -1198,13 +1206,10 @@ impl Subscriber { /// Serve the requests beneath one attached namespace on its own task. fn serve_route(&self, path: Path, generation: u64) { - let this = self.clone(); - self.tasks.push(async move { - // stop_announce is the authoritative remover: it drops the entry - // (retracting the route) once the announce refcount hits zero, - // which is what makes run_route exit, as does a limit holding it back. - this.run_route(path, generation).await; - }); + // stop_announce is the authoritative remover: it drops the entry (retracting the + // route) once the announce refcount hits zero, which is what ends the route, as does + // a limit holding it back. + self.tasks.push(RunRoute::new(self.clone(), path, generation)); } /// End every active subscription with the error that ended the session. @@ -1253,8 +1258,9 @@ impl Subscriber { } /// What the peer declared in its SETUP. See [`super::Publisher::peer`]. - pub(super) async fn peer(&self) -> cluster::Peer { - self.peer_setup.get().await.cluster + #[cfg_attr(feature = "async", kio::wrapper)] + pub(super) fn poll_peer(&self, waiter: &kio::Waiter) -> Poll { + self.peer_setup.poll_get(waiter).map(|peer| peer.cluster) } /// The announcing session's declared or assigned identity, for split-horizon. @@ -1393,19 +1399,8 @@ impl Subscriber { } /// Send SUBSCRIBE_NAMESPACE for one prefix on a bidi stream, and read what it reports - /// until the peer ends it or `live` closes. - /// The caller opens the stream, one per prefix. - /// - /// A failure here is per-prefix, so the caller decides what it means for the - /// session: [`is_protocol_violation`] separates the peer's fault (fatal) from a - /// stream of ours that simply died (survivable). - /// - /// `snapshot` holds the origin's announce consumers under the prefix short of - /// synced. It drops once the NAMESPACE messages MoQ Active Count promised are - /// in, on REQUEST_OK without the extension, or when no request goes out at all. - /// - /// [`Ended::Local`](namespaces::Ended::Local) means `live` closed, and the request is - /// off the wire by the time this returns. + /// until the peer ends it or `live` closes: [`NamespaceSubscription`] awaited. + #[cfg(test)] async fn run_subscribe_namespace( &mut self, stream: Stream, @@ -1413,172 +1408,165 @@ impl Subscriber { snapshot: crate::model::Snapshot, live: kio::Consumer<()>, ) -> Result { - let mut request = NamespaceRequest { - stream: stream.with_stats(self.origin.stats()), - messages: super::request_stream::Messages::new(self.session.clone(), self.control.client()), - prefix, - request_id: None, - live: HashMap::new(), - }; - let res = self.subscribe_namespace(&mut request, snapshot, &live).await; - - // The stream owns every advertisement it carried, so release them however it - // ends: a clean close, a decode error, the peer resetting it, or the session - // ending it. Without this each namespace keeps its refcount and the source never - // detaches. - for (path, advert) in std::mem::take(&mut request.live) { - if let Some(id) = advert.id { - let _ = self.stop_announce(path, id, false); - } - } - - if let Ok(namespaces::Ended::Local) = res { - self.unsubscribe_namespace(request).await; - } - res + let mut request = NamespaceSubscription::new(self.clone(), stream, prefix, snapshot, live); + kio::wait(|waiter| request.poll(waiter)).await } /// Send the request and read its answers until either side ends it. - async fn subscribe_namespace( + fn poll_subscribe_namespace( &mut self, request: &mut NamespaceRequest, - snapshot: crate::model::Snapshot, + phase: &mut NamespacePhase, + snapshot: &mut Option, live: &kio::Consumer<()>, - ) -> Result { - // Hidden namespaces are requested too, as on moq-lite: the session mirrors the - // peer into the origin and each local reader opts in on its own. The parameter - // fails decoding at a peer that doesn't know it, so it waits on the peer's SETUP - // to say whether it does (MoQ Hidden). - let declared = self.peer_setup.get().await; - let hidden = declared.hidden; - let prefix = request.prefix.clone(); - - let request_id = self.control.next_request_id(); - request.request_id = Some(request_id); - let stream = &mut request.stream; - - let msg = ietf::SubscribeNamespace { - authorization_token: None, - request_id, - namespace: prefix.clone(), - hidden, - }; - // Flushed while waiting for the answer, so `live` closing still ends a request the - // peer's send window holds back. - stream.writer.buffer_message(ietf::SubscribeNamespace::ID, &msg)?; + waiter: &kio::Waiter, + ) -> Poll> { + loop { + match phase { + NamespacePhase::Setup => { + // Hidden namespaces are requested too, as on moq-lite: the session mirrors the + // peer into the origin and each local reader opts in on its own. The parameter + // fails decoding at a peer that doesn't know it, so it waits on the peer's SETUP + // to say whether it does (MoQ Hidden). + let declared = ready!(self.peer_setup.poll_get(waiter)); + let hidden = declared.hidden; + let prefix = request.prefix.clone(); + + let request_id = self.control.next_request_id(); + request.request_id = Some(request_id); + + let msg = ietf::SubscribeNamespace { + authorization_token: None, + request_id, + namespace: prefix.clone(), + hidden, + }; + // Flushed while waiting for the answer, so `live` closing still ends a request the + // peer's send window holds back. + request + .stream + .writer + .buffer_message(ietf::SubscribeNamespace::ID, &msg)?; + + tracing::debug!(%prefix, "subscribe_namespace started"); + *phase = NamespacePhase::Answer { declared }; + } + NamespacePhase::Answer { declared } => { + let declared = *declared; + // Read response + let Some(framed) = ready!(Self::poll_namespace_message(request, live, waiter))? else { + return Poll::Ready(Ok(namespaces::Ended::Local)); + }; + let super::request_stream::Framed(type_id, body) = + framed.ok_or_else(|| request.messages.failed(Error::ProtocolViolation))?; + let mut data = body.decoder(self.version); + + let count = match type_id { + ietf::RequestOk::ID => ietf::RequestOk::decode_msg(&mut data, self.version)?.active, + ietf::RequestError::ID => { + let msg = ietf::RequestError::decode_msg(&mut data, self.version)?; + let err = + request::from_code(msg.error_code, request::Kind::SubscribeNamespace, self.version); + tracing::warn!(%err, reason = %msg.reason_phrase, "subscribe_namespace error"); + return Poll::Ready(Err(err)); + } + _ => return Poll::Ready(Err(Error::UnexpectedMessage)), + }; - tracing::debug!(%prefix, "subscribe_namespace started"); + tracing::debug!(prefix = %request.prefix, ?count, "subscribe_namespace ok"); - // Read response - let Some(framed) = self.next_namespace_message(request, live).await? else { - return Ok(namespaces::Ended::Local); - }; - let super::request_stream::Framed(type_id, body) = - framed.ok_or_else(|| request.messages.failed(Error::ProtocolViolation))?; - let mut data = body.decoder(self.version); + // MoQ Active Count says how many NAMESPACE messages make up the initial set. A + // count the negotiation did not promise, or a missing one it did, is the peer + // breaking the extension. Without it, the answer is the best boundary there is. + if declared.active_count != count.is_some() { + return Poll::Ready(Err(Error::ProtocolViolation)); + } + let initial = count.unwrap_or(0); + if initial == 0 { + *snapshot = None; + } - let count = match type_id { - ietf::RequestOk::ID => ietf::RequestOk::decode_msg(&mut data, self.version)?.active, - ietf::RequestError::ID => { - let msg = ietf::RequestError::decode_msg(&mut data, self.version)?; - let err = request::from_code(msg.error_code, request::Kind::SubscribeNamespace, self.version); - tracing::warn!(%err, reason = %msg.reason_phrase, "subscribe_namespace error"); - return Err(err); + // The extensions change the NAMESPACE encoding, so we can't parse one until + // the peer's SETUP says whether it negotiated them, which it has by now. + let peer = declared.cluster; + *phase = NamespacePhase::Entries(Entries { + negotiated: ietf::Negotiated { + cluster: peer.negotiated(), + }, + peer, + initial, + // Each entry routes into the origin, so a peer streaming them must not hold + // the session's other tasks. + budget: kio::coop::Budget::new(32), + }); + } + NamespacePhase::Entries(entries) => { + return self.poll_namespace_entries(request, entries, snapshot, live, waiter); + } } - _ => return Err(Error::UnexpectedMessage), - }; - - tracing::debug!(%prefix, ?count, "subscribe_namespace ok"); - - // MoQ Active Count says how many NAMESPACE messages make up the initial set. A - // count the negotiation did not promise, or a missing one it did, is the peer - // breaking the extension. Without it, the answer is the best boundary there is. - if declared.active_count != count.is_some() { - return Err(Error::ProtocolViolation); } - let initial = count.unwrap_or(0); - let snapshot = (initial > 0).then_some(snapshot); - - // The extensions change the NAMESPACE encoding, so we can't parse one until - // the peer's SETUP says whether it negotiated them. - let peer = self.peer().await; - let negotiated = ietf::Negotiated { - cluster: peer.negotiated(), - }; - - self.run_namespace_entries(request, &peer, negotiated, Initial { initial, snapshot }, live) - .await } /// The next message on a SUBSCRIBE_NAMESPACE stream: `None` once `live` closes, /// `Some(None)` once the peer finishes the stream. - async fn next_namespace_message( - &self, + fn poll_namespace_message( request: &mut NamespaceRequest, live: &kio::Consumer<()>, - ) -> Result>, Error> { + waiter: &kio::Waiter, + ) -> Poll>, Error>> { let NamespaceRequest { stream, messages, .. } = request; - kio::wait(|waiter| { - if live.poll_closed(waiter).is_ready() { - return Poll::Ready(Ok(None)); - } - ready!(stream.writer.poll_flush(waiter))?; - messages - .poll_next(&mut stream.reader, waiter) - .map(|next| next.map(Some)) - }) - .await + if live.poll_closed(waiter).is_ready() { + return Poll::Ready(Ok(None)); + } + ready!(stream.writer.poll_flush(waiter))?; + messages + .poll_next(&mut stream.reader, waiter) + .map(|next| next.map(Some)) } - /// End a request the session ended, and wait for it to leave the wire, so an + /// End a request the session ended, returning the wait for it to leave the wire, so an /// overlapping one can go out: a peer may refuse a prefix that overlaps one it still - /// holds. + /// holds. `None` when there is nothing to wait on. /// /// The stream closes (a FIN on draft-18, a reset on draft-22) and is read until the peer /// closes its half. - async fn unsubscribe_namespace(&self, request: NamespaceRequest) { + fn unsubscribe_namespace(&self, request: NamespaceRequest) -> Option { let NamespaceRequest { stream: Stream { mut writer, mut reader }, - mut messages, + messages, prefix, request_id, .. } = request; // Nothing went out, so there is nothing to end. - if request_id.is_none() { - return; - } + request_id?; // The request never reached the peer whole, so it holds nothing to wait on. if !writer.is_flushed() { tracing::debug!(%prefix, "subscribe_namespace abandoned before it was sent"); reader.abort(&Error::Cancel); writer.abort(&Error::Cancel); - return; + return None; } tracing::debug!(%prefix, "unsubscribe_namespace"); - match super::request_stream::fin_cancels(self.version) { + let writer = match super::request_stream::fin_cancels(self.version) { true => { let _ = writer.finish(); + Some(writer) } - false => writer.abort(&Error::Cancel), - } - let session = self.session.clone(); - // A peer still streaming entries must not hold the session's other tasks. - let mut budget = kio::coop::Budget::new(32); - kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() { - return Poll::Ready(()); - } - loop { - ready!(budget.poll_yield(waiter)); - match ready!(messages.poll_next(&mut reader, waiter)) { - Ok(Some(_)) => continue, - Ok(None) | Err(_) => return Poll::Ready(()), - } + false => { + writer.abort(&Error::Cancel); + None } + }; + Some(Unsubscribe { + session: self.session.clone(), + _writer: writer, + reader, + messages, + // A peer still streaming entries must not hold the session's other tasks. + budget: kio::coop::Budget::new(32), }) - .await; } /// Read NAMESPACE / NAMESPACE_DONE entries until either side ends the request. @@ -1586,30 +1574,34 @@ impl Subscriber { /// `request.live` tracks the suffixes this stream has advertised, so a repeat is /// recognized as an update rather than a second advertisement, and the caller can /// release whatever is still held when the stream ends. - async fn run_namespace_entries( + fn poll_namespace_entries( &mut self, request: &mut NamespaceRequest, - peer: &cluster::Peer, - negotiated: ietf::Negotiated, - mut initial: Initial, + entries: &mut Entries, + snapshot: &mut Option, live: &kio::Consumer<()>, - ) -> Result { - // Each entry routes into the origin, so a peer streaming them must not hold the - // session's other tasks. - let mut budget = kio::coop::Budget::new(32); + waiter: &kio::Waiter, + ) -> Poll> { + let Entries { + peer, + negotiated, + initial, + budget, + } = entries; loop { - kio::wait(|waiter| budget.poll_yield(waiter)).await; + ready!(budget.poll_yield(waiter)); // Every counted NAMESPACE is in the origin by now, dropped ones included. - if initial.initial == 0 { - initial.snapshot = None; - } - let super::request_stream::Framed(type_id, body) = match self.next_namespace_message(request, live).await? { - Some(Some(message)) => message, - None => return Ok(namespaces::Ended::Local), - // A stream ending inside the counted set cut it short. - Some(None) if initial.initial > 0 => return Err(Error::ProtocolViolation), - Some(None) => return Ok(namespaces::Ended::Peer), - }; + if *initial == 0 { + *snapshot = None; + } + let super::request_stream::Framed(type_id, body) = + match ready!(Self::poll_namespace_message(request, live, waiter))? { + Some(Some(message)) => message, + None => return Poll::Ready(Ok(namespaces::Ended::Local)), + // A stream ending inside the counted set cut it short. + Some(None) if *initial > 0 => return Poll::Ready(Err(Error::ProtocolViolation)), + Some(None) => return Poll::Ready(Ok(namespaces::Ended::Peer)), + }; let prefix = &request.prefix; let live = &mut request.live; let mut data = body.decoder(self.version); @@ -1619,11 +1611,11 @@ impl Subscriber { // relative to our root -- so the join is too, which is what everything // below wants (`create_broadcast` joins the root itself). ietf::Namespace::ID => { - let msg = ietf::Namespace::decode_body(&mut data, self.version, negotiated)?; + let msg = ietf::Namespace::decode_body(&mut data, self.version, *negotiated)?; if !data.is_empty() { - return Err(Error::WrongSize); + return Poll::Ready(Err(Error::WrongSize)); } - initial.initial = initial.initial.saturating_sub(1); + *initial = initial.saturating_sub(1); let path = prefix.join(&msg.suffix)?; let params = msg.params.unwrap_or_default(); let held = match live.entry(path.clone()) { @@ -1665,13 +1657,13 @@ impl Subscriber { Err(Error::Unauthorized) => { tracing::debug!(%path, "namespace outside the subscribe scope; ignoring"); } - Err(err) => return Err(err), + Err(err) => return Poll::Ready(Err(err)), } } } } // Only NAMESPACE may come before the counted set is in. - ietf::NamespaceDone::ID if initial.initial > 0 => return Err(Error::ProtocolViolation), + ietf::NamespaceDone::ID if *initial > 0 => return Poll::Ready(Err(Error::ProtocolViolation)), ietf::NamespaceDone::ID => { let msg = ietf::NamespaceDone::decode_msg(&mut data, self.version)?; let path = prefix.join(&msg.suffix)?; @@ -1682,12 +1674,163 @@ impl Subscriber { } _ => { tracing::warn!(type_id, "unexpected message on subscribe_namespace stream"); - return Err(Error::UnexpectedMessage); + return Poll::Ready(Err(Error::UnexpectedMessage)); + } + } + } + } +} + +/// One SUBSCRIBE_NAMESPACE for one prefix on a bidi stream, read until the peer ends it or +/// `live` closes. The caller opens the stream, one per prefix. +/// +/// A failure here is per-prefix, so the caller decides what it means for the +/// session: [`is_protocol_violation`] separates the peer's fault (fatal) from a +/// stream of ours that simply died (survivable). +/// +/// `snapshot` holds the origin's announce consumers under the prefix short of +/// synced. It drops once the NAMESPACE messages MoQ Active Count promised are +/// in, on REQUEST_OK without the extension, or when no request goes out at all. +/// +/// [`Ended::Local`](namespaces::Ended::Local) means `live` closed, and the request is +/// off the wire by the time it is ready. +struct NamespaceSubscription { + this: Subscriber, + live: kio::Consumer<()>, + state: NamespaceSubscriptionState, +} + +enum NamespaceSubscriptionState { + /// Sending the request and reading what it reports. + Running { + request: NamespaceRequest, + phase: NamespacePhase, + snapshot: Option, + }, + /// The session ended it: waiting for it to leave the wire. + Leaving(Unsubscribe), + Done, +} + +/// How far a [`NamespaceSubscription`] has come. +enum NamespacePhase { + /// Waiting on the peer's SETUP, which says how to encode the request. + Setup, + /// The request is buffered, waiting on its answer. + Answer { declared: peer::Peer }, + /// Reading NAMESPACE and NAMESPACE_DONE. + Entries(Entries), +} + +/// What reading a SUBSCRIBE_NAMESPACE's entries needs. +struct Entries { + peer: cluster::Peer, + negotiated: ietf::Negotiated, + /// NAMESPACE messages still owed before the first change, under MoQ Active Count. + initial: u64, + budget: kio::coop::Budget, +} + +impl NamespaceSubscription { + fn new( + this: Subscriber, + stream: Stream, + prefix: Path, + snapshot: crate::model::Snapshot, + live: kio::Consumer<()>, + ) -> Self { + let request = NamespaceRequest { + stream: stream.with_stats(this.origin.stats()), + messages: super::request_stream::Messages::new(this.session.clone(), this.control.client()), + prefix, + request_id: None, + live: HashMap::new(), + }; + Self { + this, + live, + state: NamespaceSubscriptionState::Running { + request, + phase: NamespacePhase::Setup, + snapshot: Some(snapshot), + }, + } + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + loop { + match &mut self.state { + NamespaceSubscriptionState::Running { + request, + phase, + snapshot, + } => { + let res = ready!( + self.this + .poll_subscribe_namespace(request, phase, snapshot, &self.live, waiter) + ); + *snapshot = None; + + // The stream owns every advertisement it carried, so release them however it + // ends: a clean close, a decode error, the peer resetting it, or the session + // ending it. Without this each namespace keeps its refcount and the source never + // detaches. + for (path, advert) in std::mem::take(&mut request.live) { + if let Some(id) = advert.id { + let _ = self.this.stop_announce(path, id, false); + } + } + + let NamespaceSubscriptionState::Running { request, .. } = + std::mem::replace(&mut self.state, NamespaceSubscriptionState::Done) + else { + unreachable!("matched above"); + }; + if let Ok(namespaces::Ended::Local) = res + && let Some(leaving) = self.this.unsubscribe_namespace(request) + { + self.state = NamespaceSubscriptionState::Leaving(leaving); + continue; + } + return Poll::Ready(res); } + NamespaceSubscriptionState::Leaving(leaving) => { + ready!(leaving.poll(waiter)); + self.state = NamespaceSubscriptionState::Done; + return Poll::Ready(Ok(namespaces::Ended::Local)); + } + NamespaceSubscriptionState::Done => panic!("NamespaceSubscription polled after it finished"), + } + } + } +} + +/// A request the session ended, read until the peer closes its half or the session ends. +struct Unsubscribe { + session: Transport, + /// Kept until then, since dropping a finished writer resets it. + _writer: Option>, + reader: Reader, + messages: super::request_stream::Messages, + budget: kio::coop::Budget, +} + +impl Unsubscribe { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + if self.session.poll_closed(waiter).is_ready() { + return Poll::Ready(()); + } + loop { + ready!(self.budget.poll_yield(waiter)); + match ready!(self.messages.poll_next(&mut self.reader, waiter)) { + Ok(Some(_)) => continue, + Ok(None) | Err(_) => return Poll::Ready(()), } } } +} +impl Subscriber { /// Handle an incoming bidi stream dispatched by the session. /// /// `peer` and `declared` are what the peer declared in its SETUP, which the dispatcher @@ -1701,10 +1844,10 @@ impl Subscriber { stream: Stream, peer: cluster::Peer, declared: Option, - ) -> Result, Error> { - let mut this = self.clone(); + ) -> Result { + let this = self.clone(); let mut data = body.decoder(this.version); - let task = match id { + let task: crate::util::BoxTask = match id { ietf::Publish::ID => { let msg = match ietf::Publish::decode_msg(&mut data, this.version) { Err(err) if super::token::refuses_request(&err) => { @@ -1722,12 +1865,13 @@ impl Subscriber { tracing::debug!(message = ?msg, "received publish"); // Owned, since the task outlives the body it was read from. let msg = msg.into_owned(); - async move { - if let Err(err) = this.serve_publish(stream, msg).await { + let mut serve = tracks::ServePublish::new(this, stream, msg); + Box::new(move |waiter: &kio::Waiter| { + if let Err(err) = ready!(serve.poll(waiter)) { tracing::debug!(%err, "publish stream error"); } - } - .boxed() + Poll::Ready(()) + }) } ietf::PublishNamespace::ID => { // A negotiated session that omits HOP_PATH fails the decode here, which @@ -1760,22 +1904,24 @@ impl Subscriber { declared, gate, }; - async move { - if let Err(err) = this.serve_publish_namespace_stream(stream, incoming).await { + let session = this.session.clone(); + let mut serve = ServePublishNamespace::new(this, stream, incoming); + Box::new(move |waiter: &kio::Waiter| { + if let Err(err) = ready!(serve.poll(waiter)) { // An advertisement update is decoded here rather than in the // dispatcher, so nothing else would surface a malformed one. The // cluster draft requires closing the session on those; a stream // the peer simply reset is not the peer's fault. if is_protocol_violation(&err) { - this.session.close( + session.close( super::error::to_session_code(&SessionError::from(&err)), err.to_string().as_ref(), ); } tracing::debug!(%err, "publish_namespace stream error"); } - } - .boxed() + Poll::Ready(()) + }) } _ => { tracing::warn!(id, "unexpected bidi stream type for subscriber"); @@ -1792,14 +1938,14 @@ impl Subscriber { stream: Stream, kind: request::Kind, err: DecodeError, - ) -> BoxFuture<'static, ()> { - let version = self.version; - async move { - if let Err(err) = super::token::reject(stream, kind, err, version).await { + ) -> crate::util::BoxTask { + let mut reject = super::token::rejection(stream, kind, err, self.version); + Box::new(move |waiter: &kio::Waiter| { + if let Err(err) = ready!(reject.poll(waiter)) { tracing::debug!(%err, "token refusal failed"); } - } - .boxed() + Poll::Ready(()) + }) } /// What the peer declared about being solicited (MoQ Solicit). @@ -1807,8 +1953,9 @@ impl Subscriber { /// Read once by the dispatch loop and handed to each stream rather than awaited per /// stream: the slot is settled by the time streams are accepted, and a stream task /// that waited on it would park forever if it never were. - pub(super) async fn solicit(&self) -> Option { - self.peer_setup.get().await.solicit + #[cfg_attr(feature = "async", kio::wrapper)] + pub(super) fn poll_solicit(&self, waiter: &kio::Waiter) -> Poll> { + self.peer_setup.poll_get(waiter).map(|peer| peer.solicit) } /// Whether an incoming PUBLISH_NAMESPACE means the peer ignored our SETUP. @@ -1840,7 +1987,8 @@ impl Subscriber { msg.track_namespace.clone(), crate::auth::Direction::Subscribe, ); - self.serve_publish_namespace_stream( + let mut serve = ServePublishNamespace::new( + self.clone(), stream, IncomingNamespace { msg, @@ -1848,237 +1996,512 @@ impl Subscriber { declared, gate, }, - ) - .await + ); + kio::wait(|waiter| serve.poll(waiter)).await } - async fn serve_publish_namespace_stream( + /// [`PublishNamespaceUpdates`] awaited. + #[cfg(test)] + async fn run_publish_namespace_updates( &mut self, - mut stream: Stream, - incoming: IncomingNamespace, - ) -> Result<(), Error> { - let IncomingNamespace { - msg, - peer, - declared, - gate, - } = incoming; - if let Some(error) = super::namespace::refusal(&msg.track_namespace, self.version) { - return super::request_stream::reject(stream, error).await; + stream: &mut Stream, + path: &Path, + held: Held, + peer: cluster::Peer, + attached: &mut Option, + admission: Admission, + ) -> Result { + let mut updates = PublishNamespaceUpdates::new(self, held, peer, admission); + kio::wait(|waiter| updates.poll(waiter, self, stream, path, attached)).await + } + + /// Retract an advertisement's route ahead of a withdrawal that waits on the peer, so a peer + /// that stops reading cannot keep it routed. `attached` makes it once only. + fn release_announce(&mut self, path: &Path, attached: &mut Option, token: bool) -> Result<(), Error> { + if let Some(id) = attached.take() { + self.stop_announce(path.clone(), id, token)?; } - let path = msg.track_namespace.to_owned(); + Ok(()) + } - if self.unsolicited_is_a_violation(declared) { - tracing::warn!(%path, "unsolicited publish_namespace from a peer that implements MoQ Solicit"); - return Err(Error::ProtocolViolation); + /// Spend a peer's request ID, closing the session on one it already used. + fn spend(&self, id: RequestId) -> Result<(), Error> { + if let Err(err) = self.control.accept(id) { + if matches!(err, Error::GoingAway) { + return Err(err); + } + self.session.clone().close( + super::error::to_session_code(&crate::SessionError::from(&err)), + "invalid request ID", + ); + return Err(err); } + Ok(()) + } - // A path that already contains our own Hop ID looped back. Reject it rather - // than attaching a source we would then have to route around. - let Some(advert) = self.route(msg.params.cluster.as_ref(), &peer) else { - tracing::debug!(%path, "dropping reflected publish_namespace"); - self.write_error(&mut stream, &Error::Unroutable, "route loops through this relay") - .await?; - let _ = stream.writer.close().await; - return Ok(()); - }; + /// Buffer REQUEST_OK on the bidi stream. + fn buffer_ok(&self, stream: &mut Stream) -> Result<(), Error> { + stream.writer.buffer_message( + ietf::RequestOk::ID, + &ietf::RequestOk { + largest: None, + active: None, + }, + ) + } - // Our authorization of the peer (the limit) admits an announce without a token, and - // refuses one it does not cover: the peer learns at once rather than publishing into - // nothing, and re-announces once its limit widens. An AUTHORIZATION TOKEN on the - // announce goes to the application to verify for this announce alone (MoQ - // request-token), even where the limit covers it, and a verified one is exempt from the - // limit while it lasts. - let token = match &msg.authorization_token { - Some(token) if gate.token_decides() => Some(token.0.clone()), - _ => None, - }; - if token.is_none() && !self.auth.within_limit(crate::auth::Direction::Subscribe, path.as_str()) { - tracing::debug!(%path, "refusing publish_namespace outside the limit"); - let reason = crate::auth::unauthorized_reason(&path); - self.write_error(&mut stream, &Error::Unauthorized, &reason).await?; - let _ = stream.writer.close().await; - return Ok(()); - } - let existing = self.state.lock().broadcasts.get(&path).map(|entry| entry._slot.clone()); - let slot: Option = Some(match existing { - Some(slot) => slot, - None => std::sync::Arc::new(self.acquire_announce()?), - }); - let (admit, token_grant) = match token { - None => (Admit::Session(slot.clone()), None), - Some(token) => { - let verdict = gate.verify(token, slot.clone()); - // The session closing or the peer withdrawing the announce withdraws the token. - let session = self.session.clone(); - let version = self.version; - let verified = kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() - || super::request_stream::poll_withdrawn( - &mut stream, - super::request_stream::fin_cancels(version), - waiter, - ) - .is_ready() - { - return Poll::Ready(None); - } - verdict.poll_grant(waiter).map(Some) - }) - .await; - let Some(verified) = verified else { - return Ok(()); - }; - let admitted = crate::auth::RequestGrant::new( - &self.runtime, - verdict, - verified, - path.clone(), - crate::auth::RequestKind::PublishNamespace, - slot.clone(), - ); - match admitted { - Ok(grant) => (Admit::Token(slot.clone()), Some(grant)), - Err(refusal) => { - self.write_error(&mut stream, &refusal.error, &refusal.reason).await?; - let _ = stream.writer.close().await; - return Ok(()); + /// Buffer the REQUEST_ERROR refusing a PUBLISH_NAMESPACE on the bidi stream that carries it. + fn buffer_error(&self, stream: &mut Stream, err: &Error, reason: &str) -> Result<(), Error> { + super::request_stream::buffer_error( + &mut stream.writer, + request::Kind::PublishNamespace, + err, + reason, + self.version, + ) + } +} + +/// Whether the session's limit stopped covering `path`, for an announce it admitted. +fn poll_limited( + auth: &crate::auth::Link, + epoch: &mut u64, + path: &Path, + waiter: &kio::Waiter, +) -> Poll { + let mut excluded = false; + while let Poll::Ready(permit) = auth.poll_permit(crate::auth::Direction::Subscribe, epoch, waiter) { + excluded = !permit.within_limit(path.as_str()); + } + match excluded { + true => Poll::Ready(crate::auth::Refusal { + error: Error::Unauthorized, + reason: crate::auth::unauthorized_reason(path), + }), + false => Poll::Pending, + } +} + +/// Serves an incoming PUBLISH_NAMESPACE on its bidi stream, from admitting it until its +/// advertisement ends on the wire. +struct ServePublishNamespace { + this: Subscriber, + /// `None` once it is taken apart to end the request. + stream: Option>, + state: ServeNamespaceState, +} + +const SERVE_STREAM: &str = "the stream is only taken apart as the request ends"; + +enum ServeNamespaceState { + /// Not polled yet. + Start(IncomingNamespace), + /// Waiting on the verdict for the announce's own token. + Verifying { + live: Box, + verdict: crate::auth::RequestVerdict, + }, + /// Admitted, running under what admitted it until either ends. + Live(Box), + /// Ending with a closed stream, once what is buffered goes out when `flush` (failing if + /// it cannot). + Close { + flush: bool, + }, + /// Waiting for the peer to acknowledge the close, whatever comes of it. + Closing { + closing: Option>, + _reader: Reader, + }, + /// Delivering the last message, then ending with `result`. + Deliver { + deliver: super::request_stream::Deliver, + result: Result<(), Error>, + }, + Done, +} + +/// An admitted PUBLISH_NAMESPACE: everything it does runs under what admitted it, its own +/// token's grant or the session's limit, either of which ending drops it wherever it waits. +struct LiveAnnounce { + path: Path, + peer: cluster::Peer, + slot: Option, + /// Taken when the route attaches. + advert: Option, + /// What the peer advertised, taken when the updates start. + cluster: Option, + /// Its own token's grant, when one admitted it. + grant: Option, + /// Whether the session's limit covers it, as of `epoch`, when no token admitted it. + epoch: u64, + /// Says the route is up. + attached: Option, + /// Says the OK went out, which decides how the peer is told. + acked: bool, + step: LiveStep, +} + +#[expect(clippy::large_enum_variant, reason = "a machine's enum is its storage")] +enum LiveStep { + /// Attaching the route and answering. + Attach, + /// Flushing the answer. + Ack, + Updates(PublishNamespaceUpdates), + /// Over, either way. + Ended, +} + +impl ServePublishNamespace { + fn new(this: Subscriber, stream: Stream, incoming: IncomingNamespace) -> Self { + Self { + this, + stream: Some(stream), + state: ServeNamespaceState::Start(incoming), + } + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + loop { + match &mut self.state { + ServeNamespaceState::Start(_) => { + let ServeNamespaceState::Start(incoming) = + std::mem::replace(&mut self.state, ServeNamespaceState::Done) + else { + unreachable!("matched above"); + }; + self.state = self.start(incoming)?; + } + ServeNamespaceState::Verifying { verdict, .. } => { + let stream = self.stream.as_mut().expect(SERVE_STREAM); + // The session closing or the peer withdrawing the announce withdraws the token. + if self.this.session.poll_closed(waiter).is_ready() + || super::request_stream::poll_withdrawn( + stream, + super::request_stream::fin_cancels(self.this.version), + waiter, + ) + .is_ready() + { + self.state = ServeNamespaceState::Done; + return Poll::Ready(Ok(())); + } + let verified = ready!(verdict.poll_grant(waiter)); + let ServeNamespaceState::Verifying { mut live, verdict } = + std::mem::replace(&mut self.state, ServeNamespaceState::Done) + else { + unreachable!("matched above"); + }; + let admitted = crate::auth::RequestGrant::new( + &self.this.runtime, + verdict, + verified, + live.path.clone(), + crate::auth::RequestKind::PublishNamespace, + live.slot.clone(), + ); + match admitted { + Ok(grant) => { + live.grant = Some(grant); + self.state = ServeNamespaceState::Live(live); + } + Err(refusal) => { + self.this.buffer_error(stream, &refusal.error, &refusal.reason)?; + self.state = ServeNamespaceState::Close { flush: true }; + } } } - } - }; - let token = token_grant.is_some(); - - // Everything the advertisement does once admitted runs under what admitted it: its - // token's grant, or the session's limit, which never holds an announce back. Either - // ending at any await, the first answer included, drops it right there, and a - // publisher whose limit widens again re-announces. `attached` says the route is up, and - // `acked` that the OK went out, which decides how the peer is told. - let mut attached = None; - let mut acked = false; - let auth = self.auth.clone(); - let mut epoch = 0; - let mut limited = |waiter: &kio::Waiter| { - let mut excluded = false; - while let Poll::Ready(permit) = auth.poll_permit(crate::auth::Direction::Subscribe, &mut epoch, waiter) { - excluded = !permit.within_limit(path.as_str()); - } - match excluded { - true => Poll::Ready(crate::auth::Refusal { - error: Error::Unauthorized, - reason: crate::auth::unauthorized_reason(&path), - }), - false => Poll::Pending, - } - }; - let mut ended = |waiter: &kio::Waiter| match &token_grant { - Some(grant) => grant.poll_ended(waiter), - None => limited(waiter), - }; - let lived = crate::auth::until( - async { - match self.start_announce(path.clone(), advert, admit) { - Ok(id) => attached = Some(id), - Err(err) => return Ok(Withdraw::Reject(err.into())), + ServeNamespaceState::Live(live) => { + let stream = self.stream.as_mut().expect(SERVE_STREAM); + // Polled first, so the announce never acts on authority that already ended. + let ended = match &live.grant { + Some(grant) => grant.poll_ended(waiter), + None => poll_limited(&self.this.auth, &mut live.epoch, &live.path, waiter), + }; + let lived = match ended { + Poll::Ready(refusal) => Err(refusal), + Poll::Pending => Ok(ready!(live.poll(&mut self.this, stream, waiter))), + }; + // Whatever the body held goes with it, before the route does. + live.step = LiveStep::Ended; + + // The route goes before anything waits on the peer, so a peer that stops reading + // cannot keep it routed. A stream that detached mid-flight (a reflected update) has + // nothing left to release. + let token = live.grant.is_some(); + self.this.release_announce(&live.path, &mut live.attached, token)?; + + let withdraw = match lived { + Ok(withdraw) => withdraw?, + Err(refusal) => { + tracing::info!(path = %live.path, err = %refusal.error, reason = %refusal.reason, "announce no longer authorized"); + if !live.acked { + self.this.buffer_error(stream, &refusal.error, &refusal.reason)?; + self.state = ServeNamespaceState::Close { flush: true }; + continue; + } + // A lapsed token cancels the request, resetting what we send and stopping what we + // read, so the publisher knows to announce again with a fresh one (moq-transport + // sections 6.4.2.3 and 9.14). It takes any OK still buffered with it. + if super::error::token_expired(&refusal.error) { + let Stream { mut reader, writer } = self.stream.take().expect(SERVE_STREAM); + reader.stop(super::error::EXPIRED_AUTH_TOKEN); + writer.reset(super::error::EXPIRED_AUTH_TOKEN); + self.state = ServeNamespaceState::Done; + return Poll::Ready(Err(refusal.error)); + } + // End the advertisement, leaving the session up: an OK cut short by the grant + // ending is still buffered, and goes out ahead of the FIN. + self.state = ServeNamespaceState::Deliver { + deliver: super::request_stream::Deliver::new(), + result: Err(refusal.error), + }; + continue; + } + }; + self.state = match withdraw { + Withdraw::Done => { + self.state = ServeNamespaceState::Done; + return Poll::Ready(Ok(())); + } + Withdraw::Close => ServeNamespaceState::Close { flush: false }, + Withdraw::Reject(refusal) => { + self.this.buffer_error(stream, &refusal.error, &refusal.reason)?; + ServeNamespaceState::Close { flush: true } + } + // Refused, then the advertisement ends, leaving the session up. + Withdraw::Cancel(refusal) => { + self.this.buffer_error(stream, &refusal.error, &refusal.reason)?; + ServeNamespaceState::Deliver { + deliver: super::request_stream::Deliver::new(), + result: Ok(()), + } + } + }; } - // Buffered whole before its first await, so however this is dropped from here, the - // advertisement can still end on the wire. - acked = true; - self.write_ok(&mut stream).await?; - - // An endpoint updates an advertisement with REQUEST_UPDATE on the stream that - // already carries it, so keep reading until the stream closes. - self.run_publish_namespace_updates( - &mut stream, - &path, - Held { - cluster: msg.params.cluster, - }, - peer, - &mut attached, - Admission { - grant: token_grant.clone(), - slot: slot.clone(), - }, - ) - .await - }, - &mut ended, - ) - .await; - - // The route goes before anything waits on the peer, so a peer that stops reading - // cannot keep it routed. A stream that detached mid-flight (a reflected update) has - // nothing left to release. - self.release_announce(&path, &mut attached, token)?; - - let withdraw = match lived { - Ok(withdraw) => withdraw?, - Err(refusal) => { - tracing::info!(%path, err = %refusal.error, reason = %refusal.reason, "announce no longer authorized"); - if !acked { - self.write_error(&mut stream, &refusal.error, &refusal.reason).await?; - let _ = stream.writer.close().await; - return Ok(()); + ServeNamespaceState::Close { flush } => { + let stream = self.stream.as_mut().expect(SERVE_STREAM); + if *flush { + ready!(stream.writer.poll_flush(waiter))?; + } + let Stream { writer, reader } = self.stream.take().expect(SERVE_STREAM); + self.state = ServeNamespaceState::Closing { + closing: writer.closing().ok(), + _reader: reader, + }; + } + ServeNamespaceState::Closing { closing, .. } => { + if let Some(closing) = closing { + let _ = ready!(closing.poll(waiter)); + } + self.state = ServeNamespaceState::Done; + return Poll::Ready(Ok(())); } - // A lapsed token cancels the request, resetting what we send and stopping what we - // read, so the publisher knows to announce again with a fresh one (moq-transport - // sections 6.4.2.3 and 9.14). It takes any OK still buffered with it. - if super::error::token_expired(&refusal.error) { - stream.reader.stop(super::error::EXPIRED_AUTH_TOKEN); - stream.writer.reset(super::error::EXPIRED_AUTH_TOKEN); - return Err(refusal.error); + ServeNamespaceState::Deliver { deliver, .. } => { + let stream = self.stream.as_mut().expect(SERVE_STREAM); + let delivered = ready!(deliver.poll(waiter, &mut stream.writer)); + let ServeNamespaceState::Deliver { result, .. } = + std::mem::replace(&mut self.state, ServeNamespaceState::Done) + else { + unreachable!("matched above"); + }; + return Poll::Ready(delivered.and(result)); } - self.cancel_announce(&mut stream).await?; - return Err(refusal.error); + ServeNamespaceState::Done => panic!("ServePublishNamespace polled after it finished"), } + } + } + + /// Decide what the request is owed before anything waits. + fn start(&mut self, incoming: IncomingNamespace) -> Result { + let IncomingNamespace { + msg, + peer, + declared, + gate, + } = incoming; + let this = &mut self.this; + let stream = self.stream.as_mut().expect(SERVE_STREAM); + if let Some(error) = super::namespace::refusal(&msg.track_namespace, this.version) { + stream.writer.buffer_message(ietf::RequestError::ID, &error)?; + return Ok(ServeNamespaceState::Deliver { + deliver: super::request_stream::Deliver::new(), + result: Ok(()), + }); + } + let path = msg.track_namespace.to_owned(); + + if this.unsolicited_is_a_violation(declared) { + tracing::warn!(%path, "unsolicited publish_namespace from a peer that implements MoQ Solicit"); + return Err(Error::ProtocolViolation); + } + + // A path that already contains our own Hop ID looped back. Reject it rather + // than attaching a source we would then have to route around. + let Some(advert) = this.route(msg.params.cluster.as_ref(), &peer) else { + tracing::debug!(%path, "dropping reflected publish_namespace"); + this.buffer_error(stream, &Error::Unroutable, "route loops through this relay")?; + return Ok(ServeNamespaceState::Close { flush: true }); }; - match withdraw { - Withdraw::Done => {} - Withdraw::Close => { - let _ = stream.writer.close().await; - } - Withdraw::Reject(refusal) => { - self.write_error(&mut stream, &refusal.error, &refusal.reason).await?; - let _ = stream.writer.close().await; + + // Our authorization of the peer (the limit) admits an announce without a token, and + // refuses one it does not cover: the peer learns at once rather than publishing into + // nothing, and re-announces once its limit widens. An AUTHORIZATION TOKEN on the + // announce goes to the application to verify for this announce alone (MoQ + // request-token), even where the limit covers it, and a verified one is exempt from the + // limit while it lasts. + let token = match &msg.authorization_token { + Some(token) if gate.token_decides() => Some(token.0.clone()), + _ => None, + }; + if token.is_none() && !this.auth.within_limit(crate::auth::Direction::Subscribe, path.as_str()) { + tracing::debug!(%path, "refusing publish_namespace outside the limit"); + let reason = crate::auth::unauthorized_reason(&path); + this.buffer_error(stream, &Error::Unauthorized, &reason)?; + return Ok(ServeNamespaceState::Close { flush: true }); + } + let existing = this.state.lock().broadcasts.get(&path).map(|entry| entry._slot.clone()); + let slot: Option = Some(match existing { + Some(slot) => slot, + None => std::sync::Arc::new(this.acquire_announce()?), + }); + let live = Box::new(LiveAnnounce { + path, + peer, + slot, + advert: Some(advert), + cluster: msg.params.cluster, + grant: None, + epoch: 0, + attached: None, + acked: false, + step: LiveStep::Attach, + }); + Ok(match token { + None => ServeNamespaceState::Live(live), + Some(token) => { + let verdict = gate.verify(token, live.slot.clone()); + ServeNamespaceState::Verifying { live, verdict } } - Withdraw::Cancel(refusal) => { - self.write_error(&mut stream, &refusal.error, &refusal.reason).await?; - self.cancel_announce(&mut stream).await?; + }) + } +} + +impl LiveAnnounce { + /// Attach the route, answer, and follow the advertisement's updates until it ends, + /// ready with how. + fn poll( + &mut self, + this: &mut Subscriber, + stream: &mut Stream, + waiter: &kio::Waiter, + ) -> Poll> { + loop { + match &mut self.step { + LiveStep::Attach => { + let advert = self.advert.take().expect("the route attaches once"); + let admit = match self.grant { + Some(_) => Admit::Token(self.slot.clone()), + None => Admit::Session(self.slot.clone()), + }; + match this.start_announce(self.path.clone(), advert, admit) { + Ok(id) => self.attached = Some(id), + Err(err) => return Poll::Ready(Ok(Withdraw::Reject(err.into()))), + } + // Buffered whole before its first wait, so however this is dropped from here, the + // advertisement can still end on the wire. + self.acked = true; + this.buffer_ok(stream)?; + self.step = LiveStep::Ack; + } + LiveStep::Ack => { + ready!(stream.writer.poll_flush(waiter))?; + // An endpoint updates an advertisement with REQUEST_UPDATE on the stream that + // already carries it, so keep reading until the stream closes. + self.step = LiveStep::Updates(PublishNamespaceUpdates::new( + this, + Held { + cluster: self.cluster.take(), + }, + self.peer, + Admission { + grant: self.grant.clone(), + slot: self.slot.clone(), + }, + )); + } + LiveStep::Updates(updates) => { + return updates.poll(waiter, this, stream, &self.path, &mut self.attached); + } + LiveStep::Ended => panic!("LiveAnnounce polled after it ended"), } } - Ok(()) } +} - /// Read advertisement updates off a live PUBLISH_NAMESPACE stream until it closes, and - /// return how it ends, which the caller writes once the route is released. - /// - /// `held` is what the peer advertised, kept current because a REQUEST_UPDATE carries - /// only what changed. Each one is answered with REQUEST_OK, or REQUEST_ERROR and a - /// closed stream when it cannot be applied, which withdraws the advertisement - /// (moq-transport Section 9.5.1). - async fn run_publish_namespace_updates( +/// Reads advertisement updates off a live PUBLISH_NAMESPACE stream until it closes, ready +/// with how it ends, which the caller writes once the route is released. +/// +/// `held` is what the peer advertised, kept current because a REQUEST_UPDATE carries +/// only what changed. Each one is answered with REQUEST_OK, or REQUEST_ERROR and a +/// closed stream when it cannot be applied, which withdraws the advertisement +/// (moq-transport Section 9.5.1). +struct PublishNamespaceUpdates { + held: Option, + peer: cluster::Peer, + admission: Admission, + // A REQUEST_UPDATE carrying a fresh token refreshes the announce's request grant (MoQ + // request-token); its verify is raced against the stream in the loop's single wait + // (never a bare wait), so a peer ending the announce ends it whatever the acceptor + // does, and the caller's guard ends it when the old grant does. Messages behind a + // pending verdict wait in `renewal`, handled in order once it resolves. + renewal: request_update::Renewal<(u64, ietf::Body)>, + messages: super::request_stream::Messages, + // Draft-22 only: the peer finished its send direction, which does not withdraw the + // advertisement there, so stop reading and keep watching the writer. + finished: bool, + // Request IDs of updates queued behind a pending verdict, already spent. + spent: std::collections::HashSet, + // Each update re-routes in the origin and is answered, so a peer streaming them must + // not hold the session's other tasks. + budget: kio::coop::Budget, + step: UpdateStep, +} + +enum UpdateStep { + /// Taking the next update, queued or off the stream. + Next, + /// Flushing the REQUEST_OK that answered one. + Ack, + /// Our send direction closed on a draft where the peer's FIN does not withdraw the + /// advertisement: waiting for its close to settle. + Closed, +} + +impl PublishNamespaceUpdates { + fn new(this: &Subscriber, held: Held, peer: cluster::Peer, admission: Admission) -> Self { + Self { + held: held.cluster, + peer, + admission, + renewal: request_update::Renewal::default(), + messages: super::request_stream::Messages::new(this.session.clone(), this.control.client()), + finished: false, + spent: std::collections::HashSet::new(), + budget: kio::coop::Budget::new(32), + step: UpdateStep::Next, + } + } + + fn poll( &mut self, + waiter: &kio::Waiter, + this: &mut Subscriber, stream: &mut Stream, path: &Path, - held: Held, - peer: cluster::Peer, attached: &mut Option, - admission: Admission, - ) -> Result { - let token_grant = &admission.grant; - let Held { cluster: mut held } = held; - // A REQUEST_UPDATE carrying a fresh token refreshes the announce's request grant (MoQ - // request-token); its verify is raced against the stream in the loop's single wait - // (never a bare await), so a peer ending the announce ends it whatever the acceptor - // does, and the caller's guard ends it when the old grant does. Messages behind a - // pending verdict wait in `renewal`, handled in order once it resolves. - let mut renewal = request_update::Renewal::<(u64, ietf::Body)>::default(); - let mut messages = super::request_stream::Messages::new(self.session.clone(), self.control.client()); - // Draft-22 only: the peer finished its send direction, which does not withdraw the - // advertisement there, so stop reading and keep watching the writer. - let mut finished = false; - + ) -> Poll> { // What one turn of the loop resolved to. enum Turn { // A whole message off the stream, or one drained from the queue. @@ -2093,39 +2516,50 @@ impl Subscriber { Renewal(request_update::Verdict), } - // Request IDs of updates queued behind a pending verdict, already spent. - let mut spent = std::collections::HashSet::new(); - // Each update re-routes in the origin and is answered, so a peer streaming them must - // not hold the session's other tasks. - let mut budget = kio::coop::Budget::new(32); + let Self { + held, + peer, + admission, + renewal, + messages, + finished, + spent, + budget, + step, + } = self; + let token_grant = &admission.grant; loop { - kio::wait(|waiter| budget.poll_yield(waiter)).await; + match step { + UpdateStep::Next => {} + UpdateStep::Ack => { + ready!(stream.writer.poll_flush(waiter))?; + *step = UpdateStep::Next; + continue; + } + UpdateStep::Closed => { + ready!(stream.writer.poll_closed(waiter))?; + return Poll::Ready(Ok(Withdraw::Done)); + } + } + + ready!(budget.poll_yield(waiter)); // Handle one buffered update before reading the next, so answers stay in order. let turn = if let Some((id, body)) = renewal.next() { Turn::Message(id, body) + } else if let Poll::Ready(result) = stream.writer.poll_closed(waiter) { + Turn::WriterClosed(result) + } else if let Poll::Ready(renewed) = renewal.poll(waiter) { + Turn::Renewal(renewed) + } else if *finished { + return Poll::Pending; } else { - kio::wait(|waiter| { - if let Poll::Ready(result) = stream.writer.poll_closed(waiter) { - return Poll::Ready(Turn::WriterClosed(result)); - } - if let Poll::Ready(renewed) = renewal.poll(waiter) { - return Poll::Ready(Turn::Renewal(renewed)); - } - // Keep reading while a verdict is pending, so a withdrawal still ends the - // announce: a buffered update must not mask a later terminal. - if !finished { - return match messages.poll_next(&mut stream.reader, waiter) { - Poll::Ready(Ok(Some(super::request_stream::Framed(id, body)))) => { - Poll::Ready(Turn::Message(id, body)) - } - Poll::Ready(Ok(None)) => Poll::Ready(Turn::Fin), - Poll::Ready(Err(err)) => Poll::Ready(Turn::Failed(err)), - Poll::Pending => Poll::Pending, - }; - } - Poll::Pending - }) - .await + // Keep reading while a verdict is pending, so a withdrawal still ends the + // announce: a buffered update must not mask a later terminal. + match ready!(messages.poll_next(&mut stream.reader, waiter)) { + Ok(Some(super::request_stream::Framed(id, body))) => Turn::Message(id, body), + Ok(None) => Turn::Fin, + Err(err) => Turn::Failed(err), + } }; let (type_id, body) = match turn { @@ -2134,75 +2568,77 @@ impl Subscriber { if id == ietf::PublishNamespaceUpdate::ID { // Its request ID is spent on arrival, so one dropped with the queue still moves // GOAWAY's cutoff past it. - let update_id = ietf::RequestId::decode(&mut body.decoder(self.version), self.version)?; - if let Err(err) = self.spend(update_id) { + let update_id = ietf::RequestId::decode(&mut body.decoder(this.version), this.version)?; + if let Err(err) = this.spend(update_id) { if matches!(err, Error::GoingAway) { - return Ok(Withdraw::Cancel(crate::auth::Refusal { + return Poll::Ready(Ok(Withdraw::Cancel(crate::auth::Refusal { error: err, reason: "going away".to_string(), - })); + }))); } - return Err(err); + return Poll::Ready(Err(err)); } spent.insert(update_id); // Decoded on arrival too, so a session-fatal update still closes the session // when a withdrawal drops the queue. A token refusal only refuses this // announce, which waits for the update's turn. - let mut data = body.decoder(self.version); - match ietf::PublishNamespaceUpdate::decode_msg(&mut data, self.version) { + let mut data = body.decoder(this.version); + match ietf::PublishNamespaceUpdate::decode_msg(&mut data, this.version) { Err(err) if super::token::refuses_request(&err) => {} msg => { msg?; } } if !data.is_empty() { - return Err(Error::WrongSize); + return Poll::Ready(Err(Error::WrongSize)); } } - match renewal.queue((id, body), self.version) { + match renewal.queue((id, body), this.version) { Ok(()) => continue, // The dispatcher closes only on is_protocol_violation, which excludes // Error::Session, so the session is closed here. Err(request_update::Overflow::Session) => { - return Err(request_update::too_many(&mut self.session)); + return Poll::Ready(Err(request_update::too_many(&mut this.session))); } // Drafts below 19 negotiate no limit, so this is a local memory guard, not a // protocol fault: end this announce, never the session. - Err(request_update::Overflow::Request) => return Ok(Withdraw::Close), + Err(request_update::Overflow::Request) => return Poll::Ready(Ok(Withdraw::Close)), } } Turn::Message(id, body) => (id, body), Turn::Fin => { - if super::request_stream::fin_cancels(self.version) { - return Ok(Withdraw::Done); + if super::request_stream::fin_cancels(this.version) { + return Poll::Ready(Ok(Withdraw::Done)); } // The advertisement outlives the peer's send direction. - finished = true; + *finished = true; continue; } // Already finished: the advertisement outlived the peer's send direction until now. - Turn::WriterClosed(result) if finished => return result.map(|()| Withdraw::Done), + Turn::WriterClosed(result) if *finished => return Poll::Ready(result.map(|()| Withdraw::Done)), Turn::WriterClosed(result) => { stream.reader.abort(&Error::Cancel); result?; - if !super::request_stream::fin_cancels(self.version) { - stream.writer.closed().await?; + if !super::request_stream::fin_cancels(this.version) { + *step = UpdateStep::Closed; + continue; } - return Ok(Withdraw::Done); + return Poll::Ready(Ok(Withdraw::Done)); } - Turn::Failed(err) => return Err(err), + Turn::Failed(err) => return Poll::Ready(Err(err)), Turn::Renewal(request_update::Verdict { verdict, grant: res }) => { // Only a token-authorized announce starts a renewal. let rg = token_grant.as_ref().expect("a renewal without a request grant"); // On accept the old grant is dropped and the deadline re-armed (REQUEST_OK). let Err(refusal) = rg.renew(verdict, res) else { - self.write_ok(stream).await?; + this.buffer_ok(stream)?; + *step = UpdateStep::Ack; continue; }; // Per draft-18 section 10.9.1, a refused or uncovered // renewal ends only this announce. Answer REQUEST_ERROR, then close the stream; // the session stays up and the old grant does not survive. - return Ok(Withdraw::Cancel(refusal)); + return Poll::Ready(Ok(Withdraw::Cancel(refusal))); } }; @@ -2210,45 +2646,45 @@ impl Subscriber { // A repeated PUBLISH_NAMESPACE lands here too: a second request on the // stream is the base draft's duplicate request ID. tracing::warn!(type_id, "unexpected message on publish_namespace stream"); - return Err(Error::UnexpectedMessage); + return Poll::Ready(Err(Error::UnexpectedMessage)); } - let mut data = body.decoder(self.version); + let mut data = body.decoder(this.version); // An update spends its own request ID, even one refused below. A queued one was // spent when it arrived. - let update_id = ietf::RequestId::decode(&mut body.decoder(self.version), self.version)?; + let update_id = ietf::RequestId::decode(&mut body.decoder(this.version), this.version)?; if !spent.remove(&update_id) - && let Err(err) = self.spend(update_id) + && let Err(err) = this.spend(update_id) { if matches!(err, Error::GoingAway) { - return Ok(Withdraw::Cancel(crate::auth::Refusal { + return Poll::Ready(Ok(Withdraw::Cancel(crate::auth::Refusal { error: err, reason: "going away".to_string(), - })); + }))); } - return Err(err); + return Poll::Ready(Err(err)); } - let msg = match ietf::PublishNamespaceUpdate::decode_msg(&mut data, self.version) { + let msg = match ietf::PublishNamespaceUpdate::decode_msg(&mut data, this.version) { // A token refusal rejects the update, which withdraws the advertisement as // any refused update does. Err(err) if super::token::refuses_request(&err) => { // Junk after the fields is fatal, which a token refusal must not hide. if !data.is_empty() { - return Err(Error::WrongSize); + return Poll::Ready(Err(Error::WrongSize)); } let refusal = crate::auth::Refusal { reason: super::token::reason(&err).to_string(), error: Error::Decode(err), }; - return Ok(Withdraw::Cancel(refusal)); + return Poll::Ready(Ok(Withdraw::Cancel(refusal))); } msg => msg?, }; // Junk inside the declared size would otherwise be applied silently, which // is the one decode path that skipped the check the others make. if !data.is_empty() { - return Err(Error::WrongSize); + return Poll::Ready(Err(Error::WrongSize)); } // Cluster parameters and a token renewal can ride the same update, so apply the @@ -2257,11 +2693,11 @@ impl Subscriber { // old one already serves and never splices the two. The parameters exist only on a // session that negotiated the extension; anywhere else they are the peer's violation. let carries_cluster = msg.hops.is_some() || msg.cost.is_some(); - held = match &held { + *held = match &*held { Some(current) => Some(msg.apply(current)), None if carries_cluster => { tracing::warn!(%path, "cluster parameters on a session that negotiated none"); - return Err(Error::ProtocolViolation); + return Poll::Ready(Err(Error::ProtocolViolation)); } None => None, }; @@ -2273,22 +2709,22 @@ impl Subscriber { // arrives here or nowhere. Ending the stream is not ours to do: a peer MAY // legitimately send a path carrying our Hop ID when a redundant sibling shares it. let applied: Result<(), Error> = match carries_cluster { - true => match self.route(held.as_ref(), &peer) { + true => match this.route(held.as_ref(), peer) { None => { if let Some(id) = attached.take() { tracing::debug!(%path, "publish_namespace now loops back; detaching"); - let _ = self.stop_announce(path.clone(), id, token_grant.is_some()); + let _ = this.stop_announce(path.clone(), id, token_grant.is_some()); } Ok(()) } Some(advert) => { tracing::debug!(%path, hops = advert.route.hops.len(), cost = ?advert.route.cost, "publish_namespace update"); match *attached { - Some(id) => self.update_announce(path.clone(), id, advert), + Some(id) => this.update_announce(path.clone(), id, advert), // Re-attach: a clean path replaced the reflected one we detached from. None => { // The grant still holds the slot it was admitted with, so reuse it. - self.start_announce(path.clone(), advert, admission.admit()) + this.start_announce(path.clone(), advert, admission.admit()) .map(|id| *attached = Some(id)) } } @@ -2301,129 +2737,38 @@ impl Subscriber { // The close is the withdrawal. if let Err(err) = applied { tracing::warn!(%path, %err, "publish_namespace update refused"); - return Ok(Withdraw::Reject(err.into())); + return Poll::Ready(Ok(Withdraw::Reject(err.into()))); } // A REQUEST_UPDATE carrying a fresh token refreshes the announce's request grant // (MoQ request-token), when the announce is token-authorized. The verify is not - // awaited here: it becomes the pending renewal raced against the deadline above. + // waited on here: it becomes the pending renewal raced against the deadline above. // The routing above is already applied, so the renewal's answer (written when its // verdict resolves) is this update's single response, cluster parameters included. if let Some(token) = msg .authorization_token - .filter(|_| self.auth.token_decides(crate::auth::Direction::Subscribe)) + .filter(|_| this.auth.token_decides(crate::auth::Direction::Subscribe)) { let Some(grant) = &token_grant else { // An announce is admitted by the session or by a token, and that never // changes, so a token here is refused rather than taken unread. - return Ok(Withdraw::Cancel(crate::auth::Refusal { + return Poll::Ready(Ok(Withdraw::Cancel(crate::auth::Refusal { error: Error::Unsupported, reason: "AUTHORIZATION TOKEN on an announce admitted without one".to_string(), - })); + }))); }; - renewal.verify(grant.verify_renewal(&self.auth, token.0)); + renewal.verify(grant.verify_renewal(&this.auth, token.0)); continue; } // No token rides this update: acknowledge it now. - self.write_ok(stream).await?; - } - } - - /// Retract an advertisement's route ahead of a withdrawal that waits on the peer, so a peer - /// that stops reading cannot keep it routed. `attached` makes it once only. - fn release_announce(&mut self, path: &Path, attached: &mut Option, token: bool) -> Result<(), Error> { - if let Some(id) = attached.take() { - self.stop_announce(path.clone(), id, token)?; - } - Ok(()) - } - - /// End an advertisement we accepted, leaving the session up: a closed stream. - async fn cancel_announce(&self, stream: &mut Stream) -> Result<(), Error> { - // An OK cut short by the grant ending is still buffered, and goes out ahead of the FIN. - super::request_stream::deliver(&mut stream.writer).await - } - - /// Reject an incoming PUBLISH. - /// - /// PUBLISH offers a single track, so honoring it means routing per - /// (namespace, track). Our model routes per namespace: a source attaches at a - /// path and serves every track under it, resolved on demand via SUBSCRIBE. - /// Accepting a PUBLISH would mean inventing a namespace-level source out of a - /// track-level offer, and that fiction then contradicts any real - /// PUBLISH_NAMESPACE for the same path. - /// - /// Declining the request rather than failing the session, since a peer using - /// a feature we don't implement is not a protocol violation. - async fn run_publish_stream( - &mut self, - mut stream: Stream, - msg: ietf::Publish<'static>, - ) -> Result<(), Error> { - if let Some(error) = super::namespace::refusal(&msg.track_namespace, self.version) { - return super::request_stream::reject(stream, error).await; - } - tracing::debug!(broadcast = %msg.track_namespace, track = %msg.track_name, "rejecting publish"); - - // We decline the method itself rather than this particular track, which would be - // UNINTERESTED. - // - // The alias the message carries is deliberately not recorded. Nothing will ever bind - // it, and a rejected request has no lifetime of ours to hang the cleanup on, so the - // entry would have to be swept asynchronously. Any data streams the publisher opened - // before reading this are dropped by the unknown-alias path instead. - let (kind, err) = (request::Kind::Publish, &Error::Unsupported); - super::request_stream::buffer_error(&mut stream.writer, kind, err, "PUBLISH is not supported", self.version)?; - // The rejection is the whole exchange, but it still has to arrive: a finish alone - // leaves the drop-time reset free to discard it before the peer acknowledges it. - super::request_stream::deliver(&mut stream.writer).await - } - - /// Send OK on the bidi stream. - /// Spend a peer's request ID, closing the session on one it already used. - fn spend(&self, id: RequestId) -> Result<(), Error> { - if let Err(err) = self.control.accept(id) { - if matches!(err, Error::GoingAway) { - return Err(err); - } - self.session.clone().close( - super::error::to_session_code(&crate::SessionError::from(&err)), - "invalid request ID", - ); - return Err(err); + this.buffer_ok(stream)?; + *step = UpdateStep::Ack; } - Ok(()) - } - - async fn write_ok(&self, stream: &mut Stream) -> Result<(), Error> { - stream - .writer - .encode_message(&ietf::RequestOk { - largest: None, - active: None, - }) - .await?; - Ok(()) - } - - /// Refuse a PUBLISH_NAMESPACE on the bidi stream that carries it. - async fn write_error( - &self, - stream: &mut Stream, - err: &Error, - reason: &str, - ) -> Result<(), Error> { - super::request_stream::write_error( - &mut stream.writer, - request::Kind::PublishNamespace, - err, - reason, - self.version, - ) - .await } +} +impl Subscriber { /// Attach the route for one newly advertised namespace, returning the id that names /// this advertisement. /// @@ -2634,134 +2979,15 @@ impl Subscriber { Ok(()) } - /// Run `tasks` to their own end, or until the session dies. + /// Ready once `tasks` run to their own end, or the session dies. /// /// A retraction does not disturb subscriptions already in flight: what ended /// takes no new work, but the work it started finishes. - async fn drain(&self, tasks: &mut TaskSet) { - let session = self.session.clone(); - kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() { - return Poll::Ready(()); - } - tasks.poll(waiter) - }) - .await - } - - /// Serve materialization requests for one announced namespace: mint a source - /// per requested path and serve its track requests until the route is - /// retracted or the session dies. Tracks in flight at a retraction run to - /// their own end, unless the announce was what authorized them. - async fn run_route(&self, path: Path, generation: u64) { - let mut broadcasts = TaskSet::owned(); - let closed_session = self.session.clone(); - let mut epoch = 0; - loop { - let next = broadcasts - .drive(|waiter| { - if closed_session.poll_closed(waiter).is_ready() { - return Poll::Ready(None); - } - // A draining peer usually stops publishing namespaces, so react - // to the GOAWAY itself; waiting for another message would leave - // the route primary until the session finally closed. - // Idempotent, since the signal stays set. - if self.going_away.poll(waiter).is_ready() { - self.drain_route(&path); - } - // A limit that no longer covers the namespace holds the route back as a - // retraction would, closing its sources. Tracks in flight end on their - // own gates, with `Unauthorized`. - let mut excluded = false; - while let Poll::Ready(permit) = - self.auth - .poll_permit(crate::auth::Direction::Subscribe, &mut epoch, waiter) - { - excluded = !permit.within_limit(path.as_str()); - } - // The route lives in the entry: stop_announce removing it retracts - // the route, and this loop ends with it. - let mut state = self.state.lock(); - let Some(entry) = state.broadcasts.get_mut(&path) else { - return Poll::Ready(None); - }; - if entry.generation != generation { - return Poll::Ready(None); - } - // A namespace its own request token authorized is exempt from the limit. - if excluded && entry.tokens == 0 && entry.dynamic.is_some() { - tracing::info!(route = %crate::path::Rooted(self.origin.root(), &path), "namespace no longer authorized"); - entry.dynamic = None; - // Dropping the namespace's token closes every source minted under it. - entry.live = Default::default(); - } - match &mut entry.dynamic { - Some(dynamic) => dynamic.poll_requested_broadcast(waiter).map(Some), - None => Poll::Ready(None), - } - }) - .await; - - let request = match next { - Some(Ok(request)) => request, - // Retracted or torn down: no request will ever arrive again, but - // the broadcasts already served keep their tracks in flight. - Some(Err(_)) | None => { - self.drain(&mut broadcasts).await; - break; - } - }; - - // The request path is absolute; the wire (and our origin handle) speak - // paths relative to the session's root. - let requested = match request.path().strip_prefix(self.origin.root()) { - Some(requested) => requested.clone(), - None => continue, - }; - // The source takes new tracks only while the peer still selects that instance. - let validity = { - let state = self.state.lock(); - let entry = state - .broadcasts - .get(&path) - .filter(|entry| entry.generation == generation); - entry - .and_then(|entry| entry.dynamic.as_ref()) - .map(|route| crate::model::broadcast::RequestValidity::for_source(route, request.path())) - .unwrap_or_else(|| { - crate::model::broadcast::RequestValidity::token(Some(kio::Producer::new(()).consume())) - }) - }; - let source = - crate::model::broadcast::SourceGuard::new(self.origin.create_source(requested.clone(), validity)); - // The handler exists before the requester sees the source, so a track it asks - // for before `run_broadcast` first polls queues instead of failing `NotFound`. - let dynamic = source.requests(); - request.serve(&*source); - - // The namespace's token closes the source as it is retracted. If it was - // retracted since the accept, close it here as that retraction would have, - // and still serve what it took on: tracks subscribed since carry on. - let route = self - .state - .lock() - .broadcasts - .get(&path) - .filter(|entry| entry.generation == generation) - .map(|entry| entry.live.consume()); - if route.is_none() { - source.close(); - } - - let mut this = self.clone(); - this.validity = dynamic.validity().clone(); - broadcasts.push(async move { - if let Err(err) = this.run_broadcast(requested, source, dynamic, route).await { - tracing::debug!(%err, "error running broadcast"); - } - }); + fn poll_drain(&self, tasks: &mut TaskSet, waiter: &kio::Waiter) -> Poll<()> { + if self.session.poll_closed(waiter).is_ready() { + return Poll::Ready(()); } + tasks.poll(waiter) } /// Re-price one attached route to a draining cost (the peer sent a GOAWAY): @@ -2784,75 +3010,17 @@ impl Subscriber { } } - /// Serve one minted source's track requests, taken through its `broadcast` handler, - /// until it closes: its namespace is retracted (`route` closes), the peer stops - /// selecting its instance, nothing holds it any more, or the session dies. + /// Serve one minted source's track requests; see [`RunBroadcast`]. + #[cfg(test)] async fn run_broadcast( &self, path: Path, source: crate::model::broadcast::SourceGuard, - mut broadcast: broadcast::Requests, + broadcast: broadcast::Requests, route: Option>, - ) -> Result<(), Error> { - let mut subscribes = TaskSet::owned(); - let closed_session = self.session.clone(); - // Answers the source's listings with the peer's tracks, while anyone lists it. - let lister = std::sync::Arc::new( - broadcast - .lister() - .expect("only the source answers its broadcast's listings"), - ); - let mut listening = None; - loop { - let next = subscribes - .drive(|waiter| { - if closed_session.poll_closed(waiter).is_ready() { - return Poll::Ready(None); - } - self.poll_listing(&path, &lister, &mut listening, waiter); - loop { - if let Poll::Ready(next) = broadcast.poll_requested_track(waiter) { - return Poll::Ready(Some(next)); - } - if route.as_ref().is_some_and(|route| route.poll_closed(waiter).is_ready()) { - source.close(); - continue; - } - // Nothing holds the source and no track is on its way: retire it, so - // what the session keeps for the path goes with the fronts that used - // it. Declined when a holder or a track got there first, so look again. - if source.poll_unheld(waiter).is_pending() { - return Poll::Pending; - } - source.close_unheld(); - } - }) - .await; - - let request = match next { - Some(Ok(request)) => request, - Some(Err(err)) => { - tracing::debug!(%err, "broadcast closed"); - drop(listening.take()); - // No new tracks, but those in flight run to their own end. - self.drain(&mut subscribes).await; - break; - } - // Session gone. - None => break, - }; - - let mut this = self.clone(); - - let path = path.clone(); - let broadcast = broadcast.clone(); - let route = route.clone(); - subscribes.push(async move { - this.run_subscribe(path, broadcast, route, request).await; - }); - } - - Ok(()) + ) { + let mut run = RunBroadcast::new(self.clone(), path, source, broadcast, route); + kio::wait(|waiter| kio::Task::poll(&mut run, waiter)).await } /// Answer the source's listings while anyone lists it, from a SUBSCRIBE_TRACKS covering @@ -2887,560 +3055,1227 @@ impl Subscriber { } } + /// Serve one track request; see [`RunSubscribe`]. + #[cfg(test)] async fn run_subscribe( &mut self, broadcast_path: Path, - // Held for the subscription's lifetime but never watched: the broadcast ending - // is a retraction, which does not disturb a subscription the scope allows. - _broadcast: broadcast::Requests, - // Open while the namespace stays announced and authorized, which its own token may - // have done in place of the scope. `None` once retracted. + broadcast: broadcast::Requests, route: Option>, request: track::Request, ) { - // Holds the session busy while the request waits on authorization or TRACK_STATUS. - // Past that, only a subscription or a fill counts. - let mut setup = Some(self.session.work().start()); - // Data streams wait on the alias bound by SUBSCRIBE_OK, so leave the model request - // pending until its immutable track metadata is known. - // Subscribe only to what our grant covers (MoQ Auth), and cancel once it no longer - // does, leaving the rest of the session alone. - let (auth, path, direction) = ( - self.auth.clone(), - broadcast_path.clone(), - crate::auth::Direction::Subscribe, - ); - let mut gate = match route { - Some(live) => crate::auth::Gate::announced(auth, path, direction, live), - None => crate::auth::Gate::new(auth, path, direction), - }; - // Nobody wanting the track any more ends a wait on a token that may never be answered. - loop { - let admitted = kio::wait(|waiter| { - if gate.poll_ready(waiter).is_ready() { - return Poll::Ready(Some(gate.poll_denied(waiter).is_pending())); - } - crate::util::poll_unwanted(&mut request.demand(), waiter).map(|_| None) - }) - .await; - match admitted { - Some(true) => break, - Some(false) => { - request.reject(Error::Unauthorized); - return; - } - None if request.reject_unwanted(Error::Cancel) => return, - // Demand returned in the gap. - None => continue, - } - } - - let track_name = request.name().to_owned(); - // Group FETCHes for cache misses: standalone, so they outlive each subscription. - let mut group_fetches = TaskSet::owned(); - - // Demand with nobody subscribing (fetches, or a TRACK_STATUS downstream) learns the - // track from TRACK_STATUS rather than a SUBSCRIBE: a finished track refuses one, and - // one made only to learn the track races the fetches it was for. So does a bounded - // subscription, which may read only history below the Largest Object the status - // reports, as a fetch would: `linger` subscribes once it reaches past that. - let mut target = match request.subscription() { - subscription - if subscription.as_ref().is_none_or(|subscription| { - subscription.ranges.is_empty() || subscription.start().is_some() && subscription.end().is_some() - }) => - { - let Some(idle) = self - .track_status(&broadcast_path, &track_name, request, &mut group_fetches, &mut gate) - .await - else { - return; - }; - setup = None; - let Some(next) = self - .linger(&broadcast_path, &track_name, idle, &mut group_fetches, &mut gate) - .await - else { - return; - }; - Target::Resume(next) - } - _ => Target::Request(request), - }; - // `subscribe_once` counts the subscription itself. - drop(setup); - loop { - let Some(idle) = self - .subscribe_once(&broadcast_path, &track_name, target, &mut group_fetches, &mut gate) - .await - else { - return; - }; - let Some(next) = self - .linger(&broadcast_path, &track_name, idle, &mut group_fetches, &mut gate) - .await - else { - return; - }; - target = Target::Resume(next); - } + let mut run = RunSubscribe::new(self.clone(), broadcast_path, broadcast, route, request); + kio::wait(|waiter| kio::Task::poll(&mut run, waiter)).await } - /// Subscribe upstream for `target`, until the subscription ends: `Some` once nobody - /// subscribes any more or the track ended, with the copy kept for [`Self::linger`], - /// `None` once it failed. + /// Subscribe upstream for `target`, until the subscription ends; see [`SubscribeOnce`]. + #[cfg(test)] async fn subscribe_once( &mut self, broadcast_path: &Path, track_name: &str, - mut target: Target, + target: Target, group_fetches: &mut TaskSet, gate: &mut crate::auth::Gate, ) -> Option { - let mut request_gate = gate.unnamed(); - let gate = &mut request_gate; - let mut token_grant = None; - let _working = self.session.work().start(); - let validity = self.validity.start(); - let mut resolved = ResolvedRanges::default(); - let subscription = target - .subscription() - .map(|subscription| match target.producer() { - Some(track) => resolved.resolve(track, &subscription), - None if needs_resolution(&subscription) => Err(Error::Unsupported), - None => Ok(subscription), - }) - .transpose(); - // A relative anchor the copy cannot settle yet counts back from the Largest Object, so - // the live join that names it goes first and the answer resolves the anchor. No Group ID - // arithmetic spells it upstream. A time anchor names no group the answer settles. - let deferred = matches!(subscription, Err(Error::Unsupported)) - && target - .subscription() - .is_some_and(|subscription| !has_time(&subscription)); - let subscription = match subscription { - Ok(subscription) => subscription, - Err(_) if deferred => None, - Err(err) => { - target.fail(err); - return None; - } - }; - let start = subscription.as_ref().and_then(|s| s.start()); - let join = subscribe_join(start, subscription.as_ref().and_then(|s| s.end()), self.version); - // A live join delivers nothing below the group SUBSCRIBE_OK names as Largest. Every - // draft-18 join is one, since its joining FETCH carries that group alone. - let live = start.is_none() || join.fetch; - - // A listed track's paused subscription is already this track's upstream copy, so - // resume it rather than open a second one. - let key: tracks::TrackKey = (broadcast_path.clone(), track_name.into()); - let adopted = match &self.want(&key) { - None => None, - Some(handoff) => { - // The handoff waits for the held stream to flush, which a peer withholding - // credit can stall, so nobody wanting the track any more ends the wait too. - enum Abandon { - Failed(Error), - Unused, - } - loop { - let adopt = kio::wait(|waiter| { - // Every wait here drives the fetches a lingering copy started. - let _ = group_fetches.poll(waiter); - if gate.poll_denied(waiter).is_ready() { - return Poll::Ready(Err(Abandon::Failed(Error::Unauthorized))); - } - if let Err(err) = validity.check(waiter) { - return Poll::Ready(Err(Abandon::Failed(err))); - } - if let Poll::Ready(adopted) = Self::poll_handoff(handoff, waiter) { - return Poll::Ready(Ok(adopted)); - } - let unwanted = match &mut target { - Target::Request(request) => crate::util::poll_unwanted(&mut request.demand(), waiter), - Target::Resume(idle) => crate::util::poll_unwanted(&mut idle.demand, waiter), - }; - unwanted.map(|_| Err(Abandon::Unused)) - }) - .await; - match adopt { - Ok(adopted) => break adopted, - Err(Abandon::Failed(err)) => { - self.unwant(key, handoff); - target.fail(err); - return None; - } - // Nobody wants the track: a resumed copy lingers again, a request is refused. - Err(Abandon::Unused) => match target { - Target::Resume(idle) => { - self.unwant(key, handoff); - return Some(idle); - } - Target::Request(ref request) if request.reject_unwanted(Error::Cancel) => { - self.unwant(key, handoff); - return None; - } - Target::Request(_) => continue, - }, - } - } - } - }; - - // The resumed subscription's alias, properties, and owed answers, when resuming one. - let mut held: Option = None; - let (request_id, mut stream, mut messages) = match adopted { - Some(upstream) => { - self.adopted(&key, upstream.request_id); - let tracks::Upstream { - token_grant: token, - stream, - messages, - request_id, - alias, - properties, - owed, - // This subscription forwards into the copy itself from here. - copy: _, - } = upstream; - if let Some(token) = &token { - gate.bind(token.clone()); - } - token_grant = token; - held = Some(tracks::Resumed { - alias, - properties, - owed, - ended: None, - }); - (request_id, stream, messages) - } - None => { - let opening = self.session.clone(); - let mut open = crate::transport::open::Open::default(); - let stream = match kio::wait(|waiter| { - let _ = group_fetches.poll(waiter); - if gate.poll_denied(waiter).is_ready() { - return Poll::Ready(Err(Error::Unauthorized)); - } - validity.check(waiter)?; - Stream::poll_open(&opening, &mut open, self.version, waiter) - }) - .await - { - Ok(stream) => stream.with_stats(self.origin.stats()), - Err(err) => { - target.fail(err); - return None; - } - }; + let mut once = SubscribeOnce::new(self, broadcast_path, track_name, target, gate); + kio::wait(|waiter| once.poll(waiter, self, broadcast_path, track_name, group_fetches)).await + } +} - let request_id = self.control.next_request_id(); - let messages = super::request_stream::Messages::new(self.session.clone(), self.control.client()); - (request_id, stream, messages) - } - }; +/// Serves materialization requests for one announced namespace: mints a source per +/// requested path and serves its track requests until the route is retracted or the +/// session dies. Tracks in flight at a retraction run to their own end, unless the +/// announce was what authorized them. +struct RunRoute { + this: Subscriber, + path: Path, + generation: u64, + broadcasts: TaskSet, + /// The auth limit's epoch last seen. + epoch: u64, + /// Retracted or torn down: no request will ever arrive again, but the broadcasts + /// already served keep their tracks in flight. + draining: bool, +} - // Register the request before writing SUBSCRIBE so SUBSCRIBE_OK can bind its alias, - // and so a fill fetch stream that overtakes it finds the subscription it answers. - let fill = kio::Producer::new(match join.fill.is_some() || join.fetch { - true => Fill::Requested, - false => Fill::Done, - }); +impl RunRoute { + fn new(this: Subscriber, path: Path, generation: u64) -> Self { + Self { + this, + path, + generation, + broadcasts: TaskSet::owned(), + epoch: 0, + draining: false, + } + } + + /// The next source requested under the route, or `None` once the route or the session + /// is gone. + fn poll_requested(&mut self, waiter: &kio::Waiter) -> Poll>> { + let this = &self.this; + let path = &self.path; + if this.session.poll_closed(waiter).is_ready() { + return Poll::Ready(None); + } + // A draining peer usually stops publishing namespaces, so react + // to the GOAWAY itself; waiting for another message would leave + // the route primary until the session finally closed. + // Idempotent, since the signal stays set. + if this.going_away.poll(waiter).is_ready() { + this.drain_route(path); + } + // A limit that no longer covers the namespace holds the route back as a + // retraction would, closing its sources. Tracks in flight end on their + // own gates, with `Unauthorized`. + let mut excluded = false; + while let Poll::Ready(permit) = + this.auth + .poll_permit(crate::auth::Direction::Subscribe, &mut self.epoch, waiter) { - let mut state = self.state.lock(); - // A resumed subscription keeps the alias it is bound to and the data streams it - // counted while it forwarded before. - let (alias, tail) = match held.is_some() { - true => state - .subscribes - .remove(&request_id) - .map(|held| (held.alias, held.tail)) - .unwrap_or_default(), - false => Default::default(), - }; - state.subscribes.insert( - request_id, - TrackState { - // A resumed copy is the session's to abort from here on. - producer: target.producer().cloned(), - alias, - tail, - ..TrackState::pending(track_name.to_owned(), broadcast_path.clone(), fill.clone()) - }, - ); + excluded = !permit.within_limit(path.as_str()); } - - // Write Subscribe message. The aggregate is read now: a subscriber can join while - // the request ID and stream were awaited. Updates follow its priority from here, but - // no draft lets one change the order. - let max_age_supported = self.peer_setup.get().await.max_age; - let aggregate = target.subscription(); - let priority = aggregate.as_ref().map(|s| s.priority).unwrap_or(0); - let max_age = - max_age_supported.then(|| super::max_age::millis(aggregate.as_ref().map_or(Duration::MAX, |s| s.budget()))); - let order = aggregate.map(|s| s.order).unwrap_or_default(); - // Without an Established subscription, a relay includes the request it holds. - let new_group = target.new_group(); - // A timed copy asks for the units that opt its objects into timestamps. - let properties_wanted = match &target { - Target::Request(_) => true, - Target::Resume(idle) => idle.properties.timescale.is_some(), - }; - // Buffered only: the wait for the answer flushes it, so the local terminals still end a - // request the peer's send window holds back. - let written = match &held { - // Resume the held subscription: its group order is the one it was made with, - // since no update can change it. Its groups wait on the answer from here. - Some(held) => { - self.rebind_alias(request_id, Alias::Resuming(request_id)); - let forward = tracks::Forward { - forward: true, - filter: Some(join.filter), - fill: join.fill, - priority: Some(priority), - max_age, - // Only a track declaring DYNAMIC_GROUPS takes one on an update. - new_group: new_group.filter(|_| held.properties.dynamic_groups), - }; - self.buffer_forward(&mut stream.writer, &forward) - } - None => stream.writer.buffer_message( - ietf::Subscribe::ID, - &ietf::Subscribe { - max_age, - new_group, - request_id, - track_namespace: broadcast_path.clone(), - track_name: track_name.into(), - subscriber_priority: super::priority::to_wire(priority), - group_order: order.into(), - filter: join.filter, - fill: join.fill, - properties_wanted, - forward: true, - rendezvous: false, - range_filters: false, - authorization_token: None, - }, - ), + // The route lives in the entry: stop_announce removing it retracts + // the route, and this loop ends with it. + let mut state = this.state.lock(); + let Some(entry) = state.broadcasts.get_mut(path) else { + return Poll::Ready(None); }; - if let Err(err) = written { - tracing::debug!(%err, "failed to encode subscribe"); - self.remove_subscribe(request_id); - self.entry_ended(&key, request_id, Err(err.clone())); - target.fail(err); - return None; + if entry.generation != self.generation { + return Poll::Ready(None); + } + // A namespace its own request token authorized is exempt from the limit. + if excluded && entry.tokens == 0 && entry.dynamic.is_some() { + tracing::info!(route = %crate::path::Rooted(this.origin.root(), path), "namespace no longer authorized"); + entry.dynamic = None; + // Dropping the namespace's token closes every source minted under it. + entry.live = Default::default(); + } + match &mut entry.dynamic { + Some(dynamic) => dynamic.poll_requested_broadcast(waiter).map(Some), + None => Poll::Ready(None), } + } +} - tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), broadcast_path), track = %track_name, "subscribe started"); +impl kio::Task for RunRoute { + type Output = (); - // Park the origin request where a session abort can reject it. The producer - // does not exist until SUBSCRIBE_OK, and dropping this task would otherwise - // end the track as `Dropped`. - let mut resumed = match target { - Target::Request(request) => { - let mut state = self.state.lock(); - let Some(held) = state.subscribes.get_mut(&request_id) else { - request.reject(Error::Cancel); - return None; - }; - held.pending = Some(request); - None + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + loop { + if self.draining { + return self.this.poll_drain(&mut self.broadcasts, waiter); } - Target::Resume(idle) => Some(idle), - }; + let request = match self.poll_requested(waiter) { + Poll::Ready(Some(Ok(request))) => request, + Poll::Ready(Some(Err(_)) | None) => { + self.draining = true; + continue; + } + // The children never end the wait; a `Ready` from them just means they're + // drained until the next submission. + Poll::Pending => { + let _ = self.broadcasts.poll(waiter); + return Poll::Pending; + } + }; - // A publisher can be serving before its SUBSCRIBE_OK reaches us, since the data - // streams are independent of the request stream. Waiting for the response alone would - // miss the local side going away in that window and leave the publisher serving a - // track nobody reads, which is the leak this whole path exists to close. The broadcast - // ending is not a local side going away: a retraction does not disturb subscriptions - // already in flight, and this one is. - enum Setup { - Response(Result, Error>), - Unused, - /// `abort` already rejected the parked request. - Gone, - } - - // A local terminal ended the wait, so the peer still holds the request to cancel. - let mut unanswered = false; - let setup = { - let mut response = std::pin::pin!(async { - kio::wait(|waiter| stream.writer.poll_flush(waiter)).await?; - self.read_answer(&mut stream, &mut messages, held.as_mut()).await - }); - loop { - let setup = kio::wait(|waiter| { - let _ = group_fetches.poll(waiter); - let response = waiter.poll_future(response.as_mut()); - if let Poll::Ready(Err(err)) = response { - return Poll::Ready(Setup::Response(Err(err))); + let this = &self.this; + let path = &self.path; + // The request path is absolute; the wire (and our origin handle) speak + // paths relative to the session's root. + let requested = match request.path().strip_prefix(this.origin.root()) { + Some(requested) => requested.clone(), + None => continue, + }; + // The source takes new tracks only while the peer still selects that instance. + let validity = { + let state = this.state.lock(); + let entry = state + .broadcasts + .get(path) + .filter(|entry| entry.generation == self.generation); + entry + .and_then(|entry| entry.dynamic.as_ref()) + .map(|route| crate::model::broadcast::RequestValidity::for_source(route, request.path())) + .unwrap_or_else(|| { + crate::model::broadcast::RequestValidity::token(Some(kio::Producer::new(()).consume())) + }) + }; + let source = + crate::model::broadcast::SourceGuard::new(this.origin.create_source(requested.clone(), validity)); + // The handler exists before the requester sees the source, so a track it asks + // for before `RunBroadcast` first polls queues instead of failing `NotFound`. + let dynamic = source.requests(); + request.serve(&*source); + + // The namespace's token closes the source as it is retracted. If it was + // retracted since the accept, close it here as that retraction would have, + // and still serve what it took on: tracks subscribed since carry on. + let route = this + .state + .lock() + .broadcasts + .get(path) + .filter(|entry| entry.generation == self.generation) + .map(|entry| entry.live.consume()); + if route.is_none() { + source.close(); + } + + let mut this = this.clone(); + this.validity = dynamic.validity().clone(); + self.broadcasts + .push(RunBroadcast::new(this, requested, source, dynamic, route)); + } + } +} + +/// Serves one minted source's track requests, taken through its `broadcast` handler, +/// until it closes: its namespace is retracted (`route` closes), the peer stops +/// selecting its instance, nothing holds it any more, or the session dies. +struct RunBroadcast { + this: Subscriber, + path: Path, + source: crate::model::broadcast::SourceGuard, + broadcast: broadcast::Requests, + route: Option>, + subscribes: TaskSet, + /// Answers the source's listings with the peer's tracks, while anyone lists it. + lister: std::sync::Arc, + listening: Option, + /// The broadcast closed: no new tracks, but those in flight run to their own end. + draining: bool, +} + +impl RunBroadcast { + fn new( + this: Subscriber, + path: Path, + source: crate::model::broadcast::SourceGuard, + broadcast: broadcast::Requests, + route: Option>, + ) -> Self { + let lister = std::sync::Arc::new( + broadcast + .lister() + .expect("only the source answers its broadcast's listings"), + ); + Self { + this, + path, + source, + broadcast, + route, + subscribes: TaskSet::owned(), + lister, + listening: None, + draining: false, + } + } + + /// The next track requested of the source, or `None` once the session is gone. + fn poll_requested(&mut self, waiter: &kio::Waiter) -> Poll>> { + if self.this.session.poll_closed(waiter).is_ready() { + return Poll::Ready(None); + } + self.this + .poll_listing(&self.path, &self.lister, &mut self.listening, waiter); + loop { + if let Poll::Ready(next) = self.broadcast.poll_requested_track(waiter) { + return Poll::Ready(Some(next)); + } + if self + .route + .as_ref() + .is_some_and(|route| route.poll_closed(waiter).is_ready()) + { + self.source.close(); + continue; + } + // Nothing holds the source and no track is on its way: retire it, so + // what the session keeps for the path goes with the fronts that used + // it. Declined when a holder or a track got there first, so look again. + if self.source.poll_unheld(waiter).is_pending() { + return Poll::Pending; + } + self.source.close_unheld(); + } + } +} + +impl kio::Task for RunBroadcast { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + loop { + if self.draining { + return self.this.poll_drain(&mut self.subscribes, waiter); + } + let next = match self.poll_requested(waiter) { + Poll::Ready(next) => next, + Poll::Pending => { + let _ = self.subscribes.poll(waiter); + return Poll::Pending; + } + }; + match next { + Some(Ok(request)) => self.subscribes.push(RunSubscribe::new( + self.this.clone(), + self.path.clone(), + self.broadcast.clone(), + self.route.clone(), + request, + )), + Some(Err(err)) => { + tracing::debug!(%err, "broadcast closed"); + drop(self.listening.take()); + self.draining = true; + } + // Session gone. + None => return Poll::Ready(()), + } + } + } +} + +/// Serves one track request: admits it, then learns the track from TRACK_STATUS or +/// subscribes to it, keeping the copy lingering between subscriptions until nobody +/// wants it. +struct RunSubscribe { + this: Subscriber, + broadcast_path: Path, + track_name: String, + /// Held for the subscription's lifetime but never watched: the broadcast ending + /// is a retraction, which does not disturb a subscription the scope allows. + _broadcast: broadcast::Requests, + /// Holds the session busy while the request waits on authorization or TRACK_STATUS. + /// Past that, only a subscription or a fill counts. + setup: Option, + /// Subscribe only to what our grant covers (MoQ Auth), and cancel once it no longer + /// does, leaving the rest of the session alone. + gate: crate::auth::Gate, + /// Group FETCHes for cache misses: standalone, so they outlive each subscription. + group_fetches: TaskSet, + step: RunStep, +} + +#[allow( + clippy::large_enum_variant, + reason = "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; an allow, since the size gap varies by target" +)] +enum RunStep { + /// Waiting for the gate to admit the request. Nobody wanting the track any more ends a + /// wait on a token that may never be answered. + Admit(Option), + Status(TrackStatus), + Linger(Linger), + Subscribe(SubscribeOnce), +} + +impl RunSubscribe { + fn new( + this: Subscriber, + broadcast_path: Path, + broadcast: broadcast::Requests, + // Open while the namespace stays announced and authorized, which its own token may + // have done in place of the scope. `None` once retracted. + route: Option>, + request: track::Request, + ) -> Self { + let setup = Some(this.session.work().start()); + let (auth, path, direction) = ( + this.auth.clone(), + broadcast_path.clone(), + crate::auth::Direction::Subscribe, + ); + let gate = match route { + Some(live) => crate::auth::Gate::announced(auth, path, direction, live), + None => crate::auth::Gate::new(auth, path, direction), + }; + Self { + this, + broadcast_path, + track_name: request.name().to_owned(), + _broadcast: broadcast, + setup, + gate, + group_fetches: TaskSet::owned(), + step: RunStep::Admit(Some(request)), + } + } +} + +impl kio::Task for RunSubscribe { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + loop { + match &mut self.step { + RunStep::Admit(request) => { + let pending = request.as_mut().expect("admitted once"); + let admitted = match self.gate.poll_ready(waiter) { + Poll::Ready(()) => Some(self.gate.poll_denied(waiter).is_pending()), + Poll::Pending => { + ready!(crate::util::poll_unwanted(&mut pending.demand(), waiter)); + None + } + }; + match admitted { + Some(true) => {} + Some(false) => { + request.take().expect("admitted once").reject(Error::Unauthorized); + return Poll::Ready(()); + } + None if pending.reject_unwanted(Error::Cancel) => return Poll::Ready(()), + // Demand returned in the gap. + None => continue, } - if let Err(err) = validity.check(waiter) { - unanswered = true; - return Poll::Ready(Setup::Response(Err(err))); + let request = request.take().expect("admitted once"); + // Data streams wait on the alias bound by SUBSCRIBE_OK, so leave the model + // request pending until its immutable track metadata is known. + // + // Demand with nobody subscribing (fetches, or a TRACK_STATUS downstream) learns + // the track from TRACK_STATUS rather than a SUBSCRIBE: a finished track refuses + // one, and one made only to learn the track races the fetches it was for. So + // does a bounded subscription, which may read only history below the Largest + // Object the status reports, as a fetch would: `Linger` subscribes once it + // reaches past that. + self.step = match request.subscription() { + subscription + if subscription.as_ref().is_none_or(|subscription| { + subscription.ranges.is_empty() + || subscription.start().is_some() && subscription.end().is_some() + }) => + { + RunStep::Status(TrackStatus::new(&self.this, request)) + } + _ => { + // `SubscribeOnce` counts the subscription itself. + self.setup = None; + RunStep::Subscribe(SubscribeOnce::new( + &self.this, + &self.broadcast_path, + &self.track_name, + Target::Request(request), + &self.gate, + )) + } + }; + } + RunStep::Status(status) => { + let Some(idle) = ready!(status.poll( + waiter, + &self.this, + &self.broadcast_path, + &self.track_name, + &mut self.group_fetches, + &mut self.gate, + )) else { + return Poll::Ready(()); + }; + self.setup = None; + self.step = RunStep::Linger(Linger::new(&self.this, idle)); + } + RunStep::Linger(linger) => { + let Some(next) = ready!(linger.poll( + waiter, + &self.this, + &self.broadcast_path, + &self.track_name, + &mut self.group_fetches, + &mut self.gate, + )) else { + return Poll::Ready(()); + }; + // `SubscribeOnce` counts the subscription itself. + self.setup = None; + self.step = RunStep::Subscribe(SubscribeOnce::new( + &self.this, + &self.broadcast_path, + &self.track_name, + Target::Resume(next), + &self.gate, + )); + } + RunStep::Subscribe(once) => { + let Some(idle) = ready!(once.poll( + waiter, + &self.this, + &self.broadcast_path, + &self.track_name, + &mut self.group_fetches, + )) else { + return Poll::Ready(()); + }; + self.step = RunStep::Linger(Linger::new(&self.this, idle)); + } + } + } + } +} + +/// How a listed track's handoff ended without a subscription to adopt. +enum Abandon { + Failed(Error), + Unused, +} + +/// What ended the wait for a SUBSCRIBE's answer. A publisher can be serving before its +/// SUBSCRIBE_OK reaches us, since the data streams are independent of the request stream. +/// Waiting for the response alone would miss the local side going away in that window and +/// leave the publisher serving a track nobody reads, which is the leak this whole path +/// exists to close. The broadcast ending is not a local side going away: a retraction does +/// not disturb subscriptions already in flight, and this one is. +enum AnswerWait { + Response(Result, Error>), + Unused, + /// `abort` already rejected the parked request. + Gone, +} + +/// One event ends an Established subscription: the last subscriber leaving, or the +/// publisher's PUBLISH_DONE. The broadcast ending does not: a retraction does not disturb +/// subscriptions already in flight. +enum ServeEnd { + Idle, + Revoked, + Done(Result, Error>), + Fill(track::Fill), + Update(RequestId, Update), + Invalid(Error), + /// An update failed to go out. + Failed(Error), +} + +/// Subscribes upstream for a target, until the subscription ends: `Some` once nobody +/// subscribes any more or the track ended, with the copy kept for [`Linger`], `None` once +/// it failed. +struct SubscribeOnce { + key: tracks::TrackKey, + /// This request's own view of the track's gate, which a resumed subscription's token + /// binds to. + gate: crate::auth::Gate, + target: Option, + token_grant: Option, + _working: crate::session::Working, + validity: crate::model::broadcast::RequestValidity, + resolved: ResolvedRanges, + /// The aggregate asked for, resolved against the copy. + subscription: Option, + /// A relative anchor the copy cannot settle yet counts back from the Largest Object, so + /// the live join that names it goes first and the answer resolves the anchor. + deferred: bool, + start: Option, + join: Join, + /// A live join delivers nothing below the group SUBSCRIBE_OK names as Largest. + live: bool, + /// A listed track's paused subscription, which this one resumes rather than open a + /// second one. + handoff: Option>, + /// The resumed subscription's alias, properties, and owed answers, when resuming one. + held: Option, + /// Set once the request has its stream. + request_id: RequestId, + stream: Option>, + messages: Option, + fill: Option>, + /// What the peer's SETUP negotiated: MAX_AGE, and compressed properties. + max_age_supported: bool, + flate: bool, + /// What the request went out with. + max_age: Option, + priority: i8, + new_group: Option, + properties_wanted: bool, + /// The lingering copy being subscribed again. + resumed: Option, + /// A local terminal ended the wait for the answer, so the peer still holds the request + /// to cancel. + unanswered: bool, + established: Option, + step: OnceStep, +} + +#[allow( + clippy::large_enum_variant, + reason = "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; an allow, since the size gap varies by target" +)] +enum OnceStep { + /// Resolving what to ask for. + Start, + /// Waiting for a listed track's paused subscription to be handed over. + Adopt, + /// Opening a fresh request stream. + Open(crate::transport::open::Open), + /// Writing the SUBSCRIBE, or the update resuming a held subscription. + Write, + /// Waiting for the answer, flushing the request meanwhile. + Answer, + /// Abandoned before the answer: cancelling the request while the fetches it drives go + /// on, then lingering again unless the session took the copy. + Abandon { cancelling: Cancelling, aborted: bool }, + /// Refused: cancelling the request, then failing. + Refuse(Cancelling), + /// The alias would not bind: cancelling the request, then failing the track with why. + Unbound { + cancelling: Cancelling, + track: track::Producer, + err: Error, + }, + /// Established, until one event ends it. + Serve, + /// The publisher ended it cleanly: waiting for the data streams it opened. + Settle { + settle: Settle, + fill: kio::Producer, + count: u64, + end: Result<(), Error>, + }, + /// Walking away from it: cancelling the request while the fetches it drives go on. + Cancel { cancelling: Cancelling, aborted: bool }, +} + +/// An Established subscription, from its answer to its end. +struct Established { + alias: u64, + properties: ietf::Properties, + /// The answer's Largest Location. + largest: Option, + timescale: Option, + track: track::Producer, + dynamic: track::Requests, + /// The copy already holds the track: it is current again up to the answer's Largest + /// Location once the join's head lands or a newer group does, since leaving cut the + /// group it was in. + resuming: Option>, + fetching: Option, + fetch_done: bool, + /// The NEW_GROUP_REQUEST last sent, the SUBSCRIBE's included. It stays outstanding + /// until the Largest Group reaches it, and joins meanwhile ask for the same group. + requested: Option, + dynamic_groups: bool, + /// The group the track's consumers want, `None` once nobody does. + wanted: Option, + /// The subscriber priority upstream holds, and the one the aggregate now asks for. + sent_max_age: Option, + wanted_max_age: Option, + sent_priority: i8, + priority: i8, + sent_bounds: WireBounds, + wanted_bounds: Result, + /// A deferred relative anchor, until the copy learns the live edge it counts back from: + /// an answer with no Largest leaves it to the first group. The bounds it then resolves + /// to go upstream as an update, absolute like any other. + unsettled: Option, + edge: track::Consumer, + /// REQUEST_UPDATEs sent and not yet answered. + unanswered: u64, + /// Our grant stopped covering the track: it is aborted rather than kept lingering. + failed: bool, + /// How the upstream subscription ended, which ends a listing entry it held. + ended: Option>, + /// The publisher ended the track cleanly: the copy lingers for its history's fills. + finished: bool, + demand: track::Demand, + /// Nobody subscribing at all (only fetches asked) needs no subscription. + subscribed: bool, + done: ReadPublishDone, +} + +impl SubscribeOnce { + fn new( + this: &Subscriber, + broadcast_path: &Path, + track_name: &str, + target: Target, + gate: &crate::auth::Gate, + ) -> Self { + Self { + key: (broadcast_path.clone(), track_name.into()), + gate: gate.unnamed(), + target: Some(target), + token_grant: None, + _working: this.session.work().start(), + validity: this.validity.start(), + resolved: ResolvedRanges::default(), + subscription: None, + deferred: false, + start: None, + join: Join::default(), + live: false, + handoff: None, + held: None, + request_id: RequestId(0), + stream: None, + messages: None, + fill: None, + max_age_supported: false, + flate: false, + max_age: None, + priority: 0, + new_group: None, + properties_wanted: false, + resumed: None, + unanswered: false, + established: None, + step: OnceStep::Start, + } + } + + fn target(&mut self) -> &mut Target { + self.target.as_mut().expect("the target is held until it is parked") + } + + fn stream(&mut self) -> Stream { + self.stream.take().expect("the request stream is open") + } + + fn poll( + &mut self, + waiter: &kio::Waiter, + this: &Subscriber, + broadcast_path: &Path, + track_name: &str, + group_fetches: &mut TaskSet, + ) -> Poll> { + loop { + match &mut self.step { + OnceStep::Start => { + let target = self.target.take().expect("started once"); + let subscription = target + .subscription() + .map(|subscription| match target.producer() { + Some(track) => self.resolved.resolve(track, &subscription), + None if needs_resolution(&subscription) => Err(Error::Unsupported), + None => Ok(subscription), + }) + .transpose(); + // A relative anchor the copy cannot settle yet counts back from the Largest + // Object, so the live join that names it goes first and the answer resolves + // the anchor. No Group ID arithmetic spells it upstream. A time anchor names + // no group the answer settles. + self.deferred = matches!(subscription, Err(Error::Unsupported)) + && target + .subscription() + .is_some_and(|subscription| !has_time(&subscription)); + let subscription = match subscription { + Ok(subscription) => subscription, + Err(_) if self.deferred => None, + Err(err) => { + target.fail(err); + return Poll::Ready(None); + } + }; + self.start = subscription.as_ref().and_then(|s| s.start()); + self.join = subscribe_join(self.start, subscription.as_ref().and_then(|s| s.end()), this.version); + // Every draft-18 join is a live one, since its joining FETCH carries that + // group alone. + self.live = self.start.is_none() || self.join.fetch; + self.subscription = subscription; + self.target = Some(target); + + // A listed track's paused subscription is already this track's upstream + // copy, so resume it rather than open a second one. + self.handoff = this.want(&self.key); + self.step = match self.handoff { + Some(_) => OnceStep::Adopt, + None => OnceStep::Open(Default::default()), + }; + } + OnceStep::Adopt => { + let handoff = self.handoff.as_ref().expect("adopting a handoff"); + // The handoff waits for the held stream to flush, which a peer withholding + // credit can stall, so nobody wanting the track any more ends the wait too. + // Every wait here drives the fetches a lingering copy started. + let _ = group_fetches.poll(waiter); + let adopt = if self.gate.poll_denied(waiter).is_ready() { + Err(Abandon::Failed(Error::Unauthorized)) + } else if let Err(err) = self.validity.check(waiter) { + Err(Abandon::Failed(err)) + } else if let Poll::Ready(adopted) = Subscriber::poll_handoff(handoff, waiter) { + Ok(adopted) + } else { + let target = self.target.as_mut().expect("the target is held until it is parked"); + ready!(match target { + Target::Request(request) => crate::util::poll_unwanted(&mut request.demand(), waiter), + Target::Resume(idle) => crate::util::poll_unwanted(&mut idle.demand, waiter), + }); + Err(Abandon::Unused) + }; + match adopt { + Ok(Some(upstream)) => { + this.adopted(&self.key, upstream.request_id); + let tracks::Upstream { + token_grant: token, + stream, + messages, + request_id, + alias, + properties, + owed, + // This subscription forwards into the copy itself from here. + copy: _, + } = upstream; + if let Some(token) = &token { + self.gate.bind(token.clone()); + } + self.token_grant = token; + self.held = Some(tracks::Resumed { + alias, + properties, + owed, + ended: None, + }); + self.request_id = request_id; + self.stream = Some(stream); + self.messages = Some(messages); + self.register(this, broadcast_path, track_name); + } + Ok(None) => self.step = OnceStep::Open(Default::default()), + Err(Abandon::Failed(err)) => { + this.unwant(self.key.clone(), handoff); + self.target + .take() + .expect("the target is held until it is parked") + .fail(err); + return Poll::Ready(None); + } + // Nobody wants the track: a resumed copy lingers again, a request is refused. + Err(Abandon::Unused) => { + match self.target.take().expect("the target is held until it is parked") { + Target::Resume(idle) => { + this.unwant(self.key.clone(), handoff); + return Poll::Ready(Some(idle)); + } + Target::Request(request) if request.reject_unwanted(Error::Cancel) => { + this.unwant(self.key.clone(), handoff); + return Poll::Ready(None); + } + target @ Target::Request(_) => self.target = Some(target), + } + } } - if gate.poll_denied(waiter).is_ready() { - unanswered = true; - return Poll::Ready(Setup::Response(Err(Error::Unauthorized))); + } + OnceStep::Open(open) => { + let _ = group_fetches.poll(waiter); + let opened = if self.gate.poll_denied(waiter).is_ready() { + Err(Error::Unauthorized) + } else if let Err(err) = self.validity.check(waiter) { + Err(err) + } else { + ready!(Stream::poll_open(&this.session, open, this.version, waiter)) + }; + match opened { + Ok(stream) => self.stream = Some(stream.with_stats(this.origin.stats())), + Err(err) => { + self.target + .take() + .expect("the target is held until it is parked") + .fail(err); + return Poll::Ready(None); + } } - // An answer that has already arrived wins over the local terminals. Both can - // be ready in one poll, and taking abandonment there would discard a response - // the publisher has already sent: if it was a rejection, the request is gone - // and cancelling it names a dead id back at a peer entitled to object. - if let Poll::Ready(res) = response { - return Poll::Ready(Setup::Response(res)); + self.request_id = this.control.next_request_id(); + self.messages = Some(super::request_stream::Messages::new( + this.session.clone(), + this.control.client(), + )); + self.register(this, broadcast_path, track_name); + } + OnceStep::Write => { + let peer = ready!(this.peer_setup.poll_get(waiter)); + self.max_age_supported = peer.max_age; + self.flate = peer.flate; + if let Some(step) = self.write(this, broadcast_path, track_name) { + return step; } - let mut state = self.state.lock(); - if let Some(idle) = &mut resumed { - // A session abort drained the subscription, and the copy with it. - if !state.subscribes.contains_key(&request_id) { - return Poll::Ready(Setup::Gone); + } + OnceStep::Answer => { + let _ = group_fetches.poll(waiter); + let stream = self.stream.as_mut().expect("the request stream is open"); + let messages = self.messages.as_mut().expect("the request stream is open"); + let response = match stream.writer.poll_flush(waiter) { + Poll::Ready(Ok(())) => this.poll_read_answer(waiter, stream, messages, self.held.as_mut()), + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Pending => Poll::Pending, + }; + let setup = if let Poll::Ready(Err(err)) = response { + AnswerWait::Response(Err(err)) + } else if let Err(err) = self.validity.check(waiter) { + self.unanswered = true; + AnswerWait::Response(Err(err)) + } else if self.gate.poll_denied(waiter).is_ready() { + self.unanswered = true; + AnswerWait::Response(Err(Error::Unauthorized)) + } else if let Poll::Ready(res) = response { + // An answer that has already arrived wins over the local terminals. Both + // can be ready in one poll, and taking abandonment there would discard a + // response the publisher has already sent: if it was a rejection, the + // request is gone and cancelling it names a dead id back at a peer + // entitled to object. + AnswerWait::Response(res) + } else { + let mut state = this.state.lock(); + match &mut self.resumed { + // A session abort drained the subscription, and the copy with it. + Some(_) if !state.subscribes.contains_key(&self.request_id) => AnswerWait::Gone, + Some(idle) => { + ready!(crate::util::poll_unwanted(&mut idle.demand, waiter)); + AnswerWait::Unused + } + None => match state + .subscribes + .get_mut(&self.request_id) + .and_then(|held| held.pending.as_mut()) + { + None => AnswerWait::Gone, + Some(pending) => { + ready!(crate::util::poll_unwanted(&mut pending.demand(), waiter)); + AnswerWait::Unused + } + }, } - return match crate::util::poll_unwanted(&mut idle.demand, waiter) { - Poll::Ready(_) => Poll::Ready(Setup::Unused), - Poll::Pending => Poll::Pending, - }; - } - let Some(pending) = state - .subscribes - .get_mut(&request_id) - .and_then(|held| held.pending.as_mut()) - else { - return Poll::Ready(Setup::Gone); }; - if crate::util::poll_unwanted(&mut pending.demand(), waiter).is_ready() { - return Poll::Ready(Setup::Unused); - } - Poll::Pending - }) - .await; - match setup { - Setup::Response(res) => break Some(res), - Setup::Gone => break None, - // Nobody holds the resumed copy any more: back to lingering. - Setup::Unused if resumed.is_some() => break None, - Setup::Unused => { - let mut state = self.state.lock(); - let Some(pending) = state - .subscribes - .get_mut(&request_id) - .and_then(|held| held.pending.take()) - else { - break None; - }; - if pending.reject_unwanted(Error::Cancel) { - break None; + // A request to abandon: a resumed copy nobody holds any more goes back to + // lingering, and a request goes once refusing it sticks. + let abandoned = match setup { + AnswerWait::Response(res) => match self.accept(this, broadcast_path, track_name, res) { + Some(step) => return step, + None => continue, + }, + AnswerWait::Gone => true, + AnswerWait::Unused if self.resumed.is_some() => true, + AnswerWait::Unused => { + let mut state = this.state.lock(); + match state + .subscribes + .get_mut(&self.request_id) + .and_then(|held| held.pending.take()) + { + None => true, + Some(pending) if pending.reject_unwanted(Error::Cancel) => true, + Some(pending) => { + if let Some(held) = state.subscribes.get_mut(&self.request_id) { + held.pending = Some(pending); + } + false + } + } } - if let Some(held) = state.subscribes.get_mut(&request_id) { - held.pending = Some(pending); + }; + if abandoned && let Some(step) = self.abandon(this, broadcast_path, track_name) { + return step; + } + } + OnceStep::Abandon { cancelling, aborted } => { + if cancelling.poll(waiter).is_pending() { + let _ = group_fetches.poll(waiter); + return Poll::Pending; + } + // A resumed copy goes back to lingering, unless the session took it. + let aborted = *aborted; + return Poll::Ready(self.resumed.take().filter(|_| !aborted)); + } + OnceStep::Refuse(cancelling) => { + ready!(cancelling.poll(waiter)); + return Poll::Ready(None); + } + OnceStep::Unbound { cancelling, track, err } => { + ready!(cancelling.poll(waiter)); + this.remove_subscribe(self.request_id); + let _ = track.clone().abort(err.clone()); + return Poll::Ready(None); + } + OnceStep::Serve => { + let end = ready!(self.poll_end(waiter, this, group_fetches)); + if let Some(idle) = self.ended(this, broadcast_path, track_name, group_fetches, end) + && let Some(step) = self.finish(this, idle) + { + return step; + } + } + OnceStep::Settle { + settle, + fill, + count, + end: _, + } => { + let live = self.established.as_mut().expect("established"); + live.poll_joining(waiter); + // Cache misses keep going meanwhile, such as one the fill settling just + // released to FETCH the joined group. + let _ = group_fetches.poll(waiter); + ready!(poll_settled(settle, waiter, fill, *count)); + let OnceStep::Settle { fill, end, .. } = std::mem::replace(&mut self.step, OnceStep::Serve) else { + unreachable!() + }; + // No tail is coming for the head. Only TRACK_ENDED makes it the whole group, + // even one that lands after this. Any other end leaves the group open + // upstream, so the head is cut short rather than passed off as the whole + // group. + if let Ok(mut fill) = fill.write() { + match &end { + Ok(()) => fill.end(), + Err(err) => fill.abort(err.clone()), } } + live.publisher_ended(end); + // The publisher already ended the request, so there is nothing to cancel. + if let Some(step) = self.finish(this, false) { + return step; + } } + OnceStep::Cancel { cancelling, aborted } => { + let _ = group_fetches.poll(waiter); + ready!(cancelling.poll(waiter)); + // A session abort took the copy too, and a revoked grant ended it. + let live = self.established.take().expect("established"); + if *aborted || live.failed { + return Poll::Ready(None); + } + return Poll::Ready(Some(Idle { + demand: live.track.demand(), + track: live.track, + dynamic: live.dynamic, + properties: live.properties, + finished: false, + })); + } + } + } + } + + /// Register the request before writing SUBSCRIBE so SUBSCRIBE_OK can bind its alias, + /// and so a fill fetch stream that overtakes it finds the subscription it answers. + fn register(&mut self, this: &Subscriber, broadcast_path: &Path, track_name: &str) { + let fill = kio::Producer::new(match self.join.fill.is_some() || self.join.fetch { + true => Fill::Requested, + false => Fill::Done, + }); + { + let mut state = this.state.lock(); + // A resumed subscription keeps the alias it is bound to and the data streams it + // counted while it forwarded before. + let (alias, tail) = match self.held.is_some() { + true => state + .subscribes + .remove(&self.request_id) + .map(|held| (held.alias, held.tail)) + .unwrap_or_default(), + false => Default::default(), + }; + state.subscribes.insert( + self.request_id, + TrackState { + // A resumed copy is the session's to abort from here on. + producer: self.target().producer().cloned(), + alias, + tail, + ..TrackState::pending(track_name.to_owned(), broadcast_path.clone(), fill.clone()) + }, + ); + } + self.fill = Some(fill); + self.step = OnceStep::Write; + } + + /// Write the SUBSCRIBE, or the update resuming a held subscription, and park the origin + /// request where a session abort can reject it. `Some` once that failed. + fn write(&mut self, this: &Subscriber, broadcast_path: &Path, track_name: &str) -> Option>> { + // The aggregate is read now: a subscriber can join while the request ID and stream + // were awaited. Updates follow its priority from here, but no draft lets one change + // the order. + let aggregate = self.target().subscription(); + self.priority = aggregate.as_ref().map(|s| s.priority).unwrap_or(0); + self.max_age = self + .max_age_supported + .then(|| super::max_age::millis(aggregate.as_ref().map_or(Duration::MAX, |s| s.budget()))); + let order = aggregate.map(|s| s.order).unwrap_or_default(); + // Without an Established subscription, a relay includes the request it holds. + self.new_group = self.target().new_group(); + // A timed copy asks for the units that opt its objects into timestamps. + self.properties_wanted = match self.target() { + Target::Request(_) => true, + Target::Resume(idle) => idle.properties.timescale.is_some(), + }; + let request_id = self.request_id; + let stream = self.stream.as_mut().expect("the request stream is open"); + // Buffered only: the wait for the answer flushes it, so the local terminals still end + // a request the peer's send window holds back. + let written = match &self.held { + // Resume the held subscription: its group order is the one it was made with, + // since no update can change it. Its groups wait on the answer from here. + Some(held) => { + this.rebind_alias(request_id, Alias::Resuming(request_id)); + let forward = tracks::Forward { + forward: true, + filter: Some(self.join.filter), + fill: self.join.fill, + priority: Some(self.priority), + max_age: self.max_age, + // Only a track declaring DYNAMIC_GROUPS takes one on an update. + new_group: self.new_group.filter(|_| held.properties.dynamic_groups), + }; + this.buffer_forward(&mut stream.writer, &forward) } + None => stream.writer.buffer_message( + ietf::Subscribe::ID, + &ietf::Subscribe { + max_age: self.max_age, + new_group: self.new_group, + request_id, + track_namespace: broadcast_path.clone(), + track_name: track_name.into(), + subscriber_priority: super::priority::to_wire(self.priority), + group_order: order.into(), + filter: self.join.filter, + fill: self.join.fill, + properties_wanted: self.properties_wanted, + forward: true, + rendezvous: false, + range_filters: false, + authorization_token: None, + }, + ), }; + if let Err(err) = written { + tracing::debug!(%err, "failed to encode subscribe"); + this.remove_subscribe(request_id); + this.entry_ended(&self.key, request_id, Err(err.clone())); + self.target + .take() + .expect("the target is held until it is parked") + .fail(err); + return Some(Poll::Ready(None)); + } - let Some(response) = setup else { - tracing::debug!( - broadcast = %crate::path::Rooted(self.origin.root(), broadcast_path), - track = %track_name, - "subscribe abandoned before it was accepted" - ); - // A held subscription is still the track's entry: pause it again, owing the - // answer to the resume as well. - if let Some(held) = held - && self.state.lock().subscribes.contains_key(&request_id) - { - let upstream = tracks::Upstream { - token_grant: token_grant.clone(), - stream, - messages, - request_id, - alias: held.alias, - properties: held.properties, - owed: held.owed + 1, - copy: resumed.as_ref().map(|idle| idle.track.clone()), + tracing::debug!(broadcast = %crate::path::Rooted(this.origin.root(), broadcast_path), track = %track_name, "subscribe started"); + + // Park the origin request where a session abort can reject it. The producer + // does not exist until SUBSCRIBE_OK, and dropping this task would otherwise + // end the track as `Dropped`. + match self.target.take().expect("the target is held until it is parked") { + Target::Request(request) => { + let mut state = this.state.lock(); + let Some(held) = state.subscribes.get_mut(&request_id) else { + request.reject(Error::Cancel); + return Some(Poll::Ready(None)); }; - self.pause(key, upstream); - return resumed; + held.pending = Some(request); } - // The publisher may already be serving before it answers. A session abort - // already rejected the parked request; dropping what remains is not a second one. - let aborted = self.remove_subscribe(request_id).is_none(); - let mut cancelling = std::pin::pin!(self.cancel_subscribe(stream)); - group_fetches - .drive(|waiter| waiter.poll_future(cancelling.as_mut())) - .await; - // A resumed copy goes back to lingering, unless the session took it. - return resumed.filter(|_| !aborted); + Target::Resume(idle) => self.resumed = Some(idle), + } + self.step = OnceStep::Answer; + None + } + + /// Walk away from a request abandoned before it was answered. `Some` once it ended here. + fn abandon(&mut self, this: &Subscriber, broadcast_path: &Path, track_name: &str) -> Option>> { + tracing::debug!( + broadcast = %crate::path::Rooted(this.origin.root(), broadcast_path), + track = %track_name, + "subscribe abandoned before it was accepted" + ); + // A held subscription is still the track's entry: pause it again, owing the + // answer to the resume as well. + if let Some(held) = self.held.take() + && this.state.lock().subscribes.contains_key(&self.request_id) + { + let upstream = tracks::Upstream { + token_grant: self.token_grant.clone(), + stream: self.stream(), + messages: self.messages.take().expect("the request stream is open"), + request_id: self.request_id, + alias: held.alias, + properties: held.properties, + owed: held.owed + 1, + copy: self.resumed.as_ref().map(|idle| idle.track.clone()), + }; + this.pause(self.key.clone(), upstream); + return Some(Poll::Ready(self.resumed.take())); + } + // The publisher may already be serving before it answers. A session abort + // already rejected the parked request; dropping what remains is not a second one. + let aborted = this.remove_subscribe(self.request_id).is_none(); + self.step = OnceStep::Abandon { + cancelling: Subscriber::cancelling(self.stream()), + aborted, }; + None + } - // SUBSCRIBE_OK commits the model's immutable metadata before the alias releases - // any data stream that arrived ahead of this control response. + /// Take in the answer. SUBSCRIBE_OK commits the model's immutable metadata before the + /// alias releases any data stream that arrived ahead of this control response. `Some` + /// once the request ended here. + fn accept( + &mut self, + this: &Subscriber, + broadcast_path: &Path, + track_name: &str, + response: Result, Error>, + ) -> Option>> { + let request_id = self.request_id; let accepted = match response { // A track that needs a property we do not understand, or carries more than we // could forward, is neither processed nor forwarded: the subscription is // cancelled and the subscribers refused. - Ok(Some(accepted)) - if refused(&accepted.properties, self.version, self.peer_setup.get().await.flate).is_some() => - { - let err = refused(&accepted.properties, self.version, self.peer_setup.get().await.flate) - .expect("checked by the guard"); - tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), broadcast_path), track = %track_name, %err, "subscribe refused"); + Ok(Some(accepted)) if refused(&accepted.properties, this.version, self.flate).is_some() => { + let err = refused(&accepted.properties, this.version, self.flate).expect("checked by the guard"); + tracing::debug!(broadcast = %crate::path::Rooted(this.origin.root(), broadcast_path), track = %track_name, %err, "subscribe refused"); // Refuse before waiting on the peer to acknowledge the cancellation. - if let Some(pending) = self.take_pending(request_id) { + if let Some(pending) = this.take_pending(request_id) { pending.reject(err.clone()); } - self.remove_subscribe(request_id); - self.entry_ended(&key, request_id, Err(err.clone())); - if let Some(idle) = resumed { + this.remove_subscribe(request_id); + this.entry_ended(&self.key, request_id, Err(err.clone())); + if let Some(idle) = self.resumed.take() { let _ = idle.track.abort(err); } - self.cancel_subscribe(stream).await; + self.step = OnceStep::Refuse(Subscriber::cancelling(self.stream())); return None; } Ok(Some(accepted)) => accepted, Ok(None) => { - if let Some(pending) = self.take_pending(request_id) { + if let Some(pending) = this.take_pending(request_id) { pending.reject(Error::UnexpectedMessage); } - self.remove_subscribe(request_id); - self.entry_ended(&key, request_id, Err(Error::UnexpectedMessage)); - if let Some(idle) = resumed { + this.remove_subscribe(request_id); + this.entry_ended(&self.key, request_id, Err(Error::UnexpectedMessage)); + if let Some(idle) = self.resumed.take() { let _ = idle.track.abort(Error::UnexpectedMessage); } - return None; + return Some(Poll::Ready(None)); } Err(err) => { tracing::debug!(%err, "subscribe response error"); // Refuse before waiting on the peer to acknowledge the cancellation. - if let Some(pending) = self.take_pending(request_id) { + if let Some(pending) = this.take_pending(request_id) { pending.reject(err.clone()); } - self.remove_subscribe(request_id); + this.remove_subscribe(request_id); // A held subscription that the publisher ended before answering ends its // entry as the PUBLISH_DONE says; one whose resume failed fails it, and is // cancelled since its PUBLISH_DONE has yet to come. - let ended = held.as_mut().and_then(|held| held.ended.take()); - let cancel = unanswered || (held.is_some() && ended.is_none()); + let ended = self.held.as_mut().and_then(|held| held.ended.take()); + let cancel = self.unanswered || (self.held.is_some() && ended.is_none()); let finished = matches!(ended, Some(Ok(()))); - self.entry_ended(&key, request_id, ended.unwrap_or_else(|| Err(err.clone()))); + this.entry_ended(&self.key, request_id, ended.unwrap_or_else(|| Err(err.clone()))); // A resumed copy whose track finished before the resume was answered cannot // place the end, so it lingers again and its reader's SUBSCRIBE learns it. - if finished && resumed.is_some() { - return resumed; + if finished && self.resumed.is_some() { + return Some(Poll::Ready(self.resumed.take())); } - if let Some(idle) = resumed { + if let Some(idle) = self.resumed.take() { let _ = idle.track.abort(err); } if cancel { - self.cancel_subscribe(stream).await; + self.step = OnceStep::Refuse(Subscriber::cancelling(self.stream())); + return None; } - return None; + return Some(Poll::Ready(None)); } }; let Accepted { @@ -3451,18 +4286,15 @@ impl Subscriber { // A resumed SUBSCRIBE opts out of the properties from draft-22, so its answer carries // none and the copy keeps the ones it learned. Before that the answer declares them // again. A pause holds these too, so a later resume still knows DYNAMIC_GROUPS. - let properties = match &resumed { - Some(idle) if held.is_none() && !properties_wanted && ietf::Filter::is_draft20(self.version) => { + let properties = match &self.resumed { + Some(idle) if self.held.is_none() && !self.properties_wanted && ietf::Filter::is_draft20(this.version) => { idle.properties.clone() } _ => properties, }; let timescale = properties.timescale; - let info = accepted_info(&properties, self.version); - // The copy already holds the track: it is current again up to the answer's Largest - // Location once the join's head lands or a newer group does, since leaving cut the - // group it was in. - let mut resuming = resumed.is_some().then(|| { + let info = accepted_info(&properties, this.version); + let resuming = self.resumed.is_some().then(|| { largest.map(|largest| track::Position { group: largest.group, frame: largest.object, @@ -3471,23 +4303,23 @@ impl Subscriber { // A fresh copy is staged until the answer is written in below, so no reader sees it // before its Largest and start. let mut staged = None; - let (mut track, dynamic) = match resumed { + let (mut track, dynamic) = match self.resumed.take() { Some(idle) => { - if !self.state.lock().subscribes.contains_key(&request_id) { + if !this.state.lock().subscribes.contains_key(&request_id) { // Aborted with the session while the answer was in hand. - return None; + return Some(Poll::Ready(None)); } (idle.track, idle.dynamic) } None => { - let Some(request) = self.take_pending(request_id) else { + let Some(request) = this.take_pending(request_id) else { // Aborted while the answer was in hand. The parked request is already rejected. - self.remove_subscribe(request_id); - return None; + this.remove_subscribe(request_id); + return Some(Poll::Ready(None)); }; // Pending until the answer is written in below, so a relay serving the copy // never reads the requested floor as where the feed starts. - let request = request.request_start(start); + let request = request.request_start(self.start); // Serves cache misses with a group FETCH. Registered before accepting, so a // miss queued meanwhile waits for it rather than failing for want of a handler. let dynamic = request.requests(); @@ -3505,90 +4337,70 @@ impl Subscriber { track.report_object(largest.group, largest.object); // A live join's Largest is where its feed is, which a relative anchor counts back // from even when a newer group lands before the joined one. - if live { + if self.live { track.report_edge(largest.group); } } // A live join starts at the publisher's edge; an absolute one where it asked. With // no Largest the join skips nothing, so every group is owed: a lower one's stream // may still arrive after a higher one's. This resolves the pending start. - let _ = match live { + let _ = match self.live { true => track.start_at(largest.map_or(0, |largest| largest.group)), - false => track.start_at(start), + false => track.start_at(self.start), }; if let Some(staged) = staged { staged.commit(); } // Resolved while the answer's Largest is still the live edge: no data stream reaches // the copy before its alias is bound below. - let settled = match deferred { + let settled = match self.deferred { true => track .subscription() - .map(|subscription| (settle_bounds(&mut resolved, &track, &subscription), subscription)), + .map(|subscription| (settle_bounds(&mut self.resolved, &track, &subscription), subscription)), false => None, }; - let mut fetching: Option = None; { - let mut state = self.state.lock(); + let mut state = this.state.lock(); if let Some(held) = state.subscribes.get_mut(&request_id) { held.producer = Some(track.clone()); held.timescale = timescale; held.fill_ok(largest); } } - self.entry_info(&key, request_id, track.info()); + this.entry_info(&self.key, request_id, track.info()); // A resumed subscription stays bound to the alias it was held with. - if held.is_some() { - self.rebind_alias(request_id, Alias::Active(request_id)); - } else if let Err(err) = self.register_alias(request_id, alias) { + if self.held.is_some() { + this.rebind_alias(request_id, Alias::Active(request_id)); + } else if let Err(err) = this.register_alias(request_id, alias) { if matches!(err, Error::Duplicate) { tracing::warn!(track_alias = %alias, "publisher reused a live track alias for another track"); - self.session.close( + this.session.close( super::error::to_session_code(&SessionError::from(&err)), err.to_string().as_ref(), ); } else { tracing::warn!(track_alias = %alias, %err, "could not bind track alias"); - self.cancel_subscribe(stream).await; + self.step = OnceStep::Unbound { + cancelling: Subscriber::cancelling(self.stream()), + track, + err, + }; + return None; } - self.remove_subscribe(request_id); + this.remove_subscribe(request_id); let _ = track.abort(err); - return None; - } - if join.fetch && largest.is_some() { - fetching = self.start_joining_fetch(request_id, &track); - } - - // One event ends the subscription: the last subscriber leaving, or the - // publisher's PUBLISH_DONE. The broadcast ending does not: a retraction - // does not disturb subscriptions already in flight. - enum End { - Idle, - Revoked, - Done(Result, Error>), - Fill(track::Fill), - Update(RequestId, Update), - Invalid(Error), - /// An update failed to go out. - Failed(Error), + return Some(Poll::Ready(None)); } + let fetching = match self.join.fetch && largest.is_some() { + true => this.start_joining_fetch(request_id, &track), + false => None, + }; - // The NEW_GROUP_REQUEST last sent, the SUBSCRIBE's included. It stays outstanding - // until the Largest Group reaches it, and joins meanwhile ask for the same group. - let mut requested = new_group; - let dynamic_groups = track.info().refresh.is_some(); - // The group the track's consumers want, `None` once nobody does. - let mut wanted = None; - // The subscriber priority upstream holds, and the one the aggregate now asks for. - let mut sent_max_age = max_age; - let mut wanted_max_age = max_age; - let mut sent_priority = priority; - let mut priority = priority; - let mut sent_bounds = (start, subscription.as_ref().and_then(|subscription| subscription.end())); + let sent_bounds = ( + self.start, + self.subscription.as_ref().and_then(|subscription| subscription.end()), + ); let mut wanted_bounds = Ok(sent_bounds); - // A deferred relative anchor, until the copy learns the live edge it counts back from: - // an answer with no Largest leaves it to the first group. The bounds it then resolves - // to go upstream as an update, absolute like any other. let mut unsettled = None; match settled { Some((Ok(Some(bounds)), _)) => wanted_bounds = Ok(bounds), @@ -3596,432 +4408,440 @@ impl Subscriber { Some((Err(err), _)) => wanted_bounds = Err(err), None => {} } - let mut edge = track.consume(); - // REQUEST_UPDATEs sent and not yet answered. - let unanswered = AtomicU64::new(0); - - let mut fetch_done = fetching.is_none(); - // Our grant stopped covering the track: it is aborted rather than kept lingering. - let mut failed = false; - // How the upstream subscription ended, which ends a listing entry it held. - let mut ended: Option> = None; - // The publisher ended the track cleanly: the copy lingers for its history's fills. - let mut finished = false; - let mut demand = track.demand(); - // Nobody subscribing at all (only fetches asked) needs no subscription. - let mut subscribed = track - .subscription() - .is_some_and(|subscription| !subscription.ranges.is_empty()); - let idle = { - let updating_track = track.clone(); - let mut done = std::pin::pin!(Self::read_publish_done( - &mut stream.reader, - self.version, - &unanswered, - &mut messages, - |ok| { - if let Some(largest) = ok.largest { - updating_track.report_object(largest.group, largest.object); - } - let mut state = self.state.lock(); - if let Some(held) = state.subscribes.get_mut(&request_id) { - held.fill_ok(ok.largest); - } + self.established = Some(Established { + alias, + properties, + largest, + timescale, + resuming, + fetch_done: fetching.is_none(), + fetching, + requested: self.new_group, + dynamic_groups: track.info().refresh.is_some(), + wanted: None, + sent_max_age: self.max_age, + wanted_max_age: self.max_age, + sent_priority: self.priority, + priority: self.priority, + sent_bounds, + wanted_bounds, + unsettled, + edge: track.consume(), + unanswered: 0, + failed: false, + ended: None, + finished: false, + demand: track.demand(), + subscribed: track + .subscription() + .is_some_and(|subscription| !subscription.ranges.is_empty()), + done: ReadPublishDone::default(), + track, + dynamic, + }); + self.step = OnceStep::Serve; + None + } + + /// The event that ends or updates the Established subscription. + fn poll_end(&mut self, waiter: &kio::Waiter, this: &Subscriber, group_fetches: &mut TaskSet) -> Poll { + let request_id = self.request_id; + let fill = self.fill.as_ref().expect("registered"); + let stream = self.stream.as_mut().expect("the request stream is open"); + let messages = self.messages.as_mut().expect("the request stream is open"); + let live = self.established.as_mut().expect("established"); + + if self.gate.poll_denied(waiter).is_ready() { + return Poll::Ready(ServeEnd::Revoked); + } + // An update the peer's send window holds back goes out from here, so it never + // holds back the ends below. + if let Poll::Ready(Err(err)) = stream.writer.poll_flush(waiter) { + return Poll::Ready(ServeEnd::Failed(err)); + } + live.poll_joining(waiter); + // An error is the track closing, which the arms below report. + if let Poll::Ready(Ok(requested)) = live.dynamic.poll_requested_fill(waiter) { + return Poll::Ready(ServeEnd::Fill(requested)); + } + let _ = group_fetches.poll(waiter); + let track = &mut live.track; + // A group past the answer's Largest also makes the copy current, so a join + // whose head never arrives cannot hold every later group back. + if let Some(largest) = live.resuming + && (poll_headed(fill, waiter).is_ready() + || largest.is_some_and(|largest| track.poll_past(largest.group, waiter).is_ready())) + { + // An absolute feed starts where it asked, whatever the answer's Largest; a + // live join's starts there. + let from = live.sent_bounds.0.filter(|_| !self.live); + match largest { + // A feed ending before the answer's Largest never shows it, so the + // answer alone makes the copy current. A live join's feed is open. + Some(largest) if !self.live && live.sent_bounds.1.is_some_and(|end| end <= largest.group) => { + track.set_live_ahead(largest, from) } - )); - loop { - let end = kio::wait(|waiter| { - if gate.poll_denied(waiter).is_ready() { - return Poll::Ready(End::Revoked); - } - // An update the peer's send window holds back goes out from here, so it never - // holds back the ends below. - if let Poll::Ready(Err(err)) = stream.writer.poll_flush(waiter) { - return Poll::Ready(End::Failed(err)); - } - if !fetch_done - && let Some(fut) = fetching.as_mut() - && let Poll::Ready(res) = fut.poll(waiter) - { - fetch_done = true; - if let Err(err) = res { - tracing::debug!(%err, "joining FETCH failed, joining at the next whole group"); - } - } - // An error is the track closing, which the arms below report. - if let Poll::Ready(Ok(requested)) = dynamic.poll_requested_fill(waiter) { - return Poll::Ready(End::Fill(requested)); - } - let _ = group_fetches.poll(waiter); - // A group past the answer's Largest also makes the copy current, so a join - // whose head never arrives cannot hold every later group back. - if let Some(largest) = resuming - && (poll_headed(&fill, waiter).is_ready() - || largest.is_some_and(|largest| track.poll_past(largest.group, waiter).is_ready())) - { - // An absolute feed starts where it asked, whatever the answer's Largest; a - // live join's starts there. - let from = sent_bounds.0.filter(|_| !live); - match largest { - // A feed ending before the answer's Largest never shows it, so the - // answer alone makes the copy current. A live join's feed is open. - Some(largest) if !live && sent_bounds.1.is_some_and(|end| end <= largest.group) => { - track.set_live_ahead(largest, from) - } - largest => track.set_live(largest, from), - } - resuming = None; - } - // The last subscriber left: the upstream subscription goes with it, as on - // lite, and the copy lingers for fetches and a returning subscriber. - while let Poll::Ready(Ok(subscription)) = track.poll_subscription_changed(waiter) { - subscribed = subscription - .as_ref() - .is_some_and(|subscription| !subscription.ranges.is_empty()); - unsettled = None; - if let Some(subscription) = subscription { - priority = subscription.priority; - match settle_bounds(&mut resolved, &track, &subscription) { - Ok(Some(bounds)) => wanted_bounds = Ok(bounds), - Ok(None) => unsettled = Some(subscription.clone()), - Err(err) => wanted_bounds = Err(err), - } - if max_age_supported { - wanted_max_age = Some(super::max_age::millis(subscription.budget())); - } - } - } - if let Some(subscription) = &unsettled { - while let Poll::Ready(Some(_)) = edge.poll_largest(waiter) {} - match settle_bounds(&mut resolved, &track, subscription) { - Ok(Some(bounds)) => { - wanted_bounds = Ok(bounds); - unsettled = None; - } - Ok(None) => {} - Err(err) => return Poll::Ready(End::Invalid(err)), - } - } - loop { - match demand.poll_wanted_changed(waiter) { - Poll::Ready(Ok(next)) => wanted = next, - Poll::Ready(Err(_)) => wanted = None, - Poll::Pending => break, - } - if wanted.is_none() { - break; - } - } - if !subscribed || wanted.is_none() { - return Poll::Ready(End::Idle); - } - if let Poll::Ready(res) = waiter.poll_future(done.as_mut()) { - return Poll::Ready(End::Done(res)); - } - // One REQUEST_UPDATE in flight at a time, which any MAX_REQUEST_UPDATES - // allows. Changes arriving meanwhile coalesce into the aggregate, read - // once its answer lands. Polled after the answers are read for that reason. - if unanswered.load(Ordering::Relaxed) > 0 { - return Poll::Pending; - } - let new_group = wanted.filter(|&value| { - dynamic_groups - && largest_group(&track, largest).is_some_and(|largest| value > largest) - && requested.is_none_or(|requested| value > requested) - }); - let bounds = match &wanted_bounds { - Ok(bounds) => *bounds, - Err(err) => return Poll::Ready(End::Invalid(err.clone())), - }; - let widened = match (bounds.0, sent_bounds.0) { - (None, Some(_)) => true, - (Some(start), Some(prior)) => start < prior, - _ => false, - } || match (bounds.1, sent_bounds.1) { - (None, Some(_)) => true, - (Some(end), Some(prior)) => end > prior, - _ => false, - }; - let ranges = if bounds != sent_bounds { - // Widen from the requested group bounds, not a fresh Next Object edge: - // that would cut the unread tail of a group already in flight. - let mut join = Join { - filter: match bounds { - (None | Some(0), None) => Filter::Unfiltered, - (start, end) => Filter::Absolute { - start: ietf::Location { - group: start.unwrap_or(0), - object: 0, - }, - end: end - .and_then(|end| end.checked_sub(1)) - .map(|group| ietf::EndLocation { group, object: None }), - }, - }, - ..Default::default() - }; - if widened && Filter::is_draft20(self.version) { - join.fill = Some(ietf::Fill { - filter: Some(match bounds.0 { - None => Filter::Relative(1), - Some(_) => join.filter, - }), - group_order: Some(GroupOrder::Ascending), - range_filters: false, - }); - } - Some(join) - } else { - None - }; - let mut update = Update { - filter: ranges.as_ref().map(|join| join.filter), - fill: ranges.and_then(|join| join.fill), - priority: (priority != sent_priority).then_some(priority), - max_age: wanted_max_age.filter(|_| wanted_max_age != sent_max_age), - new_group, - }; - let mut state = self.state.lock(); - let Some(held) = state.subscribes.get_mut(&request_id) else { - return Poll::Ready(End::Idle); - }; - // Reserve the new fill while holding the old slot. A stream that already - // owns a head must hand it off, but a reset may have discarded a pending - // fill's header, leaving nothing that could settle that request. - let prior_fill = held.fill.clone(); - let filling = if update.fill.is_some() { - match prior_fill.poll(waiter, |fill| match **fill { - Fill::Active | Fill::Ready { .. } => Poll::Pending, - _ => Poll::Ready(()), - }) { - Poll::Ready(Ok(fill)) => Some(fill), - Poll::Ready(Err(_)) => return Poll::Ready(End::Invalid(Error::Cancel)), - Poll::Pending => { - update.filter = None; - update.fill = None; - None - } - } - } else { - None - }; - if update.priority.is_none() - && update.max_age.is_none() - && update.new_group.is_none() - && update.filter.is_none() - { - return Poll::Pending; - } - let update_id = self.control.next_request_id(); - if let Some(mut filling) = filling { - // A handler may already have the old header and be waiting on its - // answer. It must keep the old state, never claim the new request. - if let Fill::Failed { sequence } = *filling { - held.headless_through = held.headless_through.max(Some(sequence)); - } - filling.cancel(); - drop(filling); - held.fill = kio::Producer::new(Fill::Requested); - held.fill_start = match update.fill.and_then(|fill| fill.filter.or(update.filter)) { - Some(Filter::Absolute { start, .. }) => Some(start), - _ => None, - }; - let previous = held.fetch_id.replace(update_id); - if let Some(previous) = previous { - state.fetches.remove(&previous); - } - state.fetches.insert(update_id, FillRequest::Update(request_id)); - } - Poll::Ready(End::Update(update_id, update)) - }) - .await; + largest => track.set_live(largest, from), + } + live.resuming = None; + } + // The last subscriber left: the upstream subscription goes with it, as on + // lite, and the copy lingers for fetches and a returning subscriber. + while let Poll::Ready(Ok(subscription)) = live.track.poll_subscription_changed(waiter) { + live.subscribed = subscription + .as_ref() + .is_some_and(|subscription| !subscription.ranges.is_empty()); + live.unsettled = None; + if let Some(subscription) = subscription { + live.priority = subscription.priority; + match settle_bounds(&mut self.resolved, &live.track, &subscription) { + Ok(Some(bounds)) => live.wanted_bounds = Ok(bounds), + Ok(None) => live.unsettled = Some(subscription.clone()), + Err(err) => live.wanted_bounds = Err(err), + } + if self.max_age_supported { + live.wanted_max_age = Some(super::max_age::millis(subscription.budget())); + } + } + } + if let Some(subscription) = &live.unsettled { + while let Poll::Ready(Some(_)) = live.edge.poll_largest(waiter) {} + match settle_bounds(&mut self.resolved, &live.track, subscription) { + Ok(Some(bounds)) => { + live.wanted_bounds = Ok(bounds); + live.unsettled = None; + } + Ok(None) => {} + Err(err) => return Poll::Ready(ServeEnd::Invalid(err)), + } + } + loop { + match live.demand.poll_wanted_changed(waiter) { + Poll::Ready(Ok(next)) => live.wanted = next, + Poll::Ready(Err(_)) => live.wanted = None, + Poll::Pending => break, + } + if live.wanted.is_none() { + break; + } + } + if !live.subscribed || live.wanted.is_none() { + return Poll::Ready(ServeEnd::Idle); + } + let updating = &live.track; + let done = live.done.poll( + waiter, + &mut stream.reader, + this.version, + &mut live.unanswered, + messages, + |ok| { + if let Some(largest) = ok.largest { + updating.report_object(largest.group, largest.object); + } + let mut state = this.state.lock(); + if let Some(held) = state.subscribes.get_mut(&request_id) { + held.fill_ok(ok.largest); + } + }, + ); + if let Poll::Ready(res) = done { + return Poll::Ready(ServeEnd::Done(res)); + } + // One REQUEST_UPDATE in flight at a time, which any MAX_REQUEST_UPDATES + // allows. Changes arriving meanwhile coalesce into the aggregate, read + // once its answer lands. Polled after the answers are read for that reason. + if live.unanswered > 0 { + return Poll::Pending; + } + let new_group = live.wanted.filter(|&value| { + live.dynamic_groups + && largest_group(&live.track, live.largest).is_some_and(|largest| value > largest) + && live.requested.is_none_or(|requested| value > requested) + }); + let bounds = match &live.wanted_bounds { + Ok(bounds) => *bounds, + Err(err) => return Poll::Ready(ServeEnd::Invalid(err.clone())), + }; + let sent_bounds = live.sent_bounds; + let widened = match (bounds.0, sent_bounds.0) { + (None, Some(_)) => true, + (Some(start), Some(prior)) => start < prior, + _ => false, + } || match (bounds.1, sent_bounds.1) { + (None, Some(_)) => true, + (Some(end), Some(prior)) => end > prior, + _ => false, + }; + let ranges = if bounds != sent_bounds { + // Widen from the requested group bounds, not a fresh Next Object edge: + // that would cut the unread tail of a group already in flight. + let mut join = Join { + filter: match bounds { + (None | Some(0), None) => Filter::Unfiltered, + (start, end) => Filter::Absolute { + start: ietf::Location { + group: start.unwrap_or(0), + object: 0, + }, + end: end + .and_then(|end| end.checked_sub(1)) + .map(|group| ietf::EndLocation { group, object: None }), + }, + }, + ..Default::default() + }; + if widened && Filter::is_draft20(this.version) { + join.fill = Some(ietf::Fill { + filter: Some(match bounds.0 { + None => Filter::Relative(1), + Some(_) => join.filter, + }), + group_order: Some(GroupOrder::Ascending), + range_filters: false, + }); + } + Some(join) + } else { + None + }; + let mut update = Update { + filter: ranges.as_ref().map(|join| join.filter), + fill: ranges.and_then(|join| join.fill), + priority: (live.priority != live.sent_priority).then_some(live.priority), + max_age: live.wanted_max_age.filter(|_| live.wanted_max_age != live.sent_max_age), + new_group, + }; + let mut state = this.state.lock(); + let Some(held) = state.subscribes.get_mut(&request_id) else { + return Poll::Ready(ServeEnd::Idle); + }; + // Reserve the new fill while holding the old slot. A stream that already + // owns a head must hand it off, but a reset may have discarded a pending + // fill's header, leaving nothing that could settle that request. + let prior_fill = held.fill.clone(); + let filling = if update.fill.is_some() { + match prior_fill.poll(waiter, |fill| match **fill { + Fill::Active | Fill::Ready { .. } => Poll::Pending, + _ => Poll::Ready(()), + }) { + Poll::Ready(Ok(fill)) => Some(fill), + Poll::Ready(Err(_)) => return Poll::Ready(ServeEnd::Invalid(Error::Cancel)), + Poll::Pending => { + update.filter = None; + update.fill = None; + None + } + } + } else { + None + }; + if update.priority.is_none() + && update.max_age.is_none() + && update.new_group.is_none() + && update.filter.is_none() + { + return Poll::Pending; + } + let update_id = this.control.next_request_id(); + if let Some(mut filling) = filling { + // A handler may already have the old header and be waiting on its + // answer. It must keep the old state, never claim the new request. + if let Fill::Failed { sequence } = *filling { + held.headless_through = held.headless_through.max(Some(sequence)); + } + filling.cancel(); + drop(filling); + held.fill = kio::Producer::new(Fill::Requested); + held.fill_start = match update.fill.and_then(|fill| fill.filter.or(update.filter)) { + Some(Filter::Absolute { start, .. }) => Some(start), + _ => None, + }; + let previous = held.fetch_id.replace(update_id); + if let Some(previous) = previous { + state.fetches.remove(&previous); + } + state.fetches.insert(update_id, FillRequest::Update(request_id)); + } + Poll::Ready(ServeEnd::Update(update_id, update)) + } - match end { - End::Invalid(err) => { - ended = Some(Err(err.clone())); - let _ = track.clone().abort(err); - break true; - } - End::Update(update_id, update) => { - let changed_range = update.filter.is_some(); - if changed_range { - sent_bounds = wanted_bounds.clone().expect("validated update bounds"); - } - let (asked, new_group, asked_age) = (update.priority, update.new_group, update.max_age); - if let Err(err) = stream.writer.buffer_message( - ietf::SubscribeUpdate::ID, - &ietf::SubscribeUpdate { - max_age: update.max_age, - request_id: update_id, - subscriber_priority: update.priority.map(super::priority::to_wire), - forward: None, - filter: update.filter, - authorization_token: None, - new_group: update.new_group, - fill: update.fill, - }, - ) { - tracing::debug!(%err, "failed to update the subscription"); - ended = Some(Err(err.clone())); - let _ = track.clone().abort(err); - break false; - } - if changed_range { - let _ = track.start_at(sent_bounds.0); - } - if asked_age.is_some() { - sent_max_age = asked_age; - } - if let Some(asked) = asked { - sent_priority = asked; - } - if let Some(value) = new_group { - requested = Some(value); - } - unanswered.fetch_add(1, Ordering::Relaxed); - } - End::Failed(err) => { - tracing::debug!(%err, "failed to update the subscription"); - ended = Some(Err(err.clone())); - let _ = track.clone().abort(err); - break false; - } - End::Revoked => { - tracing::info!(broadcast = %crate::path::Rooted(self.origin.root(), broadcast_path), track = %track_name, "subscription no longer authorized"); - let _ = track.clone().abort(Error::Unauthorized); - failed = true; - ended = Some(Err(Error::Unauthorized)); - break true; - } - End::Fill(requested) => { - let fetch = self.clone().serve_fill( - self.session.work().start(), - broadcast_path.clone(), - track_name.to_owned(), - requested, - Later::Known(timescale), - Some(Heading { request_id }), - ); - group_fetches.push(fetch); - } - End::Idle => { - tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), broadcast_path), track = %track_name, "subscribe cancelled (idle)"); - break true; - } - End::Done(res) => { - // Only TRACK_ENDED ends a listing entry quietly; a clean end of another - // kind, such as SUBSCRIPTION_ENDED, leaves the track's fate unknown. - let end = match &res { - Ok(done) => tracks::listing_end(done, self.version), - Err(err) => Err(err.clone()), - }; - ended = Some(end.clone()); - match res.map(|done| done.stream_count) { - Ok(count) => { - tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), broadcast_path), track = %track_name, "subscribe complete"); - // The publisher sends PUBLISH_DONE once every data stream it opened - // is closed, but QUIC does not order them, so some can still be on - // their way. Wait until Stream Count of their headers arrived and - // each is read to its end, or a bounded grace for any reset before - // its header (the draft says to use a timeout). The count is a hint: - // a published peer sends 0, which waits out the grace. - let held = self - .state - .lock() - .subscribes - .get(&request_id) - .map(|held| (held.tail.consume(), held.fill.clone())); - if let Some((tail, fill)) = held { - let mut settle = Settle::new(&self.runtime, tail); - kio::wait(|waiter| { - if !fetch_done - && let Some(fut) = fetching.as_mut() - && let Poll::Ready(res) = fut.poll(waiter) - { - fetch_done = true; - if let Err(err) = res { - tracing::debug!(%err, "joining FETCH failed, joining at the next whole group"); - } - } - // Cache misses keep going meanwhile, such as one the fill - // settling just released to FETCH the joined group. - let _ = group_fetches.poll(waiter); - poll_settled(&mut settle, waiter, &fill, count) - }) - .await; - // No tail is coming for the head. Only TRACK_ENDED makes it the - // whole group, even one that lands after this. Any other end - // leaves the group open upstream, so the head is cut short - // rather than passed off as the whole group. - if let Ok(mut fill) = fill.write() { - match &end { - Ok(()) => fill.end(), - Err(err) => fill.abort(err.clone()), - } - } - } - // Only TRACK_ENDED ends the track. Any other status ends just this - // subscription, so the copy fails rather than declare an end the - // track may run past: a later reader subscribes afresh to the same - // name, and its groups continue. - match end { - // A no-op once an END_OF_TRACK declared the end, whose readers - // wait for a missing group only as long as the grace. - Ok(()) => { - let _ = track.finish(); - finished = true; - } - Err(err) => { - let _ = track.clone().abort(err); - } - } - } - Err(err) => { - tracing::debug!(%err, "subscribe ended with error"); - let _ = track.clone().abort(err); + /// Act on an event of the Established subscription. `Some` once it ends, with whether + /// it went idle rather than ending upstream. + fn ended( + &mut self, + this: &Subscriber, + broadcast_path: &Path, + track_name: &str, + group_fetches: &mut TaskSet, + end: ServeEnd, + ) -> Option { + let live = self.established.as_mut().expect("established"); + match end { + ServeEnd::Invalid(err) => { + live.ended = Some(Err(err.clone())); + let _ = live.track.clone().abort(err); + Some(true) + } + ServeEnd::Update(update_id, update) => { + let stream = self.stream.as_mut().expect("the request stream is open"); + let changed_range = update.filter.is_some(); + if changed_range { + live.sent_bounds = live.wanted_bounds.clone().expect("validated update bounds"); + } + let (asked, new_group, asked_age) = (update.priority, update.new_group, update.max_age); + if let Err(err) = stream.writer.buffer_message( + ietf::SubscribeUpdate::ID, + &ietf::SubscribeUpdate { + max_age: update.max_age, + request_id: update_id, + subscriber_priority: update.priority.map(super::priority::to_wire), + forward: None, + filter: update.filter, + authorization_token: None, + new_group: update.new_group, + fill: update.fill, + }, + ) { + tracing::debug!(%err, "failed to update the subscription"); + live.ended = Some(Err(err.clone())); + let _ = live.track.clone().abort(err); + return Some(false); + } + if changed_range { + let _ = live.track.start_at(live.sent_bounds.0); + } + if asked_age.is_some() { + live.sent_max_age = asked_age; + } + if let Some(asked) = asked { + live.sent_priority = asked; + } + if let Some(value) = new_group { + live.requested = Some(value); + } + live.unanswered += 1; + None + } + ServeEnd::Failed(err) => { + tracing::debug!(%err, "failed to update the subscription"); + live.ended = Some(Err(err.clone())); + let _ = live.track.clone().abort(err); + Some(false) + } + ServeEnd::Revoked => { + tracing::info!(broadcast = %crate::path::Rooted(this.origin.root(), broadcast_path), track = %track_name, "subscription no longer authorized"); + let _ = live.track.clone().abort(Error::Unauthorized); + live.failed = true; + live.ended = Some(Err(Error::Unauthorized)); + Some(true) + } + ServeEnd::Fill(requested) => { + group_fetches.push(ServeFill::new( + this.clone(), + broadcast_path.clone(), + track_name.to_owned(), + requested, + Later::Known(live.timescale), + Some(Heading { + request_id: self.request_id, + }), + )); + None + } + ServeEnd::Idle => { + tracing::debug!(broadcast = %crate::path::Rooted(this.origin.root(), broadcast_path), track = %track_name, "subscribe cancelled (idle)"); + Some(true) + } + ServeEnd::Done(res) => { + // Only TRACK_ENDED ends a listing entry quietly; a clean end of another + // kind, such as SUBSCRIPTION_ENDED, leaves the track's fate unknown. + let end = match &res { + Ok(done) => tracks::listing_end(done, this.version), + Err(err) => Err(err.clone()), + }; + live.ended = Some(end.clone()); + match res.map(|done| done.stream_count) { + Ok(count) => { + tracing::debug!(broadcast = %crate::path::Rooted(this.origin.root(), broadcast_path), track = %track_name, "subscribe complete"); + // The publisher sends PUBLISH_DONE once every data stream it opened + // is closed, but QUIC does not order them, so some can still be on + // their way. Wait until Stream Count of their headers arrived and + // each is read to its end, or a bounded grace for any reset before + // its header (the draft says to use a timeout). The count is a hint: + // a published peer sends 0, which waits out the grace. + let held = this + .state + .lock() + .subscribes + .get(&self.request_id) + .map(|held| (held.tail.consume(), held.fill.clone())); + match held { + Some((tail, fill)) => { + self.step = OnceStep::Settle { + settle: Settle::new(&this.runtime, tail), + fill, + count, + end, + }; + return None; } + None => live.publisher_ended(end), } - // The publisher already ended the request, so there is nothing to cancel. - break false; + } + Err(err) => { + tracing::debug!(%err, "subscribe ended with error"); + let _ = live.track.clone().abort(err); } } + // The publisher already ended the request, so there is nothing to cancel. + Some(false) } - }; + } + } + /// Clean up after the Established subscription ended, `idle` when nobody wanted it any + /// more rather than its request ending. + fn finish(&mut self, this: &Subscriber, idle: bool) -> Option>> { + let request_id = self.request_id; + let live = self.established.as_mut().expect("established"); // Release an unallocated join now, or let the session cancel its committed request. - drop(fetching); + live.fetching = None; - if let Some(end) = ended { - self.entry_ended(&key, request_id, end); - } else if idle && self.is_entry(&key, request_id) && self.state.lock().subscribes.contains_key(&request_id) { + if let Some(end) = live.ended.take() { + this.entry_ended(&self.key, request_id, end); + } else if idle && this.is_entry(&self.key, request_id) && this.state.lock().subscribes.contains_key(&request_id) + { // A listed track's subscription pauses rather than ends, held again as its entry. - track.set_idle(); + let mut live = self.established.take().expect("established"); + live.track.set_idle(); let upstream = tracks::Upstream { - token_grant: token_grant.clone(), - stream, - messages, + token_grant: self.token_grant.clone(), + stream: self.stream(), + messages: self.messages.take().expect("the request stream is open"), request_id, - alias, - properties: properties.clone(), - owed: unanswered.load(Ordering::Relaxed), - copy: Some(track.clone()), + alias: live.alias, + properties: live.properties.clone(), + owed: live.unanswered, + copy: Some(live.track.clone()), }; - self.pause(key, upstream); - return Some(Idle { - demand: track.demand(), - track, - dynamic, - properties, + this.pause(self.key.clone(), upstream); + return Some(Poll::Ready(Some(Idle { + demand: live.track.demand(), + track: live.track, + dynamic: live.dynamic, + properties: live.properties, finished: false, - }); + }))); } // Clean up - let aborted = self.remove_subscribe(request_id).is_none(); + let aborted = this.remove_subscribe(request_id).is_none(); if !idle { // The publisher already ended the request, so a FIN is all we owe it, unless an // update its send window held back would go out truncated behind it. + let mut stream = self.stream(); match stream.writer.is_flushed() { true => { stream.writer.finish().ok(); @@ -4030,206 +4850,213 @@ impl Subscriber { } // A reader may still miss history the publisher holds, such as a group whose // fetch was in flight as the track ended, so a finished copy keeps serving fills. - return (finished && !aborted).then(|| Idle { - demand: track.demand(), - track, - dynamic, - properties, + let live = self.established.take().expect("established"); + return Some(Poll::Ready((live.finished && !aborted).then(|| Idle { + demand: live.track.demand(), + track: live.track, + dynamic: live.dynamic, + properties: live.properties, finished: true, - }); + }))); } // What the copy cached goes stale from here. Marked before the cancel, which waits // on the publisher: it stops serving as soon as the cancel lands, and a reader // returning in between must not take the cache as the live edge. - if !aborted && !failed { - track.set_idle(); - } + let live = self.established.as_mut().expect("established"); + if !aborted && !live.failed { + live.track.set_idle(); + } + // Refusing the copy wakes its FETCH jobs too. Keep driving their cancels while this + // stream waits for its FIN, even if the peer never ACKs it. + self.step = OnceStep::Cancel { + cancelling: Subscriber::cancelling(self.stream()), + aborted, + }; + None + } +} + +impl Established { + /// Drive the joining FETCH until it is done with. + fn poll_joining(&mut self, waiter: &kio::Waiter) { + if !self.fetch_done + && let Some(fetching) = self.fetching.as_mut() + && let Poll::Ready(res) = fetching.poll(waiter) { - // Refusing the copy wakes its FETCH jobs too. Keep driving their cancels - // while this stream waits for its FIN, even if the peer never ACKs it. - let mut cancelling = std::pin::pin!(self.cancel_subscribe(stream)); - kio::wait(|waiter| { - let _ = group_fetches.poll(waiter); - waiter.poll_future(cancelling.as_mut()) - }) - .await; + self.fetch_done = true; + if let Err(err) = res { + tracing::debug!(%err, "joining FETCH failed, joining at the next whole group"); + } } - // A session abort took the copy too, and a revoked grant ended it. - if aborted || failed { - return None; + } + + /// The publisher ended the subscription with `end`, its data streams accounted for. + fn publisher_ended(&mut self, end: Result<(), Error>) { + // Only TRACK_ENDED ends the track. Any other status ends just this + // subscription, so the copy fails rather than declare an end the + // track may run past: a later reader subscribes afresh to the same + // name, and its groups continue. + match end { + // A no-op once an END_OF_TRACK declared the end, whose readers + // wait for a missing group only as long as the grace. + Ok(()) => { + let _ = self.track.finish(); + self.finished = true; + } + Err(err) => { + let _ = self.track.clone().abort(err); + } } - Some(Idle { - demand: track.demand(), - track, - dynamic, - properties, - finished: false, - }) } +} - /// Keep a copy nobody subscribes to, cache and all, for fetches and a subscriber - /// returning soon: `Some` once one does, to subscribe again, `None` once nobody held - /// it through the linger, or the session ended. A finished copy never subscribes again. - /// - /// The session holds no upstream subscription for the copy meanwhile, so it marks the - /// copy unfed: a bounded read's holes are misses at once, which the fills answer. A - /// subscription whose ranges all lie below the Largest Object the publisher reported - /// reads only history, so it is served that way too, as a FETCH would be, rather than - /// subscribing. - async fn linger( - &mut self, - broadcast_path: &Path, - track_name: &str, - mut idle: Idle, - group_fetches: &mut TaskSet, - gate: &mut crate::auth::Gate, - ) -> Option { +/// Keeps a copy nobody subscribes to, cache and all, for fetches and a subscriber +/// returning soon: `Some` once one does, to subscribe again, `None` once nobody held +/// it through the linger, or the session ended. A finished copy never subscribes again. +/// +/// The session holds no upstream subscription for the copy meanwhile, so it marks the +/// copy unfed: a bounded read's holes are misses at once, which the fills answer. A +/// subscription whose ranges all lie below the Largest Object the publisher reported +/// reads only history, so it is served that way too, as a FETCH would be, rather than +/// subscribing. +struct Linger { + idle: Option, + /// The copy's entry in the session's registry, so a session abort ends it with its error. + id: u64, + linger: crate::time::Deadline, + subscribed: bool, + /// The grant stopped covering the copy: its fetches answer before it goes. + revoked: bool, +} + +/// What ended a wait while lingering. +enum LingerStep { + Fill(track::Fill), + Subscribe, + Expired, + Revoked, + Closed, +} + +impl Linger { + fn new(this: &Subscriber, idle: Idle) -> Self { // Registered so a session abort ends the copy with its error. let id = { - let mut state = self.state.lock(); + let mut state = this.state.lock(); state.next_lingering += 1; let id = state.next_lingering; state.lingering.insert(id, idle.track.clone()); id }; - let mut linger = crate::time::Deadline::new(&self.runtime); + let linger = crate::time::Deadline::new(&this.runtime); idle.track.set_unfed(true); - let mut subscribed = idle + let subscribed = idle .track .subscription() .is_some_and(|subscription| wants_feed(&subscription, idle.track.reported_group())); - enum Step { - Fill(track::Fill), - Subscribe, - Expired, - Revoked, - Closed, - } - let resume = loop { - let step = kio::wait(|waiter| { - if idle.track.poll_aborted(waiter).is_ready() { - return Poll::Ready(Step::Closed); - } - if gate.poll_denied(waiter).is_ready() { - return Poll::Ready(Step::Revoked); - } - if let Poll::Ready(Ok(fill)) = idle.dynamic.poll_requested_fill(waiter) { - return Poll::Ready(Step::Fill(fill)); - } + Self { + idle: Some(idle), + id, + linger, + subscribed, + revoked: false, + } + } + + /// Leave the linger, subscribing again when `resume`. + fn finish(&mut self, this: &Subscriber, resume: bool) -> Poll> { + // The session's abort took it out of the registry already, or this does. + let registered = this.state.lock().lingering.remove(&self.id).is_some(); + Poll::Ready(self.idle.take().filter(|_| resume && registered)) + } + + fn poll( + &mut self, + waiter: &kio::Waiter, + this: &Subscriber, + broadcast_path: &Path, + track_name: &str, + group_fetches: &mut TaskSet, + gate: &mut crate::auth::Gate, + ) -> Poll> { + loop { + if self.revoked { + ready!(this.poll_drain(group_fetches, waiter)); + return self.finish(this, false); + } + let idle = self.idle.as_mut().expect("lingering"); + let step = if idle.track.poll_aborted(waiter).is_ready() { + LingerStep::Closed + } else if gate.poll_denied(waiter).is_ready() { + LingerStep::Revoked + } else if let Poll::Ready(Ok(fill)) = idle.dynamic.poll_requested_fill(waiter) { + LingerStep::Fill(fill) + } else { let _ = group_fetches.poll(waiter); while let Poll::Ready(Ok(subscription)) = idle.track.poll_subscription_changed(waiter) { - subscribed = subscription + self.subscribed = subscription .as_ref() .is_some_and(|subscription| wants_feed(subscription, idle.track.reported_group())); } - if subscribed && !idle.finished { - return Poll::Ready(Step::Subscribe); - } - // Nobody holds it: let it go after the linger. A reader waiting on a fetch holds - // it too. A holder returning restarts the countdown when it next leaves. - if crate::util::poll_unwanted(&mut idle.demand, waiter).is_ready() { - if linger.deadline().is_none() { - linger.set(self.runtime.now().checked_add(track::IDLE_LINGER)); - } - if linger.poll(waiter).is_ready() { - return Poll::Ready(Step::Expired); + if self.subscribed && !idle.finished { + LingerStep::Subscribe + } else if crate::util::poll_unwanted(&mut idle.demand, waiter).is_ready() { + // Nobody holds it: let it go after the linger. A reader waiting on a fetch + // holds it too. A holder returning restarts the countdown when it next leaves. + if self.linger.deadline().is_none() { + self.linger.set(this.runtime.now().checked_add(track::IDLE_LINGER)); } + ready!(self.linger.poll(waiter)); + LingerStep::Expired } else { - linger.set(None); + self.linger.set(None); + return Poll::Pending; } - Poll::Pending - }) - .await; + }; match step { - Step::Fill(fill) => { - let fetch = self.clone().serve_fill( - self.session.work().start(), + LingerStep::Fill(fill) => { + group_fetches.push(ServeFill::new( + this.clone(), broadcast_path.clone(), track_name.to_owned(), fill, Later::Known(idle.properties.timescale), None, - ); - group_fetches.push(fetch); + )); } - Step::Subscribe => { + LingerStep::Subscribe => { // Marked before the SUBSCRIBE goes out: a hole may arrive live from here. idle.track.set_unfed(false); - break true; + return self.finish(this, true); } - Step::Expired => match idle.track.abort_unwanted(Error::Cancel) { - Ok(()) => { - tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), broadcast_path), track = %track_name, "track released (idle)"); - self.state.lock().lingering.remove(&id); - return None; - } - Err(used) => { - idle.track = used; - linger.set(None); + LingerStep::Expired => { + let mut idle = self.idle.take().expect("lingering"); + match idle.track.abort_unwanted(Error::Cancel) { + Ok(()) => { + tracing::debug!(broadcast = %crate::path::Rooted(this.origin.root(), broadcast_path), track = %track_name, "track released (idle)"); + this.state.lock().lingering.remove(&self.id); + return Poll::Ready(None); + } + Err(used) => { + idle.track = used; + self.idle = Some(idle); + self.linger.set(None); + } } - }, - Step::Revoked => { - tracing::info!(broadcast = %crate::path::Rooted(self.origin.root(), broadcast_path), track = %track_name, "lingering track no longer authorized"); - let _ = idle.track.clone().abort(Error::Unauthorized); - self.drain(group_fetches).await; - break false; - } - Step::Closed => break false, - } - }; - // The session's abort took it out of the registry already, or this does. - let registered = self.state.lock().lingering.remove(&id).is_some(); - (resume && registered).then_some(idle) - } - - /// Read the PUBLISH_DONE that ends an Established subscription, as the end it reports - /// and, for a clean end, its Stream Count. - /// - /// The publisher must send it before its FIN (draft-22 section 6.4.2.2), so a FIN - /// without one is a failed request, not a clean end. - /// Read the subscription's PUBLISH_DONE, and on the way the answers to the - /// REQUEST_UPDATEs counted in `unanswered`. Any other message is unexpected. - async fn read_publish_done( - reader: &mut Reader, - version: Version, - unanswered: &AtomicU64, - messages: &mut super::request_stream::Messages, - mut acknowledged: impl FnMut(ietf::RequestOk), - ) -> Result, Error> { - // A refused update is followed by PUBLISH_DONE, which ends the subscription. Its - // count stays, so no further update goes out meanwhile, and no answer is owed. - let mut refused = false; - loop { - let super::request_stream::Framed(type_id, body) = messages.message(reader).await?; - if type_id == ietf::PublishDone::ID { - let msg = body.message::(version)?.into_owned(); - tracing::debug!(message = ?msg, "received publish done"); - msg.end(version)?; - return Ok(msg); - } - let answer = matches!(type_id, ietf::RequestOk::ID | ietf::RequestError::ID) - && !refused && unanswered.load(Ordering::Relaxed) > 0; - if !answer { - return Err(Error::UnexpectedMessage); - } - let mut data = body.decoder(version); - match type_id { - ietf::RequestOk::ID => { - let msg = ietf::RequestOk::decode_msg(&mut data, version)?; - tracing::debug!(message = ?msg, "subscription update ok"); - acknowledged(msg); - unanswered.fetch_sub(1, Ordering::Relaxed); } - _ => { - let msg = ietf::RequestError::decode_msg(&mut data, version)?; - tracing::debug!(message = ?msg, "subscription update refused"); - refused = true; + LingerStep::Revoked => { + tracing::info!(broadcast = %crate::path::Rooted(this.origin.root(), broadcast_path), track = %track_name, "lingering track no longer authorized"); + let _ = idle.track.clone().abort(Error::Unauthorized); + self.revoked = true; } + LingerStep::Closed => return self.finish(this, false), } } } +} +impl Subscriber { /// Tell the publisher to stop serving a subscription we are walking away from. /// /// Every path that abandons an Established subscription goes through here, because @@ -4240,20 +5067,16 @@ impl Subscriber { /// cancellation (draft-22 section 6.4.2.2), so section 6.4.2.3's pair applies, an endpoint /// that has already FINed its sending direction cancels with STOP_SENDING on the /// receiving one. - async fn cancel_subscribe(&self, stream: Stream) { - let Stream { writer, mut reader } = stream; - - // STOP_SENDING needs no acknowledgement, so it goes first and the wait below covers - // only what we still have to deliver. - reader.abort(&Error::Cancel); + fn cancelling(mut stream: Stream) -> Cancelling { + // STOP_SENDING needs no acknowledgement, so it goes first and the wait covers only + // what we still have to deliver. + stream.reader.abort(&Error::Cancel); // Finishing alone would leave the writer's Drop free to RESET_STREAM, and a stream // that has sent its FIN is still retransmitting: the reset would discard what we // already wrote. Closing consumes the writer, removing that fallback, and waits for // the acknowledgement. A request the peer's send window held back resets instead. - if let Err(err) = writer.close().await { - tracing::debug!(%err, "failed to close the subscribe stream"); - } + Cancelling(StreamClose::new(stream)) } /// Send the joining FETCH that follows a draft-18 SUBSCRIBE, and keep reading its @@ -4262,423 +5085,732 @@ impl Subscriber { /// It asks for the head of the group SUBSCRIBE_OK names as Largest, and nothing below /// it: an earlier start's groups are the copy's fills, one FETCH per run of them. The /// subscribe's request id names the FETCH. The data - /// arrives on its own fetch stream and [`Self::recv_fill`] stitches it. A refusal or a + /// arrives on its own fetch stream and [`RecvFill`] stitches it. A refusal or a /// failed send settles the fill as [`Fill::Failed`], as a failed draft-22 fill does: only /// the head's group is lost, and a whole group still stands alone. A malformed answer /// closes the session. fn start_joining_fetch(&self, subscribe_id: RequestId, track: &track::Producer) -> Option { let fill = self.state.lock().subscribes.get(&subscribe_id)?.fill.clone(); - let track = track.clone(); let state = kio::Shared::::default(); - let progress = state.clone(); - let mut this = self.clone(); - let validity = self.validity.start(); - let work = async move { - let res = async { - // This subscription owns its joining fetch, so wait for stream credit. - // 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 opening = this.session.clone(); - let mut open = crate::transport::open::Open::default(); - let mut stream = kio::wait(|waiter| { - validity.check(waiter)?; - if let Poll::Ready(err) = track.poll_aborted(waiter) { - return Poll::Ready(Err(err)); - } - Stream::poll_open(&opening, &mut open, this.version, waiter) - }) - .await?; - stream = stream.with_stats(this.origin.stats()); - let fetch_id = this.control.next_request_id(); - progress.lock().committed = true; - { - let mut state = this.state.lock(); - if let Some(track) = state.subscribes.get_mut(&subscribe_id) { - track.fetch_id = Some(fetch_id); - state.fetches.insert(fetch_id, FillRequest::Joining(subscribe_id)); - } - } - // Flushed while waiting for the answer, so a cancellation still ends a request the - // peer's send window holds back. - stream.writer.buffer_message( - ietf::Fetch::ID, - &ietf::Fetch { - authorization_token: None, - request_id: fetch_id, - subscriber_priority: super::priority::to_wire( - track.subscription().map(|s| s.priority).unwrap_or(0), - ), - group_order: GroupOrder::Ascending, - fetch_type: FetchType::RelativeJoining { - subscriber_request_id: subscribe_id, - group_offset: 0, - }, - range_filters: false, - fill_timeout: false, - properties_wanted: true, - // Served by the subscription it joins, which named the instance. - }, - )?; - this.finish_joining_fetch(stream, progress).await - } - .await; - // A refusal costs only the head, but a malformed answer is the peer breaking the - // protocol, which closes the session as it would on any request. - if let Err(err) = &res - && is_protocol_violation(err) - { - this.session.close( - super::error::to_session_code(&SessionError::from(err)), - err.to_string().as_ref(), - ); - } - // Release a mid-group tail parked on the head, including one still being written: - // a request that failed after FETCH_OK cannot be trusted to finish its stream. - if res.is_err() { - let largest = this - .state - .lock() - .subscribes - .get(&subscribe_id) - .and_then(|track| track.fill_largest); - if let Ok(mut state) = fill.write() - && matches!(*state, Fill::Requested | Fill::Serving(_) | Fill::Active) - { - *state = Fill::failed(largest); - } - } - res - } - .boxed(); + let join = JoinFetch { + this: self.clone(), + subscribe_id, + track: track.clone(), + fill, + progress: state.clone(), + validity: self.validity.start(), + messages: super::request_stream::Messages::new(self.session.clone(), self.control.client()), + stream: None, + step: JoinStep::Open(Default::default()), + }; Some(JoiningTask { - work: Some(work), + work: Some(join), state, tasks: self.tasks.clone(), }) } - async fn finish_joining_fetch( + /// Learn a track nobody subscribes to from TRACK_STATUS; see [`TrackStatus`]. + #[cfg(test)] + async fn track_status( &mut self, - mut stream: Stream, - progress: kio::Shared, - ) -> Result<(), Error> { - let cancelled = |waiter: &kio::Waiter| { - progress - .poll(waiter, |state| { - state.cancelled.then_some(()).map_or(Poll::Pending, Poll::Ready) - }) - .map(drop) - }; - let mut messages = super::request_stream::Messages::new(self.session.clone(), self.control.client()); - let res = async { - let answer = { - let mut response = std::pin::pin!(async { - kio::wait(|waiter| stream.writer.poll_flush(waiter)).await?; - self.read_fetch_response(&mut stream, &mut messages).await - }); - kio::wait(|waiter| { - if cancelled(waiter).is_ready() { - return Poll::Ready(None); - } - waiter.poll_future(response.as_mut()).map(Some) - }) - .await - }; - match answer { - Some(answer) => answer?, - None => return Ok(()), - }; - // The request stays open while the subscription still needs its fill. From draft 18 - // it may carry a GOAWAY meanwhile; before that, nothing follows the answer here. - let mut finished = false; - kio::wait(|waiter| { - if !finished { - match messages.poll_end(&mut stream.reader, waiter) { - Poll::Ready(Ok(())) => finished = true, - Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), - Poll::Pending => {} - } - } - cancelled(waiter).map(Ok) - }) - .await - } - .await; - stream.reader.abort(&Error::Cancel); - let _ = stream.writer.close().await; - res + broadcast_path: &Path, + track_name: &str, + request: track::Request, + group_fetches: &mut TaskSet, + gate: &mut crate::auth::Gate, + ) -> Option { + let mut status = TrackStatus::new(self, request); + kio::wait(|waiter| status.poll(waiter, self, broadcast_path, track_name, group_fetches, gate)).await } - async fn read_subscribe_response( + /// The answer to a TRACK_STATUS: TRACK_STATUS_OK, or the publisher's refusal as an error. + fn poll_read_track_status_response( &self, - stream: &mut Stream, + reader: &mut Reader, messages: &mut super::request_stream::Messages, - ) -> Result, Error> { - // Read type_id + size + body from the stream - let super::request_stream::Framed(type_id, body) = messages.message(&mut stream.reader).await?; + waiter: &kio::Waiter, + ) -> Poll> { + let super::request_stream::Framed(type_id, body) = ready!(messages.poll_message(reader, waiter))?; let mut data = body.decoder(self.version); - match type_id { - ietf::SubscribeOk::ID => { - let msg = ietf::SubscribeOk::decode_msg(&mut data, self.version)?; - tracing::debug!(message = ?msg, "received subscribe ok"); - Ok(Some(Accepted { - alias: msg.track_alias, - properties: msg.properties, - largest: msg.largest, - })) + let code = match type_id { + id if id == ietf::TrackStatusOk::ID => { + return Poll::Ready(Ok(ietf::TrackStatusOk::decode_msg(&mut data, self.version)?)); } - ietf::RequestError::ID => { - let msg = ietf::RequestError::decode_msg(&mut data, self.version)?; - tracing::warn!(message = ?msg, "request error"); - Err(request::from_code( - msg.error_code, - request::Kind::Subscribe, - self.version, - )) + ietf::RequestError::ID if true => ietf::RequestError::decode_msg(&mut data, self.version)?.error_code, + _ => return Poll::Ready(Err(Error::UnexpectedMessage)), + }; + Poll::Ready(Err(request::from_code(code, request::Kind::TrackStatus, self.version))) + } +} + +/// Reads the PUBLISH_DONE that ends an Established subscription, as the end it reports +/// and, for a clean end, its Stream Count, and on the way the answers to the +/// REQUEST_UPDATEs counted in `unanswered`. Any other message is unexpected. +/// +/// The publisher must send it before its FIN (draft-22 section 6.4.2.2), so a FIN +/// without one is a failed request, not a clean end. +#[derive(Default)] +struct ReadPublishDone { + /// A refused update is followed by PUBLISH_DONE, which ends the subscription. Its + /// count stays, so no further update goes out meanwhile, and no answer is owed. + refused: bool, +} + +impl ReadPublishDone { + fn poll( + &mut self, + waiter: &kio::Waiter, + reader: &mut Reader, + version: Version, + unanswered: &mut u64, + messages: &mut super::request_stream::Messages, + mut acknowledged: impl FnMut(ietf::RequestOk), + ) -> Poll, Error>> { + loop { + let super::request_stream::Framed(type_id, body) = ready!(messages.poll_message(reader, waiter))?; + if type_id == ietf::PublishDone::ID { + let msg = body.message::(version)?.into_owned(); + tracing::debug!(message = ?msg, "received publish done"); + msg.end(version)?; + return Poll::Ready(Ok(msg)); + } + let answer = + matches!(type_id, ietf::RequestOk::ID | ietf::RequestError::ID) && !self.refused && *unanswered > 0; + if !answer { + return Poll::Ready(Err(Error::UnexpectedMessage)); + } + let mut data = body.decoder(version); + match type_id { + ietf::RequestOk::ID => { + let msg = ietf::RequestOk::decode_msg(&mut data, version)?; + tracing::debug!(message = ?msg, "subscription update ok"); + acknowledged(msg); + *unanswered -= 1; + } + _ => { + let msg = ietf::RequestError::decode_msg(&mut data, version)?; + tracing::debug!(message = ?msg, "subscription update refused"); + self.refused = true; + } + } + } + } +} + +/// A subscription's request stream being cancelled; see [`Subscriber::cancelling`]. +struct Cancelling(StreamClose); + +impl Cancelling { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + if let Err(err) = ready!(self.0.poll(waiter)) { + tracing::debug!(%err, "failed to close the subscribe stream"); + } + Poll::Ready(()) + } +} + +/// A draft-18 joining FETCH; see [`Subscriber::start_joining_fetch`]. +struct JoinFetch { + this: Subscriber, + subscribe_id: RequestId, + track: track::Producer, + /// The subscription's fill, which the head this FETCH brings settles. + fill: kio::Producer, + progress: kio::Shared, + validity: crate::model::broadcast::RequestValidity, + messages: super::request_stream::Messages, + stream: Option>, + step: JoinStep, +} + +enum JoinStep { + /// Opening the request stream. + Open(crate::transport::open::Open), + /// Waiting for FETCH_OK, flushing the FETCH meanwhile. + Answer, + /// The request stays open while the subscription still needs its fill. From draft 18 + /// it may carry a GOAWAY meanwhile; before that, nothing follows the answer here. + Listen { + /// The publisher finished its side. + finished: bool, + }, + /// Closing the request, then ending with how it went. + Close(StreamClose, Option>), +} + +impl JoinFetch { + /// Ready once the subscription cancelled the join. + fn poll_cancelled(progress: &kio::Shared, waiter: &kio::Waiter) -> Poll<()> { + progress + .poll(waiter, |state| { + state.cancelled.then_some(()).map_or(Poll::Pending, Poll::Ready) + }) + .map(drop) + } + + /// Abandon the request's reads and close it, ending with `res`. + fn close(&mut self, res: Result<(), Error>) { + let mut stream = self.stream.take().expect("the request stream is open"); + stream.reader.abort(&Error::Cancel); + self.step = JoinStep::Close(StreamClose::new(stream), Some(res)); + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + let res = ready!(self.poll_join(waiter)); + // A refusal costs only the head, but a malformed answer is the peer breaking the + // protocol, which closes the session as it would on any request. + if let Err(err) = &res + && is_protocol_violation(err) + { + self.this.session.close( + super::error::to_session_code(&SessionError::from(err)), + err.to_string().as_ref(), + ); + } + // Release a mid-group tail parked on the head, including one still being written: + // a request that failed after FETCH_OK cannot be trusted to finish its stream. + if res.is_err() { + let largest = self + .this + .state + .lock() + .subscribes + .get(&self.subscribe_id) + .and_then(|track| track.fill_largest); + if let Ok(mut state) = self.fill.write() + && matches!(*state, Fill::Requested | Fill::Serving(_) | Fill::Active) + { + *state = Fill::failed(largest); + } + } + Poll::Ready(res) + } + + fn poll_join(&mut self, waiter: &kio::Waiter) -> Poll> { + loop { + match &mut self.step { + JoinStep::Open(open) => { + // This subscription owns its joining fetch, so wait for stream credit. + // No ID is spent while opening a stream, so dropping a blocked open + // cannot leave a gap in the peer's sequential request IDs. + self.validity.check(waiter)?; + if let Poll::Ready(err) = self.track.poll_aborted(waiter) { + return Poll::Ready(Err(err)); + } + let stream = ready!(Stream::poll_open(&self.this.session, open, self.this.version, waiter))?; + let mut stream = stream.with_stats(self.this.origin.stats()); + let fetch_id = self.this.control.next_request_id(); + self.progress.lock().committed = true; + { + let mut state = self.this.state.lock(); + if let Some(track) = state.subscribes.get_mut(&self.subscribe_id) { + track.fetch_id = Some(fetch_id); + state.fetches.insert(fetch_id, FillRequest::Joining(self.subscribe_id)); + } + } + // Flushed while waiting for the answer, so a cancellation still ends a request the + // peer's send window holds back. + stream.writer.buffer_message( + ietf::Fetch::ID, + &ietf::Fetch { + authorization_token: None, + request_id: fetch_id, + subscriber_priority: super::priority::to_wire( + self.track.subscription().map(|s| s.priority).unwrap_or(0), + ), + group_order: GroupOrder::Ascending, + fetch_type: FetchType::RelativeJoining { + subscriber_request_id: self.subscribe_id, + group_offset: 0, + }, + range_filters: false, + fill_timeout: false, + properties_wanted: true, + // Served by the subscription it joins, which named the instance. + }, + )?; + self.stream = Some(stream); + self.step = JoinStep::Answer; + } + JoinStep::Answer => { + let stream = self.stream.as_mut().expect("the request stream is open"); + if Self::poll_cancelled(&self.progress, waiter).is_ready() { + self.close(Ok(())); + continue; + } + let answer = ready!(match stream.writer.poll_flush(waiter) { + Poll::Ready(Ok(())) => + self.this + .poll_read_fetch_response(&mut stream.reader, &mut self.messages, waiter), + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Pending => Poll::Pending, + }); + match answer { + Ok(_) => self.step = JoinStep::Listen { finished: false }, + Err(err) => self.close(Err(err)), + } + } + JoinStep::Listen { finished } => { + let stream = self.stream.as_mut().expect("the request stream is open"); + if !*finished { + match self.messages.poll_end(&mut stream.reader, waiter) { + Poll::Ready(Ok(())) => *finished = true, + Poll::Ready(Err(err)) => { + self.close(Err(err)); + continue; + } + Poll::Pending => {} + } + } + ready!(Self::poll_cancelled(&self.progress, waiter)); + self.close(Ok(())); + } + JoinStep::Close(close, res) => { + let _ = ready!(close.poll(waiter)); + return Poll::Ready(res.take().expect("polled after completion")); + } } - _ => Err(Error::UnexpectedMessage), } } +} - /// Learn a track nobody subscribes to from TRACK_STATUS, accepting an idle copy that - /// [`Self::linger`] serves fetches from and subscribes once someone does. `None` once - /// refused or abandoned: a refused status fails the query-only or fetch-only request. - /// - /// Fetches go out meanwhile into `group_fetches`, reading their objects only once the - /// status says the timescale, and failing with it. - async fn track_status( +/// Learns a track nobody subscribes to from TRACK_STATUS, accepting an idle copy that +/// [`Linger`] serves fetches from and subscribes once someone does. `None` once refused or +/// abandoned: a refused status fails the query-only or fetch-only request. +/// +/// Fetches go out meanwhile into `group_fetches`, reading their objects only once the +/// status says the timescale, and failing with it. +struct TrackStatus { + timescale: crate::util::Learning>, + step: StatusStep, +} + +#[allow( + clippy::large_enum_variant, + reason = "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; an allow, since the size gap varies by target" +)] +enum StatusStep { + Learn(LearnTrackStatus), + /// The fetches sent alongside fail with the status: letting them answer first. + Drain, +} + +impl TrackStatus { + fn new(this: &Subscriber, request: track::Request) -> Self { + let timescale = crate::util::Learning::new(); + Self { + step: StatusStep::Learn(LearnTrackStatus::new(this, request, timescale.later())), + timescale, + } + } + + fn poll( &mut self, + waiter: &kio::Waiter, + this: &Subscriber, broadcast_path: &Path, track_name: &str, - request: track::Request, group_fetches: &mut TaskSet, gate: &mut crate::auth::Gate, - ) -> Option { - let timescale = crate::util::Learning::new(); - let learned = self - .learn_track_status(broadcast_path, track_name, request, group_fetches, ×cale, gate) - .await; - match learned { - Ok(idle) => { - timescale.resolve(Ok(idle.properties.timescale)); - Some(idle) - } - Err(err) => { - // The fetches sent alongside fail with the status: let them answer first. - timescale.resolve(Err(err)); - self.drain(group_fetches).await; - None + ) -> Poll> { + loop { + match &mut self.step { + StatusStep::Learn(learn) => { + match ready!(learn.poll(waiter, this, broadcast_path, track_name, group_fetches, gate)) { + Ok(idle) => { + self.timescale.resolve(Ok(idle.properties.timescale)); + return Poll::Ready(Some(idle)); + } + Err(err) => { + self.timescale.resolve(Err(err)); + self.step = StatusStep::Drain; + } + } + } + StatusStep::Drain => { + ready!(this.poll_drain(group_fetches, waiter)); + return Poll::Ready(None); + } } } } +} + +/// [`TrackStatus`]'s exchange, failing with why the request was refused or abandoned. +struct LearnTrackStatus { + request: Option, + /// Registered before accepting: `accept` releases the request's own fetch gate, and a + /// cache-miss fetch queued in the gap would be drained as NotFound. + dynamic: Option, + demand: Option, + validity: crate::model::broadcast::RequestValidity, + /// The timescale the fetches sent alongside read their objects in. + timescale: Later>, + messages: super::request_stream::Messages, + stream: Option>, + step: LearnStep, +} + +enum LearnStep { + /// Opening the request stream. + Open(crate::transport::open::Open), + /// Waiting for the answer, flushing the TRACK_STATUS meanwhile. + Response, + /// A local terminal ended the wait: closing the request, then failing with it. + Cancelled(StreamClose, Option), + /// Checking the answer against what the peer's SETUP negotiated. + Answered(Option>), +} + +impl LearnTrackStatus { + fn new(this: &Subscriber, request: track::Request, timescale: Later>) -> Self { + Self { + dynamic: Some(request.requests()), + request: Some(request), + demand: None, + validity: this.validity.start(), + timescale, + messages: super::request_stream::Messages::new(this.session.clone(), this.control.client()), + stream: None, + step: LearnStep::Open(Default::default()), + } + } + + /// Refuse the request with `err`, and fail with it. + fn reject(&mut self, err: Error) -> Poll> { + if let Some(request) = self.request.take() { + request.reject(err.clone()); + } + Poll::Ready(Err(err)) + } - /// [`Self::track_status`]'s exchange, failing with why the request was refused or - /// abandoned. - async fn learn_track_status( + fn poll( &mut self, + waiter: &kio::Waiter, + this: &Subscriber, broadcast_path: &Path, track_name: &str, - request: track::Request, group_fetches: &mut TaskSet, - timescale: &crate::util::Learning>, gate: &mut crate::auth::Gate, - ) -> Result { - // Registered before accepting: `accept` releases the request's own fetch gate, and a - // cache-miss fetch queued in the gap would be drained as NotFound. - let dynamic = request.requests(); - let validity = self.validity.start(); - let opening = self.session.clone(); - let mut open = crate::transport::open::Open::default(); - let mut stream = match kio::wait(|waiter| { - if gate.poll_denied(waiter).is_ready() { - return Poll::Ready(Err(Error::Unauthorized)); - } - validity.check(waiter)?; - Stream::poll_open(&opening, &mut open, self.version, waiter) - }) - .await - { - Ok(stream) => stream, - Err(err) => { - request.reject(err.clone()); - return Err(err); - } - }; - stream = stream.with_stats(self.origin.stats()); - let request_id = self.control.next_request_id(); - - // Flushed while waiting for the answer, so the terminals below still end a request - // the peer's send window holds back. - let written = stream.writer.buffer_message( - ietf::TrackStatus::ID, - &ietf::TrackStatus { - authorization_token: None, - request_id, - track_namespace: broadcast_path.clone(), - track_name: track_name.into(), - properties_wanted: true, - }, - ); - if let Err(err) = written { - request.reject(err.clone()); - return Err(err); - } - - // Nobody wanting the track any more ends the wait on a peer that may never answer, - // and the session ending fails it with the session's reason. - let mut demand = request.demand(); - let closed = self.session.clone(); - let this = self.clone(); - let response = { - let mut response = std::pin::pin!(async { - kio::wait(|waiter| stream.writer.poll_flush(waiter)).await?; - self.read_track_status_response(&mut stream).await - }); - loop { - let res = kio::wait(|waiter| { - let response = waiter.poll_future(response.as_mut()); - if let Poll::Ready(Err(err)) = response { - return Poll::Ready(Some(SetupResponse::Response(Err(err)))); - } - if let Err(err) = validity.check(waiter) { - return Poll::Ready(Some(SetupResponse::Cancelled(err))); - } - if gate.poll_denied(waiter).is_ready() { - return Poll::Ready(Some(SetupResponse::Cancelled(Error::Unauthorized))); - } - if let Poll::Ready(res) = response { - return Poll::Ready(Some(SetupResponse::Response(res))); + ) -> Poll> { + loop { + match &mut self.step { + LearnStep::Open(open) => { + let opened = if gate.poll_denied(waiter).is_ready() { + Err(Error::Unauthorized) + } else if let Err(err) = self.validity.check(waiter) { + Err(err) + } else { + ready!(Stream::poll_open(&this.session, open, this.version, waiter)) + }; + let mut stream = match opened { + Ok(stream) => stream.with_stats(this.origin.stats()), + Err(err) => return self.reject(err), + }; + let request_id = this.control.next_request_id(); + + // Flushed while waiting for the answer, so the terminals below still end a + // request the peer's send window holds back. + let written = stream.writer.buffer_message( + ietf::TrackStatus::ID, + &ietf::TrackStatus { + authorization_token: None, + request_id, + track_namespace: broadcast_path.clone(), + track_name: track_name.into(), + properties_wanted: true, + }, + ); + if let Err(err) = written { + return self.reject(err); } - if let Poll::Ready(err) = closed.poll_closed(waiter) { - return Poll::Ready(Some(SetupResponse::Cancelled(Error::from_transport(err)))); + self.stream = Some(stream); + self.demand = self.request.as_ref().map(|request| request.demand()); + self.step = LearnStep::Response; + } + LearnStep::Response => { + let stream = self.stream.as_mut().expect("the request stream is open"); + let response = match stream.writer.poll_flush(waiter) { + Poll::Ready(Ok(())) => { + this.poll_read_track_status_response(&mut stream.reader, &mut self.messages, waiter) + } + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Pending => Poll::Pending, + }; + // Nobody wanting the track any more ends the wait on a peer that may never + // answer, and the session ending fails it with the session's reason. + let res = if let Poll::Ready(Err(err)) = response { + Some(SetupResponse::Response(Err(err))) + } else if let Err(err) = self.validity.check(waiter) { + Some(SetupResponse::Cancelled(err)) + } else if gate.poll_denied(waiter).is_ready() { + Some(SetupResponse::Cancelled(Error::Unauthorized)) + } else if let Poll::Ready(res) = response { + Some(SetupResponse::Response(res)) + } else if let Poll::Ready(err) = this.session.poll_closed(waiter) { + Some(SetupResponse::Cancelled(Error::from_transport(err))) + } else { + // A fetch needs nothing from the status to go out, only to read its objects. + let dynamic = self.dynamic.as_mut().expect("held until accepted"); + while let Poll::Ready(Ok(fill)) = dynamic.poll_requested_fill(waiter) { + group_fetches.push(ServeFill::new( + this.clone(), + broadcast_path.clone(), + track_name.to_owned(), + fill, + self.timescale.clone(), + None, + )); + } + let _ = group_fetches.poll(waiter); + let demand = self.demand.as_mut().expect("set with the request stream"); + ready!(crate::util::poll_unwanted(demand, waiter)); + None + }; + match res { + Some(SetupResponse::Response(response)) => { + let _ = stream.writer.finish(); + self.step = LearnStep::Answered(Some(response)); + } + Some(SetupResponse::Cancelled(err)) => { + if let Some(request) = self.request.take() { + request.reject(err.clone()); + } + stream.reader.abort(&Error::Cancel); + let stream = self.stream.take().expect("the request stream is open"); + self.step = LearnStep::Cancelled(StreamClose::new(stream), Some(err)); + } + None if self + .request + .as_ref() + .is_some_and(|request| request.reject_unwanted(Error::Cancel)) => + { + let mut stream = self.stream.take().expect("the request stream is open"); + stream.reader.abort(&Error::Cancel); + stream.writer.abort(&Error::Cancel); + return Poll::Ready(Err(Error::Cancel)); + } + // Demand returned in the gap. + None => {} } - // A fetch needs nothing from the status to go out, only to read its objects. - while let Poll::Ready(Ok(fill)) = dynamic.poll_requested_fill(waiter) { - let fetch = this.clone().serve_fill( - this.session.work().start(), - broadcast_path.clone(), - track_name.to_owned(), - fill, - timescale.later(), - None, - ); - group_fetches.push(fetch); + } + LearnStep::Cancelled(close, err) => { + let _ = ready!(close.poll(waiter)); + return Poll::Ready(Err(err.take().expect("polled after completion"))); + } + LearnStep::Answered(response) => { + // A track that needs a property we do not understand is not forwarded. + let refusal = match response { + Some(Ok(ok)) => refused( + &ok.properties, + this.version, + ready!(this.peer_setup.poll_get(waiter)).flate, + ), + _ => None, + }; + let ok = match (response.take().expect("answered once"), refusal) { + (Ok(_), Some(err)) => Err(err), + (res, _) => res, + }; + let ok = match ok { + Ok(ok) => ok, + Err(err) => { + tracing::debug!(%err, broadcast = %crate::path::Rooted(this.origin.root(), broadcast_path), track = %track_name, "track status refused"); + return self.reject(err); + } + }; + + tracing::debug!(broadcast = %crate::path::Rooted(this.origin.root(), broadcast_path), track = %track_name, "track status known"); + let request = self.request.take().expect("held until accepted"); + // Nothing feeds this copy live: what it caches is fetched, never the live edge. + // Idle from the accept, so no reader takes it as live in between. + let track = request.idle().stage(accepted_info(&ok.properties, this.version)); + // Marked received before any reader sees it, so none resolves an anchor this + // wire dropped. + track.received(); + track.anchored_cached(); + if let Some(largest) = ok.largest { + track.report_object(largest.group, largest.object); } - let _ = group_fetches.poll(waiter); - crate::util::poll_unwanted(&mut demand, waiter).map(|_| None) - }) - .await; - match res { - Some(res) => break Some(res), - None if request.reject_unwanted(Error::Cancel) => break None, - None => continue, + let track = track.commit(); + return Poll::Ready(Ok(Idle { + demand: track.demand(), + track, + dynamic: self.dynamic.take().expect("held until accepted"), + properties: ok.properties, + finished: false, + })); } } - }; - let response = match response { - Some(SetupResponse::Response(response)) => response, - Some(SetupResponse::Cancelled(err)) => { - request.reject(err.clone()); - stream.reader.abort(&Error::Cancel); - let _ = stream.writer.close().await; - return Err(err); - } - None => { - stream.reader.abort(&Error::Cancel); - stream.writer.abort(&Error::Cancel); - return Err(Error::Cancel); - } - }; - let _ = stream.writer.finish(); - let ok = match response { - // A track that needs a property we do not understand is not forwarded. - Ok(ok) => match refused(&ok.properties, self.version, self.peer_setup.get().await.flate) { - Some(err) => Err(err), - None => Ok(ok), - }, - res => res, - }; - let ok = match ok { - Ok(ok) => ok, - Err(err) => { - tracing::debug!(%err, broadcast = %crate::path::Rooted(self.origin.root(), broadcast_path), track = %track_name, "track status refused"); - request.reject(err.clone()); - return Err(err); - } - }; + } + } +} - tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), broadcast_path), track = %track_name, "track status known"); - // Nothing feeds this copy live: what it caches is fetched, never the live edge. Idle - // from the accept, so no reader takes it as live in between. - let track = request.idle().stage(accepted_info(&ok.properties, self.version)); - // Marked received before any reader sees it, so none resolves an anchor this wire dropped. - track.received(); - track.anchored_cached(); - if let Some(largest) = ok.largest { - track.report_object(largest.group, largest.object); +impl Subscriber { + /// Receive one subgroup stream into its subscription's track; see [`RecvGroup`]. + #[cfg(test)] + pub async fn recv_group( + &mut self, + stream: &mut Reader, + ) -> Result<(), Error> { + let mut recv = RecvGroup::new(); + kio::wait(|waiter| recv.poll(waiter, self, stream)).await + } + + /// Drop an OBJECT_DATAGRAM: a subscription is a stream track's, which takes no datagrams. + /// + /// A malformed datagram is still a protocol error. + pub fn recv_datagram(&self, payload: bytes::Bytes) -> Result<(), Error> { + let (datagram, _) = ietf::ObjectDatagram::decode_slice(&payload, self.version)?; + if let ietf::DatagramBody::Status(status) = datagram.body + && !super::object::status_defined(status) + { + return Err(Error::Session(SessionError::ProtocolViolation)); } - let track = track.commit(); - Ok(Idle { - demand: track.demand(), - track, - dynamic, - properties: ok.properties, - finished: false, - }) + tracing::debug!( + alias = datagram.track_alias, + group = datagram.group_id, + "dropping a datagram on a stream track's subscription" + ); + Ok(()) } +} - /// Read the answer to a TRACK_STATUS: TRACK_STATUS_OK, or the publisher's refusal as an - /// error. - async fn read_track_status_response( - &self, - stream: &mut Stream, - ) -> Result { - let mut messages = super::request_stream::Messages::new(self.session.clone(), self.control.client()); - let super::request_stream::Framed(type_id, body) = messages.message(&mut stream.reader).await?; - let mut data = body.decoder(self.version); +/// Mark where the track ends, as an END_OF_TRACK object said. +/// +/// PUBLISH_DONE carries no end location, so this is what lets a subscriber +/// learn the end before the live edge reaches it. A boundary at or below a group already +/// received is the publisher breaking its own end, which no later group can repair. A +/// marker that lands after the subscription already ended (its grace expired) changes +/// nothing. A marker carrying properties has nowhere to keep them, since the end is +/// track-level, so the track is refused rather than lose them. +fn end_track(track: &mut track::Producer, end: track::Position, properties: bool) -> Result<(), Error> { + if properties { + tracing::warn!(?end, "END_OF_TRACK properties are unsupported"); + let _ = track.clone().abort(Error::Unsupported); + return Err(Error::Unsupported); + } + // Lower groups may still be on the wire, behind the one that carried the end: the + // received copy holds readers at a hole for the grace. + if let Err(err) = track.finish_at(end, None) { + tracing::warn!(%err, ?end, "invalid END_OF_TRACK"); + let _ = track.clone().abort(Error::ProtocolViolation); + return Err(Error::ProtocolViolation); + } + Ok(()) +} - let code = match type_id { - id if id == ietf::TrackStatusOk::ID => { - return Ok(ietf::TrackStatusOk::decode_msg(&mut data, self.version)?); - } - ietf::RequestError::ID if true => ietf::RequestError::decode_msg(&mut data, self.version)?.error_code, - _ => return Err(Error::UnexpectedMessage), +/// How [`Subscriber::open_group`] resolved a subgroup stream. +enum Opened { + /// The group producer the stream writes into, and the Object ID it starts at. + Group(group::Producer, u64), + /// A standalone terminal marker, which creates no payload group, and whether it + /// carried properties. + EndOfTrack { object: u64, properties: bool }, +} + +/// How a subgroup stream ended cleanly. +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum Ended { + /// The stream finished, or carried an explicit end of group. + Group, + /// It carried an END_OF_TRACK after the group's last object, and whether that + /// marker carried properties. + Track { object: u64, properties: bool }, +} + +use super::object::{END_OF_GROUP, END_OF_TRACK}; + +/// The start of a subgroup stream's first object, peeked before its group is created. +#[derive(Debug, Clone, Copy)] +struct FirstObject { + /// The Object ID, which the first object's delta is. + id: u64, + end_of_track: bool, + has_properties: bool, +} + +/// [`FirstObject`] for a stream whose objects carry extensions, or don't. +#[derive(Debug)] +struct PeekFirst(FirstObject); + +/// The first Object ID alone, which names a compact stream's subgroup. +#[derive(Debug)] +struct PeekId(u64); + +impl<'a> Decode<'a, Version> for PeekId { + fn decode(buf: &mut Decoder<'a>, _: Version) -> Result { + Ok(Self(buf.varint()?)) + } +} + +impl<'a, const EXTENSIONS: bool> Decode<'a, Version> for PeekFirst { + fn decode(buf: &mut Decoder<'a>, version: Version) -> Result { + let id = buf.varint()?; + let has_properties = if EXTENSIONS { + let ObjectExtensionsLength(size) = ObjectExtensionsLength::decode(buf, version)?; + buf.slice(size)?; + size != 0 + } else { + false }; - Err(request::from_code(code, request::Kind::TrackStatus, self.version)) + let size = buf.varint()?; + let end_of_track = size == 0 && buf.varint()? == END_OF_TRACK; + Ok(Self(FirstObject { + id, + end_of_track, + has_properties, + })) } +} - pub async fn recv_group( +/// Receives one subgroup stream into its subscription's track. +pub(super) struct RecvGroup { + state: RecvGroupState, +} + +#[kio::machine(self.state: RecvGroupState)] +impl RecvGroup { + pub(super) fn new() -> Self { + Self { + state: RecvGroupState::Start, + } + } + + pub(super) fn poll( &mut self, + waiter: &kio::Waiter, + this: &Subscriber, stream: &mut Reader, - ) -> Result<(), Error> { - let group: ietf::GroupHeader = stream.decode().await?; + ) -> Poll> { + let group: ietf::GroupHeader = wait!(Header => stream.poll_decode(waiter))?; // Parked below every other stream until its alias resolves, so a full window stops it // before the SUBSCRIBE_OK or PUBLISH it waits on. stream.set_priority(crate::priority::RECV_PARKED); if group.sub_group_id != 0 { tracing::warn!(sub_group_id = %group.sub_group_id, "subgroup ID is not supported, dropping stream"); - return Err(Error::Unsupported); + return Poll::Ready(Err(Error::Unsupported)); } // SUBSCRIBE_OK or PUBLISH can be reordered behind this stream. Hold only the // subgroup header while waiting so the data stream cannot consume flow control. - let aliases = self.state.lock().aliases.consume(); - let mut resolving = std::pin::pin!(resolve_track_alias(&self.runtime, aliases, group.track_alias)); + let aliases = this.state.lock().aliases.consume(); + let mut resolving: ResolveAlias = ResolveAlias::new(&this.runtime, aliases, group.track_alias); // Unread meanwhile, so a reset or an eviction would otherwise go unseen, and an evicted // stream holds further eviction off until it is stopped. - let resolved = kio::wait(|waiter| match waiter.poll_future(resolving.as_mut()) { + let resolved = wait!(Resolve => match resolving.poll(waiter) { Poll::Ready(resolved) => Poll::Ready(Ok(resolved)), Poll::Pending => stream.poll_reset(waiter).map(Err), - }) - .await?; + })?; let request_id = match resolved { Ok(request_id) => request_id, // Ours: we cancelled the subscription and the publisher has not stopped yet. Err(err @ Error::Cancel) => { tracing::debug!(track_alias = %group.track_alias, "dropping group for a cancelled subscription"); - return Err(err); + return Poll::Ready(Err(err)); } // Theirs: nothing ever bound this alias. Either the publisher sent data for a // track it never acknowledged, or SUBSCRIBE_OK is more than a timeout behind. @@ -4688,23 +5820,20 @@ impl Subscriber { timeout = ?TRACK_ALIAS_TIMEOUT, "unknown track alias: no SUBSCRIBE_OK bound it" ); - return Err(err); + return Poll::Ready(Err(err)); } }; - let (mut track, timescale, fill, headless, mut reading) = { - let state = self.state.lock(); - let track = state.subscribes.get(&request_id).ok_or(Error::NotFound)?; - ( - track.producer.clone().ok_or(Error::NotFound)?, - track.timescale, - track.fill.clone(), - track.headless_through.is_some_and(|through| group.group_id <= through), - // Every data stream counts toward PUBLISH_DONE's Stream Count, even one - // dropped below, and the subscription's end waits until it is read. - Reading::open(&track.tail, Some(group.group_id), self.runtime.now()), - ) - }; + let state = this.state.lock(); + let held = state.subscribes.get(&request_id).ok_or(Error::NotFound)?; + let mut track: track::Producer = held.producer.clone().ok_or(Error::NotFound)?; + let timescale: Option = held.timescale; + let fill: kio::Producer = held.fill.clone(); + let headless: bool = held.headless_through.is_some_and(|through| group.group_id <= through); + // Every data stream counts toward PUBLISH_DONE's Stream Count, even one dropped + // below, and the subscription's end waits until it is read. + let mut reading: Reading = Reading::open(&held.tail, Some(group.group_id), this.runtime.now()); + drop(state); let (priority, order) = track .subscription() .map_or((0, Default::default()), |sub| (sub.priority, sub.order)); @@ -4716,48 +5845,36 @@ impl Subscriber { // The peek inside blocks until the publisher produces the group's first object, so // race it against the subscription going away the same way the group read below is. // Otherwise dropping the local subscriber cannot end this handler. - let passthrough = !fill.read().outstanding(); - let opened = { - let mut opening = track.clone(); - let mut open = std::pin::pin!(crate::util::sendable(self.open_group( - stream, - &mut opening, - &fill, - headless, - &group, - &mut reading - ))); - kio::wait(|waiter| { - if let Poll::Ready(err) = track.poll_aborted(waiter) { - return Poll::Ready(Err(err)); - } - waiter.poll_future(open.as_mut()) - }) - .await - }; + let passthrough: bool = !fill.read().outstanding(); + let mut opening: track::Producer = track.clone(); + let mut open: OpenGroup = OpenGroup::new(headless); + let opened = wait!(Open => match track.poll_aborted(waiter) { + Poll::Ready(err) => Poll::Ready(Err(err)), + Poll::Pending => open.poll(waiter, stream, opening, fill, group, reading), + }); let opened = match opened { // The group is at or past the end the publisher declared, which no later stream // can repair. Err(Error::Closed) => { tracing::warn!(group = group.group_id, "group past the declared end of track"); let _ = track.abort(Error::ProtocolViolation); - return Err(Error::ProtocolViolation); + return Poll::Ready(Err(Error::ProtocolViolation)); } - Err(err) => return Err(err), + Err(err) => return Poll::Ready(Err(err)), Ok(opened) => opened, }; let (producer, start) = match opened { Opened::Group(producer, start) => (producer, start), // A standalone terminal carries a Location, not a payload group. Opened::EndOfTrack { object, properties } => { - return end_track( + return Poll::Ready(end_track( &mut track, track::Position { group: group.group_id, frame: object, }, properties, - ); + )); } }; @@ -4774,29 +5891,27 @@ impl Subscriber { })?; } - // Guarded: this handler can be dropped at any await below, and a group producer + // Guarded: this handler can be dropped at any wait below, and a group producer // that dies without a terminal leaves its consumer waiting on nothing. - let mut producer = crate::recv::Group::new(producer); - let observing = producer.handle().clone(); + let mut producer: crate::recv::Group = crate::recv::Group::new(producer); + let observing: group::Handle = producer.handle().clone(); - let mut ingest = GroupIngest::new(self.version, &group, timescale); + let mut ingest: GroupIngest = GroupIngest::new(this.version, &group, timescale); ingest.expected_first = Some(start); - let res = { - kio::wait(|waiter| { - if let Poll::Ready(err) = track.poll_aborted(waiter) { - return Poll::Ready(Err(err)); - } - if let Poll::Ready(err) = observing.poll_aborted(waiter) { - return Poll::Ready(Err(err)); - } - ingest.poll(stream, &mut producer, waiter) - }) - .await - }; + let res = wait!(Ingest => { + if let Poll::Ready(err) = track.poll_aborted(waiter) { + Poll::Ready(Err(err)) + } else if let Poll::Ready(err) = observing.poll_aborted(waiter) { + Poll::Ready(Err(err)) + } else { + ingest.poll(stream, producer, waiter) + } + }); - match res { + let res = match res { Err(err @ (Error::Cancel | Error::Stream(crate::StreamError::Cancel))) => { let _ = producer.abort(err); + Ok(()) } // The peer's fault: a group we cannot carry, or one it sent malformed. Err( @@ -4810,325 +5925,255 @@ impl Subscriber { tracing::warn!(%err, group = %producer.sequence, "refusing group"); let _ = producer.abort(err.clone()); // Return the refusal to the dispatcher so it sends STOP_SENDING. - return Err(err); + Err(err) } // The track is malformed, not just this group, so no later group is trusted. Err(Error::MalformedTrack) => { tracing::warn!(group = %producer.sequence, "malformed track"); let _ = producer.abort(Error::MalformedTrack); let _ = track.abort(Error::MalformedTrack); - return Err(Error::MalformedTrack); + Err(Error::MalformedTrack) } // The peer reset the stream, the transport failed, or the track went away. Err(err) => { tracing::debug!(%err, group = %producer.sequence, "group error"); let _ = producer.abort(err); + Ok(()) } Ok(Ended::Group) => { let _ = producer.finish(); + Ok(()) } // No object past this group's last one exists, so the track ends after it. Ok(Ended::Track { object, properties }) => { let _ = producer.finish(); - return end_track( + end_track( &mut track, track::Position { group: group.group_id, frame: object, }, properties, - ); + ) } - } - - Ok(()) - } - - /// Drop an OBJECT_DATAGRAM: a subscription is a stream track's, which takes no datagrams. - /// - /// A malformed datagram is still a protocol error. - pub fn recv_datagram(&self, payload: bytes::Bytes) -> Result<(), Error> { - let (datagram, _) = ietf::ObjectDatagram::decode_slice(&payload, self.version)?; - if let ietf::DatagramBody::Status(status) = datagram.body - && !super::object::status_defined(status) - { - return Err(Error::Session(SessionError::ProtocolViolation)); - } - tracing::debug!( - alias = datagram.track_alias, - group = datagram.group_id, - "dropping a datagram on a stream track's subscription" - ); - Ok(()) + }; + // Held until the group settles above, so a frame the track or group cut short drops + // into an already-aborted group instead of reporting a loss, and the stream counts + // toward the subscription's end for as long as it is read. + drop((ingest, reading)); + Poll::Ready(res) } } -/// Mark where the track ends, as an END_OF_TRACK object said. +/// Resolves a subgroup stream to the group producer it writes into, and the Object ID it +/// starts at. /// -/// PUBLISH_DONE carries no end location, so this is what lets a subscriber -/// learn the end before the live edge reaches it. A boundary at or below a group already -/// received is the publisher breaking its own end, which no later group can repair. A -/// marker that lands after the subscription already ended (its grace expired) changes -/// nothing. A marker carrying properties has nowhere to keep them, since the end is -/// track-level, so the track is refused rather than lose them. -fn end_track(track: &mut track::Producer, end: track::Position, properties: bool) -> Result<(), Error> { - if properties { - tracing::warn!(?end, "END_OF_TRACK properties are unsupported"); - let _ = track.clone().abort(Error::Unsupported); - return Err(Error::Unsupported); - } - // Lower groups may still be on the wire, behind the one that carried the end: the - // received copy holds readers at a hole for the grace. - if let Err(err) = track.finish_at(end, None) { - tracing::warn!(%err, ?end, "invalid END_OF_TRACK"); - let _ = track.clone().abort(Error::ProtocolViolation); - return Err(Error::ProtocolViolation); - } - Ok(()) -} - -/// How [`Subscriber::open_group`] resolved a subgroup stream. -enum Opened { - /// The group producer the stream writes into, and the Object ID it starts at. - Group(group::Producer, u64), - /// A standalone terminal marker, which creates no payload group, and whether it - /// carried properties. - EndOfTrack { object: u64, properties: bool }, -} - -/// How a subgroup stream ended cleanly. -#[derive(Debug, PartialEq, Eq)] -pub(crate) enum Ended { - /// The stream finished, or carried an explicit end of group. - Group, - /// It carried an END_OF_TRACK after the group's last object, and whether that - /// marker carried properties. - Track { object: u64, properties: bool }, -} - -use super::object::{END_OF_GROUP, END_OF_TRACK}; - -/// The start of a subgroup stream's first object, peeked before its group is created. -#[derive(Debug, Clone, Copy)] -struct FirstObject { - /// The Object ID, which the first object's delta is. - id: u64, - end_of_track: bool, - has_properties: bool, +/// The first object is peeked before any producer exists, since a standalone +/// END_OF_TRACK carries no payload group. Normally the stream then starts the +/// group. While a fill is outstanding it may instead be the tail of the group the fill +/// fetch stream began, which the first Object ID decides. +struct OpenGroup { + /// Whether a fill an update replaced failed to head this group, or a later one. + headless: bool, + step: OpenStep, } -/// [`FirstObject`] for a stream whose objects carry extensions, or don't. -#[derive(Debug)] -struct PeekFirst(FirstObject); - -/// The first Object ID alone, which names a compact stream's subgroup. -#[derive(Debug)] -struct PeekId(u64); - -impl<'a> Decode<'a, Version> for PeekId { - fn decode(buf: &mut Decoder<'a>, _: Version) -> Result { - Ok(Self(buf.varint()?)) - } +enum OpenStep { + /// Peeking the compact form's first Object ID, which names its subgroup. + PeekId, + /// Handing a promised head to its live tail on the subgroup header alone. + ClaimEarly, + /// Peeking the first object. + PeekFirst, + /// Claiming the head a mid-group stream continues, from its first Object ID if it has one. + ClaimTail(Option), } -impl<'a, const EXTENSIONS: bool> Decode<'a, Version> for PeekFirst { - fn decode(buf: &mut Decoder<'a>, version: Version) -> Result { - let id = buf.varint()?; - let has_properties = if EXTENSIONS { - let ObjectExtensionsLength(size) = ObjectExtensionsLength::decode(buf, version)?; - buf.slice(size)?; - size != 0 - } else { - false - }; - let size = buf.varint()?; - let end_of_track = size == 0 && buf.varint()? == END_OF_TRACK; - Ok(Self(FirstObject { - id, - end_of_track, - has_properties, - })) +impl OpenGroup { + fn new(headless: bool) -> Self { + Self { + headless, + step: OpenStep::PeekId, + } } -} -impl Subscriber { - /// The group producer this subgroup stream writes into, and the Object ID it starts at. - /// - /// The first object is peeked before any producer exists, since a standalone - /// END_OF_TRACK carries no payload group. Normally the stream then starts the - /// group. While a fill is outstanding it may instead be the tail of the group the fill - /// fetch stream began, which the first Object ID decides. - async fn open_group( - &self, + fn poll( + &mut self, + waiter: &kio::Waiter, stream: &mut Reader, track: &mut track::Producer, fill: &kio::Producer, - // Whether a fill an update replaced failed to head this group, or a later one. - headless: bool, header: &ietf::GroupHeader, reading: &mut Reading, - ) -> Result { + ) -> Poll> { let sequence = header.group_id; // Stats (groups/frames/bytes) are counted in the model as the group is written, // through the tagged `track::Producer`. let create = |track: &mut track::Producer| track.create_group(group::Info { sequence }); - // The compact form names the subgroup by its first Object ID, so check it before the - // rest of the object can fail and open the group as subgroup 0. A stream with no - // objects names none, and a failed read fails again in the full peek below. - if header.flags.has_subgroup_object - && let Ok(Some(PeekId(id))) = stream.decode_peek_maybe::().await - && id != 0 - { - tracing::warn!(sub_group_id = id, "subgroup ID is not supported, dropping stream"); - return Err(Error::Unsupported); - } - - // The subgroup header is enough to hand a promised head to its live tail. - // Waiting for the next object would hold every later range update on a quiet group. - let outstanding = fill.read().outstanding(); - if !header.flags.first_object - && !header.flags.has_subgroup_object - && outstanding - && let Some((producer, next)) = self.claim_fill(fill, track, sequence, None, reading).await? - { - return Ok(Opened::Group(producer, next)); - } + loop { + match self.step { + OpenStep::PeekId => { + // The compact form names the subgroup by its first Object ID, so check it + // before the rest of the object can fail and open the group as subgroup 0. A + // stream with no objects names none, and a failed read fails again in the + // full peek below. + if header.flags.has_subgroup_object + && let Ok(Some(PeekId(id))) = ready!(stream.poll_decode_peek_maybe::(waiter)) + && id != 0 + { + tracing::warn!(sub_group_id = id, "subgroup ID is not supported, dropping stream"); + return Poll::Ready(Err(Error::Unsupported)); + } - let peeked = match header.flags.has_extensions { - true => stream - .decode_peek_maybe::>() - .await - .map(|peek| peek.map(|peek| peek.0)), - false => stream - .decode_peek_maybe::>() - .await - .map(|peek| peek.map(|peek| peek.0)), - }; - let first = match peeked { - Ok(first) => first, - // The header arrived, so the stream counts toward the track's end. Abort the - // group it named rather than let the track end clean without it. A fill that - // already holds the group reports it through its own producer. - Err(err) => { - if let Ok(group) = create(track) { - let _ = group.abort(err.clone()); + // The subgroup header is enough to hand a promised head to its live tail. + // Waiting for the next object would hold every later range update on a quiet + // group. + let outstanding = fill.read().outstanding(); + self.step = match !header.flags.first_object && !header.flags.has_subgroup_object && outstanding { + true => OpenStep::ClaimEarly, + false => OpenStep::PeekFirst, + }; } - return Err(err); - } - }; - if let Some(first) = first.filter(|first| first.end_of_track) { - return Ok(Opened::EndOfTrack { - object: first.id, - properties: first.has_properties, - }); - } - - // Without a fill, a relay preserves the subgroup's first Object ID and FIRST_OBJECT - // claim independently. Plain readers see only frame ordinals. A head that was owed - // and never came is the exception: the stream is the same unstitchable tail it is - // when it arrives before the fill settles, below. - if !fill.read().outstanding() { - let start = first.map_or(0, |first| first.id); - if start != 0 && (headless || fill.read().headless(sequence)) { - tracing::warn!(sequence, start, "no fill to stitch a mid-group stream onto"); - return Err(Error::Unsupported); - } - return Ok(Opened::Group(create(track)?, start)); - } - - // The first object's ID delta is its absolute Object ID (see `next_object_id`). - match first.map(|first| first.id) { - // A group delivered from its start stands alone, unless the fill already - // headed this very sequence: the publisher then served those objects twice, - // and the model has one producer per group. Refuse both unfinished pieces - // rather than cache a cut group as complete. - Some(0) => { - let headed = matches!(*fill.read(), Fill::Ready { sequence: s, .. } if s == sequence); - if headed { - tracing::warn!(sequence, "a whole group arrived for one the fill already headed"); - if let Ok(mut state) = fill.write() { - state.cancel(); + OpenStep::ClaimEarly => { + if let Some((producer, next)) = + ready!(poll_claim_fill(fill, track, sequence, None, reading, waiter))? + { + return Poll::Ready(Ok(Opened::Group(producer, next))); } - return Err(Error::Unsupported); + self.step = OpenStep::PeekFirst; } + OpenStep::PeekFirst => { + let peeked = match header.flags.has_extensions { + true => ready!(stream.poll_decode_peek_maybe::>(waiter)) + .map(|peek| peek.map(|peek| peek.0)), + false => ready!(stream.poll_decode_peek_maybe::>(waiter)) + .map(|peek| peek.map(|peek| peek.0)), + }; + let first = match peeked { + Ok(first) => first, + // The header arrived, so the stream counts toward the track's end. Abort + // the group it named rather than let the track end clean without it. A + // fill that already holds the group reports it through its own producer. + Err(err) => { + if let Ok(group) = create(track) { + let _ = group.abort(err.clone()); + } + return Poll::Ready(Err(err)); + } + }; + if let Some(first) = first.filter(|first| first.end_of_track) { + return Poll::Ready(Ok(Opened::EndOfTrack { + object: first.id, + properties: first.has_properties, + })); + } - Ok(Opened::Group(create(track)?, 0)) - } + // Without a fill, a relay preserves the subgroup's first Object ID and + // FIRST_OBJECT claim independently. Plain readers see only frame ordinals. A + // head that was owed and never came is the exception: the stream is the same + // unstitchable tail it is when it arrives before the fill settles, below. + if !fill.read().outstanding() { + let start = first.map_or(0, |first| first.id); + if start != 0 && (self.headless || fill.read().headless(sequence)) { + tracing::warn!(sequence, start, "no fill to stitch a mid-group stream onto"); + return Poll::Ready(Err(Error::Unsupported)); + } + return Poll::Ready(Ok(Opened::Group(create(track)?, start))); + } - // A group starting partway through is the tail of one the fill began, and - // without that head it has a hole at the front. - Some(start) => match self.claim_fill(fill, track, sequence, Some(start), reading).await? { - Some((producer, _)) => Ok(Opened::Group(producer, start)), - None => { - tracing::warn!(sequence, start, "no fill to stitch a mid-group stream onto"); - Err(Error::Unsupported) - } - }, + // The first object's ID delta is its absolute Object ID (see `next_object_id`). + match first.map(|first| first.id) { + // A group delivered from its start stands alone, unless the fill already + // headed this very sequence: the publisher then served those objects + // twice, and the model has one producer per group. Refuse both unfinished + // pieces rather than cache a cut group as complete. + Some(0) => { + let headed = matches!(*fill.read(), Fill::Ready { sequence: s, .. } if s == sequence); + if headed { + tracing::warn!(sequence, "a whole group arrived for one the fill already headed"); + if let Ok(mut state) = fill.write() { + state.cancel(); + } + return Poll::Ready(Err(Error::Unsupported)); + } - // A stream that ends without an object: the group is over and had nothing - // outside the fill's range, so the head it delivered is the whole group. - None => match self.claim_fill(fill, track, sequence, None, reading).await? { - Some((producer, next)) => Ok(Opened::Group(producer, next)), - None => Ok(Opened::Group(create(track)?, 0)), - }, + return Poll::Ready(Ok(Opened::Group(create(track)?, 0))); + } + // A group starting partway through is the tail of one the fill began, and + // without that head it has a hole at the front. A stream that ends without + // an object: the group is over and had nothing outside the fill's range, + // so the head it delivered is the whole group. + start => self.step = OpenStep::ClaimTail(start), + } + } + OpenStep::ClaimTail(start) => { + let claimed = ready!(poll_claim_fill(fill, track, sequence, start, reading, waiter))?; + return Poll::Ready(match (start, claimed) { + (Some(start), Some((producer, _))) => Ok(Opened::Group(producer, start)), + (Some(start), None) => { + tracing::warn!(sequence, start, "no fill to stitch a mid-group stream onto"); + Err(Error::Unsupported) + } + (None, Some((producer, next))) => Ok(Opened::Group(producer, next)), + (None, None) => Ok(Opened::Group(create(track)?, 0)), + }); + } + } } } +} - /// Take the head the fill fetch stream delivered for `sequence`, once it has finished - /// writing it. - /// - /// The model has one producer per group, so this is the handoff: the fill owns the - /// producer while it writes objects `0..next`, and the subgroup stream carrying the rest - /// picks it up here. `start` is the Object ID that stream begins at, or `None` when it - /// carries no objects at all and simply ends the group. - /// - /// Waiting is what keeps the two streams from interleaving into one producer. It ends - /// with the subscription, so a publisher that promises a fill and never delivers one - /// costs this stream and nothing else. Meanwhile this stream stops holding the - /// subscription's end open: the fill's own stream holds it if its header arrived, and - /// the grace gives up on it if not. - async fn claim_fill( - &self, - fill: &kio::Producer, - track: &track::Producer, - sequence: u64, - start: Option, - reading: &mut Reading, - ) -> Result, Error> { - reading.park(); - let settled = kio::wait(|waiter| { - if let Poll::Ready(err) = track.poll_aborted(waiter) { - return Poll::Ready(Err(err)); - } - - let settled = fill.poll(waiter, |fill| match **fill { +/// Take the head the fill fetch stream delivered for `sequence`, once it has finished +/// writing it. +/// +/// The model has one producer per group, so this is the handoff: the fill owns the +/// producer while it writes objects `0..next`, and the subgroup stream carrying the rest +/// picks it up here. `start` is the Object ID that stream begins at, or `None` when it +/// carries no objects at all and simply ends the group. +/// +/// Waiting is what keeps the two streams from interleaving into one producer. It ends +/// with the subscription, so a publisher that promises a fill and never delivers one +/// costs this stream and nothing else. Meanwhile this stream stops holding the +/// subscription's end open: the fill's own stream holds it if its header arrived, and +/// the grace gives up on it if not. +fn poll_claim_fill( + fill: &kio::Producer, + track: &track::Producer, + sequence: u64, + start: Option, + reading: &mut Reading, + waiter: &kio::Waiter, +) -> Poll, Error>> { + reading.park(); + let settled = ready!(match track.poll_aborted(waiter) { + Poll::Ready(err) => Poll::Ready(Err(err)), + Poll::Pending => fill + .poll(waiter, |fill| match **fill { Fill::Requested | Fill::Serving(_) | Fill::Active => Poll::Pending, Fill::Ready { .. } | Fill::Done | Fill::Failed { .. } | Fill::Ended => Poll::Ready(()), - }); - - match settled { - Poll::Ready(Ok(_)) => Poll::Ready(Ok(())), + }) + .map(|settled| match settled { + Ok(_) => Ok(()), // Fill is never explicitly closed, and this task holds a producer. - Poll::Ready(Err(_)) => { + Err(_) => { debug_assert!(false, "fill closed while a receive task owns it"); - Poll::Ready(Err(Error::Cancel)) + Err(Error::Cancel) } - Poll::Pending => Poll::Pending, - } - }) - .await; - // Hold the end open again before taking the head: once the fill is claimed, this - // stream is the only thing saying the group is still being read. - reading.resume(); - settled?; + }), + }); + // Hold the end open again before taking the head: once the fill is claimed, this + // stream is the only thing saying the group is still being read. + reading.resume(); + settled?; + Poll::Ready( fill.write() .map_err(|_| { debug_assert!(false, "fill closed while claiming its head"); Error::Cancel })? - .claim(sequence, start) - } + .claim(sequence, start), + ) } /// Pumps moq-transport subgroup objects from a reader into a group producer: @@ -5187,54 +6232,95 @@ impl GroupIngest { } impl Subscriber { - /// Read a fill fetch stream: the head of the group a subscription joins part way through. - /// - /// A draft-22 fill answers the FILL_PARAMETERS we sent and is named by the SUBSCRIBE's - /// Request ID. A draft-18 joining FETCH has its own request id, mapped back to that - /// subscription. Unlike a subgroup stream it needs no track alias. It writes the objects - /// into a group producer of its own and hands that to the subgroup stream carrying the - /// rest of the group; see [`Fill`]. A reset stream is the publisher's fill-failure - /// signal, and arrives here as a read error, which drops the head and the join with it. + /// Read a fill fetch stream, or a group FETCH's; see [`RecvFill`]. + #[cfg(test)] pub async fn recv_fill( &mut self, stream: &mut Reader, ) -> Result<(), Error> { - // The dispatcher peeked the stream type to get here. - let _ = stream.varint().await?; - let header: ietf::FetchHeader = stream.decode().await?; - - let group_fetch = self.state.lock().group_fetches.get(&header.request_id).cloned(); - if let Some(slot) = group_fetch { - return self.recv_group_fetch(stream, slot).await; - } - - let (subscribe_id, fill, _counted, multiple) = { - let state = self.state.lock(); - // A draft-22 fill is named by the SUBSCRIBE's request id. A draft-18 joining - // FETCH has its own id, which `fetches` maps back to that subscription. - let joined = state.fetches.get(&header.request_id).copied(); - let subscribe_id = joined.map(FillRequest::subscription).unwrap_or(header.request_id); - let track = state.subscribes.get(&subscribe_id).ok_or(Error::NotFound)?; - // Once an update owns the fill, a delayed initial fill cannot claim its slot. - if joined.is_none() && track.fetch_id.is_some() { - return Err(Error::Unsupported); - } - // 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 = (!matches!(joined, Some(FillRequest::Joining(_)))) - .then(|| Reading::open(&track.tail, None, self.runtime.now())); - ( - subscribe_id, - track.fill.clone(), - counted, - !matches!(joined, Some(FillRequest::Joining(_))), - ) - }; + let mut recv = RecvFill::new(); + kio::wait(|waiter| recv.poll(waiter, self, stream)).await + } + + /// Read the fill's objects into a group producer of its own; see [`RunFill`]. + #[cfg(test)] + async fn run_fill( + &mut self, + stream: &mut Reader, + track: &mut track::Producer, + timescale: Option, + multiple: bool, + ) -> Result { + let mut fill = RunFill::new(self.version, timescale, multiple); + kio::wait(|waiter| fill.poll(waiter, stream, track)).await + } +} + +/// Reads a fill fetch stream: the head of the group a subscription joins part way through. +/// +/// A draft-22 fill answers the FILL_PARAMETERS we sent and is named by the SUBSCRIBE's +/// Request ID. A draft-18 joining FETCH has its own request id, mapped back to that +/// subscription. Unlike a subgroup stream it needs no track alias. It writes the objects +/// into a group producer of its own and hands that to the subgroup stream carrying the +/// rest of the group; see [`Fill`]. A reset stream is the publisher's fill-failure +/// signal, and arrives here as a read error, which drops the head and the join with it. +/// A group FETCH's stream is handed to [`RecvGroupFetch`] instead. +pub(super) struct RecvFill { + state: RecvFillState, +} + +#[kio::machine(self.state: RecvFillState)] +impl RecvFill { + pub(super) fn new() -> Self { + Self { + state: RecvFillState::Start, + } + } + + pub(super) fn poll( + &mut self, + waiter: &kio::Waiter, + this: &Subscriber, + stream: &mut Reader, + ) -> Poll> { + // The dispatcher peeked the stream type to get here. + wait!(Kind => stream.poll_varint(waiter))?; + let request_id: RequestId = + wait!(Header => stream.poll_decode::(waiter).map_ok(|header| header.request_id))?; + + let group_fetch: Option = this + .state + .lock() + .group_fetches + .get(&request_id) + .cloned() + .map(|slot| RecvGroupFetch::new(this.version, slot)); + wait!(GroupFetch => match group_fetch { + Some(fetch) => return fetch.poll(waiter, stream), + None => Poll::Ready(()), + }); + + let lookup = this.state.lock(); + // A draft-22 fill is named by the SUBSCRIBE's request id. A draft-18 joining + // FETCH has its own id, which `fetches` maps back to that subscription. + let joined = lookup.fetches.get(&request_id).copied(); + let subscribe_id: RequestId = joined.map(FillRequest::subscription).unwrap_or(request_id); + let entry = lookup.subscribes.get(&subscribe_id).ok_or(Error::NotFound)?; + // Once an update owns the fill, a delayed initial fill cannot claim its slot. + if joined.is_none() && entry.fetch_id.is_some() { + return Poll::Ready(Err(Error::Unsupported)); + } + // 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 multiple: bool = !matches!(joined, Some(FillRequest::Joining(_))); + let counted: Option = multiple.then(|| Reading::open(&entry.tail, None, this.runtime.now())); + let fill: kio::Producer = entry.fill.clone(); + drop(lookup); // SUBSCRIBE_OK declares the units these object timestamps are in, and this stream can // be reordered ahead of it. Taking the fill in the same step is what refuses a second // stream for a request that asked for one fill. - let timescale = kio::wait(|waiter| { + let timescale: Option = wait!(Accept => { let accepted = fill.poll(waiter, |fill| match **fill { Fill::Requested => Poll::Pending, _ => Poll::Ready(()), @@ -5257,37 +6343,35 @@ impl Subscriber { } Poll::Pending => Poll::Pending, } - }) - .await?; - let _filling = Filling(fill.clone()); + })?; + let filling: Filling = Filling(fill.clone()); // A fill stream can overtake SUBSCRIBE_OK. The fill gate above is released only // after that response commits track Info and installs the producer. - let (track, largest) = { - let state = self.state.lock(); - let held = state.subscribes.get(&subscribe_id).ok_or(Error::NotFound)?; - (held.producer.clone().ok_or(Error::NotFound)?, held.fill_largest) + let installed = this + .state + .lock() + .subscribes + .get(&subscribe_id) + .map(|subscription| (subscription.producer.clone(), subscription.fill_largest)); + let Some((Some(producer), fill_largest)) = installed else { + // The guard cancels an unfinished head before the count lets the end go. + drop(filling); + drop(counted); + return Poll::Ready(Err(Error::NotFound)); }; + let track: track::Producer = producer; + let largest: Option = fill_largest; // Race the peer's stream against the subscription going away, the same way a // subgroup stream is served. Otherwise a peer that stalls partway through a payload // keeps this handler and its stream alive for as long as it cares to: aborting the // track does not close a group producer, since those lifecycles are independent. - let res = { - let mut serving = track.clone(); - let mut serve = std::pin::pin!(crate::util::sendable(self.run_fill( - stream, - &mut serving, - timescale, - multiple - ))); - kio::wait(|waiter| { - if let Poll::Ready(err) = track.poll_aborted(waiter) { - return Poll::Ready(Err(err)); - } - waiter.poll_future(serve.as_mut()) - }) - .await - }; + let serving: track::Producer = track.clone(); + let run: RunFill = RunFill::new(this.version, timescale, multiple); + let res = wait!(Serve => match track.poll_aborted(waiter) { + Poll::Ready(err) => Poll::Ready(Err(err)), + Poll::Pending => run.poll(waiter, stream, serving), + }); let mut head = match res { Ok(head) => head, @@ -5299,7 +6383,9 @@ impl Subscriber { if matches!(err, Error::MalformedTrack) { let _ = track.abort(Error::MalformedTrack); } - return Err(err); + drop(filling); + drop(counted); + return Poll::Ready(Err(err)); } }; @@ -5312,39 +6398,88 @@ impl Subscriber { } // The subscription can end while the head is being written, and its teardown cannot // reach a producer this task still owns. So the handoff is where that is settled. - match fill.write() { - Ok(mut state) => state.install(head), + let res = match fill.write() { + Ok(mut state) => { + state.install(head); + Ok(()) + } // Fill is never explicitly closed, and this task holds a producer. Err(_) => { debug_assert!(false, "fill closed while installing its head"); let mut head = head; head.cancel(); - return Err(Error::Cancel); + Err(Error::Cancel) } - } + }; + // Held to the end, the guard dropping first: it cancels an unfinished head before the + // count, which holds the subscription's end open while this stream is read, lets go. + drop(filling); + drop(counted); + Poll::Ready(res) + } +} - Ok(()) +/// Reads a fill's objects into a group producer of its own. +/// +/// Ready with the head for the live tail to claim, or [`Fill::Done`] when the stream +/// carried no objects at all. Aborts the producer on the way out of an error, since a +/// half-written head is a group with no end. +/// +/// Earlier groups finish at the next group's first object; the final group's head +/// waits for its live tail. A Prior Object ID Gap accounts for every missing ID. +struct RunFill { + version: Version, + timescale: Option, + /// Whether the stream may carry several groups, as a widening's fill does. + multiple: bool, + /// The group being written, the next Object ID it expects, and its producer. + head: Option<(u64, u64, crate::recv::Group)>, + prior_group: Option, + /// A group the copy already owns, read past, and the next Object ID it expects. + skipped: Option<(u64, u64)>, + budget: kio::coop::Budget, + first: bool, + step: FillStep, +} + +#[allow( + clippy::large_enum_variant, + reason = "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; an allow, since the size gap varies by target" +)] +enum FillStep { + /// Reading the next object's header. + Header, + /// Yielding before taking in the object just read. + Yield(FetchedObject), + /// Reading past an object of a group the copy already owns. + Skip(SkipPayload, u64), + /// Writing an object's payload into the head. + Payload(FetchPayloadRecv, u64), +} + +impl RunFill { + fn new(version: Version, timescale: Option, multiple: bool) -> Self { + Self { + version, + timescale, + multiple, + head: None, + prior_group: None, + skipped: None, + budget: kio::coop::Budget::new(32), + first: true, + step: FillStep::Header, + } } - /// Read the fill's objects into a group producer of its own. - /// - /// Returns the head for the live tail to claim, or [`Fill::Done`] when the stream - /// carried no objects at all. Aborts the producer on the way out of an error, since a - /// half-written head is a group with no end. - async fn run_fill( + fn poll( &mut self, + waiter: &kio::Waiter, stream: &mut Reader, track: &mut track::Producer, - timescale: Option, - multiple: bool, - ) -> Result { - let mut head: Option<(u64, u64, crate::recv::Group)> = None; - - match self - .run_fill_objects(stream, track, timescale, &mut head, multiple) - .await - { - Ok(()) => Ok(match head { + ) -> Poll> { + Poll::Ready(match ready!(self.poll_objects(waiter, stream, track)) { + Ok(()) => Ok(match self.head.take() { Some((sequence, next, producer)) => Fill::Ready { sequence, next, @@ -5353,438 +6488,679 @@ impl Subscriber { None => Fill::Done, }), Err(err) => { - if let Some((_, _, producer)) = head { + if let Some((_, _, producer)) = self.head.take() { let _ = producer.abort(err.clone()); } Err(err) } - } + }) } /// Decode fetch objects into `head`, creating its group from the first object's /// absolute IDs. - /// - /// Earlier groups finish at the next group's first object; the final group's head - /// waits for its live tail. A Prior Object ID Gap accounts for every missing ID. - async fn run_fill_objects( + fn poll_objects( &mut self, + waiter: &kio::Waiter, stream: &mut Reader, track: &mut track::Producer, - timescale: Option, - head: &mut Option<(u64, u64, crate::recv::Group)>, - multiple: bool, - ) -> Result<(), Error> { - let mut prior_group = None; - let mut skipped: Option<(u64, u64)> = None; - let mut budget = kio::coop::Budget::new(32); - let mut first = true; - while let Some(object) = decode_fetch_object(stream, std::mem::take(&mut first)).await? { - kio::wait(|waiter| budget.poll_yield(waiter)).await; - if !object.datagram && !object.subgroup_ok { - tracing::warn!("subgroup ID is not supported, dropping fill"); - return Err(Error::Unsupported); + ) -> Poll> { + loop { + match &mut self.step { + FillStep::Header => { + let Some(object) = ready!(poll_fetch_object(stream, self.first, waiter))? else { + return Poll::Ready(Ok(())); + }; + self.first = false; + self.step = FillStep::Yield(object); + } + FillStep::Yield(_) => { + ready!(self.budget.poll_yield(waiter)); + let FillStep::Yield(object) = std::mem::replace(&mut self.step, FillStep::Header) else { + unreachable!() + }; + self.step = self.take(object, track)?; + } + FillStep::Skip(skip, id) => { + ready!(skip.poll(waiter, stream))?; + let (_, next) = self.skipped.as_mut().expect("skipping a group"); + *next = id.checked_add(1).ok_or(Error::ProtocolViolation)?; + self.step = FillStep::Header; + } + FillStep::Payload(recv, id) => { + let (_, next, producer) = self.head.as_mut().expect("the head was created above"); + ready!(recv.poll(waiter, stream, producer))?; + *next = id.checked_add(1).ok_or(Error::ProtocolViolation)?; + self.step = FillStep::Header; + } } + } + } - // Object ID still keys off whether the wire Group ID field was present. - let group = resolve_fetch_group(prior_group, object.group)?; - if let (Some(prior), Some(group)) = (prior_group, group) - && group < prior - { - return Err(Error::ProtocolViolation); - } - if let Some(sequence) = group { - prior_group = Some(sequence); - } + /// Take in one object's header, opening its group when it starts one, and say how its + /// payload is read. + fn take(&mut self, object: FetchedObject, track: &mut track::Producer) -> Result { + if !object.datagram && !object.subgroup_ok { + tracing::warn!("subgroup ID is not supported, dropping fill"); + return Err(Error::Unsupported); + } - // A datagram is never a stream track's, so a FETCH stream has no group to rebuild. - if object.datagram { - tracing::debug!(?group, "a datagram group is not fetchable"); - return Err(Error::NotFound); - } + // Object ID still keys off whether the wire Group ID field was present. + let group = resolve_fetch_group(self.prior_group, object.group)?; + if let (Some(prior), Some(group)) = (self.prior_group, group) + && group < prior + { + return Err(Error::ProtocolViolation); + } + if let Some(sequence) = group { + self.prior_group = Some(sequence); + } - let raw = bytes::Bytes::from(object.properties.unwrap_or_default()); - let decoded = super::object::read(&raw, self.version, timescale)?; - let expected = head - .as_ref() - .filter(|(sequence, _, _)| group.is_none_or(|group| group == *sequence)) - .map(|(_, next, _)| *next) - .or_else(|| { - skipped - .filter(|(sequence, _)| group.is_none_or(|group| group == *sequence)) - .map(|(_, next)| next) - }) - .unwrap_or(0); - let id = match (object.group.is_some(), object.object) { - (true, Some(id)) => id, - (false, Some(delta)) => expected - .checked_sub(1) - .and_then(|prior| prior.checked_add(delta)) - .ok_or(Error::ProtocolViolation)?, - (false, None) => expected, - _ => return Err(Error::ProtocolViolation), - }; - if id < expected { - return Err(Error::ProtocolViolation); - } - if id > expected && decoded.gap.is_none_or(|gap| gap < id - expected) { - return Err(Error::Unsupported); - } - let current = head - .as_ref() - .map(|(sequence, _, _)| *sequence) - .or_else(|| skipped.map(|(sequence, _)| sequence)); - if current.is_some() && group.is_some_and(|group| Some(group) != current) && !multiple { - return Err(Error::Unsupported); - } - if current.is_none() || group.is_some_and(|group| Some(group) != current) { - if let Some((_, _, producer)) = head.take() { - producer.finish()?; - } - skipped = None; - let sequence = group.ok_or(Error::ProtocolViolation)?; - match open_fill_group(track, head, sequence) { - Ok(()) => {} - // Widening a range can fill groups the copy already owns. Drain their - // objects without replacing their producer or cutting the rest of the fill. - Err(Error::Duplicate) if multiple => skipped = Some((sequence, 0)), - Err(err) => return Err(err), - } - } - if let Some((_, next)) = skipped.as_mut() { - skip_fetch_payload(stream).await?; - *next = id.checked_add(1).ok_or(Error::ProtocolViolation)?; - continue; - } + // A datagram is never a stream track's, so a FETCH stream has no group to rebuild. + if object.datagram { + tracing::debug!(?group, "a datagram group is not fetchable"); + return Err(Error::NotFound); + } - let (_, next, producer) = head.as_mut().expect("the head was created above"); - if *next == 0 - && let Some(mut wire) = producer.wire() - { - wire.first_id = id; - producer.set_wire(wire)?; + let raw = bytes::Bytes::from(object.properties.unwrap_or_default()); + let decoded = super::object::read(&raw, self.version, self.timescale)?; + let expected = self + .head + .as_ref() + .filter(|(sequence, _, _)| group.is_none_or(|group| group == *sequence)) + .map(|(_, next, _)| *next) + .or_else(|| { + self.skipped + .filter(|(sequence, _)| group.is_none_or(|group| group == *sequence)) + .map(|(_, next)| next) + }) + .unwrap_or(0); + let id = match (object.group.is_some(), object.object) { + (true, Some(id)) => id, + (false, Some(delta)) => expected + .checked_sub(1) + .and_then(|prior| prior.checked_add(delta)) + .ok_or(Error::ProtocolViolation)?, + (false, None) => expected, + _ => return Err(Error::ProtocolViolation), + }; + if id < expected { + return Err(Error::ProtocolViolation); + } + if id > expected && decoded.gap.is_none_or(|gap| gap < id - expected) { + return Err(Error::Unsupported); + } + let current = self + .head + .as_ref() + .map(|(sequence, _, _)| *sequence) + .or_else(|| self.skipped.map(|(sequence, _)| sequence)); + if current.is_some() && group.is_some_and(|group| Some(group) != current) && !self.multiple { + return Err(Error::Unsupported); + } + if current.is_none() || group.is_some_and(|group| Some(group) != current) { + if let Some((_, _, producer)) = self.head.take() { + producer.finish()?; + } + self.skipped = None; + let sequence = group.ok_or(Error::ProtocolViolation)?; + match open_fill_group(track, &mut self.head, sequence) { + Ok(()) => {} + // Widening a range can fill groups the copy already owns. Drain their + // objects without replacing their producer or cutting the rest of the fill. + Err(Error::Duplicate) if self.multiple => self.skipped = Some((sequence, 0)), + Err(err) => return Err(err), } - self.recv_fetch_payload(stream, producer, FetchPayload { decoded, timescale, id }) - .await?; - *next = id.checked_add(1).ok_or(Error::ProtocolViolation)?; + } + if self.skipped.is_some() { + return Ok(FillStep::Skip(SkipPayload::new(), id)); } - Ok(()) + let (_, next, producer) = self.head.as_mut().expect("the head was created above"); + if *next == 0 + && let Some(mut wire) = producer.wire() + { + wire.first_id = id; + producer.set_wire(wire)?; + } + let timescale = self.timescale; + Ok(FillStep::Payload( + FetchPayloadRecv::new(FetchPayload { decoded, timescale, id }), + id, + )) } +} - /// Read one fetch object's length and payload into `producer`, after its header. - async fn recv_fetch_payload( - &self, - stream: &mut Reader, - producer: &mut crate::recv::Group, - payload: FetchPayload, - ) -> Result<(), Error> { - let mut recv = FetchPayloadRecv::new(payload); - kio::wait(|waiter| recv.poll(waiter, stream, producer)).await +impl Subscriber { + /// Serve a whole fill with one FETCH, as [`ServeFill`] serves each run. + #[cfg(test)] + async fn run_group_fetch(self, broadcast: Path, name: String, fill: track::Fill, timescale: Option) { + let run = fill.start().group..fill.end(); + let mut fetch = FetchRun::new( + self, + broadcast, + name, + std::sync::Arc::new(fill), + run, + Later::Known(timescale), + ); + kio::wait(|waiter| fetch.poll(waiter)).await } +} - /// Serve a fill a run at a time, newest run first for a descending reader, until it is - /// answered or nobody waits on it. Each run is one FETCH, except that the source's live - /// group always goes alone: a run including the group its source is still writing would - /// hold the earlier groups until it finishes, or be cut short at the Largest Object by a - /// publisher following draft-22 section 9.11. - /// - /// `heading` is the subscription's own fill, when it may be bringing a joined group's - /// head. That group waits for it instead: a FETCH racing it can cache the group first, - /// and then the tail stream, which is what ends the group for live readers, is refused - /// as a duplicate. The cache answers the group once the head lands, and a fill that - /// settles without one leaves it to a FETCH. Only that group waits, and only while a - /// reader still wants it. - /// - /// `_working` counts the fill from when it is queued, before its first poll. - async fn serve_fill( - self, - _working: crate::session::Working, +/// Serves a fill a run at a time, newest run first for a descending reader, until it is +/// answered or nobody waits on it. Each run is one FETCH, except that the source's live +/// group always goes alone: a run including the group its source is still writing would +/// hold the earlier groups until it finishes, or be cut short at the Largest Object by a +/// publisher following draft-22 section 9.11. +/// +/// `heading` is the subscription's own fill, when it may be bringing a joined group's +/// head. That group waits for it instead: a FETCH racing it can cache the group first, +/// and then the tail stream, which is what ends the group for live readers, is refused +/// as a duplicate. The cache answers the group once the head lands, and a fill that +/// settles without one leaves it to a FETCH. Only that group waits, and only while a +/// reader still wants it. +struct ServeFill { + this: Subscriber, + /// Counts the fill from when it is queued, before its first poll. + _working: crate::session::Working, + broadcast: Path, + name: String, + fill: std::sync::Arc, + timescale: Later>, + heading: Option, + /// The run fetched last, which coming back means it went unanswered. + last: Option>, + step: ServeStep, +} + +#[allow( + clippy::large_enum_variant, + reason = "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; an allow, since the size gap varies by target" +)] +enum ServeStep { + /// Picking the next run. + Next, + /// Waiting for the subscription's own fill to bring the head of a group in `run`. + Head { run: std::ops::Range, head: Head }, + /// Fetching a run. + Fetch(FetchRun), +} + +impl ServeFill { + fn new( + this: Subscriber, broadcast: Path, name: String, fill: track::Fill, timescale: Later>, heading: Option, - ) { - let fill = std::sync::Arc::new(fill); - let mut last = None; + ) -> Self { + Self { + _working: this.session.work().start(), + this, + broadcast, + name, + fill: std::sync::Arc::new(fill), + timescale, + heading, + last: None, + step: ServeStep::Next, + } + } +} + +impl kio::Task for ServeFill { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { loop { - // The groups around a pending head go first, so it holds back only its own group. - let pending = match heading.as_ref().and_then(|heading| heading.pending(&self.state)) { - Some(Head::Group(group)) => fill.next_run_except(group), - _ => None, - }; - let Some(run) = pending.or_else(|| fill.next_run()) else { - break; - }; - let mut run = match fill.live_group() { - Some(live) => first_piece(run, live, fill.order()), - None => run, - }; - if let Some(heading) = &heading - && let Some(head) = heading.pending(&self.state) - { - if let Head::Group(group) = head { - run = first_piece(run, group, fill.order()); + match &mut self.step { + ServeStep::Next => { + let fill = &self.fill; + let state = &self.this.state; + // The groups around a pending head go first, so it holds back only its own group. + let pending = match self.heading.as_ref().and_then(|heading| heading.pending(state)) { + Some(Head::Group(group)) => fill.next_run_except(group), + _ => None, + }; + let Some(run) = pending.or_else(|| fill.next_run()) else { + return Poll::Ready(()); + }; + let mut run = match fill.live_group() { + Some(live) => first_piece(run, live, fill.order()), + None => run, + }; + if let Some(heading) = &self.heading + && let Some(head) = heading.pending(state) + { + if let Head::Group(group) = head { + run = first_piece(run, group, fill.order()); + } + if head.overlaps(&run) { + self.step = ServeStep::Head { run, head }; + continue; + } + } + // Every fetch answers its run, so the same one coming back is a bug; ending the + // fill names it dropped rather than fetching it forever. + if self.last.as_ref() == Some(&run) { + debug_assert!(false, "fill run {run:?} left unanswered"); + return Poll::Ready(()); + } + self.last = Some(run.clone()); + self.step = ServeStep::Fetch(FetchRun::new( + self.this.clone(), + self.broadcast.clone(), + self.name.clone(), + fill.clone(), + run, + self.timescale.clone(), + )); } - if head.overlaps(&run) { - let waited = kio::wait(|waiter| { - if let Poll::Ready(err) = fill.poll_aborted(waiter) { - return Poll::Ready(Err(err)); + ServeStep::Head { run, head } => { + let heading = self.heading.as_ref().expect("only a heading fill waits on its head"); + let waited = ready!(match self.fill.poll_aborted(waiter) { + Poll::Ready(err) => Poll::Ready(Err(err)), + Poll::Pending if self.fill.poll_unwanted_groups(run.clone(), waiter).is_ready() => { + Poll::Ready(Ok(())) + } + Poll::Pending => heading.poll_changed(&self.this.state, *head, waiter).map(Ok), + }); + if let Err(err) = waited { + self.fill.reject_groups(run.clone(), err); + } + self.step = ServeStep::Next; + } + ServeStep::Fetch(fetch) => { + ready!(fetch.poll(waiter)); + self.step = ServeStep::Next; + } + } + } + } +} + +/// Fetches a run of groups from the publisher to fill cache misses, with one standalone +/// (or draft-22 filtered) FETCH for this subscription's track. It asks from the frame the +/// reader wants in the first group through the end of the last, so a publisher that +/// evicted a prefix can still answer. +/// +/// FETCH_OK settles what the answer covers before any group is accepted, so a refusal +/// reaches every waiting reader as the publisher's own error. The objects arrive on a +/// fetch stream, which [`RecvFill`] routes to [`RecvGroupFetch`]. +struct FetchRun { + this: Subscriber, + broadcast: Path, + name: String, + fill: std::sync::Arc, + run: std::ops::Range, + timescale: Later>, + validity: crate::model::broadcast::RequestValidity, + /// The run's last group. + last: u64, + /// The frame the reader wants in the run's first group. + start: u64, + draft20: bool, + stream: Option>, + messages: super::request_stream::Messages, + /// Where FETCH_OK hands the run to its fetch stream. + slot: kio::Producer, + /// Registered before the FETCH goes out, since its stream can overtake FETCH_OK. + registered: Option, + step: FetchRunStep, +} + +enum FetchRunStep { + /// Opening the request stream. + Open(crate::transport::open::Open), + /// Waiting for FETCH_OK, flushing the FETCH meanwhile. + Answer, + /// Waiting for the timescale the objects are stamped in. + Timescale(ietf::FetchOk), + /// Keeping the request open until its data stream finishes or every reader leaves. + Receive { + covered: std::ops::Range, + /// Whether the request stream may still carry a reset. + listening: bool, + abandoned: bool, + }, + /// Finishing the request stream. + Close(StreamClose), +} + +impl FetchRun { + fn new( + this: Subscriber, + broadcast: Path, + name: String, + fill: std::sync::Arc, + run: std::ops::Range, + timescale: Later>, + ) -> Self { + Self { + validity: this.validity.start(), + last: run.end - 1, + start: fill.frame_start(run.start), + draft20: Filter::is_draft20(this.version), + messages: super::request_stream::Messages::new(this.session.clone(), this.control.client()), + this, + broadcast, + name, + fill, + run, + timescale, + stream: None, + slot: kio::Producer::new(GroupFetch::Pending), + registered: None, + step: FetchRunStep::Open(Default::default()), + } + } + + /// Refuse what the run still owes, and cancel the request. + fn cancel(&mut self, err: Error) -> Poll<()> { + self.fill.reject_groups(self.run.clone(), err); + self.registered = None; + if let Some(stream) = self.stream.take() { + Subscriber::cancel_group_fetch(stream); + } + Poll::Ready(()) + } + + /// Refuse `groups`, then finish the request stream. + fn refuse(&mut self, groups: std::ops::Range, err: Error) { + self.fill.reject_groups(groups, err); + self.close(); + } + + /// Finish the request stream. + fn close(&mut self) { + let stream = self.stream.take().expect("the request stream is open"); + self.step = FetchRunStep::Close(StreamClose::new(stream)); + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + loop { + match &mut self.step { + FetchRunStep::Open(open) => { + // The owning subscribe/linger loop enforces its admitted gate, including a + // token-authorized announce outside the session scope. Its abort ends this job. + // The fill is already owned by waiting readers; credit backpressures it until + // their demand leaves instead of refusing accepted work. + let opened = if let Err(err) = self.validity.check(waiter) { + Err(err) + } else if let Poll::Ready(err) = self.fill.poll_aborted(waiter) { + Err(err) + } else if self.fill.poll_unwanted_groups(self.run.clone(), waiter).is_ready() { + Err(Error::Cancel) + } else if let Poll::Ready(Err(err)) = self.timescale.poll(waiter) { + // A TRACK_STATUS sent alongside that failed fails the fetch. + Err(err) + } else { + ready!(Stream::poll_open(&self.this.session, open, self.this.version, waiter)) + }; + let mut stream = match opened { + Ok(stream) => stream.with_stats(self.this.origin.stats()), + Err(err) => { + self.fill.reject_groups(self.run.clone(), err); + return Poll::Ready(()); + } + }; + let fetch_id = self.this.control.next_request_id(); + self.registered = Some(GroupFetchEntry::new(&self.this.state, fetch_id, self.slot.clone())); + + let namespace = self.broadcast.clone(); + let track = self.name.as_str().into(); + let from = ietf::Location { + group: self.run.start, + object: self.start, + }; + let fetch_type = match self.draft20 { + // No End Object includes the whole End Group. + true => FetchType::Filtered { + namespace, + track, + filter: Filter::Absolute { + start: from, + end: Some(ietf::EndLocation { + group: self.last, + object: None, + }), + }, + }, + // An End Object of 0 is the whole End Group. + false => FetchType::Standalone { + namespace, + track, + start: from, + end: ietf::Location { + group: self.last, + object: 0, + }, + }, + }; + // Flushed while waiting for the answer, so the terminals below still end a + // request the peer's send window holds back. + let written = stream.writer.buffer_message( + ietf::Fetch::ID, + &ietf::Fetch { + authorization_token: None, + request_id: fetch_id, + subscriber_priority: super::priority::to_wire(self.fill.priority()), + group_order: GroupOrder::Ascending, + fetch_type, + range_filters: false, + fill_timeout: false, + // The copy learns the track from SUBSCRIBE_OK or TRACK_STATUS_OK, so a + // FETCH_OK repeating it is waste, except that only its TIMESCALE opts the + // answer into the timestamps a timed copy needs. A FETCH sent alongside + // TRACK_STATUS cannot know yet, so it asks. + properties_wanted: match self.timescale.peek() { + Some(Ok(timescale)) => timescale.is_some(), + _ => true, + }, + }, + ); + self.stream = Some(stream); + match written { + Ok(()) => self.step = FetchRunStep::Answer, + Err(err) => { + tracing::debug!(%err, groups = ?self.run, "group fetch refused"); + self.refuse(self.run.clone(), err); + } + } + } + FetchRunStep::Answer => { + let stream = self.stream.as_mut().expect("the request stream is open"); + let response = match stream.writer.poll_flush(waiter) { + Poll::Ready(Ok(())) => { + self.this + .poll_read_fetch_response(&mut stream.reader, &mut self.messages, waiter) + } + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Pending => Poll::Pending, + }; + let res = if let Poll::Ready(Err(err)) = response { + SetupResponse::Response(Err(err)) + } else if let Err(err) = self.validity.check(waiter) { + SetupResponse::Cancelled(err) + } else if let Poll::Ready(err) = self.fill.poll_aborted(waiter) { + SetupResponse::Cancelled(err) + } else if let Poll::Ready(Err(err)) = self.timescale.poll(waiter) { + SetupResponse::Cancelled(err) + } else if let Poll::Ready(res) = response { + SetupResponse::Response(res) + } else if self.fill.poll_unwanted_groups(self.run.clone(), waiter).is_ready() { + SetupResponse::Cancelled(Error::Cancel) + } else { + return Poll::Pending; + }; + match res { + SetupResponse::Response(Ok(ok)) => self.step = FetchRunStep::Timescale(ok), + SetupResponse::Cancelled(err) => return self.cancel(err), + SetupResponse::Response(Err(err)) => { + tracing::debug!(%err, groups = ?self.run, "group fetch refused"); + self.refuse(self.run.clone(), err); + } + } + } + FetchRunStep::Timescale(_) => { + // The objects are stamped in the track's timescale, which the TRACK_STATUS sent + // alongside may still owe: they wait unread in their stream until it lands, and + // its failure fails the fetch. + let timescale = if let Err(err) = self.validity.check(waiter) { + Err(err) + } else if let Poll::Ready(err) = self.fill.poll_aborted(waiter) { + Err(err) + } else if let Poll::Ready(res) = self.timescale.poll(waiter) { + res + } else if self.fill.poll_unwanted_groups(self.run.clone(), waiter).is_ready() { + Err(Error::Cancel) + } else { + return Poll::Pending; + }; + let timescale = match timescale { + Ok(timescale) => timescale, + Err(err) => return self.cancel(err), + }; + let FetchRunStep::Timescale(ok) = std::mem::replace(&mut self.step, FetchRunStep::Answer) else { + unreachable!() + }; + if let Some(step) = self.answered(ok, timescale) { + return step; + } + } + FetchRunStep::Receive { + covered, + listening, + abandoned, + } => { + let stream = self.stream.as_mut().expect("the request stream is open"); + let slot = &self.slot; + // A publisher that fails after FETCH_OK resets the request instead and owes no + // fetch stream, so the groups it left waiting are refused. A FIN is not that: + // the fetch stream can trail it. + let reset = 'reset: { + if *listening { + match self.messages.poll_end(&mut stream.reader, waiter) { + Poll::Ready(Err(err)) => break 'reset Some(err), + Poll::Ready(Ok(())) => *listening = false, + Poll::Pending => {} + } + } + if let Poll::Ready(err) = self.fill.poll_aborted(waiter) { + // Completed history stays cached. Only this FETCH's unfinished group is + // ours to abort; closing a track never cascades into cached siblings. + let state = slot.read(); + if let GroupFetch::Receiving(Some(group)) = &*state { + group.abort_open(err.clone()); + } + *abandoned = !matches!(*state, GroupFetch::Done); + break 'reset Some(err); + } + + let unwanted = self.fill.poll_unwanted_groups(covered.clone(), waiter).is_ready(); + let _ = ready!(slot.poll(waiter, |state| match &**state { + GroupFetch::Done => Poll::Ready(()), + GroupFetch::Receiving(Some(group)) + if unwanted + && group.poll_unwanted(waiter).is_ready() + && group.abort_unwanted(Error::Cancel) => + { + *abandoned = true; + Poll::Ready(()) + } + GroupFetch::Ready(_) | GroupFetch::Receiving(None) if unwanted => { + *abandoned = true; + Poll::Ready(()) + } + _ => Poll::Pending, + })); + None + }; + let covered = covered.clone(); + if *abandoned { + if let Ok(mut state) = self.slot.write() + && !matches!(*state, GroupFetch::Done) + { + *state = GroupFetch::Cancelled; } - if fill.poll_unwanted_groups(run.clone(), waiter).is_ready() { - return Poll::Ready(Ok(())); + // A closed track answers what the run still owes with its abort, so the fill + // is not asked for the same run again; groups nobody reads are owed nothing. + if let Some(err) = reset { + self.fill.reject_groups(covered, err); } - heading.poll_changed(&self.state, head, waiter).map(Ok) - }) - .await; - if let Err(err) = waited { - fill.reject_groups(run, err); + self.registered = None; + if let Some(stream) = self.stream.take() { + Subscriber::cancel_group_fetch(stream); + } + return Poll::Ready(()); } - continue; + // A run already being written is the fetch stream's to settle; one still waiting + // for its stream is lost. + if let Some(err) = reset + && let Ok(mut state) = self.slot.write() + && matches!(*state, GroupFetch::Ready(_)) + { + *state = GroupFetch::Done; + self.fill.reject_groups(covered, err); + } + self.close(); + } + FetchRunStep::Close(close) => { + let _ = ready!(close.poll(waiter)); + return Poll::Ready(()); } } - // Every fetch answers its run, so the same one coming back is a bug; ending the - // fill names it dropped rather than fetching it forever. - if last.as_ref() == Some(&run) { - debug_assert!(false, "fill run {run:?} left unanswered"); - return; - } - last = Some(run.clone()); - self.clone() - .run_range_fetch(broadcast.clone(), name.clone(), fill.clone(), run, timescale.clone()) - .await; } } - /// Serve a whole fill with one FETCH, as [`Self::serve_fill`] serves each run. - #[cfg(test)] - async fn run_group_fetch(self, broadcast: Path, name: String, fill: track::Fill, timescale: Option) { - let run = fill.start().group..fill.end(); - self.run_range_fetch(broadcast, name, std::sync::Arc::new(fill), run, Later::Known(timescale)) - .await; - } - - /// Fetch a run of groups from the publisher to fill cache misses, with one standalone - /// (or draft-22 filtered) FETCH for this subscription's track. It asks from the frame the - /// reader wants in the first group through the end of the last, so a publisher that - /// evicted a prefix can still answer. - /// - /// FETCH_OK settles what the answer covers before any group is accepted, so a refusal - /// reaches every waiting reader as the publisher's own error. The objects arrive on a - /// fetch stream, which [`Self::recv_fill`] routes to [`Self::recv_group_fetch`]. - async fn run_range_fetch( - self, - broadcast: Path, - name: String, - fill: std::sync::Arc, - run: std::ops::Range, - timescale: Later>, - ) { - let validity = self.validity.start(); - let last = run.end - 1; - let start = fill.frame_start(run.start); - let draft20 = Filter::is_draft20(self.version); - - // The owning subscribe/linger loop enforces its admitted gate, including a - // token-authorized announce outside the session scope. Its abort ends this job. - // The fill is already owned by waiting readers; credit backpressures it until - // their demand leaves instead of refusing accepted work. - let opening = self.session.clone(); - let mut open = crate::transport::open::Open::default(); - let mut stream = match kio::wait(|waiter| { - validity.check(waiter)?; - if let Poll::Ready(err) = fill.poll_aborted(waiter) { - return Poll::Ready(Err(err)); - } - if fill.poll_unwanted_groups(run.clone(), waiter).is_ready() { - return Poll::Ready(Err(Error::Cancel)); - } - // A TRACK_STATUS sent alongside that failed fails the fetch. - if let Poll::Ready(Err(err)) = timescale.poll(waiter) { - return Poll::Ready(Err(err)); - } - Stream::poll_open(&opening, &mut open, self.version, waiter) - }) - .await - { - Ok(stream) => stream, - Err(err) => return fill.reject_groups(run, err), - }; - stream = stream.with_stats(self.origin.stats()); - let fetch_id = self.control.next_request_id(); - - // Registered before the FETCH goes out, since its stream can overtake FETCH_OK. - let slot = kio::Producer::new(GroupFetch::Pending); - let registered = GroupFetchEntry::new(&self.state, fetch_id, slot.clone()); - - let namespace = broadcast.clone(); - let track = name.as_str().into(); - let from = ietf::Location { - group: run.start, - object: start, - }; - let fetch_type = match draft20 { - // No End Object includes the whole End Group. - true => FetchType::Filtered { - namespace, - track, - filter: Filter::Absolute { - start: from, - end: Some(ietf::EndLocation { - group: last, - object: None, - }), - }, - }, - // An End Object of 0 is the whole End Group. - false => FetchType::Standalone { - namespace, - track, - start: from, - end: ietf::Location { group: last, object: 0 }, - }, - }; - // Flushed while waiting for the answer, so the terminals below still end a request the - // peer's send window holds back. - let res = stream.writer.buffer_message( - ietf::Fetch::ID, - &ietf::Fetch { - authorization_token: None, - request_id: fetch_id, - subscriber_priority: super::priority::to_wire(fill.priority()), - group_order: GroupOrder::Ascending, - fetch_type, - range_filters: false, - fill_timeout: false, - // The copy learns the track from SUBSCRIBE_OK or TRACK_STATUS_OK, so a FETCH_OK - // repeating it is waste, except that only its TIMESCALE opts the answer into the - // timestamps a timed copy needs. A FETCH sent alongside TRACK_STATUS cannot know - // yet, so it asks. - properties_wanted: match timescale.peek() { - Some(Ok(timescale)) => timescale.is_some(), - _ => true, - }, - }, - ); - let mut messages = super::request_stream::Messages::new(self.session.clone(), self.control.client()); - let res = match res { - Err(err) => SetupResponse::Response(Err(err)), - Ok(()) => { - let mut response = std::pin::pin!(async { - kio::wait(|waiter| stream.writer.poll_flush(waiter)).await?; - self.read_fetch_response(&mut stream, &mut messages).await - }); - kio::wait(|waiter| { - let response = waiter.poll_future(response.as_mut()); - if let Poll::Ready(Err(err)) = response { - return Poll::Ready(SetupResponse::Response(Err(err))); - } - if let Err(err) = validity.check(waiter) { - return Poll::Ready(SetupResponse::Cancelled(err)); - } - if let Poll::Ready(err) = fill.poll_aborted(waiter) { - return Poll::Ready(SetupResponse::Cancelled(err)); - } - if let Poll::Ready(Err(err)) = timescale.poll(waiter) { - return Poll::Ready(SetupResponse::Cancelled(err)); - } - if let Poll::Ready(res) = response { - return Poll::Ready(SetupResponse::Response(res)); - } - fill.poll_unwanted_groups(run.clone(), waiter) - .map(|_| SetupResponse::Cancelled(Error::Cancel)) - }) - .await - } - }; - - let ok = match res { - SetupResponse::Response(Ok(ok)) => ok, - SetupResponse::Cancelled(err) => { - fill.reject_groups(run, err); - drop(registered); - self.cancel_group_fetch(stream).await; - return; - } - SetupResponse::Response(Err(err)) => { - tracing::debug!(%err, groups = ?run, "group fetch refused"); - fill.reject_groups(run, err); - let _ = stream.writer.close().await; - return; - } - }; - - // The objects are stamped in the track's timescale, which the TRACK_STATUS sent - // alongside may still owe: they wait unread in their stream until it lands, and its - // failure fails the fetch. - let timescale = match kio::wait(|waiter| { - if let Err(err) = validity.check(waiter) { - return Poll::Ready(Err(err)); - } - if let Poll::Ready(err) = fill.poll_aborted(waiter) { - return Poll::Ready(Err(err)); - } - if let Poll::Ready(res) = timescale.poll(waiter) { - return Poll::Ready(res); - } - fill.poll_unwanted_groups(run.clone(), waiter) - .map(|_| Err(Error::Cancel)) - }) - .await - { - Ok(timescale) => timescale, - Err(err) => { - fill.reject_groups(run, err); - drop(registered); - self.cancel_group_fetch(stream).await; - return; - } - }; - + /// Settle what FETCH_OK covers and hand the run to its fetch stream, or end the request + /// with `Some` when the answer covers nothing it can take. + fn answered(&mut self, ok: ietf::FetchOk, timescale: Option) -> Option> { + let (run, last, start) = (self.run.clone(), self.last, self.start); // Draft-22's End Location is inclusive; older drafts spell a bounded object one past it. let mut end = ok.end_location; - if draft20 { + if self.draft20 { let Some(object) = end.object.checked_add(1) else { - fill.reject_groups(run, Error::ProtocolViolation); - let _ = stream.writer.close().await; - return; + self.refuse(run, Error::ProtocolViolation); + return None; }; end.object = object; } // Before draft-22, Object 0 covers the whole named group. Legacy answers one // group past it still pass the exclusive comparison below. - let whole_last = !draft20 && end.group == last && end.object == 0; + let whole_last = !self.draft20 && end.group == last && end.object == 0; // Draft-22's answer ends at or after its start and never past the requested end, which // a publisher following section 9.11 cuts at the Largest Object. Validate before accepting a group or letting End of // Track mark where the track ends. - if draft20 && (end.group < run.start || end.group > last || (end.group == run.start && end.object <= start)) { + if self.draft20 + && (end.group < run.start || end.group > last || (end.group == run.start && end.object <= start)) + { let err = Error::ProtocolViolation; tracing::warn!(groups = ?run, ?end, "FETCH_OK's End Location is outside the requested range"); - self.session.clone().close( + self.this.session.clone().close( super::error::to_session_code(&SessionError::from(&err)), err.to_string().as_ref(), ); - fill.reject_groups(run, err); - return; + self.fill.reject_groups(run, err); + return Some(Poll::Ready(())); } // The publisher knows where the track ends, which a range FETCH downstream needs. if ok.end_of_track { let Some(final_sequence) = end.group.checked_add(u64::from(whole_last || end.object > 0)) else { - fill.reject_groups(run, Error::ProtocolViolation); - let _ = stream.writer.close().await; - return; + self.refuse(run, Error::ProtocolViolation); + return None; }; - fill.finish_track_at(final_sequence); + self.fill.finish_track_at(final_sequence); } // Before draft-22, an empty answer opens no fetch stream at all. if !whole_last && (end.group, end.object) <= (run.start, start) { - fill.reject_groups(run, Error::NotFound); - let _ = stream.writer.close().await; - return; + self.refuse(run, Error::NotFound); + return None; } // The groups the answer covers: whole up to the End Location's group, and that one up @@ -5799,12 +7175,12 @@ impl Subscriber { // A publisher may stop at the Largest Object, so the groups past it are not answered. // They are refused for now rather than named dropped: they may still be written. if covered < last { - fill.reject_groups(covered + 1..run.end, Error::NotFound); + self.fill.reject_groups(covered + 1..run.end, Error::NotFound); } - if let Ok(mut state) = slot.write() { + if let Ok(mut state) = self.slot.write() { *state = GroupFetch::Ready(RangeFetch { - fill: fill.clone(), + fill: self.fill.clone(), timescale, first: run.start, last: covered, @@ -5812,202 +7188,244 @@ impl Subscriber { end: bound, }); } - let covered = run.start..covered + 1; // Keep the request open until its data stream finishes or every reader leaves. - // A publisher that fails after FETCH_OK resets the request instead and owes no fetch - // stream, so the groups it left waiting are refused. A FIN is not that: the fetch - // stream can trail it. - let mut open = true; - let mut abandoned = false; - let reset = kio::wait(|waiter| { - if open { - let closed = messages.poll_end(&mut stream.reader, waiter); - match closed { - Poll::Ready(Err(err)) => return Poll::Ready(Some(err)), - Poll::Ready(Ok(())) => open = false, - Poll::Pending => {} - } - } - if let Poll::Ready(err) = fill.poll_aborted(waiter) { - // Completed history stays cached. Only this FETCH's unfinished group is - // ours to abort; closing a track never cascades into cached siblings. - let state = slot.read(); - if let GroupFetch::Receiving(Some(group)) = &*state { - group.abort_open(err.clone()); - } - abandoned = !matches!(*state, GroupFetch::Done); - return Poll::Ready(Some(err)); - } + self.step = FetchRunStep::Receive { + covered: run.start..covered + 1, + listening: true, + abandoned: false, + }; + None + } +} - let unwanted = fill.poll_unwanted_groups(covered.clone(), waiter).is_ready(); - slot.poll(waiter, |state| match &**state { - GroupFetch::Done => Poll::Ready(()), - GroupFetch::Receiving(Some(group)) - if unwanted && group.poll_unwanted(waiter).is_ready() && group.abort_unwanted(Error::Cancel) => - { - abandoned = true; - Poll::Ready(()) - } - GroupFetch::Ready(_) | GroupFetch::Receiving(None) if unwanted => { - abandoned = true; - Poll::Ready(()) - } - _ => Poll::Pending, - }) - .map(|_| None) - }) - .await; - if abandoned { - if let Ok(mut state) = slot.write() - && !matches!(*state, GroupFetch::Done) - { - *state = GroupFetch::Cancelled; - } - // A closed track answers what the run still owes with its abort, so the fill is - // not asked for the same run again; groups nobody reads are owed nothing. - if let Some(err) = reset { - fill.reject_groups(covered, err); - } - drop(registered); - self.cancel_group_fetch(stream).await; - return; - } - // A run already being written is the fetch stream's to settle; one still waiting for - // its stream is lost. - if let Some(err) = reset - && let Ok(mut state) = slot.write() - && matches!(*state, GroupFetch::Ready(_)) - { - *state = GroupFetch::Done; - fill.reject_groups(covered, err); +/// A stream finishing: its writer waits for the peer to acknowledge the FIN, as +/// [`Writer::close`] does, while its reader is held until then. +struct StreamClose { + _reader: Reader, + closing: Option, Error>>, +} + +impl StreamClose { + fn new(stream: Stream) -> Self { + let Stream { writer, reader } = stream; + Self { + _reader: reader, + closing: Some(writer.closing()), } - let _ = stream.writer.close().await; } + /// Ready with how the close went. + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + let res = match self.closing.as_mut().expect("polled after completion") { + Ok(closing) => ready!(closing.poll(waiter)), + Err(_) => self.closing.take().expect("checked above").map(drop), + }; + self.closing = None; + Poll::Ready(res) + } +} + +impl Subscriber { /// Cancel a FETCH using the negotiated draft's existing cancellation signal. - async fn cancel_group_fetch(&self, mut stream: Stream) { - stream.reader.abort(&Error::Cancel); - stream.writer.abort(&Error::Cancel) + fn cancel_group_fetch(stream: Stream) { + let Stream { mut reader, writer } = stream; + reader.abort(&Error::Cancel); + writer.abort(&Error::Cancel) } - /// Read the answer to a FETCH: FETCH_OK, or the publisher's refusal as an error. - async fn read_fetch_response( + /// The answer to a FETCH: FETCH_OK, or the publisher's refusal as an error. + fn poll_read_fetch_response( &self, - stream: &mut Stream, + reader: &mut Reader, messages: &mut super::request_stream::Messages, - ) -> Result { - let super::request_stream::Framed(type_id, body) = messages.message(&mut stream.reader).await?; + waiter: &kio::Waiter, + ) -> Poll> { + let super::request_stream::Framed(type_id, body) = ready!(messages.poll_message(reader, waiter))?; let mut data = body.decoder(self.version); - match type_id { + Poll::Ready(match type_id { ietf::FetchOk::ID => Ok(ietf::FetchOk::decode_msg(&mut data, self.version)?), ietf::RequestError::ID => { let msg = ietf::RequestError::decode_msg(&mut data, self.version)?; Err(request::from_code(msg.error_code, request::Kind::Fetch, self.version)) } _ => Err(Error::UnexpectedMessage), - } + }) } /// Write a FETCH's objects into the groups of the run it answers. + #[cfg(test)] async fn recv_group_fetch( &mut self, stream: &mut Reader, slot: kio::Producer, ) -> Result<(), Error> { + let mut recv = RecvGroupFetch::new(self.version, slot); + kio::wait(|waiter| recv.poll(waiter, stream)).await + } +} + +/// Writes a FETCH's objects into the groups of the run it answers. +struct RecvGroupFetch { + slot: kio::Producer, + version: Version, + state: RecvGroupFetchState, +} + +#[kio::machine(self.state: RecvGroupFetchState)] +impl RecvGroupFetch { + fn new(version: Version, slot: kio::Producer) -> Self { + Self { + slot, + version, + state: RecvGroupFetchState::Start, + } + } + + fn poll( + &mut self, + waiter: &kio::Waiter, + stream: &mut Reader, + ) -> Poll> { // FETCH_OK can trail its own fetch stream. Taking the run in the same step is what // refuses a second stream for one request. - let taken = kio::wait(|waiter| { - match slot.poll(waiter, |state| match &**state { - GroupFetch::Pending => Poll::Pending, - _ => Poll::Ready(()), - }) { - Poll::Ready(Ok(mut state)) => { - Poll::Ready(match std::mem::replace(&mut *state, GroupFetch::Receiving(None)) { - GroupFetch::Ready(fetch) => Ok(fetch), - GroupFetch::Cancelled => { - *state = GroupFetch::Cancelled; - Err(Error::Cancel) - } - other => { - *state = other; - Err(Error::Unsupported) - } - }) + let taken = wait!(Take => match self.slot.poll(waiter, |state| match &**state { + GroupFetch::Pending => Poll::Pending, + _ => Poll::Ready(()), + }) { + Poll::Ready(Ok(mut state)) => Poll::Ready(match std::mem::replace(&mut *state, GroupFetch::Receiving(None)) { + GroupFetch::Ready(fetch) => Ok(fetch), + GroupFetch::Cancelled => { + *state = GroupFetch::Cancelled; + Err(Error::Cancel) } - Poll::Ready(Err(_)) => Poll::Ready(Err(Error::Cancel)), - Poll::Pending => Poll::Pending, - } - }) - .await; - let mut writer = RangeWriter::new(taken?, &slot); - - let res = { - let mut receive = std::pin::pin!(crate::util::sendable( - self.recv_group_fetch_objects(stream, &mut writer) - )); - kio::wait(|waiter| { - if let Poll::Ready(res) = waiter.poll_future(receive.as_mut()) { - return Poll::Ready(res); + other => { + *state = other; + Err(Error::Unsupported) } - // The request gave up on the run, or the group being written was aborted under - // it, as by eviction while the stream stalls: the rest of the run is refused - // rather than left waiting on a stream that may never resume. + }), + Poll::Ready(Err(_)) => Poll::Ready(Err(Error::Cancel)), + Poll::Pending => Poll::Pending, + }); + let mut writer: RangeWriter = RangeWriter::new(taken?, self.slot.clone()); + let objects: GroupFetchObjects = GroupFetchObjects::new(self.version); + + let res: Result<(), Error> = wait!(Receive => match objects.poll(waiter, stream, writer) { + Poll::Ready(res) => Poll::Ready(res), + // The request gave up on the run, or the group being written was aborted under + // it, as by eviction while the stream stalls: the rest of the run is refused + // rather than left waiting on a stream that may never resume. + Poll::Pending => { let mut aborted = Error::Cancel; - slot.poll(waiter, |state| match &**state { - GroupFetch::Cancelled => Poll::Ready(()), - GroupFetch::Receiving(Some(group)) => group.poll_aborted(waiter).map(|err| aborted = err), - _ => Poll::Pending, - }) - .map(|_| Err(aborted)) - }) - .await - }; + self.slot + .poll(waiter, |state| match &**state { + GroupFetch::Cancelled => Poll::Ready(()), + GroupFetch::Receiving(Some(group)) => group.poll_aborted(waiter).map(|err| aborted = err), + _ => Poll::Pending, + }) + .map(|_| Err(aborted)) + } + }); // Completion and abandonment share the slot lock, so a complete group is never // aborted between finishing it and publishing Done to the request handler. - let mut state = slot.write().ok(); + let mut state = self.slot.write().ok(); let res = writer.settle(res); if let Some(state) = state.as_mut() { **state = GroupFetch::Done; } - res + // The writer's own handle on the slot goes only once the lock does. + drop(state); + drop(writer); + Poll::Ready(res) } +} - /// Decode a FETCH's objects into the groups of its run, preserving marked gaps and their - /// properties, never past FETCH_OK's End Location. A group's objects end where the next - /// group's begin, since a FETCH carries every object of its range in order. Before - /// draft-22 the stream must also reach an object bound inside the last group. - async fn recv_group_fetch_objects( - &self, +/// Decodes a FETCH's objects into the groups of its run, preserving marked gaps and their +/// properties, never past FETCH_OK's End Location. A group's objects end where the next +/// group's begin, since a FETCH carries every object of its range in order. Before +/// draft-22 the stream must also reach an object bound inside the last group. +struct GroupFetchObjects { + version: Version, + /// The next Object ID expected in the current group. + next: u64, + prior_group: Option, + first: bool, + /// The object whose payload is being read: its group, its Object ID, and where it goes. + payload: Option<(u64, u64, FetchPayloadRead)>, +} + +/// Where a fetch object's payload goes. +#[allow( + clippy::large_enum_variant, + reason = "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; an allow, since the size gap varies by target" +)] +enum FetchPayloadRead { + Write(FetchPayloadRecv), + /// Cached already, or a datagram: read past it. + Skip(SkipPayload), +} + +impl GroupFetchObjects { + fn new(version: Version) -> Self { + Self { + version, + next: 0, + prior_group: None, + first: true, + payload: None, + } + } + + fn poll( + &mut self, + waiter: &kio::Waiter, stream: &mut Reader, - writer: &mut RangeWriter<'_>, - ) -> Result<(), Error> { + writer: &mut RangeWriter, + ) -> Poll> { let timescale = writer.fetch.timescale; - // The next Object ID expected in the current group. - let mut next = 0; - let mut prior_group = None; - let mut first = true; - while let Some(object) = decode_fetch_object(stream, std::mem::take(&mut first)).await? { + loop { + if let Some((sequence, id, payload)) = &mut self.payload { + match payload { + FetchPayloadRead::Write(recv) => { + let producer = writer.producer().expect("the object's group is being written"); + ready!(recv.poll(waiter, stream, producer))?; + } + FetchPayloadRead::Skip(skip) => ready!(skip.poll(waiter, stream))?, + } + let (sequence, id) = (*sequence, *id); + self.payload = None; + self.next = id.checked_add(1).ok_or(Error::ProtocolViolation)?; + if writer.bound(sequence).is_some_and(|end| self.next > end) { + tracing::warn!(sequence, "a group fetch continued past FETCH_OK's End Location"); + return Poll::Ready(Err(Error::ProtocolViolation)); + } + } + + let Some(object) = ready!(poll_fetch_object(stream, self.first, waiter))? else { + break; + }; + self.first = false; if !object.subgroup_ok && !object.datagram { tracing::warn!("subgroup ID is not supported, dropping group fetch"); - return Err(Error::Unsupported); + return Poll::Ready(Err(Error::Unsupported)); } - let group = resolve_fetch_group(prior_group, object.group)?; + let group = resolve_fetch_group(self.prior_group, object.group)?; if let Some(group) = group { - prior_group = Some(group); + self.prior_group = Some(group); } // The first object always names its group, so a later one leaving it out stays in it. let sequence = group.or(writer.sequence()).ok_or(Error::ProtocolViolation)?; if writer.sequence() != Some(sequence) { - next = writer.open(sequence, object.datagram)?; + self.next = writer.open(sequence, object.datagram)?; } else if object.datagram { // A datagram group has one object, so it never shares one with a stream. - return Err(Error::ProtocolViolation); + return Poll::Ready(Err(Error::ProtocolViolation)); } + let next = self.next; let id = match (object.group.is_some(), object.object) { (true, Some(id)) => id, (false, Some(delta)) => next @@ -6015,29 +7433,21 @@ impl Subscriber { .and_then(|prior| prior.checked_add(delta)) .ok_or(Error::ProtocolViolation)?, (false, None) => next, - _ => return Err(Error::ProtocolViolation), + _ => return Poll::Ready(Err(Error::ProtocolViolation)), }; if id < next { - return Err(Error::ProtocolViolation); + return Poll::Ready(Err(Error::ProtocolViolation)); } let raw = bytes::Bytes::from(object.properties.unwrap_or_default()); let decoded = super::object::read(&raw, self.version, timescale)?; if id > next && decoded.gap.is_none_or(|gap| gap < id - next) { - return Err(Error::Unsupported); - } - match writer.producer() { - Some(producer) => { - self.recv_fetch_payload(stream, producer, FetchPayload { decoded, timescale, id }) - .await? - } - // Cached already, or a datagram: read past it. - None => skip_fetch_payload(stream).await?, - } - next = id.checked_add(1).ok_or(Error::ProtocolViolation)?; - if writer.bound(sequence).is_some_and(|end| next > end) { - tracing::warn!(sequence, "a group fetch continued past FETCH_OK's End Location"); - return Err(Error::ProtocolViolation); + return Poll::Ready(Err(Error::Unsupported)); } + let payload = match writer.producer() { + Some(_) => FetchPayloadRead::Write(FetchPayloadRecv::new(FetchPayload { decoded, timescale, id })), + None => FetchPayloadRead::Skip(SkipPayload::new()), + }; + self.payload = Some((sequence, id, payload)); } // A clean FIN short of the promised end would otherwise cache a truncated group as @@ -6046,30 +7456,58 @@ impl Subscriber { // stream must reach the bounded group at all, not just its objects. if !Filter::is_draft20(self.version) && let Some(end) = writer.fetch.end - && (writer.sequence() != Some(writer.fetch.last) || next < end) + && (writer.sequence() != Some(writer.fetch.last) || self.next < end) { let sequence = writer.sequence(); + let next = self.next; tracing::warn!(?sequence, next, "a group fetch ended before FETCH_OK's End Location"); - return Err(Error::ProtocolViolation); + return Poll::Ready(Err(Error::ProtocolViolation)); } - Ok(()) + Poll::Ready(Ok(())) } } -/// Read past one fetch object's payload, for a group this fetch does not write. -async fn skip_fetch_payload(stream: &mut Reader) -> Result<(), Error> { - let mut left = stream.varint().await?; - let mut budget = kio::coop::Budget::new(32); - while left > 0 { - kio::wait(|waiter| budget.poll_yield(waiter)).await; - let max = usize::try_from(left).unwrap_or(usize::MAX); - let chunk = kio::wait(|waiter| stream.poll_read_chunk(waiter, max)) - .await? - .ok_or(Error::ProtocolViolation)?; - left -= chunk.len() as u64; +/// Reads past one fetch object's payload, for a group this fetch does not write. +struct SkipPayload { + /// The payload bytes left to read, once its length arrived. + left: Option, + budget: kio::coop::Budget, + /// This pass's budget is spent, so a read still waiting does not spend another. + passing: bool, +} + +impl SkipPayload { + fn new() -> Self { + Self { + left: None, + budget: kio::coop::Budget::new(32), + passing: false, + } + } + + fn poll( + &mut self, + waiter: &kio::Waiter, + stream: &mut Reader, + ) -> Poll> { + let mut left = match self.left { + Some(left) => left, + None => *self.left.insert(ready!(stream.poll_varint(waiter))?), + }; + while left > 0 { + if !self.passing { + ready!(self.budget.poll_yield(waiter)); + self.passing = true; + } + let max = usize::try_from(left).unwrap_or(usize::MAX); + let chunk = ready!(stream.poll_read_chunk(waiter, max))?.ok_or(Error::ProtocolViolation)?; + self.passing = false; + left -= chunk.len() as u64; + self.left = Some(left); + } + Poll::Ready(Ok(())) } - Ok(()) } struct FetchPayload { @@ -6140,47 +7578,50 @@ struct FetchedObject { properties: Option>, } -/// Decode the next fetch object header, or `None` at stream end. +/// The next fetch object header, or `None` at stream end. /// /// The `first` object on a stream has no prior to inherit from, so a field it leaves to /// the prior object is a protocol violation. -async fn decode_fetch_object( +fn poll_fetch_object( stream: &mut Reader, first: bool, -) -> Result, Error> { - Ok(match stream.decode_maybe::().await? { - None => None, - Some(ietf::FetchObject::EndOfRange { .. }) => { - tracing::warn!("a fill with an End of Range cannot be stitched"); - return Err(Error::Unsupported); - } - Some(ietf::FetchObject::Object { - subgroup, - group, - object, - priority, - properties, - }) => { - let inherits = group.is_none() - || object.is_none() - || priority.is_none() - || matches!(subgroup, ietf::FetchSubgroup::Prior | ietf::FetchSubgroup::PriorPlusOne); - if first && inherits { - tracing::warn!("the first fetch object refers to a prior object"); - return Err(Error::ProtocolViolation); + waiter: &kio::Waiter, +) -> Poll, Error>> { + Poll::Ready(Ok( + match ready!(stream.poll_decode_maybe::(waiter))? { + None => None, + Some(ietf::FetchObject::EndOfRange { .. }) => { + tracing::warn!("a fill with an End of Range cannot be stitched"); + return Poll::Ready(Err(Error::Unsupported)); } - Some(FetchedObject { + Some(ietf::FetchObject::Object { + subgroup, group, object, - datagram: subgroup == ietf::FetchSubgroup::Datagram, - subgroup_ok: matches!( - subgroup, - ietf::FetchSubgroup::Zero | ietf::FetchSubgroup::Prior | ietf::FetchSubgroup::Explicit(0) - ), + priority, properties, - }) - } - }) + }) => { + let inherits = group.is_none() + || object.is_none() + || priority.is_none() + || matches!(subgroup, ietf::FetchSubgroup::Prior | ietf::FetchSubgroup::PriorPlusOne); + if first && inherits { + tracing::warn!("the first fetch object refers to a prior object"); + return Poll::Ready(Err(Error::ProtocolViolation)); + } + Some(FetchedObject { + group, + object, + datagram: subgroup == ietf::FetchSubgroup::Datagram, + subgroup_ok: matches!( + subgroup, + ietf::FetchSubgroup::Zero | ietf::FetchSubgroup::Prior | ietf::FetchSubgroup::Explicit(0) + ), + properties, + }) + } + }, + )) } /// Absolute Group ID from a fetch object's Group ID field. @@ -6521,14 +7962,11 @@ mod tests { let session = ScriptedSession::eof(responses); let (_, recv) = session.open_bi(Default::default()).await.unwrap(); let mut reader = Reader::new(crate::transport::boxed::RecvStream::new(recv), Version::Draft18); - let result = Subscriber::read_publish_done( - &mut reader, - Version::Draft18, - &AtomicU64::new(0), - &mut crate::ietf::request_stream::Messages::new(crate::driver::Transport::new(session.clone()), true), - |_| {}, - ) - .await; + let mut messages = + crate::ietf::request_stream::Messages::new(crate::driver::Transport::new(session.clone()), true); + let mut done = ReadPublishDone::default(); + let result = + kio::wait(|waiter| done.poll(waiter, &mut reader, Version::Draft18, &mut 0, &mut messages, |_| {})).await; if clean { assert_eq!(result.unwrap().stream_count, 0); } else { @@ -10375,7 +11813,7 @@ mod tests { let (tasks, task_set) = crate::util::TaskSet::new(); std::mem::forget(task_set); - let mut subscriber = Subscriber::new( + let subscriber = Subscriber::new( crate::time::sim(), crate::driver::Transport::new(session.clone()), origin, @@ -10407,7 +11845,8 @@ mod tests { }; // Errors are surfaced to the peer on the stream, not raised as a session error. - subscriber.run_publish_stream(stream, msg).await.unwrap(); + let mut decline = subscriber.decline_publish(stream, msg); + kio::wait(|waiter| decline.poll(waiter)).await.unwrap(); kio::clock::sleep(Duration::from_millis(1)).await; assert!( @@ -13306,10 +14745,7 @@ mod stitch_tests { let stream = Stream::open(&crate::driver::Transport::new(downstream.clone()), egress) .await .unwrap(); - publisher - .handle_stream(id, ietf::Body(body.into()), stream) - .unwrap() - .await; + kio::wait(publisher.handle_stream(id, ietf::Body(body.into()), stream).unwrap()).await; let mut streams = std::collections::BTreeMap::>::new(); for (id, event) in downstream.log.trail() { if let Sent::Write(bytes) = event { diff --git a/rs/moq-net/src/ietf/subscriber/namespaces.rs b/rs/moq-net/src/ietf/subscriber/namespaces.rs index 1e855d4bc1a..c1e8a5397c8 100644 --- a/rs/moq-net/src/ietf/subscriber/namespaces.rs +++ b/rs/moq-net/src/ietf/subscriber/namespaces.rs @@ -10,7 +10,7 @@ //! which leaves the wire before the broader one goes out. The routes a replaced request //! carried go with it, and come back once the broader request reports them. -use std::future::Future; +use std::collections::VecDeque; use super::*; use crate::discovery::{Change, Discovery}; @@ -33,155 +33,205 @@ struct Request { gone: kio::Consumer<()>, } -/// A request waiting to go out. -struct Pending { - /// Close once each request it overlaps is off the wire. - blockers: Vec>, - /// Opens its stream. - opening: O, - /// Closes when the session ends it. - live: kio::Consumer<()>, - gone: kio::Producer<()>, +/// Keeps one SUBSCRIBE_NAMESPACE open for each prefix discovery decides on, each on its +/// own stream, for the session's life. Only an error ends it: a request the peer breaks +/// the protocol on. +pub(in crate::ietf) struct Namespaces { + subscriber: Subscriber, + /// Until the peer's SETUP says whether it holds overlapping requests. + setup: Option, + discovery: Option, + requests: HashMap, + /// The request discovery holds open for each prefix. + open_at: HashMap, + running: kio::Tasks, + /// Requests that ended, handled one at a time with discovery's changes between them. + done: VecDeque<(u64, Result)>, + next: u64, } -impl Subscriber { - /// Keep one SUBSCRIBE_NAMESPACE open for each prefix `discovery` decides on, each on a - /// stream from `open`, for the session's life. Only an error ends it: a request the - /// peer breaks the protocol on. - pub(in crate::ietf) async fn run_namespaces(self, discovery: Discovery, open: F) -> Result<(), Error> - where - F: Fn() -> O, - O: Future, Error>>, - { - let mut discovery = match self.peer_setup.get().await.solicit { - Some(_) => discovery, - None => discovery.exclusive(), - }; - let mut requests: HashMap = HashMap::new(); - // The request discovery holds open for each prefix. - let mut open_at: HashMap = HashMap::new(); - let mut running = futures::stream::FuturesUnordered::new(); - let mut next = 0u64; +impl Namespaces { + pub(in crate::ietf) fn new(subscriber: Subscriber, discovery: Discovery) -> Self { + Self { + subscriber, + setup: Some(discovery), + discovery: None, + requests: HashMap::new(), + open_at: HashMap::new(), + running: kio::Tasks::new(), + done: VecDeque::new(), + next: 0, + } + } - loop { - enum Event { - Change(Change), - Done(T), - } - let event = kio::wait(|waiter| { - if let Poll::Ready(change) = discovery.poll(waiter) { - return Poll::Ready(Event::Change(change)); - } - let mut cx = waiter.context(); - match futures::Stream::poll_next(std::pin::Pin::new(&mut running), &mut cx) { - Poll::Ready(Some(done)) => Poll::Ready(Event::Done(done)), - _ => Poll::Pending, - } - }) - .await; + /// Ready with the error that ends them. + pub(in crate::ietf) fn poll(&mut self, waiter: &kio::Waiter) -> Poll { + if self.setup.is_some() { + let declared = ready!(self.subscriber.peer_setup.poll_get(waiter)); + let discovery = self.setup.take().expect("checked above"); + self.discovery = Some(match declared.solicit { + Some(_) => discovery, + None => discovery.exclusive(), + }); + } + let discovery = self.discovery.as_mut().expect("decided above"); - match event { - Event::Change(Change::Open(prefix, snapshot)) => { - // Wait out the overlapping requests on their way out, such as the narrower - // ones an exclusive session just closed for this one. Those still open - // overlap it only where the peer holds both. - let blockers = requests - .values() - .filter(|request| request.live.is_none()) - .filter(|request| request.prefix.has_prefix(&prefix) || prefix.has_prefix(&request.prefix)) - .map(|request| request.gone.clone()) - .collect(); - let live = kio::Producer::new(()); - let gone = kio::Producer::new(()); - let key = next; - next += 1; - requests.insert( - key, - Request { - prefix: prefix.clone(), - live: Some(live.clone()), - gone: gone.consume(), - }, - ); - open_at.insert(prefix.clone(), key); - let pending = Pending { - blockers, - opening: open(), - live: live.consume(), - gone, - }; - running.push(self.clone().namespace_request(key, prefix, snapshot, pending)); - } - // Dropping its `live` ends the request. - Event::Change(Change::Close(prefix)) => { - if let Some(key) = open_at.remove(&prefix) - && let Some(request) = requests.get_mut(&key) - { - request.live = None; - } - } - Event::Done((key, res)) => { - let Some(request) = requests.remove(&key) else { - continue; - }; - // Ended on its own, rather than by the session. - if open_at.get(&request.prefix) == Some(&key) { - open_at.remove(&request.prefix); - discovery.ended(&request.prefix); + loop { + if let Poll::Ready(change) = discovery.poll(waiter) { + match change { + Change::Open(prefix, snapshot) => { + // Wait out the overlapping requests on their way out, such as the narrower + // ones an exclusive session just closed for this one. Those still open + // overlap it only where the peer holds both. + let blockers = self + .requests + .values() + .filter(|request| request.live.is_none()) + .filter(|request| request.prefix.has_prefix(&prefix) || prefix.has_prefix(&request.prefix)) + .map(|request| request.gone.clone()) + .collect(); + let live = kio::Producer::new(()); + let gone = kio::Producer::new(()); + let key = self.next; + self.next += 1; + self.requests.insert( + key, + Request { + prefix: prefix.clone(), + live: Some(live.clone()), + gone: gone.consume(), + }, + ); + self.open_at.insert(prefix.clone(), key); + self.running.push(Exchange { + key, + gone: Some(gone), + state: ExchangeState::Opening { + subscriber: self.subscriber.clone(), + prefix, + snapshot, + blockers, + open: crate::transport::open::Open::default(), + live: live.consume(), + }, + }); } - if let Err(err) = res { - // The peer breaking the protocol is fatal, and the driver turns - // this into the session close the draft wants. - if is_protocol_violation(&err) { - return Err(err); + // Dropping its `live` ends the request. + Change::Close(prefix) => { + if let Some(key) = self.open_at.remove(&prefix) + && let Some(request) = self.requests.get_mut(&key) + { + request.live = None; } - tracing::warn!(%err, prefix = %request.prefix, "subscribe_namespace failed, continuing without"); } } + continue; + } + + let Some((key, res)) = self.done.pop_front() else { + let done = &mut self.done; + let _ = self + .running + .poll_each(waiter, |exchange, res| done.push_back((exchange.key, res))); + if self.done.is_empty() { + return Poll::Pending; + } + continue; + }; + let Some(request) = self.requests.remove(&key) else { + continue; + }; + // Ended on its own, rather than by the session. + if self.open_at.get(&request.prefix) == Some(&key) { + self.open_at.remove(&request.prefix); + discovery.ended(&request.prefix); + } + if let Err(err) = res { + // The peer breaking the protocol is fatal, and the driver turns + // this into the session close the draft wants. + if is_protocol_violation(&err) { + return Poll::Ready(err); + } + tracing::warn!(%err, prefix = %request.prefix, "subscribe_namespace failed, continuing without"); } } } +} - /// Send one request once the ones it overlaps are off the wire, and read it until - /// either side ends it. - async fn namespace_request( - mut self, - key: u64, +/// One request: sent once the ones it overlaps are off the wire, and read until either +/// side ends it. +struct Exchange { + key: u64, + /// Dropped once it is off the wire, so an overlapping request may go out. + gone: Option>, + state: ExchangeState, +} + +#[expect(clippy::large_enum_variant, reason = "a machine's enum is its storage")] +enum ExchangeState { + /// Waiting out the requests it overlaps, then for stream credit. + Opening { + subscriber: Subscriber, prefix: Path, snapshot: crate::model::Snapshot, - pending: Pending, - ) -> (u64, Result) - where - O: Future, Error>>, - { - let Pending { - blockers, - opening, - live, - gone, - } = pending; - // The open races `live` too: a request ended before it went out holds nothing, nor - // the broader one it blocks. Without stream credit it is refused, not parked. - let mut opening = std::pin::pin!(opening); - let opened = kio::wait(|waiter| { - if live.poll_closed(waiter).is_ready() { - return Poll::Ready(None); - } - for blocker in &blockers { - if blocker.poll_closed(waiter).is_pending() { - return Poll::Pending; + /// Close once each request it overlaps is off the wire. + blockers: Vec>, + open: crate::transport::open::Open, + /// Closes when the session ends it. + live: kio::Consumer<()>, + }, + Running(NamespaceSubscription), + Done, +} + +impl kio::Task for Exchange { + type Output = Result; + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll { + let res = loop { + match &mut self.state { + ExchangeState::Opening { + subscriber, + blockers, + open, + live, + .. + } => { + // The open races `live` too: a request ended before it went out holds nothing, nor + // the broader one it blocks. Without stream credit it is refused, not parked. + if live.poll_closed(waiter).is_ready() { + break Ok(Ended::Local); + } + for blocker in blockers.iter() { + if blocker.poll_closed(waiter).is_pending() { + return Poll::Pending; + } + } + let stream = match ready!(Stream::poll_open(&subscriber.session, open, subscriber.version, waiter)) + { + Ok(stream) => stream, + Err(err) => break Err(err), + }; + let ExchangeState::Opening { + subscriber, + prefix, + snapshot, + live, + .. + } = std::mem::replace(&mut self.state, ExchangeState::Done) + else { + unreachable!("matched above"); + }; + self.state = + ExchangeState::Running(NamespaceSubscription::new(subscriber, stream, prefix, snapshot, live)); } + ExchangeState::Running(request) => break ready!(request.poll(waiter)), + ExchangeState::Done => panic!("Exchange polled after it finished"), } - waiter.poll_future(opening.as_mut()).map(Some) - }) - .await; - let res = match opened { - Some(Ok(stream)) => self.run_subscribe_namespace(stream, prefix, snapshot, live).await, - Some(Err(err)) => Err(err), - None => Ok(Ended::Local), }; + self.state = ExchangeState::Done; // Off the wire: an overlapping request may go out. - drop(gone); - (key, res) + self.gone = None; + Poll::Ready(res) } } diff --git a/rs/moq-net/src/ietf/subscriber/tracks.rs b/rs/moq-net/src/ietf/subscriber/tracks.rs index 9690aae226f..fb8c8a0c701 100644 --- a/rs/moq-net/src/ietf/subscriber/tracks.rs +++ b/rs/moq-net/src/ietf/subscriber/tracks.rs @@ -351,8 +351,8 @@ impl Subscriber { pub(super) fn listen(&self, path: Path, lister: Arc) -> Listening { let (key, failed, opening) = self.listings.lock().listen(path, lister); if let Some(opening) = opening { - let this = self.clone(); - self.tasks.push(this.run_subscribe_tracks(opening)); + let mut task = SubscribeTracks::new(self.clone(), opening); + self.tasks.push(move |waiter: &kio::Waiter| task.poll(waiter)); } Listening { listings: self.listings.clone(), @@ -361,181 +361,6 @@ impl Subscriber { } } - /// Send one SUBSCRIBE_TRACKS and read its answers until it is cancelled and the peer - /// has closed it, or it fails. - async fn run_subscribe_tracks(self, opening: Opening) { - let Opening { - key, - prefix, - blockers, - live, - gone, - } = opening; - let result = self.subscribe_tracks(&prefix, blockers, &live).await; - let mut listings = self.listings.lock(); - if let Err(err) = result - && listings - .requests - .get(&key) - .is_some_and(|request| request.live.is_some()) - { - listings.fail(&prefix, err); - } - listings.requests.remove(&key); - listings.release(); - drop(listings); - // Off the wire: an overlapping request may go out. - drop(gone); - } - - async fn subscribe_tracks( - &self, - prefix: &Path, - blockers: Vec>, - live: &kio::Consumer<()>, - ) -> Result<(), Error> { - // The peer refuses a prefix overlapping one it still holds, so wait for each - // overlapping request to leave the wire. - let proceed = kio::wait(|waiter| { - if live.poll_closed(waiter).is_ready() { - return Poll::Ready(false); - } - for blocker in &blockers { - if blocker.poll_closed(waiter).is_pending() { - return Poll::Pending; - } - } - Poll::Ready(true) - }) - .await; - if !proceed { - return Ok(()); - } - - let session = self.session.clone(); - let opened = { - let mut opening = std::pin::pin!(Stream::open(&session, self.version)); - kio::wait(|waiter| { - if live.poll_closed(waiter).is_ready() { - return Poll::Ready(None); - } - waiter.poll_future(opening.as_mut()).map(Some) - }) - .await - }; - let Some(stream) = opened else { - return Ok(()); - }; - let mut stream = stream?.with_stats(self.origin.stats()); - let request_id = self.control.next_request_id(); - stream.writer.buffer_message( - ietf::SubscribeTracks::ID, - &ietf::SubscribeTracks { - request_id, - prefix: prefix.clone(), - forward: false, - group_order: None, - properties_wanted: true, - authorization_token: None, - }, - )?; - tracing::debug!(prefix = %crate::path::Rooted(self.origin.root(), prefix), id = %request_id, "subscribe_tracks started"); - - let mut messages = Messages::new(self.session.clone(), self.control.client()); - let mut established = false; - enum Event { - Closed, - Cancelled, - Message(Option), - } - // A peer streaming skips must not hold the session's other tasks. - let mut budget = kio::coop::Budget::new(32); - loop { - kio::wait(|waiter| budget.poll_yield(waiter)).await; - let event = kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() { - return Poll::Ready(Ok(Event::Closed)); - } - if live.poll_closed(waiter).is_ready() { - return Poll::Ready(Ok(Event::Cancelled)); - } - ready!(stream.writer.poll_flush(waiter))?; - messages - .poll_next(&mut stream.reader, waiter) - .map(|next| next.map(Event::Message)) - }) - .await?; - let Framed(id, body) = match event { - Event::Closed => return Ok(()), - Event::Cancelled => break, - // The peer ended the listing, which only a cancel of ours should do. - Event::Message(None) if established => return Err(Error::Closed), - Event::Message(None) => return Err(messages.failed(Error::ProtocolViolation)), - Event::Message(Some(framed)) => framed, - }; - match id { - ietf::RequestOk::ID if !established => { - body.message::(self.version) - .map_err(|err| messages.failed(err.into()))?; - established = true; - } - ietf::RequestError::ID if !established => { - let refused = body - .message::(self.version) - .map_err(|err| messages.failed(err.into()))?; - tracing::debug!(prefix = %prefix, code = refused.error_code, reason = %refused.reason_phrase, "subscribe_tracks refused"); - return Err(request::from_code( - refused.error_code, - request::Kind::SubscribeNamespace, - self.version, - )); - } - ietf::PublishSkipped::ID if established => { - let skipped = body - .message::(self.version) - .map_err(|err| messages.failed(err.into()))?; - let namespace = prefix - .join(&skipped.suffix) - .map_err(|err| messages.failed(err.into()))?; - self.skipped(namespace, skipped.track_name.as_ref().into()); - } - _ => return Err(messages.failed(Error::UnexpectedMessage)), - } - } - - // Cancelled: close our half (FIN on draft-18, RESET_STREAM from draft-22, where a - // FIN is not a cancel), then read until the peer closes its own. The PUBLISHes it - // made stay ours. - tracing::debug!(prefix = %crate::path::Rooted(self.origin.root(), prefix), "subscribe_tracks cancelled"); - let Stream { mut writer, mut reader } = stream; - // The request never reached the peer whole, so it holds nothing to wait on. - if !writer.is_flushed() { - reader.abort(&Error::Cancel); - writer.abort(&Error::Cancel); - return Ok(()); - } - match request_stream::fin_cancels(self.version) { - true => { - let _ = writer.finish(); - } - false => writer.abort(&Error::Cancel), - } - kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() { - return Poll::Ready(()); - } - loop { - ready!(budget.poll_yield(waiter)); - match ready!(messages.poll_next(&mut reader, waiter)) { - Ok(Some(_)) => continue, - Ok(None) | Err(_) => return Poll::Ready(()), - } - } - }) - .await; - Ok(()) - } - /// A fresh subscription of ours to `name` in `namespace`, and its info once accepted. fn subscribed(&self, namespace: &Path, name: &str) -> Option<(RequestId, Option)> { let state = self.state.lock(); @@ -583,151 +408,13 @@ impl Subscriber { }; listings.insert(key.clone(), entry); drop(listings); - let this = self.clone(); - self.tasks.push(async move { - match this.open_paused(&key, &handoff).await { - Ok(Some(upstream)) => this.watch(key, upstream, handoff).await, - Ok(None) => {} - Err(err) => { - tracing::debug!(broadcast = %key.0, track = %key.1, %err, "paused subscribe failed"); - let mut listings = this.listings.lock(); - if listings.watched_by(&key, &handoff).is_some() { - listings.remove(&key, false); - listings.fail(&key.0, err); - } - } - } - }); - } - - /// Open a paused SUBSCRIBE for a skipped track and list it once its SUBSCRIBE_OK - /// arrives. `None` once its entry was let go. - async fn open_paused(&self, key: &TrackKey, handoff: &kio::Shared) -> Result, Error> { - let gone = |waiter: &kio::Waiter| { - handoff - .poll(waiter, |handoff| match handoff.gone { - true => Poll::Ready(()), - false => Poll::Pending, - }) - .is_ready() - }; - let session = self.session.clone(); - let opened = { - let mut opening = std::pin::pin!(Stream::open(&session, self.version)); - kio::wait(|waiter| { - if gone(waiter) { - return Poll::Ready(None); - } - waiter.poll_future(opening.as_mut()).map(Some) - }) - .await - }; - let Some(stream) = opened else { - return Ok(None); - }; - let mut stream = stream?.with_stats(self.origin.stats()); - let request_id = self.control.next_request_id(); - self.state.lock().subscribes.insert( - request_id, - TrackState { - ..TrackState::pending(key.1.to_string(), key.0.clone(), kio::Producer::new(Fill::Done)) - }, - ); - let mut messages = Messages::new(self.session.clone(), self.control.client()); - let accepted = async { - // Flushed while waiting for the answer, so the entry going still ends a request - // the peer's send window holds back. - stream.writer.buffer_message( - ietf::Subscribe::ID, - &ietf::Subscribe { - max_age: None, - new_group: None, - request_id, - track_namespace: key.0.clone(), - track_name: key.1.as_ref().into(), - subscriber_priority: super::super::priority::DEFAULT_WIRE, - // No update can change it, so it takes the order a reader gets by default. - group_order: track::Order::default().into(), - filter: Filter::NextObject, - fill: None, - properties_wanted: true, - forward: false, - rendezvous: false, - range_filters: false, - authorization_token: None, - }, - )?; - let mut response = std::pin::pin!(async { - kio::wait(|waiter| stream.writer.poll_flush(waiter)).await?; - self.read_subscribe_response(&mut stream, &mut messages).await - }); - let answered = kio::wait(|waiter| { - if gone(waiter) { - return Poll::Ready(None); - } - waiter.poll_future(response.as_mut()).map(Some) - }) - .await; - match answered { - Some(Ok(Some(accepted))) => Ok(Some(accepted)), - Some(Ok(None)) => Err(Error::UnexpectedMessage), - Some(Err(err)) => Err(err), - None => Ok(None), - } - } - .await; - let accepted = match accepted { - Ok(Some(accepted)) => accepted, - Ok(None) => { - self.remove_subscribe(request_id); - self.cancel_subscribe(stream).await; - return Ok(None); - } - Err(err) => { - self.remove_subscribe(request_id); - return Err(err); - } - }; - if let Some(err) = refused(&accepted.properties, self.version, self.peer_setup.get().await.flate) { - self.remove_subscribe(request_id); - self.cancel_subscribe(stream).await; - return Err(err); - } - if let Err(err) = self.bind_alias(request_id, accepted.alias) { - self.remove_subscribe(request_id); - self.cancel_subscribe(stream).await; - return Err(err); - } - let info = accepted_info(&accepted.properties, self.version); - let listed = { - let mut listings = self.listings.lock(); - match listings.watched_by(key, handoff) { - Some(entry) => { - entry.request_id = Some(request_id); - entry.info = Some(info.clone()); - for lister in listings.listers(&key.0) { - let _ = lister.add(key.1.clone(), info.clone()); - } - true - } - None => false, - } + let mut task = Skipped::Opening { + this: self.clone(), + key, + handoff, + open: OpenPaused::new(), }; - if !listed { - self.remove_subscribe(request_id); - self.cancel_subscribe(stream).await; - return Ok(None); - } - Ok(Some(Upstream { - token_grant: None, - stream, - messages, - request_id, - alias: accepted.alias, - properties: accepted.properties, - owed: 0, - copy: None, - })) + self.tasks.push(move |waiter: &kio::Waiter| task.poll(waiter)); } /// Bind the alias a held subscription's PUBLISH or SUBSCRIBE_OK chose, paused, closing @@ -750,71 +437,16 @@ impl Subscriber { Err(err) } - /// Serve a PUBLISH: hold one answering our SUBSCRIBE_TRACKS paused as its track's entry, - /// refuse a redundant one, and decline one no listing asked for. - pub(super) async fn serve_publish( - &mut self, - mut stream: Stream, + /// What to do with a PUBLISH whose token, if it decides, granted it. + fn offer( + &self, + stream: Stream, msg: ietf::Publish<'static>, - ) -> Result<(), Error> { - let token_grant = if let Some(token) = msg - .authorization_token - .clone() - .filter(|_| self.auth.token_decides(crate::auth::Direction::Subscribe)) - { - let credit = kio::Producer::<()>::default(); - let hold: crate::auth::Hold = std::sync::Arc::new(credit.consume()); - let verdict = self.auth.verify_request( - token.0, - Some(hold), - crate::auth::RequestKind::Publish, - &msg.track_namespace, - ); - let session = self.session.clone(); - let verified = kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() - || request_stream::poll_withdrawn(&mut stream, request_stream::fin_cancels(self.version), waiter) - .is_ready() - { - return Poll::Ready(None); - } - verdict.poll_grant(waiter).map(Some) - }) - .await; - let Some(verified) = verified else { - drop(verdict); - kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() || credit.poll_wanted_changed(waiter, true).is_ready() { - Poll::Ready(()) - } else { - Poll::Pending - } - }) - .await; - return Ok(()); - }; - match crate::auth::RequestGrant::new( - &self.runtime, - verdict, - verified, - msg.track_namespace.clone(), - crate::auth::RequestKind::Publish, - None, - ) { - Ok(grant) => Some(grant), - Err(refusal) => { - let (kind, err) = (request::Kind::Publish, &refusal.error); - request_stream::buffer_error(&mut stream.writer, kind, err, &refusal.reason, self.version)?; - // Dropping an unfinished writer resets it, which could discard the refusal. - return request_stream::deliver(&mut stream.writer).await; - } - } - } else { - None - }; + token_grant: Option, + flate: bool, + ) -> ServePublishState { let key: TrackKey = (msg.track_namespace.clone(), msg.track_name.as_ref().into()); let info = accepted_info(&msg.properties, self.version); - let flate = self.peer_setup.get().await.flate; let offer = { let mut listings = self.listings.lock(); // Draft-18 allows one subscription per track (section 5.1); later drafts allow @@ -885,8 +517,8 @@ impl Subscriber { } } }; - let handoff = match offer { - Offer::Unsolicited => return self.run_publish_stream(stream, msg).await, + let reject = match offer { + Offer::Unsolicited => self.decline_publish(stream, msg), Offer::Refuse(code, reason) => { tracing::debug!(broadcast = %key.0, track = %key.1, reason, "refusing publish"); let error = ietf::RequestError { @@ -894,179 +526,65 @@ impl Subscriber { reason_phrase: reason.into(), retry_interval: 0, }; - return request_stream::reject(stream, error).await; - } - Offer::Hold(handoff) => handoff, - }; - - let request_id = msg.request_id; - self.state.lock().subscribes.insert( - request_id, - TrackState { - ..TrackState::pending(key.1.to_string(), key.0.clone(), kio::Producer::new(Fill::Done)) - }, - ); - let ok = crate::auth::RequestGrant::guard(token_grant.as_ref(), async { - self.bind_alias(request_id, msg.track_alias)?; - // PUBLISH_OK is a REQUEST_OK on draft-22. FORWARD=0 holds it paused on draft-18; - // on draft-22 the PUBLISH our FORWARD=0 asked for already is. - let id = match self.version { - Version::Draft18 => ietf::PublishOk::ID, - _ => ietf::RequestOk::ID, - }; - stream.writer.buffer_message( - id, - &ietf::PublishOk { - forward: false, - subscriber_priority: None, - group_order: None, - filter: None, - new_group: None, - }, - )?; - kio::wait(|waiter| stream.writer.poll_flush(waiter)).await - }) - .await - .map_err(|refusal| refusal.error) - .and_then(|result| result); - if let Err(err) = ok { - tracing::debug!(broadcast = %key.0, track = %key.1, %err, "failed to hold publish"); - self.remove_subscribe(request_id); - let mut listings = self.listings.lock(); - if listings.watched_by(&key, &handoff).is_some() { - listings.remove(&key, false); - listings.fail(&key.0, err); + request_stream::Reject::new(stream, error) } - return Ok(()); - } - tracing::debug!(broadcast = %crate::path::Rooted(self.origin.root(), &key.0), track = %key.1, id = %request_id, "holding publish"); - let upstream = Upstream { - token_grant, - stream, - messages: Messages::new(self.session.clone(), self.control.client()), - request_id, - alias: msg.track_alias, - properties: msg.properties, - owed: 0, - copy: None, - }; - self.watch(key, upstream, handoff).await; - Ok(()) - } - - /// Watch a paused subscription until it ends, its entry is let go, or a subscription - /// takes it to resume. - async fn watch(&self, key: TrackKey, mut upstream: Upstream, handoff: kio::Shared) { - let session = self.session.clone(); - enum Step { - Closed, - Gone, - Wanted, - Message(Result, Error>), - } - let end = loop { - let step = kio::wait(|waiter| { - if session.poll_closed(waiter).is_ready() { - return Poll::Ready(Step::Closed); - } - // A held subscription is handed over only once flushed, so resumes and pauses - // can't pile up behind a peer withholding credit. - let flushed = match upstream.stream.writer.poll_flush(waiter) { - Poll::Ready(Err(err)) => return Poll::Ready(Step::Message(Err(err))), - Poll::Ready(Ok(())) => true, - Poll::Pending => false, + Offer::Hold(handoff) => { + self.state.lock().subscribes.insert( + msg.request_id, + TrackState { + ..TrackState::pending(key.1.to_string(), key.0.clone(), kio::Producer::new(Fill::Done)) + }, + ); + return ServePublishState::Accepting { + stream, + msg, + token_grant, + key, + handoff, + answered: false, }; - if let Poll::Ready(handoff) = - handoff.poll(waiter, |handoff| match handoff.gone || (handoff.wanted && flushed) { - true => Poll::Ready(()), - false => Poll::Pending, - }) { - return Poll::Ready(match handoff.gone { - true => Step::Gone, - false => Step::Wanted, - }); - } - if let Some(grant) = &upstream.token_grant - && let Poll::Ready(refusal) = grant.poll_ended(waiter) - { - return Poll::Ready(Step::Message(Err(refusal.error))); - } - upstream - .messages - .poll_next(&mut upstream.stream.reader, waiter) - .map(Step::Message) - }) - .await; - match step { - Step::Closed => return, - Step::Wanted => { - let upstream = { - let mut held = handoff.lock(); - match (held.gone, held.wanted) { - (true, _) => upstream, - // The subscription gave up before it was handed over. - (false, false) => continue, - (false, true) => { - held.upstream = Some(upstream); - return; - } - } - }; - self.remove_subscribe(upstream.request_id); - self.cancel_subscribe(upstream.stream).await; - return; - } - Step::Gone => { - self.remove_subscribe(upstream.request_id); - self.cancel_subscribe(upstream.stream).await; - return; - } - Step::Message(Ok(Some(Framed(id, body)))) => match id { - // The answer to a pause, or to a resume that paused again before it landed. - ietf::RequestOk::ID | ietf::RequestError::ID if upstream.owed > 0 => { - if let Err(err) = self.owed_answer(id, &body) { - break Err(upstream.messages.failed(err)); - } - upstream.owed -= 1; - } - ietf::PublishDone::ID => { - break body - .message::(self.version) - .map_err(|err| upstream.messages.failed(err.into())) - .and_then(|done| listing_end(&done, self.version)); - } - _ => break Err(upstream.messages.failed(Error::UnexpectedMessage)), - }, - // PUBLISH_DONE must come before the FIN (draft-22 section 6.4.2.2). - Step::Message(Ok(None)) => break Err(Error::ProtocolViolation), - Step::Message(Err(err)) => break Err(err), } }; - tracing::debug!(broadcast = %key.0, track = %key.1, ?end, "held subscription ended"); - self.remove_subscribe(upstream.request_id); - // A pause still held back would go out truncated behind a FIN. - match upstream.stream.writer.is_flushed() { - true => { - let _ = upstream.stream.writer.finish(); - } - false => upstream.stream.writer.abort(&Error::Cancel), + ServePublishState::Rejecting { + reject, + _credit: None, + _grant: token_grant, } - // A failed upstream fails the copy, as it would while forwarding. TRACK_ENDED leaves it - // idle: the end is past groups it never saw, so a returning reader's SUBSCRIBE learns - // where, as for a copy whose subscription was cancelled. - if let (Err(err), Some(copy)) = (&end, upstream.copy.take()) { - let _ = copy.abort(err.clone()); - } - let mut listings = self.listings.lock(); - if listings.watched_by(&key, &handoff).is_some() { - match end { - Ok(()) => listings.remove(&key, true), - Err(err) => { - listings.remove(&key, false); - listings.fail(&key.0, err); - } - } + } + + /// Decline a PUBLISH no listing asked for. + /// + /// PUBLISH offers a single track, so honoring it means routing per + /// (namespace, track). Our model routes per namespace: a source attaches at a + /// path and serves every track under it, resolved on demand via SUBSCRIBE. + /// Accepting a PUBLISH would mean inventing a namespace-level source out of a + /// track-level offer, and that fiction then contradicts any real + /// PUBLISH_NAMESPACE for the same path. + /// + /// Declining the request rather than failing the session, since a peer using + /// a feature we don't implement is not a protocol violation. The rejection is the + /// whole exchange, but it still has to arrive: [`request_stream::Reject`] delivers it. + pub(super) fn decline_publish( + &self, + mut stream: Stream, + msg: ietf::Publish<'static>, + ) -> request_stream::Reject { + if let Some(error) = super::super::namespace::refusal(&msg.track_namespace, self.version) { + return request_stream::Reject::new(stream, error); } + tracing::debug!(broadcast = %msg.track_namespace, track = %msg.track_name, "rejecting publish"); + + // We decline the method itself rather than this particular track, which would be + // UNINTERESTED. + // + // The alias the message carries is deliberately not recorded. Nothing will ever bind + // it, and a rejected request has no lifetime of ours to hang the cleanup on, so the + // entry would have to be swept asynchronously. Any data streams the publisher opened + // before reading this are dropped by the unknown-alias path instead. + let (kind, err) = (request::Kind::Publish, &Error::Unsupported); + let buffered = + request_stream::buffer_error(&mut stream.writer, kind, err, "PUBLISH is not supported", self.version); + request_stream::Reject::buffered(stream, buffered) } /// Ask for the paused subscription a listing holds for this track, to resume it. @@ -1148,16 +666,8 @@ impl Subscriber { } None => false, }; - let this = self.clone(); - self.tasks.push(async move { - match held { - true => this.watch(key, upstream, handoff).await, - false => { - this.remove_subscribe(upstream.request_id); - this.cancel_subscribe(upstream.stream).await; - } - } - }); + let mut task = Watch::new(self.clone(), key, upstream, handoff, held); + self.tasks.push(move |waiter: &kio::Waiter| task.poll(waiter)); } /// Pause a subscription nobody reads that is its track's entry, and hold it again. The @@ -1244,17 +754,19 @@ impl Subscriber { } /// Read the answer to a SUBSCRIBE, or to the REQUEST_UPDATE resuming `held`. - pub(super) async fn read_answer( + pub(super) fn poll_read_answer( &self, + waiter: &kio::Waiter, stream: &mut Stream, messages: &mut Messages, held: Option<&mut Resumed>, - ) -> Result, Error> { + ) -> Poll, Error>> { let Some(held) = held else { - return self.read_subscribe_response(stream, messages).await; + let Framed(id, body) = ready!(messages.poll_message(&mut stream.reader, waiter))?; + return Poll::Ready(self.subscribe_response(id, &body).map(Some)); }; loop { - let Framed(id, body) = messages.message(&mut stream.reader).await?; + let Framed(id, body) = ready!(messages.poll_message(&mut stream.reader, waiter))?; match id { // The answer to a pause that went out before this resume. ietf::RequestOk::ID | ietf::RequestError::ID if held.owed > 0 => { @@ -1264,28 +776,1007 @@ impl Subscriber { ietf::RequestOk::ID => { let ok = body.message::(self.version)?; tracing::debug!(message = ?ok, "subscription resumed"); - return Ok(Some(Accepted { + return Poll::Ready(Ok(Some(Accepted { alias: held.alias, properties: held.properties.clone(), largest: ok.largest, - })); + }))); } ietf::RequestError::ID => { let refused = body.message::(self.version)?; tracing::debug!(message = ?refused, "resume refused"); - return Err(request::from_code( + return Poll::Ready(Err(request::from_code( refused.error_code, request::Kind::Subscribe, self.version, - )); + ))); } ietf::PublishDone::ID => { let done = body.message::(self.version)?; let end = listing_end(&done, self.version); held.ended = Some(end.clone()); - return Err(end.err().unwrap_or(Error::NotFound)); + return Poll::Ready(Err(end.err().unwrap_or(Error::NotFound))); + } + _ => return Poll::Ready(Err(Error::UnexpectedMessage)), + } + } + } + + /// The answer to a SUBSCRIBE: its SUBSCRIBE_OK, or the REQUEST_ERROR refusing it. + fn subscribe_response(&self, id: u64, body: &ietf::Body) -> Result { + let mut data = body.decoder(self.version); + match id { + ietf::SubscribeOk::ID => { + let msg = ietf::SubscribeOk::decode_msg(&mut data, self.version)?; + tracing::debug!(message = ?msg, "received subscribe ok"); + Ok(Accepted { + alias: msg.track_alias, + properties: msg.properties, + largest: msg.largest, + }) + } + ietf::RequestError::ID => { + let msg = ietf::RequestError::decode_msg(&mut data, self.version)?; + tracing::warn!(message = ?msg, "request error"); + Err(request::from_code( + msg.error_code, + request::Kind::Subscribe, + self.version, + )) + } + _ => Err(Error::UnexpectedMessage), + } + } +} + +/// One SUBSCRIBE_TRACKS, sent once the overlapping requests leave the wire, then read until +/// it is cancelled and the peer has closed it, or it fails. +struct SubscribeTracks { + this: Subscriber, + key: u64, + prefix: Path, + /// The overlapping requests that must leave the wire first. + blockers: Vec>, + live: kio::Consumer<()>, + /// Dropped once it is off the wire, so an overlapping one may go out. + gone: Option>, + state: SubscribeTracksState, +} + +enum SubscribeTracksState { + /// Waiting for the overlapping requests to leave the wire. + Blocked, + /// Waiting for a stream to send it on. + Opening(crate::transport::open::Open), + /// Reading its answers until it is cancelled. + Listing { + stream: Stream, + messages: Messages, + /// Its REQUEST_OK arrived. + established: bool, + /// A peer streaming skips must not hold the session's other tasks. + budget: kio::coop::Budget, + }, + /// Cancelled: reading until the peer closes its own half. A draft-18 FIN leaves the + /// writer to reset only once that is done. + Draining { + reader: crate::coding::Reader, + _writer: Option>, + messages: Messages, + budget: kio::coop::Budget, + }, + Done, +} + +impl SubscribeTracks { + fn new(this: Subscriber, opening: Opening) -> Self { + let Opening { + key, + prefix, + blockers, + live, + gone, + } = opening; + Self { + this, + key, + prefix, + blockers, + live, + gone: Some(gone), + state: SubscribeTracksState::Blocked, + } + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + let result = ready!(self.poll_request(waiter)); + self.state = SubscribeTracksState::Done; + let mut listings = self.this.listings.lock(); + if let Err(err) = result + && listings + .requests + .get(&self.key) + .is_some_and(|request| request.live.is_some()) + { + listings.fail(&self.prefix, err); + } + listings.requests.remove(&self.key); + listings.release(); + drop(listings); + // Off the wire: an overlapping request may go out. + self.gone = None; + Poll::Ready(()) + } + + fn poll_request(&mut self, waiter: &kio::Waiter) -> Poll> { + let this = &self.this; + let prefix = &self.prefix; + loop { + match &mut self.state { + SubscribeTracksState::Blocked => { + if self.live.poll_closed(waiter).is_ready() { + return Poll::Ready(Ok(())); + } + // The peer refuses a prefix overlapping one it still holds, so wait for each + // overlapping request to leave the wire. + for blocker in &self.blockers { + if blocker.poll_closed(waiter).is_pending() { + return Poll::Pending; + } + } + self.state = SubscribeTracksState::Opening(Default::default()); + } + SubscribeTracksState::Opening(open) => { + if self.live.poll_closed(waiter).is_ready() { + return Poll::Ready(Ok(())); + } + let stream = ready!(Stream::poll_open(&this.session, open, this.version, waiter))?; + let mut stream = stream.with_stats(this.origin.stats()); + let request_id = this.control.next_request_id(); + stream.writer.buffer_message( + ietf::SubscribeTracks::ID, + &ietf::SubscribeTracks { + request_id, + prefix: prefix.clone(), + forward: false, + group_order: None, + properties_wanted: true, + authorization_token: None, + }, + )?; + tracing::debug!(prefix = %crate::path::Rooted(this.origin.root(), prefix), id = %request_id, "subscribe_tracks started"); + self.state = SubscribeTracksState::Listing { + stream, + messages: Messages::new(this.session.clone(), this.control.client()), + established: false, + budget: kio::coop::Budget::new(32), + }; + } + SubscribeTracksState::Listing { + stream, + messages, + established, + budget, + } => { + loop { + ready!(budget.poll_yield(waiter)); + if this.session.poll_closed(waiter).is_ready() { + return Poll::Ready(Ok(())); + } + if self.live.poll_closed(waiter).is_ready() { + break; + } + ready!(stream.writer.poll_flush(waiter))?; + let Framed(id, body) = match ready!(messages.poll_next(&mut stream.reader, waiter))? { + // The peer ended the listing, which only a cancel of ours should do. + None if *established => return Poll::Ready(Err(Error::Closed)), + None => return Poll::Ready(Err(messages.failed(Error::ProtocolViolation))), + Some(framed) => framed, + }; + match id { + ietf::RequestOk::ID if !*established => { + body.message::(this.version) + .map_err(|err| messages.failed(err.into()))?; + *established = true; + } + ietf::RequestError::ID if !*established => { + let refused = body + .message::(this.version) + .map_err(|err| messages.failed(err.into()))?; + tracing::debug!(prefix = %prefix, code = refused.error_code, reason = %refused.reason_phrase, "subscribe_tracks refused"); + return Poll::Ready(Err(request::from_code( + refused.error_code, + request::Kind::SubscribeNamespace, + this.version, + ))); + } + ietf::PublishSkipped::ID if *established => { + let skipped = body + .message::(this.version) + .map_err(|err| messages.failed(err.into()))?; + let namespace = prefix + .join(&skipped.suffix) + .map_err(|err| messages.failed(err.into()))?; + this.skipped(namespace, skipped.track_name.as_ref().into()); + } + _ => return Poll::Ready(Err(messages.failed(Error::UnexpectedMessage))), + } + } + + // Cancelled: close our half (FIN on draft-18, RESET_STREAM from draft-22, where + // a FIN is not a cancel), then read until the peer closes its own. The + // PUBLISHes it made stay ours. + tracing::debug!(prefix = %crate::path::Rooted(this.origin.root(), prefix), "subscribe_tracks cancelled"); + let SubscribeTracksState::Listing { + stream: Stream { writer, mut reader }, + messages, + budget, + .. + } = std::mem::replace(&mut self.state, SubscribeTracksState::Done) + else { + unreachable!() + }; + // The request never reached the peer whole, so it holds nothing to wait on. + if !writer.is_flushed() { + reader.abort(&Error::Cancel); + writer.abort(&Error::Cancel); + return Poll::Ready(Ok(())); + } + let writer = match request_stream::fin_cancels(this.version) { + true => { + let mut writer = writer; + let _ = writer.finish(); + Some(writer) + } + false => { + writer.abort(&Error::Cancel); + None + } + }; + self.state = SubscribeTracksState::Draining { + reader, + _writer: writer, + messages, + budget, + }; + } + SubscribeTracksState::Draining { + reader, + messages, + budget, + .. + } => { + if this.session.poll_closed(waiter).is_ready() { + return Poll::Ready(Ok(())); + } + loop { + ready!(budget.poll_yield(waiter)); + match ready!(messages.poll_next(reader, waiter)) { + Ok(Some(_)) => continue, + Ok(None) | Err(_) => return Poll::Ready(Ok(())), + } + } + } + SubscribeTracksState::Done => unreachable!("polled after it ended"), + } + } + } +} + +/// A skipped track's paused SUBSCRIBE: opened, then watched as the track's entry. +#[expect( + clippy::large_enum_variant, + reason = "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" +)] +enum Skipped { + Opening { + this: Subscriber, + key: TrackKey, + handoff: kio::Shared, + open: OpenPaused, + }, + Watching(Watch), + Done, +} + +impl Skipped { + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + loop { + match self { + Self::Opening { + this, + key, + handoff, + open, + } => { + let result = ready!(open.poll(waiter, this, key, handoff)); + let Self::Opening { this, key, handoff, .. } = std::mem::replace(self, Self::Done) else { + unreachable!() + }; + match result { + Ok(Some(upstream)) => *self = Self::Watching(Watch::new(this, key, upstream, handoff, true)), + Ok(None) => return Poll::Ready(()), + Err(err) => { + tracing::debug!(broadcast = %key.0, track = %key.1, %err, "paused subscribe failed"); + let mut listings = this.listings.lock(); + if listings.watched_by(&key, &handoff).is_some() { + listings.remove(&key, false); + listings.fail(&key.0, err); + } + return Poll::Ready(()); + } + } + } + Self::Watching(watch) => return watch.poll(waiter), + Self::Done => return Poll::Ready(()), + } + } + } +} + +/// Opens a paused SUBSCRIBE for a skipped track and lists it once its SUBSCRIBE_OK +/// arrives. `None` once its entry was let go. +struct OpenPaused { + state: OpenPausedState, +} + +#[expect( + clippy::large_enum_variant, + reason = "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" +)] +enum OpenPausedState { + /// Waiting for a stream to send it on. + Opening(crate::transport::open::Open), + /// Flushed while waiting for the answer, so the entry going still ends a request the + /// peer's send window holds back. + Answering { + stream: Stream, + messages: Messages, + request_id: RequestId, + }, + /// Accepted, waiting for the peer's SETUP to judge its properties. + Accepted { + stream: Stream, + messages: Messages, + request_id: RequestId, + accepted: Accepted, + }, + /// Abandoned: cancelling it before resolving with `result`. + Cancelling { + cancel: request_stream::Cancel, + result: Option, Error>>, + }, + Done, +} + +impl OpenPausedState { + /// Remove and cancel the subscription, then resolve with `result`. + fn cancel( + this: &Subscriber, + request_id: RequestId, + stream: Stream, + result: Result, Error>, + ) -> Self { + this.remove_subscribe(request_id); + Self::Cancelling { + cancel: request_stream::Cancel::new(stream), + result: Some(result), + } + } +} + +impl OpenPaused { + fn new() -> Self { + Self { + state: OpenPausedState::Opening(Default::default()), + } + } + + fn poll( + &mut self, + waiter: &kio::Waiter, + this: &Subscriber, + key: &TrackKey, + handoff: &kio::Shared, + ) -> Poll, Error>> { + let gone = |waiter: &kio::Waiter| { + handoff + .poll(waiter, |handoff| match handoff.gone { + true => Poll::Ready(()), + false => Poll::Pending, + }) + .is_ready() + }; + loop { + match &mut self.state { + OpenPausedState::Opening(open) => { + if gone(waiter) { + return Poll::Ready(Ok(None)); + } + let stream = ready!(Stream::poll_open(&this.session, open, this.version, waiter))?; + let mut stream = stream.with_stats(this.origin.stats()); + let request_id = this.control.next_request_id(); + this.state.lock().subscribes.insert( + request_id, + TrackState { + ..TrackState::pending(key.1.to_string(), key.0.clone(), kio::Producer::new(Fill::Done)) + }, + ); + let messages = Messages::new(this.session.clone(), this.control.client()); + let buffered = stream.writer.buffer_message( + ietf::Subscribe::ID, + &ietf::Subscribe { + max_age: None, + new_group: None, + request_id, + track_namespace: key.0.clone(), + track_name: key.1.as_ref().into(), + subscriber_priority: super::super::priority::DEFAULT_WIRE, + // No update can change it, so it takes the order a reader gets by default. + group_order: track::Order::default().into(), + filter: Filter::NextObject, + fill: None, + properties_wanted: true, + forward: false, + rendezvous: false, + range_filters: false, + authorization_token: None, + }, + ); + if let Err(err) = buffered { + this.remove_subscribe(request_id); + return Poll::Ready(Err(err)); + } + self.state = OpenPausedState::Answering { + stream, + messages, + request_id, + }; + } + OpenPausedState::Answering { + stream, + messages, + request_id, + } => { + let request_id = *request_id; + if gone(waiter) { + let OpenPausedState::Answering { stream, .. } = + std::mem::replace(&mut self.state, OpenPausedState::Done) + else { + unreachable!() + }; + self.state = OpenPausedState::cancel(this, request_id, stream, Ok(None)); + continue; + } + let answered = match stream.writer.poll_flush(waiter) { + Poll::Ready(Ok(())) => this.poll_read_answer(waiter, stream, messages, None), + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Pending => Poll::Pending, + }; + let accepted = match ready!(answered).and_then(|accepted| accepted.ok_or(Error::UnexpectedMessage)) + { + Ok(accepted) => accepted, + Err(err) => { + this.remove_subscribe(request_id); + return Poll::Ready(Err(err)); + } + }; + let OpenPausedState::Answering { stream, messages, .. } = + std::mem::replace(&mut self.state, OpenPausedState::Done) + else { + unreachable!() + }; + self.state = OpenPausedState::Accepted { + stream, + messages, + request_id, + accepted, + }; + } + OpenPausedState::Accepted { .. } => { + let flate = ready!(this.peer_setup.poll_get(waiter)).flate; + let OpenPausedState::Accepted { + stream, + messages, + request_id, + accepted, + } = std::mem::replace(&mut self.state, OpenPausedState::Done) + else { + unreachable!() + }; + if let Some(err) = refused(&accepted.properties, this.version, flate) { + self.state = OpenPausedState::cancel(this, request_id, stream, Err(err)); + continue; + } + if let Err(err) = this.bind_alias(request_id, accepted.alias) { + self.state = OpenPausedState::cancel(this, request_id, stream, Err(err)); + continue; + } + let info = accepted_info(&accepted.properties, this.version); + let listed = { + let mut listings = this.listings.lock(); + match listings.watched_by(key, handoff) { + Some(entry) => { + entry.request_id = Some(request_id); + entry.info = Some(info.clone()); + for lister in listings.listers(&key.0) { + let _ = lister.add(key.1.clone(), info.clone()); + } + true + } + None => false, + } + }; + if !listed { + self.state = OpenPausedState::cancel(this, request_id, stream, Ok(None)); + continue; + } + return Poll::Ready(Ok(Some(Upstream { + token_grant: None, + stream, + messages, + request_id, + alias: accepted.alias, + properties: accepted.properties, + owed: 0, + copy: None, + }))); + } + OpenPausedState::Cancelling { cancel, result } => { + ready!(cancel.poll(waiter)); + let result = result.take().expect("resolved once"); + self.state = OpenPausedState::Done; + return Poll::Ready(result); + } + OpenPausedState::Done => unreachable!("polled after it resolved"), + } + } + } +} + +/// Watches a paused subscription until it ends, its entry is let go, or a subscription +/// takes it to resume. +struct Watch { + this: Subscriber, + key: TrackKey, + /// Taken once it is handed over, cancelled, or ended. + upstream: Option, + handoff: kio::Shared, + state: WatchState, +} + +#[expect( + clippy::large_enum_variant, + reason = "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" +)] +enum WatchState { + /// Its entry was already gone when it was held, so it is cancelled at once. + Unheld, + Watching, + /// Cancelling it. The rest of the subscription, its grant and lingering copy included, + /// is released only after. + Cancelling { + cancel: request_stream::Cancel, + _rest: (Option, Messages, Option), + }, + Done, +} + +impl Watch { + /// Watch `upstream` as the entry `handoff` belongs to, or cancel it unless `held`. + fn new(this: Subscriber, key: TrackKey, upstream: Upstream, handoff: kio::Shared, held: bool) -> Self { + Self { + this, + key, + upstream: Some(upstream), + handoff, + state: match held { + true => WatchState::Watching, + false => WatchState::Unheld, + }, + } + } + + /// Remove and cancel the subscription. + fn cancel(&mut self) { + let Upstream { + token_grant, + stream, + messages, + request_id, + copy, + .. + } = self.upstream.take().expect("cancelled once"); + self.this.remove_subscribe(request_id); + self.state = WatchState::Cancelling { + cancel: request_stream::Cancel::new(stream), + _rest: (token_grant, messages, copy), + }; + } + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + enum Step { + Closed, + Gone, + Wanted, + Message(Result, Error>), + } + let end = loop { + match &mut self.state { + WatchState::Unheld => { + self.cancel(); + continue; + } + WatchState::Cancelling { cancel, .. } => { + ready!(cancel.poll(waiter)); + self.state = WatchState::Done; + return Poll::Ready(()); + } + WatchState::Done => return Poll::Ready(()), + WatchState::Watching => {} + } + let upstream = self.upstream.as_mut().expect("watched"); + let step = 'step: { + if self.this.session.poll_closed(waiter).is_ready() { + break 'step Step::Closed; + } + // A held subscription is handed over only once flushed, so resumes and pauses + // can't pile up behind a peer withholding credit. + let flushed = match upstream.stream.writer.poll_flush(waiter) { + Poll::Ready(Err(err)) => break 'step Step::Message(Err(err)), + Poll::Ready(Ok(())) => true, + Poll::Pending => false, + }; + if let Poll::Ready(handoff) = + self.handoff + .poll(waiter, |handoff| match handoff.gone || (handoff.wanted && flushed) { + true => Poll::Ready(()), + false => Poll::Pending, + }) { + break 'step match handoff.gone { + true => Step::Gone, + false => Step::Wanted, + }; + } + if let Some(grant) = &upstream.token_grant + && let Poll::Ready(refusal) = grant.poll_ended(waiter) + { + break 'step Step::Message(Err(refusal.error)); + } + Step::Message(ready!(upstream.messages.poll_next(&mut upstream.stream.reader, waiter))) + }; + match step { + Step::Closed => { + self.upstream = None; + self.state = WatchState::Done; + return Poll::Ready(()); + } + Step::Wanted => { + let mut held = self.handoff.lock(); + match (held.gone, held.wanted) { + (true, _) => {} + // The subscription gave up before it was handed over. + (false, false) => continue, + (false, true) => { + held.upstream = self.upstream.take(); + drop(held); + self.state = WatchState::Done; + return Poll::Ready(()); + } + } + drop(held); + self.cancel(); + } + Step::Gone => self.cancel(), + Step::Message(Ok(Some(Framed(id, body)))) => match id { + // The answer to a pause, or to a resume that paused again before it landed. + ietf::RequestOk::ID | ietf::RequestError::ID if upstream.owed > 0 => { + if let Err(err) = self.this.owed_answer(id, &body) { + break Err(upstream.messages.failed(err)); + } + upstream.owed -= 1; + } + ietf::PublishDone::ID => { + break body + .message::(self.this.version) + .map_err(|err| upstream.messages.failed(err.into())) + .and_then(|done| listing_end(&done, self.this.version)); + } + _ => break Err(upstream.messages.failed(Error::UnexpectedMessage)), + }, + // PUBLISH_DONE must come before the FIN (draft-22 section 6.4.2.2). + Step::Message(Ok(None)) => break Err(Error::ProtocolViolation), + Step::Message(Err(err)) => break Err(err), + } + }; + + self.state = WatchState::Done; + let key = &self.key; + let mut upstream = self.upstream.take().expect("watched"); + tracing::debug!(broadcast = %key.0, track = %key.1, ?end, "held subscription ended"); + self.this.remove_subscribe(upstream.request_id); + // A pause still held back would go out truncated behind a FIN. + match upstream.stream.writer.is_flushed() { + true => { + let _ = upstream.stream.writer.finish(); + } + false => upstream.stream.writer.abort(&Error::Cancel), + } + // A failed upstream fails the copy, as it would while forwarding. TRACK_ENDED leaves it + // idle: the end is past groups it never saw, so a returning reader's SUBSCRIBE learns + // where, as for a copy whose subscription was cancelled. + if let (Err(err), Some(copy)) = (&end, upstream.copy.take()) { + let _ = copy.abort(err.clone()); + } + let mut listings = self.this.listings.lock(); + if listings.watched_by(key, &self.handoff).is_some() { + match end { + Ok(()) => listings.remove(key, true), + Err(err) => { + listings.remove(key, false); + listings.fail(&key.0, err); } - _ => return Err(Error::UnexpectedMessage), + } + } + Poll::Ready(()) + } +} + +/// Serve a PUBLISH: hold one answering our SUBSCRIBE_TRACKS paused as its track's entry, +/// refuse a redundant one, and decline one no listing asked for. +pub(super) struct ServePublish { + this: Subscriber, + state: ServePublishState, +} + +#[expect( + clippy::large_enum_variant, + reason = "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" +)] +enum ServePublishState { + /// Waiting for the acceptor to judge the PUBLISH's token. + Verifying { + stream: Stream, + msg: ietf::Publish<'static>, + verdict: crate::auth::RequestVerdict, + /// What the acceptor holds while it decides: its credit lasts as long as this. + credit: kio::Producer<()>, + }, + /// The requester withdrew before the verdict: waiting for the acceptor to let go. + Withdrawn { + credit: kio::Producer<()>, + _stream: Stream, + }, + /// Waiting for the peer's SETUP to judge the track's properties. + Setup { + stream: Stream, + msg: ietf::Publish<'static>, + token_grant: Option, + }, + /// Refusing or declining it, holding what the request held until it is delivered. + Rejecting { + reject: request_stream::Reject, + _credit: Option>, + _grant: Option, + }, + /// Answering PUBLISH_OK, under the token's grant. + Accepting { + stream: Stream, + msg: ietf::Publish<'static>, + token_grant: Option, + key: TrackKey, + handoff: kio::Shared, + /// The alias is bound and PUBLISH_OK buffered. + answered: bool, + }, + /// Held as its track's entry. + Watching(Watch), + Done, +} + +impl ServePublish { + /// Start serving `msg`, the PUBLISH that opened `stream`. + pub(super) fn new(this: Subscriber, stream: Stream, msg: ietf::Publish<'static>) -> Self { + let state = match msg + .authorization_token + .clone() + .filter(|_| this.auth.token_decides(crate::auth::Direction::Subscribe)) + { + Some(token) => { + let credit = kio::Producer::<()>::default(); + let hold: crate::auth::Hold = std::sync::Arc::new(credit.consume()); + let verdict = this.auth.verify_request( + token.0, + Some(hold), + crate::auth::RequestKind::Publish, + &msg.track_namespace, + ); + ServePublishState::Verifying { + stream, + msg, + verdict, + credit, + } + } + None => ServePublishState::Setup { + stream, + msg, + token_grant: None, + }, + }; + Self { this, state } + } + + /// Ready once the PUBLISH is refused, declined, or its held subscription ends. + pub(super) fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + let this = &self.this; + loop { + match &mut self.state { + ServePublishState::Verifying { stream, verdict, .. } => { + let withdrawn = this.session.poll_closed(waiter).is_ready() + || request_stream::poll_withdrawn(stream, request_stream::fin_cancels(this.version), waiter) + .is_ready(); + if withdrawn { + let ServePublishState::Verifying { stream, credit, .. } = + std::mem::replace(&mut self.state, ServePublishState::Done) + else { + unreachable!() + }; + self.state = ServePublishState::Withdrawn { + credit, + _stream: stream, + }; + continue; + } + let verified = ready!(verdict.poll_grant(waiter)); + let ServePublishState::Verifying { + mut stream, + msg, + verdict, + credit, + } = std::mem::replace(&mut self.state, ServePublishState::Done) + else { + unreachable!() + }; + let granted = crate::auth::RequestGrant::new( + &this.runtime, + verdict, + verified, + msg.track_namespace.clone(), + crate::auth::RequestKind::Publish, + None, + ); + self.state = match granted { + Ok(grant) => ServePublishState::Setup { + stream, + msg, + token_grant: Some(grant), + }, + Err(refusal) => { + let (kind, err) = (request::Kind::Publish, &refusal.error); + let buffered = request_stream::buffer_error( + &mut stream.writer, + kind, + err, + &refusal.reason, + this.version, + ); + // Dropping an unfinished writer resets it, which could discard the refusal. + ServePublishState::Rejecting { + reject: request_stream::Reject::buffered(stream, buffered), + _credit: Some(credit), + _grant: None, + } + } + }; + } + ServePublishState::Withdrawn { credit, .. } => { + if this.session.poll_closed(waiter).is_ready() + || credit.poll_wanted_changed(waiter, true).is_ready() + { + self.state = ServePublishState::Done; + return Poll::Ready(Ok(())); + } + return Poll::Pending; + } + ServePublishState::Setup { .. } => { + let flate = ready!(this.peer_setup.poll_get(waiter)).flate; + let ServePublishState::Setup { + stream, + msg, + token_grant, + } = std::mem::replace(&mut self.state, ServePublishState::Done) + else { + unreachable!() + }; + self.state = this.offer(stream, msg, token_grant, flate); + } + ServePublishState::Rejecting { reject, .. } => { + let result = ready!(reject.poll(waiter)); + self.state = ServePublishState::Done; + return Poll::Ready(result); + } + ServePublishState::Accepting { + stream, + msg, + token_grant, + answered, + .. + } => { + // The grant ending at any point ends the request there, and is checked first, + // so a request never acts on authority that already ended. + let ok = match token_grant.as_ref().map(|grant| grant.poll_ended(waiter)) { + Some(Poll::Ready(refusal)) => Err(refusal.error), + _ => { + let mut answer = || { + if !std::mem::replace(answered, true) { + this.bind_alias(msg.request_id, msg.track_alias)?; + // PUBLISH_OK is a REQUEST_OK on draft-22. FORWARD=0 holds it paused on + // draft-18; on draft-22 the PUBLISH our FORWARD=0 asked for already is. + let id = match this.version { + Version::Draft18 => ietf::PublishOk::ID, + _ => ietf::RequestOk::ID, + }; + stream.writer.buffer_message( + id, + &ietf::PublishOk { + forward: false, + subscriber_priority: None, + group_order: None, + filter: None, + new_group: None, + }, + )?; + } + Ok(()) + }; + match answer() { + Ok(()) => ready!(stream.writer.poll_flush(waiter)), + Err(err) => Err(err), + } + } + }; + let ServePublishState::Accepting { + stream, + msg, + token_grant, + key, + handoff, + .. + } = std::mem::replace(&mut self.state, ServePublishState::Done) + else { + unreachable!() + }; + let request_id = msg.request_id; + if let Err(err) = ok { + tracing::debug!(broadcast = %key.0, track = %key.1, %err, "failed to hold publish"); + this.remove_subscribe(request_id); + let mut listings = this.listings.lock(); + if listings.watched_by(&key, &handoff).is_some() { + listings.remove(&key, false); + listings.fail(&key.0, err); + } + return Poll::Ready(Ok(())); + } + tracing::debug!(broadcast = %crate::path::Rooted(this.origin.root(), &key.0), track = %key.1, id = %request_id, "holding publish"); + let upstream = Upstream { + token_grant, + stream, + messages: Messages::new(this.session.clone(), this.control.client()), + request_id, + alias: msg.track_alias, + properties: msg.properties, + owed: 0, + copy: None, + }; + self.state = ServePublishState::Watching(Watch::new(this.clone(), key, upstream, handoff, true)); + } + ServePublishState::Watching(watch) => { + ready!(watch.poll(waiter)); + self.state = ServePublishState::Done; + return Poll::Ready(Ok(())); + } + ServePublishState::Done => unreachable!("polled after it resolved"), } } } @@ -1515,11 +2006,12 @@ mod tests { let id = ours.reader.varint().await.unwrap(); assert_eq!(id, ietf::Publish::ID); let body: ietf::Body = ours.reader.decode().await.unwrap(); - let mut subscriber = self.subscriber.clone(); + let subscriber = self.subscriber.clone(); let version = self.version; kio::clock::sim::spawn(async move { let msg = body.message::(version).unwrap().into_owned(); - let _ = subscriber.serve_publish(ours, msg).await; + let mut serve = ServePublish::new(subscriber, ours, msg); + let _ = kio::wait(|waiter| serve.poll(waiter)).await; }); let (id, body) = read(&mut stream).await; (stream, id, body) diff --git a/rs/moq-net/src/ietf/token.rs b/rs/moq-net/src/ietf/token.rs index 68222190458..d74d89c2789 100644 --- a/rs/moq-net/src/ietf/token.rs +++ b/rs/moq-net/src/ietf/token.rs @@ -108,15 +108,15 @@ pub(super) fn reason(err: &DecodeError) -> &'static str { /// /// Section 8.9 rejects the message rather than the session. #[cfg(feature = "async")] -pub(super) async fn reject( +pub(super) fn rejection( mut stream: Stream, kind: request::Kind, err: DecodeError, version: Version, -) -> Result<(), Error> { +) -> super::request_stream::Reject { let reason = reason(&err); - super::request_stream::buffer_error(&mut stream.writer, kind, &Error::Decode(err), reason, version)?; - super::request_stream::deliver(&mut stream.writer).await + let buffered = super::request_stream::buffer_error(&mut stream.writer, kind, &Error::Decode(err), reason, version); + super::request_stream::Reject::buffered(stream, buffered) } /// Encode `token` by value. diff --git a/rs/moq-net/src/lib.rs b/rs/moq-net/src/lib.rs index 96c10f71c14..1e04e26b1ad 100644 --- a/rs/moq-net/src/lib.rs +++ b/rs/moq-net/src/lib.rs @@ -6,12 +6,11 @@ //! protocol (the default) or the full IETF `moq-transport` protocol. //! //! ## Features -//! The default `async` feature enables async helpers, the announcement `Stream` adapter, -//! `time::run`, and -//! the Lite/IETF sessions (`Client`, `Server`, `Session`, and the session `Driver`). -//! With `default-features = false`, use the model's `poll_*` methods and -//! [`origin::Driver`] with explicit time. Origin serving stays available without -//! a native runtime; no executor or clock is chosen by the model. +//! The `async` feature, off by default, enables the generated async wrappers for the `poll_*` +//! methods, the announcement `Stream` adapter, `time::run`, and the Lite/IETF sessions +//! (`Client`, `Server`, `Session`, and the session `Driver`). Without it, use the model's +//! `poll_*` methods and [`origin::Driver`] with explicit time. Origin serving stays available +//! without a native runtime; no executor or clock is chosen by the model. //! //! ## API //! The API is built around Producer/Consumer pairs, with the hierarchy: diff --git a/rs/moq-net/src/lite/session.rs b/rs/moq-net/src/lite/session.rs index d0632dc2b28..2d8a73d43f3 100644 --- a/rs/moq-net/src/lite/session.rs +++ b/rs/moq-net/src/lite/session.rs @@ -31,7 +31,7 @@ pub(crate) struct SessionStart { pub peer_implementation: crate::setup::PeerImplementation, } -/// Server: read the peer's single SETUP message off its Setup Stream before starting +/// Server: reads the peer's single SETUP message off its Setup Stream before starting /// the session, so the caller can inspect the advertised path (and gate on it) before /// serving. lite-05+ only. /// @@ -40,17 +40,7 @@ pub(crate) struct SessionStart { /// /// Pass the result to [`start`] as its `peer_setup` so capability gating resolves /// without re-reading SETUP, and every held stream is handed back to the driver. -pub async fn accept_setup(session: &mut Transport, version: Version) -> Result { - let mut accept = SetupAccept { - version, - pending: Vec::new(), - early: Vec::new(), - setup: None, - }; - kio::wait(|waiter| accept.poll(session, waiter)).await -} - -struct SetupAccept { +pub struct SetupAccept { version: Version, pending: Vec>, early: Vec>, @@ -58,7 +48,18 @@ struct SetupAccept { } impl SetupAccept { - fn poll(&mut self, session: &mut Transport, waiter: &kio::Waiter) -> Poll> { + /// Start waiting for a `version` peer's SETUP. + pub fn new(version: Version) -> Self { + Self { + version, + pending: Vec::new(), + early: Vec::new(), + setup: None, + } + } + + /// The peer's SETUP, with the streams held until it arrived. + pub fn poll(&mut self, session: &Transport, waiter: &kio::Waiter) -> Poll> { let Self { version, pending, @@ -142,6 +143,13 @@ impl SetupAccept { } } +/// [`SetupAccept`] as a future, for tests. +#[cfg(test)] +pub(crate) async fn accept_setup(session: &Transport, version: Version) -> Result { + let mut accept = SetupAccept::new(version); + kio::wait(|waiter| accept.poll(session, waiter)).await +} + /// The peer's pre-read SETUP and the streams held until it arrived. pub struct AcceptedSetup { /// The SETUP decoded before origin binding. @@ -194,7 +202,7 @@ pub struct Config { pub our_setup: Setup, /// The peer's SETUP, when it was already read before [`start`] (e.g. a server that - /// gated on the client's path via [`accept_setup`]). Seeds the peer-setup slot so + /// gated on the client's path via [`SetupAccept`]). Seeds the peer-setup slot so /// the Setup Stream isn't expected again. `None` reads it from the wire as usual. pub peer_setup: Option, @@ -1220,7 +1228,7 @@ mod tests { .unwrap(); let session = crate::lite::test_transport::ScriptedSession::new(Vec::new()).with_incoming_unis(vec![Vec::new(), setup]); - let accepted = accept_setup(&mut crate::driver::Transport::new(session.clone()), version) + let accepted = accept_setup(&crate::driver::Transport::new(session.clone()), version) .await .unwrap(); assert_eq!(accepted.early.len(), 1); @@ -1232,7 +1240,7 @@ mod tests { let session = crate::lite::test_transport::ScriptedSession::new(Vec::new()).with_incoming_unis(vec![vec![1], vec![1]]); assert!(matches!( - accept_setup(&mut crate::driver::Transport::new(session.clone()), Version::Lite05).now_or_never(), + accept_setup(&crate::driver::Transport::new(session.clone()), Version::Lite05).now_or_never(), Some(Err(Error::ProtocolViolation)) )); } @@ -1386,7 +1394,7 @@ mod tests { async fn unknown_client_role_remains_bidirectional() { for version in [Version::Lite05, Version::Lite06, Version::Lite07] { let transport = ScriptedSession::new(Vec::new()).with_incoming_unis(vec![vec![1, 4, 1, 3, 1, 3]]); - let setup = accept_setup(&mut crate::driver::Transport::new(transport.clone()), version) + let setup = accept_setup(&crate::driver::Transport::new(transport.clone()), version) .await .unwrap(); assert_eq!(setup.setup.role, None); diff --git a/rs/moq-net/src/lite/subscriber.rs b/rs/moq-net/src/lite/subscriber.rs index c1641266fc5..a54e95cd82e 100644 --- a/rs/moq-net/src/lite/subscriber.rs +++ b/rs/moq-net/src/lite/subscriber.rs @@ -3960,7 +3960,7 @@ mod tests { let log = session.log.clone(); let peer_setup = lite::PeerSetup::new(false, crate::transport::Address::Setup); let early = if pre_read { - let accepted = lite::accept_setup(&mut crate::driver::Transport::new(session.clone()), version) + let accepted = lite::accept_setup(&crate::driver::Transport::new(session.clone()), version) .await .unwrap(); assert_eq!(accepted.early.len(), 2); diff --git a/rs/moq-net/src/model/bandwidth.rs b/rs/moq-net/src/model/bandwidth.rs index 46017b658ad..cd509137de7 100644 --- a/rs/moq-net/src/model/bandwidth.rs +++ b/rs/moq-net/src/model/bandwidth.rs @@ -167,14 +167,9 @@ impl Demand { !self.state.is_closed() && self.state.wanted() } - /// Return [`wanted`](Self::wanted) once it differs from what this handle last returned: + /// [`wanted`](Self::wanted) once it differs from what this handle last returned: /// immediately on the first call, then on each change. Errors once every producer is gone. - #[cfg(feature = "async")] - pub async fn wanted_changed(&mut self) -> Result { - kio::wait(|waiter| self.poll_wanted_changed(waiter)).await - } - - /// Poll for [`wanted`](Self::wanted) to differ from what this handle last returned. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_wanted_changed(&mut self, waiter: &kio::Waiter) -> Poll> { loop { if self.state.is_closed() { @@ -1156,7 +1151,7 @@ impl Consumer { } } - /// Poll for a bandwidth change without blocking. + /// The bandwidth estimate once it changes, or `None` when it has become unavailable. /// /// `Ok(None)` means the estimate is unavailable *for now*: the backend /// stopped reporting one, or the handle spans reconnects and is between @@ -1168,6 +1163,7 @@ impl Consumer { /// /// A backend with no bandwidth estimation at all yields no [Consumer] in the /// first place, so that case never reaches here. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_changed(&mut self, waiter: &kio::Waiter) -> Poll>> { let last = self.last; @@ -1207,18 +1203,6 @@ impl Consumer { self.last = bitrate; Poll::Ready(Ok(bitrate)) } - - /// Block until the bandwidth estimate changes, returning the new value, or - /// `None` when the estimate has become unavailable. - /// - /// # Errors - /// - /// Returns an error once the producer is closed or dropped, so a caller can - /// stop watching. See [`poll_changed`](Self::poll_changed). - #[cfg(feature = "async")] - pub async fn changed(&mut self) -> Result> { - kio::wait(|waiter| self.poll_changed(waiter)).await - } } #[cfg(test)] diff --git a/rs/moq-net/src/model/follow.rs b/rs/moq-net/src/model/follow.rs index 78d90882a38..66304339567 100644 --- a/rs/moq-net/src/model/follow.rs +++ b/rs/moq-net/src/model/follow.rs @@ -245,7 +245,7 @@ impl Following { } } - /// Wait for the next change to the instance serving the path; see + /// The next change to the instance serving the path; see /// [`Consumer::follow`](crate::origin::Consumer::follow). /// /// Returns the broadcast of each instance as it comes online or takes over, and `None` @@ -253,15 +253,8 @@ impl Following { /// instance takes it over. Returns [`Error::Closed`] once the origin closes, and any other /// resolution failure as-is, such as a handler refusing the path; the next call then waits /// for the serving route to change, even to another route of the same instance, instead of - /// asking the refusing route again. Dropping the future keeps the wait's progress for the - /// next call. - #[cfg(feature = "async")] - pub async fn available(&mut self) -> Result, Error> { - kio::wait(|waiter| self.poll_available(waiter)).await - } - - /// Poll for the next change to the instance serving the path, registering `waiter` when - /// there is none yet; see [`Self::available`]. + /// asking the refusing route again. An abandoned wait keeps its progress for the next call. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_available(&mut self, waiter: &kio::Waiter) -> Poll, Error>> { loop { // Apply every change on hand before reading an answer, so a burst of them hands out diff --git a/rs/moq-net/src/model/frame.rs b/rs/moq-net/src/model/frame.rs index ffdb2e148da..92f00eb0dac 100644 --- a/rs/moq-net/src/model/frame.rs +++ b/rs/moq-net/src/model/frame.rs @@ -1673,9 +1673,10 @@ impl Consumer { true } - /// Poll for the next chunk of bytes since the last read. + /// The next chunk of bytes since the last read. /// /// Returns `None` once the frame is finished and all bytes have been consumed. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_read_chunk(&mut self, waiter: &kio::Waiter) -> Poll>> { if self.recover.is_none() { return self.poll_read_chunk_once(waiter); @@ -1737,13 +1738,15 @@ impl Consumer { } } - /// Return the next chunk of bytes since the last read. - #[cfg(feature = "async")] - pub async fn read_chunk(&mut self) -> Result> { - kio::wait(|waiter| self.poll_read_chunk(waiter)).await + /// Return all remaining bytes once the frame is finished. + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_read_all(&mut self, waiter: &kio::Waiter) -> Poll> { + if self.recover.is_none() { + return self.poll_read_all_once(waiter); + } + self.poll_resumed(waiter, Self::poll_read_all_once) } - #[cfg(feature = "async")] fn poll_read_all_once(&mut self, waiter: &kio::Waiter) -> Poll> { if self.expired { return Poll::Ready(Err(Error::Old)); @@ -1779,18 +1782,6 @@ impl Consumer { self.stats.bytes(out.len() as u64); Poll::Ready(Ok(out)) } - - /// Return all remaining bytes, blocking until the frame is finished. - #[cfg(feature = "async")] - pub async fn read_all(&mut self) -> Result { - kio::wait(|waiter| { - if self.recover.is_none() { - return self.poll_read_all_once(waiter); - } - self.poll_resumed(waiter, Self::poll_read_all_once) - }) - .await - } } /// Poll the group channel, mapping a terminal close without an error to diff --git a/rs/moq-net/src/model/group.rs b/rs/moq-net/src/model/group.rs index 1039f909ba6..7cba1e9aa1f 100644 --- a/rs/moq-net/src/model/group.rs +++ b/rs/moq-net/src/model/group.rs @@ -2203,13 +2203,7 @@ impl Producer { err } - /// Block until the group's outcome is final: `Ok` once finished, or the abort error, - /// including eviction. - #[cfg(feature = "async")] - pub async fn closed(&self) -> Result<()> { - kio::wait(|waiter| self.handle.poll_closed(waiter)).await - } - /// Poll until a reader holds the group and the next frame, described by `frame`, fits + /// Resolves once a reader holds the group and the next frame, described by `frame`, fits /// within 16 MB past the end of the last frame its fastest reader took, so a local writer /// runs no further ahead than the live tail keeps behind that reader. /// @@ -2218,17 +2212,13 @@ impl Producer { /// for a frame larger than the window, and one whose next frame was shed has fallen off /// the group and is not waited on. Writing without waiting never fails: the group sheds /// its front instead. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_ready(&self, waiter: &kio::Waiter, frame: &frame::Info) -> Poll<()> { ready!(self.poll_paced(waiter, frame.size, super::pace::AHEAD)); // A closed group reports `Closed`, which is ready too. self.handle.state.poll_wanted_changed(waiter, false).map(|_| ()) } - /// Wait until a reader holds the group and `frame` fits; see [`Self::poll_ready`]. - pub async fn ready(&self, frame: &frame::Info) { - kio::wait(|waiter| self.poll_ready(waiter, frame)).await - } - /// Poll until the writer may take more from upstream, `partial` bytes into a frame /// it has not committed: until it is at most `window` bytes past the end of the last /// frame its fastest reader took. Ingest passes its stream's receive credit (see @@ -2287,8 +2277,9 @@ impl Producer { self.handle.consume() } - /// Poll until the group's outcome is final: `Ok` once finished, or the abort error, + /// Resolves once the group's outcome is final: `Ok` once finished, or the abort error, /// including eviction. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll> { self.handle.poll_closed(waiter) } @@ -2995,14 +2986,9 @@ impl Demand { !self.state.is_closed() && self.state.wanted() } - /// Return [`wanted`](Self::wanted) once it differs from what this handle last returned: + /// [`wanted`](Self::wanted) once it differs from what this handle last returned: /// immediately on the first call, then on each change. Errors once the group is closed. - #[cfg(feature = "async")] - pub async fn wanted_changed(&mut self) -> Result { - kio::wait(|waiter| self.poll_wanted_changed(waiter)).await - } - - /// Poll for [`wanted`](Self::wanted) to differ from what this handle last returned. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_wanted_changed(&mut self, waiter: &kio::Waiter) -> Poll> { loop { if self.state.is_closed() { @@ -3018,13 +3004,8 @@ impl Demand { } } - /// Wait until the group's outcome is final: `Ok` once finished, or the abort error. - #[cfg(feature = "async")] - pub async fn closed(&self) -> Result<()> { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - - /// Poll until the group's outcome is final: `Ok` once finished, or the abort error. + /// Resolves once the group's outcome is final: `Ok` once finished, or the abort error. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll> { settle(self.state.consume_weak().poll(waiter, |state| state.poll_outcome())) } @@ -3800,16 +3781,11 @@ impl Consumer { state.fin.unwrap_or(state.next_index) } - /// Return a consumer for the next frame for chunked reading. - #[cfg(feature = "async")] - pub async fn next_frame(&mut self) -> Result> { - kio::wait(|waiter| self.poll_next_frame(waiter)).await - } - - /// Poll for the next frame, without blocking. + /// The next frame, as a consumer for chunked reading. /// /// Returns None if the group is finished and the index is out of range, or the cursor /// passed the [`Self::with_range`] cap. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next_frame(&mut self, waiter: &kio::Waiter) -> Poll>> { let res = match self.recover.is_none() { true => { @@ -4017,11 +3993,17 @@ impl Consumer { Ok(out.filled_mut()) } - /// Poll until the group terminates, or this cursor reaches the end of its range, returning - /// its next frame index. + /// This cursor's next frame index once the group terminates, or this cursor reaches + /// the end of its range. + /// + /// This answers for the cursor, not the group: a reader that drained every frame gets the + /// clean end even if the group was aborted afterwards to release its cache, while one that + /// stopped short gets that abort. A prior [`Self::skip_to`] contributes to the index even + /// though those frames were not read. Use [`Self::frame_count`] for the producer's total. /// /// An error means the frames this cursor has yet to read are gone (evicted, cut, or /// failed upstream) and no route will continue them. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_finished(&mut self, waiter: &kio::Waiter) -> Poll> { let res = match self.recover.is_none() { true => self.poll_finished_once(waiter), @@ -4056,18 +4038,6 @@ impl Consumer { .map(|res| res.map(|()| index as u64)); self.poll_expired_if_blocked(waiter, res) } - - /// Block until the group terminates, or this cursor reaches the end of its range, returning - /// its next frame index. - /// - /// This answers for the cursor, not the group: a reader that drained every frame gets the - /// clean end even if the group was aborted afterwards to release its cache, while one that - /// stopped short gets that abort. A prior [`Self::skip_to`] contributes to the index even - /// though those frames were not read. Use [`Self::frame_count`] for the producer's total. - #[cfg(feature = "async")] - pub async fn finished(&mut self) -> Result { - kio::wait(|waiter| self.poll_finished(waiter)).await - } } impl Cursor { @@ -6276,6 +6246,7 @@ mod test { /// Every group starts cold even when a finished encoder is reused. Open groups /// keep independent windows, and aborting one does not consume another's spare. #[test] + #[cfg(feature = "async")] fn reused_encoders_start_independent_group_streams() { let pool = cache::Pool::unbounded(); let account = cache::Track::new(pool.clone(), Default::default(), kio::Weak::new()); diff --git a/rs/moq-net/src/model/listing.rs b/rs/moq-net/src/model/listing.rs index c0a2b435c5f..209a5086ccf 100644 --- a/rs/moq-net/src/model/listing.rs +++ b/rs/moq-net/src/model/listing.rs @@ -491,8 +491,9 @@ pub struct Tracks { } impl Tracks { - /// Poll for the next event: `None` once the broadcast ended, or the error that ended + /// The next event: `None` once the broadcast ended, or the error that ended /// the listing, such as an upstream refusing it. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll, Error>> { let mut queue = ready!(self.queue.poll(waiter, |queue| match queue.is_ready() { true => Poll::Ready(()), @@ -501,12 +502,6 @@ impl Tracks { Poll::Ready(queue.next()) } - /// Wait for the next event; see [`Self::poll_next`]. - #[cfg(feature = "async")] - pub async fn next(&mut self) -> Result, Error> { - kio::wait(|waiter| self.poll_next(waiter)).await - } - /// The broadcast being listed. pub fn broadcast(&self) -> &Consumer { &self.broadcast diff --git a/rs/moq-net/src/model/origin.rs b/rs/moq-net/src/model/origin.rs index 2faa52d7925..c38abc84bb4 100644 --- a/rs/moq-net/src/model/origin.rs +++ b/rs/moq-net/src/model/origin.rs @@ -3230,7 +3230,7 @@ impl Drop for DriverState { } } -/// Everything [`run_front`] owns, queued by [`Consumer::request`]. +/// Everything [`FrontServeRun`] owns, queued by [`Consumer::request`]. struct FrontTask { /// The route table the front selects from. shared: SendCell>, @@ -3412,24 +3412,65 @@ impl InFlight { /// Drives one front until it ends, then holds the tracks it left in flight until /// their last reader leaves, or until nothing owns the origin: a reader never keeps /// the driver from finishing. -async fn run_front(task: FrontTask, origin: TasksWeak) { - let mut in_flight = serve_front(task).await; - // Each keeps its copy for the readers still on their way, then lets go as an unread - // track parks, so the copy never keeps its source subscribed for nobody. Unread once - // is unread for good: the front closed its broadcast, which refuses every lookup, so - // no reader can arrive once the last one left. - kio::wait(|waiter| { +struct FrontServeRun { + state: FrontRunState, + origin: TasksWeak, +} + +#[allow(clippy::large_enum_variant, reason = "a machine's enum is its storage")] +enum FrontRunState { + /// Not polled yet: the front starts serving on its first poll, so a set that drops + /// it unpolled never starts it. + Start(FrontTask), + Serve(FrontServe), + /// Each keeps its copy for the readers still on their way, then lets go as an unread + /// track parks, so the copy never keeps its source subscribed for nobody. Unread once + /// is unread for good: the front closed its broadcast, which refuses every lookup, so + /// no reader can arrive once the last one left. + InFlight(Vec), +} + +impl FrontServeRun { + fn new(task: FrontTask, origin: TasksWeak) -> Self { + Self { + state: FrontRunState::Start(task), + origin, + } + } +} + +impl kio::Task for FrontServeRun { + type Output = (); + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll<()> { + let in_flight = loop { + match &mut self.state { + FrontRunState::Start(_) => { + let FrontRunState::Start(task) = + std::mem::replace(&mut self.state, FrontRunState::InFlight(Vec::new())) + else { + unreachable!("matched above"); + }; + self.state = FrontRunState::Serve(FrontServe::new(task)); + } + FrontRunState::Serve(serve) => { + let in_flight = ready!(serve.poll(waiter)); + self.state = FrontRunState::InFlight(in_flight); + } + FrontRunState::InFlight(in_flight) => break in_flight, + } + }; // Dropped without a release, so a reader on its way keeps the copy; it stays held // only while the logical track's state does, for an origin nobody owns. - if origin.poll_orphaned(waiter).is_ready() { - for io in &mut in_flight { + if self.origin.poll_orphaned(waiter).is_ready() { + for io in in_flight.iter_mut() { if let Some(request) = io.request.take() { request.reject(io.end.clone()); } } return Poll::Ready(()); } - for mut io in std::mem::take(&mut in_flight) { + for mut io in std::mem::take(in_flight) { if let Some(request) = &mut io.request { while let Poll::Ready(Some(_)) = request.poll_subscription_changed(waiter) {} } @@ -3466,8 +3507,7 @@ async fn run_front(task: FrontTask, origin: TasksWeak) { true => Poll::Ready(()), false => Poll::Pending, } - }) - .await; + } } /// The route a front may serve from. `resolved` is the instance the front first @@ -3489,112 +3529,159 @@ fn pick<'a>( } } +/// What [`FrontServe::poll_step`] returns: one thing that happened. +enum FrontStep { + Assigned(track::Request), + Resolved(Result), + SourceClosed(u64), + Info(Arc, u64, Result), + Ended(Arc, u64, Result<(), Error>, bool), + Demand(Arc), + Holders, + /// The broadcast closed: nobody can hold it again. + Closed, + Deadline, + Table, + /// Someone lists the broadcast, or the serving source changed: read its listing. + Listing, + /// Nobody lists the broadcast any more. + Unlisted, + Listed(Result, Error>), +} + +/// A front's in-flight upstream request. +struct FrontUpstream { + route: u64, + /// The route's instance when asked. + asked: Instance, + /// [`OriginState::next_route`] when asked. + asked_at: u64, + pending: kio::Consumer, + /// The queue's [`ServeState::renewals`], and its count when asked: a renewal + /// since released the request, whatever it resolved to and whenever. + renewals: Arc, + asked_renewals: u64, +} + /// Feeds the world's events to a [`Front`] and performs the actions it returns, -/// until the front ends; returns the tracks it ended in flight. The decisions live +/// until the front ends; resolves to the tracks it ended in flight. The decisions live /// in the machine; this only waits and executes, so nothing here decides anything /// twice. -async fn serve_front(task: FrontTask) -> Vec { - let FrontTask { - shared, - broadcast, - path, - horizon, - pin, - watch, - request, - timers, - } = task; - - /// What the wait below returns: one thing that happened. - enum Step { - Assigned(track::Request), - Resolved(Result), - SourceClosed(u64), - Info(Arc, u64, Result), - Ended(Arc, u64, Result<(), Error>, bool), - Demand(Arc), - Holders, - /// The broadcast closed: nobody can hold it again. - Closed, - Deadline, - Table, - /// Someone lists the broadcast, or the serving source changed: read its listing. - Listing, - /// Nobody lists the broadcast any more. - Unlisted, - Listed(Result, Error>), - } - - // The front serves its broadcast on demand: every track a reader names is - // handed here, and its readers read it from whichever source serves the path. - let mut dynamic = broadcast.requests(); - let lister = dynamic - .lister() - .expect("only the front answers its broadcast's listings"); - let mut listing: Option = None; - let mut front = Front::new(track::IDLE_LINGER); - let mut sources: HashMap = HashMap::new(); - // The instance each source's route served when it was asked, and the table's clock - // then: what its content is from. - let mut instances: HashMap = HashMap::new(); - let mut next_source = 0u64; - /// The in-flight upstream request. - struct Upstream { - route: u64, - /// The route's instance when asked. - asked: Instance, - /// [`OriginState::next_route`] when asked. - asked_at: u64, - pending: kio::Consumer, - /// The queue's [`ServeState::renewals`], and its count when asked: a renewal - /// since released the request, whatever it resolved to and whenever. - renewals: Arc, - asked_renewals: u64, - } - let mut upstream: Option = None; - let mut tracks: HashMap, TrackIo> = HashMap::new(); - let ready: kio::Queue> = kio::Queue::new(); - let mut deadline = crate::time::Deadline::new(&timers); - // The watch generation the last selection saw. - let mut seen = 0; - // What the front first resolved; see [`pick`]. - let mut resolved: Option = None; - // Whether a consumer held the broadcast as of the last edge fed to the machine, - // starting with the request that minted the front. - let mut held = true; - let mut events: VecDeque = VecDeque::new(); - // Once the front picks a route a session received, it asks sessions to keep - // discovering its path for as long as it lives, so the route outlives whatever - // reader found it, such as a listing that has since closed. - let mut discovered: Option = None; - // The source the broadcast's datagram lookups go to, as last handed over: datagram - // tracks are live only, so a lookup joins the serving source's own track rather than a - // copy the front would splice. - let mut datagram_source: Option = None; - broadcast.serve_datagrams_from(None); - - // Read the table for the machine: the best route and whether - // the serving source is on its way out. Also what the watch wakes for. - let select = |front: &mut Front, - sources: &HashMap, - instances: &mut HashMap, - seen: &mut u64, - resolved: &mut Option, - discovered: &mut Option| - -> Event { - let table = shared.read(); +/// +/// Declared in drop order: the state the front builds goes before the task's handles it +/// is built from, the newest first. +struct FrontServe { + /// The current wait's waiter. Each wait gets a fresh one: a reused waiter keeps the + /// registrations an earlier wait left, which wake the front spuriously. + park: kio::Park, + /// The source the broadcast's datagram lookups go to, as last handed over: datagram + /// tracks are live only, so a lookup joins the serving source's own track rather than a + /// copy the front would splice. + datagram_source: Option, + /// Once the front picks a route a session received, it asks sessions to keep + /// discovering its path for as long as it lives, so the route outlives whatever + /// reader found it, such as a listing that has since closed. + discovered: Option, + events: VecDeque, + /// Whether a consumer held the broadcast as of the last edge fed to the machine, + /// starting with the request that minted the front. + held: bool, + /// What the front first resolved; see [`pick`]. + resolved: Option, + /// The watch generation the last selection saw. + seen: u64, + deadline: crate::time::Deadline, + ready: kio::Queue>, + tracks: HashMap, TrackIo>, + upstream: Option, + next_source: u64, + /// The instance each source's route served when it was asked, and the table's clock + /// then: what its content is from. + instances: HashMap, + sources: HashMap, + front: Front, + listing: Option, + lister: broadcast::Lister, + /// The front serves its broadcast on demand: every track a reader names is + /// handed here, and its readers read it from whichever source serves the path. + dynamic: broadcast::Requests, + timers: Deadlines, + request: kio::Producer, + watch: Watch, + pin: Option, + horizon: Horizon, + path: Path, + broadcast: broadcast::Producer, + shared: SendCell>, +} + +impl FrontServe { + fn new(task: FrontTask) -> Self { + let FrontTask { + shared, + broadcast, + path, + horizon, + pin, + watch, + request, + timers, + } = task; + let dynamic = broadcast.requests(); + let lister = dynamic + .lister() + .expect("only the front answers its broadcast's listings"); + let deadline = crate::time::Deadline::new(&timers); + broadcast.serve_datagrams_from(None); + let mut serve = Self { + park: kio::Park::default(), + datagram_source: None, + discovered: None, + events: VecDeque::new(), + held: true, + resolved: None, + seen: 0, + deadline, + ready: kio::Queue::new(), + tracks: HashMap::new(), + upstream: None, + next_source: 0, + instances: HashMap::new(), + sources: HashMap::new(), + front: Front::new(track::IDLE_LINGER), + listing: None, + lister, + dynamic, + timers, + request, + watch, + pin, + horizon, + path, + broadcast, + shared, + }; + let event = serve.select(); + serve.events.push_back(event); + serve + } + + /// Read the table for the machine: the best route and whether + /// the serving source is on its way out. Also what the watch wakes for. + fn select(&mut self) -> Event { + let table = self.shared.read(); if table.closed { return Event::Closed; } // Read alongside the decision, under the lock a poke takes first. - *seen = watch.seen(); - front.retain_routes(|route| table.routes.covers(&path, route)); + self.seen = self.watch.seen(); + self.front.retain_routes(|route| table.routes.covers(&self.path, route)); // A route that took over the instance the front resolved continues it. The table // carries the front's verdict forward as each takeover lands; this also catches // one that landed before the front wrote its first, or won the path only later. - if let Some(instance) = resolved.as_mut() - && let Some(next) = request.write().ok().and_then(|mut pending| { - if let Some(best) = table.best_route(&path, horizon, &HashSet::new(), |_| true) { + if let Some(instance) = self.resolved.as_mut() + && let Some(next) = self.request.write().ok().and_then(|mut pending| { + if let Some(best) = table.best_route(&self.path, self.horizon, &HashSet::new(), |_| true) { PendingFront::continue_with(&mut pending, best); } match &pending.resolved { @@ -3602,19 +3689,25 @@ async fn serve_front(task: FrontTask) -> Vec { _ => None, } }) { - for (asked, _) in instances.values_mut() { + for (asked, _) in self.instances.values_mut() { if asked == instance { *asked = next.clone(); } } *instance = next; } - let best = pick(&table, &path, horizon, front, resolved.as_ref().or(pin.as_ref())); + let best = pick( + &table, + &self.path, + self.horizon, + &self.front, + self.resolved.as_ref().or(self.pin.as_ref()), + ); // Under the table's lock, which a request reads it under too: a request that // sees another winner from here on mints its own front rather than join this one. - if resolved.is_none() + if self.resolved.is_none() && let Some(entry) = best - && let Ok(mut pending) = request.write() + && let Ok(mut pending) = self.request.write() && pending.selected != entry.instance { pending.selected = entry.instance.clone(); @@ -3626,41 +3719,49 @@ async fn serve_front(task: FrontTask) -> Vec { local: entry.local, carried: !entry.instance.is_minted(), }); - let serving_closing = front + let serving_closing = self + .front .serving() - .and_then(|id| sources.get(&id)) + .and_then(|id| self.sources.get(&id)) .is_some_and(|source| source.is_closing()); drop(table); - if received && discovered.is_none() { - let mut table = shared.lock(); - *discovered = Some(Interest::new(&mut table, &shared, path.clone())); + if received && self.discovered.is_none() { + let mut table = self.shared.lock(); + self.discovered = Some(Interest::new(&mut table, &self.shared, self.path.clone())); } Event::Selected { best, serving_closing } - }; + } - // Whether a route's refusal is the answer: only the current winner speaks for the - // path. A retraction resolves like a rejection, and a route beaten while its - // request was pending speaks for nobody, so either one re-selects instead. - let standing = |front: &Front, route: u64, resolved: &Option| { - let table = shared.read(); - pick(&table, &path, horizon, front, resolved.as_ref().or(pin.as_ref())).is_some_and(|best| best.id == route) - }; + /// Whether a route's refusal is the answer: only the current winner speaks for the + /// path. A retraction resolves like a rejection, and a route beaten while its + /// request was pending speaks for nobody, so either one re-selects instead. + fn standing(&self, route: u64) -> bool { + let table = self.shared.read(); + pick( + &table, + &self.path, + self.horizon, + &self.front, + self.resolved.as_ref().or(self.pin.as_ref()), + ) + .is_some_and(|best| best.id == route) + } - // Whether an answer `route` gave for `asked` would serve another instance than the - // front's: the route has since moved off it, or, before the front resolved, another - // instance won the path while the request was in flight. - let stale = |front: &Front, route: u64, asked: &Instance, resolved: &Option| { - let table = shared.read(); - let Some(entry) = table.routes.covering(&path).find(|entry| entry.id == route) else { + /// Whether an answer `route` gave for `asked` would serve another instance than the + /// front's: the route has since moved off it, or, before the front resolved, another + /// instance won the path while the request was in flight. + fn stale(&self, route: u64, asked: &Instance) -> bool { + let table = self.shared.read(); + let Some(entry) = table.routes.covering(&self.path).find(|entry| entry.id == route) else { return false; }; if entry.instance != *asked { return true; } - match resolved { + match &self.resolved { Some(resolved) => resolved != asked, // A winner that took over the instance asked continues it. - None => pick(&table, &path, horizon, front, pin.as_ref()).is_some_and(|best| { + None => pick(&table, &self.path, self.horizon, &self.front, self.pin.as_ref()).is_some_and(|best| { best.instance != *asked && best .takeover @@ -3668,731 +3769,727 @@ async fn serve_front(task: FrontTask) -> Vec { .is_none_or(|takeover| takeover.instance != *asked) }), } - }; - - events.push_back(select( - &mut front, - &sources, - &mut instances, - &mut seen, - &mut resolved, - &mut discovered, - )); + } - loop { - while let Some(event) = events.pop_front() { - let refused = matches!( - &event, - Event::Resolved { - result: Err(Refusal { standing: true, .. }), - .. + /// Wait for the next thing to happen, the one place the front parks. + fn poll_step(&mut self, waiter: &kio::Waiter) -> Poll { + if let Poll::Ready(Ok(request)) = self.dynamic.poll_requested_track(waiter) { + return Poll::Ready(FrontStep::Assigned(request)); + } + if let Some(upstream) = &self.upstream + && let Poll::Ready(result) = upstream.pending.poll(waiter, |p| match &p.resolved { + Some(result) => Poll::Ready(result.clone()), + None => Poll::Pending, + }) { + return Poll::Ready(FrontStep::Resolved(match result { + Ok(resolved) => resolved, + // The queue died unresolved (its handler dropped): the route + // could not serve. + Err(_closed) => Err(Error::Unroutable), + })); + } + // The listing first: a source ends its listing before it closes, so this forwards + // the tracks it finished last before its close detaches it. + match &mut self.listing { + None => { + if self.front.serving().is_some() && matches!(self.lister.poll_wanted(waiter), Poll::Ready(Ok(()))) { + return Poll::Ready(FrontStep::Listing); } - ); - let answered = match &event { - Event::TrackInfo { track, source, .. } => Some((track.clone(), *source)), - _ => None, - }; - for action in front.step(event) { - match action { - Action::Reselect => events.push_back(select( - &mut front, - &sources, - &mut instances, - &mut seen, - &mut resolved, - &mut discovered, - )), - Action::Request { route } => { - // The entry and what it serves. - let found = { - let table = shared.read(); - table - .routes - .covering(&path) - .find(|entry| entry.id == route && entry.live()) - .map(|entry| { - let instance = (entry.instance.clone(), table.next_route); - (entry.source.clone(), entry.server.clone(), instance) - }) - }; - let Some((source, server, instance)) = found else { - events.push_back(Event::Resolved { - route, - result: Err(Refusal { - err: Error::Unroutable, - standing: false, - }), - }); - continue; - }; - if let Some(source) = source { - let id = next_source; - next_source += 1; - sources.insert(id, source); - instances.insert(id, instance); - events.push_back(Event::Resolved { route, result: Ok(id) }); - continue; + } + Some(state) => { + if self.lister.poll_unwanted(waiter).is_ready() { + return Poll::Ready(FrontStep::Unlisted); + } + if state.source != self.front.serving() { + return Poll::Ready(FrontStep::Listing); + } + if let Some(reader) = &mut state.reader + && let Poll::Ready(result) = reader.poll_next(waiter) + { + return Poll::Ready(FrontStep::Listed(result)); + } + } + } + if let Some(id) = self.front.serving() + && let Some(source) = self.sources.get(&id) + && source.poll_closed(waiter).is_ready() + { + return Poll::Ready(FrontStep::SourceClosed(id)); + } + while let Poll::Ready(Ok(name)) = self.ready.poll_pop(waiter) { + let Some(io) = self.tracks.get_mut(&name) else { continue }; + io.ready.popped(); + let waker = Waker::from(io.ready.clone()); + let waiter = io.park.hold(&Context::from_waker(&waker)); + let step = (|| { + if let Some(asked) = &mut io.query + && let Poll::Ready(result) = asked.query.poll(waiter) + { + return Some(FrontStep::Info(name.clone(), asked.source, result)); + } + follow_demand(&mut io.held, io.accepted.as_mut(), &io.weak, &io.routes, waiter); + // Settled once the copy closes: a group still open below a declared end is + // owed until then, and a session closing first means it never came. + if let Some((source, copy, _)) = &io.copy + && copy.poll_closed(waiter).is_ready() + { + let result = match copy.poll_complete(&kio::Waiter::noop()) { + Poll::Ready(result) => result, + Poll::Pending => { + // A channel cannot reopen after poll_closed reported its end. + debug_assert!(false, "closed copy did not report completion"); + Err(Error::Cancel) } - let Some(server) = server else { - events.push_back(Event::Resolved { - route, - result: Err(Refusal { - err: Error::Unroutable, - standing: standing(&front, route, &resolved), - }), - }); - continue; - }; - let mut serve = server.lock(); - if serve.closed { - // Retracted under us, or its handler dropped while the - // announcement stands: it cannot serve. A retraction - // leaves the table before it closes the server, so the - // table tells the two apart. - drop(serve); - events.push_back(Event::Resolved { + }; + let delivered = copy.latest_sequence().is_some(); + return Some(FrontStep::Ended(name.clone(), *source, result, delivered)); + } + // Watch the demand edge in whichever direction is unmet. Only `Pending` + // parks, so an answer steps and the handler looks again, polling anew when + // nothing changed. A closed track is unread: it owes one `Unwanted` if + // recorded read, and is no edge after. + let edge = match io.wanted { + true => io.weak.poll_unwanted(waiter).is_ready(), + false => matches!(io.weak.poll_wanted(waiter), Poll::Ready(Ok(()))), + }; + if edge { + return Some(FrontStep::Demand(name.clone())); + } + None + })(); + if let Some(step) = step { + // A ready answer did not park on its next edge. Revisit this track + // once the machine has consumed the event and applied its actions. + io.ready.wake_by_ref(); + return Poll::Ready(step); + } + } + // The same for whoever holds the broadcast itself. + let edge = match self.held { + true => self.broadcast.poll_unheld(waiter), + false => self.broadcast.poll_held(waiter), + }; + match edge { + Poll::Ready(Ok(())) => return Poll::Ready(FrontStep::Holders), + // Answers at once from here on, so it must end the front, not be polled again. + Poll::Ready(Err(_)) => return Poll::Ready(FrontStep::Closed), + Poll::Pending => {} + } + if self.deadline.poll(waiter).is_ready() { + return Poll::Ready(FrontStep::Deadline); + } + self.watch.poll_changed(waiter, self.seen).map(|()| FrontStep::Table) + } +} + +impl kio::Task for FrontServe { + type Output = Vec; + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + loop { + while let Some(event) = self.events.pop_front() { + let refused = matches!( + &event, + Event::Resolved { + result: Err(Refusal { standing: true, .. }), + .. + } + ); + let answered = match &event { + Event::TrackInfo { track, source, .. } => Some((track.clone(), *source)), + _ => None, + }; + for action in self.front.step(event) { + match action { + Action::Reselect => { + let event = self.select(); + self.events.push_back(event); + } + Action::Request { route } => { + // The entry and what it serves. + let found = { + let table = self.shared.read(); + table + .routes + .covering(&self.path) + .find(|entry| entry.id == route && entry.live()) + .map(|entry| { + let instance = (entry.instance.clone(), table.next_route); + (entry.source.clone(), entry.server.clone(), instance) + }) + }; + let Some((source, server, instance)) = found else { + self.events.push_back(Event::Resolved { + route, + result: Err(Refusal { + err: Error::Unroutable, + standing: false, + }), + }); + continue; + }; + if let Some(source) = source { + let id = self.next_source; + self.next_source += 1; + self.sources.insert(id, source); + self.instances.insert(id, instance); + self.events.push_back(Event::Resolved { route, result: Ok(id) }); + continue; + } + let Some(server) = server else { + self.events.push_back(Event::Resolved { + route, + result: Err(Refusal { + err: Error::Unroutable, + standing: self.standing(route), + }), + }); + continue; + }; + let mut serve = server.lock(); + if serve.closed { + // Retracted under us, or its handler dropped while the + // announcement stands: it cannot serve. A retraction + // leaves the table before it closes the server, so the + // table tells the two apart. + drop(serve); + self.events.push_back(Event::Resolved { + route, + result: Err(Refusal { + err: Error::Unroutable, + standing: self.standing(route), + }), + }); + continue; + } + // A source this route already materialized for the path + // attaches without another upstream round trip. Its session + // retires a source nothing holds, atomically with this mint, so + // one found closed is gone and the route is asked afresh. + if let Some(source) = serve + .served + .get(&self.path) + .map(|weak| weak.consume()) + .filter(|source| !source.is_closed()) + { + drop(serve); + let id = self.next_source; + self.next_source += 1; + self.sources.insert(id, source); + self.instances.insert(id, instance); + self.events.push_back(Event::Resolved { route, result: Ok(id) }); + continue; + } + let pending = match serve.requests.join(&self.path) { + Some(producer) => producer.consume(), + None => { + let producer = kio::Producer::::default(); + let consumer = producer.consume(); + match serve.requests.insert(self.path.clone(), producer) { + Ok(()) => consumer, + // No live handler behind the route: it cannot + // serve, whatever the table says. + Err(_) => { + drop(serve); + self.events.push_back(Event::Resolved { + route, + result: Err(Refusal { + err: Error::Unroutable, + standing: self.standing(route), + }), + }); + continue; + } + } + } + }; + self.upstream = Some(FrontUpstream { route, - result: Err(Refusal { - err: Error::Unroutable, - standing: standing(&front, route, &resolved), - }), + asked: instance.0, + asked_at: instance.1, + pending, + renewals: serve.renewals.clone(), + asked_renewals: serve.renewals.load(Ordering::Acquire), }); - continue; } - // A source this route already materialized for the path - // attaches without another upstream round trip. Its session - // retires a source nothing holds, atomically with this mint, so - // one found closed is gone and the route is asked afresh. - if let Some(source) = serve - .served - .get(&path) - .map(|weak| weak.consume()) - .filter(|source| !source.is_closed()) - { - drop(serve); - let id = next_source; - next_source += 1; - sources.insert(id, source); - instances.insert(id, instance); - events.push_back(Event::Resolved { route, result: Ok(id) }); - continue; - } - let pending = match serve.requests.join(&path) { - Some(producer) => producer.consume(), - None => { - let producer = kio::Producer::::default(); - let consumer = producer.consume(); - match serve.requests.insert(path.clone(), producer) { - Ok(()) => consumer, - // No live handler behind the route: it cannot - // serve, whatever the table says. - Err(_) => { - drop(serve); - events.push_back(Event::Resolved { - route, - result: Err(Refusal { - err: Error::Unroutable, - standing: standing(&front, route, &resolved), - }), - }); - continue; - } + Action::Detach { source, drain } => { + self.sources.remove(&source); + self.instances.remove(&source); + for io in self.tracks.values_mut() { + if !drain && io.query.as_ref().is_some_and(|asked| asked.source == source) { + io.query = None; + } + if !drain && io.staged.as_ref().is_some_and(|asked| asked.source == source) { + io.staged = None; } + // Any source leaving is rare: drop every parked shortcut rather + // than track which source each came from. + io.routes.unpark(); } - }; - upstream = Some(Upstream { - route, - asked: instance.0, - asked_at: instance.1, - pending, - renewals: serve.renewals.clone(), - asked_renewals: serve.renewals.load(Ordering::Acquire), - }); - } - Action::Detach { source, drain } => { - sources.remove(&source); - instances.remove(&source); - for io in tracks.values_mut() { - if !drain && io.query.as_ref().is_some_and(|asked| asked.source == source) { - io.query = None; + } + Action::Resolve => { + // The broadcast is whatever its first source serves, from the instance + // its route had when asked. + let (Some(source), Some(route)) = (self.front.serving(), self.front.serving_route()) else { + continue; + }; + let Some((instance, asked)) = self.instances.get(&source).cloned() else { + continue; + }; + // Route datagram lookups before a requester can see the broadcast, + // which would otherwise find no source to ask. + if self.datagram_source != Some(source) { + self.datagram_source = Some(source); + self.broadcast.serve_datagrams_from(self.sources.get(&source).cloned()); } - if !drain && io.staged.as_ref().is_some_and(|asked| asked.source == source) { - io.staged = None; + self.resolved = Some(instance.clone()); + if let Ok(mut pending) = self.request.write() + && pending.resolved.is_none() + { + pending.first = Some(instance.clone()); + pending.resolved = Some(Ok(instance)); + pending.via = Some((route, asked)); } - // Any source leaving is rare: drop every parked shortcut rather - // than track which source each came from. - io.routes.unpark(); - } - } - Action::Resolve => { - // The broadcast is whatever its first source serves, from the instance - // its route had when asked. - let (Some(source), Some(route)) = (front.serving(), front.serving_route()) else { - continue; - }; - let Some((instance, asked)) = instances.get(&source).cloned() else { - continue; - }; - // Route datagram lookups before a requester can see the broadcast, - // which would otherwise find no source to ask. - if datagram_source != Some(source) { - datagram_source = Some(source); - broadcast.serve_datagrams_from(sources.get(&source).cloned()); - } - resolved = Some(instance.clone()); - if let Ok(mut pending) = request.write() - && pending.resolved.is_none() - { - pending.first = Some(instance.clone()); - pending.resolved = Some(Ok(instance)); - pending.via = Some((route, asked)); } - } - Action::Query { track: name, source } => { - let Some(io) = tracks.get_mut(&name) else { continue }; - io.ready.wake_by_ref(); - let closing = sources.get(&source).is_some_and(|s| s.is_closing()); - match sources.get(&source).map(|s| s.track(&name)) { - Some(Ok(lookup)) => { - // Subscribe before the query exposes upstream demand, and query - // that same instance even if the broadcast changes meanwhile. - let sub = io.weak.subscription().map(|(demand, subscribers)| { - lookup.subscribe(io.routes.subscription(demand, subscribers)) - }); - let query = match &sub { - Some(sub) => sub.info(), - None => lookup.info().into_inner(), - }; - let copy = query.track(); - io.query = Some(Asked { + Action::Query { track: name, source } => { + let Some(io) = self.tracks.get_mut(&name) else { continue }; + io.ready.wake_by_ref(); + let closing = self.sources.get(&source).is_some_and(|s| s.is_closing()); + match self.sources.get(&source).map(|s| s.track(&name)) { + Some(Ok(lookup)) => { + // Subscribe before the query exposes upstream demand, and query + // that same instance even if the broadcast changes meanwhile. + let sub = io.weak.subscription().map(|(demand, subscribers)| { + lookup.subscribe(io.routes.subscription(demand, subscribers)) + }); + let query = match &sub { + Some(sub) => sub.info(), + None => lookup.info().into_inner(), + }; + let copy = query.track(); + io.query = Some(Asked { + source, + staged: io.routes.stage(copy.clone()), + copy, + query, + info: None, + sub, + }); + } + Some(Err(err)) => self.events.push_back(Event::TrackInfo { + track: name, source, - staged: io.routes.stage(copy.clone()), - copy, - query, - info: None, - sub, - }); + closing, + result: Err(err), + }), + None => {} } - Some(Err(err)) => events.push_back(Event::TrackInfo { - track: name, - source, - closing, - result: Err(err), - }), - None => {} } - } - Action::Splice { track: name, source } => { - let Some(io) = tracks.get_mut(&name) else { continue }; - let Some(asked) = io.staged.take() else { - continue; - }; - if asked.source != source { - continue; - } - io.splice(asked); - } - Action::Park { track: name } => { - let Some(io) = tracks.get_mut(&name) else { continue }; - io.ready.wake_by_ref(); - // Drop the copy so its source goes idle at once. A returning - // reader gets a fresh one, so nothing cached can be stale. - io.routes.park(); - io.copy = None; - io.held = None; - } - Action::Forget { track: name } => { - // A reader that looked the track up since the machine decided keeps - // it. Feed its `Wanted` edge here: the demand poll only sees the - // current level, so a reader gone before the next poll would - // otherwise leave the track unread with no linger armed. - if let Some(io) = tracks.get_mut(&name) - && !io.weak.abort_unwanted(Error::Cancel) - { - if !io.wanted { - io.wanted = true; - events.push_back(Event::Wanted { track: name }); + Action::Splice { track: name, source } => { + let Some(io) = self.tracks.get_mut(&name) else { continue }; + let Some(asked) = io.staged.take() else { + continue; + }; + if asked.source != source { + continue; } - continue; + io.splice(asked); } - // Nobody reads it, so nobody follows its copy: let the route go idle. - if let Some(io) = tracks.remove(&name) { + Action::Park { track: name } => { + let Some(io) = self.tracks.get_mut(&name) else { continue }; + io.ready.wake_by_ref(); + // Drop the copy so its source goes idle at once. A returning + // reader gets a fresh one, so nothing cached can be stale. io.routes.park(); + io.copy = None; + io.held = None; } - events.push_back(Event::Forgotten { track: name }); - } - Action::Finish { track: name } => { - if let Some(io) = tracks.get_mut(&name) { - io.end(Ok(())); + Action::Forget { track: name } => { + // A reader that looked the track up since the machine decided keeps + // it. Feed its `Wanted` edge here: the demand poll only sees the + // current level, so a reader gone before the next poll would + // otherwise leave the track unread with no linger armed. + if let Some(io) = self.tracks.get_mut(&name) + && !io.weak.abort_unwanted(Error::Cancel) + { + if !io.wanted { + io.wanted = true; + self.events.push_back(Event::Wanted { track: name }); + } + continue; + } + // Nobody reads it, so nobody follows its copy: let the route go idle. + if let Some(io) = self.tracks.remove(&name) { + io.routes.park(); + } + self.events.push_back(Event::Forgotten { track: name }); } - } - Action::Abort { track: name, err } => { - if let Some(io) = tracks.get_mut(&name) { - tracing::debug!(name = %name, %err, "aborting track"); - io.end(Err(err)); + Action::Finish { track: name } => { + if let Some(io) = self.tracks.get_mut(&name) { + io.end(Ok(())); + } } - } - Action::Arm { at } => deadline.set(at), - Action::Retire => { - // Closing refuses every holder from here on, atomically with minting - // one: a request that joined first keeps the front, and one after - // finds it closed and mints a fresh one. - let event = match broadcast.close_unheld() { - true => Event::Retired, - false => { - held = true; - Event::Held + Action::Abort { track: name, err } => { + if let Some(io) = self.tracks.get_mut(&name) { + tracing::debug!(name = %name, %err, "aborting track"); + io.end(Err(err)); } - }; - events.push_back(event); - } - Action::End { err } => { - // Leave the table before anything shows the end, whatever ended the - // front: a requester seeing it re-requests, and must mint a fresh - // front rather than join this one again. A successor already in the - // slot stays. - shared - .lock() - .fronts - .remove_if(&(path.clone(), horizon, pin.clone()), |front| { - front.request.same_channel(&request) - }); - if let Ok(mut pending) = request.write() { - pending.refused = refused; - pending.resolved.get_or_insert(Err(err.clone())); } - // Ending the broadcast only retracts it: no new requesters or - // tracks, and a newcomer at the path gets a fresh front. Tracks - // in flight carry on (moq-lite: retraction does not disturb - // subscriptions already in flight): their readers follow the copy - // they read to its end, since no front is left to replace it. - broadcast.close(); - let mut in_flight = Vec::new(); - for (_, mut io) in tracks.drain() { - let wanted = io.weak.is_wanted(); - // A reader still waiting on its source's answer is in flight - // too: serve it the copy it asked, so it ends as that copy does. - let waiting = io.staged.take().or_else(|| io.query.take()); - if io.copy.is_none() - && wanted && let Some(asked) = waiting - { - io.splice(asked); + Action::Arm { at } => self.deadline.set(at), + Action::Retire => { + // Closing refuses every holder from here on, atomically with minting + // one: a request that joined first keeps the front, and one after + // finds it closed and mints a fresh one. + let event = match self.broadcast.close_unheld() { + true => Event::Retired, + false => { + self.held = true; + Event::Held + } + }; + self.events.push_back(event); + } + Action::End { err } => { + // Leave the table before anything shows the end, whatever ended the + // front: a requester seeing it re-requests, and must mint a fresh + // front rather than join this one again. A successor already in the + // slot stays. + self.shared + .lock() + .fronts + .remove_if(&(self.path.clone(), self.horizon, self.pin.clone()), |front| { + front.request.same_channel(&self.request) + }); + if let Ok(mut pending) = self.request.write() { + pending.refused = refused; + pending.resolved.get_or_insert(Err(err.clone())); } - // Nothing in flight: unread, or nothing serving it. - if !wanted || io.copy.is_none() { - io.end(Err(err.clone())); - // An unread track the front has not parked yet still holds - // its copy, which would keep its source subscribed. - if !wanted { - io.routes.release(); + // Ending the broadcast only retracts it: no new requesters or + // tracks, and a newcomer at the path gets a fresh front. Tracks + // in flight carry on (moq-lite: retraction does not disturb + // subscriptions already in flight): their readers follow the copy + // they read to its end, since no front is left to replace it. + self.broadcast.close(); + let mut in_flight = Vec::new(); + for (_, mut io) in self.tracks.drain() { + let wanted = io.weak.is_wanted(); + // A reader still waiting on its source's answer is in flight + // too: serve it the copy it asked, so it ends as that copy does. + let waiting = io.staged.take().or_else(|| io.query.take()); + if io.copy.is_none() + && wanted && let Some(asked) = waiting + { + io.splice(asked); } - continue; + // Nothing in flight: unread, or nothing serving it. + if !wanted || io.copy.is_none() { + io.end(Err(err.clone())); + // An unread track the front has not parked yet still holds + // its copy, which would keep its source subscribed. + if !wanted { + io.routes.release(); + } + continue; + } + io.routes.conclude(); + let TrackIo { + weak, + routes, + request, + accepted, + copy, + held, + .. + } = io; + in_flight.push(InFlight { + weak, + routes, + end: err.clone(), + request, + accepted, + query: copy.map(|(_, _, query)| query), + held, + }); } - io.routes.conclude(); - let TrackIo { - weak, - routes, - request, - accepted, - copy, - held, - .. - } = io; - in_flight.push(InFlight { - weak, - routes, - end: err.clone(), - request, - accepted, - query: copy.map(|(_, _, query)| query), - held, - }); + return Poll::Ready(in_flight); } - return in_flight; } } + // A staged copy the machine did not splice was refused or is no longer wanted: + // let it go, or its source stays subscribed for nobody. + if let Some((name, source)) = answered + && let Some(io) = self.tracks.get_mut(&name) + && io.staged.as_ref().is_some_and(|asked| asked.source == source) + { + io.staged = None; + } } - // A staged copy the machine did not splice was refused or is no longer wanted: - // let it go, or its source stays subscribed for nobody. - if let Some((name, source)) = answered - && let Some(io) = tracks.get_mut(&name) - && io.staged.as_ref().is_some_and(|asked| asked.source == source) - { - io.staged = None; - } - } - if front.serving() != datagram_source { - datagram_source = front.serving(); - broadcast.serve_datagrams_from(datagram_source.and_then(|id| sources.get(&id)).cloned()); - } - - let step = kio::wait(|waiter| { - if let Poll::Ready(Ok(request)) = dynamic.poll_requested_track(waiter) { - return Poll::Ready(Step::Assigned(request)); - } - if let Some(upstream) = &upstream - && let Poll::Ready(result) = upstream.pending.poll(waiter, |p| match &p.resolved { - Some(result) => Poll::Ready(result.clone()), - None => Poll::Pending, - }) { - return Poll::Ready(Step::Resolved(match result { - Ok(resolved) => resolved, - // The queue died unresolved (its handler dropped): the route - // could not serve. - Err(_closed) => Err(Error::Unroutable), - })); + if self.front.serving() != self.datagram_source { + self.datagram_source = self.front.serving(); + self.broadcast + .serve_datagrams_from(self.datagram_source.and_then(|id| self.sources.get(&id)).cloned()); } - // The listing first: a source ends its listing before it closes, so this forwards - // the tracks it finished last before its close detaches it. - match &mut listing { - None => { - if front.serving().is_some() && matches!(lister.poll_wanted(waiter), Poll::Ready(Ok(()))) { - return Poll::Ready(Step::Listing); + + let mut park = std::mem::take(&mut self.park); + let Poll::Ready(step) = self.poll_step(park.hold(&waiter.context())) else { + self.park = park; + return Poll::Pending; + }; + // Before the step runs, so nothing it touches wakes the front for a wait that + // is over. + drop(park); + + let event = match step { + FrontStep::Assigned(request) => { + let name: Arc = request.name().into(); + let track_ready = crate::util::QueueWake::new(name.clone(), self.ready.clone()); + track_ready.wake_by_ref(); + let routes = super::resume::Producer::new(); + self.tracks.insert( + name.clone(), + TrackIo { + ready: track_ready, + park: Default::default(), + weak: request.weak(), + request: Some(request.routes(routes.consume())), + accepted: None, + routes, + query: None, + staged: None, + copy: None, + held: None, + wanted: false, + }, + ); + Event::TrackAssigned { + track: name, + now: self.timers.now(), } } - Some(state) => { - if lister.poll_unwanted(waiter).is_ready() { - return Poll::Ready(Step::Unlisted); - } - if state.source != front.serving() { - return Poll::Ready(Step::Listing); - } - if let Some(reader) = &mut state.reader - && let Poll::Ready(result) = reader.poll_next(waiter) - { - return Poll::Ready(Step::Listed(result)); + FrontStep::Resolved(result) => { + let Some(FrontUpstream { + route, + asked, + asked_at, + renewals, + asked_renewals, + .. + }) = self.upstream.take() + else { + continue; + }; + // Counted rather than read off the request, which an answer just before + // the change already took out of the queue, and rather than compared by + // instance, which may have come back since. + let released = renewals.load(Ordering::Acquire) != asked_renewals; + match result { + // An instance change released the request, the route moved to another + // instance, or another won the path, while the request was in flight: + // ask again rather than resolve requesters onto a broadcast that is + // already replaced. + _ if released => Event::Resolved { + route, + result: Err(Refusal { + err: Error::Unroutable, + standing: false, + }), + }, + Ok(_) if self.stale(route, &asked) => Event::Resolved { + route, + result: Err(Refusal { + err: Error::Unroutable, + standing: false, + }), + }, + Ok(source) => { + let id = self.next_source; + self.next_source += 1; + self.sources.insert(id, source); + self.instances.insert(id, (asked, asked_at)); + Event::Resolved { route, result: Ok(id) } + } + Err(err) => Event::Resolved { + route, + result: Err(Refusal { + err, + standing: self.standing(route), + }), + }, } } - } - if let Some(id) = front.serving() - && let Some(source) = sources.get(&id) - && source.poll_closed(waiter).is_ready() - { - return Poll::Ready(Step::SourceClosed(id)); - } - while let Poll::Ready(Ok(name)) = ready.poll_pop(waiter) { - let Some(io) = tracks.get_mut(&name) else { continue }; - io.ready.popped(); - let waker = Waker::from(io.ready.clone()); - let waiter = io.park.hold(&Context::from_waker(&waker)); - let step = (|| { - if let Some(asked) = &mut io.query - && let Poll::Ready(result) = asked.query.poll(waiter) - { - return Some(Step::Info(name.clone(), asked.source, result)); - } - follow_demand(&mut io.held, io.accepted.as_mut(), &io.weak, &io.routes, waiter); - // Settled once the copy closes: a group still open below a declared end is - // owed until then, and a session closing first means it never came. - if let Some((source, copy, _)) = &io.copy - && copy.poll_closed(waiter).is_ready() + FrontStep::SourceClosed(source) => Event::SourceClosed { source }, + FrontStep::Info(name, source, result) => { + let closing = self.sources.get(&source).is_some_and(|s| s.is_closing()); + let Some(io) = self.tracks.get_mut(&name) else { continue }; + let Some(asked) = io.query.take() else { continue }; + // An aborted copy with history still splices, so its readers drain that + // history before the outcome, as local readers do. One that delivered + // nothing is a refusal, which leaves a draining incumbent serving. + let result = match result { + Ok(info) => match asked.copy.poll_complete(&kio::Waiter::noop()) { + Poll::Ready(Err(err)) if asked.copy.latest_sequence().is_none() => Err(err), + _ => Ok(info), + }, + Err(err) => Err(err), + }; + // Staged only while the track has a reader: without one the machine + // will not splice, and a held copy would keep the source subscribed. + if let Ok(info) = &result + && io.wanted { - let result = match copy.poll_complete(&kio::Waiter::noop()) { - Poll::Ready(result) => result, - Poll::Pending => { - // A channel cannot reopen after poll_closed reported its end. - debug_assert!(false, "closed copy did not report completion"); - Err(Error::Cancel) - } - }; - let delivered = copy.latest_sequence().is_some(); - return Some(Step::Ended(name.clone(), *source, result, delivered)); + io.staged = Some(Asked { + info: Some(info.clone()), + ..asked + }); } - // Watch the demand edge in whichever direction is unmet. Only `Pending` - // parks, so an answer steps and the handler looks again, polling anew when - // nothing changed. A closed track is unread: it owes one `Unwanted` if - // recorded read, and is no edge after. - let edge = match io.wanted { - true => io.weak.poll_unwanted(waiter).is_ready(), - false => matches!(io.weak.poll_wanted(waiter), Poll::Ready(Ok(()))), - }; - if edge { - return Some(Step::Demand(name.clone())); + Event::TrackInfo { + track: name, + source, + closing, + result, } - None - })(); - if let Some(step) = step { - // A ready answer did not park on its next edge. Revisit this track - // once the machine has consumed the event and applied its actions. - io.ready.wake_by_ref(); - return Poll::Ready(step); } - } - // The same for whoever holds the broadcast itself. - let edge = match held { - true => broadcast.poll_unheld(waiter), - false => broadcast.poll_held(waiter), - }; - match edge { - Poll::Ready(Ok(())) => return Poll::Ready(Step::Holders), - // Answers at once from here on, so it must end the front, not be polled again. - Poll::Ready(Err(_)) => return Poll::Ready(Step::Closed), - Poll::Pending => {} - } - if deadline.poll(waiter).is_ready() { - return Poll::Ready(Step::Deadline); - } - watch.poll_changed(waiter, seen).map(|()| Step::Table) - }) - .await; - - let event = match step { - Step::Assigned(request) => { - let name: Arc = request.name().into(); - let track_ready = crate::util::QueueWake::new(name.clone(), ready.clone()); - track_ready.wake_by_ref(); - let routes = super::resume::Producer::new(); - tracks.insert( - name.clone(), - TrackIo { - ready: track_ready, - park: Default::default(), - weak: request.weak(), - request: Some(request.routes(routes.consume())), - accepted: None, - routes, - query: None, - staged: None, - copy: None, - held: None, - wanted: false, - }, - ); - Event::TrackAssigned { - track: name, - now: timers.now(), - } - } - Step::Resolved(result) => { - let Some(Upstream { - route, - asked, - asked_at, - renewals, - asked_renewals, - .. - }) = upstream.take() - else { - continue; - }; - // Counted rather than read off the request, which an answer just before - // the change already took out of the queue, and rather than compared by - // instance, which may have come back since. - let released = renewals.load(Ordering::Acquire) != asked_renewals; - match result { - // An instance change released the request, the route moved to another - // instance, or another won the path, while the request was in flight: - // ask again rather than resolve requesters onto a broadcast that is - // already replaced. - _ if released => Event::Resolved { - route, - result: Err(Refusal { - err: Error::Unroutable, - standing: false, - }), - }, - Ok(_) if stale(&front, route, &asked, &resolved) => Event::Resolved { - route, - result: Err(Refusal { - err: Error::Unroutable, - standing: false, - }), - }, - Ok(source) => { - let id = next_source; - next_source += 1; - sources.insert(id, source); - instances.insert(id, (asked, asked_at)); - Event::Resolved { route, result: Ok(id) } + FrontStep::Ended(name, source, result, delivered) => { + let closing = self.sources.get(&source).is_some_and(|s| s.is_closing()); + let Some(io) = self.tracks.get_mut(&name) else { continue }; + io.copy = None; + io.held = None; + Event::TrackEnded { + track: name, + source, + closing, + result, + delivered, } - Err(err) => Event::Resolved { - route, - result: Err(Refusal { - err, - standing: standing(&front, route, &resolved), - }), - }, - } - } - Step::SourceClosed(source) => Event::SourceClosed { source }, - Step::Info(name, source, result) => { - let closing = sources.get(&source).is_some_and(|s| s.is_closing()); - let Some(io) = tracks.get_mut(&name) else { continue }; - let Some(asked) = io.query.take() else { continue }; - // An aborted copy with history still splices, so its readers drain that - // history before the outcome, as local readers do. One that delivered - // nothing is a refusal, which leaves a draining incumbent serving. - let result = match result { - Ok(info) => match asked.copy.poll_complete(&kio::Waiter::noop()) { - Poll::Ready(Err(err)) if asked.copy.latest_sequence().is_none() => Err(err), - _ => Ok(info), - }, - Err(err) => Err(err), - }; - // Staged only while the track has a reader: without one the machine - // will not splice, and a held copy would keep the source subscribed. - if let Ok(info) = &result - && io.wanted - { - io.staged = Some(Asked { - info: Some(info.clone()), - ..asked - }); - } - Event::TrackInfo { - track: name, - source, - closing, - result, } - } - Step::Ended(name, source, result, delivered) => { - let closing = sources.get(&source).is_some_and(|s| s.is_closing()); - let Some(io) = tracks.get_mut(&name) else { continue }; - io.copy = None; - io.held = None; - Event::TrackEnded { - track: name, - source, - closing, - result, - delivered, - } - } - Step::Demand(name) => { - let Some(io) = tracks.get_mut(&name) else { continue }; - // Read again: the edge may have flipped back since it fired, and the next - // wait parks on it. - if io.weak.is_wanted() == io.wanted { - continue; + FrontStep::Demand(name) => { + let Some(io) = self.tracks.get_mut(&name) else { continue }; + // Read again: the edge may have flipped back since it fired, and the next + // wait parks on it. + if io.weak.is_wanted() == io.wanted { + continue; + } + io.wanted = !io.wanted; + if !io.wanted { + // Nothing will be spliced now: let go of the copies a query + // holds, or the source stays subscribed with nobody reading. + io.query = None; + io.staged = None; + } + match io.wanted { + true => Event::Wanted { track: name }, + false => Event::Unwanted { + track: name, + now: self.timers.now(), + }, + } } - io.wanted = !io.wanted; - if !io.wanted { - // Nothing will be spliced now: let go of the copies a query - // holds, or the source stays subscribed with nobody reading. - io.query = None; - io.staged = None; + FrontStep::Holders => { + // Read again: the edge may have flipped back since it fired. + if self.broadcast.is_held() == self.held { + continue; + } + self.held = !self.held; + match self.held { + true => Event::Held, + false => Event::Unheld, + } } - match io.wanted { - true => Event::Wanted { track: name }, - false => Event::Unwanted { - track: name, - now: timers.now(), - }, + FrontStep::Closed => Event::Retired, + FrontStep::Deadline => { + // Cleared here so a fired deadline cannot keep firing; the machine + // re-arms what is still parked. + self.deadline.set(None); + Event::Deadline { now: self.timers.now() } } - } - Step::Holders => { - // Read again: the edge may have flipped back since it fired. - if broadcast.is_held() == held { + FrontStep::Table => self.select(), + FrontStep::Listing => { + // A route claiming a prefix names no broadcast, so a path it serves on request + // has no inventory to list. + let covering = self.front.serving_route().is_some_and(|route| { + self.shared + .read() + .routes + .covering(&self.path) + .find(|entry| entry.id == route) + .is_some_and(|entry| entry.prefix != self.path) + }); + if covering { + self.listing = None; + self.lister.abort(Error::NotFound); + continue; + } + let state = self.listing.get_or_insert_default(); + state.source = self.front.serving(); + state.stale = state.listed.clone(); + // A closing source lists nothing more; the front moves off it soon. + state.reader = state + .source + .and_then(|id| self.sources.get(&id)) + .and_then(|source| source.tracks().ok()); continue; } - held = !held; - match held { - true => Event::Held, - false => Event::Unheld, - } - } - Step::Closed => Event::Retired, - Step::Deadline => { - // Cleared here so a fired deadline cannot keep firing; the machine - // re-arms what is still parked. - deadline.set(None); - Event::Deadline { now: timers.now() } - } - Step::Table => select( - &mut front, - &sources, - &mut instances, - &mut seen, - &mut resolved, - &mut discovered, - ), - Step::Listing => { - // A route claiming a prefix names no broadcast, so a path it serves on request - // has no inventory to list. - let covering = front.serving_route().is_some_and(|route| { - shared - .read() - .routes - .covering(&path) - .find(|entry| entry.id == route) - .is_some_and(|entry| entry.prefix != path) - }); - if covering { - listing = None; - lister.abort(Error::NotFound); + FrontStep::Unlisted => { + self.listing = None; + self.lister.clear(); continue; } - let state = listing.get_or_insert_default(); - state.source = front.serving(); - state.stale = state.listed.clone(); - // A closing source lists nothing more; the front moves off it soon. - state.reader = state - .source - .and_then(|id| sources.get(&id)) - .and_then(|source| source.tracks().ok()); - continue; - } - Step::Unlisted => { - listing = None; - lister.clear(); - continue; - } - Step::Listed(result) => { - let Some(state) = &mut listing else { continue }; - // The lister refuses only once the broadcast ended, which ends the front too. - match result { - Ok(Some(broadcast::TracksEvent::Added { name, info })) => { - if !state.stale.remove(&name) && state.listed.insert(name.clone()) { - let _ = lister.add(name, info); + FrontStep::Listed(result) => { + let Some(state) = &mut self.listing else { continue }; + // The lister refuses only once the broadcast ended, which ends the front too. + match result { + Ok(Some(broadcast::TracksEvent::Added { name, info })) => { + if !state.stale.remove(&name) && state.listed.insert(name.clone()) { + let _ = self.lister.add(name, info); + } } - } - // Forwarded as the source reported it, finished or not. - Ok(Some(broadcast::TracksEvent::Removed { name })) => { - state.stale.remove(&name); - if state.listed.remove(&name) { - let _ = lister.remove(&name); + // Forwarded as the source reported it, finished or not. + Ok(Some(broadcast::TracksEvent::Removed { name })) => { + state.stale.remove(&name); + if state.listed.remove(&name) { + let _ = self.lister.remove(&name); + } } - } - Ok(Some(broadcast::TracksEvent::Finished { name })) => { - state.stale.remove(&name); - if state.listed.remove(&name) { - let _ = lister.finish(&name); + Ok(Some(broadcast::TracksEvent::Finished { name })) => { + state.stale.remove(&name); + if state.listed.remove(&name) { + let _ = self.lister.finish(&name); + } } - } - Ok(Some(broadcast::TracksEvent::Synced)) => { - for name in state.stale.drain() { - state.listed.remove(&name); - let _ = lister.remove(&name); + Ok(Some(broadcast::TracksEvent::Synced)) => { + for name in state.stale.drain() { + state.listed.remove(&name); + let _ = self.lister.remove(&name); + } + self.lister.synced(); + } + // The source ended: the front moves to another source, or ends with it. + Ok(None) + if state + .source + .and_then(|id| self.sources.get(&id)) + .is_none_or(|source| source.is_closing()) => + { + state.reader = None + } + // The source's handler ended the listing on a live source, as a peer does + // once its broadcast ended: so does the front's, and a later reader asks again. + Ok(None) => { + self.listing = None; + self.lister.end(); + } + Err(err) => { + // Its readers fail, and a later reader asks again. + self.listing = None; + self.lister.abort(err); } - lister.synced(); - } - // The source ended: the front moves to another source, or ends with it. - Ok(None) - if state - .source - .and_then(|id| sources.get(&id)) - .is_none_or(|source| source.is_closing()) => - { - state.reader = None - } - // The source's handler ended the listing on a live source, as a peer does - // once its broadcast ended: so does the front's, and a later reader asks again. - Ok(None) => { - listing = None; - lister.end(); - } - Err(err) => { - // Its readers fail, and a later reader asks again. - listing = None; - lister.abort(err); } + continue; } - continue; - } - }; - events.push_back(event); + }; + self.events.push_back(event); + } } } @@ -5857,10 +5954,13 @@ impl Requests { self.announcement.update(route) } - /// Poll for the next requested path under this route, without blocking. + /// The next path a consumer requests under this route, as a [`Request`] to serve. /// - /// Returns [`Error::Closed`] once the origin's [`Driver`] has been dropped: - /// no request will ever arrive again, so handler loops should end. + /// Takes `&self` so a handler can serve from one task while another re-prices + /// the route; concurrent callers each receive distinct requests. Returns + /// [`Error::Closed`] once the origin's [`Driver`] has been dropped: no request will + /// ever arrive again, so handler loops should end. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_requested_broadcast(&self, waiter: &kio::Waiter) -> Poll> { let mut state = ready!(self.state.poll(waiter, |state| { if state.closed || state.requests.has_queued() { @@ -5889,16 +5989,6 @@ impl Requests { publisher: self.publisher.clone(), })) } - - /// Block until a consumer requests a path under this route, returning a - /// [`Request`] to serve. - /// - /// Takes `&self` so a handler can serve from one task while another re-prices - /// the route; concurrent callers each receive distinct requests. - #[cfg(feature = "async")] - pub async fn requested_broadcast(&self) -> Result { - kio::wait(|waiter| self.poll_requested_broadcast(waiter)).await - } } impl ServeState { @@ -6901,7 +6991,7 @@ impl Consumer { /// Follow the broadcast at `path` across the instances that serve it. /// - /// The first [`available`](Following::available) resolves the instance serving the path + /// The first [`poll_available`](Following::poll_available) resolves the instance serving the path /// once one is announced, retrying route churn but returning a selected handler's refusal; /// the call after a refusal waits for the serving route to change. Each later call returns /// another instance taking over, or `None` once the path goes offline, then whatever serves @@ -7142,7 +7232,7 @@ impl Consumer { // Released before the push: a set whose handles are gone drops the task, // and the `Watch` it carries unregisters under this same lock. drop(state); - self.tasks.push(run_front( + self.tasks.push(FrontServeRun::new( FrontTask { shared: self.shared.clone(), broadcast, @@ -7415,23 +7505,13 @@ impl Routes { event } - /// Returns the next route announcement, update, or retraction, its prefix - /// relative to this cursor's root. + /// The next route announcement, update, or retraction, its prefix relative to + /// this cursor's root, or `None` once the cursor is closed. /// /// A retraction is only delivered for a previously announced prefix, and a - /// repeated announcement for the same prefix is a metadata update. Returns - /// None if the cursor is closed. The consumer is also a [`futures::Stream`] - /// of the same events. - #[cfg(feature = "async")] - pub async fn next(&mut self) -> Option { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next event, without blocking. - /// - /// Returns `Poll::Ready(Some(_))` for an event, `Poll::Ready(None)` if the - /// cursor is closed, or `Poll::Pending` after registering `waiter` to be - /// notified when the next event arrives. + /// repeated announcement for the same prefix is a metadata update. With the + /// `async` feature, the consumer is also a `futures::Stream` of the same events. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll> { loop { let now = self.timers.now(); @@ -10178,8 +10258,8 @@ mod tests { }), timers: consumer.timers.clone(), }; - let mut front = std::pin::pin!(serve_front(task)); - assert!(front.as_mut().now_or_never().is_none(), "the front ended early"); + let mut front = kio::Pending::new(FrontServe::new(task)); + assert!((&mut front).now_or_never().is_none(), "the front ended early"); broadcast.close(); let in_flight = kio::clock::timeout(Duration::from_secs(1), front) diff --git a/rs/moq-net/src/model/pace.rs b/rs/moq-net/src/model/pace.rs index 2b0d9f6fcbd..db667bb7682 100644 --- a/rs/moq-net/src/model/pace.rs +++ b/rs/moq-net/src/model/pace.rs @@ -22,6 +22,7 @@ use std::sync::{ use std::task::Waker; use super::group::MAX_CACHE_BYTES; +#[cfg(any(test, feature = "async"))] use crate::transport::ReceiveWindow; /// The furthest a paced writer runs ahead of its group's fastest reader: half the live @@ -32,6 +33,7 @@ pub(crate) const AHEAD: u64 = MAX_CACHE_BYTES / 2; /// stream may hold under `receive`, the windows its transport grants, capped at [`AHEAD`]. /// /// A transport that reports no window never pauses a stream, so its window is moot. +#[cfg(any(test, feature = "async"))] pub(crate) fn window(receive: Option) -> u64 { receive.map_or(AHEAD, |receive| receive.stream.min(receive.connection).min(AHEAD)) } diff --git a/rs/moq-net/src/model/socket.rs b/rs/moq-net/src/model/socket.rs index 76df834ffaf..9c47c581f43 100644 --- a/rs/moq-net/src/model/socket.rs +++ b/rs/moq-net/src/model/socket.rs @@ -255,7 +255,7 @@ impl sealed::Open for TcpStream { local: Some(local), pending: None, }; - shared.tasks.push(kio::Pending::new(pump)); + shared.tasks.push(pump); Self { producer, consumer, @@ -695,33 +695,23 @@ impl Endpoint { })) } - /// Poll for the next `tcp` dial from a peer; fails with [`Error::Closed`] once the + /// The next `tcp` dial from a peer; fails with [`Error::Closed`] once the /// endpoint closed, and with [`Error::Unsupported`] while no peer can reach it (see /// [`Endpoint`]). /// /// The first call starts accepting `tcp` dials. They queue until taken, or until their /// dialer gives up, so keep accepting, and accept or reject each one. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_accept_tcp(&self, waiter: &kio::Waiter) -> Poll, Error>> { self.poll_accept(waiter) } - /// Wait for the next `tcp` dial from a peer; see [`Self::poll_accept_tcp`]. - #[cfg(feature = "async")] - pub async fn accept_tcp(&self) -> Result, Error> { - kio::wait(|waiter| self.poll_accept_tcp(waiter)).await - } - - /// Poll for the next `track` dial from a peer; see [`Self::poll_accept_tcp`]. + /// The next `track` dial from a peer; see [`Self::poll_accept_tcp`]. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_accept_track(&self, waiter: &kio::Waiter) -> Poll, Error>> { self.poll_accept(waiter) } - /// Wait for the next `track` dial from a peer; see [`Self::poll_accept_tcp`]. - #[cfg(feature = "async")] - pub async fn accept_track(&self) -> Result, Error> { - kio::wait(|waiter| self.poll_accept_track(waiter)).await - } - /// Close the endpoint; see [`Endpoint`]. pub fn close(&self) { self.shared.close(); diff --git a/rs/moq-net/src/model/track.rs b/rs/moq-net/src/model/track.rs index c7689bf2035..782f4d194d4 100644 --- a/rs/moq-net/src/model/track.rs +++ b/rs/moq-net/src/model/track.rs @@ -3778,17 +3778,10 @@ impl Producer { Ok(()) } - /// Block until the track's outcome is final: `Ok` once its declared end is reached (its - /// error, if declared with one), or the abort error, or [`Error::Dropped`] once every - /// producer left without either. - #[cfg(feature = "async")] - pub async fn closed(&self) -> Result<()> { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - - /// Poll until the track's outcome is final: `Ok` once its declared end is reached (its - /// error, if declared with one), or the abort error, or [`Error::Dropped`] once every + /// Resolves once the track's outcome is final: `Ok` once its declared end is reached + /// (its error, if declared with one), or the abort error, or [`Error::Dropped`] once every /// producer left without either. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll> { TrackState::settle(self.state.poll_ref(waiter, |state| state.poll_outcome(waiter))) } @@ -3910,16 +3903,6 @@ impl Producer { } } - /// Block until the aggregate subscription changes, then return the new value. - /// - /// Yields the most demanding request across all live subscribers, or `None` - /// once the last one drops. Used by relays to forward downstream demand - /// upstream (e.g. SUBSCRIBE_UPDATE). - #[cfg(feature = "async")] - pub async fn subscription_changed(&mut self) -> Result> { - kio::wait(|waiter| self.poll_subscription_changed(waiter)).await - } - /// A non-blocking snapshot of the current aggregate subscription, or `None` /// when there are no live subscribers. Unlike [`Self::subscription`], this /// doesn't wait for a change or advance the change cursor. @@ -3935,9 +3918,11 @@ impl Producer { snapshot_subscription(&subs, bound) } - /// Poll for a changed subscription: the - /// aggregate subscription whenever it changes, or `None` once nobody is subscribed. - /// Errors once the track is aborted. + /// The aggregate subscription once it changes: the most demanding request across all + /// live subscribers, or `None` once the last one drops. Errors once the track is aborted. + /// + /// Used by relays to forward downstream demand upstream (e.g. SUBSCRIBE_UPDATE). + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_subscription_changed(&mut self, waiter: &kio::Waiter) -> Poll>> { // Surface an abort as the stream ending. `poll_closed` parks on the closed // waiters, so per-group churn on the track state never wakes this poll. @@ -4018,23 +4003,19 @@ impl Requests { &self.name } - /// Block until readers miss a run of groups, returning a [`Fill`] to serve via - /// [`Fill::accept`] and [`Fill::drop_groups`]. + /// The next run of groups readers miss, as a [`Fill`] to serve via [`Fill::accept`] + /// and [`Fill::drop_groups`]. /// /// A relay fetches the run upstream; an origin already has its groups cached, so /// nothing reaches here. Errors once the track is aborted. - #[cfg(feature = "async")] - pub async fn requested_fill(&self) -> Result { - kio::wait(|waiter| self.poll_requested_fill(waiter)).await - } - - /// Poll for a fill to serve. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_requested_fill(&self, waiter: &kio::Waiter) -> Poll> { fill::poll_requested(&self.state, &self.fills, waiter) } /// The next fill, which a test expects to name a single group. #[cfg(test)] + #[cfg_attr(feature = "async", kio::wrapper)] pub(crate) fn poll_requested_fill_group(&self, waiter: &kio::Waiter) -> Poll> { self.poll_requested_fill(waiter).map_ok(|fill| { assert_eq!(fill.end() - fill.start().group, 1, "a one-group fill"); @@ -4042,12 +4023,6 @@ impl Requests { }) } - /// [`Self::poll_requested_fill_group`], awaited. - #[cfg(all(test, feature = "async"))] - pub(crate) async fn requested_fill_group(&self) -> Result { - kio::wait(|waiter| self.poll_requested_fill_group(waiter)).await - } - /// Watch subscriber demand without keeping the track alive. pub fn demand(&self) -> Demand { Demand::new(self.name.clone(), self.state.weak()) @@ -4431,23 +4406,18 @@ impl Demand { Some(wanted) } - /// Return [`wanted`](Self::wanted) once it differs from what this handle last returned: - /// immediately on the first call, then on each change. Errors once the track is closed - /// or aborted. - /// - /// So `while let Some(n) = demand.wanted_changed().await?` runs while anyone wants the - /// track, and `while demand.wanted_changed().await?.is_none()` waits for someone to. - #[cfg(feature = "async")] - pub async fn wanted_changed(&mut self) -> Result> { - kio::wait(|waiter| self.poll_wanted_changed(waiter)).await - } - /// The demand aggregate. fn subscriptions(&self) -> kio::Lock { self.state.read().subscriptions.clone() } - /// Poll for [`wanted`](Self::wanted) to differ from what this handle last returned. + /// [`wanted`](Self::wanted) once it differs from what this handle last returned: + /// immediately on the first call, then on each change. Errors once the track is closed + /// or aborted. + /// + /// So `while let Some(n) = demand.wanted_changed().await?` runs while anyone wants the + /// track, and `while demand.wanted_changed().await?.is_none()` waits for someone to. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_wanted_changed(&mut self, waiter: &kio::Waiter) -> Poll>> { let active = match self.poll_state(waiter) { DemandState::Closed => return Poll::Ready(Err(self.abort_reason())), @@ -4463,17 +4433,14 @@ impl Demand { Poll::Ready(Ok(next)) } - /// Block until the track's outcome is final; see [`Producer::closed`]. - #[cfg(feature = "async")] - pub async fn closed(&self) -> Result<()> { - kio::wait(|waiter| { - TrackState::settle( - self.state - .consume_weak() - .poll(waiter, |state| state.poll_outcome(waiter)), - ) - }) - .await + /// Resolves once the track's outcome is final; see [`Producer::poll_closed`]. + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll> { + TrackState::settle( + self.state + .consume_weak() + .poll(waiter, |state| state.poll_outcome(waiter)), + ) } /// The publisher's tie-break priority, as set in [`Info::priority`]. @@ -5240,7 +5207,9 @@ impl Consumer { }) } - /// Poll for the track's largest location: the current one, then each newer one, and `None` once the track ends. + /// The track's largest location, its newest group at the newest frame (0 while empty): + /// the current one, then each newer one, and `None` once the track ends. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_largest(&mut self, waiter: &kio::Waiter) -> Poll> { let mut seen = self.largest; let next = self.poll_largest_after(&mut seen, Bounds::default(), waiter); @@ -5279,12 +5248,6 @@ impl Consumer { }) } - /// The track's largest location, its newest group at the newest frame (0 while empty): the current one, then each newer one, and `None` once the track ends. - #[cfg(feature = "async")] - pub async fn largest(&mut self) -> Option { - kio::wait(|waiter| self.poll_largest(waiter)).await - } - pub(crate) fn latest_sequence(&self) -> Option { if let Some(serving) = self.serving() { return serving.latest_sequence(); @@ -8118,7 +8081,7 @@ impl Subscriber { } } - /// Poll for the next group in arrival order, without blocking. + /// The next group in arrival order. /// /// Returns each group it delivers exactly once, in the order it landed on the wire, /// which may be out of sequence due to network reordering or loss. Use @@ -8135,7 +8098,7 @@ impl Subscriber { /// Groups are semi-reliable, and the [`Delay::max`] budget is the other /// thing (alongside eviction and a moving start) that decides which of them arrive: /// one that has drifted further behind the live edge than the budget tolerates is - /// skipped rather than handed over, so a single poll walks off a whole backlog. The + /// skipped rather than handed over, so a single read walks off a whole backlog. The /// default is [`Duration::ZERO`], which takes the live /// edge and writes the rest off; raise it to read history. [`Self::with_range`] and /// [`Subscription::ranges`] are filters, not exemptions: backfill needs a budget that @@ -8165,10 +8128,9 @@ impl Subscriber { /// they cover there. One missing the frames asked for reads [`Error::Lagged`] rather /// than starting short. /// - /// Returns `Poll::Ready(Ok(Some(group)))` when a group is available, - /// `Poll::Ready(Ok(None))` when the track is finished or the ranges are read, - /// `Poll::Ready(Err(e))` when the track has been aborted, or - /// `Poll::Pending` when no group is available yet. + /// Returns `Ok(Some(group))` when a group is available, `Ok(None)` when the track is + /// finished or the ranges are read, or the error when the track has been aborted. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_recv_group(&mut self, waiter: &kio::Waiter) -> Poll>> { self.poll_group(false, waiter) } @@ -8421,19 +8383,6 @@ impl Subscriber { }) || owed()) } - /// Receive the next group in arrival order. - /// - /// Every group sequence is returned at most once, in the order it landed on the wire, which may - /// be out of sequence due to network reordering or loss. Use [`Self::ordered`] if you - /// only want groups whose sequence number is higher than any previously returned. - /// See [`Self::poll_recv_group`] for how [`Self::with_range`] applies. - /// Once a bounded range returns `Ok(None)`, later reads return it immediately too; - /// widening the ranges with [`Self::update`] resumes reading the newly covered groups. - #[cfg(feature = "async")] - pub async fn recv_group(&mut self) -> Result> { - kio::wait(|waiter| self.poll_recv_group(waiter)).await - } - /// The sequence cursor behind [`Ordered`], which owns the only public door to it. pub(super) fn poll_next_group(&mut self, waiter: &kio::Waiter) -> Poll>> { self.poll_group(true, waiter) @@ -8664,7 +8613,15 @@ impl Subscriber { } } - /// Poll for the track's declared exclusive end, without blocking. + /// The track's declared exclusive end position, or the cause on an abort. See + /// [`Producer::finish_at`] for what the position means. + /// + /// Resolves as soon as the end is known, which may be ahead of the live edge + /// when the producer finished via [`Producer::finish_at`]. This reports the declared + /// end, not that every group has arrived. A local session close without a declared + /// end returns [`Error::Closed`]. Read [`Self::poll_recv_group`] (or + /// [`Ordered::poll_next_group`]) until it yields `None` to observe the track fully drained. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_finished(&mut self, waiter: &kio::Waiter) -> Poll> { match &mut self.inner { Inner::Plain(cursor) => cursor.poll_finished(waiter), @@ -8672,19 +8629,6 @@ impl Subscriber { } } - /// Block until the track declares its end, returning the exclusive end position, or the - /// cause on an abort. See [`Producer::finish_at`] for what the position means. - /// - /// Resolves as soon as the end is known, which may be ahead of the live edge - /// when the producer finished via [`Producer::finish_at`]. This reports the declared - /// end, not that every group has arrived. A local session close without a declared - /// end returns [`Error::Closed`]. Drive [`Self::recv_group`] (or - /// [`Ordered::next_group`]) until it yields `None` to observe the track fully drained. - #[cfg(feature = "async")] - pub async fn finished(&mut self) -> Result { - kio::wait(|waiter| self.poll_finished(waiter)).await - } - /// Limit subsequent reads to these group sequences, returning `self` for chaining. /// /// `2..=5` includes groups 2 through 5; `2..5` excludes group 5, and `..0` is the @@ -8784,6 +8728,7 @@ impl Subscriber { /// the cache does not stand in for it, so a cap moved below the largest stays pending. /// Moving the bounds wakes a parked poll, which reports the largest once it is within /// them and newer than the last report. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_largest(&mut self, waiter: &kio::Waiter) -> Poll> { if self.largest.ended { return Poll::Ready(None); @@ -8820,12 +8765,6 @@ impl Subscriber { } } - /// The largest location within [`Self::with_range`]; see [`Self::poll_largest`]. - #[cfg(feature = "async")] - pub async fn largest(&mut self) -> Option { - kio::wait(|waiter| self.poll_largest(waiter)).await - } - /// The track's largest location known now, ignoring [`Self::with_range`], without waiting /// or moving [`Self::poll_largest`] on: past what the cache holds once the source /// reported a newer one. @@ -8924,8 +8863,7 @@ impl Ordered { self.inner.broadcast() } - /// Poll for the next group with a higher sequence number than any previously - /// returned, without blocking. + /// The next group with a higher sequence number than any previously returned. /// /// Honors the group range set by [`Self::with_range`]: /// a group past the cap stays in the producer's cache and becomes eligible again if @@ -8934,24 +8872,14 @@ impl Ordered { /// Later reads return `Ok(None)` immediately too; widening the ranges with /// [`Self::update`] resumes reading above the last sequence returned. /// - /// Returns `Poll::Ready(Ok(Some(group)))` when a group is available, - /// `Poll::Ready(Ok(None))` when the track is finished (including an abort after its - /// declared end settled, see [`Producer::abort`]) or the ranges are read, - /// `Poll::Ready(Err(e))` when the track has been aborted short of its end, or - /// `Poll::Pending` when no group is available yet. + /// Returns `Ok(Some(group))` when a group is available, `Ok(None)` when the track is + /// finished (including an abort after its declared end settled, see [`Producer::abort`]) + /// or the ranges are read, or the error when the track has been aborted short of its end. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next_group(&mut self, waiter: &kio::Waiter) -> Poll>> { self.inner.poll_next_group(waiter) } - /// Return the next group with a higher sequence number than any previously returned. - /// - /// Once a bounded range returns `Ok(None)`, later reads return it immediately too; - /// widening the ranges with [`Self::update`] resumes reading above the last sequence returned. - #[cfg(feature = "async")] - pub async fn next_group(&mut self) -> Result> { - kio::wait(|waiter| self.poll_next_group(waiter)).await - } - /// Limit subsequent reads to these group sequences, returning `self` for chaining. /// /// See [`Subscriber::with_range`] for how the bounds apply. An ordered read never goes @@ -8984,26 +8912,21 @@ impl Ordered { self.inner.update(subscription) } - /// Block until the track declares its end, returning the exclusive end position. + /// The exclusive end position, once the track declares its end. /// - /// See [`Subscriber::finished`]: this reports the declared end, not that every group + /// See [`Subscriber::poll_finished`]: this reports the declared end, not that every group /// has arrived. - #[cfg(feature = "async")] - pub async fn finished(&mut self) -> Result { - kio::wait(|waiter| self.inner.poll_finished(waiter)).await + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_finished(&mut self, waiter: &kio::Waiter) -> Poll> { + self.inner.poll_finished(waiter) } - /// Poll for the largest location within [`Self::with_range`]; see [`Subscriber::poll_largest`]. + /// The largest location within [`Self::with_range`]; see [`Subscriber::poll_largest`]. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_largest(&mut self, waiter: &kio::Waiter) -> Poll> { self.inner.poll_largest(waiter) } - /// The largest location within [`Self::with_range`]; see [`Subscriber::poll_largest`]. - #[cfg(feature = "async")] - pub async fn largest(&mut self) -> Option { - kio::wait(|waiter| self.poll_largest(waiter)).await - } - /// Whether `other` reads the same underlying track state. pub fn is_clone(&self, other: &Self) -> bool { self.inner.is_clone(&other.inner) @@ -9302,14 +9225,9 @@ impl Request { snapshot_subscription(&subs, bound) } - /// Block until the aggregate [`subscription`](Self::subscription) changes, - /// yielding `None` once nobody is waiting. - #[cfg(feature = "async")] - pub async fn subscription_changed(&mut self) -> Option { - kio::wait(|waiter| self.poll_subscription_changed(waiter)).await - } - - /// Poll for a changed subscription. + /// The aggregate [`subscription`](Self::subscription) once it changes, or `None` once + /// nobody is waiting. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_subscription_changed(&mut self, waiter: &kio::Waiter) -> Poll> { let state = self.state.read(); let (subs, bound) = (state.subscriptions.clone(), state.demand_bound()); diff --git a/rs/moq-net/src/model/track/fill.rs b/rs/moq-net/src/model/track/fill.rs index ff6cdc38084..c238259ecd3 100644 --- a/rs/moq-net/src/model/track/fill.rs +++ b/rs/moq-net/src/model/track/fill.rs @@ -847,13 +847,8 @@ impl Fill { self.inner.is_wanted(sequence) } - /// Wait until no reader waits for group `sequence` from this fill, so stop serving it. - #[cfg(feature = "async")] - pub async fn unwanted(&self, sequence: u64) { - kio::wait(|waiter| self.poll_unwanted(sequence, waiter)).await - } - - /// Poll until no reader waits for group `sequence` from this fill. + /// Resolves once no reader waits for group `sequence` from this fill, so stop serving it. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_unwanted(&self, sequence: u64, waiter: &kio::Waiter) -> Poll<()> { self.inner.poll_unwanted(sequence, waiter) } @@ -915,17 +910,12 @@ impl Fill { } } - /// Return [`wanted`](Self::wanted) once it differs from what this handle last returned: + /// [`wanted`](Self::wanted) once it differs from what this handle last returned: /// immediately on the first call, then on each change. /// /// So `while let Some(sequence) = fill.wanted_changed().await` serves each group as it /// becomes owed, and ends once the fill is withdrawn or answered. - #[cfg(feature = "async")] - pub async fn wanted_changed(&mut self) -> Option { - kio::wait(|waiter| self.poll_wanted_changed(waiter)).await - } - - /// Poll for [`wanted`](Self::wanted) to differ from what this handle last returned. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_wanted_changed(&mut self, waiter: &kio::Waiter) -> Poll> { // Arm every input of `wanted` before reading it, so a change in between still wakes: // the readers leaving, narrowing, or taking a group, the answers, the order, and a diff --git a/rs/moq-net/src/server.rs b/rs/moq-net/src/server.rs index b80542df110..b1b8e7249dc 100644 --- a/rs/moq-net/src/server.rs +++ b/rs/moq-net/src/server.rs @@ -1,6 +1,8 @@ //! Accepting a MoQ session, including the paused handshake that inspects the //! peer's SETUP before granting origins. +use std::task::{Poll, ready}; + use crate::driver::Transport; use crate::origin; use crate::time::{Deadlines, Instant}; @@ -8,7 +10,7 @@ use crate::transport::poll::Session as _; use crate::{ ALPN_18, ALPN_20, ALPN_22, ALPN_LITE, ALPN_LITE_03, ALPN_LITE_04, ALPN_LITE_05, ALPN_LITE_06, ALPN_LITE_07_WIP, Consume, Error, NEGOTIATED, Session, SessionError, Version, Versions, - coding::{Encode, Stream}, + coding::{Closing, Encode, Reader, Stream, Writer}, ietf, lite, setup::{self, Role}, stats, @@ -240,59 +242,54 @@ impl Server { )) } - async fn handshake_lite( + /// Pause a lite handshake with the client's SETUP, which only versions with a Setup + /// Stream carry. + fn lite_handshake( &self, runtime: Deadlines, - mut session: Transport, + session: Transport, version: lite::Version, - ) -> Result { - self.versions.select(Version::Lite(version)).ok_or(Error::Version)?; - let (path, role, origin, peer_implementation, handshake) = if version.has_setup_stream() { - // Gate on the client's SETUP: read it before serving so the - // caller can scope by the advertised path. Seeded back into - // `start` on `ok()` so PROBE gating resolves without + client_setup: Option, + ) -> Handshake { + let (path, role, origin, token, peer_implementation, handshake) = match client_setup { + // Seeded back into `start` on `ok()` so PROBE gating resolves without // re-reading the (consumed) Setup Stream. - let client_setup = lite::accept_setup(&mut session, version).await?; - ( + Some(client_setup) => ( client_setup.setup.path.clone(), client_setup.setup.role, client_setup.setup.hop, + client_setup.setup.token.clone(), client_setup.setup.implementation.clone(), PausedHandshake::LiteSetup { session, version, client_setup, }, - ) - } else { - (None, None, None, None, PausedHandshake::LiteBare { session, version }) - }; - - let auth = crate::auth::Link::new(match &handshake { - PausedHandshake::LiteBare { version, .. } | PausedHandshake::LiteSetup { version, .. } => { - version.has_auth() - } - PausedHandshake::Ietf(_) | PausedHandshake::Negotiated(_) => false, - }); - let token = match &handshake { - PausedHandshake::LiteSetup { client_setup, .. } => client_setup.setup.token.clone(), - _ => None, + ), + None => ( + None, + None, + None, + None, + None, + PausedHandshake::LiteBare { session, version }, + ), }; - Ok(Handshake { + Handshake { path, role, origin, token, peer_implementation, assigned_hop: crate::Hop::random(), - auth, + auth: crate::auth::Link::new(version.has_auth()), inner: Some(RequestInner { server: self.clone(), runtime, deadline: None, handshake, }), - }) + } } /// Perform the MoQ handshake as a server, returning the [`Session`] and its [`Driver`](crate::Driver). @@ -302,12 +299,21 @@ impl Server { /// Convenience wrapper over [`accept_request`](Self::accept_request) that /// completes the handshake immediately. Use `accept_request` when you need to /// inspect the client's advertised path before deciding what to serve. - pub async fn accept( + pub fn accept( &self, now: Instant, session: impl crate::transport::poll::Boxable, - ) -> Result<(Session, crate::Driver), Error> { - self.accept_request(now, session).await?.ok().await + ) -> impl Future> { + // Box before any protocol code runs, so it is compiled once rather than per transport. + let mut accept = Accept::new(Deadlines::new(now), Transport::new(session)); + let mut answer = None; + kio::wait(move |waiter| { + let answer = match &mut answer { + Some(answer) => answer, + None => answer.insert(Answer::new(ready!(accept.poll(self, waiter))?)), + }; + answer.poll(waiter) + }) } /// Begin the MoQ handshake, pausing once the client's request path is known so @@ -329,74 +335,21 @@ impl Server { session: impl crate::transport::poll::Boxable, ) -> impl Future> { // Box before any protocol code runs, so it is compiled once rather than per transport. - self.accept_boxed(Deadlines::new(now), Transport::new(session)) - } - - async fn accept_boxed(&self, runtime: Deadlines, session: Transport) -> Result { - let mut refused = session.clone(); - self.handshake(runtime, session) - .await - .inspect_err(|err| close(&mut refused, err, None)) + let mut accept = Accept::new(Deadlines::new(now), Transport::new(session)); + kio::wait(move |waiter| accept.poll(self, waiter)) } - async fn handshake(&self, runtime: Deadlines, mut session: Transport) -> Result { - let supported = match session.protocol() { - Some(alpn @ (ALPN_22 | ALPN_20 | ALPN_18)) => { - let draft = match alpn { - ALPN_22 => ietf::Version::Draft22, - ALPN_20 => ietf::Version::Draft20, - _ => ietf::Version::Draft18, - }; - - self.versions.select(Version::Ietf(draft)).ok_or(Error::Version)?; - return self.accept_ietf(runtime, session, draft).await; - } - Some(alpn @ (ALPN_LITE_07_WIP | ALPN_LITE_06 | ALPN_LITE_05 | ALPN_LITE_04 | ALPN_LITE_03)) => { - let version = match alpn { - ALPN_LITE_07_WIP => lite::Version::Lite07, - ALPN_LITE_06 => lite::Version::Lite06, - ALPN_LITE_05 => lite::Version::Lite05, - ALPN_LITE_04 => lite::Version::Lite04, - _ => lite::Version::Lite03, - }; - return self.handshake_lite(runtime, session, version).await; - } - // No ALPN, or `moql`: negotiate moq-lite in a bidi SETUP. - Some(ALPN_LITE) | None => self.versions.filter(&NEGOTIATED.into()).ok_or(Error::Version)?, - Some(p) => return Err(Error::UnknownAlpn(p.to_string())), - }; - - // Read the client's SETUP to choose the version; `ok()` sends the server SETUP and starts. - // Stats attach in `ok()`: a caller may still replace them on the handshake. - // The SETUP keeps lite-01's framing whichever version it offers. - let mut stream = Stream::accept(&mut session, Version::Lite(lite::Version::Lite01)).await?; - let (client, client_setup_size) = stream.reader.decode_sized::().await?; - - let version = client - .versions - .iter() - .flat_map(|v| Version::try_from(*v).ok()) - .find(|v| supported.contains(v)) - .ok_or(Error::Version)?; - let Version::Lite(version) = version else { - return Err(Error::Version); - }; - - // Answer, close the stream, and continue as that version's ALPN would. Waiting for the - // acknowledgement keeps the reply from being discarded by the reset a dropped stream sends. - if Version::Lite(version).hands_off() { - let Stream { mut writer, .. } = stream; - writer - .encode(&setup::Server { - version: Version::Lite(version).into(), - parameters: lite::Parameters::default().encode_bytes(version)?, - }) - .await?; - writer.close().await?; - return self.handshake_lite(runtime, session, version).await; - } - - Ok(Handshake { + /// Pause a lite-01/02 handshake once the client's SETUP picked the version; `ok()` sends + /// the server SETUP and starts. + fn negotiated_handshake( + &self, + runtime: Deadlines, + session: Transport, + stream: Stream, + version: lite::Version, + client_setup_size: usize, + ) -> Handshake { + Handshake { path: None, role: None, origin: None, @@ -416,20 +369,20 @@ impl Server { client_setup_size, }), }), - }) + } } - /// Read a moq-transport client's SETUP (with its request path) off its uni stream, - /// then pause. `ok()` starts the session and hands the stream back for GOAWAY. - async fn accept_ietf( + /// Pause a moq-transport handshake with the client's SETUP (and its request path). + /// `ok()` starts the session and hands the SETUP stream back for GOAWAY. + fn ietf_handshake( &self, runtime: Deadlines, - mut session: Transport, + session: Transport, version: ietf::Version, - ) -> Result { - let peer_setup = ietf::accept_setup(&mut session, version).await?; + peer_setup: ietf::PeerSetup, + ) -> Handshake { let options = &peer_setup.options; - Ok(Handshake { + Handshake { path: options.path.clone(), role: None, // A moq-transport peer only has an identity if it negotiated the MoQ @@ -451,6 +404,278 @@ impl Server { peer_setup, }), }), + } + } +} + +/// The server's side of a handshake, up to the pause where [`Handshake`] hands the client's +/// SETUP to the caller. +/// +/// A failure closes the transport with the error's code, so the peer learns why instead of +/// seeing a bare disconnect. Dropped while pending, it releases what its step holds and +/// closes nothing. +struct Accept { + /// Closed on failure. Released once resolved, so it never outlives the handshake. + refused: Option, + /// After a hand-off, the negotiation stream's receive half, kept open until the handshake + /// resolves: dropping it can stop the client's half of that stream. + held: Option>, + step: AcceptStep, +} + +/// Where [`Accept`] waits. Every step but `Done` holds the session's clock and transport, +/// which move into the [`Handshake`]. +enum AcceptStep { + /// Not polled yet: the transport's ALPN picks the path. + Start { runtime: Deadlines, session: Transport }, + /// Reading a moq-transport client's SETUP, holding the uni streams that beat it. + Ietf { + runtime: Deadlines, + session: Transport, + version: ietf::Version, + setup: ietf::SetupAccept, + }, + /// Reading a lite client's Setup Stream, holding the uni streams that beat it, or, with + /// `None`, a version that has none and pauses at once. + Lite { + runtime: Deadlines, + session: Transport, + version: lite::Version, + setup: Option, + }, + /// Waiting for the bidi stream a client without a versioned ALPN negotiates on, then + /// reading its SETUP off it. The SETUP keeps lite-01's framing whichever version it offers. + Negotiate { + runtime: Deadlines, + session: Transport, + supported: Versions, + stream: Option>, + }, + /// Flushing the server SETUP that hands the session off to the picked version's ALPN path. + HandOff { + runtime: Deadlines, + session: Transport, + version: lite::Version, + writer: Writer, + }, + /// Waiting for the client to acknowledge that reply, so a reset cannot discard it. + HandOffClose { + runtime: Deadlines, + session: Transport, + version: lite::Version, + closing: Closing, + }, + /// Resolved. + Done, +} + +impl Accept { + fn new(runtime: Deadlines, session: Transport) -> Self { + Self { + refused: Some(session.clone()), + held: None, + step: AcceptStep::Start { runtime, session }, + } + } + + /// The paused handshake, or the error the transport was closed with. + fn poll(&mut self, server: &Server, waiter: &kio::Waiter) -> Poll> { + let res = ready!(self.poll_step(server, waiter)); + // Release everything the handshake held before closing over a failure. + self.step = AcceptStep::Done; + self.held = None; + let refused = self.refused.take(); + if let (Err(err), Some(mut refused)) = (&res, refused) { + close(&mut refused, err, None); + } + Poll::Ready(res) + } + + fn take(&mut self) -> AcceptStep { + std::mem::replace(&mut self.step, AcceptStep::Done) + } + + fn poll_step(&mut self, server: &Server, waiter: &kio::Waiter) -> Poll> { + loop { + match &mut self.step { + AcceptStep::Start { .. } => { + let AcceptStep::Start { runtime, session } = self.take() else { + unreachable!("matched above"); + }; + self.step = match session.protocol() { + Some(alpn @ (ALPN_22 | ALPN_20 | ALPN_18)) => { + let version = match alpn { + ALPN_22 => ietf::Version::Draft22, + ALPN_20 => ietf::Version::Draft20, + _ => ietf::Version::Draft18, + }; + server.versions.select(Version::Ietf(version)).ok_or(Error::Version)?; + AcceptStep::Ietf { + runtime, + session, + version, + setup: ietf::SetupAccept::new(version), + } + } + Some(alpn @ (ALPN_LITE_07_WIP | ALPN_LITE_06 | ALPN_LITE_05 | ALPN_LITE_04 | ALPN_LITE_03)) => { + let version = match alpn { + ALPN_LITE_07_WIP => lite::Version::Lite07, + ALPN_LITE_06 => lite::Version::Lite06, + ALPN_LITE_05 => lite::Version::Lite05, + ALPN_LITE_04 => lite::Version::Lite04, + _ => lite::Version::Lite03, + }; + Self::lite_step(server, runtime, session, version)? + } + // No ALPN, or `moql`: negotiate moq-lite in a bidi SETUP. + Some(ALPN_LITE) | None => AcceptStep::Negotiate { + runtime, + session, + supported: server.versions.filter(&NEGOTIATED.into()).ok_or(Error::Version)?, + stream: None, + }, + Some(p) => return Poll::Ready(Err(Error::UnknownAlpn(p.to_string()))), + }; + } + AcceptStep::Ietf { session, setup, .. } => { + let peer_setup = ready!(setup.poll(session, waiter))?; + let AcceptStep::Ietf { + runtime, + session, + version, + .. + } = self.take() + else { + unreachable!("matched above"); + }; + return Poll::Ready(Ok(server.ietf_handshake(runtime, session, version, peer_setup))); + } + AcceptStep::Lite { session, setup, .. } => { + // Gate on the client's SETUP: read it before serving so the caller can + // scope by the advertised path. + let client_setup = match setup { + Some(setup) => Some(ready!(setup.poll(session, waiter))?), + None => None, + }; + let AcceptStep::Lite { + runtime, + session, + version, + .. + } = self.take() + else { + unreachable!("matched above"); + }; + return Poll::Ready(Ok(server.lite_handshake(runtime, session, version, client_setup))); + } + AcceptStep::Negotiate { + session, + supported, + stream, + .. + } => { + let stream = match stream { + Some(stream) => stream, + None => stream.insert(ready!(Stream::poll_accept( + session, + Version::Lite(lite::Version::Lite01), + waiter + ))?), + }; + // Stats attach in `ok()`: a caller may still replace them on the handshake. + let (client, client_setup_size) = ready!(stream.reader.poll_decode_sized::(waiter))?; + let version = client + .versions + .iter() + .flat_map(|v| Version::try_from(*v).ok()) + .find(|v| supported.contains(v)) + .ok_or(Error::Version)?; + let Version::Lite(version) = version else { + return Poll::Ready(Err(Error::Version)); + }; + let AcceptStep::Negotiate { + runtime, + session, + stream: Some(stream), + .. + } = self.take() + else { + unreachable!("matched above"); + }; + if !Version::Lite(version).hands_off() { + return Poll::Ready(Ok(server.negotiated_handshake( + runtime, + session, + stream, + version, + client_setup_size, + ))); + } + + // Answer, close the stream, and continue as that version's ALPN would. + let reply = setup::Server { + version: Version::Lite(version).into(), + parameters: lite::Parameters::default().encode_bytes(version)?, + }; + let Stream { mut writer, reader } = stream; + self.held = Some(reader); + writer.buffer(&reply)?; + self.step = AcceptStep::HandOff { + runtime, + session, + version, + writer, + }; + } + AcceptStep::HandOff { writer, .. } => { + ready!(writer.poll_flush(waiter))?; + let AcceptStep::HandOff { + runtime, + session, + version, + writer, + } = self.take() + else { + unreachable!("matched above"); + }; + self.step = AcceptStep::HandOffClose { + runtime, + session, + version, + closing: writer.closing()?, + }; + } + AcceptStep::HandOffClose { closing, .. } => { + ready!(closing.poll(waiter))?; + let AcceptStep::HandOffClose { + runtime, + session, + version, + .. + } = self.take() + else { + unreachable!("matched above"); + }; + self.step = Self::lite_step(server, runtime, session, version)?; + } + AcceptStep::Done => panic!("handshake polled after it resolved"), + } + } + } + + /// The step that reads a lite client's SETUP, refusing a version the server doesn't speak. + fn lite_step( + server: &Server, + runtime: Deadlines, + session: Transport, + version: lite::Version, + ) -> Result { + server.versions.select(Version::Lite(version)).ok_or(Error::Version)?; + Ok(AcceptStep::Lite { + runtime, + session, + version, + setup: version.has_setup_stream().then(|| lite::SetupAccept::new(version)), }) } } @@ -594,14 +819,15 @@ struct PausedNegotiated { } impl PausedNegotiated { - /// Send the server SETUP and start the session. - async fn ok( + /// Buffer the server SETUP, counting the client's under the final stats. The returned + /// reply sends it and starts the session. + fn reply( self, server: Server, runtime: Deadlines, peer_hop: Option, auth: crate::auth::Link, - ) -> Result<(Session, crate::Driver), Error> { + ) -> Result { let Self { session, stream, @@ -612,12 +838,64 @@ impl PausedNegotiated { server.stats.control(stats::Message::Session, client_setup_size, true); let (publish, subscribe) = server.stat_tagged_origins(); - let server_setup = setup::Server { + stream.writer.buffer(&setup::Server { version: Version::Lite(version).into(), parameters: lite::Parameters::default().encode_bytes(version)?, - }; - stream.writer.encode(&server_setup).await?; + })?; + Ok(NegotiatedReply(Some(NegotiatedStart { + server, + runtime, + session, + stream, + version, + peer_hop, + auth, + publish, + subscribe, + }))) + } +} +/// A negotiated client's server SETUP on its way out. Dropped before it lands, it resets the +/// negotiation stream. +struct NegotiatedReply(Option); + +impl NegotiatedReply { + /// The started session, once the server SETUP has hit the stream. + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + let start = self.0.as_mut().expect("reply polled after it resolved"); + let flushed = ready!(start.stream.writer.poll_flush(waiter)); + let start = self.0.take().expect("checked above"); + Poll::Ready(flushed.and_then(|()| start.start())) + } +} + +/// What a negotiated session starts with, held while its server SETUP flushes. +struct NegotiatedStart { + server: Server, + runtime: Deadlines, + session: Transport, + stream: Stream, + version: lite::Version, + peer_hop: Option, + auth: crate::auth::Link, + publish: Option, + subscribe: Option, +} + +impl NegotiatedStart { + fn start(self) -> Result<(Session, crate::Driver), Error> { + let Self { + server, + runtime, + session, + stream, + version, + peer_hop, + auth, + publish, + subscribe, + } = self; let stream = stream.with_version(version); // Pre-lite-05: no Setup Stream, so nothing to advertise or seed. // One registry, which the subscriber sends and the session hands out. @@ -795,21 +1073,89 @@ impl Handshake { /// Accept the session, returning the [`Session`] and its [`Driver`](crate::Driver). /// /// Poll or spawn the returned driver to run the session. - pub async fn ok(mut self) -> Result<(Session, crate::Driver), Error> { - let peer_hop = Some(self.assigned_hop); - let auth = self.auth.clone(); + pub async fn ok(self) -> Result<(Session, crate::Driver), Error> { + let mut answer = Answer::new(self); + kio::wait(|waiter| answer.poll(waiter)).await + } + + /// Reject the session, closing the transport with `err`'s wire code. + pub fn close(mut self, err: Error) { + let inner = self.inner.take().expect("request already responded"); + inner.close(err); + } +} + +impl RequestInner { + fn close(self, err: Error) { + self.server.stats.end(&err); + let (mut session, version) = self.handshake.transport(); + close(&mut session, &err, Some(version)); + } +} + +/// [`Handshake::ok`] as a poll: starts the session, once a negotiated client has taken the +/// server SETUP. Dropped before it resolves, it refuses the session as the deadline does. +// Fields drop in order: the reply before the guard refuses, and the handshake last. +struct Answer { + /// What starts the session, taken by the first poll. + start: Option<(Server, Deadlines, PausedHandshake)>, + /// A negotiated client's server SETUP on its way out. Boxed so the bidi SETUP write of + /// these old versions doesn't size every caller's handshake. + reply: Option>, + /// The deadline's timer, armed once the reply first waits. + timer: Option, + unanswered: Unanswered, + stats: stats::Session, + /// When to give up on the peer taking the reply. + deadline: Option, + /// The request being answered, emptied of what the start took. + handshake: Handshake, +} + +impl Answer { + fn new(mut handshake: Handshake) -> Self { let RequestInner { server, runtime, deadline, + handshake: paused, + } = handshake.inner.take().expect("request already responded"); + let stats = server.stats.clone(); + let (session, version) = paused.transport(); + Self { + start: Some((server, runtime, paused)), + reply: None, + timer: None, + unanswered: Unanswered(Some((session, version, stats.clone()))), + stats, + deadline, handshake, - } = self.inner.take().expect("request already responded"); + } + } - let stats = server.stats.clone(); - let (session, version) = handshake.transport(); - let unanswered = Unanswered(Some((session, version, stats.clone()))); - let reply = async move { - match handshake { + /// The started session, or why it failed. + fn poll(&mut self, waiter: &kio::Waiter) -> Poll> { + let res = ready!(self.poll_reply(waiter)); + // Release the reply before the guard closes over a deadline. + self.reply = None; + self.timer = None; + let Some(res) = res else { + drop(std::mem::replace(&mut self.unanswered, Unanswered(None))); + return Poll::Ready(Err(Error::Timeout)); + }; + self.unanswered.0 = None; + if let Err(err) = &res { + self.stats.end(err); + } + Poll::Ready(res) + } + + /// The reply's outcome, or `None` once the deadline passed first. + fn poll_reply(&mut self, waiter: &kio::Waiter) -> Poll>> { + if let Some((server, runtime, paused)) = self.start.take() { + let peer_hop = Some(self.handshake.assigned_hop); + let auth = self.handshake.auth.clone(); + let started = match paused { PausedHandshake::LiteBare { session, version } => { server.start_lite(runtime, session, version, None, peer_hop, auth) } @@ -822,34 +1168,27 @@ impl Handshake { server.start_lite(runtime, session, version, Some(client_setup), peer_hop, auth) } PausedHandshake::Ietf(paused) => paused.ok(server, runtime, peer_hop, auth), - // Boxed so the bidi SETUP write of these old versions doesn't size every caller's future. - PausedHandshake::Negotiated(paused) => Box::pin(paused.ok(server, runtime, peer_hop, auth)).await, - } - }; - let res = match deadline { - // The guard refuses the session on the way out. - Some(at) => kio::clock::timeout_at(at, reply).await.map_err(|_| Error::Timeout)?, - None => reply.await, - }; - unanswered.answered(); - if let Err(err) = &res { - stats.end(err); + PausedHandshake::Negotiated(paused) => match paused.reply(server, runtime, peer_hop, auth) { + Ok(reply) => { + self.reply = Some(Box::new(reply)); + return self.poll_reply(waiter); + } + Err(err) => Err(err), + }, + }; + return Poll::Ready(Some(started)); } - res - } - - /// Reject the session, closing the transport with `err`'s wire code. - pub fn close(mut self, err: Error) { - let inner = self.inner.take().expect("request already responded"); - inner.close(err); - } -} - -impl RequestInner { - fn close(self, err: Error) { - self.server.stats.end(&err); - let (mut session, version) = self.handshake.transport(); - close(&mut session, &err, Some(version)); + let reply = self.reply.as_mut().expect("answer polled after it resolved"); + // The reply is polled first, so one that is ready wins even past the deadline. + if let Poll::Ready(res) = reply.poll(waiter) { + return Poll::Ready(Some(res)); + } + let Some(at) = self.deadline else { + return Poll::Pending; + }; + let timer = self.timer.get_or_insert_with(|| kio::clock::sleep_until(at)); + ready!(waiter.poll_future(std::pin::Pin::new(timer))); + Poll::Ready(None) } } @@ -858,12 +1197,6 @@ impl RequestInner { /// took too long. struct Unanswered(Option<(Transport, Version, stats::Session)>); -impl Unanswered { - fn answered(mut self) { - self.0 = None; - } -} - impl Drop for Unanswered { fn drop(&mut self) { if let Some((mut session, version, stats)) = self.0.take() { @@ -1302,6 +1635,27 @@ mod tests { assert_eq!(failed.count(stats::Failure::Application), 1, "{failed:?}"); } + /// A handshake the caller stops waiting for before the client's SETUP arrives closes + /// nothing and lets go of the transport, so the caller's own deadline decides the close. + #[kio::test] + async fn an_abandoned_accept_closes_nothing() { + for session in [ + FakeSession::new(ALPN_22, []), + FakeSession::new(ALPN_LITE_05, []), + FakeSession::new(ALPN_LITE, []), + ] { + let observed = session.clone(); + let abandoned = kio::clock::timeout( + std::time::Duration::from_secs(1), + Server::new().accept_request(kio::clock::now(), session), + ) + .await; + assert!(abandoned.is_err(), "{:?}: accepted without a SETUP", observed.protocol); + assert_eq!(observed.close_code(), None, "{:?}", observed.protocol); + assert_eq!(Arc::strong_count(&observed.closed), 1, "{:?}", observed.protocol); + } + } + /// A rejection counts by the code the peer is told, so a "no error" close is clean /// however this side spelled it. #[kio::test] diff --git a/rs/moq-net/src/session.rs b/rs/moq-net/src/session.rs index 2ed971dfe3c..b5042245e90 100644 --- a/rs/moq-net/src/session.rs +++ b/rs/moq-net/src/session.rs @@ -1,6 +1,10 @@ //! A MoQ session handle and a snapshot of its connection statistics. -use std::{sync::Arc, task::Poll, time::Duration}; +use std::{ + sync::Arc, + task::{Poll, ready}, + time::Duration, +}; use crate::transport::poll::Session as _; use crate::transport::{Error as _, Metrics as _}; @@ -28,6 +32,20 @@ struct Ended { drain: Option>, } +/// Resolve with `read` of how the session ended, or [`kio::Closed`] once the driver dropped +/// without publishing it. +fn poll_ended( + closed: &kio::Consumer>, + waiter: &kio::Waiter, + read: impl FnOnce(&Ended) -> R, +) -> Poll> { + let ended = closed.poll(waiter, |state| match &**state { + Some(ended) => Poll::Ready(read(ended)), + None => Poll::Pending, + }); + ended.map(|ended| ended.map_err(|_| kio::Closed)) +} + /// How this side closed the transport, shared by every clone of it. /// /// Every close goes through the boxed transport, so the first one is recorded whichever @@ -191,25 +209,23 @@ impl Session { /// 01-04), or the session ended first. Peer-controlled: use it for logs and /// rollout counts, never to authorize. pub async fn peer_implementation(&self) -> Option { + kio::wait(|waiter| self.poll_peer_implementation(waiter)).await + } + + /// See [`Self::peer_implementation`]. + fn poll_peer_implementation(&self, waiter: &kio::Waiter) -> Poll> { // The slot closes only when the driver drops, which a caller holding a finished // one never does, so the transport's end is watched too. - kio::wait(|waiter| { - let slot = self.peer_implementation.poll(waiter, |slot| match &**slot { - Some(value) => Poll::Ready(value.clone()), - None => Poll::Pending, - }); - match slot { - Poll::Ready(Ok(value)) => return Poll::Ready(value), - Poll::Ready(Err(_)) => return Poll::Ready(None), - Poll::Pending => {} - } - let ended = self.closed.poll(waiter, |ended| match &**ended { - Some(_) => Poll::Ready(()), - None => Poll::Pending, - }); - ended.map(|_| None) - }) - .await + let slot = self.peer_implementation.poll(waiter, |slot| match &**slot { + Some(value) => Poll::Ready(value.clone()), + None => Poll::Pending, + }); + match slot { + Poll::Ready(Ok(value)) => return Poll::Ready(value), + Poll::Ready(Err(_)) => return Poll::Ready(None), + Poll::Pending => {} + } + poll_ended(&self.closed, waiter, |_| ()).map(|_| None) } /// Returns a consumer for the estimated send bitrate (from the congestion controller). @@ -308,16 +324,13 @@ impl Session { let closed = self.closed.clone(); drop(self); - match closed - .wait(|state| match &**state { - Some(ended) => Poll::Ready(ended.clone()), - None => Poll::Pending, + kio::wait(|waiter| { + Poll::Ready(match ready!(poll_ended(&closed, waiter, Ended::clone)) { + Ok(ended) => ended.drain.unwrap_or(Err(ended.err)), + Err(kio::Closed) => Err(Error::Cancel), }) - .await - { - Ok(ended) => ended.drain.unwrap_or(Err(ended.err)), - Err(kio::Closed) => Err(Error::Cancel), - } + }) + .await } /// Block until the transport session is closed, returning the reason. @@ -336,18 +349,18 @@ impl Session { /// completion, this resolves with [`Error::Cancel`]. A panic while polling the driver is /// a transport failure instead. pub async fn closed(&self) -> Error { - match self - .closed - .wait(|state| match &**state { - Some(ended) => Poll::Ready(ended.err.clone()), - None => Poll::Pending, - }) - .await - { - Ok(err) => err, - // The driver was dropped before it could observe the close. - Err(kio::Closed) => Error::Cancel, - } + kio::wait(|waiter| self.poll_closed(waiter)).await + } + + /// See [`Self::closed`]. + fn poll_closed(&self, waiter: &kio::Waiter) -> Poll { + Poll::Ready( + match ready!(poll_ended(&self.closed, waiter, |ended| ended.err.clone())) { + Ok(err) => err, + // The driver was dropped before it could observe the close. + Err(kio::Closed) => Error::Cancel, + }, + ) } /// Drain the peer gracefully: the handle for sending this session's single @@ -380,22 +393,16 @@ impl Session { self.goaway.consumer() } - /// Poll until the session carries no subscription or fetch, in either direction. + /// Resolves once the session carries no subscription or fetch, in either direction. /// /// Ready at once while it carries none. A track request still being set up, waiting /// on authorization or the track's info, counts too. Announcements and the other /// control streams don't, so a session can be idle while it still serves discovery. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_idle(&self, waiter: &kio::Waiter) -> Poll<()> { self.work.poll_idle(waiter) } - /// Wait until the session carries no subscription or fetch, in either direction. - /// - /// See [`poll_idle`](Self::poll_idle). - pub async fn idle(&self) { - kio::wait(|waiter| self.poll_idle(waiter)).await - } - /// The tokens this side presented and the grant they earned, plus the tokens /// the peer presents. See [`auth`]. /// diff --git a/rs/moq-net/src/time.rs b/rs/moq-net/src/time.rs index 680764439d5..cb336d86607 100644 --- a/rs/moq-net/src/time.rs +++ b/rs/moq-net/src/time.rs @@ -21,9 +21,9 @@ pub trait Driver { /// Run a driver to completion on the ambient runtime, sleeping until each deadline. #[cfg(feature = "async")] -pub async fn run(driver: D) -> Error { +pub fn run(driver: D) -> impl Future { let mut driver = Box::new(driver); - kio::clock::drive(move |now, waiter| driver.poll(now, waiter)).await + kio::clock::drive(move |now, waiter| driver.poll(now, waiter)) } /// A deadline queue attached to the deterministic test clock. diff --git a/rs/moq-net/src/transport/boxed.rs b/rs/moq-net/src/transport/boxed.rs index 4730d990ca5..d37c209c833 100644 --- a/rs/moq-net/src/transport/boxed.rs +++ b/rs/moq-net/src/transport/boxed.rs @@ -68,12 +68,6 @@ impl Session { }) } - /// Open a unidirectional stream scheduled by `options`, waiting its turn for credit. - pub(crate) async fn open_uni(&self, options: SendOptions<'_, Group>) -> Result { - let mut open = Open::default(); - kio::wait(|waiter| self.poll_open_uni(&mut open, waiter, options)).await - } - /// The handle each stream this session yields reserves its paused credit through. fn budget(&self) -> Budget { Arc::downgrade(&self.inner) diff --git a/rs/moq-net/src/util.rs b/rs/moq-net/src/util.rs index 438ce870ef5..23ea5c01c72 100644 --- a/rs/moq-net/src/util.rs +++ b/rs/moq-net/src/util.rs @@ -1,6 +1,6 @@ -//! Future boxing and the driver-owned task set. +//! The driver-owned task set. -use std::{collections::VecDeque, future::Future, pin::Pin, task::Poll}; +use std::{collections::VecDeque, task::Poll}; mod cell; mod idle; @@ -13,37 +13,10 @@ pub(crate) use idle::Idle; pub(crate) use later::{Later, Learning}; pub(crate) use wake::QueueWake; -/// Prove a future's Send bound here rather than inside every caller's future. -pub(crate) fn sendable(future: F) -> impl Future + Send { - future -} - -/// A boxed `Send` future, the same type as `futures::future::BoxFuture`. -pub(crate) type BoxFuture<'a, T> = Pin + Send + 'a>>; - -/// A boxed [`kio::Tasks`] task: hand-rolled machines and adapted futures alike. +/// A boxed [`kio::Tasks`] task. pub(crate) type BoxTask = Box Poll<()> + Send>; -/// Adapt a boxed future into a [`kio::Tasks`] task. -fn future_task(mut task: BoxFuture<'static, ()>) -> BoxTask { - Box::new(move |waiter| waiter.poll_future(task.as_mut())) -} - -/// Resolve with the error of a fallible future, parking forever on success. -/// -/// The building block for racing "watchdog" futures whose clean completion should -/// not end the race (only their failure should). -#[cfg(feature = "async")] -pub(crate) fn err_only(fut: impl Future> + Send) -> impl Future + Send { - sendable(async move { - if let Err(err) = fut.await { - return err; - } - std::future::pending().await - }) -} - -/// Cloneable handle for submitting futures to a driver-owned [`TaskSet`]. +/// Cloneable handle for submitting tasks to a driver-owned [`TaskSet`]. #[derive(Clone)] pub(crate) struct Tasks { state: kio::Shared, @@ -63,22 +36,24 @@ pub(crate) struct Keepalive { /// The submission queue between [`Tasks`] handles and their [`TaskSet`]. #[derive(Default)] struct Submissions { - queued: VecDeque>, + queued: VecDeque, /// The set dropped: submissions are discarded, as the driver that would /// have polled them has torn down. closed: bool, } impl Tasks { - /// Queue a future for polling by the associated [`TaskSet`]. + /// Queue a task for polling by the associated [`TaskSet`]. /// /// A task submitted after the set is gone is dropped: the driver has torn down, /// so the task would have been cancelled anyway. #[cfg_attr(not(feature = "async"), allow(dead_code, reason = "async-only"))] - pub fn push(&self, task: impl Future + Send + 'static) { + pub fn push(&self, mut task: impl kio::Task + Send + 'static) { let mut state = self.state.lock(); if !state.closed { - state.queued.push_back(Box::pin(task)); + state + .queued + .push_back(Box::new(move |waiter: &kio::Waiter| task.poll(waiter))); } } } @@ -111,12 +86,14 @@ pub(crate) struct TasksWeak { } impl TasksWeak { - /// Queue a future for polling by the associated [`TaskSet`], dropped if the + /// Queue a task for polling by the associated [`TaskSet`], dropped if the /// set (or every owning handle) is already gone. - pub fn push(&self, task: impl Future + Send + 'static) { + pub fn push(&self, mut task: impl kio::Task + Send + 'static) { let mut state = self.state.lock(); if !state.closed && !self.alive.is_closed() { - state.queued.push_back(Box::pin(task)); + state + .queued + .push_back(Box::new(move |waiter: &kio::Waiter| task.poll(waiter))); } } @@ -164,10 +141,11 @@ impl TaskSet { } } - /// Queue a future for polling by this set. + /// Queue a task for polling by this set. #[cfg_attr(not(feature = "async"), allow(dead_code, reason = "async-only"))] - pub fn push(&mut self, task: impl Future + Send + 'static) { - self.active.push(future_task(Box::pin(task))); + pub fn push(&mut self, mut task: impl kio::Task + Send + 'static) { + self.active + .push(Box::new(move |waiter: &kio::Waiter| task.poll(waiter))); } /// Adopt every queued submission into the active set, then poll the children. @@ -194,7 +172,7 @@ impl TaskSet { break; }; while let Some(task) = state.queued.pop_front() { - self.active.push(future_task(task)); + self.active.push(task); } // No handles left: nothing can ever submit again. if orphaned { @@ -209,25 +187,6 @@ impl TaskSet { } } - /// Poll every child while polling `f`, returning its output. - /// - /// `f` is a poll function rather than a future, so an accept loop polls the - /// transport directly and there is no accept future to cancel or rebuild when - /// a child finishes. - #[cfg(feature = "async")] - pub async fn drive(&mut self, mut f: impl FnMut(&kio::Waiter) -> Poll + Unpin) -> T { - kio::wait(|waiter| { - if let Poll::Ready(output) = f(waiter) { - return Poll::Ready(output); - } - // The children never end the drive; a `Ready` here just means they're - // drained until the next submission. - let _ = self.poll(waiter); - Poll::Pending - }) - .await - } - /// Drive submitted children until every submission handle is dropped and all /// active children finish. The async form of [`Self::poll`], for callers with /// nothing else to poll. @@ -279,42 +238,19 @@ mod tests { let nested_tasks = tasks.clone(); let outer_completed = completed.clone(); - tasks.push(async move { + tasks.push(move |_: &kio::Waiter| { outer_completed.fetch_add(1, Ordering::SeqCst); let inner_completed = outer_completed.clone(); - nested_tasks.push(async move { + nested_tasks.push(move |_: &kio::Waiter| { inner_completed.fetch_add(1, Ordering::SeqCst); + Poll::Ready(()) }); + Poll::Ready(()) }); drop(tasks); futures::executor::block_on(task_set.run()); assert_eq!(completed.load(Ordering::SeqCst), 2); } - - #[test] - fn drive_polls_children_alongside_the_accept_future() { - let (tasks, mut set) = TaskSet::new(); - let completed = Arc::new(AtomicUsize::new(0)); - - let child_completed = completed.clone(); - tasks.push(async move { - child_completed.fetch_add(1, Ordering::SeqCst); - }); - - // The driven poll only resolves after the child ran, proving drive() - // interleaves both without an executor. - let gate = completed.clone(); - let output = futures::executor::block_on(set.drive(move |waiter| { - if gate.load(Ordering::SeqCst) == 1 { - std::task::Poll::Ready(42) - } else { - waiter.waker().wake_by_ref(); - std::task::Poll::Pending - } - })); - - assert_eq!(output, 42); - } } diff --git a/rs/moq-relay/Cargo.toml b/rs/moq-relay/Cargo.toml index 21abed77acb..8fd61d1515c 100644 --- a/rs/moq-relay/Cargo.toml +++ b/rs/moq-relay/Cargo.toml @@ -68,11 +68,11 @@ hyper = { workspace = true, optional = true } # `server-auto` reclaims the socket from under the auto builder `axum_server` serves with. hyper-util = { workspace = true, features = ["server-auto"] } kio = { workspace = true, features = ["runtime", "serde"] } -moq-auth = { workspace = true, features = ["client"] } +moq-auth = { workspace = true, features = ["async", "client"] } moq-net = { workspace = true } # The flow-control limits both QUIC runtimes take, for widening a peer's connection. moq-quic = { workspace = true } -moq-stats = { workspace = true } +moq-stats = { workspace = true, features = ["async"] } moq-tokio = { workspace = true, default-features = false, features = ["watch", "tcp"] } # The /metrics text encoder; the protobuf features stay off. prometheus-client = "0.25.1" diff --git a/rs/moq-relay/tests/integration/refusals.rs b/rs/moq-relay/tests/integration/refusals.rs index dde11bd406f..c7c9993ccab 100644 --- a/rs/moq-relay/tests/integration/refusals.rs +++ b/rs/moq-relay/tests/integration/refusals.rs @@ -262,7 +262,7 @@ async fn refused_only_stats_depth_one_has_root_totals_and_shared_embedder_metric .unwrap() .unwrap(); let consumer = moq_stats::Consumer::new(broadcast, Default::default()); - let mut totals = consumer.totals().await.unwrap(); + let mut totals = consumer.totals().unwrap(); kio::clock::timeout(TIMEOUT, async { loop { let frame = totals.next().await.unwrap().unwrap(); @@ -319,7 +319,7 @@ async fn peer_zero_close_is_clean_and_application_close_counts_once() { .unwrap() .unwrap(); let consumer = moq_stats::Consumer::new(broadcast, Default::default()); - let mut totals = consumer.totals().await.unwrap(); + let mut totals = consumer.totals().unwrap(); for (index, error) in [moq_net::Error::Cancel, moq_net::Error::App(42)] .into_iter() .enumerate() diff --git a/rs/moq-room/Cargo.toml b/rs/moq-room/Cargo.toml index b94fa763709..fc03f70e41b 100644 --- a/rs/moq-room/Cargo.toml +++ b/rs/moq-room/Cargo.toml @@ -13,6 +13,15 @@ rust-version.workspace = true keywords = ["quic", "webtransport", "conferencing", "live"] categories = ["multimedia", "network-programming"] +# The self dev-dependency only turns on `async` for the tests. +[package.metadata.cargo-shear] +ignored = ["moq-room"] + +[features] +default = [] +# Generated async wrappers for the poll methods, and `chat::Subscriber::subscribe`. Enable where you `.await` them. +async = ["kio/async"] + [dependencies] kio = { workspace = true, features = ["runtime"] } moq-auth = { workspace = true } @@ -21,6 +30,8 @@ moq-net = { workspace = true } thiserror = { workspace = true } [dev-dependencies] +# Compile the generated wrappers used by the tests. +moq-room = { path = ".", features = ["async"] } moq-tokio = { path = "../moq-tokio", default-features = false, features = ["net"] } tokio = { workspace = true, features = ["macros", "rt", "test-util"] } workspace-hack = { path = "../workspace-hack" } diff --git a/rs/moq-room/MIGRATE.md b/rs/moq-room/MIGRATE.md index aa6ebd7cc57..eead40fc2f2 100644 --- a/rs/moq-room/MIGRATE.md +++ b/rs/moq-room/MIGRATE.md @@ -4,6 +4,12 @@ How to update your code for each breaking change, newest first. ## Unreleased +### Async wrappers are opt-in + +Enable the `async` Cargo feature to keep awaiting `Room::next`, `chat::Publisher::expire`, +`chat::Subscriber::recv`, and `chat::Subscriber::subscribe`. The feature is off by default; +`poll_next`, `poll_expire`, `poll_recv`, and `chat::Subscriber::new` stay available without it. + ### `moq_net::Path` is bounded `broadcast_path` returns `Result`, failing when the identity and kind exceed `moq_net::Path::MAX_BYTES`. `claims` takes `impl TryInto>` and fails with `Error::Net` on a room or identity that long. diff --git a/rs/moq-room/README.md b/rs/moq-room/README.md index b6b1b19bf85..6d4ab418aaa 100644 --- a/rs/moq-room/README.md +++ b/rs/moq-room/README.md @@ -28,6 +28,9 @@ while let Some(event) = room.next().await { } ``` +The example awaits the wrappers behind the `async` feature, which is off by default; the +`poll_*` methods work without it. + Chat uses `moq-json::window` with ten seconds of history. Drive `Publisher::expire()` alongside incoming sends (or call `poll_expire` from a poll loop) so messages expire while idle. `Subscriber::recv` returns push/pop/skip diff --git a/rs/moq-room/src/chat.rs b/rs/moq-room/src/chat.rs index 55e018f3080..7b0805e71a3 100644 --- a/rs/moq-room/src/chat.rs +++ b/rs/moq-room/src/chat.rs @@ -108,18 +108,15 @@ impl Publisher { Ok(()) } - /// Poll until the oldest retained messages expire; drive this alongside incoming sends. + /// Prune the oldest retained messages once they expire. Drive this alongside incoming sends, + /// and again after each expiry while the publisher is open. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_expire(&mut self, waiter: &kio::Waiter) -> Poll> { self.deadline.set(self.expires.front().copied()); ready!(self.deadline.poll(waiter)); Poll::Ready(self.prune()) } - /// Wait until the oldest retained messages expire; call again while the publisher is open. - pub async fn expire(&mut self) -> Result<(), Error> { - kio::wait(|waiter| self.poll_expire(waiter)).await - } - /// Finish the track, preserving the final retained window for current readers. pub fn finish(mut self) -> Result<(), Error> { Ok(self.producer.finish()?) @@ -140,8 +137,11 @@ enum State { impl Subscriber { /// Subscribe to the newest retained window on a broadcast. + #[cfg(feature = "async")] pub async fn subscribe(broadcast: &broadcast::Consumer) -> Result { - Ok(Self::new(broadcast.track(TRACK)?.subscribe(None).await?)) + let track = broadcast.track(TRACK)?; + let subscribing = track.subscribe(None).into_inner(); + Ok(Self::new(kio::wait(|waiter| subscribing.poll_ok(waiter)).await?)) } /// Read window changes from an existing subscription, starting at the newest window. @@ -151,8 +151,10 @@ impl Subscriber { } } - /// Poll the next window change and when its frame was sent, returning errors separately from - /// clean completion. + /// The next window change and when its frame was sent, or `None` on clean completion. + /// + /// A message first seen in a later window's header carries that header's time, not its send time. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_recv(&mut self, waiter: &kio::Waiter) -> Poll>, Error>> { if let State::Positioning(pending) = &mut self.state { let track = pending.as_mut().expect("positioning holds the track"); @@ -171,13 +173,6 @@ impl Subscriber { }; Poll::Ready(Ok(ready!(consumer.poll_next(waiter))?)) } - - /// Wait for the next window change and when its frame was sent, or None on clean completion. - /// - /// A message first seen in a later window's header carries that header's time, not its send time. - pub async fn recv(&mut self) -> Result>, Error> { - kio::wait(|waiter| self.poll_recv(waiter)).await - } } #[cfg(test)] diff --git a/rs/moq-room/src/lib.rs b/rs/moq-room/src/lib.rs index 2dec939f982..a2199eaf212 100644 --- a/rs/moq-room/src/lib.rs +++ b/rs/moq-room/src/lib.rs @@ -14,6 +14,10 @@ //! extracted from hang.live's roster and iroh-live's announce-bus redesign of //! `iroh-rooms`. Gossip discovery, tickets, and 1:1 Call stay in iroh-live. //! Capture and encode stay in `moq-video` / `moq-audio`. +//! +//! The `poll_*` methods are always available. Enable the `async` feature for their generated +//! async wrappers (`next`, `expire`, and `recv`) and `chat::Subscriber::subscribe`; it is off +//! by default. #![warn(missing_docs)] diff --git a/rs/moq-room/src/room.rs b/rs/moq-room/src/room.rs index 6d6cf5aa675..14f4a4f6bc8 100644 --- a/rs/moq-room/src/room.rs +++ b/rs/moq-room/src/room.rs @@ -69,7 +69,8 @@ impl Room { } } - /// Poll for the next remote camera/screen (un)announce. + /// The next remote camera/screen (un)announce, or `None` once the announce stream ends. + #[cfg_attr(feature = "async", kio::wrapper)] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll> { loop { match self.announced.poll_next(waiter) { @@ -136,11 +137,6 @@ impl Room { } Poll::Pending } - - /// Wait for the next remote camera/screen (un)announce. - pub async fn next(&mut self) -> Option { - kio::wait(|waiter| self.poll_next(waiter)).await - } } #[cfg(test)] diff --git a/rs/moq-rtc/Cargo.toml b/rs/moq-rtc/Cargo.toml index 110497ebab7..0ea7ca60820 100644 --- a/rs/moq-rtc/Cargo.toml +++ b/rs/moq-rtc/Cargo.toml @@ -29,7 +29,7 @@ hang = { workspace = true } # Drives each str0m connection over its UDP socket. kio = { workspace = true, features = ["runtime"] } moq-ice = { workspace = true } -moq-mux = { workspace = true } +moq-mux = { workspace = true, features = ["async"] } moq-net = { workspace = true } # `RedactedUrl` keeps a dialed URL's query token and userinfo out of the logs. moq-tokio = { workspace = true } diff --git a/rs/moq-rtmp/Cargo.toml b/rs/moq-rtmp/Cargo.toml index d28c5a6e31a..cb9df64851c 100644 --- a/rs/moq-rtmp/Cargo.toml +++ b/rs/moq-rtmp/Cargo.toml @@ -38,7 +38,7 @@ bytes = { workspace = true } futures = { workspace = true } hang = { workspace = true } kio = { workspace = true, features = ["runtime"] } -moq-mux = { workspace = true } +moq-mux = { workspace = true, features = ["async"] } moq-net = { workspace = true } rand = { workspace = true } rind = { workspace = true } diff --git a/rs/moq-rtsp/Cargo.toml b/rs/moq-rtsp/Cargo.toml index ab5480bc5d6..a4d7589b122 100644 --- a/rs/moq-rtsp/Cargo.toml +++ b/rs/moq-rtsp/Cargo.toml @@ -39,6 +39,8 @@ url = { workspace = true } base64 = { workspace = true } # A clock only the stall tests move, so a loaded runner never fires a deadline early. kio = { path = "../kio", features = ["manual"] } +# The tests await the exported catalog and media. +moq-mux = { path = "../moq-mux", features = ["async"] } tokio = { workspace = true, features = ["full", "test-util"] } tracing-test = { workspace = true } workspace-hack = { path = "../workspace-hack" } diff --git a/rs/moq-srt/Cargo.toml b/rs/moq-srt/Cargo.toml index 03ed9a71ff7..842a550ae82 100644 --- a/rs/moq-srt/Cargo.toml +++ b/rs/moq-srt/Cargo.toml @@ -19,7 +19,7 @@ bytes = { workspace = true } futures = { workspace = true } # The media max age an ingest declares by default. hang = { workspace = true } -moq-mux = { workspace = true } +moq-mux = { workspace = true, features = ["async"] } moq-net = { workspace = true } # Pure-Rust SRT (no libsrt / ffmpeg). SRT carries MPEG-TS, demuxed by moq-mux. srt-tokio = "0.4" diff --git a/rs/moq-stats/Cargo.toml b/rs/moq-stats/Cargo.toml index 9115c6aeb7f..86e55e19e03 100644 --- a/rs/moq-stats/Cargo.toml +++ b/rs/moq-stats/Cargo.toml @@ -12,7 +12,13 @@ rust-version.workspace = true keywords = ["quic", "http3", "webtransport", "stats", "live"] categories = ["multimedia", "network-programming", "web-programming"] +[package.metadata.docs.rs] +features = ["async"] + [features] +default = [] +# Generated async wrappers for the poll methods, and `Publisher::run`. Enable where you `.await` them. +async = ["kio/async"] # Exposes one producer tick to the producer_tick bench, through the hidden # `produce::bench` module, and moq-net's hidden counter bump. Off by default # and hidden from the docs. Not a supported API. @@ -33,6 +39,10 @@ uuid = { workspace = true } [dev-dependencies] criterion = { package = "bench-harness", path = "../bench-harness" } futures = { workspace = true } +# Compile the generated wrappers used by the tests and benches. +moq-json = { path = "../moq-json", features = ["async"] } +moq-net = { path = "../moq-net", features = ["async"] } +moq-stats = { path = ".", features = ["async"] } tokio = { workspace = true, features = ["macros", "rt", "test-util", "time"] } workspace-hack = { path = "../workspace-hack" } diff --git a/rs/moq-stats/MIGRATE.md b/rs/moq-stats/MIGRATE.md index 22bbb22bba7..e69f55662c4 100644 --- a/rs/moq-stats/MIGRATE.md +++ b/rs/moq-stats/MIGRATE.md @@ -4,6 +4,18 @@ How to update your code for each breaking change, newest first. ## Unreleased +### Async wrappers are opt-in + +Enable the `async` Cargo feature to keep awaiting the readers' `next` and `Publisher::run`: +`moq-stats = { version = "...", features = ["async"] }`. This also applies with +`default-features = false`. The feature is off by default; without it, drive `poll_next` and +`Publisher::poll_run`. + +`Consumer::{totals, presence, prefix, group}` return their reader at once instead of awaiting +the subscription handshake: drop the `.await`. The handshake finishes on the reader's first +`next`, which fails with its error, such as `NotFound` for an unlisted group. +The stats tracks and their JSON are unchanged. + ### Group counters have their own epoch type `GroupFrame::epoch` is `moq_stats::Epoch` instead of `moq_net::Instance`. diff --git a/rs/moq-stats/src/aggregate.rs b/rs/moq-stats/src/aggregate.rs index 975a601ab06..efb25f95a34 100644 --- a/rs/moq-stats/src/aggregate.rs +++ b/rs/moq-stats/src/aggregate.rs @@ -170,8 +170,9 @@ pub struct TotalsConsumer { impl TotalsConsumer { /// The next merged frame, or `None` once the announce stream ends (the /// source origin went away). - pub async fn next(&mut self) -> Result> { - kio::wait(|waiter| self.inner.poll_next(waiter)).await + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_next(&mut self, waiter: &Waiter) -> Poll>> { + self.inner.poll_next(waiter) } } @@ -184,8 +185,9 @@ pub struct GroupConsumer { impl GroupConsumer { /// The next merged frame, each tier's [`Group`] keyed by tier label as in /// [`GroupFrame::tiers`], or `None` once the announce stream ends. - pub async fn next(&mut self) -> Result>> { - kio::wait(|waiter| self.inner.poll_next(waiter)).await + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_next(&mut self, waiter: &Waiter) -> Poll>>> { + self.inner.poll_next(waiter) } } @@ -197,8 +199,9 @@ pub struct PresenceConsumer { impl PresenceConsumer { /// The next merged frame, or `None` once the announce stream ends. - pub async fn next(&mut self) -> Result> { - kio::wait(|waiter| self.inner.poll_next(waiter)).await + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_next(&mut self, waiter: &Waiter) -> Poll>> { + self.inner.poll_next(waiter) } } @@ -210,8 +213,9 @@ pub struct PrefixConsumer { impl PrefixConsumer { /// The next merged frame, or `None` once the announce stream ends. - pub async fn next(&mut self) -> Result> { - kio::wait(|waiter| self.inner.poll_next(waiter)).await + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_next(&mut self, waiter: &Waiter) -> Poll>> { + self.inner.poll_next(waiter) } } diff --git a/rs/moq-stats/src/consume.rs b/rs/moq-stats/src/consume.rs index ba493f50ed3..90b7458953c 100644 --- a/rs/moq-stats/src/consume.rs +++ b/rs/moq-stats/src/consume.rs @@ -1,6 +1,9 @@ //! The consuming half: typed readers over one published stats broadcast. +use std::task::{Poll, ready}; + use moq_net::{Path, broadcast}; +use serde::de::DeserializeOwned; use crate::{ GroupFrame, PrefixFrame, PresenceFrame, Result, TotalsFrame, group_track, prefix_track, presence_track, @@ -41,47 +44,54 @@ impl Consumer { Self { broadcast } } - /// Subscribe to the totals track, awaiting the subscription handshake. - pub async fn totals(&self) -> Result { - let name = totals_track(); + /// Subscribe to the totals track. The first read finishes the subscription + /// handshake, failing with its error. + pub fn totals(&self) -> Result { Ok(TotalsConsumer { - inner: self.subscribe(&name).await?, + inner: self.subscribe(&totals_track())?, }) } - /// Subscribe to the presence track for `label`, awaiting the - /// subscription handshake. - pub async fn presence(&self, label: impl Into) -> Result { - let name = presence_track(label); + /// Subscribe to the presence track for `label`; see [`Self::totals`]. + pub fn presence(&self, label: impl Into) -> Result { Ok(PresenceConsumer { - inner: self.subscribe(&name).await?, + inner: self.subscribe(&presence_track(label))?, }) } - /// Subscribe to the prefix track for broadcast prefix `path`, awaiting the - /// subscription handshake. - pub async fn prefix(&self, path: impl Into) -> Result { - let name = prefix_track(path); + /// Subscribe to the prefix track for broadcast prefix `path`; see + /// [`Self::totals`]. + pub fn prefix(&self, path: impl Into) -> Result { Ok(PrefixConsumer { - inner: self.subscribe(&name).await?, + inner: self.subscribe(&prefix_track(path))?, }) } - /// Subscribe to group `group`'s track, awaiting the subscription handshake. - /// Fails with [`NotFound`](moq_net::Error::NotFound) unless the group is + /// Subscribe to group `group`'s track; see [`Self::totals`]. The first read + /// fails with [`NotFound`](moq_net::Error::NotFound) unless the group is /// listed. - pub async fn group(&self, group: impl Into) -> Result { + pub fn group(&self, group: impl Into) -> Result { Ok(GroupConsumer { - inner: self.subscribe(&group_track(group)).await?, + inner: self.subscribe(&group_track(group))?, }) } - async fn subscribe(&self, name: &str) -> Result> { + fn subscribe(&self, name: &str) -> Result> { + // The reader subscribes itself, so a refused subscription fails its first read. let track = self.broadcast.track(name)?; - // The reader subscribes itself; resolving the info rides on that subscription. - let reader = moq_json::snapshot::Consumer::new(track.clone(), Default::default()); - track.info().await?; - Ok(reader) + Ok(Reader(Box::new(moq_json::snapshot::Consumer::new( + track, + Default::default(), + )))) + } +} + +/// A snapshot reader yielding only the latest value. +struct Reader(Box>); + +impl Reader { + fn poll_latest(&mut self, waiter: &kio::Waiter) -> Poll>> { + Poll::Ready(Ok(ready!(self.0.poll_latest(waiter))?.map(|value| value.value))) } } @@ -89,31 +99,33 @@ impl Consumer { /// intermediate frames a slow reader missed are collapsed, which is safe /// because the counters are cumulative. pub struct TotalsConsumer { - inner: moq_json::snapshot::Consumer, + inner: Reader, } impl TotalsConsumer { /// The next frame, or `None` once the track ends (the producer went away). - pub async fn next(&mut self) -> Result> { - Ok(self.inner.latest().await?.map(|value| value.value)) + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { + self.inner.poll_latest(waiter) } } /// A typed reader over one label's presence track; see [`TotalsConsumer`]. pub struct PresenceConsumer { - inner: moq_json::snapshot::Consumer, + inner: Reader, } impl PresenceConsumer { /// The next frame, or `None` once the track ends (the producer went away). - pub async fn next(&mut self) -> Result> { - Ok(self.inner.latest().await?.map(|value| value.value)) + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { + self.inner.poll_latest(waiter) } } /// A typed reader over one broadcast prefix's track; see [`TotalsConsumer`]. pub struct PrefixConsumer { - inner: moq_json::snapshot::Consumer, + inner: Reader, } impl PrefixConsumer { @@ -121,14 +133,15 @@ impl PrefixConsumer { /// Fails with /// [`FrameTooLarge`](moq_net::Error::FrameTooLarge) once the frame outgrows /// the cache. - pub async fn next(&mut self) -> Result> { - Ok(self.inner.latest().await?.map(|value| value.value)) + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { + self.inner.poll_latest(waiter) } } /// A typed reader over one group's track; see [`TotalsConsumer`]. pub struct GroupConsumer { - inner: moq_json::snapshot::Consumer, + inner: Reader, } impl GroupConsumer { @@ -136,8 +149,9 @@ impl GroupConsumer { /// or the group lingered out of the listing). Fails with /// [`FrameTooLarge`](moq_net::Error::FrameTooLarge) once the frame outgrows /// the cache. - pub async fn next(&mut self) -> Result> { - Ok(self.inner.latest().await?.map(|value| value.value)) + #[cfg_attr(feature = "async", kio::wrapper)] + pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll>> { + self.inner.poll_latest(waiter) } } @@ -268,7 +282,7 @@ mod tests { let broadcast = announced(&origin).await; let consumer = Consumer::new(broadcast.consume(), Config::new()); - let mut totals = consumer.totals().await.expect("subscribe totals"); + let mut totals = consumer.totals().expect("subscribe totals"); let frame = totals.next().await.expect("read").expect("frame"); assert_eq!(frame[""].publisher.bytes, 42); @@ -280,10 +294,7 @@ mod tests { assert_eq!(frame[""].sessions.active(), 1); // At depth 0 the one group is the empty one. - let mut group = consumer - .group(moq_net::Path::from_static("")) - .await - .expect("subscribe group"); + let mut group = consumer.group(moq_net::Path::from_static("")).expect("subscribe group"); let frame = group.next().await.expect("read").expect("frame"); let tier = &frame.tiers[""]; assert_eq!(tier.publisher.bytes, 50); @@ -291,19 +302,23 @@ mod tests { assert_eq!(tier.labels["acme"].active(), 1); // A presence track is produced once requested, so drive a drain. - let mut presence = std::pin::pin!(consumer.presence(moq_net::Path::from_static("acme"))); - assert!(futures::poll!(presence.as_mut()).is_pending()); + let mut presence = consumer + .presence(moq_net::Path::from_static("acme")) + .expect("subscribe presence"); + let mut next = std::pin::pin!(presence.next()); + assert!(futures::poll!(next.as_mut()).is_pending()); drive_tick().await; - let mut presence = presence.await.expect("subscribe presence"); - let frame = presence.next().await.expect("read").expect("frame"); + let frame = next.await.expect("read").expect("frame"); assert_eq!(frame.get("").expect("default tier").active(), 1); // So is a prefix track. - let mut prefix = std::pin::pin!(consumer.prefix(moq_net::Path::from_static("foo"))); - assert!(futures::poll!(prefix.as_mut()).is_pending()); + let mut prefix = consumer + .prefix(moq_net::Path::from_static("foo")) + .expect("subscribe prefix"); + let mut next = std::pin::pin!(prefix.next()); + assert!(futures::poll!(next.as_mut()).is_pending()); drive_tick().await; - let mut prefix = prefix.await.expect("subscribe prefix"); - let frame = prefix.next().await.expect("read").expect("frame"); + let frame = next.await.expect("read").expect("frame"); let tier = frame.get("").expect("default tier"); assert_eq!(tier.publisher.bytes, 50); assert_eq!(tier.children["bar"].publisher.bytes, 50); diff --git a/rs/moq-stats/src/lib.rs b/rs/moq-stats/src/lib.rs index 1aa06c1f058..c20ad0c6c68 100644 --- a/rs/moq-stats/src/lib.rs +++ b/rs/moq-stats/src/lib.rs @@ -34,6 +34,9 @@ //! [`group_track`] and [`parse_group_track`]. The full contract (paths, //! tracks, and counter semantics) is at //! . +//! +//! The `poll_*` methods are always available. Enable the `async` feature for their generated +//! async wrappers (the readers' `next`) and [`Publisher`]'s `run`; it is off by default. pub mod aggregate; pub mod consume; diff --git a/rs/moq-stats/src/produce.rs b/rs/moq-stats/src/produce.rs index f98bbdd4f6b..5069c1b23b2 100644 --- a/rs/moq-stats/src/produce.rs +++ b/rs/moq-stats/src/produce.rs @@ -4,13 +4,12 @@ use std::collections::HashMap; use std::sync::{Arc, Weak}; use std::time::Duration; -use std::task::Poll; +use std::task::{Poll, ready}; use kio::clock::Instant; use moq_net::stats::{Presence, Registry, Report, Tier, Traffic, TrafficEntry}; use moq_net::{Path, broadcast, kio, origin, track}; use serde::Serialize; -use tracing::Instrument as _; use crate::{ Epoch, Prefix, Totals, Usage, parse_prefix_track, parse_presence_track, prefix_track, presence_track, totals_track, @@ -180,7 +179,7 @@ impl Producer { registry: Registry::disabled(), _keepalive: None, }; - return (producer, Publisher(None)); + return (producer, Publisher(Run::Done)); }; // The prefix is a literal path, so its subtree claim cannot fail. @@ -196,7 +195,11 @@ impl Producer { linger, interval, }; - let publisher = Publisher(Some((task, Arc::downgrade(&keepalive), tracing::Span::current()))); + let publisher = Publisher(Run::Idle( + Box::new(task), + Arc::downgrade(&keepalive), + tracing::Span::current(), + )); let producer = Self { registry, _keepalive: Some(keepalive), @@ -214,61 +217,108 @@ impl Producer { /// The publish task of one [`Producer`], returned by [`Producer::new`]. #[must_use = "nothing is published until the publisher runs"] -pub struct Publisher(Option<(Task, Weak, tracing::Span)>); +pub struct Publisher(Run); + +/// Where a [`Publisher`] is in its run. +enum Run { + /// Not polled yet: announces on the first poll. + Idle(Box, Weak, tracing::Span), + Running(Box), + Done, +} + +/// A running publisher's drain and the interval it drains on. +struct Running { + drain: Drain, + ticker: kio::clock::Interval, + weak: Weak, + span: tracing::Span, +} impl Publisher { /// Publish until the last [`Producer`] clone drops; at once for a no-op producer. - pub async fn run(self) { - if let Some((task, weak, span)) = self.0 { - task.run(weak).instrument(span).await; + /// + /// The broadcast is announced on the first poll and ends when this resolves or + /// the publisher drops. + pub fn poll_run(&mut self, waiter: &kio::Waiter) -> Poll<()> { + if let Run::Idle(..) = self.0 { + let Run::Idle(task, weak, span) = std::mem::replace(&mut self.0, Run::Done) else { + unreachable!("matched above"); + }; + if let Some(running) = Running::start(*task, weak, span) { + self.0 = Run::Running(Box::new(running)); + } } + let Run::Running(running) = &mut self.0 else { + return Poll::Ready(()); + }; + if running.poll_run(waiter).is_pending() { + return Poll::Pending; + } + // Dropping the drain ends the broadcast, as dropping the publisher does. + self.0 = Run::Done; + Poll::Ready(()) } -} -/// Everything the publish task owns. -struct Task { - registry: Registry, - origin: origin::Producer, - prefix: Path, - node: Option, - depth: usize, - linger: Duration, - interval: Duration, + /// Publish until the last [`Producer`] clone drops; at once for a no-op producer. + /// See [`Self::poll_run`]. + #[cfg(feature = "async")] + pub async fn run(mut self) { + kio::wait(|waiter| self.poll_run(waiter)).await + } } -impl Task { - /// Publishes the stats broadcast and writes a frame per drain. Runs until - /// every [`Producer`] clone is dropped (`weak.upgrade()` returns `None`). - async fn run(self, weak: Weak) { +impl Running { + /// Announce the stats broadcast, unless every [`Producer`] clone is already + /// gone or the origin refused it. + fn start(task: Task, weak: Weak, span: tracing::Span) -> Option { // A publisher first run after its last producer dropped announces nothing, // so it cannot replace a newer producer's broadcast on the same path. if weak.strong_count() == 0 { - return; - } - let interval = self.interval; - let Some(mut drain) = Drain::new(self) else { - return; - }; - - let mut ticker = kio::clock::interval(interval); - ticker.set_missed_tick_behavior(kio::clock::MissedTickBehavior::Delay); + return None; + } + let (drain, ticker) = span.in_scope(|| { + let interval = task.interval; + let drain = Drain::new(task)?; + let mut ticker = kio::clock::interval(interval); + ticker.set_missed_tick_behavior(kio::clock::MissedTickBehavior::Delay); + Some((drain, ticker)) + })?; + Some(Self { + drain, + ticker, + weak, + span, + }) + } + /// Write a frame per drain, until every [`Producer`] clone is dropped. + fn poll_run(&mut self, waiter: &kio::Waiter) -> Poll<()> { + let _entered = self.span.enter(); loop { // Scheduled instants are at least an interval apart, even after a late tick, // so each drain gets an interval of its own to open its groups in. - let scheduled = ticker.tick().await; - - // Dropping the drain ends the broadcast, as cancelling this future does. - if weak.upgrade().is_none() { - return; + let scheduled = ready!(self.ticker.poll_tick(&mut waiter.context())); + if self.weak.upgrade().is_none() { + return Poll::Ready(()); } - - drain.collect(); - drain.publish(scheduled); + self.drain.collect(); + self.drain.publish(scheduled); } } } +/// Everything the publish task owns. +struct Task { + registry: Registry, + origin: origin::Producer, + prefix: Path, + node: Option, + depth: usize, + linger: Duration, + interval: Duration, +} + /// The publish task's state, kept across drains so a steady-state drain /// reuses its buffers instead of allocating per entry. struct Drain { @@ -3010,9 +3060,7 @@ mod tests { drop(acme); let lingering = start + Duration::from_secs(9); let stats = crate::Consumer::new(broadcast.clone(), Default::default()); - let mut acme = pumping_at(&mut drain, lingering, stats.group(moq_net::Path::from_static("acme"))) - .await - .unwrap(); + let mut acme = stats.group(moq_net::Path::from_static("acme")).unwrap(); let mut frame = pumping_at(&mut drain, lingering, acme.next()).await.unwrap().unwrap(); while !frame.tiers[""].labels.is_empty() { frame = pumping_at(&mut drain, lingering, acme.next()).await.unwrap().unwrap(); @@ -3026,9 +3074,7 @@ mod tests { ("acme".into(), false) ); assert!(pumping_at(&mut drain, lingered, acme.next()).await.unwrap().is_none()); - let mut beta_reader = pumping_at(&mut drain, lingered, stats.group(moq_net::Path::from_static("beta"))) - .await - .unwrap(); + let mut beta_reader = stats.group(moq_net::Path::from_static("beta")).unwrap(); let frame = pumping_at(&mut drain, lingered, beta_reader.next()) .await .unwrap() @@ -3068,7 +3114,7 @@ mod tests { .tier(Tier::default()) .session(moq_net::Path::from_static("acme/live")); drive_tick().await; - let mut reader = stats.group(moq_net::Path::from_static("acme")).await.expect("listed"); + let mut reader = stats.group(moq_net::Path::from_static("acme")).expect("listed"); let first = reader.next().await.unwrap().unwrap(); assert_eq!(first.tiers[""].sessions.sessions_started, 1); let old_sequence = sequence(&broadcast).await; @@ -3089,7 +3135,6 @@ mod tests { drive_tick().await; let second = stats .group(moq_net::Path::from_static("acme")) - .await .expect("listed again") .next() .await @@ -3139,12 +3184,8 @@ mod tests { let _big = registry .tier(Tier::default()) .session(moq_net::Path::from_static("big")); - let mut acme = pumping(&mut drain, stats.group(moq_net::Path::from_static("acme"))) - .await - .unwrap(); - let mut big = pumping(&mut drain, stats.group(moq_net::Path::from_static("big"))) - .await - .unwrap(); + let mut acme = stats.group(moq_net::Path::from_static("acme")).unwrap(); + let mut big = stats.group(moq_net::Path::from_static("big")).unwrap(); assert!(pumping(&mut drain, acme.next()).await.unwrap().is_some()); assert!(pumping(&mut drain, big.next()).await.unwrap().is_some()); @@ -3187,7 +3228,6 @@ mod tests { let frame = stats .group(moq_net::Path::from_static("acme")) - .await .unwrap() .next() .await @@ -3201,7 +3241,6 @@ mod tests { let frame = stats .group(moq_net::Path::from_static("beta")) - .await .unwrap() .next() .await @@ -3212,7 +3251,6 @@ mod tests { let frame = stats .group(moq_net::Path::from_static("feed")) - .await .unwrap() .next() .await diff --git a/rs/moq-tokio/Cargo.toml b/rs/moq-tokio/Cargo.toml index c8e2f335e8a..a7abbcc4092 100644 --- a/rs/moq-tokio/Cargo.toml +++ b/rs/moq-tokio/Cargo.toml @@ -16,7 +16,7 @@ categories = ["multimedia", "network-programming", "web-programming"] default = ["net", "quic", "websocket", "tcp", "uds"] # MoQ sessions over moq-net: `Client`, `Server`, `origin`, and the `connect`, `listen`, # and `transport` modules. Every transport feature turns it on, since none works without it. -net = ["dep:moq-net", "_certs"] +net = ["dep:moq-net", "moq-net/async", "_certs"] # The async QUIC driver alone (`quic::driver`): moq-quic on tokio, without moq-net or # the MoQ backend. Includes rind cryptography. driver = ["_certs", "dep:moq-congestion", "dep:moq-quic", "dep:rustc-hash"] @@ -64,7 +64,7 @@ http = { version = "1", optional = true } humantime = { workspace = true } humantime-serde = { workspace = true } if-addrs = { workspace = true, optional = true } -kio = { workspace = true, features = ["runtime"] } +kio = { workspace = true, features = ["async", "runtime"] } mdns-sd = { workspace = true, optional = true } moq-congestion = { workspace = true, optional = true } # rind's rustls provider and the `UnsupportedCpu` error, re-exported as `crypto`. Needs no moq-net. diff --git a/rs/moq-tokio/src/connection.rs b/rs/moq-tokio/src/connection.rs index 337cc4062ec..e7ceec5a7a8 100644 --- a/rs/moq-tokio/src/connection.rs +++ b/rs/moq-tokio/src/connection.rs @@ -1504,10 +1504,14 @@ impl Connection { Ok(self) } - /// Poll for the next connection status change since this handle last reported one. + /// The connection status, once it changes from what this handle last reported, or the error + /// once the loop has stopped (the give-up error, or a generic one when the handle is dropped). /// - /// `Ready(Ok(status))` on a change, `Ready(Err)` once the loop has stopped (the give-up error, - /// or a generic one when the handle is dropped), `Pending` otherwise. + /// The loop moves between `Connected`, `Disconnected`, and `Migrating` (a GOAWAY handover, where + /// the old session is still serving), so successive calls report changes rather than a fixed + /// alternation; a status that flips and flips back before the caller polls is reported once. + /// This tracks the *current* state, not every edge. + #[kio::wrapper] pub fn poll_status(&mut self, waiter: &kio::Waiter) -> Poll> { let last = self.last_reported; let status = match ready!(self.state.poll(waiter, |state| match state.status { @@ -1522,17 +1526,6 @@ impl Connection { Poll::Ready(Ok(status)) } - /// Wait until the connection status changes from what this handle last reported. - /// - /// Returns the current [`Status`]. The loop moves between `Connected`, - /// `Disconnected`, and `Migrating` (a GOAWAY handover, where the old session is - /// still serving), so successive calls report changes rather than a fixed - /// alternation; a status that flips and flips back before the caller polls is - /// reported once. This tracks the *current* state, not every edge. - pub async fn status(&mut self) -> crate::Result { - kio::wait(|waiter| self.poll_status(waiter)).await - } - /// Whether a session is currently connected. /// /// The synchronous read behind [`status`](Self::status), for callers that just want the current @@ -1616,11 +1609,10 @@ impl Connection { self.viewers.clone() } - /// Poll whether the connection loop has stopped. - /// - /// `Ready(Err)` if it permanently gave up (reconnect timeout exceeded, an auth - /// rejection, or the session ending in one-shot mode), `Ready(Ok(()))` if stopped - /// by dropping the handle, `Pending` while it's still running. + /// Resolves once the connection loop stops: an error if it permanently gave up (reconnect + /// timeout exceeded, an auth rejection, or the session ending in one-shot mode), `Ok(())` if + /// stopped by dropping the handle. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll> { ready!(self.state.poll_closed(waiter)); Poll::Ready(match &self.state.read().error { @@ -1629,11 +1621,6 @@ impl Connection { }) } - /// Wait until the connection loop stops. - pub async fn closed(&self) -> crate::Result<()> { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - /// The auth handle every session of this connection is built on: its tokens are /// presented on each reconnect, so a route naming one keeps it, and its grant follows /// the live session. It ends with the connection: the abort's or [`Error::MoqNet`]'s diff --git a/rs/moq-tokio/src/quic/demux.rs b/rs/moq-tokio/src/quic/demux.rs index fc2af50649a..5cc4c6b4518 100644 --- a/rs/moq-tokio/src/quic/demux.rs +++ b/rs/moq-tokio/src/quic/demux.rs @@ -147,8 +147,9 @@ pub struct WebRtc { } impl WebRtc { - /// Poll for the next datagram; `None` once every [`Demux`] clone is gone and - /// every listening socket it split is released. + /// The next datagram, or `None` once every [`Demux`] clone is gone and every listening socket + /// it split is released. + #[kio::wrapper] pub fn poll_recv(&self, waiter: &kio::Waiter) -> Poll> { if let Poll::Ready(mut queue) = self.shared.queue.poll(waiter, |queue| match queue.is_empty() { true => Poll::Pending, @@ -163,13 +164,7 @@ impl WebRtc { } } - /// Receive the next datagram; `None` once the listener and every [`Demux`] - /// clone are gone; see [`Self::poll_recv`]. - pub async fn recv(&self) -> Option { - kio::wait(|waiter| self.poll_recv(waiter)).await - } - - /// Poll to send `transmit` from the shared port, with GSO and ECN as it asks. + /// Send `transmit` from the shared port, with GSO and ECN as it asks. /// /// Its `src_ip` is kept as given. On a wildcard bind with several local /// addresses, set it to the [`Datagram::local`] being answered; `None` lets the @@ -178,6 +173,7 @@ impl WebRtc { /// A datagram sent once the listener has released its socket is lost, as UDP /// allows. Errors are those of the endpoint's own sends: an unreachable destination /// comes back, as [`AsyncUdpSocket::try_send`] describes, and anything else is logged. + #[kio::wrapper] pub fn poll_send(&self, waiter: &kio::Waiter, transmit: &Transmit<'_>) -> Poll> { // Any member of a reuseport group answers from the same port. let sockets = self.sockets.lock().unwrap(); @@ -195,11 +191,6 @@ impl WebRtc { socket.slot.poll_send(&transmit, &mut waiter.context()) } - /// Send `transmit` from the shared port; see [`Self::poll_send`]. - pub async fn send(&self, transmit: &Transmit<'_>) -> io::Result<()> { - kio::wait(|waiter| self.poll_send(waiter, transmit)).await - } - /// Route every datagram from `remote` to this feed, whatever its first byte, /// until the returned [`Pin`] drops. pub fn pin(&self, remote: SocketAddr) -> Pin { diff --git a/rs/moq-tokio/src/quic/driver/connection.rs b/rs/moq-tokio/src/quic/driver/connection.rs index 234f23b538a..2c57bc7c3df 100644 --- a/rs/moq-tokio/src/quic/driver/connection.rs +++ b/rs/moq-tokio/src/quic/driver/connection.rs @@ -258,7 +258,7 @@ impl Drop for ConnectionDriver { pub struct Connection(ConnectionRef); impl Connection { - /// Initiate a new outgoing unidirectional stream, scheduled by `options`. + /// Open a new outgoing unidirectional stream, scheduled by `options`. /// /// Streams are cheap and instantaneous to open unless blocked by flow control. As a /// consequence, the peer won't be notified that a stream has been opened until the stream is @@ -266,11 +266,7 @@ impl Connection { /// /// # Panics /// - when `options.group` belongs to another connection - pub async fn open_uni(&self, options: send::Options<'_>) -> Result { - kio::wait(|waiter| self.poll_open_uni(waiter, options)).await - } - - /// See [`Self::open_uni`]. + #[kio::wrapper] pub fn poll_open_uni(&self, waiter: &kio::Waiter, options: send::Options<'_>) -> Poll> { let (conn, id) = ready!(poll_open(waiter, &self.0, options, Dir::Uni))?; Poll::Ready(Ok(send::Stream::new(conn, id))) @@ -282,7 +278,7 @@ impl Connection { send::Group::new(self.0.0.clone(), priority) } - /// Initiate a new outgoing bidirectional stream, its sending half scheduled by `options`. + /// Open a new outgoing bidirectional stream, its sending half scheduled by `options`. /// /// Streams are cheap and instantaneous to open unless blocked by flow control. As a /// consequence, the peer won't be notified that a stream has been opened until the stream is @@ -293,11 +289,7 @@ impl Connection { /// /// # Panics /// - when `options.group` belongs to another connection - pub async fn open_bi(&self, options: send::Options<'_>) -> Result<(send::Stream, recv::Stream), Error> { - kio::wait(|waiter| self.poll_open_bi(waiter, options)).await - } - - /// See [`Self::open_bi`]. + #[kio::wrapper] pub fn poll_open_bi( &self, waiter: &kio::Waiter, @@ -307,22 +299,18 @@ impl Connection { Poll::Ready(Ok((send::Stream::new(conn.clone(), id), recv::Stream::new(conn, id)))) } - /// Accept the next incoming uni-directional stream + /// The next incoming unidirectional stream. /// /// # Cancel safety /// /// This method is cancellation safe. If this does not resolve, no streams were consumed. - pub async fn accept_uni(&self) -> Result { - kio::wait(|waiter| self.poll_accept_uni(waiter)).await - } - - /// See [`Self::accept_uni`]. + #[kio::wrapper] pub fn poll_accept_uni(&self, waiter: &kio::Waiter) -> Poll> { let (conn, id) = ready!(poll_accept(waiter, &self.0, Dir::Uni))?; Poll::Ready(Ok(recv::Stream::new(conn, id))) } - /// Accept the next incoming bidirectional stream + /// The next incoming bidirectional stream. /// /// **Important Note**: The `Connection` that calls [`open_bi()`] must write to its [`send::Stream`] /// before the other `Connection` is able to `accept_bi()`. Calling [`open_bi()`] then @@ -334,11 +322,7 @@ impl Connection { /// # Cancel safety /// /// This method is cancellation safe. If this does not resolve, no streams were consumed. - pub async fn accept_bi(&self) -> Result<(send::Stream, recv::Stream), Error> { - kio::wait(|waiter| self.poll_accept_bi(waiter)).await - } - - /// See [`Self::accept_bi`]. + #[kio::wrapper] pub fn poll_accept_bi(&self, waiter: &kio::Waiter) -> Poll> { let (conn, id) = ready!(poll_accept(waiter, &self.0, Dir::Bi))?; Poll::Ready(Ok((send::Stream::new(conn.clone(), id), recv::Stream::new(conn, id)))) @@ -360,16 +344,12 @@ impl Connection { Poll::Ready(datagram.map_err(|err| state.lost_as(err))) } - /// Wait for the connection to be closed for any reason + /// The error the connection closed with, for any reason. /// /// Despite the return type's name, closed connections are often not an error condition at the /// application layer. Cases that might be routine include [`Error::LocallyClosed`] /// and [`Error::ApplicationClosed`]. - pub async fn closed(&self) -> Error { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - - /// Wait for the connection to close, returning why. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll { let mut state = self.0.state.lock(); let err = ready!(state.inner.poll_closed(waiter)); diff --git a/rs/moq-tokio/src/quic/web/recv.rs b/rs/moq-tokio/src/quic/web/recv.rs index 70027b2a510..262c75c5483 100644 --- a/rs/moq-tokio/src/quic/web/recv.rs +++ b/rs/moq-tokio/src/quic/web/recv.rs @@ -83,12 +83,8 @@ impl Stream { self.stop_mut(code); } - /// Wait until the stream ends after its bytes are read, the peer resets it, or the session closes. - pub async fn closed(&mut self) -> Result<(), Error> { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - - /// Poll until the stream ends after its bytes are read, the peer resets it, or the session closes. + /// Resolves once the stream ends after its bytes are read, the peer resets it, or the session closes. + #[kio::wrapper] pub fn poll_closed(&mut self, waiter: &kio::Waiter) -> Poll> { let close = &self.close; let Some(inner) = &mut self.inner else { diff --git a/rs/moq-tokio/src/quic/web/send.rs b/rs/moq-tokio/src/quic/web/send.rs index f72f0b88681..9422c39a4a0 100644 --- a/rs/moq-tokio/src/quic/web/send.rs +++ b/rs/moq-tokio/src/quic/web/send.rs @@ -112,7 +112,9 @@ impl Stream { self.state = State::Reset(stopped); } - /// Poll until all data was acknowledged, the peer stopped the stream, or the session closed. + /// Resolves once all data was acknowledged, the peer stopped the stream (with its code), or the + /// session closed. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll, Error>> { let result = match &self.state { State::Open(stream) => stream.poll_closed(waiter), @@ -129,11 +131,6 @@ impl Stream { }) } - /// Wait until all data was acknowledged, the peer stopped the stream, or the session closed. - pub async fn closed(&self) -> Result, Error> { - kio::wait(|w| self.poll_closed(w)).await - } - /// Finish sending, retaining a handle to acknowledgements and peer cancellation. pub fn finish(mut self) -> Finished { self.finish_mut(); @@ -276,14 +273,11 @@ pub struct Finished { park: kio::Park, } impl Finished { - /// Poll until the peer acknowledges or stops the stream, or the session closes. + /// Resolves once the peer acknowledges or stops the stream (with its code), or the session closes. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll, Error>> { self.stream.poll_closed(waiter) } - /// Wait until the peer acknowledges or stops the stream, or the session closes. - pub async fn closed(&self) -> Result, Error> { - self.stream.closed().await - } /// Reset the stream while still delivering its WebTransport header. pub fn reset(self, code: u32) { self.stream.reset(code); diff --git a/rs/moq-tokio/src/quic/web/session.rs b/rs/moq-tokio/src/quic/web/session.rs index 7fccf5c88ef..47b2b6f25ec 100644 --- a/rs/moq-tokio/src/quic/web/session.rs +++ b/rs/moq-tokio/src/quic/web/session.rs @@ -212,27 +212,8 @@ impl Session { .await } - /// Accept the next unidirectional stream after its WebTransport header. - pub async fn accept_uni(&self) -> Result { - kio::wait(|waiter| self.poll_accept_uni(waiter)).await - } - - /// Accept the next bidirectional stream after its WebTransport header. - pub async fn accept_bi(&self) -> Result<(SendStream, RecvStream), SessionError> { - kio::wait(|waiter| self.poll_accept_bi(waiter)).await - } - - /// Open a unidirectional stream scheduled by `options`. - pub async fn open_uni(&self, options: send::Options<'_>) -> Result { - kio::wait(|waiter| self.poll_open_uni(waiter, options)).await - } - - /// Open a bidirectional stream scheduled by `options`. - pub async fn open_bi(&self, options: send::Options<'_>) -> Result<(SendStream, RecvStream), SessionError> { - kio::wait(|waiter| self.poll_open_bi(waiter, options)).await - } - - /// Poll for the next unidirectional stream after its WebTransport header. + /// The next unidirectional stream, after its WebTransport header. + #[kio::wrapper] pub fn poll_accept_uni(&self, waiter: &Waiter) -> Poll> { if let Some(accept) = &self.accept { poll_accept_uni_shared(accept, waiter).map_err(|e| self.map_error(e)) @@ -244,7 +225,8 @@ impl Session { } } - /// Poll for the next bidirectional stream after its WebTransport header. + /// The next bidirectional stream, after its WebTransport header. + #[kio::wrapper] pub fn poll_accept_bi(&self, waiter: &Waiter) -> Poll> { if let Some(accept) = &self.accept { poll_accept_bi_shared(accept, waiter).map_err(|e| self.map_error(e)) @@ -261,7 +243,8 @@ impl Session { } } - /// Poll to open a unidirectional stream and write its WebTransport header. + /// Open a unidirectional stream scheduled by `options`, writing its WebTransport header. + #[kio::wrapper] pub fn poll_open_uni(&self, waiter: &Waiter, options: send::Options<'_>) -> Poll> { self.opening.poll(waiter, |opening, cx| { if opening.uni.is_none() { @@ -286,7 +269,8 @@ impl Session { }) } - /// Poll to open a bidirectional stream and write its WebTransport header. + /// Open a bidirectional stream scheduled by `options`, writing its WebTransport header. + #[kio::wrapper] pub fn poll_open_bi( &self, waiter: &Waiter, @@ -319,15 +303,11 @@ impl Session { }) } - /// Receive the next datagram with its WebTransport session header removed. - pub async fn recv_datagram(&self) -> Result { - kio::wait(|waiter| self.poll_recv_datagram(waiter)).await - } - - /// Poll for a datagram with its WebTransport session header removed. + /// The next datagram, with its WebTransport session header removed. /// /// A datagram that does not name this session, such as another session's on the same /// connection, is dropped: datagrams are unreliable, so a stray one must not end the session. + #[kio::wrapper] pub fn poll_recv_datagram(&self, waiter: &Waiter) -> Poll> { loop { let mut datagram = @@ -498,19 +478,6 @@ impl Session { } } - /// Wait until the session is closed, returning the error. See [`crate::quic::driver::connection::Connection::closed`]. - /// - /// If the peer sent a `CloseWebTransportSession` capsule, the returned error will be - /// [`WebTransportError::Closed`] with the code and reason from the capsule. - /// - /// Unlike [`crate::quic::driver::connection::Connection::closed`], this does **not** return early when - /// [`close()`](Self::close) has been called. It waits for the underlying QUIC - /// connection to shut down, ensuring the `CloseWebTransportSession` capsule has - /// been delivered. Use [`close_reason()`](Self::close_reason) for a non-blocking check. - pub async fn closed(&self) -> SessionError { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - /// Return why the session was closed, or None if it's not closed. See /// [`crate::quic::driver::connection::Connection::close_reason`]. pub fn close_reason(&self) -> Option { @@ -577,7 +544,16 @@ impl Session { } } - /// Poll until the session's connection has closed. + /// The error the session closed with. See [`crate::quic::driver::connection::Connection::closed`]. + /// + /// If the peer sent a `CloseWebTransportSession` capsule, the returned error will be + /// [`WebTransportError::Closed`] with the code and reason from the capsule. + /// + /// Unlike [`crate::quic::driver::connection::Connection::closed`], this does **not** return early when + /// [`close()`](Self::close) has been called. It waits for the underlying QUIC + /// connection to shut down, ensuring the `CloseWebTransportSession` capsule has + /// been delivered. Use [`close_reason()`](Self::close_reason) for a non-blocking check. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &Waiter) -> Poll { self.conn.poll_closed(waiter).map(|e| self.map_error(e)) } diff --git a/rs/moq-tokio/src/webrtc/channel.rs b/rs/moq-tokio/src/webrtc/channel.rs index 29efc11f548..97fa6869b54 100644 --- a/rs/moq-tokio/src/webrtc/channel.rs +++ b/rs/moq-tokio/src/webrtc/channel.rs @@ -187,17 +187,12 @@ pub(crate) struct Reader { impl Reader { /// The next record the driver read off the channel. + #[cfg_attr(test, kio::wrapper)] pub fn poll_recv(&mut self, waiter: &kio::Waiter) -> Poll> { let record = ready!(self.rx.poll_recv(&mut waiter.context())).unwrap_or(Err(qmux::Error::Closed))?; self.unread.fetch_sub(record.len() as u64, Ordering::Relaxed); Poll::Ready(Ok(record)) } - - /// Wait for [`Self::poll_recv`]. - #[cfg(test)] - pub async fn recv(&mut self) -> Result { - kio::wait(|waiter| self.poll_recv(waiter)).await - } } #[cfg(test)] diff --git a/rs/moq-transcode/Cargo.toml b/rs/moq-transcode/Cargo.toml index 8bbfebd8399..a2fe4db53ae 100644 --- a/rs/moq-transcode/Cargo.toml +++ b/rs/moq-transcode/Cargo.toml @@ -35,7 +35,7 @@ v4l2 = ["moq-video/v4l2"] [dependencies] bytes = { workspace = true } hang = { workspace = true } -kio = { workspace = true } +kio = { workspace = true, features = ["async"] } moq-mux = { workspace = true } moq-net = { workspace = true } moq-video = { workspace = true } diff --git a/rs/moq-transcode/src/active.rs b/rs/moq-transcode/src/active.rs index a9789b0fb71..17c8eab48e2 100644 --- a/rs/moq-transcode/src/active.rs +++ b/rs/moq-transcode/src/active.rs @@ -365,14 +365,7 @@ struct Seen { impl Consumer { /// The next rendition to report, or `None` once the transcoder is gone. - pub async fn next(&mut self) -> Option { - kio::wait(|waiter| self.poll_next(waiter)).await - } - - /// Poll for the next rendition to report, without blocking. - /// - /// Returns `Poll::Ready(Some(_))` for an update, `Poll::Ready(None)` once the - /// transcoder is gone, or `Poll::Pending` after registering `waiter`. + #[kio::wrapper] pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll> { let update = { let seen = &mut self.seen; diff --git a/rs/moq-uring/Cargo.toml b/rs/moq-uring/Cargo.toml index 67a133e14ff..0b8966e2dac 100644 --- a/rs/moq-uring/Cargo.toml +++ b/rs/moq-uring/Cargo.toml @@ -30,12 +30,12 @@ futures = "0.3" http = "1" io-uring = { workspace = true } # The worker is a kio clock backend: its turn's instant is the ambient clock. -kio = { workspace = true, features = ["runtime"] } +kio = { workspace = true, features = ["async", "runtime"] } libc = { workspace = true } moq-congestion = { workspace = true, optional = true } # rind's rustls provider and the `UnsupportedCpu` error the transport crates share. moq-crypto = { workspace = true, features = ["provider"], optional = true } -moq-net = { workspace = true } +moq-net = { workspace = true, features = ["async"] } moq-quic = { workspace = true, features = [], optional = true } moq-sock = { workspace = true } rustc-hash = { workspace = true } diff --git a/rs/moq-uring/src/quic/driver/connection.rs b/rs/moq-uring/src/quic/driver/connection.rs index 849a33213a7..144c6ebc1c2 100644 --- a/rs/moq-uring/src/quic/driver/connection.rs +++ b/rs/moq-uring/src/quic/driver/connection.rs @@ -266,11 +266,6 @@ impl Connection { self.shared.close_code(error_code.into_inner(), reason); } - /// Wait for the connection to close, returning why. - pub async fn closed(&self) -> Error { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - /// The connection's statistics. pub fn stats(&self) -> moq_quic::connection::Stats { self.shared.conn.borrow().stats() @@ -472,13 +467,15 @@ pub(crate) async fn establish(shared: Shared, deadline: Option) -> Resu impl Connection { /// Accept the next unidirectional stream the peer opens. - pub async fn accept_uni(&self) -> Result { - kio::wait(|waiter| self.poll_accept_uni(waiter)).await + #[kio::wrapper] + pub fn poll_accept_uni(&self, waiter: &kio::Waiter) -> Poll> { + accept_uni(&self.shared, waiter) } /// Accept the next bidirectional stream the peer opens. - pub async fn accept_bi(&self) -> Result<(send::Stream, recv::Stream), Error> { - kio::wait(|waiter| self.poll_accept_bi(waiter)).await + #[kio::wrapper] + pub fn poll_accept_bi(&self, waiter: &kio::Waiter) -> Poll> { + accept_bi(&self.shared, waiter) } /// Open a unidirectional stream scheduled by `options`, once the peer's stream credit @@ -486,8 +483,9 @@ impl Connection { /// /// # Panics /// - when `options.group` belongs to another connection - pub async fn open_uni(&self, options: send::Options<'_>) -> Result { - kio::wait(|waiter| self.poll_open_uni(waiter, options)).await + #[kio::wrapper] + pub fn poll_open_uni(&self, waiter: &kio::Waiter, options: send::Options<'_>) -> Poll> { + open(&self.shared, Dir::Uni, options, waiter).map_ok(|(send, _)| send) } /// Open a bidirectional stream scheduled by `options`, once the peer's stream credit @@ -495,31 +493,7 @@ impl Connection { /// /// # Panics /// - when `options.group` belongs to another connection - pub async fn open_bi(&self, options: send::Options<'_>) -> Result<(send::Stream, recv::Stream), Error> { - kio::wait(|waiter| self.poll_open_bi(waiter, options)).await - } - - /// Receive the next datagram. - pub async fn recv_datagram(&self) -> Result { - kio::wait(|waiter| self.poll_recv_datagram(waiter)).await - } - - /// [`Self::accept_uni`], as a poll. - pub fn poll_accept_uni(&self, waiter: &kio::Waiter) -> Poll> { - accept_uni(&self.shared, waiter) - } - - /// [`Self::accept_bi`], as a poll. - pub fn poll_accept_bi(&self, waiter: &kio::Waiter) -> Poll> { - accept_bi(&self.shared, waiter) - } - - /// [`Self::open_uni`], as a poll. - pub fn poll_open_uni(&self, waiter: &kio::Waiter, options: send::Options<'_>) -> Poll> { - open(&self.shared, Dir::Uni, options, waiter).map_ok(|(send, _)| send) - } - - /// [`Self::open_bi`], as a poll. + #[kio::wrapper] pub fn poll_open_bi( &self, waiter: &kio::Waiter, @@ -528,12 +502,14 @@ impl Connection { open(&self.shared, Dir::Bi, options, waiter).map_ok(|(send, recv)| (send, recv.expect("bidirectional"))) } - /// [`Self::recv_datagram`], as a poll. + /// Receive the next datagram. + #[kio::wrapper] pub fn poll_recv_datagram(&self, waiter: &kio::Waiter) -> Poll> { recv_datagram(&self.shared, waiter) } - /// [`Self::closed`], as a poll. + /// Why the connection closed, once it has. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll { poll_closed(&self.shared, waiter) } diff --git a/rs/moq-uring/src/quic/driver/endpoint.rs b/rs/moq-uring/src/quic/driver/endpoint.rs index 95eb33c7305..914d776b1af 100644 --- a/rs/moq-uring/src/quic/driver/endpoint.rs +++ b/rs/moq-uring/src/quic/driver/endpoint.rs @@ -149,11 +149,7 @@ impl Endpoint { /// Fails immediately on an endpoint built without a /// [`server`](Config::server) configuration, and with the socket's error /// once the endpoint has died, its worker included. - pub async fn accept(&self) -> Result { - kio::wait(|waiter| self.poll_accept(waiter)).await - } - - /// The next handshaken incoming connection, or the endpoint's terminal error. + #[kio::wrapper] pub fn poll_accept(&self, waiter: &kio::Waiter) -> Poll> { if let Some(err) = self.inner.closed() { return Poll::Ready(Err(err)); diff --git a/rs/moq-uring/src/quic/driver/recv.rs b/rs/moq-uring/src/quic/driver/recv.rs index 962cacab753..ecb4f1e2a2b 100644 --- a/rs/moq-uring/src/quic/driver/recv.rs +++ b/rs/moq-uring/src/quic/driver/recv.rs @@ -78,12 +78,8 @@ impl Stream { /// Segment boundaries do not correspond to peer writes, and hence cannot be used as framing. /// /// Cancel-safe. - pub async fn read_chunk(&mut self, max: usize) -> Result, Error> { - kio::wait(|waiter| self.poll_read_chunk(waiter, max)).await - } - - /// [`Self::read_chunk`], as a poll. - fn poll_read_chunk(&mut self, waiter: &kio::Waiter, max: usize) -> Poll, Error>> { + #[kio::wrapper] + pub fn poll_read_chunk(&mut self, waiter: &kio::Waiter, max: usize) -> Poll, Error>> { self.poll_read_generic(waiter, |chunks, waiter| match chunks.poll_next(waiter, max) { Poll::Ready(Ok(Some(chunk))) => ReadStatus::Readable(chunk.bytes), Poll::Ready(res) => (None, res.err()).into(), @@ -192,21 +188,17 @@ impl Stream { self.drop_stop = code; } - /// Completes when the stream has been reset by the peer or otherwise closed. + /// Resolves once the stream has been reset by the peer or otherwise closed. /// /// Yields `Some` with the reset error code when the stream is reset by the peer, and `None` /// once every byte the peer finished it with has been read, after which it is no longer /// meaningful for the stream to be reset. Fails with [`Error::Evicted`] once the stream was - /// evicted to free a stuck receive window, until it is stopped. + /// evicted to free a stuck receive window, until it is stopped. It takes no bytes, so a lease + /// in progress is left as it was. /// /// Cancel-safe. - pub async fn received_reset(&mut self) -> Result, Error> { - kio::wait(|waiter| self.poll_received_reset(waiter)).await - } - - /// Poll form of [`Self::received_reset`]. It takes no bytes, so a lease in progress is - /// left as it was. - pub(crate) fn poll_received_reset(&mut self, waiter: &kio::Waiter) -> Poll, Error>> { + #[kio::wrapper] + pub fn poll_received_reset(&mut self, waiter: &kio::Waiter) -> Poll, Error>> { if let Some(code) = self.reset { return Poll::Ready(Ok(Some(code))); } diff --git a/rs/moq-uring/src/quic/driver/send.rs b/rs/moq-uring/src/quic/driver/send.rs index 7777d2b2f12..423439589e8 100644 --- a/rs/moq-uring/src/quic/driver/send.rs +++ b/rs/moq-uring/src/quic/driver/send.rs @@ -156,12 +156,8 @@ impl Stream { self.0.set_group(group); } - /// Completes when the peer stops the stream. See [`Finished::closed()`]. - pub async fn closed(&self) -> Result, Error> { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - - /// Wait for the stream to end, retaining the peer's stop code; see [`Self::closed`]. + /// Resolves once the peer stops the stream, with its stop code. See [`Finished::poll_closed`]. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll, Error>> { poll_closed(&self.0.shared, self.0.id, waiter) } @@ -204,17 +200,13 @@ impl Finished { self.0.set_group(group); } - /// Completes when the peer stops the stream or acknowledges all of its data. + /// Resolves once the peer stops the stream or acknowledges all of its data. /// /// Yields `Some` with the stop error code if the peer stops the stream, and `None` once the /// peer acknowledges receipt of all stream data (although not necessarily the processing of /// it). Fails only when the connection ends first, since then the caller cannot read success /// as "every byte arrived". - pub async fn closed(&self) -> Result, Error> { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - - /// Wait for the stream to end, retaining the peer's stop code; see [`Self::closed`]. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll, Error>> { poll_closed(&self.0.shared, self.0.id, waiter) } diff --git a/rs/moq-uring/src/quic/web.rs b/rs/moq-uring/src/quic/web.rs index 925033a1124..aea49b43106 100644 --- a/rs/moq-uring/src/quic/web.rs +++ b/rs/moq-uring/src/quic/web.rs @@ -537,7 +537,8 @@ struct Opening { } impl Session { - /// Poll for a unidirectional stream after its session header. + /// The next unidirectional stream, after its session header. + #[kio::wrapper] pub fn poll_accept_uni(&self, waiter: &kio::Waiter) -> Poll> { let Some(web) = &self.web else { return self @@ -569,7 +570,8 @@ impl Session { }) } - /// Poll for a bidirectional stream after its session header. + /// The next bidirectional stream, after its session header. + #[kio::wrapper] pub fn poll_accept_bi(&self, waiter: &kio::Waiter) -> Poll> { let Some(web) = &self.web else { return self.conn.poll_accept_bi(waiter).map_ok(|(send, recv)| { @@ -606,7 +608,8 @@ impl Session { }) } - /// Poll to open a unidirectional stream scheduled by `options`. + /// Open a unidirectional stream scheduled by `options`. + #[kio::wrapper] pub fn poll_open_uni( &self, waiter: &kio::Waiter, @@ -635,7 +638,8 @@ impl Session { }) } - /// Poll to open a bidirectional stream scheduled by `options`. + /// Open a bidirectional stream scheduled by `options`. + #[kio::wrapper] pub fn poll_open_bi( &self, waiter: &kio::Waiter, @@ -669,10 +673,11 @@ impl Session { }) } - /// Poll for the next datagram after its session header. + /// The next datagram, with its session header removed. /// /// A datagram that does not name this session, such as another session's on the same /// connection, is dropped: datagrams are unreliable, so a stray one must not end the session. + #[kio::wrapper] pub fn poll_recv_datagram(&self, waiter: &kio::Waiter) -> Poll> { loop { let datagram = ready!(self.conn.poll_recv_datagram(waiter)).map_err(|e| self.map_err(e))?; @@ -715,7 +720,8 @@ impl Session { pub fn close(&self, code: u32, reason: &str) { close(&self.conn, self.web.as_ref(), code, reason); } - /// Poll until the session's connection has closed. + /// Why the session's connection closed, once it has. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll { let err = ready!(self.conn.poll_closed(waiter)); Poll::Ready( @@ -733,30 +739,6 @@ impl Session { pub fn connection(&self) -> &Connection { &self.conn } - /// Accept the next unidirectional stream after its session header. - pub async fn accept_uni(&self) -> Result { - kio::wait(|w| self.poll_accept_uni(w)).await - } - /// Accept the next bidirectional stream after its session header. - pub async fn accept_bi(&self) -> Result<(SendStream, RecvStream), Error> { - kio::wait(|w| self.poll_accept_bi(w)).await - } - /// Open a unidirectional stream scheduled by `options`. - pub async fn open_uni(&self, options: driver_send::Options<'_>) -> Result { - kio::wait(|w| self.poll_open_uni(w, options)).await - } - /// Open a bidirectional stream scheduled by `options`. - pub async fn open_bi(&self, options: driver_send::Options<'_>) -> Result<(SendStream, RecvStream), Error> { - kio::wait(|w| self.poll_open_bi(w, options)).await - } - /// Receive a datagram with its session header removed. - pub async fn recv_datagram(&self) -> Result { - kio::wait(|w| self.poll_recv_datagram(w)).await - } - /// Wait until the session's connection has closed. - pub async fn closed(&self) -> Error { - kio::wait(|w| self.poll_closed(w)).await - } } impl moq_net::transport::poll::Session for Session { diff --git a/rs/moq-uring/src/quic/web/recv.rs b/rs/moq-uring/src/quic/web/recv.rs index 1a0a39e7c2b..495420b3af5 100644 --- a/rs/moq-uring/src/quic/web/recv.rs +++ b/rs/moq-uring/src/quic/web/recv.rs @@ -82,12 +82,8 @@ impl Stream { self.stop_mut(code); } - /// Wait until the stream ends after its bytes are read, the peer resets it, or the session closes. - pub async fn closed(&mut self) -> Result<(), Error> { - kio::wait(|waiter| self.poll_closed(waiter)).await - } - - /// Poll until the stream ends after its bytes are read, the peer resets it, or the session closes. + /// Resolves once the stream ends after its bytes are read, the peer resets it, or the session closes. + #[kio::wrapper] pub fn poll_closed(&mut self, waiter: &kio::Waiter) -> Poll> { let close = &self.close; let Some(inner) = &mut self.inner else { diff --git a/rs/moq-uring/src/quic/web/send.rs b/rs/moq-uring/src/quic/web/send.rs index 08f6c95f73e..ba0df503eb9 100644 --- a/rs/moq-uring/src/quic/web/send.rs +++ b/rs/moq-uring/src/quic/web/send.rs @@ -105,7 +105,9 @@ impl Stream { } } - /// Poll until all data was acknowledged, the peer stopped the stream, or the session closed. + /// Resolves once all data was acknowledged, the peer stopped the stream (with its code), or the + /// session closed. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll, Error>> { let result = match &self.state { State::Open(stream) => stream.poll_closed(waiter), @@ -122,11 +124,6 @@ impl Stream { }) } - /// Wait until all data was acknowledged, the peer stopped the stream, or the session closed. - pub async fn closed(&self) -> Result, Error> { - kio::wait(|w| self.poll_closed(w)).await - } - /// Finish sending, retaining a handle to acknowledgements and peer cancellation. pub fn finish(mut self) -> Finished { self.finish_mut(); @@ -262,14 +259,11 @@ pub struct Finished { park: kio::Park, } impl Finished { - /// Poll until the peer acknowledges or stops the stream, or the session closes. + /// Resolves once the peer acknowledges or stops the stream (with its code), or the session closes. + #[kio::wrapper] pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll, Error>> { self.stream.poll_closed(waiter) } - /// Wait until the peer acknowledges or stops the stream, or the session closes. - pub async fn closed(&self) -> Result, Error> { - self.stream.closed().await - } /// Reset the stream while still delivering its WebTransport header. pub fn reset(self, code: u32) { self.stream.reset(code); diff --git a/rs/moq-uring/src/udp.rs b/rs/moq-uring/src/udp.rs index fb6ae374889..cd0bde4e0f6 100644 --- a/rs/moq-uring/src/udp.rs +++ b/rs/moq-uring/src/udp.rs @@ -812,9 +812,9 @@ impl Socket { self.shared.shard } - /// A received packet, or the socket's terminal error, registering `waiter` - /// while neither is available. Queued packets drain before an error + /// A received packet, or the socket's terminal error. Queued packets drain before an error /// surfaces. + #[kio::wrapper] pub fn poll_recv(&self, waiter: &kio::Waiter) -> Poll> { let mut rx = self.shared.rx.borrow_mut(); if let Some(queued) = rx.queue.pop_front() { @@ -862,13 +862,8 @@ impl Socket { Some((to, io::Error::from_raw_os_error(code))) } - /// Await [`poll_recv`](Self::poll_recv). - pub async fn recv(&self) -> io::Result { - kio::wait(|waiter| self.poll_recv(waiter)).await - } - - /// A free send-staging buffer, registering `waiter` while the pool is - /// drained. Backpressure lives here: the pool caps in-flight sends. + /// A free send-staging buffer. Backpressure lives here: the pool caps in-flight sends. + #[kio::wrapper] pub fn poll_acquire(&self, waiter: &kio::Waiter) -> Poll> { let mut tx = self.shared.tx.borrow_mut(); if let Some(code) = tx.error { @@ -906,11 +901,6 @@ impl Socket { waiter.register(&mut tx.waiters); Poll::Pending } - - /// Await [`poll_acquire`](Self::poll_acquire). - pub async fn acquire(&self) -> io::Result { - kio::wait(|waiter| self.poll_acquire(waiter)).await - } } impl Drop for Socket { diff --git a/rs/moq-video/Cargo.toml b/rs/moq-video/Cargo.toml index 768d2ac3bfc..8745fcd9b5f 100644 --- a/rs/moq-video/Cargo.toml +++ b/rs/moq-video/Cargo.toml @@ -126,7 +126,7 @@ kio = { workspace = true, features = ["runtime"] } # libvpx bindings for the `vpx` feature. The library itself comes from the build # host; see the feature above. libvpx-native-sys = { version = "5.0.17", optional = true } -moq-mux = { workspace = true } +moq-mux = { workspace = true, features = ["async"] } moq-net = { workspace = true } # Vendored, statically linked software H.264 fallback (no system dependency). # Hardware encoders are preferred, but an `openh264` build falls back when diff --git a/sh/rs/coroutines.sh b/sh/rs/coroutines.sh index d8cbc994a32..02ac6db652e 100755 --- a/sh/rs/coroutines.sh +++ b/sh/rs/coroutines.sh @@ -1,69 +1,45 @@ #!/usr/bin/env bash # Prove a poll-only library build compiles no coroutines: no `async fn` body, # `async` block, or `async` closure in the crate's own MIR. A library's `async` -# feature carries its wrappers, so a build without it should have none. Each -# `--allow NAME` excuses the `async fn NAME` bodies still waiting on a poll -# conversion, and fails once none is left, so the list only shrinks. +# feature carries its wrappers, so a build without it should have none. Name a +# PACKAGE by its `Cargo.toml` path when another dependency shares its name. # -# usage: coroutines.sh PACKAGE [--allow NAME]... [-- CARGO_ARGS...] +# usage: coroutines.sh PACKAGE [-- CARGO_ARGS...] set -euo pipefail if [[ "${GITHUB_ACTIONS:-}" != true && "${LOCAL_SUPERVISED:-}" != 1 ]]; then exec bash "$(dirname "$0")/../local.sh" bash "$0" "$@" fi -usage='usage: coroutines.sh PACKAGE [--allow NAME]... [-- CARGO_ARGS...]' -if [[ $# == 0 ]]; then +usage='usage: coroutines.sh PACKAGE [-- CARGO_ARGS...]' +if [[ $# == 0 || ($# -ge 2 && "$2" != --) ]]; then echo "$usage" >&2 exit 2 fi package=$1 shift -allowed=() -while [[ $# != 0 ]]; do - case "$1" in - --allow) - [[ $# -ge 2 && "$2" =~ ^[A-Za-z_][A-Za-z0-9_]*$ ]] || { - echo "$usage" >&2 - exit 2 - } - allowed+=("$2") - shift 2 - ;; - --) - shift - break - ;; - *) - echo "$usage" >&2 - exit 2 - ;; - esac -done +if [[ $# != 0 ]]; then + shift +fi cd "$(git rev-parse --show-toplevel)" mir=$(mktemp) trap 'rm -f "$mir"' EXIT # Only the final crate gets `--emit`, so its dependencies build as usual. -cargo rustc --locked -p "$package" --lib "$@" -- --emit "mir=$mir" +if [[ $package == */Cargo.toml ]]; then + select=(--manifest-path "$package") +else + select=(-p "$package") +fi +cargo rustc --locked "${select[@]}" --lib "$@" -- --emit "mir=$mir" -# rustc names each coroutine by its body, as `{async fn body of queue::Queue::pop()}`, -# sometimes without the module path, so an allowed name matches as the last segment. +# rustc names each coroutine by its body, as `{async fn body of queue::Queue::pop()}`. found=$(grep -oE '\{async (fn body of [^{}]*|block@[^{}]*|closure@[^{}]*)\}' "$mir" | sort -u || true) -unexpected=$found -for name in ${allowed[@]+"${allowed[@]}"}; do - pattern="^\{async fn body of ([A-Za-z0-9_<>, ]+::)*${name}[<(]" - if ! grep -qE -- "$pattern" <<<"$found"; then - echo "error: $package no longer compiles async fn $name; remove its --allow" >&2 - exit 1 - fi - unexpected=$(grep -vE -- "$pattern" <<<"$unexpected" || true) -done -if [[ -n "$unexpected" ]]; then +if [[ -n "$found" ]]; then echo "error: $package compiles coroutines without its async feature:" >&2 - echo "$unexpected" >&2 + echo "$found" >&2 echo "Gate each behind the feature, or generate it from a poll with #[kio::wrapper]." >&2 exit 1 fi -echo "$package: no coroutines beyond ${#allowed[@]} allowed async fn names" +echo "$package: no coroutines" diff --git a/test/drill/mutations/subscriber-max-age-ignored.patch b/test/drill/mutations/subscriber-max-age-ignored.patch index 2b4732aa115..ab941992b22 100644 --- a/test/drill/mutations/subscriber-max-age-ignored.patch +++ b/test/drill/mutations/subscriber-max-age-ignored.patch @@ -4,10 +4,10 @@ # Preserve negotiation and decoding, but ignore the requested budget when serving. --- a/rs/moq-net/src/ietf/publisher.rs +++ b/rs/moq-net/src/ietf/publisher.rs -@@ -2028,9 +2028,6 @@ - // A live SUBSCRIBE's NEW_GROUP_REQUEST asks nothing more: subscribing is itself a join, +@@ -1887,9 +1887,6 @@ // which the track's refresh ratio judges. - let mut subscription = serving_subscription(msg.subscriber_priority).with_order(msg.group_order.to_order()); + let mut subscription: Subscription = + serving_subscription(msg.subscriber_priority).with_order(msg.group_order.to_order()); - if let Some(max_age) = msg.max_age { - subscription.delay.max = Some(Duration::from_millis(max_age)); - }