pleme_redis/

worker_pool.rs

//! Generic worker pool for processing Redis-backed job queues
//!
//! This module provides a generic worker pool that can process any job type
//! that implements the `Job` trait. Services provide their own `JobExecutor`
//! and `JobStore` implementations for business logic and persistence.

use crate::job_queue::{JobQueue, QueueError};
use chrono::Utc;
use std::sync::Arc;
use std::time::Duration;
use thiserror::Error;
use tokio::task::JoinHandle;
use uuid::Uuid;

/// Job executor trait - services implement this for business logic
///
/// # Example
/// ```rust
/// use pleme_redis::worker_pool::JobExecutor;
/// use async_trait::async_trait;
///
/// struct MyJobExecutor;
///
/// #[async_trait]
/// impl JobExecutor<MyJob> for MyJobExecutor {
///     async fn execute(&self, job: &MyJob) -> Result<serde_json::Value, ExecutorError> {
///         // Business logic here
///         Ok(serde_json::json!({"status": "success"}))
///     }
/// }
/// ```
#[async_trait::async_trait]
pub trait JobExecutor<J>: Send + Sync
where
    J: Send + Sync,
{
    /// Execute a job and return results
    ///
    /// # Arguments
    /// * `job` - Job to execute
    ///
    /// # Returns
    /// * `Ok(Value)` - Job results (serializable JSON)
    /// * `Err(ExecutorError)` - Execution failed
    async fn execute(&self, job: &J) -> Result<serde_json::Value, ExecutorError>;
}

/// Job store trait - services implement this for database persistence
///
/// This trait abstracts database operations so the worker pool can work
/// with any persistence layer (PostgreSQL, MongoDB, etc.).
#[async_trait::async_trait]
pub trait JobStore<J>: Send + Sync
where
    J: Send + Sync,
{
    /// Atomically claim next available job
    ///
    /// Should use SELECT FOR UPDATE SKIP LOCKED or equivalent
    /// to ensure only one worker claims a job.
    ///
    /// # Arguments
    /// * `worker_id` - Worker identifier for tracking
    ///
    /// # Returns
    /// * `Some(Job)` - Job successfully claimed
    /// * `None` - No jobs available
    async fn claim_next_job(&self, worker_id: &str) -> Result<Option<J>, StoreError>;

    /// Load job by ID
    async fn load_job(&self, job_id: Uuid) -> Result<J, StoreError>;

    /// Check if job was soft-deleted (for cancellation support)
    async fn is_job_deleted(&self, job_id: Uuid) -> Result<bool, StoreError>;

    /// Mark job as running
    async fn mark_job_running(&self, job_id: Uuid) -> Result<(), StoreError>;

    /// Mark job as completed with results
    async fn mark_job_completed(&self, job_id: Uuid, results: serde_json::Value) -> Result<(), StoreError>;

    /// Schedule job retry with exponential backoff
    async fn schedule_retry(
        &self,
        job_id: Uuid,
        retry_count: i32,
        next_retry_at: chrono::DateTime<Utc>,
        error_message: &str,
        error_details: serde_json::Value,
    ) -> Result<(), StoreError>;

    /// Mark job as permanently failed
    async fn mark_job_failed(
        &self,
        job_id: Uuid,
        error_message: &str,
        error_details: serde_json::Value,
    ) -> Result<(), StoreError>;

    /// Find orphaned jobs (PENDING/QUEUED without worker)
    ///
    /// Used for startup recovery
    async fn find_orphaned_jobs(&self) -> Result<Vec<(Uuid, String, i32, i32)>, StoreError>;

    /// Find stale jobs (RUNNING without updates)
    ///
    /// Used by stale job checker
    async fn find_stale_jobs(&self, threshold_seconds: i64) -> Result<Vec<J>, StoreError>;

    /// Count claimable jobs (for monitoring)
    async fn count_claimable_jobs(&self) -> Result<i64, StoreError>;
}

/// Worker pool configuration
#[derive(Debug, Clone)]
pub struct WorkerPoolConfig {
    /// Number of concurrent workers
    pub worker_count: usize,

    /// Redis queue key
    pub queue_key: String,

    /// Worker poll timeout (seconds)
    pub poll_timeout: u64,

    /// Maximum retries per job (default: 3)
    pub max_retries: i32,

    /// Base delay for exponential backoff (seconds, default: 60)
    pub base_retry_delay: i64,

    /// Maximum delay for exponential backoff (seconds, default: 3600)
    pub max_retry_delay: i64,

