ProjectNautilus/flight/src/scheduler/mod.rs

use anyhow::Result;
use log::trace;
use nautilus_common::on_drop::on_drop;
use std::any::type_name;
use std::fmt::Debug;
use std::ops::Deref;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{Receiver, Sender, channel};
use std::thread;
use std::thread::{Scope, sleep};
use std::time::{Duration, Instant};

#[derive(Clone, Debug)]
pub struct TaskHandle<Message, Data> {
    pub sender: Sender<Message>,
    data: Data,
}

impl<Message, Data> Deref for TaskHandle<Message, Data> {
    type Target = Data;

    fn deref(&self) -> &Self::Target {
        &self.data
    }
}

#[allow(dead_code)]
pub trait Task {
    type Message;
    type Data;

    fn get_data(&self) -> Self::Data;
    fn run(self, receiver: Receiver<Self::Message>, running: Arc<AtomicBool>);
}

pub trait CyclicTask {
    type Message;
    type Data;

    fn get_data(&self) -> Self::Data;
    fn step(&mut self, receiver: &Receiver<Self::Message>, step_time: Instant);
}

impl<F> CyclicTask for F
where
    F: FnMut(),
{
    type Message = ();
    type Data = ();

    fn get_data(&self) -> Self::Data {}

    fn step(&mut self, _receiver: &Receiver<Self::Message>, _step_time: Instant) {
        self();
    }
}

pub struct Scheduler<'s, 'e>
where
    'e: 's,
{
    scope: &'s Scope<'s, 'e>,
    running: Arc<AtomicBool>,
}

impl<'s> Scheduler<'s, '_> {
    pub fn scope<'env, F, R>(running: Arc<AtomicBool>, f: F) -> R
    where
        F: FnOnce(Scheduler<'_, 'env>) -> R,
    {
        trace!("Scheduler::scope(running: {running:?}, f)");
        thread::scope(|scope: &Scope| {
            let running_result = running.clone();
            // This will automatically set running to false when it drops
            // This means that if the function returns any side branches
            // checking running will shut down
            let _shutdown_threads = on_drop(move || running.store(false, Ordering::Relaxed));

            f(Scheduler {
                scope,
                running: running_result,
            })
        })
    }

    // pub fn inner_scope<'env, F, R>(&self, f: F) -> R
    // where
    //     F: FnOnce(Scheduler<'_, 'env>) -> R,
    //     's: 'env,
    // {
    //     trace!("Scheduler::inner_scope(self)");
    //     thread::scope(|scope: &Scope<'_, 'env>| {
    //         f(Scheduler {
    //             scope,
    //             running: self.running.clone(),
    //         })
    //     })
    // }

    #[allow(dead_code)]
    pub fn run<T>(
        &self,
        name: impl Into<String>,
        task: T,
    ) -> Result<TaskHandle<T::Message, T::Data>>
    where
        T: Task + Send + Debug + 's,
        T::Message: Send,
    {
        let name = name.into();
        trace!(
            "Scheduler::run<T={}>(name: {name}, task: {task:?})",
            type_name::<T>()
        );
        let running = self.running.clone();
        let (sender, receiver) = channel::<T::Message>();
        let data = task.get_data();
        let _ = thread::Builder::new()
            .name(name.clone())
            .spawn_scoped(self.scope, move || {
                task.run(receiver, running);
            })?;
        Ok(TaskHandle { sender, data })
    }

    pub fn run_cyclic<T>(
        &self,
        name: impl Into<String>,
        mut task: T,
        frequency: u64,
    ) -> Result<TaskHandle<T::Message, T::Data>>
    where
        T: CyclicTask + Send + 's,
        T::Message: Send,
    {
        let name = name.into();
        trace!(
            "Scheduler::run_cyclic<T={}>(name: {name}, task, frequency: {frequency})",
            type_name::<T>()
        );
        let running = self.running.clone();
        let (sender, receiver) = channel::<T::Message>();
        let data = task.get_data();
        let _ = thread::Builder::new()
            .name(name.clone())
            .spawn_scoped(self.scope, move || {
                let period = Duration::from_nanos(1_000_000_000 / frequency);
                let mut cycle_start_time = Instant::now();
                while running.load(Ordering::Relaxed) {
                    task.step(&receiver, cycle_start_time);

                    cycle_start_time += period;
                    let sleep_duration = cycle_start_time - Instant::now();
                    if sleep_duration > Duration::ZERO {
                        sleep(sleep_duration);
                    }
                }
            })?;
        Ok(TaskHandle { sender, data })
    }
}