Coverage for agentos/core/resource_manager.py: 0%
288 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 10:19 +0800
1"""AgentOS Resource Manager — production-grade resource lifecycle & cleanup.
3Manages:
4- Async resource acquisition/release (connections, pools, files)
5- Ordered shutdown (LIFO — last acquired, first released)
6- Finalizer registry for guaranteed cleanup
7- Health-checked resource pools
8- Leak detection
10Design: ~330 lines, zero external deps beyond stdlib + asyncio.
11"""
13from __future__ import annotations
15import asyncio
16import logging
17import time
18from abc import ABC, abstractmethod
19from collections.abc import AsyncIterator, Callable
20from contextlib import asynccontextmanager
21from dataclasses import dataclass
22from enum import StrEnum
23from typing import (
24 Any,
25 TypeVar,
26)
28logger = logging.getLogger(__name__)
30T = TypeVar("T")
33# ============================================================================
34# Data types
35# ============================================================================
38class ResourceState(StrEnum):
39 CREATED = "created"
40 ACTIVE = "active"
41 CLOSING = "closing"
42 CLOSED = "closed"
43 ERROR = "error"
46class ResourceType(StrEnum):
47 CONNECTION = "connection"
48 POOL = "pool"
49 FILE = "file"
50 LOCK = "lock"
51 SESSION = "session"
52 OTHER = "other"
55@dataclass
56class ResourceInfo:
57 """Metadata about a managed resource."""
59 name: str
60 resource_type: ResourceType = ResourceType.OTHER
61 state: ResourceState = ResourceState.CREATED
62 acquired_at: float = 0.0
63 released_at: float = 0.0
64 leak_warn_threshold: float = 300.0 # Warn if held > 5 min
67# ============================================================================
68# Abstract Resource
69# ============================================================================
72class AbstractResource(ABC):
73 """Interface for managed resources."""
75 @abstractmethod
76 async def close(self) -> None:
77 """Release the underlying resource."""
79 async def health_check(self) -> bool:
80 """Optional health check. Default: assume healthy if not closed."""
81 return True
83 def on_leak_detected(self) -> None:
84 """Callback when resource appears to be leaked."""
85 logger.warning("Resource leak suspected: %s", self)
88# ============================================================================
89# Managed Resource wrapper
90# ============================================================================
93class ManagedResource(AbstractResource):
94 """Wraps any async-closeable object into a managed resource."""
96 def __init__(
97 self,
98 resource: Any,
99 close_fn: Callable[[Any], Any] | None = None,
100 name: str = "unnamed",
101 resource_type: ResourceType = ResourceType.OTHER,
102 ):
103 self._resource = resource
104 self._close_fn = (
105 close_fn or getattr(resource, "close", None) or getattr(resource, "aclose", None)
106 )
107 self.info = ResourceInfo(name=name, resource_type=resource_type)
108 self.info.acquired_at = time.monotonic()
109 self.info.state = ResourceState.ACTIVE
111 @property
112 def raw(self) -> Any:
113 """Access the underlying resource object."""
114 return self._resource
116 async def close(self) -> None:
117 if self.info.state == ResourceState.CLOSED:
118 return
120 self.info.state = ResourceState.CLOSING
121 try:
122 if self._close_fn is not None:
123 result = self._close_fn()
124 if asyncio.iscoroutine(result):
125 await result
126 except Exception as exc:
127 self.info.state = ResourceState.ERROR
128 logger.error("Failed to close resource %s: %s", self.info.name, exc)
129 raise
130 else:
131 self.info.state = ResourceState.CLOSED
132 self.info.released_at = time.monotonic()
134 async def health_check(self) -> bool:
135 if self.info.state == ResourceState.CLOSED:
136 return False
137 health_fn = getattr(self._resource, "health_check", None) or getattr(
138 self._resource, "ping", None
139 )
140 if health_fn is not None:
141 try:
142 result = health_fn()
143 if asyncio.iscoroutine(result):
144 result = await result
145 return bool(result)
146 except Exception:
147 return False
148 return True
151# ============================================================================
152# Resource Pool
153# ============================================================================
156class ResourcePool(AbstractResource):
157 """Generic async resource pool with health-checked lend/return.
159 Usage:
160 pool = ResourcePool(factory=create_db_conn, max_size=10)
161 async with pool.acquire() as conn:
162 await conn.query(...)
163 """
165 def __init__(
166 self,
167 factory: Callable[[], Any],
168 max_size: int = 10,
169 min_size: int = 0,
170 idle_timeout: float = 300.0,
171 health_check_interval: float = 30.0,
172 name: str = "pool",
173 ):
174 self._factory = factory
175 self._max_size = max_size
176 self._min_size = min_size
177 self._idle_timeout = idle_timeout
178 self._health_check_interval = health_check_interval
179 self._name = name
181 self._available: asyncio.Queue = asyncio.Queue(maxsize=max_size)
182 self._in_use: dict[int, Any] = {} # id(resource) → resource
183 self._total: int = 0
184 self._lock = asyncio.Lock()
185 self._closed = False
186 self._health_task: asyncio.Task | None = None
188 async def _prefill(self):
189 """Pre-create min_size connections."""
190 for _ in range(self._min_size):
191 resource = self._factory()
192 if asyncio.iscoroutine(resource):
193 resource = await resource
194 await self._available.put(resource)
195 self._total += 1
197 async def start(self):
198 """Initialize the pool."""
199 await self._prefill()
200 if self._health_check_interval > 0:
201 self._health_task = asyncio.ensure_future(self._health_loop())
203 @asynccontextmanager
204 async def acquire(self) -> AsyncIterator[Any]:
205 """Acquire a resource from the pool. Returns via context manager."""
206 resource = await self._get()
207 try:
208 yield resource
209 finally:
210 await self._return(resource)
212 async def _get(self) -> Any:
213 if self._closed:
214 raise RuntimeError(f"ResourcePool '{self._name}' is closed")
216 # Try to get from available
217 try:
218 resource = self._available.get_nowait()
219 self._in_use[id(resource)] = resource
220 return resource
221 except asyncio.QueueEmpty:
222 pass
224 # Try to create new
225 async with self._lock:
226 if self._total < self._max_size:
227 resource = self._factory()
228 if asyncio.iscoroutine(resource):
229 resource = await resource
230 self._total += 1
231 self._in_use[id(resource)] = resource
232 return resource
234 # Wait for one to become available
235 resource = await self._available.get()
236 self._in_use[id(resource)] = resource
237 return resource
239 async def _return(self, resource: Any):
240 self._in_use.pop(id(resource), None)
241 if not self._closed:
242 await self._available.put(resource)
244 async def _health_loop(self):
245 """Periodic health check, remove dead connections."""
246 while not self._closed:
247 await asyncio.sleep(self._health_check_interval)
248 # Drain and re-check
249 healthy: list[Any] = []
250 while not self._available.empty():
251 try:
252 resource = self._available.get_nowait()
253 if self._is_healthy(resource):
254 healthy.append(resource)
255 else:
256 self._total -= 1
257 logger.debug("Pool '%s': removed unhealthy connection", self._name)
258 except asyncio.QueueEmpty:
259 break
261 for resource in healthy:
262 await self._available.put(resource)
264 # Top up to min_size
265 while self._total < self._min_size and not self._closed:
266 async with self._lock:
267 if self._total < self._max_size:
268 resource = self._factory()
269 if asyncio.iscoroutine(resource):
270 resource = await resource
271 await self._available.put(resource)
272 self._total += 1
274 def _is_healthy(self, resource: Any) -> bool:
275 """Check if a resource is healthy."""
276 health_fn = getattr(resource, "health_check", None) or getattr(resource, "ping", None)
277 if health_fn is None:
278 return True
279 try:
280 result = health_fn()
281 if asyncio.iscoroutine(result):
282 return True # Can't check async in sync context
283 return bool(result)
284 except Exception:
285 return False
287 def _close_resource(self, resource: Any, errors: list[Exception]):
288 """Attempt to close a single resource, collecting errors."""
289 close_fn = (
290 getattr(resource, "close", None)
291 or getattr(resource, "aclose", None)
292 or (resource.get("close") if isinstance(resource, dict) else None)
293 or (resource.get("aclose") if isinstance(resource, dict) else None)
294 )
295 if close_fn is not None:
296 try:
297 result = close_fn()
298 if asyncio.iscoroutine(result):
299 asyncio.ensure_future(self._await_close(result, errors))
300 except Exception as exc:
301 errors.append(exc)
303 @staticmethod
304 async def _await_close(coro, errors: list[Exception]):
305 try:
306 await coro
307 except Exception as exc:
308 errors.append(exc)
310 async def close(self) -> None:
311 """Close all resources in the pool."""
312 if self._closed:
313 return
314 self._closed = True
316 if self._health_task:
317 self._health_task.cancel()
318 try:
319 await self._health_task
320 except asyncio.CancelledError:
321 pass
323 # Close all resources
324 errors: list[Exception] = []
326 # Close in-use
327 for rid, resource in list(self._in_use.items()):
328 self._close_resource(resource, errors)
329 self._in_use.clear()
331 # Close available
332 while not self._available.empty():
333 try:
334 resource = self._available.get_nowait()
335 self._close_resource(resource, errors)
336 except asyncio.QueueEmpty:
337 break
338 except Exception as exc:
339 errors.append(exc)
341 if errors:
342 logger.error("Pool '%s': %d errors during close", self._name, len(errors))
344 async def health_check(self) -> bool:
345 return not self._closed
347 @property
348 def stats(self) -> dict[str, Any]:
349 """Pool statistics."""
350 return {
351 "name": self._name,
352 "total": self._total,
353 "in_use": len(self._in_use),
354 "available": self._available.qsize(),
355 "max_size": self._max_size,
356 "closed": self._closed,
357 }
360# ============================================================================
361# Resource Manager (central registry)
362# ============================================================================
365class ResourceManager:
366 """Central resource lifecycle manager.
368 Tracks all resources, ensures ordered shutdown (LIFO),
369 provides leak detection.
371 Usage:
372 rm = ResourceManager()
373 conn = await rm.register("db", managed_pg_conn)
374 # ... use conn ...
375 await rm.shutdown() # closes all in reverse order
376 """
378 def __init__(self, leak_warn_threshold: float = 300.0):
379 self._resources: list[tuple[str, AbstractResource]] = []
380 self._lock = asyncio.Lock()
381 self._leak_threshold = leak_warn_threshold
382 self._shutting_down = False
383 self._finalizers: list[Callable[[], Any]] = []
385 async def register(self, name: str, resource: AbstractResource) -> AbstractResource:
386 """Register a resource for lifecycle management."""
387 async with self._lock:
388 if self._shutting_down:
389 raise RuntimeError(
390 "ResourceManager is shutting down — cannot register new resources"
391 )
392 self._resources.append((name, resource))
393 return resource
395 async def unregister(self, name: str) -> AbstractResource | None:
396 """Remove a resource from management (without closing)."""
397 async with self._lock:
398 for i, (n, r) in enumerate(self._resources):
399 if n == name:
400 self._resources.pop(i)
401 return r
402 return None
404 def add_finalizer(self, fn: Callable[[], Any]):
405 """Register a finalizer to run during shutdown."""
406 self._finalizers.append(fn)
408 async def get(self, name: str) -> AbstractResource | None:
409 """Find a managed resource by name."""
410 for n, r in self._resources:
411 if n == name:
412 return r
413 return None
415 async def shutdown(self, timeout: float = 30.0) -> list[str]:
416 """Ordered shutdown — LIFO order, with timeout per resource.
418 Returns list of resource names that failed to close.
419 """
420 self._shutting_down = True
421 failures: list[str] = []
423 # Close resources in reverse order
424 async with self._lock:
425 resources = list(reversed(self._resources))
426 self._resources.clear()
428 for name, resource in resources:
429 try:
430 await asyncio.wait_for(resource.close(), timeout=timeout)
431 logger.debug("Closed resource: %s", name)
432 except TimeoutError:
433 failures.append(f"{name} (timeout)")
434 logger.error("Resource '%s' close timed out after %.0fs", name, timeout)
435 except Exception as exc:
436 failures.append(f"{name} ({exc})")
437 logger.error("Failed to close resource '%s': %s", name, exc)
439 # Run finalizers
440 for fn in self._finalizers:
441 try:
442 result = fn()
443 if asyncio.iscoroutine(result):
444 await result
445 except Exception as exc:
446 logger.error("Finalizer failed: %s", exc)
448 return failures
450 async def health_report(self) -> dict[str, bool]:
451 """Health check all registered resources."""
452 report: dict[str, bool] = {}
453 async with self._lock:
454 for name, resource in self._resources:
455 try:
456 report[name] = await resource.health_check()
457 except Exception:
458 report[name] = False
459 return report
461 def check_leaks(self) -> list[str]:
462 """Check for resources held beyond leak threshold."""
463 now = time.monotonic()
464 leaks: list[str] = []
465 for name, resource in self._resources:
466 if hasattr(resource, "info") and hasattr(resource.info, "acquired_at"):
467 age = now - resource.info.acquired_at
468 if age > self._leak_threshold and resource.info.state != ResourceState.CLOSED:
469 leaks.append(f"{name} (held {age:.0f}s)")
470 return leaks
472 @property
473 def size(self) -> int:
474 return len(self._resources)
477# ============================================================================
478# Global singleton
479# ============================================================================
481_global_resource_manager: ResourceManager | None = None
484def get_resource_manager() -> ResourceManager:
485 """Get or create the global ResourceManager singleton."""
486 global _global_resource_manager
487 if _global_resource_manager is None:
488 _global_resource_manager = ResourceManager()
489 return _global_resource_manager