Distributed Systems

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Optional, List, Any, Dict
from datetime import datetime
import hashlib
import json
import asyncio

@dataclass
class VectorClock:
    """Track causality in distributed systems"""
    clocks: Dict[str, int] = field(default_factory=dict)
    
    def increment(self, node_id: str) -> None:
        self.clocks[node_id] = self.clocks.get(node_id, 0) + 1
    
    def merge(self, other: 'VectorClock') -> 'VectorClock':
        """Merge two vector clocks (take max of each)"""
        merged = VectorClock()
        all_nodes = set(self.clocks.keys()) | set(other.clocks.keys())
        for node in all_nodes:
            merged.clocks[node] = max(
                self.clocks.get(node, 0),
                other.clocks.get(node, 0)
            )
        return merged
    
    def is_concurrent_with(self, other: 'VectorClock') -> bool:
        """Check if two events are concurrent (neither happened before the other)"""
        self_greater = any(
            self.clocks.get(k, 0) > other.clocks.get(k, 0)
            for k in self.clocks
        )
        other_greater = any(
            other.clocks.get(k, 0) > self.clocks.get(k, 0)
            for k in other.clocks
        )
        return self_greater and other_greater


@dataclass
class ReplicatedValue:
    """Last-Write-Wins Register for eventual consistency"""
    value: Any
    timestamp: float
    node_id: str
    vector_clock: VectorClock
    
    def merge(self, other: 'ReplicatedValue') -> 'ReplicatedValue':
        """Merge with conflict resolution"""
        # Check for concurrent updates
        if self.vector_clock.is_concurrent_with(other.vector_clock):
            # Conflict! Use timestamp + node_id as tiebreaker
            if (self.timestamp, self.node_id) > (other.timestamp, other.node_id):
                winner = self
            else:
                winner = other
        elif self.timestamp > other.timestamp:
            winner = self
        else:
            winner = other
            
        return ReplicatedValue(
            value=winner.value,
            timestamp=winner.timestamp,
            node_id=winner.node_id,
            vector_clock=self.vector_clock.merge(other.vector_clock)
        )


class EventuallyConsistentStore:
    """Distributed key-value store with eventual consistency"""
    
    def __init__(self, node_id: str, replicas: List['EventuallyConsistentStore'] = None):
        self.node_id = node_id
        self.data: Dict[str, ReplicatedValue] = {}
        self.vector_clock = VectorClock()
        self.replicas = replicas or []
        self.pending_sync: List[tuple] = []
    
    def write(self, key: str, value: Any) -> None:
        """Write with vector clock tracking"""
        self.vector_clock.increment(self.node_id)
        
        replicated = ReplicatedValue(
            value=value,
            timestamp=datetime.now().timestamp(),
            node_id=self.node_id,
            vector_clock=VectorClock(clocks=self.vector_clock.clocks.copy())
        )
        
        # Local write
        if key in self.data:
            self.data[key] = self.data[key].merge(replicated)
        else:
            self.data[key] = replicated
        
        # Queue for async replication
        self.pending_sync.append((key, replicated))
    
    def read(self, key: str) -> Optional[Any]:
        """Read (may return stale data)"""
        if key in self.data:
            return self.data[key].value
        return None
    
    async def sync_to_replicas(self) -> None:
        """Background sync to other replicas"""
        while self.pending_sync:
            key, value = self.pending_sync.pop(0)
            for replica in self.replicas:
                try:
                    await replica.receive_sync(key, value)
                except Exception as e:
                    # Re-queue failed sync
                    self.pending_sync.append((key, value))
                    print(f"Sync failed to {replica.node_id}: {e}")
    
    async def receive_sync(self, key: str, incoming: ReplicatedValue) -> None:
        """Receive and merge update from another replica"""
        if key in self.data:
            self.data[key] = self.data[key].merge(incoming)
        else:
            self.data[key] = incoming
        
        # Update our vector clock
        self.vector_clock = self.vector_clock.merge(incoming.vector_clock)


