1use std::collections::HashMap;
37use std::future::Future;
38use std::pin::Pin;
39use std::sync::{Arc, RwLock};
40use std::time::Duration;
41
42use async_trait::async_trait;
43use chrono::{DateTime, SecondsFormat, Utc};
44use serde::de::DeserializeOwned;
45use serde::{Deserialize, Serialize};
46
47use crate::cache::Cache;
48use crate::config::QueueConfig;
49use crate::error::{CoreError, CoreResult};
50use crate::events::{Event, Events};
51use crate::parameters::Parameters;
52use crate::registry::Registry;
53use crate::{Db, Module, ModuleId};
54
55mod database;
56pub mod migrations;
57#[cfg(feature = "redis")]
58pub mod redis;
59mod worker;
60
61pub use database::DatabaseQueue;
62pub use worker::{WorkOptions, Worker};
63
64pub struct QueueModule;
67
68impl Module for QueueModule {
69 fn id(&self) -> ModuleId {
70 ModuleId::new(migrations::MODULE_ID)
71 }
72 fn migrations(&self) -> crate::MigrationSet {
73 migrations::migrations()
74 }
75}
76
77pub const DEFAULT_QUEUE: &str = "default";
79
80#[derive(Debug, Clone, PartialEq, Eq)]
82pub enum JobError {
83 Retry(String),
85 Fail(String),
87}
88
89impl JobError {
90 pub fn retry(message: impl std::fmt::Display) -> Self {
92 Self::Retry(message.to_string())
93 }
94
95 pub fn fail(message: impl std::fmt::Display) -> Self {
97 Self::Fail(message.to_string())
98 }
99
100 pub fn message(&self) -> &str {
101 match self {
102 Self::Retry(m) | Self::Fail(m) => m,
103 }
104 }
105}
106
107impl std::fmt::Display for JobError {
108 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109 f.write_str(self.message())
110 }
111}
112
113impl std::error::Error for JobError {}
114
115impl From<CoreError> for JobError {
116 fn from(e: CoreError) -> Self {
117 Self::Retry(e.to_string())
118 }
119}
120
121impl From<sqlx::Error> for JobError {
122 fn from(e: sqlx::Error) -> Self {
123 Self::Retry(e.to_string())
124 }
125}
126
127#[derive(Clone)]
129pub struct JobCx {
130 db: Db,
131 events: Events,
132 cache: Cache,
133 parameters: Parameters,
134 id: Option<i64>,
135 attempt: u32,
136 queue: String,
137}
138
139impl JobCx {
140 pub fn db(&self) -> &Db {
141 &self.db
142 }
143 pub fn events(&self) -> &Events {
144 &self.events
145 }
146 pub fn cache(&self) -> &Cache {
147 &self.cache
148 }
149 pub fn parameters(&self) -> &Parameters {
150 &self.parameters
151 }
152 pub fn id(&self) -> Option<i64> {
154 self.id
155 }
156 pub fn attempt(&self) -> u32 {
158 self.attempt
159 }
160 pub fn queue(&self) -> &str {
161 &self.queue
162 }
163}
164
165#[async_trait]
168pub trait Job: Serialize + DeserializeOwned + Send + Sync + 'static {
169 const NAME: &'static str;
171
172 async fn run(&self, cx: &JobCx) -> Result<(), JobError>;
173
174 fn queue(&self) -> String {
176 DEFAULT_QUEUE.to_string()
177 }
178
179 fn tries(&self) -> u32 {
182 3
183 }
184
185 fn backoff(&self, attempt: u32) -> Duration {
188 Duration::from_secs(10u64 << attempt.saturating_sub(1).min(8))
189 }
190
191 fn timeout(&self) -> Option<Duration> {
194 Some(Duration::from_secs(60))
195 }
196
197 fn retry_until(&self) -> Option<DateTime<Utc>> {
200 None
201 }
202
203 fn unique_key(&self) -> Option<String> {
206 None
207 }
208}
209
210#[derive(Debug, Clone)]
212pub struct JobMeta {
213 pub tries: u32,
214 pub timeout: Option<Duration>,
215 pub retry_until: Option<DateTime<Utc>>,
216}
217
218type RunFn = dyn Fn(serde_json::Value, JobCx) -> Pin<Box<dyn Future<Output = Result<(), JobError>> + Send>>
219 + Send
220 + Sync;
221type MetaFn = dyn Fn(&serde_json::Value) -> Result<JobMeta, String> + Send + Sync;
222type BackoffFn = dyn Fn(&serde_json::Value, u32) -> Duration + Send + Sync;
223
224#[derive(Clone)]
227pub struct JobReg {
228 name: &'static str,
229 run: Arc<RunFn>,
230 meta: Arc<MetaFn>,
231 backoff: Arc<BackoffFn>,
232}
233
234impl std::fmt::Debug for JobReg {
235 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
236 f.debug_struct("JobReg").field("name", &self.name).finish()
237 }
238}
239
240impl JobReg {
241 pub fn of<T: Job>() -> Self {
243 Self {
244 name: T::NAME,
245 run: Arc::new(|payload, cx| {
246 Box::pin(async move {
247 let job: T = serde_json::from_value(payload)
248 .map_err(|e| JobError::Fail(format!("payload does not decode: {e}")))?;
249 job.run(&cx).await
250 })
251 }),
252 meta: Arc::new(|payload| {
253 let job: T = serde_json::from_value(payload.clone())
254 .map_err(|e| format!("payload does not decode: {e}"))?;
255 Ok(JobMeta {
256 tries: job.tries(),
257 timeout: job.timeout(),
258 retry_until: job.retry_until(),
259 })
260 }),
261 backoff: Arc::new(|payload, attempt| {
262 serde_json::from_value::<T>(payload.clone())
263 .map(|job| job.backoff(attempt))
264 .unwrap_or(Duration::from_secs(10))
265 }),
266 }
267 }
268
269 pub fn with_handler<T, F, Fut>(handler: F) -> Self
274 where
275 T: Job,
276 F: Fn(T, JobCx) -> Fut + Send + Sync + 'static,
277 Fut: Future<Output = Result<(), JobError>> + Send + 'static,
278 {
279 let handler = Arc::new(handler);
280 let mut reg = Self::of::<T>();
281 reg.run = Arc::new(move |payload, cx| {
282 let handler = Arc::clone(&handler);
283 Box::pin(async move {
284 let job: T = serde_json::from_value(payload)
285 .map_err(|e| JobError::Fail(format!("payload does not decode: {e}")))?;
286 handler(job, cx).await
287 })
288 });
289 reg
290 }
291
292 pub fn name(&self) -> &'static str {
293 self.name
294 }
295}
296
297pub trait QueueRegistry {
299 fn add_job<T: Job>(&mut self);
300}
301
302impl QueueRegistry for Registry {
303 fn add_job<T: Job>(&mut self) {
304 self.add(JobReg::of::<T>());
305 }
306}
307
308#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
310pub struct Queued {
311 pub id: i64,
312 pub queue: String,
313 pub name: String,
314 pub payload: serde_json::Value,
315 pub attempts: u32,
317 pub available_at: String,
318 pub reserved_at: Option<String>,
319 pub reserved_by: Option<String>,
320 pub created_at: String,
321}
322
323#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
325pub struct FailedJob {
326 pub id: i64,
327 pub queue: String,
328 pub name: String,
329 pub payload: serde_json::Value,
330 pub error: String,
331 pub failed_at: String,
332}
333
334#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
336pub struct QueueCount {
337 pub queue: String,
338 pub pending: u64,
340 pub reserved: u64,
342}
343
344#[derive(Debug, Clone, Copy, PartialEq, Eq)]
346pub enum Pushed {
347 Queued(i64),
349 Ran,
351 Duplicate,
353}
354
355#[async_trait]
358pub trait QueueDriver: Send + Sync {
359 fn name(&self) -> &'static str;
360 async fn push(
362 &self,
363 queue: &str,
364 name: &str,
365 payload: &serde_json::Value,
366 available_at: DateTime<Utc>,
367 unique_key: Option<&str>,
368 ) -> CoreResult<Option<i64>>;
369 async fn claim(
372 &self,
373 queues: &[String],
374 retry_after: Duration,
375 worker: &str,
376 ) -> CoreResult<Option<Queued>>;
377 async fn release(&self, id: i64, available_at: DateTime<Utc>) -> CoreResult<()>;
379 async fn delete(&self, id: i64) -> CoreResult<()>;
381 async fn fail(&self, job: &Queued, error: &str) -> CoreResult<i64>;
383 async fn counts(&self) -> CoreResult<Vec<QueueCount>>;
384 async fn failed(&self) -> CoreResult<Vec<FailedJob>>;
385 async fn retry(&self, id: Option<i64>) -> CoreResult<usize>;
388 async fn forget(&self, id: Option<i64>) -> CoreResult<usize>;
390 async fn clear(&self, queue: Option<&str>) -> CoreResult<usize>;
392 async fn ping(&self) -> CoreResult<()> {
394 Ok(())
395 }
396}
397
398#[derive(Debug, Clone, Copy, PartialEq, Eq)]
400pub enum Mode {
401 Worker,
403 External,
405 Inline,
407}
408
409impl Mode {
410 pub fn parse(s: &str) -> CoreResult<Self> {
411 match s {
412 "worker" => Ok(Self::Worker),
413 "external" => Ok(Self::External),
414 "inline" => Ok(Self::Inline),
415 other => Err(CoreError::Config(format!(
416 "queue.mode `{other}` is not one of worker, external, inline"
417 ))),
418 }
419 }
420}
421
422#[derive(Clone)]
424pub struct Services {
425 pub db: Db,
426 pub events: Events,
427 pub cache: Cache,
428 pub parameters: Parameters,
429}
430
431struct Inner {
432 driver: Box<dyn QueueDriver>,
433 jobs: RwLock<HashMap<&'static str, JobReg>>,
434 services: Services,
435 mode: Mode,
436 queues: Vec<String>,
437 retry_after: Duration,
438 poll: Duration,
439 workers: usize,
440 wake: tokio::sync::Notify,
442}
443
444#[derive(Clone)]
446pub struct Queue {
447 inner: Arc<Inner>,
448}
449
450impl std::fmt::Debug for Queue {
451 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
452 f.debug_struct("Queue")
453 .field("driver", &self.inner.driver.name())
454 .field("mode", &self.inner.mode)
455 .field("jobs", &self.job_names().len())
456 .finish_non_exhaustive()
457 }
458}
459
460#[derive(Debug, Clone)]
462pub struct Options {
463 pub mode: Mode,
464 pub queues: Vec<String>,
466 pub retry_after: Duration,
468 pub poll: Duration,
470 pub workers: usize,
472}
473
474impl Options {
475 pub fn from_config(config: &QueueConfig) -> CoreResult<Self> {
477 let mut queues = config.queues.clone();
478 if queues.is_empty() {
479 queues.push(DEFAULT_QUEUE.to_string());
480 }
481 Ok(Self {
482 mode: Mode::parse(&config.mode)?,
483 queues,
484 retry_after: Duration::from_secs(config.retry_after_seconds.max(1)),
485 poll: Duration::from_secs(config.poll_seconds.max(1)),
486 workers: config.workers,
487 })
488 }
489}
490
491impl Default for Options {
492 fn default() -> Self {
493 Self {
494 mode: Mode::Worker,
495 queues: vec![DEFAULT_QUEUE.to_string()],
496 retry_after: Duration::from_secs(90),
497 poll: Duration::from_secs(1),
498 workers: 2,
499 }
500 }
501}
502
503impl Queue {
504 pub fn new(
506 driver: impl QueueDriver + 'static,
507 jobs: impl IntoIterator<Item = JobReg>,
508 services: Services,
509 options: Options,
510 ) -> Self {
511 let jobs = jobs.into_iter().map(|j| (j.name, j)).collect();
512 Self {
513 inner: Arc::new(Inner {
514 driver: Box::new(driver),
515 jobs: RwLock::new(jobs),
516 services,
517 mode: options.mode,
518 queues: options.queues,
519 retry_after: options.retry_after,
520 poll: options.poll,
521 workers: options.workers,
522 wake: tokio::sync::Notify::new(),
523 }),
524 }
525 }
526
527 pub fn inline(jobs: impl IntoIterator<Item = JobReg>, services: Services) -> Self {
530 Self::new(
531 NoQueue,
532 jobs,
533 services,
534 Options {
535 mode: Mode::Inline,
536 ..Options::default()
537 },
538 )
539 }
540
541 pub fn mode(&self) -> Mode {
542 self.inner.mode
543 }
544
545 pub fn driver_name(&self) -> &'static str {
546 self.inner.driver.name()
547 }
548
549 pub fn job_names(&self) -> Vec<&'static str> {
551 let jobs = self.inner.jobs.read().unwrap_or_else(|e| e.into_inner());
552 let mut names: Vec<&'static str> = jobs.keys().copied().collect();
553 names.sort_unstable();
554 names
555 }
556
557 pub fn register(&self, reg: JobReg) {
560 self.inner
561 .jobs
562 .write()
563 .unwrap_or_else(|e| e.into_inner())
564 .insert(reg.name, reg);
565 }
566
567 pub(crate) fn job(&self, name: &str) -> Option<JobReg> {
568 self.inner
569 .jobs
570 .read()
571 .unwrap_or_else(|e| e.into_inner())
572 .get(name)
573 .cloned()
574 }
575
576 pub(crate) fn services(&self) -> &Services {
577 &self.inner.services
578 }
579
580 pub(crate) fn driver(&self) -> &dyn QueueDriver {
581 self.inner.driver.as_ref()
582 }
583
584 pub(crate) fn queues(&self) -> &[String] {
585 &self.inner.queues
586 }
587
588 pub(crate) fn retry_after(&self) -> Duration {
589 self.inner.retry_after
590 }
591
592 pub(crate) fn poll(&self) -> Duration {
593 self.inner.poll
594 }
595
596 pub(crate) fn wake(&self) -> &tokio::sync::Notify {
597 &self.inner.wake
598 }
599
600 fn check<T: Job>(&self) -> CoreResult<()> {
601 if self.job(T::NAME).is_some() {
602 Ok(())
603 } else {
604 Err(CoreError::Validation(format!(
605 "job `{}` is not registered; add `registry.add_job::<{}>()` to a module",
606 T::NAME,
607 std::any::type_name::<T>()
608 )))
609 }
610 }
611
612 pub async fn push<T: Job>(&self, job: &T) -> CoreResult<Pushed> {
614 self.push_at(job, Utc::now()).await
615 }
616
617 pub async fn push_in<T: Job>(&self, job: &T, delay: Duration) -> CoreResult<Pushed> {
619 let delay = chrono::Duration::from_std(delay).unwrap_or(chrono::Duration::MAX);
620 self.push_at(job, Utc::now() + delay).await
621 }
622
623 pub async fn push_at<T: Job>(&self, job: &T, when: DateTime<Utc>) -> CoreResult<Pushed> {
625 self.check::<T>()?;
626 let payload = serde_json::to_value(job)
627 .map_err(|e| CoreError::Internal(format!("job `{}` cannot serialize: {e}", T::NAME)))?;
628 if self.inner.mode == Mode::Inline {
629 let reg = self.job(T::NAME).expect("checked above");
632 let cx = self.context(None, 1, &job.queue());
633 return match (reg.run)(payload, cx).await {
634 Ok(()) => Ok(Pushed::Ran),
635 Err(e) => Err(CoreError::Internal(format!(
636 "job `{}` failed: {e}",
637 T::NAME
638 ))),
639 };
640 }
641 let unique = job.unique_key();
642 let stored = self
643 .inner
644 .driver
645 .push(&job.queue(), T::NAME, &payload, when, unique.as_deref())
646 .await?;
647 match stored {
648 Some(id) => {
649 self.inner.wake.notify_waiters();
650 tracing::debug!(job = T::NAME, id, "job queued");
651 Ok(Pushed::Queued(id))
652 }
653 None => Ok(Pushed::Duplicate),
654 }
655 }
656
657 pub async fn push_named(
660 &self,
661 name: &str,
662 payload: &serde_json::Value,
663 queue: &str,
664 ) -> CoreResult<Pushed> {
665 let Some(reg) = self.job(name) else {
666 return Err(CoreError::Validation(format!(
667 "job `{name}` is not registered"
668 )));
669 };
670 if self.inner.mode == Mode::Inline {
671 let cx = self.context(None, 1, queue);
672 return match (reg.run)(payload.clone(), cx).await {
673 Ok(()) => Ok(Pushed::Ran),
674 Err(e) => Err(CoreError::Internal(format!("job `{name}` failed: {e}"))),
675 };
676 }
677 let stored = self
678 .inner
679 .driver
680 .push(queue, name, payload, Utc::now(), None)
681 .await?;
682 match stored {
683 Some(id) => {
684 self.inner.wake.notify_waiters();
685 Ok(Pushed::Queued(id))
686 }
687 None => Ok(Pushed::Duplicate),
688 }
689 }
690
691 pub(crate) fn context(&self, id: Option<i64>, attempt: u32, queue: &str) -> JobCx {
692 let s = &self.inner.services;
693 JobCx {
694 db: s.db.clone(),
695 events: s.events.clone(),
696 cache: s.cache.clone(),
697 parameters: s.parameters.clone(),
698 id,
699 attempt,
700 queue: queue.to_string(),
701 }
702 }
703
704 pub async fn counts(&self) -> CoreResult<Vec<QueueCount>> {
706 self.inner.driver.counts().await
707 }
708
709 pub async fn failed(&self) -> CoreResult<Vec<FailedJob>> {
711 self.inner.driver.failed().await
712 }
713
714 pub async fn retry(&self, id: Option<i64>) -> CoreResult<usize> {
716 let moved = self.inner.driver.retry(id).await?;
717 if moved > 0 {
718 self.inner.wake.notify_waiters();
719 }
720 Ok(moved)
721 }
722
723 pub async fn forget(&self, id: Option<i64>) -> CoreResult<usize> {
725 self.inner.driver.forget(id).await
726 }
727
728 pub async fn clear(&self, queue: Option<&str>) -> CoreResult<usize> {
730 self.inner.driver.clear(queue).await
731 }
732
733 pub async fn work(
736 &self,
737 options: WorkOptions,
738 stop: tokio::sync::watch::Receiver<bool>,
739 ) -> CoreResult<usize> {
740 Worker::new(self.clone(), options).run(stop).await
741 }
742
743 pub fn spawn_workers(
746 &self,
747 stop: tokio::sync::watch::Receiver<bool>,
748 ) -> Vec<tokio::task::JoinHandle<()>> {
749 (0..self.inner.workers)
750 .map(|n| {
751 let queue = self.clone();
752 let stop = stop.clone();
753 tokio::spawn(async move {
754 let options = WorkOptions::default().name(format!("worker-{n}"));
755 if let Err(e) = queue.work(options, stop).await {
756 tracing::error!(worker = n, error = %e, "worker stopped");
757 }
758 })
759 })
760 .collect()
761 }
762}
763
764struct NoQueue;
766
767#[async_trait]
768impl QueueDriver for NoQueue {
769 fn name(&self) -> &'static str {
770 "inline"
771 }
772 async fn push(
773 &self,
774 _queue: &str,
775 _name: &str,
776 _payload: &serde_json::Value,
777 _available_at: DateTime<Utc>,
778 _unique_key: Option<&str>,
779 ) -> CoreResult<Option<i64>> {
780 Ok(None)
781 }
782 async fn claim(
783 &self,
784 _queues: &[String],
785 _retry_after: Duration,
786 _worker: &str,
787 ) -> CoreResult<Option<Queued>> {
788 Ok(None)
789 }
790 async fn release(&self, _id: i64, _available_at: DateTime<Utc>) -> CoreResult<()> {
791 Ok(())
792 }
793 async fn delete(&self, _id: i64) -> CoreResult<()> {
794 Ok(())
795 }
796 async fn fail(&self, _job: &Queued, _error: &str) -> CoreResult<i64> {
797 Ok(0)
798 }
799 async fn counts(&self) -> CoreResult<Vec<QueueCount>> {
800 Ok(Vec::new())
801 }
802 async fn failed(&self) -> CoreResult<Vec<FailedJob>> {
803 Ok(Vec::new())
804 }
805 async fn retry(&self, _id: Option<i64>) -> CoreResult<usize> {
806 Ok(0)
807 }
808 async fn forget(&self, _id: Option<i64>) -> CoreResult<usize> {
809 Ok(0)
810 }
811 async fn clear(&self, _queue: Option<&str>) -> CoreResult<usize> {
812 Ok(0)
813 }
814}
815
816#[derive(Debug, Clone, Serialize, Deserialize)]
818pub struct JobFailed {
819 pub id: i64,
820 pub name: String,
821 pub queue: String,
822 pub attempts: u32,
823 pub error: String,
824 pub payload: serde_json::Value,
826}
827
828impl Event for JobFailed {
829 const NAME: &'static str = "queue.job_failed";
830}
831
832pub async fn from_config(
835 config: &QueueConfig,
836 jobs: impl IntoIterator<Item = JobReg>,
837 services: Services,
838) -> CoreResult<Queue> {
839 let options = Options::from_config(config)?;
840 match config.store.as_str() {
841 "auto" | "database" => Ok(Queue::new(
842 DatabaseQueue::new(services.db.clone()),
843 jobs,
844 services,
845 options,
846 )),
847 #[cfg(feature = "redis")]
848 "redis" => {
849 let url = config.url.as_deref().ok_or_else(|| {
850 CoreError::Config("queue.store = \"redis\" needs queue.url".to_string())
851 })?;
852 let driver = redis::connect(url, &config.prefix).await?;
853 driver.ping().await?;
854 Ok(Queue::new(driver, jobs, services, options))
855 }
856 #[cfg(not(feature = "redis"))]
857 "redis" => Err(CoreError::Config(
858 "queue.store = \"redis\" needs the `redis` feature of laterite-core".to_string(),
859 )),
860 other => Err(CoreError::Config(format!(
861 "queue.store `{other}` is not one of auto, database, redis"
862 ))),
863 }
864}
865
866pub fn time_text(at: DateTime<Utc>) -> String {
867 at.to_rfc3339_opts(SecondsFormat::Micros, true)
868}
869
870#[cfg(test)]
871mod tests {
872 use super::*;
873
874 #[test]
875 fn the_default_backoff_doubles_from_ten_seconds() {
876 #[derive(Serialize, Deserialize)]
877 struct Noop;
878 #[async_trait]
879 impl Job for Noop {
880 const NAME: &'static str = "test::noop";
881 async fn run(&self, _cx: &JobCx) -> Result<(), JobError> {
882 Ok(())
883 }
884 }
885 assert_eq!(Noop.backoff(1), Duration::from_secs(10));
886 assert_eq!(Noop.backoff(2), Duration::from_secs(20));
887 assert_eq!(Noop.backoff(4), Duration::from_secs(80));
888 assert_eq!(Noop.backoff(40), Duration::from_secs(2560));
889 assert_eq!(Noop.tries(), 3);
890 }
891
892 #[test]
893 fn a_mode_parses_or_is_refused() {
894 assert_eq!(Mode::parse("inline").unwrap(), Mode::Inline);
895 assert!(Mode::parse("sometimes").is_err());
896 }
897}