Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixes an issue where token_bucket::last_refill would update even if there were no tokens to add. #4103

Merged
merged 1 commit into from
Feb 6, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 21 additions & 2 deletions nano/core_test/utility.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,25 @@ TEST (rate, unlimited)
ASSERT_EQ (bucket.largest_burst (), static_cast<size_t> (1e9));
}

TEST (rate, busy_spin)
{
// Bucket should refill at a rate of 1 token per second
nano::rate::token_bucket bucket (1, 1);

// Run a very tight loop for 5 seconds + a bit of wiggle room
int counter = 0;
for (auto start = std::chrono::steady_clock::now (), now = start; now < start + std::chrono::milliseconds{ 5500 }; now = std::chrono::steady_clock::now ())
{
if (bucket.try_consume ())
{
++counter;
}
}

// Bucket starts fully refilled, therefore we see 1 additional request
ASSERT_EQ (counter, 6);
}

TEST (optional_ptr, basic)
{
struct valtype
Expand Down Expand Up @@ -164,7 +183,7 @@ TEST (thread_pool_alarm, one)
}
condition.notify_one ();
});
nano::unique_lock<nano::mutex> unique (mutex);
nano::unique_lock<nano::mutex> unique{ mutex };
condition.wait (unique, [&] () { return !!done; });
}

Expand All @@ -184,7 +203,7 @@ TEST (thread_pool_alarm, many)
condition.notify_one ();
});
}
nano::unique_lock<nano::mutex> unique (mutex);
nano::unique_lock<nano::mutex> unique{ mutex };
condition.wait (unique, [&] () { return count == 50; });
}

Expand Down
26 changes: 15 additions & 11 deletions nano/lib/rate_limiting.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,15 @@

#include <limits>

nano::rate::token_bucket::token_bucket (size_t max_token_count_a, size_t refill_rate_a)
nano::rate::token_bucket::token_bucket (std::size_t max_token_count_a, std::size_t refill_rate_a)
{
reset (max_token_count_a, refill_rate_a);
}

bool nano::rate::token_bucket::try_consume (unsigned tokens_required_a)
{
debug_assert (tokens_required_a <= 1e9);
nano::lock_guard<nano::mutex> lk (bucket_mutex);
nano::lock_guard<nano::mutex> guard{ mutex };
refill ();
bool possible = current_size >= tokens_required_a;
if (possible)
Expand All @@ -27,32 +27,36 @@ bool nano::rate::token_bucket::try_consume (unsigned tokens_required_a)
// Keep track of smallest observed bucket size so burst size can be computed (for tests and stats)
smallest_size = std::min (smallest_size, current_size);

return possible || refill_rate == 1e9;
return possible || refill_rate == unlimited_rate_sentinel;
}

void nano::rate::token_bucket::refill ()
{
auto now (std::chrono::steady_clock::now ());
auto tokens_to_add = static_cast<size_t> (std::chrono::duration_cast<std::chrono::nanoseconds> (now - last_refill).count () / 1e9 * refill_rate);
current_size = std::min (current_size + tokens_to_add, max_token_count);
last_refill = std::chrono::steady_clock::now ();
std::size_t tokens_to_add = static_cast<std::size_t> (std::chrono::duration_cast<std::chrono::nanoseconds> (now - last_refill).count () / 1e9 * refill_rate);
// Only update if there are any tokens to add
if (tokens_to_add > 0)
{
current_size = std::min (current_size + tokens_to_add, max_token_count);
last_refill = std::chrono::steady_clock::now ();
}
}

size_t nano::rate::token_bucket::largest_burst () const
std::size_t nano::rate::token_bucket::largest_burst () const
{
nano::lock_guard<nano::mutex> lk (bucket_mutex);
nano::lock_guard<nano::mutex> guard{ mutex };
return max_token_count - smallest_size;
}

void nano::rate::token_bucket::reset (size_t max_token_count_a, size_t refill_rate_a)
void nano::rate::token_bucket::reset (std::size_t max_token_count_a, std::size_t refill_rate_a)
{
nano::lock_guard<nano::mutex> lk (bucket_mutex);
nano::lock_guard<nano::mutex> guard{ mutex };

// A token count of 0 indicates unlimited capacity. We use 1e9 as
// a sentinel, allowing largest burst to still be computed.
if (max_token_count_a == 0 || refill_rate_a == 0)
{
refill_rate_a = max_token_count_a = static_cast<size_t> (1e9);
refill_rate_a = max_token_count_a = unlimited_rate_sentinel;
}
max_token_count = smallest_size = current_size = max_token_count_a;
refill_rate = refill_rate_a;
Expand Down
30 changes: 18 additions & 12 deletions nano/lib/rate_limiting.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,34 +25,40 @@ namespace rate
public:
/**
* Set up a token bucket.
* @param max_token_count_a Maximum number of tokens in this bucket, which limits bursts.
* @param refill_rate_a Token refill rate, which limits the long term rate (tokens per seconds)
* @param max_token_count Maximum number of tokens in this bucket, which limits bursts.
* @param refill_rate Token refill rate, which limits the long term rate (tokens per seconds)
*/
token_bucket (size_t max_token_count_a, size_t refill_rate_a);
token_bucket (std::size_t max_token_count, std::size_t refill_rate);

/**
* Determine if an operation of cost \p tokens_required_a is possible, and deduct from the
* Determine if an operation of cost \p tokens_required is possible, and deduct from the
* bucket if that's the case.
* The default cost is 1 token, but resource intensive operations may request
* more tokens to be available.
*/
bool try_consume (unsigned tokens_required_a = 1);
bool try_consume (unsigned tokens_required = 1);

/** Returns the largest burst observed */
size_t largest_burst () const;
std::size_t largest_burst () const;

/** Update the max_token_count and/or refill_rate_a parameters */
void reset (size_t max_token_count_a, size_t refill_rate_a);
void reset (std::size_t max_token_count, std::size_t refill_rate);

private:
void refill ();
size_t max_token_count;
size_t refill_rate;
size_t current_size{ 0 };

private:
std::size_t max_token_count;
std::size_t refill_rate;

std::size_t current_size{ 0 };
/** The minimum observed bucket size, from which the largest burst can be derived */
size_t smallest_size{ 0 };
std::size_t smallest_size{ 0 };
std::chrono::steady_clock::time_point last_refill;
mutable nano::mutex bucket_mutex;

mutable nano::mutex mutex;

static std::size_t constexpr unlimited_rate_sentinel{ static_cast<std::size_t> (1e9) };
};
}
}