use std::ops::ControlFlow;
use std::sync::atomic::Ordering;
use tokio::io::{AsyncRead, AsyncWrite};
use crate::h2::error::{ErrorCode, H2Error};
use crate::h2::stream_state::StreamState;
use super::*;
impl<T: AsyncRead + AsyncWrite + Unpin + Send + 'static> Driver<T> {
pub(super) async fn run(mut self) -> Result<(), H2Error> {
let result = self.event_loop().await;
self.closed.store(true, Ordering::Release);
let final_err = match &result {
Ok(()) => H2Error::Connection {
code: ErrorCode::NoError,
reason: "connection closed".into(),
},
Err(e) => e.duplicate(),
};
for (_, mut actor) in self.streams.drain() {
if actor.remote_done {
actor.deliver_ok();
} else {
actor.deliver_err(final_err.duplicate());
}
}
while let Some(cmd) = self.pending.pop_front() {
if let Some(sink) = cmd.into_sink() {
send_err_to_sink(sink, final_err.duplicate());
}
}
while let Ok(cmd) = self.command_rx.try_recv() {
if let Some(sink) = cmd.into_sink() {
send_err_to_sink(sink, final_err.duplicate());
}
}
result
}
pub(super) async fn event_loop(&mut self) -> Result<(), H2Error> {
let mut sweep_tick = tokio::time::interval(std::time::Duration::from_millis(100));
sweep_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut flush_tick = tokio::time::interval(std::time::Duration::from_millis(1));
flush_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
if self.shutdown_started && self.streams.is_empty() && self.goaway_drained().await? {
return Ok(());
}
if self.read_buffered_frames().await?.is_break() {
return Ok(());
}
self.run_ready_work().await?;
let Some(tick_fired) = self.wait_event(&mut sweep_tick, &mut flush_tick).await? else {
return Ok(());
};
self.after_event(tick_fired).await?;
}
}
async fn run_ready_work(&mut self) -> Result<(), H2Error> {
self.take_queued_commands().await?;
self.take_queued_body_chunks().await?;
self.drain_pending().await?;
self.try_drain_pending().await?;
Ok(())
}
async fn after_event(&mut self, tick_fired: bool) -> Result<(), H2Error> {
self.take_queued_commands().await?;
self.try_drain_pending().await?;
if tick_fired {
self.sweep_cancelled_streams().await?;
self.check_write_stall()?;
}
Ok(())
}
pub(super) async fn on_command_received(
&mut self,
cmd: Option<DriverCommand>,
) -> Result<(), H2Error> {
match cmd {
Some(cmd) => self.on_command(cmd).await,
None => {
self.shutdown_started = true;
Ok(())
}
}
}
async fn read_buffered_frames(&mut self) -> Result<ControlFlow<()>, H2Error> {
for _ in 0..128 {
if !self.reader.buffered() || self.writer.saturated() {
break;
}
let frame = self.reader.next().await;
if self.on_read(frame).await?.is_break() {
return Ok(ControlFlow::Break(()));
}
}
Ok(ControlFlow::Continue(()))
}
pub(super) async fn on_read(
&mut self,
frame: Result<Option<Frame>, H2Error>,
) -> Result<ControlFlow<()>, H2Error> {
let Some(frame) = frame? else {
if self.shutdown_started {
return Ok(ControlFlow::Break(()));
}
return Err(H2Error::Connection {
code: ErrorCode::NoError,
reason: "peer closed connection".into(),
});
};
self.on_inbound_frame(frame).await?;
self.drain_pending().await?;
Ok(ControlFlow::Continue(()))
}
async fn take_queued_body_chunks(&mut self) -> Result<(), H2Error> {
for _ in 0..self.body_chunk_rx.len() {
if self.writer.congested() {
break;
}
let Ok(chunk) = self.body_chunk_rx.try_recv() else {
break;
};
self.on_body_chunk(chunk).await?;
}
Ok(())
}
pub(super) fn accepting_commands(&self) -> bool {
!self.shutdown_started
&& self.pending.is_empty()
&& (!self.writer.congested() || self.output.stalled())
}
async fn take_queued_commands(&mut self) -> Result<(), H2Error> {
for _ in 0..self.command_rx.len() {
if !self.pending.is_empty() || self.writer.congested() {
break;
}
let Ok(next) = self.command_rx.try_recv() else {
break;
};
self.on_command(next).await?;
}
Ok(())
}
pub(super) async fn sweep_cancelled_streams(&mut self) -> Result<(), H2Error> {
self.pending.retain(|cmd| !cmd.is_cancelled());
let to_cancel: Vec<u32> = self
.streams
.iter()
.filter(|(_, actor)| {
!actor.state.is_closed()
&& actor
.response_tx
.as_ref()
.is_some_and(ResponseSink::is_cancelled)
})
.map(|(id, _)| id)
.collect();
for sid in to_cancel {
tracing::debug!(
target: "leyline::h2",
stream_id = sid,
"caller cancelled — RST_STREAM(CANCEL) to release the slot"
);
self.writer.write_rst_stream(sid, ErrorCode::Cancel).await?;
let err = H2Error::Stream {
stream_id: sid,
code: ErrorCode::Cancel,
};
self.fail_stream(sid, err);
}
Ok(())
}
pub(super) fn reject_after_goaway(&self) -> Result<(), H2Error> {
if self.peer_goaway_last_stream.is_some() || self.output.stalled() {
return Err(H2Error::Stream {
stream_id: 0,
code: ErrorCode::RefusedStream,
});
}
Ok(())
}
pub(super) fn peer_ready(&self) -> bool {
self.peer_greeted
}
pub(super) fn check_stream_capacity(&self) -> Result<(), H2Error> {
if let Some(limit) = self.peer_settings.max_concurrent_streams
&& self.streams.len() >= limit as usize
&& self.active_stream_count() >= limit
{
return Err(H2Error::Connection {
code: ErrorCode::RefusedStream,
reason: "MAX_CONCURRENT_STREAMS exceeded".into(),
});
}
if self.next_stream_id > 0x7FFF_FFFF {
return Err(H2Error::Connection {
code: ErrorCode::NoError,
reason: "stream ID space exhausted".into(),
});
}
Ok(())
}
pub(super) fn admit_new_stream(&self) -> Result<(), H2Error> {
self.reject_after_goaway()?;
self.check_stream_capacity()
}
pub(super) fn alloc_stream_id(&mut self) -> u32 {
let stream_id = self.next_stream_id;
self.next_stream_id = stream_id + 2;
stream_id
}
pub(super) fn effective_send_window(&self, stream_id: u32) -> usize {
let conn = self.conn_send_window.max(0) as usize;
let stream = self
.streams
.get(stream_id)
.map(|i| i.send_window.max(0) as usize)
.unwrap_or(0);
conn.min(stream)
}
pub(super) fn park_stream(&mut self, stream_id: u32, pending: PendingSend) {
if let Some(actor) = self.streams.get_mut(stream_id) {
actor.pending_send = Some(pending);
if !self.buffered_pending.contains(&stream_id) {
self.buffered_pending.push_back(stream_id);
}
}
}
pub(super) async fn finish_remote(&mut self, stream_id: u32) -> Result<(), H2Error> {
if self
.streams
.get(stream_id)
.is_some_and(StreamActor::length_mismatch)
{
self.writer
.write_rst_stream(stream_id, ErrorCode::ProtocolError)
.await?;
self.fail_stream(
stream_id,
H2Error::Stream {
stream_id,
code: ErrorCode::ProtocolError,
},
);
return Ok(());
}
let local_open = self
.streams
.get(stream_id)
.is_some_and(|a| !a.state.is_closed());
if local_open {
self.writer
.write_rst_stream(stream_id, ErrorCode::NoError)
.await?;
}
if let Some(actor) = self.streams.get_mut(stream_id)
&& !actor.stalled.is_empty()
{
actor.remote_done = true;
return Ok(());
}
self.complete_stream(stream_id);
Ok(())
}
pub(super) fn has_stalled(&self) -> bool {
self.stalled > 0
}
pub(super) async fn flush_stalled(&mut self) -> Result<(), H2Error> {
let mut drained = Vec::new();
for (sid, actor) in self.streams.iter_mut() {
let Some(ResponseSink::StreamingEx { body_tx, .. }) = actor.response_tx.as_ref() else {
continue;
};
while let Some(chunk) = actor.stalled.pop_front() {
self.stalled -= 1;
match body_tx.try_send(Ok(chunk)) {
Ok(()) => {}
Err(mpsc::error::TrySendError::Full(Ok(chunk))) => {
actor.stalled.push_front(chunk);
self.stalled += 1;
break;
}
Err(_) => break,
}
}
if actor.stalled.is_empty() {
drained.push((sid, actor.remote_done));
}
}
for (sid, done) in drained {
if done {
self.complete_stream(sid);
} else {
self.maybe_top_up_stream_window(sid).await?;
}
}
Ok(())
}
pub(super) fn complete_stream(&mut self, stream_id: u32) {
if let Some(mut actor) = self.streams.remove(stream_id) {
self.stalled = self.stalled.saturating_sub(actor.stalled.len());
actor.deliver_ok();
}
self.buffered_pending.retain(|&s| s != stream_id);
}
pub(super) fn fail_stream(&mut self, stream_id: u32, err: H2Error) {
if let Some(mut actor) = self.streams.remove(stream_id) {
self.stalled = self.stalled.saturating_sub(actor.stalled.len());
actor.deliver_err(err);
}
self.buffered_pending.retain(|&s| s != stream_id);
}
pub(super) fn active_stream_count(&self) -> u32 {
self.streams
.values()
.filter(|a| {
matches!(
a.state,
StreamState::Open
| StreamState::HalfClosedLocal
| StreamState::HalfClosedRemote
)
})
.count() as u32
}
}