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