dexscreen 0.0.1__py3-none-any.whl → 0.0.4__py3-none-any.whl

dexscreen/utils/middleware.py

@@ -0,0 +1,363 @@
+"""
+Middleware utilities for request tracking and correlation ID propagation
+"""
+
+import inspect
+import time
+from functools import wraps
+from typing import Any, Callable, Optional, TypeVar
+
+from .logging_config import (
+    generate_correlation_id,
+    get_contextual_logger,
+    get_correlation_id,
+    set_correlation_id,
+    with_correlation_id,
+)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class RequestTracker:
+    """Tracks request lifecycle and propagates correlation IDs"""
+
+    def __init__(self):
+        self.contextual_logger = get_contextual_logger(__name__)
+        self.active_requests: dict[str, dict[str, Any]] = {}
+
+    def start_request(self, operation: str, context: Optional[dict[str, Any]] = None) -> str:
+        """Start tracking a new request and return correlation ID"""
+        correlation_id = generate_correlation_id()
+        set_correlation_id(correlation_id)
+
+        request_info = {
+            "operation": operation,
+            "correlation_id": correlation_id,
+            "start_time": time.time(),
+            "context": context or {},
+            "status": "active",
+        }
+
+        self.active_requests[correlation_id] = request_info
+
+        track_context = {
+            "operation": "request_start",
+            "correlation_id": correlation_id,
+            "tracked_operation": operation,
+            "active_requests_count": len(self.active_requests),
+            **request_info["context"],
+        }
+
+        self.contextual_logger.info("Starting request tracking for %s", operation, context=track_context)
+
+        return correlation_id
+
+    def end_request(
+        self,
+        correlation_id: Optional[str] = None,
+        status: str = "completed",
+        result_context: Optional[dict[str, Any]] = None,
+    ) -> Optional[dict[str, Any]]:
+        """End request tracking and return request info"""
+        if correlation_id is None:
+            correlation_id = get_correlation_id()
+
+        if not correlation_id or correlation_id not in self.active_requests:
+            self.contextual_logger.warning(
+                "Attempted to end request tracking with unknown correlation ID: %s",
+                correlation_id,
+                context={"operation": "request_end_unknown", "correlation_id": correlation_id},
+            )
+            return None
+
+        request_info = self.active_requests.pop(correlation_id)
+        request_info["end_time"] = time.time()
+        request_info["duration"] = request_info["end_time"] - request_info["start_time"]
+        request_info["status"] = status
+        request_info["result_context"] = result_context or {}
+
+        track_context = {
+            "operation": "request_end",
+            "correlation_id": correlation_id,
+            "tracked_operation": request_info["operation"],
+            "duration": request_info["duration"],
+            "status": status,
+            "active_requests_count": len(self.active_requests),
+            **request_info["context"],
+            **request_info["result_context"],
+        }
+
+        log_level = "info" if status == "completed" else "warning" if status == "failed" else "error"
+        log_method = getattr(self.contextual_logger, log_level)
+
+        log_method(
+            "Completed request tracking for %s (%.3fs, %s)",
+            request_info["operation"],
+            request_info["duration"],
+            status,
+            context=track_context,
+        )
+
+        return request_info
+
+    def get_active_requests(self) -> dict[str, dict[str, Any]]:
+        """Get all currently active requests"""
+        return self.active_requests.copy()
+
+    def log_active_requests(self):
+        """Log information about currently active requests"""
+        if not self.active_requests:
+            self.contextual_logger.debug("No active requests", context={"operation": "active_requests_check"})
+            return
+
+        current_time = time.time()
+        requests_info = []
+
+        for correlation_id, info in self.active_requests.items():
+            duration = current_time - info["start_time"]
+            requests_info.append(
+                {
+                    "correlation_id": correlation_id,
+                    "operation": info["operation"],
+                    "duration": duration,
+                }
+            )
+
+        active_context = {
+            "operation": "active_requests_summary",
+            "active_count": len(self.active_requests),
+            "requests": requests_info,
+        }
+
+        self.contextual_logger.info(
+            "%d active requests currently being tracked", len(self.active_requests), context=active_context
+        )
+
+
+# Global request tracker instance
+_request_tracker = RequestTracker()
+
+
+def track_request(operation: str, include_args: bool = False, include_result: bool = False):
+    """
+    Decorator to automatically track request lifecycle with correlation IDs.
+
+    Args:
+        operation: Name of the operation being tracked
+        include_args: Whether to include function arguments in context
+        include_result: Whether to include function result in context
+    """
+
+    def decorator(func: F) -> F:
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            # Build context
+            context: dict[str, Any] = {"function_name": func.__name__}
+            if include_args:
+                context["args"] = _sanitize_args(args, kwargs)
+
+            # Start tracking
+            correlation_id = _request_tracker.start_request(operation, context)
+
+            try:
+                result = func(*args, **kwargs)
+
+                # Build result context
+                result_context = {}
+                if include_result:
+                    result_context["result"] = _sanitize_result(result)
+
+                _request_tracker.end_request(correlation_id, "completed", result_context)
+                return result
+
+            except Exception as e:
+                error_context = {
+                    "error_type": type(e).__name__,
+                    "error_message": str(e),
+                }
+                _request_tracker.end_request(correlation_id, "failed", error_context)
+                raise
+
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            # Build context
+            context: dict[str, Any] = {"function_name": func.__name__}
+            if include_args:
+                context["args"] = _sanitize_args(args, kwargs)
+
+            # Start tracking
+            correlation_id = _request_tracker.start_request(operation, context)
+
+            try:
+                result = await func(*args, **kwargs)
+
+                # Build result context
+                result_context = {}
+                if include_result:
+                    result_context["result"] = _sanitize_result(result)
+
+                _request_tracker.end_request(correlation_id, "completed", result_context)
+                return result
+
+            except Exception as e:
+                error_context = {
+                    "error_type": type(e).__name__,
+                    "error_message": str(e),
+                }
+                _request_tracker.end_request(correlation_id, "failed", error_context)
+                raise
+
+        if inspect.iscoroutinefunction(func):
+            return async_wrapper  # type: ignore[return-value]
+        else:
+            return sync_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def _sanitize_args(args: tuple, kwargs: dict) -> dict[str, Any]:
+    """Sanitize function arguments for logging"""
+    sanitized = {}
+
+    # Limit args to first 3 to avoid log bloat
+    if args:
+        sanitized["positional"] = [_sanitize_value(arg) for arg in args[:3]]
+        if len(args) > 3:
+            sanitized["positional_truncated"] = len(args)
+
+    # Sanitize kwargs
+    if kwargs:
+        sanitized["keyword"] = {k: _sanitize_value(v) for k, v in list(kwargs.items())[:5]}
+        if len(kwargs) > 5:
+            sanitized["keyword_truncated"] = len(kwargs)
+
+    return sanitized
+
+
+def _sanitize_value(value: Any) -> Any:
+    """Sanitize a single value for logging"""
+    if isinstance(value, str):
+        # Truncate long strings
+        return value[:100] + "..." if len(value) > 100 else value
+    elif isinstance(value, (list, tuple)):
+        # Show length for collections
+        return f"<{type(value).__name__} of {len(value)} items>"
+    elif isinstance(value, dict):
+        return f"<dict with {len(value)} keys>"
+    elif hasattr(value, "__dict__"):
+        # Object with attributes
+        return f"<{type(value).__name__} instance>"
+    else:
+        return value
+
+
+def _sanitize_result(result: Any) -> Any:
+    """Sanitize function result for logging"""
+    if isinstance(result, (list, tuple)):
+        return {"type": type(result).__name__, "length": len(result)}
+    elif isinstance(result, dict):
+        return {"type": "dict", "keys": len(result)}
+    elif hasattr(result, "__dict__"):
+        return {"type": type(result).__name__}
+    else:
+        return _sanitize_value(result)
+
+
+class CorrelationMiddleware:
+    """Middleware for automatic correlation ID management in HTTP requests"""
+
+    def __init__(self):
+        self.contextual_logger = get_contextual_logger(__name__)
+
+    def wrap_http_client(self, client_class):
+        """Wrap an HTTP client class to add correlation ID headers"""
+        original_request = client_class.request
+        original_request_async = getattr(client_class, "request_async", None)
+
+        @with_correlation_id()
+        def wrapped_request(self, method, url, **kwargs):
+            correlation_id = get_correlation_id()
+            if correlation_id:
+                # Add correlation ID to headers
+                headers = kwargs.get("headers", {})
+                headers["X-Correlation-ID"] = correlation_id
+                kwargs["headers"] = headers
+
+                middleware_context = {
+                    "operation": "http_request_correlation",
+                    "method": method,
+                    "url": url[:100] + "..." if len(url) > 100 else url,
+                    "correlation_id": correlation_id,
+                    "has_existing_headers": bool(kwargs.get("headers", {})),
+                }
+
+                self.contextual_logger.debug(
+                    "Adding correlation ID to HTTP %s request", method, context=middleware_context
+                )
+
+            return original_request(self, method, url, **kwargs)
+
+        if original_request_async:
+
+            @with_correlation_id()
+            async def wrapped_request_async(self, method, url, **kwargs):
+                correlation_id = get_correlation_id()
+                if correlation_id:
+                    # Add correlation ID to headers
+                    headers = kwargs.get("headers", {})
+                    headers["X-Correlation-ID"] = correlation_id
+                    kwargs["headers"] = headers
+
+                    middleware_context = {
+                        "operation": "async_http_request_correlation",
+                        "method": method,
+                        "url": url[:100] + "..." if len(url) > 100 else url,
+                        "correlation_id": correlation_id,
+                        "has_existing_headers": bool(kwargs.get("headers", {})),
+                    }
+
+                    self.contextual_logger.debug(
+                        "Adding correlation ID to async HTTP %s request", method, context=middleware_context
+                    )
+
+                return await original_request_async(self, method, url, **kwargs)
+
+            client_class.request_async = wrapped_request_async
+
+        client_class.request = wrapped_request
+        return client_class
+
+
+# Global middleware instance
+_correlation_middleware = CorrelationMiddleware()
+
+
+def get_request_tracker() -> RequestTracker:
+    """Get the global request tracker instance"""
+    return _request_tracker
+
+
+def get_correlation_middleware() -> CorrelationMiddleware:
+    """Get the global correlation middleware instance"""
+    return _correlation_middleware
+
+
+def auto_track_requests(operation_prefix: str = "api"):
+    """
+    Class decorator to automatically add request tracking to all public methods.
+
+    Args:
+        operation_prefix: Prefix for operation names (method names will be appended)
+    """
+
+    def decorator(cls):
+        for attr_name in dir(cls):
+            if not attr_name.startswith("_"):  # Only public methods
+                attr = getattr(cls, attr_name)
+                if callable(attr):
+                    operation_name = f"{operation_prefix}_{attr_name}"
+                    tracked_method = track_request(operation_name)(attr)
+                    setattr(cls, attr_name, tracked_method)
+        return cls
+
+    return decorator

