token_bucket.cpp

Go to the documentation of this file.

// BSD 3-Clause License
// Copyright (c) 2024, 🍀☀🌕🌥 🌊
// See the LICENSE file in the project root for full license information.
 
#include <kcenon/thread/core/token_bucket.h>
 
#include <algorithm>
#include <thread>
 
namespace kcenon::thread
{
    token_bucket::token_bucket(std::size_t tokens_per_second, std::size_t burst_size)
        : tokens_(static_cast<std::int64_t>(burst_size) * PRECISION_FACTOR)
        , max_tokens_(static_cast<std::int64_t>(burst_size) * PRECISION_FACTOR)
        , refill_rate_(static_cast<double>(tokens_per_second) * PRECISION_FACTOR / 1e9)
        , last_refill_(std::chrono::steady_clock::now().time_since_epoch().count())
    {
    }
 
    auto token_bucket::refill() -> void
    {
        auto now = std::chrono::steady_clock::now().time_since_epoch().count();
        auto last = last_refill_.load(std::memory_order_acquire);
 
        // Calculate elapsed nanoseconds
        auto elapsed_ns = now - last;
        if (elapsed_ns <= 0)
        {
            return;  // No time passed, nothing to refill
        }
 
        // Try to update last_refill_ atomically
        if (!last_refill_.compare_exchange_weak(
                last, now,
                std::memory_order_acq_rel,
                std::memory_order_relaxed))
        {
            // Another thread updated, our calculation is stale
            return;
        }
 
        // Calculate tokens to add
        double rate = refill_rate_.load(std::memory_order_relaxed);
        auto new_tokens = static_cast<std::int64_t>(elapsed_ns * rate);
 
        if (new_tokens <= 0)
        {
            return;
        }
 
        // Add tokens (capped at max)
        std::int64_t max = max_tokens_.load(std::memory_order_relaxed);
        std::int64_t current = tokens_.load(std::memory_order_relaxed);
        std::int64_t updated = std::min(current + new_tokens, max);
 
        // Relaxed CAS is fine here since we're just accumulating
        tokens_.compare_exchange_weak(
            current, updated,
            std::memory_order_relaxed,
            std::memory_order_relaxed);
    }
 
    auto token_bucket::try_acquire(std::size_t tokens) -> bool
    {
        // First, refill based on elapsed time
        refill();
 
        // Scale requested tokens by precision factor
        std::int64_t needed = static_cast<std::int64_t>(tokens) * PRECISION_FACTOR;
 
        // CAS loop to atomically decrement tokens
        std::int64_t current = tokens_.load(std::memory_order_acquire);
        while (current >= needed)
        {
            if (tokens_.compare_exchange_weak(
                    current, current - needed,
                    std::memory_order_acq_rel,
                    std::memory_order_acquire))
            {
                return true;  // Successfully acquired
            }
            // CAS failed, current has been updated, retry
        }
 
        return false;  // Insufficient tokens
    }
 
    auto token_bucket::try_acquire_for(
        std::size_t tokens,
        std::chrono::milliseconds timeout) -> bool
    {
        auto deadline = std::chrono::steady_clock::now() + timeout;
 
        // Start with small backoff, increase exponentially
        auto backoff = std::chrono::microseconds{1};
        constexpr auto max_backoff = std::chrono::milliseconds{1};
 
        while (std::chrono::steady_clock::now() < deadline)
        {
            if (try_acquire(tokens))
            {
                return true;
            }
 
            // Sleep with exponential backoff
            std::this_thread::sleep_for(backoff);
 
            // Double backoff, cap at max
            backoff = std::min(
                backoff * 2,
                std::chrono::duration_cast<std::chrono::microseconds>(max_backoff));
        }
 
        // Final attempt after loop
        return try_acquire(tokens);
    }
 
    auto token_bucket::available_tokens() const -> std::size_t
    {
        // Need non-const refill, but we're returning a snapshot anyway
        const_cast<token_bucket*>(this)->refill();
 
        std::int64_t current = tokens_.load(std::memory_order_acquire);
        if (current <= 0)
        {
            return 0;
        }
        return static_cast<std::size_t>(current / PRECISION_FACTOR);
    }
 
    auto token_bucket::time_until_available(std::size_t tokens) const
        -> std::chrono::nanoseconds
    {
        const_cast<token_bucket*>(this)->refill();
 
        std::int64_t needed = static_cast<std::int64_t>(tokens) * PRECISION_FACTOR;
        std::int64_t current = tokens_.load(std::memory_order_acquire);
 
        if (current >= needed)
        {
            return std::chrono::nanoseconds{0};
        }
 
        // Calculate deficit
        std::int64_t deficit = needed - current;
 
        // Time = deficit / rate (rate is in tokens per nanosecond)
        double rate = refill_rate_.load(std::memory_order_relaxed);
        if (rate <= 0)
        {
            // Infinite wait if rate is zero
            return std::chrono::nanoseconds::max();
        }
 
        auto wait_ns = static_cast<std::int64_t>(deficit / rate);
        return std::chrono::nanoseconds{wait_ns};
    }
 
    auto token_bucket::set_rate(std::size_t tokens_per_second) -> void
    {
        // Refill with current rate before changing
        refill();
 
        // Update rate (tokens per nanosecond, scaled by precision)
        double new_rate = static_cast<double>(tokens_per_second) * PRECISION_FACTOR / 1e9;
        refill_rate_.store(new_rate, std::memory_order_release);
    }
 
    auto token_bucket::set_burst_size(std::size_t burst_size) -> void
    {
        std::int64_t new_max = static_cast<std::int64_t>(burst_size) * PRECISION_FACTOR;
        max_tokens_.store(new_max, std::memory_order_release);
 
        // Cap current tokens if they exceed new max
        std::int64_t current = tokens_.load(std::memory_order_acquire);
        while (current > new_max)
        {
            if (tokens_.compare_exchange_weak(
                    current, new_max,
                    std::memory_order_acq_rel,
                    std::memory_order_acquire))
            {
                break;
            }
        }
    }
 
    auto token_bucket::get_rate() const -> std::size_t
    {
        double rate = refill_rate_.load(std::memory_order_acquire);
        // Convert back: rate * 1e9 / PRECISION_FACTOR
        return static_cast<std::size_t>(rate * 1e9 / PRECISION_FACTOR);
    }
 
    auto token_bucket::get_burst_size() const -> std::size_t
    {
        std::int64_t max = max_tokens_.load(std::memory_order_acquire);
        return static_cast<std::size_t>(max / PRECISION_FACTOR);
    }
 
    auto token_bucket::reset() -> void
    {
        std::int64_t max = max_tokens_.load(std::memory_order_acquire);
        tokens_.store(max, std::memory_order_release);
        last_refill_.store(
            std::chrono::steady_clock::now().time_since_epoch().count(),
            std::memory_order_release);
    }
 
} // namespace kcenon::thread