From 38fc0c41a7e8df64ec52c900ab2aaa4daa497083 Mon Sep 17 00:00:00 2001 From: Fabian Terhorst Date: Sat, 27 Jul 2024 17:34:40 +0200 Subject: [PATCH] tweak(server): rate limiter improvements use eastl lru cache to limit the keys to 4000. --- .../include/KeyedRateLimiter.h | 129 ++++++++++++------ code/tests/EastlAllocator.cpp | 13 ++ 2 files changed, 103 insertions(+), 39 deletions(-) create mode 100644 code/tests/EastlAllocator.cpp diff --git a/code/components/citizen-server-instance/include/KeyedRateLimiter.h b/code/components/citizen-server-instance/include/KeyedRateLimiter.h index 6d5f7c7813..9c24bc44ae 100644 --- a/code/components/citizen-server-instance/include/KeyedRateLimiter.h +++ b/code/components/citizen-server-instance/include/KeyedRateLimiter.h @@ -3,15 +3,80 @@ #include #include -#include #include #include #include +#include "EASTL/bonus/lru_cache.h" + +namespace eastl +{ +template <> +struct hash +{ + size_t operator()(const net::PeerAddress& address) const + { + auto sockaddr = address.GetSocketAddress(); + + if (sockaddr->sa_family == AF_INET) + { + auto in = (sockaddr_in*)sockaddr; + return std::hash()(in->sin_addr.s_addr); + } + else if (sockaddr->sa_family == AF_INET6) + { + auto in6 = (sockaddr_in6*)sockaddr; + return std::hash()(std::string_view{ reinterpret_cast(&in6->sin6_addr), sizeof(in6->sin6_addr) }); + } + + return std::hash()(address.GetHost()); + } +}; + +template <> +struct hash +{ + size_t operator()(const std::string& str) const + { + return std::hash()(str); + } +}; +} + namespace fx { +// lru cache requires a default ctor for the value to compile +struct Bucket +{ + folly::TokenBucket bucket; + + Bucket(double genRate, double burstSize) + : bucket(genRate, burstSize) + { + } + + Bucket(): bucket {1, 1} + { + } + + void Reset(double genRate, double burstSize) + { + bucket.reset(genRate, burstSize); + } + + bool Consume(double toConsume) + { + return bucket.consume(toConsume); + } + + void ReturnTokens(double tokensToReturn) + { + bucket.returnTokens(tokensToReturn); + } +}; + template struct KeyMangler { @@ -47,7 +112,9 @@ template class KeyedRateLimiter { private: - using TBucket = folly::TokenBucket; + static constexpr size_t MAX_KEYS = 4000; + + using TBucket = Bucket; using TMangledKey = std::result_of_t(TKey)>; public: @@ -73,7 +140,7 @@ class KeyedRateLimiter for (auto& bucket : m_buckets) { - bucket.second.reset(genRate, burstSize); + bucket.second.first.Reset(genRate, burstSize); } } } @@ -84,14 +151,8 @@ class KeyedRateLimiter auto mangled = KeyMangler()(key); - auto it = m_buckets.find(mangled); - - if (it == m_buckets.end()) - { - it = m_buckets.emplace(mangled, TBucket{ m_genRate, m_burstSize }).first; - } - - it->second.returnTokens(n); + TBucket& bucket = m_buckets.get(mangled); + bucket.ReturnTokens(n); } bool Consume(const TKey& key, double n = 1.0, bool* isCooldown = nullptr) @@ -102,11 +163,11 @@ class KeyedRateLimiter if (Cooldown) { - auto cit = m_cooldowns.find(mangled); + eastl::optional cooldown = m_cooldowns.at(mangled); - if (cit != m_cooldowns.end()) + if (cooldown != eastl::nullopt) { - if (std::chrono::high_resolution_clock::now().time_since_epoch() <= cit->second) + if (std::chrono::high_resolution_clock::now().time_since_epoch() <= *cooldown) { if (isCooldown) { @@ -116,33 +177,21 @@ class KeyedRateLimiter return false; } - m_cooldowns.erase(cit); + m_cooldowns.erase(mangled); } } - auto it = m_buckets.find(mangled); - - if (it == m_buckets.end()) - { - it = m_buckets.emplace(mangled, TBucket{m_genRate, m_burstSize}).first; - } - - bool valid = it->second.consume(n); + TBucket& bucket = m_buckets.get(mangled); + bool valid = bucket.Consume(n); if (Cooldown && !valid) { - it = m_cooldownBuckets.find(mangled); - - if (it == m_cooldownBuckets.end()) - { - it = m_cooldownBuckets.emplace(mangled, TBucket{ m_genRate * 1.5, m_burstSize * 1.5 }).first; - } - - bool cooldownValid = it->second.consume(n); + bucket = m_cooldownBuckets.get(mangled); + bool cooldownValid = bucket.Consume(n); if (!cooldownValid) { - m_cooldowns[mangled] = std::chrono::duration_cast(std::chrono::high_resolution_clock::now().time_since_epoch() + std::chrono::seconds(15)); + m_cooldowns.insert_or_assign(mangled, std::chrono::duration_cast(std::chrono::high_resolution_clock::now().time_since_epoch() + std::chrono::seconds(15))); } } @@ -154,22 +203,24 @@ class KeyedRateLimiter std::unique_lock lock(m_mutex); auto mangled = KeyMangler()(key); - auto it = m_buckets.find(mangled); - if (it != m_buckets.end()) + if (m_buckets.contains(mangled)) { - it->second.reset(m_genRate, m_burstSize); + m_buckets.get(mangled).Reset(m_genRate, m_burstSize); } } - private: - std::unordered_map m_buckets; - std::unordered_map m_cooldownBuckets; - std::unordered_map m_cooldowns; std::mutex m_mutex; double m_genRate; double m_burstSize; + + eastl::function m_createTBucket = [this](TMangledKey) { return TBucket{ m_genRate, m_burstSize }; }; + eastl::function m_createTCooldownBucket = [this](TMangledKey) { return TBucket{ m_genRate * 1.5, m_burstSize * 1.5 }; }; + + eastl::lru_cache m_buckets {MAX_KEYS, eastl::allocator(EASTL_LRUCACHE_DEFAULT_NAME), m_createTBucket}; + eastl::lru_cache m_cooldownBuckets {MAX_KEYS, eastl::allocator(EASTL_LRUCACHE_DEFAULT_NAME), m_createTCooldownBucket}; + eastl::lru_cache m_cooldowns {MAX_KEYS}; }; struct RateLimiterDefaults diff --git a/code/tests/EastlAllocator.cpp b/code/tests/EastlAllocator.cpp new file mode 100644 index 0000000000..b32a56d311 --- /dev/null +++ b/code/tests/EastlAllocator.cpp @@ -0,0 +1,13 @@ +#include + +// both operators are required to be defined for eastl and the CitiTest module uses eastl in some tests. + +void* operator new[](size_t size, const char* pName, int flags, unsigned debugFlags, const char* file, int line) +{ + return ::operator new[](size); +} + +void* operator new[](size_t size, size_t alignment, size_t alignmentOffset, const char* pName, int flags, unsigned debugFlags, const char* file, int line) +{ + return ::operator new[](size); +}