Health Check DB Query — Load Balancer 503 Outage

from dataclasses import dataclass, field
from enum import Enum
from typing import List, Optional
import time
import threading
from io.thecodeforge.loadbalancer.health import HealthChecker
from io.thecodeforge.loadbalancer.algorithms import LoadBalancingAlgorithm

class BackendState(Enum):
    HEALTHY = "healthy"
    UNHEALTHY = "unhealthy"
    DRAINING = "draining"

@dataclass
class BackendServer:
    host: str
    port: int
    weight: int = 1
    state: BackendState = BackendState.HEALTHY
    active_connections: int = 0
    last_health_check: float = 0.0
    consecutive_failures: int = 0
    
    @property
    def address(self) -> str:
        return f"{self.host}:{self.port}"
    
    def is_available(self) -> bool:
        return self.state in (BackendState.HEALTHY, BackendState.DRAINING)


class LoadBalancer:
    """
    Production-grade load balancer with health checking,
    connection draining, and multiple routing algorithms.
    """
    
    def __init__( health_check_interval: float = 5.0):
        self.backends: List[BackendServer] = []
        self.algorithm = algorithm
        self.health_checker = HealthChecker(interval=health_check_interval)
        self._lock = threading.Lock()
        self._minimum_healthy_hosts: int = 0
    
    def add_backend(self, host: str, port: int, weight: int = 1) -> BackendServer:
        """
        Register a new backend server with the load balancer.
        """
        with self._lock:
            backend = BackendServer(host=host, port=port, weight=weight)
            self.backends.append(backend)
            self.health_checker.register(backend)
            return backend
    
    def remove_backend(self, backend: BackendServer) -> None:
        """
        Gracefully remove a backend with connection draining.
        """
        with self._lock:
            backend.state = BackendState.DRAINING
            self.health_checker.unregister(backend)
    
    def select_backend(self) -> Optional[BackendServer]:
        """
        Select next backend using configured algorithm.
        Respects minimum_healthy_hosts threshold.
        """
        with self._lock:
            available = [b for b in self.backends if b.is_available()]
            healthy = [b for b in available if b.state == BackendState.HEALTHY]
            
            if len(healthy) < self._minimum_healthy_hosts and available:
                return self.algorithm.select(available)
            
            if not healthy:
                return None
            
            return self.algorithm.select(healthy)
    
    def set_minimum_healthy_hosts(self, count: int) -> None:
        """
        Configure minimum healthy backends before routing stops.
        Prevents complete pool exhaustion during failures.
        """
        self._minimum_healthy_hosts = count


# Example usage
from io.thecodeforge.loadbalancer.algorithms import RoundRobinAlgorithm

lb = LoadBalancer(algorithm=RoundRobinAlgorithm(), health_check_interval=5.0)
lb.add_backend("10.0.1.10", 8080)
lb.add_backend("10.0.1.11", 8080)
lb.addself, algorithm: LoadBalancingAlgorithm,_backend("10.0.1.12", 8080)
lb.set_minimum_healthy_hosts(1)

for i in range(6):
    backend = lb.select_backend()
    if backend:
        print(f"Request {i} -> {backend.address}")

from abc import ABC, abstractmethod
from typing import Optional, Dict, Any
from io.thecodeforge.loadbalancer.core import BackendServer
from io.thecodeforge.loadbalancer.models import Request, Connection

class LoadBalancerType(ABC):
    """
    Abstract base for Layer 4 and Layer 7 load balancer types.
    """
    
    @abstractmethod
    def route(self, request: Any) -> Optional[BackendServer]:
        pass


class Layer4LoadBalancer(LoadBalancerType):
    """
    Transport-layer load balancer.
    Routes based on source/destination IP and port.
    Does not inspect packet contents.
    """
    
    def __init__(self, balancer):
        self.balancer = balancer
    
    def route(self, connection: Connection) -> Optional[BackendServer]:
        """
        Route based on 5-tuple: src_ip, src_port, dst_ip, dst_port, protocol.
        """
        backend = self.balancer.select_backend()
        if backend:
            backend.active_connections += 1
        return backend
    
    def get_capabilities(self) -> Dict[str, bool]:
        return {
            "protocol_agnostic": True,
            "url_routing": False,
            "header_inspection": False,
            "cookie_persistence": False,
            "content_based_routing": False,
            "ssl_termination": False,
            "websocket_support": True,
            "latency_overhead": "minimal"
        }


