rate_limiter.h

Go to the documentation of this file.

// BSD 3-Clause License
// Copyright (c) 2021-2025, 🍀☀🌕🌥 🌊
// See the LICENSE file in the project root for full license information.
 
#pragma once
 
#include <chrono>
#include <mutex>
#include <shared_mutex>
#include <unordered_map>
#include <string>
#include <string_view>
#include <atomic>
#include <cstddef>
#include <cstdint>
 
namespace kcenon::network {
 
struct rate_limiter_config {
    size_t max_requests_per_second = 100;
 
    size_t burst_size = 20;
 
    std::chrono::seconds window = std::chrono::seconds(1);
 
    bool auto_cleanup = true;
 
    std::chrono::seconds stale_timeout = std::chrono::seconds(300);
};
 
class rate_limiter {
public:
    explicit rate_limiter(rate_limiter_config config = {})
        : config_(std::move(config))
        , last_cleanup_(std::chrono::steady_clock::now()) {}
 
    [[nodiscard]] bool allow(std::string_view client_id) {
        std::unique_lock<std::shared_mutex> lock(mutex_);
 
        auto now = std::chrono::steady_clock::now();
 
        // Periodic cleanup
        if (config_.auto_cleanup) {
            maybe_cleanup(now);
        }
 
        std::string key(client_id);
        auto& bucket = buckets_[key];
 
        // Initialize new bucket
        if (bucket.last_refill.time_since_epoch().count() == 0) {
            bucket.tokens = static_cast<double>(config_.burst_size);
            bucket.last_refill = now;
        }
 
        // Refill tokens based on elapsed time
        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
            now - bucket.last_refill);
 
        double tokens_to_add = elapsed.count() *
            (static_cast<double>(config_.max_requests_per_second) / 1000.0);
 
        bucket.tokens = std::min(
            static_cast<double>(config_.burst_size),
            bucket.tokens + tokens_to_add);
        bucket.last_refill = now;
 
        // Try to consume a token
        if (bucket.tokens >= 1.0) {
            bucket.tokens -= 1.0;
            return true;
        }
 
        return false;
    }
 
    [[nodiscard]] bool allow(std::string_view client_id, std::string_view session_id) {
        if (session_id.empty()) {
            return allow(client_id);
        }
        std::string composite_key;
        composite_key.reserve(client_id.size() + 1 + session_id.size());
        composite_key.append(client_id);
        composite_key.push_back(':');
        composite_key.append(session_id);
        return allow(std::string_view(composite_key));
    }
 
    [[nodiscard]] bool would_allow(std::string_view client_id) const {
        std::shared_lock<std::shared_mutex> lock(mutex_);
 
        auto it = buckets_.find(std::string(client_id));
        if (it == buckets_.end()) {
            return true; // New client would have full bucket
        }
 
        auto now = std::chrono::steady_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
            now - it->second.last_refill);
 
        double tokens_to_add = elapsed.count() *
            (static_cast<double>(config_.max_requests_per_second) / 1000.0);
 
        double available = std::min(
            static_cast<double>(config_.burst_size),
            it->second.tokens + tokens_to_add);
 
        return available >= 1.0;
    }
 
    [[nodiscard]] double remaining_tokens(std::string_view client_id) const {
        std::shared_lock<std::shared_mutex> lock(mutex_);
 
        auto it = buckets_.find(std::string(client_id));
        if (it == buckets_.end()) {
            return static_cast<double>(config_.burst_size);
        }
 
        auto now = std::chrono::steady_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
            now - it->second.last_refill);
 
        double tokens_to_add = elapsed.count() *
            (static_cast<double>(config_.max_requests_per_second) / 1000.0);
 
        return std::min(
            static_cast<double>(config_.burst_size),
            it->second.tokens + tokens_to_add);
    }
 
    void reset(std::string_view client_id) {
        std::unique_lock<std::shared_mutex> lock(mutex_);
        buckets_.erase(std::string(client_id));
    }
 
    void reset_all() {
        std::unique_lock<std::shared_mutex> lock(mutex_);
        buckets_.clear();
    }
 
    [[nodiscard]] size_t client_count() const {
        std::shared_lock<std::shared_mutex> lock(mutex_);
        return buckets_.size();
    }
 
    void set_config(rate_limiter_config config) {
        std::unique_lock<std::shared_mutex> lock(mutex_);
        config_ = std::move(config);
    }
 
    [[nodiscard]] rate_limiter_config config() const {
        std::shared_lock<std::shared_mutex> lock(mutex_);
        return config_;
    }
 
private:
    struct bucket {
        double tokens = 0.0;
        std::chrono::steady_clock::time_point last_refill{};
    };
 
    void maybe_cleanup(std::chrono::steady_clock::time_point now) {
        // Cleanup at most once per minute
        auto since_cleanup = std::chrono::duration_cast<std::chrono::seconds>(
            now - last_cleanup_);
 
        if (since_cleanup < std::chrono::seconds(60)) {
            return;
        }
 
        last_cleanup_ = now;
 
        // Remove stale entries
        for (auto it = buckets_.begin(); it != buckets_.end();) {
            auto since_refill = std::chrono::duration_cast<std::chrono::seconds>(
                now - it->second.last_refill);
 
            if (since_refill > config_.stale_timeout) {
                it = buckets_.erase(it);
            } else {
                ++it;
            }
        }
    }
 
    rate_limiter_config config_;
    mutable std::shared_mutex mutex_;
    std::unordered_map<std::string, bucket> buckets_;
    std::chrono::steady_clock::time_point last_cleanup_;
};
 
