Concurrency LLDs: Task Scheduler, Thread Pool, Producer-Consumer & Mini Kafka

import heapq
import itertools
import threading
import time
from abc import ABC, abstractmethod
from concurrent.futures import ThreadPoolExecutor
from enum import IntEnum, Enum
from typing import Callable


class TaskPriority(IntEnum):        # smaller number = runs first
    HIGH = 1
    MEDIUM = 2
    LOW = 3


class TaskStatus(Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"
    CANCELLED = "CANCELLED"


class Task(ABC):
    _ids = itertools.count(1)

    def __init__(self, priority: TaskPriority, run_at: float, job: Callable[[], None]):
        self.task_id = next(Task._ids)
        self.priority = priority
        self.next_run_time = run_at
        self.job = job                       # Java's Runnable = any Python callable

    def execute(self) -> None:
        self.job()

    @abstractmethod
    def on_complete(self) -> bool:
        """Return True if the task must be queued again."""


class OneTimeTask(Task):
    def on_complete(self) -> bool:
        return False                         # done forever


class RecurringTask(Task):
    def __init__(self, priority, interval: float, job, first_run: float | None = None):
        super().__init__(priority, first_run or time.time(), job)
        self.interval = interval

    def on_complete(self) -> bool:
        self.next_run_time = time.time() + self.interval   # e.g. alarm: same time tomorrow
        return True                                        # re-queue me


class ScheduledTask:
    """Wrapper that knows how to ORDER tasks: time first, then priority."""

    def __init__(self, task: Task):
        self.task = task

    def __lt__(self, other: "ScheduledTask") -> bool:
        a, b = self.task, other.task
        if a.next_run_time != b.next_run_time:
            return a.next_run_time < b.next_run_time       # earlier time wins
        return a.priority < b.priority                     # same time: priority decides


class TaskScheduler:
    def __init__(self, workers: int = 3, tick: float = 0.1):
        self._queue: list[ScheduledTask] = []              # min-heap
        self._lock = threading.Lock()
        self.registry: dict[int, TaskStatus] = {}
        self._workers = ThreadPoolExecutor(max_workers=workers)
        self._tick = tick
        self._running = False

    def submit(self, task: Task) -> int:
        with self._lock:
            heapq.heappush(self._queue, ScheduledTask(task))
            self.registry[task.task_id] = TaskStatus.PENDING
        return task.task_id

    def cancel(self, task_id: int) -> None:
        with self._lock:
            self.registry[task_id] = TaskStatus.CANCELLED

    def start(self) -> None:
        self._running = True
        threading.Thread(target=self._ticker, daemon=True).start()

    def _ticker(self) -> None:                             # wakes up every 100 ms
        while self._running:
            self._poll()
            time.sleep(self._tick)

    def _poll(self) -> None:
        now = time.time()
        with self._lock:
            while self._queue and self._queue[0].task.next_run_time <= now:
                st = heapq.heappop(self._queue)
                if self.registry[st.task.task_id] != TaskStatus.CANCELLED:
                    self._workers.submit(self._run, st.task)   # hand to a worker

    def _run(self, task: Task) -> None:                    # runs on a worker thread
        self.registry[task.task_id] = TaskStatus.RUNNING
        try:
            task.execute()
            self.registry[task.task_id] = TaskStatus.COMPLETED
            if task.on_complete():                         # recurring -> back in queue
                self.submit(task)
        except Exception:
            self.registry[task.task_id] = TaskStatus.FAILED

    def shutdown(self) -> None:
        self._running = False
        self._workers.shutdown(wait=True)


# usage
if __name__ == "__main__":
    s = TaskScheduler()
    s.start()
    now = time.time()
    s.submit(OneTimeTask(TaskPriority.LOW, now + 0.3, lambda: print("low  @0.3s")))
    s.submit(OneTimeTask(TaskPriority.HIGH, now + 0.3, lambda: print("high @0.3s")))
    rid = s.submit(RecurringTask(TaskPriority.MEDIUM, 0.25, lambda: print("tick (recurring)")))
    time.sleep(0.8)
    s.cancel(rid)
    time.sleep(0.3)
    s.shutdown()
# tick (recurring)
# high @0.3s      <- same time as "low", but HIGH priority is taken first
# low  @0.3s
# tick (recurring)
# tick (recurring)

import queue
import threading
import time
from abc import ABC, abstractmethod
from typing import Callable

_STOP = object()                     # "poison pill": tells a worker to exit its loop


class ThreadPool(ABC):
    @abstractmethod
    def submit(self, task: Callable[[], None]) -> None: ...
    @abstractmethod
    def shutdown(self) -> None: ...


class Worker(threading.Thread):
    def __init__(self, task_queue: queue.Queue, name: str):
        super().__init__(name=name, daemon=True)
        self.task_queue = task_queue      # a REFERENCE to the shared queue

    def run(self) -> None:
        while True:                       # loop forever: a worker never quits on its own
            task = self.task_queue.get()  # BLOCKS (sleeps) while the queue is empty
            if task is _STOP:
                break
            try:
                task()
            except Exception as e:        # one bad task must not kill the worker
                print(f"{self.name} task failed: {e}")


class FixedThreadPool(ThreadPool):
    def __init__(self, pool_size: int, queue_capacity: int):
        self.task_queue: queue.Queue = queue.Queue(maxsize=queue_capacity)  # BOUNDED
        self.workers = [Worker(self.task_queue, f"worker-{i}") for i in range(pool_size)]
        self.is_shutdown = False
        for w in self.workers:
            w.start()                     # start() -> calls run() on a new thread

    def submit(self, task: Callable[[], None]) -> None:
        if self.is_shutdown:
            raise RuntimeError("pool is shut down")
        self.task_queue.put(task)         # BLOCKS while the queue is full (back-pressure)

    def shutdown(self) -> None:
        self.is_shutdown = True
        for _ in self.workers:
            self.task_queue.put(_STOP)    # one pill per worker, after the real tasks
        for w in self.workers:
            w.join()


# usage
if __name__ == "__main__":
    pool = FixedThreadPool(pool_size=3, queue_capacity=10)
    for n in range(6):
        pool.submit(lambda n=n: (time.sleep(0.05),
                                 print(f"task {n} on {threading.current_thread().name}")))
    pool.shutdown()
# task 0 on worker-2
# task 2 on worker-0
# task 1 on worker-1
# task 3 on worker-2   <- only 3 threads, reused for 6 tasks (order may vary)
# ...

import queue
import threading
import time
from abc import ABC, abstractmethod


class Buffer(ABC):                         # "a piece of memory" - not necessarily a queue
    @abstractmethod
    def put(self, item) -> None: ...
    @abstractmethod
    def take(self): ...


class BoundedBuffer(Buffer):
    def __init__(self, capacity: int):
        self._q: queue.Queue = queue.Queue(maxsize=capacity)   # FIFO + thread-safe

    def put(self, item) -> None:
        self._q.put(item)                  # blocks if full

    def take(self):
        return self._q.get()               # blocks if empty


class Producer(threading.Thread):
    def __init__(self, buffer: Buffer, count: int):
        super().__init__(daemon=True)
        self.buffer, self.count = buffer, count

    def run(self) -> None:
        for i in range(self.count):
            self.buffer.put(i)
            print(f"produced {i}")
            time.sleep(0.01)               # simulate work
        self.buffer.put(None)              # signal "no more items"


class Consumer(threading.Thread):
    def __init__(self, buffer: Buffer):
        super().__init__(daemon=True)
        self.buffer = buffer
        self.is_running = True

    def run(self) -> None:
        while self.is_running:             # WHILE, not IF: keep consuming
            item = self.buffer.take()
            if item is None:
                break
            print(f"  consumed {item}")
            time.sleep(0.03)               # consumer is slower than producer


# usage
if __name__ == "__main__":
    buf = BoundedBuffer(capacity=2)
    p, c = Producer(buf, 4), Consumer(buf)
    p.start(); c.start(); p.join(); c.join()
# produced 0 / consumed 0 / produced 1 / produced 2 ... producer waits when buffer is full

import threading
import zlib
from dataclasses import dataclass


@dataclass(frozen=True)
class Message:
    key: str        # decides the partition
    value: str


class Partition:
    """The physical unit: an append-only log of messages."""

    def __init__(self, pid: int):
        self.pid = pid
        self.log: list[Message] = []
        self._lock = threading.Lock()     # Java used a ReadWriteLock

    def append(self, msg: Message) -> int:
        with self._lock:                  # two writers must not grab the same offset
            self.log.append(msg)
            return len(self.log) - 1      # offset of the new message

    def read_from(self, offset: int) -> list[Message]:
        with self._lock:
            return self.log[offset:]      # reading never removes anything


class Topic:
    """Logical name grouping several partitions (aggregation)."""

    def __init__(self, name: str, partitions: int):
        self.name = name
        self.partitions = [Partition(i) for i in range(partitions)]

    def partition_for(self, key: str) -> Partition:
        # stable hash: same key -> same partition -> per-key ordering
        return self.partitions[zlib.crc32(key.encode()) % len(self.partitions)]


class Broker:
    def __init__(self):
        self.topics: dict[str, Topic] = {}

    def create_topic(self, name: str, partitions: int) -> Topic:
        self.topics[name] = Topic(name, partitions)
        return self.topics[name]


class Producer:
    def __init__(self, broker: Broker):
        self.broker = broker

    def send(self, topic: str, key: str, value: str) -> tuple[int, int]:
        p = self.broker.topics[topic].partition_for(key)
        return p.pid, p.append(Message(key, value))


class Consumer:
    """Each consumer remembers its own offset per partition, so it can replay."""

    def __init__(self, broker: Broker, topic: str):
        self.topic = broker.topics[topic]
        self.offsets = {p.pid: 0 for p in self.topic.partitions}

    def poll(self) -> list[Message]:
        out = []
        for p in self.topic.partitions:
            msgs = p.read_from(self.offsets[p.pid])
            self.offsets[p.pid] += len(msgs)
            out.extend(msgs)
        return out

    def seek_to_beginning(self) -> None:              # replay old events
        self.offsets = {pid: 0 for pid in self.offsets}


# usage
if __name__ == "__main__":
    kafka = Broker()
    kafka.create_topic("lld", partitions=3)
    prod = Producer(kafka)
    for i, user in enumerate(["ajay", "david", "ajay", "sumit"]):
        print(prod.send("lld", user, f"event-{i}"))   # (partition, offset)
    # (1, 0) (0, 0) (1, 1) (1, 2)  <- both "ajay" events land in partition 1
    cons = Consumer(kafka, "lld")
    print([m.value for m in cons.poll()])              # ['event-1', 'event-0', 'event-2', 'event-3']
    print(cons.poll())                                 # [] - nothing new
    cons.seek_to_beginning()
    print(len(cons.poll()))                            # 4 - replayed