sark-h2 0.13.0

use std::marker::PhantomData;
use std::pin::Pin;

use dope::DriverContext;
use dope::driver::token::Token;
use dope::manifold::Outcome;
use dope::manifold::listener::{
    application::{Application, ApplicationHooks},
    egress::SlotEgress,
    state::{EgressCtx, State},
};
use dope_fiber::abi::Fiber;
use dope_fiber::raw::task::RootWaker;
use dope_net::link::slot::Slot;
use dope_net::wire::Wire;
use dope_net::wire::identity::Identity;
use o3::buffer::RetainBytes;

use super::Config;
use super::connection::{ConnectionState, Dispatch, EventSink, Limits, Request, Response};
use super::driver;
use super::scheduler::{Resumed, Scheduler, StartContext, Started};
use crate::conn::{self, ConfigError, Conn, ConnError, ValidatedConfig};
use crate::frame::{ErrorCode, HEADER_LEN};
use crate::role::ServerRole;
use crate::stream::StreamId;

const SERVER_EVENT_CAPACITY: usize = 64;

fn protocol_config(config: Config) -> conn::Config {
    let mut protocol = conn::Config::default();
    // `chunk` retains the complete transport chunk before parsing it, so this
    // bound must fit both one receive buffer and one maximum-sized peer frame.
    protocol.inbound_capacity = config
        .receive_buffer_bytes
        .max(HEADER_LEN + protocol.local_settings.max_frame_size as usize);
    let frame_headroom = HEADER_LEN + protocol.local_settings.max_frame_size as usize;
    protocol.outbound_capacity = config
        .max_outbound_bytes
        .saturating_add(frame_headroom)
        .max(protocol.initial_outbound_capacity::<ServerRole>());
    // The driver drains and resumes parsing on overload. A bounded event burst
    // therefore avoids provisioning the generic 8K-event queue per connection.
    protocol.event_capacity = SERVER_EVENT_CAPACITY;
    protocol
}

pub trait Handler: 'static {
    type Fut<'h>: Fiber<'h, Output = Response> + 'h
    where
        Self: 'h;

    fn request<'h>(&'h self, request: Request) -> Self::Fut<'h>;
}

pub trait SyncHandler: 'static {
    fn request(&self, request: Request) -> Response;
}

impl<F> SyncHandler for F
where
    F: Fn(Request) -> Response + 'static,
{
    fn request(&self, request: Request) -> Response {
        self(request)
    }
}

trait ConnectionContainer: Default + 'static {
    fn connection(&mut self) -> &mut ConnectionState;
}

fn flush_into<'d, W, C>(
    slot: &mut Slot<'d, W, State<C>>,
    egress: &mut EgressCtx<'_, 'd, '_>,
    driver: &mut DriverContext<'_, 'd>,
    close_after: bool,
) where
    W: Wire,
    C: ConnectionContainer,
{
    let send_token = slot.token();
    let mut write_buffer = egress.write_buf_for(slot);
    if close_after {
        slot.set_close_after();
    }
    let state = slot.state.conn.connection();
    match state.connection.take_write(&mut write_buffer) {
        crate::conn::OutboundWrite::Buffered(written) => {
            slot.submit_buffered(write_buffer, written, send_token, driver);
        }
        crate::conn::OutboundWrite::Split { prefix_len, body } => {
            slot.submit_split_shared(write_buffer, prefix_len, body, send_token, driver);
        }
    }
}

fn finish_ingest<'d, W, C>(
    slot: &mut Slot<'d, W, State<C>>,
    egress: &mut EgressCtx<'_, 'd, '_>,
    driver: &mut DriverContext<'_, 'd>,
    error: Option<ConnError>,
) -> Outcome
where
    W: Wire,
    C: ConnectionContainer,
{
    if let Some(error) = error {
        let state = slot.state.conn.connection();
        let _ = state.connection.goaway(ErrorCode::from(&error), b"");
        flush_into(slot, egress, driver, true);
    } else {
        flush_connection(slot, egress, driver);
    }
    Outcome::Ok
}

