counters.py

"""
Scalable Like Counter System
=============================
Features:
  - Sharded counters with per-shard locks (no global bottleneck)
  - O(1) reads via running_shard_total
  - Dynamic shard scaling (16 → 32 → 128 → 1024)
  - Worker pool for parallel event processing
  - Duplicate-like prevention (in-memory; swap for Redis in prod)
  - Periodic aggregator flush to simulated DB
  - Shard monitor that auto-scales on queue depth
"""

import random
import threading
import time
from queue import Queue


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

INITIAL_SHARDS  = 16
NUM_WORKERS     = 4       # parallel event-processor threads
FLUSH_INTERVAL  = 10      # seconds between aggregator flushes
MONITOR_INTERVAL = 5      # seconds between shard-monitor checks

# Queue-depth thresholds → target shard count
SCALE_THRESHOLDS = [
    (500,    32),
    (2_000,  128),
    (10_000, 1_024),
]


# ---------------------------------------------------------------------------
# DynamicShardedCounter
# ---------------------------------------------------------------------------

class DynamicShardedCounter:
    """
    Thread-safe sharded counter with:
      - Per-shard locks  → parallel increments with no global contention
      - running_total    → O(1) reads at all times
      - rescale()        → zero-loss resharding under a single resize lock
    """

    def __init__(self, initial_shards: int = INITIAL_SHARDS):
        self._num_shards  = initial_shards
        self._shards      = [0] * initial_shards
        self._shard_locks = [threading.Lock() for _ in range(initial_shards)]

        self._running_total = 0
        self._total_lock    = threading.Lock()

        # RLock so rescale() can call increment() internally if needed
        self._resize_lock = threading.RLock()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def increment(self, user_id: int) -> None:
        """Increment the counter for a given user (write path)."""
        with self._resize_lock:                       # blocks during resize
            shard = hash(user_id) % self._num_shards
            with self._shard_locks[shard]:
                self._shards[shard] += 1
            with self._total_lock:
                self._running_total += 1              # O(1) running total

    def read(self) -> int:
        """Return current like count — O(1)."""
        with self._total_lock:
            return self._running_total

    def flush(self) -> int:
        """
        Drain all shard values and reset running_total.
        Returns the batch count to be committed to the DB.
        Called by the aggregator every FLUSH_INTERVAL seconds.
        """
        with self._resize_lock:
            batch = 0
            for i in range(self._num_shards):
                with self._shard_locks[i]:
                    batch += self._shards[i]
                    self._shards[i] = 0
            with self._total_lock:
                self._running_total = 0
            return batch

    def rescale(self, new_shard_count: int) -> None:
        """
        Resize shard array without losing unflushed counts.

        Strategy:
          1. Acquire resize_lock exclusively (pauses all increments).
          2. Sum up all live shard values (carried).
          3. Rebuild shards + locks at new size.
          4. Redistribute carried counts evenly across new shards.
          5. Release lock → writes resume immediately.
        """
        with self._resize_lock:
            old_count = self._num_shards
            if new_shard_count <= old_count:
                return  # never downscale in this impl

            # Step 1: capture live (unflushed) counts
            carried = sum(self._shards[i] for i in range(old_count))

            # Step 2: rebuild
            self._num_shards  = new_shard_count
            self._shards      = [0] * new_shard_count
            self._shard_locks = [threading.Lock() for _ in range(new_shard_count)]

            # Step 3: redistribute carried counts without rounding loss
            base, remainder = divmod(carried, new_shard_count)
            for i in range(new_shard_count):
                self._shards[i] = base + (1 if i < remainder else 0)

            print(
                f"[Resharding] {old_count} → {new_shard_count} shards | "
                f"Carried {carried} unflushed likes"
            )

    @property
    def num_shards(self) -> int:
        return self._num_shards


# ---------------------------------------------------------------------------
# Duplicate-like guard
# ---------------------------------------------------------------------------

class LikeDeduplicator:
    """
    In-memory seen set.  In production replace with Redis SETNX + TTL.
    """

    def __init__(self):
        self._seen: set = set()
        self._lock = threading.Lock()

    def is_duplicate(self, user_id: int, post_id: int) -> bool:
        key = (user_id, post_id)
        with self._lock:
            if key in self._seen:
                return True
            self._seen.add(key)
            return False


