vishal maurya

LLD Series: In-Memory Pub/Sub (Kafka-lite) - LLD

Overview

Part of the Interview Prep series. Dotted-underline words are glossary terms: hover for the meaning or open the Glossary. Track your progress on the Interview Prep tracker. Series index: Start here.

1. Problem statement and how it is asked

"Design an in-memory publish-subscribe system. Publishers publish messages to topics; subscribers subscribe to topics and receive messages. Delivery should be asynchronous, a slow subscriber must not block others, and ordering must be preserved. Later: let consumers resume from an offset." Asked at: Visa, Mastercard, PayU, AmEx (event/payment-notification flavours), NPCI (switch -> event fan-out). Often phrased as "design a message queue / Kafka-lite / event bus".

2. Requirements

Functional: create topic; publish(topic, key, payload); subscribe/unsubscribe; push delivery to each subscriber (fan-out: every subscriber gets every message); pull consumers with poll/commit/seek (offset-based replay); new subscriber chooses "from beginning" or "latest". Non-functional: thread-safe, per-subscriber ordering equals publish (log) order, slow-consumer isolation, retry + dead-letter for poison messages, no user code executed under a lock. Assumptions: single process, one partition per topic, unbounded log (memory), at-least-once delivery per subscriber. Out of scope: persistence, partitions/consumer groups (discussed in follow-ups), network, auth, retention/compaction.

3. Entities and class diagram

 Broker 1 ---- * Topic 1 ---- * Message (offset, key, payload, ts)
                 | append-only log
                 | * PushSubscription  (own offset, own single-thread executor, DLQ)  -> Subscriber (functional iface)
 Broker ------- creates PullConsumer (own position + committed offset) reads Topic log

+-----------------+        +---------------------------+      +-------------------------+
| Broker          |        | Topic                     |      | PushSubscription        |
|-----------------|        |---------------------------|      |-------------------------|
| topics: CHM     |1     * | log: List<Message>        |1   * | offset (volatile)       |
| createTopic()   |------->| subs: CopyOnWriteArrayList|----->| exec: 1 thread          |
| publish()       |        | append(), read(from,max)  |      | scheduled: AtomicBool   |
| subscribe()     |        +---------------------------+      | signal(), drain()       |
| pullConsumer()  |                                           | deliver() retry -> DLQ  |
+-----------------+                                           +-------------------------+
                                  PullConsumer: poll(max) / commit() / seek() / rollbackToCommitted()

4. Design patterns and why

  • Observer / Pub-Sub: the core; Topic is the subject, subscriptions the observers.
  • Append-only log + offsets (not per-subscriber queues): storage is O(messages) not O(messages x subscribers); replay and late joiners come free. Rejected: a BlockingQueue per subscriber (copies, no replay, unbounded memory per slow consumer, needs offsets bolted on later).
  • Producer-consumer with a "signal + drain" loop: publisher only appends and signals; consumer thread drains from its own offset. Rejected: publisher calling subscribers inline (slow subscriber blocks publisher); thread-per-message (breaks ordering).
  • Single-thread executor per subscriber: gives ordering per subscriber with isolation across subscribers. Rejected: shared thread pool (ordering lost unless keyed).
  • Strategy-ready: Subscriber is a functional interface; retry/DLQ policy can be extracted.

5. Code

All files in default package; compile with javac *.java && java Main.

Message.java

import java.time.Instant;
public final class Message {
    private final long offset; private final String key; private final String payload; private final Instant ts;
    public Message(long offset, String key, String payload) { this.offset=offset; this.key=key; this.payload=payload; this.ts=Instant.now(); }
    public long offset(){return offset;} public String key(){return key;} public String payload(){return payload;} public Instant timestamp(){return ts;}
    @Override public String toString(){return "Message{offset="+offset+", payload="+payload+"}";}
}

Subscriber.java

@FunctionalInterface
public interface Subscriber { void onMessage(Message m) throws Exception; }

Topic.java

