use std::sync::{Arc, Mutex};
use std::time::Duration;
use rainier_cache::{LockGuard, LockManager};
use rainier_container::Container;
use rainier_support::{BoxFuture, Result};
use crate::job::{Job, JobContext, JobRegistry, QueuedJob};
use crate::queue::Queue;
pub struct SyncQueue {
registry: Arc<JobRegistry>,
container: Arc<Container>,
}
impl SyncQueue {
pub fn new(registry: Arc<JobRegistry>, container: Arc<Container>) -> Self {
Self { registry, container }
}
}
impl Queue for SyncQueue {
fn name(&self) -> &str {
"sync"
}
fn push<'a>(&'a self, mut job: QueuedJob) -> BoxFuture<'a, Result<String>> {
Box::pin(async move {
job.attempts = 1;
let id = job.id.clone();
let context = Arc::new(JobContext::new(
Arc::clone(&self.container),
job.id.clone(),
job.queue.clone(),
1,
1,
));
self.registry.run(&job, context).await?;
Ok(id)
})
}
fn reserve<'a>(&'a self, _queue: &'a str) -> BoxFuture<'a, Result<Option<QueuedJob>>> {
Box::pin(async { Ok(None) })
}
fn acknowledge<'a>(&'a self, _job: &'a QueuedJob) -> BoxFuture<'a, Result<()>> {
Box::pin(async { Ok(()) })
}
fn release<'a>(&'a self, _job: &'a QueuedJob, _delay: Duration) -> BoxFuture<'a, Result<()>> {
Box::pin(async { Ok(()) })
}
fn fail<'a>(&'a self, _job: &'a QueuedJob, _error: &'a str) -> BoxFuture<'a, Result<()>> {
Box::pin(async { Ok(()) })
}
fn size<'a>(&'a self, _queue: &'a str) -> BoxFuture<'a, Result<u64>> {
Box::pin(async { Ok(0) })
}
fn clear<'a>(&'a self, _queue: &'a str) -> BoxFuture<'a, Result<u64>> {
Box::pin(async { Ok(0) })
}
}
pub struct PendingDispatch<'a> {
manager: &'a QueueManager,
job: QueuedJob,
unique: Option<(String, Duration)>,
}
impl<'a> PendingDispatch<'a> {
pub fn on_queue(mut self, queue: impl Into<String>) -> Self {
self.job = self.job.on_queue(queue);
self
}
pub fn delay(mut self, delay: Duration) -> Self {
self.job = self.job.delayed_by(delay);
self
}
pub fn tries(mut self, tries: u32) -> Self {
self.job.max_attempts = tries.max(1);
self
}
pub async fn send(self) -> Result<Option<String>> {
let mut job = self.job;
let claim = match &self.unique {
Some((key, ttl)) => {
let claim = self.manager.claim(key, *ttl).await?;
if matches!(claim, Claim::AlreadyQueued) {
tracing::debug!(key, "dropping a duplicate dispatch");
return Ok(None);
}
if matches!(claim, Claim::Held(_)) {
job.unique_key = Some(key.clone());
}
claim
}
None => Claim::NotEnforced,
};
let id = self.manager.push(job).await?;
if let Claim::Held(guard) = claim {
guard.keep();
}
Ok(Some(id))
}
}
enum Claim {
Held(LockGuard),
AlreadyQueued,
NotEnforced,
}
pub struct QueueManager {
queue: Arc<dyn Queue>,
registry: Arc<JobRegistry>,
recorded: Option<Mutex<Vec<QueuedJob>>>,
locks: Option<LockManager>,
}
impl QueueManager {
pub fn new(queue: Arc<dyn Queue>, registry: Arc<JobRegistry>) -> Self {
Self { queue, registry, recorded: None, locks: None }
}
pub fn with_locks(mut self, locks: LockManager) -> Self {
self.locks = Some(locks);
self
}
pub fn fake() -> Self {
Self {
queue: Arc::new(SyncQueue::new(
Arc::new(JobRegistry::new()),
Arc::new(Container::new()),
)),
registry: Arc::new(JobRegistry::new()),
recorded: Some(Mutex::new(Vec::new())),
locks: None,
}
}
pub fn is_faking(&self) -> bool {
self.recorded.is_some()
}
pub fn queue(&self) -> &Arc<dyn Queue> {
&self.queue
}
pub fn registry(&self) -> &Arc<JobRegistry> {
&self.registry
}
pub async fn dispatch<J: Job>(&self, job: J) -> Result<Option<String>> {
self.pending(job)?.send().await
}
fn unique_key<J: Job>(job: &J) -> Option<(String, Duration)> {
job.unique_id().map(|id| (format!("queue:unique:{}:{id}", J::NAME), J::UNIQUE_FOR))
}
async fn claim(&self, key: &str, ttl: Duration) -> Result<Claim> {
let Some(locks) = &self.locks else {
tracing::warn!(
key,
"a job declares `unique_id` but the queue has no lock manager; dispatching without deduplicating"
);
return Ok(Claim::NotEnforced);
};
Ok(match locks.lock(key, ttl).acquire().await? {
Some(guard) => Claim::Held(guard),
None => Claim::AlreadyQueued,
})
}
pub async fn release_uniqueness(&self, key: &str) -> Result<bool> {
let Some(locks) = &self.locks else { return Ok(false) };
locks.lock(key.to_string(), Duration::from_secs(1)).force_release().await
}
pub fn pending<J: Job>(&self, job: J) -> Result<PendingDispatch<'_>> {
let unique = Self::unique_key(&job);
Ok(PendingDispatch { manager: self, job: QueuedJob::from_job(&job)?, unique })
}
pub async fn dispatch_on<J: Job>(
&self,
queue: impl Into<String>,
job: J,
) -> Result<Option<String>> {
self.pending(job)?.on_queue(queue).send().await
}
pub async fn dispatch_after<J: Job>(&self, delay: Duration, job: J) -> Result<Option<String>> {
self.pending(job)?.delay(delay).send().await
}
pub async fn dispatch_now<J: Job>(&self, job: J, container: Arc<Container>) -> Result<()> {
let queued = QueuedJob::from_job(&job)?;
let context =
Arc::new(JobContext::new(container, queued.id.clone(), queued.queue.clone(), 1, 1));
self.registry.run(&queued, context).await
}
async fn push(&self, job: QueuedJob) -> Result<String> {
if let Some(recorded) = &self.recorded {
let id = job.id.clone();
recorded.lock().expect("recorder lock poisoned").push(job);
return Ok(id);
}
self.queue.push(job).await
}
pub fn pushed<J: Job>(&self) -> Vec<QueuedJob> {
let Some(recorded) = &self.recorded else {
return Vec::new();
};
recorded
.lock()
.expect("recorder lock poisoned")
.iter()
.filter(|job| job.name == J::NAME)
.cloned()
.collect()
}
pub fn all_pushed(&self) -> Vec<QueuedJob> {
match &self.recorded {
Some(recorded) => recorded.lock().expect("recorder lock poisoned").clone(),
None => Vec::new(),
}
}
pub fn assert_pushed<J: Job>(&self) {
self.require_faking("assert_pushed");
assert!(
!self.pushed::<J>().is_empty(),
"expected a `{}` job to have been queued. Queued: {:?}",
J::NAME,
self.all_pushed().iter().map(|job| job.name.clone()).collect::<Vec<_>>()
);
}
pub fn assert_pushed_times<J: Job>(&self, times: usize) {
self.require_faking("assert_pushed_times");
let actual = self.pushed::<J>().len();
assert_eq!(
actual,
times,
"expected `{}` to be queued {times} time(s), but it was queued {actual}",
J::NAME
);
}
pub fn assert_not_pushed<J: Job>(&self) {
self.require_faking("assert_not_pushed");
assert!(
self.pushed::<J>().is_empty(),
"expected no `{}` job to be queued, but one was",
J::NAME
);
}
pub fn assert_pushed_on<J: Job>(&self, queue: &str) {
self.require_faking("assert_pushed_on");
let queues: Vec<String> = self.pushed::<J>().iter().map(|job| job.queue.clone()).collect();
assert!(
queues.iter().any(|q| q == queue),
"expected `{}` to be queued on `{queue}`, but it was queued on {queues:?}",
J::NAME
);
}
fn require_faking(&self, method: &str) {
assert!(
self.is_faking(),
"`{method}` needs a faking manager — build it with `QueueManager::fake()`, \
otherwise nothing is recorded and the assertion is meaningless"
);
}
}
impl std::fmt::Debug for QueueManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("QueueManager")
.field("driver", &self.queue.name())
.field("jobs", &self.registry.names())
.field("faking", &self.is_faking())
.finish()
}
}
const _: fn() = || {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<QueueManager>();
};
#[cfg(test)]
mod unique_tests {
use super::*;
use crate::queue::MemoryQueue;
use rainier_cache::MemoryCache;
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize)]
struct RebuildIndex;
#[async_trait::async_trait]
impl Job for RebuildIndex {
const NAME: &'static str = "search.rebuild";
fn unique_id(&self) -> Option<String> {
Some(String::new())
}
async fn handle(&self, _: &JobContext) -> Result<()> {
Ok(())
}
}
#[derive(Serialize, Deserialize)]
struct SendInvoice {
id: u64,
}
#[async_trait::async_trait]
impl Job for SendInvoice {
const NAME: &'static str = "billing.invoice";
fn unique_id(&self) -> Option<String> {
Some(self.id.to_string())
}
async fn handle(&self, _: &JobContext) -> Result<()> {
Ok(())
}
}
#[derive(Serialize, Deserialize)]
struct Ordinary;
#[async_trait::async_trait]
impl Job for Ordinary {
const NAME: &'static str = "test.ordinary";
async fn handle(&self, _: &JobContext) -> Result<()> {
Ok(())
}
}
fn manager() -> (QueueManager, Arc<MemoryQueue>) {
let queue = Arc::new(MemoryQueue::new());
let manager = QueueManager::new(Arc::clone(&queue) as Arc<_>, Arc::new(JobRegistry::new()))
.with_locks(rainier_cache::LockManager::new(Arc::new(MemoryCache::new())));
(manager, queue)
}
#[tokio::test]
async fn a_second_identical_dispatch_is_dropped() {
let (manager, queue) = manager();
assert!(manager.dispatch(RebuildIndex).await.unwrap().is_some());
assert!(manager.dispatch(RebuildIndex).await.unwrap().is_none());
assert_eq!(queue.size("default").await.unwrap(), 1, "one copy, not two");
}
#[tokio::test]
async fn different_ids_are_different_jobs() {
let (manager, queue) = manager();
assert!(manager.dispatch(SendInvoice { id: 1 }).await.unwrap().is_some());
assert!(manager.dispatch(SendInvoice { id: 2 }).await.unwrap().is_some());
assert!(manager.dispatch(SendInvoice { id: 1 }).await.unwrap().is_none());
assert_eq!(queue.size("default").await.unwrap(), 2);
}
#[tokio::test]
async fn two_job_types_sharing_an_id_do_not_collide() {
let (manager, queue) = manager();
assert!(manager.dispatch(RebuildIndex).await.unwrap().is_some());
assert!(manager.dispatch(SendInvoice { id: 0 }).await.unwrap().is_some());
assert_eq!(queue.size("default").await.unwrap(), 2);
}
#[tokio::test]
async fn an_ordinary_job_is_never_deduplicated() {
let (manager, queue) = manager();
for _ in 0..3 {
assert!(manager.dispatch(Ordinary).await.unwrap().is_some());
}
assert_eq!(queue.size("default").await.unwrap(), 3);
}
#[tokio::test]
async fn the_queued_job_carries_the_key_so_a_worker_can_release_it() {
let (manager, queue) = manager();
manager.dispatch(RebuildIndex).await.unwrap();
let job = queue.reserve("default").await.unwrap().expect("a job");
assert_eq!(job.unique_key.as_deref(), Some("queue:unique:search.rebuild:"));
}
#[tokio::test]
async fn without_a_lock_manager_a_unique_job_is_dispatched_anyway() {
let queue = Arc::new(MemoryQueue::new());
let manager = QueueManager::new(Arc::clone(&queue) as Arc<_>, Arc::new(JobRegistry::new()));
assert!(manager.dispatch(RebuildIndex).await.unwrap().is_some());
assert!(manager.dispatch(RebuildIndex).await.unwrap().is_some());
assert_eq!(queue.size("default").await.unwrap(), 2);
}
#[tokio::test]
async fn releasing_the_lock_lets_the_next_one_in() {
let (manager, _queue) = manager();
manager.dispatch(RebuildIndex).await.unwrap();
assert!(manager.dispatch(RebuildIndex).await.unwrap().is_none());
manager.release_uniqueness("queue:unique:search.rebuild:").await.unwrap();
assert!(
manager.dispatch(RebuildIndex).await.unwrap().is_some(),
"once the first has run, another may be queued"
);
}
#[tokio::test]
async fn a_configured_dispatch_is_deduplicated_too() {
let (manager, queue) = manager();
assert!(manager
.pending(RebuildIndex)
.unwrap()
.on_queue("high")
.send()
.await
.unwrap()
.is_some());
assert!(manager
.pending(RebuildIndex)
.unwrap()
.on_queue("high")
.send()
.await
.unwrap()
.is_none());
assert_eq!(queue.size("high").await.unwrap(), 1);
}
#[tokio::test]
async fn a_pending_dispatch_that_is_dropped_leaves_no_lock() {
let (manager, _queue) = manager();
let abandoned = manager.pending(RebuildIndex).unwrap().on_queue("high");
drop(abandoned);
assert!(manager.dispatch(RebuildIndex).await.unwrap().is_some());
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::queue::MemoryQueue;
use rainier_support::Error;
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicU32, Ordering};
type Runs = Arc<AtomicU32>;
fn counting_container() -> (Arc<Container>, Runs) {
let runs: Runs = Arc::new(AtomicU32::new(0));
let container = Arc::new(Container::new());
container.instance(Arc::clone(&runs));
(container, runs)
}
#[derive(Serialize, Deserialize)]
struct Ping;
#[async_trait::async_trait]
impl Job for Ping {
const NAME: &'static str = "test.ping";
async fn handle(&self, context: &JobContext) -> Result<()> {
if let Ok(runs) = context.resolve::<Runs>() {
runs.fetch_add(1, Ordering::SeqCst);
}
Ok(())
}
}
#[derive(Serialize, Deserialize)]
struct Boom;
#[async_trait::async_trait]
impl Job for Boom {
const NAME: &'static str = "test.boom";
async fn handle(&self, _: &JobContext) -> Result<()> {
Err(Error::internal("exploded"))
}
}
fn registry() -> Arc<JobRegistry> {
Arc::new(JobRegistry::new().with::<Ping>().with::<Boom>())
}
fn manager(queue: Arc<dyn Queue>) -> QueueManager {
QueueManager::new(queue, registry())
}
#[tokio::test]
async fn dispatching_enqueues_the_job() {
let queue = Arc::new(MemoryQueue::new());
let manager = manager(Arc::clone(&queue) as Arc<dyn Queue>);
let id = manager.dispatch(Ping).await.unwrap();
assert!(id.is_some_and(|id| !id.is_empty()));
assert_eq!(queue.size("default").await.unwrap(), 1);
}
#[tokio::test]
async fn dispatch_options_apply() {
let queue = Arc::new(MemoryQueue::new());
let manager = manager(Arc::clone(&queue) as Arc<dyn Queue>);
manager
.pending(Ping)
.unwrap()
.on_queue("mail")
.delay(Duration::from_secs(60))
.tries(5)
.send()
.await
.unwrap();
let pending = queue.pending("mail");
assert_eq!(pending.len(), 1);
assert_eq!(pending[0].max_attempts, 5);
assert!(!pending[0].is_available(), "the delay should hold it back");
}
#[tokio::test]
async fn dispatch_on_and_after_are_shorthands() {
let queue = Arc::new(MemoryQueue::new());
let manager = manager(Arc::clone(&queue) as Arc<dyn Queue>);
manager.dispatch_on("mail", Ping).await.unwrap();
manager.dispatch_after(Duration::from_secs(30), Ping).await.unwrap();
assert_eq!(queue.size("mail").await.unwrap(), 1);
assert!(!queue.pending("default")[0].is_available());
}
#[tokio::test]
async fn the_sync_driver_runs_the_job_immediately() {
let (container, runs) = counting_container();
let manager = manager(Arc::new(SyncQueue::new(registry(), container)));
manager.dispatch(Ping).await.unwrap();
assert_eq!(runs.load(Ordering::SeqCst), 1, "it ran during dispatch");
}
#[tokio::test]
async fn the_sync_driver_surfaces_a_failure_to_the_dispatcher() {
let sync = Arc::new(SyncQueue::new(registry(), Arc::new(Container::new())));
let manager = manager(sync);
let err = manager.dispatch(Boom).await.unwrap_err();
assert!(err.message().contains("exploded"), "{}", err.message());
}
#[tokio::test]
async fn dispatch_now_bypasses_the_queue() {
let (container, runs) = counting_container();
let queue = Arc::new(MemoryQueue::new());
let manager = manager(Arc::clone(&queue) as Arc<dyn Queue>);
manager.dispatch_now(Ping, container).await.unwrap();
assert_eq!(runs.load(Ordering::SeqCst), 1);
assert_eq!(queue.size("default").await.unwrap(), 0, "nothing was queued");
}
#[tokio::test]
async fn a_fake_records_instead_of_queueing() {
let manager = QueueManager::fake();
manager.dispatch(Ping).await.unwrap();
manager.dispatch_on("mail", Ping).await.unwrap();
manager.assert_pushed::<Ping>();
manager.assert_pushed_times::<Ping>(2);
manager.assert_pushed_on::<Ping>("mail");
manager.assert_not_pushed::<Boom>();
assert_eq!(manager.all_pushed().len(), 2);
}
#[tokio::test]
#[should_panic(expected = "needs a faking manager")]
async fn assertions_refuse_to_pass_vacuously() {
let queue = Arc::new(MemoryQueue::new());
let manager = manager(queue as Arc<dyn Queue>);
manager.dispatch(Ping).await.unwrap();
manager.assert_not_pushed::<Ping>();
}
#[tokio::test]
#[should_panic(expected = "test.boom")]
async fn a_missing_dispatch_reports_what_was_queued() {
let manager = QueueManager::fake();
manager.dispatch(Ping).await.unwrap();
manager.assert_pushed::<Boom>();
}
}