Skip to main content

wae_scheduler/
task.rs

1use async_trait::async_trait;
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4use std::sync::{
5    Arc,
6    atomic::{AtomicBool, AtomicU64, Ordering},
7};
8use tokio::sync::RwLock;
9use wae_types::WaeResult;
10
11/// 任务 ID 类型
12pub type TaskId = String;
13
14/// 任务状态
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16pub enum TaskState {
17    /// 待执行
18    Pending,
19    /// 运行中
20    Running,
21    /// 已暂停
22    Paused,
23    /// 已完成
24    Completed,
25    /// 已取消
26    Cancelled,
27    /// 执行失败
28    Failed,
29}
30
31impl std::fmt::Display for TaskState {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        match self {
34            TaskState::Pending => write!(f, "pending"),
35            TaskState::Running => write!(f, "running"),
36            TaskState::Paused => write!(f, "paused"),
37            TaskState::Completed => write!(f, "completed"),
38            TaskState::Cancelled => write!(f, "cancelled"),
39            TaskState::Failed => write!(f, "failed"),
40        }
41    }
42}
43
44/// 定时任务 trait
45///
46/// 定义定时任务的核心接口,所有需要被调度的任务都需要实现此 trait。
47#[async_trait]
48pub trait ScheduledTask: Send + Sync {
49    /// 执行任务
50    ///
51    /// # 返回值
52    ///
53    /// 返回 `Ok(())` 表示任务执行成功,返回 `Err` 表示任务执行失败。
54    async fn execute(&self) -> WaeResult<()>;
55
56    /// 获取任务名称
57    fn name(&self) -> &str;
58}
59
60/// 任务句柄
61///
62/// 用于管理和查询任务状态。
63#[derive(Debug, Clone)]
64pub struct TaskHandle {
65    /// 任务 ID
66    pub id: TaskId,
67    /// 任务名称
68    pub name: String,
69    /// 任务状态
70    state: Arc<RwLock<TaskState>>,
71    /// 执行次数计数器
72    execution_count: Arc<AtomicU64>,
73    /// 最后执行时间
74    last_execution: Arc<RwLock<Option<DateTime<Utc>>>>,
75    /// 最后错误信息
76    last_error: Arc<RwLock<Option<String>>>,
77    /// 取消标志
78    cancelled: Arc<AtomicBool>,
79}
80
81impl TaskHandle {
82    /// 创建新的任务句柄
83    pub fn new(id: TaskId, name: String) -> Self {
84        Self {
85            id,
86            name,
87            state: Arc::new(RwLock::new(TaskState::Pending)),
88            execution_count: Arc::new(AtomicU64::new(0)),
89            last_execution: Arc::new(RwLock::new(None)),
90            last_error: Arc::new(RwLock::new(None)),
91            cancelled: Arc::new(AtomicBool::new(false)),
92        }
93    }
94
95    /// 获取任务状态
96    pub async fn state(&self) -> TaskState {
97        *self.state.read().await
98    }
99
100    /// 获取执行次数
101    pub fn execution_count(&self) -> u64 {
102        self.execution_count.load(Ordering::SeqCst)
103    }
104
105    /// 获取最后执行时间
106    pub async fn last_execution(&self) -> Option<DateTime<Utc>> {
107        *self.last_execution.read().await
108    }
109
110    /// 获取最后错误信息
111    pub async fn last_error(&self) -> Option<String> {
112        self.last_error.read().await.clone()
113    }
114
115    /// 检查任务是否已取消
116    pub fn is_cancelled(&self) -> bool {
117        self.cancelled.load(Ordering::SeqCst)
118    }
119
120    /// 更新状态
121    pub async fn set_state(&self, state: TaskState) {
122        *self.state.write().await = state;
123    }
124
125    /// 记录执行
126    pub async fn record_execution(&self) {
127        self.execution_count.fetch_add(1, Ordering::SeqCst);
128        *self.last_execution.write().await = Some(Utc::now());
129    }
130
131    /// 记录错误
132    pub async fn record_error(&self, error: String) {
133        *self.last_error.write().await = Some(error);
134    }
135
136    /// 取消任务
137    pub fn cancel(&self) {
138        self.cancelled.store(true, Ordering::SeqCst);
139    }
140}