Make Token-Bucket Admission Safe Across Threads

Problem Statement

In high-throughput distributed healthcare platforms, background jobs and external API webhooks frequently compete for restricted upstream endpoints. To prevent overwhelming downstream services, incoming requests must pass through a strict rate-limiting admission controller.

Implement a thread-safe TokenBucket class that regulates admissions using the standard token-bucket algorithm. Rather than running a background worker thread to continuously deposit tokens, the bucket should refill tokens lazily whenever admission is attempted, using a monotonic clock.

Class Interface

Implement the following methods:

class TokenBucket:
    def __init__(self, capacity: int, refill_rate: float):
        """
        capacity: maximum number of tokens the bucket can hold (integer > 0).
        refill_rate: tokens added per second (float > 0.0).
        The bucket starts full at initialization.
        """
        pass

    def try_acquire(self, tokens: int = 1) -> bool:
        """
        Attempts to consume `tokens` immediately without blocking.
        Returns True if tokens were acquired, False otherwise.
        """
        pass

    def acquire(self, tokens: int = 1, timeout: float = None) -> bool:
        """
        Blocks until `tokens` are available or until `timeout` seconds elapse.
        If `timeout` is None, blocks indefinitely until tokens can be acquired.
        Returns True if acquired within the timeout, False if the timeout expired.
        """
        pass

Requirements

1. Thread Safety: Multiple threads will call try_acquire and acquire concurrently. All state mutations must be synchronized to prevent race conditions such as over-allocation or lost updates.
2. Monotonic Time: Time calculations must rely on a steady, monotonic clock (time.monotonic() in Python) to avoid skew from system wall-clock adjustments.
3. Efficient Waiting: In acquire, threads should wait on a condition variable or equivalent synchronization primitive rather than busy-waiting / polling in a loop.
4. Refill Cap: Tokens accumulate continuously proportional to elapsed time, but the bucket must never exceed capacity.

Constraints

• 1 <= capacity = 0.0).

Example

bucket = TokenBucket(capacity=5, refill_rate=2.0)  # Starts with 5.0 tokens

# Immediate consumption
bucket.try_acquire(3)  # Returns True (2.0 tokens remain)
bucket.try_acquire(3)  # Returns False (only 2.0 tokens available)

# Blocking consumption with timeout
# 1 token needed to reach 3; at 2.0 tokens/sec, requires 0.5s wait
bucket.acquire(3, timeout=0.2)  # Returns False after 0.2s (timed out)
bucket.acquire(3, timeout=1.0)  # Waits ~0.5s, acquires, returns True