1use crate::error::{QueueError, QueueResult};
4use crate::job::{Job, JobData, JobId, JobPriority, JobState, JobStatus};
5use armature_log::{debug, info};
6use chrono::{DateTime, Utc};
7use redis::{AsyncCommands, Client, aio::ConnectionManager};
8use std::time::Duration;
9
10const SCAN_COUNT: usize = 500;
12
13const MOVE_DELAYED_SCRIPT: &str = r#"
25local delayed_key = KEYS[1]
26local prefix = ARGV[1]
27local now = ARGV[2]
28
29local job_ids = redis.call('ZRANGEBYSCORE', delayed_key, '-inf', now)
30local promoted = 0
31
32for _, job_id in ipairs(job_ids) do
33 local job_json = redis.call('GET', prefix .. ':job:' .. job_id)
34 if job_json then
35 -- Claim the job atomically; skip if another pass already took it.
36 if redis.call('ZREM', delayed_key, job_id) == 1 then
37 local pname = 'normal'
38 local pscore = -1
39 local ok, job = pcall(cjson.decode, job_json)
40 if ok and type(job) == 'table' and job.priority then
41 local p = job.priority
42 if p == 'Low' then pname = 'low'; pscore = 0
43 elseif p == 'Normal' then pname = 'normal'; pscore = -1
44 elseif p == 'High' then pname = 'high'; pscore = -2
45 elseif p == 'Critical' then pname = 'critical'; pscore = -3
46 end
47 end
48 redis.call('ZADD', prefix .. ':pending:' .. pname, pscore, job_id)
49 promoted = promoted + 1
50 end
51 end
52end
53
54return promoted
55"#;
56
57const DEQUEUE_POP_SCRIPT: &str = r#"
70local prefix = ARGV[1]
71for i = 1, #KEYS do
72 while true do
73 local popped = redis.call('ZPOPMIN', KEYS[i], 1)
74 if not popped or not popped[1] then
75 break
76 end
77 local job_id = popped[1]
78 local job_json = redis.call('GET', prefix .. ':job:' .. job_id)
79 if job_json then
80 return {job_id, job_json}
81 end
82 -- Job body expired between enqueue and dequeue: the id is discarded
83 -- (already popped) and we keep draining this same queue.
84 end
85end
86return nil
87"#;
88
89#[derive(Debug, Clone)]
91pub struct QueueConfig {
92 pub redis_url: String,
94
95 pub queue_name: String,
97
98 pub key_prefix: String,
100
101 pub max_size: usize,
103
104 pub retention_time: Duration,
106}
107
108impl QueueConfig {
109 pub fn new(redis_url: impl Into<String>, queue_name: impl Into<String>) -> Self {
111 let queue_name = queue_name.into();
112 Self {
113 redis_url: redis_url.into(),
114 key_prefix: format!("armature:queue:{}", queue_name),
115 queue_name,
116 max_size: 0,
117 retention_time: Duration::from_secs(86400), }
119 }
120
121 pub fn with_key_prefix(mut self, prefix: impl Into<String>) -> Self {
123 self.key_prefix = prefix.into();
124 self
125 }
126
127 pub fn with_max_size(mut self, max_size: usize) -> Self {
129 self.max_size = max_size;
130 self
131 }
132
133 pub fn with_retention_time(mut self, retention_time: Duration) -> Self {
135 self.retention_time = retention_time;
136 self
137 }
138
139 fn key(&self, suffix: &str) -> String {
141 format!("{}:{}", self.key_prefix, suffix)
142 }
143}
144
145#[derive(Clone)]
147pub struct Queue {
148 connection: ConnectionManager,
149 config: QueueConfig,
150}
151
152impl Queue {
153 pub async fn new(
155 redis_url: impl Into<String>,
156 queue_name: impl Into<String>,
157 ) -> QueueResult<Self> {
158 let config = QueueConfig::new(redis_url, queue_name);
159 Self::with_config(config).await
160 }
161
162 pub async fn with_config(config: QueueConfig) -> QueueResult<Self> {
164 info!("Initializing job queue: {}", config.queue_name);
165 debug!(
166 "Queue config - prefix: {}, max_size: {}",
167 config.key_prefix, config.max_size
168 );
169
170 let client = Client::open(config.redis_url.as_str())
171 .map_err(|e| QueueError::Config(e.to_string()))?;
172
173 let connection = ConnectionManager::new(client).await?;
174
175 info!("Job queue '{}' ready", config.queue_name);
176 Ok(Self { connection, config })
177 }
178
179 pub async fn enqueue(&self, job_type: impl Into<String>, data: JobData) -> QueueResult<JobId> {
181 let job_type = job_type.into();
182 debug!(
183 "Enqueueing job: {} on queue '{}'",
184 job_type, self.config.queue_name
185 );
186 let job = Job::new(&self.config.queue_name, &job_type, data);
187 self.enqueue_job(job).await
188 }
189
190 pub async fn enqueue_in(
197 &self,
198 delay: chrono::Duration,
199 job_type: impl Into<String>,
200 data: JobData,
201 ) -> QueueResult<JobId> {
202 let job = Job::new(&self.config.queue_name, job_type, data).schedule_after(delay);
203 self.enqueue_job(job).await
204 }
205
206 pub async fn enqueue_at(
213 &self,
214 when: DateTime<Utc>,
215 job_type: impl Into<String>,
216 data: JobData,
217 ) -> QueueResult<JobId> {
218 let job = Job::new(&self.config.queue_name, job_type, data).schedule_at(when);
219 self.enqueue_job(job).await
220 }
221
222 pub async fn enqueue_job(&self, job: Job) -> QueueResult<JobId> {
224 if self.config.max_size > 0 {
229 let size = self.backlog_size().await?;
230 if size >= self.config.max_size {
231 return Err(QueueError::QueueFull);
232 }
233 }
234
235 let job_id = job.id;
236 let mut conn = self.connection.clone();
237
238 let job_json =
240 serde_json::to_string(&job).map_err(|e| QueueError::Serialization(e.to_string()))?;
241
242 let job_key = self.config.key(&format!("job:{}", job_id));
244 let _: () = conn
245 .set_ex(&job_key, job_json, self.config.retention_time.as_secs())
246 .await?;
247
248 if job.is_ready() {
250 let queue_key = self.priority_queue_key(job.priority);
251 let score = -(job.priority as i64); let _: () = conn.zadd(&queue_key, job_id.to_string(), score).await?;
253 } else {
254 let delayed_key = self.config.key("delayed");
256 let score = job.scheduled_at.unwrap().timestamp();
257 let _: () = conn.zadd(&delayed_key, job_id.to_string(), score).await?;
258 }
259
260 Ok(job_id)
261 }
262
263 pub async fn dequeue(&self) -> QueueResult<Option<Job>> {
265 self.move_delayed_jobs().await?;
266
267 let mut conn = self.connection.clone();
268
269 let script = redis::Script::new(DEQUEUE_POP_SCRIPT);
277 let popped: Option<(String, String)> = script
278 .key(self.priority_queue_key(JobPriority::Critical))
279 .key(self.priority_queue_key(JobPriority::High))
280 .key(self.priority_queue_key(JobPriority::Normal))
281 .key(self.priority_queue_key(JobPriority::Low))
282 .arg(&self.config.key_prefix)
283 .invoke_async(&mut conn)
284 .await?;
285
286 let Some((job_id_str, job_json)) = popped else {
287 return Ok(None);
288 };
289 let job_id = job_id_str
292 .parse::<JobId>()
293 .map_err(|e| QueueError::Deserialization(e.to_string()))?;
294 let mut job: Job = serde_json::from_str(&job_json)
295 .map_err(|e| QueueError::Deserialization(e.to_string()))?;
296
297 job.start_processing();
298 let job_key = self.config.key(&format!("job:{}", job_id));
299 let processing_key = self.config.key("processing");
300 let updated_json =
301 serde_json::to_string(&job).map_err(|e| QueueError::Serialization(e.to_string()))?;
302
303 let _: () = redis::pipe()
306 .set_ex(&job_key, updated_json, self.config.retention_time.as_secs())
307 .ignore()
308 .zadd(&processing_key, job_id.to_string(), Utc::now().timestamp())
309 .ignore()
310 .query_async(&mut conn)
311 .await?;
312
313 Ok(Some(job))
314 }
315
316 pub async fn complete(&self, job_id: JobId) -> QueueResult<()> {
318 if let Some(mut job) = self.get_job(job_id).await? {
323 job.complete();
324 let job_key = self.config.key(&format!("job:{}", job_id));
325 let processing_key = self.config.key("processing");
326 let job_json = serde_json::to_string(&job)
327 .map_err(|e| QueueError::Serialization(e.to_string()))?;
328
329 let mut conn = self.connection.clone();
330 let _: () = redis::pipe()
331 .set_ex(&job_key, job_json, self.config.retention_time.as_secs())
332 .ignore()
333 .zrem(&processing_key, job_id.to_string())
334 .ignore()
335 .query_async(&mut conn)
336 .await?;
337 } else {
338 self.remove_from_processing(job_id).await?;
339 }
340 Ok(())
341 }
342
343 pub async fn fail(&self, job_id: JobId, error: String) -> QueueResult<()> {
345 if let Some(mut job) = self.get_job(job_id).await? {
349 job.fail(error);
350
351 let job_key = self.config.key(&format!("job:{}", job_id));
352 let processing_key = self.config.key("processing");
353 let job_json = serde_json::to_string(&job)
354 .map_err(|e| QueueError::Serialization(e.to_string()))?;
355 let mut conn = self.connection.clone();
356
357 if job.status.state == JobState::Failed && job.can_retry() {
358 let retry_at = Utc::now() + job.backoff_delay();
360 job.scheduled_at = Some(retry_at);
361 let job_json = serde_json::to_string(&job)
362 .map_err(|e| QueueError::Serialization(e.to_string()))?;
363
364 let delayed_key = self.config.key("delayed");
366 let _: () = redis::pipe()
367 .set_ex(&job_key, job_json, self.config.retention_time.as_secs())
368 .ignore()
369 .zadd(&delayed_key, job_id.to_string(), retry_at.timestamp())
370 .ignore()
371 .zrem(&processing_key, job_id.to_string())
372 .ignore()
373 .query_async(&mut conn)
374 .await?;
375 } else {
376 let dead_key = self.config.key("dead");
378 let _: () = redis::pipe()
379 .set_ex(&job_key, job_json, self.config.retention_time.as_secs())
380 .ignore()
381 .zadd(&dead_key, job_id.to_string(), Utc::now().timestamp())
382 .ignore()
383 .zrem(&processing_key, job_id.to_string())
384 .ignore()
385 .query_async(&mut conn)
386 .await?;
387 }
388 } else {
389 self.remove_from_processing(job_id).await?;
390 }
391 Ok(())
392 }
393
394 pub async fn requeue(&self, job: &Job) -> QueueResult<()> {
406 let mut job = job.clone();
407
408 job.status = JobStatus::pending();
410 job.started_at = None;
411 job.attempts = job.attempts.saturating_sub(1);
412 self.save_job(&job).await?;
413
414 let mut conn = self.connection.clone();
415 let queue_key = self.priority_queue_key(job.priority);
416 let score = -(job.priority as i64);
417 let _: () = conn.zadd(&queue_key, job.id.to_string(), score).await?;
418
419 self.remove_from_processing(job.id).await?;
420 Ok(())
421 }
422
423 pub async fn get_job(&self, job_id: JobId) -> QueueResult<Option<Job>> {
425 let mut conn = self.connection.clone();
426 let job_key = self.config.key(&format!("job:{}", job_id));
427
428 let job_json: Option<String> = conn.get(&job_key).await?;
429
430 if let Some(json) = job_json {
431 let job: Job = serde_json::from_str(&json)
432 .map_err(|e| QueueError::Deserialization(e.to_string()))?;
433 Ok(Some(job))
434 } else {
435 Ok(None)
436 }
437 }
438
439 async fn save_job(&self, job: &Job) -> QueueResult<()> {
441 let mut conn = self.connection.clone();
442 let job_key = self.config.key(&format!("job:{}", job.id));
443 let job_json =
444 serde_json::to_string(job).map_err(|e| QueueError::Serialization(e.to_string()))?;
445
446 let _: () = conn
447 .set_ex(&job_key, job_json, self.config.retention_time.as_secs())
448 .await?;
449 Ok(())
450 }
451
452 pub async fn size(&self) -> QueueResult<usize> {
454 let mut conn = self.connection.clone();
455
456 let mut pipe = redis::pipe();
459 for priority in [
460 JobPriority::Critical,
461 JobPriority::High,
462 JobPriority::Normal,
463 JobPriority::Low,
464 ] {
465 pipe.zcard(self.priority_queue_key(priority));
466 }
467
468 let counts: Vec<usize> = pipe.query_async(&mut conn).await?;
469 Ok(counts.iter().sum())
470 }
471
472 pub async fn backlog_size(&self) -> QueueResult<usize> {
481 let mut conn = self.connection.clone();
482
483 let mut pipe = redis::pipe();
484 for priority in [
485 JobPriority::Critical,
486 JobPriority::High,
487 JobPriority::Normal,
488 JobPriority::Low,
489 ] {
490 pipe.zcard(self.priority_queue_key(priority));
491 }
492 pipe.zcard(self.config.key("delayed"));
493 pipe.zcard(self.config.key("processing"));
494
495 let counts: Vec<usize> = pipe.query_async(&mut conn).await?;
496 Ok(counts.iter().sum())
497 }
498
499 pub async fn processing_len(&self) -> QueueResult<usize> {
501 let mut conn = self.connection.clone();
502 let processing_key = self.config.key("processing");
503 let count: usize = conn.zcard(&processing_key).await?;
504 Ok(count)
505 }
506
507 async fn move_delayed_jobs(&self) -> QueueResult<()> {
515 let mut conn = self.connection.clone();
516 let delayed_key = self.config.key("delayed");
517 let now = Utc::now().timestamp();
518
519 let earliest: Vec<(String, i64)> = conn.zrange_withscores(&delayed_key, 0, 0).await?;
525 match earliest.first() {
526 Some((_, score)) if *score <= now => {}
527 _ => return Ok(()),
528 }
529
530 let script = redis::Script::new(MOVE_DELAYED_SCRIPT);
531 let _: i64 = script
532 .key(&delayed_key)
533 .arg(&self.config.key_prefix)
534 .arg(now)
535 .invoke_async(&mut conn)
536 .await?;
537
538 Ok(())
539 }
540
541 async fn remove_from_processing(&self, job_id: JobId) -> QueueResult<()> {
543 let mut conn = self.connection.clone();
544 let processing_key = self.config.key("processing");
545 let _: () = conn.zrem(&processing_key, job_id.to_string()).await?;
546 Ok(())
547 }
548
549 fn priority_queue_key(&self, priority: JobPriority) -> String {
551 self.config
552 .key(&format!("pending:{:?}", priority).to_lowercase())
553 }
554
555 pub async fn clear(&self) -> QueueResult<()> {
557 let mut conn = self.connection.clone();
558 let pattern = format!("{}:*", self.config.key_prefix);
559
560 let mut cursor: u64 = 0;
563 loop {
564 let (next, keys): (u64, Vec<String>) = redis::cmd("SCAN")
565 .arg(cursor)
566 .arg("MATCH")
567 .arg(&pattern)
568 .arg("COUNT")
569 .arg(SCAN_COUNT)
570 .query_async(&mut conn)
571 .await?;
572
573 if !keys.is_empty() {
574 let _: () = redis::cmd("UNLINK")
575 .arg(&keys)
576 .query_async(&mut conn)
577 .await?;
578 }
579
580 cursor = next;
581 if cursor == 0 {
582 break;
583 }
584 }
585
586 Ok(())
587 }
588}
589
590#[cfg(test)]
591mod tests {
592 use super::*;
593
594 #[test]
595 fn test_queue_config() {
596 let config = QueueConfig::new("redis://localhost:6379", "test");
597 assert_eq!(config.queue_name, "test");
598 assert!(config.key_prefix.contains("test"));
599 }
600
601 #[test]
606 fn test_lua_priority_mapping_matches_rust() {
607 for priority in [
608 JobPriority::Low,
609 JobPriority::Normal,
610 JobPriority::High,
611 JobPriority::Critical,
612 ] {
613 let rust_score = -(priority as i64);
615 let variant = format!("{priority:?}");
617 let name = variant.to_lowercase();
619
620 let expected = format!("'{variant}' then pname = '{name}'; pscore = {rust_score}");
621 assert!(
622 MOVE_DELAYED_SCRIPT.contains(&expected),
623 "script missing mapping for {priority:?}: expected `{expected}`"
624 );
625 }
626 }
627
628 #[test]
629 fn test_priority_queue_key_layout() {
630 let config = QueueConfig::new("redis://localhost:6379", "jobs");
632 for (priority, name) in [
633 (JobPriority::Low, "low"),
634 (JobPriority::Normal, "normal"),
635 (JobPriority::High, "high"),
636 (JobPriority::Critical, "critical"),
637 ] {
638 let key = config.key(&format!("pending:{priority:?}").to_lowercase());
639 assert_eq!(key, format!("{}:pending:{}", config.key_prefix, name));
640 }
641 }
642
643 #[test]
644 fn test_dequeue_script_scans_all_keys() {
645 assert!(DEQUEUE_POP_SCRIPT.contains("for i = 1, #KEYS do"));
647 assert!(DEQUEUE_POP_SCRIPT.contains("ZPOPMIN"));
648 }
649
650 #[tokio::test]
652 #[ignore = "requires a running Redis instance"]
653 async fn test_move_delayed_promotes_due_jobs() {
654 use crate::job::Job;
655
656 let queue = Queue::new("redis://localhost:6379", "test_move_delayed")
657 .await
658 .unwrap();
659 queue.clear().await.unwrap();
660
661 let past = Utc::now() - chrono::Duration::seconds(30);
663 let job = Job::new("test_move_delayed", "task", serde_json::json!({}))
664 .with_priority(JobPriority::High)
665 .schedule_at(past);
666 queue.enqueue_job(job).await.unwrap();
667
668 assert_eq!(queue.size().await.unwrap(), 0);
670
671 let dequeued = queue.dequeue().await.unwrap();
673 assert!(dequeued.is_some());
674 assert_eq!(dequeued.unwrap().priority, JobPriority::High);
675
676 queue.clear().await.unwrap();
677 }
678
679 #[test]
680 fn test_priority_queue_key() {
681 let config = QueueConfig::new("redis://localhost:6379", "test");
682 assert!(config.key("pending:high").contains("high"));
683 }
684
685 #[test]
686 fn test_queue_config_with_custom_prefix() {
687 let config = QueueConfig::new("redis://localhost:6379", "myqueue").with_key_prefix("app");
688 assert!(config.key_prefix.contains("app"));
689 }
690
691 #[test]
692 fn test_queue_config_default_retention() {
693 let config = QueueConfig::new("redis://localhost:6379", "test");
694 assert_eq!(config.retention_time, Duration::from_secs(86400)); }
696
697 #[test]
698 fn test_queue_config_custom_retention() {
699 let retention = Duration::from_secs(3600);
700 let config =
701 QueueConfig::new("redis://localhost:6379", "test").with_retention_time(retention);
702 assert_eq!(config.retention_time, retention);
703 }
704
705 #[test]
706 fn test_queue_config_default_max_size() {
707 let config = QueueConfig::new("redis://localhost:6379", "test");
708 assert_eq!(config.max_size, 0); }
710
711 #[test]
712 fn test_queue_config_custom_max_size() {
713 let config = QueueConfig::new("redis://localhost:6379", "test").with_max_size(1000);
714 assert_eq!(config.max_size, 1000);
715 }
716
717 #[test]
718 fn test_queue_key_generation() {
719 let config = QueueConfig::new("redis://localhost:6379", "jobs");
720
721 let pending_key = config.key("pending:normal");
722 let processing_key = config.key("processing");
723 let completed_key = config.key("completed");
724
725 assert!(pending_key.contains("jobs"));
726 assert!(processing_key.contains("jobs"));
727 assert!(completed_key.contains("jobs"));
728 }
729
730 #[test]
731 fn test_queue_config_clone() {
732 let config1 = QueueConfig::new("redis://localhost:6379", "test");
733 let config2 = config1.clone();
734
735 assert_eq!(config1.queue_name, config2.queue_name);
736 assert_eq!(config1.redis_url, config2.redis_url);
737 }
738
739 #[test]
740 fn test_queue_config_different_queues() {
741 let config1 = QueueConfig::new("redis://localhost:6379", "queue1");
742 let config2 = QueueConfig::new("redis://localhost:6379", "queue2");
743
744 assert_ne!(config1.key_prefix, config2.key_prefix);
745 }
746
747 #[test]
748 fn test_queue_config_key_consistency() {
749 let config = QueueConfig::new("redis://localhost:6379", "test");
750
751 let key1 = config.key("pending");
752 let key2 = config.key("pending");
753
754 assert_eq!(key1, key2);
755 }
756
757 #[test]
758 fn test_queue_config_builder_pattern() {
759 let config = QueueConfig::new("redis://localhost:6379", "test")
760 .with_key_prefix("app")
761 .with_retention_time(Duration::from_secs(7200))
762 .with_max_size(500);
763
764 assert!(config.key_prefix.contains("app"));
765 assert_eq!(config.retention_time, Duration::from_secs(7200));
766 assert_eq!(config.max_size, 500);
767 }
768
769 #[test]
770 fn test_queue_config_redis_url() {
771 let url = "redis://user:pass@host:6380/2";
772 let config = QueueConfig::new(url, "test");
773 assert_eq!(config.redis_url, url);
774 }
775
776 #[test]
777 fn test_queue_config_key_with_empty_suffix() {
778 let config = QueueConfig::new("redis://localhost:6379", "test");
779 let key = config.key("");
780 assert!(key.contains("test"));
781 }
782
783 #[test]
784 fn test_queue_config_key_with_special_characters() {
785 let config = QueueConfig::new("redis://localhost:6379", "test");
786 let key = config.key("pending:high:priority");
787 assert!(key.contains("pending:high:priority"));
788 }
789
790 #[test]
791 fn test_queue_config_multiple_prefixes() {
792 let config1 =
793 QueueConfig::new("redis://localhost:6379", "app1").with_key_prefix("production");
794 let config2 =
795 QueueConfig::new("redis://localhost:6379", "app2").with_key_prefix("development");
796
797 let key1 = config1.key("jobs");
798 let key2 = config2.key("jobs");
799
800 assert_ne!(key1, key2);
801 }
802
803 #[test]
804 fn test_queue_config_unlimited_max_size() {
805 let config = QueueConfig::new("redis://localhost:6379", "test").with_max_size(0);
806 assert_eq!(config.max_size, 0);
807 }
808
809 #[test]
810 fn test_queue_config_large_retention() {
811 let week = Duration::from_secs(7 * 24 * 3600);
812 let config = QueueConfig::new("redis://localhost:6379", "test").with_retention_time(week);
813 assert_eq!(config.retention_time, week);
814 }
815}