mirror of
https://github.com/Lumither/nocturne.git
synced 2026-10-17 18:36:21 +00:00
feat(scheduler) remove terminate() etc. impl
This commit is contained in:
10 files changed
+91
-409
No files matched your search
@@ -7,6 +7,7 @@ use std::{
|
||||
use crate::scheduler::error::SchedulerError::MutexPoison;
|
||||
|
||||
// todo: refactor with thiserror
|
||||
#[allow(dead_code)]
|
||||
pub enum SchedulerError {
|
||||
IdNotFound,
|
||||
IdAlreadyRunning,
|
||||
|
||||
@@ -1,145 +1,56 @@
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
sync::Arc,
|
||||
};
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use crate::scheduler::{
|
||||
error::{SchedulerError, SchedulerError::IdNotFound},
|
||||
runner::Runner,
|
||||
tasks::CronTask,
|
||||
};
|
||||
use crate::scheduler::{error::SchedulerError, tasks::CronTask, utils::execute};
|
||||
|
||||
use tokio::sync::RwLock;
|
||||
use uuid::Uuid;
|
||||
use futures::future::join_all;
|
||||
use tracing::warn;
|
||||
|
||||
pub mod error;
|
||||
mod runner;
|
||||
mod sleeper;
|
||||
pub mod task_func;
|
||||
pub mod tasks;
|
||||
mod utils;
|
||||
|
||||
pub type TaskPool = HashMap<String, Arc<Box<dyn CronTask>>>;
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct Scheduler {
|
||||
running: RwLock<HashSet<Uuid>>,
|
||||
parking_pool: RwLock<HashMap<Uuid, Arc<Box<dyn CronTask>>>>,
|
||||
runner: Runner,
|
||||
task_pool: TaskPool,
|
||||
}
|
||||
|
||||
impl Scheduler {
|
||||
/// Start all tasks in the scheduler
|
||||
pub async fn start(&self) -> Result<(), SchedulerError> {
|
||||
self.runner.run_all().await
|
||||
pub async fn start(&self) {
|
||||
dbg!("entering");
|
||||
|
||||
let fut = self
|
||||
.task_pool
|
||||
.clone()
|
||||
.into_iter()
|
||||
.map(|(name, task)| execute(name, task))
|
||||
.collect::<Vec<_>>();
|
||||
// self.runner.run_all().await
|
||||
join_all(fut).await;
|
||||
}
|
||||
|
||||
/// Terminate all tasks in the scheduler
|
||||
pub async fn halt(&self) -> Result<(), SchedulerError> {
|
||||
for id in self.running.read().await.iter() {
|
||||
let func = &self.runner.terminate_remove(*id).await?;
|
||||
self.running.write().await.remove(id);
|
||||
self.parking_pool.write().await.insert(*id, func.clone());
|
||||
pub fn insert(&mut self, task: (&str, Box<dyn CronTask>)) -> Result<&Self, SchedulerError> {
|
||||
let (k, v) = task;
|
||||
if self.task_pool.insert(k.into(), v.into()).is_some() {
|
||||
warn!("task list duplicated key detected: {}", k);
|
||||
}
|
||||
Ok(())
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// Run a task with `id`
|
||||
pub async fn run(&self, id: &Uuid) -> Result<(), SchedulerError> {
|
||||
if self.running.read().await.contains(id) {
|
||||
Ok(())
|
||||
} else if self.parking_pool.read().await.contains_key(id) {
|
||||
let func = self.parking_pool.write().await.remove(id).unwrap();
|
||||
self.running.write().await.insert(*id);
|
||||
self.runner.add_with_id(func, id).await?;
|
||||
self.runner.run_id(id).await
|
||||
} else {
|
||||
Err(IdNotFound)
|
||||
pub fn insert_list(
|
||||
&mut self,
|
||||
tasks: Vec<(&str, Box<dyn CronTask>)>,
|
||||
) -> Result<&Self, SchedulerError> {
|
||||
for task in tasks {
|
||||
self.insert(task)?;
|
||||
}
|
||||
}
|
||||
|
||||
/// Stop a task with `id`
|
||||
pub async fn stop(&self, id: &Uuid) -> Result<(), SchedulerError> {
|
||||
if self.running.read().await.contains(id) {
|
||||
let func = self.runner.terminate_remove(*id).await?;
|
||||
self.running.write().await.remove(id);
|
||||
self.parking_pool.write().await.insert(*id, func.clone());
|
||||
Ok(())
|
||||
} else if self.parking_pool.read().await.get(id).is_some() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(IdNotFound)
|
||||
}
|
||||
}
|
||||
|
||||
// pub fn now(&self, id: &Uuid) -> Result<(), SchedulerError> {
|
||||
// }
|
||||
|
||||
pub async fn insert(&self, task: Box<dyn CronTask>) -> Result<Uuid, SchedulerError> {
|
||||
let id = self.runner.add(task.into()).await?;
|
||||
self.running.write().await.insert(id);
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
pub async fn insert_list(
|
||||
&self,
|
||||
tasks: Vec<Box<dyn CronTask>>,
|
||||
) -> Result<Vec<Uuid>, SchedulerError> {
|
||||
let ids = self
|
||||
.runner
|
||||
.add_list(tasks.into_iter().map(|t| t.into()).collect())
|
||||
.await?;
|
||||
self.running.write().await.extend(&ids);
|
||||
Ok(ids)
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
// #[cfg(test)]
|
||||
// mod tests {
|
||||
// use std::{error::Error, thread, time::Duration};
|
||||
//
|
||||
// use crate::scheduler::{
|
||||
// tasks::{async_basic::AsyncBasic, basic::BasicTask},
|
||||
// Scheduler,
|
||||
// };
|
||||
//
|
||||
// use uuid::Uuid;
|
||||
//
|
||||
// async fn execution(scheduler: Scheduler, id: Uuid) -> Result<(), Box<dyn Error>> {
|
||||
// for i in 1..=2 {
|
||||
// println!("Loop {} start", i);
|
||||
// println!("Run {}", i);
|
||||
// scheduler.run(&id).await?;
|
||||
// println!("Sleep {}", i);
|
||||
// thread::sleep(Duration::from_secs(4));
|
||||
// println!("Sleep {} finished", i);
|
||||
// scheduler.stop(&id)?;
|
||||
// println!("Stop {}", i);
|
||||
// thread::sleep(Duration::from_secs(2));
|
||||
// println!("Loop {} finished", i);
|
||||
// }
|
||||
// Ok(())
|
||||
// }
|
||||
//
|
||||
// #[tokio::test]
|
||||
// async fn test_run() -> Result<(), Box<dyn Error>> {
|
||||
// let scheduler = Scheduler::default();
|
||||
// let task = BasicTask::new(|| println!("test print"), "* * * * * *")?;
|
||||
// let id = scheduler.insert(Box::new(task)).await?;
|
||||
// scheduler.start().await?;
|
||||
// execution(scheduler, id)
|
||||
// }
|
||||
//
|
||||
// #[tokio::test]
|
||||
// async fn test_run_async() -> Result<(), Box<dyn Error>> {
|
||||
// let scheduler = Scheduler::default();
|
||||
// let task = AsyncBasic::new(
|
||||
// || Box::pin(async { println!("test async print") }),
|
||||
// "* * * * * *",
|
||||
// )?;
|
||||
// let id = scheduler.insert(Box::new(task)).await?;
|
||||
// scheduler.start().await?;
|
||||
// execution(scheduler, id)
|
||||
// }
|
||||
// }
|
||||
@@ -1,186 +0,0 @@
|
||||
use crate::scheduler::{
|
||||
error::{SchedulerError, SchedulerError::IdNotFound},
|
||||
sleeper::SkippableSleeper,
|
||||
tasks::CronTask,
|
||||
};
|
||||
use futures::future;
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
},
|
||||
};
|
||||
|
||||
use tokio::runtime::Runtime;
|
||||
use tokio::sync::RwLock;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct TaskStatus {
|
||||
sleeper: RwLock<SkippableSleeper>,
|
||||
is_running: AtomicBool,
|
||||
expected_running: AtomicBool,
|
||||
}
|
||||
|
||||
pub struct TaskEntity {
|
||||
func: Arc<Box<dyn CronTask>>,
|
||||
status: TaskStatus,
|
||||
}
|
||||
|
||||
impl TaskEntity {
|
||||
pub fn new(func: Arc<Box<dyn CronTask>>) -> Self {
|
||||
let next_execution: SkippableSleeper = func.get_next_execution().unwrap().into();
|
||||
Self {
|
||||
func,
|
||||
status: TaskStatus {
|
||||
sleeper: next_execution.into(),
|
||||
is_running: false.into(),
|
||||
expected_running: true.into(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_blocking(&self) {
|
||||
dbg!("entered");
|
||||
dbg!("chkpt 1");
|
||||
loop {
|
||||
dbg!("chkpt 2");
|
||||
// deadlock spotted
|
||||
self.status.sleeper.read().await.start().await;
|
||||
dbg!("chkpt 3");
|
||||
if self.status.expected_running.load(Ordering::SeqCst) {
|
||||
dbg!("chkpt 4");
|
||||
self.func.call().await;
|
||||
dbg!("chkpt 5");
|
||||
if let Some(next_execution) = self.func.get_next_execution() {
|
||||
dbg!("chkpt 6a");
|
||||
*(self.status.sleeper.write().await) = next_execution.into();
|
||||
} else {
|
||||
dbg!("chkpt 6b");
|
||||
self.status.is_running.store(false, Ordering::SeqCst);
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn terminate(&self) {
|
||||
self.status.expected_running.store(false, Ordering::SeqCst);
|
||||
self.status.is_running.store(false, Ordering::SeqCst);
|
||||
self.status.sleeper.read().await.skip();
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Runner {
|
||||
task_list: RwLock<HashMap<Uuid, Arc<TaskEntity>>>,
|
||||
rt: Runtime,
|
||||
}
|
||||
|
||||
impl Runner {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
task_list: HashMap::new().into(),
|
||||
rt: Runtime::new().unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_all(&self) -> Result<(), SchedulerError> {
|
||||
let list = &self.task_list.read().await;
|
||||
self.rt.block_on(async {
|
||||
for (_, entity) in list.iter() {
|
||||
self.run_entity(entity.clone()).await
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn run_id(&self, id: &Uuid) -> Result<(), SchedulerError> {
|
||||
if let Some(entity) = self.task_list.read().await.get(id).cloned() {
|
||||
self.rt.block_on(async { self.run_entity(entity).await });
|
||||
Ok(())
|
||||
} else {
|
||||
Err(IdNotFound)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_entity(&self, entity: Arc<TaskEntity>) {
|
||||
if !entity.status.is_running.load(Ordering::SeqCst)
|
||||
&& entity.status.expected_running.load(Ordering::SeqCst)
|
||||
{
|
||||
entity.status.is_running.store(true, Ordering::SeqCst);
|
||||
let fut = {
|
||||
let entity_clone = entity.clone();
|
||||
async move { entity_clone.run_blocking().await }
|
||||
};
|
||||
dbg!("spawn");
|
||||
tokio::spawn(fut);
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn add_with_id(
|
||||
&self,
|
||||
task: Arc<Box<dyn CronTask>>,
|
||||
id: &Uuid,
|
||||
) -> Result<(), SchedulerError> {
|
||||
self.task_list
|
||||
.write()
|
||||
.await
|
||||
.insert(*id, TaskEntity::new(task).into());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn add(&self, task: Arc<Box<dyn CronTask>>) -> Result<Uuid, SchedulerError> {
|
||||
let id = Uuid::new_v4();
|
||||
self.add_with_id(task, &id).await?;
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
pub async fn add_list_with_id(
|
||||
&self,
|
||||
tasks: Vec<(&Uuid, Arc<Box<dyn CronTask>>)>,
|
||||
) -> Result<(), SchedulerError> {
|
||||
let fut = tasks
|
||||
.into_iter()
|
||||
.map(async |(id, entity)| self.add_with_id(entity, id).await)
|
||||
.collect::<Vec<_>>();
|
||||
future::join_all(fut).await.into_iter().collect()
|
||||
}
|
||||
|
||||
pub async fn add_list(
|
||||
&self,
|
||||
tasks: Vec<Arc<Box<dyn CronTask>>>,
|
||||
) -> Result<Vec<Uuid>, SchedulerError> {
|
||||
let fut = tasks
|
||||
.into_iter()
|
||||
.map(async |task| {
|
||||
let id = Uuid::new_v4();
|
||||
self.add_with_id(task, &id).await.map(move |_| id)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
future::join_all(fut).await.into_iter().collect()
|
||||
}
|
||||
|
||||
pub async fn terminate_remove(
|
||||
&self,
|
||||
id: Uuid,
|
||||
) -> Result<Arc<Box<dyn CronTask>>, SchedulerError> {
|
||||
let mut list = self.task_list.write().await;
|
||||
match list.get(&id) {
|
||||
None => Err(IdNotFound),
|
||||
Some(task) => {
|
||||
task.terminate().await;
|
||||
let func = list.remove(&id).unwrap().func.clone();
|
||||
Ok(func.clone())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Runner {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
@@ -1,81 +0,0 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::{sync::Arc, time::Duration};
|
||||
use tokio::sync::Notify;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SkippableSleeper {
|
||||
duration: Duration,
|
||||
skipped: Arc<AtomicBool>,
|
||||
notify: Arc<Notify>,
|
||||
}
|
||||
|
||||
impl SkippableSleeper {
|
||||
pub fn new(duration: Duration) -> Self {
|
||||
Self {
|
||||
duration,
|
||||
skipped: Arc::new(AtomicBool::new(false)),
|
||||
notify: Arc::new(Notify::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn start(&self) {
|
||||
if self.skipped.load(Ordering::Acquire) {
|
||||
return;
|
||||
}
|
||||
|
||||
let sleep = tokio::time::sleep(self.duration);
|
||||
let notify = self.notify.notified();
|
||||
|
||||
tokio::select! {
|
||||
_ = sleep => {}
|
||||
_ = notify => {
|
||||
self.skipped.store(true, Ordering::Release);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn skip(&self) {
|
||||
self.notify.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
impl From<DateTime<Utc>> for SkippableSleeper {
|
||||
fn from(value: DateTime<Utc>) -> Self {
|
||||
let now = Utc::now();
|
||||
let duration = value.signed_duration_since(now);
|
||||
let duration = duration.to_std().unwrap_or(Duration::ZERO);
|
||||
Self::new(duration)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::time::{sleep, Duration};
|
||||
|
||||
#[tokio::test]
|
||||
async fn skip_sleeper() {
|
||||
let sleeper = SkippableSleeper::new(Duration::from_secs(10));
|
||||
let sleeper_clone = sleeper.clone();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
println!("Sleeping...");
|
||||
sleeper_clone.start().await;
|
||||
println!("Woke up!");
|
||||
});
|
||||
|
||||
println!("Start waiting...");
|
||||
sleep(Duration::from_secs(4)).await;
|
||||
println!("Skipping sleep...");
|
||||
sleeper.skip();
|
||||
|
||||
handle.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn normal_sleeper() {
|
||||
let sleeper = SkippableSleeper::new(Duration::from_secs(4));
|
||||
sleeper.start().await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use crate::scheduler::{error::SchedulerError, tasks::CronTask};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use tokio::time::{sleep, Sleep};
|
||||
|
||||
pub fn duration_from_now(future: DateTime<Utc>) -> Duration {
|
||||
future
|
||||
.signed_duration_since(Utc::now())
|
||||
.to_std()
|
||||
.unwrap_or(Duration::ZERO)
|
||||
}
|
||||
|
||||
pub fn sleeper_from_now(future: DateTime<Utc>) -> Sleep {
|
||||
sleep(duration_from_now(future))
|
||||
}
|
||||
|
||||
#[allow(unused_variables)]
|
||||
pub async fn execute(name: String, func: Arc<Box<dyn CronTask>>) -> Result<(), SchedulerError> {
|
||||
// todo: stats
|
||||
while let Some(next_execution) = func.get_next_execution() {
|
||||
let sleeper = sleeper_from_now(next_execution);
|
||||
tokio::pin!(sleeper);
|
||||
tokio::select! {
|
||||
_ = sleeper => {
|
||||
func.call().await
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Reference in new issue
Block a user