Distributed Rate Limiter

Problem Overview

Design and implement a distributed rate limiter using a token bucket algorithm with lazy token refill. You are given two pre-implemented classes:

1. DistributedCache: A distributed key-value store (e.g., Redis) with get() and put() methods
2. TokenBucket: A data class representing a user's token bucket state

Your task is to implement two functions:

1. refill_token_bucket(bucket, current_time): Calculate and apply lazy token refill based on elapsed time
2. allow_request(cache, user_id, tokens_requested): Check if a user has enough tokens and consume them if so

This pattern is commonly used in:

• API rate limiting at scale (millions of users)
• Preventing abuse in distributed systems
• Fair resource allocation across services
• GPU inference quota management

Given Classes

import time
from dataclasses import dataclass

class DistributedCache:
    """
    A distributed cache (e.g., Redis) with atomic get/put operations.
    Assume this is already implemented and handles serialization.
    """
    def get(self, key: str) -> object | None:
        """Get value by key. Returns None if key doesn't exist."""
        pass

    def put(self, key: str, value: object) -> None:
        """Store value with key."""
        pass

@dataclass
class TokenBucket:
    """Represents the state of a user's token bucket."""
    tokens: float           # Current number of tokens
    last_refill_time: float # Timestamp of last refill (Unix time)
    capacity: int           # Maximum tokens the bucket can hold
    refill_rate: float      # Tokens added per second

Example Usage

cache = DistributedCache()

# User makes first request
result = allow_request(cache, "user_123", tokens_requested=1)
print(result)  # True (new bucket created with full capacity)

# User makes many requests quickly (exhaust remaining 99 tokens)
for _ in range(99):
    allow_request(cache, "user_123", tokens_requested=1)

# All 100 tokens consumed, next request fails
result = allow_request(cache, "user_123", tokens_requested=1)
print(result)  # False (out of tokens)

# Wait for refill...
time.sleep(1)  # Wait 1 second (assuming refill_rate=10/sec)

# Now user has tokens again
result = allow_request(cache, "user_123", tokens_requested=1)
print(result)  # True (10 new tokens available)

---

Part 1: Lazy Token Refill

Problem Statement

Implement the refill_token_bucket function that calculates how many tokens should be added based on elapsed time since the last refill.

Why lazy refill? Instead of running a background process to constantly refill all user buckets, we calculate the refill on-demand when a request arrives. This is more efficient for systems with millions of users where most are inactive at any given time.

def refill_token_bucket(bucket: TokenBucket, current_time: float) -> TokenBucket:
    """
    Calculate and apply lazy token refill based on elapsed time.

    Args:
        bucket: The current token bucket state
        current_time: Current Unix timestamp (seconds)

    Returns:
        Updated TokenBucket with refilled tokens and updated timestamp
    """
    pass

Requirements

