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 { pub sender: Sender, data: Data, } impl Deref for TaskHandle { 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, running: Arc); } pub trait CyclicTask { type Message; type Data; fn get_data(&self) -> Self::Data; fn step(&mut self, receiver: &Receiver, step_time: Instant); } impl CyclicTask for F where F: FnMut(), { type Message = (); type Data = (); fn get_data(&self) -> Self::Data {} fn step(&mut self, _receiver: &Receiver, _step_time: Instant) { self(); } } pub struct Scheduler<'s, 'e> where 'e: 's, { scope: &'s Scope<'s, 'e>, running: Arc, } impl<'s> Scheduler<'s, '_> { pub fn scope<'env, F, R>(running: Arc, 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( &self, name: impl Into, task: T, ) -> Result> where T: Task + Send + Debug + 's, T::Message: Send, { let name = name.into(); trace!( "Scheduler::run(name: {name}, task: {task:?})", type_name::() ); let running = self.running.clone(); let (sender, receiver) = channel::(); 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( &self, name: impl Into, mut task: T, frequency: u64, ) -> Result> where T: CyclicTask + Send + 's, T::Message: Send, { let name = name.into(); trace!( "Scheduler::run_cyclic(name: {name}, task, frequency: {frequency})", type_name::() ); let running = self.running.clone(); let (sender, receiver) = channel::(); 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 }) } }