Microservices Architecture

import asyncio
import random
import time
import hashlib
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Callable, Any
from enum import Enum
from abc import ABC, abstractmethod
import aiohttp
import logging

logger = logging.getLogger(__name__)

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

@dataclass
class ServiceInstance:
    service_name: str
    instance_id: str
    host: str
    port: int
    metadata: Dict[str, str] = field(default_factory=dict)
    status: ServiceStatus = ServiceStatus.HEALTHY
    last_heartbeat: float = field(default_factory=time.time)
    weight: int = 100
    
    @property
    def address(self) -> str:
        return f"{self.host}:{self.port}"

class ServiceRegistry:
    """In-memory service registry with health checking"""
    
    def __init__(self, heartbeat_timeout: float = 30.0):
        self.services: Dict[str, Dict[str, ServiceInstance]] = {}
        self.heartbeat_timeout = heartbeat_timeout
        self._cleanup_task = None
    
    def register(self, instance: ServiceInstance) -> None:
        """Register a service instance"""
        if instance.service_name not in self.services:
            self.services[instance.service_name] = {}
        
        self.services[instance.service_name][instance.instance_id] = instance
        logger.info(f"Registered {instance.service_name}/{instance.instance_id} at {instance.address}")
    
    def deregister(self, service_name: str, instance_id: str) -> None:
        """Deregister a service instance"""
        if service_name in self.services:
            if instance_id in self.services[service_name]:
                del self.services[service_name][instance_id]
                logger.info(f"Deregistered {service_name}/{instance_id}")
    
    def heartbeat(self, service_name: str, instance_id: str) -> bool:
        """Update heartbeat for an instance"""
        if service_name in self.services:
            if instance_id in self.services[service_name]:
                instance = self.services[service_name][instance_id]
                instance.last_heartbeat = time.time()
                instance.status = ServiceStatus.HEALTHY
                return True
        return False
    
    def get_instances(
        self, 
        service_name: str, 
        healthy_only: bool = True
    ) -> List[ServiceInstance]:
        """Get all instances of a service"""
        if service_name not in self.services:
            return []
        
        instances = list(self.services[service_name].values())
        
        if healthy_only:
            instances = [i for i in instances if i.status == ServiceStatus.HEALTHY]
        
        return instances
    
    async def start_cleanup(self) -> None:
        """Start background task to clean up dead instances"""
        while True:
            await asyncio.sleep(10)
            self._cleanup_expired()
    
    def _cleanup_expired(self) -> None:
        """Remove instances that haven't sent heartbeat"""
        now = time.time()
        for service_name in list(self.services.keys()):
            for instance_id in list(self.services[service_name].keys()):
                instance = self.services[service_name][instance_id]
                if now - instance.last_heartbeat > self.heartbeat_timeout:
                    instance.status = ServiceStatus.UNHEALTHY
                    logger.warning(f"Instance {service_name}/{instance_id} marked unhealthy")


class LoadBalancer(ABC):
    """Abstract load balancer"""
    
    @abstractmethod
    def select(self, instances: List[ServiceInstance]) -> Optional[ServiceInstance]:
        pass

class RoundRobinLoadBalancer(LoadBalancer):
    def __init__(self):
        self.counters: Dict[str, int] = {}
    
    def select(self, instances: List[ServiceInstance]) -> Optional[ServiceInstance]:
        if not instances:
            return None
        
        service_name = instances[0].service_name
        if service_name not in self.counters:
            self.counters[service_name] = 0
        
        index = self.counters[service_name] % len(instances)
        self.counters[service_name] += 1
        
        return instances[index]

class WeightedRoundRobinLoadBalancer(LoadBalancer):
    def __init__(self):
        self.current_weights: Dict[str, Dict[str, int]] = {}
    
    def select(self, instances: List[ServiceInstance]) -> Optional[ServiceInstance]:
        if not instances:
            return None
        
        service_name = instances[0].service_name
        
        if service_name not in self.current_weights:
            self.current_weights[service_name] = {}
        
        weights = self.current_weights[service_name]
        
        # Initialize weights
        for instance in instances:
            if instance.instance_id not in weights:
                weights[instance.instance_id] = 0
        
        # Add original weights
        for instance in instances:
            weights[instance.instance_id] += instance.weight
        
        # Select instance with highest current weight
        max_weight = -1
        selected = None
        for instance in instances:
            if weights[instance.instance_id] > max_weight:
                max_weight = weights[instance.instance_id]
                selected = instance
        
        # Decrease selected instance's weight
        if selected:
            total_weight = sum(i.weight for i in instances)
            weights[selected.instance_id] -= total_weight
        
        return selected

