Lock-Free Token Bucket with AtomicLong and CAS

Your learning outcome for this lesson:

You will refactor the token bucket to a lock-free implementation using AtomicLong with CAS loops.

1. The Problem with Simple Operations: Beyond Locks

You might think that a simple operation like tokens-- is atomic, but it's not. Even if you declare the tokens variable as volatile to ensure its latest value is visible across threads, the increment/decrement operation itself is a sequence of three steps:

1. Read the current value.
2. Modify the value (add or subtract one).
3. Write the new value back.

Another thread can interfere between these steps, leading to lost updates—the very race condition we solved with locks.

To solve this without locks, we need a way to perform this read-modify-write sequence as a single, indivisible (atomic) hardware operation. This is where Atomic classes and Compare-and-Swap come in.

2. Compare-and-Swap (CAS): The Heart of Lock-Free Algorithms

Java's java.util.concurrent.atomic package, with classes like AtomicLong and AtomicInteger, provides methods that are guaranteed to be atomic. The most fundamental of these is compareAndSet(), which implements the CAS mechanism.

A CAS operation is an optimistic approach to updates:

1. It reads a value from memory.
2. It performs some computation to create a new value.
3. It then tries to set the new value only if the value in memory has not changed since it was first read.

If another thread changed the value in the meantime, the CAS operation fails, and the current thread typically retries the entire process in a loop. This "try-fail-retry" cycle happens without blocking and is often called a spin-loop or CAS loop.

The key idea is optimistic concurrency: we assume that contention (two threads trying to update at the same time) is rare. Threads proceed without acquiring a lock, and if a conflict occurs, one thread simply retries. In many scenarios, this is much more efficient than the pessimistic approach of locking, where threads are blocked even if there's no contention.

3. Refactoring the Token Bucket

Let's refactor our TokenBucket class. We'll replace the ReentrantLock and double availableTokens with AtomicLong. Since AtomicLong works with whole numbers, we'll adjust our design to use long for token counts. We also need to handle the lastRefillTimestamp in a thread-safe manner, so we'll make that an AtomicLong as well.

The core of our refactoring involves replacing the single locked critical section with two distinct, lock-free operations:

1. A lock-free refill operation.
2. A lock-free consumption operation.

The Race Condition in Naive Consumption

Before we write the correct code, it's vital to understand the wrong way to do it. As we've discussed, this seemingly safe code contains a race condition:

// THIS CODE IS WRONG
if (tokens.get() > 0) {
    tokens.decrementAndGet(); // Another thread can sneak in between get() and decrementAndGet()
    return true;
}

This exact issue is a common pitfall. Let's look at a real-world example of a developer encountering this problem.

Implementing the Lock-Free Logic

To implement the TokenBucket correctly without locks, we'll follow these steps in the tryConsume() method:

1. Refill:

   • We need to ensure only one thread refills the bucket at a time for a given period. We can achieve this by performing a CAS operation on lastRefillTimestamp.
   • The first thread to successfully update the timestamp "wins" the right to perform the refill. It calculates the tokens to add and updates availableTokens using another CAS loop.
   • Other threads that fail the timestamp CAS simply skip the refill logic, knowing another thread is handling it.

2. Consume:

   • After the potential refill, every thread will attempt to consume a token using the "optimistic locking" pattern we saw in the video and Stack Overflow answer.
   • It will read the current number of tokens, check if there are any available, and then use compareAndSet to decrement the count, retrying if the value was changed by another thread.

4. Your Implementation Task

It's time to apply these concepts. Below is the skeleton of our lock-free TokenBucket. Your task is to replace the ReentrantLock and double fields from the previous lesson with AtomicLongs and implement the new tryConsume() method based on the CAS patterns we've discussed.

Note: To keep the focus on the concurrency pattern, we're using long for token counts and calculating refills based on whole seconds.

import java.util.concurrent.atomic.AtomicLong;

public class LockFreeTokenBucket {

    private final long capacity;
    private final long refillRate; // tokens per second

