services/priority_worker_pool.py

"""
Modular Priority-Tier Worker Pool

A self-contained, plug-and-play worker pool for processing async jobs
with priority-tier scheduling. Can be used in any Python application.

Usage:
    from services.priority_worker_pool import PriorityWorkerPool, WorkerConfig
    
    # Define your job processor function
    async def process_my_job(job, session):
        # Process job and return updated job
        job.status = "completed"
        job.output_data = {"result": "done"}
        return job
    
    # Configure and start pool
    pool = PriorityWorkerPool(
        database_url="sqlite+aiosqlite:///./my_db.db",
        job_model=MyJobModel,
        job_processor=process_my_job,
        config=WorkerConfig(fast_workers=5, medium_workers=5, slow_workers=5)
    )
    await pool.start()

Environment Variables (optional):
    FAST_WORKERS: Number of fast workers (default: 5)
    MEDIUM_WORKERS: Number of medium workers (default: 5)
    SLOW_WORKERS: Number of slow workers (default: 5)
    FAST_INTERVAL: Fast tier polling interval in seconds (default: 5)
    MEDIUM_INTERVAL: Medium tier polling interval in seconds (default: 30)
    SLOW_INTERVAL: Slow tier polling interval in seconds (default: 60)

Dependencies:
    sqlalchemy[asyncio]>=2.0.0
    aiosqlite (for SQLite) or asyncpg (for PostgreSQL)

Job Model Requirements:
    Your job model must have these columns:
    - job_id: str (unique identifier)
    - status: str (queued, processing, completed, failed, cancelled)
    - priority: str (fast, medium, slow)
    - next_process_at: datetime (nullable, for rescheduling)
    - retry_count: int (default 0)
    - created_at: datetime
    - started_at: datetime (nullable)
    - completed_at: datetime (nullable)
    - error_message: str (nullable)
"""
import asyncio
import logging
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Optional, List, Callable, Any, TypeVar, Generic
from sqlalchemy import select, or_, and_
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession

logger = logging.getLogger(__name__)

# Generic type for job model
JobType = TypeVar('JobType')


@dataclass
class WorkerConfig:
    """Configuration for the worker pool."""
    fast_workers: int = 5
    medium_workers: int = 5
    slow_workers: int = 5
    fast_interval: int = 2      # seconds
    medium_interval: int = 10   # seconds
    slow_interval: int = 15     # seconds
    max_retries: int = 60       # Max retry attempts before failing
    job_per_api_key: int = 1    # Max concurrent jobs per API key
    
    @classmethod
    def from_env(cls) -> 'WorkerConfig':
        """Create config from environment variables."""
        return cls(
            fast_workers=int(os.getenv("FAST_WORKERS", "5")),
            medium_workers=int(os.getenv("MEDIUM_WORKERS", "5")),
            slow_workers=int(os.getenv("SLOW_WORKERS", "5")),
            fast_interval=int(os.getenv("FAST_INTERVAL", "5")),
            medium_interval=int(os.getenv("MEDIUM_INTERVAL", "30")),
            slow_interval=int(os.getenv("SLOW_INTERVAL", "60")),
            job_per_api_key=int(os.getenv("JOB_PER_API_KEY", "1")),
        )


@dataclass
class PriorityMapping:
    """Maps job types to priority tiers."""
    mappings: dict = field(default_factory=dict)
    
    def get_priority(self, job_type: str, default: str = "fast") -> str:
        """Get priority for a job type."""
        return self.mappings.get(job_type, default)
    
    def get_interval(self, priority: str, config: WorkerConfig) -> int:
        """Get polling interval for a priority tier."""
        if priority == "fast":
            return config.fast_interval
        elif priority == "medium":
            return config.medium_interval
        else:
            return config.slow_interval


class JobProcessor(ABC, Generic[JobType]):
    """Abstract base class for job processors."""
    
    @abstractmethod
    async def process(self, job: JobType, session: AsyncSession) -> JobType:
        """
        Process a job and return the updated job.
        
        Args:
            job: The job to process
            session: Database session for updates
            
        Returns:
            The updated job with new status/output
        """
        pass
    
    @abstractmethod
    async def check_status(self, job: JobType, session: AsyncSession) -> JobType:
        """
        Check status of an in-progress job (for async third-party operations).
        
        Args:
            job: The job to check
            session: Database session for updates
            
        Returns:
            The updated job. Set next_process_at to reschedule if not done.
        """
        pass


