klieo-bus-memory 0.5.0

//! In-process `JobQueue` implementation.
//!
//! Backed by per-queue [`PerQueue`] state guarded by a single
//! `tokio::sync::Mutex`. Each `PerQueue` separates ready and leased ids
//! into two `VecDeque`s with the [`JobEntry`] payload kept in a shared
//! `HashMap` (W1.A3 / round-2 HIGH — old impl scanned the full
//! VecDeque per claim).
//!
//! Lease expiry is **lazy** — checked only at `claim()` time when the
//! ready queue is empty. Each `claim` increments `attempts`, sets
//! `lease_until = now + ttl`. `ack` removes the entry. `nak(delay)`
//! either resets the lease + sets `not_before` for back-off or, if
//! `attempts >= max_attempts`, removes the entry and publishes the
//! payload to `<queue>.dlq` with `dlq_reason: "max_attempts_exceeded"`.
//! `dead_letter(reason)` removes the entry and publishes to
//! `<queue>.dlq` with the caller-supplied reason.
//!
//! Dedup-CAS via the shared [`crate::kv::MemoryKv`]: when
//! `Job.dedup_key` is set, `enqueue` writes an idempotency record via
//! `KvStore::cas(... expected=None)` before queueing. A `CasConflict`
//! response means the same dedup key has already been enqueued; the
//! existing `JobId` is read back and returned without queueing.

use crate::kv::MemoryKv;
use crate::pubsub::MemoryPubsub;
use async_trait::async_trait;
use bytes::Bytes;
use klieo_core::bus::{ClaimHandle, ClaimHandleImpl, ClaimedJob, Job, JobQueue, Lease, LeaseImpl};
use klieo_core::error::BusError;
use klieo_core::ids::JobId;
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;

#[derive(Debug)]
struct JobEntry {
    payload: Bytes,
    attempts: u32,
    max_attempts: u32,
    /// `Some(deadline)` — entry is claimed; the lease expires at this
    /// time. `None` — no active lease.
    lease_until: Option<Instant>,
    /// `Some(deadline)` — claim is suppressed until this time
    /// (post-`nak(delay)`).
    not_before: Option<Instant>,
}

#[derive(Default)]
struct PerQueue {
    /// Source-of-truth payload + lease metadata, keyed by JobId.
    entries: HashMap<JobId, JobEntry>,
    /// FIFO of ids currently eligible to be claimed (may include
    /// back-off-suppressed entries — `claim` skips and re-pushes them).
    ready: VecDeque<JobId>,
    /// Ids currently leased. `claim` only sweeps this for expiry when
    /// `ready` is empty.
    leased: HashSet<JobId>,
}

#[derive(Default)]
pub(crate) struct State {
    queues: HashMap<String, PerQueue>,
}

/// In-process `JobQueue` impl.
pub struct MemoryJobQueue {
    state: Arc<Mutex<State>>,
    pub(crate) pubsub: Arc<MemoryPubsub>,
    kv: Arc<MemoryKv>,
    next_id: AtomicU64,
}

impl MemoryJobQueue {
    /// Build atop the supplied pubsub (for DLQ routing) and KV (for
    /// dedup-CAS).
    pub fn new(pubsub: Arc<MemoryPubsub>, kv: Arc<MemoryKv>) -> Self {
        Self {
            state: Arc::new(Mutex::new(State::default())),
            pubsub,
            kv,
            next_id: AtomicU64::new(1),
        }
    }

    fn fresh_id(&self) -> JobId {
        let n = self.next_id.fetch_add(1, Ordering::Relaxed);
        JobId(format!("mem-{n}"))
    }
}

/// Sweep expired leases on `pq` back into the ready queue. Caller holds
/// the queue lock.
fn sweep_expired_leases(pq: &mut PerQueue, now: Instant) {
    let expired: Vec<JobId> = pq
        .leased
        .iter()
        .filter(|id| {
            pq.entries
                .get(*id)
                .and_then(|e| e.lease_until)
                .map_or(true, |t| t <= now)
        })
        .cloned()
        .collect();
    for id in expired {
        pq.leased.remove(&id);
        if let Some(entry) = pq.entries.get_mut(&id) {
            entry.lease_until = None;
        }
        pq.ready.push_back(id);
    }
}

