two_phase_commit.py

"""
Two-Phase Commit Coordinator for Distributed Transactions

This module implements a two-phase commit protocol coordinator that manages
distributed transactions across multiple participants. It handles the prepare
and commit/abort phases while managing timeouts and failures gracefully.
"""

import asyncio
import logging
from enum import Enum
from typing import Dict, List, Optional, Set, Tuple, Union
from dataclasses import dataclass
from abc import ABC, abstractmethod
import time
import uuid

# Configure logging
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


class TransactionState(Enum):
    """Enumeration of transaction states"""
    PREPARING = "preparing"
    PREPARED = "prepared"
    COMMITTING = "committing"
    COMMITTED = "committed"
    ABORTING = "aborting"
    ABORTED = "aborted"
    FAILED = "failed"


class ParticipantResponse(Enum):
    """Enumeration of participant responses"""
    PREPARED = "prepared"
    ABORTED = "aborted"
    COMMITTED = "committed"
    FAILED = "failed"


@dataclass
class TransactionContext:
    """Context for a distributed transaction"""
    transaction_id: str
    participants: List[str]
    created_at: float
    timeout: float = 30.0  # Default timeout in seconds


class ParticipantInterface(ABC):
    """Abstract interface for transaction participants"""
    
    @abstractmethod
    async def prepare(self, transaction_id: str) -> ParticipantResponse:
        """Prepare phase - check if participant can commit"""
        pass
    
    @abstractmethod
    async def commit(self, transaction_id: str) -> ParticipantResponse:
        """Commit phase - actually commit the transaction"""
        pass
    
    @abstractmethod
    async def abort(self, transaction_id: str) -> ParticipantResponse:
        """Abort phase - rollback the transaction"""
        pass


class NetworkParticipant(ParticipantInterface):
    """
    Network-based participant implementation
    
    This would typically communicate over HTTP/gRPC/other protocols
    """
    
    def __init__(self, participant_id: str, endpoint: str):
        self.participant_id = participant_id
        self.endpoint = endpoint
    
    async def prepare(self, transaction_id: str) -> ParticipantResponse:
        """Simulate network prepare call"""
        try:
            # In a real implementation, this would make an HTTP/gRPC call
            logger.info(f"Preparing participant {self.participant_id} for transaction {transaction_id}")
            # Simulate network delay
            await asyncio.sleep(0.1)
            return ParticipantResponse.PREPARED
        except Exception as e:
            logger.error(f"Prepare failed for participant {self.participant_id}: {e}")
            return ParticipantResponse.FAILED
    
    async def commit(self, transaction_id: str) -> ParticipantResponse:
        """Simulate network commit call"""
        try:
            logger.info(f"Committing participant {self.participant_id} for transaction {transaction_id}")
            await asyncio.sleep(0.1)
            return ParticipantResponse.COMMITTED
        except Exception as e:
            logger.error(f"Commit failed for participant {self.participant_id}: {e}")
            return ParticipantResponse.FAILED
    
    async def abort(self, transaction_id: str) -> ParticipantResponse:
        """Simulate network abort call"""
        try:
            logger.info(f"Aborting participant {self.participant_id} for transaction {transaction_id}")
            await asyncio.sleep(0.1)
            return ParticipantResponse.ABORTED
        except Exception as e:
            logger.error(f"Abort failed for participant {self.participant_id}: {e}")
            return ParticipantResponse.FAILED


