diff --git a/src/test/util_tests.cpp b/src/test/util_tests.cpp index e26234f8696..15c54d412ce 100644 --- a/src/test/util_tests.cpp +++ b/src/test/util_tests.cpp @@ -25,6 +25,7 @@ #include #include #include +#include #include #include @@ -1930,4 +1931,135 @@ BOOST_AUTO_TEST_CASE(gib_string_literal_test) BOOST_CHECK_EQUAL(32_GiB, 32768_MiB); } +BOOST_AUTO_TEST_CASE(token_bucket_initial_value) +{ + // Initial value is clamped to cap + util::TokenBucket b1(/*rate=*/1, /*value=*/100, /*cap=*/10); + BOOST_CHECK_EQUAL(b1.value(), 10); + + // Initial value below cap is kept as-is + util::TokenBucket b2(/*rate=*/1, /*value=*/5, /*cap=*/10); + BOOST_CHECK_EQUAL(b2.value(), 5); +} + +BOOST_AUTO_TEST_CASE(token_bucket_first_increment) +{ + // First increment establishes the time baseline but does not refill + util::TokenBucket b(/*rate=*/100, /*value=*/0, /*cap=*/1000); + b.increment(NodeClock::time_point{10s}); + BOOST_CHECK_EQUAL(b.value(), 0); + + // Second increment refills based on elapsed time + b.increment(NodeClock::time_point{15s}); + BOOST_CHECK_EQUAL(b.value(), 500); // 100/s * 5s +} + +BOOST_AUTO_TEST_CASE(token_bucket_refill_caps) +{ + util::TokenBucket b(/*rate=*/10, /*value=*/90, /*cap=*/100); + b.increment(NodeClock::time_point{1s}); + b.increment(NodeClock::time_point{100s}); // would add 990, but cap is 100 + BOOST_CHECK_EQUAL(b.value(), 100); +} + +BOOST_AUTO_TEST_CASE(token_bucket_time_backwards) +{ + util::TokenBucket b(/*rate=*/10, /*value=*/50, /*cap=*/200); + b.increment(NodeClock::time_point{10s}); + b.increment(NodeClock::time_point{5s}); // backwards, no change + BOOST_CHECK_EQUAL(b.value(), 50); + b.increment(NodeClock::time_point{15s}); // forwards takes backwards into account + BOOST_CHECK_EQUAL(b.value(), 150); +} + +BOOST_AUTO_TEST_CASE(token_bucket_decrement_no_debt) +{ + // Default debt=0: returns false at exactly 0 + util::TokenBucket b(/*rate=*/1, /*value=*/3, /*cap=*/10); + BOOST_CHECK(b.decrement(1)); // 3 -> 2 + BOOST_CHECK(b.decrement(1)); // 2 -> 1 + BOOST_CHECK(!b.decrement(1)); // 1 -> 0, at floor + BOOST_CHECK_EQUAL(b.value(), 0); + BOOST_CHECK(!b.decrement(1)); // 0 -> -1, despite being at floor + BOOST_CHECK_EQUAL(b.value(), -1); +} + +BOOST_AUTO_TEST_CASE(token_bucket_decrement_with_debt) +{ + util::TokenBucket b(/*rate=*/1, /*value=*/2, /*cap=*/10); + BOOST_CHECK(b.decrement(1, -3)); // 2 -> 1 + BOOST_CHECK(b.decrement(1, -3)); // 1 -> 0 + BOOST_CHECK(b.decrement(1, -3)); // 0 -> -1, still above -3 + BOOST_CHECK(b.decrement(1, -3)); // -1 -> -2, still above -3 + BOOST_CHECK(!b.decrement(1, -3)); // -2 -> -3, at floor + BOOST_CHECK_EQUAL(b.value(), -3); +} + +BOOST_AUTO_TEST_CASE(token_bucket_drain_and_refill) +{ + util::TokenBucket b(/*rate=*/10, /*value=*/20, /*cap=*/100); + b.decrement(20); // drain to 0 + BOOST_CHECK_EQUAL(b.value(), 0); + + b.increment(NodeClock::time_point{1s}); + b.increment(NodeClock::time_point{4s}); // +30 + BOOST_CHECK_EQUAL(b.value(), 30); +} + + +BOOST_AUTO_TEST_CASE(token_bucket_first_increment_at_epoch) +{ + // The first increment establishes the baseline (no refill) even when it + // lands exactly on the clock epoch; later increments then refill normally. + util::TokenBucket b(/*rate=*/100, /*value=*/0, /*cap=*/1000); + b.increment(NodeClock::time_point{0s}); + BOOST_CHECK_EQUAL(b.value(), 0); + b.increment(NodeClock::time_point{5s}); + BOOST_CHECK_EQUAL(b.value(), 500); // 100/s * 5s +} + +BOOST_AUTO_TEST_CASE(token_bucket_at_cap_advances_baseline) +{ + util::TokenBucket b(/*rate=*/10, /*value=*/100, /*cap=*/100); + BOOST_CHECK_EQUAL(b.value(), 100); // already at cap + b.increment(NodeClock::time_point{1s}); // baseline established at 1s + b.increment(NodeClock::time_point{100s}); // 99s spent at the cap; baseline -> 100s + BOOST_CHECK_EQUAL(b.value(), 100); + + b.decrement(100); // drain to 0 + BOOST_CHECK_EQUAL(b.value(), 0); + + // refill doesn't "bank" the extra 99s we were at cap + b.increment(NodeClock::time_point{101s}); + BOOST_CHECK_EQUAL(b.value(), 10); + + // And when real time genuinely elapses, a single increment refills straight + // back to the cap immediately. + b.increment(NodeClock::time_point{200s}); // 99s elapsed -> +990, clamped to cap + BOOST_CHECK_EQUAL(b.value(), 100); +} + +BOOST_AUTO_TEST_CASE(token_bucket_fractional_refill) +{ + // Sub-second elapsed time accumulates fractional tokens via double math. + util::TokenBucket b(/*rate=*/10, /*value=*/0, /*cap=*/100); + b.increment(NodeClock::time_point{1s}); + b.increment(NodeClock::time_point{1250ms}); // 10/s * 0.25s = 2.5 + BOOST_CHECK_EQUAL(b.value(), 2.5); +} + +BOOST_AUTO_TEST_CASE(token_bucket_refill_from_debt) +{ + // Refilling from a negative (debt) balance accrues normally and still + // clamps to the cap rather than to debt + increment. + util::TokenBucket b(/*rate=*/10, /*value=*/0, /*cap=*/100); + BOOST_CHECK(!b.decrement(50)); // -> -50, below floor 0 + BOOST_CHECK_EQUAL(b.value(), -50); + b.increment(NodeClock::time_point{1s}); // baseline + b.increment(NodeClock::time_point{4s}); // +30 -> -20 + BOOST_CHECK_EQUAL(b.value(), -20); + b.increment(NodeClock::time_point{100s}); // +960 but clamped to cap + BOOST_CHECK_EQUAL(b.value(), 100); +} + BOOST_AUTO_TEST_SUITE_END() diff --git a/src/util/tokenbucket.h b/src/util/tokenbucket.h new file mode 100644 index 00000000000..ae8e8b53ad9 --- /dev/null +++ b/src/util/tokenbucket.h @@ -0,0 +1,69 @@ +// Copyright (c) The Bitcoin Core developers +// Distributed under the MIT software license, see the accompanying +// file COPYING or http://www.opensource.org/licenses/mit-license.php. + +#ifndef BITCOIN_UTIL_TOKENBUCKET_H +#define BITCOIN_UTIL_TOKENBUCKET_H + +#include + +namespace util { + +/** A token bucket rate limiter. + * + * Tokens are added at a steady rate (m_rate per second) up to a capacity + * cap (m_cap). Tokens are removed by calling decrement(), which returns + * false if the bucket is emptied. + * + * Typical usage: + * bucket.increment(now); // refill based on elapsed time + * if (bucket.value() >= 1) bucket.decrement(1); // consume a token + */ +template +class TokenBucket +{ +public: + using clock = Clock; + using time_point = typename Clock::time_point; + using duration = typename Clock::duration; + + const double m_rate{1}; //!< Tokens added per second + const double m_cap{0}; //!< Maximum token balance + + /** @param rate Tokens added per second. + * @param value Initial token balance (clamped to cap). + * @param cap Maximum token balance. */ + TokenBucket(double rate, double value, double cap) : m_rate{rate}, m_cap{cap}, m_value{std::min(value, cap)} {} + + /** Refill tokens based on elapsed time since last call. No refill + * occurs on the first call (establishes the time baseline). */ + void increment(const time_point& now) + { + if (now > m_last_updated) { + if (m_value < m_cap && m_last_updated > MIN_TIME) { + double inc = m_rate * std::chrono::duration_cast(now - m_last_updated).count(); + m_value = std::min(m_cap, m_value + inc); + } + } + m_last_updated = now; + } + + /** Consume n tokens. Returns false if the balance dropped to/below the given floor. */ + bool decrement(double n = 1.0, double floor = 0.0) + { + m_value -= n; + return (m_value > floor); + } + + /** Current token balance. */ + double value() const { return m_value; } + +private: + static constexpr time_point MIN_TIME{time_point::min()}; + time_point m_last_updated{MIN_TIME}; + double m_value{0}; +}; + +} // namespace util + +#endif // BITCOIN_UTIL_TOKENBUCKET_H