/src/wasmtime/crates/fiber/src/lib.rs

Source (jump to first uncovered line)
use std::any::Any;
use std::cell::Cell;
use std::io;
use std::marker::PhantomData;
use std::panic::{self, AssertUnwindSafe};

#[cfg(windows)]
mod windows;
#[cfg(windows)]
use windows as imp;

#[cfg(unix)]
mod unix;
#[cfg(unix)]
use unix as imp;

pub struct Fiber<'a, Resume, Yield, Return> {
    inner: imp::Fiber,
    done: Cell<bool>,
    _phantom: PhantomData<&'a (Resume, Yield, Return)>,
}

pub struct Suspend<Resume, Yield, Return> {
    inner: imp::Suspend,
    _phantom: PhantomData<(Resume, Yield, Return)>,
}

enum RunResult<Resume, Yield, Return> {
    Executing,
    Resuming(Resume),
    Yield(Yield),
    Returned(Return),
    Panicked(Box<dyn Any + Send>),
}

impl<'a, Resume, Yield, Return> Fiber<'a, Resume, Yield, Return> {
    /// Creates a new fiber which will execute `func` on a new native stack of
    /// size `stack_size`.
    ///
    /// This function returns a `Fiber` which, when resumed, will execute `func`
    /// to completion. When desired the `func` can suspend itself via
    /// `Fiber::suspend`.
    pub fn new(
        stack_size: usize,
        func: impl FnOnce(Resume, &Suspend<Resume, Yield, Return>) -> Return + 'a,
    ) -> io::Result<Fiber<'a, Resume, Yield, Return>> {
        Ok(Fiber {
            inner: imp::Fiber::new(stack_size, func)?,
            done: Cell::new(false),
            _phantom: PhantomData,
        })
    }

    /// Creates a new fiber with existing stack space that will execute `func`.
    ///
    /// This function returns a `Fiber` which, when resumed, will execute `func`
    /// to completion. When desired the `func` can suspend itself via
    /// `Fiber::suspend`.
    ///
    /// # Safety
    ///
    /// The caller must properly allocate the stack space with a guard page and
    /// make the pages accessible for correct behavior.
    pub unsafe fn new_with_stack(
        top_of_stack: *mut u8,
        func: impl FnOnce(Resume, &Suspend<Resume, Yield, Return>) -> Return + 'a,
    ) -> io::Result<Fiber<'a, Resume, Yield, Return>> {
        Ok(Fiber {
            inner: imp::Fiber::new_with_stack(top_of_stack, func)?,
            done: Cell::new(false),
            _phantom: PhantomData,
        })
    }

    /// Resumes execution of this fiber.
    ///
    /// This function will transfer execution to the fiber and resume from where
    /// it last left off.
    ///
    /// Returns `true` if the fiber finished or `false` if the fiber was
    /// suspended in the middle of execution.
    ///
    /// # Panics
    ///
    /// Panics if the current thread is already executing a fiber or if this
    /// fiber has already finished.
    ///
    /// Note that if the fiber itself panics during execution then the panic
    /// will be propagated to this caller.
    pub fn resume(&self, val: Resume) -> Result<Return, Yield> {
        assert!(!self.done.replace(true), "cannot resume a finished fiber");
        let result = Cell::new(RunResult::Resuming(val));
        self.inner.resume(&result);
        match result.into_inner() {
            RunResult::Resuming(_) | RunResult::Executing => unreachable!(),
            RunResult::Yield(y) => {
                self.done.set(false);
                Err(y)
            }
            RunResult::Returned(r) => Ok(r),
            RunResult::Panicked(payload) => std::panic::resume_unwind(payload),
        }
    }

    /// Returns whether this fiber has finished executing.
    pub fn done(&self) -> bool {
        self.done.get()
    }
}

impl<Resume, Yield, Return> Suspend<Resume, Yield, Return> {
    /// Suspend execution of a currently running fiber.
    ///
    /// This function will switch control back to the original caller of
    /// `Fiber::resume`. This function will then return once the `Fiber::resume`
    /// function is called again.
    ///
    /// # Panics
    ///
    /// Panics if the current thread is not executing a fiber from this library.
    pub fn suspend(&self, value: Yield) -> Resume {
        self.inner
            .switch::<Resume, Yield, Return>(RunResult::Yield(value))
    }