class Layer7LoadBalancer(LoadBalancerType):
    """
    Application-layer load balancer.
    Routes based on HTTP headers, URL path, hostname, cookies.
    """
    
    def __init__(self, balancer):
        self.balancer = balancer
        self.rules: list = []
    
    def add_routing_rule(self, condition: callable, target_pool: str) -> None:
        """
        Add content-based routing rule.
        """
        self.rules.append({"condition": condition, "pool": target_pool})
    
    def route(self, request: Request) -> Optional[BackendServer]:
        """
        Route based on HTTP request content.
        """
        for rule in self.rules:
            if rule["condition"](request):
                pool = rule["pool"]
                return self.balancer.select_backend_from_pool(pool)
        
        return self.balancer.select_backend()
    
    def get_capabilities(self) -> Dict[str, bool]:
        return {
            "protocol_agnostic": False,
            "url_routing": True,
            "header_inspection": True,
            "cookie_persistence": True,
            "content_based_routing": True,
            "ssl_termination": True,
            "websocket_support": True,
            "latency_overhead": "moderate"
        }


# Example: Layer 7 routing rules
from io.thecodeforge.loadbalancer.algorithms import WeightedRoundRobin

l7 = Layer7LoadBalancer(balancer=WeightedRoundRobin())

# Route API traffic to API servers
l7.add_routing_rule(
    condition=lambda req: req.path.startswith("/api/"),
    target_pool="api-servers"
)

# Route static assets to CDN-backed servers
l7.add_routing_rule(
    condition=lambda req: req.path.startswith("/static/"),
    target_pool="static-servers"
)

# Route by hostname
l7.add_routing_rule(
    condition=lambda req: req.host == "api.example.com",
    target_pool="api-servers"
)

from abc import ABC, abstractmethod
from typing import List, Optional
import hashlib
import random
from io.thecodeforge.loadbalancer.core import BackendServer


class LoadBalancingAlgorithm(ABC):
    """
    Abstract base for all load balancing algorithms.
    """
    
    @abstractmethod
    def select(self, backends: List[BackendServer]) -> Optional[BackendServer]:
        pass


class RoundRobinAlgorithm(LoadBalancingAlgorithm):
    """
    Distributes requests sequentially across all healthy backends.
    Simple and fair when servers have equal capacity.
    """
    
    def __init__(self):
        self._index = 0
    
    def select(self, backends: List[BackendServer]) -> Optional[BackendServer]:
        if not backends:
            return None
        backend = backends[self._index % len(backends)]
        self._index += 1
        return backend


class WeightedRoundRobinAlgorithm(LoadBalancingAlgorithm):
    """
    Distributes requests proportionally based on server weights.
    Higher weight servers receive proportionally more requests.
    """
    
    def __init__(self):
        self._current_weights: dict = {}
        self._index = 0
    
    def select(self, backends: List[BackendServer]) -> Optional[BackendServer]:
        if not backends:
            return None
        
        total_weight = sum(b.weight for b in backends)
        
        for backend in backends:
            addr = backend.address
            if addr not in self._current_weights:
                self._current_weights[addr] = 0
            self._current_weights[addr] += backend.weight
        
        selected = max(backends, key=lambda b: self._current_weights[b.address])
        self._current_weights[selected.address] -= total_weight
        return selected


class LeastConnectionsAlgorithm(LoadBalancingAlgorithm):
    """
    Routes to the server with the fewest active connections.
    Best for workloads with variable request durations.
    """
    
    def select(self, backends: List[BackendServer]) -> Optional[BackendServer]:
        if not backends:
            return None
        return min(backends, key=lambda b: b.active_connections)


class IpHashAlgorithm(LoadBalancingAlgorithm):
    """
    Routes based on hash of client IP address.
    Provides session affinity without cookies.
    """
    
    def __init__(self, client_ip_getter: callable = None):
        self._get_client_ip = client_ip_getter or (lambda: "127.0.0.1")
    
    def select(self, backends: List[BackendServer]) -> Optional[BackendServer]:
        if not backends:
            return None
        
        client_ip = self._get_client_ip()
        hash_value = int(hashlib.md5(client_ip.encode()).hexdigest(), 16)
        index = hash_value % len(backends)
        return backends[index]


