Coverage for agentos/core/background.py: 97%
241 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 11:37 +0800
1"""
2Production-grade background task execution framework.
4Supports:
5- Fire-and-forget tasks
6- Scheduled tasks (cron-like + interval)
7- Task queues with worker pools
8- Retry policies with backoff
9- Task state tracking (pending/running/done/failed/cancelled)
10- Concurrency limiting (semaphore)
11- Graceful shutdown
12- Task result storage
14Copyright 2026 AgentOS. All rights reserved.
15"""
17from __future__ import annotations
19import asyncio
20import random
21import time
22import uuid
23from dataclasses import dataclass, field
24from enum import Enum, auto
25from typing import (
26 Any, Awaitable, Callable, Dict, Generic, List, Optional, Tuple, TypeVar,
27 Union,
28)
29import logging
31logger = logging.getLogger("agentos.background")
33T = TypeVar("T")
35# ---------------------------------------------------------------------------
36# Task State
37# ---------------------------------------------------------------------------
39class TaskState(Enum):
40 PENDING = auto()
41 RUNNING = auto()
42 DONE = auto()
43 FAILED = auto()
44 CANCELLED = auto()
45 RETRYING = auto()
48@dataclass
49class TaskResult(Generic[T]):
50 task_id: str
51 state: TaskState
52 result: Optional[T] = None
53 error: Optional[Exception] = None
54 started_at: Optional[float] = None
55 finished_at: Optional[float] = None
56 retries: int = 0
57 metadata: Dict[str, Any] = field(default_factory=dict)
59 @property
60 def duration(self) -> Optional[float]:
61 if self.started_at and self.finished_at:
62 return self.finished_at - self.started_at
63 return None
66# ---------------------------------------------------------------------------
67# Retry Policy
68# ---------------------------------------------------------------------------
70@dataclass
71class RetryPolicy:
72 max_retries: int = 3
73 base_delay: float = 1.0
74 max_delay: float = 60.0
75 backoff_factor: float = 2.0
76 jitter: bool = True
77 retry_on: Tuple[type, ...] = (Exception,)
79 def delay(self, attempt: int) -> float:
80 d = min(self.base_delay * (self.backoff_factor ** attempt), self.max_delay)
81 if self.jitter:
82 d *= 0.5 + random.uniform(0, 0.5)
83 return d
86# ---------------------------------------------------------------------------
87# Schedule
88# ---------------------------------------------------------------------------
90@dataclass(frozen=True)
91class CronSchedule:
92 """Simple cron-like schedule (minute hour day_of_month month day_of_week)."""
93 minute: str = "*"
94 hour: str = "*"
95 day_of_month: str = "*"
96 month: str = "*"
97 day_of_week: str = "*"
100@dataclass(frozen=True)
101class IntervalSchedule:
102 """Run every N seconds."""
103 seconds: float
104 align_to_start: bool = True # drift correction
107Schedule = Union[CronSchedule, IntervalSchedule, float]
110# ---------------------------------------------------------------------------
111# Task Definition
112# ---------------------------------------------------------------------------
114@dataclass
115class TaskDef:
116 func: Callable[..., Any]
117 args: tuple = ()
118 kwargs: Dict[str, Any] = field(default_factory=dict)
119 task_id: str = ""
120 retry_policy: Optional[RetryPolicy] = None
121 timeout: Optional[float] = None
122 schedule: Optional[Schedule] = None
123 name: str = ""
125 def __post_init__(self):
126 if not self.task_id:
127 self.task_id = uuid.uuid4().hex[:12]
128 if not self.name:
129 self.name = self.func.__name__ if hasattr(self.func, "__name__") else self.task_id
132# ---------------------------------------------------------------------------
133# Background Executor
134# ---------------------------------------------------------------------------
136class TaskQueue:
137 """Bounded async task queue with priority."""
139 def __init__(self, maxsize: int = 0):
140 self._queue: asyncio.Queue = asyncio.Queue(maxsize=maxsize)
142 async def put(self, task: TaskDef) -> None:
143 await self._queue.put(task)
145 async def get(self) -> TaskDef:
146 return await self._queue.get()
148 def task_done(self) -> None:
149 self._queue.task_done()
151 async def join(self) -> None:
152 await self._queue.join()
154 @property
155 def qsize(self) -> int:
156 return self._queue.qsize()
158 @property
159 def empty(self) -> bool:
160 return self._queue.empty()
163class BackgroundExecutor:
164 """Main background task execution engine.
166 Usage:
167 executor = BackgroundExecutor(max_concurrency=10)
168 result = await executor.submit(my_func, arg1, arg2, retry=RetryPolicy(3))
170 @executor.scheduled(every=60)
171 async def cleanup_job():
172 ...
174 await executor.start()
175 # ... app runs ...
176 await executor.shutdown()
177 """
179 def __init__(
180 self,
181 max_concurrency: int = 10,
182 queue_size: int = 0,
183 ):
184 self._max_concurrency = max_concurrency
185 self._semaphore = asyncio.Semaphore(max_concurrency)
186 self._queue = TaskQueue(maxsize=queue_size)
187 self._results: Dict[str, TaskResult] = {}
188 self._scheduled: List[Tuple[Schedule, TaskDef]] = []
189 self._workers: List[asyncio.Task] = []
190 self._scheduler_task: Optional[asyncio.Task] = None
191 self._running = False
192 self._shutting_down = False
193 self._accepting = True
194 self._cleanup_interval = 3600.0 # auto-clean results older than this
196 # -- Lifecycle --
198 async def start(self, num_workers: int = 4) -> None:
199 if self._running:
200 return
201 self._running = True
202 self._workers = [
203 asyncio.create_task(self._worker(i), name=f"bg-worker-{i}")
204 for i in range(num_workers)
205 ]
206 if self._scheduled:
207 self._scheduler_task = asyncio.create_task(
208 self._scheduler(), name="bg-scheduler"
209 )
210 logger.info(
211 "BackgroundExecutor started: workers=%d, scheduled=%d",
212 num_workers, len(self._scheduled),
213 )
215 async def shutdown(self, timeout: float = 30.0) -> None:
216 if not self._running:
217 return
218 self._accepting = False
219 logger.info("BackgroundExecutor shutting down (draining queue)...")
221 # Cancel scheduler
222 if self._scheduler_task:
223 self._scheduler_task.cancel()
224 try:
225 await self._scheduler_task
226 except asyncio.CancelledError:
227 pass
229 # Wait for queue to drain while workers are still running
230 try:
231 await asyncio.wait_for(self._queue.join(), timeout=timeout)
232 except asyncio.TimeoutError:
233 logger.warning("Queue drain timed out after %.1fs", timeout)
235 # Now stop workers
236 self._shutting_down = True
237 for worker in self._workers:
238 worker.cancel()
239 results = await asyncio.gather(*self._workers, return_exceptions=True)
240 for r in results:
241 if isinstance(r, Exception) and not isinstance(r, asyncio.CancelledError):
242 logger.error("Worker error during shutdown: %s", r)
244 self._workers.clear()
245 self._running = False
246 self._shutting_down = False
247 logger.info("BackgroundExecutor shut down")
249 # -- Submission --
251 async def submit(
252 self,
253 func: Callable[..., Awaitable[T]],
254 *args: Any,
255 retry_policy: Optional[RetryPolicy] = None,
256 timeout: Optional[float] = None,
257 task_id: Optional[str] = None,
258 **kwargs: Any,
259 ) -> TaskResult[T]:
260 """Submit a task and wait for its result."""
261 if not self._accepting:
262 raise RuntimeError("Executor is shutting down, not accepting tasks")
263 td = TaskDef(
264 func=func, args=args, kwargs=kwargs,
265 task_id=task_id or uuid.uuid4().hex[:12],
266 retry_policy=retry_policy, timeout=timeout,
267 name=func.__name__ if hasattr(func, "__name__") else "anonymous",
268 )
269 async with self._semaphore:
270 return await self._execute(td)
272 async def submit_async(self, func: Callable, *args: Any,
273 retry_policy: Optional[RetryPolicy] = None,
274 timeout: Optional[float] = None,
275 task_id: Optional[str] = None,
276 **kwargs: Any) -> str:
277 """Fire-and-forget: enqueue and return task_id immediately."""
278 td = TaskDef(
279 func=func, args=args, kwargs=kwargs,
280 task_id=task_id or uuid.uuid4().hex[:12],
281 retry_policy=retry_policy, timeout=timeout,
282 name=func.__name__ if hasattr(func, "__name__") else "anonymous",
283 )
284 await self._queue.put(td)
285 return td.task_id
287 def scheduled(self, every: Schedule, retry_policy: Optional[RetryPolicy] = None):
288 """Decorator for recurring scheduled tasks."""
289 def decorator(fn):
290 td = TaskDef(
291 func=fn, retry_policy=retry_policy, schedule=every,
292 )
293 self._scheduled.append((every, td))
294 return fn
295 return decorator
297 # -- Result access --
299 def get_result(self, task_id: str) -> Optional[TaskResult]:
300 return self._results.get(task_id)
302 async def wait_for(self, task_id: str, timeout: Optional[float] = None) -> TaskResult:
303 deadline = time.monotonic() + timeout if timeout else None
304 while True:
305 result = self._results.get(task_id)
306 if result and result.state not in (TaskState.PENDING, TaskState.RUNNING, TaskState.RETRYING):
307 return result
308 if deadline and time.monotonic() > deadline:
309 raise asyncio.TimeoutError(f"Task {task_id} did not complete within {timeout}s")
310 await asyncio.sleep(0.05)
312 # -- Internal --
314 async def _worker(self, worker_id: int):
315 while not self._shutting_down:
316 try:
317 td = await self._queue.get()
318 except asyncio.CancelledError:
319 break
320 async with self._semaphore:
321 try:
322 await self._execute(td)
323 finally:
324 self._queue.task_done()
326 async def _execute(self, td: TaskDef) -> TaskResult:
327 result = TaskResult(task_id=td.task_id, state=TaskState.PENDING)
328 self._results[td.task_id] = result
330 attempts = 1 + (td.retry_policy.max_retries if td.retry_policy else 0)
331 for attempt in range(attempts):
332 result.state = TaskState.RUNNING if attempt == 0 else TaskState.RETRYING
333 result.started_at = time.monotonic()
334 try:
335 coro = td.func(*td.args, **td.kwargs)
336 if td.timeout is not None:
337 value = await asyncio.wait_for(coro, timeout=td.timeout)
338 else:
339 value = await coro
340 result.result = value
341 result.state = TaskState.DONE
342 result.finished_at = time.monotonic()
343 return result
344 except asyncio.CancelledError:
345 result.state = TaskState.CANCELLED
346 result.finished_at = time.monotonic()
347 raise
348 except Exception as exc:
349 if td.retry_policy:
350 if not isinstance(exc, td.retry_policy.retry_on):
351 result.state = TaskState.FAILED
352 result.error = exc
353 result.finished_at = time.monotonic()
354 return result
355 if attempt < td.retry_policy.max_retries:
356 delay = td.retry_policy.delay(attempt)
357 logger.debug(
358 "Task %s retry %d/%d in %.1fs: %s",
359 td.task_id, attempt + 1, td.retry_policy.max_retries,
360 delay, exc,
361 )
362 await asyncio.sleep(delay)
363 result.retries += 1
364 continue
365 result.state = TaskState.FAILED
366 result.error = exc
367 result.finished_at = time.monotonic()
368 return result
370 result.state = TaskState.FAILED
371 result.finished_at = time.monotonic()
372 return result
374 async def _scheduler(self):
375 """Run scheduled tasks at their intervals."""
376 last_runs: Dict[str, float] = {}
377 while not self._shutting_down:
378 now = time.monotonic()
379 for schedule, td in self._scheduled:
380 next_run = self._next_run(schedule, last_runs.get(td.task_id, 0), now)
381 if next_run is not None and now >= next_run:
382 last_runs[td.task_id] = now
383 await self._queue.put(TaskDef(
384 func=td.func, args=td.args, kwargs=td.kwargs,
385 task_id=f"{td.task_id}-{int(now)}",
386 retry_policy=td.retry_policy, timeout=td.timeout,
387 name=td.name,
388 ))
389 await asyncio.sleep(0.1)
391 @staticmethod
392 def _next_run(schedule: Schedule, last: float, now: float) -> Optional[float]:
393 if isinstance(schedule, (int, float)):
394 nxt = last + schedule
395 return nxt if nxt <= now else None
396 if isinstance(schedule, IntervalSchedule):
397 if last == 0:
398 return now # first run: fire immediately
399 return last + schedule.seconds
400 # CronSchedule — simplified: just interval-based for now
401 return None
404# ---------------------------------------------------------------------------
405# Global singleton
406# ---------------------------------------------------------------------------
408_default_executor: Optional[BackgroundExecutor] = None
411def get_background_executor() -> BackgroundExecutor:
412 global _default_executor
413 if _default_executor is None:
414 _default_executor = BackgroundExecutor()
415 return _default_executor
418def set_background_executor(executor: BackgroundExecutor) -> None:
419 global _default_executor
420 _default_executor = executor