armdb 0.14.2

use std::marker::PhantomData;
use std::mem::size_of;
use std::ops::Range;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use armour_core::{GetRefs, GetType, Typ};

use crate::armour::{CollectionMeta, Db};
use crate::durability::Fixed;
use crate::sync::{Mutex, lock};
use crate::{
    DbError, DbResult, FixedConfig, LeaseToken, LeasedEntry, Queue, QueueConfig, QueueKey,
    QueueSpec, TtlHook, TtlPolicy, ZeroTree,
};

use super::retention::{FixedHeaderHook, HeaderRetention};
use super::wait::{QueueHeaderRow, QueueScan, QueueWait};

const FIXED_PAYLOAD_BYTES: usize = 4096;
const FIXED_ROW_BYTES: usize = 32 + FIXED_PAYLOAD_BYTES;

/// Debug-only synchronization point for deterministic queue admission tests.
///
/// It is absent from release builds and is never configured by production
/// code. The barrier runs immediately after the lock-free depth check and
/// before the row write, making the documented soft-cap race observable.
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
#[doc(hidden)]
pub mod test_support {
    use std::sync::{Arc, Barrier, LazyLock, Mutex};

    static PUSH_ADMISSION_BARRIER: LazyLock<Mutex<Option<Arc<Barrier>>>> =
        LazyLock::new(|| Mutex::new(None));

    pub fn set_fixed_queue_push_admission_barrier(barrier: Option<Arc<Barrier>>) {
        *PUSH_ADMISSION_BARRIER
            .lock()
            .expect("queue test barrier lock") = barrier;
    }

    pub(crate) fn wait_after_push_admission() {
        let barrier = PUSH_ADMISSION_BARRIER
            .lock()
            .expect("queue test barrier lock")
            .clone();
        if let Some(barrier) = barrier {
            barrier.wait();
        }
    }
}

/// Fixed-size metadata stored beside every queue payload.
#[repr(C)]
#[derive(
    Clone,
    Copy,
    zerocopy::FromBytes,
    zerocopy::IntoBytes,
    zerocopy::Immutable,
    zerocopy::KnownLayout,
)]
pub(crate) struct QueueHeader {
    pub(crate) lease_deadline_ms: u64,
    pub(crate) dead_at_ms: u64,
    pub(crate) enqueued_at_ms: u64,
    pub(crate) lease_epoch: u32,
    pub(crate) attempts: u16,
    pub(crate) flags: u16,
}

/// FixedStore representation of one fixed-payload queue row.
///
/// The payload's exact type is selected by `N`, while fixed backing storage
/// keeps every FixedStore slot the same size.
#[repr(C)]
#[derive(
    Clone,
    Copy,
    zerocopy::FromBytes,
    zerocopy::IntoBytes,
    zerocopy::Immutable,
    zerocopy::KnownLayout,
)]
struct QueueFixedStorage {
    header: QueueHeader,
    payload: [u8; FIXED_PAYLOAD_BYTES],
}

#[repr(transparent)]
#[derive(zerocopy::FromBytes, zerocopy::IntoBytes, zerocopy::Immutable, zerocopy::KnownLayout)]
pub(crate) struct QueueFixedRow<N: QueueSpec> {
    storage: QueueFixedStorage,
    _queue: PhantomData<fn() -> N>,
}

impl<N: QueueSpec> Copy for QueueFixedRow<N> {}

impl<N: QueueSpec> Clone for QueueFixedRow<N> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<N> GetType for QueueFixedRow<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
    const TYPE: Typ = Typ::Custom("QueueFixedRow", &[N::Payload::TYPE]);
}

impl<N> GetRefs for QueueFixedRow<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
}

impl<N> CollectionMeta for QueueFixedRow<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
    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> QueueFixedRow<N> {
    pub(crate) fn new<P>(payload: P, enqueued_at_ms: u64) -> DbResult<Self>
    where
        P: zerocopy::IntoBytes + zerocopy::Immutable,
    {
        let payload = payload.as_bytes();
        if payload.len() > FIXED_PAYLOAD_BYTES {
            return Err(DbError::Client("queue fixed payload exceeds 4096 bytes"));
        }
        let mut storage = QueueFixedStorage {
            header: QueueHeader {
                lease_deadline_ms: 0,
                dead_at_ms: 0,
                enqueued_at_ms,
                lease_epoch: 0,
                attempts: 0,
                flags: 0,
            },
            payload: [0; FIXED_PAYLOAD_BYTES],
        };
        storage.payload[..payload.len()].copy_from_slice(payload);
        Ok(Self {
            storage,
            _queue: PhantomData,
        })
    }

    pub(crate) fn header(&self) -> QueueHeader {
        self.storage.header
    }

    pub(crate) fn with_header(mut self, header: QueueHeader) -> Self {
        self.storage.header = header;
        self
    }

    fn payload(&self) -> DbResult<N::Payload>
    where
        N::Payload: zerocopy::FromBytes,
    {
        zerocopy::FromBytes::read_from_bytes(&self.storage.payload[..size_of::<N::Payload>()])
            .map_err(|_| DbError::Internal("queue fixed payload size mismatch"))
    }
}

