Coverage for agentos/tools/connection_pool.py: 0%
234 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-08 20:40 +0800
1"""
2Connection Pooling & Resource Management for AgentOS.
3Generic connection pool, rate limiter, resource quota manager, and health-checked pools.
4"""
6import threading
7import time
8from collections import deque
9from collections.abc import Callable
10from dataclasses import dataclass
11from typing import Any, Generic, TypeVar
13T = TypeVar("T")
16# ============================================================================
17# ConnectionPool
18# ============================================================================
21@dataclass
22class _PooledConn(Generic[T]):
23 conn: T
24 created_at: float
25 last_used: float
26 borrowed: bool = False
29class ConnectionPool(Generic[T]):
30 """Thread-safe generic connection pool with health checks and idle eviction.
32 Supports: min/max sizing, health validation, auto-reconnect, idle timeout.
33 """
35 def __init__(
36 self,
37 factory: Callable[[], T],
38 health_check: Callable[[T], bool] | None = None,
39 closer: Callable[[T], None] | None = None,
40 min_size: int = 2,
41 max_size: int = 20,
42 max_idle: int = 10,
43 idle_timeout: float = 300.0,
44 checkout_timeout: float = 30.0,
45 ):
46 self._factory = factory
47 self._health_check = health_check
48 self._closer = closer
49 self._min_size = min_size
50 self._max_size = max_size
51 self._max_idle = max_idle
52 self._idle_timeout = idle_timeout
53 self._checkout_timeout = checkout_timeout
54 self._pool: deque[_PooledConn[T]] = deque()
55 self._lock = threading.RLock()
56 self._condition = threading.Condition(self._lock)
57 self._total_created: int = 0
58 self._total_borrowed: int = 0
59 self._total_returned: int = 0
60 self._total_failed_health: int = 0
61 self._closed: bool = False
63 def _create(self) -> _PooledConn[T]:
64 conn = self._factory()
65 self._total_created += 1
66 now = time.monotonic()
67 return _PooledConn(conn=conn, created_at=now, last_used=now)
69 def _validate(self, pc: _PooledConn[T]) -> bool:
70 if self._health_check is None:
71 return True
72 try:
73 ok = self._health_check(pc.conn)
74 if not ok:
75 self._total_failed_health += 1
76 return ok
77 except Exception:
78 self._total_failed_health += 1
79 return False
81 def acquire(self, timeout: float | None = None) -> T:
82 """Borrow a connection from the pool. Blocks until available or timeout."""
83 if timeout is None:
84 timeout = self._checkout_timeout
86 with self._condition:
87 deadline = time.monotonic() + timeout
89 while True:
90 if self._closed:
91 raise RuntimeError("ConnectionPool is closed")
93 # Find a valid idle connection; evict unhealthy ones
94 unhealthy: list[_PooledConn[T]] = []
95 for pc in self._pool:
96 if pc.borrowed:
97 continue
98 if self._validate(pc):
99 pc.borrowed = True
100 pc.last_used = time.monotonic()
101 self._total_borrowed += 1
102 return pc.conn
103 else:
104 unhealthy.append(pc)
106 for pc in unhealthy:
107 if pc in self._pool:
108 self._pool.remove(pc)
109 self._close_conn(pc)
111 # Try to create a new one if under max
112 active = sum(1 for pc in self._pool if pc.borrowed)
113 if active + len([pc for pc in self._pool if not pc.borrowed]) < self._max_size:
114 pc = self._create()
115 pc.borrowed = True
116 self._total_borrowed += 1
117 self._pool.append(pc)
118 return pc.conn
120 # Wait for a connection to be returned
121 remaining = deadline - time.monotonic()
122 if remaining <= 0:
123 raise TimeoutError(f"Timed out waiting for connection after {timeout}s")
124 self._condition.wait(timeout=min(remaining, 1.0))
126 def release(self, conn: T) -> None:
127 """Return a connection to the pool."""
128 with self._condition:
129 for pc in self._pool:
130 if pc.conn is conn:
131 pc.borrowed = False
132 pc.last_used = time.monotonic()
133 self._total_returned += 1
134 self._condition.notify()
135 return
136 # Connection not in pool — close it
137 self._close_raw(conn)
138 self._condition.notify()
140 def _close_conn(self, pc: _PooledConn[T]) -> None:
141 self._close_raw(pc.conn)
143 def _close_raw(self, conn: T) -> None:
144 if self._closer:
145 try:
146 self._closer(conn)
147 except Exception:
148 pass
150 def warm_up(self) -> int:
151 """Pre-create connections up to min_size. Returns number created."""
152 count = 0
153 with self._condition:
154 idle = sum(1 for pc in self._pool if not pc.borrowed)
155 needed = self._min_size - idle
156 for _ in range(needed):
157 self._pool.append(self._create())
158 count += 1
159 return count
161 def evict_idle(self) -> int:
162 """Remove idle connections past timeout. Returns number evicted."""
163 now = time.monotonic()
164 count = 0
165 with self._condition:
166 # Preserve min_size
167 idle = [pc for pc in self._pool if not pc.borrowed]
168 to_keep = self._min_size
169 old_first = sorted(idle, key=lambda pc: pc.last_used)
170 for pc in old_first[to_keep:]:
171 if now - pc.last_used > self._idle_timeout:
172 self._pool.remove(pc)
173 self._close_conn(pc)
174 count += 1
175 return count
177 def close(self) -> None:
178 """Close all connections and shut down the pool."""
179 with self._condition:
180 self._closed = True
181 for pc in self._pool:
182 self._close_conn(pc)
183 self._pool.clear()
184 self._condition.notify_all()
186 @property
187 def stats(self) -> dict[str, Any]:
188 with self._lock:
189 active = sum(1 for pc in self._pool if pc.borrowed)
190 idle = sum(1 for pc in self._pool if not pc.borrowed)
191 return {
192 "total_created": self._total_created,
193 "total_borrowed": self._total_borrowed,
194 "total_returned": self._total_returned,
195 "active": active,
196 "idle": idle,
197 "total": len(self._pool),
198 "failed_health_checks": self._total_failed_health,
199 "capacity": self._max_size,
200 }
202 def __enter__(self):
203 return self
205 def __exit__(self, *args):
206 self.close()
209# ============================================================================
210# RateLimiter
211# ============================================================================
214class RateLimiter:
215 """Thread-safe token bucket rate limiter with burst support."""
217 def __init__(self, rate: float, burst: int = 1):
218 """rate: tokens per second. burst: max tokens accumulated."""
219 self._rate = rate
220 self._burst = burst
221 self._tokens: float = burst
222 self._last_refill: float = time.monotonic()
223 self._lock = threading.Lock()
224 self._total_acquired: int = 0
225 self._total_rejected: int = 0
227 def acquire(self, count: int = 1, timeout: float | None = None) -> bool:
228 """Try to acquire N tokens. Blocks up to timeout if not enough."""
229 deadline = time.monotonic() + timeout if timeout else None
231 with self._lock:
232 while True:
233 self._refill()
234 if self._tokens >= count:
235 self._tokens -= count
236 self._total_acquired += count
237 return True
239 if deadline and time.monotonic() >= deadline:
240 self._total_rejected += count
241 return False
243 # Wait for refill
244 wait_time = (count - self._tokens) / self._rate
245 if deadline:
246 wait_time = min(wait_time, deadline - time.monotonic())
247 if wait_time <= 0:
248 self._total_rejected += count
249 return False
251 # Release lock during wait
252 self._lock.release()
253 try:
254 time.sleep(wait_time)
255 finally:
256 self._lock.acquire()
258 def try_acquire(self, count: int = 1) -> bool:
259 """Non-blocking attempt to acquire tokens."""
260 with self._lock:
261 self._refill()
262 if self._tokens >= count:
263 self._tokens -= count
264 self._total_acquired += count
265 return True
266 self._total_rejected += count
267 return False
269 def _refill(self) -> None:
270 now = time.monotonic()
271 elapsed = now - self._last_refill
272 self._tokens = min(self._burst, self._tokens + elapsed * self._rate)
273 self._last_refill = now
275 @property
276 def available(self) -> float:
277 with self._lock:
278 self._refill()
279 return self._tokens
281 @property
282 def stats(self) -> dict[str, Any]:
283 with self._lock:
284 self._refill()
285 return {
286 "rate": self._rate,
287 "burst": self._burst,
288 "tokens_available": round(self._tokens, 2),
289 "total_acquired": self._total_acquired,
290 "total_rejected": self._total_rejected,
291 }
294# ============================================================================
295# ResourceQuota
296# ============================================================================
299class ResourceQuota:
300 """Track and enforce resource usage quotas per component."""
302 def __init__(self, global_limit: int = 1024):
303 self._global_limit = global_limit
304 self._allocations: dict[str, int] = {}
305 self._lock = threading.Lock()
307 def allocate(self, component: str, amount: int = 1) -> bool:
308 """Try to allocate resources. Returns True if successful."""
309 with self._lock:
310 current_total = sum(self._allocations.values())
311 if current_total + amount > self._global_limit:
312 return False
313 self._allocations[component] = self._allocations.get(component, 0) + amount
314 return True
316 def release(self, component: str, amount: int = 1) -> None:
317 with self._lock:
318 current = self._allocations.get(component, 0)
319 self._allocations[component] = max(0, current - amount)
321 def set_limit(self, component: str, limit: int) -> None:
322 with self._lock:
323 current = self._allocations.get(component, 0)
324 if current > limit:
325 self._allocations[component] = limit
327 def get_usage(self, component: str) -> int:
328 with self._lock:
329 return self._allocations.get(component, 0)
331 @property
332 def total_used(self) -> int:
333 with self._lock:
334 return sum(self._allocations.values())
336 @property
337 def remaining(self) -> int:
338 return self._global_limit - self.total_used
340 @property
341 def stats(self) -> dict[str, Any]:
342 with self._lock:
343 return {
344 "global_limit": self._global_limit,
345 "total_used": self.total_used,
346 "remaining": self.remaining,
347 "allocations": dict(self._allocations),
348 }
351# ============================================================================
352# Convenience Functions
353# ============================================================================
356def create_connection_pool(
357 factory: Callable[[], T],
358 health_check: Callable[[T], bool] | None = None,
359 closer: Callable[[T], None] | None = None,
360 min_size: int = 2,
361 max_size: int = 20,
362) -> ConnectionPool[T]:
363 """Create a thread-safe connection pool."""
364 return ConnectionPool(
365 factory,
366 health_check=health_check,
367 closer=closer,
368 min_size=min_size,
369 max_size=max_size,
370 )
373def create_rate_limiter(rate: float, burst: int = 10) -> RateLimiter:
374 """Create a token bucket rate limiter."""
375 return RateLimiter(rate=rate, burst=burst)
378def create_resource_quota(global_limit: int = 1024) -> ResourceQuota:
379 """Create a resource quota manager."""
380 return ResourceQuota(global_limit=global_limit)