class P2CLeastConnectionsAlgorithm(LoadBalancingAlgorithm):
    """
    Power of Two Choices: randomly pick two backends,
    then route to the one with fewer connections.
    Near-optimal load distribution with O(1) selection.
    """
    
    def select(self, backends: List[BackendServer]) -> Optional[BackendServer]:
        if not backends:
            return None
        if len(backends) == 1:
            return backends[0]
        
        a, b = random.sample(backends, 2)
        return a if a.active_connections <= b.active_connections else b


# Algorithm comparison
algorithms = {
    "Round Robin": "Simple sequential distribution. Assumes equal server capacity.",
    "Weighted Round Robin": "Proportional distribution based on server weight. For heterogeneous pools.",
    "Least Connections": "Routes to fewest active connections. Best for variable-duration requests.",
    "IP Hash": "Consistent routing by client IP. Provides session affinity without cookies.",
    "P2C Least Connections": "Near-optimal distribution with O(1) complexity. Used by Envoy and gRPC."
}

import time
import threading
import requests
from dataclasses import dataclass
from enum import Enum
from typing import Callable, Optional
from io.thecodeforge.loadbalancer.core import BackendServer, BackendState


class HealthCheckType(Enum):
    HTTP = "http"
    TCP = "tcp"
    GRPC = "grpc"


@dataclass
class HealthCheckConfig:
    check_type: HealthCheckType
    path: str = "/health"
    port: Optional[int] = None
    interval_seconds: float = 5.0
    timeout_seconds: float = 2.0
    healthy_threshold: int = 2
    unhealthy_threshold: int = 3
    expected_status_codes: list = None
    
    def __post_init__(self):
        if self.expected_status_codes is None:
            self.expected_status_codes = [200]


class HealthChecker:
    """
    Production health checker with configurable thresholds,
    grace periods, and passive failure detection.
    """
    
    def __init__(self, config: HealthCheckConfig = None):
        self.config = config or HealthCheckConfig(
            check_type=HealthCheckType.HTTP,
            path="/health"
        )
        self._backends: dict = {}
        self._running = False
        self._thread: Optional[threading.Thread] = None
    
    def register(self, backend: BackendServer) -> None:
        """
        Register a backend for health checking.
        """
        self._backends[backend.address] = {
            "backend": backend,
            "consecutive_successes": 0,
            "consecutive_failures": 0,
            "last_check_time": 0.0
        }
    
    def unregister(self, backend: BackendServer) -> None:
        """
        Remove a backend from health checking.
        """
        self._backends.pop(backend.address, None)
    
    def check_http(self, backend: BackendServer) -> bool:
        """
        Perform HTTP health check against backend.
        """
        port = self.config.port or backend.port
        url = f"http://{backend.host}:{port}{self.config.path}"
        
        try:
            response = requests.get(
                url,
                timeout=self.config.timeout_seconds
            )
            return response.status_code in self.config.expected_status_codes
        except (requests.ConnectionError, requests.Timeout):
            return False
    
    def check_tcp(self, backend: BackendServer) -> bool:
        """
        Perform TCP connection check against backend.
        """
        import socket
        port = self.config.port or backend.port
        
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(self.config.timeout_seconds)
            result = sock.connect_ex((backend.host, port))
            sock.close()
            return result == 0
        except socket.error:
            return False
    
    def run_check(self, backend: BackendServer) -> bool:
        """
        Execute health check and update backend state based on thresholds.
        """
        if self.config.check_type == HealthCheckType.HTTP:
            is_healthy = self.check_http(backend)
        else:
            is_healthy = self.check_tcp(backend)
        
        entry = self._backends.get(backend.address)
        if not entry:
            return is_healthy
        
        if is_healthy:
            entry["consecutive_failures"] = 0
            entry["consecutive_successes"] += 1
            
            if entry["consecutive_successes"] >= self.config.healthy_threshold:
                backend.state = BackendState.HEALTHY
                backend.consecutive_failures = 0
        else:
            entry["consecutive_successes"] = 0
            entry["consecutive_failures"] += 1
            
            if entry["consecutive_failures"] >= self.config.unhealthy_threshold:
                backend.state = BackendState.UNHEALTHY
                backend.consecutive_failures = entry["consecutive_failures"]
        
        entry["last_check_time"] = time.time()
        return is_healthy
    
    def start(self) -> None:
        """
        Start background health checking thread.
        """
        self._running = True
        self._thread = threading.Thread(target=self._check_loop, daemon=True)
        self._thread.start()
    
    def stop(self) -> None:
        """
        Stop background health checking.
        """
        self._running = False
        if self._thread:
            self._thread.join(timeout=10.0)
    
    def _check_loop(self) -> None:
        """
        Continuous health check loop.
        """
        while self._running:
            for entry in list(self._backends.values()):
                backend = entry["backend"]
                if backend.state != BackendState.DRAINING:
                    self.run_check(backend)
            time.sleep(self.config.interval_seconds)


