Allow DeterministicExecutor to park until forbid_parking is called

This allows us to perform async setup such as talking to the database.
This commit is contained in:
Nathan Sobo
2021-07-20 12:22:02 -06:00
parent f6263768b3
commit 64927abbc4
4 changed files with 89 additions and 40 deletions
+69 -37
View File
@@ -2,7 +2,6 @@ use anyhow::{anyhow, Result};
use async_task::Runnable;
pub use async_task::Task;
use backtrace::{Backtrace, BacktraceFmt, BytesOrWideString};
use futures::task::noop_waker;
use parking_lot::Mutex;
use rand::prelude::*;
use smol::{channel, prelude::*, Executor};
@@ -20,6 +19,7 @@ use std::{
thread,
time::Duration,
};
use waker_fn::waker_fn;
use crate::{platform, util};
@@ -45,18 +45,26 @@ struct DeterministicState {
seed: u64,
scheduled: Vec<(Runnable, Backtrace)>,
spawned_from_foreground: Vec<(Runnable, Backtrace)>,
forbid_parking: bool,
}
pub struct Deterministic(Arc<Mutex<DeterministicState>>);
pub struct Deterministic {
state: Arc<Mutex<DeterministicState>>,
parker: Mutex<parking::Parker>,
}
impl Deterministic {
fn new(seed: u64) -> Self {
Self(Arc::new(Mutex::new(DeterministicState {
rng: StdRng::seed_from_u64(seed),
seed,
scheduled: Default::default(),
spawned_from_foreground: Default::default(),
})))
Self {
state: Arc::new(Mutex::new(DeterministicState {
rng: StdRng::seed_from_u64(seed),
seed,
scheduled: Default::default(),
spawned_from_foreground: Default::default(),
forbid_parking: false,
})),
parker: Default::default(),
}
}
pub fn spawn_from_foreground<F, T>(&self, future: F) -> Task<T>
@@ -66,7 +74,8 @@ impl Deterministic {
{
let backtrace = Backtrace::new_unresolved();
let scheduled_once = AtomicBool::new(false);
let state = self.0.clone();
let state = self.state.clone();
let unparker = self.parker.lock().unparker();
let (runnable, task) = async_task::spawn_local(future, move |runnable| {
let mut state = state.lock();
let backtrace = backtrace.clone();
@@ -75,6 +84,7 @@ impl Deterministic {
} else {
state.spawned_from_foreground.push((runnable, backtrace));
}
unparker.unpark();
});
runnable.schedule();
task
@@ -86,9 +96,12 @@ impl Deterministic {
F: 'static + Send + Future<Output = T>,
{
let backtrace = Backtrace::new_unresolved();
let state = self.0.clone();
let state = self.state.clone();
let unparker = self.parker.lock().unparker();
let (runnable, task) = async_task::spawn(future, move |runnable| {
state.lock().scheduled.push((runnable, backtrace.clone()));
let mut state = state.lock();
state.scheduled.push((runnable, backtrace.clone()));
unparker.unpark();
});
runnable.schedule();
task
@@ -109,36 +122,45 @@ impl Deterministic {
{
smol::pin!(future);
let waker = noop_waker();
let unparker = self.parker.lock().unparker();
let waker = waker_fn(move || {
unparker.unpark();
});
let mut cx = Context::from_waker(&waker);
let mut trace = Trace::default();
for _ in 0..max_ticks {
let runnable = {
let state = &mut *self.0.lock();
let runnable_count = state.scheduled.len() + state.spawned_from_foreground.len();
let ix = state.rng.gen_range(0..=runnable_count);
if ix < state.scheduled.len() {
let (_, backtrace) = &state.scheduled[ix];
trace.record(&state, backtrace.clone());
state.scheduled.remove(ix).0
} else if ix < runnable_count {
let (_, backtrace) = &state.spawned_from_foreground[0];
trace.record(&state, backtrace.clone());
state.spawned_from_foreground.remove(0).0
} else {
if let Poll::Ready(result) = future.as_mut().poll(&mut cx) {
return Some(result);
let mut state = self.state.lock();
let runnable_count = state.scheduled.len() + state.spawned_from_foreground.len();
let ix = state.rng.gen_range(0..=runnable_count);
if ix < state.scheduled.len() {
let (_, backtrace) = &state.scheduled[ix];
trace.record(&state, backtrace.clone());
let runnable = state.scheduled.remove(ix).0;
drop(state);
runnable.run();
} else if ix < runnable_count {
let (_, backtrace) = &state.spawned_from_foreground[0];
trace.record(&state, backtrace.clone());
let runnable = state.spawned_from_foreground.remove(0).0;
drop(state);
runnable.run();
} else {
drop(state);
if let Poll::Ready(result) = future.as_mut().poll(&mut cx) {
return Some(result);
}
let state = &mut *self.state.lock();
if state.scheduled.is_empty() && state.spawned_from_foreground.is_empty() {
if state.forbid_parking {
panic!("deterministic executor parked after a call to forbid_parking");
}
if state.scheduled.is_empty() && state.spawned_from_foreground.is_empty() {
panic!("detected non-determinism in deterministic executor");
} else {
continue;
}
drop(state);
self.parker.lock().park();
}
};
runnable.run();
continue;
}
}
None
@@ -311,11 +333,21 @@ impl Foreground {
Self::Platform { .. } => panic!("can't call this method on a platform executor"),
Self::Test(_) => panic!("can't call this method on a test executor"),
Self::Deterministic(executor) => {
let state = &mut *executor.0.lock();
let state = &mut *executor.state.lock();
state.rng = StdRng::seed_from_u64(state.seed);
}
}
}
pub fn forbid_parking(&self) {
match self {
Self::Platform { .. } => panic!("can't call this method on a platform executor"),
Self::Test(_) => panic!("can't call this method on a test executor"),
Self::Deterministic(executor) => {
executor.state.lock().forbid_parking = true;
}
}
}
}
impl Background {
@@ -363,7 +395,7 @@ impl Background {
smol::block_on(async move { util::timeout(timeout, future).await.ok() })
}
Self::Deterministic(executor) => {
let max_ticks = executor.0.lock().rng.gen_range(1..=1000);
let max_ticks = executor.state.lock().rng.gen_range(1..=1000);
executor.block_on(max_ticks, future)
}
}