fn resume_egress<'d, W, C>(
    slot: &mut Slot<'d, W, State<C>>,
    limits: Limits,
    egress: &mut EgressCtx<'_, 'd, '_>,
    driver: &mut DriverContext<'_, 'd>,
) where
    W: Wire,
    C: ConnectionContainer,
{
    let state = slot.state.conn.connection();
    state.pump_pending(limits, false);
    if state.connection.outbound_len() != 0 {
        flush_connection(slot, egress, driver);
    }
}

#[derive(Default)]
pub struct SyncConnState {
    state: ConnectionState,
}

impl SyncConnState {
    fn new(protocol: ValidatedConfig<ServerRole>) -> Self {
        Self {
            state: ConnectionState::new(Conn::from_config(protocol)),
        }
    }
}

impl ConnectionContainer for SyncConnState {
    fn connection(&mut self) -> &mut ConnectionState {
        &mut self.state
    }
}

pub struct SyncApp<'h, H: SyncHandler, W: Wire = Identity> {
    user: &'h H,
    limits: Limits,
    protocol: ValidatedConfig<ServerRole>,
    wire: PhantomData<fn() -> W>,
}

impl<'h, H: SyncHandler, W: Wire> SyncApp<'h, H, W> {
    pub fn new(user: &'h H, config: Config) -> Result<Self, ConfigError> {
        Ok(Self {
            user,
            limits: config.into(),
            protocol: ValidatedConfig::new(protocol_config(config))?,
            wire: PhantomData,
        })
    }

    pub fn handler(&self) -> &H {
        self.user
    }
}

struct SyncSink<'h, H: SyncHandler> {
    user: &'h H,
}

impl<H: SyncHandler> EventSink for SyncSink<'_, H> {
    fn request(&mut self, _stream_id: StreamId, request: Request) -> Dispatch {
        Dispatch::Response(self.user.request(request))
    }

    fn reset(&mut self, _stream_id: StreamId) {}
}

struct SyncTransport<'a, 'h, H: SyncHandler> {
    user: &'h H,
    limits: Limits,
    state: &'a mut ConnectionState,
}

impl<H: SyncHandler> driver::Transport for SyncTransport<'_, '_, H> {
    fn connection(&mut self) -> &mut Conn<ServerRole> {
        &mut self.state.connection
    }

    fn drain_events(&mut self) -> usize {
        self.state
            .drain_events(self.limits, &mut SyncSink { user: self.user })
    }
}

impl<'d, H: SyncHandler + 'd, W: Wire> Application<'d> for SyncApp<'d, H, W> {
    type Conn = SyncConnState;
    type Wire = W;
    type Hooks = Self;

    fn connection(self: Pin<&Self>) -> Self::Conn {
        SyncConnState::new(self.protocol)
    }
}

impl<'d, H: SyncHandler + 'd, W: Wire> ApplicationHooks<'d, SyncApp<'d, H, W>>
    for SyncApp<'d, H, W>
{
    fn chunk<R: RetainBytes>(
        app: Pin<&mut SyncApp<'d, H, W>>,
        slot: &mut Slot<'d, W, State<SyncConnState>>,
        mut egress: EgressCtx<'_, 'd, '_>,
        chunk: R,
        driver: &mut DriverContext<'_, 'd>,
    ) -> Outcome {
        let this = app.get_mut();
        let state = &mut slot.state.conn.state;
        let error = driver::Driver::new(&mut SyncTransport {
            user: this.user,
            limits: this.limits,
            state,
        })
        .ingest_retained(chunk.into_retained())
        .err();
        finish_ingest(slot, &mut egress, driver, error)
    }

    fn send(
        app: Pin<&mut SyncApp<'d, H, W>>,
        slot: &mut Slot<'d, W, State<SyncConnState>>,
        mut egress: EgressCtx<'_, 'd, '_>,
        _sent: usize,
        driver: &mut DriverContext<'_, 'd>,
    ) {
        resume_egress(slot, app.get_mut().limits, &mut egress, driver);
    }

    fn close(
        _app: Pin<&mut SyncApp<'d, H, W>>,
        slot: &mut Slot<'d, W, State<SyncConnState>>,
        _egress: EgressCtx<'_, 'd, '_>,
    ) {
        slot.state.conn.state.close();
    }
}

