Improve nano::rate::token_bucket

nano/lib/rate_limiting.cpp

@@ -13,18 +13,19 @@ nano::rate::token_bucket::token_bucket (std::size_t max_token_count_a, std::size
 	reset (max_token_count_a, refill_rate_a);
 }
 
-bool nano::rate::token_bucket::try_consume (unsigned tokens_required_a)
+bool nano::rate::token_bucket::try_consume (unsigned tokens_required)
 {
-	debug_assert (tokens_required_a <= 1e9);
+	debug_assert (tokens_required <= 1e9);
+
 	refill ();
-	bool possible = current_size >= tokens_required_a;
+
+	// Keep track of largest observed bucket size so burst size can be computed (for tests and stats)
+	largest_size = std::max (largest_size, current_size);
+
+	bool possible = current_size >= tokens_required;
 	if (possible)
 	{
-		current_size -= tokens_required_a;
-	}
-	else if (tokens_required_a == 1e9)
-	{
-		current_size = 0;
+		current_size -= tokens_required;
 	}
 
 	// Keep track of smallest observed bucket size so burst size can be computed (for tests and stats)
@@ -35,32 +36,41 @@ bool nano::rate::token_bucket::try_consume (unsigned tokens_required_a)
 
 void nano::rate::token_bucket::refill ()
 {
-	auto now (std::chrono::steady_clock::now ());
+	auto now = 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
+	// Only update if there are 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 ();
+		last_refill = now;
 	}
 }
 
 void nano::rate::token_bucket::reset (std::size_t max_token_count_a, std::size_t refill_rate_a)
 {
 	// 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)
+	if (max_token_count_a == 0)
+	{
+		// Unlimited capacity
+		max_token_count_a = unlimited_rate_sentinel;
+	}
+	if (refill_rate_a == 0)
 	{
-		refill_rate_a = max_token_count_a = unlimited_rate_sentinel;
+		// Unlimited rate
+		refill_rate_a = unlimited_rate_sentinel;
 	}
-	max_token_count = smallest_size = current_size = max_token_count_a;
+
+	max_token_count = smallest_size = max_token_count_a;
 	refill_rate = refill_rate_a;
+	current_size = max_token_count < unlimited_rate_sentinel ? max_token_count : 0;
 	last_refill = std::chrono::steady_clock::now ();
 }
 
 std::size_t nano::rate::token_bucket::largest_burst () const
 {
-	return max_token_count - smallest_size;
+	debug_assert (largest_size >= smallest_size);
+	return largest_size - smallest_size;
 }
 
 std::size_t nano::rate::token_bucket::size () const

nano/lib/rate_limiting.hpp

@@ -25,8 +25,8 @@ class token_bucket
 public:
 	/**
 	 * Set up a token bucket.
-	 * @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)
+	 * @param max_token_count Maximum number of tokens in this bucket, which limits bursts. 0 is unlimited.
+	 * @param refill_rate Token refill rate, which limits the long term rate (tokens per seconds). 0 is unlimited (everything passes).
 	 */
 	token_bucket (std::size_t max_token_count, std::size_t refill_rate);
 
@@ -41,20 +41,21 @@ class token_bucket
 	/** Update the max_token_count and/or refill_rate_a parameters */
 	void reset (std::size_t max_token_count, std::size_t refill_rate);
 
+	/** Returns the current number of tokens in the bucket */
+	std::size_t size () const;
+
 	/** Returns the largest burst observed */
 	std::size_t largest_burst () const;
-	std::size_t size () const;
 
-private:
 	void refill ();
 
 private:
-	std::size_t max_token_count;
-	std::size_t refill_rate;
+	std::size_t max_token_count{ 0 };
+	std::size_t refill_rate{ 0 };
 
 	std::size_t current_size{ 0 };
-	/** The minimum observed bucket size, from which the largest burst can be derived */
-	std::size_t smallest_size{ 0 };
+	std::size_t smallest_size{ 0 }; // The minimum observed bucket size, from which the largest burst can be derived
+	std::size_t largest_size{ 0 }; // The largest observed bucket size, from which the largest burst can be derived
 	std::chrono::steady_clock::time_point last_refill;
 
 	static std::size_t constexpr unlimited_rate_sentinel{ static_cast<std::size_t> (1e9) };
@@ -70,8 +71,7 @@ class rate_limiter final
 	rate_limiter (std::size_t limit, double burst_ratio = 1.0);
 
 	bool should_pass (std::size_t buffer_size);
-	void reset (std::size_t limit, double burst_ratio = 1.0);
-
+	void reset (std::size_t limit, double burst_ratio);
 	std::size_t size () const;
 
 private: