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
10namespace 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 */
22template <typename Clock>
23class TokenBucket
24{
25public:
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    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    void increment(const time_point& now)
41    {
42        if (now > m_last_updated) {
43            if (m_value < m_cap && m_last_updated > MIN_TIME) {
44                double inc = m_rate * std::chrono::duration_cast<SecondsDouble>(now - m_last_updated).count();
45                m_value = std::min(m_cap, m_value + inc);
46            }
47        }
48        m_last_updated = now;
49    }
50
51    /** Consume n tokens. Returns false if the balance dropped to/below the given floor. */
52    bool decrement(double n = 1.0, double floor = 0.0)
53    {
54        m_value -= n;
55        return (m_value > floor);
56    }
57
58    /** Current token balance. */
59    double value() const { return m_value; }
60
61private:
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
70