klukai_types/tripwire/
mod.rs

1//! Provides [Tripwire], which completes when the program is requested to shutdown.
2//! It allows for graceful shutdown of various asynchronous tasks.
3
4#![warn(missing_docs)]
5#![deny(clippy::await_holding_lock)]
6
7mod preempt;
8mod signalstream;
9
10pub use preempt::{Outcome, PreemptibleFuture, PreemptibleFutureExt, TimeoutFutureExt};
11
12use futures::stream::Stream;
13use futures_util::stream::{Select, select};
14use std::{
15    future::Future,
16    pin::Pin,
17    task::{Context, Poll},
18};
19use tokio::{
20    signal::unix::{SignalKind, signal},
21    sync::{mpsc, watch},
22};
23use tokio_stream::wrappers::{ReceiverStream, WatchStream};
24use tracing::{debug, warn};
25
26use crate::tripwire::signalstream::SignalStream;
27
28/// A `Future` that completes once the program is requested to shutdown. This
29/// gives a chance for other `Futures` to do additional cleanup before being
30/// `drop`ped.
31#[derive(Debug)]
32pub struct Tripwire {
33    subscription: WatchStream<TripwireState>,
34    subscription_rx: watch::Receiver<TripwireState>,
35}
36
37impl Tripwire {
38    /// Create a new [Tripwire]. The returned worker will trip the tripwire when
39    /// dropped, and must be spawned on a runtime to listen for SIGTERM/SIGINT
40    pub fn new<S>(stream: S) -> (Self, TripwireWorker<S>)
41    where
42        S: Stream<Item = ()>,
43    {
44        let (sender, receiver) = watch::channel(TripwireState::Running);
45        let tripwire = Tripwire {
46            subscription: WatchStream::new(receiver.clone()),
47            subscription_rx: receiver,
48        };
49        let worker = TripwireWorker {
50            subscription: sender,
51            stream,
52        };
53        (tripwire, worker)
54    }
55
56    /// Create a new [Tripwire] listens for a `()` being sent to the returned sender.
57    /// Graceful shutdown can then be awaited through the [TripwireWorker]
58    pub fn new_simple() -> (Self, TripwireWorker<ReceiverStream<()>>, mpsc::Sender<()>) {
59        let (tx, rx) = mpsc::channel::<()>(1);
60        let (tw, w) = Self::new(ReceiverStream::new(rx));
61        (tw, w, tx)
62    }
63
64    /// Listen for SIGTERM and SIGINT
65    pub fn new_signals() -> (Self, TripwireWorker<Select<SignalStream, SignalStream>>) {
66        let sigterms = SignalStream::new(signal(SignalKind::terminate()).unwrap());
67        let sigints = SignalStream::new(signal(SignalKind::interrupt()).unwrap());
68        Self::new(select(sigterms, sigints))
69    }
70
71    /// Returns an Arc of the current [TripwireState]
72    pub fn state(&self) -> TripwireState {
73        *self.subscription_rx.borrow()
74    }
75
76    /// Returns true if we're shutting down
77    pub fn is_shutting_down(&self) -> bool {
78        matches!(self.state(), TripwireState::ShuttingDown)
79    }
80
81    /// Returns a [Future] that completes when this wire is tripped
82    pub fn tripwired<F>(self, inner: F) -> Tripwired<F> {
83        Tripwired {
84            tripwire: self,
85            inner,
86        }
87    }
88}
89
90impl Clone for Tripwire {
91    fn clone(&self) -> Self {
92        Tripwire {
93            subscription: WatchStream::new(self.subscription_rx.clone()),
94            subscription_rx: self.subscription_rx.clone(),
95        }
96    }
97}
98
99impl Future for Tripwire {
100    type Output = ();
101
102    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
103        loop {
104            match futures::ready!(Pin::new(&mut self.subscription).poll_next(cx)) {
105                Some(TripwireState::Running) => {}
106                Some(TripwireState::ShuttingDown) => return Poll::Ready(()),
107                None => return Poll::Ready(()),
108            }
109        }
110    }
111}
112
113/// A [Future] that completes when the program is requested to shutdown.
114pub struct Tripwired<F> {
115    tripwire: Tripwire,
116    inner: F,
117}
118
119impl<F> Future for Tripwired<F>
120where
121    F: Future<Output = ()> + Unpin,
122{
123    type Output = F::Output;
124
125    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
126        if let Poll::Ready(()) = Pin::new(&mut self.tripwire).poll(cx) {
127            return Poll::Ready(());
128        }
129
130        Pin::new(&mut self.inner).poll(cx)
131    }
132}
133
134impl<S> Stream for Tripwired<S>
135where
136    S: Stream + Unpin,
137{
138    type Item = S::Item;
139
140    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
141        if let Poll::Ready(()) = Pin::new(&mut self.tripwire).poll(cx) {
142            return Poll::Ready(None);
143        }
144
145        Pin::new(&mut self.inner).poll_next(cx)
146    }
147}
148
149/// Whether this trip wire has been activated
150#[derive(Copy, Clone, Debug)]
151#[allow(missing_docs)]
152pub enum TripwireState {
153    Running,
154    ShuttingDown,
155}
156
157/// Trips the [Tripwire] when receiving anything from a stream
158/// (used for signals), or when dropping.
159pub struct TripwireWorker<S> {
160    subscription: watch::Sender<TripwireState>,
161    stream: S,
162}
163
164impl<S> Future for TripwireWorker<S>
165where
166    S: Stream + Unpin,
167{
168    type Output = ();
169
170    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
171        match futures::ready!(Pin::new(&mut self.stream).poll_next(cx)) {
172            Some(_) => "either SIGTERM or SIGINT",
173            None => return Poll::Ready(()),
174        };
175
176        println!(); // cleaner logs! (new line after ^C)
177        debug!("TripwireWorker tripped");
178
179        if let Err(error) = self.subscription.send(TripwireState::ShuttingDown) {
180            warn!("all subscription handles have been cancelled?: {:?}", error);
181        }
182        debug!("tripwire worker's subscription has been updated");
183        Poll::Ready(())
184    }
185}
klukai_types/tripwire/mod.rs

klukai_types/tripwire/
mod.rs
