Coverage for agentos/tools/memory_optimizer.py: 32%
213 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 12:29 +0800
1"""
2Memory Optimization Tools for AgentOS.
3Object pooling, LRU caching, memory monitoring, and smart caching with TTL.
4"""
6import threading
7import time
8from collections import OrderedDict
9from collections.abc import Callable
10from dataclasses import dataclass
11from typing import Any, Generic, TypeVar
13T = TypeVar("T")
16# ============================================================================
17# ObjectPool
18# ============================================================================
21class ObjectPool(Generic[T]):
22 """Thread-safe object pool with auto-expiry and size limits.
24 Reuses pre-allocated objects instead of creating/destroying them repeatedly.
25 """
27 def __init__(
28 self,
29 factory: Callable[[], T],
30 max_size: int = 100,
31 max_idle: int = 30,
32 idle_timeout: float = 300.0,
33 ):
34 self._factory = factory
35 self._max_size = max_size
36 self._max_idle = max_idle
37 self._idle_timeout = idle_timeout
38 self._pool: list[_PooledItem[T]] = []
39 self._lock = threading.Lock()
40 self._created: int = 0
41 self._borrowed: int = 0
42 self._returned: int = 0
44 def acquire(self) -> T:
45 """Borrow an object from the pool or create a new one."""
46 with self._lock:
47 now = time.monotonic()
48 self._evict_expired(now)
50 if self._pool:
51 item = self._pool.pop()
52 item.idle = False
53 self._borrowed += 1
54 return item.obj
56 if self._created < self._max_size:
57 self._created += 1
58 self._borrowed += 1
59 return self._factory()
61 # Pool fully allocated; create temporary object outside pool
62 self._borrowed += 1
63 return self._factory()
65 def release(self, obj: T) -> None:
66 """Return an object to the pool for reuse."""
67 with self._lock:
68 self._returned += 1
69 self._evict_expired(time.monotonic())
71 if len(self._pool) < self._max_idle:
72 self._pool.append(_PooledItem(obj=obj, idle=True, acquired_at=time.monotonic()))
73 # else: discard excess objects
75 def _evict_expired(self, now: float) -> None:
76 self._pool[:] = [item for item in self._pool if now - item.acquired_at < self._idle_timeout]
78 @property
79 def stats(self) -> dict[str, Any]:
80 with self._lock:
81 return {
82 "created": self._created,
83 "borrowed": self._borrowed,
84 "returned": self._returned,
85 "idle": len(self._pool),
86 "active": self._borrowed - self._returned,
87 }
89 def __len__(self) -> int:
90 return len(self._pool)
93@dataclass
94class _PooledItem(Generic[T]):
95 obj: T
96 idle: bool
97 acquired_at: float
100# ============================================================================
101# LRUCache
102# ============================================================================
105class LRUCache(Generic[T]):
106 """Thread-safe LRU cache with capacity limit and optional TTL."""
108 def __init__(self, capacity: int = 1024, ttl: float | None = None):
109 self._capacity = capacity
110 self._ttl = ttl
111 self._cache: OrderedDict[str, _CacheEntry[T]] = OrderedDict()
112 self._lock = threading.Lock()
113 self._hits: int = 0
114 self._misses: int = 0
115 self._evictions: int = 0
117 def get(self, key: str) -> T | None:
118 with self._lock:
119 entry = self._cache.get(key)
120 if entry is None:
121 self._misses += 1
122 return None
123 if self._ttl and time.monotonic() - entry.timestamp > self._ttl:
124 del self._cache[key]
125 self._misses += 1
126 self._evictions += 1
127 return None
128 self._cache.move_to_end(key)
129 self._hits += 1
130 return entry.value
132 def put(self, key: str, value: T) -> None:
133 with self._lock:
134 if key in self._cache:
135 self._cache.move_to_end(key)
136 self._cache[key] = _CacheEntry(value=value, timestamp=time.monotonic())
137 return
138 if len(self._cache) >= self._capacity:
139 self._cache.popitem(last=False)
140 self._evictions += 1
141 self._cache[key] = _CacheEntry(value=value, timestamp=time.monotonic())
143 def remove(self, key: str) -> bool:
144 with self._lock:
145 if key in self._cache:
146 del self._cache[key]
147 return True
148 return False
150 def clear(self) -> None:
151 with self._lock:
152 self._cache.clear()
154 @property
155 def hit_rate(self) -> float:
156 total = self._hits + self._misses
157 return self._hits / total if total > 0 else 0.0
159 @property
160 def stats(self) -> dict[str, Any]:
161 with self._lock:
162 return {
163 "size": len(self._cache),
164 "capacity": self._capacity,
165 "hits": self._hits,
166 "misses": self._misses,
167 "hit_rate": round(self.hit_rate, 4),
168 "evictions": self._evictions,
169 }
171 def __len__(self) -> int:
172 return len(self._cache)
174 def __contains__(self, key: str) -> bool:
175 return key in self._cache
178@dataclass
179class _CacheEntry(Generic[T]):
180 value: T
181 timestamp: float
184# ============================================================================
185# SmartCache
186# ============================================================================
189class SmartCache(Generic[T]):
190 """Multi-tier cache with compute-on-miss and automatic invalidation."""
192 def __init__(self, compute: Callable[[str], T], capacity: int = 1024, ttl: float = 300.0):
193 self._lru = LRUCache[T](capacity=capacity, ttl=ttl)
194 self._compute = compute
195 self._lock = threading.Lock()
197 def get(self, key: str) -> T:
198 """Get from cache or compute and cache on miss."""
199 cached = self._lru.get(key)
200 if cached is not None:
201 return cached
202 with self._lock:
203 # Double-check after acquiring lock
204 cached = self._lru.get(key)
205 if cached is not None:
206 return cached
207 value = self._compute(key)
208 self._lru.put(key, value)
209 return value
211 def prefetch(self, keys: list[str]) -> int:
212 """Pre-compute and cache values for a list of keys. Returns count cached."""
213 count = 0
214 for key in keys:
215 if key not in self._lru:
216 try:
217 self.get(key)
218 count += 1
219 except Exception:
220 pass
221 return count
223 def invalidate(self, key: str) -> bool:
224 return self._lru.remove(key)
226 def invalidate_pattern(self, pattern: str) -> int:
227 """Invalidate all keys containing pattern substring. Returns count removed."""
228 count = 0
229 for key in list(self._lru._cache.keys()):
230 if pattern in key:
231 if self._lru.remove(key):
232 count += 1
233 return count
235 def clear(self) -> None:
236 self._lru.clear()
238 @property
239 def stats(self) -> dict[str, Any]:
240 return self._lru.stats
242 def __len__(self) -> int:
243 return len(self._lru)
246# ============================================================================
247# MemoryMonitor
248# ============================================================================
251class MemoryMonitor:
252 """Monitor per-component memory usage with high-water mark tracking."""
254 _singleton = None
255 _lock = threading.Lock()
257 def __new__(cls):
258 if cls._singleton is None:
259 with cls._lock:
260 if cls._singleton is None:
261 cls._singleton = super().__new__(cls)
262 cls._singleton._initialized = False
263 return cls._singleton
265 def __init__(self):
266 if self._initialized:
267 return
268 self._initialized = True
269 self._components: dict[str, _ComponentMetrics] = {}
270 self._lock = threading.Lock()
272 def register(self, name: str) -> None:
273 with self._lock:
274 if name not in self._components:
275 self._components[name] = _ComponentMetrics(name=name)
277 def record_alloc(self, name: str, size_bytes: int) -> None:
278 with self._lock:
279 comp = self._components.get(name)
280 if comp:
281 comp.current_bytes += size_bytes
282 comp.total_allocations += 1
283 comp.peak_bytes = max(comp.peak_bytes, comp.current_bytes)
285 def record_free(self, name: str, size_bytes: int) -> None:
286 with self._lock:
287 comp = self._components.get(name)
288 if comp:
289 comp.current_bytes = max(0, comp.current_bytes - size_bytes)
291 def snapshot(self) -> dict[str, dict[str, Any]]:
292 with self._lock:
293 return {name: comp.to_dict() for name, comp in self._components.items()}
295 def alert(self, name: str, threshold_bytes: int) -> bool:
296 """Check if a component exceeds memory threshold."""
297 with self._lock:
298 comp = self._components.get(name)
299 if comp:
300 return comp.current_bytes > threshold_bytes
301 return False
303 @property
304 def total_current(self) -> int:
305 with self._lock:
306 return sum(c.current_bytes for c in self._components.values())
309@dataclass
310class _ComponentMetrics:
311 name: str
312 current_bytes: int = 0
313 peak_bytes: int = 0
314 total_allocations: int = 0
316 def to_dict(self) -> dict[str, Any]:
317 return {
318 "name": self.name,
319 "current_bytes": self.current_bytes,
320 "peak_bytes": self.peak_bytes,
321 "total_allocations": self.total_allocations,
322 }
325# ============================================================================
326# Convenience Functions
327# ============================================================================
330def create_object_pool(
331 factory: Callable[[], T],
332 max_size: int = 100,
333 max_idle: int = 30,
334 idle_timeout: float = 300.0,
335) -> ObjectPool[T]:
336 """Create a thread-safe object pool."""
337 return ObjectPool(factory, max_size=max_size, max_idle=max_idle, idle_timeout=idle_timeout)
340def create_lru_cache(capacity: int = 1024, ttl: float | None = None) -> LRUCache[Any]:
341 """Create a thread-safe LRU cache."""
342 return LRUCache(capacity=capacity, ttl=ttl)
345def create_smart_cache(
346 compute: Callable[[str], T],
347 capacity: int = 1024,
348 ttl: float = 300.0,
349) -> SmartCache[T]:
350 """Create a smart cache with compute-on-miss."""
351 return SmartCache(compute, capacity=capacity, ttl=ttl)
354def get_memory_monitor() -> MemoryMonitor:
355 """Get the singleton memory monitor."""
356 return MemoryMonitor()