# Example: configure health checks
config = HealthCheckConfig(
    check_type=HealthCheckType.HTTP,
    path="/health/ready",
    interval_seconds=5.0,
    timeout_seconds=2.0,
    healthy_threshold=2,
    unhealthy_threshold=3,
    expected_status_codes=[200, 204]
)

checker = HealthChecker(config=config)

import hashlib
import time
from typing import Dict, Optional
from dataclasses import dataclass
from io.thecodeforge.loadbalancer.core import BackendServer


@dataclass
class SessionEntry:
    backend_address: str
    created_at: float
    last_used: float
    ttl_seconds: float
    
    def is_expired(self) -> bool:
        return time.time() - self.last_used > self.ttl_seconds


class SessionPersistenceManager:
    """
    Manages session affinity between clients and backends.
    Supports cookie-based and IP-based persistence.
    """
    
    def __init__(self, ttl_seconds: int = 3600, cookie_name: str = "SERVERID"):
        self._sessions: Dict[str, SessionEntry] = {}
        self._ttl = ttl_seconds
        self._cookie_name = cookie_name
    
    def get_backend_for_client(
        self,
        client_id: str,
        healthy_backends: list
    ) -> Optional[BackendServer]:
        """
        Look up persisted backend for client.
        Returns None if session expired or backend unhealthy.
        """
        entry = self._sessions.get(client_id)
        
        if entry is None or entry.is_expired():
            return None
        
        for backend in healthy_backends:
            if backend.address == entry.backend_address:
                entry.last_used = time.time()
                return backend
        
        del self._sessions[client_id]
        return None
    
    def persist_session(self, client_id: str, backend: BackendServer) -> None:
        """
        Create or update session affinity for a client.
        """
        self._sessions[client_id] = SessionEntry(
            backend_address=backend.address,
            created_at=time.time(),
            last_used=time.time(),
            ttl_seconds=self._ttl
        )
    
    def extract_client_id_from_cookie(self, cookies: dict) -> Optional[str]:
        """
        Extract client identifier from load balancer cookie.
        """
        return cookies.get(self._cookie_name)
    
    def create_session_cookie(self, client_id: str, backend: BackendServer) -> dict:
        """
        Create cookie header for session persistence.
        """
        return {
            "name": self._cookie_name,
            "value": backend.address,
            "max_age": self._ttl,
            "path": "/",
            "http_only": True,
            "secure": True
        }
    
    def cleanup_expired(self) -> int:
        """
        Remove expired session entries.
        Returns count of removed entries.
        """
        expired = [
            k for k, v in self._sessions.items()
            if v.is_expired()
        ]
        for key in expired:
            del self._sessions[key]
        return len(expired)
    
    @property
    def active_sessions(self) -> int:
        return len(self._sessions)


class ConsistentHashPersistence:
    """
    IP-based persistence using consistent hashing.
    Minimizes redistribution when backends are added or removed.
    """
    
    def __init__(self, virtual_nodes: int = 150):
        self._virtual_nodes = virtual_nodes
        self._ring: Dict[int, str] = {}
    
    def _hash(self, key: str) -> int:
        return int(hashlib.md5(key.encode()).hexdigest(), 16)
    
    def add_backend(self, backend: BackendServer) -> None:
        for i in range(self._virtual_nodes):
            vnode_key = f"{backend.address}#{i}"
            hash_val = self._hash(vnode_key)
            self._ring[hash_val] = backend.address
    
    def remove_backend(self, backend: BackendServer) -> None:
        for i in range(self._virtual_nodes):
            vnode_key = f"{backend.address}#{i}"
            hash_val = self._hash(vnode_key)
            self._ring.pop(hash_val, None)
    
    def get_backend(self, client_ip: str) -> Optional[str]:
        if not self._ring:
            return None
        
        hash_val = self._hash(client_ip)
        sorted_hashes = sorted(self._ring.keys())
        
        for h in sorted_hashes:
            if h >= hash_val:
                return self._ring[h]
        
        return self._ring[sorted_hashes[0]]