1use std::collections::HashMap;
4use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
5use std::sync::{Arc, RwLock};
6use std::thread;
7use std::time::{Duration, SystemTime, UNIX_EPOCH};
8
9use crate::model::{Backend, DbError, DbPool, ModelRow, Value};
10use crate::scheduler::CronSchedule;
11
12type Handler = Arc<dyn Fn(&str) -> Result<(), String> + Send + Sync>;
13type DeadHook = Arc<dyn Fn(&JobInfo) + Send + Sync>;
14
15#[derive(Debug, Clone, Default)]
30pub struct EnqueueOptions {
31 queue: Option<String>,
32 priority: i64,
33 run_at: Option<SystemTime>,
34 unique_key: Option<String>,
35 max_retries: Option<u32>,
36}
37
38impl EnqueueOptions {
39 pub fn new() -> Self {
40 EnqueueOptions::default()
41 }
42
43 pub fn queue(mut self, name: &str) -> Self {
46 self.queue = Some(name.to_string());
47 self
48 }
49
50 pub fn priority(mut self, priority: i64) -> Self {
52 self.priority = priority;
53 self
54 }
55
56 pub fn run_at(mut self, at: SystemTime) -> Self {
58 self.run_at = Some(at);
59 self
60 }
61
62 pub fn run_in(mut self, delay: Duration) -> Self {
64 self.run_at = Some(SystemTime::now() + delay);
65 self
66 }
67
68 pub fn unique(mut self, key: &str) -> Self {
73 self.unique_key = Some(key.to_string());
74 self
75 }
76
77 pub fn max_retries(mut self, n: u32) -> Self {
79 self.max_retries = Some(n);
80 self
81 }
82}
83
84#[derive(Debug, Clone, PartialEq, Eq)]
86pub struct JobInfo {
87 pub id: String,
88 pub job_type: String,
89 pub payload: String,
90 pub queue: String,
91 pub priority: i64,
92 pub status: String,
94 pub attempts: i64,
96 pub max_retries: i64,
97 pub run_at: i64,
99 pub created_at: i64,
100 pub last_error: Option<String>,
101 pub unique_key: Option<String>,
102}
103
104impl JobInfo {
105 fn from_row(row: &ModelRow) -> Result<JobInfo, DbError> {
106 Ok(JobInfo {
107 id: row.get("id")?,
108 job_type: row.get("job_type")?,
109 payload: row.get("payload")?,
110 queue: row.get("queue")?,
111 priority: row.get("priority")?,
112 status: row.get("status")?,
113 attempts: row.get("attempts")?,
114 max_retries: row.get("max_retries")?,
115 run_at: row.get("next_run_at")?,
116 created_at: row.get("created_at")?,
117 last_error: row.get("last_error")?,
118 unique_key: row.get("unique_key")?,
119 })
120 }
121}
122
123const COLUMNS: &str = "id, job_type, payload, queue, priority, status, attempts, max_retries, next_run_at, created_at, last_error, unique_key";
124
125struct Schedule {
128 name: String,
129 cron: CronSchedule,
130 job_type: String,
131 payload: String,
132 options: EnqueueOptions,
133}
134
135pub struct PersistentJobQueue {
175 pool: DbPool,
176 handlers: RwLock<HashMap<String, Handler>>,
177 default_max_retries: AtomicU32,
178 poll_interval_ms: AtomicU64,
179 initial_backoff_ms: AtomicU64,
180 backoff_multiplier: AtomicU32,
181 stop: Arc<AtomicBool>,
182 id_counter: AtomicU64,
183 queues: RwLock<Vec<String>>,
184 schedules: RwLock<Vec<Arc<Schedule>>>,
185 on_dead: RwLock<Vec<DeadHook>>,
186}
187
188fn portable(backend: Backend, sql: &str) -> String {
190 #[cfg(feature = "model-postgres")]
191 if backend == Backend::Postgres {
192 let mut out = String::with_capacity(sql.len() + 8);
193 let mut n = 0;
194 for c in sql.chars() {
195 if c == '?' {
196 n += 1;
197 out.push_str(&format!("${n}"));
198 } else {
199 out.push(c);
200 }
201 }
202 return out;
203 }
204 let _ = backend;
205 sql.to_string()
206}
207
208fn now_epoch() -> i64 {
209 SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs() as i64).unwrap_or(0)
210}
211
212fn epoch_of(at: SystemTime) -> i64 {
213 at.duration_since(UNIX_EPOCH).map(|d| d.as_secs() as i64).unwrap_or(0)
214}
215
216impl PersistentJobQueue {
217 pub async fn new(pool: DbPool) -> Result<Self, DbError> {
222 let queue = PersistentJobQueue {
223 pool,
224 handlers: RwLock::new(HashMap::new()),
225 default_max_retries: AtomicU32::new(3),
226 poll_interval_ms: AtomicU64::new(500),
227 initial_backoff_ms: AtomicU64::new(500),
228 backoff_multiplier: AtomicU32::new(2),
229 stop: Arc::new(AtomicBool::new(false)),
230 id_counter: AtomicU64::new(0),
231 queues: RwLock::new(Vec::new()),
232 schedules: RwLock::new(Vec::new()),
233 on_dead: RwLock::new(Vec::new()),
234 };
235 queue.ensure_tables().await?;
236 queue.recover_interrupted().await?;
237 Ok(queue)
238 }
239
240 fn sql(&self, sql: &str) -> String {
241 portable(self.pool.backend(), sql)
242 }
243
244 async fn ensure_tables(&self) -> Result<(), DbError> {
245 self.pool
246 .execute(
247 "CREATE TABLE IF NOT EXISTS rws_jobs \
248 (id VARCHAR(64) PRIMARY KEY, job_type TEXT NOT NULL, payload TEXT NOT NULL, \
249 status VARCHAR(16) NOT NULL DEFAULT 'pending', attempts INTEGER NOT NULL DEFAULT 0, \
250 max_retries INTEGER NOT NULL DEFAULT 3, next_run_at BIGINT NOT NULL, \
251 created_at BIGINT NOT NULL, last_error TEXT, \
252 queue VARCHAR(64) NOT NULL DEFAULT 'default', priority BIGINT NOT NULL DEFAULT 0, \
253 unique_key VARCHAR(191))",
254 &[],
255 )
256 .await?;
257 for (column, definition) in [
259 ("queue", "VARCHAR(64) NOT NULL DEFAULT 'default'"),
260 ("priority", "BIGINT NOT NULL DEFAULT 0"),
261 ("unique_key", "VARCHAR(191)"),
262 ] {
263 if self.pool.query_rows(&format!("SELECT {column} FROM rws_jobs WHERE 1 = 0"), &[]).await.is_err() {
264 self.pool.execute(&format!("ALTER TABLE rws_jobs ADD COLUMN {column} {definition}"), &[]).await?;
265 }
266 }
267 self.pool
268 .execute(
269 "CREATE TABLE IF NOT EXISTS rws_job_schedules (name VARCHAR(191) PRIMARY KEY, next_run_at BIGINT NOT NULL)",
270 &[],
271 )
272 .await?;
273 Ok(())
274 }
275
276 async fn recover_interrupted(&self) -> Result<(), DbError> {
277 self.pool.execute("UPDATE rws_jobs SET status = 'pending' WHERE status = 'running'", &[]).await?;
278 Ok(())
279 }
280
281 pub fn max_retries(self, n: u32) -> Self {
285 self.default_max_retries.store(n, Ordering::Relaxed);
286 self
287 }
288
289 pub fn poll_interval(self, interval: Duration) -> Self {
291 self.poll_interval_ms.store(interval.as_millis() as u64, Ordering::Relaxed);
292 self
293 }
294
295 pub fn backoff(self, initial: Duration, multiplier: u32) -> Self {
299 self.initial_backoff_ms.store(initial.as_millis() as u64, Ordering::Relaxed);
300 self.backoff_multiplier.store(multiplier.max(1), Ordering::Relaxed);
301 self
302 }
303
304 pub fn queues(self, names: &[&str]) -> Self {
307 *self.queues.write().unwrap_or_else(|e| e.into_inner()) = names.iter().map(|n| n.to_string()).collect();
308 self
309 }
310
311 pub fn schedule(self, name: &str, cron: &str, job_type: &str, payload: &str) -> Result<Self, String> {
317 self.schedule_with(name, cron, job_type, payload, EnqueueOptions::new())
318 }
319
320 pub fn schedule_with(self, name: &str, cron: &str, job_type: &str, payload: &str, options: EnqueueOptions) -> Result<Self, String> {
322 let cron = CronSchedule::parse(cron)?;
323 self.schedules.write().unwrap_or_else(|e| e.into_inner()).push(Arc::new(Schedule {
324 name: name.to_string(),
325 cron,
326 job_type: job_type.to_string(),
327 payload: payload.to_string(),
328 options,
329 }));
330 Ok(self)
331 }
332
333 pub fn on_dead(&self, hook: impl Fn(&JobInfo) + Send + Sync + 'static) {
336 self.on_dead.write().unwrap_or_else(|e| e.into_inner()).push(Arc::new(hook));
337 }
338
339 pub fn register(&self, job_type: &str, handler: impl Fn(&str) -> Result<(), String> + Send + Sync + 'static) {
343 self.handlers.write().unwrap_or_else(|e| e.into_inner()).insert(job_type.to_string(), Arc::new(handler));
344 }
345
346 fn generate_id(&self) -> String {
347 let nanos = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_nanos() as u64).unwrap_or(0);
348 let count = self.id_counter.fetch_add(1, Ordering::Relaxed);
349 format!("{:x}-{:x}", nanos, count)
350 }
351
352 pub async fn enqueue(&self, job_type: &str, payload: &str) -> Result<String, DbError> {
354 self.enqueue_with(job_type, payload, EnqueueOptions::new()).await
355 }
356
357 pub async fn enqueue_with_retries(&self, job_type: &str, payload: &str, max_retries: u32) -> Result<String, DbError> {
359 self.enqueue_with(job_type, payload, EnqueueOptions::new().max_retries(max_retries)).await
360 }
361
362 pub async fn enqueue_with(&self, job_type: &str, payload: &str, options: EnqueueOptions) -> Result<String, DbError> {
366 if let Some(key) = &options.unique_key {
367 let rows = self
368 .pool
369 .query_rows(
370 &self.sql("SELECT id FROM rws_jobs WHERE job_type = ? AND unique_key = ? AND status IN ('pending', 'running')"),
371 &[Value::Text(job_type.to_string()), Value::Text(key.clone())],
372 )
373 .await?;
374 if let Some(row) = rows.first() {
375 return row.get("id");
376 }
377 }
378 let id = self.generate_id();
379 let now = now_epoch();
380 let max_retries = options.max_retries.unwrap_or_else(|| self.default_max_retries.load(Ordering::Relaxed));
381 self.pool
382 .execute(
383 &self.sql(
384 "INSERT INTO rws_jobs (id, job_type, payload, status, attempts, max_retries, next_run_at, created_at, queue, priority, unique_key) \
385 VALUES (?, ?, ?, 'pending', 0, ?, ?, ?, ?, ?, ?)",
386 ),
387 &[
388 Value::Text(id.clone()),
389 Value::Text(job_type.to_string()),
390 Value::Text(payload.to_string()),
391 Value::Int(max_retries as i64),
392 Value::Int(options.run_at.map(epoch_of).unwrap_or(now)),
393 Value::Int(now),
394 Value::Text(options.queue.clone().unwrap_or_else(|| "default".to_string())),
395 Value::Int(options.priority),
396 options.unique_key.clone().map(Value::Text).unwrap_or(Value::Null),
397 ],
398 )
399 .await?;
400 Ok(id)
401 }
402
403 async fn claim_next(&self) -> Result<Option<JobInfo>, DbError> {
408 let queues = self.queues.read().unwrap_or_else(|e| e.into_inner()).clone();
409 let mut params = vec![Value::Int(now_epoch())];
410 let mut filter = String::new();
411 if !queues.is_empty() {
412 filter = format!(" AND queue IN ({})", vec!["?"; queues.len()].join(", "));
413 params.extend(queues.into_iter().map(Value::Text));
414 }
415 let rows = self
416 .pool
417 .query_rows(
418 &self.sql(&format!(
419 "SELECT {COLUMNS} FROM rws_jobs WHERE status = 'pending' AND next_run_at <= ?{filter} \
420 ORDER BY priority DESC, next_run_at ASC, created_at ASC LIMIT 1"
421 )),
422 ¶ms,
423 )
424 .await?;
425 let Some(row) = rows.into_iter().next() else {
426 return Ok(None);
427 };
428 let job = JobInfo::from_row(&row)?;
429 let claimed = self
430 .pool
431 .execute(&self.sql("UPDATE rws_jobs SET status = 'running' WHERE id = ? AND status = 'pending'"), &[Value::Text(job.id.clone())])
432 .await?;
433 if claimed == 0 {
434 return Ok(None); }
436 Ok(Some(job))
437 }
438
439 async fn mark_done(&self, id: &str) -> Result<(), DbError> {
440 self.pool.execute(&self.sql("DELETE FROM rws_jobs WHERE id = ?"), &[Value::Text(id.to_string())]).await?;
441 Ok(())
442 }
443
444 async fn mark_failed(&self, job: &JobInfo, attempts: i64, error: &str) -> Result<(), DbError> {
445 self.pool
446 .execute(
447 &self.sql("UPDATE rws_jobs SET status = 'failed', attempts = ?, last_error = ? WHERE id = ?"),
448 &[Value::Int(attempts), Value::Text(error.to_string()), Value::Text(job.id.clone())],
449 )
450 .await?;
451 let dead = JobInfo { status: "failed".into(), attempts, last_error: Some(error.to_string()), ..job.clone() };
452 let hooks = self.on_dead.read().unwrap_or_else(|e| e.into_inner()).clone();
453 for hook in hooks {
454 hook(&dead);
455 }
456 Ok(())
457 }
458
459 async fn reschedule(&self, id: &str, attempts: i64, error: &str, delay: Duration) -> Result<(), DbError> {
460 let next_run_at = now_epoch() + delay.as_secs() as i64;
461 self.pool
462 .execute(
463 &self.sql("UPDATE rws_jobs SET status = 'pending', attempts = ?, last_error = ?, next_run_at = ? WHERE id = ?"),
464 &[Value::Int(attempts), Value::Text(error.to_string()), Value::Int(next_run_at), Value::Text(id.to_string())],
465 )
466 .await?;
467 Ok(())
468 }
469
470 pub async fn tick(&self) -> Result<bool, DbError> {
475 let Some(job) = self.claim_next().await? else {
476 return Ok(false);
477 };
478
479 let handler = self.handlers.read().unwrap_or_else(|e| e.into_inner()).get(&job.job_type).cloned();
480 let result = match handler {
481 Some(h) => h(&job.payload),
482 None => Err(format!("no handler registered for job_type '{}'", job.job_type)),
483 };
484
485 match result {
486 Ok(()) => self.mark_done(&job.id).await?,
487 Err(e) => {
488 let next_attempts = job.attempts + 1;
492 if next_attempts > job.max_retries {
493 eprintln!(
494 "[jobs] persistent job '{}' ({}) failed permanently after {} attempt(s): {}",
495 job.job_type, job.id, next_attempts, e
496 );
497 self.mark_failed(&job, next_attempts, &e).await?;
498 } else {
499 let initial_ms = self.initial_backoff_ms.load(Ordering::Relaxed);
500 let multiplier = self.backoff_multiplier.load(Ordering::Relaxed).max(1) as u64;
501 let delay = Duration::from_millis(initial_ms.saturating_mul(multiplier.saturating_pow(job.attempts as u32)));
502 eprintln!(
503 "[jobs] persistent job '{}' ({}) failed (attempt {}/{}): {} — retrying in {:?}",
504 job.job_type,
505 job.id,
506 next_attempts,
507 job.max_retries + 1,
508 e,
509 delay
510 );
511 self.reschedule(&job.id, next_attempts, &e, delay).await?;
512 }
513 }
514 }
515 Ok(true)
516 }
517
518 pub async fn tick_schedules(&self) -> Result<usize, DbError> {
523 self.tick_schedules_at(now_epoch()).await
524 }
525
526 async fn tick_schedules_at(&self, now: i64) -> Result<usize, DbError> {
527 let schedules = self.schedules.read().unwrap_or_else(|e| e.into_inner()).clone();
528 let mut enqueued = 0;
529 for schedule in schedules {
530 let next_after = |t: i64| schedule.cron.next_after(t.max(0) as u64).map(|n| n as i64);
531 let rows = self
532 .pool
533 .query_rows(&self.sql("SELECT next_run_at FROM rws_job_schedules WHERE name = ?"), &[Value::Text(schedule.name.clone())])
534 .await?;
535 let Some(row) = rows.first() else {
536 if let Some(next) = next_after(now) {
539 let _ = self
540 .pool
541 .execute(&self.sql("INSERT INTO rws_job_schedules (name, next_run_at) VALUES (?, ?)"), &[Value::Text(schedule.name.clone()), Value::Int(next)])
542 .await;
543 }
544 continue;
545 };
546 let due: i64 = row.get("next_run_at")?;
547 if due > now {
548 continue;
549 }
550 let Some(next) = next_after(now) else { continue };
552 let claimed = self
553 .pool
554 .execute(
555 &self.sql("UPDATE rws_job_schedules SET next_run_at = ? WHERE name = ? AND next_run_at = ?"),
556 &[Value::Int(next), Value::Text(schedule.name.clone()), Value::Int(due)],
557 )
558 .await?;
559 if claimed == 1 {
560 self.enqueue_with(&schedule.job_type, &schedule.payload, schedule.options.clone()).await?;
561 enqueued += 1;
562 }
563 }
564 Ok(enqueued)
565 }
566
567 pub async fn jobs(&self, status: Option<&str>, limit: usize) -> Result<Vec<JobInfo>, DbError> {
570 let (filter, params) = match status {
571 Some(status) => (" WHERE status = ?", vec![Value::Text(status.to_string())]),
572 None => ("", Vec::new()),
573 };
574 let rows = self
575 .pool
576 .query_rows(&self.sql(&format!("SELECT {COLUMNS} FROM rws_jobs{filter} ORDER BY priority DESC, next_run_at ASC, created_at ASC LIMIT {limit}")), ¶ms)
577 .await?;
578 rows.iter().map(JobInfo::from_row).collect()
579 }
580
581 pub async fn failed_jobs(&self, limit: usize) -> Result<Vec<JobInfo>, DbError> {
583 self.jobs(Some("failed"), limit).await
584 }
585
586 pub async fn retry(&self, id: &str) -> Result<bool, DbError> {
589 let changed = self
590 .pool
591 .execute(
592 &self.sql("UPDATE rws_jobs SET status = 'pending', attempts = 0, next_run_at = ? WHERE id = ? AND status = 'failed'"),
593 &[Value::Int(now_epoch()), Value::Text(id.to_string())],
594 )
595 .await?;
596 Ok(changed > 0)
597 }
598
599 pub async fn delete(&self, id: &str) -> Result<bool, DbError> {
601 let changed = self
602 .pool
603 .execute(&self.sql("DELETE FROM rws_jobs WHERE id = ? AND status <> 'running'"), &[Value::Text(id.to_string())])
604 .await?;
605 Ok(changed > 0)
606 }
607
608 pub fn start(self: Arc<Self>, workers: usize) -> Vec<thread::JoinHandle<()>> {
615 assert!(workers > 0, "PersistentJobQueue requires at least one worker");
616 (0..workers)
617 .map(|index| {
618 let queue = Arc::clone(&self);
619 thread::spawn(move || {
620 let rt = tokio::runtime::Builder::new_current_thread()
621 .enable_all()
622 .build()
623 .expect("failed to start tokio runtime for PersistentJobQueue worker");
624 rt.block_on(async {
625 let mut last_schedule_check = 0;
626 loop {
627 if queue.stop.load(Ordering::Relaxed) {
628 break;
629 }
630 if index == 0 && now_epoch() != last_schedule_check {
631 last_schedule_check = now_epoch();
632 if let Err(e) = queue.tick_schedules().await {
633 eprintln!("[jobs] schedule check error: {}", e);
634 }
635 }
636 let poll_interval = Duration::from_millis(queue.poll_interval_ms.load(Ordering::Relaxed));
637 match queue.tick().await {
638 Ok(true) => {}
639 Ok(false) => tokio::time::sleep(poll_interval).await,
640 Err(e) => {
641 eprintln!("[jobs] persistent queue poll error: {}", e);
642 tokio::time::sleep(poll_interval).await;
643 }
644 }
645 }
646 });
647 })
648 })
649 .collect()
650 }
651
652 pub fn stop(&self) {
656 self.stop.store(true, Ordering::Relaxed);
657 }
658
659 pub fn pool(&self) -> &DbPool {
661 &self.pool
662 }
663}
664
665#[cfg(test)]
666mod tests;