# Usage example
async def demo_eventual_consistency():
    # Three replicas
    replica_a = EventuallyConsistentStore("node-a")
    replica_b = EventuallyConsistentStore("node-b")
    replica_c = EventuallyConsistentStore("node-c")
    
    # Cross-reference replicas
    replica_a.replicas = [replica_b, replica_c]
    replica_b.replicas = [replica_a, replica_c]
    replica_c.replicas = [replica_a, replica_b]
    
    # Concurrent writes (simulating network partition)
    replica_a.write("user:1:name", "Alice")
    replica_b.write("user:1:name", "Alicia")  # Conflict!
    
    # Before sync: different values
    print(f"A sees: {replica_a.read('user:1:name')}")  # Alice
    print(f"B sees: {replica_b.read('user:1:name')}")  # Alicia
    
    # After sync: converges to same value
    await replica_a.sync_to_replicas()
    await replica_b.sync_to_replicas()
    
    print(f"After sync - A: {replica_a.read('user:1:name')}")
    print(f"After sync - B: {replica_b.read('user:1:name')}")
    # Both will show the same value (winner based on timestamp + node_id)

import asyncio
import random
from enum import Enum
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Callable
from datetime import datetime

class NodeState(Enum):
    FOLLOWER = "follower"
    CANDIDATE = "candidate"
    LEADER = "leader"

@dataclass
class LogEntry:
    term: int
    index: int
    command: str
    data: dict

@dataclass
class RaftNode:
    """Simplified Raft consensus node"""
    node_id: str
    state: NodeState = NodeState.FOLLOWER
    current_term: int = 0
    voted_for: Optional[str] = None
    log: List[LogEntry] = field(default_factory=list)
    commit_index: int = 0
    last_applied: int = 0
    
    # Leader state
    next_index: Dict[str, int] = field(default_factory=dict)
    match_index: Dict[str, int] = field(default_factory=dict)
    
    # Cluster info
    peers: List[str] = field(default_factory=list)
    
    # Timers
    election_timeout: float = 0
    heartbeat_interval: float = 0.15  # 150ms
    
    def __post_init__(self):
        self.reset_election_timeout()
    
    def reset_election_timeout(self):
        """Randomized election timeout (150-300ms typically)"""
        self.election_timeout = random.uniform(0.15, 0.30)
    
    async def start_election(self) -> bool:
        """Candidate starts election"""
        self.state = NodeState.CANDIDATE
        self.current_term += 1
        self.voted_for = self.node_id
        votes = 1  # Vote for self
        
        # Request votes from all peers
        vote_requests = []
        for peer_id in self.peers:
            vote_requests.append(
                self.request_vote(peer_id)
            )
        
        results = await asyncio.gather(*vote_requests, return_exceptions=True)
        votes += sum(1 for r in results if r is True)
        
        # Check if won election
        quorum = (len(self.peers) + 1) // 2 + 1
        if votes >= quorum:
            await self.become_leader()
            return True
        
        return False
    
    async def become_leader(self):
        """Transition to leader state"""
        self.state = NodeState.LEADER
        print(f"[{self.node_id}] Became leader for term {self.current_term}")
        
        # Initialize leader state
        for peer_id in self.peers:
            self.next_index[peer_id] = len(self.log) + 1
            self.match_index[peer_id] = 0
        
        # Start heartbeat
        asyncio.create_task(self.send_heartbeats())
    
    async def send_heartbeats(self):
        """Leader sends periodic heartbeats"""
        while self.state == NodeState.LEADER:
            for peer_id in self.peers:
                asyncio.create_task(self.append_entries(peer_id, []))
            await asyncio.sleep(self.heartbeat_interval)
    
    async def request_vote(self, peer_id: str) -> bool:
        """Request vote from a peer"""
        # In real implementation, this would be RPC
        last_log_index = len(self.log)
        last_log_term = self.log[-1].term if self.log else 0
        
        # Simulate RPC response
        # Peer grants vote if:
        # 1. Candidate's term >= peer's term
        # 2. Peer hasn't voted for someone else
        # 3. Candidate's log is at least as up-to-date
        return True  # Simplified
    
    async def append_entries(
        self, 
        peer_id: str, 
        entries: List[LogEntry]
    ) -> bool:
        """AppendEntries RPC (heartbeat + log replication)"""
        prev_log_index = self.next_index.get(peer_id, 1) - 1
        prev_log_term = self.log[prev_log_index - 1].term if prev_log_index > 0 else 0
        
        # Simulate RPC
        success = True  # In real implementation, check response
        
        if success and entries:
            self.match_index[peer_id] = prev_log_index + len(entries)
            self.next_index[peer_id] = self.match_index[peer_id] + 1
            
            # Update commit index
            self.update_commit_index()
        
        return success
    
    def update_commit_index(self):
        """Leader updates commit index based on replication"""
        for n in range(self.commit_index + 1, len(self.log) + 1):
            replicated_count = 1  # Count self
            for peer_id in self.peers:
                if self.match_index.get(peer_id, 0) >= n:
                    replicated_count += 1
            
            quorum = (len(self.peers) + 1) // 2 + 1
            if replicated_count >= quorum:
                if self.log[n - 1].term == self.current_term:
                    self.commit_index = n
    
    def propose(self, command: str, data: dict) -> Optional[LogEntry]:
        """Client proposes a command (only leader can accept)"""
        if self.state != NodeState.LEADER:
            return None  # Redirect to leader
        
        entry = LogEntry(
            term=self.current_term,
            index=len(self.log) + 1,
            command=command,
            data=data
        )
        self.log.append(entry)
        
        # Replicate to followers (async)
        asyncio.create_task(self.replicate_log())
        
        return entry
    
    async def replicate_log(self):
        """Replicate latest log entries to all followers"""
        for peer_id in self.peers:
            next_idx = self.next_index.get(peer_id, 1)
            entries_to_send = self.log[next_idx - 1:]
            
            if entries_to_send:
                await self.append_entries(peer_id, entries_to_send)


