136 lines
5.5 KiB
Python
136 lines
5.5 KiB
Python
"""
|
|
Resource Tracking Middleware
|
|
Tracks CPU, memory, and I/O usage per request for admin debugging.
|
|
"""
|
|
import psutil
|
|
import time
|
|
import threading
|
|
import logging
|
|
from django.utils.deprecation import MiddlewareMixin
|
|
from django.core.cache import cache
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ResourceTrackingMiddleware(MiddlewareMixin):
|
|
"""
|
|
Middleware to track resource usage per request.
|
|
Stores metrics in cache with request ID for retrieval.
|
|
Only tracks for authenticated admin/developer users.
|
|
"""
|
|
thread_local = threading.local()
|
|
|
|
def process_request(self, request):
|
|
"""Start tracking resources for this request"""
|
|
# Only track if user is authenticated
|
|
if not hasattr(request, 'user') or not request.user.is_authenticated:
|
|
return None
|
|
|
|
# Check if user is admin/developer
|
|
if not (hasattr(request.user, 'is_admin_or_developer') and request.user.is_admin_or_developer()):
|
|
return None
|
|
|
|
# Check if debug tracking is enabled via header (set by frontend)
|
|
debug_enabled = request.headers.get('X-Debug-Resource-Tracking', '').lower() == 'true'
|
|
if not debug_enabled:
|
|
return None
|
|
|
|
try:
|
|
# Generate request ID
|
|
request_id = f"req_{int(time.time() * 1000000)}"
|
|
request.resource_tracking_id = request_id
|
|
|
|
# Get initial process stats
|
|
process = psutil.Process()
|
|
initial_cpu_times = process.cpu_times()
|
|
initial_memory = process.memory_info()
|
|
initial_io = process.io_counters() if hasattr(process, 'io_counters') else None
|
|
|
|
# Store initial state
|
|
self.thread_local.start_time = time.time()
|
|
self.thread_local.initial_cpu_times = initial_cpu_times
|
|
self.thread_local.initial_memory = initial_memory
|
|
self.thread_local.initial_io = initial_io
|
|
self.thread_local.process = process
|
|
self.thread_local.request_id = request_id
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Error starting resource tracking: {str(e)}")
|
|
# Don't break the request if tracking fails
|
|
|
|
return None
|
|
|
|
def process_response(self, request, response):
|
|
"""Calculate and store resource usage for this request"""
|
|
if not hasattr(request, 'resource_tracking_id'):
|
|
return response
|
|
|
|
try:
|
|
# Calculate elapsed time
|
|
elapsed_time = time.time() - self.thread_local.start_time
|
|
|
|
# Get final process stats
|
|
process = self.thread_local.process
|
|
final_cpu_times = process.cpu_times()
|
|
final_memory = process.memory_info()
|
|
final_io = process.io_counters() if hasattr(process, 'io_counters') else None
|
|
|
|
# Calculate CPU usage (user + system time)
|
|
cpu_user_time = (final_cpu_times.user - self.thread_local.initial_cpu_times.user) * 1000 # ms
|
|
cpu_system_time = (final_cpu_times.system - self.thread_local.initial_cpu_times.system) * 1000 # ms
|
|
cpu_total_time = cpu_user_time + cpu_system_time
|
|
|
|
# Calculate memory delta
|
|
memory_delta = final_memory.rss - self.thread_local.initial_memory.rss
|
|
|
|
# Calculate I/O
|
|
io_read = 0
|
|
io_write = 0
|
|
if final_io and self.thread_local.initial_io:
|
|
io_read = final_io.read_bytes - self.thread_local.initial_io.read_bytes
|
|
io_write = final_io.write_bytes - self.thread_local.initial_io.write_bytes
|
|
|
|
# Get system-wide stats
|
|
cpu_percent = psutil.cpu_percent(interval=0.1)
|
|
memory = psutil.virtual_memory()
|
|
|
|
# Store metrics in cache (expire after 5 minutes)
|
|
metrics = {
|
|
'request_id': request.resource_tracking_id,
|
|
'path': request.path,
|
|
'method': request.method,
|
|
'elapsed_time_ms': round(elapsed_time * 1000, 2),
|
|
'cpu': {
|
|
'user_time_ms': round(cpu_user_time, 2),
|
|
'system_time_ms': round(cpu_system_time, 2),
|
|
'total_time_ms': round(cpu_total_time, 2),
|
|
'system_percent': round(cpu_percent, 2),
|
|
},
|
|
'memory': {
|
|
'delta_bytes': memory_delta,
|
|
'delta_mb': round(memory_delta / (1024**2), 2),
|
|
'final_rss_mb': round(final_memory.rss / (1024**2), 2),
|
|
'system_used_percent': round(memory.percent, 2),
|
|
},
|
|
'io': {
|
|
'read_bytes': io_read,
|
|
'read_mb': round(io_read / (1024**2), 2),
|
|
'write_bytes': io_write,
|
|
'write_mb': round(io_write / (1024**2), 2),
|
|
},
|
|
'timestamp': time.time(),
|
|
}
|
|
|
|
# Store in cache with 5 minute expiry
|
|
cache.set(f"resource_tracking_{request.resource_tracking_id}", metrics, 300)
|
|
|
|
# Add request ID to response header
|
|
response['X-Resource-Tracking-ID'] = request.resource_tracking_id
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Error calculating resource tracking: {str(e)}")
|
|
# Don't break the response if tracking fails
|
|
|
|
return response
|
|
|