class PriorityWorker(Generic[JobType]):
    """Worker that processes jobs of a specific priority tier."""
    
    def __init__(
        self,
        worker_id: int,
        priority: str,
        poll_interval: int,
        session_maker: async_sessionmaker,
        job_model: type,
        job_processor: JobProcessor[JobType],
        max_retries: int = 60,
        wake_event: Optional[asyncio.Event] = None,
        config: Optional[WorkerConfig] = None
    ):
        self.worker_id = worker_id
        self.priority = priority
        self.poll_interval = poll_interval
        self.session_maker = session_maker
        self.job_model = job_model
        self.job_processor = job_processor
        self.max_retries = max_retries
        self._running = False
        self._current_job_id: Optional[str] = None
        self._wake_event = wake_event  # Event to wake worker immediately when new jobs arrive
        self._config = config or WorkerConfig.from_env()
    
    async def start(self):
        """Start the worker polling loop."""
        self._running = True
        logger.debug(f"Worker {self.worker_id} ({self.priority}) started, polling every {self.poll_interval}s")
        asyncio.create_task(self._poll_loop())
    
    async def stop(self):
        """Stop the worker."""
        self._running = False
        logger.info(f"Worker {self.worker_id} ({self.priority}) stopped")
    
    async def _poll_loop(self):
        """Main polling loop with optimized scheduling.
        
        Optimizations:
        - When no jobs are found, sleep for poll_interval before checking again
        - When a job is processed, immediately check for the next job (no waiting)
        - This ensures first job starts immediately when queue was empty
        - This ensures next job starts immediately after current job finishes
        """
        while self._running:
            job_found = False
            try:
                job_found = await self._process_one_job()
            except Exception as e:
                logger.error(f"Worker {self.worker_id}: Error in poll loop: {e}")
            
            # Only sleep if no job was found - otherwise immediately look for next job
            if not job_found:
                # Wait on event with timeout - allows immediate wake-up when new job arrives
                if self._wake_event:
                    try:
                        # Wait for event or timeout (whichever comes first)
                        await asyncio.wait_for(
                            self._wake_event.wait(),
                            timeout=self.poll_interval
                        )
                        # Clear event after waking (we'll check for jobs)
                        self._wake_event.clear()
                    except asyncio.TimeoutError:
                        pass  # Normal timeout, check for jobs
                else:
                    await asyncio.sleep(self.poll_interval)
    
    async def _process_one_job(self) -> bool:
        """Find and process one job.
        
        Enforces constraints:
        1. Only one job per user can be in processing state at a time
        2. Total processing jobs limited to JOB_PER_API_KEY * number of API keys
        
        Returns:
            True if a job was found and processed, False if no jobs available
        """
        async with self.session_maker() as session:
            from sqlalchemy import func
            
            now = datetime.utcnow()
            
            # Get number of API keys for capacity calculation
            try:
                from services.api_key_manager import get_key_count
                num_api_keys = get_key_count()
                max_processing = self._config.job_per_api_key * num_api_keys
            except ImportError:
                max_processing = 10  # Default fallback
            
            # Check if we're at max processing capacity (only for new jobs being picked up)
            count_query = select(func.count()).where(
                self.job_model.status == "processing"
            )
            count_result = await session.execute(count_query)
            current_processing = count_result.scalar() or 0
            
            # Query for jobs matching this priority tier
            query = select(self.job_model).where(
                and_(
                    self.job_model.priority == self.priority,
                    self.job_model.status.in_(["queued", "processing"]),
                    or_(
                        self.job_model.next_process_at.is_(None),
                        self.job_model.next_process_at <= now
                    )
                )
            ).order_by(self.job_model.created_at).limit(1)
            
            result = await session.execute(query)
            job = result.scalar_one_or_none()
            
            if not job:
                return False
            
            # For queued jobs, apply the constraints
            if job.status == "queued":
                # Constraint 1: Check if this user already has a job in processing
                user_processing_query = select(func.count()).where(
                    and_(
                        self.job_model.user_id == job.user_id,
                        self.job_model.status == "processing"
                    )
                )
                user_result = await session.execute(user_processing_query)
                user_processing_count = user_result.scalar() or 0
                
                if user_processing_count > 0:
                    logger.debug(f"Worker {self.worker_id}: User {job.user_id} already has a job processing, skipping")
                    return False
                
                # Constraint 2: Check if we're at max total processing capacity
                if current_processing >= max_processing:
                    logger.debug(f"Worker {self.worker_id}: At max capacity ({current_processing}/{max_processing}), skipping new job")
                    return False
            
            self._current_job_id = job.job_id
            
            try:
                await self._process_job(session, job)
                return True
            except Exception as e:
                logger.error(f"Worker {self.worker_id}: Error processing job {job.job_id}: {e}")
                job.status = "failed"
                job.error_message = str(e)
                job.completed_at = datetime.utcnow()
                await session.commit()
                return True  # Job was found, even though it failed
            finally:
                self._current_job_id = None
    
    async def _process_job(self, session: AsyncSession, job: JobType):
        """Process a single job."""
        logger.info(f"Worker {self.worker_id}: Processing job {job.job_id} (status: {job.status})")
        
        from sqlalchemy import update
        
        if job.status == "queued":
            # New job - try to claim it atomically
            # Set next_process_at to future to prevent others from picking it up while we process
            next_check = datetime.utcnow() + timedelta(seconds=self.poll_interval * 2)
            
            stmt = (
                update(self.job_model)
                .where(
                    self.job_model.job_id == job.job_id,
                    self.job_model.status == "queued"
                )
                .values(
                    status="processing",
                    started_at=datetime.utcnow(),
                    next_process_at=next_check
                )
            )
            result = await session.execute(stmt)
            await session.commit()
            
            if result.rowcount == 0:
                logger.info(f"Worker {self.worker_id}: Failed to claim job {job.job_id} (already taken)")
                return

            # We claimed it. Refresh and process.
            await session.refresh(job)
            job = await self.job_processor.process(job, session)
            
        else:
            # Already processing - try to claim for status check
            # Ensure we only pick it up if next_process_at matches (or is null/past)
            # But the SELECT already filtered for that.
            # We just need to ensure no one else grabbed it between SELECT and UPDATE.
            
            # Update next_process_at to future to lock it for this check
            next_check = datetime.utcnow() + timedelta(seconds=self.poll_interval * 2)
            
            stmt = (
                update(self.job_model)
                .where(
                    self.job_model.job_id == job.job_id,
                    or_(
                        self.job_model.next_process_at.is_(None),
                        self.job_model.next_process_at <= datetime.utcnow()
                    )
                )
                .values(next_process_at=next_check)
            )
            result = await session.execute(stmt)
            await session.commit()
            
            if result.rowcount == 0:
                logger.info(f"Worker {self.worker_id}: Failed to claim job {job.job_id} for check (already taken)")
                return
                
            await session.refresh(job)
            job = await self.job_processor.check_status(job, session)
        
        # Handle retry limit
        if job.status == "processing" and job.retry_count > self.max_retries:
            job.status = "failed"
            job.error_message = f"Max retries ({self.max_retries}) exceeded"
            job.completed_at = datetime.utcnow()
        
        # Handle credit finalization for jobs with reserved credits
        if job.status in ("completed", "failed", "cancelled"):
            await self._handle_job_credits(session, job)
        
        await session.commit()
    
    async def _handle_job_credits(self, session: AsyncSession, job: JobType):
        """Handle credit finalization when job reaches terminal state."""
        # Check if job has credits_reserved attribute (credit-enabled jobs)
        if not hasattr(job, 'credits_reserved') or job.credits_reserved <= 0:
            return
        
        try:
            from services.credit_service.credit_manager import handle_job_completion
            await handle_job_completion(session, job)
        except ImportError:
            # Credit service not available - skip
            logger.debug(f"Credit service not available for job {job.job_id}")
        except Exception as e:
            logger.error(f"Error handling credits for job {job.job_id}: {e}")