    /// Threshold for detecting stale RUNNING jobs in seconds (default: 7200 = 2 hours)
    pub stale_job_threshold_seconds: i64,

    /// How often to check for stale jobs in seconds (default: 300 = 5 minutes)
    pub stale_check_interval_seconds: u64,
}

impl Default for WorkerPoolConfig {
    fn default() -> Self {
        Self {
            worker_count: 5,
            queue_key: "jobs:queue".to_string(),
            poll_timeout: 30,
            max_retries: 3,
            base_retry_delay: 60,
            max_retry_delay: 3600,
            stale_job_threshold_seconds: 7200,
            stale_check_interval_seconds: 300,
        }
    }
}

/// Generic worker pool for processing jobs
pub struct WorkerPool<J, E, S>
where
    J: Send + Sync + 'static,
    E: JobExecutor<J> + 'static,
    S: JobStore<J> + 'static,
{
    config: WorkerPoolConfig,
    redis_queue: Arc<tokio::sync::Mutex<JobQueue>>,
    executor: Arc<E>,
    store: Arc<S>,
    workers: Vec<JoinHandle<()>>,
    _phantom: std::marker::PhantomData<J>,
}

impl<J, E, S> WorkerPool<J, E, S>
where
    J: Send + Sync + 'static + crate::job::Job,
    E: JobExecutor<J> + 'static,
    S: JobStore<J> + 'static,
{
    /// Create new worker pool
    pub fn new(
        config: WorkerPoolConfig,
        redis_queue: JobQueue,
        executor: E,
        store: S,
    ) -> Self {
        Self {
            config,
            redis_queue: Arc::new(tokio::sync::Mutex::new(redis_queue)),
            executor: Arc::new(executor),
            store: Arc::new(store),
            workers: Vec::new(),
            _phantom: std::marker::PhantomData,
        }
    }

    /// Start all workers
    ///
    /// Spawns `worker_count` concurrent workers that poll the queue and process jobs.
    /// Also spawns a dedicated stale job checker that runs periodically.
    pub async fn start(&mut self) -> Result<(), WorkerError> {
        tracing::info!(
            "🚀 [WORKER_POOL] Starting worker pool with {} workers + 1 stale job checker",
            self.config.worker_count
        );

        // CRITICAL: Recover orphaned jobs from database before starting workers
        self.recover_orphaned_jobs().await?;

        // Spawn job processing workers
        for worker_id in 0..self.config.worker_count {
            let worker = Worker::new(
                worker_id,
                self.config.clone(),
                self.redis_queue.clone(),
                self.executor.clone(),
                self.store.clone(),
            );

            let handle = tokio::spawn(async move {
                if let Err(e) = worker.run().await {
                    tracing::error!("❌ [WORKER_{}] Worker failed: {}", worker_id, e);
                }
            });

            self.workers.push(handle);
        }

        // Spawn dedicated stale job checker
        let checker = StaleJobChecker::new(
            self.config.clone(),
            self.redis_queue.clone(),
            self.store.clone(),
        );

        let handle = tokio::spawn(async move {
            if let Err(e) = checker.run().await {
                tracing::error!("❌ [STALE_CHECKER] Stale job checker failed: {}", e);
            }
        });

        self.workers.push(handle);

        tracing::info!(
            "✅ [WORKER_POOL] All workers started successfully - stale_threshold: {}s",
            self.config.stale_job_threshold_seconds
        );

        Ok(())
    }

    /// Recover orphaned jobs from database and re-enqueue to Redis
    async fn recover_orphaned_jobs(&mut self) -> Result<(), WorkerError> {
        tracing::info!("🔍 [RECOVERY] Scanning database for orphaned jobs to recover...");

        let orphaned_jobs = self
            .store
            .find_orphaned_jobs()
            .await
            .map_err(|e| WorkerError::StoreError(e.to_string()))?;

        if orphaned_jobs.is_empty() {
            tracing::info!("✅ [RECOVERY] No orphaned jobs found - queue is clean");
            return Ok(());
        }

        tracing::info!(
            "📋 [RECOVERY] Found {} orphaned jobs to recover",
            orphaned_jobs.len()
        );

        let mut recovered = 0;
        let mut failed_recovery = 0;

        for (job_id, job_type, retry_count, max_retries) in orphaned_jobs {
            tracing::debug!(
                "🔄 [RECOVERY] Processing job {} (type: {}, retry: {}/{})",
                job_id,
                job_type,
                retry_count,
                max_retries
            );

            match self.redis_queue.lock().await.enqueue(job_id, job_type.clone()).await {
                Ok(_) => {
                    recovered += 1;
                    tracing::info!(
                        "✅ [RECOVERY] Re-enqueued job {} to Redis",
                        job_id
                    );
                }
                Err(e) => {
                    failed_recovery += 1;
                    tracing::error!(
                        "❌ [RECOVERY] Failed to re-enqueue job {}: {}",
                        job_id,
                        e
                    );
                }
            }
        }

        tracing::info!(
            "✅ [RECOVERY] Recovery complete - recovered: {}, failed: {}",
            recovered,
            failed_recovery
        );

        Ok(())
    }

    /// Stop all workers
    pub async fn stop(&mut self) {
        tracing::info!("Stopping worker pool");

        for handle in self.workers.drain(..) {
            handle.abort();
        }
    }

    /// Wait for all workers to complete
    pub async fn wait(self) -> Result<(), WorkerError> {
        for handle in self.workers {
            handle
                .await
                .map_err(|e| WorkerError::WorkerFailed(e.to_string()))?;
        }
        Ok(())
    }
}

