lelet_utils/
lib.rs

1use std::future::Future;
2use std::marker::PhantomData;
3use std::mem;
4use std::pin::Pin;
5use std::process::abort;
6use std::sync::atomic::AtomicBool;
7use std::sync::atomic::Ordering;
8use std::sync::Arc;
9use std::task::Context;
10use std::task::Poll;
11use std::task::RawWaker;
12use std::task::RawWakerVTable;
13use std::task::Waker;
14use std::thread;
15use std::thread::Thread;
16
17/// Call [`abort`] when `f` panic
18///
19/// [`abort`]: https://doc.rust-lang.org/std/process/fn.abort.html
20pub fn abort_on_panic<R>(f: impl FnOnce() -> R) -> R {
21    struct Bomb;
22
23    impl Drop for Bomb {
24        fn drop(&mut self) {
25            abort();
26        }
27    }
28
29    let bomb = Bomb;
30
31    let r = f();
32
33    mem::forget(bomb);
34
35    r
36}
37
38/// Future that will yield multiple times
39#[derive(Debug)]
40pub struct Yields(pub usize);
41
42impl Future for Yields {
43    type Output = ();
44
45    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
46        if self.0 == 0 {
47            Poll::Ready(())
48        } else {
49            self.0 -= 1;
50            cx.waker().wake_by_ref();
51            Poll::Pending
52        }
53    }
54}
55
56/// Block current thread until f is complete
57pub fn block_on<F: Future>(f: F) -> F::Output {
58    static VTABLE: RawWakerVTable = RawWakerVTable::new(
59        |clone_me| {
60            let clone_me = unsafe { Arc::from_raw(clone_me as *const Unparker) };
61            mem::forget(clone_me.clone());
62            RawWaker::new(Arc::into_raw(clone_me) as *const (), &VTABLE)
63        },
64        |wake_me| unsafe { Arc::from_raw(wake_me as *const Unparker) }.unpark(),
65        |wake_me_by_ref| unsafe { &*(wake_me_by_ref as *const Unparker) }.unpark(),
66        |drop_me| drop(unsafe { Arc::from_raw(drop_me as *const Unparker) }),
67    );
68
69    let parker = Parker::new();
70
71    let waker = {
72        let raw = RawWaker::new(Arc::into_raw(parker.unparker()) as *const (), &VTABLE);
73        unsafe { Waker::from_raw(raw) }
74    };
75
76    let mut f = f;
77    let mut cx = Context::from_waker(&waker);
78    loop {
79        match unsafe { Pin::new_unchecked(&mut f) }.poll(&mut cx) {
80            Poll::Pending => parker.park(),
81            Poll::Ready(val) => return val,
82        }
83    }
84}
85
86struct ParkerInner {
87    parked: AtomicBool,
88    thread: Thread,
89
90    // !Send + !Sync
91    _marker: PhantomData<*mut ()>,
92}
93
94impl ParkerInner {
95    fn new() -> ParkerInner {
96        ParkerInner {
97            parked: AtomicBool::new(false),
98            thread: thread::current(),
99            _marker: PhantomData,
100        }
101    }
102
103    fn park(&self) {
104        // wait while the flag is set
105        while !self.parked.load(Ordering::Relaxed) {
106            thread::park();
107        }
108
109        // consume the flag
110        self.parked.store(false, Ordering::Relaxed);
111    }
112
113    fn unpark(&self) {
114        // set the flag
115        self.parked.store(true, Ordering::Relaxed);
116
117        self.thread.unpark();
118    }
119}
120
121/// A thread parking primitive
122pub struct Parker(Arc<ParkerInner>);
123
124impl Parker {
125    /// Create new Parker
126    pub fn new() -> Parker {
127        Parker(Arc::new(ParkerInner::new()))
128    }
129
130    /// Park current thread
131    pub fn park(&self) {
132        self.0.park();
133    }
134
135    /// Return the associated unparker
136    pub fn unparker(&self) -> Arc<Unparker> {
137        let cloned = self.0.clone();
138        unsafe { mem::transmute(cloned) }
139    }
140}
141
142/// Unparker for the associated parked thread
143pub struct Unparker(ParkerInner);
144
145impl Unparker {
146    /// Unpark the associated parked thread
147    pub fn unpark(&self) {
148        self.0.unpark();
149    }
150}
151
152unsafe impl Send for Unparker {}
153unsafe impl Sync for Unparker {}
lelet_utils/lib.rs

lelet_utils/
lib.rs
