control_plane/controller.py

import asyncio
import json
import time
import logging
import numpy as np
from typing import Dict, List, Any, Tuple
from sqlalchemy.ext.asyncio import AsyncSession
from database import async_session
from models import AuditLog, ReliabilityEvent

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("reliability-controller")

SERVICES = ["service-a", "service-b", "service-c"]
SLIDING_WINDOW_SEC = 10.0
DEFAULT_RPS = 100.0
MIN_RPS = 1.0

# Store state in memory
class ServiceState:
    def __init__(self, name: str):
        self.name = name
        self.current_limit = DEFAULT_RPS
        self.circuit_state = "CLOSED" # CLOSED, OPEN, HALF_OPEN
        self.last_circuit_change = time.time()
        # For anti-flapping
        self.healthy_ticks = 0
        self.history_latencies = []
        self.history_errors = []
        self.ema_latency = 0.0
        self.ema_error = 0.0

service_states = {s: ServiceState(s) for s in SERVICES}

async def publish_event(redis_client, event_type: str, service: str, severity: str, message: str, db_session: AsyncSession):
    """Saves event to Postgres and broadcasts to Redis Pub/Sub."""
    logger.info(f"[{severity}] Service {service}: {message}")
    
    event_dict = {
        "event_type": event_type,
        "service": service,
        "severity": severity,
        "message": message,
        "timestamp": time.time()
    }
    
    # Save to PostgreSQL
    event_db = ReliabilityEvent(
        event_type=event_type,
        service=service,
        severity=severity,
        message=message
    )
    db_session.add(event_db)
    
    # Publish to Redis
    await redis_client.publish("reliability_events", json.dumps(event_dict))

async def save_audit_log(service: str, action: str, old_val: str, new_val: str, reason: str, metrics: dict, db_session: AsyncSession):
    """Saves a controller decision to the database."""
    log = AuditLog(
        service=service,
        action=action,
        old_value=old_val,
        new_value=new_val,
        reason=reason,
        trigger_metrics=metrics
    )
    db_session.add(log)

async def evaluate_service_metrics(redis_client, service: str) -> Dict[str, Any]:
    """Fetches telemetry logs from Redis and computes statistics."""
    now = time.time()
    cutoff = now - SLIDING_WINDOW_SEC
    
    # Fetch records in sliding window
    records_raw = await redis_client.zrangebyscore(f"metrics:raw:{service}", cutoff, now)
    
    total_reqs = 0
    errors = 0
    latencies = []
    
    for r in records_raw:
        try:
            data = json.loads(r)
            total_reqs += 1
            latencies.append(data.get("latency_ms", 0.0) / 1000.0) # convert to seconds
            status = int(data.get("status_code", 200))
            if status in [500, 502, 504] or data.get("downstream_error", False):
                errors += 1
        except Exception:
            pass
            
    if total_reqs == 0:
        # Check if service is responsive via health checks
        try:
            # We can check if any recent ping is there, if not, treat as unreachable
            is_unreachable = True
        except Exception:
            is_unreachable = True
            
        return {
            "total_requests": 0,
            "error_rate": 0.0,
            "p50": 0.0,
            "p95": 0.0,
            "p99": 0.0,
            "rps": 0.0,
            "status": "UNREACHABLE"
        }
        
    latencies = sorted(latencies)
    p50 = float(np.percentile(latencies, 50)) if latencies else 0.0
    p95 = float(np.percentile(latencies, 95)) if latencies else 0.0
    p99 = float(np.percentile(latencies, 99)) if latencies else 0.0
    
    error_rate = float(errors) / total_reqs
    rps = float(total_reqs) / SLIDING_WINDOW_SEC
    
    status = "HEALTHY"
    if error_rate > 0.10 or p95 > 0.50:
        status = "CRITICAL"
    elif error_rate > 0.02 or p95 > 0.20:
        status = "DEGRADED"
        
    return {
        "total_requests": total_reqs,
        "error_rate": error_rate,
        "p50": p50,
        "p95": p95,
        "p99": p99,
        "rps": rps,
        "status": status,
        "latencies_raw": latencies
    }

