Branch data     Line data    Source code

       1                 :             : // Copyright (c) The Bitcoin Core developers
       2                 :             : // Distributed under the MIT software license, see the accompanying
       3                 :             : // file COPYING or http://www.opensource.org/licenses/mit-license.php.
       4                 :             : 
       5                 :             : #ifndef BITCOIN_UTIL_TOKENBUCKET_H
       6                 :             : #define BITCOIN_UTIL_TOKENBUCKET_H
       7                 :             : 
       8                 :             : #include <util/time.h>
       9                 :             : 
      10                 :             : namespace util {
      11                 :             : 
      12                 :             : /** A token bucket rate limiter.
      13                 :             :  *
      14                 :             :  * Tokens are added at a steady rate (m_rate per second) up to a capacity
      15                 :             :  * cap (m_cap). Tokens are removed by calling decrement(), which returns
      16                 :             :  * false if the bucket is emptied.
      17                 :             :  *
      18                 :             :  * Typical usage:
      19                 :             :  *   bucket.increment(now);   // refill based on elapsed time
      20                 :             :  *   if (bucket.value() >= 1) bucket.decrement(1);  // consume a token
      21                 :             :  */
      22                 :             : template <typename Clock>
      23                 :             : class TokenBucket
      24                 :             : {
      25                 :             : public:
      26                 :             :     using clock = Clock;
      27                 :             :     using time_point = typename Clock::time_point;
      28                 :             :     using duration = typename Clock::duration;
      29                 :             : 
      30                 :             :     const double m_rate{1};     //!< Tokens added per second
      31                 :             :     const double m_cap{0};      //!< Maximum token balance
      32                 :             : 
      33                 :             :     /** @param rate   Tokens added per second.
      34                 :             :      *  @param value  Initial token balance (clamped to cap).
      35                 :             :      *  @param cap    Maximum token balance. */
      36                 :        2512 :     TokenBucket(double rate, double value, double cap) : m_rate{rate}, m_cap{cap}, m_value{std::min(value, cap)} {}
      37                 :             : 
      38                 :             :     /** Refill tokens based on elapsed time since last call. No refill
      39                 :             :      *  occurs on the first call (establishes the time baseline). */
      40                 :      334960 :     void increment(const time_point& now)
      41                 :             :     {
      42         [ +  + ]:      334960 :         if (now > m_last_updated) {
      43   [ +  +  +  + ]:      333027 :             if (m_value < m_cap && m_last_updated > MIN_TIME) {
      44         [ +  + ]:      231700 :                 double inc = m_rate * std::chrono::duration_cast<SecondsDouble>(now - m_last_updated).count();
      45         [ +  + ]:      236697 :                 m_value = std::min(m_cap, m_value + inc);
      46                 :             :             }
      47                 :             :         }
      48                 :      334960 :         m_last_updated = now;
      49                 :      334960 :     }
      50                 :             : 
      51                 :             :     /** Consume n tokens. Returns false if the balance dropped to/below the given floor. */
      52                 :       46680 :     bool decrement(double n = 1.0, double floor = 0.0)
      53                 :             :     {
      54   [ +  -  +  -  :       46679 :         m_value -= n;
          +  -  +  -  +  
          -  +  -  +  -  
             +  -  +  - ]
      55         [ +  + ]:       46669 :         return (m_value > floor);
      56                 :             :     }
      57                 :             : 
      58                 :             :     /** Current token balance. */
      59   [ -  +  -  +  :       97889 :     double value() const { return m_value; }
          +  -  -  +  +  
             -  +  +  +  
           + ][ +  -  +  
          -  +  -  +  -  
          +  -  +  -  +  
          -  +  -  +  -  
          +  -  +  -  +  
          -  +  -  +  -  
          +  -  +  -  +  
                -  +  - ]
      60                 :             : 
      61                 :             : private:
      62                 :             :     static constexpr time_point MIN_TIME{time_point::min()};
      63                 :             :     time_point m_last_updated{MIN_TIME};
      64                 :             :     double m_value{0};
      65                 :             : };
      66                 :             : 
      67                 :             : } // namespace util
      68                 :             : 
      69                 :             : #endif // BITCOIN_UTIL_TOKENBUCKET_H