62 lines
1.6 KiB
C++
62 lines
1.6 KiB
C++
#ifndef RATE_LIMITER_HPP
|
|
#define RATE_LIMITER_HPP
|
|
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <chrono>
|
|
#include <mutex>
|
|
#include <queue>
|
|
|
|
namespace simple_object_storage {
|
|
|
|
class RateLimiter {
|
|
public:
|
|
RateLimiter(int max_requests, std::chrono::seconds window)
|
|
: max_requests_(max_requests), window_(window) {}
|
|
|
|
bool is_allowed(const std::string& key) {
|
|
std::lock_guard<std::mutex> lock(mutex_);
|
|
|
|
auto now = std::chrono::steady_clock::now();
|
|
auto& bucket = buckets_[key];
|
|
|
|
// Clean up old requests
|
|
while (!bucket.empty() && (now - bucket.front()) > window_) {
|
|
bucket.pop();
|
|
}
|
|
|
|
// Check if we're under the limit
|
|
if (bucket.size() < max_requests_) {
|
|
bucket.push(now);
|
|
return true;
|
|
}
|
|
|
|
return false;
|
|
}
|
|
|
|
// New method: check if over the limit WITHOUT incrementing
|
|
bool is_over_limit(const std::string& key) {
|
|
std::lock_guard<std::mutex> lock(mutex_);
|
|
auto now = std::chrono::steady_clock::now();
|
|
auto& bucket = buckets_[key];
|
|
while (!bucket.empty() && (now - bucket.front()) > window_) {
|
|
bucket.pop();
|
|
}
|
|
return bucket.size() >= max_requests_;
|
|
}
|
|
|
|
void reset(const std::string& key) {
|
|
std::lock_guard<std::mutex> lock(mutex_);
|
|
buckets_.erase(key);
|
|
}
|
|
|
|
private:
|
|
int max_requests_;
|
|
std::chrono::seconds window_;
|
|
std::unordered_map<std::string, std::queue<std::chrono::steady_clock::time_point>> buckets_;
|
|
std::mutex mutex_;
|
|
};
|
|
|
|
} // namespace simple_object_storage
|
|
|
|
#endif |