dexscreen/utils/ratelimit.py

@@ -3,12 +3,9 @@ import collections
 import threading
 import time
 from collections import deque
+from typing import Any
 
-
-class RateLimitError(Exception):
-    """Raised when rate limit is exceeded"""
-
-    pass
+from .logging_config import get_contextual_logger, with_correlation_id
 
 
 class RateLimiter:
@@ -21,12 +18,75 @@ class RateLimiter:
         self.sync_lock = threading.Lock()
         self.async_lock = asyncio.Lock()
 
+        # Enhanced logging
+        self.contextual_logger = get_contextual_logger(__name__)
+
+        # Rate limiting statistics
+        self.stats = {
+            "total_requests": 0,
+            "blocked_requests": 0,
+            "total_wait_time": 0.0,
+            "max_wait_time": 0.0,
+            "average_wait_time": 0.0,
+            "calls_in_current_window": 0,
+            "window_start_time": None,
+        }
+
+        init_context = {
+            "max_calls": max_calls,
+            "period": period,
+            "rate_per_second": max_calls / period,
+        }
+
+        self.contextual_logger.debug("RateLimiter initialized", context=init_context)
+
+    @with_correlation_id()
     def __enter__(self):
         with self.sync_lock:
+            self.stats["total_requests"] += 1
             sleep_time = self.get_sleep_time()
 
