use std::os::fd::{AsRawFd, RawFd}; use std::{cell::RefCell, collections::VecDeque, io, rc::Rc, task::Poll}; use ntex_neon::driver::{DriverApi, Handler, Interest}; use ntex_neon::{syscall, Runtime}; use ntex_util::channel::oneshot::Sender; use slab::Slab; use socket2::SockAddr; #[derive(Clone)] pub(crate) struct ConnectOps(Rc); #[derive(Debug)] enum Change { Readable, Writable, Error(io::Error), } struct ConnectOpsBatcher { feed: VecDeque<(usize, Change)>, inner: Rc, } struct Item { fd: RawFd, sender: Sender>, } struct ConnectOpsInner { api: DriverApi, connects: RefCell>, } impl ConnectOps { pub(crate) fn current() -> Self { Runtime::value(|rt| { let mut inner = None; rt.driver().register(|api| { let ops = Rc::new(ConnectOpsInner { api, connects: RefCell::new(Slab::new()), }); inner = Some(ops.clone()); Box::new(ConnectOpsBatcher { inner: ops, feed: VecDeque::new(), }) }); ConnectOps(inner.unwrap()) }) } pub(crate) fn connect( &self, fd: RawFd, addr: SockAddr, sender: Sender>, ) -> io::Result { let result = syscall!(break libc::connect(fd, addr.as_ptr(), addr.len())); if let Poll::Ready(res) = result { res?; } let item = Item { fd, sender }; let id = self.0.connects.borrow_mut().insert(item); self.0.api.register(fd, id, Interest::Writable); Ok(id) } } impl Handler for ConnectOpsBatcher { fn readable(&mut self, id: usize) { log::debug!("connect-fd is readable {:?}", id); self.feed.push_back((id, Change::Readable)); } fn writable(&mut self, id: usize) { log::debug!("connect-fd is writable {:?}", id); self.feed.push_back((id, Change::Writable)); } fn error(&mut self, id: usize, err: io::Error) { self.feed.push_back((id, Change::Error(err))); } fn commit(&mut self) { if self.feed.is_empty() { return; } log::debug!("Commit connect driver changes, num: {:?}", self.feed.len()); let mut connects = self.inner.connects.borrow_mut(); for (id, change) in self.feed.drain(..) { if connects.contains(id) { let item = connects.remove(id); match change { Change::Readable => unreachable!(), Change::Writable => { let mut err: libc::c_int = 0; let mut err_len = std::mem::size_of::() as libc::socklen_t; let res = syscall!(libc::getsockopt( item.fd.as_raw_fd(), libc::SOL_SOCKET, libc::SO_ERROR, &mut err as *mut _ as *mut _, &mut err_len )); let res = if err == 0 { res.map(|_| ()) } else { Err(io::Error::from_raw_os_error(err)) }; self.inner.api.unregister_all(item.fd); let _ = item.sender.send(res); } Change::Error(err) => { let _ = item.sender.send(Err(err)); self.inner.api.unregister_all(item.fd); } } } } } }