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    type WriteBuf = crate::WriteBufArr;
149
150    #[inline(always)]
151    fn has_pending(&self) -> bool {
152        Self::has_pending(self)
153    }
154
155    #[inline(always)]
156    fn tick(
157        &mut self,
158        core: &mut PoolCore<Self::Transport, Self::CodecLayer, Self::WriteBuf, Self::SlotUser, B>,
159        driver: &mut B::Driver,
160    ) {
161        while let Some(fd) = self.accept_pending.pop_front() {
162            let _ = T::apply_stream_socket_opts(&fd, &self.socket_opts);
163            let Some((slot_idx, conn_id)) = core.allocate_slot(driver, fd) else {
164                continue;
165            };
166            let user = core.slots[slot_idx].user.get_mut();
167            if let Some(future) = self.protocol.on_register(conn_id, &mut user.protocol_state) {
168                user.long_future = Some(LongFuture {
169                    generation: conn_id.generation(),
170                    future,
171                });
172                WakeQueue::push(conn_id);
173            }
174        }
175        self.accept_flush_rearm::<B>(driver);
176        core.flush_recv_rearm(driver);
177    }
178
179    #[inline(always)]
180    fn on_accept_event(&mut self, token: B::Token, result: i32, more: bool) {
181        Self::on_accept_event::<B>(self, token, result, more)
182    }
183}
184
185#[derive(Clone, Debug)]
186pub struct TcpServerConfig {
187    pub max_conn: usize,
188    pub bind: SocketAddr,
189    pub backlog: i32,
190    pub socket_opts: dope_transport::config::StreamSocketConfig,
191    pub shutdown: Option<Arc<AtomicBool>>,
192    pub cpu_id: Option<u16>,
193}
194
195impl<H, F, B> Pool<Server<H, Tcp>, F, B>
196where
197    H: ServerProtocol,
198    F: Profile,
199    B: PoolDriver,
200{
201    pub fn build(protocol: H, cfg: TcpServerConfig) -> io::Result<Runtime<Self>> {
202        let shutdown = cfg.shutdown.clone();
203        let max_conn = cfg.max_conn;
204        let cpu_id = cfg.cpu_id;
205        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);
206
207        let driver_cfg =
208            <<B as Backend>::Config as dope_core::driver::DriverConfig>::for_tcp_profile(
209                F::RING_ENTRIES,
210                max_conn,
211                crate::pool::slot::INGRESS_BUF_CAP,
212            )
213            .with_defer_taskrun(F::DEFER_TASKRUN)
214            .with_cpu_id(cpu_id);
215        let mut driver = <B as Backend>::new_driver(driver_cfg)?;
216        let reactor = Self::build_raw(protocol, cfg, &mut driver, 0)?;
217        Ok(Runtime::new(reactor, driver).with_shutdown(shutdown))
218    }
219
220    pub fn build_raw(
221        protocol: H,
222        cfg: TcpServerConfig,
223        driver: &mut B::Driver,
224        slot_id_base: usize,
225    ) -> io::Result<Self> {
226        let TcpServerConfig {
227            max_conn,
228            bind,
229            backlog,
230            socket_opts,
231            shutdown: _,
232            cpu_id: _,
233        } = cfg;
234        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);
235
236        WakeQueue::reserve(max_conn);
237
238        let listener_cfg = dope_transport::config::ListenerSocketConfig {
239            reuseport: dope_transport::config::SocketToggle::Enabled,
240            fastopen_backlog: None,
241        };
242
243        B::set_current(driver as *mut _);
244
245        let listener_result = Tcp::bind_listener(bind, backlog, &listener_cfg);
246
247        B::clear_current();
248
249        let listener_fd = listener_result?;
250        let core = PoolCore::<Tcp, H::CodecLayer, crate::WriteBufArr, ServerSlotState<H>, B>::new(
251            max_conn,
252            slot_id_base,
253        )?;
254        let accept_slot_id = slot_id_base + max_conn;
255
256        let direction = Server {
257            protocol,
258            listener_fd,
259            accept_slot: AcceptSlot::default(),
260            accept_pending: VecDeque::with_capacity(ACCEPT_PENDING_INIT_CAP),
261            socket_opts,
262            accept_slot_id,
263            transport: PhantomData,
264        };
265
266        let mut reactor = Self {
267            direction,
268            core,
269            profile: PhantomData,
270        };
271        reactor.direction.accept_arm::<B>(driver);
272        if !reactor.direction.accept_slot.armed && !reactor.direction.accept_slot.pending_rearm {
273            return Err(io::Error::other("dope-net: failed to arm multishot accept"));
274        }
275        Ok(reactor)
276    }
277}
278
279impl<H, T, F, B> Pool<Server<H, T>, F, B>
280where
281    H: ServerProtocol,
282    T: Transport,
283    F: Profile,
284    B: PoolDriver,
285{
286    pub fn slot_id_range(&self) -> std::ops::Range<usize> {
287        let base = self.core.slot_id_base;
288        base..(self.direction.accept_slot_id + 1)
289    }
290
291    pub fn conn_slot_count(&self) -> usize {
292        self.core.slots.len()
293    }
294
295    pub fn protocol(&self) -> &H {
296        &self.direction.protocol
297    }
298
299    pub fn protocol_mut(&mut self) -> &mut H {
300        &mut self.direction.protocol
301    }
302
303    pub fn listener_fd(&self) -> &FdHandle {
304        &self.direction.listener_fd
305    }
306}
307
308impl<H, T, F, B> super::PoolTick<B> for Pool<Server<H, T>, F, B>
309where
310    H: ServerProtocol,
311    T: Transport,
312    F: Profile,
313    B: PoolDriver,
314{
315    #[inline(always)]
316    fn tick(&mut self, driver: &mut B::Driver) {
317        self.direction.tick(&mut self.core, driver);
318        self.drain_ready_queue(driver);
319    }
320}