Skip to content
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,10 +14,12 @@ mod lock;
mod waiter;

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

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
4 changes: 3 additions & 1 deletion rs/kio/src/producer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,9 @@ impl<T> Producer<T> {
}
}

fn poll_unused(&self, waiter: &Waiter) -> Poll<Option<()>> {
/// Poll-based variant of [`Self::unused`]: `Ready(Some(()))` when no consumers
/// remain, `Ready(None)` if the channel closed first, else `Pending`.
pub fn poll_unused(&self, waiter: &Waiter) -> Poll<Option<()>> {
let mut state = self.state.lock();
if state.closed {
return Poll::Ready(None);
Expand Down
117 changes: 117 additions & 0 deletions rs/moq-native/tests/broadcast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,123 @@ async fn broadcast_moq_lite_05_timestamps_webtransport() {
lite05_timestamp_roundtrip("https").await;
}

/// Lite05 FETCH round-trip: retrieve a past group by sequence without holding a
/// subscription, exercising the FETCH / FETCH_OK control flow and per-frame
/// timestamp decoding on the fetch stream.
async fn lite05_fetch_roundtrip(scheme: &str) {
use moq_native::moq_net::{Timescale, Timestamp};

let pub_origin = Origin::random().produce();
let mut broadcast = pub_origin.create_broadcast("test").expect("failed to create broadcast");
let mut track = broadcast
.create_track("video", moq_net::TrackInfo::default().with_timescale(Timescale::MICRO))
.expect("failed to create track");

// A group with a few timestamped frames (middle PTS goes backwards, so the
// fetch stream carries a negative zigzag delta too).
let timestamps_us = [10_000u64, 30_000, 20_000];
let mut group = track.append_group().expect("failed to append group"); // seq 0
for &us in &timestamps_us {
let payload = format!("frame@{us}").into_bytes();
let frame = moq_native::moq_net::Frame {
size: payload.len() as u64,
timestamp: Some(Timestamp::new(us, Timescale::MICRO).unwrap()),
};
let mut writer = group.create_frame(frame).expect("failed to create frame");
writer
.write(bytes::Bytes::from(payload))
.expect("failed to write frame");
writer.finish().expect("failed to finish frame");
}
group.finish().expect("failed to finish group");

let mut server_config = moq_native::ServerConfig::default();
server_config.bind = Some("[::]:0".to_string());
server_config.tls.generate = vec!["localhost".into()];
server_config.version = vec!["moq-lite-05-wip".parse().unwrap()];
let mut server = server_config.init().expect("failed to init server");
let addr = server.local_addr().expect("failed to get local addr");

let sub_origin = Origin::random().produce();
let mut announcements = sub_origin.consume().announced();

let mut client_config = moq_native::ClientConfig::default();
client_config.tls.disable_verify = Some(true);
client_config.version = vec!["moq-lite-05-wip".parse().unwrap()];
let client = client_config.init().expect("failed to init client");
let url: url::Url = format!("{scheme}://localhost:{}", addr.port()).parse().unwrap();

let server_handle = tokio::spawn(async move {
let request = server.accept().await.expect("no incoming connection");
let session = request.with_publisher(pub_origin.clone()).ok().await?;
let _broadcast = broadcast;
let _track = track;
let _ = session.closed().await;
Ok::<_, anyhow::Error>(())
});

let client = client.with_consumer(sub_origin);
let session = tokio::time::timeout(TIMEOUT, client.connect(url))
.await
.expect("client connect timed out")
.expect("client connect failed");

let (path, bc) = tokio::time::timeout(TIMEOUT, announcements.next())
.await
.expect("announce timed out")
.expect("origin closed");
assert_eq!(path.as_str(), "test");
let bc = bc.broadcast().expect("expected announce, got unannounce");

// Fetch group 0 directly, without subscribing. No live producer holds the group
// on the client, so this issues a wire FETCH upstream.
let mut group_sub = tokio::time::timeout(TIMEOUT, async {
bc.track("video").unwrap().fetch(0, None).unwrap().await
})
.await
.expect("fetch timed out")
.expect("fetch failed");
assert_eq!(group_sub.sequence, 0);

for &expected_us in &timestamps_us {
let mut frame_sub = tokio::time::timeout(TIMEOUT, group_sub.next_frame())
.await
.expect("next_frame timed out")
.expect("next_frame failed")
.expect("group closed prematurely");

let ts = frame_sub
.timestamp
.expect("Lite05 fetch must carry per-frame timestamps");
assert_eq!(ts.scale(), Timescale::MICRO);
assert_eq!(ts.value(), expected_us);

let payload = frame_sub.read_all().await.expect("failed to read frame");
assert_eq!(payload, bytes::Bytes::from(format!("frame@{expected_us}")));
}

// The fetched group ends cleanly (stream FIN → no more frames).
let end = tokio::time::timeout(TIMEOUT, group_sub.next_frame())
.await
.expect("next_frame timed out")
.expect("next_frame failed");
assert!(end.is_none(), "group should finish after its frames");

drop(session);
server_handle
.await
.expect("server task panicked")
.expect("server task failed");
}

#[tracing_test::traced_test]
#[tokio::test]
async fn broadcast_moq_lite_05_fetch_webtransport() {
// WebTransport only: Lite05Wip isn't advertised over ALPN, so raw QUIC (moqt://)
// can't negotiate it (same reason the other Lite05 tests are https-only).
lite05_fetch_roundtrip("https").await;
}

/// On Lite05 a publisher that doesn't advertise a timescale still works:
/// SUBSCRIBE_OK carries `timescale = 0` and neither side encodes a
/// per-frame timestamp byte. Subscribers receive `frame.timestamp = None`.
Expand Down
Loading
Loading