async def run_reliability_loop(redis_client):
    """Main Reliability Controller Loop running every 2 seconds."""
    logger.info("Reliability control loop started.")
    
    # Make sure default Redis keys are set
    for s in SERVICES:
        await redis_client.set(f"limit:capacity:{s}", str(DEFAULT_RPS))
        await redis_client.set(f"limit:rate:{s}", str(DEFAULT_RPS))
        await redis_client.set(f"circuit_breaker:state:{s}", "CLOSED")
        
    while True:
        await asyncio.sleep(2.0)
        
        async with async_session() as db_session:
            try:
                metrics_summary = {}
                
                # Step 1: Collect Metrics
                for s in SERVICES:
                    metrics_summary[s] = await evaluate_service_metrics(redis_client, s)
                
                # Sync states in Redis
                for s in SERVICES:
                    metrics = metrics_summary[s]
                    state = service_states[s]
                    
                    # Update EMA (smooth calculations to prevent flapping)
                    alpha = 0.3
                    state.ema_latency = (alpha * metrics["p95"]) + ((1 - alpha) * state.ema_latency)
                    state.ema_error = (alpha * metrics["error_rate"]) + ((1 - alpha) * state.ema_error)
                    
                    # Store history for anomaly detection
                    state.history_latencies.append(metrics["p95"])
                    state.history_errors.append(metrics["error_rate"])
                    if len(state.history_latencies) > 30: # 1 minute history
                        state.history_latencies.pop(0)
                        state.history_errors.pop(0)
                        
                    # Anomaly Detection: Latency Spike Check
                    if len(state.history_latencies) >= 5:
                        mean = np.mean(state.history_latencies[:-1])
                        std = np.std(state.history_latencies[:-1])
                        if std > 0.01 and metrics["p95"] > mean + 3 * std:
                            await publish_event(
                                redis_client, "ANOMALY_DETECTED", s, "WARNING",
                                f"Anomaly Detected: Latency Spike! Current P95={metrics['p95']*1000:.1f}ms exceeds threshold. Mean={mean*1000:.1f}ms, Std={std*1000:.1f}ms",
                                db_session
                            )
                            # Quick clamp limit to mitigate immediate load
                            old_limit = state.current_limit
                            state.current_limit = max(MIN_RPS, state.current_limit * 0.4)
                            await redis_client.set(f"limit:capacity:{s}", str(state.current_limit))
                            await redis_client.set(f"limit:rate:{s}", str(state.current_limit))
                            await save_audit_log(
                                s, "ADAPTIVE_LIMIT_CLAMPED", str(old_limit), str(state.current_limit),
                                "Anomaly detection triggered immediate clamp due to latency spike",
                                metrics, db_session
                            )

                # Step 2: Evaluate State Machines & Backpressure propagation
                # We analyze from downstream (C) to upstream (A)
                for s in reversed(SERVICES):
                    metrics = metrics_summary[s]
                    state = service_states[s]
                    now = time.time()
                    
                    # Circuit Breaker Logic
                    if state.circuit_state == "CLOSED":
                        # If critical error rate or latency is exceeded, open circuit
                        if metrics["error_rate"] >= 0.10 or metrics["p95"] >= 0.50:
                            state.circuit_state = "OPEN"
                            state.last_circuit_change = now
                            await redis_client.set(f"circuit_breaker:state:{s}", "OPEN")
                            await publish_event(
                                redis_client, "CIRCUIT_OPENED", s, "CRITICAL",
                                f"Circuit breaker OPENED. Error Rate={metrics['error_rate']*100:.1f}%, P95 Latency={metrics['p95']*1000:.1f}ms",
                                db_session
                            )
                            # Instantly drop limits to minimum
                            old_limit = state.current_limit
                            state.current_limit = MIN_RPS
                            await redis_client.set(f"limit:capacity:{s}", str(MIN_RPS))
                            await redis_client.set(f"limit:rate:{s}", str(MIN_RPS))
                            await save_audit_log(
                                s, "CIRCUIT_OPEN", str(old_limit), str(MIN_RPS),
                                "Circuit opened due to critical failures. Rate limits clamped.",
                                metrics, db_session
                            )
                    
                    elif state.circuit_state == "OPEN":
                        # Stay open for 8 seconds, then transition to HALF_OPEN
                        if now - state.last_circuit_change > 8.0:
                            state.circuit_state = "HALF_OPEN"
                            state.last_circuit_change = now
                            await redis_client.set(f"circuit_breaker:state:{s}", "HALF_OPEN")
                            await publish_event(
                                redis_client, "CIRCUIT_HALF_OPEN", s, "WARNING",
                                "Circuit transitioned to HALF-OPEN. Probing downstream connectivity.",
                                db_session
                            )
                            
                    elif state.circuit_state == "HALF_OPEN":
                        # If errors still persist in Half-Open, return to OPEN
                        if metrics["error_rate"] > 0.05:
                            state.circuit_state = "OPEN"
                            state.last_circuit_change = now
                            await redis_client.set(f"circuit_breaker:state:{s}", "OPEN")
                            await publish_event(
                                redis_client, "CIRCUIT_REOPENED", s, "CRITICAL",
                                f"HALF-OPEN probe failed. Circuit re-opened. Error Rate={metrics['error_rate']*100:.1f}%",
                                db_session
                            )
                        # If healthy over multiple evaluations, close the circuit
                        elif metrics["total_requests"] >= 1 and metrics["error_rate"] <= 0.01:
                            state.healthy_ticks += 1
                            if state.healthy_ticks >= 2: # Sustained health (4 seconds)
                                state.circuit_state = "CLOSED"
                                state.healthy_ticks = 0
                                state.last_circuit_change = now
                                await redis_client.set(f"circuit_breaker:state:{s}", "CLOSED")
                                await publish_event(
                                    redis_client, "CIRCUIT_CLOSED", s, "INFO",
                                    "Circuit closed. Normal operations fully restored.",
                                    db_session
                                )
                                # Gradually scale up limits
                                old_limit = state.current_limit
                                state.current_limit = 20.0
                                await redis_client.set(f"limit:capacity:{s}", str(state.current_limit))
                                await redis_client.set(f"limit:rate:{s}", str(state.current_limit))
                                await save_audit_log(
                                    s, "CIRCUIT_CLOSED_RECOVERY", str(old_limit), str(state.current_limit),
                                    "Circuit closed. Starting gradual load recovery.",
                                    metrics, db_session
                                )
                    
                    # Adaptive Throttling and Backpressure logic for CLOSED circuit
                    if state.circuit_state == "CLOSED":
                        if metrics["status"] == "CRITICAL":
                            # Clamp limit down by 60%
                            old_limit = state.current_limit
                            state.current_limit = max(MIN_RPS, state.current_limit * 0.4)
                            await redis_client.set(f"limit:capacity:{s}", str(state.current_limit))
                            await redis_client.set(f"limit:rate:{s}", str(state.current_limit))
                            await save_audit_log(
                                s, "LIMIT_CLAMPED_CRITICAL", str(old_limit), str(state.current_limit),
                                f"Critical condition. P95={metrics['p95']*1000:.1f}ms, Error Rate={metrics['error_rate']*100:.1f}%",
                                metrics, db_session
                            )
                            
                        elif metrics["status"] == "DEGRADED":
                            # Clamp limit down by 30%
                            old_limit = state.current_limit
                            state.current_limit = max(MIN_RPS, state.current_limit * 0.7)
                            await redis_client.set(f"limit:capacity:{s}", str(state.current_limit))
                            await redis_client.set(f"limit:rate:{s}", str(state.current_limit))
                            await save_audit_log(
                                s, "LIMIT_CLAMPED_DEGRADED", str(old_limit), str(state.current_limit),
                                f"Degraded condition. P95={metrics['p95']*1000:.1f}ms, Error Rate={metrics['error_rate']*100:.1f}%",
                                metrics, db_session
                            )
                            
                        elif metrics["status"] == "HEALTHY":
                            # If capacity is not at default, gradually restore
                            if state.current_limit < DEFAULT_RPS:
                                old_limit = state.current_limit
                                # Incremental increase
                                state.current_limit = min(DEFAULT_RPS, state.current_limit + 20.0)
                                await redis_client.set(f"limit:capacity:{s}", str(state.current_limit))
                                await redis_client.set(f"limit:rate:{s}", str(state.current_limit))
                                await save_audit_log(
                                    s, "LIMIT_RESTORED_STEP", str(old_limit), str(state.current_limit),
                                    "Gradual self-healing recovery step.",
                                    metrics, db_session
                                )

                # Step 3: Coordinate Backpressure Propagation
                # If service C degrades/opens, we immediately reduce limits of B and A
                # If B degrades/opens, we immediately reduce limits of A
                
                # Propagation C -> B -> A
                if service_states["service-c"].circuit_state in ["OPEN", "HALF_OPEN"] or service_states["service-c"].current_limit < DEFAULT_RPS:
                    c_health_ratio = service_states["service-c"].current_limit / DEFAULT_RPS
                    # Clamp B to at most B's current capacity scaled by C's health, or match C's limit
                    target_b_limit = min(service_states["service-b"].current_limit, service_states["service-c"].current_limit)
                    if service_states["service-b"].current_limit > target_b_limit:
                        old = service_states["service-b"].current_limit
                        service_states["service-b"].current_limit = max(MIN_RPS, target_b_limit)
                        await redis_client.set(f"limit:capacity:service-b", str(service_states["service-b"].current_limit))
                        await redis_client.set(f"limit:rate:service-b", str(service_states["service-b"].current_limit))
                        await save_audit_log(
                            "service-b", "BACKPRESSURE_PROPAGATED", str(old), str(service_states["service-b"].current_limit),
                            "Backpressure propagated from degraded service-c",
                            metrics_summary["service-b"], db_session
                        )
                        
                if service_states["service-b"].circuit_state in ["OPEN", "HALF_OPEN"] or service_states["service-b"].current_limit < DEFAULT_RPS:
                    target_a_limit = min(service_states["service-a"].current_limit, service_states["service-b"].current_limit)
                    if service_states["service-a"].current_limit > target_a_limit:
                        old = service_states["service-a"].current_limit
                        service_states["service-a"].current_limit = max(MIN_RPS, target_a_limit)
                        await redis_client.set(f"limit:capacity:service-a", str(service_states["service-a"].current_limit))
                        await redis_client.set(f"limit:rate:service-a", str(service_states["service-a"].current_limit))
                        await save_audit_log(
                            "service-a", "BACKPRESSURE_PROPAGATED", str(old), str(service_states["service-a"].current_limit),
                            "Backpressure propagated from degraded service-b",
                            metrics_summary["service-a"], db_session
                        )
                        
                # Commit all database changes (audit logs & events)
                await db_session.commit()
                
            except Exception as e:
                logger.error(f"Error in reliability control loop: {e}")
                await db_session.rollback()