# Leader election simulation
async def simulate_raft():
    nodes = [
        RaftNode(node_id="node-1", peers=["node-2", "node-3"]),
        RaftNode(node_id="node-2", peers=["node-1", "node-3"]),
        RaftNode(node_id="node-3", peers=["node-1", "node-2"]),
    ]
    
    # Simulate node-1 winning election
    leader = nodes[0]
    await leader.start_election()
    
    # Propose a command
    if leader.state == NodeState.LEADER:
        entry = leader.propose("SET", {"key": "name", "value": "Alice"})
        print(f"Proposed: {entry}")

import asyncio
from enum import Enum
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any
from datetime import datetime, timedelta
import uuid

class TxState(Enum):
    PENDING = "pending"
    PREPARED = "prepared"
    COMMITTED = "committed"
    ABORTED = "aborted"

@dataclass
class Transaction:
    tx_id: str
    state: TxState = TxState.PENDING
    participants: List[str] = field(default_factory=list)
    votes: Dict[str, bool] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.now)
    timeout: timedelta = timedelta(seconds=30)

class TwoPhaseCommitCoordinator:
    """Coordinates distributed transactions using 2PC"""
    
    def __init__(self, coordinator_id: str):
        self.coordinator_id = coordinator_id
        self.transactions: Dict[str, Transaction] = {}
        self.participants: Dict[str, 'TransactionParticipant'] = {}
    
    def register_participant(self, participant: 'TransactionParticipant'):
        self.participants[participant.participant_id] = participant
    
    async def begin_transaction(
        self, 
        participant_ids: List[str],
        operations: Dict[str, Any]
    ) -> str:
        """Start a new distributed transaction"""
        tx_id = f"tx-{uuid.uuid4().hex[:8]}"
        tx = Transaction(tx_id=tx_id, participants=participant_ids)
        self.transactions[tx_id] = tx
        
        print(f"[Coordinator] Starting transaction {tx_id}")
        
        try:
            # Phase 1: Prepare
            prepared = await self.prepare_phase(tx_id, operations)
            
            if prepared:
                # Phase 2: Commit
                await self.commit_phase(tx_id)
                return tx_id
            else:
                # Phase 2: Abort
                await self.abort_phase(tx_id)
                raise Exception("Transaction aborted - not all participants prepared")
                
        except asyncio.TimeoutError:
            await self.abort_phase(tx_id)
            raise Exception("Transaction timed out")
    
    async def prepare_phase(
        self, 
        tx_id: str, 
        operations: Dict[str, Any]
    ) -> bool:
        """Phase 1: Ask all participants to prepare"""
        tx = self.transactions[tx_id]
        print(f"[Coordinator] Phase 1: PREPARE for {tx_id}")
        
        prepare_tasks = []
        for participant_id in tx.participants:
            participant = self.participants.get(participant_id)
            if participant:
                op = operations.get(participant_id, {})
                prepare_tasks.append(
                    participant.prepare(tx_id, op)
                )
        
        # Wait for all votes with timeout
        try:
            results = await asyncio.wait_for(
                asyncio.gather(*prepare_tasks, return_exceptions=True),
                timeout=tx.timeout.total_seconds()
            )
            
            # Check all votes
            for participant_id, result in zip(tx.participants, results):
                if isinstance(result, Exception):
                    tx.votes[participant_id] = False
                else:
                    tx.votes[participant_id] = result
            
            all_prepared = all(tx.votes.values())
            tx.state = TxState.PREPARED if all_prepared else TxState.PENDING
            
            return all_prepared
            
        except asyncio.TimeoutError:
            return False
    
    async def commit_phase(self, tx_id: str):
        """Phase 2: Tell all participants to commit"""
        tx = self.transactions[tx_id]
        print(f"[Coordinator] Phase 2: COMMIT for {tx_id}")
        
        commit_tasks = []
        for participant_id in tx.participants:
            participant = self.participants.get(participant_id)
            if participant:
                commit_tasks.append(participant.commit(tx_id))
        
        await asyncio.gather(*commit_tasks, return_exceptions=True)
        tx.state = TxState.COMMITTED
        print(f"[Coordinator] Transaction {tx_id} COMMITTED")
    
    async def abort_phase(self, tx_id: str):
        """Phase 2: Tell all participants to abort"""
        tx = self.transactions[tx_id]
        print(f"[Coordinator] Phase 2: ABORT for {tx_id}")
        
        abort_tasks = []
        for participant_id in tx.participants:
            participant = self.participants.get(participant_id)
            if participant:
                abort_tasks.append(participant.abort(tx_id))
        
        await asyncio.gather(*abort_tasks, return_exceptions=True)
        tx.state = TxState.ABORTED
        print(f"[Coordinator] Transaction {tx_id} ABORTED")


