diff --git a/include/cgimap/rate_limiter.hpp b/include/cgimap/rate_limiter.hpp index 96b9dafe..562e977c 100644 --- a/include/cgimap/rate_limiter.hpp +++ b/include/cgimap/rate_limiter.hpp @@ -42,7 +42,7 @@ class memcached_rate_limiter : public rate_limiter { void update(const std::string &key, int bytes, bool moderator) override; private: - memcached_st *ptr; + memcached_st *ptr = nullptr; struct state; }; diff --git a/src/rate_limiter.cpp b/src/rate_limiter.cpp index 78c7a69f..8b300a3a 100644 --- a/src/rate_limiter.cpp +++ b/src/rate_limiter.cpp @@ -8,8 +8,10 @@ */ #include +#include #include +#include "cgimap/logger.hpp" #include "cgimap/options.hpp" #include "cgimap/rate_limiter.hpp" @@ -28,21 +30,23 @@ struct memcached_rate_limiter::state { memcached_rate_limiter::memcached_rate_limiter( const boost::program_options::variables_map &options) { - if (options.count("memcache") && (ptr = memcached_create(nullptr)) != nullptr) { - memcached_server_st *server_list; + + if (!options.count("memcache")) + return; + + if ((ptr = memcached_create(nullptr)) != nullptr) { memcached_behavior_set(ptr, MEMCACHED_BEHAVIOR_NO_BLOCK, 1); memcached_behavior_set(ptr, MEMCACHED_BEHAVIOR_BINARY_PROTOCOL, 1); // memcached_behavior_set(ptr, MEMCACHED_BEHAVIOR_SUPPORT_CAS, 1); - server_list = - memcached_servers_parse(options["memcache"].as().c_str()); + const auto server = options["memcache"].as(); + memcached_server_st * server_list = memcached_servers_parse(server.c_str()); memcached_server_push(ptr, server_list); - memcached_server_list_free(server_list); - } else { - ptr = nullptr; + + logger::message(fmt::format("memcached rate limiting enabled ({})", server)); } } @@ -52,22 +56,22 @@ memcached_rate_limiter::~memcached_rate_limiter() { } std::tuple memcached_rate_limiter::check(const std::string &key, bool moderator) { + uint32_t bytes_served = 0; - std::string mc_key; state *sp; size_t length; uint32_t flags; - memcached_return error; + memcached_return_t error; - mc_key = "cgimap:" + key; - auto bytes_per_sec = global_settings::get_ratelimiter_ratelimit(moderator); + const auto mc_key = "cgimap:" + key; + const auto bytes_per_sec = global_settings::get_ratelimiter_ratelimit(moderator); if (ptr && (sp = (state *)memcached_get(ptr, mc_key.data(), mc_key.size(), &length, &flags, &error)) != nullptr) { assert(length == sizeof(state)); - int64_t elapsed = time(nullptr) - sp->last_update; + const int64_t elapsed = time(nullptr) - sp->last_update; if (elapsed * bytes_per_sec < sp->bytes_served) { bytes_served = sp->bytes_served - elapsed * bytes_per_sec; @@ -76,7 +80,7 @@ std::tuple memcached_rate_limiter::check(const std::string &key, bool free(sp); } - auto max_bytes = global_settings::get_ratelimiter_maxdebt(moderator); + const auto max_bytes = global_settings::get_ratelimiter_maxdebt(moderator); if (bytes_served < max_bytes) { return {false, 0}; } else { @@ -86,51 +90,63 @@ std::tuple memcached_rate_limiter::check(const std::string &key, bool } void memcached_rate_limiter::update(const std::string &key, int bytes, bool moderator) { - if (ptr) { - time_t now = time(nullptr); - std::string mc_key; - state *sp; - size_t length; - uint32_t flags; - memcached_return error; - mc_key = "cgimap:" + key; - auto bytes_per_sec = global_settings::get_ratelimiter_ratelimit(moderator); + if (!ptr) + return; - retry: + const auto now = time(nullptr); - if (ptr && - (sp = (state *)memcached_get(ptr, mc_key.data(), mc_key.size(), &length, - &flags, &error)) != nullptr) { - assert(length == sizeof(state)); + state *sp; + size_t length; + uint32_t flags; + memcached_return_t error; - int64_t elapsed = now - sp->last_update; + const auto mc_key = "cgimap:" + key; + const auto bytes_per_sec = global_settings::get_ratelimiter_ratelimit(moderator); - sp->last_update = now; + // upper limit in memcached for relative TTL values + // anything bigger is considered as absolute timestamp + constexpr auto REALTIME_MAXDELTA = 60L*60L*24L*30L; - if (elapsed * bytes_per_sec < sp->bytes_served) { - sp->bytes_served = sp->bytes_served - elapsed * bytes_per_sec + bytes; - } else { - sp->bytes_served = bytes; - } + // calculate number of seconds after which the memcached entry is guaranteed + // to be irrelevant (adding a bit of headroom). + const auto memcached_expiration = std::min(REALTIME_MAXDELTA, + 2L * global_settings::get_ratelimiter_maxdebt(moderator) / bytes_per_sec); - // should use CAS but it's a right pain so we'll wing it for now... - memcached_replace(ptr, mc_key.data(), mc_key.size(), (char *)sp, - sizeof(state), 0, 0); +retry: - free(sp); + if (!ptr) + return; + + if ((sp = (state *)memcached_get(ptr, mc_key.data(), mc_key.size(), &length, + &flags, &error)) != nullptr) { + assert(length == sizeof(state)); + + const int64_t elapsed = now - sp->last_update; + + sp->last_update = now; + + if (elapsed * bytes_per_sec < sp->bytes_served) { + sp->bytes_served = sp->bytes_served - elapsed * bytes_per_sec + bytes; } else { - state s; + sp->bytes_served = bytes; + } - s.last_update = now; - s.bytes_served = bytes; + // should use CAS but it's a right pain so we'll wing it for now... + auto rc = memcached_replace(ptr, mc_key.data(), mc_key.size(), (char *)sp, + sizeof(state), memcached_expiration, 0); + free(sp); - if (memcached_add(ptr, mc_key.data(), mc_key.size(), (char *)&s, - sizeof(state), 0, 0) == MEMCACHED_NOTSTORED) { - goto retry; - } + } else { + state s; + + s.last_update = now; + s.bytes_served = bytes; + + auto rc = memcached_add(ptr, mc_key.data(), mc_key.size(), (char *)&s, sizeof(state), memcached_expiration, 0); + + if (rc == MEMCACHED_NOTSTORED) { + goto retry; } } - - return; }