AWS S3.

# storage/erasure_coding.py

"""
Erasure Coding for 11 Nines Durability

Uses Reed-Solomon coding to achieve extreme durability
with only 1.6x storage overhead (vs 3x for replication).
"""

from dataclasses import dataclass
from typing import List, Optional, Tuple
import hashlib


@dataclass
class Shard:
    """A single shard of erasure-coded data."""
    shard_id: int
    data: bytes
    checksum: str
    storage_node: str
    availability_zone: str


@dataclass
class ErasureConfig:
    """Configuration for erasure coding."""
    data_shards: int = 10      # K: minimum shards to reconstruct
    parity_shards: int = 6     # M: redundant shards
    # Total = 16 shards, can lose up to 6


class ErasureCodingService:
    """
    Erasure coding service for S3 durability.
    
    Uses Reed-Solomon (10, 6) coding:
    - Split data into 10 data shards
    - Compute 6 parity shards
    - Can reconstruct from any 10 of 16 shards
    
    Applies concepts:
    - Week 1: Replication and partitioning
    - Week 1: Data durability strategies
    """
    
    def __init__(self, config: ErasureConfig, storage_nodes: List):
        self.config = config
        self.storage_nodes = storage_nodes
        self.total_shards = config.data_shards + config.parity_shards
    
    async def encode_and_store(
        self, 
        object_key: str, 
        data: bytes
    ) -> List[Shard]:
        """
        Encode data and distribute shards across AZs.
        
        Steps:
        1. Calculate checksum of original data
        2. Split into data shards
        3. Compute parity shards
        4. Distribute across AZs (ensuring AZ diversity)
        5. Verify (bracketing) before confirming
        """
        # Step 1: Original checksum
        original_checksum = self._compute_checksum(data)
        
        # Step 2: Split into data shards
        data_shards = self._split_data(data, self.config.data_shards)
        
        # Step 3: Compute parity shards
        parity_shards = self._compute_parity(data_shards)
        
        # Step 4: Select storage nodes (AZ-aware)
        all_shards = data_shards + parity_shards
        shard_placements = self._select_nodes(len(all_shards))
        
        # Step 5: Write shards to storage
        stored_shards = []
        for i, (shard_data, (node, az)) in enumerate(
            zip(all_shards, shard_placements)
        ):
            checksum = self._compute_checksum(shard_data)
            shard = Shard(
                shard_id=i,
                data=shard_data,
                checksum=checksum,
                storage_node=node,
                availability_zone=az
            )
            await self._write_shard(shard)
            stored_shards.append(shard)
        
        # Step 6: BRACKETING - verify we can reconstruct
        if not await self._verify_reconstruction(
            stored_shards, original_checksum
        ):
            raise DurabilityError("Bracketing verification failed")
        
        return stored_shards
    
    async def read_and_decode(
        self, 
        shards: List[Shard]
    ) -> bytes:
        """
        Read shards and reconstruct original data.
        
        Can reconstruct from any K (10) of N (16) shards.
        """
        # Read available shards
        available = []
        for shard in shards:
            try:
                data = await self._read_shard(shard)
                if self._verify_checksum(data, shard.checksum):
                    available.append((shard.shard_id, data))
            except Exception:
                continue  # Shard unavailable
        
        # Need at least K shards
        if len(available) < self.config.data_shards:
            raise DurabilityError(
                f"Only {len(available)} shards available, "
                f"need {self.config.data_shards}"
            )
        
        # Reconstruct using Reed-Solomon
        return self._reconstruct(available)
    
    def _select_nodes(self, count: int) -> List[Tuple[str, str]]:
        """
        Select storage nodes ensuring AZ diversity.
        
        Critical: No more than ceil(count/3) shards per AZ
        This ensures we can lose an entire AZ and still reconstruct.
        """
        placements = []
        az_counts = {}
        max_per_az = (count + 2) // 3  # Ceiling division
        
        for node in self.storage_nodes:
            az = node.availability_zone
            if az_counts.get(az, 0) < max_per_az:
                placements.append((node.id, az))
                az_counts[az] = az_counts.get(az, 0) + 1
            if len(placements) == count:
                break
        
        return placements
    
    async def _verify_reconstruction(
        self, 
        shards: List[Shard], 
        expected_checksum: str
    ) -> bool:
        """
        Bracketing: Verify data can be reconstructed.
        
        This is critical - we don't confirm write success
        until we've proven the data is recoverable.
        """
        try:
            reconstructed = await self.read_and_decode(shards)
            actual_checksum = self._compute_checksum(reconstructed)
            return actual_checksum == expected_checksum
        except Exception:
            return False
    
    def _compute_checksum(self, data: bytes) -> str:
        return hashlib.sha256(data).hexdigest()

# metadata/consistency.py

"""
Strong Consistency for S3 Metadata

Ensures read-after-write consistency for all operations.
"""

