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
11pub type TaskId = String;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16pub enum TaskState {
17 Pending,
19 Running,
21 Paused,
23 Completed,
25 Cancelled,
27 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#[async_trait]
48pub trait ScheduledTask: Send + Sync {
49 async fn execute(&self) -> WaeResult<()>;
55
56 fn name(&self) -> &str;
58}
59
60#[derive(Debug, Clone)]
64pub struct TaskHandle {
65 pub id: TaskId,
67 pub name: String,
69 state: Arc<RwLock<TaskState>>,
71 execution_count: Arc<AtomicU64>,
73 last_execution: Arc<RwLock<Option<DateTime<Utc>>>>,
75 last_error: Arc<RwLock<Option<String>>>,
77 cancelled: Arc<AtomicBool>,
79}
80
81impl TaskHandle {
82 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 pub async fn state(&self) -> TaskState {
97 *self.state.read().await
98 }
99
100 pub fn execution_count(&self) -> u64 {
102 self.execution_count.load(Ordering::SeqCst)
103 }
104
105 pub async fn last_execution(&self) -> Option<DateTime<Utc>> {
107 *self.last_execution.read().await
108 }
109
110 pub async fn last_error(&self) -> Option<String> {
112 self.last_error.read().await.clone()
113 }
114
115 pub fn is_cancelled(&self) -> bool {
117 self.cancelled.load(Ordering::SeqCst)
118 }
119
120 pub async fn set_state(&self, state: TaskState) {
122 *self.state.write().await = state;
123 }
124
125 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 pub async fn record_error(&self, error: String) {
133 *self.last_error.write().await = Some(error);
134 }
135
136 pub fn cancel(&self) {
138 self.cancelled.store(true, Ordering::SeqCst);
139 }
140}