class ConsistentHashLoadBalancer(LoadBalancer):
    """Consistent hashing for sticky sessions"""
    
    def __init__(self, replicas: int = 100):
        self.replicas = replicas
        self.ring: Dict[int, ServiceInstance] = {}
        self.sorted_keys: List[int] = []
    
    def _hash(self, key: str) -> int:
        return int(hashlib.md5(key.encode()).hexdigest(), 16)
    
    def _build_ring(self, instances: List[ServiceInstance]) -> None:
        self.ring.clear()
        for instance in instances:
            for i in range(self.replicas):
                key = f"{instance.instance_id}:{i}"
                hash_val = self._hash(key)
                self.ring[hash_val] = instance
        self.sorted_keys = sorted(self.ring.keys())
    
    def select(
        self, 
        instances: List[ServiceInstance], 
        key: str = None
    ) -> Optional[ServiceInstance]:
        if not instances:
            return None
        
        self._build_ring(instances)
        
        if not key:
            key = str(random.random())
        
        hash_val = self._hash(key)
        
        # Find first node >= hash
        for ring_key in self.sorted_keys:
            if ring_key >= hash_val:
                return self.ring[ring_key]
        
        return self.ring[self.sorted_keys[0]]


class ServiceDiscoveryClient:
    """Client-side service discovery with load balancing"""
    
    def __init__(
        self,
        registry: ServiceRegistry,
        load_balancer: LoadBalancer = None,
        cache_ttl: float = 30.0
    ):
        self.registry = registry
        self.load_balancer = load_balancer or RoundRobinLoadBalancer()
        self.cache_ttl = cache_ttl
        self.cache: Dict[str, tuple] = {}  # (instances, timestamp)
    
    def discover(self, service_name: str) -> Optional[ServiceInstance]:
        """Discover and select a service instance"""
        instances = self._get_instances(service_name)
        if not instances:
            raise ServiceNotFoundException(f"No instances found for {service_name}")
        
        return self.load_balancer.select(instances)
    
    def _get_instances(self, service_name: str) -> List[ServiceInstance]:
        """Get instances with caching"""
        if service_name in self.cache:
            instances, timestamp = self.cache[service_name]
            if time.time() - timestamp < self.cache_ttl:
                return instances
        
        instances = self.registry.get_instances(service_name)
        self.cache[service_name] = (instances, time.time())
        return instances
    
    async def call(
        self,
        service_name: str,
        path: str,
        method: str = "GET",
        **kwargs
    ) -> Any:
        """Make HTTP call to discovered service"""
        instance = self.discover(service_name)
        url = f"http://{instance.address}{path}"
        
        async with aiohttp.ClientSession() as session:
            async with session.request(method, url, **kwargs) as response:
                return await response.json()

class ServiceNotFoundException(Exception):
    pass


# Self-registering service
class SelfRegisteringService:
    """Base class for services that self-register"""
    
    def __init__(
        self,
        service_name: str,
        host: str,
        port: int,
        registry: ServiceRegistry,
        heartbeat_interval: float = 10.0
    ):
        self.instance = ServiceInstance(
            service_name=service_name,
            instance_id=f"{host}:{port}:{random.randint(1000, 9999)}",
            host=host,
            port=port
        )
        self.registry = registry
        self.heartbeat_interval = heartbeat_interval
        self._heartbeat_task = None
    
    async def start(self) -> None:
        """Register and start heartbeat"""
        self.registry.register(self.instance)
        self._heartbeat_task = asyncio.create_task(self._send_heartbeats())
    
    async def stop(self) -> None:
        """Deregister and stop heartbeat"""
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
        self.registry.deregister(
            self.instance.service_name,
            self.instance.instance_id
        )
    
    async def _send_heartbeats(self) -> None:
        while True:
            await asyncio.sleep(self.heartbeat_interval)
            self.registry.heartbeat(
                self.instance.service_name,
                self.instance.instance_id
            )


