Coverage for agentos/core/background.py: 97%
242 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 12:29 +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 logging
21import random
22import time
23import uuid
24from collections.abc import Awaitable, Callable
25from dataclasses import dataclass, field
26from enum import Enum, auto
27from typing import (
28 Any,
29 Generic,
30 TypeVar,
31 Union,
32)
34logger = logging.getLogger("agentos.background")
36T = TypeVar("T")
38# ---------------------------------------------------------------------------
39# Task State
40# ---------------------------------------------------------------------------
43class TaskState(Enum):
44 PENDING = auto()
45 RUNNING = auto()
46 DONE = auto()
47 FAILED = auto()
48 CANCELLED = auto()
49 RETRYING = auto()
52@dataclass
53class TaskResult(Generic[T]):
54 task_id: str
55 state: TaskState
56 result: T | None = None
57 error: Exception | None = None
58 started_at: float | None = None
59 finished_at: float | None = None
60 retries: int = 0
61 metadata: dict[str, Any] = field(default_factory=dict)
63 @property
64 def duration(self) -> float | None:
65 if self.started_at and self.finished_at:
66 return self.finished_at - self.started_at
67 return None
70# ---------------------------------------------------------------------------
71# Retry Policy
72# ---------------------------------------------------------------------------
75@dataclass
76class RetryPolicy:
77 max_retries: int = 3
78 base_delay: float = 1.0
79 max_delay: float = 60.0
80 backoff_factor: float = 2.0
81 jitter: bool = True
82 retry_on: tuple[type, ...] = (Exception,)
84 def delay(self, attempt: int) -> float:
85 d = min(self.base_delay * (self.backoff_factor**attempt), self.max_delay)
86 if self.jitter:
87 d *= 0.5 + random.uniform(0, 0.5)
88 return d
91# ---------------------------------------------------------------------------
92# Schedule
93# ---------------------------------------------------------------------------
96@dataclass(frozen=True)
97class CronSchedule:
98 """Simple cron-like schedule (minute hour day_of_month month day_of_week)."""
100 minute: str = "*"
101 hour: str = "*"
102 day_of_month: str = "*"
103 month: str = "*"
104 day_of_week: str = "*"
107@dataclass(frozen=True)
108class IntervalSchedule:
109 """Run every N seconds."""
111 seconds: float
112 align_to_start: bool = True # drift correction
115Schedule = Union[CronSchedule, IntervalSchedule, float]
118# ---------------------------------------------------------------------------
119# Task Definition
120# ---------------------------------------------------------------------------
123@dataclass
124class TaskDef:
125 func: Callable[..., Any]
126 args: tuple = ()
127 kwargs: dict[str, Any] = field(default_factory=dict)
128 task_id: str = ""
129 retry_policy: RetryPolicy | None = None
130 timeout: float | None = None
131 schedule: Schedule | None = None
132 name: str = ""
134 def __post_init__(self):
135 if not self.task_id:
136 self.task_id = uuid.uuid4().hex[:12]
137 if not self.name:
138 self.name = self.func.__name__ if hasattr(self.func, "__name__") else self.task_id
141# ---------------------------------------------------------------------------
142# Background Executor
143# ---------------------------------------------------------------------------
146class TaskQueue:
147 """Bounded async task queue with priority."""
149 def __init__(self, maxsize: int = 0):
150 self._queue: asyncio.Queue = asyncio.Queue(maxsize=maxsize)
152 async def put(self, task: TaskDef) -> None:
153 await self._queue.put(task)
155 async def get(self) -> TaskDef:
156 return await self._queue.get()
158 def task_done(self) -> None:
159 self._queue.task_done()
161 async def join(self) -> None:
162 await self._queue.join()
164 @property
165 def qsize(self) -> int:
166 return self._queue.qsize()
168 @property
169 def empty(self) -> bool:
170 return self._queue.empty()
173class BackgroundExecutor:
174 """Main background task execution engine.
176 Usage:
177 executor = BackgroundExecutor(max_concurrency=10)
178 result = await executor.submit(my_func, arg1, arg2, retry=RetryPolicy(3))
180 @executor.scheduled(every=60)
181 async def cleanup_job():
182 ...
184 await executor.start()
185 # ... app runs ...
186 await executor.shutdown()
187 """
189 def __init__(
190 self,
191 max_concurrency: int = 10,
192 queue_size: int = 0,
193 ):
194 self._max_concurrency = max_concurrency
195 self._semaphore = asyncio.Semaphore(max_concurrency)
196 self._queue = TaskQueue(maxsize=queue_size)
197 self._results: dict[str, TaskResult] = {}
198 self._scheduled: list[tuple[Schedule, TaskDef]] = []
199 self._workers: list[asyncio.Task] = []
200 self._scheduler_task: asyncio.Task | None = None
201 self._running = False
202 self._shutting_down = False
203 self._accepting = True
204 self._cleanup_interval = 3600.0 # auto-clean results older than this
206 # -- Lifecycle --
208 async def start(self, num_workers: int = 4) -> None:
209 if self._running:
210 return
211 self._running = True
212 self._workers = [
213 asyncio.create_task(self._worker(i), name=f"bg-worker-{i}") for i in range(num_workers)
214 ]
215 if self._scheduled:
216 self._scheduler_task = asyncio.create_task(self._scheduler(), name="bg-scheduler")
217 logger.info(
218 "BackgroundExecutor started: workers=%d, scheduled=%d",
219 num_workers,
220 len(self._scheduled),
221 )
223 async def shutdown(self, timeout: float = 30.0) -> None:
224 if not self._running:
225 return
226 self._accepting = False
227 logger.info("BackgroundExecutor shutting down (draining queue)...")
229 # Cancel scheduler
230 if self._scheduler_task:
231 self._scheduler_task.cancel()
232 try:
233 await self._scheduler_task
234 except asyncio.CancelledError:
235 pass
237 # Wait for queue to drain while workers are still running
238 try:
239 await asyncio.wait_for(self._queue.join(), timeout=timeout)
240 except TimeoutError:
241 logger.warning("Queue drain timed out after %.1fs", timeout)
243 # Now stop workers
244 self._shutting_down = True
245 for worker in self._workers:
246 worker.cancel()
247 results = await asyncio.gather(*self._workers, return_exceptions=True)
248 for r in results:
249 if isinstance(r, Exception) and not isinstance(r, asyncio.CancelledError):
250 logger.error("Worker error during shutdown: %s", r)
252 self._workers.clear()
253 self._running = False
254 self._shutting_down = False
255 logger.info("BackgroundExecutor shut down")
257 # -- Submission --
259 async def submit(
260 self,
261 func: Callable[..., Awaitable[T]],
262 *args: Any,
263 retry_policy: RetryPolicy | None = None,
264 timeout: float | None = None,
265 task_id: str | None = None,
266 **kwargs: Any,
267 ) -> TaskResult[T]:
268 """Submit a task and wait for its result."""
269 if not self._accepting:
270 raise RuntimeError("Executor is shutting down, not accepting tasks")
271 td = TaskDef(
272 func=func,
273 args=args,
274 kwargs=kwargs,
275 task_id=task_id or uuid.uuid4().hex[:12],
276 retry_policy=retry_policy,
277 timeout=timeout,
278 name=func.__name__ if hasattr(func, "__name__") else "anonymous",
279 )
280 async with self._semaphore:
281 return await self._execute(td)
283 async def submit_async(
284 self,
285 func: Callable,
286 *args: Any,
287 retry_policy: RetryPolicy | None = None,
288 timeout: float | None = None,
289 task_id: str | None = None,
290 **kwargs: Any,
291 ) -> str:
292 """Fire-and-forget: enqueue and return task_id immediately."""
293 td = TaskDef(
294 func=func,
295 args=args,
296 kwargs=kwargs,
297 task_id=task_id or uuid.uuid4().hex[:12],
298 retry_policy=retry_policy,
299 timeout=timeout,
300 name=func.__name__ if hasattr(func, "__name__") else "anonymous",
301 )
302 await self._queue.put(td)
303 return td.task_id
305 def scheduled(self, every: Schedule, retry_policy: RetryPolicy | None = None):
306 """Decorator for recurring scheduled tasks."""
308 def decorator(fn):
309 td = TaskDef(
310 func=fn,
311 retry_policy=retry_policy,
312 schedule=every,
313 )
314 self._scheduled.append((every, td))
315 return fn
317 return decorator
319 # -- Result access --
321 def get_result(self, task_id: str) -> TaskResult | None:
322 return self._results.get(task_id)
324 async def wait_for(self, task_id: str, timeout: float | None = None) -> TaskResult:
325 deadline = time.monotonic() + timeout if timeout else None
326 while True:
327 result = self._results.get(task_id)
328 if result and result.state not in (
329 TaskState.PENDING,
330 TaskState.RUNNING,
331 TaskState.RETRYING,
332 ):
333 return result
334 if deadline and time.monotonic() > deadline:
335 raise TimeoutError(f"Task {task_id} did not complete within {timeout}s")
336 await asyncio.sleep(0.05)
338 # -- Internal --
340 async def _worker(self, worker_id: int):
341 while not self._shutting_down:
342 try:
343 td = await self._queue.get()
344 except asyncio.CancelledError:
345 break
346 async with self._semaphore:
347 try:
348 await self._execute(td)
349 finally:
350 self._queue.task_done()
352 async def _execute(self, td: TaskDef) -> TaskResult:
353 result = TaskResult(task_id=td.task_id, state=TaskState.PENDING)
354 self._results[td.task_id] = result
356 attempts = 1 + (td.retry_policy.max_retries if td.retry_policy else 0)
357 for attempt in range(attempts):
358 result.state = TaskState.RUNNING if attempt == 0 else TaskState.RETRYING
359 result.started_at = time.monotonic()
360 try:
361 coro = td.func(*td.args, **td.kwargs)
362 if td.timeout is not None:
363 value = await asyncio.wait_for(coro, timeout=td.timeout)
364 else:
365 value = await coro
366 result.result = value
367 result.state = TaskState.DONE
368 result.finished_at = time.monotonic()
369 return result
370 except asyncio.CancelledError:
371 result.state = TaskState.CANCELLED
372 result.finished_at = time.monotonic()
373 raise
374 except Exception as exc:
375 if td.retry_policy:
376 if not isinstance(exc, td.retry_policy.retry_on):
377 result.state = TaskState.FAILED
378 result.error = exc
379 result.finished_at = time.monotonic()
380 return result
381 if attempt < td.retry_policy.max_retries:
382 delay = td.retry_policy.delay(attempt)
383 logger.debug(
384 "Task %s retry %d/%d in %.1fs: %s",
385 td.task_id,
386 attempt + 1,
387 td.retry_policy.max_retries,
388 delay,
389 exc,
390 )
391 await asyncio.sleep(delay)
392 result.retries += 1
393 continue
394 result.state = TaskState.FAILED
395 result.error = exc
396 result.finished_at = time.monotonic()
397 return result
399 result.state = TaskState.FAILED
400 result.finished_at = time.monotonic()
401 return result
403 async def _scheduler(self):
404 """Run scheduled tasks at their intervals."""
405 last_runs: dict[str, float] = {}
406 while not self._shutting_down:
407 now = time.monotonic()
408 for schedule, td in self._scheduled:
409 next_run = self._next_run(schedule, last_runs.get(td.task_id, 0), now)
410 if next_run is not None and now >= next_run:
411 last_runs[td.task_id] = now
412 await self._queue.put(
413 TaskDef(
414 func=td.func,
415 args=td.args,
416 kwargs=td.kwargs,
417 task_id=f"{td.task_id}-{int(now)}",
418 retry_policy=td.retry_policy,
419 timeout=td.timeout,
420 name=td.name,
421 )
422 )
423 await asyncio.sleep(0.1)
425 @staticmethod
426 def _next_run(schedule: Schedule, last: float, now: float) -> float | None:
427 if isinstance(schedule, (int, float)):
428 nxt = last + schedule
429 return nxt if nxt <= now else None
430 if isinstance(schedule, IntervalSchedule):
431 if last == 0:
432 return now # first run: fire immediately
433 return last + schedule.seconds
434 # CronSchedule — simplified: just interval-based for now
435 return None
438# ---------------------------------------------------------------------------
439# Global singleton
440# ---------------------------------------------------------------------------
442_default_executor: BackgroundExecutor | None = None
445def get_background_executor() -> BackgroundExecutor:
446 global _default_executor
447 if _default_executor is None:
448 _default_executor = BackgroundExecutor()
449 return _default_executor
452def set_background_executor(executor: BackgroundExecutor) -> None:
453 global _default_executor
454 _default_executor = executor