From a02588b16ffd5d243b87e0ef54c0c0c0f2e0d99f Mon Sep 17 00:00:00 2001
From: Zygo Blaxell <bees@furryterror.org>
Date: Tue, 10 Dec 2024 00:33:22 -0500
Subject: [PATCH] time: add more methods to support dynamic rate throttling

 * Allow RateLimiter to change rate after construction.
 * Check range of rate argument in constructor.
 * Atomic increment for RateEstimator.

Signed-off-by: Zygo Blaxell <bees@furryterror.org>
---
 include/crucible/time.h |  7 ++++++-
 lib/time.cc             | 27 +++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/include/crucible/time.h b/include/crucible/time.h
index cad7fe3..5aa195f 100644
--- a/include/crucible/time.h
+++ b/include/crucible/time.h
@@ -34,7 +34,7 @@ namespace crucible {
 		double	m_rate;
 		double	m_burst;
 		double  m_tokens = 0.0;
-		mutex	m_mutex;
+		mutable mutex m_mutex;
 
 		void update_tokens();
 		RateLimiter() = delete;
@@ -45,6 +45,8 @@ namespace crucible {
 		double sleep_time(double cost = 1.0);
 		bool is_ready();
 		void borrow(double cost = 1.0);
+		void rate(double new_rate);
+		double rate() const;
 	};
 
 	class RateEstimator {
@@ -88,6 +90,9 @@ namespace crucible {
 		// Read count
 		uint64_t count() const;
 
+		/// Increment count (like update(count() + more), but atomic)
+		void increment(uint64_t more = 1);
+
 		// Convert counts to chrono types
 		chrono::high_resolution_clock::time_point time_point(uint64_t absolute_count) const;
 		chrono::duration<double> duration(uint64_t relative_count) const;
diff --git a/lib/time.cc b/lib/time.cc
index 99001b6..4f991e5 100644
--- a/lib/time.cc
+++ b/lib/time.cc
@@ -98,12 +98,16 @@ namespace crucible {
 		m_rate(rate),
 		m_burst(burst)
 	{
+		THROW_CHECK1(invalid_argument, m_rate, m_rate > 0);
+		THROW_CHECK1(invalid_argument, m_burst, m_burst >= 0);
 	}
 
 	RateLimiter::RateLimiter(double rate) :
 		m_rate(rate),
 		m_burst(rate)
 	{
+		THROW_CHECK1(invalid_argument, m_rate, m_rate > 0);
+		THROW_CHECK1(invalid_argument, m_burst, m_burst >= 0);
 	}
 
 	void
@@ -119,6 +123,7 @@ namespace crucible {
 	double
 	RateLimiter::sleep_time(double cost)
 	{
+		THROW_CHECK1(invalid_argument, m_rate, m_rate > 0);
 		borrow(cost);
 		unique_lock<mutex> lock(m_mutex);
 		update_tokens();
@@ -154,6 +159,21 @@ namespace crucible {
 		m_tokens -= cost;
 	}
 
+	void
+	RateLimiter::rate(double const new_rate)
+	{
+		THROW_CHECK1(invalid_argument, new_rate, new_rate > 0);
+		unique_lock<mutex> lock(m_mutex);
+		m_rate = new_rate;
+	}
+
+	double
+	RateLimiter::rate() const
+	{
+		unique_lock<mutex> lock(m_mutex);
+		return m_rate;
+	}
+
 	RateEstimator::RateEstimator(double min_delay, double max_delay) :
 		m_min_delay(min_delay),
 		m_max_delay(max_delay)
@@ -202,6 +222,13 @@ namespace crucible {
 		}
 	}
 
+	void
+	RateEstimator::increment(const uint64_t more)
+	{
+		unique_lock<mutex> lock(m_mutex);
+		return update_unlocked(m_last_count + more);
+	}
+
 	uint64_t
 	RateEstimator::count() const
 	{