feat(backend) scheduler rewrite (deadlock for now)

This commit is contained in:
lumither committed 2025-03-05 13:52:07 -08:00
1 parent 0c6325d5f2
commit 5e8b1ff35d
8 files changed
+266 -138

No files matched your search

+6 -3
View File
@@ -3,7 +3,6 @@ mod cron;
use crate::{ use crate::{
modules::{ modules::{
Module,
blog::{ blog::{
api::{ api::{
get_page_count, get_page_count,
@@ -12,11 +11,12 @@ use crate::{
}, },
// cron::check_update, // cron::check_update,
}, },
Module,
}, },
scheduler::tasks::CronTask, scheduler::tasks::{async_basic::AsyncBasic, CronTask},
}; };
use axum::{Router, routing::get}; use axum::{routing::get, Router};
use sqlx::PgPool; use sqlx::PgPool;
const MOUNT_POINT: &str = "/blog"; const MOUNT_POINT: &str = "/blog";
@@ -47,6 +47,9 @@ impl Module for Blog {
// BasicTask::new(move || check_update::task(&db_handler), "0/15 * * * * *") // BasicTask::new(move || check_update::task(&db_handler), "0/15 * * * * *")
// .unwrap() // .unwrap()
// .to_task(), // .to_task(),
AsyncBasic::new(|| Box::pin(async { println!("test") }), "* * * * * *")
.unwrap()
.to_task(),
] ]
} }
} }
+5 -8
View File
@@ -5,11 +5,12 @@ use std::num::ParseIntError;
use crate::{ use crate::{
constants::config::server::{default_value, var_name}, constants::config::server::{default_value, var_name},
scheduler::{Scheduler, tasks::CronTask}, scheduler::{tasks::CronTask, Scheduler},
}; };
use macros::error_panic; use macros::error_panic;
use axum::Router; use axum::Router;
use tokio::net::TcpListener;
use tracing::{info, warn}; use tracing::{info, warn};
// todo: merge `get_server_router` & `get_mount_point`, make Option // todo: merge `get_server_router` & `get_mount_point`, make Option
@@ -51,13 +52,7 @@ impl ModuleTree {
Err(_) => default_value::PORT, Err(_) => default_value::PORT,
}; };
let listener = match tokio::net::TcpListener::bind(format!( let listener = match TcpListener::bind(format!("{}:{}", default_value::HOST, port)).await {
"{}:{}",
default_value::HOST,
port
))
.await
{
Ok(listener) => { Ok(listener) => {
info!("server started on {}:{}", default_value::HOST, port); info!("server started on {}:{}", default_value::HOST, port);
listener listener
@@ -77,6 +72,8 @@ impl ModuleTree {
app = app.nest(module.get_mount_point(), module.get_server_router()); app = app.nest(module.get_mount_point(), module.get_server_router());
} }
scheduler.start().unwrap();
axum::serve(listener, app).await.unwrap_or_else(|e| { axum::serve(listener, app).await.unwrap_or_else(|e| {
error_panic!("failed to start axum: {}", e.to_string()); error_panic!("failed to start axum: {}", e.to_string());
}); });
+47 -122
View File
@@ -1,172 +1,94 @@
use std::{ use std::{
collections::HashMap, collections::{HashMap, HashSet},
sync::{ sync::{Arc, RwLock},
atomic::{AtomicBool, Ordering}, Arc,
RwLock,
},
thread,
}; };
use crate::scheduler::{ use crate::scheduler::{
scheduler_error::{ runner::Runner,
SchedulerError, SchedulerError::IdAlreadyRunning, SchedulerError::IdNotFound, scheduler_error::{SchedulerError, SchedulerError::IdNotFound},
},
sleeper::SkippableSleeper,
tasks::CronTask, tasks::CronTask,
}; };
use tracing::log::error;
use uuid::Uuid; use uuid::Uuid;
mod runner;
pub mod scheduler_error; pub mod scheduler_error;
mod sleeper; mod sleeper;
mod task_func; mod task_func;
pub mod tasks; pub mod tasks;
struct TaskStatus {
sleeper: Option<Arc<SkippableSleeper>>,
is_alive: AtomicBool,
}
#[derive(Default)] #[derive(Default)]
pub struct Scheduler { pub struct Scheduler {
tasks: Arc<RwLock<HashMap<Uuid, Arc<Box<dyn CronTask>>>>>, running: Arc<RwLock<HashSet<Uuid>>>,
running: Arc<RwLock<HashMap<Uuid, Arc<TaskStatus>>>>, parking_pool: Arc<RwLock<HashMap<Uuid, Arc<Box<dyn CronTask>>>>>,
runner: Runner,
} }
impl Scheduler { impl Scheduler {
/// Start all tasks in the scheduler /// Start all tasks in the scheduler
pub fn start(&self) -> Result<(), SchedulerError> { pub fn start(&self) -> Result<(), SchedulerError> {
for (k, _v) in self.tasks.read()?.iter() { self.runner.run_all()
self.run_non_blocking(k)?
}
Ok(())
} }
/// Terminate all tasks in the scheduler /// Terminate all tasks in the scheduler
pub fn halt() {} pub fn halt(&self) -> Result<(), SchedulerError> {
for id in self.running.read()?.iter() {
fn run_task( let func = &self.runner.terminate_remove(*id)?;
running: Arc<RwLock<HashMap<Uuid, Arc<TaskStatus>>>>, self.running.write()?.remove(id);
id: Uuid, self.parking_pool.write()?.insert(*id, func.clone());
task: Arc<Box<dyn CronTask>>,
) -> Result<(), SchedulerError> {
if running.read()?.get(&id).is_some() {
return Err(IdAlreadyRunning);
} else {
running.write()?.insert(
id,
TaskStatus {
sleeper: None,
is_alive: AtomicBool::new(true),
}
.into(),
);
}
loop {
if let Some(next_execution) = task.get_next_execution() {
// guarantee is alive
let task_status = {
let running_reading_guard = running.read()?;
running_reading_guard.get(&id).cloned().clone()
};
if let Some(task_status) = task_status {
if !&task_status.is_alive.load(Ordering::SeqCst) {
{
let mut running_lock_guard = running.write()?;
running_lock_guard.remove(&id);
}
return Ok(());
} else {
let sleeper = Arc::from(SkippableSleeper::from(next_execution));
{
let mut running_lock_guard = running.write()?;
running_lock_guard.insert(
id,
TaskStatus {
sleeper: Some(sleeper.clone()),
is_alive: AtomicBool::new(true),
}
.into(),
);
}
sleeper.start();
task.call();
}
}
} else {
error!("unexpected error ");
break;
}
} }
Ok(()) Ok(())
} }
/// Run a task with `id` /// Run a task with `id`
pub fn run(&self, id: &Uuid) -> Result<(), SchedulerError> { pub fn run(&self, id: &Uuid) -> Result<(), SchedulerError> {
let id = id.to_owned(); if self.running.read()?.contains(id) {
if let Some(task) = self.tasks.read()?.get(&id) {
Self::run_task(self.running.clone(), id, task.clone())
} else {
Err(IdNotFound)
}
}
pub fn run_non_blocking(&self, id: &Uuid) -> Result<(), SchedulerError> {
let id = id.to_owned();
if let Some(task) = self.tasks.read()?.get(&id) {
let running_clone = self.running.clone();
let task_clone = task.clone();
thread::spawn(move || Self::run_task(running_clone, id, task_clone));
Ok(()) Ok(())
} else if self.parking_pool.read()?.contains_key(id) {
let func = self.parking_pool.write()?.remove(id).unwrap();
self.running.write()?.insert(*id);
self.runner.add_with_id(func, id)
} else { } else {
Err(IdNotFound) Err(IdNotFound)
} }
} }
fn get_running(&self, id: &Uuid) -> Option<Arc<TaskStatus>> {
let ret = self.running.read().unwrap().get(id).cloned();
ret
}
/// Stop a task with `id` /// Stop a task with `id`
pub fn stop(&self, id: &Uuid) -> Result<(), SchedulerError> { pub fn stop(&self, id: &Uuid) -> Result<(), SchedulerError> {
match self.get_running(id) { dbg!("stop() called");
None => Err(IdNotFound), if self.running.read()?.contains(id) {
Some(status) => { dbg!("b1");
status.is_alive.store(false, Ordering::SeqCst); let func = self.runner.terminate_remove(*id)?;
if let Some(s) = &status.sleeper { dbg!("b1a");
s.skip(); self.running.write()?.remove(id);
} else { dbg!("b1b");
return Err(IdAlreadyRunning); self.parking_pool.write()?.insert(*id, func.clone());
} dbg!("b1c");
Ok(()) Ok(())
} } else if self.parking_pool.read()?.get(id).is_some() {
dbg!("b2");
Ok(())
} else {
dbg!("b3");
Err(IdNotFound)
} }
} }
pub fn now(&self, id: &Uuid) -> Result<(), SchedulerError> { // pub fn now(&self, id: &Uuid) -> Result<(), SchedulerError> {
match self.get_running(id) { // }
None => Err(IdNotFound),
Some(status) => {
if let Some(s) = &status.sleeper {
s.skip()
} else {
return Err(IdAlreadyRunning);
}
Ok(())
}
}
}
pub fn insert(&self, task: Box<dyn CronTask>) -> Result<Uuid, SchedulerError> { pub fn insert(&self, task: Box<dyn CronTask>) -> Result<Uuid, SchedulerError> {
let id = Uuid::new_v4(); let id = self.runner.add(task.into())?;
self.tasks.write()?.insert(id, Arc::new(task)); self.running.write()?.insert(id);
Ok(id) Ok(id)
} }
pub fn insert_list(&self, tasks: Vec<Box<dyn CronTask>>) -> Result<Vec<Uuid>, SchedulerError> { pub fn insert_list(&self, tasks: Vec<Box<dyn CronTask>>) -> Result<Vec<Uuid>, SchedulerError> {
tasks.into_iter().map(|t| self.insert(t)).collect() let ids = self
.runner
.add_list(tasks.into_iter().map(|t| t.into()).collect())?;
self.running.write()?.extend(&ids);
Ok(ids)
} }
pub fn new() -> Self { pub fn new() -> Self {
@@ -187,11 +109,14 @@ mod tests {
let scheduler = Scheduler::default(); let scheduler = Scheduler::default();
let task = BasicTask::new(|| println!("test print"), "* * * * * *")?; let task = BasicTask::new(|| println!("test print"), "* * * * * *")?;
let id = scheduler.insert(Box::new(task))?; let id = scheduler.insert(Box::new(task))?;
scheduler.start()?;
for i in 1..=2 { for i in 1..=2 {
println!("Loop {} start", i); println!("Loop {} start", i);
println!("Run {}", i); println!("Run {}", i);
scheduler.run_non_blocking(&id)?; scheduler.run(&id)?;
println!("Sleep {}", i);
thread::sleep(Duration::from_secs(4)); thread::sleep(Duration::from_secs(4));
println!("Sleep {} finished", i);
scheduler.stop(&id)?; scheduler.stop(&id)?;
println!("Stop {}", i); println!("Stop {}", i);
thread::sleep(Duration::from_secs(2)); thread::sleep(Duration::from_secs(2));
+152
View File
@@ -0,0 +1,152 @@
use std::{
collections::HashMap,
sync::{
atomic::{AtomicBool, Ordering}, Arc,
RwLock,
},
};
use crate::scheduler::{
scheduler_error::{SchedulerError, SchedulerError::IdNotFound},
sleeper::SkippableSleeper,
tasks::CronTask,
};
use tokio::runtime::Runtime;
use uuid::Uuid;
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) {
loop {
if self.status.expected_running.load(Ordering::SeqCst) {
self.status.sleeper.read().unwrap().start();
self.func.call().await;
if let Some(next_execution) = self.func.get_next_execution() {
*(self.status.sleeper.write().unwrap()) = next_execution.into();
} else {
self.status.is_running.store(false, Ordering::SeqCst);
break;
}
} else {
break;
}
}
}
pub fn terminate(&self) {
self.status.expected_running.store(false, Ordering::SeqCst);
self.status.is_running.store(false, Ordering::SeqCst);
self.status.sleeper.read().unwrap().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 fn run_all(&self) -> Result<(), SchedulerError> {
let list = &self.task_list.read()?;
self.rt.block_on(async {
for (_, entity) in list.iter() {
if !entity.status.is_running.load(Ordering::SeqCst)
&& entity.status.expected_running.load(Ordering::SeqCst)
{
entity.status.is_running.store(true, Ordering::SeqCst);
tokio::spawn({
let entity_clone = entity.clone();
async move { entity_clone.run_blocking().await }
});
}
}
});
Ok(())
}
pub fn add_with_id(
&self,
task: Arc<Box<dyn CronTask>>,
id: &Uuid,
) -> Result<(), SchedulerError> {
self.task_list
.write()?
.insert(*id, TaskEntity::new(task).into());
Ok(())
}
pub fn add(&self, task: Arc<Box<dyn CronTask>>) -> Result<Uuid, SchedulerError> {
let id = Uuid::new_v4();
self.add_with_id(task, &id)?;
Ok(id)
}
pub fn add_list_with_id(
&self,
tasks: Vec<(&Uuid, Arc<Box<dyn CronTask>>)>,
) -> Result<(), SchedulerError> {
tasks
.into_iter()
.try_for_each(|(id, task)| self.add_with_id(task, id))
}
pub fn add_list(
&self,
tasks: Vec<Arc<Box<dyn CronTask>>>,
) -> Result<Vec<Uuid>, SchedulerError> {
tasks
.into_iter()
.map(|task| {
let id = Uuid::new_v4();
self.add_with_id(task, &id).map(move |_| id)
})
.collect::<Result<Vec<_>, _>>()
}
pub fn terminate_remove(&self, id: Uuid) -> Result<Arc<Box<dyn CronTask>>, SchedulerError> {
match self.task_list.read()?.get(&id) {
None => Err(IdNotFound),
Some(task) => {
task.terminate();
let func = &self.task_list.write()?.remove(&id).unwrap().func;
Ok(func.clone())
}
}
}
}
impl Default for Runner {
fn default() -> Self {
Self::new()
}
}
+12
View File
@@ -1,3 +1,15 @@
use std::pin::Pin;
pub trait TaskFunc: Fn() + Send + Sync + 'static {} pub trait TaskFunc: Fn() + Send + Sync + 'static {}
impl<T> TaskFunc for T where T: Fn() + Send + Sync + Clone + 'static {} impl<T> TaskFunc for T where T: Fn() + Send + Sync + Clone + 'static {}
pub trait AsyncTaskFunc:
Fn() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync + 'static
{
}
impl<T> AsyncTaskFunc for T where
T: Fn() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync + Clone + 'static
{
}
@@ -0,0 +1,37 @@
use std::{error::Error, pin::Pin, str::FromStr};
use crate::{scheduler::task_func::AsyncTaskFunc, scheduler::tasks::CronTask};
use chrono::{DateTime, Utc};
use cron::Schedule;
pub struct AsyncBasic {
schedule: Schedule,
func: Box<dyn AsyncTaskFunc>,
}
impl AsyncBasic {
pub fn new<F>(func: F, cron_expr: &str) -> Result<Self, Box<dyn Error>>
where
F: AsyncTaskFunc,
{
Ok(AsyncBasic {
schedule: Schedule::from_str(cron_expr)?,
func: Box::new(func),
})
}
pub fn to_task(self) -> Box<dyn CronTask> {
Box::new(self)
}
}
impl CronTask for AsyncBasic {
fn get_next_execution(&self) -> Option<DateTime<Utc>> {
self.schedule.upcoming(Utc).next()
}
fn call(&self) -> Pin<Box<(dyn Future<Output = ()> + Send)>> {
Box::pin((self.func)())
}
}
+3 -4
View File
@@ -1,5 +1,4 @@
use std::error::Error; use std::{error::Error, pin::Pin, str::FromStr};
use std::str::FromStr;
use crate::scheduler::{task_func::TaskFunc, tasks::CronTask}; use crate::scheduler::{task_func::TaskFunc, tasks::CronTask};
@@ -32,7 +31,7 @@ impl CronTask for BasicTask {
self.schedule.upcoming(Utc).next() self.schedule.upcoming(Utc).next()
} }
fn call(&self) { fn call(&self) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
(self.func)() Box::pin(async { (self.func)() })
} }
} }
+4 -1
View File
@@ -1,8 +1,11 @@
use std::pin::Pin;
use chrono::{DateTime, Utc}; use chrono::{DateTime, Utc};
pub mod async_basic;
pub mod basic; pub mod basic;
pub trait CronTask: Send + Sync { pub trait CronTask: Send + Sync {
fn get_next_execution(&self) -> Option<DateTime<Utc>>; fn get_next_execution(&self) -> Option<DateTime<Utc>>;
fn call(&self); fn call(&self) -> Pin<Box<dyn Future<Output = ()> + Send + '_>>;
} }