Thread-safe designs: limiter, queue, transfers

Part of the free Low-Level Design course on LogicWiz, in the chapter: Concurrency in LLD. The class-diagram exercises run in the app.

Thread-safe designs: limiter, queue, transfers

"Your rate limiter passed every test. Then one bot sent eight threads at it, and it let through three times the limit."

Priya, at a design review


The rate limiter that let a bot through

A rate limiter decides whether a request may go ahead. SwiftKart limits each API key to a fixed number of requests. The usual design is a token bucket. Think of a bucket that holds up to 1,000 tokens. Each request takes one token. Tokens drip back in at a steady rate. An empty bucket means the request is refused.

Here is the part that decides, with no protection yet. The Bucket interface has one method, tryTake(now). The time arrives as a number of nanoseconds, so tests can control it.

public final class UnsafeBucket implements Bucket {
    private final long capacity;
    private final double perNano;         // tokens added per nanosecond
    private double tokens;
    private long last;                    // time of the last refill

    public UnsafeBucket(long capacity, double perSecond, long now) {
        this.capacity = capacity;
        this.perNano = perSecond / 1_000_000_000.0;
        this.tokens = capacity;           // a new bucket starts full
        this.last = now;
    }

    public boolean tryTake(long now) {
        if (now > last) {                 // add tokens for the time that passed
            tokens = Math.min(capacity, tokens + (now - last) * perNano);
            last = now;
        }
        if (tokens < 1) return false;     // check
        tokens -= 1;                      // act
        return true;
    }
}

The race is in two places. The check and the act are separate steps, so two threads can both see one token and both take it. The refill reads tokens and last, then writes both, so one thread can overwrite the other's work.

To prove it, we need a test that starts threads at the same moment, many times. A barrier makes all threads wait at one gate and start together. The clock stays frozen, so no tokens refill, and exactly 1,000 requests may pass.

static int oneRound(Supplier<Bucket> make, ExecutorService pool) throws Exception {
    Bucket bucket = make.get();                   // 1000 tokens, clock frozen
    CyclicBarrier go = new CyclicBarrier(8);      // all 8 threads start together
    AtomicInteger allowed = new AtomicInteger();
    List<Future<?>> done = new ArrayList<>();
    for (int t = 0; t < 8; t++) {
        done.add(pool.submit(() -> {
            go.await();
            for (int i = 0; i < 1000; i++) {
                if (bucket.tryTake(1_000L)) allowed.incrementAndGet();
            }
            return null;
        }));
    }
    for (Future<?> f : done) f.get();
    return allowed.get();                         // must be exactly 1000
}

In three runs of 2,000 rounds, the unsafe bucket gave a wrong count in 292, 346 and 297 rounds. In the worst round of each run it let through 3,649, 3,747 and 3,943 requests, not 1,000. With a smaller test (a bucket of 50 and 20 tries per thread), 5,000 rounds showed the bug only 3 times in one run and never in two others. A race may need many tries to show. A run with no failure proves nothing.

Fix it with a lock, or with one atomic swap

The first fix is a lock around the whole method. The check, the act and the refill become one critical section.

private final ReentrantLock lock = new ReentrantLock();

public boolean tryTake(long now) {
    lock.lock();
    try {
        if (now > last) {
            tokens = Math.min(capacity, tokens + (now - last) * perNano);
            last = now;
        }
        if (tokens < 1) return false;
        tokens -= 1;
        return true;
    } finally {
        lock.unlock();
    }
}

The second fix uses no lock. tokens and last must change together, so put them in one immutable object and replace it in a single step with compareAndSet. If another thread replaced it first, the step fails and the loop reads again.

private record State(double tokens, long last) { }
private final AtomicReference<State> state = new AtomicReference<>(new State(1000, 1_000L));

