use crate::{JobEntry, JobId, JobState, QueueBackend, QueueError, QueueResult, QueueStats, JobResult, QueueConfig};
use async_trait::async_trait;
use dashmap::DashMap;
use parking_lot::RwLock;
use std::collections::BinaryHeap;
use std::cmp::Ordering;
use std::sync::Arc;
use tokio::time::Instant;
use chrono::Utc;
#[derive(Debug, Clone)]
struct PriorityJobEntry {
entry: JobEntry,
enqueue_time: Instant,
}
impl PartialEq for PriorityJobEntry {
fn eq(&self, other: &Self) -> bool {
self.entry.priority() == other.entry.priority() &&
self.entry.run_at() == other.entry.run_at() &&
self.enqueue_time == other.enqueue_time
}
}
impl Eq for PriorityJobEntry {}
impl PartialOrd for PriorityJobEntry {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for PriorityJobEntry {
fn cmp(&self, other: &Self) -> Ordering {
match self.entry.priority().cmp(&other.entry.priority()) {
Ordering::Equal => {
match other.entry.run_at().cmp(&self.entry.run_at()) {
Ordering::Equal => other.enqueue_time.cmp(&self.enqueue_time),
other_ord => other_ord,
}
}
priority_ord => priority_ord,
}
}
}
pub struct MemoryBackend {
config: QueueConfig,
jobs: DashMap<JobId, JobEntry>,
pending_queue: Arc<RwLock<BinaryHeap<PriorityJobEntry>>>,
stats: Arc<RwLock<QueueStats>>,
}
impl MemoryBackend {
pub fn new(config: QueueConfig) -> Self {
Self {
config,
jobs: DashMap::new(),
pending_queue: Arc::new(RwLock::new(BinaryHeap::new())),
stats: Arc::new(RwLock::new(QueueStats::default())),
}
}
fn update_stats(&self, old_state: Option<JobState>, new_state: JobState) {
let mut stats = self.stats.write();
if let Some(old) = old_state {
match old {
JobState::Pending => stats.pending_jobs = stats.pending_jobs.saturating_sub(1),
JobState::Processing => stats.processing_jobs = stats.processing_jobs.saturating_sub(1),
JobState::Completed => stats.completed_jobs = stats.completed_jobs.saturating_sub(1),
JobState::Failed => stats.failed_jobs = stats.failed_jobs.saturating_sub(1),
JobState::Dead => stats.dead_jobs = stats.dead_jobs.saturating_sub(1),
}
} else {
stats.total_jobs += 1;
}
match new_state {
JobState::Pending => stats.pending_jobs += 1,
JobState::Processing => stats.processing_jobs += 1,
JobState::Completed => stats.completed_jobs += 1,
JobState::Failed => stats.failed_jobs += 1,
JobState::Dead => stats.dead_jobs += 1,
}
}
fn get_next_ready_job(&self) -> Option<JobEntry> {
let mut queue = self.pending_queue.write();
let now = Utc::now();
while let Some(priority_entry) = queue.peek() {
if !self.jobs.contains_key(&priority_entry.entry.id()) {
queue.pop();
continue;
}
if priority_entry.entry.is_ready() {
let priority_entry = queue.pop().unwrap();
return Some(priority_entry.entry);
} else if priority_entry.entry.run_at() > now {
break;
} else {
queue.pop();
}
}
None
}
}
#[async_trait]
impl QueueBackend for MemoryBackend {
async fn enqueue(&self, job: JobEntry) -> QueueResult<JobId> {
let job_id = job.id();
if *self.config.get_max_queue_size() > 0 && self.jobs.len() >= *self.config.get_max_queue_size() {
return Err(QueueError::Configuration(
format!("Queue size limit exceeded: {}", *self.config.get_max_queue_size())
));
}
self.update_stats(None, job.state().clone());
if job.state() == &JobState::Pending {
let priority_entry = PriorityJobEntry {
entry: job.clone(),
enqueue_time: Instant::now(),
};
self.pending_queue.write().push(priority_entry);
}
self.jobs.insert(job_id, job);
Ok(job_id)
}
async fn dequeue(&self) -> QueueResult<Option<JobEntry>> {
if let Some(mut job) = self.get_next_ready_job() {
let old_state = job.state().clone();
job.mark_processing();
self.update_stats(Some(old_state), job.state().clone());
self.jobs.insert(job.id(), job.clone());
Ok(Some(job))
} else {
Ok(None)
}
}
async fn complete(&self, job_id: JobId, result: JobResult<()>) -> QueueResult<()> {
if let Some(mut job_entry) = self.jobs.get_mut(&job_id) {
let old_state = job_entry.state().clone();
match result {
Ok(_) => {
job_entry.mark_completed();
self.update_stats(Some(old_state), job_entry.state().clone());
}
Err(error) => {
let error_message = error.to_string();
job_entry.mark_failed(error_message);
let new_state = job_entry.state().clone();
self.update_stats(Some(old_state), new_state.clone());
if new_state == JobState::Failed {
let priority_entry = PriorityJobEntry {
entry: job_entry.clone(),
enqueue_time: Instant::now(),
};
self.pending_queue.write().push(priority_entry);
}
}
}
Ok(())
} else {
Err(QueueError::JobNotFound(job_id.to_string()))
}
}
async fn get_job(&self, job_id: JobId) -> QueueResult<Option<JobEntry>> {
Ok(self.jobs.get(&job_id).map(|entry| entry.clone()))
}
async fn get_jobs_by_state(&self, state: JobState, limit: Option<usize>) -> QueueResult<Vec<JobEntry>> {
let mut jobs: Vec<JobEntry> = self.jobs
.iter()
.filter(|entry| entry.state() == &state)
.map(|entry| entry.clone())
.collect();
jobs.sort_by(|a, b| a.created_at.cmp(&b.created_at));
if let Some(limit) = limit {
jobs.truncate(limit);
}
Ok(jobs)
}
async fn remove_job(&self, job_id: JobId) -> QueueResult<bool> {
if let Some((_, job)) = self.jobs.remove(&job_id) {
let mut stats = self.stats.write();
match job.state() {
JobState::Pending => stats.pending_jobs = stats.pending_jobs.saturating_sub(1),
JobState::Processing => stats.processing_jobs = stats.processing_jobs.saturating_sub(1),
JobState::Completed => stats.completed_jobs = stats.completed_jobs.saturating_sub(1),
JobState::Failed => stats.failed_jobs = stats.failed_jobs.saturating_sub(1),
JobState::Dead => stats.dead_jobs = stats.dead_jobs.saturating_sub(1),
}
stats.total_jobs = stats.total_jobs.saturating_sub(1);
Ok(true)
} else {
Ok(false)
}
}
async fn clear(&self) -> QueueResult<()> {
self.jobs.clear();
self.pending_queue.write().clear();
*self.stats.write() = QueueStats::default();
Ok(())
}
async fn stats(&self) -> QueueResult<QueueStats> {
Ok(self.stats.read().clone())
}
async fn requeue_job(&self, job_id: JobId, _job: JobEntry) -> QueueResult<bool> {
if let Some(mut existing_job) = self.jobs.get_mut(&job_id) {
if existing_job.state() == &JobState::Dead {
existing_job.reset_for_retry();
if existing_job.is_ready() {
let priority_entry = PriorityJobEntry {
entry: existing_job.clone(),
enqueue_time: Instant::now(),
};
self.pending_queue.write().push(priority_entry);
}
let mut stats = self.stats.write();
stats.dead_jobs = stats.dead_jobs.saturating_sub(1);
stats.pending_jobs = stats.pending_jobs.saturating_add(1);
Ok(true)
} else {
Ok(false)
}
} else {
Ok(false)
}
}
async fn clear_jobs_by_state(&self, state: JobState) -> QueueResult<u64> {
let mut count = 0u64;
let mut stats = self.stats.write();
self.jobs.retain(|_, job| {
if job.state() == &state {
count += 1;
match state {
JobState::Pending => stats.pending_jobs = stats.pending_jobs.saturating_sub(1),
JobState::Processing => stats.processing_jobs = stats.processing_jobs.saturating_sub(1),
JobState::Completed => stats.completed_jobs = stats.completed_jobs.saturating_sub(1),
JobState::Failed => stats.failed_jobs = stats.failed_jobs.saturating_sub(1),
JobState::Dead => stats.dead_jobs = stats.dead_jobs.saturating_sub(1),
}
false } else {
true }
});
if state == JobState::Pending {
let mut pending_queue = self.pending_queue.write();
pending_queue.retain(|priority_entry| {
priority_entry.entry.state() != &JobState::Pending
});
}
stats.total_jobs = stats.total_jobs.saturating_sub(count);
Ok(count)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Job, Priority};
use serde::{Deserialize, Serialize};
use std::time::Duration;
#[derive(Debug, Clone, Serialize, Deserialize)]
struct TestJob {
id: u32,
message: String,
}
#[async_trait]
impl Job for TestJob {
async fn execute(&self) -> JobResult<()> {
Ok(())
}
fn job_type(&self) -> &'static str {
"test"
}
}
#[tokio::test]
async fn test_memory_backend_basic_operations() {
let backend = MemoryBackend::new(QueueConfig::default());
let job = TestJob {
id: 1,
message: "test job".to_string(),
};
let entry = JobEntry::new(job, Some(Priority::Normal), None).unwrap();
let job_id = entry.id();
backend.enqueue(entry).await.unwrap();
let stats = backend.stats().await.unwrap();
assert_eq!(stats.pending_jobs, 1);
assert_eq!(stats.total_jobs, 1);
let dequeued = backend.dequeue().await.unwrap().unwrap();
assert_eq!(dequeued.id(), job_id);
assert_eq!(dequeued.state(), &JobState::Processing);
backend.complete(job_id, Ok(())).await.unwrap();
let stats = backend.stats().await.unwrap();
assert_eq!(stats.completed_jobs, 1);
assert_eq!(stats.processing_jobs, 0);
assert_eq!(stats.pending_jobs, 0);
}
#[tokio::test]
async fn test_priority_ordering() {
let backend = MemoryBackend::new(QueueConfig::default());
let low_job = TestJob { id: 1, message: "low".to_string() };
let high_job = TestJob { id: 2, message: "high".to_string() };
let normal_job = TestJob { id: 3, message: "normal".to_string() };
let low_entry = JobEntry::new(low_job, Some(Priority::Low), None).unwrap();
let high_entry = JobEntry::new(high_job, Some(Priority::High), None).unwrap();
let normal_entry = JobEntry::new(normal_job, Some(Priority::Normal), None).unwrap();
backend.enqueue(low_entry).await.unwrap();
backend.enqueue(high_entry).await.unwrap();
backend.enqueue(normal_entry).await.unwrap();
let first = backend.dequeue().await.unwrap().unwrap();
assert_eq!(first.priority(), Priority::High);
let second = backend.dequeue().await.unwrap().unwrap();
assert_eq!(second.priority(), Priority::Normal);
let third = backend.dequeue().await.unwrap().unwrap();
assert_eq!(third.priority(), Priority::Low);
}
#[tokio::test]
async fn test_ghost_job_cleanup() {
let config = crate::config::QueueConfigBuilder::testing().build().expect("Failed to build config");
let backend = MemoryBackend::new(config);
let job = TestJob { id: 1, message: "ghost test".to_string() };
let entry = JobEntry::new(job, Some(Priority::Normal), None).unwrap();
let job_id = backend.enqueue(entry).await.unwrap();
assert!(backend.jobs.contains_key(&job_id));
assert_eq!(backend.pending_queue.read().len(), 1);
backend.jobs.remove(&job_id);
assert_eq!(backend.pending_queue.read().len(), 1);
let result = backend.dequeue().await.unwrap();
assert!(result.is_none());
assert_eq!(backend.pending_queue.read().len(), 0);
}
#[tokio::test]
async fn test_multiple_ghost_jobs_cleanup() {
let config = crate::config::QueueConfigBuilder::testing().build().expect("Failed to build config");
let backend = MemoryBackend::new(config);
let mut job_ids = Vec::new();
for i in 1..=5 {
let job = TestJob { id: i, message: format!("ghost test {}", i) };
let entry = JobEntry::new(job, Some(Priority::Normal), None).unwrap();
let job_id = backend.enqueue(entry).await.unwrap();
job_ids.push(job_id);
}
assert_eq!(backend.pending_queue.read().len(), 5);
assert_eq!(backend.jobs.len(), 5);
for &job_id in &job_ids[0..3] {
backend.jobs.remove(&job_id);
}
assert_eq!(backend.pending_queue.read().len(), 5);
assert_eq!(backend.jobs.len(), 2);
let result = backend.dequeue().await.unwrap().unwrap();
assert_eq!(result.payload.get("id").unwrap().as_u64().unwrap(), 4);
assert_eq!(backend.pending_queue.read().len(), 1);
}
#[tokio::test]
async fn test_delayed_job() {
let backend = MemoryBackend::new(QueueConfig::default());
let job = TestJob {
id: 1,
message: "delayed job".to_string(),
};
let delay = Duration::from_millis(100);
let entry = JobEntry::new(job, None, Some(delay)).unwrap();
backend.enqueue(entry).await.unwrap();
let result = backend.dequeue().await.unwrap();
assert!(result.is_none());
tokio::time::sleep(delay + Duration::from_millis(10)).await;
let result = backend.dequeue().await.unwrap();
assert!(result.is_some());
}
#[tokio::test]
async fn test_job_failure_and_retry() {
let backend = MemoryBackend::new(QueueConfig::default());
let job = TestJob {
id: 1,
message: "failing job".to_string(),
};
let entry = JobEntry::new(job, None, None).unwrap();
let job_id = entry.id();
backend.enqueue(entry).await.unwrap();
let _job_entry = backend.dequeue().await.unwrap().unwrap();
let error = Box::new(std::io::Error::new(std::io::ErrorKind::Other, "test error"));
backend.complete(job_id, Err(error)).await.unwrap();
let stats = backend.stats().await.unwrap();
assert_eq!(stats.failed_jobs, 1);
assert_eq!(stats.processing_jobs, 0);
let failed_jobs = backend.get_jobs_by_state(JobState::Failed, None).await.unwrap();
assert_eq!(failed_jobs.len(), 1);
assert_eq!(failed_jobs[0].attempts(), 1);
}
#[tokio::test]
async fn test_queue_size_limit() {
let config = crate::config::QueueConfigBuilder::new()
.max_queue_size(2)
.build().expect("Failed to build config");
let backend = MemoryBackend::new(config);
for i in 1..=2 {
let job = TestJob { id: i, message: format!("job {}", i) };
let entry = JobEntry::new(job, None, None).unwrap();
backend.enqueue(entry).await.unwrap();
}
let job = TestJob { id: 3, message: "overflow job".to_string() };
let entry = JobEntry::new(job, None, None).unwrap();
let result = backend.enqueue(entry).await;
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), QueueError::Configuration(_)));
}
}