later 0.0.49

use crate::{
    encoder,
    id::{Id, IdOf},
    models::{DelayedStage, Job, RecurringJob, RequeuedStage, Stage, StageName},
    storage::{RangeOrder, RangePage, Storage, StorageOperation},
    JobId, RecurringJobId, UtcDateTime,
};
use serde::{de::DeserializeOwned, Serialize};

pub(crate) struct Persist {
    pub(crate) inner: Box<dyn Storage>,
    key_prefix: String,
}

pub(crate) struct Config<'s> {
    inner: &'s dyn Storage,
    key_prefix: String,
}

#[cfg(feature = "prometheus")]
pub(crate) struct QueueDepth {
    pub(crate) state: QueueState,
    pub(crate) count: usize,
}

#[cfg(feature = "prometheus")]
pub(crate) enum QueueState {
    Delayed,
    DependencyWaiting,
    Ready,
    Leased,
    RetryWaiting,
}

#[cfg(feature = "prometheus")]
impl QueueState {
    pub(crate) fn metric_name(&self) -> &'static str {
        match self {
            Self::Delayed => "delayed",
            Self::DependencyWaiting => "dependency_waiting",
            Self::Ready => "ready",
            Self::Leased => "leased",
            Self::RetryWaiting => "retry_waiting",
        }
    }
}

impl Persist {
    fn requires_stage_index(stage: &Stage) -> bool {
        match stage {
            Stage::Delayed(_) | Stage::Requeued(_) => true,
            // Only ever read back to feed the Prometheus queue-depth gauge.
            // With the dashboard feature, `later_jobs_index` already tracks
            // these stages exactly and the gauge reads it (see
            // `queue_depths`), so the extra add/remove per transition - two
            // more writes on a million-row key - is pure overhead. Without
            // the dashboard there is no index, so the key is still needed.
            Stage::Waiting(_) | Stage::Enqueued(_) | Stage::Running(_) => {
                cfg!(all(feature = "prometheus", not(feature = "dashboard")))
            }
            Stage::Success(_) | Stage::Failed(_) => false,
        }
    }

    pub fn new(storage: Box<dyn Storage>, key_prefix: String) -> Self {
        Self {
            inner: storage,
            key_prefix,
        }
    }

