pistonite_cu/cli/ctrlc.rs
1use std::sync::atomic::{AtomicU8, Ordering};
2use std::sync::{Arc, LazyLock, Mutex};
3
4/// New handlers that will be added by the ctrlc thread
5/// once ctrlc is pressed
6static CTRLC_HANDLERS_NEW: Mutex<Vec<Box<dyn FnMut() + Send>>> = Mutex::new(Vec::new());
7/// Stack of CtrlC frames to signal
8static CTRLC_SIGNAL_STACK: Mutex<Vec<CtrlcFrame>> = Mutex::new(Vec::new());
9/// Thread safe init lock
10static INIT_ONCE: LazyLock<Result<(), String>> = LazyLock::new(|| {
11 let mut handlers = vec![];
12 let set_result = ctrlc::try_set_handler(move || {
13 // populate new handlers
14 if let Ok(mut new_handlers) = CTRLC_HANDLERS_NEW.lock() {
15 handlers.extend(new_handlers.drain(..));
16 }
17 // signal the stack
18 let mut signalled = false;
19 if let Ok(stack) = CTRLC_SIGNAL_STACK.lock() {
20 if let Some(frame) = stack.last() {
21 signalled = true;
22 frame.signal.signal();
23 if let Some(f) = &frame.on_signal {
24 f(frame.signal.clone())
25 }
26 }
27 }
28
29 // note we are not holding any lock when invoking user-defined handlers
30
31 // if user did not set any global handler or action frames, then we terminate
32 if !signalled && handlers.is_empty() {
33 std::process::exit(1);
34 }
35 for handler in handlers.iter_mut().rev() {
36 handler();
37 }
38 });
39 match set_result {
40 Err(ctrlc::Error::MultipleHandlers) => {
41 Err("failed to set ctrl-c handler: a handler is already set using the `ctrlc` crate. please set with cu::cli instead (see documentation for more information)".to_string())
42 }
43 Err(other_error) => {
44 Err(format!("failed to set ctrl-c handler: {other_error}"))
45 }
46 Ok(_) => Ok(())
47 }
48});
49
50/// Add a global handler to handle Ctrl-C signals
51///
52/// See [Handling Ctrl-C](fn@crate::cli::ctrlc_frame).
53///
54/// Note that this is only available with feature `cli` - since
55/// you should not be adding a global handler from a library.
56#[cfg(feature = "cli")]
57pub fn add_global_ctrlc_handler<F: FnMut() + Send + 'static>(handler: F) -> cu::Result<()> {
58 {
59 let Ok(mut handlers) = CTRLC_HANDLERS_NEW.lock() else {
60 cu::bail!("global ctrl-c handler vector is poisoned");
61 };
62 handlers.push(Box::new(handler));
63 }
64 if let Err(e) = &*INIT_ONCE {
65 cu::bail!("{e}");
66 }
67 Ok(())
68}
69
70/// # Handling Ctrl-C
71///
72/// The [`ctrlc`](https://docs.rs/ctrlc) crate provides a low-level cross-platform
73/// way to set up a handler for `Ctrl-C`. `cu` builds wrappers around it so it can
74/// be intergrated with other internals, such as prompts.
75///
76/// # Action Frames
77/// Most of the time, custom Ctrl-C behavior is for executing some long running
78/// tasks and give the user the ability to abort it.
79/// The [`execute`](CtrlcBuilder::execute) function pushes a new frame
80/// to an internal stack, then executes the task synchronously on the same thread.
81/// **The task itself is responsible for periodically checking the received
82/// signal object if it has been aborted (by calling `ctrlc.check?`).
83///
84/// The frame will only be removed from the stack when the task returns.
85/// Note that different threads can spawn Ctrl-C frames, and those frame
86/// may not finish in order. When user hits `Ctrl-C`, only the most recently
87/// added frame (the top-most of the stack) will be notified.
88/// It will continue to notify the frame until the task ends and removes
89/// the frame from the stask.
90///
91/// # Async Behavior
92/// The async version of the API, [`co_execute`](CtrlcBuilder::co_execute),
93/// is exactly the same, other than it takes an async closure instead.
94/// Even though in the async case, we can let the runtime drop the future
95/// to abort it without explicit checks, the explicit checks make it clearer
96/// and easier to reason about program states when aborting.
97///
98/// If you prefer automatically cancellation, consider this pattern:
99#[cfg_attr(not(feature = "coroutine"), doc = "```rust,ignore")]
100#[cfg_attr(feature = "coroutine", doc = "```rust,no_run")]
101/// # use pistonite_cu as cu;
102/// # async fn main_() -> cu::Result<()> {
103/// let (send, mut recv) = tokio::sync::mpsc::unbounded_channel();
104/// let result = cu::cli::ctrlc_frame()
105/// .on_signal(move |_| { let _ = send.send(()); })
106/// .co_execute(async move |ctrlc| {
107/// let waiter = async move {
108/// loop {
109/// if recv.recv().await.is_none() {
110/// return;
111/// }
112/// if ctrlc.should_abort() {
113/// return;
114/// }
115/// }
116/// };
117/// tokio::select! {
118/// result = my_long_running_task() => {
119/// return result;
120/// }
121/// _ = waiter => {
122/// // does not matter what the value is -
123/// // co_execute will ensure None is returned
124/// // when aborted
125/// return Ok(0);
126/// }
127/// }
128/// }).await?;
129/// match result {
130/// Some(x) => cu::info!("valud is: {x}"),
131/// None => cu::error!("aborted!"),
132/// }
133/// # Ok(()) }
134///
135/// async fn my_long_running_task() -> cu::Result<i32> {
136/// // your task here can get aborted without checking the signal
137/// // which can be dangerous and have unintended effects!
138/// Ok(42)
139/// }
140/// ```
141///
142/// # Return Value
143/// The return value of the inner `cu::Result<T>` will be transformed into
144/// `cu::Result<Option<T>>`, where:
145/// - `Ok(Some(value))`: the action was not aborted, and produced a value.
146/// - `Ok(None)`: the action was aborted.
147/// - The definition of "aborted" can be customized by setting the [`abort_threshold()`](CtrlcBuilder::abort_threshold)
148/// on the builder
149/// - `Err(e)`: either the inner task returned error, or some other error happened in the framework
150/// (for example, when joining the task thread/future)
151///
152/// ```rust,no_run
153/// # use pistonite_cu as cu;
154/// use std::thread;
155/// use std::time::Duration;
156///
157/// match cu::cli::ctrlc_frame().execute(|ctrlc| {
158/// for _ in 0..10 {
159/// cu::print!("please press Ctrl-C");
160/// thread::sleep(Duration::from_secs(2));
161/// ctrlc.check()?; // returns Err if signaled
162/// }
163/// cu::Ok(42)
164/// }) {
165/// Ok(None) => cu::info!("was aborted!"),
166/// Ok(Some(n)) => cu::info!("was finished: {n}"),
167/// Err(e) => cu::error!("error: {e:?}"),
168/// }
169/// ```
170///
171/// # Fallback
172/// If the `Ctrl-C` framework failed to initialize, this may fallback to simply running
173/// the task without the ability to receive signals.
174///
175/// # Global Handlers
176/// Action frames should be used whenever possible.
177/// If you do need a global handler for `Ctrl-C`, use
178/// [`cu::cli::add_global_ctrlc_handler`]. This is because the `ctrlc` crate
179/// only allows one global handler.
180///
181/// Global handlers are called after the action frame is notified (if any),
182/// and called in reverse order of registration.
183/// The underlying handler is lazily set up whenever a global handler
184/// or action frame is added. If there are no longer any action frames
185/// and there are no global handlers, the underlying handler
186/// will call `std::process::exit(1)` to terminate.
187#[inline(always)]
188pub fn ctrlc_frame() -> CtrlcBuilder {
189 CtrlcBuilder::default()
190}
191
192/// Builder for a new frame for handling `Ctrl-C` signals.
193/// The canonical factory function for this is `cu::cli::ctrlc_frame`.
194/// See [Handling Ctrl-C](ctrlc_frame).
195pub struct CtrlcBuilder {
196 abort_threshold: u8,
197 on_signal: Option<OnSignalFn>,
198}
199impl Default for CtrlcBuilder {
200 #[inline(always)]
201 fn default() -> Self {
202 Self {
203 abort_threshold: 1,
204 on_signal: None,
205 }
206 }
207}
208impl CtrlcBuilder {
209 /// Set the number of `Ctrl-C` signals required for the task
210 /// to be considered aborted (where the return will be `Ok(None`).
211 ///
212 /// Default is 1
213 #[inline(always)]
214 pub fn abort_threshold(mut self, threshold: u8) -> Self {
215 self.abort_threshold = threshold;
216 self
217 }
218
219 /// Set a function to be called when `Ctrl-C` signal is received.
220 ///
221 /// The function will be executed on the `Ctrl-C` signal handling thread,
222 /// not the thread that runs the task, and is
223 #[inline(always)]
224 pub fn on_signal<F: Fn(CtrlcSignal) + Send + 'static>(mut self, f: F) -> Self {
225 self.on_signal = Some(Box::new(f));
226 self
227 }
228
229 /// Execute the task
230 pub fn execute<T, F>(self, f: F) -> cu::Result<Option<T>>
231 where
232 F: FnOnce(CtrlcSignal) -> cu::Result<T>,
233 {
234 let signal = CtrlcSignal::new(self.abort_threshold);
235 let Some(ctrlc_frame_scope) = CtrlcFrame::push_scope(signal.clone(), self.on_signal) else {
236 return f(signal).map(Some);
237 };
238 if let Err(e) = &*INIT_ONCE {
239 cu::bail!("{e}");
240 }
241 let result = f(signal.clone());
242 if signal.should_abort() {
243 return Ok(None);
244 }
245 drop(ctrlc_frame_scope); // suppress unused warning
246 result.map(Some)
247 }
248
249 #[cfg(feature = "coroutine")]
250 pub async fn co_execute<T, TFuture, F>(self, f: F) -> cu::Result<Option<T>>
251 where
252 TFuture: Future<Output = cu::Result<T>>,
253 F: FnOnce(CtrlcSignal) -> TFuture,
254 {
255 let signal = CtrlcSignal::new(self.abort_threshold);
256 let Some(ctrlc_frame_scope) = CtrlcFrame::push_scope(signal.clone(), self.on_signal) else {
257 return f(signal).await.map(Some);
258 };
259 if let Err(e) = &*INIT_ONCE {
260 cu::bail!("{e}");
261 }
262 let result = f(signal.clone()).await;
263 if signal.should_abort() {
264 return Ok(None);
265 }
266 drop(ctrlc_frame_scope); // suppress unused warning
267 result.map(Some)
268 }
269}
270
271type OnSignalFn = Box<dyn Fn(CtrlcSignal) + Send>;
272struct CtrlcFrame {
273 id: usize,
274 signal: CtrlcSignal,
275 on_signal: Option<OnSignalFn>,
276}
277struct CtrlcScope(usize);
278/// Signal passed into a task executing inside a `Ctrl-C` action frame,
279/// for it to check if `Ctrl-C` has been signaled
280#[derive(Clone)]
281pub struct CtrlcSignal {
282 signaled_times: Arc<AtomicU8>,
283 abort_threshold: u8,
284}
285impl CtrlcFrame {
286 pub fn push_scope(signal: CtrlcSignal, on_signal: Option<OnSignalFn>) -> Option<CtrlcScope> {
287 let Ok(mut signal_stack) = CTRLC_SIGNAL_STACK.lock() else {
288 cu::trace!("failed to register new ctrl-c frame");
289 return None;
290 };
291 let id = crate::next_atomic_usize();
292 signal_stack.push(Self {
293 id,
294 signal,
295 on_signal,
296 });
297 Some(CtrlcScope(id))
298 }
299}
300impl Drop for CtrlcScope {
301 fn drop(&mut self) {
302 if let Ok(mut signal_stack) = CTRLC_SIGNAL_STACK.lock() {
303 signal_stack.retain(|x| x.id != self.0);
304 }
305 }
306}
307
308impl CtrlcSignal {
309 fn new(abort_threshold: u8) -> Self {
310 Self {
311 signaled_times: Arc::new(AtomicU8::new(0)),
312 abort_threshold,
313 }
314 }
315 /// Return an `Err` if `Ctrl-C` has been signaled
316 pub fn check(&self) -> cu::Result<()> {
317 if self.should_abort() {
318 cu::bail!("interrupted")
319 }
320 Ok(())
321 }
322 /// Return `true` if `Ctrl-C` has been signaled at least
323 /// the same number of times as the abort_threshold
324 pub fn should_abort(&self) -> bool {
325 self.signaled_times() >= self.abort_threshold
326 }
327 /// Return `true` if `Ctrl-C` has been signaled at least once
328 pub fn signaled(&self) -> bool {
329 self.signaled_times() > 0
330 }
331 /// Get the number of times `Ctrl-C` has been signaled
332 pub fn signaled_times(&self) -> u8 {
333 self.signaled_times.load(Ordering::Acquire)
334 }
335 /// Programmatically trigger the signal. Note that this does not
336 /// send actual signal or keyboard events
337 pub fn signal(&self) {
338 self.signaled_times.fetch_add(1, Ordering::SeqCst);
339 }
340}