snerd_rust/
queue.rs

1use chrono::Utc;
2use std::collections::{HashMap, HashSet, BinaryHeap};
3use tokio::sync::Semaphore;
4use std::sync::{Arc, Mutex};
5use std::time::Duration;
6use tokio::sync::RwLock;
7use serde_json::json;
8
9use crate::file_store::FileStore;
10use crate::rate_limiter::RateLimiter;
11use crate::task::{RetryableTask, PriorityTask, ProgressMessage};
12use tokio::sync::broadcast;
13
14pub type TaskHandler = Arc<dyn Fn(String) -> Result<(), String> + Send + Sync>;
15pub type MaxRetryHandler = Arc<dyn Fn(String) -> Result<(), String> + Send + Sync>;
16
17struct ExecutingGuard {
18    executing_tasks: Arc<Mutex<HashSet<String>>>,
19    task_id: String,
20}
21
22impl Drop for ExecutingGuard {
23    fn drop(&mut self) {
24        if let Ok(mut executing) = self.executing_tasks.lock() {
25            executing.remove(&self.task_id);
26        }
27    }
28}
29
30#[derive(Clone)]
31pub struct SnerdQueue {
32    pub name: String,
33    pub file_store: FileStore,
34    pub rate_limiter: RateLimiter,
35    task_handlers: Arc<RwLock<HashMap<String, TaskHandler>>>,
36    max_retry_handlers: Arc<RwLock<HashMap<String, MaxRetryHandler>>>,
37    active_hashes: Arc<Mutex<HashSet<String>>>,
38    executing_tasks: Arc<Mutex<HashSet<String>>>,
39    worker_semaphore: Arc<Semaphore>,
40    pub progress_tx: broadcast::Sender<ProgressMessage>,
41}
42
43impl SnerdQueue {
44    pub fn new(name: &str, file_store: FileStore, rate_limiter: RateLimiter) -> Self {
45        let mut initial_hashes = HashSet::new();
46        if let Ok(tasks) = file_store.read_tasks() {
47            for task in tasks {
48                if task.deleted_at.is_none() {
49                    if let Some(hash) = task.payload_hash {
50                        initial_hashes.insert(hash);
51                    }
52                }
53            }
54        }
55
56        let (progress_tx, _) = broadcast::channel(1024);
57        Self {
58            name: name.to_string(),
59            file_store,
60            rate_limiter,
61            task_handlers: Arc::new(RwLock::new(HashMap::new())),
62            max_retry_handlers: Arc::new(RwLock::new(HashMap::new())),
63            active_hashes: Arc::new(Mutex::new(initial_hashes)),
64            executing_tasks: Arc::new(Mutex::new(HashSet::new())),
65            worker_semaphore: Arc::new(Semaphore::new(100)), // Limit to 100 concurrent tasks
66            progress_tx,
67        }
68    }
69
70    pub fn subscribe_progress(&self) -> broadcast::Receiver<ProgressMessage> {
71        self.progress_tx.subscribe()
72    }
73
74    pub fn yield_progress(&self, task_id: &str, data: &str) {
75        let _ = self.progress_tx.send(ProgressMessage {
76            task_id: task_id.to_string(),
77            data: data.to_string(),
78        });
79    }
80
81    pub async fn register_task_handler<F>(&self, task_type: &str, handler: F)
82    where
83        F: Fn(String) -> Result<(), String> + Send + Sync + 'static,
84    {
85        self.task_handlers
86            .write()
87            .await
88            .insert(task_type.to_string(), Arc::new(handler));
89    }
90
91    pub async fn register_max_retry_handler<F>(&self, task_type: &str, handler: F)
92    where
93        F: Fn(String) -> Result<(), String> + Send + Sync + 'static,
94    {
95        self.max_retry_handlers
96            .write()
97            .await
98            .insert(task_type.to_string(), Arc::new(handler));
99    }
100
101    pub fn enqueue(&self, mut task: RetryableTask) -> std::io::Result<()> {
102        if let Some(ref hash) = task.payload_hash {
103            if let Ok(mut hashes) = self.active_hashes.lock() {
104                if hashes.contains(hash) {
105                    return Ok(());
106                }
107                hashes.insert(hash.clone());
108            }
109        }
110        task.deleted_at = None;
111        self.file_store.save_task(&task)?;
112
113        if task.execute_at <= Utc::now() && task.retry_after_time <= Utc::now() {
114            if let Some(ref group) = task.rate_limit_group {
115                if let Some(limit) = task.max_per_minute {
116                    match self.rate_limiter.check_and_increment(group, limit) {
117                        Ok(true) => {}
118                        Ok(false) | Err(_) => {
119                            task.retry_after_time = Utc::now() + chrono::Duration::seconds(60);
120                            let _ = self.file_store.save_task(&task);
121                            return Ok(());
122                        }
123                    }
124                }
125            }
126
127            // Lock check before executing
128            if let Ok(mut executing) = self.executing_tasks.lock() {
129                if executing.contains(&task.task_id) {
130                    return Ok(());
131                }
132                executing.insert(task.task_id.clone());
133            }
134
135            let q = self.clone();
136            tokio::spawn(async move {
137                q.execute_task(task).await;
138            });
139        }
140        Ok(())
141    }
142
143    pub async fn start_processor(&self, interval: Duration) {
144        let q = self.clone();
145        tokio::spawn(async move {
146            let mut interval_timer = tokio::time::interval(interval);
147            loop {
148                interval_timer.tick().await;
149                q.process_due_tasks().await;
150            }
151        });
152    }
153
154    pub async fn process_due_tasks(&self) {
155        let tasks = match self.file_store.read_tasks() {
156            Ok(t) => t,
157            Err(_) => return,
158        };
159
160        let now = Utc::now();
161        let mut heap = BinaryHeap::new();
162        
163        for task in tasks {
164            if task.execute_at <= now && task.retry_after_time <= now && task.deleted_at.is_none() {
165                heap.push(PriorityTask(task));
166            }
167        }
168
169        let available = self.worker_semaphore.available_permits();
170        for _ in 0..available {
171            if let Some(PriorityTask(mut task)) = heap.pop() {
172                if let Some(ref group) = task.rate_limit_group {
173                    if let Some(limit) = task.max_per_minute {
174                        match self.rate_limiter.check_and_increment(group, limit) {
175                            Ok(true) => {}
176                            Ok(false) | Err(_) => {
177                                task.retry_after_time = now + chrono::Duration::seconds(60);
178                                let _ = self.file_store.save_task(&task);
179                                continue;
180                            }
181                        }
182                    }
183                }
184
185                // Lock check before executing
186                if let Ok(mut executing) = self.executing_tasks.lock() {
187                    // Double-check against the latest state in the file store to avoid TOCTOU race conditions
188                    if let Ok(Some(latest_task)) = self.file_store.get_latest_task(&task.task_id) {
189                        if latest_task.execute_at > now || latest_task.retry_after_time > now || latest_task.deleted_at.is_some() {
190                            continue;
191                        }
192                    } else {
193                        // Task was deleted completely
194                        continue;
195                    }
196
197                    if executing.contains(&task.task_id) {
198                        continue;
199                    }
200                    executing.insert(task.task_id.clone());
201                }
202
203                if let Ok(permit) = self.worker_semaphore.clone().try_acquire_owned() {
204                    let q = self.clone();
205                    tokio::spawn(async move {
206                        let _p = permit;
207                        q.execute_task(task).await;
208                    });
209                }
210            } else {
211                break;
212            }
213        }
214    }
215
216    async fn execute_task(&self, mut task: RetryableTask) {
217        // Drop guard guarantees removal from executing_tasks
218        let _guard = ExecutingGuard {
219            executing_tasks: Arc::clone(&self.executing_tasks),
220            task_id: task.task_id.clone(),
221        };
222
223        // Build the execution result — either via webhook HTTP call or local handler
224        let result: Result<(), String> = if let Some(ref url) = task.webhook_url.clone() {
225            // --- Webhook Path ---
226            let payload = json!({
227                "taskId": task.task_id,
228                "taskType": task.task_type,
229                "data": task.task_data,
230            });
231            let url = url.clone();
232            tokio::task::spawn_blocking(move || {
233                let rt = tokio::runtime::Handle::current();
234                rt.block_on(async {
235                    let mut client_builder = reqwest::Client::builder();
236                    if let Some(secs) = task.max_execution_seconds {
237                        client_builder = client_builder.timeout(std::time::Duration::from_secs(secs));
238                    }
239                    let client = client_builder.build().unwrap_or_else(|_| reqwest::Client::new());
240                    
241                    match client
242                        .post(&url)
243                        .header("Content-Type", "application/json")
244                        .header("X-SnerdMQ-Event", "Execute")
245                        .json(&payload)
246                        .send()
247                        .await
248                    {
249                        Ok(resp) if resp.status().is_success() => Ok(()),
250                        Ok(resp) => Err(format!("Webhook returned non-2xx status: {}", resp.status())),
251                        Err(e) => {
252                            if e.is_timeout() {
253                                Err(format!("Webhook execution timed out after {} seconds", task.max_execution_seconds.unwrap_or(0)))
254                            } else {
255                                Err(format!("Webhook request failed: {}", e))
256                            }
257                        }
258                    }
259                })
260            })
261            .await
262            .unwrap_or_else(|e| Err(format!("Webhook task panic: {:?}", e)))
263        } else {
264            // --- Local Handler Path ---
265            let handler = {
266                let handlers = self.task_handlers.read().await;
267                handlers.get(&task.task_type).cloned()
268            };
269            if let Some(h) = handler {
270                let task_data = task.task_data.clone();
271                let fut = tokio::task::spawn_blocking(move || h(task_data));
272                
273                if let Some(secs) = task.max_execution_seconds {
274                    match tokio::time::timeout(std::time::Duration::from_secs(secs), fut).await {
275                        Ok(Ok(res)) => res,
276                        Ok(Err(e)) => Err(format!("Task panic: {:?}", e)),
277                        Err(_) => Err(format!("Task execution timed out after {} seconds", secs)),
278                    }
279                } else {
280                    fut.await.unwrap_or_else(|e| Err(format!("Task panic: {:?}", e)))
281                }
282            } else {
283                return; // No handler and no webhook — nothing to do
284            }
285        };
286
287        match result {
288            Ok(_) => {
289                let mut rescheduled = false;
290                if let Some(ref cron_expr) = task.cron_expression {
291                    use cron::Schedule;
292                    use std::str::FromStr;
293                    if let Ok(schedule) = Schedule::from_str(cron_expr) {
294                        if let Some(next) = schedule.upcoming(Utc).next() {
295                            task.execute_at = next;
296                            task.retry_count = 0;
297                            task.last_error_obj = None;
298                            task.last_job_error = None;
299                            let _ = self.file_store.save_task(&task);
300                            rescheduled = true;
301                        }
302                    }
303                }
304
305                if !rescheduled {
306                    let _ = self.file_store.delete_task(&task.task_id);
307                    if let Some(ref hash) = task.payload_hash {
308                        if let Ok(mut hashes) = self.active_hashes.lock() {
309                            hashes.remove(hash);
310                        }
311                    }
312                }
313            }
314            Err(e) => {
315                if task.retry_count < task.max_retries {
316                    task.update_retry_config(Some(e));
317                    let _ = self.file_store.save_task(&task);
318                } else {
319                    // Max retries reached — fire DLQ webhook or local max retry handler
320                    if let Some(ref url) = task.webhook_url.clone() {
321                        let payload = json!({
322                            "taskId": task.task_id,
323                            "taskType": task.task_type,
324                            "data": task.task_data,
325                        });
326                        let url = url.clone();
327                        tokio::spawn(async move {
328                            let _ = reqwest::Client::new()
329                                .post(&url)
330                                .header("Content-Type", "application/json")
331                                .header("X-SnerdMQ-Event", "MaxRetriesReached")
332                                .json(&payload)
333                                .send()
334                                .await;
335                        });
336                    } else {
337                        let max_handler = {
338                            let max_handlers = self.max_retry_handlers.read().await;
339                            max_handlers.get(&task.task_type).cloned()
340                        };
341                        if let Some(mh) = max_handler {
342                            let max_data = task.task_data.clone();
343                            let _ = tokio::task::spawn_blocking(move || mh(max_data)).await;
344                        }
345                    }
346
347                    let _ = self.file_store.delete_task(&task.task_id);
348                    if let Some(ref hash) = task.payload_hash {
349                        if let Ok(mut hashes) = self.active_hashes.lock() {
350                            hashes.remove(hash);
351                        }
352                    }
353                }
354            }
355        }
356    }
357}
snerd_rust/queue.rs

snerd_rust/
queue.rs
