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