tako_rs_server_pt/
shutdown.rs

1use std::io;
2
3/// Worker bind results, so the parent can fail when every worker failed to bind instead of
4/// waiting for a signal with no listener up.
5#[derive(Default)]
6struct BindStatus {
7  succeeded: std::sync::atomic::AtomicUsize,
8  failed: std::sync::atomic::AtomicUsize,
9  first_err: std::sync::Mutex<Option<io::Error>>,
10  notify: tokio::sync::Notify,
11}
12
13/// Stops the workers of [`spawn_per_thread`](crate::spawn_per_thread); the signal is sticky, so
14/// a worker that starts waiting after [`trigger`](Self::trigger) still sees it.
15#[derive(Clone)]
16pub struct PerThreadShutdown {
17  pub(crate) inner: tokio_util::sync::CancellationToken,
18  bind_status: std::sync::Arc<BindStatus>,
19}
20
21impl PerThreadShutdown {
22  pub(crate) fn new() -> Self {
23    Self {
24      inner: tokio_util::sync::CancellationToken::new(),
25      bind_status: std::sync::Arc::default(),
26    }
27  }
28
29  /// Tells every worker to stop accepting and drain; later calls do nothing.
30  pub fn trigger(&self) {
31    self.inner.cancel();
32  }
33
34  /// Resolves once shutdown has been triggered.
35  pub async fn notified(&self) {
36    self.inner.cancelled().await;
37  }
38
39  pub(crate) fn report_bind_success(&self) {
40    self
41      .bind_status
42      .succeeded
43      .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
44    self.bind_status.notify.notify_waiters();
45  }
46
47  /// Records a worker's bind failure; only the first error is kept for the caller.
48  pub(crate) fn report_bind_failure(&self, err: io::Error) {
49    {
50      let mut guard = self.bind_status.first_err.lock().unwrap();
51      if guard.is_none() {
52        *guard = Some(err);
53      }
54    }
55    self
56      .bind_status
57      .failed
58      .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
59    self.bind_status.notify.notify_waiters();
60  }
61
62  /// Waits until one worker has bound, or returns the first bind error once all `total` failed.
63  pub async fn wait_for_bind_outcome(&self, total: usize) -> io::Result<()> {
64    use std::sync::atomic::Ordering;
65
66    loop {
67      // Arm the notified future BEFORE reading state so a wake fired
68      // between the load and the await is not lost.
69      let notified = self.bind_status.notify.notified();
70      tokio::pin!(notified);
71      notified.as_mut().enable();
72
73      let succ = self.bind_status.succeeded.load(Ordering::SeqCst);
74      let fail = self.bind_status.failed.load(Ordering::SeqCst);
75
76      if succ > 0 {
77        return Ok(());
78      }
79      if succ + fail >= total {
80        let err = self
81          .bind_status
82          .first_err
83          .lock()
84          .unwrap()
85          .take()
86          .unwrap_or_else(|| {
87            io::Error::other(format!("all {total} per-thread workers failed to bind"))
88          });
89        return Err(err);
90      }
91
92      notified.await;
93    }
94  }
95}
tako_rs_server_pt/shutdown.rs

tako_rs_server_pt/
shutdown.rs
