moq_uring/quic/noq/

stream.rs

//! Stream handles: thin, direct calls into the shared noq-proto connection.
//!
//! Single-threaded sans-IO means a write goes straight into noq's send
//! queue (no staging copy) and a read comes straight out of its reassembly
//! buffer; the handles just kick the driver so egress reaches the wire and
//! park on the per-stream waiter lists the driver wakes.

use std::task::{Context, Poll};

use bytes::{Buf, Bytes, BytesMut};
use moq_noq_proto::{StreamId, VarInt};

use super::super::Error;
use super::{End, Shared};

/// An outgoing stream. Dropping it unfinished resets it with code 0.
pub struct SendStream {
	shared: Shared,
	id: StreamId,
	park: kio::Park,
	/// The FIN went out; further writes are refused and `poll_closed` waits
	/// for the acknowledgement.
	fin: bool,
	/// We reset the stream; it is as closed as it will ever be.
	reset: bool,
}

impl SendStream {
	pub(crate) fn new(shared: Shared, id: StreamId) -> Self {
		// The driver records how this stream ends only while a handle holds
		// it, so the handle is what announces itself.
		shared.track(id);
		Self {
			shared,
			id,
			park: kio::Park::default(),
			fin: false,
			reset: false,
		}
	}

	/// The QUIC stream id, which the WebTransport layer uses as the session id.
	pub(crate) fn id(&self) -> u64 {
		self.id.into()
	}

	/// Whether the send side is already terminated, so [`Drop`] would do
	/// nothing. The WebTransport wrapper asks before mapping a reset of its
	/// own, since a finished stream must not be reset instead.
	pub(crate) fn ended(&self) -> bool {
		self.fin || self.reset
	}

	/// Queue as much of `buf` as noq will take right now, without parking.
	/// Best-effort, for the close path where nobody is left to poll.
	pub(crate) fn try_write(&mut self, buf: &[u8]) -> usize {
		if self.fin || self.reset {
			return 0;
		}
		let n = self
			.shared
			.conn
			.borrow_mut()
			.send_stream(self.id)
			.write(buf)
			.unwrap_or(0);
		if n > 0 {
			self.shared.kick();
		}
		n
	}

	/// [`reset`](web_transport_trait::poll::SendStream::reset) with a
	/// full-width code, for the WebTransport HTTP/3 error mapping.
	pub(crate) fn reset_code(&mut self, code: u64) {
		if self.reset {
			return;
		}
		// Err means the stream is already gone, which is what we wanted.
		let _ = self
			.shared
			.conn
			.borrow_mut()
			.send_stream(self.id)
			.reset(VarInt::from_u64(code).unwrap_or(VarInt::MAX));
		self.reset = true;
		self.shared.kick();
	}
}

impl web_transport_trait::poll::SendStream for SendStream {
	type Error = Error;

	fn poll_write(&mut self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, Self::Error>> {
		let waiter = self.park.hold(cx);
		if self.fin || self.reset {
			return Poll::Ready(Err(Error::Quic("stream already finished".to_string())));
		}
		if let Some(err) = self.shared.closed() {
			return Poll::Ready(Err(err));
		}
		let result = self.shared.conn.borrow_mut().send_stream(self.id).write(buf);
		match result {
			Ok(n) => {
				self.shared.kick();
				Poll::Ready(Ok(n))
			}
			// No capacity right now; the driver wakes us when noq reports
			// the stream writable.
			Err(moq_noq_proto::WriteError::Blocked) => {
				self.shared.park_writable(self.id, waiter);
				Poll::Pending
			}
			Err(moq_noq_proto::WriteError::Stopped(code)) => Poll::Ready(Err(Error::Stop(code.into_inner()))),
			Err(moq_noq_proto::WriteError::ClosedStream) => {
				Poll::Ready(Err(Error::Quic("stream already finished".to_string())))
			}
		}
	}

	fn set_priority(&mut self, order: u8) {
		// The trait (like W3C sendOrder) sends HIGHER values first, and so
		// does noq.
		let _ = self
			.shared
			.conn
			.borrow_mut()
			.send_stream(self.id)
			.set_priority(i32::from(order));
	}

	fn finish(&mut self) -> Result<(), Self::Error> {
		if self.fin || self.reset {
			return Ok(());
		}
		match self.shared.conn.borrow_mut().send_stream(self.id).finish() {
			Ok(()) => {}
			// A STOP_SENDING beat us here. Carry the code like `poll_write`
			// does, or `moq_net::Error::from_transport` cannot decode a
			// routine cancellation.
			Err(moq_noq_proto::FinishError::Stopped(code)) => {
				self.reset = true;
				return Err(Error::Stop(code.into_inner()));
			}
			// Already finished or reset, so the FIN it wanted is out.
			Err(moq_noq_proto::FinishError::ClosedStream) => {}
		}
		self.fin = true;
		self.shared.kick();
		Ok(())
	}

	fn reset(&mut self, code: u32) {
		self.reset_code(u64::from(code));
	}

	fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
		let waiter = self.park.hold(cx);
		if self.reset {
			return Poll::Ready(Ok(()));
		}
		// noq reports a send stream's end as an event, which the driver
		// records: an acknowledged FIN, or the peer's STOP_SENDING.
		match self.shared.ended(self.id) {
			Some(End::Stopped(code)) => Poll::Ready(Err(Error::Stop(code))),
			Some(End::Delivered) => Poll::Ready(Ok(())),
			// The end never came. If the connection died first the caller
			// cannot read success as "every byte arrived".
			None => match self.shared.closed() {
				Some(err) => Poll::Ready(Err(err)),
				None => {
					self.shared.park_finishing(self.id, waiter);
					Poll::Pending
				}
			},
		}
	}
}