class TwoPhaseCommitCoordinator:
    """
    Two-Phase Commit Coordinator
    
    Manages distributed transactions across multiple participants using
    the two-phase commit protocol. Handles timeouts, failures, and
    maintains transaction state.
    """
    
    def __init__(self, default_timeout: float = 30.0):
        self.default_timeout = default_timeout
        self.transactions: Dict[str, TransactionContext] = {}
        self.transaction_states: Dict[str, TransactionState] = {}
        self.participants: Dict[str, ParticipantInterface] = {}
        self._lock = asyncio.Lock()
    
    async def register_participant(self, participant: ParticipantInterface) -> None:
        """
        Register a participant with the coordinator
        
        Args:
            participant: Participant to register
        """
        async with self._lock:
            participant_id = getattr(participant, 'participant_id', str(uuid.uuid4()))
            self.participants[participant_id] = participant
            logger.info(f"Registered participant: {participant_id}")
    
    async def begin_transaction(
        self, 
        participants: List[str], 
        timeout: Optional[float] = None
    ) -> str:
        """
        Begin a new distributed transaction
        
        Args:
            participants: List of participant IDs
            timeout: Transaction timeout in seconds
            
        Returns:
            Transaction ID
            
        Raises:
            ValueError: If participants list is empty or contains invalid participants
        """
        if not participants:
            raise ValueError("Participants list cannot be empty")
        
        # Validate participants exist
        missing_participants = set(participants) - set(self.participants.keys())
        if missing_participants:
            raise ValueError(f"Unknown participants: {missing_participants}")
        
        transaction_id = str(uuid.uuid4())
        timeout = timeout or self.default_timeout
        
        async with self._lock:
            self.transactions[transaction_id] = TransactionContext(
                transaction_id=transaction_id,
                participants=participants,
                created_at=time.time(),
                timeout=timeout
            )
            self.transaction_states[transaction_id] = TransactionState.PREPARING
        
        logger.info(f"Started transaction {transaction_id} with participants {participants}")
        return transaction_id
    
    async def prepare_transaction(self, transaction_id: str) -> bool:
        """
        Execute the prepare phase of the two-phase commit protocol
        
        Args:
            transaction_id: ID of the transaction to prepare
            
        Returns:
            True if all participants prepared successfully, False otherwise
        """
        async with self._lock:
            if transaction_id not in self.transactions:
                raise ValueError(f"Unknown transaction: {transaction_id}")
            
            if self.transaction_states[transaction_id] != TransactionState.PREPARING:
                raise RuntimeError(f"Transaction {transaction_id} is not in PREPARING state")
        
        context = self.transactions[transaction_id]
        participants = context.participants
        
        # Check for timeout
        if time.time() - context.created_at > context.timeout:
            await self._abort_transaction(transaction_id)
            return False
        
        # Prepare all participants concurrently
        prepare_tasks = [
            self._prepare_participant(transaction_id, participant_id)
            for participant_id in participants
        ]
        
        try:
            results = await asyncio.wait_for(
                asyncio.gather(*prepare_tasks, return_exceptions=True),
                timeout=context.timeout
            )
        except asyncio.TimeoutError:
            logger.error(f"Prepare phase timeout for transaction {transaction_id}")
            await self._abort_transaction(transaction_id)
            return False
        
        # Check results
        prepared_count = 0
        failed_count = 0
        
        for i, result in enumerate(results):
            participant_id = participants[i]
            if isinstance(result, Exception):
                logger.error(f"Prepare failed for participant {participant_id}: {result}")
                failed_count += 1
            elif result == ParticipantResponse.PREPARED:
                prepared_count += 1
            elif result == ParticipantResponse.FAILED:
                failed_count += 1
            else:
                logger.warning(f"Unexpected prepare response from {participant_id}: {result}")
                failed_count += 1
        
        # If any participant failed to prepare, abort the transaction
        if failed_count > 0 or prepared_count != len(participants):
            logger.info(f"Prepare phase failed for transaction {transaction_id}. Aborting.")
            await self._abort_transaction(transaction_id)
            return False
        
        # All participants prepared successfully
        async with self._lock:
            self.transaction_states[transaction_id] = TransactionState.PREPARED
        
        logger.info(f"Prepare phase completed successfully for transaction {transaction_id}")
        return True
    
    async def commit_transaction(self, transaction_id: str) -> bool:
        """
        Execute the commit phase of the two-phase commit protocol
        
        Args:
            transaction_id: ID of the transaction to commit
            
        Returns:
            True if transaction was committed successfully, False otherwise
        """
        async with self._lock:
            if transaction_id not in self.transactions:
                raise ValueError(f"Unknown transaction: {transaction_id}")
            
            if self.transaction_states[transaction_id] != TransactionState.PREPARED:
                raise RuntimeError(f"Transaction {transaction_id} is not in PREPARED state")
            
            self.transaction_states[transaction_id] = TransactionState.COMMITTING
        
        context = self.transactions[transaction_id]
        participants = context.participants
        
        # Commit all participants concurrently
        commit_tasks = [
            self._commit_participant(transaction_id, participant_id)
            for participant_id in participants
        ]
        
        try:
            results = await asyncio.wait_for(
                asyncio.gather(*commit_tasks, return_exceptions=True),
                timeout=context.timeout
            )
        except asyncio.TimeoutError:
            logger.error(f"Commit phase timeout for transaction {transaction_id}")
            # Even on timeout, we consider the transaction committed
            # In a real system, you'd want to handle this more carefully
            async with self._lock:
                self.transaction_states[transaction_id] = TransactionState.FAILED
            return False
        
        # Check results
        committed_count = 0
        failed_count = 0
        
        for i, result in enumerate(results):
            participant_id = participants[i]
            if isinstance(result, Exception):
                logger.error(f"Commit failed for participant {participant_id}: {result}")
                failed_count += 1
            elif result == ParticipantResponse.COMMITTED:
                committed_count += 1
            elif result == ParticipantResponse.FAILED:
                failed_count += 1
            else:
                logger.warning(f"Unexpected commit response from {participant_id}: {result}")
                failed_count += 1
        
        # Update transaction state
        async with self._lock:
            if failed_count > 0:
                self.transaction_states[transaction_id] = TransactionState.FAILED
                return False
            else:
                self.transaction_states[transaction_id] = TransactionState.COMMITTED
        
        logger.info(f"Commit phase completed for transaction {transaction_id}")
        return True
    
    async def abort_transaction(self, transaction_id: str) -> bool:
        """
        Abort a transaction
        
        Args:
            transaction_id: ID of the transaction to abort
            
        Returns:
            True if transaction was aborted successfully, False otherwise
        """
        return await self._abort_transaction(transaction_id)
    
    async def get_transaction_status(self, transaction_id: str) -> Optional[TransactionState]:
        """
        Get the current status of a transaction
        
        Args:
            transaction_id: ID of the transaction
            
        Returns:
            Current transaction state or None if transaction doesn't exist
        """
        async with self._lock:
            return self.transaction_states.get(transaction_id)
    
    async def _prepare_participant(
        self, 
        transaction_id: str, 
        participant_id: str
    ) -> ParticipantResponse:
        """Prepare a single participant"""
        try:
            participant = self.participants[participant_id]
            return await participant.prepare(transaction_id)
        except Exception as e:
            logger.error(f"Error preparing participant {participant_id}: {e}")
            return ParticipantResponse.FAILED
    
    async def _commit_participant(
        self, 
        transaction_id: str, 
        participant_id: str
    ) -> ParticipantResponse:
        """Commit a single participant"""
        try:
            participant = self.participants[participant_id]
            return await participant.commit(transaction_id)
        except Exception as e:
            logger.error(f"Error committing participant {participant_id}: {e}")
            return ParticipantResponse.FAILED
    
    async def _abort_transaction(self, transaction_id: str) -> bool:
        """Abort a transaction and notify all participants"""
        async with self._lock:
            if transaction_id not in self.transactions:
                return False
            
            self.transaction_states[transaction_id] = TransactionState.ABORTING
        
        context = self.transactions[transaction_id]
        participants = context.participants
        
        # Abort all participants concurrently
        abort_tasks = [
            self._abort_participant(transaction_id, participant_id)
            for participant_id in participants
        ]
        
        try:
            results = await asyncio.gather(*abort_tasks, return_exceptions=True)
        except Exception as e:
            logger.error(f"Error during abort phase for transaction {transaction_id}: {e}")
            async with self._lock:
                self.transaction_states[transaction_id] = TransactionState.FAILED
            return False
        
        # Update transaction state
        async with self._lock:
            self.transaction_states[transaction_id] = TransactionState.ABORTED
        
        logger.info(f"Transaction {transaction_id} aborted")
        return True
    
    async def _abort_participant(
        self, 
        transaction_id: str, 
        participant_id: str
    ) -> ParticipantResponse:
        """Abort a single participant"""
        try:
            participant = self.participants[participant_id]
            return await participant.abort(transaction_id)
        except Exception as e:
            logger.error(f"Error aborting participant {participant_id}: {e}")
            return ParticipantResponse.FAILED


