token_bucket.hpp

Go to the documentation of this file.

#pragma once
 
/// @file userver/utils/token_bucket.hpp
/// @brief @copybrief utils::TokenBucket
 
#include <atomic>
#include <chrono>
 
USERVER_NAMESPACE_BEGIN
 
namespace utils {
 
/// @ingroup userver_universal userver_concurrency
///
/// Thread safe ratelimiter
class TokenBucket final {
 public:
  using TimePoint = std::chrono::steady_clock::time_point;
  using Duration = std::chrono::steady_clock::duration;
 
  /// Token bucket refill policy
  struct RefillPolicy {
    /// Refill amount (zero disables refills)
    size_t amount{1};
    /// Refill interval (zero makes bucket to instantly refill)
    Duration interval{Duration::max()};
  };
 
  /// Create an initially always empty token bucket
  TokenBucket() noexcept;
 
  /// Create a token bucket with max_size tokens and a specified refill policy
  TokenBucket(size_t max_size, RefillPolicy policy);
 
  /// Start with max_size tokens and add 1 token each
  /// single_token_update_interval up to max_size.
  /// Zero duration means "no limit".
  [[deprecated]] TokenBucket(size_t max_size,
                             Duration single_token_update_interval);
 
  /// Create an initially unbounded token bucket (largest size, instant refill)
  static TokenBucket MakeUnbounded() noexcept;
 
  TokenBucket(const TokenBucket&) = delete;
  TokenBucket(TokenBucket&&) noexcept;
  TokenBucket& operator=(const TokenBucket&) = delete;
  TokenBucket& operator=(TokenBucket&&) noexcept;
 
  bool IsUnbounded() const;
 
  /// Get current token limit (might be inaccurate as the result is stale)
  size_t GetMaxSizeApprox() const;
 
  /// Get current refill amount (might be inaccurate as the result is stale)
  size_t GetRefillAmountApprox() const;
 
  /// Get current refill interval (might be inaccurate as the result is stale)
  Duration GetRefillIntervalApprox() const;
 
  /// Get rate (tokens per second)
  double GetRatePs() const;
 
  /// Get current token count (might be inaccurate as the result is stale)
  size_t GetTokensApprox();
 
  /// Set max token count
  void SetMaxSize(size_t max_size);
 
  /// Add 1 token each token_update_interval.
  /// Zero duration means "no limit".
  [[deprecated]] void SetUpdateInterval(Duration single_token_update_interval);
 
  /// Set refill policy for the bucket
  void SetRefillPolicy(RefillPolicy policy);
 
  /// Set refill policy to "instant refill".
  ///
  /// Obtain does not deplete the bucket in this mode.
  /// Equivalent to `amount=1, interval=zero()` policy.
  void SetInstantRefillPolicy();
 
  /// @returns true if token was successfully obtained
  [[nodiscard]] bool Obtain();
 
  /// @return true if the requested number of tokens was successfully obtained
  [[nodiscard]] bool ObtainAll(size_t count);
 
  /// Get rate for specified update interval (updates per second)
  static double GetRatePs(Duration interval);
 
 private:
  void Update();
 
  std::atomic<size_t> max_size_;
  std::atomic<size_t> token_refill_amount_;
  std::atomic<Duration> token_refill_interval_;
  std::atomic<size_t> tokens_;
  std::atomic<TimePoint> last_update_;
};
 
}  // namespace utils
 
USERVER_NAMESPACE_END