Agent Zero
a26405f15e
- app/core/resilience.py: CircuitBreaker (CLOSED/OPEN/HALF_OPEN), retry_db, redis_call_with_fallback, InMemoryRateLimiter, CircuitBreakerMiddleware - app/core/auth.py: get_session_data now falls back to PostgreSQL sessions table when Redis is unavailable - app/core/permissions.py: get_cached_permissions falls back to direct DB resolution when Redis circuit is open - app/core/rate_limit.py: check_rate_limit falls back to in-memory limiter when Redis is down; reset_rate_limit clears both Redis and in-memory - app/core/middleware.py: CSRF validation uses get_session_data (Redis+DB fallback); sliding session TTL is best-effort during outage - app/core/db/__init__.py: get_db() wraps session creation with retry_db for transient connection errors; records circuit breaker success/failure - app/deps.py: refresh_session_ttl wrapped in try/except for Redis outage - app/main.py: CircuitBreakerMiddleware registered (returns 503 when DB circuit is OPEN, skips health/metrics endpoints) - app/config.py: Added resilience settings (thresholds, cooldown, retries) - tests/test_resilience.py: 30 tests covering all patterns 30/30 resilience tests pass. No regressions in plugin lifecycle tests.
301 lines
9.2 KiB
Python
301 lines
9.2 KiB
Python
"""Resilience patterns: circuit breaker, DB retry, Redis graceful degradation.
|
|
|
|
Provides:
|
|
- CircuitBreaker: tracks failures per service, opens after threshold, auto-recovers
|
|
- retry_db: decorator that retries DB operations on transient connection errors
|
|
- redis_call_with_fallback: helper that catches Redis errors and falls back
|
|
- InMemoryRateLimiter: process-local rate limiter for Redis outage fallback
|
|
- CircuitBreakerMiddleware: ASGI middleware that fail-fasts on open circuits
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from collections import defaultdict
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any, TypeVar
|
|
|
|
from fastapi import HTTPException, status
|
|
from starlette.responses import JSONResponse
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
T = TypeVar("T")
|
|
|
|
# -- Circuit Breaker ----------------------------------------------------------
|
|
|
|
|
|
class CircuitBreaker:
|
|
"""Circuit breaker for a single service.
|
|
|
|
States:
|
|
- CLOSED: normal operation, requests pass through
|
|
- OPEN: service is down, requests fail fast with 503
|
|
- HALF_OPEN: cooldown expired, one probe request is allowed through
|
|
"""
|
|
|
|
CLOSED = "closed"
|
|
OPEN = "open"
|
|
HALF_OPEN = "half_open"
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
failure_threshold: int = 5,
|
|
window_seconds: int = 30,
|
|
cooldown_seconds: int = 60,
|
|
) -> None:
|
|
self.name = name
|
|
self.failure_threshold = failure_threshold
|
|
self.window_seconds = window_seconds
|
|
self.cooldown_seconds = cooldown_seconds
|
|
self._state = self.CLOSED
|
|
self._failures: list[float] = []
|
|
self._opened_at: float = 0.0
|
|
self._half_open_probe_in_flight = False
|
|
self._lock = asyncio.Lock()
|
|
|
|
@property
|
|
def state(self) -> str:
|
|
if self._state == self.OPEN:
|
|
if time.monotonic() - self._opened_at >= self.cooldown_seconds:
|
|
return self.HALF_OPEN
|
|
return self._state
|
|
|
|
async def can_proceed(self) -> bool:
|
|
current_state = self.state
|
|
if current_state == self.CLOSED:
|
|
return True
|
|
if current_state == self.HALF_OPEN:
|
|
async with self._lock:
|
|
if not self._half_open_probe_in_flight:
|
|
self._half_open_probe_in_flight = True
|
|
return True
|
|
return False
|
|
return False
|
|
|
|
async def record_success(self) -> None:
|
|
was_open = self._state in (self.OPEN, self.HALF_OPEN)
|
|
self._state = self.CLOSED
|
|
self._failures.clear()
|
|
self._half_open_probe_in_flight = False
|
|
if was_open:
|
|
logger.info("Circuit breaker '%s' recovered - state: CLOSED", self.name)
|
|
|
|
async def record_failure(self) -> None:
|
|
now = time.monotonic()
|
|
self._failures.append(now)
|
|
cutoff = now - self.window_seconds
|
|
self._failures = [t for t in self._failures if t >= cutoff]
|
|
|
|
if self._state == self.HALF_OPEN:
|
|
self._state = self.OPEN
|
|
self._opened_at = now
|
|
self._half_open_probe_in_flight = False
|
|
logger.warning("Circuit breaker '%s' probe failed - state: OPEN", self.name)
|
|
return
|
|
|
|
if len(self._failures) >= self.failure_threshold:
|
|
self._state = self.OPEN
|
|
self._opened_at = now
|
|
logger.error(
|
|
"Circuit breaker '%s' tripped - %d failures in %ds - state: OPEN",
|
|
self.name,
|
|
len(self._failures),
|
|
self.window_seconds,
|
|
)
|
|
|
|
def reset(self) -> None:
|
|
self._state = self.CLOSED
|
|
self._failures.clear()
|
|
self._opened_at = 0.0
|
|
self._half_open_probe_in_flight = False
|
|
|
|
|
|
_circuits: dict[str, CircuitBreaker] = {}
|
|
|
|
|
|
def get_circuit(name: str) -> CircuitBreaker:
|
|
if name not in _circuits:
|
|
_circuits[name] = CircuitBreaker(name)
|
|
return _circuits[name]
|
|
|
|
|
|
def reset_all_circuits() -> None:
|
|
for cb in _circuits.values():
|
|
cb.reset()
|
|
|
|
|
|
# -- DB Retry -----------------------------------------------------------------
|
|
|
|
_DB_RETRYABLE_EXC = (ConnectionError, OSError, asyncio.TimeoutError)
|
|
|
|
|
|
def _is_transient_db_error(exc: Exception) -> bool:
|
|
if isinstance(exc, _DB_RETRYABLE_EXC):
|
|
return True
|
|
try:
|
|
from sqlalchemy.exc import OperationalError, DBAPIError
|
|
if isinstance(exc, OperationalError):
|
|
return True
|
|
if isinstance(exc, DBAPIError):
|
|
cause = exc.__cause__ or exc.orig
|
|
if cause and isinstance(cause, (ConnectionError, OSError)):
|
|
return True
|
|
cause_str = str(cause) if cause else ""
|
|
if any(kw in cause_str.lower() for kw in (
|
|
"connection", "timeout", "refused", "reset",
|
|
"broken pipe", "server closed",
|
|
)):
|
|
return True
|
|
except ImportError:
|
|
pass
|
|
return False
|
|
|
|
|
|
async def retry_db(
|
|
func: Callable[..., Awaitable[T]],
|
|
*args: Any,
|
|
max_retries: int = 3,
|
|
base_delay: float = 0.1,
|
|
**kwargs: Any,
|
|
) -> T:
|
|
last_exc: Exception | None = None
|
|
for attempt in range(max_retries):
|
|
try:
|
|
result = await func(*args, **kwargs)
|
|
if attempt > 0:
|
|
logger.info("DB operation succeeded on retry %d", attempt + 1)
|
|
return result
|
|
except Exception as exc:
|
|
last_exc = exc
|
|
if not _is_transient_db_error(exc):
|
|
raise
|
|
if attempt < max_retries - 1:
|
|
delay = base_delay * (2 ** attempt)
|
|
logger.warning(
|
|
"DB transient error (attempt %d/%d): %s - retrying in %.2fs",
|
|
attempt + 1, max_retries, exc, delay,
|
|
)
|
|
await asyncio.sleep(delay)
|
|
else:
|
|
logger.error("DB operation failed after %d retries: %s", max_retries, exc)
|
|
|
|
await get_circuit("db").record_failure()
|
|
raise HTTPException(
|
|
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
|
detail={"detail": "Database temporarily unavailable", "code": "db_unavailable"},
|
|
)
|
|
|
|
|
|
# -- Redis Fallback -----------------------------------------------------------
|
|
|
|
|
|
class RedisUnavailableError(Exception):
|
|
pass
|
|
|
|
|
|
async def redis_call_with_fallback(
|
|
redis_op: Callable[..., Awaitable[T]],
|
|
fallback: Callable[..., Awaitable[T]] | None = None,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> T | None:
|
|
circuit = get_circuit("redis")
|
|
if not await circuit.can_proceed():
|
|
logger.warning("Redis circuit OPEN - using fallback")
|
|
if fallback is not None:
|
|
return await fallback(*args, **kwargs)
|
|
return None
|
|
|
|
try:
|
|
result = await redis_op(*args, **kwargs)
|
|
await circuit.record_success()
|
|
return result
|
|
except Exception as exc:
|
|
logger.warning("Redis operation failed: %s - using fallback", exc)
|
|
await circuit.record_failure()
|
|
if fallback is not None:
|
|
return await fallback(*args, **kwargs)
|
|
return None
|
|
|
|
|
|
# -- In-Memory Rate Limiter (Redis fallback) ----------------------------------
|
|
|
|
|
|
class InMemoryRateLimiter:
|
|
"""Process-local sliding-window rate limiter for Redis outage fallback."""
|
|
|
|
def __init__(self) -> None:
|
|
self._counts: dict[str, list[float]] = defaultdict(list)
|
|
|
|
async def check(self, key: str, max_attempts: int, window_seconds: int) -> tuple[bool, int]:
|
|
now = time.monotonic()
|
|
cutoff = now - window_seconds
|
|
self._counts[key] = [t for t in self._counts[key] if t >= cutoff]
|
|
if len(self._counts[key]) >= max_attempts:
|
|
retry_after = int(window_seconds - (now - self._counts[key][0]))
|
|
return False, max(retry_after, 1)
|
|
self._counts[key].append(now)
|
|
return True, 0
|
|
|
|
async def reset(self, key: str) -> None:
|
|
self._counts.pop(key, None)
|
|
|
|
def clear_all(self) -> None:
|
|
self._counts.clear()
|
|
|
|
|
|
_inmemory_limiter = InMemoryRateLimiter()
|
|
|
|
|
|
def get_inmemory_limiter() -> InMemoryRateLimiter:
|
|
return _inmemory_limiter
|
|
|
|
|
|
# -- Circuit Breaker Middleware -----------------------------------------------
|
|
|
|
|
|
class CircuitBreakerMiddleware:
|
|
"""ASGI middleware: returns 503 when DB circuit is OPEN."""
|
|
|
|
SKIP_PATHS = {
|
|
"/api/v1/health",
|
|
"/api/v1/health/live",
|
|
"/api/v1/health/ready",
|
|
"/api/v1/metrics",
|
|
"/docs",
|
|
"/redoc",
|
|
"/openapi.json",
|
|
}
|
|
|
|
def __init__(self, app) -> None:
|
|
self.app = app
|
|
|
|
async def __call__(self, scope, receive, send) -> None:
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
path = scope.get("path", "")
|
|
if path in self.SKIP_PATHS or path.startswith("/docs") or path.startswith("/redoc"):
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
db_circuit = get_circuit("db")
|
|
if not await db_circuit.can_proceed():
|
|
response = JSONResponse(
|
|
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
|
content={
|
|
"detail": "Service temporarily unavailable",
|
|
"code": "circuit_open",
|
|
"service": "db",
|
|
},
|
|
)
|
|
await response(scope, receive, send)
|
|
return
|
|
|
|
await self.app(scope, receive, send)
|