Fix ratelimiter logic (#390)

Author: hhvrc

include/RateLimiter.h

@@ -4,6 +4,7 @@
 #include "SimpleMutex.h"
 
 #include <cstdint>
+#include <deque>
 #include <vector>
 
 namespace OpenShock {
@@ -33,6 +34,6 @@ namespace OpenShock {
     int64_t m_nextSlot;
     int64_t m_nextCleanup;
     std::vector<Limit> m_limits;
-    std::vector<int64_t> m_requests;
+    std::deque<int64_t> m_requests;
   };
 }  // namespace OpenShock

src/RateLimiter.cpp

@@ -23,24 +23,22 @@ OpenShock::RateLimiter::~RateLimiter()
 
 void OpenShock::RateLimiter::addLimit(uint32_t durationMs, uint16_t count)
 {
-  m_mutex.lock(portMAX_DELAY);
+  OpenShock::ScopedLock lock__(&m_mutex);
 
   // Insert sorted
   m_limits.insert(std::upper_bound(m_limits.begin(), m_limits.end(), durationMs, [](int64_t durationMs, const Limit& limit) { return durationMs < limit.durationMs; }), {durationMs, count});
 
   m_nextSlot    = 0;
   m_nextCleanup = 0;
-
-  m_mutex.unlock();
 }
 
 void OpenShock::RateLimiter::clearLimits()
 {
-  m_mutex.lock(portMAX_DELAY);
+  OpenShock::ScopedLock lock__(&m_mutex);
 
   m_limits.clear();
-
-  m_mutex.unlock();
+  m_nextSlot    = 0;
+  m_nextCleanup = 0;
 }
 
 bool OpenShock::RateLimiter::tryRequest()
@@ -57,40 +55,49 @@ bool OpenShock::RateLimiter::tryRequest()
     return true;
   }
 
+  // Cleanup based on longest limit
   if (m_nextCleanup <= now) {
     int64_t longestLimit = m_limits.back().durationMs;
     int64_t expiresAt    = now - longestLimit;
 
-    auto nextToExpire = std::find_if(m_requests.begin(), m_requests.end(), [expiresAt](int64_t requestedAtMs) { return requestedAtMs > expiresAt; });
-    if (nextToExpire != m_requests.end()) {
-      m_requests.erase(m_requests.begin(), nextToExpire);
-    }
+    // erase everything that’s expired
+    auto firstAlive = std::upper_bound(m_requests.begin(), m_requests.end(), expiresAt);
+    m_requests.erase(m_requests.begin(), firstAlive);
 
-    m_nextCleanup = m_requests.front() + longestLimit;
+    if (!m_requests.empty()) {
+      m_nextCleanup = m_requests.front() + longestLimit;
+    } else {
+      // nothing to clean until we add a new request
+      m_nextCleanup = now + longestLimit;
+    }
   }
 
   if (m_nextSlot > now) {
     return false;
   }
 
-  // Check if we've exceeded any limits, starting with the highest limit first
+  // Check if we've exceeded any limits, starting from the largest duration
   for (std::size_t i = m_limits.size(); i > 0;) {
     const auto& limit = m_limits[--i];
 
     // Calculate the window start time
     int64_t windowStart = now - limit.durationMs;
 
-    // Check how many requests are inside the limit window
+    // Check how many requests are inside the limit window and track earliest in-window element
     std::size_t insideWindow = 0;
-    for (int64_t request : m_requests) {
-      if (request > windowStart) {
-        insideWindow++;
-      }
+    auto it = m_requests.rbegin();
+    for (; it != m_requests.rend(); ++it) {
+      if (*it < windowStart) break;
+      ++insideWindow;
     }
 
     // If the window is full, set the wait time until its available, and reject the request
     if (insideWindow >= limit.count) {
-      m_nextSlot = m_requests.back() + limit.durationMs;
+      auto firstInWindow = (it == m_requests.rend())
+        ? m_requests.begin()
+        : it.base();
+
+      m_nextSlot = *firstInWindow + limit.durationMs;
       return false;
     }
   }
@@ -102,20 +109,17 @@ bool OpenShock::RateLimiter::tryRequest()
 }
 void OpenShock::RateLimiter::clearRequests()
 {
-  m_mutex.lock(portMAX_DELAY);
+  OpenShock::ScopedLock lock__(&m_mutex);
 
   m_requests.clear();
-
-  m_mutex.unlock();
+  m_nextSlot    = 0;
+  m_nextCleanup = 0;
 }
 
 void OpenShock::RateLimiter::blockFor(int64_t blockForMs)
 {
-  int64_t blockUntil = OpenShock::millis() + blockForMs;
-
-  m_mutex.lock(portMAX_DELAY);
+  OpenShock::ScopedLock lock__(&m_mutex);
 
+  int64_t blockUntil = OpenShock::millis() + blockForMs;
   m_nextSlot = std::max(m_nextSlot, blockUntil);
-
-  m_mutex.unlock();
 }