Yuriy Yuriev
163 lines
4.8 KiB
Python
163 lines
4.8 KiB
Python
"""Background task manager for async task tracking."""
|
|
|
|
import asyncio
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, Coroutine
|
|
|
|
from core.constants import TaskStatus
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class TaskInfo:
|
|
"""Information about a running task."""
|
|
task_id: str
|
|
task_type: str
|
|
coroutine: Coroutine = None
|
|
status: TaskStatus = TaskStatus.RUNNING
|
|
created_at: float = field(default_factory=lambda: asyncio.get_running_loop().time())
|
|
metadata: dict = field(default_factory=dict)
|
|
|
|
|
|
class BackgroundTaskManager:
|
|
"""
|
|
Manages background async tasks with tracking and cancellation support.
|
|
|
|
Features:
|
|
- Start and track async tasks
|
|
- Cancel individual or all tasks
|
|
- Get task status and metadata
|
|
- Prevent duplicate tasks
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._tasks: Dict[str, asyncio.Task] = {}
|
|
self._task_info: Dict[str, TaskInfo] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def start_task(
|
|
self,
|
|
task_id: str,
|
|
coro: Coroutine,
|
|
task_type: str = "unknown",
|
|
metadata: dict = None
|
|
) -> bool:
|
|
"""
|
|
Start a background task.
|
|
|
|
Args:
|
|
task_id: Unique task identifier
|
|
coro: Coroutine to run
|
|
task_type: Type of task
|
|
metadata: Additional task metadata
|
|
|
|
Returns:
|
|
True if started successfully, False if already running
|
|
"""
|
|
async with self._lock:
|
|
if task_id in self._tasks and not self._tasks[task_id].done():
|
|
logger.warning(f"Task {task_id} is already running")
|
|
return False
|
|
|
|
task = asyncio.create_task(coro)
|
|
self._tasks[task_id] = task
|
|
|
|
self._task_info[task_id] = TaskInfo(
|
|
task_id=task_id,
|
|
task_type=task_type,
|
|
coroutine=coro,
|
|
metadata=metadata or {}
|
|
)
|
|
|
|
task.add_done_callback(lambda t, tid=task_id: self._on_task_done(tid, t))
|
|
|
|
logger.info(f"Started background task: {task_id} ({task_type})")
|
|
return True
|
|
|
|
def _on_task_done(self, task_id: str, task: asyncio.Task) -> None:
|
|
if task.cancelled():
|
|
logger.info(f"Task {task_id} cancelled")
|
|
return
|
|
exc = task.exception()
|
|
if exc:
|
|
logger.error(f"Task {task_id} failed: {exc}")
|
|
else:
|
|
logger.info(f"Task {task_id} completed successfully")
|
|
|
|
async def cancel_task(self, task_id: str) -> bool:
|
|
"""
|
|
Cancel a specific task.
|
|
|
|
Args:
|
|
task_id: Task identifier to cancel
|
|
|
|
Returns:
|
|
True if task was cancelled, False if not found
|
|
"""
|
|
async with self._lock:
|
|
task = self._tasks.get(task_id)
|
|
if task and not task.done():
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
logger.info(f"Task cancelled: {task_id}")
|
|
return True
|
|
return False
|
|
|
|
async def cancel_all(self) -> int:
|
|
"""
|
|
Cancel all running tasks.
|
|
|
|
Returns:
|
|
Number of tasks cancelled
|
|
"""
|
|
cancelled = 0
|
|
|
|
async with self._lock:
|
|
for task in list(self._tasks.values()):
|
|
if not task.done():
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
cancelled += 1
|
|
|
|
self._tasks.clear()
|
|
self._task_info.clear()
|
|
|
|
logger.info(f"All tasks cancelled: {cancelled}")
|
|
return cancelled
|
|
|
|
def get_active_tasks(self) -> Dict[str, dict]:
|
|
"""
|
|
Get all active tasks with their metadata.
|
|
|
|
Returns:
|
|
Dictionary of active tasks
|
|
"""
|
|
active = {}
|
|
for task_id, task in self._tasks.items():
|
|
if not task.done():
|
|
info = self._task_info.get(task_id)
|
|
active[task_id] = {
|
|
"info": info.metadata if info else {},
|
|
"status": TaskStatus.RUNNING.value,
|
|
"type": info.task_type if info else "unknown",
|
|
"created_at": info.created_at if info else 0
|
|
}
|
|
return active
|
|
|
|
@property
|
|
def active_count(self) -> int:
|
|
"""Number of currently active tasks."""
|
|
return sum(1 for t in self._tasks.values() if not t.done())
|
|
|
|
def has_active_tasks(self) -> bool:
|
|
"""Check if there are any active tasks."""
|
|
return self.active_count > 0 |