use std::marker::PhantomData;
use std::ops::Range;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime};
use armour_core::{GetRefs, GetType, Typ};
use rapira::Rapira;
use crate::armour::{CollectionMeta, Db};
use crate::durability::Fixed;
use crate::sync::{Mutex, lock};
use crate::{
DbError, DbResult, LeaseToken, LeasedEntry, NoHook, Queue, QueueConfig, QueueKey, QueueSpec,
RapiraCodec, TtlHook, TtlPolicy, VarTypedTree, ZeroTree,
};
use super::fixed::QueueHeader;
use super::retention::{HeaderRetention, QueueRetentionPass, VarHeaderHook, VarRetention};
use super::wait::{QueueHeaderRow, QueueScan, QueueWait};
#[repr(transparent)]
#[derive(zerocopy::FromBytes, zerocopy::IntoBytes, zerocopy::Immutable, zerocopy::KnownLayout)]
pub(crate) struct QueueVarHeaderRow<N: QueueSpec> {
header: QueueHeader,
_queue: PhantomData<fn() -> N>,
}
impl<N: QueueSpec> Copy for QueueVarHeaderRow<N> {}
impl<N: QueueSpec> Clone for QueueVarHeaderRow<N> {
fn clone(&self) -> Self {
*self
}
}
impl<N: QueueSpec> GetType for QueueVarHeaderRow<N> {
const TYPE: Typ = Typ::Custom("QueueVarHeaderRow", &[]);
}
impl<N: QueueSpec> GetRefs for QueueVarHeaderRow<N> {}
impl<N: QueueSpec> CollectionMeta for QueueVarHeaderRow<N> {
type SelfId = QueueKey;
const NAME: &'static str = N::NAME;
const TTL: Option<TtlPolicy<Self>> = Some(TtlPolicy {
revision: 1,
deadline: |header, key| {
(key.class() == crate::DEAD_CLASS && header.header().dead_at_ms != 0)
.then_some(header.header().dead_at_ms)
},
});
}
impl<N: QueueSpec> QueueVarHeaderRow<N> {
pub(crate) fn new(enqueued_at_ms: u64) -> Self {
Self {
header: QueueHeader {
lease_deadline_ms: 0,
dead_at_ms: 0,
enqueued_at_ms,
lease_epoch: 0,
attempts: 0,
flags: 0,
},
_queue: PhantomData,
}
}
pub(crate) fn header(&self) -> QueueHeader {
self.header
}
pub(crate) fn with_header(mut self, header: QueueHeader) -> Self {
self.header = header;
self
}
}
impl<N: QueueSpec> super::wait::QueueHeaderRow for QueueVarHeaderRow<N> {
fn lease_deadline_ms(&self) -> u64 {
self.header().lease_deadline_ms
}
}
pub(crate) type QueueVarHeaderTree<N> = ZeroTree<
QueueKey,
32,
QueueVarHeaderRow<N>,
TtlHook<QueueVarHeaderRow<N>, VarHeaderHook<N>>,
Fixed,
>;
pub(crate) struct QueueVarPayload<N: QueueSpec> {
pub(crate) bytes: Vec<u8>,
pub(crate) _queue: PhantomData<fn() -> N>,
}
impl<N: QueueSpec> Clone for QueueVarPayload<N> {
fn clone(&self) -> Self {
Self {
bytes: self.bytes.clone(),
_queue: PhantomData,
}
}
}
impl<N: QueueSpec> Rapira for QueueVarPayload<N> {
const STATIC_SIZE: Option<usize> = <Vec<u8> as Rapira>::STATIC_SIZE;
const MIN_SIZE: usize = <Vec<u8> as Rapira>::MIN_SIZE;
fn size(&self) -> usize {
self.bytes.size()
}
fn check_bytes(slice: &mut &[u8]) -> rapira::Result<()> {
<Vec<u8> as Rapira>::check_bytes(slice)
}
fn from_slice(slice: &mut &[u8]) -> rapira::Result<Self> {
Ok(Self {
bytes: Vec::<u8>::from_slice(slice)?,
_queue: PhantomData,
})
}
fn convert_to_bytes(&self, slice: &mut [u8], cursor: &mut usize) {
self.bytes.convert_to_bytes(slice, cursor);
}
}
impl<N> GetType for QueueVarPayload<N>
where
N: QueueSpec,
N::Payload: GetType,
{
const TYPE: Typ = Typ::Custom("QueueVarPayload", &[N::Payload::TYPE]);
}
impl<N> GetRefs for QueueVarPayload<N>
where
N: QueueSpec,
N::Payload: GetType,
{
}
impl<N> CollectionMeta for QueueVarPayload<N>
where
N: QueueSpec,
N::Payload: GetType,
{
type SelfId = rend::u64_be;
const NAME: &'static str = N::PAYLOAD_NAME;
}
pub(crate) type QueuePayloadTree<N> = VarTypedTree<
rend::u64_be,
QueueVarPayload<N>,
RapiraCodec,
TtlHook<QueueVarPayload<N>, NoHook>,
>;
#[doc(hidden)]
pub struct VarQueueInner<N>
where
N: QueueSpec,
N::Payload: GetType,
{
pub(crate) headers: Arc<QueueVarHeaderTree<N>>,
pub(crate) payloads: Arc<QueuePayloadTree<N>>,
pub(crate) config: QueueConfig,
pub(crate) seq: Arc<crate::armour::seq::SeqGen>,
pub(crate) ids: Mutex<Range<u64>>,
pub(crate) live_count: Arc<AtomicU64>,
pub(crate) dead_count: Arc<AtomicU64>,
pub(crate) retention: Arc<VarRetention<N>>,
pub(crate) runtime: Arc<crate::txn::TxRuntime>,
pub(crate) wait: Arc<QueueWait>,
pub(crate) payload_writes: Arc<AtomicU64>,
pub(crate) instance_id: u64,
}
impl<N> VarQueueInner<N>
where
N: QueueSpec,
N::Payload: GetType,
{
pub(crate) fn seed_counts(
headers: &Arc<QueueVarHeaderTree<N>>,
retention: &Arc<VarRetention<N>>,
) -> DbResult<u64> {
let mut live_count = 0;
let mut dead_count = 0;
for (key, _) in headers.iter() {
if key.class() == crate::LIVE_CLASS {
live_count += 1;
} else {
dead_count += 1;
}
}
retention.header.set_dead_count(dead_count);
Ok(live_count)
}
fn open(
db: &Db,
config: QueueConfig,
headers: Arc<QueueVarHeaderTree<N>>,
payloads: Arc<QueuePayloadTree<N>>,
retention: Arc<VarRetention<N>>,
live_count: u64,
) -> DbResult<Self> {
let recheck = config.wakeup_recheck;
Ok(Self {
headers,
payloads,
config,
seq: Arc::clone(&db.seq),
ids: Mutex::new(0..0),
live_count: Arc::new(AtomicU64::new(live_count)),
dead_count: retention.header.dead_count(),
retention,
runtime: Arc::clone(&db.tx_runtime),
wait: QueueWait::new(db.queue_wake_service(), recheck),
payload_writes: Arc::new(AtomicU64::new(0)),
instance_id: super::next_queue_instance_id(),
})
}
pub(crate) fn next_id(&self) -> DbResult<u64> {
let mut ids = lock(&self.ids);
if ids.is_empty() {
*ids = self.seq.reserve_block(N::NAME, crate::ID_BLOCK)?;
}
ids.next()
.ok_or(DbError::Internal("queue id block was empty"))
}
pub(crate) fn next_id_in_tx(&self) -> DbResult<u64> {
let mut ids = lock(&self.ids);
if ids.is_empty() {
*ids = self.seq.reserve_block_in_tx(N::NAME, crate::ID_BLOCK)?;
}
ids.next()
.ok_or(DbError::Internal("queue id block was empty"))
}
}
impl<N> super::QueueOpen<N> for crate::VarQueueMode
where
N: QueueSpec<Mode = crate::VarQueueMode>,
N::Payload: Rapira + GetType + Send + Sync,
{
fn open(db: &Db, config: QueueConfig) -> DbResult<Arc<Queue<N>>> {
let fixed_config = config.fixed_config();
let payload_config = config.var_payload_config();
let header_retention = HeaderRetention::new();
let retention = VarRetention::new(Arc::clone(&header_retention));
let mut live_count = 0;
let (headers, payloads) = db.open_managed_pair(
|meta| {
QueueVarHeaderTree::<N>::open_with_hook_inner(
db.tree_path(meta.name),
fixed_config,
TtlHook::new(VarHeaderHook::new(Arc::clone(&retention))),
)
},
crate::access::ZeroWire::<32>,
&[],
|meta| {
QueuePayloadTree::<N>::open_hooked_inner(
db.tree_path(meta.name),
payload_config,
RapiraCodec,
TtlHook::new(crate::NoHook),
)
},
crate::access::CodecWire(RapiraCodec),
&[],
|headers, payloads| {
VarRetention::orphan_pass_pair(headers, payloads)?;
live_count = VarQueueInner::<N>::seed_counts(headers, &retention)?;
Ok(())
},
crate::armour::db::ManagedCollectionAccess::Internal,
crate::armour::db::ManagedCollectionAccess::Internal,
)?;
retention.bind(&headers, &payloads);
let queue_retention: Arc<dyn QueueRetentionPass> = retention.clone();
let inner = VarQueueInner::open(db, config, headers, payloads, retention, live_count)?;
let queue = Arc::new(Queue { inner });
db.register_queue_retention(queue_retention);
Ok(queue)
}
}
impl<N> Queue<N>
where
N: QueueSpec<Mode = crate::VarQueueMode>,
N::Payload: Rapira + GetType + Send + Sync,
{
fn var_push(&self, payload: N::Payload, not_before: Option<SystemTime>) -> DbResult<u64> {
self.inner.retention.process_before_queue_operation();
let _producer = self.inner.retention.producer_permit();
if self
.inner
.config
.max_depth
.is_some_and(|limit| self.inner.live_count.load(Ordering::Acquire) >= limit)
{
return Err(DbError::QueueFull);
}
let now = SystemTime::now();
let visible_at_ms = super::fixed::FixedQueueInner::<N>::now_ms(not_before.unwrap_or(now))?;
if visible_at_ms >= (1 << 63) {
return Err(DbError::Client("queue visibility exceeds 63-bit domain"));
}
let id = self.inner.next_id()?;
let payload_key = rend::u64_be::from_native(id);
self.inner.payloads.insert(
&payload_key,
&QueueVarPayload {
bytes: rapira::serialize(&payload),
_queue: PhantomData,
},
)?;
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
self.inner.retention.wait_after_payload();
self.inner.payload_writes.fetch_add(1, Ordering::Release);
let row = QueueVarHeaderRow::<N>::new(super::fixed::FixedQueueInner::<N>::now_ms(now)?);
self.inner
.headers
.insert(&QueueKey::live(visible_at_ms, id), &row)?;
self.inner.live_count.fetch_add(1, Ordering::Release);
self.inner.wait.notify_visible_at(visible_at_ms);
Ok(id)
}
fn var_len(&self) -> u64 {
self.inner.live_count.load(Ordering::Acquire)
}
fn var_dead_len(&self) -> u64 {
self.inner.dead_count.load(Ordering::Acquire)
}
fn var_lease(
&self,
n: usize,
timeout: Option<Duration>,
) -> DbResult<Vec<LeasedEntry<N::Payload>>> {
self.inner.retention.process_before_queue_operation();
let now_ms = super::fixed::FixedQueueInner::<N>::now_ms(SystemTime::now())?;
let timeout_ms = u64::try_from(
timeout
.unwrap_or(self.inner.config.lease_timeout)
.as_millis(),
)
.map_err(|_| DbError::Client("queue lease timeout exceeds u64 milliseconds"))?;
if timeout_ms == 0 {
return Err(DbError::Client(
"queue lease timeout must be at least one millisecond",
));
}
let lease_deadline_ms = now_ms.checked_add(timeout_ms).ok_or(DbError::Client(
"queue lease deadline exceeds u64 milliseconds",
))?;
let mut leased = Vec::new();
for (key, candidate) in self.inner.headers.iter() {
if leased.len() == n || key.class() != crate::LIVE_CLASS || key.visible_at_ms() > now_ms
{
break;
}
if candidate.lease_deadline_ms() > now_ms {
continue;
}
let claimed = match self.inner.headers.atomic(&key, |shard| {
let Some(row) = shard.try_get(&key)? else {
return Ok(None);
};
let mut header = row.header();
if header.lease_deadline_ms > now_ms {
return Ok(None);
}
header.lease_epoch = super::next_lease_epoch(header.lease_epoch);
header.lease_deadline_ms = lease_deadline_ms;
let claimed = row.with_header(header);
shard.put(&key, &claimed)?;
Ok(Some(claimed))
}) {
Ok(claimed) => claimed,
Err(error) => return super::partial_lease_batch(leased, N::NAME, key, error),
};
if let Some(row) = claimed {
let header = row.header();
let payload = match self
.inner
.payloads
.get_or_err(&rend::u64_be::from_native(key.id()))
.and_then(|payload| {
rapira::deserialize(&payload.bytes)
.map_err(|_| DbError::CorruptedEntry { offset: 0 })
}) {
Ok(payload) => payload,
Err(error) => return super::partial_lease_batch(leased, N::NAME, key, error),
};
leased.push(LeasedEntry {
key,
id: key.id(),
attempts: header.attempts,
payload,
token: LeaseToken::new(self.inner.instance_id, key, header.lease_epoch),
});
}
}
Ok(leased)
}
fn var_ack(&self, token: LeaseToken) -> DbResult<()> {
crate::txn::Tx::from_runtime(Arc::clone(&self.inner.runtime))
.var_ack_queue(self, token)?
.commit()
}
fn var_nack(&self, token: LeaseToken) -> DbResult<()> {
crate::txn::Tx::from_runtime(Arc::clone(&self.inner.runtime))
.var_nack_queue(self, token)?
.commit()
}
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub fn debug_payload_write_count(&self) -> u64 {
self.inner.payload_writes.load(Ordering::Acquire)
}
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub fn debug_payload_len(&self) -> usize {
self.inner.payloads.len()
}
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub fn debug_header_shard(&self, key: &QueueKey) -> usize {
self.inner.headers.shard_for(key)
}
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub fn debug_fail_next_payload_delete(&self) {
self.inner.retention.fail_next_payload_delete();
}
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub fn debug_fail_payload_deletes(&self, count: usize) {
self.inner.retention.fail_payload_deletes(count);
}
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub fn debug_pause_after_payload(
&self,
) -> (
std::sync::Arc<std::sync::Barrier>,
std::sync::Arc<std::sync::Barrier>,
) {
self.inner.retention.pause_after_payload()
}
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub fn debug_force_full_orphan_pass(&self) -> DbResult<()> {
self.inner.retention.force_full_orphan_pass()
}
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub fn debug_mark_full_orphan_pass_pending(&self) {
self.inner.retention.mark_full_orphan_pass_pending();
}
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub fn debug_notify_maintenance_wait(&self) -> std::sync::mpsc::Receiver<()> {
self.inner.retention.notify_maintenance_wait()
}
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub fn debug_insert_orphan_payload(&self, id: u64, payload: N::Payload) -> DbResult<()> {
self.inner.payloads.insert(
&rend::u64_be::from_native(id),
&QueueVarPayload {
bytes: rapira::serialize(&payload),
_queue: PhantomData,
},
)
}
}
impl<N> crate::queue::QueueHandle<N> for crate::VarQueueMode
where
N: QueueSpec<Mode = crate::VarQueueMode>,
N::Payload: Rapira + GetType + Send + Sync,
{
fn push(
queue: &Queue<N>,
payload: N::Payload,
not_before: Option<SystemTime>,
) -> DbResult<u64> {
queue.var_push(payload, not_before)
}
fn len(queue: &Queue<N>) -> u64 {
queue.var_len()
}
fn dead_len(queue: &Queue<N>) -> u64 {
queue.var_dead_len()
}
fn lease(
queue: &Queue<N>,
n: usize,
timeout: Option<Duration>,
) -> DbResult<Vec<LeasedEntry<N::Payload>>> {
queue.var_lease(n, timeout)
}
fn ack(queue: &Queue<N>, token: LeaseToken) -> DbResult<()> {
queue.var_ack(token)
}
fn nack(queue: &Queue<N>, token: LeaseToken) -> DbResult<()> {
queue.var_nack(token)
}
fn scan_due(queue: &Queue<N>) -> DbResult<QueueScan> {
QueueScan::rows(
queue.inner.headers.iter(),
queue.inner.config.wakeup_recheck,
)
}
fn wait(queue: &Queue<N>) -> &QueueWait {
&queue.inner.wait
}
fn stats(queue: &Queue<N>) -> crate::QueueStats {
crate::QueueStats {
compaction_bytes: 0,
payload_writes: queue.inner.payload_writes.load(Ordering::Acquire),
}
}
}
#[cfg(test)]
mod post_log_tests {
use super::*;
use crate::queue;
use armour_core::GetType;
use rapira::Rapira;
#[derive(Clone, GetType, Rapira)]
pub struct Payload(String);
queue!(
PinnedJobs,
payload = Payload,
name = "queue_var_post_log_pinned",
var
);
#[test]
fn post_log_failure_pins_var_producer_until_runtime_drops() {
let dir = tempfile::tempdir().unwrap();
let db = crate::armour::Db::open_test(dir.path()).unwrap();
let queue = db
.open_queue::<PinnedJobs>(QueueConfig::balanced().build().unwrap())
.unwrap();
let retention = Arc::clone(&queue.inner.retention);
let fail = crate::txn::failpoints::ApplyFailGuard::at(0);
let error = db
.tx()
.push(&queue, Payload("body".into()), None)
.unwrap()
.commit()
.unwrap_err();
drop(fail);
assert!(matches!(error, DbError::TxFailedAfterLog { .. }));
assert_eq!(retention.producer_count(), 1);
drop(queue);
drop(db);
assert_eq!(retention.producer_count(), 0);
}
}