    pub fn config(&self) -> Config<'_> {
        Config {
            inner: self.inner.as_ref(),
            key_prefix: self.key_prefix.clone(),
        }
    }

    #[cfg(feature = "dashboard")]
    pub(crate) fn key_prefix(&self) -> &str {
        &self.key_prefix
    }

    #[tracing::instrument(skip(self), level = "trace")]
    pub async fn get_job(&self, id: JobId) -> anyhow::Result<Option<Job>> {
        let id = IdOf::SavedJob(id).get_id(&self.key_prefix);
        self.get_of_type(id).await
    }

    /// The storage key a job's encoded state is kept under.
    pub(crate) fn job_key(&self, id: &JobId) -> String {
        IdOf::SavedJob(id.clone())
            .get_id(&self.key_prefix)
            .to_string()
    }

    #[tracing::instrument(skip(self), level = "trace")]
    pub async fn get_recurring_job(
        &self,
        id: RecurringJobId,
    ) -> anyhow::Result<Option<RecurringJob>> {
        let id = IdOf::SavedRecurringJob(id).get_id(&self.key_prefix);
        self.get_of_type(id).await
    }

    pub async fn get_of_type<T>(&self, id: Id) -> anyhow::Result<Option<T>>
    where
        T: DeserializeOwned,
    {
        self.inner
            .get(&id.to_string())
            .await?
            .map(|bytes| encoder::decode::<T>(&bytes))
            .transpose()
    }

    #[cfg(all(test, feature = "dashboard"))]
    pub async fn save_job(&self, job: &Job) -> anyhow::Result<()> {
        self.inner
            .apply(Self::job_operations(&self.key_prefix, job)?)
            .await
    }

    pub(crate) fn job_operations(
        key_prefix: &str,
        job: &Job,
    ) -> anyhow::Result<Vec<StorageOperation>> {
        let id = IdOf::SavedJob(job.id.clone()).get_id(key_prefix);
        let encoded_id = encoder::encode(&job.id)?;
        let mut operations = vec![StorageOperation::Set {
            key: id.to_string(),
            value: encoder::encode(job)?,
        }];
        if let Some(previous_stage) = job.previous_stages.last() {
            if Self::requires_stage_index(previous_stage) {
                operations.push(StorageOperation::RangeRemove {
                    key: Self::polling_range_key_for(key_prefix, previous_stage.get_name()),
                    value: encoded_id.clone(),
                });
            }
        }
        if Self::requires_stage_index(&job.stage) {
            operations.push(StorageOperation::RangeAdd {
                key: Self::polling_range_key_for(key_prefix, job.stage.get_name()),
                value: encoded_id.clone(),
            });
        }
        if let Stage::Waiting(waiting) = &job.stage {
            let key = IdOf::ContinuationOf(waiting.parent_id.clone()).get_id(key_prefix);
            operations.push(StorageOperation::RangeAdd {
                key: key.to_string(),
                value: encoded_id,
            });
        }

        Ok(operations)
    }

    #[tracing::instrument(skip(self))]
    pub async fn save_recurring_job(&self, job: &RecurringJob) -> anyhow::Result<()> {
        let is_new = self.get_recurring_job(job.id.clone()).await?.is_none();
        let id = IdOf::SavedRecurringJob(job.id.clone()).get_id(&self.key_prefix);
        self.save(id, job).await?;

        // Only index a recurring job the first time it's registered — an
        // upsert must not add a duplicate entry to the listing the
        // recurring-job poller iterates.
        if is_new {
            let encoded_id = encoder::encode(&job.id)?;
            self.inner
                .apply(vec![StorageOperation::RangeAdd {
                    key: IdOf::AllRecurringJobs.get_id(&self.key_prefix).to_string(),
                    value: encoded_id,
                }])
                .await?;
        }

        Ok(())
    }

    /// Atomically creates a recurring job definition only if `job.id` isn't
    /// already registered. Returns `true` if this call created it, `false`
    /// if another concurrent caller (for example, another process
    /// registering the same brand-new identifier at process startup) won
    /// the race and created it first - in which case this call changed
    /// nothing, and the caller should defer to whatever that other writer
    /// wrote rather than also treating itself as the one true first
    /// registration.
    #[tracing::instrument(skip(self))]
    pub async fn create_recurring_job_if_absent(&self, job: &RecurringJob) -> anyhow::Result<bool> {
        let id = IdOf::SavedRecurringJob(job.id.clone()).get_id(&self.key_prefix);
        let created = self
            .inner
            .set_if_absent(&id.to_string(), &encoder::encode(job)?)
            .await?;
        if created {
            let encoded_id = encoder::encode(&job.id)?;
            self.inner
                .apply(vec![StorageOperation::RangeAdd {
                    key: IdOf::AllRecurringJobs.get_id(&self.key_prefix).to_string(),
                    value: encoded_id,
                }])
                .await?;
        }
        Ok(created)
    }

    /// Pages through every registered recurring job definition, oldest
    /// first. Used by the recurring-job poller.
    pub async fn all_recurring_jobs_page(
        &self,
        cursor: Option<i64>,
        limit: usize,
    ) -> anyhow::Result<RangePage> {
        self.inner
            .range_page(
                &IdOf::AllRecurringJobs.get_id(&self.key_prefix).to_string(),
                cursor,
                limit,
                RangeOrder::OldestFirst,
            )
            .await
    }

    #[tracing::instrument(skip(self))]
    pub async fn save<T>(&self, id: Id, item: T) -> anyhow::Result<()>
    where
        T: Serialize + std::fmt::Debug,
    {
        self.inner
            .set(&id.to_string(), &encoder::encode(item)?)
            .await
    }

    #[cfg(all(test, feature = "dashboard"))]
    pub async fn expire_job(
        &self,
        job_id: JobId,
        duration: std::time::Duration,
    ) -> anyhow::Result<()> {
        self.inner
            .apply(Self::expire_job_operations(
                &self.key_prefix,
                job_id,
                duration,
            )?)
            .await
    }

    pub(crate) fn expire_job_operations(
        key_prefix: &str,
        job_id: JobId,
        duration: std::time::Duration,
    ) -> anyhow::Result<Vec<StorageOperation>> {
        let key = IdOf::SavedJob(job_id).get_id(key_prefix);
        let ttl_seconds = usize::try_from(duration.as_secs())?;
        Ok(vec![StorageOperation::Expire {
            key: key.to_string(),
            ttl_seconds,
        }])
    }

    pub async fn get_continuation_jobs(&self, job: &Job) -> anyhow::Result<Vec<Job>> {
        let key = IdOf::ContinuationOf(job.id.clone()).get_id(&self.key_prefix);
        let mut cursor = None;
        let mut jobs = Vec::new();
        loop {
            let page = self
                .inner
                .range_page(&key.to_string(), cursor, 100, RangeOrder::OldestFirst)
                .await?;
            for item in page.items {
                let job_id = encoder::decode::<JobId>(&item.value)?;
                if let Some(job) = self.get_job(job_id).await? {
                    jobs.push(job);
                }
            }
            let Some(next_cursor) = page.next_cursor else {
                return Ok(jobs);
            };
            cursor = Some(next_cursor);
        }
    }

    pub async fn clear_continuation_jobs(&self, job: &Job) -> anyhow::Result<()> {
        let key = IdOf::ContinuationOf(job.id.clone()).get_id(&self.key_prefix);
        self.inner
            .apply(vec![StorageOperation::RangeClear {
                key: key.to_string(),
            }])
            .await
    }

    pub fn polling_range_key(&self, stage_name: String) -> String {
        Self::polling_range_key_for(&self.key_prefix, stage_name)
    }

    fn polling_range_key_for(key_prefix: &str, stage_name: String) -> String {
        IdOf::JobsInStagesId(stage_name)
            .get_id(key_prefix)
            .to_string()
    }

    pub async fn remove_job_from_poll_range(
        &self,
        id: &JobId,
        stage_name: &str,
    ) -> anyhow::Result<()> {
        self.inner
            .apply(vec![StorageOperation::RangeRemove {
                key: self.polling_range_key(stage_name.to_string()),
                value: encoder::encode(id)?,
            }])
            .await
    }

    pub async fn delayed_jobs_page(&self, cursor: Option<i64>) -> anyhow::Result<RangePage> {
        self.jobs_to_poll_page(&DelayedStage::get_name(), cursor)
            .await
    }

    pub async fn requeued_jobs_page(&self, cursor: Option<i64>) -> anyhow::Result<RangePage> {
        self.jobs_to_poll_page(&RequeuedStage::get_name(), cursor)
            .await
    }

    #[cfg(feature = "prometheus")]
    pub(crate) async fn queue_depths(&self) -> anyhow::Result<Vec<QueueDepth>> {
        use crate::models::{EnqueuedStage, RunningStage, WaitingStage};

        let states = [
            (QueueState::Delayed, DelayedStage::get_name()),
            (QueueState::DependencyWaiting, WaitingStage::get_name()),
            (QueueState::Ready, EnqueuedStage::get_name()),
            (QueueState::Leased, RunningStage::get_name()),
            (QueueState::RetryWaiting, RequeuedStage::get_name()),
        ];
        // One covering-index query instead of a `range_count` per state:
        // counting the enqueued stage's range key meant scanning every
        // queued job (2-4s per call at a 1.3M-job backlog), every 5
        // seconds, on a pool the workers share.
        #[cfg(feature = "dashboard")]
        let indexed = self.inner.job_index_stage_counts(&self.key_prefix).await?;
        let mut depths = Vec::with_capacity(states.len());
        for (state, stage_name) in states {
            #[cfg(feature = "dashboard")]
            let count = indexed.get(&stage_name).copied().unwrap_or(0);
            #[cfg(not(feature = "dashboard"))]
            let count = self
                .inner
                .range_count(&Self::polling_range_key_for(&self.key_prefix, stage_name))
                .await?;
            depths.push(QueueDepth { state, count });
        }
        Ok(depths)
    }

    async fn jobs_to_poll_page(
        &self,
        stage_name: &str,
        cursor: Option<i64>,
    ) -> anyhow::Result<RangePage> {
        self.inner
            .range_page(
                &self.polling_range_key(stage_name.to_string()),
                cursor,
                100,
                RangeOrder::OldestFirst,
            )
            .await
    }
}

