Coverage for agentos/core/resource_manager.py: 81%
288 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 00:18 +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
18import weakref
19from abc import ABC, abstractmethod
20from contextlib import asynccontextmanager
21from dataclasses import dataclass, field
22from enum import Enum
23from typing import (
24 Any,
25 AsyncIterator,
26 Callable,
27 Dict,
28 List,
29 Optional,
30 Set,
31 Tuple,
32 Type,
33 TypeVar,
34)
36logger = logging.getLogger(__name__)
38T = TypeVar("T")
41# ============================================================================
42# Data types
43# ============================================================================
45class ResourceState(str, Enum):
46 CREATED = "created"
47 ACTIVE = "active"
48 CLOSING = "closing"
49 CLOSED = "closed"
50 ERROR = "error"
53class ResourceType(str, Enum):
54 CONNECTION = "connection"
55 POOL = "pool"
56 FILE = "file"
57 LOCK = "lock"
58 SESSION = "session"
59 OTHER = "other"
62@dataclass
63class ResourceInfo:
64 """Metadata about a managed resource."""
66 name: str
67 resource_type: ResourceType = ResourceType.OTHER
68 state: ResourceState = ResourceState.CREATED
69 acquired_at: float = 0.0
70 released_at: float = 0.0
71 leak_warn_threshold: float = 300.0 # Warn if held > 5 min
74# ============================================================================
75# Abstract Resource
76# ============================================================================
78class AbstractResource(ABC):
79 """Interface for managed resources."""
81 @abstractmethod
82 async def close(self) -> None:
83 """Release the underlying resource."""
85 async def health_check(self) -> bool:
86 """Optional health check. Default: assume healthy if not closed."""
87 return True
89 def on_leak_detected(self) -> None:
90 """Callback when resource appears to be leaked."""
91 logger.warning("Resource leak suspected: %s", self)
94# ============================================================================
95# Managed Resource wrapper
96# ============================================================================
98class ManagedResource(AbstractResource):
99 """Wraps any async-closeable object into a managed resource."""
101 def __init__(
102 self,
103 resource: Any,
104 close_fn: Optional[Callable[[Any], Any]] = None,
105 name: str = "unnamed",
106 resource_type: ResourceType = ResourceType.OTHER,
107 ):
108 self._resource = resource
109 self._close_fn = close_fn or getattr(resource, "close", None) or getattr(resource, "aclose", None)
110 self.info = ResourceInfo(name=name, resource_type=resource_type)
111 self.info.acquired_at = time.monotonic()
112 self.info.state = ResourceState.ACTIVE
114 @property
115 def raw(self) -> Any:
116 """Access the underlying resource object."""
117 return self._resource
119 async def close(self) -> None:
120 if self.info.state == ResourceState.CLOSED:
121 return
123 self.info.state = ResourceState.CLOSING
124 try:
125 if self._close_fn is not None:
126 result = self._close_fn()
127 if asyncio.iscoroutine(result):
128 await result
129 except Exception as exc:
130 self.info.state = ResourceState.ERROR
131 logger.error("Failed to close resource %s: %s", self.info.name, exc)
132 raise
133 else:
134 self.info.state = ResourceState.CLOSED
135 self.info.released_at = time.monotonic()
137 async def health_check(self) -> bool:
138 if self.info.state == ResourceState.CLOSED:
139 return False
140 health_fn = getattr(self._resource, "health_check", None) or getattr(self._resource, "ping", None)
141 if health_fn is not None:
142 try:
143 result = health_fn()
144 if asyncio.iscoroutine(result):
145 result = await result
146 return bool(result)
147 except Exception:
148 return False
149 return True
152# ============================================================================
153# Resource Pool
154# ============================================================================
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: Optional[asyncio.Task] = 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(
343 "Pool '%s': %d errors during close", self._name, len(errors)
344 )
346 async def health_check(self) -> bool:
347 return not self._closed
349 @property
350 def stats(self) -> Dict[str, Any]:
351 """Pool statistics."""
352 return {
353 "name": self._name,
354 "total": self._total,
355 "in_use": len(self._in_use),
356 "available": self._available.qsize(),
357 "max_size": self._max_size,
358 "closed": self._closed,
359 }
362# ============================================================================
363# Resource Manager (central registry)
364# ============================================================================
366class ResourceManager:
367 """Central resource lifecycle manager.
369 Tracks all resources, ensures ordered shutdown (LIFO),
370 provides leak detection.
372 Usage:
373 rm = ResourceManager()
374 conn = await rm.register("db", managed_pg_conn)
375 # ... use conn ...
376 await rm.shutdown() # closes all in reverse order
377 """
379 def __init__(self, leak_warn_threshold: float = 300.0):
380 self._resources: List[Tuple[str, AbstractResource]] = []
381 self._lock = asyncio.Lock()
382 self._leak_threshold = leak_warn_threshold
383 self._shutting_down = False
384 self._finalizers: List[Callable[[], Any]] = []
386 async def register(self, name: str, resource: AbstractResource) -> AbstractResource:
387 """Register a resource for lifecycle management."""
388 async with self._lock:
389 if self._shutting_down:
390 raise RuntimeError("ResourceManager is shutting down — cannot register new resources")
391 self._resources.append((name, resource))
392 return resource
394 async def unregister(self, name: str) -> Optional[AbstractResource]:
395 """Remove a resource from management (without closing)."""
396 async with self._lock:
397 for i, (n, r) in enumerate(self._resources):
398 if n == name:
399 self._resources.pop(i)
400 return r
401 return None
403 def add_finalizer(self, fn: Callable[[], Any]):
404 """Register a finalizer to run during shutdown."""
405 self._finalizers.append(fn)
407 async def get(self, name: str) -> Optional[AbstractResource]:
408 """Find a managed resource by name."""
409 for n, r in self._resources:
410 if n == name:
411 return r
412 return None
414 async def shutdown(self, timeout: float = 30.0) -> List[str]:
415 """Ordered shutdown — LIFO order, with timeout per resource.
417 Returns list of resource names that failed to close.
418 """
419 self._shutting_down = True
420 failures: List[str] = []
422 # Close resources in reverse order
423 async with self._lock:
424 resources = list(reversed(self._resources))
425 self._resources.clear()
427 for name, resource in resources:
428 try:
429 await asyncio.wait_for(resource.close(), timeout=timeout)
430 logger.debug("Closed resource: %s", name)
431 except asyncio.TimeoutError:
432 failures.append(f"{name} (timeout)")
433 logger.error("Resource '%s' close timed out after %.0fs", name, timeout)
434 except Exception as exc:
435 failures.append(f"{name} ({exc})")
436 logger.error("Failed to close resource '%s': %s", name, exc)
438 # Run finalizers
439 for fn in self._finalizers:
440 try:
441 result = fn()
442 if asyncio.iscoroutine(result):
443 await result
444 except Exception as exc:
445 logger.error("Finalizer failed: %s", exc)
447 return failures
449 async def health_report(self) -> Dict[str, bool]:
450 """Health check all registered resources."""
451 report: Dict[str, bool] = {}
452 async with self._lock:
453 for name, resource in self._resources:
454 try:
455 report[name] = await resource.health_check()
456 except Exception:
457 report[name] = False
458 return report
460 def check_leaks(self) -> List[str]:
461 """Check for resources held beyond leak threshold."""
462 now = time.monotonic()
463 leaks: List[str] = []
464 for name, resource in self._resources:
465 if hasattr(resource, "info") and hasattr(resource.info, "acquired_at"):
466 age = now - resource.info.acquired_at
467 if age > self._leak_threshold and resource.info.state != ResourceState.CLOSED:
468 leaks.append(f"{name} (held {age:.0f}s)")
469 return leaks
471 @property
472 def size(self) -> int:
473 return len(self._resources)
476# ============================================================================
477# Global singleton
478# ============================================================================
480_global_resource_manager: Optional[ResourceManager] = None
483def get_resource_manager() -> ResourceManager:
484 """Get or create the global ResourceManager singleton."""
485 global _global_resource_manager
486 if _global_resource_manager is None:
487 _global_resource_manager = ResourceManager()
488 return _global_resource_manager