1031 lines · 32.3 KB · Rust
RawDownload
use anyhow::{anyhow, Result};use async_task::Runnable;use futures::channel::mpsc;use smol::{channel, prelude::*, Executor};use std::{    any::Any,    fmt::{self, Display},    marker::PhantomData,    mem,    panic::Location,    pin::Pin,    rc::Rc,    sync::Arc,    task::{Context, Poll},    thread,    time::Duration,}; use crate::{    platform::{self, Dispatcher},    util, AppContext,}; pub enum Foreground {    Platform {        dispatcher: Arc<dyn platform::Dispatcher>,        _not_send_or_sync: PhantomData<Rc<()>>,    },    #[cfg(any(test, feature = "test-support"))]    Deterministic {        cx_id: usize,        executor: Arc<Deterministic>,    },} pub enum Background {    #[cfg(any(test, feature = "test-support"))]    Deterministic { executor: Arc<Deterministic> },    Production {        executor: Arc<smol::Executor<'static>>,        _stop: channel::Sender<()>,    },} type AnyLocalFuture = Pin<Box<dyn 'static + Future<Output = Box<dyn Any + 'static>>>>;type AnyFuture = Pin<Box<dyn 'static + Send + Future<Output = Box<dyn Any + Send + 'static>>>>;type AnyTask = async_task::Task<Box<dyn Any + Send + 'static>>;type AnyLocalTask = async_task::Task<Box<dyn Any + 'static>>; #[must_use]pub enum Task<T> {    Ready(Option<T>),    Local {        any_task: AnyLocalTask,        result_type: PhantomData<T>,    },    Send {        any_task: AnyTask,        result_type: PhantomData<T>,    },} unsafe impl<T: Send> Send for Task<T> {} #[cfg(any(test, feature = "test-support"))]struct DeterministicState {    rng: rand::prelude::StdRng,    seed: u64,    scheduled_from_foreground: collections::HashMap<usize, Vec<ForegroundRunnable>>,    scheduled_from_background: Vec<BackgroundRunnable>,    forbid_parking: bool,    block_on_ticks: std::ops::RangeInclusive<usize>,    now: std::time::Instant,    next_timer_id: usize,    pending_timers: Vec<(usize, std::time::Instant, postage::barrier::Sender)>,    waiting_backtrace: Option<backtrace::Backtrace>,    next_runnable_id: usize,    poll_history: Vec<ExecutorEvent>,    previous_poll_history: Option<Vec<ExecutorEvent>>,    enable_runnable_backtraces: bool,    runnable_backtraces: collections::HashMap<usize, backtrace::Backtrace>,} #[derive(Copy, Clone, Debug, PartialEq, Eq)]pub enum ExecutorEvent {    PollRunnable { id: usize },    EnqueuRunnable { id: usize },} #[cfg(any(test, feature = "test-support"))]struct ForegroundRunnable {    id: usize,    runnable: Runnable,    main: bool,} #[cfg(any(test, feature = "test-support"))]struct BackgroundRunnable {    id: usize,    runnable: Runnable,} #[cfg(any(test, feature = "test-support"))]pub struct Deterministic {    state: Arc<parking_lot::Mutex<DeterministicState>>,    parker: parking_lot::Mutex<parking::Parker>,} #[must_use]pub enum Timer {    Production(smol::Timer),    #[cfg(any(test, feature = "test-support"))]    Deterministic(DeterministicTimer),} #[cfg(any(test, feature = "test-support"))]pub struct DeterministicTimer {    rx: postage::barrier::Receiver,    id: usize,    state: Arc<parking_lot::Mutex<DeterministicState>>,} #[cfg(any(test, feature = "test-support"))]impl Deterministic {    pub fn new(seed: u64) -> Arc<Self> {        use rand::prelude::*;         Arc::new(Self {            state: Arc::new(parking_lot::Mutex::new(DeterministicState {                rng: StdRng::seed_from_u64(seed),                seed,                scheduled_from_foreground: Default::default(),                scheduled_from_background: Default::default(),                forbid_parking: false,                block_on_ticks: 0..=1000,                now: std::time::Instant::now(),                next_timer_id: Default::default(),                pending_timers: Default::default(),                waiting_backtrace: None,                next_runnable_id: 0,                poll_history: Default::default(),                previous_poll_history: Default::default(),                enable_runnable_backtraces: false,                runnable_backtraces: Default::default(),            })),            parker: Default::default(),        })    }     pub fn execution_history(&self) -> Vec<ExecutorEvent> {        self.state.lock().poll_history.clone()    }     pub fn set_previous_execution_history(&self, history: Option<Vec<ExecutorEvent>>) {        self.state.lock().previous_poll_history = history;    }     pub fn enable_runnable_backtrace(&self) {        self.state.lock().enable_runnable_backtraces = true;    }     pub fn runnable_backtrace(&self, runnable_id: usize) -> backtrace::Backtrace {        let mut backtrace = self.state.lock().runnable_backtraces[&runnable_id].clone();        backtrace.resolve();        backtrace    }     pub fn build_background(self: &Arc<Self>) -> Arc<Background> {        Arc::new(Background::Deterministic {            executor: self.clone(),        })    }     pub fn build_foreground(self: &Arc<Self>, id: usize) -> Rc<Foreground> {        Rc::new(Foreground::Deterministic {            cx_id: id,            executor: self.clone(),        })    }     fn spawn_from_foreground(        &self,        cx_id: usize,        future: AnyLocalFuture,        main: bool,    ) -> AnyLocalTask {        let state = self.state.clone();        let id;        {            let mut state = state.lock();            id = util::post_inc(&mut state.next_runnable_id);            if state.enable_runnable_backtraces {                state                    .runnable_backtraces                    .insert(id, backtrace::Backtrace::new_unresolved());            }        }         let unparker = self.parker.lock().unparker();        let (runnable, task) = async_task::spawn_local(future, move |runnable| {            let mut state = state.lock();            state.push_to_history(ExecutorEvent::EnqueuRunnable { id });            state                .scheduled_from_foreground                .entry(cx_id)                .or_default()                .push(ForegroundRunnable { id, runnable, main });            unparker.unpark();        });        runnable.schedule();        task    }     fn spawn(&self, future: AnyFuture) -> AnyTask {        let state = self.state.clone();        let id;        {            let mut state = state.lock();            id = util::post_inc(&mut state.next_runnable_id);            if state.enable_runnable_backtraces {                state                    .runnable_backtraces                    .insert(id, backtrace::Backtrace::new_unresolved());            }        }         let unparker = self.parker.lock().unparker();        let (runnable, task) = async_task::spawn(future, move |runnable| {            let mut state = state.lock();            state                .poll_history                .push(ExecutorEvent::EnqueuRunnable { id });            state                .scheduled_from_background                .push(BackgroundRunnable { id, runnable });            unparker.unpark();        });        runnable.schedule();        task    }     fn run<'a>(        &self,        cx_id: usize,        main_future: Pin<Box<dyn 'a + Future<Output = Box<dyn Any>>>>,    ) -> Box<dyn Any> {        use std::sync::atomic::{AtomicBool, Ordering::SeqCst};         let woken = Arc::new(AtomicBool::new(false));         let state = self.state.clone();        let id;        {            let mut state = state.lock();            id = util::post_inc(&mut state.next_runnable_id);            if state.enable_runnable_backtraces {                state                    .runnable_backtraces                    .insert(id, backtrace::Backtrace::new_unresolved());            }        }         let unparker = self.parker.lock().unparker();        let (runnable, mut main_task) = unsafe {            async_task::spawn_unchecked(main_future, move |runnable| {                let state = &mut *state.lock();                state                    .scheduled_from_foreground                    .entry(cx_id)                    .or_default()                    .push(ForegroundRunnable {                        id: util::post_inc(&mut state.next_runnable_id),                        runnable,                        main: true,                    });                unparker.unpark();            })        };        runnable.schedule();         loop {            if let Some(result) = self.run_internal(woken.clone(), Some(&mut main_task)) {                return result;            }             if !woken.load(SeqCst) {                self.state.lock().will_park();            }             woken.store(false, SeqCst);            self.parker.lock().park();        }    }     pub fn run_until_parked(&self) {        use std::sync::atomic::AtomicBool;        let woken = Arc::new(AtomicBool::new(false));        self.run_internal(woken, None);    }     fn run_internal(        &self,        woken: Arc<std::sync::atomic::AtomicBool>,        mut main_task: Option<&mut AnyLocalTask>,    ) -> Option<Box<dyn Any>> {        use rand::prelude::*;        use std::sync::atomic::Ordering::SeqCst;         let unparker = self.parker.lock().unparker();        let waker = waker_fn::waker_fn(move || {            woken.store(true, SeqCst);            unparker.unpark();        });         let mut cx = Context::from_waker(&waker);        loop {            let mut state = self.state.lock();             if state.scheduled_from_foreground.is_empty()                && state.scheduled_from_background.is_empty()            {                if let Some(main_task) = main_task {                    if let Poll::Ready(result) = main_task.poll(&mut cx) {                        return Some(result);                    }                }                 return None;            }             if !state.scheduled_from_background.is_empty() && state.rng.gen() {                let background_len = state.scheduled_from_background.len();                let ix = state.rng.gen_range(0..background_len);                let background_runnable = state.scheduled_from_background.remove(ix);                state.push_to_history(ExecutorEvent::PollRunnable {                    id: background_runnable.id,                });                drop(state);                background_runnable.runnable.run();            } else if !state.scheduled_from_foreground.is_empty() {                let available_cx_ids = state                    .scheduled_from_foreground                    .keys()                    .copied()                    .collect::<Vec<_>>();                let cx_id_to_run = *available_cx_ids.iter().choose(&mut state.rng).unwrap();                let scheduled_from_cx = state                    .scheduled_from_foreground                    .get_mut(&cx_id_to_run)                    .unwrap();                let foreground_runnable = scheduled_from_cx.remove(0);                if scheduled_from_cx.is_empty() {                    state.scheduled_from_foreground.remove(&cx_id_to_run);                }                state.push_to_history(ExecutorEvent::PollRunnable {                    id: foreground_runnable.id,                });                 drop(state);                 foreground_runnable.runnable.run();                if let Some(main_task) = main_task.as_mut() {                    if foreground_runnable.main {                        if let Poll::Ready(result) = main_task.poll(&mut cx) {                            return Some(result);                        }                    }                }            }        }    }     fn block<F, T>(&self, future: &mut F, max_ticks: usize) -> Option<T>    where        F: Unpin + Future<Output = T>,    {        use rand::prelude::*;         let unparker = self.parker.lock().unparker();        let waker = waker_fn::waker_fn(move || {            unparker.unpark();        });         let mut cx = Context::from_waker(&waker);        for _ in 0..max_ticks {            let mut state = self.state.lock();            let runnable_count = state.scheduled_from_background.len();            let ix = state.rng.gen_range(0..=runnable_count);            if ix < state.scheduled_from_background.len() {                let background_runnable = state.scheduled_from_background.remove(ix);                state.push_to_history(ExecutorEvent::PollRunnable {                    id: background_runnable.id,                });                drop(state);                background_runnable.runnable.run();            } else {                drop(state);                if let Poll::Ready(result) = future.poll(&mut cx) {                    return Some(result);                }                let mut state = self.state.lock();                if state.scheduled_from_background.is_empty() {                    state.will_park();                    drop(state);                    self.parker.lock().park();                }                 continue;            }        }         None    }     pub fn timer(&self, duration: Duration) -> Timer {        let (tx, rx) = postage::barrier::channel();        let mut state = self.state.lock();        let wakeup_at = state.now + duration;        let id = util::post_inc(&mut state.next_timer_id);        match state            .pending_timers            .binary_search_by_key(&wakeup_at, |e| e.1)        {            Ok(ix) | Err(ix) => state.pending_timers.insert(ix, (id, wakeup_at, tx)),        }        let state = self.state.clone();        Timer::Deterministic(DeterministicTimer { rx, id, state })    }     pub fn now(&self) -> std::time::Instant {        let state = self.state.lock();        state.now    }     pub fn advance_clock(&self, duration: Duration) {        let new_now = self.state.lock().now + duration;        loop {            self.run_until_parked();            let mut state = self.state.lock();             if let Some((_, wakeup_time, _)) = state.pending_timers.first() {                let wakeup_time = *wakeup_time;                if wakeup_time <= new_now {                    let timer_count = state                        .pending_timers                        .iter()                        .take_while(|(_, t, _)| *t == wakeup_time)                        .count();                    state.now = wakeup_time;                    let timers_to_wake = state                        .pending_timers                        .drain(0..timer_count)                        .collect::<Vec<_>>();                    drop(state);                    drop(timers_to_wake);                    continue;                }            }             break;        }         self.state.lock().now = new_now;    }     pub fn start_waiting(&self) {        self.state.lock().waiting_backtrace = Some(backtrace::Backtrace::new_unresolved());    }     pub fn finish_waiting(&self) {        self.state.lock().waiting_backtrace.take();    }     pub fn forbid_parking(&self) {        use rand::prelude::*;         let mut state = self.state.lock();        state.forbid_parking = true;        state.rng = StdRng::seed_from_u64(state.seed);    }     pub fn allow_parking(&self) {        use rand::prelude::*;         let mut state = self.state.lock();        state.forbid_parking = false;        state.rng = StdRng::seed_from_u64(state.seed);    }     pub async fn simulate_random_delay(&self) {        use rand::prelude::*;        use smol::future::yield_now;        if self.state.lock().rng.gen_bool(0.2) {            let yields = self.state.lock().rng.gen_range(1..=10);            for _ in 0..yields {                yield_now().await;            }        }    } 

Lines 1–500 of 1031.