+            rate_limit_context = {
+                "operation": "sync_rate_limit_enter",
+                "sleep_time": sleep_time,
+                "calls_in_window": len(self.calls),
+                "max_calls": self.max_calls,
+                "period": self.period,
+                "will_block": sleep_time > 0,
+            }
+
             if sleep_time > 0:
+                self.stats["blocked_requests"] += 1
+                self.stats["total_wait_time"] += sleep_time
+                self.stats["max_wait_time"] = max(self.stats["max_wait_time"], sleep_time)
+
+                # Update average wait time
+                if self.stats["blocked_requests"] > 0:
+                    self.stats["average_wait_time"] = self.stats["total_wait_time"] / self.stats["blocked_requests"]
+
+                rate_limit_context["blocking_duration"] = sleep_time
+
+                self.contextual_logger.warning(
+                    "Rate limit exceeded, sleeping for %.3fs (calls: %d/%d)",
+                    sleep_time,
+                    len(self.calls),
+                    self.max_calls,
+                    context=rate_limit_context,
+                )
+
+                start_time = time.time()
                 time.sleep(sleep_time)
+                actual_sleep = time.time() - start_time
+
+                if abs(actual_sleep - sleep_time) > 0.1:  # More than 100ms difference
+                    self.contextual_logger.debug(
+                        "Sleep time deviation: expected %.3fs, actual %.3fs",
+                        sleep_time,
+                        actual_sleep,
+                        context={"expected_sleep": sleep_time, "actual_sleep": actual_sleep},
+                    )
+            else:
+                self.contextual_logger.debug("Rate limit check passed", context=rate_limit_context)
 
             return self
 