impl<'s> Config<'s> {
    fn one_year_ago() -> UtcDateTime {
        let now = chrono::Utc::now();
        chrono::Duration::try_days(365)
            .and_then(|duration| now.checked_sub_signed(duration))
            .unwrap_or(now)
    }

    pub async fn poll_delayed_jobs_last_run(&self) -> anyhow::Result<UtcDateTime> {
        self.get(
            IdOf::ConfigDateLastPolledForDelayedJobs.get_id(&self.key_prefix),
            Self::one_year_ago(),
        )
        .await
    }

    pub async fn poll_requeued_jobs_last_run(&self) -> anyhow::Result<UtcDateTime> {
        self.get(
            IdOf::ConfigDateLastPolledForReqdJobs.get_id(&self.key_prefix),
            Self::one_year_ago(),
        )
        .await
    }

    pub async fn poll_delayed_jobs_last_run_set(&self) -> anyhow::Result<()> {
        self.set(
            IdOf::ConfigDateLastPolledForDelayedJobs.get_id(&self.key_prefix),
            chrono::Utc::now(),
        )
        .await
    }

    pub async fn poll_requeued_jobs_last_run_set(&self) -> anyhow::Result<()> {
        self.set(
            IdOf::ConfigDateLastPolledForReqdJobs.get_id(&self.key_prefix),
            chrono::Utc::now(),
        )
        .await
    }