#[async_trait]
impl JobQueue for MemoryJobQueue {
    async fn enqueue(&self, queue: &str, job: Job) -> Result<JobId, BusError> {
        let id = self.fresh_id();

        // Dedup-CAS: write idempotency record before queueing.
        if let Some(dedup_key) = &job.dedup_key {
            use klieo_core::KvStore;
            let bucket = format!("dedup.{queue}");
            let id_bytes: Bytes = id.0.clone().into_bytes().into();
            match self.kv.cas(&bucket, dedup_key, id_bytes, None).await {
                Ok(_) => {}
                Err(BusError::CasConflict { .. }) => {
                    // Duplicate — read the existing JobId.
                    let existing = self.kv.get(&bucket, dedup_key).await?.ok_or_else(|| {
                        BusError::Permanent("dedup record missing after CAS conflict".into())
                    })?;
                    let id_str = String::from_utf8(existing.value.to_vec())
                        .map_err(|e| BusError::Permanent(format!("dedup id bytes: {e}")))?;
                    return Ok(JobId(id_str));
                }
                Err(e) => return Err(e),
            }
        }

        let max_attempts = job.max_attempts.unwrap_or(5);
        let mut g = self.state.lock().await;
        let pq = g.queues.entry(queue.into()).or_default();
        pq.entries.insert(
            id.clone(),
            JobEntry {
                payload: job.payload,
                attempts: 0,
                max_attempts,
                lease_until: None,
                not_before: None,
            },
        );
        pq.ready.push_back(id.clone());
        Ok(id)
    }

    async fn claim(
        &self,
        queue: &str,
        _worker_id: &str,
        lease_ttl: Duration,
    ) -> Result<Option<ClaimedJob>, BusError> {
        let now = Instant::now();
        let mut g = self.state.lock().await;
        let pq = g.queues.entry(queue.into()).or_default();

        // Fast path: try ready queue. Skip + re-push entries still in
        // their back-off window. We scan at most `ready.len()` per call
        // — bounded by current ready depth, independent of leased depth
        // (W1.A3 / round-2 HIGH).
        if let Some(claim) = try_claim_from_ready(pq, now, lease_ttl) {
            drop(g);
            return Ok(Some(build_claimed_job(
                self.state.clone(),
                self.pubsub.clone(),
                queue.to_string(),
                claim,
                lease_ttl,
            )));
        }

        // Slow path: ready was empty (or fully back-off-suppressed).
        // Sweep expired leases back to ready and try once more.
        sweep_expired_leases(pq, now);
        if let Some(claim) = try_claim_from_ready(pq, now, lease_ttl) {
            drop(g);
            return Ok(Some(build_claimed_job(
                self.state.clone(),
                self.pubsub.clone(),
                queue.to_string(),
                claim,
                lease_ttl,
            )));
        }

        Ok(None)
    }
}

/// Pop the front-most ready id whose `not_before` has elapsed and
/// commit it as leased. Returns the id + payload + max_attempts so the
/// caller can build a `ClaimedJob` outside the lock. Skipped (still
/// back-off-suppressed) ids are pushed to the back of `ready` to
/// preserve FIFO progress.
fn try_claim_from_ready(
    pq: &mut PerQueue,
    now: Instant,
    lease_ttl: Duration,
) -> Option<ClaimSnapshot> {
    let mut rotated = 0;
    let initial_len = pq.ready.len();
    while rotated < initial_len {
        let Some(id) = pq.ready.pop_front() else {
            return None;
        };
        let claimable = pq
            .entries
            .get(&id)
            .map(|e| e.not_before.is_none_or(|t| t <= now))
            .unwrap_or(false);

        if !claimable {
            pq.ready.push_back(id);
            rotated += 1;
            continue;
        }

        let entry = pq.entries.get_mut(&id).expect("ready id without entry");
        entry.attempts += 1;
        entry.lease_until = Some(now + lease_ttl);
        entry.not_before = None;
        let snapshot = ClaimSnapshot {
            id: id.clone(),
            payload: entry.payload.clone(),
            max_attempts: entry.max_attempts,
        };
        pq.leased.insert(id);
        return Some(snapshot);
    }
    None
}

/// Per-claim snapshot of the fields needed to build a `ClaimedJob`.
struct ClaimSnapshot {
    id: JobId,
    payload: Bytes,
    max_attempts: u32,
}