@@ -34,12 +94,58 @@ class RateLimiter:
         with self.sync_lock:
             self._clear_calls()
 
+            # Update current window stats
+            self.stats["calls_in_current_window"] = len(self.calls)
+            if self.calls and self.stats["window_start_time"] is None:
+                self.stats["window_start_time"] = self.calls[0]
+
+    @with_correlation_id()
     async def __aenter__(self):
         async with self.async_lock:
+            self.stats["total_requests"] += 1
             sleep_time = self.get_sleep_time()
 
+            rate_limit_context = {
+                "operation": "async_rate_limit_enter",
+                "sleep_time": sleep_time,
+                "calls_in_window": len(self.calls),
+                "max_calls": self.max_calls,
+                "period": self.period,
+                "will_block": sleep_time > 0,
+            }
+
             if sleep_time > 0:
+                self.stats["blocked_requests"] += 1
+                self.stats["total_wait_time"] += sleep_time
+                self.stats["max_wait_time"] = max(self.stats["max_wait_time"], sleep_time)
+
+                # Update average wait time
+                if self.stats["blocked_requests"] > 0:
+                    self.stats["average_wait_time"] = self.stats["total_wait_time"] / self.stats["blocked_requests"]
+
+                rate_limit_context["blocking_duration"] = sleep_time
+
+                self.contextual_logger.warning(
+                    "Async rate limit exceeded, sleeping for %.3fs (calls: %d/%d)",
+                    sleep_time,
+                    len(self.calls),
+                    self.max_calls,
+                    context=rate_limit_context,
+                )
+
+                start_time = time.time()
                 await asyncio.sleep(sleep_time)
+                actual_sleep = time.time() - start_time
+
+                if abs(actual_sleep - sleep_time) > 0.1:  # More than 100ms difference
+                    self.contextual_logger.debug(
+                        "Async sleep time deviation: expected %.3fs, actual %.3fs",
+                        sleep_time,
+                        actual_sleep,
+                        context={"expected_sleep": sleep_time, "actual_sleep": actual_sleep},
+                    )
+            else:
+                self.contextual_logger.debug("Async rate limit check passed", context=rate_limit_context)
 
             return self
 
@@ -47,19 +153,117 @@ class RateLimiter:
         async with self.async_lock:
             self._clear_calls()
 
