1use crate::job::{JobRegistry, ReservedJob};
9use crate::queue::Queue;
10use rustlavel_core::events::{self, Event};
11use rustlavel_core::{Json, Result};
12use std::sync::Arc;
13use std::sync::atomic::{AtomicBool, Ordering};
14use std::time::{Duration, Instant};
15use tokio::sync::watch;
16
17pub const DEFAULT_POLL_INTERVAL: Duration = Duration::from_millis(500);
19
20pub const DEFAULT_BACKOFF_CAP: Duration = Duration::from_secs(3600);
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum Outcome {
27 Processed,
29 Released,
31 Failed,
33}
34
35#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
37pub struct WorkerStats {
38 pub processed: u64,
39 pub released: u64,
40 pub failed: u64,
41}
42
43impl WorkerStats {
44 pub fn attempts(&self) -> u64 {
46 self.processed + self.released + self.failed
47 }
48
49 fn record(&mut self, outcome: Outcome) {
50 match outcome {
51 Outcome::Processed => self.processed += 1,
52 Outcome::Released => self.released += 1,
53 Outcome::Failed => self.failed += 1,
54 }
55 }
56
57 fn merge(&mut self, other: WorkerStats) {
58 self.processed += other.processed;
59 self.released += other.released;
60 self.failed += other.failed;
61 }
62}
63
64#[derive(Clone)]
68pub struct Shutdown {
69 flag: Arc<AtomicBool>,
70 sender: Arc<watch::Sender<bool>>,
71}
72
73impl Default for Shutdown {
74 fn default() -> Self {
75 Shutdown::new()
76 }
77}
78
79impl Shutdown {
80 pub fn new() -> Self {
81 let (sender, _) = watch::channel(false);
82 Shutdown { flag: Arc::new(AtomicBool::new(false)), sender: Arc::new(sender) }
83 }
84
85 pub fn signal(&self) {
87 self.flag.store(true, Ordering::SeqCst);
90 let _ = self.sender.send(true);
91 }
92
93 pub fn is_signalled(&self) -> bool {
94 self.flag.load(Ordering::SeqCst)
95 }
96
97 pub fn on_ctrl_c(&self) {
99 let signal = self.clone();
100 tokio::spawn(async move {
101 if tokio::signal::ctrl_c().await.is_ok() {
102 rustlavel_core::info!("shutting down: finishing the job in hand");
103 signal.signal();
104 }
105 });
106 }
107
108 pub async fn wait_for(&self, duration: Duration) -> bool {
111 let mut receiver = self.sender.subscribe();
112 if self.is_signalled() {
113 return true;
114 }
115
116 tokio::select! {
117 _ = receiver.changed() => true,
118 _ = tokio::time::sleep(duration) => self.is_signalled(),
119 }
120 }
121}
122
123#[derive(Debug, Clone)]
125pub struct WorkerOptions {
126 pub queue: String,
128 pub poll_interval: Duration,
130 pub max_jobs: Option<u64>,
133 pub backoff_cap: Duration,
135}
136
137impl Default for WorkerOptions {
138 fn default() -> Self {
139 WorkerOptions {
140 queue: crate::job::DEFAULT_QUEUE.to_string(),
141 poll_interval: DEFAULT_POLL_INTERVAL,
142 max_jobs: None,
143 backoff_cap: DEFAULT_BACKOFF_CAP,
144 }
145 }
146}
147
148#[derive(Clone)]
150pub struct Worker {
151 queue: Arc<dyn Queue>,
152 registry: Arc<JobRegistry>,
153 options: WorkerOptions,
154 progress: Option<Arc<dyn crate::progress::ProgressStore>>,
156}
157
158impl Worker {
159 pub fn new(queue: Arc<dyn Queue>, registry: Arc<JobRegistry>) -> Self {
160 rustlavel_http::panic::install_hook();
164
165 Worker { queue, registry, options: WorkerOptions::default(), progress: None }
166 }
167
168 pub fn reporting_to(mut self, store: Arc<dyn crate::progress::ProgressStore>) -> Self {
172 self.progress = Some(store);
173 self
174 }
175
176 pub fn on_queue(mut self, queue: impl Into<String>) -> Self {
177 self.options.queue = queue.into();
178 self
179 }
180
181 pub fn poll_interval(mut self, interval: Duration) -> Self {
182 self.options.poll_interval = interval;
183 self
184 }
185
186 pub fn max_jobs(mut self, count: u64) -> Self {
188 self.options.max_jobs = Some(count);
189 self
190 }
191
192 pub fn backoff_cap(mut self, cap: Duration) -> Self {
193 self.options.backoff_cap = cap;
194 self
195 }
196
197 pub fn options(&self) -> &WorkerOptions {
198 &self.options
199 }
200
201 pub async fn run(&self, shutdown: Shutdown) -> Result<WorkerStats> {
208 let mut stats = WorkerStats::default();
209
210 loop {
211 if shutdown.is_signalled() {
212 break;
213 }
214 if self.options.max_jobs.is_some_and(|max| stats.attempts() >= max) {
215 break;
216 }
217
218 if let Some(outcome) = self.run_once().await? {
219 stats.record(outcome);
220 } else if shutdown.wait_for(self.options.poll_interval).await {
221 break;
222 }
223 }
224
225 Ok(stats)
226 }
227
228 pub async fn run_once(&self) -> Result<Option<Outcome>> {
230 let Some(reserved) = self.queue.pop(&self.options.queue).await? else {
231 return Ok(None);
232 };
233 Ok(Some(self.process(reserved).await?))
234 }
235
236 pub async fn process(&self, reserved: ReservedJob) -> Result<Outcome> {
238 let started = Instant::now();
239 let result = self.run_handler(&reserved).await;
240 let elapsed = started.elapsed();
241
242 match result {
243 Ok(()) => {
244 self.queue.delete(&reserved).await?;
245
246 if let Some(next) = reserved.job.next_in_chain() {
250 self.queue.push(next).await?;
251 }
252 self.close_progress(&reserved, None).await;
253
254 record(
255 "queue.processed",
256 self.queue.driver(),
257 &reserved,
258 elapsed,
259 |event| event,
260 );
261 Ok(Outcome::Processed)
262 }
263 Err(error) => {
264 let dead = reserved.is_last_attempt();
265
266 if dead {
267 self.queue.fail(&reserved, &error).await?;
268 self.close_progress(&reserved, Some(&error)).await;
269 rustlavel_core::error!(
270 "job `{}` failed after {} attempt(s) and was moved to failed_jobs: {error}",
271 reserved.job.name,
272 reserved.attempts
273 );
274 } else {
275 let delay = backoff(
276 reserved.job.retry_after,
277 reserved.attempts,
278 self.options.backoff_cap,
279 );
280 self.queue.release(&reserved, delay).await?;
281 rustlavel_core::warn!(
282 "job `{}` failed on attempt {} of {} and will retry in {}s: {error}",
283 reserved.job.name,
284 reserved.attempts,
285 reserved.job.max_tries,
286 delay.as_secs()
287 );
288 }
289
290 record("queue.failed", self.queue.driver(), &reserved, elapsed, |event| {
294 event.with("exception", error.as_str()).with("dead_lettered", dead)
295 });
296
297 Ok(if dead { Outcome::Failed } else { Outcome::Released })
298 }
299 }
300 }
301
302 async fn close_progress(&self, reserved: &ReservedJob, failure: Option<&str>) {
309 let Some(store) = &self.progress else { return };
310
311 let last = store.get(&reserved.id).await.ok().flatten();
312 let closed = crate::progress::Progress {
313 percent: if failure.is_none() { 100 } else { last.as_ref().map_or(0, |p| p.percent) },
314 stage: match failure {
315 None => "finished".to_string(),
316 Some(_) => last.as_ref().map_or_else(|| "failed".to_string(), |p| p.stage.clone()),
317 },
318 updated_at: crate::time::unix_now().max(0) as u64,
319 finished: true,
320 failed: failure.map(str::to_string),
321 };
322
323 let _ = store.set(&reserved.id, closed.clone()).await;
324
325 let chain_is_over = failure.is_some() || reserved.job.then.is_empty();
326 if let (Some(chain), true) = (&reserved.job.chain, chain_is_over) {
327 let _ = store.set(&format!("chain:{chain}"), closed).await;
328 }
329 }
330
331 async fn run_handler(&self, reserved: &ReservedJob) -> std::result::Result<(), String> {
339 let ctx = crate::progress::JobContext::new(
340 reserved.id.clone(),
341 reserved.job.chain.clone(),
342 self.progress.clone(),
343 );
344 let future = self.registry.run_with(&reserved.job.name, reserved.job.payload.clone(), ctx);
345
346 match rustlavel_http::panic::catch(future).await {
347 Ok(Ok(())) => Ok(()),
348 Ok(Err(error)) => Err(error.to_string()),
349 Err(payload) => {
350 let location = rustlavel_http::panic::take_location()
351 .map(|at| format!(" at {}:{}:{}", at.file, at.line, at.column))
352 .unwrap_or_default();
353 Err(format!("the job panicked{location}: {payload}"))
354 }
355 }
356 }
357}
358
359pub async fn run_pool(worker: Worker, count: usize, shutdown: Shutdown) -> Result<WorkerStats> {
365 let mut handles = Vec::with_capacity(count.max(1));
366
367 for _ in 0..count.max(1) {
368 let worker = worker.clone();
369 let shutdown = shutdown.clone();
370 handles.push(tokio::spawn(async move { worker.run(shutdown).await }));
371 }
372
373 let mut total = WorkerStats::default();
374 for handle in handles {
375 let stats = handle
376 .await
377 .map_err(|e| rustlavel_core::Error::msg(format!("a worker task did not finish: {e}")))??;
378 total.merge(stats);
379 }
380
381 Ok(total)
382}
383
384pub fn backoff(retry_after: Duration, attempts: u32, cap: Duration) -> Duration {
390 let doublings = attempts.saturating_sub(1).min(20);
392 let factor = 1u32 << doublings;
393
394 retry_after.saturating_mul(factor).min(cap)
395}
396
397fn record(
399 kind: &'static str,
400 driver: &'static str,
401 reserved: &ReservedJob,
402 elapsed: Duration,
403 extend: impl FnOnce(Event) -> Event,
404) {
405 if !events::has_subscribers() {
406 return;
407 }
408
409 let event = Event::new(kind)
410 .took(elapsed)
411 .with("job", reserved.job.name.as_str())
412 .with("queue", reserved.job.queue.as_str())
413 .with("id", reserved.id.as_str())
414 .with("driver", driver)
415 .with("attempt", reserved.attempts)
416 .with("max_tries", reserved.job.max_tries);
417
418 extend(event).dispatch();
419}
420
421pub fn outcome_to_json(outcome: Outcome) -> Json {
423 Json::from(match outcome {
424 Outcome::Processed => "processed",
425 Outcome::Released => "released",
426 Outcome::Failed => "failed",
427 })
428}
429
430#[cfg(test)]
431mod tests {
432 use super::*;
433 use crate::job::{Job, QueuedJob};
434 use crate::memory::MemoryQueue;
435 use crate::queue::QueueExt;
436 use crate::tests_support::{
437 CountingJob, FailingJob, FlakyJob, PanickingJob, SlowJob, counter, registry,
438 };
439
440 fn worker() -> (Arc<MemoryQueue>, Worker) {
442 let queue = Arc::new(MemoryQueue::new());
443 let worker = Worker::new(queue.clone(), Arc::new(registry()))
444 .poll_interval(Duration::from_millis(5));
445 (queue, worker)
446 }
447
448 #[tokio::test]
449 async fn a_worker_runs_a_job_and_removes_it() {
450 let (queue, worker) = worker();
451 let tally = counter();
452 queue.dispatch(&CountingJob::new(tally, 7)).await.unwrap();
453
454 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Processed));
455
456 assert_eq!(tally.runs(), 1);
457 assert_eq!(tally.total(), 7);
458 assert_eq!(queue.size("default").await.unwrap(), 0);
459 assert!(queue.failed_jobs().await.unwrap().is_empty());
460 }
461
462 #[tokio::test]
463 async fn an_empty_queue_gives_the_worker_nothing_to_do() {
464 let (_queue, worker) = worker();
465 assert_eq!(worker.run_once().await.unwrap(), None);
466 }
467
468 #[tokio::test]
469 async fn a_worker_only_takes_from_the_queue_it_was_pointed_at() {
470 let (queue, worker) = worker();
471 let tally = counter();
472 queue
473 .push_on("emails", CountingJob::new(tally, 1).to_queued())
474 .await
475 .unwrap();
476
477 assert_eq!(worker.run_once().await.unwrap(), None, "nothing on `default`");
478
479 let emails = worker.clone().on_queue("emails");
480 assert_eq!(emails.run_once().await.unwrap(), Some(Outcome::Processed));
481 assert_eq!(tally.runs(), 1);
482 }
483
484 #[tokio::test]
485 async fn a_failing_job_is_retried_and_then_dead_lettered() {
486 let (queue, worker) = worker();
487 let tally = counter();
488 queue.dispatch(&FailingJob::new(tally, 3)).await.unwrap();
489
490 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Released));
492 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Released));
493 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Failed));
494 assert_eq!(worker.run_once().await.unwrap(), None, "nothing is left to run");
495
496 assert_eq!(tally.runs(), 3, "it ran exactly as many times as it had tries");
497 assert_eq!(queue.size("default").await.unwrap(), 0);
498
499 let failed = queue.failed_jobs().await.unwrap();
500 assert_eq!(failed.len(), 1);
501 assert_eq!(failed[0].name, "failing");
502 assert_eq!(failed[0].attempts, 3);
503 assert!(failed[0].error.contains("attempt 3 went wrong"), "{}", failed[0].error);
504 }
505
506 #[tokio::test]
507 async fn a_job_that_recovers_never_reaches_the_dead_letter_store() {
508 let (queue, worker) = worker();
509 let tally = counter();
510 queue.dispatch(&FlakyJob::new(tally, 3, 5)).await.unwrap();
511
512 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Released));
513 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Released));
514 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Processed));
515
516 assert_eq!(tally.runs(), 3);
517 assert!(queue.failed_jobs().await.unwrap().is_empty());
518 }
519
520 #[tokio::test]
521 async fn a_job_with_one_try_is_buried_on_its_first_failure() {
522 let (queue, worker) = worker();
523 let tally = counter();
524 queue.dispatch(&FailingJob::new(tally, 1)).await.unwrap();
525
526 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Failed));
527 assert_eq!(queue.failed_jobs().await.unwrap().len(), 1);
528 }
529
530 #[tokio::test]
531 async fn a_panicking_job_fails_cleanly_instead_of_killing_the_worker() {
532 let (queue, worker) = worker();
533 let tally = counter();
534 let healthy = counter();
535
536 queue.dispatch(&PanickingJob::new(tally, 1)).await.unwrap();
537 queue.dispatch(&CountingJob::new(healthy, 1)).await.unwrap();
538
539 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Failed));
540
541 let failed = queue.failed_jobs().await.unwrap();
542 assert_eq!(failed.len(), 1);
543 assert!(
544 failed[0].error.contains("this job panicked on purpose"),
545 "the panic message should survive: {}",
546 failed[0].error
547 );
548
549 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Processed));
551 assert_eq!(healthy.runs(), 1);
552 }
553
554 #[tokio::test]
555 async fn a_job_with_no_registered_handler_fails_rather_than_vanishing() {
556 let (queue, worker) = worker();
557 queue.push(QueuedJob::new("never-registered", Json::Null).with_tries(1)).await.unwrap();
558
559 assert_eq!(worker.run_once().await.unwrap(), Some(Outcome::Failed));
560
561 let failed = queue.failed_jobs().await.unwrap();
562 assert!(failed[0].error.contains("no handler is registered"), "{}", failed[0].error);
563 }
564
565 #[tokio::test]
566 async fn max_jobs_stops_the_loop_without_a_shutdown_signal() {
567 let (queue, worker) = worker();
568 let tally = counter();
569 for _ in 0..5 {
570 queue.dispatch(&CountingJob::new(tally, 1)).await.unwrap();
571 }
572
573 let stats = worker.max_jobs(2).run(Shutdown::new()).await.unwrap();
574
575 assert_eq!(stats.processed, 2);
576 assert_eq!(tally.runs(), 2);
577 assert_eq!(queue.size("default").await.unwrap(), 3, "the rest are still waiting");
578 }
579
580 #[tokio::test]
581 async fn an_already_signalled_shutdown_starts_no_work_at_all() {
582 let (queue, worker) = worker();
583 let tally = counter();
584 queue.dispatch(&CountingJob::new(tally, 1)).await.unwrap();
585
586 let shutdown = Shutdown::new();
587 shutdown.signal();
588
589 let stats = worker.run(shutdown).await.unwrap();
590
591 assert_eq!(stats, WorkerStats::default());
592 assert_eq!(tally.runs(), 0);
593 assert_eq!(queue.size("default").await.unwrap(), 1);
594 }
595
596 #[tokio::test]
597 async fn shutdown_lets_the_job_in_hand_finish_first() {
598 let (queue, worker) = worker();
599 let tally = counter();
600 queue.dispatch(&SlowJob::new(tally, 150)).await.unwrap();
601 queue.dispatch(&CountingJob::new(tally, 1)).await.unwrap();
602
603 let shutdown = Shutdown::new();
604 let signal = shutdown.clone();
605 tokio::spawn(async move {
606 tokio::time::sleep(Duration::from_millis(40)).await;
608 signal.signal();
609 });
610
611 let stats = worker.run(shutdown).await.unwrap();
612
613 assert_eq!(stats.processed, 1, "the slow job ran to completion");
614 assert_eq!(tally.runs(), 1);
615 assert_eq!(queue.size("default").await.unwrap(), 1, "the second job was left alone");
616 }
617
618 #[tokio::test]
619 async fn an_idle_worker_wakes_up_for_shutdown_instead_of_waiting_out_its_poll() {
620 let (_queue, worker) = worker();
621 let worker = worker.poll_interval(Duration::from_secs(30));
622
623 let shutdown = Shutdown::new();
624 let signal = shutdown.clone();
625 tokio::spawn(async move {
626 tokio::time::sleep(Duration::from_millis(20)).await;
627 signal.signal();
628 });
629
630 let started = Instant::now();
631 worker.run(shutdown).await.unwrap();
632
633 assert!(
634 started.elapsed() < Duration::from_secs(5),
635 "an idle worker should stop as soon as it is asked, not after its poll interval"
636 );
637 }
638
639 #[tokio::test]
640 async fn several_workers_share_one_queue_and_run_each_job_once() {
641 let (queue, worker) = worker();
642 let tally = counter();
643 for _ in 0..40 {
644 queue.dispatch(&CountingJob::new(tally, 1)).await.unwrap();
645 }
646
647 let shutdown = Shutdown::new();
648 let signal = shutdown.clone();
649 let drain = queue.clone();
650 tokio::spawn(async move {
651 while drain.size("default").await.unwrap_or(0) > 0 {
652 tokio::time::sleep(Duration::from_millis(5)).await;
653 }
654 signal.signal();
655 });
656
657 let stats = run_pool(worker, 4, shutdown).await.unwrap();
658
659 assert_eq!(stats.processed, 40);
660 assert_eq!(tally.runs(), 40, "no job ran twice and none was skipped");
661 assert_eq!(tally.total(), 40);
662 }
663
664 #[tokio::test]
665 async fn the_worker_reports_what_it_did_on_the_event_bus() {
666 let (queue, worker) = worker();
670 let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
671 let sink = Arc::clone(&seen);
672
673 events::subscribe(move |event: &Event| {
674 if event.field("job").and_then(Json::as_str) == Some("failing") {
675 sink.lock()
676 .unwrap()
677 .push((event.kind, event.field("attempt").and_then(Json::as_i64)));
678 }
679 });
680
681 queue.dispatch(&FailingJob::new(counter(), 2)).await.unwrap();
682 worker.run_once().await.unwrap();
683 worker.run_once().await.unwrap();
684
685 let seen = seen.lock().unwrap().clone();
686 assert!(seen.contains(&("queue.pushed", None)), "{seen:?}");
687 assert!(seen.contains(&("queue.failed", Some(1))), "{seen:?}");
688 assert!(seen.contains(&("queue.failed", Some(2))), "{seen:?}");
689
690 events::clear_subscribers();
691 }
692
693 #[test]
694 fn backoff_doubles_per_attempt_and_then_stops_at_the_cap() {
695 let base = Duration::from_secs(10);
696 let cap = Duration::from_secs(300);
697
698 assert_eq!(backoff(base, 1, cap), Duration::from_secs(10));
699 assert_eq!(backoff(base, 2, cap), Duration::from_secs(20));
700 assert_eq!(backoff(base, 3, cap), Duration::from_secs(40));
701 assert_eq!(backoff(base, 4, cap), Duration::from_secs(80));
702 assert_eq!(backoff(base, 10, cap), cap, "the cap wins eventually");
703 assert_eq!(backoff(base, 5000, cap), cap, "and never overflows");
704 }
705
706 #[test]
707 fn backoff_of_zero_stays_zero_so_a_test_can_retry_immediately() {
708 assert_eq!(backoff(Duration::ZERO, 9, DEFAULT_BACKOFF_CAP), Duration::ZERO);
709 }
710
711 #[test]
712 fn stats_add_up_across_a_pool() {
713 let mut stats = WorkerStats { processed: 1, released: 2, failed: 3 };
714 stats.merge(WorkerStats { processed: 10, released: 20, failed: 30 });
715
716 assert_eq!(stats, WorkerStats { processed: 11, released: 22, failed: 33 });
717 assert_eq!(stats.attempts(), 66);
718 }
719
720 #[test]
721 fn outcomes_render_for_a_dashboard() {
722 assert_eq!(outcome_to_json(Outcome::Processed).as_str(), Some("processed"));
723 assert_eq!(outcome_to_json(Outcome::Released).as_str(), Some("released"));
724 assert_eq!(outcome_to_json(Outcome::Failed).as_str(), Some("failed"));
725 }
726}