1#[cfg(test)]
26mod tests;
27
28use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
29use std::sync::mpsc::{self, Sender};
30use std::sync::Arc;
31use std::thread;
32use std::time::Duration;
33
34pub trait Job: Send {
41 fn run(&self) -> Result<(), String>;
44
45 fn name(&self) -> &str {
47 "job"
48 }
49}
50
51impl<F> Job for F
52where
53 F: Fn() -> Result<(), String> + Send,
54{
55 fn run(&self) -> Result<(), String> {
56 self()
57 }
58}
59
60struct Envelope {
61 job: Box<dyn Job>,
62 attempt: u32,
63}
64
65pub struct JobQueue {
76 sender: Sender<Envelope>,
77 workers: Vec<thread::JoinHandle<()>>,
78 max_retries: Arc<AtomicU32>,
79 initial_backoff_ms: Arc<AtomicU64>,
80 backoff_multiplier: Arc<AtomicU32>,
81}
82
83impl JobQueue {
84 pub fn new(workers: usize) -> Self {
90 assert!(workers > 0, "JobQueue requires at least one worker");
91
92 let (sender, receiver) = mpsc::channel::<Envelope>();
93 let receiver = Arc::new(std::sync::Mutex::new(receiver));
94
95 let max_retries = Arc::new(AtomicU32::new(3));
96 let initial_backoff_ms = Arc::new(AtomicU64::new(500));
97 let backoff_multiplier = Arc::new(AtomicU32::new(2));
98
99 let mut handles = Vec::with_capacity(workers);
100 for _ in 0..workers {
101 let receiver = Arc::clone(&receiver);
102 let max_retries = Arc::clone(&max_retries);
103 let initial_backoff_ms = Arc::clone(&initial_backoff_ms);
104 let backoff_multiplier = Arc::clone(&backoff_multiplier);
105 handles.push(thread::spawn(move || {
106 worker_loop(receiver, max_retries, initial_backoff_ms, backoff_multiplier);
107 }));
108 }
109
110 JobQueue {
111 sender,
112 workers: handles,
113 max_retries,
114 initial_backoff_ms,
115 backoff_multiplier,
116 }
117 }
118
119 pub fn max_retries(self, n: u32) -> Self {
121 self.max_retries.store(n, Ordering::Relaxed);
122 self
123 }
124
125 pub fn backoff(self, initial: Duration, multiplier: u32) -> Self {
128 self.initial_backoff_ms.store(initial.as_millis() as u64, Ordering::Relaxed);
129 self.backoff_multiplier.store(multiplier.max(1), Ordering::Relaxed);
130 self
131 }
132
133 pub fn submit(&self, job: impl Job + 'static) {
135 let _ = self.sender.send(Envelope { job: Box::new(job), attempt: 0 });
136 }
137
138 pub fn join(self) {
142 let JobQueue { sender, workers, .. } = self;
143 drop(sender);
144 for handle in workers {
145 let _ = handle.join();
146 }
147 }
148}
149
150fn worker_loop(
151 receiver: Arc<std::sync::Mutex<mpsc::Receiver<Envelope>>>,
152 max_retries: Arc<AtomicU32>,
153 initial_backoff_ms: Arc<AtomicU64>,
154 backoff_multiplier: Arc<AtomicU32>,
155) {
156 loop {
157 let received = {
158 let lock = receiver.lock().unwrap();
159 lock.recv()
160 };
161 let mut envelope = match received {
162 Ok(envelope) => envelope,
163 Err(_) => break, };
165
166 loop {
167 match envelope.job.run() {
168 Ok(()) => break,
169 Err(e) => {
170 let max = max_retries.load(Ordering::Relaxed);
171 if envelope.attempt >= max {
172 eprintln!(
173 "[jobs] '{}' failed permanently after {} attempt(s): {}",
174 envelope.job.name(),
175 envelope.attempt + 1,
176 e
177 );
178 break;
179 }
180
181 let initial_ms = initial_backoff_ms.load(Ordering::Relaxed);
182 let multiplier = backoff_multiplier.load(Ordering::Relaxed).max(1) as u64;
183 let delay = Duration::from_millis(
184 initial_ms.saturating_mul(multiplier.saturating_pow(envelope.attempt)),
185 );
186 eprintln!(
187 "[jobs] '{}' failed (attempt {}/{}): {} — retrying in {:?}",
188 envelope.job.name(),
189 envelope.attempt + 1,
190 max + 1,
191 e,
192 delay
193 );
194 thread::sleep(delay);
195 envelope.attempt += 1;
196 }
197 }
198 }
199 }
200}
201
202#[cfg(any(feature = "model-sqlite", feature = "model-postgres", feature = "model-mysql"))]
234pub struct PersistentJobQueue {
235 pool: crate::model::DbPool,
236 handlers: std::sync::RwLock<std::collections::HashMap<String, Arc<dyn Fn(&str) -> Result<(), String> + Send + Sync>>>,
237 default_max_retries: AtomicU32,
238 poll_interval_ms: AtomicU64,
239 initial_backoff_ms: AtomicU64,
240 backoff_multiplier: AtomicU32,
241 stop: Arc<std::sync::atomic::AtomicBool>,
242 id_counter: AtomicU64,
243}
244
245#[cfg(any(feature = "model-sqlite", feature = "model-postgres", feature = "model-mysql"))]
246impl PersistentJobQueue {
247 pub async fn new(pool: crate::model::DbPool) -> Result<Self, crate::model::DbError> {
251 let queue = PersistentJobQueue {
252 pool,
253 handlers: std::sync::RwLock::new(std::collections::HashMap::new()),
254 default_max_retries: AtomicU32::new(3),
255 poll_interval_ms: AtomicU64::new(500),
256 initial_backoff_ms: AtomicU64::new(500),
257 backoff_multiplier: AtomicU32::new(2),
258 stop: Arc::new(std::sync::atomic::AtomicBool::new(false)),
259 id_counter: AtomicU64::new(0),
260 };
261 queue.ensure_table().await?;
262 queue.recover_interrupted().await?;
263 Ok(queue)
264 }
265
266 async fn ensure_table(&self) -> Result<(), crate::model::DbError> {
267 self.pool
268 .execute(
269 "CREATE TABLE IF NOT EXISTS rws_jobs \
270 (id TEXT PRIMARY KEY, job_type TEXT NOT NULL, payload TEXT NOT NULL DEFAULT '', \
271 status TEXT NOT NULL DEFAULT 'pending', attempts INTEGER NOT NULL DEFAULT 0, \
272 max_retries INTEGER NOT NULL DEFAULT 3, next_run_at INTEGER NOT NULL, \
273 created_at INTEGER NOT NULL, last_error TEXT)",
274 &[],
275 )
276 .await?;
277 Ok(())
278 }
279
280 async fn recover_interrupted(&self) -> Result<(), crate::model::DbError> {
281 self.pool
282 .execute("UPDATE rws_jobs SET status = 'pending' WHERE status = 'running'", &[])
283 .await?;
284 Ok(())
285 }
286
287 pub fn max_retries(self, n: u32) -> Self {
291 self.default_max_retries.store(n, Ordering::Relaxed);
292 self
293 }
294
295 pub fn poll_interval(self, interval: Duration) -> Self {
297 self.poll_interval_ms.store(interval.as_millis() as u64, Ordering::Relaxed);
298 self
299 }
300
301 pub fn backoff(self, initial: Duration, multiplier: u32) -> Self {
305 self.initial_backoff_ms.store(initial.as_millis() as u64, Ordering::Relaxed);
306 self.backoff_multiplier.store(multiplier.max(1), Ordering::Relaxed);
307 self
308 }
309
310 pub fn register(&self, job_type: &str, handler: impl Fn(&str) -> Result<(), String> + Send + Sync + 'static) {
314 self.handlers.write().unwrap().insert(job_type.to_string(), Arc::new(handler));
315 }
316
317 fn generate_id(&self) -> String {
318 let nanos = std::time::SystemTime::now()
319 .duration_since(std::time::UNIX_EPOCH)
320 .map(|d| d.as_nanos() as u64)
321 .unwrap_or(0);
322 let count = self.id_counter.fetch_add(1, Ordering::Relaxed);
323 format!("{:x}-{:x}", nanos, count)
324 }
325
326 fn now_epoch() -> i64 {
327 std::time::SystemTime::now()
328 .duration_since(std::time::UNIX_EPOCH)
329 .map(|d| d.as_secs() as i64)
330 .unwrap_or(0)
331 }
332
333 pub async fn enqueue(&self, job_type: &str, payload: &str) -> Result<String, crate::model::DbError> {
336 let max_retries = self.default_max_retries.load(Ordering::Relaxed);
337 self.enqueue_with_retries(job_type, payload, max_retries).await
338 }
339
340 pub async fn enqueue_with_retries(
342 &self,
343 job_type: &str,
344 payload: &str,
345 max_retries: u32,
346 ) -> Result<String, crate::model::DbError> {
347 use crate::model::Value;
348 let id = self.generate_id();
349 let now = Self::now_epoch();
350 self.pool
351 .execute(
352 "INSERT INTO rws_jobs (id, job_type, payload, status, attempts, max_retries, next_run_at, created_at) \
353 VALUES (?, ?, ?, 'pending', 0, ?, ?, ?)",
354 &[
355 Value::Text(id.clone()),
356 Value::Text(job_type.to_string()),
357 Value::Text(payload.to_string()),
358 Value::Int(max_retries as i64),
359 Value::Int(now),
360 Value::Int(now),
361 ],
362 )
363 .await?;
364 Ok(id)
365 }
366
367 async fn claim_next(&self) -> Result<Option<(String, String, String, i64, i64)>, crate::model::DbError> {
372 use crate::model::Value;
373 let now = Self::now_epoch();
374 let rows = self
375 .pool
376 .query_rows(
377 "SELECT id, job_type, payload, attempts, max_retries FROM rws_jobs \
378 WHERE status = 'pending' AND next_run_at <= ? ORDER BY next_run_at ASC LIMIT 1",
379 &[Value::Int(now)],
380 )
381 .await?;
382 let row = match rows.into_iter().next() {
383 Some(r) => r,
384 None => return Ok(None),
385 };
386 let id: String = row.get("id")?;
387 let claimed = self
388 .pool
389 .execute(
390 "UPDATE rws_jobs SET status = 'running' WHERE id = ? AND status = 'pending'",
391 &[Value::Text(id.clone())],
392 )
393 .await?;
394 if claimed == 0 {
395 return Ok(None); }
397 Ok(Some((id, row.get("job_type")?, row.get("payload")?, row.get("attempts")?, row.get("max_retries")?)))
398 }
399
400 async fn mark_done(&self, id: &str) -> Result<(), crate::model::DbError> {
401 use crate::model::Value;
402 self.pool.execute("DELETE FROM rws_jobs WHERE id = ?", &[Value::Text(id.to_string())]).await?;
403 Ok(())
404 }
405
406 async fn mark_failed(&self, id: &str, error: &str) -> Result<(), crate::model::DbError> {
407 use crate::model::Value;
408 self.pool
409 .execute(
410 "UPDATE rws_jobs SET status = 'failed', last_error = ? WHERE id = ?",
411 &[Value::Text(error.to_string()), Value::Text(id.to_string())],
412 )
413 .await?;
414 Ok(())
415 }
416
417 async fn reschedule(&self, id: &str, attempts: i64, error: &str, delay: Duration) -> Result<(), crate::model::DbError> {
418 use crate::model::Value;
419 let next_run_at = Self::now_epoch() + delay.as_secs() as i64;
420 self.pool
421 .execute(
422 "UPDATE rws_jobs SET status = 'pending', attempts = ?, last_error = ?, next_run_at = ? WHERE id = ?",
423 &[Value::Int(attempts), Value::Text(error.to_string()), Value::Int(next_run_at), Value::Text(id.to_string())],
424 )
425 .await?;
426 Ok(())
427 }
428
429 pub async fn tick(&self) -> Result<bool, crate::model::DbError> {
434 let (id, job_type, payload, attempts, max_retries) = match self.claim_next().await? {
435 Some(claimed) => claimed,
436 None => return Ok(false),
437 };
438
439 let handler = self.handlers.read().unwrap().get(&job_type).cloned();
440 let result = match handler {
441 Some(h) => h(&payload),
442 None => Err(format!("no handler registered for job_type '{}'", job_type)),
443 };
444
445 match result {
446 Ok(()) => self.mark_done(&id).await?,
447 Err(e) => {
448 let next_attempts = attempts + 1;
452 if next_attempts > max_retries {
453 eprintln!(
454 "[jobs] persistent job '{}' ({}) failed permanently after {} attempt(s): {}",
455 job_type, id, next_attempts, e
456 );
457 self.mark_failed(&id, &e).await?;
458 } else {
459 let initial_ms = self.initial_backoff_ms.load(Ordering::Relaxed);
460 let multiplier = self.backoff_multiplier.load(Ordering::Relaxed).max(1) as u64;
461 let delay = Duration::from_millis(
462 initial_ms.saturating_mul(multiplier.saturating_pow(attempts as u32)),
463 );
464 eprintln!(
465 "[jobs] persistent job '{}' ({}) failed (attempt {}/{}): {} — retrying in {:?}",
466 job_type, id, next_attempts, max_retries + 1, e, delay
467 );
468 self.reschedule(&id, next_attempts, &e, delay).await?;
469 }
470 }
471 }
472 Ok(true)
473 }
474
475 pub fn start(self: Arc<Self>, workers: usize) -> Vec<thread::JoinHandle<()>> {
481 assert!(workers > 0, "PersistentJobQueue requires at least one worker");
482 (0..workers)
483 .map(|_| {
484 let queue = Arc::clone(&self);
485 thread::spawn(move || {
486 let rt = tokio::runtime::Builder::new_current_thread()
487 .enable_all()
488 .build()
489 .expect("failed to start tokio runtime for PersistentJobQueue worker");
490 rt.block_on(async {
491 loop {
492 if queue.stop.load(Ordering::Relaxed) {
493 break;
494 }
495 let poll_interval = Duration::from_millis(queue.poll_interval_ms.load(Ordering::Relaxed));
496 match queue.tick().await {
497 Ok(true) => {}
498 Ok(false) => tokio::time::sleep(poll_interval).await,
499 Err(e) => {
500 eprintln!("[jobs] persistent queue poll error: {}", e);
501 tokio::time::sleep(poll_interval).await;
502 }
503 }
504 }
505 });
506 })
507 })
508 .collect()
509 }
510
511 pub fn stop(&self) {
515 self.stop.store(true, Ordering::Relaxed);
516 }
517}