use reifydb_codec::row::pod::EncodedPodRow;
use reifydb_core::{
interface::{
catalog::{
id::QueueId,
queue::{QueueItemState, decode_queue_item_state},
},
store::SingleVersionRange,
},
key::queue::QueueItemStateKey,
};
use reifydb_test_harness::engine::TestEngine;
use reifydb_transaction::transaction::Transaction;
use reifydb_value::value::{Value, datetime::DateTime, frame::frame::Frame};
const ONE_PARTITION: &str = "CREATE QUEUE test::jobs { id: int4 } WITH { fifo: { partitions: 1 } }";
fn engine_with_queue(declaration: &str) -> TestEngine {
let t = TestEngine::new();
t.admin("CREATE NAMESPACE test");
t.admin(declaration);
t
}
fn queue_id(t: &TestEngine, name: &str) -> QueueId {
let catalog = t.inner().catalog();
let mut query_txn = t.inner().begin_query(TestEngine::identity()).unwrap();
let mut txn = Transaction::Query(&mut query_txn);
let namespace = catalog.find_namespace_by_name(&mut txn, "test").unwrap().unwrap();
catalog.find_queue_by_name(&mut txn, namespace.id(), name).unwrap().unwrap().id
}
fn state_of(t: &TestEngine, queue: QueueId) -> QueueItemState {
let store = t.inner().single().read_store();
SingleVersionRange::range_batch(&store, QueueItemStateKey::queue_scan(queue).encode(), 1024)
.unwrap()
.items
.iter()
.map(|item| decode_queue_item_state(EncodedPodRow::view(&item.bytes)).unwrap())
.next()
.expect("the queue must hold exactly one item")
}
fn claim_one(t: &TestEngine, worker: &str) -> String {
let frames = t.command(&format!(r#"CALL queue::claim("{worker}", "test::jobs", 1, duration::seconds(30))"#));
match frames[0].columns.iter().find(|c| c.name == "token").unwrap().data.get_value(0) {
Value::Utf8(t) => t,
other => panic!("token must be Utf8, got {other:?}"),
}
}
fn deadline_of(frames: &[Frame]) -> DateTime {
match frames[0].columns.iter().find(|c| c.name == "deadline").unwrap().data.get_value(0) {
Value::DateTime(d) => d,
other => panic!("deadline must be DateTime, got {other:?}"),
}
}
#[test]
fn test_extend_moves_the_deadline_of_a_live_lease() {
let t = engine_with_queue(ONE_PARTITION);
t.command("INSERT test::jobs [{ id: 1 }]");
let queue = queue_id(&t, "jobs");
t.mock_clock().set_nanos(1_000);
let token = claim_one(&t, "w1");
assert_eq!(state_of(&t, queue).lease_deadline, Some(DateTime::from_nanos(1_000 + 30 * 1_000_000_000)));
t.mock_clock().set_nanos(20 * 1_000_000_000);
let frames = t.command(&format!(r#"CALL queue::extend("{token}", duration::seconds(60))"#));
let expected = DateTime::from_nanos(80 * 1_000_000_000);
assert_eq!(deadline_of(&frames), expected);
assert_eq!(state_of(&t, queue).lease_deadline, Some(expected), "the durable record must carry the extension");
}
#[test]
fn test_extend_never_shortens_a_deadline() {
let t = engine_with_queue(ONE_PARTITION);
t.command("INSERT test::jobs [{ id: 1 }]");
let queue = queue_id(&t, "jobs");
t.mock_clock().set_nanos(0);
let token = claim_one(&t, "w1");
let original = state_of(&t, queue).lease_deadline.unwrap();
let frames = t.command(&format!(r#"CALL queue::extend("{token}", duration::seconds(5))"#));
assert_eq!(deadline_of(&frames), original, "a shorter request must return the existing deadline");
assert_eq!(state_of(&t, queue).lease_deadline, Some(original));
}
#[test]
fn test_extend_fails_hard_once_the_item_has_been_acked() {
let t = engine_with_queue(ONE_PARTITION);
t.command("INSERT test::jobs [{ id: 1 }]");
let token = claim_one(&t, "w1");
t.command(&format!(r#"CALL queue::ack("{token}")"#));
let err = t.command_err(&format!(r#"CALL queue::extend("{token}", duration::seconds(60))"#));
assert!(err.contains("QUEUE_002"), "{err}");
assert!(err.contains("not leased"), "{err}");
}
#[test]
fn test_extend_fails_hard_for_an_attempt_that_is_no_longer_current() {
let t = engine_with_queue(ONE_PARTITION);
t.command("INSERT test::jobs [{ id: 1 }]");
let token = claim_one(&t, "w1");
let superseded = token.replace(":1:w1", ":2:w1");
let err = t.command_err(&format!(r#"CALL queue::extend("{superseded}", duration::seconds(60))"#));
assert!(err.contains("QUEUE_002"), "{err}");
assert!(err.contains("later attempt"), "{err}");
}
#[test]
fn test_extend_fails_hard_for_an_item_with_no_scheduling_state() {
let t = engine_with_queue(ONE_PARTITION);
t.command("INSERT test::jobs [{ id: 1 }]");
let queue = queue_id(&t, "jobs");
let err =
t.command_err(&format!(r#"CALL queue::extend("qt1:{}:0:9999:1:w1", duration::seconds(60))"#, queue.0));
assert!(err.contains("QUEUE_002"), "{err}");
}
#[test]
fn test_extend_rejects_a_malformed_token() {
let t = engine_with_queue(ONE_PARTITION);
let err = t.command_err(r#"CALL queue::extend("nonsense", duration::seconds(60))"#);
assert!(err.contains("QUEUE_003"), "{err}");
}
#[test]
fn test_extend_is_rejected_outside_a_command_transaction() {
let t = engine_with_queue(ONE_PARTITION);
let err = t.query_err(r#"CALL queue::extend("qt1:1:0:1:1:w1", duration::seconds(60))"#);
assert!(err.contains("must run in a command transaction"), "{err}");
}