class TransactionParticipant:
    """Participant in a distributed transaction"""
    
    def __init__(self, participant_id: str):
        self.participant_id = participant_id
        self.prepared_data: Dict[str, Any] = {}
        self.committed_data: Dict[str, Any] = {}
        self.undo_log: Dict[str, Any] = {}
    
    async def prepare(self, tx_id: str, operation: Dict) -> bool:
        """Prepare to commit - acquire locks, validate, write to redo log"""
        print(f"[{self.participant_id}] Preparing {tx_id}: {operation}")
        
        try:
            # Simulate validation and lock acquisition
            await asyncio.sleep(0.1)  # Simulate I/O
            
            # Write to redo log (would be durable in real implementation)
            self.prepared_data[tx_id] = operation
            
            # Save undo information
            key = operation.get("key")
            if key:
                self.undo_log[tx_id] = {
                    "key": key,
                    "old_value": self.committed_data.get(key)
                }
            
            print(f"[{self.participant_id}] Voted YES for {tx_id}")
            return True
            
        except Exception as e:
            print(f"[{self.participant_id}] Voted NO for {tx_id}: {e}")
            return False
    
    async def commit(self, tx_id: str):
        """Commit the prepared transaction"""
        print(f"[{self.participant_id}] Committing {tx_id}")
        
        operation = self.prepared_data.pop(tx_id, {})
        
        # Apply the operation
        key = operation.get("key")
        value = operation.get("value")
        if key:
            self.committed_data[key] = value
        
        # Clear undo log
        self.undo_log.pop(tx_id, None)
    
    async def abort(self, tx_id: str):
        """Abort - rollback using undo log"""
        print(f"[{self.participant_id}] Aborting {tx_id}")
        
        # Rollback using undo log
        undo = self.undo_log.pop(tx_id, None)
        if undo:
            key = undo["key"]
            old_value = undo["old_value"]
            if old_value is not None:
                self.committed_data[key] = old_value
            elif key in self.committed_data:
                del self.committed_data[key]
        
        self.prepared_data.pop(tx_id, None)


# Usage example
async def demo_2pc():
    coordinator = TwoPhaseCommitCoordinator("coordinator-1")
    
    # Register participants (e.g., database shards)
    orders_db = TransactionParticipant("orders-db")
    inventory_db = TransactionParticipant("inventory-db")
    
    coordinator.register_participant(orders_db)
    coordinator.register_participant(inventory_db)
    
    # Execute distributed transaction
    try:
        tx_id = await coordinator.begin_transaction(
            participant_ids=["orders-db", "inventory-db"],
            operations={
                "orders-db": {"key": "order:123", "value": {"product": "laptop", "qty": 1}},
                "inventory-db": {"key": "product:laptop:stock", "value": 99}
            }
        )
        print(f"Transaction {tx_id} completed successfully!")
        
    except Exception as e:
        print(f"Transaction failed: {e}")

asyncio.run(demo_2pc())