#[derive(Default)]
pub struct ConnState {
    state: ConnectionState,
    task_head: Option<u32>,
}

impl ConnState {
    fn new(protocol: ValidatedConfig<ServerRole>) -> Self {
        Self {
            state: ConnectionState::new(Conn::from_config(protocol)),
            task_head: None,
        }
    }
}

impl ConnectionContainer for ConnState {
    fn connection(&mut self) -> &mut ConnectionState {
        &mut self.state
    }
}

pub struct App<'d, H: Handler + 'd, W: Wire = Identity> {
    user: &'d H,
    limits: Limits,
    protocol: ValidatedConfig<ServerRole>,
    scheduler: Scheduler<'d, H::Fut<'d>>,
    wire: PhantomData<fn() -> W>,
}

impl<'d, H: Handler + 'd, W: Wire> App<'d, H, W> {
    pub fn new(user: &'d H, config: Config) -> Result<Self, ConfigError> {
        if config.max_handler_tasks == 0 {
            return Err(ConfigError::ZeroCapacity("handler task"));
        }
        let scheduler = Scheduler::try_with_capacity(config.max_handler_tasks)
            .ok_or(ConfigError::HandlerCapacityOverflow)?;
        Ok(Self {
            user,
            limits: config.into(),
            protocol: ValidatedConfig::new(protocol_config(config))?,
            scheduler,
            wire: PhantomData,
        })
    }

    pub fn handler(&self) -> &H {
        self.user
    }

    fn drain_events(
        &mut self,
        slot: &mut Slot<'d, W, State<ConnState>>,
        driver: &mut DriverContext<'_, 'd>,
    ) -> usize {
        let connection_id = slot.token();
        let parent = RootWaker::from_ready(slot.driver(), slot.ready_key());
        let ConnState { state, task_head } = &mut slot.state.conn;
        let mut sink = AsyncSink {
            user: self.user,
            scheduler: &mut self.scheduler,
            connection_id,
            task_head,
            parent,
            driver,
        };
        state.drain_events(self.limits, &mut sink)
    }
}

struct AsyncSink<'a, 'turn, 'd, H>
where
    H: Handler + 'd,
{
    user: &'d H,
    scheduler: &'a mut Scheduler<'d, H::Fut<'d>>,
    connection_id: Token,
    task_head: &'a mut Option<u32>,
    parent: RootWaker<'d>,
    driver: &'a mut DriverContext<'turn, 'd>,
}

impl<'d, H> EventSink for AsyncSink<'_, '_, 'd, H>
where
    H: Handler + 'd,
{
    fn request(&mut self, stream_id: StreamId, request: Request) -> Dispatch {
        match self.scheduler.start(
            self.user.request(request),
            StartContext {
                connection_id: self.connection_id,
                stream_id,
                task_head: self.task_head,
                parent: self.parent,
                driver: self.driver,
            },
        ) {
            Started::Ready(response) => Dispatch::Response(response),
            Started::Pending => Dispatch::Pending,
            Started::Refused => Dispatch::Reset(ErrorCode::RefusedStream),
            Started::Failed => Dispatch::Reset(ErrorCode::InternalError),
        }
    }

    fn reset(&mut self, stream_id: StreamId) {
        self.scheduler
            .cancel(self.task_head, self.connection_id, stream_id);
    }
}

struct AsyncTransport<'a, 'turn, 'd, H: Handler + 'd, W: Wire> {
    app: &'a mut App<'d, H, W>,
    slot: &'a mut Slot<'d, W, State<ConnState>>,
    driver: &'a mut DriverContext<'turn, 'd>,
}

