1use std::sync::{Arc, Mutex};
5use std::time::Duration;
6
7use rainier_cache::{LockGuard, LockManager};
8use rainier_container::Container;
9use rainier_support::{BoxFuture, Result};
10
11use crate::job::{Job, JobContext, JobRegistry, QueuedJob};
12use crate::queue::Queue;
13
14pub struct SyncQueue {
22 registry: Arc<JobRegistry>,
23 container: Arc<Container>,
24}
25
26impl SyncQueue {
27 pub fn new(registry: Arc<JobRegistry>, container: Arc<Container>) -> Self {
29 Self { registry, container }
30 }
31}
32
33impl Queue for SyncQueue {
34 fn name(&self) -> &str {
35 "sync"
36 }
37
38 fn push<'a>(&'a self, mut job: QueuedJob) -> BoxFuture<'a, Result<String>> {
39 Box::pin(async move {
40 job.attempts = 1;
41 let id = job.id.clone();
42
43 let context = Arc::new(JobContext::new(
44 Arc::clone(&self.container),
45 job.id.clone(),
46 job.queue.clone(),
47 1,
50 1,
51 ));
52
53 self.registry.run(&job, context).await?;
54 Ok(id)
55 })
56 }
57
58 fn reserve<'a>(&'a self, _queue: &'a str) -> BoxFuture<'a, Result<Option<QueuedJob>>> {
59 Box::pin(async { Ok(None) })
61 }
62
63 fn acknowledge<'a>(&'a self, _job: &'a QueuedJob) -> BoxFuture<'a, Result<()>> {
64 Box::pin(async { Ok(()) })
65 }
66
67 fn release<'a>(&'a self, _job: &'a QueuedJob, _delay: Duration) -> BoxFuture<'a, Result<()>> {
68 Box::pin(async { Ok(()) })
69 }
70
71 fn fail<'a>(&'a self, _job: &'a QueuedJob, _error: &'a str) -> BoxFuture<'a, Result<()>> {
72 Box::pin(async { Ok(()) })
73 }
74
75 fn size<'a>(&'a self, _queue: &'a str) -> BoxFuture<'a, Result<u64>> {
76 Box::pin(async { Ok(0) })
77 }
78
79 fn clear<'a>(&'a self, _queue: &'a str) -> BoxFuture<'a, Result<u64>> {
80 Box::pin(async { Ok(0) })
81 }
82}
83
84pub struct PendingDispatch<'a> {
88 manager: &'a QueueManager,
89 job: QueuedJob,
90 unique: Option<(String, Duration)>,
96}
97
98impl<'a> PendingDispatch<'a> {
99 pub fn on_queue(mut self, queue: impl Into<String>) -> Self {
101 self.job = self.job.on_queue(queue);
102 self
103 }
104
105 pub fn delay(mut self, delay: Duration) -> Self {
107 self.job = self.job.delayed_by(delay);
108 self
109 }
110
111 pub fn tries(mut self, tries: u32) -> Self {
113 self.job.max_attempts = tries.max(1);
114 self
115 }
116
117 pub async fn send(self) -> Result<Option<String>> {
123 let mut job = self.job;
124
125 let claim = match &self.unique {
126 Some((key, ttl)) => {
127 let claim = self.manager.claim(key, *ttl).await?;
128 if matches!(claim, Claim::AlreadyQueued) {
129 tracing::debug!(key, "dropping a duplicate dispatch");
130 return Ok(None);
131 }
132 if matches!(claim, Claim::Held(_)) {
136 job.unique_key = Some(key.clone());
137 }
138 claim
139 }
140 None => Claim::NotEnforced,
141 };
142
143 let id = self.manager.push(job).await?;
144
145 if let Claim::Held(guard) = claim {
149 guard.keep();
150 }
151
152 Ok(Some(id))
153 }
154}
155
156enum Claim {
158 Held(LockGuard),
160 AlreadyQueued,
162 NotEnforced,
165}
166
167pub struct QueueManager {
171 queue: Arc<dyn Queue>,
172 registry: Arc<JobRegistry>,
173 recorded: Option<Mutex<Vec<QueuedJob>>>,
175 locks: Option<LockManager>,
181}
182
183impl QueueManager {
184 pub fn new(queue: Arc<dyn Queue>, registry: Arc<JobRegistry>) -> Self {
186 Self { queue, registry, recorded: None, locks: None }
187 }
188
189 pub fn with_locks(mut self, locks: LockManager) -> Self {
200 self.locks = Some(locks);
201 self
202 }
203
204 pub fn fake() -> Self {
209 Self {
210 queue: Arc::new(SyncQueue::new(
211 Arc::new(JobRegistry::new()),
212 Arc::new(Container::new()),
213 )),
214 registry: Arc::new(JobRegistry::new()),
215 recorded: Some(Mutex::new(Vec::new())),
216 locks: None,
217 }
218 }
219
220 pub fn is_faking(&self) -> bool {
222 self.recorded.is_some()
223 }
224
225 pub fn queue(&self) -> &Arc<dyn Queue> {
227 &self.queue
228 }
229
230 pub fn registry(&self) -> &Arc<JobRegistry> {
232 &self.registry
233 }
234
235 pub async fn dispatch<J: Job>(&self, job: J) -> Result<Option<String>> {
237 self.pending(job)?.send().await
238 }
239
240 fn unique_key<J: Job>(job: &J) -> Option<(String, Duration)> {
243 job.unique_id().map(|id| (format!("queue:unique:{}:{id}", J::NAME), J::UNIQUE_FOR))
246 }
247
248 async fn claim(&self, key: &str, ttl: Duration) -> Result<Claim> {
250 let Some(locks) = &self.locks else {
251 tracing::warn!(
254 key,
255 "a job declares `unique_id` but the queue has no lock manager; dispatching without deduplicating"
256 );
257 return Ok(Claim::NotEnforced);
258 };
259
260 Ok(match locks.lock(key, ttl).acquire().await? {
261 Some(guard) => Claim::Held(guard),
262 None => Claim::AlreadyQueued,
263 })
264 }
265
266 pub async fn release_uniqueness(&self, key: &str) -> Result<bool> {
278 let Some(locks) = &self.locks else { return Ok(false) };
279
280 locks.lock(key.to_string(), Duration::from_secs(1)).force_release().await
281 }
282
283 pub fn pending<J: Job>(&self, job: J) -> Result<PendingDispatch<'_>> {
289 let unique = Self::unique_key(&job);
290 Ok(PendingDispatch { manager: self, job: QueuedJob::from_job(&job)?, unique })
291 }
292
293 pub async fn dispatch_on<J: Job>(
295 &self,
296 queue: impl Into<String>,
297 job: J,
298 ) -> Result<Option<String>> {
299 self.pending(job)?.on_queue(queue).send().await
300 }
301
302 pub async fn dispatch_after<J: Job>(&self, delay: Duration, job: J) -> Result<Option<String>> {
304 self.pending(job)?.delay(delay).send().await
305 }
306
307 pub async fn dispatch_now<J: Job>(&self, job: J, container: Arc<Container>) -> Result<()> {
312 let queued = QueuedJob::from_job(&job)?;
313 let context =
314 Arc::new(JobContext::new(container, queued.id.clone(), queued.queue.clone(), 1, 1));
315 self.registry.run(&queued, context).await
316 }
317
318 async fn push(&self, job: QueuedJob) -> Result<String> {
319 if let Some(recorded) = &self.recorded {
320 let id = job.id.clone();
321 recorded.lock().expect("recorder lock poisoned").push(job);
322 return Ok(id);
323 }
324 self.queue.push(job).await
325 }
326
327 pub fn pushed<J: Job>(&self) -> Vec<QueuedJob> {
331 let Some(recorded) = &self.recorded else {
332 return Vec::new();
333 };
334 recorded
335 .lock()
336 .expect("recorder lock poisoned")
337 .iter()
338 .filter(|job| job.name == J::NAME)
339 .cloned()
340 .collect()
341 }
342
343 pub fn all_pushed(&self) -> Vec<QueuedJob> {
345 match &self.recorded {
346 Some(recorded) => recorded.lock().expect("recorder lock poisoned").clone(),
347 None => Vec::new(),
348 }
349 }
350
351 pub fn assert_pushed<J: Job>(&self) {
358 self.require_faking("assert_pushed");
359 assert!(
360 !self.pushed::<J>().is_empty(),
361 "expected a `{}` job to have been queued. Queued: {:?}",
362 J::NAME,
363 self.all_pushed().iter().map(|job| job.name.clone()).collect::<Vec<_>>()
364 );
365 }
366
367 pub fn assert_pushed_times<J: Job>(&self, times: usize) {
373 self.require_faking("assert_pushed_times");
374 let actual = self.pushed::<J>().len();
375 assert_eq!(
376 actual,
377 times,
378 "expected `{}` to be queued {times} time(s), but it was queued {actual}",
379 J::NAME
380 );
381 }
382
383 pub fn assert_not_pushed<J: Job>(&self) {
389 self.require_faking("assert_not_pushed");
390 assert!(
391 self.pushed::<J>().is_empty(),
392 "expected no `{}` job to be queued, but one was",
393 J::NAME
394 );
395 }
396
397 pub fn assert_pushed_on<J: Job>(&self, queue: &str) {
403 self.require_faking("assert_pushed_on");
404 let queues: Vec<String> = self.pushed::<J>().iter().map(|job| job.queue.clone()).collect();
405 assert!(
406 queues.iter().any(|q| q == queue),
407 "expected `{}` to be queued on `{queue}`, but it was queued on {queues:?}",
408 J::NAME
409 );
410 }
411
412 fn require_faking(&self, method: &str) {
413 assert!(
414 self.is_faking(),
415 "`{method}` needs a faking manager — build it with `QueueManager::fake()`, \
416 otherwise nothing is recorded and the assertion is meaningless"
417 );
418 }
419}
420
421impl std::fmt::Debug for QueueManager {
422 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
423 f.debug_struct("QueueManager")
424 .field("driver", &self.queue.name())
425 .field("jobs", &self.registry.names())
426 .field("faking", &self.is_faking())
427 .finish()
428 }
429}
430
431const _: fn() = || {
435 fn assert_send_sync<T: Send + Sync>() {}
436 assert_send_sync::<QueueManager>();
437};
438
439#[cfg(test)]
440mod unique_tests {
441 use super::*;
442 use crate::queue::MemoryQueue;
443 use rainier_cache::MemoryCache;
444 use serde::{Deserialize, Serialize};
445
446 #[derive(Serialize, Deserialize)]
447 struct RebuildIndex;
448
449 #[async_trait::async_trait]
450 impl Job for RebuildIndex {
451 const NAME: &'static str = "search.rebuild";
452 fn unique_id(&self) -> Option<String> {
454 Some(String::new())
455 }
456 async fn handle(&self, _: &JobContext) -> Result<()> {
457 Ok(())
458 }
459 }
460
461 #[derive(Serialize, Deserialize)]
462 struct SendInvoice {
463 id: u64,
464 }
465
466 #[async_trait::async_trait]
467 impl Job for SendInvoice {
468 const NAME: &'static str = "billing.invoice";
469 fn unique_id(&self) -> Option<String> {
471 Some(self.id.to_string())
472 }
473 async fn handle(&self, _: &JobContext) -> Result<()> {
474 Ok(())
475 }
476 }
477
478 #[derive(Serialize, Deserialize)]
479 struct Ordinary;
480
481 #[async_trait::async_trait]
482 impl Job for Ordinary {
483 const NAME: &'static str = "test.ordinary";
484 async fn handle(&self, _: &JobContext) -> Result<()> {
485 Ok(())
486 }
487 }
488
489 fn manager() -> (QueueManager, Arc<MemoryQueue>) {
490 let queue = Arc::new(MemoryQueue::new());
491 let manager = QueueManager::new(Arc::clone(&queue) as Arc<_>, Arc::new(JobRegistry::new()))
492 .with_locks(rainier_cache::LockManager::new(Arc::new(MemoryCache::new())));
493 (manager, queue)
494 }
495
496 #[tokio::test]
497 async fn a_second_identical_dispatch_is_dropped() {
498 let (manager, queue) = manager();
499
500 assert!(manager.dispatch(RebuildIndex).await.unwrap().is_some());
501 assert!(manager.dispatch(RebuildIndex).await.unwrap().is_none());
502
503 assert_eq!(queue.size("default").await.unwrap(), 1, "one copy, not two");
504 }
505
506 #[tokio::test]
507 async fn different_ids_are_different_jobs() {
508 let (manager, queue) = manager();
509
510 assert!(manager.dispatch(SendInvoice { id: 1 }).await.unwrap().is_some());
511 assert!(manager.dispatch(SendInvoice { id: 2 }).await.unwrap().is_some());
512 assert!(manager.dispatch(SendInvoice { id: 1 }).await.unwrap().is_none());
513
514 assert_eq!(queue.size("default").await.unwrap(), 2);
515 }
516
517 #[tokio::test]
518 async fn two_job_types_sharing_an_id_do_not_collide() {
519 let (manager, queue) = manager();
522
523 assert!(manager.dispatch(RebuildIndex).await.unwrap().is_some());
524 assert!(manager.dispatch(SendInvoice { id: 0 }).await.unwrap().is_some());
525
526 assert_eq!(queue.size("default").await.unwrap(), 2);
527 }
528
529 #[tokio::test]
530 async fn an_ordinary_job_is_never_deduplicated() {
531 let (manager, queue) = manager();
532
533 for _ in 0..3 {
534 assert!(manager.dispatch(Ordinary).await.unwrap().is_some());
535 }
536 assert_eq!(queue.size("default").await.unwrap(), 3);
537 }
538
539 #[tokio::test]
540 async fn the_queued_job_carries_the_key_so_a_worker_can_release_it() {
541 let (manager, queue) = manager();
542 manager.dispatch(RebuildIndex).await.unwrap();
543
544 let job = queue.reserve("default").await.unwrap().expect("a job");
545 assert_eq!(job.unique_key.as_deref(), Some("queue:unique:search.rebuild:"));
546 }
547
548 #[tokio::test]
549 async fn without_a_lock_manager_a_unique_job_is_dispatched_anyway() {
550 let queue = Arc::new(MemoryQueue::new());
553 let manager = QueueManager::new(Arc::clone(&queue) as Arc<_>, Arc::new(JobRegistry::new()));
554
555 assert!(manager.dispatch(RebuildIndex).await.unwrap().is_some());
556 assert!(manager.dispatch(RebuildIndex).await.unwrap().is_some());
557
558 assert_eq!(queue.size("default").await.unwrap(), 2);
559 }
560
561 #[tokio::test]
562 async fn releasing_the_lock_lets_the_next_one_in() {
563 let (manager, _queue) = manager();
564
565 manager.dispatch(RebuildIndex).await.unwrap();
566 assert!(manager.dispatch(RebuildIndex).await.unwrap().is_none());
567
568 manager.release_uniqueness("queue:unique:search.rebuild:").await.unwrap();
569
570 assert!(
571 manager.dispatch(RebuildIndex).await.unwrap().is_some(),
572 "once the first has run, another may be queued"
573 );
574 }
575
576 #[tokio::test]
577 async fn a_configured_dispatch_is_deduplicated_too() {
578 let (manager, queue) = manager();
580
581 assert!(manager
582 .pending(RebuildIndex)
583 .unwrap()
584 .on_queue("high")
585 .send()
586 .await
587 .unwrap()
588 .is_some());
589 assert!(manager
590 .pending(RebuildIndex)
591 .unwrap()
592 .on_queue("high")
593 .send()
594 .await
595 .unwrap()
596 .is_none());
597
598 assert_eq!(queue.size("high").await.unwrap(), 1);
599 }
600
601 #[tokio::test]
602 async fn a_pending_dispatch_that_is_dropped_leaves_no_lock() {
603 let (manager, _queue) = manager();
606
607 let abandoned = manager.pending(RebuildIndex).unwrap().on_queue("high");
608 drop(abandoned);
609
610 assert!(manager.dispatch(RebuildIndex).await.unwrap().is_some());
611 }
612}
613
614#[cfg(test)]
615mod tests {
616 use super::*;
617 use crate::queue::MemoryQueue;
618 use rainier_support::Error;
619 use serde::{Deserialize, Serialize};
620 use std::sync::atomic::{AtomicU32, Ordering};
621
622 type Runs = Arc<AtomicU32>;
633
634 fn counting_container() -> (Arc<Container>, Runs) {
635 let runs: Runs = Arc::new(AtomicU32::new(0));
636 let container = Arc::new(Container::new());
637 container.instance(Arc::clone(&runs));
638 (container, runs)
639 }
640
641 #[derive(Serialize, Deserialize)]
642 struct Ping;
643
644 #[async_trait::async_trait]
645 impl Job for Ping {
646 const NAME: &'static str = "test.ping";
647 async fn handle(&self, context: &JobContext) -> Result<()> {
648 if let Ok(runs) = context.resolve::<Runs>() {
651 runs.fetch_add(1, Ordering::SeqCst);
652 }
653 Ok(())
654 }
655 }
656
657 #[derive(Serialize, Deserialize)]
658 struct Boom;
659
660 #[async_trait::async_trait]
661 impl Job for Boom {
662 const NAME: &'static str = "test.boom";
663 async fn handle(&self, _: &JobContext) -> Result<()> {
664 Err(Error::internal("exploded"))
665 }
666 }
667
668 fn registry() -> Arc<JobRegistry> {
669 Arc::new(JobRegistry::new().with::<Ping>().with::<Boom>())
670 }
671
672 fn manager(queue: Arc<dyn Queue>) -> QueueManager {
673 QueueManager::new(queue, registry())
674 }
675
676 #[tokio::test]
677 async fn dispatching_enqueues_the_job() {
678 let queue = Arc::new(MemoryQueue::new());
679 let manager = manager(Arc::clone(&queue) as Arc<dyn Queue>);
680
681 let id = manager.dispatch(Ping).await.unwrap();
682 assert!(id.is_some_and(|id| !id.is_empty()));
683 assert_eq!(queue.size("default").await.unwrap(), 1);
684 }
685
686 #[tokio::test]
687 async fn dispatch_options_apply() {
688 let queue = Arc::new(MemoryQueue::new());
689 let manager = manager(Arc::clone(&queue) as Arc<dyn Queue>);
690
691 manager
692 .pending(Ping)
693 .unwrap()
694 .on_queue("mail")
695 .delay(Duration::from_secs(60))
696 .tries(5)
697 .send()
698 .await
699 .unwrap();
700
701 let pending = queue.pending("mail");
702 assert_eq!(pending.len(), 1);
703 assert_eq!(pending[0].max_attempts, 5);
704 assert!(!pending[0].is_available(), "the delay should hold it back");
705 }
706
707 #[tokio::test]
708 async fn dispatch_on_and_after_are_shorthands() {
709 let queue = Arc::new(MemoryQueue::new());
710 let manager = manager(Arc::clone(&queue) as Arc<dyn Queue>);
711
712 manager.dispatch_on("mail", Ping).await.unwrap();
713 manager.dispatch_after(Duration::from_secs(30), Ping).await.unwrap();
714
715 assert_eq!(queue.size("mail").await.unwrap(), 1);
716 assert!(!queue.pending("default")[0].is_available());
717 }
718
719 #[tokio::test]
720 async fn the_sync_driver_runs_the_job_immediately() {
721 let (container, runs) = counting_container();
722 let manager = manager(Arc::new(SyncQueue::new(registry(), container)));
723
724 manager.dispatch(Ping).await.unwrap();
725
726 assert_eq!(runs.load(Ordering::SeqCst), 1, "it ran during dispatch");
727 }
728
729 #[tokio::test]
730 async fn the_sync_driver_surfaces_a_failure_to_the_dispatcher() {
731 let sync = Arc::new(SyncQueue::new(registry(), Arc::new(Container::new())));
733 let manager = manager(sync);
734
735 let err = manager.dispatch(Boom).await.unwrap_err();
736 assert!(err.message().contains("exploded"), "{}", err.message());
737 }
738
739 #[tokio::test]
740 async fn dispatch_now_bypasses_the_queue() {
741 let (container, runs) = counting_container();
742 let queue = Arc::new(MemoryQueue::new());
743 let manager = manager(Arc::clone(&queue) as Arc<dyn Queue>);
744
745 manager.dispatch_now(Ping, container).await.unwrap();
746
747 assert_eq!(runs.load(Ordering::SeqCst), 1);
748 assert_eq!(queue.size("default").await.unwrap(), 0, "nothing was queued");
749 }
750
751 #[tokio::test]
752 async fn a_fake_records_instead_of_queueing() {
753 let manager = QueueManager::fake();
754
755 manager.dispatch(Ping).await.unwrap();
756 manager.dispatch_on("mail", Ping).await.unwrap();
757
758 manager.assert_pushed::<Ping>();
759 manager.assert_pushed_times::<Ping>(2);
760 manager.assert_pushed_on::<Ping>("mail");
761 manager.assert_not_pushed::<Boom>();
762 assert_eq!(manager.all_pushed().len(), 2);
763 }
764
765 #[tokio::test]
766 #[should_panic(expected = "needs a faking manager")]
767 async fn assertions_refuse_to_pass_vacuously() {
768 let queue = Arc::new(MemoryQueue::new());
769 let manager = manager(queue as Arc<dyn Queue>);
770 manager.dispatch(Ping).await.unwrap();
771 manager.assert_not_pushed::<Ping>();
772 }
773
774 #[tokio::test]
775 #[should_panic(expected = "test.boom")]
776 async fn a_missing_dispatch_reports_what_was_queued() {
777 let manager = QueueManager::fake();
778 manager.dispatch(Ping).await.unwrap();
779 manager.assert_pushed::<Boom>();
780 }
781}