import java.util.*;
import java.util.concurrent.CopyOnWriteArrayList;
/** Append-only log. One lock guards append so offsets are dense and ordered. */
public class Topic {
    private final String name;
    private final List<Message> log = new ArrayList<>();
    private final List<PushSubscription> subs = new CopyOnWriteArrayList<>();
    public Topic(String name){this.name=name;}
    public String name(){return name;}
    public Message append(String key, String payload) {
        Message m;
        synchronized (log) { m = new Message(log.size(), key, payload); log.add(m); }
        for (PushSubscription s : subs) s.signal();   // outside the lock: never call user code under it
        return m;
    }
    public List<Message> read(long from, int max) {
        synchronized (log) {
            if (from >= log.size()) return Collections.emptyList();
            return new ArrayList<>(log.subList((int) from, (int) Math.min(log.size(), from + max)));
        }
    }
    public long size(){ synchronized (log){ return log.size(); } }
    void addSubscription(PushSubscription s){subs.add(s);}
    void removeSubscription(PushSubscription s){subs.remove(s);}
}

PushSubscription.java

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;
/** Push consumer: own offset + single-thread executor => per-subscriber ordering, isolation from slow peers. */
public class PushSubscription {
    private final String id; private final Topic topic; private final Subscriber handler;
    private final ExecutorService exec;
    private final AtomicBoolean scheduled = new AtomicBoolean(false);
    private volatile long offset;               // next offset to deliver; only touched by the drain thread
    private final int maxRetries;
    private final List<Message> deadLetters = new CopyOnWriteArrayList<>();
    public PushSubscription(String id, Topic topic, Subscriber handler, long startOffset, int maxRetries) {
        this.id=id; this.topic=topic; this.handler=handler; this.offset=startOffset; this.maxRetries=maxRetries;
        this.exec = Executors.newSingleThreadExecutor(r -> { Thread t = new Thread(r, "sub-"+id); t.setDaemon(true); return t; });
    }
    void signal() {
        if (scheduled.compareAndSet(false, true)) exec.execute(this::drain);
    }
    private void drain() {
        try {
            List<Message> batch;
            while (!(batch = topic.read(offset, 100)).isEmpty()) {
                for (Message m : batch) { deliver(m); offset = m.offset() + 1; }
            }
        } finally {
            scheduled.set(false);
            if (topic.size() > offset) signal();   // close the race: message arrived after last read
        }
    }
    private void deliver(Message m) {
        for (int attempt = 1; attempt <= maxRetries; attempt++) {
            try { handler.onMessage(m); return; }
            catch (Exception e) { /* retry */ }
        }
        deadLetters.add(m);   // poison message must not block the partition forever
    }
    public List<Message> deadLetters(){return deadLetters;}
    public long offset(){return offset;}
    public String id(){return id;}
    public void close(){ topic.removeSubscription(this); exec.shutdown(); }
    public boolean awaitIdle(long ms) throws InterruptedException {
        long end = System.currentTimeMillis()+ms;
        while (System.currentTimeMillis()<end) { if (!scheduled.get() && offset>=topic.size()) return true; Thread.sleep(5);}
        return false;
    }
}

PullConsumer.java

import java.util.*;
import java.util.concurrent.atomic.AtomicLong;
/** Kafka-style: consumer owns the cursor. poll() doesn't advance committed offset; commit() does. */
public class PullConsumer {
    private final Topic topic; private final AtomicLong position = new AtomicLong(); private final AtomicLong committed = new AtomicLong();
    public PullConsumer(Topic t, long start){topic=t; position.set(start); committed.set(start);}
    public List<Message> poll(int max){
        List<Message> r = topic.read(position.get(), max);
        position.addAndGet(r.size()); return r;
    }
    public void commit(){ committed.set(position.get()); }
    public void seek(long o){ position.set(o); }
    public void rollbackToCommitted(){ position.set(committed.get()); }
    public long position(){return position.get();} public long committed(){return committed.get();}
}