class PriorityWorkerPool(Generic[JobType]):
    """
    Modular priority-tier worker pool.
    
    Can be used with any job model that follows the required schema.
    """
    
    def __init__(
        self,
        database_url: str,
        job_model: type,
        job_processor: JobProcessor[JobType],
        config: Optional[WorkerConfig] = None
    ):
        """
        Initialize the worker pool.
        
        Args:
            database_url: SQLAlchemy async database URL
            job_model: Your ORM model class for jobs
            job_processor: Instance of JobProcessor to handle jobs
            config: Worker configuration (uses env vars if not provided)
        """
        self.database_url = database_url
        self.job_model = job_model
        self.job_processor = job_processor
        self.config = config or WorkerConfig.from_env()
        
        self.engine = create_async_engine(database_url, echo=False)
        self.session_maker = async_sessionmaker(
            self.engine,
            class_=AsyncSession,
            expire_on_commit=False
        )
        self.workers: List[PriorityWorker] = []
        self._running = False
        
        # Wake events for each priority tier - allows immediate job notification
        self._wake_events: dict[str, asyncio.Event] = {
            "fast": asyncio.Event(),
            "medium": asyncio.Event(),
            "slow": asyncio.Event()
        }
    
    async def start(self):
        """Start all workers."""
        self._running = True
        worker_id = 0
        
        # Create fast workers
        for i in range(self.config.fast_workers):
            worker = PriorityWorker(
                worker_id=worker_id,
                priority="fast",
                poll_interval=self.config.fast_interval,
                session_maker=self.session_maker,
                job_model=self.job_model,
                job_processor=self.job_processor,
                max_retries=self.config.max_retries,
                wake_event=self._wake_events["fast"],
                config=self.config
            )
            self.workers.append(worker)
            await worker.start()
            worker_id += 1
        
        # Create medium workers
        for i in range(self.config.medium_workers):
            worker = PriorityWorker(
                worker_id=worker_id,
                priority="medium",
                poll_interval=self.config.medium_interval,
                session_maker=self.session_maker,
                job_model=self.job_model,
                job_processor=self.job_processor,
                max_retries=self.config.max_retries,
                wake_event=self._wake_events["medium"],
                config=self.config
            )
            self.workers.append(worker)
            await worker.start()
            worker_id += 1
        
        # Create slow workers
        for i in range(self.config.slow_workers):
            worker = PriorityWorker(
                worker_id=worker_id,
                priority="slow",
                poll_interval=self.config.slow_interval,
                session_maker=self.session_maker,
                job_model=self.job_model,
                job_processor=self.job_processor,
                max_retries=self.config.max_retries,
                wake_event=self._wake_events["slow"],
                config=self.config
            )
            self.workers.append(worker)
            await worker.start()
            worker_id += 1
        
        total = self.config.fast_workers + self.config.medium_workers + self.config.slow_workers
        logger.info(
            f"PriorityWorkerPool started with {total} workers: "
            f"{self.config.fast_workers} fast, {self.config.medium_workers} medium, {self.config.slow_workers} slow"
        )
    
    def notify_new_job(self, priority: str):
        """
        Wake sleeping workers of the specified priority tier.
        Call this when a new job is created to start processing immediately.
        
        Args:
            priority: Priority tier ("fast", "medium", or "slow")
        """
        if priority in self._wake_events:
            self._wake_events[priority].set()
            logger.debug(f"Notified {priority} workers of new job")
    
    async def stop(self):
        """Stop all workers and refund orphaned jobs."""
        self._running = False
        
        # Refund credits for any jobs that were processing when server stopped
        await self._refund_orphaned_jobs()
        
        for worker in self.workers:
            await worker.stop()
        logger.info("PriorityWorkerPool stopped")
    
    async def _refund_orphaned_jobs(self):
        """Refund credits for jobs abandoned during shutdown."""
        try:
            from services.credit_service.credit_manager import refund_orphaned_jobs
            async with self.session_maker() as session:
                refund_count = await refund_orphaned_jobs(session)
                if refund_count > 0:
                    logger.info(f"Shutdown: Refunded {refund_count} orphaned job(s)")
        except ImportError:
            logger.debug("Credit service not available for orphaned job refunds")
        except Exception as e:
            logger.error(f"Error refunding orphaned jobs during shutdown: {e}")


# Convenience functions for priority mapping
def get_priority_for_job_type(job_type: str, mappings: dict) -> str:
    """Get priority tier for a job type using provided mappings."""
    return mappings.get(job_type, "fast")


def get_interval_for_priority(priority: str, config: Optional[WorkerConfig] = None) -> int:
    """Get polling interval for a priority tier."""
    cfg = config or WorkerConfig.from_env()
    if priority == "fast":
        return cfg.fast_interval
    elif priority == "medium":
        return cfg.medium_interval
    else:
        return cfg.slow_interval