gwr_engine/

executor.rs

// Copyright (c) 2023 Graphcore Ltd. All rights reserved.

use std::cell::{Cell, RefCell};
use std::future::Future;
use std::mem;
use std::pin::Pin;
use std::rc::Rc;
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};

use gwr_track::entity::Entity;
use rand::SeedableRng;
use rand::rngs::StdRng;
use rand::seq::SliceRandom;

use crate::time::clock::Clock;
use crate::time::simtime::SimTime;
use crate::types::SimResult;

fn no_op(_: *const ()) {}

unsafe fn drop_task(data: *const ()) {
    unsafe {
        drop(Rc::from_raw(data as *const Task));
    }
}

static VTABLE: RawWakerVTable = RawWakerVTable::new(clone_raw_waker, wake_task, no_op, drop_task);

fn task_raw_waker(task: Rc<Task>) -> RawWaker {
    let ptr = Rc::into_raw(task) as *const ();
    RawWaker::new(ptr, &VTABLE)
}

fn waker_for_task(task: Rc<Task>) -> Waker {
    unsafe { Waker::from_raw(task_raw_waker(task)) }
}

unsafe fn clone_raw_waker(data: *const ()) -> RawWaker {
    unsafe {
        // Tasks are always wrapped in a reference counter to allow them to be shared
        // read-only. The input `data` pointer is borrowed — we must not decrement its
        // refcount, so we mem::forget the reconstructed Rc rather than letting it drop.
        let rc_task = Rc::from_raw(data as *const Task);
        let clone = rc_task.clone();
        mem::forget(rc_task);
        let ptr = Rc::into_raw(clone) as *const ();
        RawWaker::new(ptr, &VTABLE)
    }
}

unsafe fn wake_task(data: *const ()) {
    unsafe {
        // Tasks are always wrapped in a reference counter to allow them to be shared
        // read-only.
        let rc_task = Rc::from_raw(data as *const Task);
        let cloned = rc_task.clone();
        rc_task.executor_state.new_tasks.borrow_mut().push(cloned);
    }
}

struct Task {
    future: RefCell<Option<Pin<Box<dyn Future<Output = SimResult>>>>>,
    executor_state: Rc<ExecutorState>,
}

impl Task {
    pub fn new(
        future: impl Future<Output = SimResult> + 'static,
        executor_state: Rc<ExecutorState>,
    ) -> Task {
        Task {
            future: RefCell::new(Some(Box::pin(future))),
            executor_state,
        }
    }

    fn poll(&self, context: &mut Context) -> Poll<SimResult> {
        let mut future_slot = self.future.borrow_mut();
        let Some(future) = future_slot.as_mut() else {
            return Poll::Ready(Ok(()));
        };

        let poll_result = future.as_mut().poll(context);
        if poll_result.is_ready() {
            future_slot.take();
        }

        poll_result
    }
}

struct ExecutorState {
    task_queue: RefCell<Vec<Rc<Task>>>,
    new_tasks: RefCell<Vec<Rc<Task>>>,
    time: RefCell<SimTime>,
    randomize_task_order: Cell<bool>,
    task_order_rng: RefCell<StdRng>,
}

impl ExecutorState {
    pub fn new(top: &Rc<Entity>) -> Self {
        Self {
            task_queue: RefCell::new(Vec::new()),
            new_tasks: RefCell::new(Vec::new()),
            time: RefCell::new(SimTime::new(top)),
            randomize_task_order: Cell::new(false),
            task_order_rng: RefCell::new(StdRng::seed_from_u64(rand::random())),
        }
    }
}

/// Single-threaded executor
///
/// This is a thin-wrapper (using [`Rc`]) around the real executor, so that this
/// struct can be cloned and passed around.
///
/// See the [module documentation] for more details.
///
/// [module documentation]: index.html
#[derive(Clone)]
pub struct Executor {
    state: Rc<ExecutorState>,
}

impl Executor {
    pub fn run(&self, finished: &Rc<RefCell<bool>>) -> SimResult {
        loop {
            self.step(finished)?;
            if *finished.borrow() {
                break;
            }

            if self.state.new_tasks.borrow().is_empty() {
                if self.state.time.borrow().can_exit() {
                    break;
                }

                if let Some(wakers) = self.state.time.borrow_mut().advance_time() {
                    // No events left, advance time
                    for task_waker in wakers.into_iter() {
                        task_waker.waker.wake();
                    }
                } else {
                    break;
                }
            }
        }
        Ok(())
    }

    pub fn step(&self, finished: &Rc<RefCell<bool>>) -> SimResult {
        // Append new tasks created since the last step into the task queue
        let mut task_queue = self.state.task_queue.borrow_mut();
        task_queue.append(&mut self.state.new_tasks.borrow_mut());
        if self.state.randomize_task_order.get() {
            task_queue.shuffle(&mut *self.state.task_order_rng.borrow_mut());
        }

        // Loop over all tasks, polling them. If a task is not ready, add it to the
        // pending tasks.
        for task in task_queue.drain(..) {
            if *finished.borrow() {
                break;
            }

            // Dummy waker and context (not used as we poll all tasks)
            let waker = waker_for_task(task.clone());
            let mut context = Context::from_waker(&waker);

            match task.poll(&mut context) {
                Poll::Ready(Err(e)) => {
                    // Error - return early
                    return Err(e);
                }
                Poll::Ready(Ok(())) => {
                    // Otherwise, drop task as it is complete
                }
                Poll::Pending => {
                    // Task will have parked itself waiting somewhere
                }
            }
        }
        Ok(())
    }

    #[must_use]
    pub fn get_clock(&self, freq_mhz: f64) -> Clock {
        self.state.time.borrow_mut().get_clock(freq_mhz)
    }

    #[must_use]
    pub fn time_now_ns(&self) -> f64 {
        self.state.time.borrow().time_now_ns()
    }

    pub fn set_randomize_task_order(&self, randomize: bool) {
        self.state.randomize_task_order.set(randomize);
    }

    pub fn set_task_order_seed(&self, seed: u64) {
        *self.state.task_order_rng.borrow_mut() = StdRng::seed_from_u64(seed);
    }
}

/// `Spawner` spawns new futures into the executor.
#[derive(Clone)]
pub struct Spawner {
    state: Rc<ExecutorState>,
}

impl Spawner {
    pub fn spawn(&self, future: impl Future<Output = SimResult> + 'static) {
        self.state
            .new_tasks
            .borrow_mut()
            .push(Rc::new(Task::new(future, self.state.clone())));
    }
}

#[must_use]
pub fn new_executor_and_spawner(top: &Rc<Entity>) -> (Executor, Spawner) {
    let state = Rc::new(ExecutorState::new(top));
    (
        Executor {
            state: state.clone(),
        },
        Spawner { state },
    )
}