ntex_server/net/
accept.rs

1#![allow(clippy::missing_panics_doc)]
2use std::time::{Duration, Instant};
3use std::{cell::Cell, fmt, io, sync::Arc, sync::mpsc, thread};
4use std::{collections::VecDeque, num::NonZeroUsize};
5
6use ntex_rt::System;
7use ntex_util::{future::Either, time::Millis, time::sleep};
8use polling::{Event, Events, Poller};
9
10use super::socket::{Connection, Listener, SocketAddr};
11use super::{Server, ServerStatus, Token};
12
13const EXIT_TIMEOUT: Duration = Duration::from_millis(100);
14const ERR_TIMEOUT: Duration = Duration::from_millis(500);
15const ERR_SLEEP_TIMEOUT: Millis = Millis(525);
16
17/// Command sent to the accept loop.
18#[derive(Debug)]
19pub enum AcceptorCommand {
20    /// Stops accepting and closes the listeners, then signals the sender.
21    Stop(oneshot::Sender<()>),
22    /// Stops accepting and closes the listeners immediately.
23    Terminate,
24    /// Stops accepting connections.
25    Pause,
26    /// Resumes accepting connections.
27    Resume,
28    /// Re-registers listeners that were paused after an accept error.
29    Timer,
30}
31
32#[derive(Debug)]
33struct ServerSocketInfo {
34    addr: SocketAddr,
35    token: Token,
36    sock: Listener,
37    registered: Cell<bool>,
38    timeout: Cell<Option<Instant>>,
39}
40
41/// Handle for sending commands to the accept loop.
42#[derive(Debug, Clone)]
43pub struct AcceptNotify(Arc<Poller>, mpsc::Sender<AcceptorCommand>);
44
45impl AcceptNotify {
46    fn new(waker: Arc<Poller>, tx: mpsc::Sender<AcceptorCommand>) -> Self {
47        AcceptNotify(waker, tx)
48    }
49
50    /// Sends a command and wakes the accept loop.
51    pub fn send(&self, cmd: AcceptorCommand) {
52        let _ = self.1.send(cmd);
53        let _ = self.0.notify();
54    }
55}
56
57/// Accept loop for network listeners.
58///
59/// The loop runs on its own thread and passes accepted connections to the
60/// server.
61pub struct AcceptLoop {
62    name: String,
63    testing: bool,
64    notify: AcceptNotify,
65    inner: Option<(mpsc::Receiver<AcceptorCommand>, Arc<Poller>)>,
66    status_handler: Option<Box<dyn FnMut(ServerStatus) + Send>>,
67}
68
69impl Default for AcceptLoop {
70    fn default() -> Self {
71        Self::new()
72    }
73}
74
75impl AcceptLoop {
76    /// Creates an accept loop.
77    ///
78    /// # Panics
79    ///
80    /// Panics if the poller cannot be created.
81    pub fn new() -> AcceptLoop {
82        // Create a poller instance
83        let poll = Arc::new(
84            Poller::new()
85                .map_err(|e| panic!("Cannot create Poller {e}"))
86                .unwrap(),
87        );
88
89        let (tx, rx) = mpsc::channel();
90        let notify = AcceptNotify::new(poll.clone(), tx);
91
92        AcceptLoop {
93            notify,
94            name: "ntex:accept".to_string(),
95            inner: Some((rx, poll)),
96            testing: false,
97            status_handler: None,
98        }
99    }
100
101    /// Sets the server name.
102    ///
103    /// The accept thread is named `{name}:accept`.
104    pub fn name<T: AsRef<str>>(&mut self, name: T) {
105        self.name = format!("{}:accept", name.as_ref());
106    }
107
108    /// Returns a handle for sending commands to the loop.
109    pub fn notify(&self) -> AcceptNotify {
110        self.notify.clone()
111    }
112
113    /// Sets the handler that receives accept status changes.
114    ///
115    /// See [`ServerBuilder::status_handler`](super::ServerBuilder::status_handler).
116    pub fn set_status_handler<F>(&mut self, f: F)
117    where
118        F: FnMut(ServerStatus) + Send + 'static,
119    {
120        self.status_handler = Some(Box::new(f));
121    }
122
123    /// Enables test mode, which skips the exit delay on stop.
124    pub fn testing(&mut self) {
125        self.testing = true;
126    }
127
128    /// Starts the accept loop on a new thread.
129    ///
130    /// Blocks until the listeners are registered.
131    pub fn start(mut self, socks: Vec<(Token, Listener)>, srv: Server) {
132        let (tx, rx_start) = oneshot::channel();
133        let (rx, poll) = self
134            .inner
135            .take()
136            .expect("AcceptLoop cannot be used multiple times");
137
138        Accept::start(
139            tx,
140            rx,
141            poll,
142            socks,
143            srv,
144            self.name.clone(),
145            self.notify.clone(),
146            self.testing,
147            self.status_handler.take(),
148        );
149
150        let _ = rx_start.recv();
151    }
152}
153
154impl fmt::Debug for AcceptLoop {
155    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
156        f.debug_struct("AcceptLoop")
157            .field("name", &self.name)
158            .field("notify", &self.notify)
159            .field("inner", &self.inner)
160            .field("status_handler", &self.status_handler.is_some())
161            .finish()
162    }
163}
164
165struct Accept {
166    name: String,
167    poller: Arc<Poller>,
168    rx: mpsc::Receiver<AcceptorCommand>,
169    tx: Option<oneshot::Sender<()>>,
170    sockets: Vec<ServerSocketInfo>,
171    srv: Server,
172    notify: AcceptNotify,
173    testing: bool,
174    backpressure: bool,
175    backlog: VecDeque<Connection>,
176    status_handler: Option<Box<dyn FnMut(ServerStatus) + Send>>,
177}
178
179impl Accept {
180    #[allow(clippy::too_many_arguments)]
181    fn start(
182        tx: oneshot::Sender<()>,
183        rx: mpsc::Receiver<AcceptorCommand>,
184        poller: Arc<Poller>,
185        socks: Vec<(Token, Listener)>,
186        srv: Server,
187        name: String,
188        notify: AcceptNotify,
189        testing: bool,
190        status_handler: Option<Box<dyn FnMut(ServerStatus) + Send>>,
191    ) {
192        log::info!("Starting {name:?} accept loop");
193
194        // start accept thread
195        let sys = System::current();
196        let _ = thread::Builder::new().name(name.clone()).spawn(move || {
197            System::set_current(sys);
198            Accept::new(
199                name,
200                tx,
201                rx,
202                poller,
203                socks,
204                srv,
205                notify,
206                testing,
207                status_handler,
208            )
209            .poll();
210        });
211    }
212
213    #[allow(clippy::too_many_arguments)]
214    fn new(
215        name: String,
216        tx: oneshot::Sender<()>,
217        rx: mpsc::Receiver<AcceptorCommand>,
218        poller: Arc<Poller>,
219        socks: Vec<(Token, Listener)>,
220        srv: Server,
221        notify: AcceptNotify,
222        testing: bool,
223        status_handler: Option<Box<dyn FnMut(ServerStatus) + Send>>,
224    ) -> Accept {
225        let mut sockets = Vec::new();
226        for (hnd_token, lst) in socks {
227            sockets.push(ServerSocketInfo {
228                addr: lst.local_addr(),
229                sock: lst,
230                token: hnd_token,
231                registered: Cell::new(false),
232                timeout: Cell::new(None),
233            });
234        }
235
236        Accept {
237            name,
238            poller,
239            rx,
240            sockets,
241            notify,
242            srv,
243            testing,
244            status_handler,
245            tx: Some(tx),
246            backpressure: true,
247            backlog: VecDeque::new(),
248        }
249    }
250
251    fn update_status(&mut self, st: ServerStatus) {
252        if let Some(ref mut hnd) = self.status_handler {
253            (*hnd)(st);
254        }
255    }
256
257    fn poll(mut self) {
258        // Create storage for events
259        let mut events = Events::with_capacity(NonZeroUsize::new(512).unwrap());
260
261        // notify start
262        for idx in 0..self.sockets.len() {
263            self.add_source(idx);
264        }
265        if let Some(tx) = self.tx.take() {
266            thread::sleep(Duration::from_millis(25));
267            let _ = tx.send(());
268        }
269
270        loop {
271            for idx in 0..self.sockets.len() {
272                if self.sockets[idx].registered.get() {
273                    let readd = self.accept(idx);
274                    if readd {
275                        self.add_source(idx);
276                    }
277                }
278            }
279
280            if let Either::Right(rx) = self.process_cmd() {
281                // cleanup
282                for info in self.sockets.drain(..) {
283                    info.sock.remove_source();
284                }
285                log::info!("Accept loop {:?} has been stopped", self.name);
286
287                if let Some(rx) = rx {
288                    if !self.testing {
289                        thread::sleep(EXIT_TIMEOUT);
290                    }
291                    let _ = rx.send(());
292                }
293
294                break;
295            }
296
297            events.clear();
298            if let Err(e) = self.poller.wait(&mut events, None) {
299                assert!(
300                    e.kind() == io::ErrorKind::Interrupted,
301                    "Cannot wait for events in poller: {e}"
302                );
303            }
304        }
305    }
306
307    fn add_source(&self, idx: usize) {
308        let info = &self.sockets[idx];
309
310        loop {
311            // try to register poller source
312            let result = if info.registered.get() {
313                self.poller.modify(&info.sock, Event::readable(idx))
314            } else {
315                unsafe { self.poller.add(&info.sock, Event::readable(idx)) }
316            };
317            if let Err(err) = result {
318                if err.kind() == io::ErrorKind::WouldBlock {
319                    continue;
320                }
321                log::error!("Cannot register socket listener: {err}");
322
323                // sleep after error
324                info.timeout.set(Some(Instant::now() + ERR_TIMEOUT));
325
326                let notify = self.notify.clone();
327                System::current().handle().spawn(async move {
328                    sleep(ERR_SLEEP_TIMEOUT).await;
329                    notify.send(AcceptorCommand::Timer);
330                });
331            } else {
332                info.registered.set(true);
333            }
334
335            break;
336        }
337    }
338
339    fn remove_source(&self, key: usize) {
340        let info = &self.sockets[key];
341
342        let result = if info.registered.get() {
343            self.poller.modify(&info.sock, Event::none(key))
344        } else {
345            return;
346        };
347
348        // stop listening for incoming connections
349        if let Err(err) = result {
350            log::error!("Cannot stop socket listener for {} err: {}", info.addr, err);
351        }
352    }
353
354    fn process_timer(&mut self) {
355        let now = Instant::now();
356        for key in 0..self.sockets.len() {
357            let info = &mut self.sockets[key];
358            if let Some(inst) = info.timeout.get()
359                && now > inst
360                && !self.backpressure
361            {
362                log::info!("Resuming socket listener on {} after timeout", info.addr);
363                info.timeout.take();
364                self.add_source(key);
365            }
366        }
367    }
368
369    fn process_cmd(&mut self) -> Either<(), Option<oneshot::Sender<()>>> {
370        loop {
371            match self.rx.try_recv() {
372                Ok(cmd) => match cmd {
373                    AcceptorCommand::Stop(rx) => {
374                        if !self.backpressure {
375                            log::info!("Stopping accept loop {:?}", self.name);
376                            self.backpressure(true);
377                        }
378                        break Either::Right(Some(rx));
379                    }
380                    AcceptorCommand::Terminate => {
381                        log::info!("Stopping accept loop {:?}", self.name);
382                        self.backpressure(true);
383                        break Either::Right(None);
384                    }
385                    AcceptorCommand::Pause => {
386                        if !self.backpressure {
387                            log::info!("Pausing accept loop {:?}", self.name);
388                            self.backpressure(true);
389                        }
390                    }
391                    AcceptorCommand::Resume => {
392                        if self.backpressure {
393                            log::info!("Resuming accept loop {:?}", self.name);
394                            self.backpressure(false);
395                        }
396                    }
397                    AcceptorCommand::Timer => {
398                        self.process_timer();
399                    }
400                },
401                Err(err) => {
402                    break match err {
403                        mpsc::TryRecvError::Empty => Either::Left(()),
404                        mpsc::TryRecvError::Disconnected => {
405                            log::error!("Dropping accept loop {:?}", self.name);
406                            self.backpressure(true);
407                            Either::Right(None)
408                        }
409                    };
410                }
411            }
412        }
413    }
414
415    fn backpressure(&mut self, on: bool) {
416        self.update_status(if on {
417            ServerStatus::NotReady
418        } else {
419            ServerStatus::Ready
420        });
421
422        if self.backpressure && !on {
423            // handle backlog
424            while let Some(msg) = self.backlog.pop_front() {
425                if let Err(msg) = self.srv.process(msg) {
426                    log::trace!("Server is unavailable");
427                    self.backlog.push_front(msg);
428                    return;
429                }
430            }
431
432            // re-enable acceptors
433            self.backpressure = false;
434            for (key, info) in self.sockets.iter().enumerate() {
435                if info.timeout.get().is_none() {
436                    // socket with timeout will re-register itself after timeout
437                    log::info!(
438                        "Resuming socket listener on {} after back-pressure",
439                        info.addr
440                    );
441                    self.add_source(key);
442                }
443            }
444        } else if !self.backpressure && on {
445            self.backpressure = true;
446            for key in 0..self.sockets.len() {
447                // disable err timeout
448                let info = &mut self.sockets[key];
449                if info.timeout.take().is_none() {
450                    log::info!("Stopping socket listener on {}", info.addr);
451                    self.remove_source(key);
452                }
453            }
454        }
455    }
456
457    fn accept(&mut self, token: usize) -> bool {
458        loop {
459            if let Some(info) = self.sockets.get_mut(token) {
460                match info.sock.accept() {
461                    Ok(Some(io)) => {
462                        let msg = Connection {
463                            io,
464                            token: info.token,
465                        };
466                        if let Err(msg) = self.srv.process(msg) {
467                            log::trace!("Server is unavailable");
468                            self.backlog.push_back(msg);
469                            self.backpressure(true);
470                            return false;
471                        }
472                    }
473                    Ok(None) => return true,
474                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return true,
475                    Err(ref e) if connection_error(e) => (),
476                    Err(e) => {
477                        log::error!("Error accepting socket: {e}");
478
479                        // sleep after error
480                        info.timeout.set(Some(Instant::now() + ERR_TIMEOUT));
481
482                        let notify = self.notify.clone();
483                        System::current().handle().spawn(async move {
484                            sleep(ERR_SLEEP_TIMEOUT).await;
485                            notify.send(AcceptorCommand::Timer);
486                        });
487                        return false;
488                    }
489                }
490            }
491        }
492    }
493}
494
495/// This function defines errors that are per-connection. Which basically
496/// means that if we get this error from `accept()` system call it means
497/// next connection might be ready to be accepted.
498///
499/// All other errors will incur a timeout before next `accept()` is performed.
500/// The timeout is useful to handle resource exhaustion errors like ENFILE
501/// and EMFILE. Otherwise, could enter into tight loop.
502fn connection_error(e: &io::Error) -> bool {
503    matches!(
504        e.kind(),
505        io::ErrorKind::ConnectionRefused
506            | io::ErrorKind::ConnectionAborted
507            | io::ErrorKind::ConnectionReset
508    )
509}
ntex_server/net/accept.rs

ntex_server/net/
accept.rs
