kio: rework Producer::poll/wait to a read-only predicate that returns a Mut

[kixelated]

Summary

Producer::poll / Producer::wait were a self-triggering footgun. This reworks them so the closure is a read-only predicate and a satisfied poll hands back a Mut.

The old versions handed the closure a &mut Mut<T> and tracked a modified flag via DerefMut. On a Poll::Pending, the poll woke the consumer waiters. The trap: merely reading through a mutable path trips DerefMut. The canonical case is requests.pop() on an empty Vec — pop() takes &mut, so an empty (no-op) poll still flags the state modified, wakes a waiter in the very list the poller just registered itself into, re-polls, and spins into an infinite busy loop.

Approach

Producer::poll (and its async sibling Producer::wait) now evaluates a read-only predicate over a Ref. While the predicate returns Poll::Pending nothing is mutated, so it can't spuriously flag the state modified. On Poll::Ready it upgrades the held lock guard to a Mut and hands it back, so the caller mutates atomically — the lock is never dropped between the predicate and the mutation, so there's no TOCTOU even with multiple drainers.

let mut state = ready!(self.poll(waiter, |state| {
    if state.requests.is_empty() { Poll::Pending } else { Poll::Ready(()) }
}))?;
Poll::Ready(Ok(state.requests.pop().expect("predicate guaranteed a request")))

Naming: both channel sides now expose a poll()/wait() pair. They differ in contract by design — Consumer::poll's closure returns the value (Poll<R>), while Producer::poll's closure is a bare Poll<()> predicate that gates handing back a Mut.

Changes

• kio: Producer::poll / Producer::wait reworked to the read-only-predicate-returns-Mut shape. Weak::poll_write / Weak::wait had no users and are removed. The read-only Consumer::poll / Consumer::wait are unchanged. Dropped test = false from the [lib] so unit tests run, and added a test for the new Producer::poll.
• moq-net: the two callers, BroadcastDynamic::poll_requested_track and OriginConsumer::poll_announced, now check "queue non-empty" read-only, then pop/take through the returned Mut (the latter tidied to use ready!).

API note

This is a breaking change to the published kio crate's surface (Producer::poll/wait change signature; the Weak methods are removed). No other workspace crate used the removed/changed APIs, and moq-net's public surface is unchanged — the affected helpers are private. Targeting main per request.

Test plan

• cargo test -p kio (2 passed, incl. the new poll_gates_on_predicate_then_writes)
• cargo test -p kio -p moq-net (346 passed; incl. broadcast::test::requests, requested_unused, origin announce tests)
• cargo clippy -p kio -p moq-net -p moq-native clean
• cargo build for kio-dependent crates (moq-net, moq-native, hang, moq-mux, moq-cli, moq-bench)

(Written by Claude)

[coderabbitai]

Walkthrough

The PR replaces the old Mut-closure-based polling API (Producer::poll, Producer::wait, Weak::poll_write, Weak::wait) in the kio crate with a predicate-based alternative: Producer::poll_write_when and its async counterpart Producer::write_when. The new API evaluates a predicate over a read-locked &Ref<'_, T> and upgrades to Mut<'_, T> only when the predicate returns Ready, registering the waiter on Pending. Weak<T> loses its two polling methods entirely. Callers in moq-net — BroadcastDynamic::poll_requested_track and OriginConsumer::poll_announced — are updated to use the new predicate pattern. The breaking change is documented in rs/kio/CHANGELOG.md.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title accurately summarizes the main change: reworking Producer::poll/wait to use a read-only predicate that returns a Mut instead of a mutable closure.
Description check	✅ Passed	The description comprehensively explains the problem, solution, and changes made, all directly related to the changeset across multiple files.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

✨ Simplify code

• Create PR with simplified code
• Commit simplified code in branch claude/kio-mut-poll-footgun-t0k9zm

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

rs/moq-net/src/model/origin.rs (1)

882-894: ⚡ Quick win

Use ready! in poll_announced poll plumbing for consistency with project rules.

This block manually matches Poll where the guideline expects ready!(...) composition.

Suggested refactor

 	pub fn poll_announced(&mut self, waiter: &kio::Waiter) -> Poll<Option<OriginAnnounce>> {
-		let res = self.state.poll_write_when(waiter, |state| {
+		let mut state = match ready!(self.state.poll_write_when(waiter, |state| {
 			if state.pending.is_empty() {
 				Poll::Pending
 			} else {
 				Poll::Ready(())
 			}
-		});
-		match res {
-			Poll::Ready(Ok(mut state)) => Poll::Ready(Some(state.take().expect("predicate guaranteed an update"))),
-			// Closed: discard the Ref so its MutexGuard doesn't escape this call.
-			Poll::Ready(Err(_)) => Poll::Ready(None),
-			Poll::Pending => Poll::Pending,
-		}
+		})) {
+			Ok(state) => state,
+			// Closed: discard the Ref so its MutexGuard doesn't escape this call.
+			Err(_) => return Poll::Ready(None),
+		};
+		Poll::Ready(Some(state.take().expect("predicate guaranteed an update")))
 	}

As per coding guidelines, in poll_* plumbing, use ready!(...) to collapse Poll matching.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@rs/moq-net/src/model/origin.rs` around lines 882 - 894, In the poll_announced
method, replace the manual match on res (which handles Poll::Ready and
Poll::Pending cases) with the ready! macro for consistency with project
guidelines. Use ready! to unwrap the Poll result, which will automatically
propagate Poll::Pending and return the state value for further processing. Then
simplify the remaining match statement to only handle the Ok and Err cases from
the inner Result, eliminating the explicit Poll::Pending and Poll::Ready
wrapping in the match arms.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@rs/moq-net/src/model/origin.rs`:
- Around line 882-894: In the poll_announced method, replace the manual match on
res (which handles Poll::Ready and Poll::Pending cases) with the ready! macro
for consistency with project guidelines. Use ready! to unwrap the Poll result,
which will automatically propagate Poll::Pending and return the state value for
further processing. Then simplify the remaining match statement to only handle
the Ok and Err cases from the inner Result, eliminating the explicit
Poll::Pending and Poll::Ready wrapping in the match arms.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: b1f2b83d-c0f8-4853-9363-45e29f78c61e

📥 Commits

Reviewing files that changed from the base of the PR and between 15e2e88 and 0b4a085.

📒 Files selected for processing (5)

• rs/kio/CHANGELOG.md
• rs/kio/src/producer.rs
• rs/kio/src/weak.rs
• rs/moq-net/src/model/broadcast.rs
• rs/moq-net/src/model/origin.rs

💤 Files with no reviewable changes (1)

• rs/kio/src/weak.rs