ntex_net/polling/
reactor.rs

1use std::os::fd::{AsRawFd, BorrowedFd, RawFd};
2use std::{cell::Cell, cell::UnsafeCell, fmt, io, net, rc::Rc};
3use std::{collections::VecDeque, num::NonZeroUsize, time::Duration};
4
5#[cfg(unix)]
6use std::os::unix::net::UnixStream as OsUnixStream;
7
8use ::ntex_polling::{Event, Events, Notifier, PollMode, Poller};
9use ntex_io::Io;
10use ntex_rt::{DriverType, Notify, PollResult, Runtime};
11use ntex_service::cfg::SharedCfg;
12use socket2::{Protocol, SockAddr, Socket, Type};
13
14use super::{TcpStream, UnixStream, stream::StreamOps};
15use crate::channel::Receiver;
16
17pub trait Handler {
18    /// Submitted interest
19    fn event(&mut self, id: usize, event: Event);
20
21    /// Operation submission has failed
22    fn error(&mut self, id: usize, err: io::Error);
23
24    /// Process deferred work after events and before the driver waits.
25    ///
26    /// Called even when no events were received.
27    fn tick(&mut self);
28
29    /// Cleanup before drop
30    fn cleanup(&mut self);
31}
32
33type HandlerItem = Box<dyn Handler>;
34
35enum Change {
36    Error {
37        batch: usize,
38        user_data: u32,
39        error: io::Error,
40    },
41}
42
43#[derive(Debug)]
44/// Polling reactor api.
45pub struct ReactorApi {
46    id: usize,
47    batch: u64,
48    poll: Rc<Poller>,
49    changes: Rc<UnsafeCell<VecDeque<Change>>>,
50}
51
52impl ReactorApi {
53    /// Attach an fd to the driver.
54    ///
55    /// `fd` must be attached to the driver before using register/unregister
56    /// methods.
57    pub fn attach(&self, fd: RawFd, id: u32, event: Event) {
58        self.attach_with_mode(fd, id, event, PollMode::Oneshot);
59    }
60
61    /// Attach an fd to the driver with specific mode.
62    ///
63    /// `fd` must be attached to the driver before using register/unregister
64    /// methods.
65    pub fn attach_with_mode(&self, fd: RawFd, id: u32, mut event: Event, mode: PollMode) {
66        event.key = (u64::from(id) | self.batch) as usize;
67        if let Err(err) = unsafe { self.poll.add_with_mode(fd, event, mode) } {
68            self.change(Change::Error {
69                batch: self.id,
70                user_data: id,
71                error: err,
72            });
73        }
74    }
75
76    /// Detach an fd from the driver.
77    pub fn detach(&self, fd: RawFd, id: u32) {
78        if let Err(err) = self.poll.delete(unsafe { BorrowedFd::borrow_raw(fd) }) {
79            self.change(Change::Error {
80                batch: self.id,
81                user_data: id,
82                error: err,
83            });
84        }
85    }
86
87    /// Register interest for specified file descriptor.
88    pub fn modify(&self, fd: RawFd, id: u32, event: Event) {
89        self.modify_with_mode(fd, id, event, PollMode::Oneshot);
90    }
91
92    /// Register interest for specified file descriptor.
93    pub fn modify_with_mode(&self, fd: RawFd, id: u32, mut event: Event, mode: PollMode) {
94        event.key = (u64::from(id) | self.batch) as usize;
95
96        let result = self
97            .poll
98            .modify_with_mode(unsafe { BorrowedFd::borrow_raw(fd) }, event, mode);
99        self.check(id, result);
100    }
101
102    /// Register interest for specified file descriptor, the change may be
103    /// deferred to the next poll.
104    ///
105    /// On kqueue the change is submitted together with the next wait, saving
106    /// a syscall. A failure of a deferred change is reported as a readable
107    /// and writable event.
108    pub fn modify_with_mode_deferred(&self, fd: RawFd, id: u32, mut event: Event, mode: PollMode) {
109        event.key = (u64::from(id) | self.batch) as usize;
110
111        let result =
112            self.poll
113                .modify_with_mode_deferred(unsafe { BorrowedFd::borrow_raw(fd) }, event, mode);
114        self.check(id, result);
115    }
116
117    /// Whether the poller supports level-triggered events.
118    pub fn supports_level(&self) -> bool {
119        self.poll.supports_level()
120    }
121
122    fn check(&self, id: u32, result: io::Result<()>) {
123        if let Err(err) = result {
124            self.change(Change::Error {
125                batch: self.id,
126                user_data: id,
127                error: err,
128            });
129        }
130    }
131
132    fn change(&self, ev: Change) {
133        unsafe { (*self.changes.get()).push_back(ev) };
134    }
135}
136
137/// Polling reactor.
138///
139/// Uses `epoll` or `kqueue`, depending on the platform.
140pub struct Reactor {
141    poll: Rc<Poller>,
142    capacity: usize,
143    changes: Rc<UnsafeCell<VecDeque<Change>>>,
144    hid: Cell<u64>,
145    #[allow(clippy::box_collection)]
146    handlers: Cell<Option<Box<Vec<HandlerItem>>>>,
147}
148
149impl Reactor {
150    const BATCH: u64 = 48;
151    const BATCH_MASK: u64 = 0xFFFF_0000_0000_0000;
152    const DATA_MASK: u64 = 0x0000_FFFF_FFFF_FFFF;
153
154    pub fn new() -> io::Result<Self> {
155        Reactor::with_capacity(2048)
156    }
157
158    pub fn with_capacity(io_queue_capacity: u32) -> io::Result<Self> {
159        log::trace!("New poll driver");
160
161        Ok(Self {
162            hid: Cell::new(0),
163            poll: Rc::new(Poller::new()?),
164            capacity: io_queue_capacity as usize,
165            changes: Rc::new(UnsafeCell::new(VecDeque::with_capacity(32))),
166            handlers: Cell::new(Some(Box::new(Vec::default()))),
167        })
168    }
169
170    /// Reactor type
171    pub const fn tp(&self) -> DriverType {
172        DriverType::Poll
173    }
174
175    /// Register updates handler
176    pub fn register<F>(&self, f: F)
177    where
178        F: FnOnce(ReactorApi) -> Box<dyn Handler>,
179    {
180        let id = self.hid.get();
181        let mut handlers = self
182            .handlers
183            .take()
184            .expect("Cannot register handler during event handling");
185
186        let api = ReactorApi {
187            id: id as usize,
188            batch: id << Self::BATCH,
189            poll: self.poll.clone(),
190            changes: self.changes.clone(),
191        };
192        handlers.push(f(api));
193        self.hid.set(id + 1);
194        self.handlers.set(Some(handlers));
195    }
196
197    fn apply_changes(&self, handlers: &mut [HandlerItem]) {
198        while let Some(op) = unsafe { (*self.changes.get()).pop_front() } {
199            match op {
200                Change::Error {
201                    batch,
202                    user_data,
203                    error,
204                } => handlers[batch].error(user_data as usize, error),
205            }
206        }
207    }
208
209    fn tick(&self, handlers: &mut [HandlerItem]) {
210        loop {
211            self.apply_changes(handlers);
212            for h in handlers.iter_mut() {
213                h.tick();
214            }
215            // Ticks can fail poller operations, and handling those errors can
216            // enqueue more deferred work. Finish both before waiting.
217            if unsafe { (*self.changes.get()).is_empty() } {
218                break;
219            }
220        }
221    }
222}
223
224impl AsRawFd for Reactor {
225    fn as_raw_fd(&self) -> RawFd {
226        self.poll.as_raw_fd()
227    }
228}
229
230impl fmt::Debug for Reactor {
231    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
232        f.debug_struct("Reactor")
233            .field("poll", &self.poll)
234            .field("capacity", &self.capacity)
235            .field("hid", &self.hid)
236            .finish()
237    }
238}
239
240impl crate::Reactor for Reactor {
241    fn tcp_connect(&self, addr: net::SocketAddr, cfg: SharedCfg) -> Receiver<Io> {
242        let addr = SockAddr::from(addr);
243        let result = Socket::new(addr.domain(), Type::STREAM, Some(Protocol::TCP))
244            .and_then(crate::helpers::prep_tcp_socket)
245            .and_then(crate::helpers::prep_socket)
246            .map(move |sock| (addr, sock));
247
248        match result {
249            Err(err) => Receiver::new(Err(err)),
250            Ok((addr, sock)) => {
251                super::connect::ConnectOps::get(self).connect(sock, &addr, cfg, false)
252            }
253        }
254    }
255
256    fn unix_connect(&self, addr: std::path::PathBuf, cfg: SharedCfg) -> Receiver<Io> {
257        let result = SockAddr::unix(addr).and_then(|addr| {
258            Socket::new(addr.domain(), Type::STREAM, None)
259                .and_then(crate::helpers::prep_socket)
260                .map(move |sock| (addr, sock))
261        });
262
263        match result {
264            Err(err) => Receiver::new(Err(err)),
265            Ok((addr, sock)) => {
266                super::connect::ConnectOps::get(self).connect(sock, &addr, cfg, true)
267            }
268        }
269    }
270
271    fn from_tcp_stream(&self, stream: net::TcpStream, cfg: SharedCfg) -> io::Result<Io> {
272        stream.set_nodelay(true)?;
273
274        Ok(Io::new(
275            TcpStream(
276                crate::helpers::prep_socket(Socket::from(stream))?,
277                StreamOps::get(self),
278            ),
279            cfg,
280        ))
281    }
282
283    #[cfg(unix)]
284    fn from_unix_stream(&self, stream: OsUnixStream, cfg: SharedCfg) -> io::Result<Io> {
285        Ok(Io::new(
286            UnixStream(
287                crate::helpers::prep_socket(Socket::from(stream))?,
288                StreamOps::get(self),
289            ),
290            cfg,
291        ))
292    }
293}
294
295impl ntex_rt::Driver for Reactor {
296    /// Poll the driver and handle completed entries.
297    fn run(&self, rt: &Runtime) -> io::Result<()> {
298        let mut events = if self.capacity == 0 {
299            Events::new()
300        } else {
301            Events::with_capacity(NonZeroUsize::new(self.capacity).unwrap())
302        };
303
304        let result = loop {
305            let timeout = match rt.poll() {
306                PollResult::Pending => None,
307                PollResult::PollAgain => Some(Duration::ZERO),
308                PollResult::Ready => break Ok(()),
309            };
310            // Runtime tasks can queue cleanup without receiving a stream event.
311            let mut handlers = self.handlers.take().unwrap();
312            self.tick(&mut handlers);
313            self.handlers.set(Some(handlers));
314
315            events.clear();
316            self.poll.wait(&mut events, timeout)?;
317            // tasks woken until the runtime is polled do not need to notify
318            rt.awake();
319
320            let mut handlers = self.handlers.take().unwrap();
321            for event in events.iter() {
322                let key = event.key as u64;
323                let batch = ((key & Self::BATCH_MASK) >> Self::BATCH) as usize;
324                handlers[batch].event((key & Self::DATA_MASK) as usize, event);
325            }
326            self.tick(&mut handlers);
327            self.handlers.set(Some(handlers));
328        };
329
330        for mut h in self.handlers.take().unwrap().into_iter() {
331            h.cleanup();
332        }
333        result
334    }
335
336    /// Get notification handle
337    fn handle(&self) -> Box<dyn Notify> {
338        Box::new(NotifyHandle::new(self.poll.notifier()))
339    }
340
341    /// Clear handlers
342    fn clear(&self) {}
343}
344
345#[derive(Clone, Debug)]
346/// A notify handle to the inner driver.
347pub(crate) struct NotifyHandle {
348    notifier: Notifier,
349}
350
351impl NotifyHandle {
352    fn new(notifier: Notifier) -> Self {
353        Self { notifier }
354    }
355}
356
357impl Notify for NotifyHandle {
358    /// Notify the driver
359    fn notify(&self) -> io::Result<()> {
360        self.notifier.notify()
361    }
362}
363
364#[cfg(test)]
365mod tests {
366    use std::cell::RefCell;
367
368    use super::*;
369
370    struct Recorder(Rc<RefCell<Vec<(usize, io::ErrorKind)>>>);
371
372    impl Handler for Recorder {
373        fn event(&mut self, _: usize, _: Event) {}
374
375        fn error(&mut self, id: usize, err: io::Error) {
376            self.0.borrow_mut().push((id, err.kind()));
377        }
378
379        fn tick(&mut self) {}
380
381        fn cleanup(&mut self) {}
382    }
383
384    #[test]
385    fn reactor_info() {
386        let reactor = Reactor::with_capacity(0).unwrap();
387        assert_eq!(reactor.tp(), DriverType::Poll);
388        assert!(reactor.as_raw_fd() >= 0);
389        let s = format!("{reactor:?}");
390        assert!(s.contains("Reactor") && s.contains("capacity"), "{s}");
391    }
392
393    /// Failed poller operations are reported to the handler that submitted them.
394    #[test]
395    fn api_errors_are_reported() {
396        let reactor = Reactor::new().unwrap();
397        let errors = Rc::new(RefCell::new(Vec::new()));
398        let mut api = None;
399        reactor.register(|a| {
400            api = Some(a);
401            Box::new(Recorder(errors.clone()))
402        });
403        let api = api.unwrap();
404
405        // the socket is not attached
406        let (sock, _peer) = OsUnixStream::pair().unwrap();
407        api.modify(sock.as_raw_fd(), 1, Event::readable(0));
408        api.detach(sock.as_raw_fd(), 2);
409        assert!(errors.borrow().is_empty());
410
411        let mut handlers = reactor.handlers.take().unwrap();
412        reactor.apply_changes(&mut handlers);
413        reactor.handlers.set(Some(handlers));
414
415        let ids: Vec<_> = errors.borrow().iter().map(|(id, _)| *id).collect();
416        assert_eq!(ids, vec![1, 2]);
417    }
418
419    #[test]
420    fn tick_drains_errors_and_their_deferred_work() {
421        struct TickErrors {
422            api: ReactorApi,
423            socket: OsUnixStream,
424            next: Option<u32>,
425            errors: Rc<RefCell<Vec<usize>>>,
426        }
427
428        impl Handler for TickErrors {
429            fn event(&mut self, _: usize, _: Event) {
430                panic!("unexpected socket event");
431            }
432
433            fn error(&mut self, id: usize, _: io::Error) {
434                self.errors.borrow_mut().push(id);
435                if id == 1 {
436                    self.next = Some(2);
437                }
438            }
439
440            fn tick(&mut self) {
441                if let Some(id) = self.next.take() {
442                    // The socket is not attached, so the operation fails.
443                    self.api
444                        .modify(self.socket.as_raw_fd(), id, Event::readable(0));
445                }
446            }
447
448            fn cleanup(&mut self) {}
449        }
450
451        let reactor = Reactor::new().unwrap();
452        let (socket, _peer) = OsUnixStream::pair().unwrap();
453        let errors = Rc::new(RefCell::new(Vec::new()));
454        reactor.register(|api| {
455            Box::new(TickErrors {
456                api,
457                socket,
458                next: Some(1),
459                errors: errors.clone(),
460            })
461        });
462
463        let mut handlers = reactor.handlers.take().unwrap();
464        reactor.tick(&mut handlers);
465        reactor.handlers.set(Some(handlers));
466
467        assert_eq!(*errors.borrow(), vec![1, 2]);
468        assert!(unsafe { (*reactor.changes.get()).is_empty() });
469    }
470}
ntex_net/polling/reactor.rs

ntex_net/polling/
reactor.rs
