Skip to main content

dope_session/pool/server/
mod.rs

1pub(crate) mod poll;
2pub(crate) mod recv;
3pub(crate) mod state;
4pub(crate) mod wake;
5pub(crate) mod write;
6pub(crate) mod write_stream;
7
8pub(crate) use self::state::ServerSlotState;
9
10use self::wake::WakeQueue;
11use super::Direction;
12use super::Pool;
13use super::core::PoolCore;
14use crate::protocol::server::{LongFuture, ServerProtocol};
15use dope_core::FdHandle;
16use dope_core::driver::{
17    Backend, BackendToken, DriverConfig as _, DriverLifecycle, DriverOps, PoolDriver,
18};
19use dope_core::profile::Profile;
20use dope_runtime::runtime::Runtime;
21use dope_transport::{Tcp, Transport};
22use std::collections::VecDeque;
23use std::io;
24use std::marker::PhantomData;
25use std::net::SocketAddr;
26use std::sync::Arc;
27use std::sync::atomic::AtomicBool;
28
29const ACCEPT_PENDING_INIT_CAP: usize = 1024;
30
31#[derive(Default)]
32pub(crate) struct AcceptSlot {
33    pub(super) epoch_seq: u32,
34    pub(super) epoch: u32,
35    pub(super) armed: bool,
36    pub(super) canceling: bool,
37    pub(super) pending_rearm: bool,
38}
39
40pub struct Server<H: ServerProtocol, T: Transport = Tcp> {
41    pub(crate) protocol: H,
42    pub(crate) listener_fd: FdHandle,
43    pub(crate) accept_slot: AcceptSlot,
44    pub(crate) accept_pending: VecDeque<FdHandle>,
45    pub(crate) socket_opts: dope_transport::config::StreamSocketConfig,
46    pub(crate) accept_slot_id: usize,
47    transport: PhantomData<T>,
48}
49
50impl<H: ServerProtocol, T: Transport> Server<H, T> {
51    #[inline(always)]
52    pub fn protocol(&self) -> &H {
53        &self.protocol
54    }
55
56    #[inline(always)]
57    pub fn protocol_mut(&mut self) -> &mut H {
58        &mut self.protocol
59    }
60
61    #[inline(always)]
62    pub fn listener_fd(&self) -> &FdHandle {
63        &self.listener_fd
64    }
65
66    #[cold]
67    #[inline(never)]
68    pub(crate) fn accept_arm<B: PoolDriver>(&mut self, driver: &mut B::Driver) {
69        if self.accept_slot.armed || self.accept_slot.canceling {
70            self.accept_slot.pending_rearm = false;
71            return;
72        }
73        let token = <B::Token as BackendToken>::next_for_slot(
74            self.accept_slot_id,
75            &mut self.accept_slot.epoch_seq,
76        );
77        let (_, epoch) = token.parts();
78        let fd = self.listener_fd.as_fd();
79        let armed = driver.arm_with_retry(|d| d.arm_accept_multi(token, fd));
80        self.accept_slot.epoch = epoch;
81        if armed {
82            self.accept_slot.armed = true;
83            self.accept_slot.pending_rearm = false;
84        } else {
85            self.accept_slot.pending_rearm = true;
86        }
87    }
88
89    #[inline(always)]
90    pub(crate) fn accept_flush_rearm<B: PoolDriver>(&mut self, driver: &mut B::Driver) {
91        if self.accept_slot.pending_rearm {
92            self.accept_arm::<B>(driver);
93        }
94    }
95
96    pub(crate) fn on_accept_event<B: PoolDriver>(
97        &mut self,
98        token: B::Token,
99        result: i32,
100        more: bool,
101    ) {
102        let (_, epoch) = token.parts();
103        if self.accept_slot.epoch != epoch {
104            return;
105        }
106
107        if result >= 0 {
108            let fd = dope_core::take_fd(result);
109            self.accept_pending.push_back(fd);
110            if !more {
111                self.accept_slot.armed = false;
112                self.accept_slot.pending_rearm = !self.accept_slot.canceling;
113            }
114            return;
115        }
116
117        let errno = -result;
118        if errno == libc::EINVAL {
119            panic!("dope invariant violated: accept-multi unsupported on required uring platform");
120        }
121        if Tcp::is_accept_retryable(errno) {
122            if !more {
123                self.accept_slot.armed = false;
124                self.accept_slot.pending_rearm = !self.accept_slot.canceling;
125            }
126            return;
127        }
128        if errno == libc::ECANCELED {
129            self.accept_slot.armed = false;
130            self.accept_slot.pending_rearm = !self.accept_slot.canceling;
131            self.accept_slot.canceling = false;
132            return;
133        }
134        self.accept_slot.armed = false;
135        self.accept_slot.pending_rearm = !self.accept_slot.canceling;
136    }
137
138    #[inline(always)]
139    pub(crate) fn has_pending(&self) -> bool {
140        WakeQueue::is_pending() || !self.accept_pending.is_empty() || self.accept_slot.pending_rearm
141    }
142}
143
144impl<H: ServerProtocol, T: Transport, B: PoolDriver> Direction<B> for Server<H, T> {
145    type Transport = T;
146    type SlotUser = ServerSlotState<H>;
147    type CodecLayer = H::CodecLayer;
148
149    #[inline(always)]
150    fn has_pending(&self) -> bool {
151        Self::has_pending(self)
152    }
153
154    #[inline(always)]
155    fn tick(
156        &mut self,
157        core: &mut PoolCore<Self::Transport, Self::CodecLayer, Self::SlotUser, B>,
158        driver: &mut B::Driver,
159    ) {
160        while let Some(fd) = self.accept_pending.pop_front() {
161            let _ = T::apply_stream_socket_opts(&fd, &self.socket_opts);
162            let Some((slot_idx, conn_id)) = core.allocate_slot(driver, fd) else {
163                continue;
164            };
165            let user = core.slots[slot_idx].user.get_mut();
166            if let Some(future) = self.protocol.on_register(conn_id, &mut user.protocol_state) {
167                user.long_future = Some(LongFuture {
168                    generation: conn_id.generation(),
169                    future,
170                });
171                WakeQueue::push(conn_id);
172            }
173        }
174        self.accept_flush_rearm::<B>(driver);
175        core.flush_recv_rearm(driver);
176    }
177
178    #[inline(always)]
179    fn on_accept_event(&mut self, token: B::Token, result: i32, more: bool) {
180        Self::on_accept_event::<B>(self, token, result, more)
181    }
182}
183
184#[derive(Clone, Debug)]
185pub struct TcpServerConfig {
186    pub max_conn: usize,
187    pub bind: SocketAddr,
188    pub backlog: i32,
189    pub socket_opts: dope_transport::config::StreamSocketConfig,
190    pub shutdown: Option<Arc<AtomicBool>>,
191    pub cpu_id: Option<u16>,
192}
193
194impl<H, F, B> Pool<Server<H, Tcp>, F, B>
195where
196    H: ServerProtocol,
197    F: Profile,
198    B: PoolDriver,
199{
200    pub fn build(protocol: H, cfg: TcpServerConfig) -> io::Result<Runtime<Self>> {
201        let shutdown = cfg.shutdown.clone();
202        let max_conn = cfg.max_conn;
203        let cpu_id = cfg.cpu_id;
204        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);
205
206        let driver_cfg =
207            <<B as Backend>::Config as dope_core::driver::DriverConfig>::for_tcp_profile(
208                F::RING_ENTRIES,
209                max_conn,
210                crate::pool::slot::INGRESS_BUF_CAP,
211            )
212            .with_defer_taskrun(F::DEFER_TASKRUN)
213            .with_cpu_id(cpu_id);
214        let mut driver = <B as Backend>::new_driver(driver_cfg)?;
215        let reactor = Self::build_raw(protocol, cfg, &mut driver, 0)?;
216        Ok(Runtime::new(reactor, driver).with_shutdown(shutdown))
217    }
218
219    pub fn build_raw(
220        protocol: H,
221        cfg: TcpServerConfig,
222        driver: &mut B::Driver,
223        slot_id_base: usize,
224    ) -> io::Result<Self> {
225        let TcpServerConfig {
226            max_conn,
227            bind,
228            backlog,
229            socket_opts,
230            shutdown: _,
231            cpu_id: _,
232        } = cfg;
233        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);
234
235        WakeQueue::reserve(max_conn);
236
237        let listener_cfg = dope_transport::config::ListenerSocketConfig {
238            reuseport: dope_transport::config::SocketToggle::Enabled,
239            fastopen_backlog: None,
240        };
241
242        B::set_current(driver as *mut _);
243
244        let listener_result = Tcp::bind_listener(bind, backlog, &listener_cfg);
245
246        B::clear_current();
247
248        let listener_fd = listener_result?;
249        let core =
250            PoolCore::<Tcp, H::CodecLayer, ServerSlotState<H>, B>::new(max_conn, slot_id_base)?;
251        let accept_slot_id = slot_id_base + max_conn;
252
253        let direction = Server {
254            protocol,
255            listener_fd,
256            accept_slot: AcceptSlot::default(),
257            accept_pending: VecDeque::with_capacity(ACCEPT_PENDING_INIT_CAP),
258            socket_opts,
259            accept_slot_id,
260            transport: PhantomData,
261        };
262
263        let mut reactor = Self {
264            direction,
265            core,
266            profile: PhantomData,
267        };
268        reactor.direction.accept_arm::<B>(driver);
269        if !reactor.direction.accept_slot.armed && !reactor.direction.accept_slot.pending_rearm {
270            return Err(io::Error::other("dope-net: failed to arm multishot accept"));
271        }
272        Ok(reactor)
273    }
274}
275
276impl<H, T, F, B> Pool<Server<H, T>, F, B>
277where
278    H: ServerProtocol,
279    T: Transport,
280    F: Profile,
281    B: PoolDriver,
282{
283    pub fn slot_id_range(&self) -> std::ops::Range<usize> {
284        let base = self.core.slot_id_base;
285        base..(self.direction.accept_slot_id + 1)
286    }
287
288    pub fn conn_slot_count(&self) -> usize {
289        self.core.slots.len()
290    }
291
292    pub fn protocol(&self) -> &H {
293        &self.direction.protocol
294    }
295
296    pub fn protocol_mut(&mut self) -> &mut H {
297        &mut self.direction.protocol
298    }
299
300    pub fn listener_fd(&self) -> &FdHandle {
301        &self.direction.listener_fd
302    }
303}
304
305impl<H, T, F, B> super::PoolTick<B> for Pool<Server<H, T>, F, B>
306where
307    H: ServerProtocol,
308    T: Transport,
309    F: Profile,
310    B: PoolDriver,
311{
312    #[inline(always)]
313    fn tick(&mut self, driver: &mut B::Driver) {
314        self.direction.tick(&mut self.core, driver);
315        self.drain_ready_queue(driver);
316    }
317}