Skip to main content

tako_rs_server/builder/
handle.rs

1use std::future::Future;
2use std::time::Duration;
3
4/// Background-task handle returned by every `spawn_*` method.
5///
6/// Drop semantics: dropping the handle does **not** stop the server. Call
7/// [`ServerHandle::shutdown`] (or [`ServerHandle::trigger`] + `.join().await`)
8/// so the drain logic in the underlying `serve_*_with_shutdown` runs.
9///
10/// Runtime-agnostic — the `done` signal is fired by an `async` wrapper around
11/// the underlying `serve_*` future, so the same `ServerHandle` works whether
12/// the spawned task lives on the tokio runtime or the compio runtime.
13pub struct ServerHandle {
14  pub(crate) shutdown: tokio_util::sync::CancellationToken,
15  pub(crate) done: tokio_util::sync::CancellationToken,
16  pub(crate) drain_timeout: Duration,
17}
18
19impl ServerHandle {
20  /// Trigger graceful shutdown without awaiting completion.
21  pub fn trigger(&self) {
22    self.shutdown.cancel();
23  }
24
25  /// Await the spawned task's completion (without triggering shutdown).
26  ///
27  /// Returns when the underlying `serve_*` future resolves — typically
28  /// because [`ServerHandle::trigger`] / [`ServerHandle::shutdown`] was called
29  /// or because the listener errored fatally.
30  pub async fn join(&self) {
31    self.done.cancelled().await;
32  }
33
34  /// Trigger graceful shutdown and await the drain.
35  ///
36  /// The `_timeout` argument is kept for API symmetry with the original
37  /// builder; the actual drain bound is the `drain_timeout` on the
38  /// [`ServerConfig`](crate::ServerConfig) that was handed to the builder, enforced inside
39  /// `serve_*_with_shutdown`.
40  pub async fn shutdown(self, _timeout: Duration) {
41    self.shutdown.cancel();
42    self.done.cancelled().await;
43  }
44
45  /// Returns the drain timeout the underlying `serve_*` will honor.
46  #[inline]
47  pub fn drain_timeout(&self) -> Duration {
48    self.drain_timeout
49  }
50}
51
52impl std::fmt::Debug for ServerHandle {
53  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54    f.debug_struct("ServerHandle")
55      .field("drain_timeout", &self.drain_timeout)
56      .finish_non_exhaustive()
57  }
58}
59
60/// Convenience: await `signal_a` *or* `signal_b`, whichever fires first.
61pub async fn either<A, B>(a: A, b: B)
62where
63  A: Future<Output = ()>,
64  B: Future<Output = ()>,
65{
66  use futures_util::future::Either;
67  let a = std::pin::pin!(a);
68  let b = std::pin::pin!(b);
69  match futures_util::future::select(a, b).await {
70    Either::Left(_) | Either::Right(_) => {}
71  }
72}