+            # Update current window stats
+            self.stats["calls_in_current_window"] = len(self.calls)
+            if self.calls and self.stats["window_start_time"] is None:
+                self.stats["window_start_time"] = self.calls[0]
+
     def get_sleep_time(self) -> float:
+        """Calculate how long to sleep before allowing the next call"""
         if len(self.calls) >= self.max_calls:
             until = time.time() + self.period - self._timespan
-            return until - time.time()
+            sleep_time = until - time.time()
+
+            # Log when rate limit calculations result in significant wait times
+            if sleep_time > 1.0:  # More than 1 second
+                sleep_context = {
+                    "operation": "calculate_sleep_time",
+                    "calculated_sleep": sleep_time,
+                    "calls_in_window": len(self.calls),
+                    "window_timespan": self._timespan,
+                    "max_calls": self.max_calls,
+                    "period": self.period,
+                    "utilization_percent": (len(self.calls) / self.max_calls) * 100,
+                }
+
+                self.contextual_logger.debug(
+                    "Rate limit calculation: need to sleep %.3fs (window %.1f%% full)",
+                    sleep_time,
+                    sleep_context["utilization_percent"],
+                    context=sleep_context,
+                )
+
+            return max(0, sleep_time)  # Ensure non-negative
 
         return 0
 
     def _clear_calls(self):
-        self.calls.append(time.time())
+        """Add current call and remove expired calls from the sliding window"""
+        current_time = time.time()
+        calls_before = len(self.calls)
+
+        self.calls.append(current_time)
 
-        while self._timespan >= self.period:
+        # Remove expired calls
+        while len(self.calls) > 1 and self._timespan >= self.period:
             self.calls.popleft()
 
+        calls_after = len(self.calls)
+        expired_calls = calls_before + 1 - calls_after  # +1 for the new call we just added
+
+        if expired_calls > 0:
+            clear_context = {
+                "operation": "clear_expired_calls",
+                "expired_calls": expired_calls,
+                "calls_remaining": calls_after,
+                "window_timespan": self._timespan if len(self.calls) > 1 else 0,
+                "current_utilization": (calls_after / self.max_calls) * 100,
+            }
+
+            self.contextual_logger.debug(
+                "Cleared %d expired calls, %d calls remaining (%.1f%% capacity)",
+                expired_calls,
+                calls_after,
+                clear_context["current_utilization"],
+                context=clear_context,
+            )
+
     @property
     def _timespan(self) -> float:
+        """Get the time span of the current sliding window"""
+        if len(self.calls) < 2:
+            return 0
         return self.calls[-1] - self.calls[0]
+
+    def get_rate_limit_stats(self) -> dict[str, Any]:
+        """Get comprehensive rate limiting statistics"""
+        time.time()
+
+        # Calculate current rate
+        current_rate = 0.0
+        if len(self.calls) > 1:
+            window_duration = min(self._timespan, self.period)
+            if window_duration > 0:
+                current_rate = len(self.calls) / window_duration
+
+        stats = self.stats.copy()
+        stats.update(
+            {
+                "current_calls_in_window": len(self.calls),
+                "current_window_timespan": self._timespan,
+                "current_rate_per_second": current_rate,
+                "configured_max_rate": self.max_calls / self.period,
+                "capacity_utilization_percent": (len(self.calls) / self.max_calls) * 100,
+                "next_sleep_time": self.get_sleep_time(),
+                "is_rate_limited": len(self.calls) >= self.max_calls,
+                "efficiency_ratio": (self.stats["total_requests"] - self.stats["blocked_requests"])
+                / max(1, self.stats["total_requests"]),
+            }
+        )
+
+        return stats
+
+    def log_stats(self, operation: str = "rate_limit_stats"):
+        """Log current rate limiting statistics"""
+        stats = self.get_rate_limit_stats()
+
+        stats_context = {"operation": operation, **stats}
+
+        self.contextual_logger.info(
+            "Rate limiter stats: %d/%d requests, %.1f%% blocked, %.3fs avg wait",
+            stats["total_requests"],
+            stats["max_calls"],
+            (stats["blocked_requests"] / max(1, stats["total_requests"])) * 100,
+            stats["average_wait_time"],
+            context=stats_context,
+        )