from dataclasses import dataclass
from typing import Optional, List
from datetime import datetime
import asyncio


@dataclass
class MetadataEntry:
    """Object metadata entry."""
    bucket: str
    key: str
    version_id: str
    etag: str
    size: int
    storage_class: str
    last_modified: datetime
    shard_locations: List[str]


class StrongConsistencyMetadataStore:
    """
    Metadata store with strong read-after-write consistency.
    
    Key mechanism: synchronous replication to a majority before
    acknowledging the write, and majority reads (R + W > N) so a
    read always overlaps the latest write.
    
    Applies concepts:
    - Week 5: Strong vs eventual consistency
    - Week 5: Consensus protocols
    """
    
    def __init__(self, replicas: List, quorum_size: int, node_id: str):
        self.replicas = replicas
        self.quorum_size = quorum_size  # Majority
        self.node_id = node_id  # this (primary) metadata node
    
    async def put_metadata(
        self, 
        entry: MetadataEntry
    ) -> bool:
        """
        Store metadata with strong consistency.
        
        Only returns success after majority of replicas
        have durably stored the metadata.
        """
        # Generate monotonic version
        entry.version_id = self._generate_version_id()
        
        # Write to all replicas in parallel
        write_tasks = [
            self._write_to_replica(replica, entry)
            for replica in self.replicas
        ]
        
        # Wait for quorum (majority)
        results = await asyncio.gather(
            *write_tasks, 
            return_exceptions=True
        )
        
        successes = sum(1 for r in results if r is True)
        
        if successes >= self.quorum_size:
            # Quorum achieved - write is durable and consistent
            return True
        else:
            # Rollback on failure
            await self._rollback(entry)
            return False
    
    async def get_metadata(
        self, 
        bucket: str, 
        key: str
    ) -> Optional[MetadataEntry]:
        """
        Read metadata - guaranteed to return latest version.
        
        A majority write can leave a minority of replicas behind,
        so a single "nearest" replica could be stale. Read a majority
        and keep the newest version (any two majorities overlap).
        """
        results = await asyncio.gather(
            *(self._read_from_replica(r, bucket, key) for r in self.replicas),
            return_exceptions=True
        )
        responses = [r for r in results if not isinstance(r, Exception)]
        if len(responses) < self.quorum_size:
            raise QuorumUnavailableError(bucket, key)
        
        entries = [e for e in responses if e is not None]
        return max(entries, key=lambda e: e.version_id, default=None)
    
    async def list_objects(
        self, 
        bucket: str, 
        prefix: str
    ) -> List[MetadataEntry]:
        """
        List objects - strongly consistent.
        
        After a PUT, the object is immediately visible in LIST.
        """
        # Same rule as get_metadata: use a replica confirmed current
        # (via the witness / a quorum check), not just the nearest one
        replica = await self._select_current_replica(bucket, prefix)
        
        # Read from partition responsible for this prefix
        partition = self._get_partition(bucket, prefix)
        
        return await self._list_from_partition(replica, partition, prefix)
    
    def _generate_version_id(self) -> str:
        """Generate monotonically increasing version ID."""
        # Zero-padded timestamp + per-node sequence: sorts in write order
        # (a random uuid4 would not, so max() could pick an old write)
        import time
        self._seq = getattr(self, "_seq", 0) + 1
        return f"{time.time_ns():020d}-{self.node_id}-{self._seq:08d}"

# index/partition_manager.py

"""
Index Partitioning for S3

Dynamically partitions index to handle 500+ trillion objects.
"""

from dataclasses import dataclass
from typing import List, Optional, Dict
import hashlib


@dataclass
class Partition:
    """An index partition."""
    partition_id: str
    prefix_range: tuple  # (start_prefix, end_prefix)
    request_rate: float
    node_assignment: str