impl Drop for SendStream {
	fn drop(&mut self) {
		self.shared.forget_send(self.id);
		if !self.fin && !self.reset {
			let _ = self
				.shared
				.conn
				.borrow_mut()
				.send_stream(self.id)
				.reset(VarInt::from_u32(0));
			self.shared.kick();
		}
	}
}

impl std::fmt::Debug for SendStream {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		f.debug_struct("SendStream").field("id", &self.id).finish()
	}
}

/// How far [`RecvStream::poll_closed`] reads ahead of the application before
/// it waits for the backlog to drain.
const READ_AHEAD: usize = 64 * 1024;
/// How much to ask for per read-ahead chunk.
const READ_CHUNK: usize = 8 * 1024;

/// One read out of noq's reassembly buffer.
enum Read {
	/// Bytes, at most as many as were asked for.
	Chunk(Bytes),
	/// Every byte up to the FIN has been handed over.
	Finished,
	/// Nothing buffered; the driver wakes us when that changes.
	Blocked,
	/// The peer reset the stream with this code.
	Reset(u64),
}

/// An incoming stream. Dropping it unfinished sends STOP_SENDING with code 0.
pub struct RecvStream {
	shared: Shared,
	id: StreamId,
	park: kio::Park,
	/// Every byte up to the FIN was read out of noq; reads report the end
	/// once `backlog` is drained too.
	finished: bool,
	/// We stopped the stream; no more reads matter.
	stopped: bool,
	/// Bytes `poll_closed` read ahead, handed to `poll_read` before noq's.
	backlog: BytesMut,
}

impl RecvStream {
	pub(crate) fn new(shared: Shared, id: StreamId) -> Self {
		Self {
			shared,
			id,
			park: kio::Park::default(),
			finished: false,
			stopped: false,
			backlog: BytesMut::new(),
		}
	}

	/// Whether the read side is already terminated, so [`Drop`] would do
	/// nothing.
	pub(crate) fn ended(&self) -> bool {
		self.finished || self.stopped
	}

	/// [`stop`](web_transport_trait::poll::RecvStream::stop) with a
	/// full-width code, for the WebTransport HTTP/3 error mapping.
	pub(crate) fn stop_code(&mut self, code: u64) {
		self.backlog.clear();
		if self.stopped || self.finished {
			return;
		}
		// Err means the stream is already gone, which is what we wanted.
		let _ = self
			.shared
			.conn
			.borrow_mut()
			.recv_stream(self.id)
			.stop(VarInt::from_u64(code).unwrap_or(VarInt::MAX));
		self.stopped = true;
		self.shared.kick();
	}

