Skip to main content

ntex_server/
server.rs

1#![allow(clippy::missing_panics_doc)]
2use std::sync::{Arc, atomic::AtomicBool, atomic::Ordering};
3use std::task::{Context, Poll, ready};
4use std::{future::Future, io, pin::Pin};
5
6use async_channel::Sender;
7use ntex_rt::signals::Signal;
8
9use crate::manager::ServerCommand;
10
11#[derive(Debug)]
12pub(crate) struct ServerShared {
13    pub(crate) paused: AtomicBool,
14}
15
16/// Controller and completion future for a running server.
17///
18/// Clones can pause, resume, stop, or submit items to the server. Awaiting a
19/// `Server` resolves when the server has stopped.
20#[derive(Debug)]
21pub struct Server<T> {
22    shared: Arc<ServerShared>,
23    cmd: Sender<ServerCommand<T>>,
24    stop: Option<oneshot::AsyncReceiver<()>>,
25}
26
27impl<T> Server<T> {
28    pub(crate) fn new(cmd: Sender<ServerCommand<T>>, shared: Arc<ServerShared>) -> Self {
29        Server {
30            cmd,
31            shared,
32            stop: None,
33        }
34    }
35
36    /// Creates a network server builder with no application configuration.
37    pub fn builder() -> crate::net::ServerBuilder {
38        crate::net::ServerBuilder::default()
39    }
40
41    pub(crate) fn signal(&self, sig: Signal) {
42        let _ = self.cmd.try_send(ServerCommand::Signal(sig));
43    }
44
45    /// Submits an item to the worker pool.
46    ///
47    /// Returns the item unchanged if the server is paused or cannot accept it.
48    pub fn process(&self, item: T) -> Result<(), T> {
49        if self.shared.paused.load(Ordering::Acquire) {
50            Err(item)
51        } else if let Err(e) = self.cmd.try_send(ServerCommand::Item(item)) {
52            if let ServerCommand::Item(item) = e.into_inner() {
53                Err(item)
54            } else {
55                panic!()
56            }
57        } else {
58            Ok(())
59        }
60    }
61
62    /// Pauses processing new items.
63    ///
64    /// For network servers, pending connections may be dropped. Existing
65    /// connections remain active.
66    pub fn pause(&self) -> impl Future<Output = ()> + use<T> {
67        let (tx, rx) = oneshot::channel();
68        let _ = self.cmd.try_send(ServerCommand::Pause(tx));
69        async move {
70            let _ = rx.await;
71        }
72    }
73
74    /// Resumes processing new items.
75    pub fn resume(&self) -> impl Future<Output = ()> + use<T> {
76        let (tx, rx) = oneshot::channel();
77        let _ = self.cmd.try_send(ServerCommand::Resume(tx));
78        async move {
79            let _ = rx.await;
80        }
81    }
82
83    /// Stops processing new items and shuts down all workers.
84    ///
85    /// If `graceful` is `true`, workers are given time to finish active work
86    /// before they are stopped.
87    pub fn stop(&self, graceful: bool) -> impl Future<Output = ()> + use<T> {
88        let (tx, rx) = oneshot::channel();
89        let _ = self.cmd.try_send(ServerCommand::Stop {
90            graceful,
91            completion: Some(tx),
92        });
93        async move {
94            let _ = rx.await;
95        }
96    }
97}
98
99impl<T> Clone for Server<T> {
100    fn clone(&self) -> Self {
101        Self {
102            cmd: self.cmd.clone(),
103            shared: self.shared.clone(),
104            stop: None,
105        }
106    }
107}
108
109impl<T> Future for Server<T> {
110    type Output = io::Result<()>;
111
112    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
113        let this = self.get_mut();
114
115        if this.stop.is_none() {
116            let (tx, rx) = oneshot::async_channel();
117            if this.cmd.try_send(ServerCommand::NotifyStopped(tx)).is_err() {
118                return Poll::Ready(Ok(()));
119            }
120            this.stop = Some(rx);
121        }
122
123        let _ = ready!(Pin::new(this.stop.as_mut().unwrap()).poll(cx));
124
125        Poll::Ready(Ok(()))
126    }
127}