120 lines
3.5 KiB
Rust
120 lines
3.5 KiB
Rust
use std::collections::HashMap;
|
|
use tokio_cron_scheduler::{Job, JobScheduler};
|
|
use tracing::{error, info, instrument};
|
|
use uuid::Uuid;
|
|
|
|
#[derive(Clone)]
|
|
pub struct ScheduleJob {
|
|
pub job_id: HashMap<(u64, u64), Uuid>,
|
|
pub scheduler: JobScheduler,
|
|
}
|
|
|
|
impl ScheduleJob {
|
|
#[instrument()]
|
|
pub async fn start_cron_scheduler() -> Result<Self, ()> {
|
|
let scheduler = JobScheduler::new().await;
|
|
let mut future_self = match scheduler {
|
|
Ok(scheduler) => ScheduleJob {
|
|
job_id: Default::default(),
|
|
scheduler,
|
|
},
|
|
Err(e) => {
|
|
error!("Error starting cron scheduler: {:?}", e);
|
|
return Err(());
|
|
}
|
|
};
|
|
future_self.scheduler.set_shutdown_handler(Box::new(|| {
|
|
Box::pin(async {
|
|
info!("Cron scheduler stopped");
|
|
})
|
|
}));
|
|
future_self.scheduler.shutdown_on_ctrl_c();
|
|
match future_self.scheduler.start().await {
|
|
Ok(_) => {
|
|
info!("Cron scheduler started");
|
|
Ok(future_self)
|
|
}
|
|
Err(e) => {
|
|
error!("Error starting cron scheduler: {:?}", e);
|
|
Err(())
|
|
}
|
|
}
|
|
}
|
|
|
|
#[instrument(skip(self))]
|
|
pub async fn stop_cron_scheduler(&mut self) -> &mut Self {
|
|
for (server_id, channel_id) in self.job_id.keys() {
|
|
match self
|
|
.scheduler
|
|
.remove(self.job_id.get(&(*server_id, *channel_id)).unwrap())
|
|
.await
|
|
{
|
|
Ok(_) => {
|
|
info!("Cron job removed");
|
|
}
|
|
Err(e) => {
|
|
error!("Error removing cron job: {:?}", e);
|
|
}
|
|
}
|
|
}
|
|
match self.scheduler.shutdown().await {
|
|
Ok(_) => {
|
|
info!("Cron scheduler stopped");
|
|
self
|
|
}
|
|
Err(e) => {
|
|
error!("Error stopping cron scheduler: {:?}", e);
|
|
self
|
|
}
|
|
}
|
|
}
|
|
|
|
#[instrument(skip(self))]
|
|
pub async fn add_cron_job(
|
|
&mut self,
|
|
server_id: u64,
|
|
channel_id: u64,
|
|
cron_expression: String,
|
|
) -> Result<Uuid, ()> {
|
|
let job = match Job::new_async(cron_expression.as_str(), |uuid, _l| {
|
|
Box::pin(async move {
|
|
info!("Cron job fired: {:?}", uuid);
|
|
})
|
|
}) {
|
|
Ok(job) => job,
|
|
Err(e) => {
|
|
error!("Error creating cron job: {:?}", e);
|
|
return Err(());
|
|
}
|
|
};
|
|
|
|
match self.scheduler.add(job).await {
|
|
Ok(job_uid) => {
|
|
info!("Cron job added");
|
|
self.job_id.insert((server_id, channel_id), job_uid);
|
|
Ok(job_uid)
|
|
}
|
|
Err(e) => {
|
|
error!("Error adding cron job: {:?}", e);
|
|
Err(())
|
|
}
|
|
}
|
|
}
|
|
#[instrument(skip(self))]
|
|
pub async fn stop_scheduled_job(&mut self, server_id: u64, channel_id: u64) {
|
|
match self.job_id.remove(&(server_id, channel_id)) {
|
|
Some(job_uid) => match self.scheduler.remove(&job_uid).await {
|
|
Ok(_) => {
|
|
info!("Cron job removed");
|
|
}
|
|
Err(e) => {
|
|
error!("Error removing cron job: {:?}", e);
|
|
}
|
|
},
|
|
None => {
|
|
error!("Cron job not found");
|
|
}
|
|
}
|
|
}
|
|
}
|