dope-session 0.2.3

Thin io_uring adaptor with "Manifolds"
Documentation
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;
        }
    }
}