compliance-scanner-agent/compliance-agent/src/werkbank/queue.rs

//! The Mongo-backed Werkbank job queue (WB-02).
//!
//! A pull queue: the control plane [`enqueue`](JobQueue::enqueue)s jobs; a runner
//! [`lease`](JobQueue::lease)s the oldest queued job it can run (matched by
//! executor + labels), [`heartbeat`](JobQueue::heartbeat)s while it works, and
//! [`complete`](JobQueue::complete)s it. Leases carry a visibility timeout: if a
//! runner dies mid-job its heartbeats stop, the lease expires, and
//! [`sweep_expired`](JobQueue::sweep_expired) returns the job to `queued` (or
//! `expired` once it has been retried too many times).
//!
//! All state transitions are single atomic Mongo updates guarded by the lease
//! token, so two runners can never both own a job. Every operation takes an
//! explicit `now` so the queue's time-dependent behaviour is deterministically
//! testable.

use std::time::Duration;

use chrono::{DateTime, Utc};
use mongodb::bson::{doc, Bson, DateTime as BsonDateTime};
use mongodb::error::{ErrorKind, WriteFailure};
use mongodb::options::ReturnDocument;
use mongodb::Collection;

use compliance_core::models::werkbank::{
    Executor, HeartbeatAck, Job, JobRecord, JobResult, JobStatus, LeasedJob,
};

use crate::database::Database;
use crate::error::AgentError;

/// The non-terminal states a job can be swept or cancelled from.
const ACTIVE_STATES: [&str; 2] = ["leased", "running"];
/// Every terminal state (no further transitions).
const TERMINAL_STATES: [&str; 4] = ["succeeded", "failed", "expired", "cancelled"];

/// What a visibility-timeout sweep did.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct SweepOutcome {
    /// Expired-lease jobs returned to `queued` for another runner.
    pub requeued: u64,
    /// Jobs that had exhausted their attempts and were marked `expired`.
    pub expired: u64,
}

/// The Mongo-backed job queue.
pub struct JobQueue {
    coll: Collection<JobRecord>,
}

impl JobQueue {
    /// Build a queue over a tenant database's `werkbank_jobs` collection.
    pub fn new(db: &Database) -> Self {
        Self {
            coll: db.werkbank_jobs(),
        }
    }

    /// Enqueue a job. Idempotent by job id: a job that is already present is a
    /// no-op. Returns `true` if this call inserted it, `false` if it existed.
    pub async fn enqueue(&self, job: Job, now: DateTime<Utc>) -> Result<bool, AgentError> {
        let record = JobRecord::queued(job, now);
        match self.coll.insert_one(&record).await {
            Ok(_) => Ok(true),
            Err(e) if is_duplicate_key(&e) => Ok(false),
            Err(e) => Err(e.into()),
        }
    }

    /// Atomically lease the oldest `queued` job this runner can run — matched by
    /// executor and by labels (every label the job requires must be one the
    /// runner advertises). Returns the job plus a lease token, or `None` if
    /// nothing is runnable.
    pub async fn lease(
        &self,
        runner_id: &str,
        executor: Executor,
        runner_labels: &[String],
        lease_ttl: Duration,
        now: DateTime<Utc>,
    ) -> Result<Option<LeasedJob>, AgentError> {
        let token = uuid::Uuid::new_v4().to_string();
        let expires = bson_dt(now + ttl(lease_ttl));
        let executor_bson = mongodb::bson::to_bson(&executor).unwrap_or(Bson::Null);

        let filter = doc! {
            "status": "queued",
            "cancel_requested": { "$ne": true },
            "job.executor": executor_bson,
            // Every label the job requires must be in the runner's set — i.e. the
            // job has no label that is not offered by the runner. Absent/empty
            // job labels match any runner.
            "job.labels": { "$not": { "$elemMatch": { "$nin": runner_labels.to_vec() } } },
        };
        let update = doc! {
            "$set": {
                "status": "leased",
                "lease_token": &token,
                "leased_by": runner_id,
                "lease_expires_at": expires,
                "heartbeat_at": bson_dt(now),
                "updated_at": bson_dt(now),
            },
            "$inc": { "attempts": 1 },
        };

        let record = self
            .coll
            .find_one_and_update(filter, update)
            .sort(doc! { "created_at": 1 }) // FIFO
            .return_document(ReturnDocument::After)
            .await?;
        Ok(record.map(|r| LeasedJob {
            job: r.job,
            lease_token: token,
        }))
    }

    /// Extend a lease and report whether the job has been asked to cancel.
    /// Transitions the job to `running` on the first heartbeat. Returns `None`
    /// when the lease is no longer valid (token mismatch, or the job is already
    /// terminal) — the runner should then abandon the work.
    pub async fn heartbeat(
        &self,
        job_id: &str,
        lease_token: &str,
        lease_ttl: Duration,
        now: DateTime<Utc>,
    ) -> Result<Option<HeartbeatAck>, AgentError> {
        let filter = doc! {
            "job.id": job_id,
            "lease_token": lease_token,
            "status": { "$in": ACTIVE_STATES.to_vec() },
        };
        let update = doc! {
            "$set": {
                "status": "running",
                "lease_expires_at": bson_dt(now + ttl(lease_ttl)),
                "heartbeat_at": bson_dt(now),
                "updated_at": bson_dt(now),
            },
        };
        let record = self
            .coll
            .find_one_and_update(filter, update)
            .return_document(ReturnDocument::After)
            .await?;
        Ok(record.map(|r| HeartbeatAck {
            cancelled: r.cancel_requested,
        }))
    }