/// Individual worker
struct Worker<J, E, S>
where
    J: Send + Sync,
    E: JobExecutor<J>,
    S: JobStore<J>,
{
    id: usize,
    config: WorkerPoolConfig,
    redis_queue: Arc<tokio::sync::Mutex<JobQueue>>,
    executor: Arc<E>,
    store: Arc<S>,
    current_backoff_ms: u64,
    max_backoff_ms: u64,
    consecutive_lock_failures: u32,
    _phantom: std::marker::PhantomData<J>,
}

impl<J, E, S> Worker<J, E, S>
where
    J: Send + Sync + crate::job::Job,
    E: JobExecutor<J>,
    S: JobStore<J>,
{
    fn new(
        id: usize,
        config: WorkerPoolConfig,
        redis_queue: Arc<tokio::sync::Mutex<JobQueue>>,
        executor: Arc<E>,
        store: Arc<S>,
    ) -> Self {
        Self {
            id,
            config,
            redis_queue,
            executor,
            store,
            current_backoff_ms: 100,
            max_backoff_ms: 60_000,
            consecutive_lock_failures: 0,
            _phantom: std::marker::PhantomData,
        }
    }

    async fn run(mut self) -> Result<(), WorkerError> {
        let worker_id = format!("worker_{}", self.id);
        tracing::info!(
            "🟢 [WORKER_{}] Started with ID '{}' - atomic database-driven job claiming",
            self.id,
            worker_id
        );

        loop {
            match self.store.claim_next_job(&worker_id).await {
                Ok(Some(job)) => {
                    let job_id = job.id();
                    let job_type_str = job.job_type();

                    tracing::info!(
                        "📦 [WORKER_{}] Atomically claimed job {} (type: {}, retry: {}/{})",
                        self.id,
                        job_id,
                        job_type_str,
                        job.retry_count(),
                        job.max_retries()
                    );

                    // Get lock TTL from configuration
                    let lock_ttl_seconds = job
                        .configuration()
                        .and_then(|config| config.get("redis_lock_ttl_seconds"))
                        .and_then(|v| v.as_u64())
                        .unwrap_or(172800); // Default: 48 hours

                    // Check distributed lock
                    let mut queue = self.redis_queue.lock().await;
                    let lock_acquired = match queue
                        .acquire_job_type_lock(&job_type_str, job.parameters(), lock_ttl_seconds)
                        .await
                    {
                        Ok(acquired) => acquired,
                        Err(e) => {
                            tracing::error!(
                                "❌ [WORKER_{}] Failed to check lock for job {}: {}",
                                self.id,
                                job_id,
                                e
                            );
                            true // Fail-open
                        }
                    };
                    drop(queue);

                    if !lock_acquired {
                        self.consecutive_lock_failures += 1;
                        tracing::warn!(
                            "🔒 [WORKER_{}] Job {} skipped - lock held",
                            self.id,
                            job_id
                        );

                        // Apply exponential backoff
                        tokio::time::sleep(Duration::from_millis(self.current_backoff_ms)).await;
                        self.current_backoff_ms = std::cmp::min(self.current_backoff_ms * 2, self.max_backoff_ms);
                        continue;
                    }

                    // Lock acquired - reset backoff
                    if self.consecutive_lock_failures > 0 {
                        self.consecutive_lock_failures = 0;
                        self.current_backoff_ms = 100;
                    }

                    // Process job with retry logic
                    if let Err(e) = self.process_job_with_retry(job).await {
                        tracing::error!(
                            "❌ [WORKER_{}] Failed to process job {}: {}",
                            self.id,
                            job_id,
                            e
                        );
                    }
                }
                Ok(None) => {
                    // No jobs available - sleep and check for orphaned jobs
                    tracing::trace!(
                        "⏱️  [WORKER_{}] No claimable jobs - sleeping 5s",
                        self.id
                    );
                    tokio::time::sleep(Duration::from_secs(5)).await;

                    if let Err(e) = self.check_orphaned_jobs().await {
                        tracing::error!(
                            "❌ [WORKER_{}] Failed to check orphaned jobs: {}",
                            self.id,
                            e
                        );
                    }
                }
                Err(e) => {
                    tracing::error!("❌ [WORKER_{}] Failed to claim job: {}", self.id, e);
                    tokio::time::sleep(Duration::from_secs(1)).await;
                }
            }
        }
    }

    async fn process_job_with_retry(&self, job: J) -> Result<(), WorkerError> {
        let job_id = job.id();

        tracing::info!(
            "🎬 [WORKER_{}] Starting job {} (retry: {}/{})",
            self.id,
            job_id,
            job.retry_count(),
            job.max_retries()
        );

        // Mark as running
        self.store
            .mark_job_running(job_id)
            .await
            .map_err(|e| WorkerError::StoreError(e.to_string()))?;

        // Execute job
        match self.executor.execute(&job).await {
            Ok(results) => {
                // Success
                self.store
                    .mark_job_completed(job_id, results)
                    .await
                    .map_err(|e| WorkerError::StoreError(e.to_string()))?;

                tracing::info!(
                    "✅ [WORKER_{}] Job {} completed successfully",
                    self.id,
                    job_id
                );
            }
            Err(e) => {
                // Execution failed
                let error_message = e.to_string();

                if job.can_retry() {
                    // Schedule retry
                    let next_retry_at = job.calculate_next_retry_at();
                    let error_details = self.build_error_details(&error_message, "EXECUTION_FAILED", &job);

                    self.store
                        .schedule_retry(
                            job_id,
                            job.retry_count() + 1,
                            next_retry_at,
                            &error_message,
                            error_details,
                        )
                        .await
                        .map_err(|e| WorkerError::StoreError(e.to_string()))?;

                    tracing::warn!(
                        "⚠️  [WORKER_{}] Job {} failed (attempt {}/{}), will retry at {}",
                        self.id,
                        job_id,
                        job.retry_count() + 1,
                        job.max_retries(),
                        next_retry_at
                    );
                } else {
                    // Max retries exceeded
                    let error_details = self.build_error_details(
                        &error_message,
                        "MAX_RETRIES_EXCEEDED",
                        &job,
                    );

                    self.store
                        .mark_job_failed(job_id, &error_message, error_details)
                        .await
                        .map_err(|e| WorkerError::StoreError(e.to_string()))?;

                    tracing::error!(
                        "❌ [WORKER_{}] Job {} PERMANENTLY FAILED after {} attempts",
                        self.id,
                        job_id,
                        job.retry_count()
                    );
                }
            }
        }

        Ok(())
    }

    fn build_error_details(&self, error_message: &str, error_code: &str, job: &J) -> serde_json::Value {
        serde_json::json!({
            "error_code": error_code,
            "error_message": error_message,
            "worker_id": self.id,
            "job_type": job.job_type(),
            "retry_count": job.retry_count(),
            "max_retries": job.max_retries(),
            "timestamp": Utc::now().to_rfc3339(),
        })
    }

    async fn check_orphaned_jobs(&self) -> Result<(), WorkerError> {
        let count = self
            .store
            .count_claimable_jobs()
            .await
            .map_err(|e| WorkerError::StoreError(e.to_string()))?;

        if count > 0 {
            tracing::info!(
                "🔄 [WORKER_{}] Found {} claimable jobs",
                self.id,
                count
            );
        }

        Ok(())
    }
}

