token_bucket.h

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.
 
#pragma once
 
#include <atomic>
#include <chrono>
#include <cstddef>
 
namespace kcenon::thread
{
    class token_bucket
    {
    public:
        token_bucket(std::size_t tokens_per_second, std::size_t burst_size);
 
        ~token_bucket() = default;
 
        // Non-copyable, non-movable for thread safety
        token_bucket(const token_bucket&) = delete;
        token_bucket& operator=(const token_bucket&) = delete;
        token_bucket(token_bucket&&) = delete;
        token_bucket& operator=(token_bucket&&) = delete;
 
        [[nodiscard]] auto try_acquire(std::size_t tokens = 1) -> bool;
 
        [[nodiscard]] auto try_acquire_for(
            std::size_t tokens,
            std::chrono::milliseconds timeout) -> bool;
 
        [[nodiscard]] auto available_tokens() const -> std::size_t;
 
        [[nodiscard]] auto time_until_available(std::size_t tokens) const
            -> std::chrono::nanoseconds;
 
        auto set_rate(std::size_t tokens_per_second) -> void;
 
        auto set_burst_size(std::size_t burst_size) -> void;
 
        [[nodiscard]] auto get_rate() const -> std::size_t;
 
        [[nodiscard]] auto get_burst_size() const -> std::size_t;
 
        auto reset() -> void;
 
    private:
        auto refill() -> void;
 
        std::atomic<std::int64_t> tokens_;
 
        std::atomic<std::int64_t> max_tokens_;
 
        std::atomic<double> refill_rate_;
 
        std::atomic<std::chrono::steady_clock::time_point::rep> last_refill_;
 
        static constexpr std::int64_t PRECISION_FACTOR = 1000;
    };
 
} // namespace kcenon::thread