From fc59674153a4bca08573d09b77e065c35e443f10 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Piotr=20Wo=CC=81jcik?=
 <3044353+pwojcikdev@users.noreply.github.com>
Date: Mon, 4 Mar 2024 20:06:19 +0100
Subject: [PATCH] Single threaded rate limiter variant

---
 nano/lib/rate_limiting.cpp      |  4 ----
 nano/lib/rate_limiting.hpp      |  2 --
 nano/node/bandwidth_limiter.cpp | 29 +++++++++++++++++++++++++----
 nano/node/bandwidth_limiter.hpp | 20 ++++++++++++++++++--
 4 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/nano/lib/rate_limiting.cpp b/nano/lib/rate_limiting.cpp
index de6d0fde57..2b5d0d7bb4 100644
--- a/nano/lib/rate_limiting.cpp
+++ b/nano/lib/rate_limiting.cpp
@@ -12,7 +12,6 @@ nano::rate::token_bucket::token_bucket (std::size_t max_token_count_a, std::size
 bool nano::rate::token_bucket::try_consume (unsigned tokens_required_a)
 {
 	debug_assert (tokens_required_a <= 1e9);
-	nano::lock_guard<nano::mutex> guard{ mutex };
 	refill ();
 	bool possible = current_size >= tokens_required_a;
 	if (possible)
@@ -44,14 +43,11 @@ void nano::rate::token_bucket::refill ()
 
 std::size_t nano::rate::token_bucket::largest_burst () const
 {
-	nano::lock_guard<nano::mutex> guard{ mutex };
 	return max_token_count - smallest_size;
 }
 
 void nano::rate::token_bucket::reset (std::size_t max_token_count_a, std::size_t refill_rate_a)
 {
-	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)
diff --git a/nano/lib/rate_limiting.hpp b/nano/lib/rate_limiting.hpp
index 2c392963e4..fec74336f4 100644
--- a/nano/lib/rate_limiting.hpp
+++ b/nano/lib/rate_limiting.hpp
@@ -56,8 +56,6 @@ namespace rate
 		std::size_t smallest_size{ 0 };
 		std::chrono::steady_clock::time_point last_refill;
 
-		mutable nano::mutex mutex;
-
 		static std::size_t constexpr unlimited_rate_sentinel{ static_cast<std::size_t> (1e9) };
 	};
 }
diff --git a/nano/node/bandwidth_limiter.cpp b/nano/node/bandwidth_limiter.cpp
index 344cb6a129..e2e3f146f2 100644
--- a/nano/node/bandwidth_limiter.cpp
+++ b/nano/node/bandwidth_limiter.cpp
@@ -2,24 +2,45 @@
 #include <nano/node/bandwidth_limiter.hpp>
 
 /*
- * bandwidth_limiter
+ * bandwidth_limiter_st
  */
 
-nano::bandwidth_limiter::bandwidth_limiter (std::size_t limit_a, double burst_ratio_a) :
+nano::bandwidth_limiter_st::bandwidth_limiter_st (std::size_t limit_a, double burst_ratio_a) :
 	bucket (static_cast<std::size_t> (limit_a * burst_ratio_a), limit_a)
 {
 }
 
-bool nano::bandwidth_limiter::should_pass (std::size_t message_size_a)
+bool nano::bandwidth_limiter_st::should_pass (std::size_t message_size_a)
 {
 	return bucket.try_consume (nano::narrow_cast<unsigned int> (message_size_a));
 }
 
-void nano::bandwidth_limiter::reset (std::size_t limit_a, double burst_ratio_a)
+void nano::bandwidth_limiter_st::reset (std::size_t limit_a, double burst_ratio_a)
 {
 	bucket.reset (static_cast<std::size_t> (limit_a * burst_ratio_a), limit_a);
 }
 
+/*
+ * bandwidth_limiter_mt
+ */
+
+nano::bandwidth_limiter_mt::bandwidth_limiter_mt (std::size_t limit_a, double burst_ratio_a) :
+	bucket (static_cast<std::size_t> (limit_a * burst_ratio_a), limit_a)
+{
+}
+
+bool nano::bandwidth_limiter_mt::should_pass (std::size_t message_size_a)
+{
+	nano::lock_guard<nano::mutex> guard{ mutex };
+	return bucket.try_consume (nano::narrow_cast<unsigned int> (message_size_a));
+}
+
+void nano::bandwidth_limiter_mt::reset (std::size_t limit_a, double burst_ratio_a)
+{
+	nano::lock_guard<nano::mutex> guard{ mutex };
+	bucket.reset (static_cast<std::size_t> (limit_a * burst_ratio_a), limit_a);
+}
+
 /*
  * outbound_bandwidth_limiter
  */
diff --git a/nano/node/bandwidth_limiter.hpp b/nano/node/bandwidth_limiter.hpp
index 5cb13f0eba..365dfa1a5a 100644
--- a/nano/node/bandwidth_limiter.hpp
+++ b/nano/node/bandwidth_limiter.hpp
@@ -21,11 +21,11 @@ nano::bandwidth_limit_type to_bandwidth_limit_type (nano::transport::traffic_typ
 /**
  * Class that tracks and manages bandwidth limits for IO operations
  */
-class bandwidth_limiter final
+class bandwidth_limiter_st final // Single-threaded
 {
 public:
 	// initialize with limit 0 = unbounded
-	bandwidth_limiter (std::size_t limit, double burst_ratio);
+	bandwidth_limiter_st (std::size_t limit, double burst_ratio);
 
 	bool should_pass (std::size_t buffer_size);
 	void reset (std::size_t limit, double burst_ratio);
@@ -34,6 +34,22 @@ class bandwidth_limiter final
 	nano::rate::token_bucket bucket;
 };
 
+class bandwidth_limiter_mt final // Multi-threaded
+{
+public:
+	// initialize with limit 0 = unbounded
+	bandwidth_limiter_mt (std::size_t limit, double burst_ratio);
+
+	bool should_pass (std::size_t buffer_size);
+	void reset (std::size_t limit, double burst_ratio);
+
+private:
+	nano::rate::token_bucket bucket;
+	nano::mutex mutex;
+};
+
+using bandwidth_limiter = bandwidth_limiter_mt;
+
 class outbound_bandwidth_limiter final
 {
 public: // Config