    fn execute(
        inner: imp::Suspend,
        initial: Resume,
        func: impl FnOnce(Resume, &Suspend<Resume, Yield, Return>) -> Return,
    ) {
        let suspend = Suspend {
            inner,
            _phantom: PhantomData,
        };
        let result = panic::catch_unwind(AssertUnwindSafe(|| (func)(initial, &suspend)));
        suspend.inner.switch::<Resume, Yield, Return>(match result {
            Ok(result) => RunResult::Returned(result),
            Err(panic) => RunResult::Panicked(panic),
        });
    }
}

impl<A, B, C> Drop for Fiber<'_, A, B, C> {
    fn drop(&mut self) {
        debug_assert!(self.done.get(), "fiber dropped without finishing");
    }
}

#[cfg(test)]
mod tests {
    use super::Fiber;
    use std::cell::Cell;
    use std::panic::{self, AssertUnwindSafe};
    use std::rc::Rc;

    #[test]
    fn small_stacks() {
        Fiber::<(), (), ()>::new(0, |_, _| {})
            .unwrap()
            .resume(())
            .unwrap();
        Fiber::<(), (), ()>::new(1, |_, _| {})
            .unwrap()
            .resume(())
            .unwrap();
    }

    #[test]
    fn smoke() {
        let hit = Rc::new(Cell::new(false));
        let hit2 = hit.clone();
        let fiber = Fiber::<(), (), ()>::new(1024 * 1024, move |_, _| {
            hit2.set(true);
        })
        .unwrap();
        assert!(!hit.get());
        fiber.resume(()).unwrap();
        assert!(hit.get());
    }

    #[test]
    fn suspend_and_resume() {
        let hit = Rc::new(Cell::new(false));
        let hit2 = hit.clone();
        let fiber = Fiber::<(), (), ()>::new(1024 * 1024, move |_, s| {
            s.suspend(());
            hit2.set(true);
            s.suspend(());
        })
        .unwrap();
        assert!(!hit.get());
        assert!(fiber.resume(()).is_err());
        assert!(!hit.get());
        assert!(fiber.resume(()).is_err());
        assert!(hit.get());
        assert!(fiber.resume(()).is_ok());
        assert!(hit.get());
    }

    #[test]
    fn backtrace_traces_to_host() {
        #[inline(never)] // try to get this to show up in backtraces
        fn look_for_me() {
            run_test();
        }
        fn assert_contains_host() {
            let trace = backtrace::Backtrace::new();
            println!("{:?}", trace);
            assert!(
                trace
                .frames()
                .iter()
                .flat_map(|f| f.symbols())
                .filter_map(|s| Some(s.name()?.to_string()))
                .any(|s| s.contains("look_for_me"))
                // TODO: apparently windows unwind routines don't unwind through fibers, so this will always fail. Is there a way we can fix that?
                || cfg!(windows)
            );
        }

        fn run_test() {
            let fiber = Fiber::<(), (), ()>::new(1024 * 1024, move |(), s| {
                assert_contains_host();
                s.suspend(());
                assert_contains_host();
                s.suspend(());
                assert_contains_host();
            })
            .unwrap();
            assert!(fiber.resume(()).is_err());
            assert!(fiber.resume(()).is_err());
            assert!(fiber.resume(()).is_ok());
        }

        look_for_me();
    }

    #[test]
    fn panics_propagated() {
        let a = Rc::new(Cell::new(false));
        let b = SetOnDrop(a.clone());
        let fiber = Fiber::<(), (), ()>::new(1024 * 1024, move |(), _s| {
            drop(&b);
            panic!();
        })
        .unwrap();
        assert!(panic::catch_unwind(AssertUnwindSafe(|| fiber.resume(()))).is_err());
        assert!(a.get());

        struct SetOnDrop(Rc<Cell<bool>>);

        impl Drop for SetOnDrop {
            fn drop(&mut self) {
                self.0.set(true);
            }
        }
    }

    #[test]
    fn suspend_and_resume_values() {
        let fiber = Fiber::new(1024 * 1024, move |first, s| {
            assert_eq!(first, 2.0);
            assert_eq!(s.suspend(4), 3.0);
            "hello".to_string()
        })
        .unwrap();
        assert_eq!(fiber.resume(2.0), Err(4));
        assert_eq!(fiber.resume(3.0), Ok("hello".to_string()));
    }
}

Coverage Report

Created: 2021-03-22 08:29