    pub async fn poll_recurring_jobs_last_run(&self) -> anyhow::Result<UtcDateTime> {
        self.get(
            IdOf::ConfigDateLastPolledForRecurringJobs.get_id(&self.key_prefix),
            Self::one_year_ago(),
        )
        .await
    }

    pub async fn poll_recurring_jobs_last_run_set(&self) -> anyhow::Result<()> {
        self.set(
            IdOf::ConfigDateLastPolledForRecurringJobs.get_id(&self.key_prefix),
            chrono::Utc::now(),
        )
        .await
    }

    pub async fn poll_expired_storage_last_run(&self) -> anyhow::Result<UtcDateTime> {
        self.get(
            IdOf::ConfigDateLastPolledForExpiredStorage.get_id(&self.key_prefix),
            Self::one_year_ago(),
        )
        .await
    }

    pub async fn poll_expired_storage_last_run_set(&self) -> anyhow::Result<()> {
        self.set(
            IdOf::ConfigDateLastPolledForExpiredStorage.get_id(&self.key_prefix),
            chrono::Utc::now(),
        )
        .await
    }

    pub async fn poll_stuck_jobs_last_run(&self) -> anyhow::Result<UtcDateTime> {
        self.get(
            IdOf::ConfigDateLastPolledForStuckJobs.get_id(&self.key_prefix),
            Self::one_year_ago(),
        )
        .await
    }

    pub async fn poll_stuck_jobs_last_run_set(&self) -> anyhow::Result<()> {
        self.set(
            IdOf::ConfigDateLastPolledForStuckJobs.get_id(&self.key_prefix),
            chrono::Utc::now(),
        )
        .await
    }

    pub async fn poll_vacuum_last_run(&self) -> anyhow::Result<UtcDateTime> {
        self.get(
            IdOf::ConfigDateLastPolledForVacuum.get_id(&self.key_prefix),
            Self::one_year_ago(),
        )
        .await
    }

    pub async fn poll_vacuum_last_run_set(&self) -> anyhow::Result<()> {
        self.set(
            IdOf::ConfigDateLastPolledForVacuum.get_id(&self.key_prefix),
            chrono::Utc::now(),
        )
        .await
    }

    async fn get<T: DeserializeOwned>(&self, key: Id, default: T) -> anyhow::Result<T> {
        Ok(self
            .inner
            .get(&key.to_string())
            .await?
            .map(|bytes| encoder::decode::<T>(&bytes))
            .transpose()?
            .unwrap_or(default))
    }