    private final AtomicLong availableTokens;
    private final AtomicLong lastRefillTimestamp;

    private static final long NANOS_IN_SECOND = 1_000_000_000L;

    public LockFreeTokenBucket(long capacity, long refillRate) {
        this.capacity = capacity;
        this.refillRate = refillRate;
        this.availableTokens = new AtomicLong(capacity);
        this.lastRefillTimestamp = new AtomicLong(System.nanoTime());
    }

    /**
     * Attempts to consume one token from the bucket in a lock-free manner.
     * This method is thread-safe.
     *
     * @return {@code true} if a token was consumed, {@code false} otherwise.
     */
    public boolean tryConsume() {
        refill();

        // --- CONSUMPTION LOGIC ---
        // Implement the consumption CAS loop here.
        // 1. Get current tokens.
        // 2. If tokens are less than 1, return false.
        // 3. Try to decrement the token count using compareAndSet.
        // 4. Repeat if CAS fails.
        long currentTokens;
        do {
            currentTokens = availableTokens.get();
            if (currentTokens < 1) {
                return false;
            }
        } while (!availableTokens.compareAndSet(currentTokens, currentTokens - 1));

        return true;
    }

    private void refill() {
        long now = System.nanoTime();
        long lastRefill = lastRefillTimestamp.get();
        long elapsedNanos = now - lastRefill;

        if (elapsedNanos < NANOS_IN_SECOND) {
            return; // No need to refill yet
        }

        // --- REFILL LOGIC ---
        // Use CAS to ensure only one thread refills.
        // 1. Attempt to update the lastRefillTimestamp from its old value to `now`.
        if (lastRefillTimestamp.compareAndSet(lastRefill, now)) {
            // 2. If successful, this thread won the race. Calculate tokens to add.
            long secondsElapsed = elapsedNanos / NANOS_IN_SECOND;
            long tokensToAdd = secondsElapsed * refillRate;

            if (tokensToAdd > 0) {
                // 3. Add the new tokens using another CAS loop to avoid exceeding capacity.
                long currentTokens;
                long newTokens;
                do {
                    currentTokens = availableTokens.get();
                    // Ensure we don't exceed capacity.
                    newTokens = Math.min(capacity, currentTokens + tokensToAdd);
                    // No need to update if we are not adding any tokens
                    if (newTokens == currentTokens) break;
                } while (!availableTokens.compareAndSet(currentTokens, newTokens));
            }
        }
        // If the initial CAS fails, another thread is already handling the refill, so we do nothing.
    }
}

Instructions:

1. Create a new class LockFreeTokenBucket in your project.
2. The code for the tryConsume() consumption loop and the refill() method are provided above. Study them carefully to understand how the CAS loops for both consumption and refilling work.
3. Notice how the refill logic uses a CAS on the timestamp to "elect" a leader, which then performs the update to the token count. This is a common and powerful lock-free pattern.
4. Later, we will make this class conform to the RateLimiter interface you designed earlier.

Conclusion

Congratulations! You have now implemented a sophisticated, lock-free token bucket. By moving from a blocking lock to non-blocking atomic operations, you've taken a significant step toward building a high-performance system.

Key Takeaways:

• Lock-free algorithms use atomic hardware instructions like Compare-and-Swap (CAS) to manage concurrent state changes without blocking threads.
• The java.util.concurrent.atomic package provides classes like AtomicLong that expose these atomic operations.
• The common pattern for lock-free updates is an "optimistic" CAS loop: read a value, compute a new one, and atomically set it if the original value hasn't changed. If it has, retry.
• While powerful, lock-free code can be more complex than its locked equivalent, often requiring careful handling of multiple atomic variables to ensure correctness (as seen in our refill method).

Next Steps

We now have two implementations of our in-memory token bucket: one lock-based and one lock-free. Which one is actually better? And by how much?

In the next lesson, we will explain the performance trade-offs between lock-based and lock-free concurrency models. We'll discuss when to choose one over the other by looking at factors like contention levels and hardware characteristics. This will complete our tour of advanced in-memory implementations.