HttpRateLimiter.hpp

Go to the documentation of this file.

#pragma once
#ifndef _KURLYK_HTTP_RATE_LIMITER_HPP_INCLUDED
#define _KURLYK_HTTP_RATE_LIMITER_HPP_INCLUDED

 
namespace kurlyk {

    class HttpRateLimiter {
    public:

        long create_limit(long requests_per_period, long period_ms) {
            std::lock_guard<std::mutex> lock(m_mutex);
            long id = m_next_id++;
            m_limits[id] = LimitData{
                requests_per_period,
                period_ms,
                0,
                std::chrono::steady_clock::now()
            };
            return id;
        }

        bool remove_limit(long limit_id) {
            std::lock_guard<std::mutex> lock(m_mutex);
            return m_limits.erase(limit_id) > 0;
        }

        bool allow_request(long general_rate_limit_id, long specific_rate_limit_id) {
            std::lock_guard<std::mutex> lock(m_mutex);
            auto general_it = m_limits.find(general_rate_limit_id);
            auto specific_it = m_limits.find(specific_rate_limit_id);
            if (general_it == m_limits.end() &&
                specific_it == m_limits.end()) return true;
 
            auto now = std::chrono::steady_clock::now();
            bool general_limit_allowed = true;
            bool specific_limit_allowed = true;
 
            if (general_it != m_limits.end()) {
                auto& limit_data = general_it->second;
                general_limit_allowed = check_limit(limit_data, now);
            }
 
            if (specific_it != m_limits.end()) {
                auto& limit_data = specific_it->second;
                specific_limit_allowed = check_limit(limit_data, now);
            }
 
            if (general_limit_allowed && specific_limit_allowed) {
                // All limits allow the request, now update counts
                if (general_it != m_limits.end()) {
                    auto& limit_data = general_it->second;
                    update_limit(limit_data, now);
                }
                if (specific_it != m_limits.end()) {
                    auto& limit_data = specific_it->second;
                    update_limit(limit_data, now);
                }
                return true;
            }
 
            // Request is not allowed under one or more limits
            return false;
        }

        template<typename Duration = std::chrono::milliseconds>
        Duration time_until_next_allowed(long general_rate_limit_id, long specific_rate_limit_id) {
            std::lock_guard<std::mutex> lock(m_mutex);
            auto now = std::chrono::steady_clock::now();
            Duration max_delay{0};
 
            auto it = m_limits.find(general_rate_limit_id);
            if (it != m_limits.end()) {
                max_delay = std::max(max_delay, time_until_limit_allows<Duration>(it->second, now));
            }
 
            it = m_limits.find(specific_rate_limit_id);
            if (it != m_limits.end()) {
                max_delay = std::max(max_delay, time_until_limit_allows<Duration>(it->second, now));
            }
 
            return max_delay;
        }

        template<typename Duration = std::chrono::milliseconds>
        Duration time_until_any_limit_allows() {
            std::lock_guard<std::mutex> lock(m_mutex);
            auto now = std::chrono::steady_clock::now();
 
            Duration min_delay = Duration::max();
 
            for (const auto& pair : m_limits) {
                const auto& limit = pair.second;
                Duration delay = time_until_limit_allows<Duration>(limit, now);
                if (delay.count() > 0 && delay < min_delay) {
                    min_delay = delay;
                }
            }
 
            return (min_delay == Duration::max()) ? Duration{0} : min_delay;
        }
 
    private:
        using time_point_t = std::chrono::steady_clock::time_point;

        struct LimitData {
            long requests_per_period;   
            long period_ms;             
            long count;                 
            time_point_t start_time;    
        };
 
        std::mutex m_mutex;             
        long m_next_id = 1;             
        std::unordered_map<long, LimitData> m_limits; 

        bool check_limit(LimitData& limit_data, const time_point_t& now) {
            auto elapsed_time = std::chrono::duration_cast<std::chrono::milliseconds>(now - limit_data.start_time);
 
            // If period has elapsed, the count will reset upon update
            if (elapsed_time.count() >= limit_data.period_ms) {
                return true; // Period elapsed, limit can be reset
            }
 
            // Check if under limit or if no limit is set (0 requests per period means unlimited)
            if (limit_data.count < limit_data.requests_per_period ||
                limit_data.requests_per_period == 0) {
                return true;
            }
            return false;
        }

        void update_limit(LimitData& limit_data, const time_point_t& now) {
            auto elapsed_time = std::chrono::duration_cast<std::chrono::milliseconds>(now - limit_data.start_time);
 
            // Reset count if period has elapsed
            if (elapsed_time.count() >= limit_data.period_ms) {
                limit_data.start_time = now;
                limit_data.count = 0;
            }
 
            ++limit_data.count;
        }

        template<typename Duration>
        Duration time_until_limit_allows(const LimitData& limit_data, const time_point_t& now) const {
            if (limit_data.requests_per_period == 0) return Duration{0};
            auto elapsed = std::chrono::duration_cast<Duration>(now - limit_data.start_time);
            auto period_duration = std::chrono::duration_cast<Duration>(std::chrono::milliseconds(limit_data.period_ms));
            if (elapsed >= period_duration ||
                limit_data.count < limit_data.requests_per_period) {
                return Duration{0};
            }
            return period_duration - elapsed;
        }
    }; // HttpRateLimiter
 
} // namespace kurlyk
 
#endif // _KURLYK_HTTP_RATE_LIMITER_HPP_INCLUDED