kailash 0.9.15__py3-none-any.whl → 0.9.17__py3-none-any.whl

kailash/monitoring/asyncsql_metrics.py

@@ -0,0 +1,275 @@
+"""
+Prometheus metrics integration for AsyncSQL lock contention monitoring.
+
+This module provides easy-to-use Prometheus metrics for monitoring AsyncSQL
+per-pool locking performance and contention patterns.
+"""
+
+import time
+from typing import Optional, Dict, Any
+from contextlib import asynccontextmanager
+
+try:
+    import prometheus_client
+    PROMETHEUS_AVAILABLE = True
+except ImportError:
+    PROMETHEUS_AVAILABLE = False
+
+
+class AsyncSQLMetrics:
+    """Prometheus metrics collector for AsyncSQL lock contention monitoring."""
+    
+    def __init__(self, enabled: bool = True, registry: Optional[prometheus_client.CollectorRegistry] = None):
+        """
+        Initialize AsyncSQL metrics collector.
+        
+        Args:
+            enabled: Whether to collect metrics (disabled if prometheus_client not available)
+            registry: Custom Prometheus registry (uses default if None)
+        """
+        self.enabled = enabled and PROMETHEUS_AVAILABLE
+        self.registry = registry or prometheus_client.REGISTRY
+        
+        if not self.enabled:
+            return
+        
+        # Lock acquisition counter
+        self.lock_acquisition_counter = prometheus_client.Counter(
+            'asyncsql_lock_acquisitions_total',
+            'Total number of AsyncSQL lock acquisitions',
+            ['pool_key', 'status'],  # status: success, timeout, error
+            registry=self.registry
+        )
+        
+        # Lock wait time histogram  
+        self.lock_wait_time_histogram = prometheus_client.Histogram(
+            'asyncsql_lock_wait_seconds',
+            'Time spent waiting for AsyncSQL locks',
+            ['pool_key'],
+            buckets=(0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, float('inf')),
+            registry=self.registry
+        )
+        
+        # Active locks gauge
+        self.active_locks_gauge = prometheus_client.Gauge(
+            'asyncsql_active_locks',
+            'Number of currently active AsyncSQL locks',
+            ['pool_key'],
+            registry=self.registry
+        )
+        
+        # Pool operations counter
+        self.pool_operations_counter = prometheus_client.Counter(
+            'asyncsql_pool_operations_total',
+            'Total number of AsyncSQL pool operations',
+            ['pool_key', 'operation'],  # operation: create, cleanup, acquire, release
+            registry=self.registry
+        )
+        
+        # Lock contention summary
+        self.lock_contention_summary = prometheus_client.Summary(
+            'asyncsql_lock_contention_seconds',
+            'Summary of AsyncSQL lock contention patterns',
+            ['pool_key'],
+            registry=self.registry
+        )
+    
+    def record_lock_acquisition(self, pool_key: str, status: str, wait_time: float = 0.0):
+        """
+        Record a lock acquisition event.
+        
+        Args:
+            pool_key: The pool key for the lock
+            status: 'success', 'timeout', or 'error'
+            wait_time: Time spent waiting for the lock in seconds
+        """
+        if not self.enabled:
+            return
+        
+        self.lock_acquisition_counter.labels(pool_key=pool_key, status=status).inc()
+        
+        if wait_time > 0:
+            self.lock_wait_time_histogram.labels(pool_key=pool_key).observe(wait_time)
+            self.lock_contention_summary.labels(pool_key=pool_key).observe(wait_time)
+    
+    def set_active_locks(self, pool_key: str, count: int):
+        """
+        Update the count of active locks for a pool.
+        
+        Args:
+            pool_key: The pool key
+            count: Number of active locks
+        """
+        if not self.enabled:
+            return
+        
+        self.active_locks_gauge.labels(pool_key=pool_key).set(count)
+    
+    def record_pool_operation(self, pool_key: str, operation: str):
+        """
+        Record a pool operation event.
+        
+        Args:
+            pool_key: The pool key
+            operation: 'create', 'cleanup', 'acquire', 'release'
+        """
+        if not self.enabled:
+            return
+        
+        self.pool_operations_counter.labels(pool_key=pool_key, operation=operation).inc()
+    
+    @asynccontextmanager
+    async def timed_lock_acquisition(self, pool_key: str):
+        """
+        Context manager to time lock acquisition and automatically record metrics.
+        
+        Usage:
+            async with metrics.timed_lock_acquisition('my_pool_key'):
+                # Lock acquisition logic here
+                async with some_lock:
+                    # Work while holding lock
+                    pass
+        """
+        start_time = time.time()
+        status = 'error'
+        
+        try:
+            yield
+            status = 'success'
+        except Exception as e:
+            if 'timeout' in str(e).lower():
+                status = 'timeout'
+            else:
+                status = 'error'
+            raise
+        finally:
+            wait_time = time.time() - start_time
+            self.record_lock_acquisition(pool_key, status, wait_time)
+
+
+# Global metrics instance (can be overridden)
+_global_metrics: Optional[AsyncSQLMetrics] = None
+
+
+def get_global_metrics() -> Optional[AsyncSQLMetrics]:
+    """Get the global AsyncSQL metrics instance."""
+    global _global_metrics
+    if _global_metrics is None and PROMETHEUS_AVAILABLE:
+        _global_metrics = AsyncSQLMetrics()
+    return _global_metrics
+
+
+def set_global_metrics(metrics: Optional[AsyncSQLMetrics]):
+    """Set the global AsyncSQL metrics instance."""
+    global _global_metrics
+    _global_metrics = metrics
+
+
+def enable_metrics(registry: Optional[prometheus_client.CollectorRegistry] = None) -> AsyncSQLMetrics:
+    """
+    Enable global AsyncSQL metrics collection.
+    
+    Args:
+        registry: Custom Prometheus registry (uses default if None)
+    
+    Returns:
+        The configured metrics instance
+    """
+    metrics = AsyncSQLMetrics(enabled=True, registry=registry)
+    set_global_metrics(metrics)
+    return metrics
+
+
+def disable_metrics():
+    """Disable global AsyncSQL metrics collection."""
+    set_global_metrics(None)
+
+
+# Convenience functions for manual metric recording
+def record_lock_acquisition(pool_key: str, status: str, wait_time: float = 0.0):
+    """Record a lock acquisition event using global metrics."""
+    metrics = get_global_metrics()
+    if metrics:
+        metrics.record_lock_acquisition(pool_key, status, wait_time)
+
+
+def record_pool_operation(pool_key: str, operation: str):
+    """Record a pool operation event using global metrics."""
+    metrics = get_global_metrics()
+    if metrics:
+        metrics.record_pool_operation(pool_key, operation)
+
+
+def set_active_locks(pool_key: str, count: int):
+    """Update active locks count using global metrics."""
+    metrics = get_global_metrics()
+    if metrics:
+        metrics.set_active_locks(pool_key, count)
+
+
+# Integration example for AsyncSQLDatabaseNode
+def integrate_with_async_sql():
+    """
+    Example of how to integrate metrics with AsyncSQLDatabaseNode.
+    
+    This would typically be called during AsyncSQL initialization or
+    through a configuration setting.
+    """
+    if not PROMETHEUS_AVAILABLE:
+        return None
+    
+    # Enable metrics
+    metrics = enable_metrics()
+    
+    # Example: monkey-patch AsyncSQL methods to include metrics
+    # (This is just an example - actual integration would be cleaner)
+    from kailash.nodes.data.async_sql import AsyncSQLDatabaseNode
+    
+    # Store original methods
+    original_get_pool_creation_lock = AsyncSQLDatabaseNode._get_pool_creation_lock
+    original_acquire_lock = AsyncSQLDatabaseNode._acquire_pool_lock_with_timeout
+    
+    @classmethod  
+    def instrumented_get_pool_creation_lock(cls, pool_key: str):
+        """Instrumented version that records pool operations."""
+        record_pool_operation(pool_key, 'acquire')
+        return original_get_pool_creation_lock(pool_key)
+    
+    @classmethod
+    async def instrumented_acquire_lock(cls, pool_key: str, timeout: float = 5.0):
+        """Instrumented version that records lock acquisitions."""
+        async with metrics.timed_lock_acquisition(pool_key):
+            async with original_acquire_lock(pool_key, timeout):
+                yield
+    
+    # Apply instrumentation
+    AsyncSQLDatabaseNode._get_pool_creation_lock = instrumented_get_pool_creation_lock
+    AsyncSQLDatabaseNode._acquire_pool_lock_with_timeout = instrumented_acquire_lock
+    
+    return metrics
+
+
+if __name__ == "__main__":
+    # Example usage
+    print("AsyncSQL Metrics Module")
+    print(f"Prometheus available: {PROMETHEUS_AVAILABLE}")
+    
+    if PROMETHEUS_AVAILABLE:
+        # Enable metrics
+        metrics = enable_metrics()
+        
+        # Simulate some metrics
+        metrics.record_lock_acquisition('test_pool_1', 'success', 0.005)
+        metrics.record_lock_acquisition('test_pool_1', 'success', 0.003)
+        metrics.record_lock_acquisition('test_pool_2', 'timeout', 5.0)
+        metrics.set_active_locks('test_pool_1', 2)
+        metrics.record_pool_operation('test_pool_1', 'create')
+        
+        print("Metrics recorded successfully")
+        print("Access metrics at: http://localhost:8000/metrics")
+        print("(Start prometheus_client HTTP server to view metrics)")
+        
+        # Start metrics server (for testing)
+        # prometheus_client.start_http_server(8000)
+    else:
+        print("Install prometheus_client to enable metrics: pip install prometheus_client")