Line data    Source code

       1             : #include "limiter.h"
       2             : 
       3             : #include <algorithm>
       4             : #include <cmath>
       5             : #include <numeric>
       6             : 
       7             : namespace datadog {
       8             : namespace tracing {
       9             : 
      10         631 : Limiter::Limiter(const Clock& clock, int max_tokens, double refresh_rate,
      11         631 :                  int tokens_per_refresh)
      12         631 :     : clock_(clock),
      13         631 :       num_tokens_(max_tokens),
      14         631 :       max_tokens_(max_tokens),
      15         631 :       tokens_per_refresh_(tokens_per_refresh),
      16         631 :       previous_rates_(9, 1.0) {
      17             :   // calculate refresh interval: (1/rate) * tokens per refresh as nanoseconds
      18         631 :   refresh_interval_ = std::chrono::duration_cast<std::chrono::nanoseconds>(
      19         631 :                           std::chrono::duration_cast<std::chrono::nanoseconds>(
      20           0 :                               std::chrono::seconds(1)) /
      21         631 :                           refresh_rate) *
      22        1262 :                       tokens_per_refresh_;
      23             : 
      24         631 :   auto now = clock_().tick;
      25         631 :   next_refresh_ = now + refresh_interval_;
      26         631 :   current_period_ = std::chrono::time_point_cast<std::chrono::seconds>(now);
      27         631 :   previous_rates_sum_ =
      28         631 :       std::accumulate(previous_rates_.begin(), previous_rates_.end(), 0.0);
      29         631 : }
      30             : 
      31         625 : Limiter::Limiter(const Clock& clock, double allowed_per_second)
      32         625 :     : Limiter(clock, int(std::ceil(allowed_per_second)), allowed_per_second,
      33         625 :               1) {}
      34             : 
      35       26374 : Limiter::Result Limiter::allow() { return allow(1); }
      36             : 
      37       26374 : Limiter::Result Limiter::allow(int tokens_requested) {
      38       26374 :   auto now = clock_().tick;
      39             : 
      40             :   // update effective rate calculations
      41       26374 :   auto intervals = std::chrono::duration_cast<std::chrono::seconds>(
      42           0 :                        std::chrono::time_point_cast<std::chrono::seconds>(now) -
      43       26374 :                        current_period_)
      44       26374 :                        .count();
      45       26374 :   if (intervals > 0) {
      46          11 :     if (std::size_t(intervals) >= previous_rates_.size()) {
      47           1 :       std::fill(previous_rates_.begin() + 1, previous_rates_.end(), 1.0);
      48             :     } else {
      49          20 :       std::move_backward(previous_rates_.begin(),
      50          10 :                          previous_rates_.end() - intervals,
      51             :                          previous_rates_.end());
      52          10 :       if (num_requested_ > 0) {
      53          10 :         previous_rates_[intervals - 1] =
      54          10 :             double(num_allowed_) / double(num_requested_);
      55             :       } else {
      56           0 :         previous_rates_[intervals - 1] = 1.0;
      57             :       }
      58          10 :       if (intervals - 2 > 0) {
      59           0 :         std::fill(previous_rates_.begin(),
      60           0 :                   previous_rates_.begin() + intervals - 2, 1.0);
      61             :       }
      62             :     }
      63          11 :     previous_rates_sum_ =
      64          11 :         std::accumulate(previous_rates_.begin(), previous_rates_.end(), 0.0);
      65          11 :     num_allowed_ = 0;
      66          11 :     num_requested_ = 0;
      67          11 :     current_period_ = now;
      68             :   }
      69             : 
      70       26374 :   num_requested_++;
      71             :   // refill "tokens"
      72       26374 :   if (now >= next_refresh_) {
      73             :     auto intervals =
      74        9195 :         (now - next_refresh_).count() / refresh_interval_.count() + 1;
      75        9195 :     if (intervals > 0) {
      76        9195 :       next_refresh_ += refresh_interval_ * intervals;
      77        9195 :       num_tokens_ += intervals * tokens_per_refresh_;
      78        9195 :       if (num_tokens_ > max_tokens_) {
      79         111 :         num_tokens_ = max_tokens_;
      80             :       }
      81             :     }
      82             :   }
      83             :   // determine if allowed or not
      84       26374 :   bool allowed = false;
      85       26374 :   if (num_tokens_ >= tokens_requested) {
      86       25272 :     allowed = true;
      87       25272 :     num_allowed_++;
      88       25272 :     num_tokens_ -= tokens_requested;
      89             :   }
      90             : 
      91             :   // `effective_rate` is guaranteed to be between 0.0 and 1.0.
      92             :   double effective_rate =
      93       52748 :       (previous_rates_sum_ + double(num_allowed_) / double(num_requested_)) /
      94       26374 :       (previous_rates_.size() + 1);
      95             : 
      96       26374 :   return {allowed, *Rate::from(effective_rate)};
      97             : }
      98             : 
      99             : }  // namespace tracing
     100             : }  // namespace datadog