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}