# Example usage and testing
async def main():
    """Example usage of the Two-Phase Commit Coordinator"""
    
    # Create coordinator
    coordinator = TwoPhaseCommitCoordinator(default_timeout=10.0)
    
    # Create participants
    participant1 = NetworkParticipant("participant_1", "http://localhost:8001")
    participant2 = NetworkParticipant("participant_2", "http://localhost:8002")
    
    # Register participants
    await coordinator.register_participant(participant1)
    await coordinator.register_participant(participant2)
    
    try:
        # Begin transaction
        transaction_id = await coordinator.begin_transaction(
            participants=["participant_1", "participant_2"],
            timeout=5.0
        )
        
        # Prepare phase
        prepare_success = await coordinator.prepare_transaction(transaction_id)
        
        if prepare_success:
            # Commit phase
            commit_success = await coordinator.commit_transaction(transaction_id)
            if commit_success:
                print(f"Transaction {transaction_id} committed successfully")
            else:
                print(f"Transaction {transaction_id} commit failed")
        else:
            print(f"Transaction {transaction_id} prepare failed")
        
        # Check final status
        status = await coordinator.get_transaction_status(transaction_id)
        print(f"Final transaction status: {status}")
        
    except Exception as e:
        logger.error(f"Transaction failed: {e}")


if __name__ == "__main__":
    asyncio.run(main())