This guide provides quick reference for the resource management patterns implemented to fix leaks.
Always use this pattern:
session = self.get_session()
try:
# Your database operations here
session.add(object)
session.commit()
return result
except Exception:
session.rollback()
raise
finally:
session.close()❌ DON'T:
with self.get_session() as session: # Missing proper error handling
session.add(object)
session.commit()Always create session lazily and clean up:
async def _get_session(self):
if self.session is None or self.session.closed:
async with self._session_lock:
if self.session is None or self.session.closed:
self.session = aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=30)
)
return self.session
async def cleanup(self):
if self.session and not self.session.closed:
await self.session.close()
await asyncio.sleep(0.25) # Allow connections to close
self.session = NonePrevent unbounded growth:
def __init__(self, max_size=1000):
self.items = []
self.max_size = max_size
def add_item(self, item):
self.items.append(item)
if len(self.items) > self.max_size:
self.items = self.items[-self.max_size:] # Keep only recentImplement expiring caches:
def __init__(self):
self.cache = {}
self.timestamps = {}
self.ttl_seconds = 3600 # 1 hour
def set(self, key, value):
self._clean_expired()
self.cache[key] = value
self.timestamps[key] = datetime.utcnow()
def _clean_expired(self):
now = datetime.utcnow()
expired = [
k for k, t in self.timestamps.items()
if (now - t).total_seconds() > self.ttl_seconds
]
for k in expired:
self.cache.pop(k, None)
self.timestamps.pop(k, None)Track and cleanup background tasks:
def __init__(self):
self._background_tasks = set()
async def start_background_work(self):
task = asyncio.create_task(self._worker())
self._background_tasks.add(task)
task.add_done_callback(self._background_tasks.discard)
async def cleanup(self):
for task in self._background_tasks:
if not task.done():
task.cancel()
if self._background_tasks:
await asyncio.gather(*self._background_tasks, return_exceptions=True)Properly handle WebSocket lifecycle:
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
active_connections.append(websocket)
try:
while True:
try:
await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
except asyncio.TimeoutError:
await websocket.send_json({"type": "ping"})
except WebSocketDisconnect:
logger.info("Client disconnected")
except Exception as e:
logger.error(f"WebSocket error: {e}")
finally:
if websocket in active_connections:
active_connections.remove(websocket)Broadcasting with cleanup:
async def broadcast(message):
dead_connections = []
for conn in active_connections:
try:
await conn.send_json(message)
except Exception:
dead_connections.append(conn)
for conn in dead_connections:
if conn in active_connections:
active_connections.remove(conn)Always close figures:
try:
plt.figure(figsize=(12, 8))
# Your plotting code
plt.savefig(output_path)
plt.close('all')
except Exception as e:
logger.error(f"Plot error: {e}")
finally:
plt.close('all') # Ensure cleanup even on error# BAD
session = self.get_session()
session.add(object)
session.commit()
# Session never closed!# BAD
self.history.append(event) # Grows forever!# BAD
asyncio.create_task(worker()) # Task created but never cleaned up!# BAD
try:
while True:
await websocket.receive_text()
except WebSocketDisconnect:
pass # Connection never removed from active list!When adding new features, ensure:
- Database sessions are closed in finally blocks
- Collections have maximum size limits
- Caches have TTL/expiration
- Background tasks are tracked and cancelled
- WebSocket connections are removed on disconnect
- Files/handles are closed after use
- Heavy objects (matplotlib figures) are explicitly freed
- Cleanup methods are implemented for classes with resources
- Exceptions don't prevent cleanup (use finally)
import psutil
import asyncio
async def test_memory_leak():
process = psutil.Process()
initial = process.memory_info().rss / 1024 / 1024 # MB
# Run your code 1000 times
for i in range(1000):
await your_function()
final = process.memory_info().rss / 1024 / 1024 # MB
growth = final - initial
print(f"Memory growth: {growth:.2f} MB")
assert growth < 50, f"Memory leak detected: {growth:.2f} MB growth"async def test_resource_cleanup():
agent = YourAgent()
# Use the agent
await agent.execute(task)
# Cleanup
await agent.cleanup()
# Verify resources are released
assert agent.session is None
assert len(agent._background_tasks) == 0
assert agent.cache == {}-
Memory Usage:
process = psutil.Process() memory_mb = process.memory_info().rss / 1024 / 1024
-
Database Connection Pool:
pool_size = engine.pool.size() checked_out = engine.pool.checkedout()
-
Active WebSocket Connections:
active_count = len(active_connections)
-
Background Tasks:
task_count = len(background_tasks)
- Check: Are WebSocket connections being cleaned up?
- Check: Are database sessions being closed?
- Check: Are aiohttp sessions being closed?
- Check: Are all database operations using try/finally?
- Check: Are exceptions being handled properly?
- Check: Is rollback being called on errors?
- Check: Are collections bounded?
- Check: Are caches expiring old entries?
- Check: Are matplotlib figures being closed?
- Check: Are background tasks being cancelled?
- Check: Are background tasks being tracked?
- Check: Is cleanup() being called?
- Check: Are tasks being cancelled and awaited?
Last Updated: 2024-11-06 Status: ✅ Production Ready Maintainer: Alpha Evolve Team