class PartitionManager:
    """
    Manages index partitions for S3.
    
    Key features:
    - Dynamic splitting based on load
    - Prefix-based range partitioning
    - Automatic rebalancing
    
    Applies concepts:
    - Week 1: Hash vs range partitioning
    - Week 1: Hot key handling
    """
    
    # Thresholds for partition operations
    SPLIT_THRESHOLD_RPS = 5500  # GET requests/second
    MIN_PARTITION_SIZE = 1000   # Objects
    
    def __init__(self, partition_store):
        self.partition_store = partition_store
    
    def get_partition_for_key(
        self, 
        bucket: str, 
        key: str
    ) -> Partition:
        """
        Find the partition responsible for a key.
        
        Uses hierarchical lookup:
        1. Hash bucket to bucket partition
        2. Binary search within bucket for key prefix
        """
        # Level 1: Bucket partition
        bucket_hash = self._hash_bucket(bucket)
        bucket_partition = self._get_bucket_partition(bucket_hash)
        
        # Level 2: Key prefix partition
        key_partition = self._find_key_partition(bucket_partition, key)
        
        return key_partition
    
    async def check_and_split(
        self, 
        partition: Partition
    ) -> List[Partition]:
        """
        Check if partition needs splitting due to high load.
        
        Triggered when request rate exceeds threshold.
        """
        if partition.request_rate < self.SPLIT_THRESHOLD_RPS:
            return [partition]  # No split needed
        
        # Split partition in half by prefix
        start, end = partition.prefix_range
        mid = self._find_split_point(partition)
        
        new_partitions = [
            Partition(
                partition_id=f"{partition.partition_id}-a",
                prefix_range=(start, mid),
                request_rate=partition.request_rate / 2,
                node_assignment=self._select_node()
            ),
            Partition(
                partition_id=f"{partition.partition_id}-b",
                prefix_range=(mid, end),
                request_rate=partition.request_rate / 2,
                node_assignment=self._select_node()
            )
        ]
        
        # Atomic partition update
        await self.partition_store.atomic_split(partition, new_partitions)
        
        return new_partitions
    
    def _hash_bucket(self, bucket: str) -> int:
        """Hash bucket name to partition number."""
        return int(hashlib.md5(bucket.encode()).hexdigest(), 16) % 10000
    
    def _find_split_point(self, partition: Partition) -> str:
        """
        Find optimal split point for partition.
        
        Analyzes key distribution to split evenly.
        """
        # Sample keys in partition
        keys = self.partition_store.sample_keys(partition, n=1000)
        
        # Find median prefix
        keys.sort()
        median_key = keys[len(keys) // 2]
        
        # Extract common prefix
        return self._extract_prefix(median_key)

# lifecycle/policy_engine.py

"""
S3 Lifecycle Policy Engine

Automatically transitions objects between storage classes
based on age and access patterns.
"""

from dataclasses import dataclass
from typing import List, Optional
from datetime import datetime, timedelta
from enum import Enum


class StorageClass(Enum):
    STANDARD = "STANDARD"
    STANDARD_IA = "STANDARD_IA"
    INTELLIGENT_TIERING = "INTELLIGENT_TIERING"
    GLACIER_INSTANT = "GLACIER_IR"
    GLACIER_FLEXIBLE = "GLACIER"
    GLACIER_DEEP = "DEEP_ARCHIVE"


@dataclass
class LifecycleRule:
    """A lifecycle transition rule."""
    rule_id: str
    prefix_filter: Optional[str]
    tag_filter: Optional[dict]
    transitions: List[dict]  # [{days: 30, storage_class: "STANDARD_IA"}]
    expiration_days: Optional[int]


class LifecyclePolicyEngine:
    """
    Evaluates and applies lifecycle policies to objects.
    
    Common pattern:
    1. Standard for 30 days (frequent access)
    2. Standard-IA for 60 days (occasional access)
    3. Glacier for 1 year (rare access)
    4. Delete after 7 years (compliance)
    
    Applies concepts:
    - Week 4: Cache tiering strategies
    - Week 8: Data lifecycle management
    """
    
    def __init__(self, metadata_store, storage_service):
        self.metadata = metadata_store
        self.storage = storage_service
    
    async def evaluate_object(
        self, 
        bucket: str, 
        key: str, 
        rules: List[LifecycleRule]
    ) -> Optional[StorageClass]:
        """
        Determine if object should transition.
        """
        obj_metadata = await self.metadata.get(bucket, key)
        
        for rule in rules:
            if not self._matches_filter(obj_metadata, rule):
                continue
            
            age_days = (datetime.utcnow() - obj_metadata.last_modified).days
            
            # Check expiration
            if rule.expiration_days and age_days >= rule.expiration_days:
                return "DELETE"
            
            # Check transitions (in reverse order - most aggressive first)
            for transition in sorted(
                rule.transitions, 
                key=lambda t: t['days'], 
                reverse=True
            ):
                if age_days >= transition['days']:
                    target_class = StorageClass(transition['storage_class'])
                    if target_class != obj_metadata.storage_class:
                        return target_class
        
        return None  # No transition needed
    
    async def apply_transition(
        self, 
        bucket: str, 
        key: str,
        target_class: StorageClass
    ) -> bool:
        """
        Transition object to new storage class.
        
        For Glacier classes, this involves:
        1. Reading object from current storage
        2. Re-encoding for archive (more aggressive erasure coding)
        3. Writing to archive storage
        4. Updating metadata
        """
        if target_class == "DELETE":
            return await self._delete_object(bucket, key)
        
        # Get current object location
        current_shards = await self.metadata.get_shard_locations(bucket, key)
        
        # Transition to new storage class
        new_shards = await self.storage.transition(
            current_shards, 
            target_class
        )
        
        # Update metadata atomically
        await self.metadata.update_storage_class(
            bucket, key, target_class, new_shards
        )
        
        return True