Coverage for agentos/core/lifecycle.py: 99%
169 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 11:37 +0800
1"""AgentOS Lifecycle — graceful startup/shutdown with ordered hooks.
3Provides enterprise-grade process lifecycle management:
4- Ordered startup hooks with timeout and health gate
5- Ordered shutdown hooks with grace period
6- SIGTERM/SIGINT graceful shutdown integration
7- Component-level health registration
8- Liveness/readiness probe support
10Design: ~370 lines, zero external deps beyond stdlib + asyncio.
11"""
13from __future__ import annotations
15import asyncio
16import logging
17import signal
18import time
19from dataclasses import dataclass, field
20from enum import Enum
21from typing import Any, Callable, Dict, List, Optional
23logger = logging.getLogger(__name__)
26# ============================================================================
27# Types
28# ============================================================================
30class LifecyclePhase(str, Enum):
31 """Ordered lifecycle phases during startup."""
33 CONFIG = "config" # Configuration loading
34 INFRA = "infra" # DB, Redis, message queues
35 SECURITY = "security" # Auth, encryption, certs
36 SERVICES = "services" # Internal services
37 MIDDLEWARE = "middleware" # Middleware pipeline
38 API = "api" # HTTP/gRPC server
39 READY = "ready" # Final readiness signal
42class ComponentStatus(str, Enum):
43 """Component health status."""
45 UNINITIALIZED = "uninitialized"
46 STARTING = "starting"
47 HEALTHY = "healthy"
48 DEGRADED = "degraded"
49 UNHEALTHY = "unhealthy"
50 SHUTTING_DOWN = "shutting_down"
51 STOPPED = "stopped"
54@dataclass
55class LifecycleHook:
56 """A single startup or shutdown hook with metadata."""
58 name: str
59 phase: LifecyclePhase
60 fn: Callable[[], Any] # or async callable
61 timeout_seconds: float = 30.0
62 is_async: bool = False
63 critical: bool = True # Fail startup if critical hook fails
64 weight: int = 50 # Ordering within same phase (lower = first)
65 retries: int = 0
66 retry_delay: float = 1.0
69@dataclass
70class ComponentHealth:
71 """Health status of a single component."""
73 name: str
74 status: ComponentStatus = ComponentStatus.UNINITIALIZED
75 phase: Optional[LifecyclePhase] = None
76 message: str = ""
77 error: Optional[str] = None
78 started_at: Optional[float] = None
79 duration_ms: float = 0.0
82@dataclass
83class LifecycleReport:
84 """Full lifecycle status report."""
86 overall_status: ComponentStatus = ComponentStatus.UNINITIALIZED
87 phase: str = ""
88 components: Dict[str, ComponentHealth] = field(default_factory=dict)
89 startup_duration_ms: float = 0.0
90 shutdown_remaining_hooks: int = 0
92 @property
93 def is_healthy(self) -> bool:
94 return self.overall_status == ComponentStatus.HEALTHY
96 @property
97 def is_ready(self) -> bool:
98 return self.overall_status in (
99 ComponentStatus.HEALTHY, ComponentStatus.DEGRADED
100 )
103# ============================================================================
104# Lifecycle Manager
105# ============================================================================
107class LifecycleManager:
108 """Orchestrates ordered startup and graceful shutdown.
110 Usage:
111 lm = LifecycleManager()
113 @lm.on_startup(phase=LifecyclePhase.INFRA)
114 async def init_db():
115 ...
117 @lm.on_shutdown
118 async def close_db():
119 ...
121 async with lm:
122 # Application runs here
123 ...
125 Signal handling (SIGTERM, SIGINT) integrated automatically.
126 """
128 PHASE_ORDER: List[LifecyclePhase] = [
129 LifecyclePhase.CONFIG,
130 LifecyclePhase.INFRA,
131 LifecyclePhase.SECURITY,
132 LifecyclePhase.SERVICES,
133 LifecyclePhase.MIDDLEWARE,
134 LifecyclePhase.API,
135 LifecyclePhase.READY,
136 ]
138 def __init__(
139 self,
140 grace_period: float = 30.0,
141 startup_timeout: float = 120.0,
142 ):
143 self._startup_hooks: List[LifecycleHook] = []
144 self._shutdown_hooks: List[LifecycleHook] = []
145 self._health: Dict[str, ComponentHealth] = {}
146 self._status = ComponentStatus.UNINITIALIZED
147 self._start_time: Optional[float] = None
148 self._grace_period = grace_period
149 self._startup_timeout = startup_timeout
150 self._shutdown_event = asyncio.Event()
151 self._ready_event = asyncio.Event()
153 # ── Registration ──────────────────────────────────────────────────────
155 def on_startup(
156 self,
157 name: Optional[str] = None,
158 *,
159 phase: LifecyclePhase = LifecyclePhase.SERVICES,
160 critical: bool = True,
161 timeout_seconds: float = 30.0,
162 weight: int = 50,
163 retries: int = 0,
164 retry_delay: float = 1.0,
165 ) -> Callable:
166 """Decorator: register a startup hook."""
167 def decorator(fn):
168 hook_name = name or fn.__name__
169 is_async = asyncio.iscoroutinefunction(fn)
170 self._startup_hooks.append(LifecycleHook(
171 name=hook_name,
172 phase=phase,
173 fn=fn,
174 timeout_seconds=timeout_seconds,
175 is_async=is_async,
176 critical=critical,
177 weight=weight,
178 retries=retries,
179 retry_delay=retry_delay,
180 ))
181 self._health[hook_name] = ComponentHealth(
182 name=hook_name, phase=phase
183 )
184 return fn
185 return decorator
187 def on_shutdown(
188 self,
189 name: Optional[str] = None,
190 *,
191 timeout_seconds: float = 10.0,
192 weight: int = 50,
193 ) -> Callable:
194 """Decorator: register a shutdown hook (reverse order)."""
195 def decorator(fn):
196 hook_name = name or fn.__name__
197 is_async = asyncio.iscoroutinefunction(fn)
198 self._shutdown_hooks.append(LifecycleHook(
199 name=hook_name,
200 phase=LifecyclePhase.READY, # irrelevant for shutdown
201 fn=fn,
202 timeout_seconds=timeout_seconds,
203 is_async=is_async,
204 critical=False,
205 weight=weight,
206 ))
207 return fn
208 return decorator
210 # ── Startup ───────────────────────────────────────────────────────────
212 async def start(self) -> LifecycleReport:
213 """Execute all startup hooks in phase/weight order."""
214 self._start_time = time.perf_counter()
215 self._status = ComponentStatus.STARTING
217 # Sort: phase order first, then weight within phase
218 phase_idx = {p: i for i, p in enumerate(self.PHASE_ORDER)}
219 sorted_hooks = sorted(
220 self._startup_hooks,
221 key=lambda h: (phase_idx.get(h.phase, 99), h.weight),
222 )
224 for hook in sorted_hooks:
225 ok = await self._execute_hook(hook, is_startup=True)
226 if not ok and hook.critical:
227 self._status = ComponentStatus.UNHEALTHY
228 return self.report()
230 self._status = ComponentStatus.HEALTHY
231 self._ready_event.set()
233 return self.report()
235 async def _execute_hook(
236 self, hook: LifecycleHook, is_startup: bool = True
237 ) -> bool:
238 """Execute a single hook with timeout, retries, and health tracking."""
239 health = self._health.get(hook.name) or ComponentHealth(name=hook.name)
240 health.status = ComponentStatus.STARTING if is_startup else ComponentStatus.SHUTTING_DOWN
241 health.started_at = time.time()
243 attempt = 0
244 last_error = None
246 while attempt <= hook.retries:
247 t0 = time.perf_counter()
248 try:
249 if hook.is_async:
250 await asyncio.wait_for(hook.fn(), timeout=hook.timeout_seconds)
251 else:
252 loop = asyncio.get_event_loop()
253 await asyncio.wait_for(
254 loop.run_in_executor(None, hook.fn),
255 timeout=hook.timeout_seconds,
256 )
257 health.duration_ms = (time.perf_counter() - t0) * 1000
258 health.status = ComponentStatus.HEALTHY if is_startup else ComponentStatus.STOPPED
259 health.message = "OK"
260 logger.info(f"[lifecycle] {hook.name}: OK ({health.duration_ms:.0f}ms)")
261 return True
263 except asyncio.TimeoutError:
264 last_error = f"Timeout after {hook.timeout_seconds}s"
265 health.error = last_error
266 logger.warning(f"[lifecycle] {hook.name}: {last_error}")
267 except Exception as e:
268 last_error = str(e)
269 health.error = last_error
270 logger.warning(f"[lifecycle] {hook.name}: {last_error}")
272 attempt += 1
273 if attempt <= hook.retries:
274 await asyncio.sleep(hook.retry_delay)
276 health.status = ComponentStatus.UNHEALTHY
277 health.duration_ms = (time.perf_counter() - t0) * 1000
278 return False
280 # ── Shutdown ──────────────────────────────────────────────────────────
282 async def shutdown(self, signal_name: str = "") -> LifecycleReport:
283 """Execute all shutdown hooks in reverse registration order."""
284 if self._status == ComponentStatus.STOPPED:
285 return self.report()
287 self._status = ComponentStatus.SHUTTING_DOWN
288 self._shutdown_event.set()
289 logger.info(f"[lifecycle] Shutting down gracefully{f' ({signal_name})' if signal_name else ''}")
291 # Reverse order for shutdown (LIFO — last started, first stopped)
292 for hook in reversed(self._shutdown_hooks):
293 await self._execute_hook(hook, is_startup=False)
295 self._status = ComponentStatus.STOPPED
296 return self.report()
298 # ── Signal integration ────────────────────────────────────────────────
300 def setup_signal_handlers(self, loop: Optional[asyncio.AbstractEventLoop] = None):
301 """Register SIGTERM/SIGINT handlers on the event loop."""
302 if loop is None:
303 loop = asyncio.get_event_loop()
305 for sig in (signal.SIGTERM, signal.SIGINT):
306 try:
307 loop.add_signal_handler(
308 sig,
309 lambda s=sig: asyncio.ensure_future(
310 self.shutdown(signal.Signals(s).name)
311 ),
312 )
313 except (NotImplementedError, RuntimeError):
314 # Windows or non-main-thread — fallback to signal.signal
315 signal.signal(sig, lambda s, f: asyncio.ensure_future(
316 self.shutdown(signal.Signals(s).name)
317 ))
319 # ── Probes ────────────────────────────────────────────────────────────
321 def is_ready(self) -> bool:
322 """Readiness probe: is the service ready to accept requests?"""
323 return self._ready_event.is_set()
325 def is_live(self) -> bool:
326 """Liveness probe: is the process alive (not hung)?"""
327 return self._status not in (ComponentStatus.STOPPED, ComponentStatus.UNHEALTHY)
329 # ── Report ────────────────────────────────────────────────────────────
331 def report(self) -> LifecycleReport:
332 """Generate a full lifecycle status report."""
333 startup_ms = 0.0
334 if self._start_time:
335 startup_ms = (time.perf_counter() - self._start_time) * 1000
337 return LifecycleReport(
338 overall_status=self._status,
339 phase=self._status.value,
340 components=dict(self._health),
341 startup_duration_ms=startup_ms,
342 shutdown_remaining_hooks=len([
343 h for h in self._shutdown_hooks
344 if self._health.get(h.name, ComponentHealth(name=h.name)).status
345 not in (ComponentStatus.STOPPED,)
346 ]),
347 )
349 # ── Context manager ───────────────────────────────────────────────────
351 async def __aenter__(self):
352 """Async context manager: start lifecycle."""
353 self.setup_signal_handlers()
354 await self.start()
355 return self
357 async def __aexit__(self, exc_type, exc_val, exc_tb):
358 """Async context manager: graceful shutdown."""
359 await self.shutdown()
360 return False # Don't suppress exceptions
363# ============================================================================
364# Singleton helper
365# ============================================================================
367_default_lifecycle: Optional[LifecycleManager] = None
370def get_lifecycle(
371 grace_period: float = 30.0,
372 startup_timeout: float = 120.0,
373) -> LifecycleManager:
374 """Get or create the global LifecycleManager singleton."""
375 global _default_lifecycle
376 if _default_lifecycle is None:
377 _default_lifecycle = LifecycleManager(
378 grace_period=grace_period,
379 startup_timeout=startup_timeout,
380 )
381 return _default_lifecycle