async_ctrlc/
lib.rs

1// Copyright 2019 kennytm. Licensed under MIT OR Apache-2.0.
2
3//! `async-ctrlc` is an async wrapper of the `ctrlc` crate.
4
5use ctrlc::set_handler;
6pub use ctrlc::Error;
7#[cfg(feature = "stream")]
8use futures_core::stream::Stream;
9use std::{
10    future::Future,
11    pin::Pin,
12    ptr::null_mut,
13    sync::atomic::{AtomicBool, AtomicPtr, Ordering},
14    task::{Context, Poll, Waker},
15};
16
17// TODO: Replace this with `AtomicOptionBox<Waker>`
18// after https://github.com/jorendorff/atomicbox/pull/3 is merged.
19static WAKER: AtomicPtr<Waker> = AtomicPtr::new(null_mut());
20static ACTIVE: AtomicBool = AtomicBool::new(false);
21
22/// A future which is fulfilled when the program receives the Ctrl+C signal.
23#[derive(Debug)]
24pub struct CtrlC {
25    _private: (),
26}
27
28impl Future for CtrlC {
29    type Output = ();
30    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
31        if ACTIVE.swap(false, Ordering::SeqCst) {
32            Poll::Ready(())
33        } else {
34            let new_waker = Box::new(cx.waker().clone());
35            let old_waker_ptr = WAKER.swap(Box::into_raw(new_waker), Ordering::SeqCst);
36            if !old_waker_ptr.is_null() {
37                unsafe { Box::from_raw(old_waker_ptr) };
38            }
39            Poll::Pending
40        }
41    }
42}
43
44#[cfg(feature = "stream")]
45impl Stream for CtrlC {
46    type Item = ();
47    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
48        self.poll(cx).map(Some)
49    }
50}
51
52impl CtrlC {
53    /// Creates a new `CtrlC` future.
54    ///
55    /// There should be at most one `CtrlC` instance in the whole program. The
56    /// second call to `Ctrl::new()` would return an error.
57    pub fn new() -> Result<Self, Error> {
58        set_handler(|| {
59            ACTIVE.store(true, Ordering::SeqCst);
60            let waker_ptr = WAKER.swap(null_mut(), Ordering::SeqCst);
61            if !waker_ptr.is_null() {
62                unsafe { Box::from_raw(waker_ptr) }.wake();
63            }
64        })?;
65        Ok(CtrlC { _private: () })
66    }
67}
68
69#[cfg(unix)]
70#[test]
71fn test_unix() {
72    use async_std::{future::timeout, task::block_on};
73    use libc::{getpid, kill, SIGINT};
74    use std::{
75        thread::{sleep, spawn},
76        time::Duration,
77    };
78
79    let thread = spawn(|| unsafe {
80        sleep(Duration::from_millis(100));
81        kill(getpid(), SIGINT);
82    });
83
84    let c = CtrlC::new().unwrap();
85    let result = block_on(async move {
86        let i = 1;
87        timeout(Duration::from_millis(500), c).await.unwrap();
88        i + 2
89    });
90    assert_eq!(result, 3);
91
92    thread.join().unwrap();
93}
async_ctrlc/lib.rs

async_ctrlc/
lib.rs