Broker.java

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
public class Broker {
    private final ConcurrentMap<String, Topic> topics = new ConcurrentHashMap<>();
    public Topic createTopic(String n){ return topics.computeIfAbsent(n, Topic::new); }
    private Topic get(String n){ Topic t = topics.get(n); if (t==null) throw new IllegalArgumentException("No such topic: "+n); return t; }
    public Message publish(String topic, String key, String payload){ return get(topic).append(key, payload); }
    public PushSubscription subscribe(String topic, String id, Subscriber s, boolean fromBeginning){
        Topic t = get(topic);
        PushSubscription p = new PushSubscription(id, t, s, fromBeginning ? 0 : t.size(), 3);
        t.addSubscription(p); p.signal(); return p;
    }
    public PullConsumer pullConsumer(String topic, long start){ return new PullConsumer(get(topic), start); }
}

Main.java

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;
public class Main {
    public static void main(String[] a) throws Exception {
        Broker b = new Broker(); b.createTopic("orders");
        b.publish("orders", "k", "old-1");   // published before any subscriber
        List<String> fast = new CopyOnWriteArrayList<>(), slow = new CopyOnWriteArrayList<>();
        PushSubscription f = b.subscribe("orders","fast", m -> fast.add(m.payload()), true);
        PushSubscription s = b.subscribe("orders","slow", m -> { Thread.sleep(20); slow.add(m.payload()); }, false);
        // 4 concurrent publishers x 25 msgs
        ExecutorService pool = Executors.newFixedThreadPool(4); CountDownLatch l = new CountDownLatch(4);
        for (int p=0;p<4;p++){ final int pp=p; pool.execute(() -> { for(int i=0;i<25;i++) b.publish("orders","k","p"+pp+"-"+i); l.countDown(); }); }
        l.await(); pool.shutdown();
        f.awaitIdle(5000); s.awaitIdle(10000);
        System.out.println("fast (from beginning) got " + fast.size() + ", first=" + fast.get(0));
        System.out.println("slow (latest only) got " + slow.size());
        // ordering check: per-publisher sequence must be increasing in each subscriber
        System.out.println("per-publisher order preserved (fast): " + ordered(fast));
        System.out.println("per-publisher order preserved (slow): " + ordered(slow));
        // retry + DLQ
        AtomicInteger tries = new AtomicInteger();
        b.createTopic("pay");
        PushSubscription bad = b.subscribe("pay","bad", m -> { if (m.payload().equals("poison")) { tries.incrementAndGet(); throw new RuntimeException("x"); } }, true);
        b.publish("pay","k","ok"); b.publish("pay","k","poison"); b.publish("pay","k","ok2");
        bad.awaitIdle(2000);
        System.out.println("poison attempts=" + tries + ", DLQ=" + bad.deadLetters() + ", offset=" + bad.offset());
        // pull consumer with commit/replay
        PullConsumer pc = b.pullConsumer("pay", 0);
        System.out.println("poll(2): " + pc.poll(2)); pc.commit();
        System.out.println("poll(5): " + pc.poll(5) + " (not committed)");
        pc.rollbackToCommitted();
        System.out.println("after rollback, position=" + pc.position() + ", replay: " + pc.poll(5));
        try { b.publish("nope","k","x"); } catch (IllegalArgumentException e) { System.out.println("error: " + e.getMessage()); }
    }
    static boolean ordered(List<String> xs){
        Map<String,Integer> last = new HashMap<>();
        for (String x : xs){ if(!x.startsWith("p")) continue; String[] p=x.split("-"); int i=Integer.parseInt(p[1]);
            Integer prev=last.put(p[0], i); if(prev!=null && prev>=i) return false; }
        return true;
    }
}

Lock granularity: Topic.log monitor is held only for the O(1) append / O(k) copy, never while invoking subscriber code. PushSubscription.offset is written only by its single drain thread (volatile for visibility). scheduled AtomicBoolean ensures at most one drain task is queued; the finally re-check closes the lost-wakeup race (message appended after last read but before scheduled=false).