impl<N: QueueSpec> super::wait::QueueHeaderRow for QueueFixedRow<N> {
    fn lease_deadline_ms(&self) -> u64 {
        self.header().lease_deadline_ms
    }
}

pub(crate) type QueueHeaderTree<N> = ZeroTree<
    QueueKey,
    FIXED_ROW_BYTES,
    QueueFixedRow<N>,
    TtlHook<QueueFixedRow<N>, FixedHeaderHook<N>>,
    Fixed,
>;

#[doc(hidden)]
pub struct FixedQueueInner<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
    pub(crate) headers: Arc<QueueHeaderTree<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) runtime: Arc<crate::txn::TxRuntime>,
    pub(crate) wait: Arc<QueueWait>,
    pub(crate) instance_id: u64,
}

impl<N> FixedQueueInner<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
    /// Count live and dead rows on disk.
    ///
    /// Must run **before** the header collection is published: publication both
    /// makes it visible to the TTL sweeper and arms `FixedHeaderHook`, so a
    /// sweep of an expired dead row could otherwise decrement a counter that is
    /// still zero (saturating to zero) and then have this seed overwrite it —
    /// leaving `dead_len` permanently over-reported for the session.
    pub(crate) fn seed_counts(
        headers: &Arc<QueueHeaderTree<N>>,
        retention: &Arc<HeaderRetention>,
    ) -> DbResult<u64> {
        let mut live_count = 0;
        let mut dead_count = 0;
        for (key, _) in headers.iter() {
            match key.class() {
                crate::LIVE_CLASS => live_count += 1,
                crate::DEAD_CLASS => dead_count += 1,
                _ => unreachable!("QueueKey only has two classes"),
            }
        }
        retention.set_dead_count(dead_count);
        Ok(live_count)
    }

    pub(crate) fn open(
        db: &Db,
        config: QueueConfig,
        headers: Arc<QueueHeaderTree<N>>,
        retention: Arc<HeaderRetention>,
        live_count: u64,
    ) -> DbResult<Self> {
        let recheck = config.wakeup_recheck;
        Ok(Self {
            headers,
            config,
            seq: Arc::clone(&db.seq),
            ids: Mutex::new(0..0),
            live_count: Arc::new(AtomicU64::new(live_count)),
            dead_count: retention.dead_count(),
            runtime: Arc::clone(&db.tx_runtime),
            wait: QueueWait::new(db.queue_wake_service(), recheck),
            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"))
    }

    pub(crate) fn now_ms(time: SystemTime) -> DbResult<u64> {
        let millis = time
            .duration_since(UNIX_EPOCH)
            .map_err(|_| DbError::Client("queue time is before unix epoch"))?
            .as_millis();
        u64::try_from(millis).map_err(|_| DbError::Client("queue time exceeds u64 milliseconds"))
    }
}

impl<N> Queue<N>
where
    N: QueueSpec<Mode = crate::FixedQueueMode>,
    N::Payload: Copy + zerocopy::IntoBytes + zerocopy::FromBytes + zerocopy::Immutable + GetType,
{
    /// Enqueue a fixed-size payload for immediate or delayed delivery.
    fn fixed_push(&self, payload: N::Payload, not_before: Option<SystemTime>) -> DbResult<u64> {
        let inner = &self.inner;
        if inner
            .config
            .max_depth
            .is_some_and(|limit| inner.live_count.load(Ordering::Acquire) >= limit)
        {
            return Err(DbError::QueueFull);
        }
        #[cfg(all(debug_assertions, feature = "queue-test-support"))]
        test_support::wait_after_push_admission();
        let now = SystemTime::now();
        let visible_at_ms = 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 = inner.next_id()?;
        let row = QueueFixedRow::<N>::new(payload, FixedQueueInner::<N>::now_ms(now)?)?;
        inner
            .headers
            .insert(&QueueKey::live(visible_at_ms, id), &row)?;
        inner.live_count.fetch_add(1, Ordering::Release);
        inner.wait.notify_visible_at(visible_at_ms);
        Ok(id)
    }

    /// Number of live rows, including currently leased rows.
    #[allow(clippy::len_without_is_empty)] // Queue intentionally has no empty-state API before leasing lands.
    fn fixed_len(&self) -> u64 {
        self.inner.live_count.load(Ordering::Acquire)
    }

    /// Number of retained dead-letter rows.
    fn fixed_dead_len(&self) -> u64 {
        self.inner.dead_count.load(Ordering::Acquire)
    }

    /// Claim up to `n` due rows. Claims are deliberately independent per row.
    fn fixed_lease(
        &self,
        n: usize,
        timeout: Option<Duration>,
    ) -> DbResult<Vec<LeasedEntry<N::Payload>>> {
        let now_ms = 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 {
                break;
            }
            if key.visible_at_ms() > now_ms {
                break;
            }
            // Stage 1 is lock-free: a live lease is skipped straight off the
            // iterator. The SeqLock makes the read consistent, and a stale
            // answer costs at most one shard lock we would have taken anyway —
            // the guarded claim below is what makes the decision correct.
            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 row.payload() {
                    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 fixed_renew(&self, token: LeaseToken, timeout: Option<Duration>) -> DbResult<()> {
        if !token.belongs_to(self.inner.instance_id) || token.key().class() != crate::LIVE_CLASS {
            return Err(DbError::LeaseExpired);
        }
        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 deadline = self.inner.headers.atomic(&token.key(), |shard| {
            let Some(row) = shard.try_get(&token.key())? else {
                return Err(DbError::LeaseExpired);
            };
            let mut header = row.header();
            // Sample under the shard lock: time spent waiting for another
            // writer cannot turn an expired lease into a successful renewal.
            let now_ms = FixedQueueInner::<N>::now_ms(SystemTime::now())?;
            if header.lease_epoch != token.epoch() || header.lease_deadline_ms <= now_ms {
                return Err(DbError::LeaseExpired);
            }
            let deadline = now_ms
                .checked_add(timeout_ms)
                .ok_or(DbError::Client(
                    "queue lease deadline exceeds u64 milliseconds",
                ))?
                .max(header.lease_deadline_ms);
            if deadline != header.lease_deadline_ms {
                header.lease_deadline_ms = deadline;
                shard.put(&token.key(), &row.with_header(header))?;
            }
            Ok(deadline)
        })?;
        // Waiters rescan the header deadline; the physical key stays fixed.
        self.inner.wait.notify_visible_at(deadline);
        Ok(())
    }

    /// Acknowledge a lease through the queue's private T5 runtime.
    fn fixed_ack(&self, token: LeaseToken) -> DbResult<()> {
        crate::txn::Tx::from_runtime(Arc::clone(&self.inner.runtime))
            .fixed_ack_queue(self, token)?
            .commit()
    }

    /// Retry or dead-letter a lease through the queue's private T5 runtime.
    fn fixed_nack(&self, token: LeaseToken) -> DbResult<()> {
        crate::txn::Tx::from_runtime(Arc::clone(&self.inner.runtime))
            .fixed_nack_queue(self, token)?
            .commit()
    }
}

impl<N> crate::queue::QueueHandle<N> for crate::FixedQueueMode
where
    N: QueueSpec<Mode = crate::FixedQueueMode>,
    N::Payload: Copy + zerocopy::IntoBytes + zerocopy::FromBytes + zerocopy::Immutable + GetType,
{
    fn push(
        queue: &Queue<N>,
        payload: N::Payload,
        not_before: Option<SystemTime>,
    ) -> DbResult<u64> {
        queue.fixed_push(payload, not_before)
    }
    fn len(queue: &Queue<N>) -> u64 {
        queue.fixed_len()
    }
    fn dead_len(queue: &Queue<N>) -> u64 {
        queue.fixed_dead_len()
    }
    fn lease(
        queue: &Queue<N>,
        n: usize,
        timeout: Option<Duration>,
    ) -> DbResult<Vec<LeasedEntry<N::Payload>>> {
        queue.fixed_lease(n, timeout)
    }
    fn renew(queue: &Queue<N>, token: LeaseToken, timeout: Option<Duration>) -> DbResult<()> {
        queue.fixed_renew(token, timeout)
    }
    fn ack(queue: &Queue<N>, token: LeaseToken) -> DbResult<()> {
        queue.fixed_ack(token)
    }
    fn nack(queue: &Queue<N>, token: LeaseToken) -> DbResult<()> {
        queue.fixed_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<N>) -> crate::QueueStats {
        crate::QueueStats::default()
    }
}

impl<N> crate::queue::QueueOpen<N> for crate::FixedQueueMode
where
    N: QueueSpec<Mode = crate::FixedQueueMode>,
    N::Payload: Copy + zerocopy::IntoBytes + zerocopy::FromBytes + zerocopy::Immutable + GetType,
{
    fn open(db: &Db, config: QueueConfig) -> DbResult<Arc<Queue<N>>> {
        open_fixed_queue(db, config)
    }
}

pub(crate) fn open_fixed_queue<N>(db: &Db, config: QueueConfig) -> DbResult<Arc<Queue<N>>>
where
    N: QueueSpec<Mode = crate::FixedQueueMode>,
    N::Payload: Copy + zerocopy::IntoBytes + zerocopy::FromBytes + zerocopy::Immutable + GetType,
{
    if size_of::<N::Payload>() > FIXED_PAYLOAD_BYTES {
        return Err(DbError::Client("queue fixed payload exceeds 4096 bytes"));
    }
    let fixed_config: FixedConfig = config.fixed_config();
    let retention = HeaderRetention::new();
    let mut live_count = 0;
    let headers = db.open_zero_tree_fixed_internal::<QueueFixedRow<N>, FIXED_ROW_BYTES, _>(
        fixed_config,
        FixedHeaderHook::new(Arc::clone(&retention)),
        &[],
        |headers| {
            live_count = FixedQueueInner::<N>::seed_counts(headers, &retention)?;
            Ok(())
        },
    )?;
    let inner = FixedQueueInner::open(db, config, headers, retention, live_count)?;
    Ok(Arc::new(Queue { inner }))
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;
    use std::sync::Arc;
    use std::thread;

    use crate::{QueueConfig, queue};

    use super::*;

    #[derive(
        Clone,
        Copy,
        GetType,
        zerocopy::FromBytes,
        zerocopy::IntoBytes,
        zerocopy::Immutable,
        zerocopy::KnownLayout,
    )]
    #[repr(transparent)]
    pub struct AllocatorJob(u64);

    queue!(
        AllocatorJobs,
        payload = AllocatorJob,
        name = "queue_fixed_allocator_jobs"
    );

    #[test]
    fn concurrent_queue_local_allocator_refills_without_reusing_ids_after_reopen() {
        let dir = tempfile::tempdir().expect("tempdir");
        let db = Db::open_test(dir.path()).expect("open database");
        let queue = db
            .open_queue::<AllocatorJobs>(QueueConfig::balanced().build().expect("config"))
            .expect("open queue");

        const THREADS: usize = 8;
        const IDS_PER_THREAD: usize = 300;
        let workers: Vec<_> = (0..THREADS)
            .map(|_| {
                let queue = Arc::clone(&queue);
                thread::spawn(move || {
                    (0..IDS_PER_THREAD)
                        .map(|_| queue.inner.next_id().expect("allocate id"))
                        .collect::<Vec<_>>()
                })
            })
            .collect();

        let ids: Vec<_> = workers
            .into_iter()
            .flat_map(|worker| worker.join().expect("join worker"))
            .collect();
        let unique: HashSet<_> = ids.iter().copied().collect();
        assert_eq!(unique.len(), ids.len(), "allocator reused an id");
        assert!(
            ids.len() > crate::ID_BLOCK as usize,
            "must refill an id block"
        );
        let greatest = *ids.iter().max().expect("allocated ids");

        drop(queue);
        drop(db);

        let reopened = Db::open_test(dir.path()).expect("reopen database");
        let queue = reopened
            .open_queue::<AllocatorJobs>(QueueConfig::balanced().build().expect("config"))
            .expect("reopen queue");
        assert!(
            queue.inner.next_id().expect("allocate after reopen") > greatest,
            "reopened allocator reused a previously published id"
        );
    }

    #[test]
    fn watermark_sync_failure_rejects_push_without_creating_a_queue_row() {
        use crate::fixed::shard::{SyncSite, fail_next_sync};

        let dir = tempfile::tempdir().expect("tempdir");
        let db = Db::open_test(dir.path()).expect("open database");
        let queue = db
            .open_queue::<AllocatorJobs>(QueueConfig::balanced().build().expect("config"))
            .expect("open queue");

        queue.inner.seq.flush().expect("reset sync interval");
        let failure = fail_next_sync(SyncSite::Write);
        assert!(queue.push(AllocatorJob(1), None).is_err());
        drop(failure);
        assert_eq!(queue.len(), 0, "failed reservation must not increase depth");
        assert_eq!(
            queue.inner.headers.iter().count(),
            0,
            "failed reservation must not create a header row"
        );
    }
}