Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion rs/kio/src/consumer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use std::{
use crate::{
Counts, State, Weak,
lock::*,
producer::{Producer, Ref},
producer::{Mut, Producer, Ref},
waiter::*,
};

Expand Down Expand Up @@ -51,6 +51,21 @@ impl<T> Consumer<T> {
Poll::Pending
}

/// Acquire write access to the shared state from the consumer side.
///
/// Unlike [`poll`](Self::poll), this never registers a waiter; it simply locks the
/// state for mutation. Returns `Err(`[`Ref`]`)` if the channel is already closed.
/// Mirrors [`Producer::write`], for the rare case where a consumer also feeds shared
/// state back to producers (e.g. a request queue).
pub fn write(&self) -> Result<Mut<'_, T>, Ref<'_, T>> {
let state = self.state.lock();
if state.closed {
Err(Ref { state })
} else {
Ok(Mut::new(state))
}
}

/// Poll for channel closure, registering the waiter if still open.
pub fn poll_closed(&self, waiter: &Waiter) -> Poll<()> {
let mut state = self.state.lock();
Expand Down
146 changes: 146 additions & 0 deletions rs/kio/src/future.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,146 @@
use std::{
ops::{Deref, DerefMut},
pin::Pin,
task::{Context, Poll},
};

use crate::Waiter;

/// A pollable computation backed by kio channels.
///
/// Implementors write only [`Self::poll`], registering the [`Waiter`] with the
/// channels they read. Wrap the value in [`Pending`] to get a real
/// [`std::future::Future`].
///
/// This exists because a kio [`Waiter`] holds the strong `Arc<Waker>` while the
/// channel's [`crate::WaiterList`] keeps only a `Weak`. A bare
/// [`std::future::Future`] would have to stash the strong `Waiter` in a field and
/// replace it every poll (or lose its wakeup); [`Pending`] does that once so each
/// implementor doesn't have to.
pub trait Future: Unpin {
type Output;

/// Poll for the output, registering `waiter` with the relevant channels if not
/// yet ready.
///
/// Takes `&self`: kio channels poll immutably, so a pollable can be driven
/// through a shared borrow (e.g. while it lives inside an `&self`-borrowed enum).
/// Carry any per-poll mutable state in a kio channel or a [`std::cell`] type.
fn poll(&self, waiter: &Waiter) -> Poll<Self::Output>;
}

/// Adapts a kio [`Future`] into a [`std::future::Future`], retaining the strong
/// [`Waiter`] between polls so its weak registration stays live.
///
/// Derefs to the inner value, so any inherent methods you define on it are
/// reachable through the pending handle (e.g. a non-blocking `poll`, or an
/// `update`).
pub struct Pending<F> {
inner: F,
// Retain the previous waiter so its Weak registration survives until the next
// poll replaces it (see [`crate::WaiterList`]).
waiter: Option<Waiter>,
}

impl<F> Pending<F> {
/// Wrap a [`Future`] so it can be `.await`ed.
pub fn new(inner: F) -> Self {
Self { inner, waiter: None }
}

/// Consume the wrapper, returning the inner value.
pub fn into_inner(self) -> F {
self.inner
}
}

impl<F> Deref for Pending<F> {
type Target = F;

fn deref(&self) -> &F {
&self.inner
}
}

impl<F> DerefMut for Pending<F> {
fn deref_mut(&mut self) -> &mut F {
&mut self.inner
}
}

impl<F: Future> std::future::Future for Pending<F> {
type Output = F::Output;

fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
// Replacing drops the previous waiter, killing its Weak ref in the list so
// the inner poll's register call can recycle the slot (see `WaiterList`).
// `Pending<F>` is `Unpin` (F is, via the trait bound), so this deref is sound.
let this = &mut *self;
this.waiter = Some(Waiter::new(cx.waker().clone()));
Future::poll(&this.inner, this.waiter.as_ref().unwrap())
}
}

#[cfg(test)]
mod test {
use super::*;
use crate::Producer;

/// A pollable that waits for the channel value to reach a threshold, with an
/// inherent method reachable through `Pending`'s `DerefMut`.
struct AtLeast {
consumer: crate::Consumer<u64>,
threshold: u64,
}

impl AtLeast {
fn bump_threshold(&mut self) {
self.threshold += 1;
}
}

impl Future for AtLeast {
type Output = u64;

fn poll(&self, waiter: &Waiter) -> Poll<u64> {
let threshold = self.threshold;
match self.consumer.poll(waiter, |v| {
let current = **v;
if current >= threshold {
Poll::Ready(current)
} else {
Poll::Pending
}
}) {
Poll::Ready(Ok(v)) => Poll::Ready(v),
_ => Poll::Pending,
}
}
}

#[test]
fn pending_derefs_and_drives() {
use std::task::Waker;

let producer = Producer::new(0u64);
let mut pending = Pending::new(AtLeast {
consumer: producer.consume(),
threshold: 5,
});

// Inherent method on the inner reached via DerefMut.
pending.bump_threshold(); // threshold now 6

// The kio-level poll (reached through Deref) is pending until the value catches up.
assert!(Future::poll(&*pending, &Waiter::noop()).is_pending());

if let Ok(mut v) = producer.write() {
*v = 6;
}

// The std Future resolves once the threshold is met.
let mut cx = Context::from_waker(Waker::noop());
let mut pending = std::pin::pin!(pending);
assert_eq!(std::future::Future::poll(pending.as_mut(), &mut cx), Poll::Ready(6));
}
}
2 changes: 2 additions & 0 deletions rs/kio/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,15 @@ mod lock;
mod waiter;

mod consumer;
mod future;
mod producer;
mod weak;

#[cfg(test)]
mod tests;

pub use consumer::Consumer;
pub use future::{Future, Pending};
pub use producer::{Mut, Producer, Ref};
pub use waiter::{Waiter, WaiterList, wait};
pub use weak::Weak;
Expand Down
8 changes: 8 additions & 0 deletions rs/moq-net/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,11 @@ pub enum Error {
#[error("not found")]
NotFound,

/// A broadcast was requested that is neither announced nor served by a dynamic
/// router, so there is no route to it.
#[error("unroutable")]
Unroutable,

#[error("wrong frame size")]
WrongSize,

Expand Down Expand Up @@ -123,6 +128,9 @@ impl Error {
Self::Closed => 25,
Self::CacheFull => 26,
Self::FrameTooLarge => 27,
// 28 (Decompress) and 29 (TimestampMismatch) are reserved on the dev branch;
// keep Unroutable at 30 so the wire code is identical across branches.
Self::Unroutable => 30,
Self::App(app) => *app as u32 + 64,
Self::Remote(code) => *code,
}
Expand Down
Loading