Implement a Distributed Rate Limiter

Problem Statement

In a distributed infrastructure, multiple worker nodes must enforce API usage quotas against a shared backend. Design and implement a token bucket rate limiter that operates against a distributed key-value store supporting atomic Compare-And-Swap (CAS) operations.

The rate limiter tracks token consumption per key (e.g., customer or IP). Each key has a fixed maximum token capacity and replenishes at refill_rate_per_sec tokens per second. When a request arrives, the limiter computes accrued tokens since the last recorded timestamp, checks whether enough tokens are available, and updates the shared store via CAS.

Shared Key-Value Store Interface

You are provided an abstract KeyValueStore with the following operations:

class KeyValueStore:
    def get(self, key: str) -> tuple[dict | None, int]:
        """
        Retrieves the value and current version for the key.
        Returns (None, 0) if the key does not exist.
        """
        pass

    def cas(self, key: str, expected_version: int, new_value: dict, ttl_seconds: int) -> bool:
        """
        Atomically updates key to new_value if the current version equals expected_version.
        Returns True on success, False if a version mismatch occurs.
        """
        pass

Implementation Requirements

Implement the DistributedRateLimiter class:

class DistributedRateLimiter:
    def __init__(
        self,
        store: KeyValueStore,
        capacity: int,
        refill_rate_per_sec: float,
        max_cas_retries: int = 3,
        ttl_seconds: int = 3600,
    ):
        pass

    def allow_request(self, key: str, timestamp_ms: int, tokens: int = 1) -> bool:
        pass

• Token Bucket Behavior: Tokens replenish continuously: new_tokens = min(capacity, current_tokens + (elapsed_sec * refill_rate_per_sec)).
• Clock Skew Handling: If timestamp_ms < last_updated_ms (due to out-of-order events or local clock drift), assume elapsed_sec = 0 and do not reduce available tokens or roll back last_updated_ms.
• Concurrency Control: If another node modifies the key between get and cas, cas returns False. The limiter must re-fetch the latest state and retry up to max_cas_retries times. If all retries fail due to contention, return False.
• If available tokens after replenishment are strictly less than tokens, immediately return False without attempting CAS.

Constraints

• 1 <= capacity <= 10^6
• 0.001 <= refill_rate_per_sec <= 10^6
• 1 <= tokens <= capacity
• 0 <= max_cas_retries <= 10
• timestamp_ms is a non-negative 64-bit integer.

Examples

Example 1

store = InMemoryStore()
limiter = DistributedRateLimiter(store, capacity=5, refill_rate_per_sec=1.0)

# At t=0 ms: Key is uninitialized; bucket starts full (5 tokens). Consumes 3 tokens.
assert limiter.allow_request("user_1", timestamp_ms=0, tokens=3) == True

# At t=1000 ms: 1 second elapsed -> +1 token. Tokens = min(5, 2 + 1) = 3 tokens.
# Requesting 4 tokens fails (3 < 4). Store is not updated.
assert limiter.allow_request("user_1", timestamp_ms=1000, tokens=4) == False

# At t=1000 ms: Requesting 2 tokens succeeds. Remaining tokens = 1.
assert limiter.allow_request("user_1", timestamp_ms=1000, tokens=2) == True

Example 2: Clock Skew

# Continuing from Example 1 (last_updated_ms = 1000, tokens = 1.0)
# Out-of-order request arrives with timestamp_ms = 800
assert limiter.allow_request("user_1", timestamp_ms=800, tokens=1) == True
# Elapsed time is treated as 0 (not negative). Consumes 1 token, 0 remain.
# last_updated_ms remains 1000 in the store.