cja/jobs/

worker.rs

use std::time::Duration;

use thiserror::Error;
use tokio_util::sync::CancellationToken;
use tracing::Span;

use crate::app_state::AppState as AS;

use super::registry::JobRegistry;

/// Default maximum number of retry attempts before a job is moved to the dead letter queue.
///
/// With exponential backoff (`2^(error_count+1)` seconds), 20 retries means:
/// - First retry after 2 seconds
/// - Last retry after ~6 days
/// - Total retry window of ~12 days
pub const DEFAULT_MAX_RETRIES: i32 = 20;

/// Default lock timeout duration (2 hours).
///
/// Jobs locked for longer than this duration will be considered abandoned and
/// made available for other workers to pick up. This handles cases where a worker
/// crashes or becomes unresponsive while processing a job.
pub const DEFAULT_LOCK_TIMEOUT: Duration = Duration::from_secs(2 * 60 * 60);

pub(super) type RunJobResult = Result<RunJobSuccess, JobError>;

#[derive(Debug)]
pub(super) struct RunJobSuccess(JobFromDB);

#[derive(Debug, sqlx::FromRow)]
pub struct JobFromDB {
    pub job_id: uuid::Uuid,
    pub name: String,
    pub payload: serde_json::Value,
    pub priority: i32,
    pub run_at: chrono::DateTime<chrono::Utc>,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub context: String,
    pub error_count: i32,
    pub last_error_message: Option<String>,
    pub last_failed_at: Option<chrono::DateTime<chrono::Utc>>,
}

#[derive(Debug, Error)]
#[error("JobError(id:${}) ${1}", self.0.job_id)]
pub(crate) struct JobError(JobFromDB, color_eyre::Report);

struct Worker<AppState: AS, R: JobRegistry<AppState>> {
    id: uuid::Uuid,
    state: AppState,
    registry: R,
    sleep_duration: Duration,
    max_retries: i32,
    cancellation_token: CancellationToken,
    lock_timeout: Duration,
}

impl<AppState: AS, R: JobRegistry<AppState>> Worker<AppState, R> {
    fn new(
        state: AppState,
        registry: R,
        sleep_duration: Duration,
        max_retries: i32,
        cancellation_token: CancellationToken,
        lock_timeout: Duration,
    ) -> Self {
        Self {
            id: uuid::Uuid::new_v4(),
            state,
            registry,
            sleep_duration,
            max_retries,
            cancellation_token,
            lock_timeout,
        }
    }

