use std::pin::Pin;
use dambi::{Bytes, BytesMut};
use crate::protocol::server::{
LongFuture, PendingFuture, RequestStreamHandle, ResponseChunkStream, ServerProtocol,
};
pub(crate) enum SlotIngress {
Idle,
Spill {
buf: Box<BytesMut>,
filled: u32,
target: u32,
},
RequestStream(Box<RequestStreamHandle>),
}
impl SlotIngress {
#[inline(always)]
pub(crate) fn is_idle(&self) -> bool {
matches!(self, SlotIngress::Idle)
}
#[inline(always)]
pub(crate) fn is_spill(&self) -> bool {
matches!(self, SlotIngress::Spill { .. })
}
}
pub(crate) enum SlotEgress {
Idle,
Stream(Box<StreamState>),
Gather { iov_count: u8 },
}
impl SlotEgress {
#[inline(always)]
pub(crate) fn is_idle(&self) -> bool {
matches!(self, SlotEgress::Idle)
}
#[inline(always)]
pub(crate) fn is_stream(&self) -> bool {
matches!(self, SlotEgress::Stream(_))
}
#[inline(always)]
pub(crate) fn gather_iov_count(&self) -> Option<u8> {
match self {
SlotEgress::Gather { iov_count } => Some(*iov_count),
_ => None,
}
}
}
pub struct ServerSlotState<H: ServerProtocol> {
pub(crate) protocol_state: H::ConnState,
pub(crate) pending_future: Option<PendingFuture>,
pub(crate) long_future: Option<LongFuture>,
pub(crate) ingress: SlotIngress,
pub(crate) egress: SlotEgress,
}
impl<H: ServerProtocol> Default for ServerSlotState<H> {
fn default() -> Self {
Self {
protocol_state: H::ConnState::default(),
pending_future: None,
long_future: None,
ingress: SlotIngress::Idle,
egress: SlotEgress::Idle,
}
}
}
impl<H: ServerProtocol> Drop for ServerSlotState<H> {
fn drop(&mut self) {
if let SlotIngress::RequestStream(handle) =
std::mem::replace(&mut self.ingress, SlotIngress::Idle)
{
handle.close();
}
}
}
pub(crate) struct StreamState {
pub(crate) stream: Pin<Box<dyn ResponseChunkStream>>,
pub(crate) cur_chunk: Option<Bytes>,
pub(crate) cur_chunk_offset: u32,
pub(crate) eof: bool,
}
impl StreamState {
#[cold]
pub(crate) fn new(stream: Pin<Box<dyn ResponseChunkStream>>) -> Self {
Self {
stream,
cur_chunk: None,
cur_chunk_offset: 0,
eof: false,
}
}
pub(crate) fn advance(&mut self, n: u32) {
let Some(chunk) = self.cur_chunk.as_ref() else {
return;
};
let chunk_len = chunk.len() as u32;
let unsent = chunk_len.saturating_sub(self.cur_chunk_offset);
let consume = n.min(unsent);
self.cur_chunk_offset += consume;
if self.cur_chunk_offset >= chunk_len {
self.cur_chunk = None;
self.cur_chunk_offset = 0;
}
}
}