impl<'d, H: Handler + 'd, W: Wire> driver::Transport for AsyncTransport<'_, '_, 'd, H, W> {
    fn connection(&mut self) -> &mut Conn<ServerRole> {
        &mut self.slot.state.conn.state.connection
    }

    fn drain_events(&mut self) -> usize {
        self.app.drain_events(self.slot, self.driver)
    }
}

impl<'d, H: Handler + 'd, W: Wire> Application<'d> for App<'d, H, W> {
    type Conn = ConnState;
    type Wire = W;
    type Hooks = Self;

    fn connection(self: Pin<&Self>) -> Self::Conn {
        ConnState::new(self.protocol)
    }
}

impl<'d, H: Handler + 'd, W: Wire> ApplicationHooks<'d, App<'d, H, W>> for App<'d, H, W> {
    fn chunk<R: RetainBytes>(
        app: Pin<&mut App<'d, H, W>>,
        slot: &mut Slot<'d, W, State<ConnState>>,
        mut egress: EgressCtx<'_, 'd, '_>,
        chunk: R,
        driver: &mut DriverContext<'_, 'd>,
    ) -> Outcome {
        let this = app.get_mut();
        let error = driver::Driver::new(&mut AsyncTransport {
            app: this,
            slot,
            driver,
        })
        .ingest_retained(chunk.into_retained())
        .err();
        finish_ingest(slot, &mut egress, driver, error)
    }

    fn send(
        app: Pin<&mut App<'d, H, W>>,
        slot: &mut Slot<'d, W, State<ConnState>>,
        mut egress: EgressCtx<'_, 'd, '_>,
        _sent: usize,
        driver: &mut DriverContext<'_, 'd>,
    ) {
        resume_egress(slot, app.get_mut().limits, &mut egress, driver);
    }

    fn activate(
        app: Pin<&mut App<'d, H, W>>,
        slot: &mut Slot<'d, W, State<ConnState>>,
        mut egress: EgressCtx<'_, 'd, '_>,
        driver: &mut DriverContext<'_, 'd>,
    ) {
        let this = app.get_mut();
        let parent = RootWaker::from_ready(slot.driver(), slot.ready_key());
        let connection_id = slot.token();
        let ConnState { state, task_head } = &mut slot.state.conn;
        this.scheduler.route_ready(parent);
        while let Some(key) = this.scheduler.pop_ready(connection_id) {
            match this.scheduler.resume(key, connection_id, task_head, driver) {
                Resumed::Ready(Some(stream_id), response) => {
                    state.begin_response(stream_id, response, this.limits);
                }
                Resumed::Ready(None, _) => {
                    debug_assert!(false, "registered task must be releasable");
                }
                Resumed::Failed(Some(stream_id)) => {
                    let _ = state
                        .connection
                        .reset_stream(stream_id, ErrorCode::InternalError);
                }
                Resumed::Pending | Resumed::Failed(None) | Resumed::Stale => {}
            }
        }
        if state.connection.outbound_len() != 0 {
            flush_connection(slot, &mut egress, driver);
        }
    }

    fn close(
        app: Pin<&mut App<'d, H, W>>,
        slot: &mut Slot<'d, W, State<ConnState>>,
        _egress: EgressCtx<'_, 'd, '_>,
    ) {
        let this = app.get_mut();
        let state = &mut slot.state.conn;
        this.scheduler.close(&mut state.task_head);
        state.state.close();
    }
}

fn flush_connection<'d, W, C>(
    slot: &mut Slot<'d, W, State<C>>,
    egress: &mut EgressCtx<'_, 'd, '_>,
    driver: &mut DriverContext<'_, 'd>,
) where
    W: Wire,
    C: ConnectionContainer,
{
    let close_after = slot.state.conn.connection().connection.goaway_sent();
    flush_into(slot, egress, driver, close_after);
}