# ---------------------------------------------------------------------------
# LikeSystem — wires everything together
# ---------------------------------------------------------------------------

class LikeSystem:
    def __init__(self):
        self.event_queue  = Queue()
        self.counter      = DynamicShardedCounter(INITIAL_SHARDS)
        self.deduplicator = LikeDeduplicator()

        self._db_like_count = 0       # simulated persistent DB
        self._db_lock       = threading.Lock()

        self._running = False

    # ------------------------------------------------------------------
    # Public entry points
    # ------------------------------------------------------------------

    def like(self, user_id: int, post_id: int) -> bool:
        """
        Called by the web layer when a user taps ♥.
        Returns False if the like is a duplicate (already counted).
        """
        if self.deduplicator.is_duplicate(user_id, post_id):
            return False
        self.event_queue.put((user_id, post_id))
        return True

    def get_like_count(self) -> int:
        """O(1) read: in-flight shards + DB snapshot."""
        with self._db_lock:
            return self._db_like_count + self.counter.read()

    # ------------------------------------------------------------------
    # Background threads
    # ------------------------------------------------------------------

    def _event_processor(self) -> None:
        """Worker: pulls events off the queue → increments a shard."""
        while self._running:
            try:
                user_id, post_id = self.event_queue.get(timeout=1)
                self.counter.increment(user_id)
                self.event_queue.task_done()
            except Exception:
                pass  # timeout — loop and check _running

    def _aggregator(self) -> None:
        """Flush shards → DB every FLUSH_INTERVAL seconds."""
        while self._running:
            time.sleep(FLUSH_INTERVAL)
            batch = self.counter.flush()
            with self._db_lock:
                self._db_like_count += batch
            print(
                f"[Aggregator] Flushed {batch:,} likes | "
                f"DB total: {self._db_like_count:,} | "
                f"Queue depth: {self.event_queue.qsize():,} | "
                f"Shards: {self.counter.num_shards}"
            )

    def _shard_monitor(self) -> None:
        """Auto-scale shards based on queue depth."""
        while self._running:
            time.sleep(MONITOR_INTERVAL)
            depth = self.event_queue.qsize()
            for threshold, target in SCALE_THRESHOLDS:
                if depth > threshold and self.counter.num_shards < target:
                    self.counter.rescale(target)
                    break

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def start(self) -> None:
        self._running = True

        for _ in range(NUM_WORKERS):
            threading.Thread(
                target=self._event_processor, daemon=True, name="EventProcessor"
            ).start()

        threading.Thread(
            target=self._aggregator, daemon=True, name="Aggregator"
        ).start()

        threading.Thread(
            target=self._shard_monitor, daemon=True, name="ShardMonitor"
        ).start()

        print(
            f"[LikeSystem] Started | "
            f"shards={INITIAL_SHARDS} | workers={NUM_WORKERS}"
        )

    def stop(self) -> None:
        self._running = False


# ---------------------------------------------------------------------------
# Traffic simulation
# ---------------------------------------------------------------------------

def simulate_normal_traffic(system: LikeSystem) -> None:
    """Steady low-volume traffic."""
    while True:
        user_id = random.randint(1, 100_000)
        system.like(user_id, post_id=123)
        time.sleep(0.01)


def simulate_viral_spike(system: LikeSystem) -> None:
    """
    Fires after 20 s to simulate a post going viral.
    Hammers the system with 500 likes/s to trigger auto-rescaling.
    """
    time.sleep(20)
    print("\n[Simulation] *** VIRAL SPIKE STARTING ***\n")
    while True:
        user_id = random.randint(1, 10_000_000)
        system.like(user_id, post_id=123)
        time.sleep(0.002)


def print_stats(system: LikeSystem) -> None:
    """Periodically print a live read of the like count."""
    while True:
        time.sleep(3)
        print(
            f"[Read]  Like count (O1): {system.get_like_count():,} | "
            f"Queue: {system.event_queue.qsize():,} | "
            f"Shards: {system.counter.num_shards}"
        )


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    system = LikeSystem()
    system.start()

    threading.Thread(target=simulate_normal_traffic, args=(system,), daemon=True).start()
    threading.Thread(target=simulate_viral_spike,    args=(system,), daemon=True).start()
    threading.Thread(target=print_stats,             args=(system,), daemon=True).start()

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n[LikeSystem] Shutting down...")
        system.stop()