148 lines
5.2 KiB
Python
148 lines
5.2 KiB
Python
"""
|
|
Task Queue Manager for sequential execution of scheduled downloads.
|
|
Ensures only one task runs at a time.
|
|
"""
|
|
|
|
import asyncio
|
|
from typing import Optional, Dict
|
|
from datetime import datetime
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TaskQueue:
|
|
"""
|
|
Manages sequential execution of scheduled download tasks.
|
|
Ensures only one task runs at a time.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.queue: asyncio.Queue = asyncio.Queue()
|
|
self.current_task: Optional[str] = None # Current task_id being executed
|
|
self.is_processing: bool = False
|
|
self.worker_task: Optional[asyncio.Task] = None
|
|
self._execute_callback = None
|
|
|
|
def set_execute_callback(self, callback):
|
|
"""
|
|
Set the callback function to execute tasks.
|
|
|
|
Args:
|
|
callback: Async function that takes task_id and executes it
|
|
"""
|
|
self._execute_callback = callback
|
|
|
|
async def add_task(self, task_id: str, priority: int = 0):
|
|
"""
|
|
Add a task to the queue.
|
|
|
|
Args:
|
|
task_id: UUID of the scheduled task
|
|
priority: 0 = scheduled (normal), 1 = manual "Run Now" (higher priority)
|
|
"""
|
|
task_info = {
|
|
'task_id': task_id,
|
|
'priority': priority,
|
|
'queued_at': datetime.now()
|
|
}
|
|
|
|
# For priority tasks, we need to reorder the queue
|
|
if priority > 0:
|
|
# Get all items from queue
|
|
items = []
|
|
while not self.queue.empty():
|
|
try:
|
|
items.append(await asyncio.wait_for(self.queue.get(), timeout=0.1))
|
|
except asyncio.TimeoutError:
|
|
break
|
|
|
|
# Add new priority task
|
|
await self.queue.put(task_info)
|
|
|
|
# Re-add other items
|
|
for item in items:
|
|
await self.queue.put(item)
|
|
|
|
logger.info(f"Priority task {task_id} added to front of queue (queue_size={self.queue.qsize()})")
|
|
else:
|
|
await self.queue.put(task_info)
|
|
logger.info(f"Task {task_id} added to queue (priority={priority}, queue_size={self.queue.qsize()})")
|
|
|
|
# Start worker if not already running
|
|
if not self.is_processing:
|
|
await self.start_worker()
|
|
|
|
async def start_worker(self):
|
|
"""Start the queue worker if not already running"""
|
|
if self.worker_task is None or self.worker_task.done():
|
|
self.worker_task = asyncio.create_task(self._process_queue())
|
|
logger.info("Queue worker started")
|
|
|
|
async def _process_queue(self):
|
|
"""Process tasks from queue sequentially"""
|
|
self.is_processing = True
|
|
logger.info("Queue worker processing started")
|
|
|
|
while True:
|
|
try:
|
|
# Wait for next task (with timeout to allow graceful shutdown)
|
|
try:
|
|
task_info = await asyncio.wait_for(
|
|
self.queue.get(),
|
|
timeout=60.0
|
|
)
|
|
except asyncio.TimeoutError:
|
|
# Check if queue is empty
|
|
if self.queue.empty():
|
|
logger.info("Queue empty, worker stopping")
|
|
break
|
|
continue
|
|
|
|
task_id = task_info['task_id']
|
|
self.current_task = task_id
|
|
|
|
logger.info(f"Executing task {task_id} from queue (queue_size={self.queue.qsize()})")
|
|
|
|
# Execute the task (this will block until download completes)
|
|
try:
|
|
if self._execute_callback:
|
|
await self._execute_callback(task_id)
|
|
logger.info(f"Task {task_id} completed successfully")
|
|
else:
|
|
logger.error(f"No execute callback set, cannot run task {task_id}")
|
|
except Exception as e:
|
|
logger.error(f"Task {task_id} failed: {e}", exc_info=True)
|
|
finally:
|
|
self.current_task = None
|
|
self.queue.task_done()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Queue worker error: {e}", exc_info=True)
|
|
|
|
self.is_processing = False
|
|
logger.info("Queue worker stopped")
|
|
|
|
def get_queue_status(self) -> Dict:
|
|
"""Get current queue status"""
|
|
return {
|
|
'current_task': self.current_task,
|
|
'queue_size': self.queue.qsize(),
|
|
'is_processing': self.is_processing
|
|
}
|
|
|
|
async def stop(self):
|
|
"""Stop the queue worker gracefully"""
|
|
logger.info("Stopping queue worker...")
|
|
if self.worker_task and not self.worker_task.done():
|
|
# Wait for current task to complete
|
|
try:
|
|
await asyncio.wait_for(self.worker_task, timeout=300) # 5 minute timeout
|
|
except asyncio.TimeoutError:
|
|
logger.warning("Queue worker did not stop within timeout, cancelling")
|
|
self.worker_task.cancel()
|
|
logger.info("Queue worker stopped")
|
|
|
|
|
|
# Global queue instance
|
|
task_queue = TaskQueue() |