LLD Case Study: Thread-Safe LRU Cache

from __future__ import annotations
import threading
from typing import Optional, TypeVar, Generic, Callable
from datetime import datetime, timedelta

K = TypeVar("K")
V = TypeVar("V")


class Node(Generic[K, V]):
    __slots__ = ("key", "value", "prev", "next", "expires_at")

    def __init__(self, key: K, value: V, ttl_seconds: Optional[float] = None):
        self.key = key
        self.value = value
        self.prev: Optional[Node[K, V]] = None
        self.next: Optional[Node[K, V]] = None
        self.expires_at: Optional[datetime] = (
            datetime.utcnow() + timedelta(seconds=ttl_seconds)
            if ttl_seconds else None
        )

    def is_expired(self) -> bool:
        return self.expires_at is not None and datetime.utcnow() > self.expires_at


class LRUCache(Generic[K, V]):
    """
    Thread-safe LRU cache with optional TTL and loader support.

    Invariants:
      1. len(map) == node count between sentinels
      2. head.next is MRU; tail.prev is LRU
      3. len(map) <= capacity after every put()
      4. N.prev.next == N and N.next.prev == N for all nodes
    """

    def __init__(
        self,
        capacity: int,
        loader: Optional[Callable[[K], V]] = None,
        ttl_seconds: Optional[float] = None,
    ):
        if capacity < 1:
            raise ValueError("Capacity must be at least 1")
        self._capacity = capacity
        self._loader = loader
        self._ttl = ttl_seconds
        self._map: dict[K, Node[K, V]] = {}
        self._lock = threading.Lock()

        # Sentinel nodes eliminate null-check edge cases
        self._head: Node = Node.__new__(Node)  # type: ignore[assignment]
        self._tail: Node = Node.__new__(Node)  # type: ignore[assignment]
        self._head.prev = None
        self._head.next = self._tail
        self._tail.prev = self._head
        self._tail.next = None

    # ── Public API ──────────────────────────────────────────────────────────

    def get(self, key: K) -> Optional[V]:
        with self._lock:
            node = self._map.get(key)
            if node is None:
                return self._load(key)   # cache miss
            if node.is_expired():
                self._remove_node(node)
                del self._map[key]
                return self._load(key)
            self._move_to_head(node)
            return node.value

    def put(self, key: K, value: V) -> None:
        with self._lock:
            node = self._map.get(key)
            if node:
                node.value = value
                if self._ttl:
                    node.expires_at = datetime.utcnow() + timedelta(seconds=self._ttl)
                self._move_to_head(node)
            else:
                new_node = Node(key, value, self._ttl)
                self._map[key] = new_node
                self._insert_at_head(new_node)
                if len(self._map) > self._capacity:
                    evicted = self._evict_tail()
                    del self._map[evicted.key]

    def delete(self, key: K) -> bool:
        with self._lock:
            node = self._map.pop(key, None)
            if node:
                self._remove_node(node)
                return True
            return False

    def __len__(self) -> int:
        with self._lock:
            return len(self._map)

    # ── Private helpers (must always be called with _lock held) ─────────────

    def _insert_at_head(self, node: Node) -> None:
        """O(1): insert node immediately after sentinel head."""
        node.prev = self._head
        node.next = self._head.next
        self._head.next.prev = node   # invariant 4
        self._head.next = node

    def _remove_node(self, node: Node) -> None:
        """O(1): unlink node from list. Caller must delete from map."""
        node.prev.next = node.next
        node.next.prev = node.prev

    def _move_to_head(self, node: Node) -> None:
        """O(1): move existing node to MRU position."""
        self._remove_node(node)
        self._insert_at_head(node)

    def _evict_tail(self) -> Node:
        """O(1): remove and return the LRU node (tail.prev)."""
        lru = self._tail.prev
        self._remove_node(lru)
        return lru

    def _load(self, key: K) -> Optional[V]:
        """Call loader if configured. Lock is held by caller."""
        if self._loader is None:
            return None
        value = self._loader(key)   # NOTE: lock is held during load — use singleflight in prod
        if value is not None:
            new_node = Node(key, value, self._ttl)
            self._map[key] = new_node
            self._insert_at_head(new_node)
            if len(self._map) > self._capacity:
                evicted = self._evict_tail()
                del self._map[evicted.key]
        return value


# ── Segmented LRU — eliminates single-lock bottleneck ──────────────────────

class SegmentedLRUCache(Generic[K, V]):
    """
    Partitions keys across N independent LRU segments.
    Concurrent gets/puts on different segments never block each other.
    """

    def __init__(self, capacity: int, num_segments: int = 16,
                 ttl_seconds: Optional[float] = None):
        segment_capacity = max(1, capacity // num_segments)
        self._segments = [
            LRUCache(segment_capacity, ttl_seconds=ttl_seconds)
            for _ in range(num_segments)
        ]
        self._num_segments = num_segments

    def _segment(self, key: K) -> LRUCache[K, V]:
        return self._segments[hash(key) % self._num_segments]

    def get(self, key: K) -> Optional[V]:
        return self._segment(key).get(key)

    def put(self, key: K, value: V) -> None:
        self._segment(key).put(key, value)

    def delete(self, key: K) -> bool:
        return self._segment(key).delete(key)