    async fn set<T: Serialize>(&self, key: Id, value: T) -> anyhow::Result<()> {
        self.inner
            .set(&key.to_string(), &encoder::encode(&value)?)
            .await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::{EnqueuedStage, JobConfig, Stage};
    #[cfg(feature = "prometheus")]
    use crate::storage::memory::MemoryStorage;

    #[test]
    fn a_job_keeps_no_global_list_of_all_jobs() -> anyhow::Result<()> {
        let job = Job {
            id: JobId("no-global-list".to_string()),
            payload_type: "test".to_string(),
            payload: Vec::new(),
            config: JobConfig::default(),
            stage: Stage::Enqueued(EnqueuedStage {
                date: chrono::Utc::now(),
            }),
            previous_stages: Vec::new(),
            recurring_job_id: None,
            topic: None,
            partition: None,
            sequence: None,
        };
        // An `all-jobs` list had one row per job ever created, written at
        // creation and expired with the job, read only by a rebuild tool that
        // no longer exists - two statements per job and millions of rows.
        let running = job.clone().transition();
        let success = running.clone().transition_success()?;
        let mut operations = Persist::job_operations("test", &job)?;
        operations.extend(Persist::job_operations("test", &running)?);
        operations.extend(Persist::job_operations("test", &success)?);
        operations.extend(Persist::expire_job_operations(
            "test",
            success.id.clone(),
            std::time::Duration::from_secs(3600),
        )?);
        let touches_a_global_list = operations.iter().any(|operation| match operation {
            StorageOperation::RangeAdd { key, .. }
            | StorageOperation::RangeRemove { key, .. }
            | StorageOperation::RangeExpire { key, .. }
            | StorageOperation::RangeClear { key } => key.ends_with("all-jobs"),
            StorageOperation::Set { .. }
            | StorageOperation::Delete { .. }
            | StorageOperation::Expire { .. } => false,
        });
        assert!(!touches_a_global_list);
        Ok(())
    }

    #[cfg(all(feature = "prometheus", feature = "dashboard"))]
    #[test]
    fn dashboard_builds_skip_the_redundant_stage_range_writes() -> anyhow::Result<()> {
        let job = |stage: Stage| Job {
            id: JobId("stage-key-job".to_string()),
            payload_type: "test".to_string(),
            payload: Vec::new(),
            config: JobConfig::default(),
            stage,
            previous_stages: Vec::new(),
            recurring_job_id: None,
            topic: None,
            partition: None,
            sequence: None,
        };
        let stage_key_writes = |job: &Job| -> anyhow::Result<usize> {
            Ok(Persist::job_operations("k", job)?
                .iter()
                .filter(|operation| {
                    matches!(operation, StorageOperation::RangeAdd { key, .. }
                        if key.contains("-stage-"))
                })
                .count())
        };
        let enqueued = job(Stage::Enqueued(EnqueuedStage {
            date: chrono::Utc::now(),
        }));
        assert_eq!(
            stage_key_writes(&enqueued)?,
            0,
            "enqueued must not write a stage range key when the index tracks it"
        );
        let requeued = enqueued.transition().transition_failure("x".to_string())?;
        let crate::models::FailureTransition::Retry(requeued) = requeued else {
            anyhow::bail!("expected a retry");
        };
        assert_eq!(
            stage_key_writes(&requeued)?,
            1,
            "requeued is still polled from its stage range key"
        );
        Ok(())
    }

    #[cfg(feature = "prometheus")]
    #[tokio::test]
    async fn queue_depth_tracks_each_active_job_state() -> anyhow::Result<()> {
        let persist = Persist::new(
            Box::new(MemoryStorage::new()),
            "queue-depth-test".to_string(),
        );
        let enqueued = Job {
            id: JobId("queue-depth-job".to_string()),
            payload_type: "test".to_string(),
            payload: Vec::new(),
            config: JobConfig::default(),
            stage: Stage::Enqueued(EnqueuedStage {
                date: chrono::Utc::now(),
            }),
            previous_stages: Vec::new(),
            recurring_job_id: None,
            topic: None,
            partition: None,
            sequence: None,
        };
        persist
            .inner
            .apply(Persist::job_operations("queue-depth-test", &enqueued)?)
            .await?;
        #[cfg(feature = "dashboard")]
        persist
            .inner
            .job_index_upsert(
                "queue-depth-test",
                crate::stats::job_index_row(&enqueued, None),
            )
            .await?;
        let enqueued_depths = persist
            .queue_depths()
            .await?
            .into_iter()
            .map(|depth| (depth.state.metric_name(), depth.count))
            .collect::<Vec<_>>();

        let running = enqueued.transition();
        persist
            .inner
            .apply(Persist::job_operations("queue-depth-test", &running)?)
            .await?;
        #[cfg(feature = "dashboard")]
        persist
            .inner
            .job_index_upsert(
                "queue-depth-test",
                crate::stats::job_index_row(&running, None),
            )
            .await?;
        let running_depths = persist
            .queue_depths()
            .await?
            .into_iter()
            .map(|depth| (depth.state.metric_name(), depth.count))
            .collect::<Vec<_>>();

        let success = running.transition_success()?;
        persist
            .inner
            .apply(Persist::job_operations("queue-depth-test", &success)?)
            .await?;
        #[cfg(feature = "dashboard")]
        persist
            .inner
            .job_index_upsert(
                "queue-depth-test",
                crate::stats::job_index_row(&success, None),
            )
            .await?;
        let success_depths = persist
            .queue_depths()
            .await?
            .into_iter()
            .map(|depth| (depth.state.metric_name(), depth.count))
            .collect::<Vec<_>>();

        assert_eq!(
            (enqueued_depths, running_depths, success_depths),
            (
                vec![
                    ("delayed", 0),
                    ("dependency_waiting", 0),
                    ("ready", 1),
                    ("leased", 0),
                    ("retry_waiting", 0),
                ],
                vec![
                    ("delayed", 0),
                    ("dependency_waiting", 0),
                    ("ready", 0),
                    ("leased", 1),
                    ("retry_waiting", 0),
                ],
                vec![
                    ("delayed", 0),
                    ("dependency_waiting", 0),
                    ("ready", 0),
                    ("leased", 0),
                    ("retry_waiting", 0),
                ],
            )
        );
        Ok(())
    }
}