	/// Take up to `max` bytes out of noq's reassembly buffer.
	///
	/// Reading is what returns the peer's flow control credit, so a read that
	/// owes it a frame kicks the driver.
	///
	/// Takes the shared state rather than `&mut self` so a caller can hold a
	/// waiter from `self.park` across the read.
	fn read(shared: &Shared, id: StreamId, max: usize) -> Read {
		let mut conn = shared.conn.borrow_mut();
		let mut recv = conn.recv_stream(id);
		let mut chunks = match recv.read(true) {
			Ok(chunks) => chunks,
			// The stream is gone, so everything it held is already ours.
			Err(_) => return Read::Finished,
		};
		let read = match chunks.next(max) {
			Ok(Some(chunk)) => Read::Chunk(chunk.bytes),
			Ok(None) => Read::Finished,
			Err(moq_noq_proto::ReadError::Blocked) => Read::Blocked,
			Err(moq_noq_proto::ReadError::Reset(code)) => Read::Reset(code.into_inner()),
		};
		let transmit = chunks.finalize().should_transmit();
		drop(conn);
		if transmit {
			shared.kick();
		}
		read
	}

	/// Move up to `dst.len()` read-ahead bytes out of the backlog.
	///
	/// Wakes a `poll_closed` parked at the read-ahead cap: the room it was
	/// waiting for is what this just made.
	fn drain(&mut self, dst: &mut [u8]) -> usize {
		let n = dst.len().min(self.backlog.len());
		dst[..n].copy_from_slice(&self.backlog[..n]);
		self.backlog.advance(n);
		if n > 0 {
			self.shared.wake_readable(self.id);
		}
		n
	}
}

impl web_transport_trait::poll::RecvStream for RecvStream {
	type Error = Error;

	fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll<Result<Option<usize>, Self::Error>> {
		let waiter = self.park.hold(cx);
		if dst.is_empty() {
			return Poll::Ready(Ok(Some(0)));
		}
		if !self.backlog.is_empty() {
			return Poll::Ready(Ok(Some(self.drain(dst))));
		}
		if self.finished {
			return Poll::Ready(Ok(None));
		}
		match Self::read(&self.shared, self.id, dst.len()) {
			Read::Chunk(bytes) => {
				let n = bytes.len().min(dst.len());
				dst[..n].copy_from_slice(&bytes[..n]);
				Poll::Ready(Ok(Some(n)))
			}
			Read::Finished => {
				self.finished = true;
				Poll::Ready(Ok(None))
			}
			Read::Blocked => {
				if let Some(err) = self.shared.closed() {
					return Poll::Ready(Err(err));
				}
				self.shared.park_readable(self.id, waiter);
				Poll::Pending
			}
			Read::Reset(code) => Poll::Ready(Err(Error::Reset(code))),
		}
	}

	fn stop(&mut self, code: u32) {
		// Giving up on the read side abandons whatever was read ahead, even
		// when the FIN is already in and only the backlog is left.
		self.stop_code(u64::from(code));
	}

	fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
		let waiter = self.park.hold(cx);
		if self.finished || self.stopped {
			return Poll::Ready(Ok(()));
		}
		// The FIN sits behind whatever the peer sent before it, and noq only
		// reports the stream finished once that is read out. Waiting on
		// readability alone would park behind bytes nobody is reading, so read
		// ahead into the backlog `poll_read` serves first: this watch resolves
		// without the application draining the stream, and without losing what
		// it might still want.
		loop {
			if self.backlog.len() >= READ_AHEAD {
				// Enough held: reading further would let the peer send more
				// still, so the memory bound wins over watch liveness here.
				// `drain` wakes this once the application takes some, and a
				// stream nobody reads past the cap keeps its watch pending
				// until the connection's idle timeout.
				self.shared.park_readable(self.id, waiter);
				return Poll::Pending;
			}
			match Self::read(&self.shared, self.id, READ_CHUNK) {
				Read::Chunk(bytes) => self.backlog.extend_from_slice(&bytes),
				Read::Finished => {
					self.finished = true;
					return Poll::Ready(Ok(()));
				}
				Read::Blocked => {
					if let Some(err) = self.shared.closed() {
						return Poll::Ready(Err(err));
					}
					self.shared.park_readable(self.id, waiter);
					return Poll::Pending;
				}
				Read::Reset(code) => return Poll::Ready(Err(Error::Reset(code))),
			}
		}
	}
}

impl Drop for RecvStream {
	fn drop(&mut self) {
		self.shared.forget_recv(self.id);
		if !self.finished && !self.stopped {
			let _ = self
				.shared
				.conn
				.borrow_mut()
				.recv_stream(self.id)
				.stop(VarInt::from_u32(0));
			self.shared.kick();
		}
	}
}

impl std::fmt::Debug for RecvStream {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		f.debug_struct("RecvStream").field("id", &self.id).finish()
	}
}