# Usage example
async def main():
    # Create registry
    registry = ServiceRegistry()
    
    # Register some services
    registry.register(ServiceInstance(
        service_name="order-service",
        instance_id="order-1",
        host="10.0.0.1",
        port=8080,
        weight=100
    ))
    registry.register(ServiceInstance(
        service_name="order-service",
        instance_id="order-2",
        host="10.0.0.2",
        port=8080,
        weight=50  # Less capacity
    ))
    
    # Create discovery client
    client = ServiceDiscoveryClient(
        registry,
        load_balancer=WeightedRoundRobinLoadBalancer()
    )
    
    # Discover and call service
    for _ in range(10):
        instance = client.discover("order-service")
        print("Selected:", instance.address)

import asyncio
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional, Callable
from enum import Enum
from datetime import datetime
import uuid
import logging

logger = logging.getLogger(__name__)

class SagaStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    COMPENSATING = "compensating"
    COMPENSATED = "compensated"
    FAILED = "failed"

@dataclass
class SagaStep:
    name: str
    action: Callable
    compensation: Callable
    timeout: float = 30.0
    retries: int = 3

@dataclass
class StepResult:
    step_name: str
    success: bool
    data: Any = None
    error: Optional[str] = None

@dataclass
class Saga:
    saga_id: str
    name: str
    steps: List[SagaStep]
    status: SagaStatus = SagaStatus.PENDING
    current_step: int = 0
    context: Dict[str, Any] = field(default_factory=dict)
    results: List[StepResult] = field(default_factory=list)
    created_at: datetime = field(default_factory=datetime.now)
    completed_at: Optional[datetime] = None

class SagaOrchestrator:
    """
    Central saga orchestrator for distributed transactions.
    Handles step execution, compensation, and state management.
    """
    
    def __init__(self, store: 'SagaStore' = None):
        self.store = store or InMemorySagaStore()
    
    async def execute(
        self,
        name: str,
        steps: List[SagaStep],
        context: Dict[str, Any] = None
    ) -> Saga:
        """Execute a saga"""
        saga = Saga(
            saga_id=str(uuid.uuid4()),
            name=name,
            steps=steps,
            context=context or {}
        )
        
        await self.store.save(saga)
        
        try:
            saga.status = SagaStatus.RUNNING
            await self._execute_steps(saga)
            saga.status = SagaStatus.COMPLETED
            
        except SagaStepFailed as e:
            logger.error(f"Saga {saga.saga_id} step failed: {e}")
            saga.status = SagaStatus.COMPENSATING
            await self._compensate(saga)
            saga.status = SagaStatus.COMPENSATED
            
        except Exception as e:
            logger.error(f"Saga {saga.saga_id} failed unexpectedly: {e}")
            saga.status = SagaStatus.FAILED
        
        finally:
            saga.completed_at = datetime.now()
            await self.store.save(saga)
        
        return saga
    
    async def _execute_steps(self, saga: Saga) -> None:
        """Execute saga steps in order"""
        for i, step in enumerate(saga.steps):
            saga.current_step = i
            await self.store.save(saga)
            
            result = await self._execute_step_with_retry(saga, step)
            saga.results.append(result)
            
            if not result.success:
                raise SagaStepFailed(step.name, result.error)
            
            # Store step result in context for next steps
            saga.context[f"{step.name}_result"] = result.data
    
    async def _execute_step_with_retry(
        self, 
        saga: Saga, 
        step: SagaStep
    ) -> StepResult:
        """Execute a step with retries"""
        last_error = None
        
        for attempt in range(step.retries):
            try:
                async with asyncio.timeout(step.timeout):
                    result = await step.action(saga.context)
                    return StepResult(
                        step_name=step.name,
                        success=True,
                        data=result
                    )
                    
            except asyncio.TimeoutError:
                last_error = f"Step timed out after {step.timeout}s"
                
            except Exception as e:
                last_error = str(e)
                
            if attempt < step.retries - 1:
                await asyncio.sleep(2 ** attempt)  # Exponential backoff
        
        return StepResult(
            step_name=step.name,
            success=False,
            error=last_error
        )
    
    async def _compensate(self, saga: Saga) -> None:
        """Execute compensating actions in reverse order"""
        # Only compensate completed steps
        completed_steps = saga.results[:saga.current_step]
        
        for i in range(len(completed_steps) - 1, -1, -1):
            step = saga.steps[i]
            
            try:
                logger.info(f"Compensating step: {step.name}")
                async with asyncio.timeout(step.timeout):
                    await step.compensation(saga.context)
                    
            except Exception as e:
                logger.error(f"Compensation failed for {step.name}: {e}")
                # Continue compensating other steps