1. Calculate elapsed time since last_refill_time
2. Add elapsed_time * refill_rate tokens to the bucket
3. Cap tokens at capacity (don't exceed maximum)
4. Update last_refill_time to current_time
5. Handle edge case: if current_time < last_refill_time (clock skew), don't remove tokens

Test Cases

# Test 1: Basic refill after 1 second
bucket = TokenBucket(tokens=50, last_refill_time=1000.0, capacity=100, refill_rate=10.0)
updated = refill_token_bucket(bucket, current_time=1001.0)
assert updated.tokens == 60  # 50 + (1 * 10)
assert updated.last_refill_time == 1001.0

# Test 2: Refill capped at capacity
bucket = TokenBucket(tokens=95, last_refill_time=1000.0, capacity=100, refill_rate=10.0)
updated = refill_token_bucket(bucket, current_time=1001.0)
assert updated.tokens == 100  # Capped at capacity, not 105

# Test 3: No time elapsed
bucket = TokenBucket(tokens=50, last_refill_time=1000.0, capacity=100, refill_rate=10.0)
updated = refill_token_bucket(bucket, current_time=1000.0)
assert updated.tokens == 50  # No change

# Test 4: Large time gap (user inactive for a while)
bucket = TokenBucket(tokens=0, last_refill_time=1000.0, capacity=100, refill_rate=10.0)
updated = refill_token_bucket(bucket, current_time=1100.0)  # 100 seconds later
assert updated.tokens == 100  # Fully refilled to capacity

---

Part 2: Allow Request Function

Problem Statement

Implement the allow_request function that checks if a user has enough tokens and consumes them atomically.

def allow_request(
    cache: DistributedCache,
    user_id: str,
    tokens_requested: int = 1,
    capacity: int = 100,
    refill_rate: float = 10.0
) -> bool:
    """
    Check if user has enough tokens. If yes, consume them and return True.
    If no, return False without consuming any tokens.

    Args:
        cache: Distributed cache instance
        user_id: Unique identifier for the user
        tokens_requested: Number of tokens to consume (default: 1)
        capacity: Maximum bucket capacity for new users
        refill_rate: Tokens per second refill rate for new users

    Returns:
        True if request is allowed (tokens consumed), False otherwise
    """
    pass

Requirements

1. Build the cache key from user_id (e.g., f"rate_limit:{user_id}")
2. If user doesn't exist in cache, create a new bucket with full capacity
3. Refill tokens based on elapsed time (call refill_token_bucket)
4. Check if tokens >= tokens_requested
5. If yes: deduct tokens, update cache, return True
6. If no: update cache with refilled state (so next request benefits), return False

Test Cases

# Test 1: New user gets full capacity
cache = DistributedCache()
assert allow_request(cache, "new_user") == True

# Test 2: Exhaust tokens
cache = DistributedCache()
for _ in range(100):
    assert allow_request(cache, "user_a") == True  # First 100 succeed

assert allow_request(cache, "user_a") == False  # 101st fails

# Test 3: Different users have separate buckets
cache = DistributedCache()
for _ in range(100):
    allow_request(cache, "user_x")

assert allow_request(cache, "user_x") == False  # user_x exhausted
assert allow_request(cache, "user_y") == True   # user_y still has tokens

# Test 4: Request multiple tokens at once
cache = DistributedCache()
assert allow_request(cache, "user_b", tokens_requested=50) == True   # 50 tokens left
assert allow_request(cache, "user_b", tokens_requested=50) == True   # 0 tokens left
assert allow_request(cache, "user_b", tokens_requested=1) == False   # Exhausted

---

Part 3: Concurrency Considerations

Problem Statement

In a distributed system, multiple servers may handle requests from the same user simultaneously. The basic implementation has a race condition:

Server A: GET user_123 -> bucket(tokens=10)
Server B: GET user_123 -> bucket(tokens=10)
Server A: tokens - 1 = 9, PUT user_123 -> bucket(tokens=9)
Server B: tokens - 1 = 9, PUT user_123 -> bucket(tokens=9)
# Both requests allowed, but only 1 token deducted!

Discussion Questions

1. How would you prevent this race condition?

   • Distributed locks (Redis SETNX/Redlock)
   • Atomic operations (Redis INCR/DECR, Lua scripts)
   • Optimistic locking with version numbers

2. What are the tradeoffs of each approach?

3. How would you handle lock failures or timeouts?

Example: Atomic Operation with Lua Script

If using Redis, you can execute the entire check-and-deduct as an atomic Lua script:

# Pseudocode for Redis Lua script approach
RATE_LIMIT_SCRIPT = """
local key = KEYS[1]
local tokens_requested = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local refill_rate = tonumber(ARGV[3])
local current_time = tonumber(ARGV[4])

local bucket = redis.call('HGETALL', key)
-- Parse bucket, refill, check, deduct, save
-- All in one atomic operation
"""

def allow_request_atomic(redis_client, user_id: str, tokens: int = 1) -> bool:
    result = redis_client.eval(
        RATE_LIMIT_SCRIPT,
        keys=[f"rate_limit:{user_id}"],
        args=[tokens, 100, 10.0, time.time()]
    )
    return result == 1

Requirements for Discussion

• Explain why the basic implementation has a race condition
• Propose at least two solutions with their tradeoffs
• Discuss which solution you would use in production and why

---

Follow-Up Discussion Topics

Comparison of Concurrency Solutions

Approach	Latency	Correctness	Complexity	Best For
No locking	Lowest	Approximate	Simplest	Soft limits, high scale
Distributed lock	Higher	Exact	Medium	Strict limits, lower QPS
Lua script	Low	Exact	Medium	Redis users, high QPS
Optimistic (CAS)	Variable	Exact	Higher	Low contention per user

Production Enhancements

1. Sliding Window vs Token Bucket

   • Token bucket allows bursts up to capacity
   • Sliding window provides smoother rate limiting
   • Hybrid: sliding window log with token bucket for burst control

2. Multi-Tier Rate Limits

   • Per-second, per-minute, per-hour limits
   • Check all tiers, fail if any exceeded

3. Rate Limit Headers

   • Return X-RateLimit-Remaining, X-RateLimit-Reset
   • Helps clients implement backoff

4. Graceful Degradation

   • If cache is down, allow requests (fail open) or reject (fail closed)?
   • Consider local in-memory fallback with sync

5. Monitoring and Alerting

   • Track rejection rates per user/endpoint
   • Alert on sudden spikes in rejections (possible attack)