use crate::error::{QueueError, QueueResult};
use crate::job::{Job, JobData, JobId, JobPriority, JobState, JobStatus};
use armature_log::{debug, info};
use chrono::{DateTime, Utc};
use redis::{AsyncCommands, Client, aio::ConnectionManager};
use std::time::Duration;
const SCAN_COUNT: usize = 500;
const MOVE_DELAYED_SCRIPT: &str = r#"
local delayed_key = KEYS[1]
local prefix = ARGV[1]
local now = ARGV[2]
local job_ids = redis.call('ZRANGEBYSCORE', delayed_key, '-inf', now)
local promoted = 0
for _, job_id in ipairs(job_ids) do
local job_json = redis.call('GET', prefix .. ':job:' .. job_id)
if job_json then
-- Claim the job atomically; skip if another pass already took it.
if redis.call('ZREM', delayed_key, job_id) == 1 then
local pname = 'normal'
local pscore = -1
local ok, job = pcall(cjson.decode, job_json)
if ok and type(job) == 'table' and job.priority then
local p = job.priority
if p == 'Low' then pname = 'low'; pscore = 0
elseif p == 'Normal' then pname = 'normal'; pscore = -1
elseif p == 'High' then pname = 'high'; pscore = -2
elseif p == 'Critical' then pname = 'critical'; pscore = -3
end
end
redis.call('ZADD', prefix .. ':pending:' .. pname, pscore, job_id)
promoted = promoted + 1
end
end
end
return promoted
"#;
const DEQUEUE_POP_SCRIPT: &str = r#"
local prefix = ARGV[1]
for i = 1, #KEYS do
while true do
local popped = redis.call('ZPOPMIN', KEYS[i], 1)
if not popped or not popped[1] then
break
end
local job_id = popped[1]
local job_json = redis.call('GET', prefix .. ':job:' .. job_id)
if job_json then
return {job_id, job_json}
end
-- Job body expired between enqueue and dequeue: the id is discarded
-- (already popped) and we keep draining this same queue.
end
end
return nil
"#;
#[derive(Debug, Clone)]
pub struct QueueConfig {
pub redis_url: String,
pub queue_name: String,
pub key_prefix: String,
pub max_size: usize,
pub retention_time: Duration,
}
impl QueueConfig {
pub fn new(redis_url: impl Into<String>, queue_name: impl Into<String>) -> Self {
let queue_name = queue_name.into();
Self {
redis_url: redis_url.into(),
key_prefix: format!("armature:queue:{}", queue_name),
queue_name,
max_size: 0,
retention_time: Duration::from_secs(86400), }
}
pub fn with_key_prefix(mut self, prefix: impl Into<String>) -> Self {
self.key_prefix = prefix.into();
self
}
pub fn with_max_size(mut self, max_size: usize) -> Self {
self.max_size = max_size;
self
}
pub fn with_retention_time(mut self, retention_time: Duration) -> Self {
self.retention_time = retention_time;
self
}
fn key(&self, suffix: &str) -> String {
format!("{}:{}", self.key_prefix, suffix)
}
}
#[derive(Clone)]
pub struct Queue {
connection: ConnectionManager,
config: QueueConfig,
}
impl Queue {
pub async fn new(
redis_url: impl Into<String>,
queue_name: impl Into<String>,
) -> QueueResult<Self> {
let config = QueueConfig::new(redis_url, queue_name);
Self::with_config(config).await
}
pub async fn with_config(config: QueueConfig) -> QueueResult<Self> {
info!("Initializing job queue: {}", config.queue_name);
debug!(
"Queue config - prefix: {}, max_size: {}",
config.key_prefix, config.max_size
);
let client = Client::open(config.redis_url.as_str())
.map_err(|e| QueueError::Config(e.to_string()))?;
let connection = ConnectionManager::new(client).await?;
info!("Job queue '{}' ready", config.queue_name);
Ok(Self { connection, config })
}
pub async fn enqueue(&self, job_type: impl Into<String>, data: JobData) -> QueueResult<JobId> {
let job_type = job_type.into();
debug!(
"Enqueueing job: {} on queue '{}'",
job_type, self.config.queue_name
);
let job = Job::new(&self.config.queue_name, &job_type, data);
self.enqueue_job(job).await
}
pub async fn enqueue_in(
&self,
delay: chrono::Duration,
job_type: impl Into<String>,
data: JobData,
) -> QueueResult<JobId> {
let job = Job::new(&self.config.queue_name, job_type, data).schedule_after(delay);
self.enqueue_job(job).await
}
pub async fn enqueue_at(
&self,
when: DateTime<Utc>,
job_type: impl Into<String>,
data: JobData,
) -> QueueResult<JobId> {
let job = Job::new(&self.config.queue_name, job_type, data).schedule_at(when);
self.enqueue_job(job).await
}
pub async fn enqueue_job(&self, job: Job) -> QueueResult<JobId> {
if self.config.max_size > 0 {
let size = self.backlog_size().await?;
if size >= self.config.max_size {
return Err(QueueError::QueueFull);
}
}
let job_id = job.id;
let mut conn = self.connection.clone();
let job_json =
serde_json::to_string(&job).map_err(|e| QueueError::Serialization(e.to_string()))?;
let job_key = self.config.key(&format!("job:{}", job_id));
let _: () = conn
.set_ex(&job_key, job_json, self.config.retention_time.as_secs())
.await?;
if job.is_ready() {
let queue_key = self.priority_queue_key(job.priority);
let score = -(job.priority as i64); let _: () = conn.zadd(&queue_key, job_id.to_string(), score).await?;
} else {
let delayed_key = self.config.key("delayed");
let score = job.scheduled_at.unwrap().timestamp();
let _: () = conn.zadd(&delayed_key, job_id.to_string(), score).await?;
}
Ok(job_id)
}
pub async fn dequeue(&self) -> QueueResult<Option<Job>> {
self.move_delayed_jobs().await?;
let mut conn = self.connection.clone();
let script = redis::Script::new(DEQUEUE_POP_SCRIPT);
let popped: Option<(String, String)> = script
.key(self.priority_queue_key(JobPriority::Critical))
.key(self.priority_queue_key(JobPriority::High))
.key(self.priority_queue_key(JobPriority::Normal))
.key(self.priority_queue_key(JobPriority::Low))
.arg(&self.config.key_prefix)
.invoke_async(&mut conn)
.await?;
let Some((job_id_str, job_json)) = popped else {
return Ok(None);
};
let job_id = job_id_str
.parse::<JobId>()
.map_err(|e| QueueError::Deserialization(e.to_string()))?;
let mut job: Job = serde_json::from_str(&job_json)
.map_err(|e| QueueError::Deserialization(e.to_string()))?;
job.start_processing();
let job_key = self.config.key(&format!("job:{}", job_id));
let processing_key = self.config.key("processing");
let updated_json =
serde_json::to_string(&job).map_err(|e| QueueError::Serialization(e.to_string()))?;
let _: () = redis::pipe()
.set_ex(&job_key, updated_json, self.config.retention_time.as_secs())
.ignore()
.zadd(&processing_key, job_id.to_string(), Utc::now().timestamp())
.ignore()
.query_async(&mut conn)
.await?;
Ok(Some(job))
}
pub async fn complete(&self, job_id: JobId) -> QueueResult<()> {
if let Some(mut job) = self.get_job(job_id).await? {
job.complete();
let job_key = self.config.key(&format!("job:{}", job_id));
let processing_key = self.config.key("processing");
let job_json = serde_json::to_string(&job)
.map_err(|e| QueueError::Serialization(e.to_string()))?;
let mut conn = self.connection.clone();
let _: () = redis::pipe()
.set_ex(&job_key, job_json, self.config.retention_time.as_secs())
.ignore()
.zrem(&processing_key, job_id.to_string())
.ignore()
.query_async(&mut conn)
.await?;
} else {
self.remove_from_processing(job_id).await?;
}
Ok(())
}
pub async fn fail(&self, job_id: JobId, error: String) -> QueueResult<()> {
if let Some(mut job) = self.get_job(job_id).await? {
job.fail(error);
let job_key = self.config.key(&format!("job:{}", job_id));
let processing_key = self.config.key("processing");
let job_json = serde_json::to_string(&job)
.map_err(|e| QueueError::Serialization(e.to_string()))?;
let mut conn = self.connection.clone();
if job.status.state == JobState::Failed && job.can_retry() {
let retry_at = Utc::now() + job.backoff_delay();
job.scheduled_at = Some(retry_at);
let job_json = serde_json::to_string(&job)
.map_err(|e| QueueError::Serialization(e.to_string()))?;
let delayed_key = self.config.key("delayed");
let _: () = redis::pipe()
.set_ex(&job_key, job_json, self.config.retention_time.as_secs())
.ignore()
.zadd(&delayed_key, job_id.to_string(), retry_at.timestamp())
.ignore()
.zrem(&processing_key, job_id.to_string())
.ignore()
.query_async(&mut conn)
.await?;
} else {
let dead_key = self.config.key("dead");
let _: () = redis::pipe()
.set_ex(&job_key, job_json, self.config.retention_time.as_secs())
.ignore()
.zadd(&dead_key, job_id.to_string(), Utc::now().timestamp())
.ignore()
.zrem(&processing_key, job_id.to_string())
.ignore()
.query_async(&mut conn)
.await?;
}
} else {
self.remove_from_processing(job_id).await?;
}
Ok(())
}
pub async fn requeue(&self, job: &Job) -> QueueResult<()> {
let mut job = job.clone();
job.status = JobStatus::pending();
job.started_at = None;
job.attempts = job.attempts.saturating_sub(1);
self.save_job(&job).await?;
let mut conn = self.connection.clone();
let queue_key = self.priority_queue_key(job.priority);
let score = -(job.priority as i64);
let _: () = conn.zadd(&queue_key, job.id.to_string(), score).await?;
self.remove_from_processing(job.id).await?;
Ok(())
}
pub async fn get_job(&self, job_id: JobId) -> QueueResult<Option<Job>> {
let mut conn = self.connection.clone();
let job_key = self.config.key(&format!("job:{}", job_id));
let job_json: Option<String> = conn.get(&job_key).await?;
if let Some(json) = job_json {
let job: Job = serde_json::from_str(&json)
.map_err(|e| QueueError::Deserialization(e.to_string()))?;
Ok(Some(job))
} else {
Ok(None)
}
}
async fn save_job(&self, job: &Job) -> QueueResult<()> {
let mut conn = self.connection.clone();
let job_key = self.config.key(&format!("job:{}", job.id));
let job_json =
serde_json::to_string(job).map_err(|e| QueueError::Serialization(e.to_string()))?;
let _: () = conn
.set_ex(&job_key, job_json, self.config.retention_time.as_secs())
.await?;
Ok(())
}
pub async fn size(&self) -> QueueResult<usize> {
let mut conn = self.connection.clone();
let mut pipe = redis::pipe();
for priority in [
JobPriority::Critical,
JobPriority::High,
JobPriority::Normal,
JobPriority::Low,
] {
pipe.zcard(self.priority_queue_key(priority));
}
let counts: Vec<usize> = pipe.query_async(&mut conn).await?;
Ok(counts.iter().sum())
}
pub async fn backlog_size(&self) -> QueueResult<usize> {
let mut conn = self.connection.clone();
let mut pipe = redis::pipe();
for priority in [
JobPriority::Critical,
JobPriority::High,
JobPriority::Normal,
JobPriority::Low,
] {
pipe.zcard(self.priority_queue_key(priority));
}
pipe.zcard(self.config.key("delayed"));
pipe.zcard(self.config.key("processing"));
let counts: Vec<usize> = pipe.query_async(&mut conn).await?;
Ok(counts.iter().sum())
}
pub async fn processing_len(&self) -> QueueResult<usize> {
let mut conn = self.connection.clone();
let processing_key = self.config.key("processing");
let count: usize = conn.zcard(&processing_key).await?;
Ok(count)
}
async fn move_delayed_jobs(&self) -> QueueResult<()> {
let mut conn = self.connection.clone();
let delayed_key = self.config.key("delayed");
let now = Utc::now().timestamp();
let earliest: Vec<(String, i64)> = conn.zrange_withscores(&delayed_key, 0, 0).await?;
match earliest.first() {
Some((_, score)) if *score <= now => {}
_ => return Ok(()),
}
let script = redis::Script::new(MOVE_DELAYED_SCRIPT);
let _: i64 = script
.key(&delayed_key)
.arg(&self.config.key_prefix)
.arg(now)
.invoke_async(&mut conn)
.await?;
Ok(())
}
async fn remove_from_processing(&self, job_id: JobId) -> QueueResult<()> {
let mut conn = self.connection.clone();
let processing_key = self.config.key("processing");
let _: () = conn.zrem(&processing_key, job_id.to_string()).await?;
Ok(())
}
fn priority_queue_key(&self, priority: JobPriority) -> String {
self.config
.key(&format!("pending:{:?}", priority).to_lowercase())
}
pub async fn clear(&self) -> QueueResult<()> {
let mut conn = self.connection.clone();
let pattern = format!("{}:*", self.config.key_prefix);
let mut cursor: u64 = 0;
loop {
let (next, keys): (u64, Vec<String>) = redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(&pattern)
.arg("COUNT")
.arg(SCAN_COUNT)
.query_async(&mut conn)
.await?;
if !keys.is_empty() {
let _: () = redis::cmd("UNLINK")
.arg(&keys)
.query_async(&mut conn)
.await?;
}
cursor = next;
if cursor == 0 {
break;
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_queue_config() {
let config = QueueConfig::new("redis://localhost:6379", "test");
assert_eq!(config.queue_name, "test");
assert!(config.key_prefix.contains("test"));
}
#[test]
fn test_lua_priority_mapping_matches_rust() {
for priority in [
JobPriority::Low,
JobPriority::Normal,
JobPriority::High,
JobPriority::Critical,
] {
let rust_score = -(priority as i64);
let variant = format!("{priority:?}");
let name = variant.to_lowercase();
let expected = format!("'{variant}' then pname = '{name}'; pscore = {rust_score}");
assert!(
MOVE_DELAYED_SCRIPT.contains(&expected),
"script missing mapping for {priority:?}: expected `{expected}`"
);
}
}
#[test]
fn test_priority_queue_key_layout() {
let config = QueueConfig::new("redis://localhost:6379", "jobs");
for (priority, name) in [
(JobPriority::Low, "low"),
(JobPriority::Normal, "normal"),
(JobPriority::High, "high"),
(JobPriority::Critical, "critical"),
] {
let key = config.key(&format!("pending:{priority:?}").to_lowercase());
assert_eq!(key, format!("{}:pending:{}", config.key_prefix, name));
}
}
#[test]
fn test_dequeue_script_scans_all_keys() {
assert!(DEQUEUE_POP_SCRIPT.contains("for i = 1, #KEYS do"));
assert!(DEQUEUE_POP_SCRIPT.contains("ZPOPMIN"));
}
#[tokio::test]
#[ignore = "requires a running Redis instance"]
async fn test_move_delayed_promotes_due_jobs() {
use crate::job::Job;
let queue = Queue::new("redis://localhost:6379", "test_move_delayed")
.await
.unwrap();
queue.clear().await.unwrap();
let past = Utc::now() - chrono::Duration::seconds(30);
let job = Job::new("test_move_delayed", "task", serde_json::json!({}))
.with_priority(JobPriority::High)
.schedule_at(past);
queue.enqueue_job(job).await.unwrap();
assert_eq!(queue.size().await.unwrap(), 0);
let dequeued = queue.dequeue().await.unwrap();
assert!(dequeued.is_some());
assert_eq!(dequeued.unwrap().priority, JobPriority::High);
queue.clear().await.unwrap();
}
#[test]
fn test_priority_queue_key() {
let config = QueueConfig::new("redis://localhost:6379", "test");
assert!(config.key("pending:high").contains("high"));
}
#[test]
fn test_queue_config_with_custom_prefix() {
let config = QueueConfig::new("redis://localhost:6379", "myqueue").with_key_prefix("app");
assert!(config.key_prefix.contains("app"));
}
#[test]
fn test_queue_config_default_retention() {
let config = QueueConfig::new("redis://localhost:6379", "test");
assert_eq!(config.retention_time, Duration::from_secs(86400)); }
#[test]
fn test_queue_config_custom_retention() {
let retention = Duration::from_secs(3600);
let config =
QueueConfig::new("redis://localhost:6379", "test").with_retention_time(retention);
assert_eq!(config.retention_time, retention);
}
#[test]
fn test_queue_config_default_max_size() {
let config = QueueConfig::new("redis://localhost:6379", "test");
assert_eq!(config.max_size, 0); }
#[test]
fn test_queue_config_custom_max_size() {
let config = QueueConfig::new("redis://localhost:6379", "test").with_max_size(1000);
assert_eq!(config.max_size, 1000);
}
#[test]
fn test_queue_key_generation() {
let config = QueueConfig::new("redis://localhost:6379", "jobs");
let pending_key = config.key("pending:normal");
let processing_key = config.key("processing");
let completed_key = config.key("completed");
assert!(pending_key.contains("jobs"));
assert!(processing_key.contains("jobs"));
assert!(completed_key.contains("jobs"));
}
#[test]
fn test_queue_config_clone() {
let config1 = QueueConfig::new("redis://localhost:6379", "test");
let config2 = config1.clone();
assert_eq!(config1.queue_name, config2.queue_name);
assert_eq!(config1.redis_url, config2.redis_url);
}
#[test]
fn test_queue_config_different_queues() {
let config1 = QueueConfig::new("redis://localhost:6379", "queue1");
let config2 = QueueConfig::new("redis://localhost:6379", "queue2");
assert_ne!(config1.key_prefix, config2.key_prefix);
}
#[test]
fn test_queue_config_key_consistency() {
let config = QueueConfig::new("redis://localhost:6379", "test");
let key1 = config.key("pending");
let key2 = config.key("pending");
assert_eq!(key1, key2);
}
#[test]
fn test_queue_config_builder_pattern() {
let config = QueueConfig::new("redis://localhost:6379", "test")
.with_key_prefix("app")
.with_retention_time(Duration::from_secs(7200))
.with_max_size(500);
assert!(config.key_prefix.contains("app"));
assert_eq!(config.retention_time, Duration::from_secs(7200));
assert_eq!(config.max_size, 500);
}
#[test]
fn test_queue_config_redis_url() {
let url = "redis://user:pass@host:6380/2";
let config = QueueConfig::new(url, "test");
assert_eq!(config.redis_url, url);
}
#[test]
fn test_queue_config_key_with_empty_suffix() {
let config = QueueConfig::new("redis://localhost:6379", "test");
let key = config.key("");
assert!(key.contains("test"));
}
#[test]
fn test_queue_config_key_with_special_characters() {
let config = QueueConfig::new("redis://localhost:6379", "test");
let key = config.key("pending:high:priority");
assert!(key.contains("pending:high:priority"));
}
#[test]
fn test_queue_config_multiple_prefixes() {
let config1 =
QueueConfig::new("redis://localhost:6379", "app1").with_key_prefix("production");
let config2 =
QueueConfig::new("redis://localhost:6379", "app2").with_key_prefix("development");
let key1 = config1.key("jobs");
let key2 = config2.key("jobs");
assert_ne!(key1, key2);
}
#[test]
fn test_queue_config_unlimited_max_size() {
let config = QueueConfig::new("redis://localhost:6379", "test").with_max_size(0);
assert_eq!(config.max_size, 0);
}
#[test]
fn test_queue_config_large_retention() {
let week = Duration::from_secs(7 * 24 * 3600);
let config = QueueConfig::new("redis://localhost:6379", "test").with_retention_time(week);
assert_eq!(config.retention_time, week);
}
}