fn build_claimed_job(
    state: Arc<Mutex<State>>,
    pubsub: Arc<MemoryPubsub>,
    queue: String,
    snapshot: ClaimSnapshot,
    lease_ttl: Duration,
) -> ClaimedJob {
    let id = snapshot.id.clone();
    ClaimedJob {
        id: id.clone(),
        payload: snapshot.payload,
        lease: Lease::new(Box::new(JobLease {
            state: state.clone(),
            queue: queue.clone(),
            job_id: id.clone(),
            ttl: lease_ttl,
        })),
        claim: ClaimHandle::new(Box::new(MemoryClaimHandle {
            state,
            pubsub,
            queue,
            job_id: id,
            max_attempts: snapshot.max_attempts,
        })),
    }
}

/// Lease handle on a claimed job — extends `lease_until` on heartbeat.
struct JobLease {
    state: Arc<Mutex<State>>,
    queue: String,
    job_id: JobId,
    ttl: Duration,
}

#[async_trait]
impl LeaseImpl for JobLease {
    async fn heartbeat(&self) -> Result<(), BusError> {
        let mut g = self.state.lock().await;
        let pq = g
            .queues
            .get_mut(&self.queue)
            .ok_or_else(|| BusError::NotFound(self.queue.clone()))?;
        let entry = pq
            .entries
            .get_mut(&self.job_id)
            .ok_or_else(|| BusError::NotFound(format!("job {}", self.job_id.0)))?;
        entry.lease_until = Some(Instant::now() + self.ttl);
        Ok(())
    }
}

/// Claim resolution handle. Holds the state needed by `ack`, `nak`, and
/// `dead_letter` (queue name, job id, max-attempts threshold, plus
/// references to the shared queue state and pubsub for DLQ routing).
pub(crate) struct MemoryClaimHandle {
    pub(crate) state: Arc<Mutex<State>>,
    pub(crate) pubsub: Arc<MemoryPubsub>,
    pub(crate) queue: String,
    pub(crate) job_id: JobId,
    pub(crate) max_attempts: u32,
}

#[async_trait]
impl ClaimHandleImpl for MemoryClaimHandle {
    async fn ack(self: Box<Self>) -> Result<(), BusError> {
        let mut g = self.state.lock().await;
        if let Some(pq) = g.queues.get_mut(&self.queue) {
            pq.entries.remove(&self.job_id);
            pq.leased.remove(&self.job_id);
            // The id may still be in `ready` if a heartbeat/lease-sweep
            // race put it there; scrub it.
            pq.ready.retain(|id| id != &self.job_id);
        }
        Ok(())
    }

    async fn nak(self: Box<Self>, delay: Duration) -> Result<(), BusError> {
        let payload_to_dlq;
        {
            let mut g = self.state.lock().await;
            let pq = match g.queues.get_mut(&self.queue) {
                Some(q) => q,
                None => return Ok(()),
            };
            let entry = match pq.entries.get_mut(&self.job_id) {
                Some(e) => e,
                None => return Ok(()),
            };
            if entry.attempts >= self.max_attempts {
                let removed = pq
                    .entries
                    .remove(&self.job_id)
                    .expect("entry exists, just checked");
                pq.leased.remove(&self.job_id);
                pq.ready.retain(|id| id != &self.job_id);
                payload_to_dlq = Some(removed.payload);
            } else {
                entry.lease_until = None;
                entry.not_before = Some(Instant::now() + delay);
                pq.leased.remove(&self.job_id);
                pq.ready.push_back(self.job_id.clone());
                payload_to_dlq = None;
            }
        }
        if let Some(payload) = payload_to_dlq {
            use klieo_core::bus::Headers;
            use klieo_core::Pubsub;
            let mut headers = Headers::new();
            headers.insert("dlq_reason".into(), "max_attempts_exceeded".into());
            let subject = format!("{}.dlq", self.queue);
            self.pubsub.publish(&subject, payload, headers).await?;
        }
        Ok(())
    }