Real output (compiled with javac, run with java)

fast (from beginning) got 101, first=old-1
slow (latest only) got 100
per-publisher order preserved (fast): true
per-publisher order preserved (slow): true
poison attempts=3, DLQ=[Message{offset=1, payload=poison}], offset=3
poll(2): [Message{offset=0, payload=ok}, Message{offset=1, payload=poison}]
poll(5): [Message{offset=2, payload=ok2}] (not committed)
after rollback, position=2, replay: [Message{offset=2, payload=ok2}]
error: No such topic: nope

6. Key flows

Publisher            Topic                     PushSubscription (thread sub-X)      Subscriber
   | publish(k,p)      |                              |                                  |
   |------------------>| synchronized: offset=size, add                                  |
   |                   |-- signal() for each sub ---->| CAS scheduled false->true         |
   |<-- Message -------|                              | exec.execute(drain)               |
                                                      | read(offset,100) -> batch         |
                                                      |-- onMessage(m) ------------------>|
                                                      |   (fail x3 -> DLQ) ; offset=m+1   |
                                                      | loop until empty; scheduled=false |
                                                      | if size>offset -> signal() again  |
Pull:  poll(n) -> read(position,n), position+=n ; commit() -> committed=position ; rollbackToCommitted() -> replay

7. Extensibility

  • Partitions: Topic holds N logs, choose by hash(key) % N (PartitionStrategy interface); ordering per key retained.
  • Consumer groups: introduce ConsumerGroup owning one committed offset per partition; members share partitions.
  • Filtering: wrap handler with a Predicate<Message>/decorator subscriber.
  • Retry policy: extract RetryPolicy (backoff) injected into PushSubscription.
  • Retention: a trimmer that advances a log base offset (log becomes base + deque). None require modifying Subscriber or Broker.publish.

8. Concurrency and edge cases

  • Concurrent publishers: dense, unique offsets because append is under one lock; the order seen by every subscriber is that log order.
  • Lost wakeup: handled by the finally re-check.
  • Slow subscriber: lags only itself; the log grows (memory) -> bounded retention / backpressure in real systems.
  • Poison message: bounded retries then DLQ, offset advances (no head-of-line blocking forever).
  • Subscribe while publishing: fromBeginning=false uses topic.size() as start; a message racing exactly may be seen or missed - document as "latest means from subscription time"; fromBeginning is the lossless choice.
  • Unknown topic -> IllegalArgumentException; null payloads should be rejected; close() stops the executor.
  • Handler exceptions never kill the drain thread (caught in deliver).

9. Follow-ups and model answers

  • Exactly-once? Not possible end-to-end by the broker alone: at-least-once + idempotent consumer (dedupe on message id / offset) or transactional offset commit with the side-effect.
  • Make it distributed/persistent? Log segments on disk (append-only, sequential writes, mmap/page cache, index by offset), partitions across brokers with leader/follower replication (ISR), offsets stored in an internal topic, ZooKeeper/Raft for metadata, consumer-group rebalance. That is Kafka.
  • Backpressure? Bounded log with retention + slow consumers either lose data (retention) or publisher gets blocked/errors; pull model naturally gives consumer-side backpressure.
  • Ordering across subscribers? Only total order per partition; none across partitions.
  • Why not pushing from the publisher thread? Latency coupling and failures leak into publisher.
  • Complexity: publish O(1) amortised + O(S) signals (S subscribers); read O(k); memory O(N) messages.

10. Common mistakes and rubric

Mistakes: calling subscribers under a lock or on publisher thread; one shared queue so messages are consumed by only one subscriber (competing consumers confusion); thread pool per message losing order; no way to replay; HashMap for topics without concurrency; ignoring exceptions in handlers killing the thread; unbounded retry loop. Rubric (10): requirements + fan-out vs competing semantics clarified (2); log+offset model (2); ordering and isolation correct (2); thread-safety reasoning incl. wakeup race (2); retry/DLQ/pull-commit extras and distributed discussion (2).