Get most tests passing when respecting wake order for foreground tasks in Deterministic executor

Co-Authored-By: Max Brunsfeld <maxbrunsfeld@gmail.com>
This commit is contained in:
Nathan Sobo 2022-01-24 17:13:56 -07:00
parent b7314ef2aa
commit 469ee554a0

View file

@ -73,7 +73,7 @@ unsafe impl<T: Send> Send for Task<T> {}
struct DeterministicState { struct DeterministicState {
rng: StdRng, rng: StdRng,
seed: u64, seed: u64,
scheduled_from_foreground: HashMap<usize, Vec<ScheduledForeground>>, scheduled_from_foreground: HashMap<usize, Vec<ForegroundRunnable>>,
scheduled_from_background: Vec<Runnable>, scheduled_from_background: Vec<Runnable>,
forbid_parking: bool, forbid_parking: bool,
block_on_ticks: RangeInclusive<usize>, block_on_ticks: RangeInclusive<usize>,
@ -82,9 +82,9 @@ struct DeterministicState {
waiting_backtrace: Option<Backtrace>, waiting_backtrace: Option<Backtrace>,
} }
enum ScheduledForeground { struct ForegroundRunnable {
MainFuture, runnable: Runnable,
Runnable(Runnable), main: bool,
} }
pub struct Deterministic { pub struct Deterministic {
@ -123,7 +123,12 @@ impl Deterministic {
}) })
} }
fn spawn_from_foreground(&self, cx_id: usize, future: AnyLocalFuture) -> AnyLocalTask { fn spawn_from_foreground(
&self,
cx_id: usize,
future: AnyLocalFuture,
main: bool,
) -> AnyLocalTask {
let state = self.state.clone(); let state = self.state.clone();
let unparker = self.parker.lock().unparker(); let unparker = self.parker.lock().unparker();
let (runnable, task) = async_task::spawn_local(future, move |runnable| { let (runnable, task) = async_task::spawn_local(future, move |runnable| {
@ -132,7 +137,7 @@ impl Deterministic {
.scheduled_from_foreground .scheduled_from_foreground
.entry(cx_id) .entry(cx_id)
.or_default() .or_default()
.push(ScheduledForeground::Runnable(runnable)); .push(ForegroundRunnable { runnable, main });
unparker.unpark(); unparker.unpark();
}); });
runnable.schedule(); runnable.schedule();
@ -151,10 +156,12 @@ impl Deterministic {
task task
} }
fn run(&self, cx_id: usize, mut future: AnyLocalFuture) -> Box<dyn Any> { fn run(&self, cx_id: usize, main_future: AnyLocalFuture) -> Box<dyn Any> {
let woken = Arc::new(AtomicBool::new(false)); let woken = Arc::new(AtomicBool::new(false));
let mut main_task = self.spawn_from_foreground(cx_id, main_future, true);
loop { loop {
if let Some(result) = self.run_internal(cx_id, woken.clone(), &mut future) { if let Some(result) = self.run_internal(woken.clone(), Some(&mut main_task)) {
return result; return result;
} }
@ -167,44 +174,20 @@ impl Deterministic {
} }
} }
fn run_until_parked(&self, cx_id: usize) { fn run_until_parked(&self) {
let woken = Arc::new(AtomicBool::new(false)); let woken = Arc::new(AtomicBool::new(false));
let mut future = any_local_future(std::future::pending::<()>()); self.run_internal(woken, None);
self.run_internal(cx_id, woken, &mut future);
} }
fn run_internal( fn run_internal(
&self, &self,
cx_id: usize,
woken: Arc<AtomicBool>, woken: Arc<AtomicBool>,
future: &mut AnyLocalFuture, mut main_task: Option<&mut AnyLocalTask>,
) -> Option<Box<dyn Any>> { ) -> Option<Box<dyn Any>> {
let unparker = self.parker.lock().unparker(); let unparker = self.parker.lock().unparker();
let scheduled_main_future = Arc::new(AtomicBool::new(true)); let waker = waker_fn(move || {
self.state woken.store(true, SeqCst);
.lock() unparker.unpark();
.scheduled_from_foreground
.entry(cx_id)
.or_default()
.insert(0, ScheduledForeground::MainFuture);
let waker = waker_fn({
let state = self.state.clone();
let scheduled_main_future = scheduled_main_future.clone();
move || {
woken.store(true, SeqCst);
if !scheduled_main_future.load(SeqCst) {
scheduled_main_future.store(true, SeqCst);
state
.lock()
.scheduled_from_foreground
.entry(cx_id)
.or_default()
.push(ScheduledForeground::MainFuture);
}
unparker.unpark();
}
}); });
let mut cx = Context::from_waker(&waker); let mut cx = Context::from_waker(&waker);
@ -234,25 +217,21 @@ impl Deterministic {
.scheduled_from_foreground .scheduled_from_foreground
.get_mut(&cx_id_to_run) .get_mut(&cx_id_to_run)
.unwrap(); .unwrap();
let runnable = scheduled_from_cx.remove(0); let foreground_runnable = scheduled_from_cx.remove(0);
if scheduled_from_cx.is_empty() { if scheduled_from_cx.is_empty() {
state.scheduled_from_foreground.remove(&cx_id_to_run); state.scheduled_from_foreground.remove(&cx_id_to_run);
} }
drop(state); drop(state);
match runnable {
ScheduledForeground::MainFuture => { foreground_runnable.runnable.run();
scheduled_main_future.store(false, SeqCst); if let Some(main_task) = main_task.as_mut() {
if let Poll::Ready(result) = future.poll(&mut cx) { if foreground_runnable.main {
if let Poll::Ready(result) = main_task.poll(&mut cx) {
return Some(result); return Some(result);
} }
} }
ScheduledForeground::Runnable(runnable) => {
runnable.run();
}
} }
} else {
return None;
} }
} }
} }
@ -364,7 +343,7 @@ impl Foreground {
let future = any_local_future(future); let future = any_local_future(future);
let any_task = match self { let any_task = match self {
Self::Deterministic { cx_id, executor } => { Self::Deterministic { cx_id, executor } => {
executor.spawn_from_foreground(*cx_id, future) executor.spawn_from_foreground(*cx_id, future, false)
} }
Self::Platform { dispatcher, .. } => { Self::Platform { dispatcher, .. } => {
fn spawn_inner( fn spawn_inner(
@ -448,8 +427,8 @@ impl Foreground {
pub fn advance_clock(&self, duration: Duration) { pub fn advance_clock(&self, duration: Duration) {
match self { match self {
Self::Deterministic { cx_id, executor } => { Self::Deterministic { executor, .. } => {
executor.run_until_parked(*cx_id); executor.run_until_parked();
let mut state = executor.state.lock(); let mut state = executor.state.lock();
state.now += duration; state.now += duration;