public boolean tryTake(long now) {
    while (true) {
        State s = state.get();
        double tokens = s.tokens();
        long last = s.last();
        if (now > last) {
            tokens = Math.min(capacity, tokens + (now - last) * perNano);
            last = now;
        }
        if (tokens < 1) return false;
        if (state.compareAndSet(s, new State(tokens - 1, last))) return true;
    }
}

Both versions were right in all 2,000 rounds of all three runs. Which should you choose? The lock is easier to read, and the Rate Limiter problem in the problem library uses it. The atomic version never makes a thread wait, but the loop may run twice, so the code inside must have no side effects. Start with the lock in an interview, and name the atomic version as an option.

A bounded queue by hand: wait and notify

The double booking lesson used ArrayBlockingQueue. Interviewers ask you to write one yourself. It holds at most capacity items. put waits when the queue is full. take waits when it is empty.

First, two tools. Every Java object has a built-in lock called its monitor, and synchronized takes it. Inside a synchronized method, wait() does two things at once: it releases the monitor and puts the thread to sleep in the object's wait set, the group of sleeping threads. notifyAll() wakes every thread in that set. notify() wakes one of them. A woken thread does not run at once. It must get the monitor back first, one thread at a time.

Picture a shelf with three slots and a bench beside it. A packer who finds the shelf full sits on the bench. A rider who finds it empty also sits on the bench. Whenever the shelf changes, someone rings a bell. The comparison stops here: one bench holds both kinds of people, and the bell cannot pick which kind should wake.

public synchronized void put(T item) throws InterruptedException {
    while (items.size() == capacity) {
        wait();                       // full: sleep until someone takes an item
    }
    items.addLast(item);
    notifyAll();                      // an item exists now: wake the waiters
}

public synchronized T take() throws InterruptedException {
    while (items.isEmpty()) {
        wait();                       // empty: sleep until someone puts an item
    }
    T item = items.removeFirst();
    notifyAll();                      // a place is free now: wake the waiters
    return item;
}

Two mistakes are classic. The first is if instead of while. Arjun's first draft used if. Between the wake-up and the thread's turn, another thread may have taken the item. Java also allows a spurious wake-up, where a thread wakes with no notify at all. So the thread must check the condition again. In a test, two consumers waited on an empty queue and then one item arrived. With if, one consumer got the item and the other crashed with NoSuchElementException. With while, the other one correctly went back to waiting.

The second mistake is notify instead of notifyAll. The bench holds both producers and consumers, and notify wakes one of them, with no way to choose which. It may wake a consumer when a producer was needed. That consumer finds the queue empty and sleeps again, and the wake-up is used up. The harder test used two producers and two consumers on a queue of size 1. With notifyAll, all 20 runs finished. With notify, none finished within three seconds, and one run left for ten seconds was still stuck.

The same queue with ReentrantLock and Condition

A Condition is a bench for one reason only. You make one from a lock with lock.newCondition(). Now the packers sit on a notFull bench and the riders on a notEmpty bench. await() and signal() are the wait and notify of a condition. signal now wakes the right kind of thread, so one wake-up is enough, because each put adds one item and each take frees one place. You still check in a while loop.

private final ReentrantLock lock = new ReentrantLock();
private final Condition notFull = lock.newCondition();
private final Condition notEmpty = lock.newCondition();

public void put(T item) throws InterruptedException {
    lock.lock();
    try {
        while (items.size() == capacity) notFull.await();
        items.addLast(item);
        notEmpty.signal();            // wake one rider, not a packer
    } finally {
        lock.unlock();
    }
}
// take() is the mirror: await on notEmpty, then notFull.signal()

All 20 runs of the hard case finished. In another test, four producers each put 5,000 numbers through a queue of size 3 while four consumers took them. The consumers counted 20,000 items with a sum of 50,010,000, which is exactly right, and the queue never held more than 3. In real code, use ArrayBlockingQueue. Write this one only when asked.

Draw the design below. The starter has a plain list with no protection.

Moving money between two wallets

Each Wallet has an id, a balance and its own lock. A transfer must change two wallets as one step, so it needs both locks. The obvious code locks from first and then to.