    async fn dead_letter(self: Box<Self>, reason: &str) -> Result<(), BusError> {
        let payload = {
            let mut g = self.state.lock().await;
            let pq = match g.queues.get_mut(&self.queue) {
                Some(q) => q,
                None => return Ok(()),
            };
            let removed = match pq.entries.remove(&self.job_id) {
                Some(e) => e,
                None => return Ok(()),
            };
            pq.leased.remove(&self.job_id);
            pq.ready.retain(|id| id != &self.job_id);
            removed.payload
        };
        use klieo_core::bus::Headers;
        use klieo_core::Pubsub;
        let mut headers = Headers::new();
        headers.insert("dlq_reason".into(), reason.into());
        let subject = format!("{}.dlq", self.queue);
        self.pubsub.publish(&subject, payload, headers).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::kv::MemoryKv;
    use crate::pubsub::MemoryPubsub;
    use klieo_core::bus::Job;
    use klieo_core::ids::DurableName;
    use klieo_core::{JobQueue, Pubsub};
    use std::sync::Arc;
    use std::time::Duration;

    fn build_jq() -> MemoryJobQueue {
        MemoryJobQueue::new(Arc::new(MemoryPubsub::new()), Arc::new(MemoryKv::new()))
    }

    #[tokio::test]
    async fn claim_returns_none_when_empty() {
        let jq = build_jq();
        let claimed = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
        assert!(claimed.is_none());
    }

    #[tokio::test]
    async fn enqueue_then_claim_returns_job() {
        let jq = build_jq();
        let id = jq
            .enqueue("q", Job::new(Bytes::from_static(b"payload")))
            .await
            .unwrap();
        let claimed = jq
            .claim("q", "w1", Duration::from_secs(1))
            .await
            .unwrap()
            .unwrap();
        assert_eq!(claimed.id, id);
        assert_eq!(claimed.payload, Bytes::from_static(b"payload"));
    }

    #[tokio::test]
    async fn ack_removes_job() {
        let jq = build_jq();
        jq.enqueue("q", Job::new(Bytes::from_static(b"x")))
            .await
            .unwrap();
        let claimed = jq
            .claim("q", "w1", Duration::from_secs(1))
            .await
            .unwrap()
            .unwrap();
        claimed.ack().await.unwrap();
        let none = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
        assert!(none.is_none());
    }

    #[tokio::test]
    async fn lease_expiry_redelivers_job() {
        let jq = build_jq();
        jq.enqueue("q", Job::new(Bytes::from_static(b"x")))
            .await
            .unwrap();
        let claim1 = jq
            .claim("q", "w1", Duration::from_millis(20))
            .await
            .unwrap()
            .unwrap();
        // Drop without ack — lease should expire.
        drop(claim1);
        tokio::time::sleep(Duration::from_millis(40)).await;
        let claim2 = jq
            .claim("q", "w2", Duration::from_secs(1))
            .await
            .unwrap()
            .unwrap();
        assert_eq!(claim2.payload, Bytes::from_static(b"x"));
    }

    #[tokio::test]
    async fn nak_with_delay_suppresses_claim_until_deadline() {
        let jq = build_jq();
        jq.enqueue("q", Job::new(Bytes::from_static(b"x")))
            .await
            .unwrap();
        let claim1 = jq
            .claim("q", "w1", Duration::from_secs(1))
            .await
            .unwrap()
            .unwrap();
        claim1.nak(Duration::from_millis(40)).await.unwrap();
        // Immediately re-claim — should be suppressed.
        let none = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
        assert!(
            none.is_none(),
            "should be suppressed during not_before window"
        );
        tokio::time::sleep(Duration::from_millis(60)).await;
        let some = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
        assert!(
            some.is_some(),
            "should be reclaimable after not_before elapses"
        );
    }

    #[tokio::test]
    async fn nak_after_max_attempts_routes_to_dlq() {
        let jq = build_jq();
        let pubsub_handle = jq.pubsub.clone();
        let mut dlq_sub = pubsub_handle
            .subscribe("q.dlq", DurableName::new("dlq"))
            .await
            .unwrap();

        let mut job = Job::new(Bytes::from_static(b"x"));
        job.max_attempts = Some(2);
        jq.enqueue("q", job).await.unwrap();

        for _ in 0..2 {
            let claim = jq
                .claim("q", "w1", Duration::from_secs(1))
                .await
                .unwrap()
                .unwrap();
            claim.nak(Duration::from_millis(0)).await.unwrap();
        }

        // After max_attempts naks, the job should be in the DLQ and the
        // queue should be empty.
        let none = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
        assert!(none.is_none(), "queue should be empty after DLQ routing");

        use tokio_stream::StreamExt;
        let dlq_msg = tokio::time::timeout(Duration::from_millis(200), dlq_sub.next())
            .await
            .expect("DLQ delivery should arrive within 200ms")
            .expect("stream did not close")
            .expect("no error");
        assert_eq!(dlq_msg.payload, Bytes::from_static(b"x"));
        assert_eq!(
            dlq_msg.headers.get("dlq_reason").map(|s| s.as_str()),
            Some("max_attempts_exceeded")
        );
    }

    #[tokio::test]
    async fn dead_letter_removes_job_and_publishes_with_reason() {
        let jq = build_jq();
        use klieo_core::Pubsub;
        let mut dlq_sub = jq
            .pubsub
            .subscribe("q.dlq", klieo_core::ids::DurableName::new("dlq"))
            .await
            .unwrap();
        jq.enqueue("q", Job::new(Bytes::from_static(b"x")))
            .await
            .unwrap();
        let claim = jq
            .claim("q", "w1", Duration::from_secs(1))
            .await
            .unwrap()
            .unwrap();
        claim.dead_letter("explicit").await.unwrap();
        let none = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
        assert!(none.is_none(), "queue should be empty after DLQ");
        use tokio_stream::StreamExt;
        let msg = tokio::time::timeout(Duration::from_millis(200), dlq_sub.next())
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        assert_eq!(msg.payload, Bytes::from_static(b"x"));
        assert_eq!(
            msg.headers.get("dlq_reason").map(|s| s.as_str()),
            Some("explicit")
        );
    }

    #[tokio::test]
    async fn dedup_key_returns_existing_id_on_duplicate_enqueue() {
        let jq = build_jq();
        let mut j1 = Job::new(Bytes::from_static(b"x"));
        j1.dedup_key = Some("d1".into());
        let id1 = jq.enqueue("q", j1).await.unwrap();
        let mut j2 = Job::new(Bytes::from_static(b"y"));
        j2.dedup_key = Some("d1".into());
        let id2 = jq.enqueue("q", j2).await.unwrap();
        assert_eq!(id1, id2, "duplicate dedup_key returns existing id");
        // Only one job actually enqueued.
        let _claimed = jq
            .claim("q", "w", Duration::from_secs(1))
            .await
            .unwrap()
            .unwrap();
        let none = jq.claim("q", "w", Duration::from_secs(1)).await.unwrap();
        assert!(none.is_none(), "second enqueue must not have queued a job");
    }

    #[tokio::test]
    async fn dedup_key_distinct_keys_enqueue_independently() {
        let jq = build_jq();
        let mut j1 = Job::new(Bytes::from_static(b"x"));
        j1.dedup_key = Some("a".into());
        let id1 = jq.enqueue("q", j1).await.unwrap();
        let mut j2 = Job::new(Bytes::from_static(b"y"));
        j2.dedup_key = Some("b".into());
        let id2 = jq.enqueue("q", j2).await.unwrap();
        assert_ne!(id1, id2);

        // Both jobs must actually be queued, not just have distinct ids.
        let claim1 = jq
            .claim("q", "w", Duration::from_secs(1))
            .await
            .unwrap()
            .expect("first job should be claimable");
        let claim2 = jq
            .claim("q", "w", Duration::from_secs(1))
            .await
            .unwrap()
            .expect("second job should be claimable");
        // Order is FIFO — first enqueued comes first.
        assert_eq!(claim1.id, id1);
        assert_eq!(claim2.id, id2);
        assert_eq!(claim1.payload, Bytes::from_static(b"x"));
        assert_eq!(claim2.payload, Bytes::from_static(b"y"));
    }

    /// W1.A3 / round-2 HIGH: with many leased entries, claim() must not
    /// scan them on the fast path. Asserts FIFO-ordering of the ready
    /// queue when most entries are leased — old impl walked the full
    /// VecDeque per call, new impl drains `ready` directly.
    #[tokio::test]
    async fn claim_skips_leased_entries_fast_path() {
        let jq = build_jq();
        // Enqueue 100 jobs.
        for i in 0..100u8 {
            jq.enqueue("q", Job::new(Bytes::from(vec![i])))
                .await
                .unwrap();
        }
        // Claim and hold the first 50.
        let mut held = Vec::new();
        for _ in 0..50 {
            held.push(
                jq.claim("q", "w", Duration::from_secs(10))
                    .await
                    .unwrap()
                    .unwrap(),
            );
        }
        // The 51st claim must succeed without scanning the 50 leased.
        let claim51 = jq
            .claim("q", "w", Duration::from_secs(10))
            .await
            .unwrap()
            .expect("51st claim returns next ready entry");
        assert_eq!(claim51.payload, Bytes::from(vec![50]));
        // Keep `held` alive across the assertion to keep leases active.
        drop(held);
    }
}