control_plane/main.py

import os
import asyncio
import json
import time
import httpx
from typing import List, Optional
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import redis.asyncio as redis
from sqlalchemy import select, desc
from database import init_db, async_session
from models import AuditLog, ReliabilityEvent, BenchmarkResult
from controller import run_reliability_loop, service_states

app = FastAPI(title="PulseGuard Control Plane")

# CORS middleware for React UI development
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
redis_client = redis.from_url(REDIS_URL, decode_responses=True)

# Services downstream mappings for REST forwarding
SERVICES_URLS = {
    "service-a": "http://service-a:8000",
    "service-b": "http://service-b:8000",
    "service-c": "http://service-c:8000"
}

# WebSocket connections storage
class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except Exception:
                pass

manager = ConnectionManager()

# Background task to stream states via WebSocket
async def websocket_streamer_task():
    while True:
        await asyncio.sleep(0.5)
        if not manager.active_connections:
            continue
            
        try:
            # Gather state from Redis and memory
            stats = {}
            for s in ["service-a", "service-b", "service-c"]:
                capacity = await redis_client.get(f"limit:capacity:{s}") or "100"
                rate = await redis_client.get(f"limit:rate:{s}") or "100"
                cb_state = await redis_client.get(f"circuit_breaker:state:{s}") or "CLOSED"
                
                # Fetch recent metrics summary from memory state / redis
                from controller import evaluate_service_metrics
                metrics = await evaluate_service_metrics(redis_client, s)
                
                stats[s] = {
                    "limit_capacity": float(capacity),
                    "limit_rate": float(rate),
                    "circuit_state": cb_state,
                    "status": metrics["status"],
                    "rps": metrics["rps"],
                    "error_rate": metrics["error_rate"],
                    "p50": metrics["p50"],
                    "p95": metrics["p95"],
                    "p99": metrics["p99"],
                    "total_requests": metrics["total_requests"]
                }
            
            # Fetch latest 10 events
            async with async_session() as session:
                query_events = select(ReliabilityEvent).order_by(desc(ReliabilityEvent.timestamp)).limit(10)
                res_events = await session.execute(query_events)
                events = res_events.scalars().all()
                events_list = [
                    {
                        "id": e.id,
                        "timestamp": e.timestamp.isoformat() if e.timestamp else "",
                        "event_type": e.event_type,
                        "service": e.service,
                        "severity": e.severity,
                        "message": e.message
                    } for e in events
                ]
                
                query_audit = select(AuditLog).order_by(desc(AuditLog.timestamp)).limit(10)
                res_audit = await session.execute(query_audit)
                audits = res_audit.scalars().all()
                audits_list = [
                    {
                        "id": a.id,
                        "timestamp": a.timestamp.isoformat() if a.timestamp else "",
                        "service": a.service,
                        "action": a.action,
                        "old_value": a.old_value,
                        "new_value": a.new_value,
                        "reason": a.reason
                    } for a in audits
                ]
                
            payload = {
                "services": stats,
                "events": events_list,
                "audits": audits_list,
                "timestamp": time.time()
            }
            
            await manager.broadcast(json.dumps(payload))
        except Exception as e:
            # logger.error(f"Error in websocket streamer: {e}")
            pass

@app.on_event("startup")
async def startup_event():
    # 1. Initialize Postgres Database Tables
    await init_db()
    
    # 2. Launch Reliability Controller Loop
    asyncio.create_task(run_reliability_loop(redis_client))
    
    # 3. Launch WebSocket Broadcasting Task
    asyncio.create_task(websocket_streamer_task())

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            # Keep connection alive, listen for any client messages if needed
            data = await websocket.receive_text()
    except WebSocketDisconnect:
        manager.disconnect(websocket)

# REST APIs

@app.get("/api/events")
async def get_events(limit: int = 50):
    async with async_session() as session:
        query = select(ReliabilityEvent).order_by(desc(ReliabilityEvent.timestamp)).limit(limit)
        res = await session.execute(query)
        events = res.scalars().all()
        return events

@app.get("/api/audit-logs")
async def get_audit_logs(limit: int = 50):
    async with async_session() as session:
        query = select(AuditLog).order_by(desc(AuditLog.timestamp)).limit(limit)
        res = await session.execute(query)
        logs = res.scalars().all()
        return logs

@app.get("/api/chaos/status")
async def get_chaos_status():
    status = {}
    async with httpx.AsyncClient(timeout=2.0) as client:
        for s, url in SERVICES_URLS.items():
            try:
                res = await client.get(f"{url}/chaos/status")
                status[s] = res.json()
            except Exception:
                status[s] = {"error": "unreachable"}
    return status

class ChaosConfigPayload(BaseModel):
    delay: Optional[float] = 0.0
    error_rate: Optional[float] = 0.0
    crash: Optional[bool] = False
    cpu_load: Optional[float] = 0.0
    memory_load: Optional[int] = 0
    packet_loss: Optional[float] = 0.0

