Coverage for agentos/core/resource_manager.py: 81%
287 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 11:37 +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 contextlib import asynccontextmanager
20from dataclasses import dataclass
21from enum import Enum
22from typing import (
23 Any,
24 AsyncIterator,
25 Callable,
26 Dict,
27 List,
28 Optional,
29 Tuple,
30 TypeVar,
31)
33logger = logging.getLogger(__name__)
35T = TypeVar("T")
38# ============================================================================
39# Data types
40# ============================================================================
42class ResourceState(str, Enum):
43 CREATED = "created"
44 ACTIVE = "active"
45 CLOSING = "closing"
46 CLOSED = "closed"
47 ERROR = "error"
50class ResourceType(str, Enum):
51 CONNECTION = "connection"
52 POOL = "pool"
53 FILE = "file"
54 LOCK = "lock"
55 SESSION = "session"
56 OTHER = "other"
59@dataclass
60class ResourceInfo:
61 """Metadata about a managed resource."""
63 name: str
64 resource_type: ResourceType = ResourceType.OTHER
65 state: ResourceState = ResourceState.CREATED
66 acquired_at: float = 0.0
67 released_at: float = 0.0
68 leak_warn_threshold: float = 300.0 # Warn if held > 5 min
71# ============================================================================
72# Abstract Resource
73# ============================================================================
75class AbstractResource(ABC):
76 """Interface for managed resources."""
78 @abstractmethod
79 async def close(self) -> None:
80 """Release the underlying resource."""
82 async def health_check(self) -> bool:
83 """Optional health check. Default: assume healthy if not closed."""
84 return True
86 def on_leak_detected(self) -> None:
87 """Callback when resource appears to be leaked."""
88 logger.warning("Resource leak suspected: %s", self)
91# ============================================================================
92# Managed Resource wrapper
93# ============================================================================
95class ManagedResource(AbstractResource):
96 """Wraps any async-closeable object into a managed resource."""
98 def __init__(
99 self,
100 resource: Any,
101 close_fn: Optional[Callable[[Any], Any]] = None,
102 name: str = "unnamed",
103 resource_type: ResourceType = ResourceType.OTHER,
104 ):
105 self._resource = resource
106 self._close_fn = close_fn or getattr(resource, "close", None) or getattr(resource, "aclose", None)
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(self._resource, "ping", None)
138 if health_fn is not None:
139 try:
140 result = health_fn()
141 if asyncio.iscoroutine(result):
142 result = await result
143 return bool(result)
144 except Exception:
145 return False
146 return True
149# ============================================================================
150# Resource Pool
151# ============================================================================
153class ResourcePool(AbstractResource):
154 """Generic async resource pool with health-checked lend/return.
156 Usage:
157 pool = ResourcePool(factory=create_db_conn, max_size=10)
158 async with pool.acquire() as conn:
159 await conn.query(...)
160 """
162 def __init__(
163 self,
164 factory: Callable[[], Any],
165 max_size: int = 10,
166 min_size: int = 0,
167 idle_timeout: float = 300.0,
168 health_check_interval: float = 30.0,
169 name: str = "pool",
170 ):
171 self._factory = factory
172 self._max_size = max_size
173 self._min_size = min_size
174 self._idle_timeout = idle_timeout
175 self._health_check_interval = health_check_interval
176 self._name = name
178 self._available: asyncio.Queue = asyncio.Queue(maxsize=max_size)
179 self._in_use: Dict[int, Any] = {} # id(resource) → resource
180 self._total: int = 0
181 self._lock = asyncio.Lock()
182 self._closed = False
183 self._health_task: Optional[asyncio.Task] = None
185 async def _prefill(self):
186 """Pre-create min_size connections."""
187 for _ in range(self._min_size):
188 resource = self._factory()
189 if asyncio.iscoroutine(resource):
190 resource = await resource
191 await self._available.put(resource)
192 self._total += 1
194 async def start(self):
195 """Initialize the pool."""
196 await self._prefill()
197 if self._health_check_interval > 0:
198 self._health_task = asyncio.ensure_future(self._health_loop())
200 @asynccontextmanager
201 async def acquire(self) -> AsyncIterator[Any]:
202 """Acquire a resource from the pool. Returns via context manager."""
203 resource = await self._get()
204 try:
205 yield resource
206 finally:
207 await self._return(resource)
209 async def _get(self) -> Any:
210 if self._closed:
211 raise RuntimeError(f"ResourcePool '{self._name}' is closed")
213 # Try to get from available
214 try:
215 resource = self._available.get_nowait()
216 self._in_use[id(resource)] = resource
217 return resource
218 except asyncio.QueueEmpty:
219 pass
221 # Try to create new
222 async with self._lock:
223 if self._total < self._max_size:
224 resource = self._factory()
225 if asyncio.iscoroutine(resource):
226 resource = await resource
227 self._total += 1
228 self._in_use[id(resource)] = resource
229 return resource
231 # Wait for one to become available
232 resource = await self._available.get()
233 self._in_use[id(resource)] = resource
234 return resource
236 async def _return(self, resource: Any):
237 self._in_use.pop(id(resource), None)
238 if not self._closed:
239 await self._available.put(resource)
241 async def _health_loop(self):
242 """Periodic health check, remove dead connections."""
243 while not self._closed:
244 await asyncio.sleep(self._health_check_interval)
245 # Drain and re-check
246 healthy: List[Any] = []
247 while not self._available.empty():
248 try:
249 resource = self._available.get_nowait()
250 if self._is_healthy(resource):
251 healthy.append(resource)
252 else:
253 self._total -= 1
254 logger.debug("Pool '%s': removed unhealthy connection", self._name)
255 except asyncio.QueueEmpty:
256 break
258 for resource in healthy:
259 await self._available.put(resource)
261 # Top up to min_size
262 while self._total < self._min_size and not self._closed:
263 async with self._lock:
264 if self._total < self._max_size:
265 resource = self._factory()
266 if asyncio.iscoroutine(resource):
267 resource = await resource
268 await self._available.put(resource)
269 self._total += 1
271 def _is_healthy(self, resource: Any) -> bool:
272 """Check if a resource is healthy."""
273 health_fn = getattr(resource, "health_check", None) or getattr(resource, "ping", None)
274 if health_fn is None:
275 return True
276 try:
277 result = health_fn()
278 if asyncio.iscoroutine(result):
279 return True # Can't check async in sync context
280 return bool(result)
281 except Exception:
282 return False
284 def _close_resource(self, resource: Any, errors: List[Exception]):
285 """Attempt to close a single resource, collecting errors."""
286 close_fn = (
287 getattr(resource, "close", None)
288 or getattr(resource, "aclose", None)
289 or (resource.get("close") if isinstance(resource, dict) else None)
290 or (resource.get("aclose") if isinstance(resource, dict) else None)
291 )
292 if close_fn is not None:
293 try:
294 result = close_fn()
295 if asyncio.iscoroutine(result):
296 asyncio.ensure_future(self._await_close(result, errors))
297 except Exception as exc:
298 errors.append(exc)
300 @staticmethod
301 async def _await_close(coro, errors: List[Exception]):
302 try:
303 await coro
304 except Exception as exc:
305 errors.append(exc)
307 async def close(self) -> None:
308 """Close all resources in the pool."""
309 if self._closed:
310 return
311 self._closed = True
313 if self._health_task:
314 self._health_task.cancel()
315 try:
316 await self._health_task
317 except asyncio.CancelledError:
318 pass
320 # Close all resources
321 errors: List[Exception] = []
323 # Close in-use
324 for rid, resource in list(self._in_use.items()):
325 self._close_resource(resource, errors)
326 self._in_use.clear()
328 # Close available
329 while not self._available.empty():
330 try:
331 resource = self._available.get_nowait()
332 self._close_resource(resource, errors)
333 except asyncio.QueueEmpty:
334 break
335 except Exception as exc:
336 errors.append(exc)
338 if errors:
339 logger.error(
340 "Pool '%s': %d errors during close", self._name, len(errors)
341 )
343 async def health_check(self) -> bool:
344 return not self._closed
346 @property
347 def stats(self) -> Dict[str, Any]:
348 """Pool statistics."""
349 return {
350 "name": self._name,
351 "total": self._total,
352 "in_use": len(self._in_use),
353 "available": self._available.qsize(),
354 "max_size": self._max_size,
355 "closed": self._closed,
356 }
359# ============================================================================
360# Resource Manager (central registry)
361# ============================================================================
363class ResourceManager:
364 """Central resource lifecycle manager.
366 Tracks all resources, ensures ordered shutdown (LIFO),
367 provides leak detection.
369 Usage:
370 rm = ResourceManager()
371 conn = await rm.register("db", managed_pg_conn)
372 # ... use conn ...
373 await rm.shutdown() # closes all in reverse order
374 """
376 def __init__(self, leak_warn_threshold: float = 300.0):
377 self._resources: List[Tuple[str, AbstractResource]] = []
378 self._lock = asyncio.Lock()
379 self._leak_threshold = leak_warn_threshold
380 self._shutting_down = False
381 self._finalizers: List[Callable[[], Any]] = []
383 async def register(self, name: str, resource: AbstractResource) -> AbstractResource:
384 """Register a resource for lifecycle management."""
385 async with self._lock:
386 if self._shutting_down:
387 raise RuntimeError("ResourceManager is shutting down — cannot register new resources")
388 self._resources.append((name, resource))
389 return resource
391 async def unregister(self, name: str) -> Optional[AbstractResource]:
392 """Remove a resource from management (without closing)."""
393 async with self._lock:
394 for i, (n, r) in enumerate(self._resources):
395 if n == name:
396 self._resources.pop(i)
397 return r
398 return None
400 def add_finalizer(self, fn: Callable[[], Any]):
401 """Register a finalizer to run during shutdown."""
402 self._finalizers.append(fn)
404 async def get(self, name: str) -> Optional[AbstractResource]:
405 """Find a managed resource by name."""
406 for n, r in self._resources:
407 if n == name:
408 return r
409 return None
411 async def shutdown(self, timeout: float = 30.0) -> List[str]:
412 """Ordered shutdown — LIFO order, with timeout per resource.
414 Returns list of resource names that failed to close.
415 """
416 self._shutting_down = True
417 failures: List[str] = []
419 # Close resources in reverse order
420 async with self._lock:
421 resources = list(reversed(self._resources))
422 self._resources.clear()
424 for name, resource in resources:
425 try:
426 await asyncio.wait_for(resource.close(), timeout=timeout)
427 logger.debug("Closed resource: %s", name)
428 except asyncio.TimeoutError:
429 failures.append(f"{name} (timeout)")
430 logger.error("Resource '%s' close timed out after %.0fs", name, timeout)
431 except Exception as exc:
432 failures.append(f"{name} ({exc})")
433 logger.error("Failed to close resource '%s': %s", name, exc)
435 # Run finalizers
436 for fn in self._finalizers:
437 try:
438 result = fn()
439 if asyncio.iscoroutine(result):
440 await result
441 except Exception as exc:
442 logger.error("Finalizer failed: %s", exc)
444 return failures
446 async def health_report(self) -> Dict[str, bool]:
447 """Health check all registered resources."""
448 report: Dict[str, bool] = {}
449 async with self._lock:
450 for name, resource in self._resources:
451 try:
452 report[name] = await resource.health_check()
453 except Exception:
454 report[name] = False
455 return report
457 def check_leaks(self) -> List[str]:
458 """Check for resources held beyond leak threshold."""
459 now = time.monotonic()
460 leaks: List[str] = []
461 for name, resource in self._resources:
462 if hasattr(resource, "info") and hasattr(resource.info, "acquired_at"):
463 age = now - resource.info.acquired_at
464 if age > self._leak_threshold and resource.info.state != ResourceState.CLOSED:
465 leaks.append(f"{name} (held {age:.0f}s)")
466 return leaks
468 @property
469 def size(self) -> int:
470 return len(self._resources)
473# ============================================================================
474# Global singleton
475# ============================================================================
477_global_resource_manager: Optional[ResourceManager] = None
480def get_resource_manager() -> ResourceManager:
481 """Get or create the global ResourceManager singleton."""
482 global _global_resource_manager
483 if _global_resource_manager is None:
484 _global_resource_manager = ResourceManager()
485 return _global_resource_manager