Skip to main content

dope_session/protocol/
server.rs

1use crate::{CodecLayer, SlotGen, SlotId};
2use std::pin::Pin;
3
4use crate::Framer;
5
6#[derive(Clone, Copy, PartialEq, Eq, Debug)]
7pub enum RouteRequestBodyKind {
8    Inline,
9    Spilled,
10    Stream,
11}
12
13pub trait ServerProtocol: 'static {
14    type Framer: Framer;
15
16    type ConnState: Default + 'static;
17
18    type CodecLayer: CodecLayer;
19
20    const SUPPORTS_PARK: bool = false;
21
22    fn framer(&mut self) -> &mut Self::Framer;
23
24    fn handle<'buf>(
25        &mut self,
26        req: &'buf [u8],
27        head: <Self::Framer as Framer>::Head<'buf>,
28        write: &mut [u8],
29    ) -> HandlerAction;
30
31    fn handle_park<'buf>(
32        &mut self,
33        conn_id: SlotId,
34        conn_state: &mut Self::ConnState,
35        req: &'buf [u8],
36        head: <Self::Framer as Framer>::Head<'buf>,
37        write: &mut [u8],
38    ) -> ParkAction {
39        let _ = (conn_id, conn_state, req, head, write);
40        unreachable!("ServerProtocol::handle_park called with SUPPORTS_PARK = true but no override")
41    }
42
43    fn oversize_response(&self) -> &'static [u8] {
44        b""
45    }
46
47    fn route_request_body_kind(
48        &self,
49        head: &<Self::Framer as Framer>::Head<'_>,
50    ) -> RouteRequestBodyKind {
51        let _ = head;
52        RouteRequestBodyKind::Inline
53    }
54
55    fn handle_request_stream<'buf>(
56        &mut self,
57        head_bytes: &'buf [u8],
58        head: <Self::Framer as Framer>::Head<'buf>,
59        body_stream: RequestBodyStream,
60        write: &mut [u8],
61    ) -> RequestStreamAction {
62        let _ = (head_bytes, head, body_stream, write);
63        unreachable!(
64            "ServerProtocol::handle_request_stream called without an override; \
65             a route returned RouteRequestBodyKind::Stream but the \
66             handler does not implement the stream dispatch path"
67        )
68    }
69
70    fn on_register(
71        &mut self,
72        conn_id: SlotId,
73        conn_state: &mut Self::ConnState,
74    ) -> Option<Pin<Box<dyn std::future::Future<Output = ()> + 'static>>> {
75        let _ = (conn_id, conn_state);
76        None
77    }
78
79    fn drain_outbound(&mut self, conn_state: &mut Self::ConnState, write: &mut [u8]) -> u32 {
80        let _ = (conn_state, write);
81        0
82    }
83
84    fn close_pending(&self, conn_state: &Self::ConnState) -> bool {
85        let _ = conn_state;
86        false
87    }
88}
89
90pub enum ParkAction {
91    Send { written: u16, close_after: bool },
92    Park,
93    Close(&'static [u8]),
94}
95
96pub enum RequestStreamAction {
97    Pending(dambi::InlineFuture<'static, PendingResponse, 64>),
98    Close(&'static [u8]),
99}
100
101pub enum HandlerAction {
102    Send {
103        written: u16,
104        close_after: bool,
105    },
106    SendStatic {
107        hdr_written: u32,
108        body: &'static [u8],
109        close_after: bool,
110    },
111
112    SendVectored {
113        iovs: [dope_core::IoVec; 4],
114        iov_count: u8,
115        cookie: *const u8,
116        total_bytes: u32,
117        close_after: bool,
118    },
119    SendStream {
120        hdr_written: u32,
121        stream: Pin<Box<dyn ResponseChunkStream>>,
122        close_after: bool,
123    },
124    Pending(dambi::InlineFuture<'static, PendingResponse, 64>),
125    Close(&'static [u8]),
126}
127
128pub trait ResponseChunkStream: 'static {
129    fn poll_chunk(
130        self: Pin<&mut Self>,
131        cx: &mut std::task::Context<'_>,
132    ) -> std::task::Poll<Option<dambi::Bytes>>;
133}
134
135pub struct RequestBodyStream {
136    inner: std::rc::Rc<std::cell::RefCell<RequestStreamInner>>,
137}
138
139#[derive(Default)]
140pub(crate) struct RequestStreamInner {
141    pub(crate) pending: std::collections::VecDeque<dambi::Bytes>,
142    pub(crate) waker: Option<std::task::Waker>,
143    pub(crate) received: u32,
144    pub(crate) target: u32,
145    pub(crate) eof: bool,
146}
147
148impl RequestBodyStream {
149    pub(crate) fn new(target: u32) -> (Self, RequestStreamHandle) {
150        let inner = std::rc::Rc::new(std::cell::RefCell::new(RequestStreamInner {
151            pending: std::collections::VecDeque::new(),
152            waker: None,
153            received: 0,
154            target,
155            eof: target == 0,
156        }));
157        (
158            Self {
159                inner: inner.clone(),
160            },
161            RequestStreamHandle { inner },
162        )
163    }
164
165    pub fn poll_chunk(
166        &mut self,
167        cx: &mut std::task::Context<'_>,
168    ) -> std::task::Poll<Option<dambi::Bytes>> {
169        let mut state = self.inner.borrow_mut();
170        if let Some(chunk) = state.pending.pop_front() {
171            return std::task::Poll::Ready(Some(chunk));
172        }
173        if state.eof {
174            return std::task::Poll::Ready(None);
175        }
176        state.waker = Some(cx.waker().clone());
177        std::task::Poll::Pending
178    }
179}
180
181pub(crate) struct RequestStreamHandle {
182    pub(crate) inner: std::rc::Rc<std::cell::RefCell<RequestStreamInner>>,
183}
184
185const REQUEST_STREAM_QUEUE_CAP: usize = 8 * 1024 * 1024;
186
187impl RequestStreamHandle {
188    pub(crate) fn push_slice(&self, src: &[u8]) -> usize {
189        let mut state = self.inner.borrow_mut();
190        if src.is_empty() {
191            return 0;
192        }
193        let queued: usize = state.pending.iter().map(|b| b.len()).sum();
194        if queued >= REQUEST_STREAM_QUEUE_CAP {
195            return 0;
196        }
197        let needed = (state.target.saturating_sub(state.received)) as usize;
198        let take = src.len().min(needed);
199        if take > 0 {
200            state
201                .pending
202                .push_back(dambi::Bytes::copy_from_slice(&src[..take]));
203            state.received = state.received.saturating_add(take as u32);
204        }
205        if state.received >= state.target {
206            state.eof = true;
207        }
208        if let Some(waker) = state.waker.take() {
209            drop(state);
210            waker.wake();
211        }
212        take
213    }
214
215    pub(crate) fn close(&self) {
216        let mut state = self.inner.borrow_mut();
217        state.eof = true;
218        if let Some(waker) = state.waker.take() {
219            drop(state);
220            waker.wake();
221        }
222    }
223
224    pub(crate) fn is_complete(&self) -> bool {
225        let state = self.inner.borrow();
226        state.eof
227    }
228}
229
230type ResponseSerializer = Box<dyn FnOnce(&mut [u8]) -> PendingOutcome + 'static>;
231
232pub struct PendingResponse {
233    pub serializer: ResponseSerializer,
234}
235
236pub enum PendingOutcome {
237    Inline {
238        written: u16,
239        close_after: bool,
240    },
241    Static {
242        hdr_written: u32,
243        body: &'static [u8],
244        close_after: bool,
245    },
246    Stream {
247        hdr_written: u32,
248        stream: Pin<Box<dyn ResponseChunkStream>>,
249        close_after: bool,
250    },
251}
252
253pub(crate) struct PendingFuture {
254    pub(crate) generation: SlotGen,
255    pub(crate) future: dambi::InlineFuture<'static, PendingResponse, 64>,
256}
257
258pub(crate) struct LongFuture {
259    pub(crate) generation: SlotGen,
260    pub(crate) future: Pin<Box<dyn std::future::Future<Output = ()> + 'static>>,
261}