use std::collections::HashSet;
use std::sync::{Arc, Barrier};
use std::thread;
use std::time::{Duration, SystemTime};
use armdb::armour::Db;
use armdb::{CollectionMeta, DbError, FixedConfig, NoHook, QueueConfig, queue};
use armour_core::GetType;
#[derive(
Clone,
Copy,
Debug,
PartialEq,
Eq,
GetType,
zerocopy::FromBytes,
zerocopy::IntoBytes,
zerocopy::Immutable,
zerocopy::KnownLayout,
)]
#[repr(transparent)]
pub struct Job(u64);
queue!(Jobs, payload = Job, name = "queue_tx_jobs");
queue!(OtherJobs, payload = Job, name = "queue_tx_other_jobs");
#[derive(
Clone,
Copy,
Debug,
PartialEq,
Eq,
GetType,
zerocopy::FromBytes,
zerocopy::IntoBytes,
zerocopy::Immutable,
zerocopy::KnownLayout,
)]
#[repr(transparent)]
struct Effect(u64);
impl CollectionMeta for Effect {
type SelfId = [u8; 8];
const NAME: &'static str = "queue_tx_effect";
}
fn open_queue(config: QueueConfig) -> (tempfile::TempDir, Arc<Db>, Arc<armdb::Queue<Jobs>>) {
let dir = tempfile::tempdir().unwrap();
let db = Arc::new(Db::open_test(dir.path()).unwrap());
let queue = db.open_queue::<Jobs>(config).unwrap();
(dir, db, queue)
}
#[test]
fn transaction_queue_methods_take_arc_queue_handles() {
let _: fn(
armdb::Tx<'static>,
&'static Arc<armdb::Queue<Jobs>>,
Job,
Option<SystemTime>,
) -> armdb::DbResult<armdb::Tx<'static>> = armdb::Tx::push::<Jobs>;
let _: fn(
armdb::Tx<'static>,
&'static Arc<armdb::Queue<Jobs>>,
armdb::LeaseToken,
) -> armdb::DbResult<armdb::Tx<'static>> = armdb::Tx::ack::<Jobs>;
let _: fn(
armdb::Tx<'static>,
&'static Arc<armdb::Queue<Jobs>>,
armdb::LeaseToken,
) -> armdb::DbResult<armdb::Tx<'static>> = armdb::Tx::nack::<Jobs>;
}
#[test]
fn expiry_releases_once_and_stale_ack_cannot_delete_the_new_attempt() {
let (_dir, _db, queue) = open_queue(
QueueConfig::balanced()
.lease_timeout(Duration::from_millis(1))
.build()
.unwrap(),
);
queue.push(Job(7), None).unwrap();
let first = queue.lease(1, None).unwrap().pop().unwrap();
thread::sleep(Duration::from_millis(3));
let second = queue.lease(1, None).unwrap().pop().unwrap();
assert_ne!(first.token(), second.token());
assert!(matches!(
queue.ack(first.token()),
Err(DbError::LeaseExpired)
));
queue.ack(second.token()).unwrap();
assert_eq!(queue.len(), 0);
}
#[test]
fn lease_token_cannot_ack_a_row_in_another_queue() {
let dir = tempfile::tempdir().unwrap();
let db = Arc::new(Db::open_test(dir.path()).unwrap());
let first_queue = db
.open_queue::<Jobs>(QueueConfig::balanced().build().unwrap())
.unwrap();
let second_queue = db
.open_queue::<OtherJobs>(QueueConfig::balanced().build().unwrap())
.unwrap();
let visible_at = SystemTime::UNIX_EPOCH;
first_queue.push(Job(1), Some(visible_at)).unwrap();
second_queue.push(Job(2), Some(visible_at)).unwrap();
let first = first_queue.lease(1, None).unwrap().pop().unwrap();
let second = second_queue.lease(1, None).unwrap().pop().unwrap();
assert_eq!(first.token().key(), second.token().key());
assert_eq!(first.token().epoch(), second.token().epoch());
assert!(matches!(
second_queue.ack(first.token()),
Err(DbError::LeaseExpired)
));
assert_eq!(second_queue.len(), 1);
second_queue.ack(second.token()).unwrap();
}
#[test]
fn stale_nack_after_release_and_missing_nack_return_lease_expired() {
let (_dir, _db, queue) = open_queue(
QueueConfig::balanced()
.lease_timeout(Duration::from_millis(1))
.build()
.unwrap(),
);
queue.push(Job(7), None).unwrap();
let first = queue.lease(1, None).unwrap().pop().unwrap();
thread::sleep(Duration::from_millis(3));
let second = queue.lease(1, None).unwrap().pop().unwrap();
assert!(matches!(
queue.nack(first.token()),
Err(DbError::LeaseExpired)
));
queue.ack(second.token()).unwrap();
assert!(matches!(
queue.nack(second.token()),
Err(DbError::LeaseExpired)
));
}
#[test]
fn nack_invalidates_the_token_before_the_retry_is_leased_again() {
let (_dir, _db, queue) = open_queue(
QueueConfig::balanced()
.max_attempts(2)
.backoff_secs(&[0])
.build()
.unwrap(),
);
queue.push(Job(7), None).unwrap();
let leased = queue.lease(1, None).unwrap().pop().unwrap();
queue.nack(leased.token()).unwrap();
assert!(matches!(
queue.ack(leased.token()),
Err(DbError::LeaseExpired)
));
}
#[test]
fn sub_millisecond_lease_timeouts_reject_without_claiming() {
assert!(
QueueConfig::balanced()
.lease_timeout(Duration::from_nanos(1))
.build()
.is_err()
);
let (_dir, _db, queue) = open_queue(QueueConfig::balanced().build().unwrap());
queue.push(Job(7), None).unwrap();
assert!(queue.lease(1, Some(Duration::from_nanos(1))).is_err());
assert_eq!(queue.len(), 1);
assert_eq!(queue.lease(1, None).unwrap().len(), 1);
}
#[test]
fn guard_failure_rolls_back_effect_and_acknowledgement_together() {
let (_dir, db, queue) = open_queue(
QueueConfig::balanced()
.lease_timeout(Duration::from_millis(1))
.build()
.unwrap(),
);
let effects = db
.open_zero_tree_fixed::<Effect, 8, _>(FixedConfig::test(), NoHook, &[])
.unwrap();
queue.push(Job(7), None).unwrap();
let leased = queue.lease(1, None).unwrap().pop().unwrap();
let result = db
.tx()
.put(&effects, &1_u64.to_be_bytes(), &Effect(1))
.unwrap()
.ack(&queue, leased.token())
.unwrap()
.guard(&effects, &2_u64.to_be_bytes(), |_| {
Err(DbError::Client("abort queue tx"))
})
.unwrap()
.commit();
assert!(matches!(result, Err(DbError::Client("abort queue tx"))));
assert_eq!(effects.get(&1_u64.to_be_bytes()), None);
assert_eq!(queue.len(), 1);
thread::sleep(Duration::from_millis(3));
assert_eq!(queue.lease(1, None).unwrap().len(), 1);
}
#[test]
fn transactional_push_publishes_one_row_and_count_only_after_commit() {
let (_dir, db, queue) = open_queue(QueueConfig::balanced().build().unwrap());
let tx = db.tx().push(&queue, Job(9), None).unwrap();
assert_eq!(queue.len(), 0);
tx.commit().unwrap();
assert_eq!(queue.len(), 1);
assert_eq!(
*queue.lease(1, None).unwrap().pop().unwrap().payload(),
Job(9)
);
}
#[test]
fn transaction_pushes_respect_max_depth_before_commit() {
let (_dir, db, queue) = open_queue(QueueConfig::balanced().max_depth(Some(1)).build().unwrap());
let result = db
.tx()
.push(&queue, Job(1), None)
.unwrap()
.push(&queue, Job(2), None);
assert!(matches!(result, Err(DbError::QueueFull)));
assert_eq!(queue.len(), 0);
}
#[test]
fn concurrent_leases_claim_each_ready_row_at_most_once() {
let (_dir, _db, queue) = open_queue(QueueConfig::balanced().build().unwrap());
for id in 0..64 {
queue.push(Job(id), None).unwrap();
}
let left = Arc::clone(&queue);
let right = Arc::clone(&queue);
let a = thread::spawn(move || left.lease(64, None).unwrap());
let b = thread::spawn(move || right.lease(64, None).unwrap());
let ids_a: HashSet<_> = a.join().unwrap().into_iter().map(|row| row.id()).collect();
let ids_b: HashSet<_> = b.join().unwrap().into_iter().map(|row| row.id()).collect();
assert!(ids_a.is_disjoint(&ids_b));
assert_eq!(ids_a.len() + ids_b.len(), 64);
}
#[test]
fn multiple_consumers_ack_every_job_exactly_once() {
const CONSUMERS: usize = 4;
const JOBS: u64 = 64;
let (_dir, _db, queue) = open_queue(QueueConfig::balanced().build().unwrap());
for id in 0..JOBS {
queue.push(Job(id), None).unwrap();
}
let barrier = Arc::new(Barrier::new(CONSUMERS));
let consumers: Vec<_> = (0..CONSUMERS)
.map(|_| {
let queue = Arc::clone(&queue);
let barrier = Arc::clone(&barrier);
thread::spawn(move || {
barrier.wait();
let mut consumed = Vec::new();
while let Some(entry) = queue.lease(1, None).unwrap().pop() {
consumed.push(entry.payload().0);
queue.ack(entry.token()).unwrap();
}
consumed
})
})
.collect();
let consumed: Vec<_> = consumers
.into_iter()
.flat_map(|consumer| consumer.join().unwrap())
.collect();
let unique: HashSet<_> = consumed.iter().copied().collect();
assert_eq!(consumed.len(), JOBS as usize, "every job is acknowledged");
assert_eq!(unique.len(), JOBS as usize, "no job is delivered twice");
assert_eq!(unique, (0..JOBS).collect());
assert_eq!(queue.len(), 0);
}
#[test]
fn stress_eight_consumers_ack_1024_jobs_exactly_once() {
const CONSUMERS: usize = 8;
const JOBS: u64 = 1_024;
let (_dir, _db, queue) = open_queue(QueueConfig::balanced().build().unwrap());
for id in 0..JOBS {
queue.push(Job(id), None).unwrap();
}
let barrier = Arc::new(Barrier::new(CONSUMERS));
let consumers: Vec<_> = (0..CONSUMERS)
.map(|_| {
let queue = Arc::clone(&queue);
let barrier = Arc::clone(&barrier);
thread::spawn(move || {
barrier.wait();
let mut consumed = Vec::new();
while let Some(entry) = queue.lease(1, None).unwrap().pop() {
consumed.push(entry.payload().0);
queue.ack(entry.token()).unwrap();
}
consumed
})
})
.collect();
let consumed: Vec<_> = consumers
.into_iter()
.flat_map(|consumer| consumer.join().unwrap())
.collect();
let unique: HashSet<_> = consumed.iter().copied().collect();
assert_eq!(consumed.len(), JOBS as usize, "every job is acknowledged");
assert_eq!(unique.len(), JOBS as usize, "no job is delivered twice");
assert_eq!(unique, (0..JOBS).collect());
assert_eq!(queue.len(), 0);
}
#[test]
fn terminal_nack_moves_to_dead_without_another_backoff() {
let (_dir, _db, queue) = open_queue(
QueueConfig::balanced()
.max_attempts(1)
.backoff_secs(&[60])
.build()
.unwrap(),
);
queue.push(Job(7), None).unwrap();
let leased = queue.lease(1, None).unwrap().pop().unwrap();
queue.nack(leased.token()).unwrap();
assert_eq!(queue.len(), 0);
assert_eq!(queue.dead_len(), 1);
assert!(queue.lease(1, None).unwrap().is_empty());
}
#[test]
fn overflowing_retry_backoff_leaves_the_leased_row_and_counts_unchanged() {
let (_dir, _db, queue) = open_queue(
QueueConfig::balanced()
.max_attempts(2)
.backoff_secs(&[u64::MAX])
.build()
.unwrap(),
);
queue.push(Job(7), None).unwrap();
let leased = queue.lease(1, None).unwrap().pop().unwrap();
assert!(queue.nack(leased.token()).is_err());
assert_eq!(queue.len(), 1);
assert_eq!(queue.dead_len(), 0);
assert!(queue.lease(1, None).unwrap().is_empty());
queue.ack(leased.token()).unwrap();
}
#[test]
fn queue_keeps_root_lease_after_db_is_dropped() {
let dir = tempfile::tempdir().unwrap();
let db = Db::open_test(dir.path()).unwrap();
let queue = db
.open_queue::<Jobs>(QueueConfig::balanced().build().unwrap())
.unwrap();
drop(db);
assert!(matches!(
Db::open_test(dir.path()),
Err(DbError::Locked { .. })
));
drop(queue);
Db::open_test(dir.path()).unwrap();
}