use std::{ future::Future, io, pin::Pin, task::{Context, Poll}, thread, time::Duration, }; use actix_rt::time::{sleep, Instant, Sleep}; use log::error; use tokio::sync::mpsc::UnboundedReceiver; use crate::server::Server; use crate::socket::Listener; use crate::worker::{Conn, WorkerHandleAccept}; struct ServerSocketInfo { token: usize, lst: Listener, } /// poll instance of the server. pub(crate) struct Accept { handles: Vec, sockets: Box<[ServerSocketInfo]>, rx: UnboundedReceiver, srv: Server, next: usize, avail: Availability, state: AcceptState, timeout: Pin>, } enum AcceptState { Accept, Pause, Stop, } pub(crate) enum Interest { Pause, Resume, Stop, WorkerIndex(usize), Worker(WorkerHandleAccept), } impl Accept { pub(crate) fn start( socks: Vec<(usize, Listener)>, rx: UnboundedReceiver, srv: Server, handles: Vec, ) { // Accept runs in its own thread and tokio runtime thread::Builder::new() .name("actix-server accept loop".to_owned()) .spawn(move || { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() .block_on(async { let accept = Self::new(socks, rx, srv, handles); accept.await }); }) .unwrap(); } fn new( socks: Vec<(usize, Listener)>, rx: UnboundedReceiver, srv: Server, handles: Vec, ) -> Self { let sockets = socks .into_iter() .map(|(token, lst)| ServerSocketInfo { token, lst }) .collect(); let mut avail = Availability::default(); // Assume all handles are avail at construct time. avail.set_available_all(&handles); Accept { handles, sockets, rx, srv, next: 0, avail, state: AcceptState::Accept, timeout: Box::pin(sleep(Duration::from_millis(500))), } } fn poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<()> { let len = self.sockets.len(); let mut idx = 0; while idx < len { 'socket: loop { if !self.avail.available() { return Poll::Pending; } let socket = &mut self.sockets[idx]; match socket.lst.poll_accept(cx) { Poll::Ready(Ok(io)) => { let conn = Conn { io, token: socket.token, }; self.accept_one(conn); } Poll::Ready(Err(ref e)) if connection_error(e) => continue 'socket, Poll::Ready(Err(ref e)) => { error!("Error accepting connection: {}", e); let deadline = Instant::now() + Duration::from_millis(500); self.timeout.as_mut().reset(deadline); let _ = self.timeout.as_mut().poll(cx); break 'socket; } Poll::Pending => break 'socket, }; } idx += 1; } Poll::Pending } // Send connection to worker and handle error. fn send_connection(&mut self, conn: Conn) -> Result<(), Conn> { let next = self.next(); match next.send(conn) { Ok(_) => { // Increment counter of WorkerHandle. // Set worker to unavailable with it hit max (Return false). if !next.incr_counter() { let idx = next.idx(); self.avail.set_available(idx, false); } self.set_next(); Ok(()) } Err(conn) => { // Worker thread is error and could be gone. // Remove worker handle and notify `ServerBuilder`. self.remove_next(); if self.handles.is_empty() { error!("No workers"); // All workers are gone and Conn is nowhere to be sent. // Treat this situation as Ok and drop Conn. return Ok(()); } else if self.handles.len() <= self.next { self.next = 0; } Err(conn) } } } fn accept_one(&mut self, mut conn: Conn) { loop { let next = self.next(); let idx = next.idx(); if self.avail.get_available(idx) { match self.send_connection(conn) { Ok(_) => return, Err(c) => conn = c, } } else { self.avail.set_available(idx, false); self.set_next(); if !self.avail.available() { while let Err(c) = self.send_connection(conn) { conn = c; } return; } } } } #[inline(always)] fn next(&self) -> &WorkerHandleAccept { &self.handles[self.next] } /// Set next worker handle that would accept connection. #[inline(always)] fn set_next(&mut self) { self.next = (self.next + 1) % self.handles.len(); } /// Remove next worker handle that fail to accept connection. fn remove_next(&mut self) { let handle = self.handles.swap_remove(self.next); let idx = handle.idx(); // A message is sent to `ServerBuilder` future to notify it a new worker // should be made. self.srv.worker_faulted(idx); self.avail.set_available(idx, false); } } impl Future for Accept { type Output = (); fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); while let Poll::Ready(Some(interest)) = this.rx.poll_recv(cx) { match interest { Interest::WorkerIndex(idx) => { this.avail.set_available(idx, true); } Interest::Worker(handle) => { this.avail.set_available(handle.idx(), true); this.handles.push(handle); } Interest::Pause => this.state = AcceptState::Pause, Interest::Resume => this.state = AcceptState::Accept, Interest::Stop => this.state = AcceptState::Stop, } } match this.state { AcceptState::Accept => this.poll_accept(cx), AcceptState::Pause => Poll::Pending, AcceptState::Stop => Poll::Ready(()), } } } /// This function defines errors that are per-connection. Which basically /// means that if we get this error from `accept()` system call it means /// next connection might be ready to be accepted. /// /// All other errors will incur a timeout before next `accept()` is performed. /// The timeout is useful to handle resource exhaustion errors like ENFILE /// and EMFILE. Otherwise, could enter into tight loop. fn connection_error(e: &io::Error) -> bool { e.kind() == io::ErrorKind::ConnectionRefused || e.kind() == io::ErrorKind::ConnectionAborted || e.kind() == io::ErrorKind::ConnectionReset } /// Array of u128 with every bit as marker for a worker handle's availability. struct Availability([u128; 4]); impl Default for Availability { fn default() -> Self { Self([0; 4]) } } impl Availability { /// Check if any worker handle is available #[inline(always)] fn available(&self) -> bool { self.0.iter().any(|a| *a != 0) } /// Check if worker handle is available by index #[inline(always)] fn get_available(&self, idx: usize) -> bool { let (offset, idx) = Self::offset(idx); self.0[offset] & (1 << idx as u128) != 0 } /// Set worker handle available state by index. fn set_available(&mut self, idx: usize, avail: bool) { let (offset, idx) = Self::offset(idx); let off = 1 << idx as u128; if avail { self.0[offset] |= off; } else { self.0[offset] &= !off } } /// Set all worker handle to available state. /// This would result in a re-check on all workers' availability. fn set_available_all(&mut self, handles: &[WorkerHandleAccept]) { handles.iter().for_each(|handle| { self.set_available(handle.idx(), true); }) } /// Get offset and adjusted index of given worker handle index. fn offset(idx: usize) -> (usize, usize) { if idx < 128 { (0, idx) } else if idx < 128 * 2 { (1, idx - 128) } else if idx < 128 * 3 { (2, idx - 128 * 2) } else if idx < 128 * 4 { (3, idx - 128 * 3) } else { panic!("Max WorkerHandle count is 512") } } } #[cfg(test)] mod test { use super::Availability; fn single(aval: &mut Availability, idx: usize) { aval.set_available(idx, true); assert!(aval.available()); aval.set_available(idx, true); aval.set_available(idx, false); assert!(!aval.available()); aval.set_available(idx, false); assert!(!aval.available()); } fn multi(aval: &mut Availability, mut idx: Vec) { idx.iter().for_each(|idx| aval.set_available(*idx, true)); assert!(aval.available()); while let Some(idx) = idx.pop() { assert!(aval.available()); aval.set_available(idx, false); } assert!(!aval.available()); } #[test] fn availability() { let mut aval = Availability::default(); single(&mut aval, 1); single(&mut aval, 128); single(&mut aval, 256); single(&mut aval, 511); let idx = (0..511).filter(|i| i % 3 == 0 && i % 5 == 0).collect(); multi(&mut aval, idx); multi(&mut aval, (0..511).collect()) } #[test] #[should_panic] fn overflow() { let mut aval = Availability::default(); single(&mut aval, 512); } #[test] fn pin_point() { let mut aval = Availability::default(); aval.set_available(438, true); aval.set_available(479, true); assert_eq!(aval.0[3], 1 << (438 - 384) | 1 << (479 - 384)); } }