Skip to main content

dope_session/protocol/
server.rs

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