class connection_limiter {
public:
    explicit connection_limiter(size_t max_connections = 1000)
        : max_connections_(max_connections)
        , current_connections_(0) {}
 
    [[nodiscard]] bool can_accept() const noexcept {
        return current_connections_.load(std::memory_order_acquire) < max_connections_;
    }
 
    [[nodiscard]] bool try_accept() noexcept {
        size_t current = current_connections_.load(std::memory_order_acquire);
        while (current < max_connections_) {
            if (current_connections_.compare_exchange_weak(
                    current, current + 1,
                    std::memory_order_acq_rel,
                    std::memory_order_acquire)) {
                return true;
            }
        }
        return false;
    }
 
    void on_connect() noexcept {
        current_connections_.fetch_add(1, std::memory_order_release);
    }
 
    void on_disconnect() noexcept {
        size_t prev = current_connections_.fetch_sub(1, std::memory_order_release);
        // Prevent underflow (shouldn't happen in correct usage)
        if (prev == 0) {
            current_connections_.fetch_add(1, std::memory_order_release);
        }
    }
 
    [[nodiscard]] size_t current() const noexcept {
        return current_connections_.load(std::memory_order_acquire);
    }
 
    [[nodiscard]] size_t max() const noexcept {
        return max_connections_;
    }
 
    void set_max(size_t max_connections) noexcept {
        max_connections_ = max_connections;
    }
 
    [[nodiscard]] size_t available() const noexcept {
        size_t current = current_connections_.load(std::memory_order_acquire);
        return (current < max_connections_) ? (max_connections_ - current) : 0;
    }
 
    [[nodiscard]] bool at_capacity() const noexcept {
        return current_connections_.load(std::memory_order_acquire) >= max_connections_;
    }
 
private:
    size_t max_connections_;
    std::atomic<size_t> current_connections_;
};
 
class connection_guard {
public:
    explicit connection_guard(connection_limiter& limiter)
        : limiter_(&limiter)
        , accepted_(limiter.try_accept()) {}
 
    connection_guard(const connection_guard&) = delete;
    connection_guard& operator=(const connection_guard&) = delete;
 
    connection_guard(connection_guard&& other) noexcept
        : limiter_(other.limiter_)
        , accepted_(other.accepted_) {
        other.accepted_ = false;
    }
 
    connection_guard& operator=(connection_guard&& other) noexcept {
        if (this != &other) {
            release();
            limiter_ = other.limiter_;
            accepted_ = other.accepted_;
            other.accepted_ = false;
        }
        return *this;
    }
 
    ~connection_guard() {
        release();
    }
 
    [[nodiscard]] bool accepted() const noexcept {
        return accepted_;
    }
 
    explicit operator bool() const noexcept {
        return accepted_;
    }
 
    void release() noexcept {
        if (accepted_ && limiter_) {
            limiter_->on_disconnect();
            accepted_ = false;
        }
    }
 
private:
    connection_limiter* limiter_;
    bool accepted_;
};
 
class per_client_connection_limiter {
public:
    explicit per_client_connection_limiter(
            size_t max_per_client = 10,
            size_t max_total = 1000)
        : max_per_client_(max_per_client)
        , total_limiter_(max_total) {}
 
    [[nodiscard]] bool try_accept(std::string_view client_id) {
        // First check total limit
        if (!total_limiter_.can_accept()) {
            return false;
        }
 
        std::unique_lock<std::mutex> lock(mutex_);
        std::string key(client_id);
 
        auto& count = client_connections_[key];
        if (count >= max_per_client_) {
            return false;
        }
 
        if (!total_limiter_.try_accept()) {
            return false;
        }
 
        ++count;
        return true;
    }
 
    void release(std::string_view client_id) {
        std::unique_lock<std::mutex> lock(mutex_);
        std::string key(client_id);
 
        auto it = client_connections_.find(key);
        if (it != client_connections_.end() && it->second > 0) {
            --it->second;
            if (it->second == 0) {
                client_connections_.erase(it);
            }
            total_limiter_.on_disconnect();
        }
    }
 
    [[nodiscard]] size_t client_connections(std::string_view client_id) const {
        std::unique_lock<std::mutex> lock(mutex_);
        auto it = client_connections_.find(std::string(client_id));
        return (it != client_connections_.end()) ? it->second : 0;
    }
 
    [[nodiscard]] size_t total_connections() const noexcept {
        return total_limiter_.current();
    }
 
private:
    size_t max_per_client_;
    connection_limiter total_limiter_;
    mutable std::mutex mutex_;
    std::unordered_map<std::string, size_t> client_connections_;
};
 
} // namespace kcenon::network