/// Dedicated stale job checker
struct StaleJobChecker<J, S>
where
    J: Send + Sync,
    S: JobStore<J>,
{
    config: WorkerPoolConfig,
    redis_queue: Arc<tokio::sync::Mutex<JobQueue>>,
    store: Arc<S>,
    _phantom: std::marker::PhantomData<J>,
}

impl<J, S> StaleJobChecker<J, S>
where
    J: Send + Sync + crate::job::Job,
    S: JobStore<J>,
{
    fn new(
        config: WorkerPoolConfig,
        redis_queue: Arc<tokio::sync::Mutex<JobQueue>>,
        store: Arc<S>,
    ) -> Self {
        Self {
            config,
            redis_queue,
            store,
            _phantom: std::marker::PhantomData,
        }
    }

    async fn run(self) -> Result<(), WorkerError> {
        tracing::info!(
            "🔍 [STALE_CHECKER] Started - checking every {}s for jobs running longer than {}s",
            self.config.stale_check_interval_seconds,
            self.config.stale_job_threshold_seconds
        );

        let check_interval = Duration::from_secs(self.config.stale_check_interval_seconds);

        loop {
            tokio::time::sleep(check_interval).await;

            if let Err(e) = self.check_and_recover_stale_jobs().await {
                tracing::error!("❌ [STALE_CHECKER] Failed to check stale jobs: {}", e);
            }
        }
    }

    async fn check_and_recover_stale_jobs(&self) -> Result<(), WorkerError> {
        let stale_jobs = self
            .store
            .find_stale_jobs(self.config.stale_job_threshold_seconds)
            .await
            .map_err(|e| WorkerError::StoreError(e.to_string()))?;

        if stale_jobs.is_empty() {
            tracing::debug!("✓ [STALE_CHECKER] No stale jobs found");
            return Ok(());
        }

        tracing::warn!(
            "⚠️  [STALE_CHECKER] Found {} stale jobs - recovering them",
            stale_jobs.len()
        );

        for job in stale_jobs {
            let job_id = job.id();

            if job.can_retry() {
                // Schedule retry
                let next_retry_at = job.calculate_next_retry_at();
                let error_message = format!(
                    "Job was stuck in RUNNING state - likely due to worker crash"
                );
                let error_details = serde_json::json!({
                    "error_code": "STALE_JOB_RECOVERY",
                    "reason": "Job automatically recovered from RUNNING state",
                    "threshold_seconds": self.config.stale_job_threshold_seconds,
                });

                self.store
                    .schedule_retry(
                        job_id,
                        job.retry_count() + 1,
                        next_retry_at,
                        &error_message,
                        error_details,
                    )
                    .await
                    .map_err(|e| WorkerError::StoreError(e.to_string()))?;

                tracing::info!(
                    "✅ [STALE_CHECKER] Job {} scheduled for retry",
                    job_id
                );
            } else {
                // Max retries exceeded
                let error_message = format!(
                    "Job failed permanently: stuck in RUNNING state after {} retries",
                    job.retry_count()
                );
                let error_details = serde_json::json!({
                    "error_code": "STALE_JOB_MAX_RETRIES_EXCEEDED",
                    "reason": "Job stuck in RUNNING state and exhausted all retries",
                });

                self.store
                    .mark_job_failed(job_id, &error_message, error_details)
                    .await
                    .map_err(|e| WorkerError::StoreError(e.to_string()))?;

                tracing::error!(
                    "🔴 [STALE_CHECKER] Job {} marked as PERMANENTLY FAILED",
                    job_id
                );
            }
        }

        Ok(())
    }
}