    #[tracing::instrument(
        name = "worker.run_job",
        skip(self, job),
        fields(
            job.id = %job.job_id,
            job.name = job.name,
            job.priority = job.priority,
            job.run_at = %job.run_at,
            job.created_at = %job.created_at,
            job.context = job.context,
            job.error_count = job.error_count,
            worker.id = %self.id,
        )
        err,
    )]
    async fn run_job(&self, job: &JobFromDB) -> color_eyre::Result<()> {
        self.registry
            .run_job(job, self.state.clone(), self.cancellation_token.clone())
            .await
    }

    pub(crate) async fn run_next_job(&self, job: JobFromDB) -> color_eyre::Result<RunJobResult> {
        let job_result = self.run_job(&job).await;

        if let Err(e) = job_result {
            // Extract error message from color_eyre::Report
            let error_message = format!("{e}");

            // Check if job has exceeded max retries
            if job.error_count >= self.max_retries {
                // Move job to dead letter queue
                tracing::error!(
                    worker.id = %self.id,
                    job_id = %job.job_id,
                    error_count = job.error_count,
                    max_retries = self.max_retries,
                    "Job permanently failed - moved to dead letter queue"
                );

                let mut tx = self.state.db().begin().await?;

                sqlx::query(
                    "INSERT INTO dead_letter_jobs (original_job_id, name, payload, context, priority, error_count, last_error_message, created_at)
                     VALUES ($1, $2, $3, $4, $5, $6, $7, $8)",
                )
                .bind(job.job_id)
                .bind(&job.name)
                .bind(&job.payload)
                .bind(&job.context)
                .bind(job.priority)
                .bind(job.error_count)
                .bind(&error_message)
                .bind(job.created_at)
                .execute(&mut *tx)
                .await?;

                sqlx::query("DELETE FROM jobs WHERE job_id = $1 AND locked_by = $2")
                    .bind(job.job_id)
                    .bind(self.id.to_string())
                    .execute(&mut *tx)
                    .await?;

                tx.commit().await?;

                return Ok(Err(JobError(job, e)));
            }

            // Job is under max retries - requeue with exponential backoff
            tracing::warn!(
                worker.id = %self.id,
                job_id = %job.job_id,
                error_count = job.error_count,
                retry_attempt = job.error_count + 1,
                "Job failed, retry #{}",
                job.error_count + 1
            );

            sqlx::query(
                "
                UPDATE jobs
                SET locked_by = NULL,
                    locked_at = NULL,
                    error_count = error_count + 1,
                    last_error_message = $3,
                    last_failed_at = NOW(),
                    run_at = NOW() + (POWER(2, error_count + 1)) * interval '1 second'
                WHERE job_id = $1 AND locked_by = $2
                ",
            )
            .bind(job.job_id)
            .bind(self.id.to_string())
            .bind(error_message)
            .execute(self.state.db())
            .await?;

            return Ok(Err(JobError(job, e)));
        }

        sqlx::query(
            "
                DELETE FROM jobs
                WHERE job_id = $1 AND locked_by = $2
                ",
        )
        .bind(job.job_id)
        .bind(self.id.to_string())
        .execute(self.state.db())
        .await?;

        Ok(Ok(RunJobSuccess(job)))
    }

    #[tracing::instrument(
        name = "worker.fetch_next_job",
        skip(self),
        fields(
            worker.id = %self.id,
            job.id,
            job.name,
            lock_timeout_secs = self.lock_timeout.as_secs(),
        ),
        err,
    )]
    #[allow(clippy::cast_possible_wrap)]
    async fn fetch_next_job(&self) -> color_eyre::Result<Option<JobFromDB>> {
        // Cast is safe: lock timeouts are typically hours, not approaching i64::MAX seconds
        let lock_timeout_secs = self.lock_timeout.as_secs() as i64;

        let job = sqlx::query_as::<_, JobFromDB>(
            "
            UPDATE jobs
            SET LOCKED_BY = $1, LOCKED_AT = NOW()
            WHERE job_id = (
                SELECT job_id
                FROM jobs
                WHERE run_at <= NOW()
                  AND (
                    locked_by IS NULL
                    OR locked_at < NOW() - ($2 || ' seconds')::interval
                  )
                ORDER BY priority DESC, created_at ASC
                LIMIT 1
                FOR UPDATE SKIP LOCKED
            )
            RETURNING job_id, name, payload, priority, run_at, created_at, context, error_count, last_error_message, last_failed_at
            ",
        )
        .bind(self.id.to_string())
        .bind(lock_timeout_secs.to_string())
        .fetch_optional(self.state.db())
        .await?;

        if let Some(job) = &job {
            let span = Span::current();
            span.record("job.id", job.job_id.to_string());
            span.record("job.name", &job.name);
        }

        Ok(job)
    }

    #[tracing::instrument(
        name = "worker.tick",
        skip(self),
        fields(
            worker.id = %self.id,
        ),
    )]
    async fn tick(&self) -> color_eyre::Result<()> {
        let job = self.fetch_next_job().await?;

        let Some(job) = job else {
            let duration = self.sleep_duration;
            tracing::debug!(worker.id =% self.id, ?duration, "No Job to Run, sleeping for requested duration");

            tokio::time::sleep(duration).await;

            return Ok(());
        };

        let result = self.run_next_job(job).await?;

        match result {
            Ok(RunJobSuccess(job)) => {
                tracing::info!(worker.id =% self.id, job_id =% job.job_id, "Job Ran");
            }
            Err(job_error) => {
                tracing::error!(
                    worker.id =% self.id,
                    job_id =% job_error.0.job_id,
                    error_count =% job_error.0.error_count,
                    error_msg =% job_error.1,
                    "Job Errored"
                );
            }
        }

        Ok(())
    }
}

/// Release database locks held by a worker.
///
/// This should be called during graceful shutdown to immediately release any job locks
/// held by this worker, rather than waiting for the 2-hour lock timeout.
async fn cleanup_worker_locks<AppState: AS, R: JobRegistry<AppState>>(
    worker: &Worker<AppState, R>,
) -> color_eyre::Result<()> {
    tracing::info!(worker_id = %worker.id, "Releasing database locks");

    let result = sqlx::query(
        "UPDATE jobs
         SET locked_by = NULL, locked_at = NULL
         WHERE locked_by = $1",
    )
    .bind(worker.id.to_string())
    .execute(worker.state.db())
    .await?;

    tracing::info!(
        worker_id = %worker.id,
        locks_released = result.rows_affected(),
        "Database locks released"
    );

    Ok(())
}

