#![allow(dead_code)]
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use bytes::{Buf, Bytes};
use tokio::sync::{mpsc, oneshot};
use h3::quic::{self, ConnectionErrorIncoming, StreamErrorIncoming, StreamId, WriteBuf};
use crate::buffer::{SendAccounting, TerminalCell, WriteCompletion, WriteOutcome};
use crate::driver::{BidiHandoff, ConnShared, DriverCommand, RecvHandoff, SendHandoff};
use crate::error::{internal_stream_error, ConnTerminal, RecvEnd, SendEnd};
fn stream_id(id: u64) -> StreamId {
StreamId::try_from(id).expect("worker allocates only valid QUIC stream ids")
}
fn conn_terminal_stream_err(term: &Arc<ConnTerminal>) -> StreamErrorIncoming {
StreamErrorIncoming::ConnectionErrorIncoming {
connection_error: term.to_h3(),
}
}
pub struct H3RecvStream<B: Buf> {
id: u64,
bytes: mpsc::Receiver<Bytes>,
terminal: TerminalCell<RecvEnd>,
resume: Arc<AtomicBool>,
blocked: Arc<AtomicBool>,
cmd_tx: mpsc::UnboundedSender<DriverCommand<B>>,
terminal_seen: bool,
stop_sent: bool,
}
impl<B: Buf> H3RecvStream<B> {
pub(crate) fn from_handoff(h: RecvHandoff<B>) -> Self {
h.cleanup.disarm();
H3RecvStream {
id: h.id,
bytes: h.bytes,
terminal: h.terminal,
resume: h.resume,
blocked: h.blocked,
cmd_tx: h.cmd_tx,
terminal_seen: false,
stop_sent: false,
}
}
fn signal_resume(&self) {
if self.blocked.swap(false, Ordering::AcqRel) && !self.resume.swap(true, Ordering::Relaxed)
{
let _ = self.cmd_tx.send(DriverCommand::RecvResume { id: self.id });
}
}
fn resolve_terminal(
&mut self,
end: RecvEnd,
) -> Poll<Result<Option<Bytes>, StreamErrorIncoming>> {
self.terminal_seen = true;
match end.to_h3() {
None => Poll::Ready(Ok(None)),
Some(err) => Poll::Ready(Err(err)),
}
}
}
impl<B: Buf> quic::RecvStream for H3RecvStream<B> {
type Buf = Bytes;
fn poll_data(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Option<Self::Buf>, StreamErrorIncoming>> {
match self.bytes.poll_recv(cx) {
Poll::Ready(Some(b)) => {
self.signal_resume();
Poll::Ready(Ok(Some(b)))
}
Poll::Ready(None) => {
match self.terminal.poll(cx) {
Poll::Ready(end) => self.resolve_terminal(end),
Poll::Pending => Poll::Ready(Err(internal_stream_error(
"recv byte channel closed without a published terminal",
))),
}
}
Poll::Pending => {
match self.terminal.poll(cx) {
Poll::Ready(end) => {
if let Poll::Ready(Some(b)) = self.bytes.poll_recv(cx) {
self.signal_resume();
return Poll::Ready(Ok(Some(b)));
}
self.resolve_terminal(end)
}
Poll::Pending => Poll::Pending,
}
}
}
}
fn stop_sending(&mut self, error_code: u64) {
self.stop_sent = true;
let _ = self.cmd_tx.send(DriverCommand::StopSending {
id: self.id,
code: error_code,
});
}
fn recv_id(&self) -> StreamId {
stream_id(self.id)
}
}
impl<B: Buf> Drop for H3RecvStream<B> {
fn drop(&mut self) {
if self.stop_sent || self.terminal_seen || self.terminal.get().is_some() {
return;
}
let _ = self.cmd_tx.send(DriverCommand::StopSending {
id: self.id,
code: 0,
});
}
}
pub struct H3SendStream<B: Buf> {
id: u64,
status: TerminalCell<SendEnd>,
cmd_tx: mpsc::UnboundedSender<DriverCommand<B>>,
stash: Option<WriteBuf<B>>,
write_completion: WriteCompletion<SendEnd>,
send_gen: Option<u64>,
finish_completion: Option<oneshot::Receiver<Result<(), SendEnd>>>,
finish_result: Option<Result<(), SendEnd>>,
finalized: bool,
local_terminal: Option<SendEnd>,
send_accounting: Arc<SendAccounting>,
}
impl<B: Buf> H3SendStream<B> {
pub(crate) fn from_handoff(h: SendHandoff<B>) -> Self {
h.cleanup.disarm();
H3SendStream {
id: h.id,
status: h.status,
cmd_tx: h.cmd_tx,
stash: None,
write_completion: WriteCompletion::new(),
send_gen: None,
finish_completion: None,
finish_result: None,
finalized: false,
local_terminal: None,
send_accounting: h.send_accounting,
}
}
fn terminal_now(&self, cx: &mut Context<'_>) -> Option<SendEnd> {
if let Some(end) = &self.local_terminal {
return Some(end.clone());
}
match self.status.poll(cx) {
Poll::Ready(end) => Some(end),
Poll::Pending => None,
}
}
fn terminal_now_noctx(&self) -> Option<SendEnd> {
self.local_terminal.clone().or_else(|| self.status.get())
}
fn sticky_or_internal(&self, cx: &mut Context<'_>, msg: &'static str) -> StreamErrorIncoming {
match self.terminal_now(cx) {
Some(end) => end.to_h3(),
None => internal_stream_error(msg),
}
}
#[cfg(test)]
pub(crate) fn write_generation(&self) -> u64 {
self.write_completion.generation()
}
fn resolve_write(
&self,
outcome: WriteOutcome<SendEnd>,
cx: &mut Context<'_>,
) -> Result<(), StreamErrorIncoming> {
match outcome {
WriteOutcome::Done(result) => result.map_err(|e| e.to_h3()),
WriteOutcome::Cancelled => {
Err(self.sticky_or_internal(cx, "send completion cancelled without a terminal"))
}
}
}
fn sticky_send_end_or_internal(&self, cx: &mut Context<'_>, msg: &'static str) -> SendEnd {
self.terminal_now(cx)
.unwrap_or_else(|| SendEnd::Conn(Arc::new(ConnTerminal::Internal(msg))))
}
}
impl<B: Buf> quic::SendStream<B> for H3SendStream<B> {
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), StreamErrorIncoming>> {
if let Some(generation) = self.send_gen {
match self.write_completion.poll(generation, cx) {
Poll::Ready(outcome) => {
self.send_gen = None;
return Poll::Ready(self.resolve_write(outcome, cx));
}
Poll::Pending => return Poll::Pending,
}
}
if let Some(end) = self.terminal_now(cx) {
return Poll::Ready(Err(end.to_h3()));
}
let buf = match self.stash.take() {
None => return Poll::Ready(Ok(())),
Some(buf) => buf,
};
let bytes = buf.remaining();
let permit = match self.send_accounting.try_reserve(bytes) {
Some(permit) => permit,
None => {
self.send_accounting.register_waiter(self.id, cx.waker());
match self.send_accounting.try_reserve(bytes) {
Some(permit) => permit,
None => {
self.stash = Some(buf);
return Poll::Pending;
}
}
}
};
self.send_accounting.unregister_waiter(self.id);
let generation = self.write_completion.begin();
let done = self.write_completion.completer(generation);
if self
.cmd_tx
.send(DriverCommand::Send {
id: self.id,
buf,
done,
permit: Some(permit),
})
.is_err()
{
return Poll::Ready(Err(
self.sticky_or_internal(cx, "send channel closed without a terminal")
));
}
self.send_gen = Some(generation);
match self.write_completion.poll(generation, cx) {
Poll::Ready(outcome) => {
self.send_gen = None;
Poll::Ready(self.resolve_write(outcome, cx))
}
Poll::Pending => Poll::Pending,
}
}
fn send_data<T: Into<WriteBuf<B>>>(&mut self, data: T) -> Result<(), StreamErrorIncoming> {
if self.stash.is_some() {
return Err(internal_stream_error(
"send_data called while a previous write is still pending poll_ready",
));
}
self.stash = Some(data.into());
Ok(())
}
fn poll_finish(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), StreamErrorIncoming>> {
if let Some(result) = &self.finish_result {
return Poll::Ready(result.clone().map_err(|e| e.to_h3()));
}
if self.finish_completion.is_some() {
match Pin::new(self.finish_completion.as_mut().unwrap()).poll(cx) {
Poll::Ready(Ok(result)) => {
self.finish_completion = None;
self.finish_result = Some(result.clone());
return Poll::Ready(result.map_err(|e| e.to_h3()));
}
Poll::Ready(Err(_)) => {
self.finish_completion = None;
let end = self.sticky_send_end_or_internal(
cx,
"finish completion cancelled without a terminal",
);
self.finish_result = Some(Err(end.clone()));
return Poll::Ready(Err(end.to_h3()));
}
Poll::Pending => return Poll::Pending,
}
}
if self.finalized {
return Poll::Ready(match self.terminal_now(cx) {
Some(end) => Err(end.to_h3()),
None => Ok(()),
});
}
if let Some(end) = self.terminal_now(cx) {
self.finalized = true;
self.finish_result = Some(Err(end.clone()));
return Poll::Ready(Err(end.to_h3()));
}
let (done_tx, done_rx) = oneshot::channel();
self.finalized = true;
if self
.cmd_tx
.send(DriverCommand::Finish {
id: self.id,
done: done_tx,
})
.is_err()
{
let end =
self.sticky_send_end_or_internal(cx, "finish channel closed without a terminal");
self.finish_result = Some(Err(end.clone()));
return Poll::Ready(Err(end.to_h3()));
}
self.finish_completion = Some(done_rx);
match Pin::new(self.finish_completion.as_mut().unwrap()).poll(cx) {
Poll::Ready(Ok(result)) => {
self.finish_completion = None;
self.finish_result = Some(result.clone());
Poll::Ready(result.map_err(|e| e.to_h3()))
}
Poll::Ready(Err(_)) => {
self.finish_completion = None;
let end = self.sticky_send_end_or_internal(
cx,
"finish completion cancelled without a terminal",
);
self.finish_result = Some(Err(end.clone()));
Poll::Ready(Err(end.to_h3()))
}
Poll::Pending => Poll::Pending,
}
}
fn reset(&mut self, reset_code: u64) {
if self.finalized {
return;
}
self.finalized = true;
if self.status.get().is_none() {
self.local_terminal = Some(SendEnd::Reset {
error_code: reset_code,
});
}
let _ = self.cmd_tx.send(DriverCommand::Reset {
id: self.id,
code: reset_code,
});
}
fn send_id(&self) -> StreamId {
stream_id(self.id)
}
}
impl<B: Buf> Drop for H3SendStream<B> {
fn drop(&mut self) {
self.send_accounting.unregister_waiter(self.id);
if self.finalized || self.terminal_now_noctx().is_some() {
return;
}
self.finalized = true;
let (done_tx, _done_rx) = oneshot::channel();
let _ = self.cmd_tx.send(DriverCommand::Finish {
id: self.id,
done: done_tx,
});
}
}
pub struct H3Stream<B: Buf> {
send: H3SendStream<B>,
recv: H3RecvStream<B>,
}
impl<B: Buf> H3Stream<B> {
pub(crate) fn from_handoff(h: BidiHandoff<B>) -> Self {
H3Stream {
send: H3SendStream::from_handoff(h.send),
recv: H3RecvStream::from_handoff(h.recv),
}
}
}
impl<B: Buf> quic::SendStream<B> for H3Stream<B> {
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), StreamErrorIncoming>> {
self.send.poll_ready(cx)
}
fn send_data<T: Into<WriteBuf<B>>>(&mut self, data: T) -> Result<(), StreamErrorIncoming> {
self.send.send_data(data)
}
fn poll_finish(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), StreamErrorIncoming>> {
self.send.poll_finish(cx)
}
fn reset(&mut self, reset_code: u64) {
self.send.reset(reset_code)
}
fn send_id(&self) -> StreamId {
self.send.send_id()
}
}
impl<B: Buf> quic::RecvStream for H3Stream<B> {
type Buf = Bytes;
fn poll_data(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Option<Self::Buf>, StreamErrorIncoming>> {
self.recv.poll_data(cx)
}
fn stop_sending(&mut self, error_code: u64) {
self.recv.stop_sending(error_code)
}
fn recv_id(&self) -> StreamId {
self.recv.recv_id()
}
}
impl<B: Buf> quic::BidiStream<B> for H3Stream<B> {
type SendStream = H3SendStream<B>;
type RecvStream = H3RecvStream<B>;
fn split(self) -> (Self::SendStream, Self::RecvStream) {
(self.send, self.recv)
}
}
pub struct StreamOpener<B: Buf> {
cmd_tx: mpsc::UnboundedSender<DriverCommand<B>>,
shared: Arc<ConnShared>,
pending_bidi: Option<oneshot::Receiver<Result<BidiHandoff<B>, Arc<ConnTerminal>>>>,
pending_uni: Option<oneshot::Receiver<Result<SendHandoff<B>, Arc<ConnTerminal>>>>,
}
impl<B: Buf> StreamOpener<B> {
pub(crate) fn from_parts(
cmd_tx: mpsc::UnboundedSender<DriverCommand<B>>,
shared: Arc<ConnShared>,
) -> Self {
StreamOpener {
cmd_tx,
shared,
pending_bidi: None,
pending_uni: None,
}
}
fn submit_terminal(&self) -> StreamErrorIncoming {
match self.shared.conn_terminal.get() {
Some(term) => conn_terminal_stream_err(&term),
None => internal_stream_error("open declined without a published terminal"),
}
}
}
impl<B: Buf> Clone for StreamOpener<B> {
fn clone(&self) -> Self {
StreamOpener {
cmd_tx: self.cmd_tx.clone(),
shared: Arc::clone(&self.shared),
pending_bidi: None,
pending_uni: None,
}
}
}
impl<B: Buf> quic::OpenStreams<B> for StreamOpener<B> {
type BidiStream = H3Stream<B>;
type SendStream = H3SendStream<B>;
fn poll_open_bidi(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Self::BidiStream, StreamErrorIncoming>> {
if self.pending_bidi.is_none() {
if let Some(term) = self.shared.conn_terminal.get() {
return Poll::Ready(Err(conn_terminal_stream_err(&term)));
}
let (reply_tx, reply_rx) = oneshot::channel();
if self
.cmd_tx
.send(DriverCommand::OpenBidi { reply: reply_tx })
.is_err()
{
return Poll::Ready(Err(self.submit_terminal()));
}
self.pending_bidi = Some(reply_rx);
}
match Pin::new(self.pending_bidi.as_mut().unwrap()).poll(cx) {
Poll::Ready(Ok(Ok(handoff))) => {
self.pending_bidi = None;
Poll::Ready(Ok(H3Stream::from_handoff(handoff)))
}
Poll::Ready(Ok(Err(term))) => {
self.pending_bidi = None;
Poll::Ready(Err(conn_terminal_stream_err(&term)))
}
Poll::Ready(Err(_)) => {
self.pending_bidi = None;
Poll::Ready(Err(self.submit_terminal()))
}
Poll::Pending => Poll::Pending,
}
}
fn poll_open_send(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Self::SendStream, StreamErrorIncoming>> {
if self.pending_uni.is_none() {
if let Some(term) = self.shared.conn_terminal.get() {
return Poll::Ready(Err(conn_terminal_stream_err(&term)));
}
let (reply_tx, reply_rx) = oneshot::channel();
if self
.cmd_tx
.send(DriverCommand::OpenUni { reply: reply_tx })
.is_err()
{
return Poll::Ready(Err(self.submit_terminal()));
}
self.pending_uni = Some(reply_rx);
}
match Pin::new(self.pending_uni.as_mut().unwrap()).poll(cx) {
Poll::Ready(Ok(Ok(handoff))) => {
self.pending_uni = None;
Poll::Ready(Ok(H3SendStream::from_handoff(handoff)))
}
Poll::Ready(Ok(Err(term))) => {
self.pending_uni = None;
Poll::Ready(Err(conn_terminal_stream_err(&term)))
}
Poll::Ready(Err(_)) => {
self.pending_uni = None;
Poll::Ready(Err(self.submit_terminal()))
}
Poll::Pending => Poll::Pending,
}
}
fn close(&mut self, code: h3::error::Code, reason: &[u8]) {
let _ = self.cmd_tx.send(DriverCommand::Close {
code: code.value(),
reason: Bytes::copy_from_slice(reason),
});
}
}
pub struct Connection<B: Buf> {
accept_bidi_rx: mpsc::Receiver<BidiHandoff<B>>,
accept_uni_rx: mpsc::Receiver<RecvHandoff<B>>,
accept_terminal_bidi: TerminalCell<Arc<ConnTerminal>>,
accept_terminal_uni: TerminalCell<Arc<ConnTerminal>>,
accept_bidi_resume: Arc<AtomicBool>,
accept_uni_resume: Arc<AtomicBool>,
opener: StreamOpener<B>,
}
impl<B: Buf> Connection<B> {
#[allow(clippy::too_many_arguments)]
pub(crate) fn from_parts(
accept_bidi_rx: mpsc::Receiver<BidiHandoff<B>>,
accept_uni_rx: mpsc::Receiver<RecvHandoff<B>>,
accept_terminal_bidi: TerminalCell<Arc<ConnTerminal>>,
accept_terminal_uni: TerminalCell<Arc<ConnTerminal>>,
accept_bidi_resume: Arc<AtomicBool>,
accept_uni_resume: Arc<AtomicBool>,
opener: StreamOpener<B>,
) -> Self {
Connection {
accept_bidi_rx,
accept_uni_rx,
accept_terminal_bidi,
accept_terminal_uni,
accept_bidi_resume,
accept_uni_resume,
opener,
}
}
fn signal_accept_bidi_resume(&self) {
if !self.accept_bidi_resume.swap(true, Ordering::Relaxed) {
let _ = self.opener.cmd_tx.send(DriverCommand::AcceptBidiResume);
}
}
fn signal_accept_uni_resume(&self) {
if !self.accept_uni_resume.swap(true, Ordering::Relaxed) {
let _ = self.opener.cmd_tx.send(DriverCommand::AcceptUniResume);
}
}
}
impl<B: Buf> quic::OpenStreams<B> for Connection<B> {
type BidiStream = H3Stream<B>;
type SendStream = H3SendStream<B>;
fn poll_open_bidi(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Self::BidiStream, StreamErrorIncoming>> {
self.opener.poll_open_bidi(cx)
}
fn poll_open_send(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Self::SendStream, StreamErrorIncoming>> {
self.opener.poll_open_send(cx)
}
fn close(&mut self, code: h3::error::Code, reason: &[u8]) {
self.opener.close(code, reason)
}
}
impl<B: Buf> quic::Connection<B> for Connection<B> {
type RecvStream = H3RecvStream<B>;
type OpenStreams = StreamOpener<B>;
fn poll_accept_recv(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Self::RecvStream, ConnectionErrorIncoming>> {
match self.accept_uni_rx.poll_recv(cx) {
Poll::Ready(Some(handoff)) => {
self.signal_accept_uni_resume();
Poll::Ready(Ok(H3RecvStream::from_handoff(handoff)))
}
Poll::Ready(None) => match self.accept_terminal_uni.poll(cx) {
Poll::Ready(term) => Poll::Ready(Err(term.to_h3())),
Poll::Pending => Poll::Ready(Err(ConnectionErrorIncoming::InternalError(
"uni accept channel closed without a published terminal".to_string(),
))),
},
Poll::Pending => match self.accept_terminal_uni.poll(cx) {
Poll::Ready(term) => {
if let Poll::Ready(Some(handoff)) = self.accept_uni_rx.poll_recv(cx) {
self.signal_accept_uni_resume();
return Poll::Ready(Ok(H3RecvStream::from_handoff(handoff)));
}
Poll::Ready(Err(term.to_h3()))
}
Poll::Pending => Poll::Pending,
},
}
}
fn poll_accept_bidi(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Self::BidiStream, ConnectionErrorIncoming>> {
match self.accept_bidi_rx.poll_recv(cx) {
Poll::Ready(Some(handoff)) => {
self.signal_accept_bidi_resume();
Poll::Ready(Ok(H3Stream::from_handoff(handoff)))
}
Poll::Ready(None) => match self.accept_terminal_bidi.poll(cx) {
Poll::Ready(term) => Poll::Ready(Err(term.to_h3())),
Poll::Pending => Poll::Ready(Err(ConnectionErrorIncoming::InternalError(
"bidi accept channel closed without a published terminal".to_string(),
))),
},
Poll::Pending => match self.accept_terminal_bidi.poll(cx) {
Poll::Ready(term) => {
if let Poll::Ready(Some(handoff)) = self.accept_bidi_rx.poll_recv(cx) {
self.signal_accept_bidi_resume();
return Poll::Ready(Ok(H3Stream::from_handoff(handoff)));
}
Poll::Ready(Err(term.to_h3()))
}
Poll::Pending => Poll::Pending,
},
}
}
fn opener(&self) -> Self::OpenStreams {
self.opener.clone()
}
}
impl<B: Buf> Drop for Connection<B> {
fn drop(&mut self) {
let _ = self.opener.cmd_tx.send(DriverCommand::ConnectionDropped);
}
}
fn _assert_h3_traits<B: Buf>() {
fn is_connection<B: Buf, T: quic::Connection<B>>() {}
fn is_open_streams<B: Buf, T: quic::OpenStreams<B>>() {}
fn is_bidi_stream<B: Buf, T: quic::BidiStream<B>>() {}
fn is_send_stream<B: Buf, T: quic::SendStream<B>>() {}
fn is_recv_stream<T: quic::RecvStream>() {}
is_connection::<B, Connection<B>>();
is_open_streams::<B, StreamOpener<B>>();
is_bidi_stream::<B, H3Stream<B>>();
is_send_stream::<B, H3SendStream<B>>();
is_recv_stream::<H3RecvStream<B>>();
}
#[cfg(test)]
mod tests {
use super::*;
use crate::error::CloseOrigin;
use h3::quic::{Connection as _, OpenStreams as _, RecvStream as _, SendStream as _};
use std::task::{RawWaker, RawWakerVTable, Waker};
fn noop_cx() -> Context<'static> {
Context::from_waker(noop_waker_ref())
}
fn noop_waker_ref() -> &'static Waker {
static VTABLE: RawWakerVTable = RawWakerVTable::new(
|_| RawWaker::new(std::ptr::null(), &VTABLE),
|_| {},
|_| {},
|_| {},
);
static WAKER: std::sync::OnceLock<Waker> = std::sync::OnceLock::new();
WAKER.get_or_init(|| unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) })
}
fn flag_waker(flag: Arc<AtomicBool>) -> Waker {
let ptr = Arc::into_raw(flag) as *const ();
unsafe { Waker::from_raw(RawWaker::new(ptr, &FLAG_VTABLE)) }
}
static FLAG_VTABLE: RawWakerVTable = RawWakerVTable::new(
|p| unsafe {
let arc = Arc::from_raw(p as *const AtomicBool);
let cloned = arc.clone();
std::mem::forget(arc);
RawWaker::new(Arc::into_raw(cloned) as *const (), &FLAG_VTABLE)
},
|p| unsafe {
let arc = Arc::from_raw(p as *const AtomicBool);
arc.store(true, std::sync::atomic::Ordering::SeqCst);
},
|p| unsafe {
let arc = Arc::from_raw(p as *const AtomicBool);
arc.store(true, std::sync::atomic::Ordering::SeqCst);
std::mem::forget(arc);
},
|p| unsafe {
drop(Arc::from_raw(p as *const AtomicBool));
},
);
#[allow(clippy::type_complexity)]
fn recv_channel() -> (
mpsc::Sender<Bytes>,
TerminalCell<RecvEnd>,
Arc<AtomicBool>,
Arc<AtomicBool>,
H3RecvStream<Bytes>,
mpsc::UnboundedReceiver<DriverCommand<Bytes>>,
) {
let (btx, brx) = mpsc::channel(4);
let (ctx, crx) = mpsc::unbounded_channel();
let terminal = TerminalCell::new();
let resume = Arc::new(AtomicBool::new(false));
let blocked = Arc::new(AtomicBool::new(false));
let recv = H3RecvStream::from_handoff(RecvHandoff {
id: 0,
bytes: brx,
terminal: terminal.clone(),
resume: Arc::clone(&resume),
blocked: Arc::clone(&blocked),
cmd_tx: ctx.clone(),
cleanup: crate::driver::HandoffCleanup::new(0, true, ctx),
});
(btx, terminal, resume, blocked, recv, crx)
}
fn send_half(
id: u64,
) -> (
TerminalCell<SendEnd>,
H3SendStream<Bytes>,
mpsc::UnboundedReceiver<DriverCommand<Bytes>>,
) {
send_half_with(id, SendAccounting::new(None))
}
fn send_half_with(
id: u64,
accounting: Arc<SendAccounting>,
) -> (
TerminalCell<SendEnd>,
H3SendStream<Bytes>,
mpsc::UnboundedReceiver<DriverCommand<Bytes>>,
) {
let (ctx, crx) = mpsc::unbounded_channel();
let status = TerminalCell::new();
let send = H3SendStream::from_handoff(SendHandoff {
id,
status: status.clone(),
cmd_tx: ctx.clone(),
send_accounting: accounting,
cleanup: crate::driver::HandoffCleanup::new(id, false, ctx),
});
(status, send, crx)
}
fn wbuf(payload: &'static [u8]) -> WriteBuf<Bytes> {
WriteBuf::from(h3::proto::frame::Frame::Data(Bytes::from_static(payload)))
}
fn wire_len(payload: &'static [u8]) -> usize {
wbuf(payload).remaining()
}
#[test]
fn sf6_enqueue_failure_rolls_back_reserved_bytes() {
let acct = SendAccounting::new(Some(1024));
let (status, mut send, crx) = send_half_with(0, Arc::clone(&acct));
drop(crx);
status.set(SendEnd::Reset { error_code: 9 });
let mut cx = noop_cx();
send.send_data(wbuf(b"hello")).unwrap();
assert_eq!(acct.resident(), 0, "nothing reserved until the flush");
match send.poll_ready(&mut cx) {
Poll::Ready(Err(_)) => {}
other => panic!("expected terminal error on closed channel, got {other:?}"),
}
assert_eq!(
acct.resident(),
0,
"a failed enqueue must not leak the reserved bytes"
);
}
#[test]
fn poll_data_delivers_buffered_bytes_before_terminal() {
let (btx, terminal, _resume, _blocked, mut recv, _crx) = recv_channel();
btx.try_send(Bytes::from_static(b"hi")).unwrap();
terminal.set(RecvEnd::Fin);
let mut cx = noop_cx();
match recv.poll_data(&mut cx) {
Poll::Ready(Ok(Some(b))) => assert_eq!(&b[..], b"hi"),
other => panic!("expected buffered bytes first, got {other:?}"),
}
assert!(matches!(recv.poll_data(&mut cx), Poll::Ready(Ok(None))));
}
#[test]
fn poll_data_maps_fin_reset_conn() {
let mut cx = noop_cx();
{
let (_btx, terminal, _r, _blocked, mut recv, _c) = recv_channel();
terminal.set(RecvEnd::Fin);
assert!(matches!(recv.poll_data(&mut cx), Poll::Ready(Ok(None))));
}
{
let (_btx, terminal, _r, _blocked, mut recv, _c) = recv_channel();
terminal.set(RecvEnd::Reset { error_code: 42 });
match recv.poll_data(&mut cx) {
Poll::Ready(Err(StreamErrorIncoming::StreamTerminated { error_code })) => {
assert_eq!(error_code, 42)
}
other => panic!("expected StreamTerminated, got {other:?}"),
}
}
{
let (_btx, terminal, _r, _blocked, mut recv, _c) = recv_channel();
terminal.set(RecvEnd::Conn(Arc::new(ConnTerminal::Timeout)));
match recv.poll_data(&mut cx) {
Poll::Ready(Err(StreamErrorIncoming::ConnectionErrorIncoming {
connection_error: ConnectionErrorIncoming::Timeout,
})) => {}
other => panic!("expected ConnectionErrorIncoming::Timeout, got {other:?}"),
}
}
}
#[test]
fn poll_data_closed_channel_without_terminal_is_internal_error() {
let (btx, _terminal, _r, _blocked, mut recv, _c) = recv_channel();
drop(btx); let mut cx = noop_cx();
match recv.poll_data(&mut cx) {
Poll::Ready(Err(StreamErrorIncoming::ConnectionErrorIncoming {
connection_error: ConnectionErrorIncoming::InternalError(_),
})) => {}
other => panic!("expected InternalError, got {other:?}"),
}
}
#[test]
fn recv_resume_gated_when_worker_never_blocked() {
let (btx, _terminal, resume, blocked, mut recv, mut crx) = recv_channel();
assert!(!blocked.load(Ordering::Relaxed));
btx.try_send(Bytes::from_static(b"a")).unwrap();
btx.try_send(Bytes::from_static(b"b")).unwrap();
let mut cx = noop_cx();
assert!(matches!(recv.poll_data(&mut cx), Poll::Ready(Ok(Some(_)))));
assert!(matches!(recv.poll_data(&mut cx), Poll::Ready(Ok(Some(_)))));
assert!(!resume.load(Ordering::Relaxed));
assert!(
crx.try_recv().is_err(),
"must not emit RecvResume when worker never blocked"
);
}
#[test]
fn recv_resume_sent_once_when_worker_blocked() {
let (btx, _terminal, resume, blocked, mut recv, mut crx) = recv_channel();
blocked.store(true, Ordering::Release);
btx.try_send(Bytes::from_static(b"a")).unwrap();
btx.try_send(Bytes::from_static(b"b")).unwrap();
let mut cx = noop_cx();
assert!(matches!(recv.poll_data(&mut cx), Poll::Ready(Ok(Some(_)))));
assert!(resume.load(Ordering::Relaxed));
assert!(
!blocked.load(Ordering::Relaxed),
"park flag must be cleared"
);
match crx.try_recv() {
Ok(DriverCommand::RecvResume { id: 0 }) => {}
other => panic!("expected one RecvResume, got {other:?}"),
}
assert!(matches!(recv.poll_data(&mut cx), Poll::Ready(Ok(Some(_)))));
assert!(crx.try_recv().is_err(), "must not resend RecvResume");
}
#[test]
fn recv_drop_enqueues_stop_sending_zero() {
let (_btx, _terminal, _r, _blocked, recv, mut crx) = recv_channel();
drop(recv);
match crx.try_recv() {
Ok(DriverCommand::StopSending { id: 0, code: 0 }) => {}
other => panic!("expected StopSending(0), got {other:?}"),
}
}
#[test]
fn recv_drop_after_terminal_does_not_stop_send() {
let (_btx, terminal, _r, _blocked, recv, mut crx) = recv_channel();
terminal.set(RecvEnd::Fin);
drop(recv);
assert!(
crx.try_recv().is_err(),
"terminal recv must not stop-send on drop"
);
}
#[test]
fn send_data_single_slot_errors_on_double_stash() {
let (_status, mut send, _crx) = send_half(0);
assert!(send.send_data(wbuf(b"one")).is_ok());
match send.send_data(wbuf(b"two")) {
Err(StreamErrorIncoming::ConnectionErrorIncoming {
connection_error: ConnectionErrorIncoming::InternalError(_),
}) => {}
other => panic!("expected InternalError on double stash, got {other:?}"),
}
}
#[test]
fn poll_ready_returns_recorded_completion_once_then_sticky() {
let (status, mut send, mut crx) = send_half(0);
let mut cx = noop_cx();
assert!(matches!(send.poll_ready(&mut cx), Poll::Ready(Ok(()))));
send.send_data(wbuf(b"body")).unwrap();
assert!(matches!(send.poll_ready(&mut cx), Poll::Pending));
let done = match crx.try_recv() {
Ok(DriverCommand::Send { id: 0, done, .. }) => done,
other => panic!("expected Send, got {other:?}"),
};
done.complete(Ok(()));
status.set(SendEnd::Stopped { error_code: 7 });
assert!(matches!(send.poll_ready(&mut cx), Poll::Ready(Ok(()))));
match send.poll_ready(&mut cx) {
Poll::Ready(Err(StreamErrorIncoming::StreamTerminated { error_code: 7 })) => {}
other => panic!("expected sticky StreamTerminated, got {other:?}"),
}
}
#[test]
fn poll_ready_reuses_one_completion_cell_across_writes() {
let (_status, mut send, mut crx) = send_half(0);
let mut cx = noop_cx();
const K: u64 = 6;
for expected_gen in 1..=K {
send.send_data(wbuf(b"chunk")).unwrap();
assert!(matches!(send.poll_ready(&mut cx), Poll::Pending));
assert_eq!(
send.write_generation(),
expected_gen,
"one generation bump per write — the cell is reused, not reallocated"
);
let done = match crx.try_recv() {
Ok(DriverCommand::Send { id: 0, done, .. }) => done,
other => panic!("expected Send, got {other:?}"),
};
done.complete(Ok(()));
assert!(matches!(send.poll_ready(&mut cx), Poll::Ready(Ok(()))));
assert!(matches!(send.poll_ready(&mut cx), Poll::Ready(Ok(()))));
}
assert_eq!(
send.write_generation(),
K,
"one cell reused for all K writes"
);
}
#[test]
fn sf6_unlimited_accounting_tracks_and_releases_bytes() {
let acct = SendAccounting::new(None);
let (_status, mut send, mut crx) = send_half_with(0, Arc::clone(&acct));
let mut cx = noop_cx();
assert_eq!(acct.resident(), 0);
send.send_data(wbuf(b"hello")).unwrap();
assert!(matches!(send.poll_ready(&mut cx), Poll::Pending));
let hello = wire_len(b"hello");
assert_eq!(acct.resident(), hello, "reserved on admission");
let (done, permit) = match crx.try_recv() {
Ok(DriverCommand::Send { done, permit, .. }) => (done, permit),
other => panic!("expected Send, got {other:?}"),
};
assert!(permit.is_some(), "front end carries a byte permit (SF-6)");
assert_eq!(permit.as_ref().unwrap().bytes(), hello);
assert_eq!(acct.resident(), hello);
done.complete(Ok(()));
drop(permit);
assert_eq!(acct.resident(), 0, "released once the permit dropped");
assert!(matches!(send.poll_ready(&mut cx), Poll::Ready(Ok(()))));
}
#[test]
fn sf6_capped_accounting_parks_then_admits_on_release() {
let hello = wire_len(b"hello");
let x = wire_len(b"x");
let acct = SendAccounting::new(Some(hello));
let (_sa, mut send_a, mut crx_a) = send_half_with(0, Arc::clone(&acct));
let (_sb, mut send_b, mut crx_b) = send_half_with(4, Arc::clone(&acct));
let mut cx_a = noop_cx();
send_a.send_data(wbuf(b"hello")).unwrap();
assert!(matches!(send_a.poll_ready(&mut cx_a), Poll::Pending));
assert_eq!(acct.resident(), hello);
let cmd_a = crx_a.try_recv().expect("A admitted");
let woken = Arc::new(AtomicBool::new(false));
let waker = flag_waker(woken.clone());
let mut cx_b = Context::from_waker(&waker);
send_b.send_data(wbuf(b"x")).unwrap();
assert!(matches!(send_b.poll_ready(&mut cx_b), Poll::Pending));
assert!(crx_b.try_recv().is_err(), "B must not enqueue over the cap");
assert_eq!(
acct.resident(),
hello,
"B's bytes not reserved while parked"
);
drop(cmd_a);
assert_eq!(acct.resident(), 0, "A released");
assert!(
woken.load(std::sync::atomic::Ordering::SeqCst),
"release woke B"
);
assert!(matches!(send_b.poll_ready(&mut cx_b), Poll::Pending));
assert_eq!(acct.resident(), x, "B admitted after A freed capacity");
match crx_b.try_recv() {
Ok(DriverCommand::Send { id: 4, permit, .. }) => {
assert_eq!(permit.as_ref().unwrap().bytes(), x);
}
other => panic!("expected B's Send after release, got {other:?}"),
}
}
#[test]
fn poll_ready_unapplied_send_resolves_via_sticky_terminal() {
let (status, mut send, mut crx) = send_half(0);
let mut cx = noop_cx();
send.send_data(wbuf(b"body")).unwrap();
assert!(matches!(send.poll_ready(&mut cx), Poll::Pending));
let done = match crx.try_recv() {
Ok(DriverCommand::Send { id: 0, done, .. }) => done,
other => panic!("expected Send, got {other:?}"),
};
drop(done); status.set(SendEnd::Stopped { error_code: 9 });
match send.poll_ready(&mut cx) {
Poll::Ready(Err(StreamErrorIncoming::StreamTerminated { error_code: 9 })) => {}
other => panic!("expected sticky StreamTerminated, got {other:?}"),
}
}
#[test]
fn poll_finish_idempotent_one_finish() {
let (_status, mut send, mut crx) = send_half(0);
let mut cx = noop_cx();
assert!(matches!(send.poll_finish(&mut cx), Poll::Pending));
let done = match crx.try_recv() {
Ok(DriverCommand::Finish { id: 0, done }) => done,
other => panic!("expected Finish, got {other:?}"),
};
assert!(matches!(send.poll_finish(&mut cx), Poll::Pending));
assert!(crx.try_recv().is_err(), "must not enqueue a second Finish");
done.send(Ok(())).unwrap();
assert!(matches!(send.poll_finish(&mut cx), Poll::Ready(Ok(()))));
assert!(matches!(send.poll_finish(&mut cx), Poll::Ready(Ok(()))));
assert!(crx.try_recv().is_err());
}
#[test]
fn poll_finish_failure_is_retained_not_success() {
let (_status, mut send, crx) = send_half(0);
drop(crx); let mut cx = noop_cx();
match send.poll_finish(&mut cx) {
Poll::Ready(Err(StreamErrorIncoming::ConnectionErrorIncoming {
connection_error: ConnectionErrorIncoming::InternalError(_),
})) => {}
other => panic!("expected InternalError on first poll, got {other:?}"),
}
match send.poll_finish(&mut cx) {
Poll::Ready(Err(StreamErrorIncoming::ConnectionErrorIncoming {
connection_error: ConnectionErrorIncoming::InternalError(_),
})) => {}
other => panic!("finalized failure must not become Ok, got {other:?}"),
}
}
#[test]
fn reset_enqueues_once_and_finalizes() {
let (_status, mut send, mut crx) = send_half(4);
send.reset(7);
match crx.try_recv() {
Ok(DriverCommand::Reset { id: 4, code: 7 }) => {}
other => panic!("expected Reset(7), got {other:?}"),
}
send.reset(9);
assert!(crx.try_recv().is_err(), "must not enqueue a second Reset");
let mut cx = noop_cx();
match send.poll_finish(&mut cx) {
Poll::Ready(Err(StreamErrorIncoming::StreamTerminated { error_code: 7 })) => {}
other => panic!("expected sticky reset terminal, got {other:?}"),
}
}
#[test]
fn send_drop_enqueues_graceful_finish() {
let (_status, send, mut crx) = send_half(0);
drop(send);
match crx.try_recv() {
Ok(DriverCommand::Finish { id: 0, .. }) => {}
other => panic!("expected graceful Finish on drop, got {other:?}"),
}
}
#[test]
fn send_drop_after_finalize_does_not_finish() {
let (_status, mut send, mut crx) = send_half(0);
send.reset(3);
let _ = crx.try_recv(); drop(send);
assert!(
crx.try_recv().is_err(),
"finalized send must not finish on drop"
);
}
#[test]
fn dropped_recv_handoff_enqueues_stop_sending() {
let (_btx, brx) = mpsc::channel(1);
let (ctx, mut crx) = mpsc::unbounded_channel::<DriverCommand<Bytes>>();
let handoff = RecvHandoff {
id: 8,
bytes: brx,
terminal: TerminalCell::new(),
resume: Arc::new(AtomicBool::new(false)),
blocked: Arc::new(AtomicBool::new(false)),
cmd_tx: ctx.clone(),
cleanup: crate::driver::HandoffCleanup::new(8, true, ctx),
};
drop(handoff); match crx.try_recv() {
Ok(DriverCommand::StopSending { id: 8, code: 0 }) => {}
other => panic!("expected StopSending on dropped handoff, got {other:?}"),
}
}
#[test]
fn dropped_send_handoff_enqueues_finish() {
let (ctx, mut crx) = mpsc::unbounded_channel::<DriverCommand<Bytes>>();
let handoff = SendHandoff {
id: 8,
status: TerminalCell::new(),
cmd_tx: ctx.clone(),
send_accounting: SendAccounting::new(None),
cleanup: crate::driver::HandoffCleanup::new(8, false, ctx),
};
drop(handoff);
match crx.try_recv() {
Ok(DriverCommand::Finish { id: 8, .. }) => {}
other => panic!("expected graceful Finish on dropped handoff, got {other:?}"),
}
}
#[test]
fn converted_handoff_disarms_guard() {
let (_btx, _terminal, _resume, _blocked, recv, mut crx) = recv_channel();
assert!(
crx.try_recv().is_err(),
"conversion must not fire the guard"
);
drop(recv);
assert!(
matches!(crx.try_recv(), Ok(DriverCommand::StopSending { .. })),
"stream Drop (not the disarmed guard) enqueues cleanup"
);
}
fn opener() -> (
StreamOpener<Bytes>,
mpsc::UnboundedReceiver<DriverCommand<Bytes>>,
Arc<ConnShared>,
) {
let (ctx, crx) = mpsc::unbounded_channel();
let shared = ConnShared::new(None);
(
StreamOpener::from_parts(ctx, Arc::clone(&shared)),
crx,
shared,
)
}
#[test]
fn stream_opener_submit_helper_resolves_terminal_when_conn_terminal_preset() {
let (mut op, mut crx, shared) = opener();
shared.conn_terminal.set(Arc::new(ConnTerminal::AppClose {
origin: CloseOrigin::Peer,
error_code: 0x101,
reason: Bytes::new(),
}));
let mut cx = noop_cx();
match op.poll_open_bidi(&mut cx) {
Poll::Ready(Err(StreamErrorIncoming::ConnectionErrorIncoming {
connection_error: ConnectionErrorIncoming::ApplicationClose { error_code: 0x101 },
})) => {}
_ => panic!("expected preset terminal resolution"),
}
assert!(
crx.try_recv().is_err(),
"must not submit under a preset terminal"
);
}
#[test]
fn cloned_opener_has_fresh_pending_slots() {
let (mut op, mut crx, _shared) = opener();
let mut cx = noop_cx();
assert!(matches!(op.poll_open_bidi(&mut cx), Poll::Pending));
assert!(op.pending_bidi.is_some());
assert!(matches!(crx.try_recv(), Ok(DriverCommand::OpenBidi { .. })));
let clone = op.clone();
assert!(clone.pending_bidi.is_none());
assert!(clone.pending_uni.is_none());
}
#[test]
fn opener_open_bidi_resolves_handoff_into_stream() {
let (mut op, mut crx, _shared) = opener();
let mut cx = noop_cx();
assert!(matches!(op.poll_open_bidi(&mut cx), Poll::Pending));
let reply = match crx.try_recv() {
Ok(DriverCommand::OpenBidi { reply }) => reply,
other => panic!("expected OpenBidi, got {other:?}"),
};
let (_btx, brx) = mpsc::channel(1);
let (ictx, _icrx) = mpsc::unbounded_channel();
let handoff = BidiHandoff {
send: SendHandoff {
id: 0,
status: TerminalCell::new(),
cmd_tx: ictx.clone(),
send_accounting: SendAccounting::new(None),
cleanup: crate::driver::HandoffCleanup::new(0, false, ictx.clone()),
},
recv: RecvHandoff {
id: 0,
bytes: brx,
terminal: TerminalCell::new(),
resume: Arc::new(AtomicBool::new(false)),
blocked: Arc::new(AtomicBool::new(false)),
cmd_tx: ictx.clone(),
cleanup: crate::driver::HandoffCleanup::new(0, true, ictx),
},
};
reply.send(Ok(handoff)).ok().expect("deliver handoff");
match op.poll_open_bidi(&mut cx) {
Poll::Ready(Ok(_stream)) => {}
_ => panic!("expected resolved H3Stream"),
}
assert!(op.pending_bidi.is_none(), "slot cleared after resolution");
}
#[allow(clippy::type_complexity)]
fn connection() -> (
Connection<Bytes>,
mpsc::Sender<BidiHandoff<Bytes>>,
TerminalCell<Arc<ConnTerminal>>,
Arc<AtomicBool>,
mpsc::UnboundedReceiver<DriverCommand<Bytes>>,
) {
let (btx, brx) = mpsc::channel(4);
let (_utx, urx) = mpsc::channel(4);
let (ctx, crx) = mpsc::unbounded_channel();
let at_bidi = TerminalCell::new();
let at_uni = TerminalCell::new();
let rb = Arc::new(AtomicBool::new(false));
let ru = Arc::new(AtomicBool::new(false));
let shared = ConnShared::new(None);
let opener = StreamOpener::from_parts(ctx, shared);
let conn = Connection::from_parts(
brx,
urx,
at_bidi.clone(),
at_uni,
Arc::clone(&rb),
ru,
opener,
);
(conn, btx, at_bidi, rb, crx)
}
fn make_bidi_handoff() -> BidiHandoff<Bytes> {
let (_btx, brx) = mpsc::channel(1);
let (ictx, _icrx) = mpsc::unbounded_channel();
BidiHandoff {
send: SendHandoff {
id: 0,
status: TerminalCell::new(),
cmd_tx: ictx.clone(),
send_accounting: SendAccounting::new(None),
cleanup: crate::driver::HandoffCleanup::new(0, false, ictx.clone()),
},
recv: RecvHandoff {
id: 0,
bytes: brx,
terminal: TerminalCell::new(),
resume: Arc::new(AtomicBool::new(false)),
blocked: Arc::new(AtomicBool::new(false)),
cmd_tx: ictx.clone(),
cleanup: crate::driver::HandoffCleanup::new(0, true, ictx),
},
}
}
#[test]
fn poll_accept_bidi_delivers_then_maps_terminal() {
let (mut conn, btx, at_bidi, rb, mut crx) = connection();
let mut cx = noop_cx();
btx.try_send(make_bidi_handoff()).unwrap();
match conn.poll_accept_bidi(&mut cx) {
Poll::Ready(Ok(_stream)) => {}
_ => panic!("expected accepted stream"),
}
assert!(rb.load(Ordering::Relaxed));
match crx.try_recv() {
Ok(DriverCommand::AcceptBidiResume) => {}
other => panic!("expected AcceptBidiResume, got {other:?}"),
}
at_bidi.set(Arc::new(ConnTerminal::Timeout));
match conn.poll_accept_bidi(&mut cx) {
Poll::Ready(Err(ConnectionErrorIncoming::Timeout)) => {}
_ => panic!("expected Timeout"),
}
}
#[test]
fn poll_accept_bidi_sealing_recheck_yields_queued_stream_before_terminal() {
let (mut conn, btx, at_bidi, _rb, _crx) = connection();
let mut cx = noop_cx();
btx.try_send(make_bidi_handoff()).unwrap();
at_bidi.set(Arc::new(ConnTerminal::Timeout));
match conn.poll_accept_bidi(&mut cx) {
Poll::Ready(Ok(_stream)) => {}
_ => panic!("expected queued stream ahead of terminal"),
}
}
#[test]
fn connection_drop_enqueues_connection_dropped() {
let (conn, _btx, _at, _rb, mut crx) = connection();
drop(conn);
match crx.try_recv() {
Ok(DriverCommand::ConnectionDropped) => {}
other => panic!("expected ConnectionDropped, got {other:?}"),
}
}
}