/// Worker errors
#[derive(Error, Debug)]
pub enum WorkerError {
    #[error("Store error: {0}")]
    StoreError(String),

    #[error("Queue error: {0}")]
    QueueError(String),

    #[error("Worker failed: {0}")]
    WorkerFailed(String),

    #[error("Execution error: {0}")]
    ExecutionError(String),
}

impl From<QueueError> for WorkerError {
    fn from(err: QueueError) -> Self {
        WorkerError::QueueError(err.to_string())
    }
}

/// Store errors
#[derive(Error, Debug)]
pub enum StoreError {
    #[error("Database error: {0}")]
    DatabaseError(String),

    #[error("Job not found: {0}")]
    JobNotFound(Uuid),

    #[error("Serialization error: {0}")]
    SerializationError(String),
}

/// Executor errors
#[derive(Error, Debug)]
pub enum ExecutorError {
    #[error("Execution failed: {0}")]
    ExecutionFailed(String),

    #[error("Job was cancelled")]
    JobCancelled,

    #[error("Provider error: {0}")]
    ProviderError(String),

    #[error("Validation error: {0}")]
    ValidationError(String),
}

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

    #[test]
    fn test_worker_pool_config_defaults() {
        let config = WorkerPoolConfig::default();

        assert_eq!(config.worker_count, 5);
        assert_eq!(config.queue_key, "jobs:queue");
        assert_eq!(config.poll_timeout, 30);
        assert_eq!(config.max_retries, 3);
        assert_eq!(config.base_retry_delay, 60);
        assert_eq!(config.max_retry_delay, 3600);
        assert_eq!(config.stale_job_threshold_seconds, 7200);
        assert_eq!(config.stale_check_interval_seconds, 300);
    }
}