use std::sync::atomic::{AtomicBool, Ordering}; use std::{cell::Cell, cell::RefCell, collections::VecDeque, rc::Rc, sync::Arc}; use async_channel::{unbounded, Receiver, Sender}; use ntex_rt::System; use ntex_util::{future::join_all, time::sleep, time::Millis}; use crate::server::ServerShared; use crate::signals::Signal; use crate::{Server, ServerConfiguration, Worker, WorkerId, WorkerPool, WorkerStatus}; const STOP_DELAY: Millis = Millis(500); const RESTART_DELAY: Millis = Millis(250); #[derive(Clone)] pub(crate) struct ServerManager(Rc>); #[derive(Debug)] pub(crate) enum ServerCommand { Item(T), Pause(oneshot::Sender<()>), Resume(oneshot::Sender<()>), Signal(Signal), Stop { graceful: bool, completion: Option>, }, NotifyStopped(oneshot::Sender<()>), Worker(Update), } #[derive(Debug)] pub(crate) enum Update { Available(Worker), Unavailable(Worker), } struct Inner { id: Cell, factory: F, cfg: WorkerPool, shared: Arc, stopping: Cell, stop_notify: RefCell>>, cmd: Sender>, } impl ServerManager { pub(crate) fn start(cfg: WorkerPool, factory: F) -> Server { log::info!("Starting {} workers", cfg.num); let (tx, rx) = unbounded(); let no_signals = cfg.no_signals; let shared = Arc::new(ServerShared { paused: AtomicBool::new(true), }); let mgr = ServerManager(Rc::new(Inner { cfg, factory, id: Cell::new(WorkerId::default()), shared: shared.clone(), stopping: Cell::new(false), stop_notify: RefCell::new(Vec::new()), cmd: tx.clone(), })); // handle cmd let _ = ntex_rt::spawn(handle_cmd(mgr.clone(), rx)); // start workers for _ in 0..mgr.0.cfg.num { start_worker(mgr.clone()); } let srv = Server::new(tx, shared); // handle signals if !no_signals { crate::signals::start(srv.clone()); } srv } pub(crate) fn factory(&self) -> F { self.0.factory.clone() } pub(crate) fn next_id(&self) -> WorkerId { let mut id = self.0.id.get(); let next_id = id.next(); self.0.id.set(id); next_id } pub(crate) fn pause(&self) { self.0.shared.paused.store(true, Ordering::Release); self.0.factory.paused(); } pub(crate) fn resume(&self) { self.0.shared.paused.store(false, Ordering::Release); self.0.factory.resumed(); } fn available(&self, wrk: Worker) { let _ = self .0 .cmd .try_send(ServerCommand::Worker(Update::Available(wrk))); } fn unavailable(&self, wrk: Worker) { let _ = self .0 .cmd .try_send(ServerCommand::Worker(Update::Unavailable(wrk))); } fn add_stop_notify(&self, tx: oneshot::Sender<()>) { self.0.stop_notify.borrow_mut().push(tx); } fn stopping(&self) -> bool { self.0.stopping.get() } } fn start_worker(mgr: ServerManager) { let _ = ntex_rt::spawn(async move { let id = mgr.next_id(); let mut wrk = Worker::start(id, mgr.factory()); loop { match wrk.status() { WorkerStatus::Available => mgr.available(wrk.clone()), WorkerStatus::Unavailable => mgr.unavailable(wrk.clone()), WorkerStatus::Failed => { mgr.unavailable(wrk); sleep(RESTART_DELAY).await; if !mgr.stopping() { wrk = Worker::start(id, mgr.factory()); } else { return; } } } wrk.wait_for_status().await; } }); } struct HandleCmdState { next: usize, backlog: VecDeque, workers: Vec>, mgr: ServerManager, } impl HandleCmdState { fn new(mgr: ServerManager) -> Self { Self { next: 0, backlog: VecDeque::new(), workers: Vec::with_capacity(mgr.0.cfg.num), mgr, } } fn process(&mut self, mut item: F::Item) { loop { if !self.workers.is_empty() { if self.next > self.workers.len() { self.next = self.workers.len() - 1; } match self.workers[self.next].send(item) { Ok(()) => { self.next = (self.next + 1) % self.workers.len(); break; } Err(i) => { if !self.mgr.0.stopping.get() { log::trace!("Worker failed while processing item"); } item = i; self.workers.remove(self.next); } } } else { if !self.mgr.0.stopping.get() { log::error!("No workers"); self.backlog.push_back(item); self.mgr.pause(); } break; } } } fn update_workers(&mut self, upd: Update) { match upd { Update::Available(worker) => { self.workers.push(worker); if self.workers.len() == 1 { self.mgr.resume(); } else { self.workers.sort(); } } Update::Unavailable(worker) => { if let Ok(idx) = self.workers.binary_search(&worker) { self.workers.remove(idx); } if self.workers.is_empty() { self.mgr.pause(); } } } // handle backlog if !self.backlog.is_empty() && !self.workers.is_empty() { while let Some(item) = self.backlog.pop_front() { // handle worker failure if let Err(item) = self.workers[0].send(item) { self.backlog.push_back(item); self.workers.remove(0); break; } } } } async fn stop(&mut self, graceful: bool, completion: Option>) { self.mgr.0.stopping.set(true); // stop server self.mgr.0.factory.stop().await; // stop workers if !self.workers.is_empty() { let timeout = self.mgr.0.cfg.shutdown_timeout; if graceful && !timeout.is_zero() { let futs: Vec<_> = self .workers .iter() .map(|worker| worker.stop(timeout)) .collect(); let _ = join_all(futs).await; } else { self.workers.iter().for_each(|worker| { let _ = worker.stop(Millis::ZERO); }); } } // notify sender if let Some(tx) = completion { let _ = tx.send(()); } let notify = std::mem::take(&mut *self.mgr.0.stop_notify.borrow_mut()); for tx in notify { let _ = tx.send(()); } // stop system if server was spawned if self.mgr.0.cfg.stop_runtime { sleep(STOP_DELAY).await; System::current().stop(); } } } struct DropHandle(ServerManager); impl Drop for DropHandle { fn drop(&mut self) { self.0 .0.stopping.set(true); self.0 .0.factory.terminate(); } } async fn handle_cmd( mgr: ServerManager, rx: Receiver>, ) { let _drop_hnd = DropHandle(mgr.clone()); let mut state = HandleCmdState::new(mgr); loop { let item = if let Ok(item) = rx.recv().await { item } else { return; }; match item { ServerCommand::Item(item) => state.process(item), ServerCommand::Worker(upd) => state.update_workers(upd), ServerCommand::Pause(tx) => { state.mgr.pause(); let _ = tx.send(()); } ServerCommand::Resume(tx) => { state.mgr.resume(); let _ = tx.send(()); } ServerCommand::NotifyStopped(tx) => state.mgr.add_stop_notify(tx), ServerCommand::Stop { graceful, completion, } => { state.stop(graceful, completion).await; return; } ServerCommand::Signal(sig) => { // Signals support // Handle `SIGINT`, `SIGTERM`, `SIGQUIT` signals and stop ntex system match sig { Signal::Int => { log::info!("SIGINT received, exiting"); state.stop(false, None).await; return; } Signal::Term => { log::info!("SIGTERM received, stopping"); state.stop(true, None).await; return; } Signal::Quit => { log::info!("SIGQUIT received, exiting"); state.stop(false, None).await; return; } _ => (), } } } } }