queue_manager.py

import asyncio
import logging
from typing import Optional
from models import Job

logger = logging.getLogger(__name__)

class JobQueue:
    def __init__(self, max_size: int = 1000):
        self.queue = asyncio.PriorityQueue(maxsize=max_size)
        self.paused = asyncio.Event()
        self.paused.set() # 預設為啟動狀態 (未暫停)
        
    async def push(self, job: Job):
        """
        將任務推入隊列。
        如果隊列已滿，拋出 asyncio.QueueFull (由呼叫者處理以回傳 503)。
        """
        try:
            self.queue.put_nowait(job)
            logger.info(f"Job {job.submission_id} added to queue. Priority: {job.priority}, Queue size: {self.queue.qsize()}")
        except asyncio.QueueFull:
            logger.warning(f"Job queue is full! Rejected submission {job.submission_id}")
            raise

    async def pop(self) -> Job:
        """
        從隊列取出任務。
        這是一個阻塞操作，直到有任務可用且隊列未暫停。
        """
        # 等待暫停解除
        await self.paused.wait()
        
        job = await self.queue.get()
        logger.info(f"Job {job.submission_id} popped from queue. Priority: {job.priority}")
        return job

    def pause(self):
        """暫停隊列處理 (Worker 將無法獲取新任務)"""
        self.paused.clear()
        logger.info("Queue paused")

    def resume(self):
        """恢復隊列處理"""
        self.paused.set()
        logger.info("Queue resumed")

    def clear(self) -> int:
        """清空隊列，回傳被移除的任務數量"""
        count = 0
        try:
            while not self.queue.empty():
                self.queue.get_nowait()
                count += 1
        except asyncio.QueueEmpty:
            pass
        logger.info(f"Queue cleared. Removed {count} jobs.")
        return count

    def is_paused(self) -> bool:
        return not self.paused.is_set()

    def size(self) -> int:
        return self.queue.qsize()

    def empty(self) -> bool:
        return self.queue.empty()

# 全域單例
# 將 Queue 大小調大以容納大量批次測試 (100 -> 20000)
job_queue = JobQueue(max_size=20000)