    /// Record a job's terminal result. Guarded by the lease token and only from
    /// an active (`leased`/`running`) state, so it is idempotent — a duplicate or
    /// late submission after the job already finished matches nothing. Returns
    /// `true` if this call recorded the result.
    pub async fn complete(
        &self,
        job_id: &str,
        lease_token: &str,
        result: &JobResult,
        now: DateTime<Utc>,
    ) -> Result<bool, AgentError> {
        let status = result.status.unwrap_or(JobStatus::Failed);
        let status_bson = mongodb::bson::to_bson(&status).unwrap_or(Bson::String("failed".into()));
        let result_bson =
            mongodb::bson::to_bson(result).map_err(|e| AgentError::Other(e.to_string()))?;

        let filter = doc! {
            "job.id": job_id,
            "lease_token": lease_token,
            "status": { "$in": ACTIVE_STATES.to_vec() },
        };
        let update = doc! {
            "$set": {
                "status": status_bson,
                "result": result_bson,
                "lease_token": Bson::Null,
                "lease_expires_at": Bson::Null,
                "updated_at": bson_dt(now),
            },
        };
        let res = self.coll.update_one(filter, update).await?;
        Ok(res.modified_count == 1)
    }

    /// Request cancellation of a job. A still-`queued` job is cancelled outright;
    /// an in-flight one is flagged so the runner sees it on its next heartbeat and
    /// tears down. Returns `true` if a non-terminal job matched.
    pub async fn cancel(&self, job_id: &str, now: DateTime<Utc>) -> Result<bool, AgentError> {
        let filter = doc! {
            "job.id": job_id,
            "status": { "$nin": TERMINAL_STATES.to_vec() },
        };
        // Pipeline update: flag cancellation, and if still queued flip straight to
        // cancelled (nothing is running it).
        let pipeline = vec![doc! {
            "$set": {
                "cancel_requested": true,
                "status": {
                    "$cond": [ { "$eq": ["$status", "queued"] }, "cancelled", "$status" ]
                },
                "updated_at": bson_dt(now),
            }
        }];
        let res = self.coll.update_one(filter, pipeline).await?;
        Ok(res.matched_count == 1)
    }

    /// Sweep leases whose visibility timeout has elapsed: return them to `queued`
    /// for another runner, or mark them `expired` once they have been leased
    /// `max_attempts` times. This is what makes a crashed runner's job recover.
    pub async fn sweep_expired(
        &self,
        now: DateTime<Utc>,
        max_attempts: u32,
        // (kept explicit rather than a const so callers can tune retry policy)
    ) -> Result<SweepOutcome, AgentError> {
        let now_bson = bson_dt(now);
        let max = i64::from(max_attempts);

        let requeue = self
            .coll
            .update_many(
                doc! {
                    "status": { "$in": ACTIVE_STATES.to_vec() },
                    "lease_expires_at": { "$lt": &now_bson },
                    "attempts": { "$lt": max },
                },
                doc! { "$set": {
                    "status": "queued",
                    "lease_token": Bson::Null,
                    "leased_by": Bson::Null,
                    "lease_expires_at": Bson::Null,
                    "updated_at": &now_bson,
                } },
            )
            .await?;

        let expire = self
            .coll
            .update_many(
                doc! {
                    "status": { "$in": ACTIVE_STATES.to_vec() },
                    "lease_expires_at": { "$lt": &now_bson },
                    "attempts": { "$gte": max },
                },
                doc! { "$set": {
                    "status": "expired",
                    "lease_token": Bson::Null,
                    "lease_expires_at": Bson::Null,
                    "updated_at": &now_bson,
                } },
            )
            .await?;

        Ok(SweepOutcome {
            requeued: requeue.modified_count,
            expired: expire.modified_count,
        })
    }

    /// Fetch a job record by job id (inspection / control-plane reads).
    pub async fn get(&self, job_id: &str) -> Result<Option<JobRecord>, AgentError> {
        Ok(self.coll.find_one(doc! { "job.id": job_id }).await?)
    }
}

/// A `chrono::Duration` for a lease TTL, saturating rather than panicking on an
/// absurd input (`chrono::Duration::seconds` panics past its internal bound).
fn ttl(d: Duration) -> chrono::Duration {
    let secs = i64::try_from(d.as_secs()).unwrap_or(i64::MAX);
    chrono::Duration::try_seconds(secs).unwrap_or(chrono::Duration::MAX)
}

/// A chrono instant as a BSON date (so Mongo stores/compares it as a real date).
fn bson_dt(dt: DateTime<Utc>) -> BsonDateTime {
    BsonDateTime::from_chrono(dt)
}

/// Whether a Mongo error is a duplicate-key (E11000) violation — a job with this
/// id is already enqueued.
fn is_duplicate_key(e: &mongodb::error::Error) -> bool {
    match &*e.kind {
        ErrorKind::Write(WriteFailure::WriteError(we)) => we.code == 11000,
        _ => false,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn ttl_saturates_and_converts() {
        assert_eq!(ttl(Duration::from_secs(30)), chrono::Duration::seconds(30));
        // An absurd TTL saturates instead of panicking.
        assert_eq!(ttl(Duration::from_secs(u64::MAX)), chrono::Duration::MAX);
    }

    #[test]
    fn state_constants_are_disjoint() {
        for s in ACTIVE_STATES {
            assert!(
                !TERMINAL_STATES.contains(&s),
                "{s} cannot be both active and terminal"
            );
        }
    }
}