static boolean naive(Wallet from, Wallet to, long amount) {
    from.lock.lock();
    try {
        to.lock.lock();                // thread 2 holds this lock and wants ours
        try {
            if (!from.withdraw(amount)) return false;
            to.deposit(amount);
            return true;
        } finally { to.lock.unlock(); }
    } finally { from.lock.unlock(); }
}

In a test, thread 1 moved paise from A to B while thread 2 moved from B to A, 20,000 times each. In all three runs the program hung in the first round, and Java's deadlock detector reported two stuck threads. With no lock at all, the total of the two wallets was wrong in 41, 87 and 73 of 200 rounds. Both bugs are real.

The first fix is lock ordering. Every transfer locks the wallet with the lower id first, whatever the direction. No thread can then hold the higher one while waiting for the lower one.

Wallet first  = from.id < to.id ? from : to;
Wallet second = (first == from) ? to : from;
first.lock.lock();
try {
    second.lock.lock();
    try { /* withdraw from, deposit to */ } finally { second.lock.unlock(); }
} finally { first.lock.unlock(); }

The second fix is tryLock. Take the first lock, try the second for a short time, and if that fails, release the first and try again after a random pause. The random pause stops two threads from retrying in step with each other.

for (int attempt = 0; attempt < 50; attempt++) {
    if (from.lock.tryLock(1, TimeUnit.MILLISECONDS)) {
        try {
            if (to.lock.tryLock(1, TimeUnit.MILLISECONDS)) {
                try { return move(from, to, amount); } finally { to.lock.unlock(); }
            }
        } finally { from.lock.unlock(); }
    }
    Thread.sleep(ThreadLocalRandom.current().nextInt(1, 5));
}
return false;                          // gave up: the caller can try again

Ordering and tryLock both ended with the right total and no hang in 200 rounds, in all three runs. Ordering needs a fixed id for every lock. tryLock works when you cannot order, but it needs a retry plan.

Why a plain flag can hang a program

SwiftKart's notification worker runs until a flag says stop. This code looks right.

static boolean running = true;          // not volatile

Thread worker = new Thread(() -> {
    long spins = 0;
    while (running) { spins++; }         // nothing here changes running
});
worker.start();
Thread.sleep(500);
running = false;                         // the worker may never see this

In three runs, the worker was still running two seconds later. Each thread may keep its own copy of a value, and the compiler may read running once, because nothing in the loop changes it. Visibility is the question of when one thread sees another thread's write. Marking the field volatile makes every read see the latest write. With it, the worker stopped at once. A lock also gives visibility: what a thread wrote before releasing a lock is seen by the next thread that takes it.

Java does not promise that the plain version hangs. It does not promise that it stops. When a System.out.print call was added inside the loop, the worker stopped, because printing takes a lock. That hid the bug and did not fix it. Remember what the thread safety basics lesson said: volatile gives visibility, not atomicity.

⚠️ Warning: A test that passes does not prove that code is safe. Run races many times, start the threads together, and describe what you saw.

📌 Summary: Protect check-and-act with one lock or one atomic swap. Wait in a while loop and wake with notifyAll, or use Conditions. Lock wallets in id order. Use volatile for flags.

What Arjun learns next

Arjun can now build the three designs. The next lesson is about the interview itself: when to raise concurrency before the interviewer does, how to explain a design in 90 seconds, and the standard questions with short answers.

Check your understanding

  1. In a token bucket, two threads both see one token left, and both let their request through. What is the cause?
  2. A hand-written bounded queue uses `wait()` and `notify()`. With two producers and two consumers, it sometimes hangs for ever. What is the best fix?
  3. Thread 1 transfers from wallet A to B and thread 2 from B to A. Each locks `from` first, then `to`. Which change removes the deadlock without tryLock?
  4. A worker loops `while (running)` on a plain boolean. Main sets `running = false`, but the worker keeps going. What explains this and what is the fix?