snerd_rust/
queue.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3use tokio::sync::RwLock;
4use std::time::Duration;
5use chrono::Utc;
6
7use crate::file_store::FileStore;
8use crate::task::RetryableTask;
9
10pub type TaskHandler = Arc<dyn Fn(String) -> Result<(), String> + Send + Sync>;
11pub type MaxRetryHandler = Arc<dyn Fn(String) -> Result<(), String> + Send + Sync>;
12
13#[derive(Clone)]
14pub struct SnerdQueue {
15    pub name: String,
16    pub file_store: FileStore,
17    task_handlers: Arc<RwLock<HashMap<String, TaskHandler>>>,
18    max_retry_handlers: Arc<RwLock<HashMap<String, MaxRetryHandler>>>,
19}
20
21impl SnerdQueue {
22    pub fn new(name: &str, file_store: FileStore) -> Self {
23        Self {
24            name: name.to_string(),
25            file_store,
26            task_handlers: Arc::new(RwLock::new(HashMap::new())),
27            max_retry_handlers: Arc::new(RwLock::new(HashMap::new())),
28        }
29    }
30
31    pub async fn register_task_handler<F>(&self, task_type: &str, handler: F)
32    where
33        F: Fn(String) -> Result<(), String> + Send + Sync + 'static,
34    {
35        self.task_handlers
36            .write()
37            .await
38            .insert(task_type.to_string(), Arc::new(handler));
39    }
40
41    pub async fn register_max_retry_handler<F>(&self, task_type: &str, handler: F)
42    where
43        F: Fn(String) -> Result<(), String> + Send + Sync + 'static,
44    {
45        self.max_retry_handlers
46            .write()
47            .await
48            .insert(task_type.to_string(), Arc::new(handler));
49    }
50
51    pub fn enqueue(&self, mut task: RetryableTask) -> std::io::Result<()> {
52        task.deleted_at = None;
53        self.file_store.save_task(&task)?;
54
55        if task.retry_after_time <= Utc::now() {
56            let q = self.clone();
57            tokio::spawn(async move {
58                q.execute_task(task).await;
59            });
60        }
61        Ok(())
62    }
63
64    pub async fn start_processor(&self, interval: Duration) {
65        let q = self.clone();
66        tokio::spawn(async move {
67            let mut interval_timer = tokio::time::interval(interval);
68            loop {
69                interval_timer.tick().await;
70                q.process_due_tasks().await;
71            }
72        });
73    }
74
75    pub async fn process_due_tasks(&self) {
76        let tasks = match self.file_store.read_tasks() {
77            Ok(t) => t,
78            Err(_) => return,
79        };
80
81        let now = Utc::now();
82        for task in tasks {
83            if task.retry_after_time <= now && task.deleted_at.is_none() {
84                let q = self.clone();
85                tokio::spawn(async move {
86                    q.execute_task(task).await;
87                });
88            }
89        }
90    }
91
92    async fn execute_task(&self, mut task: RetryableTask) {
93        let handler = {
94            let handlers = self.task_handlers.read().await;
95            handlers.get(&task.task_type).cloned()
96        };
97
98        if let Some(h) = handler {
99            // In a real production system with blocking synchronous handlers,
100            // we should spawn them using spawn_blocking to avoid starving the executor.
101            let task_data = task.task_data.clone();
102            
103            let result = tokio::task::spawn_blocking(move || {
104                h(task_data)
105            }).await.unwrap_or_else(|e| Err(format!("Task panic: {:?}", e)));
106
107            match result {
108                Ok(_) => {
109                    let _ = self.file_store.delete_task(&task.task_id);
110                }
111                Err(e) => {
112                    if task.retry_count < task.max_retries {
113                        task.update_retry_config(Some(e));
114                        let _ = self.file_store.save_task(&task);
115                    } else {
116                        // Max retries reached
117                        let max_handler = {
118                            let max_handlers = self.max_retry_handlers.read().await;
119                            max_handlers.get(&task.task_type).cloned()
120                        };
121
122                        if let Some(mh) = max_handler {
123                            let max_data = task.task_data.clone();
124                            let _ = tokio::task::spawn_blocking(move || {
125                                mh(max_data)
126                            }).await;
127                        }
128
129                        let _ = self.file_store.delete_task(&task.task_id);
130                    }
131                }
132            }
133        }
134    }
135}
snerd_rust/queue.rs

snerd_rust/
queue.rs
