tako_rs_server/builder/
handle.rs

1use std::net::SocketAddr;
2use std::sync::Arc;
3use std::sync::Mutex;
4use std::time::Duration;
5
6/// A clonable server startup or background-task failure.
7#[derive(Debug, Clone)]
8pub struct ServerError(pub(crate) Arc<dyn std::error::Error + Send + Sync>);
9
10impl std::fmt::Display for ServerError {
11  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
12    self.0.fmt(f)
13  }
14}
15
16impl std::error::Error for ServerError {
17  fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
18    Some(self.0.as_ref())
19  }
20}
21
22impl From<tako_rs_core::types::BoxError> for ServerError {
23  fn from(error: tako_rs_core::types::BoxError) -> Self {
24    Self(Arc::from(error))
25  }
26}
27
28/// Handle to a spawned server on either runtime. Dropping it does not stop the server; call
29/// [`ServerHandle::shutdown`], or [`ServerHandle::trigger`] and then `join`, to drain it.
30pub struct ServerHandle {
31  pub(crate) shutdown: tokio_util::sync::CancellationToken,
32  pub(crate) done: tokio_util::sync::CancellationToken,
33  pub(crate) abort: tokio_util::sync::CancellationToken,
34  pub(crate) drain_timeout: Duration,
35  pub(crate) result: Arc<Mutex<Option<Result<(), ServerError>>>>,
36  pub(crate) local_addr: Option<SocketAddr>,
37}
38
39impl ServerHandle {
40  /// Stops accepting and starts draining without waiting for it.
41  pub fn trigger(&self) {
42    self.shutdown.cancel();
43  }
44
45  /// Waits until the server stops, after a shutdown or a fatal listener error, without
46  /// triggering shutdown.
47  pub async fn join(&self) {
48    self.done.cancelled().await;
49  }
50
51  /// Awaits completion and returns startup, I/O, panic, or forced-abort errors.
52  pub async fn result(&self) -> Result<(), ServerError> {
53    self.join().await;
54    self
55      .result
56      .lock()
57      .unwrap_or_else(std::sync::PoisonError::into_inner)
58      .as_ref()
59      .expect("completed task has a result")
60      .clone()
61  }
62
63  /// The bound address, including the port chosen for a port-zero bind; `None` for Unix and vsock.
64  pub fn local_addr(&self) -> Option<SocketAddr> {
65    self.local_addr
66  }
67
68  /// Triggers shutdown and waits at most the smaller of `timeout` and the
69  /// configured drain timeout before aborting remaining server work.
70  pub async fn shutdown(self, timeout: Duration) {
71    self.trigger();
72    let deadline = async {
73      #[cfg(not(feature = "compio"))]
74      tokio::time::sleep(timeout.min(self.drain_timeout)).await;
75      #[cfg(feature = "compio")]
76      compio::time::sleep(timeout.min(self.drain_timeout)).await;
77    };
78    let finished = std::pin::pin!(self.done.cancelled());
79    let deadline = std::pin::pin!(deadline);
80    if let futures_util::future::Either::Right(_) =
81      futures_util::future::select(finished, deadline).await
82    {
83      self.abort.cancel();
84      self.done.cancelled().await;
85    }
86  }
87
88  /// Waits for Ctrl+C or SIGTERM, then drains the server.
89  pub async fn shutdown_on_signal(self) -> std::io::Result<()> {
90    tako_rs_core::server_support::shutdown_signal().await?;
91    let timeout = self.drain_timeout;
92    self.shutdown(timeout).await;
93    Ok(())
94  }
95
96  /// `ServerConfig::drain_timeout`, which also caps [`ServerHandle::shutdown`].
97  #[inline]
98  pub fn drain_timeout(&self) -> Duration {
99    self.drain_timeout
100  }
101}
102
103impl std::fmt::Debug for ServerHandle {
104  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105    f.debug_struct("ServerHandle")
106      .field("drain_timeout", &self.drain_timeout)
107      .finish_non_exhaustive()
108  }
109}
tako_rs_server/builder/handle.rs

tako_rs_server/builder/
handle.rs