class SagaStepFailed(Exception):
    def __init__(self, step_name: str, error: str):
        self.step_name = step_name
        self.error = error
        super().__init__(f"Step '{step_name}' failed: {error}")


class SagaStore(ABC):
    @abstractmethod
    async def save(self, saga: Saga) -> None:
        pass
    
    @abstractmethod
    async def get(self, saga_id: str) -> Optional[Saga]:
        pass

class InMemorySagaStore(SagaStore):
    def __init__(self):
        self.sagas: Dict[str, Saga] = {}
    
    async def save(self, saga: Saga) -> None:
        self.sagas[saga.saga_id] = saga
    
    async def get(self, saga_id: str) -> Optional[Saga]:
        return self.sagas.get(saga_id)


# Example: Order processing saga
async def create_order_saga(order_data: dict) -> Saga:
    orchestrator = SagaOrchestrator()
    
    # Define saga steps
    steps = [
        SagaStep(
            name="create_order",
            action=lambda ctx: create_order(ctx["order_data"]),
            compensation=lambda ctx: cancel_order(ctx["create_order_result"]["order_id"]),
            timeout=10.0
        ),
        SagaStep(
            name="reserve_inventory",
            action=lambda ctx: reserve_inventory(
                ctx["create_order_result"]["order_id"],
                ctx["order_data"]["items"]
            ),
            compensation=lambda ctx: release_inventory(
                ctx["create_order_result"]["order_id"]
            ),
            timeout=15.0
        ),
        SagaStep(
            name="charge_payment",
            action=lambda ctx: charge_payment(
                ctx["create_order_result"]["order_id"],
                ctx["order_data"]["amount"]
            ),
            compensation=lambda ctx: refund_payment(
                ctx["charge_payment_result"]["transaction_id"]
            ),
            timeout=30.0,
            retries=3
        ),
        SagaStep(
            name="ship_order",
            action=lambda ctx: initiate_shipping(
                ctx["create_order_result"]["order_id"]
            ),
            compensation=lambda ctx: cancel_shipping(
                ctx["ship_order_result"]["shipment_id"]
            ),
            timeout=20.0
        )
    ]
    
    return await orchestrator.execute(
        name="order_processing",
        steps=steps,
        context={"order_data": order_data}
    )


# Service functions (would be actual service calls)
async def create_order(order_data: dict) -> dict:
    return {"order_id": str(uuid.uuid4()), "status": "created"}

async def cancel_order(order_id: str) -> None:
    logger.info(f"Cancelling order {order_id}")

async def reserve_inventory(order_id: str, items: list) -> dict:
    return {"reservation_id": str(uuid.uuid4())}

async def release_inventory(order_id: str) -> None:
    logger.info(f"Releasing inventory for order {order_id}")

async def charge_payment(order_id: str, amount: float) -> dict:
    return {"transaction_id": str(uuid.uuid4()), "amount": amount}

async def refund_payment(transaction_id: str) -> None:
    logger.info(f"Refunding transaction {transaction_id}")

async def initiate_shipping(order_id: str) -> dict:
    return {"shipment_id": str(uuid.uuid4())}

async def cancel_shipping(shipment_id: str) -> None:
    logger.info(f"Cancelling shipment {shipment_id}")


# Usage
async def main():
    result = await create_order_saga({
        "items": [{"product_id": "123", "quantity": 2}],
        "amount": 99.99,
        "customer_id": "cust_123"
    })
    
    print(f"Saga status: {result.status}")
    print(f"Results: {result.results}")