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