Line	Count	Source (jump to first uncovered line)
1		use std::any::Any;
2		use std::cell::Cell;
3		use std::io;
4		use std::marker::PhantomData;
5		use std::panic::{self, AssertUnwindSafe};
6
7		#[cfg(windows)]
8		mod windows;
9		#[cfg(windows)]
10		use windows as imp;
11
12		#[cfg(unix)]
13		mod unix;
14		#[cfg(unix)]
15		use unix as imp;
16
17		pub struct Fiber<'a, Resume, Yield, Return> {
18		inner: imp::Fiber,
19		done: Cell<bool>,
20		_phantom: PhantomData<&'a (Resume, Yield, Return)>,
21		}
22
23		pub struct Suspend<Resume, Yield, Return> {
24		inner: imp::Suspend,
25		_phantom: PhantomData<(Resume, Yield, Return)>,
26		}
27
28		enum RunResult<Resume, Yield, Return> {
29		Executing,
30		Resuming(Resume),
31		Yield(Yield),
32		Returned(Return),
33		Panicked(Box<dyn Any + Send>),
34		}
35
36		impl<'a, Resume, Yield, Return> Fiber<'a, Resume, Yield, Return> {
37		/// Creates a new fiber which will execute `func` on a new native stack of
38		/// size `stack_size`.
39		///
40		/// This function returns a `Fiber` which, when resumed, will execute `func`
41		/// to completion. When desired the `func` can suspend itself via
42		/// `Fiber::suspend`.
43		pub fn new(
44		stack_size: usize,
45		func: impl FnOnce(Resume, &Suspend<Resume, Yield, Return>) -> Return + 'a,
46		) -> io::Result<Fiber<'a, Resume, Yield, Return>> {
47		Ok(Fiber {
48		inner: imp::Fiber::new(stack_size, func)?,
49		done: Cell::new(false),
50		_phantom: PhantomData,
51		})
52		}
53
54		/// Creates a new fiber with existing stack space that will execute `func`.
55		///
56		/// This function returns a `Fiber` which, when resumed, will execute `func`
57		/// to completion. When desired the `func` can suspend itself via
58		/// `Fiber::suspend`.
59		///
60		/// # Safety
61		///
62		/// The caller must properly allocate the stack space with a guard page and
63		/// make the pages accessible for correct behavior.
64		pub unsafe fn new_with_stack(
65		top_of_stack: *mut u8,
66		func: impl FnOnce(Resume, &Suspend<Resume, Yield, Return>) -> Return + 'a,
67		) -> io::Result<Fiber<'a, Resume, Yield, Return>> {
68		Ok(Fiber {
69		inner: imp::Fiber::new_with_stack(top_of_stack, func)?,
70		done: Cell::new(false),
71		_phantom: PhantomData,
72		})
73		}
74
75		/// Resumes execution of this fiber.
76		///
77		/// This function will transfer execution to the fiber and resume from where
78		/// it last left off.
79		///
80		/// Returns `true` if the fiber finished or `false` if the fiber was
81		/// suspended in the middle of execution.
82		///
83		/// # Panics
84		///
85		/// Panics if the current thread is already executing a fiber or if this
86		/// fiber has already finished.
87		///
88		/// Note that if the fiber itself panics during execution then the panic
89		/// will be propagated to this caller.
90	0	pub fn resume(&self, val: Resume) -> Result<Return, Yield> {
91	0	assert!(!self.done.replace(true), "cannot resume a finished fiber");
92	0	let result = Cell::new(RunResult::Resuming(val));
93	0	self.inner.resume(&result);
94	0	match result.into_inner() {
95	0	RunResult::Resuming(_) \| RunResult::Executing => unreachable!(),
96	0	RunResult::Yield(y) => {
97	0	self.done.set(false);
98	0	Err(y)
99		}
100	0	RunResult::Returned(r) => Ok(r),
101	0	RunResult::Panicked(payload) => std::panic::resume_unwind(payload),
102		}
103	0	}
104
105		/// Returns whether this fiber has finished executing.
106		pub fn done(&self) -> bool {
107		self.done.get()
108		}
109		}
110
111		impl<Resume, Yield, Return> Suspend<Resume, Yield, Return> {
112		/// Suspend execution of a currently running fiber.
113		///
114		/// This function will switch control back to the original caller of
115		/// `Fiber::resume`. This function will then return once the `Fiber::resume`
116		/// function is called again.
117		///
118		/// # Panics
119		///
120		/// Panics if the current thread is not executing a fiber from this library.
121		pub fn suspend(&self, value: Yield) -> Resume {
122		self.inner
123		.switch::<Resume, Yield, Return>(RunResult::Yield(value))
124		}
125
126		fn execute(
127		inner: imp::Suspend,
128		initial: Resume,
129		func: impl FnOnce(Resume, &Suspend<Resume, Yield, Return>) -> Return,
130		) {
131		let suspend = Suspend {
132		inner,
133		_phantom: PhantomData,
134		};
135		let result = panic::catch_unwind(AssertUnwindSafe(\|\| (func)(initial, &suspend)));
136		suspend.inner.switch::<Resume, Yield, Return>(match result {
137		Ok(result) => RunResult::Returned(result),
138		Err(panic) => RunResult::Panicked(panic),
139		});
140		}
141		}
142
143		impl<A, B, C> Drop for Fiber<'_, A, B, C> {
144		fn drop(&mut self) {
145		debug_assert!(self.done.get(), "fiber dropped without finishing");
146		}
147		}
148
149		#[cfg(test)]
150		mod tests {
151		use super::Fiber;
152		use std::cell::Cell;
153		use std::panic::{self, AssertUnwindSafe};
154		use std::rc::Rc;
155
156		#[test]
157		fn small_stacks() {
158		Fiber::<(), (), ()>::new(0, \|_, _\| {})
159		.unwrap()
160		.resume(())
161		.unwrap();
162		Fiber::<(), (), ()>::new(1, \|_, _\| {})
163		.unwrap()
164		.resume(())
165		.unwrap();
166		}
167
168		#[test]
169		fn smoke() {
170		let hit = Rc::new(Cell::new(false));
171		let hit2 = hit.clone();
172		let fiber = Fiber::<(), (), ()>::new(1024 * 1024, move \|_, _\| {
173		hit2.set(true);
174		})
175		.unwrap();
176		assert!(!hit.get());
177		fiber.resume(()).unwrap();
178		assert!(hit.get());
179		}
180
181		#[test]
182		fn suspend_and_resume() {
183		let hit = Rc::new(Cell::new(false));
184		let hit2 = hit.clone();
185		let fiber = Fiber::<(), (), ()>::new(1024 * 1024, move \|_, s\| {
186		s.suspend(());
187		hit2.set(true);
188		s.suspend(());
189		})
190		.unwrap();
191		assert!(!hit.get());
192		assert!(fiber.resume(()).is_err());
193		assert!(!hit.get());
194		assert!(fiber.resume(()).is_err());
195		assert!(hit.get());
196		assert!(fiber.resume(()).is_ok());
197		assert!(hit.get());
198		}
199
200		#[test]
201		fn backtrace_traces_to_host() {
202		#[inline(never)] // try to get this to show up in backtraces
203		fn look_for_me() {
204		run_test();
205		}
206		fn assert_contains_host() {
207		let trace = backtrace::Backtrace::new();
208		println!("{:?}", trace);
209		assert!(
210		trace
211		.frames()
212		.iter()
213		.flat_map(\|f\| f.symbols())
214		.filter_map(\|s\| Some(s.name()?.to_string()))
215		.any(\|s\| s.contains("look_for_me"))
216		// TODO: apparently windows unwind routines don't unwind through fibers, so this will always fail. Is there a way we can fix that?
217		\|\| cfg!(windows)
218		);
219		}
220
221		fn run_test() {
222		let fiber = Fiber::<(), (), ()>::new(1024 * 1024, move \|(), s\| {
223		assert_contains_host();
224		s.suspend(());
225		assert_contains_host();
226		s.suspend(());
227		assert_contains_host();
228		})
229		.unwrap();
230		assert!(fiber.resume(()).is_err());
231		assert!(fiber.resume(()).is_err());
232		assert!(fiber.resume(()).is_ok());
233		}
234
235		look_for_me();
236		}
237
238		#[test]
239		fn panics_propagated() {
240		let a = Rc::new(Cell::new(false));
241		let b = SetOnDrop(a.clone());
242		let fiber = Fiber::<(), (), ()>::new(1024 * 1024, move \|(), _s\| {
243		drop(&b);
244		panic!();
245		})
246		.unwrap();
247		assert!(panic::catch_unwind(AssertUnwindSafe(\|\| fiber.resume(()))).is_err());
248		assert!(a.get());
249
250		struct SetOnDrop(Rc<Cell<bool>>);
251
252		impl Drop for SetOnDrop {
253		fn drop(&mut self) {
254		self.0.set(true);
255		}
256		}
257		}
258
259		#[test]
260		fn suspend_and_resume_values() {
261		let fiber = Fiber::new(1024 * 1024, move \|first, s\| {
262		assert_eq!(first, 2.0);
263		assert_eq!(s.suspend(4), 3.0);
264		"hello".to_string()
265		})
266		.unwrap();
267		assert_eq!(fiber.resume(2.0), Err(4));
268		assert_eq!(fiber.resume(3.0), Ok("hello".to_string()));
269		}
270		}