89 lines
3.2 KiB
Python
89 lines
3.2 KiB
Python
# ================================================================
|
|
# shared/monitoring/decorators.py
|
|
# ================================================================
|
|
"""
|
|
Decorators for monitoring and metrics
|
|
"""
|
|
|
|
import time
|
|
import logging
|
|
import functools
|
|
from typing import Callable, Any, Optional
|
|
from .metrics import get_metrics_collector
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def track_execution_time(metric_name: str, service_name: str,
|
|
labels: Optional[dict] = None):
|
|
"""Decorator to track function execution time"""
|
|
def decorator(func: Callable) -> Callable:
|
|
@functools.wraps(func)
|
|
async def async_wrapper(*args, **kwargs) -> Any:
|
|
start_time = time.time()
|
|
try:
|
|
result = await func(*args, **kwargs)
|
|
duration = time.time() - start_time
|
|
|
|
metrics_collector = get_metrics_collector(service_name)
|
|
if metrics_collector:
|
|
metrics_collector.observe_histogram(metric_name, duration, labels)
|
|
|
|
return result
|
|
except Exception as e:
|
|
duration = time.time() - start_time
|
|
logger.error(f"Function {func.__name__} failed after {duration:.2f}s: {e}")
|
|
raise
|
|
|
|
@functools.wraps(func)
|
|
def sync_wrapper(*args, **kwargs) -> Any:
|
|
start_time = time.time()
|
|
try:
|
|
result = func(*args, **kwargs)
|
|
duration = time.time() - start_time
|
|
|
|
metrics_collector = get_metrics_collector(service_name)
|
|
if metrics_collector:
|
|
metrics_collector.observe_histogram(metric_name, duration, labels)
|
|
|
|
return result
|
|
except Exception as e:
|
|
duration = time.time() - start_time
|
|
logger.error(f"Function {func.__name__} failed after {duration:.2f}s: {e}")
|
|
raise
|
|
|
|
# Return appropriate wrapper based on function type
|
|
import asyncio
|
|
if asyncio.iscoroutinefunction(func):
|
|
return async_wrapper
|
|
else:
|
|
return sync_wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
def count_calls(metric_name: str, service_name: str,
|
|
labels: Optional[dict] = None):
|
|
"""Decorator to count function calls"""
|
|
def decorator(func: Callable) -> Callable:
|
|
@functools.wraps(func)
|
|
async def async_wrapper(*args, **kwargs) -> Any:
|
|
metrics_collector = get_metrics_collector(service_name)
|
|
if metrics_collector:
|
|
metrics_collector.increment_counter(metric_name, labels=labels)
|
|
return await func(*args, **kwargs)
|
|
|
|
@functools.wraps(func)
|
|
def sync_wrapper(*args, **kwargs) -> Any:
|
|
metrics_collector = get_metrics_collector(service_name)
|
|
if metrics_collector:
|
|
metrics_collector.increment_counter(metric_name, labels=labels)
|
|
return func(*args, **kwargs)
|
|
|
|
# Return appropriate wrapper based on function type
|
|
import asyncio
|
|
if asyncio.iscoroutinefunction(func):
|
|
return async_wrapper
|
|
else:
|
|
return sync_wrapper
|
|
|
|
return decorator |