@app.post("/api/chaos/inject/{service}")
async def inject_chaos(service: str, payload: ChaosConfigPayload):
    if service not in SERVICES_URLS:
        return {"error": "Invalid service name"}
        
    url = SERVICES_URLS[service]
    async with httpx.AsyncClient(timeout=5.0) as client:
        try:
            res = await client.post(f"{url}/chaos/inject", json=payload.dict())
            return res.json()
        except Exception as e:
            return {"error": f"Failed to forward to {service}: {str(e)}"}

@app.post("/api/chaos/clear/{service}")
async def clear_chaos(service: str):
    if service not in SERVICES_URLS:
        return {"error": "Invalid service name"}
        
    url = SERVICES_URLS[service]
    async with httpx.AsyncClient(timeout=5.0) as client:
        try:
            res = await client.post(f"{url}/chaos/clear")
            return res.json()
        except Exception as e:
            return {"error": f"Failed to forward clear to {service}: {str(e)}"}

@app.post("/api/generator/rps")
async def set_generator_rps(rps: int = Query(..., ge=0)):
    await redis_client.set("load_generator:target_rps", str(rps))
    event = {
        "event_type": "TRAFFIC_GENERATOR_CHANGED",
        "service": "load-generator",
        "severity": "INFO",
        "message": f"Traffic generator rate set to {rps} RPS",
        "timestamp": time.time()
    }
    await redis_client.publish("reliability_events", json.dumps(event))
    return {"status": "success", "rps": rps}


# Resilience Benchmarking
async def run_benchmark_task(test_name: str, has_protection: bool):
    """Executes a structured benchmark test and saves outcomes."""
    logger.info(f"Starting resilience benchmark: {test_name} (Protection={has_protection})")
    
    # Store previous protection state, turn on/off controller
    # For simulation, if has_protection=False we can manually set service limits to infinite
    if not has_protection:
        for s in ["service-a", "service-b", "service-c"]:
            await redis_client.set(f"limit:capacity:{s}", "999999")
            await redis_client.set(f"limit:rate:{s}", "999999")
            await redis_client.set(f"circuit_breaker:state:{s}", "CLOSED")
            
    # Set high traffic mode in Redis (RPS = 200)
    await redis_client.set("load_generator:target_rps", "200")
    
    # Let run for 5 seconds under healthy conditions
    await asyncio.sleep(5.0)
    
    # Inject chaos on Service C
    async with httpx.AsyncClient(timeout=3.0) as client:
        # Crash Service C or delay it heavily
        await client.post(f"{SERVICES_URLS['service-c']}/chaos/inject", json={"delay": 1.5, "error_rate": 0.3})
        
    # Observe failure propagation for 12 seconds
    await asyncio.sleep(12.0)
    
    # Recover Service C
    async with httpx.AsyncClient(timeout=3.0) as client:
        await client.post(f"{SERVICES_URLS['service-c']}/chaos/clear")
        
    # Wait for gradual recovery
    await asyncio.sleep(10.0)
    
    # Set traffic back to normal
    await redis_client.set("load_generator:target_rps", "50")
    
    # Calculate performance over this benchmark duration
    # Fetch Service A metrics
    from controller import evaluate_service_metrics
    metrics = await evaluate_service_metrics(redis_client, "service-a")
    
    # Reset limit defaults if protection was disabled
    if not has_protection:
        for s in ["service-a", "service-b", "service-c"]:
            await redis_client.set(f"limit:capacity:{s}", "100")
            await redis_client.set(f"limit:rate:{s}", "100")
            
    # Save results to Postgres
    async with async_session() as session:
        result = BenchmarkResult(
            test_name=test_name,
            duration_seconds=27,
            total_requests=metrics["total_requests"],
            success_rate=1.0 - metrics["error_rate"],
            p50_latency=metrics["p50"],
            p95_latency=metrics["p95"],
            p99_latency=metrics["p99"],
            has_protection=has_protection,
            notes=f"Simulated heavy delay on C with RPS=200. Protection active={has_protection}"
        )
        session.add(result)
        
        # Log event
        event = ReliabilityEvent(
            event_type="BENCHMARK_COMPLETED",
            service="control-plane",
            severity="INFO",
            message=f"Benchmark '{test_name}' completed. Success Rate: {(1.0-metrics['error_rate'])*100:.1f}%, P95 Latency: {metrics['p95']*1000:.1f}ms"
        )
        session.add(event)
        await session.commit()
        
    logger.info(f"Benchmark completed: {test_name}")

@app.post("/api/benchmark/start")
async def start_benchmark(test_name: str = "Chaos Propagation Test", has_protection: bool = True):
    asyncio.create_task(run_benchmark_task(test_name, has_protection))
    return {"status": "benchmark started in background"}

@app.get("/api/benchmarks")
async def get_benchmarks():
    async with async_session() as session:
        query = select(BenchmarkResult).order_by(desc(BenchmarkResult.timestamp)).limit(20)
        res = await session.execute(query)
        results = res.scalars().all()
        return results