/// Start a job worker that processes jobs from the queue.
///
/// The worker will continuously poll for jobs and execute them using the provided registry.
/// Jobs are executed with automatic retry logic on failure.
///
/// # Arguments
///
/// * `app_state` - The application state containing database connection and configuration
/// * `registry` - The job registry that maps job names to their implementations
/// * `sleep_duration` - How long to sleep when no jobs are available
/// * `max_retries` - Maximum number of times to retry a failed job before moving to the dead letter queue (default: 20)
/// * `shutdown_token` - Cancellation token for graceful shutdown. When cancelled, the worker
///   will stop accepting new jobs and release database locks before exiting.
/// * `lock_timeout` - How long a job can be locked before it's considered abandoned and
///   becomes available for other workers (default: 2 hours)
///
/// # Retry Behavior
///
/// When a job fails:
/// - The error count is incremented
/// - The error message and timestamp are recorded
/// - The job is requeued with exponential backoff: delay = `2^(error_count + 1)` seconds
///   (first retry: 2s, second: 4s, third: 8s, fourth: 16s, etc.)
/// - If `error_count` >= `max_retries`, the job is moved to the dead letter queue
///
/// # Graceful Shutdown
///
/// When the `shutdown_token` is cancelled:
/// - The worker stops polling for new jobs
/// - Any currently executing job is allowed to complete
/// - Database locks are released immediately (instead of waiting for the lock timeout)
///
/// # Lock Timeout
///
/// If a worker crashes or becomes unresponsive while processing a job, the job will remain
/// locked in the database. The `lock_timeout` parameter controls how long to wait before
/// considering such jobs abandoned. After the timeout expires, any worker can pick up the
/// job and retry it.
///
/// # Example
///
/// ```ignore
/// use std::time::Duration;
/// use tokio_util::sync::CancellationToken;
///
/// let shutdown_token = CancellationToken::new();
/// let worker_token = shutdown_token.clone();
///
/// // Start worker with graceful shutdown support and lock timeout
/// tokio::spawn(async move {
///     cja::jobs::worker::job_worker(
///         app_state,
///         registry,
///         Duration::from_secs(60),      // poll every 60s when idle
///         20,                            // max 20 retries
///         worker_token,                  // for graceful shutdown
///         Duration::from_secs(2 * 3600), // 2 hour lock timeout
///     ).await.unwrap();
/// });
///
/// // Later, trigger shutdown
/// shutdown_token.cancel();
/// ```
pub async fn job_worker<AppState: AS>(
    app_state: AppState,
    registry: impl JobRegistry<AppState>,
    sleep_duration: Duration,
    max_retries: i32,
    shutdown_token: CancellationToken,
    lock_timeout: Duration,
) -> color_eyre::Result<()> {
    let worker = Worker::new(
        app_state,
        registry,
        sleep_duration,
        max_retries,
        shutdown_token.clone(),
        lock_timeout,
    );

    loop {
        tokio::select! {
            result = worker.tick() => {
                result?;
            }
            () = shutdown_token.cancelled() => {
                tracing::info!(worker_id = %worker.id, "Job worker shutdown requested");
                break;
            }
        }
    }

    cleanup_worker_locks(&worker).await?;
    tracing::info!(worker_id = %worker.id, "Job worker shutdown complete");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::app_state::AppState;
    use crate::impl_job_registry;
    use crate::jobs::Job;
    use crate::server::cookies::CookieKey;

    #[derive(Clone)]
    struct TestAppState {
        db: sqlx::PgPool,
        cookie_key: CookieKey,
    }

    impl AppState for TestAppState {
        fn db(&self) -> &sqlx::PgPool {
            &self.db
        }

        fn version(&self) -> &'static str {
            "test"
        }

        fn cookie_key(&self) -> &CookieKey {
            &self.cookie_key
        }
    }

    #[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
    struct TestJob {
        id: String,
    }

    #[async_trait::async_trait]
    impl Job<TestAppState> for TestJob {
        const NAME: &'static str = "TestJob";

        async fn run(&self, _app_state: TestAppState) -> color_eyre::Result<()> {
            Ok(())
        }
    }

    impl_job_registry!(TestAppState, TestJob);

    /// Test that `fetch_next_job` picks up a job with a stale lock (lock older than timeout)
    #[sqlx::test]
    async fn test_fetch_next_job_picks_up_stale_locked_job(db: sqlx::PgPool) {
        let app_state = TestAppState {
            db: db.clone(),
            cookie_key: CookieKey::generate(),
        };

        let job_id = uuid::Uuid::new_v4();
        let stale_worker_id = "crashed-worker";

        // Insert a job locked 120 seconds ago
        sqlx::query(
            "INSERT INTO jobs (job_id, name, payload, priority, run_at, created_at, context, error_count, locked_by, locked_at)
             VALUES ($1, $2, $3, $4, NOW(), NOW(), $5, $6, $7, NOW() - interval '120 seconds')",
        )
        .bind(job_id)
        .bind("TestJob")
        .bind(serde_json::json!({"id": "stale-lock-test"}))
        .bind(0)
        .bind("test-stale-lock")
        .bind(0)
        .bind(stale_worker_id)
        .execute(&db)
        .await
        .unwrap();

        // Create a worker with 60 second lock timeout
        let worker = Worker::new(
            app_state,
            Jobs,
            Duration::from_secs(1),
            20,
            CancellationToken::new(),
            Duration::from_secs(60), // 60 second timeout
        );

        // fetch_next_job should pick up the stale locked job
        let fetched = worker.fetch_next_job().await.unwrap();
        assert!(fetched.is_some());
        assert_eq!(fetched.unwrap().job_id, job_id);
    }

    /// Test that `fetch_next_job` does NOT pick up a job with a fresh lock
    #[sqlx::test]
    async fn test_fetch_next_job_skips_recently_locked_job(db: sqlx::PgPool) {
        let app_state = TestAppState {
            db: db.clone(),
            cookie_key: CookieKey::generate(),
        };

        let job_id = uuid::Uuid::new_v4();
        let active_worker_id = "active-worker";

        // Insert a job locked only 10 seconds ago
        sqlx::query(
            "INSERT INTO jobs (job_id, name, payload, priority, run_at, created_at, context, error_count, locked_by, locked_at)
             VALUES ($1, $2, $3, $4, NOW(), NOW(), $5, $6, $7, NOW() - interval '10 seconds')",
        )
        .bind(job_id)
        .bind("TestJob")
        .bind(serde_json::json!({"id": "recent-lock-test"}))
        .bind(0)
        .bind("test-recent-lock")
        .bind(0)
        .bind(active_worker_id)
        .execute(&db)
        .await
        .unwrap();

        // Create a worker with 1 hour lock timeout
        let worker = Worker::new(
            app_state,
            Jobs,
            Duration::from_secs(1),
            20,
            CancellationToken::new(),
            Duration::from_secs(3600), // 1 hour timeout
        );

        // fetch_next_job should NOT pick up the recently locked job
        let fetched = worker.fetch_next_job().await.unwrap();
        assert!(fetched.is_none());
    }

    /// Test that unlocked jobs are picked up before stale locked jobs (by priority)
    #[sqlx::test]
    async fn test_fetch_next_job_prefers_unlocked_by_priority(db: sqlx::PgPool) {
        let app_state = TestAppState {
            db: db.clone(),
            cookie_key: CookieKey::generate(),
        };

        let unlocked_job_id = uuid::Uuid::new_v4();
        let stale_locked_job_id = uuid::Uuid::new_v4();

        // Insert unlocked job with higher priority
        sqlx::query(
            "INSERT INTO jobs (job_id, name, payload, priority, run_at, created_at, context, error_count)
             VALUES ($1, $2, $3, $4, NOW(), NOW(), $5, $6)",
        )
        .bind(unlocked_job_id)
        .bind("TestJob")
        .bind(serde_json::json!({"id": "unlocked"}))
        .bind(10) // Higher priority
        .bind("test-unlocked")
        .bind(0)
        .execute(&db)
        .await
        .unwrap();

        // Insert stale locked job with lower priority
        sqlx::query(
            "INSERT INTO jobs (job_id, name, payload, priority, run_at, created_at, context, error_count, locked_by, locked_at)
             VALUES ($1, $2, $3, $4, NOW(), NOW(), $5, $6, $7, NOW() - interval '120 seconds')",
        )
        .bind(stale_locked_job_id)
        .bind("TestJob")
        .bind(serde_json::json!({"id": "stale-locked"}))
        .bind(5) // Lower priority
        .bind("test-stale")
        .bind(0)
        .bind("crashed-worker")
        .execute(&db)
        .await
        .unwrap();

        // Create a worker with 60 second lock timeout
        let worker = Worker::new(
            app_state,
            Jobs,
            Duration::from_secs(1),
            20,
            CancellationToken::new(),
            Duration::from_secs(60),
        );

        // Should pick the higher priority unlocked job first
        let fetched = worker.fetch_next_job().await.unwrap();
        assert!(fetched.is_some());
        assert_eq!(fetched.unwrap().job_id, unlocked_job_id);
    }
}