diff --git a/tdutils/td/utils/FloodControlFast.h b/tdutils/td/utils/FloodControlFast.h index 22f206efd..59f32f55d 100644 --- a/tdutils/td/utils/FloodControlFast.h +++ b/tdutils/td/utils/FloodControlFast.h @@ -7,54 +7,9 @@ #pragma once #include "td/utils/common.h" -#include "td/utils/TimedStat.h" namespace td { -class FloodControlBucket { - public: - FloodControlBucket(double duration, double count) - : max_capacity_(count), speed_(count / duration), volume_(max_capacity_) { - } - - void add_event(double now, double size = 1) { - CHECK(now >= wakeup_at_); - update_volume(now); - if (volume_ >= size) { - volume_ -= size; - return; - } - size -= volume_; - volume_ = 0; - wakeup_at_ = volume_at_ + size / speed_; - volume_at_ = wakeup_at_; - } - double get_wakeup_at() const { - return wakeup_at_; - } - - void clear_events() { - volume_ = max_capacity_; - volume_at_ = 0; - wakeup_at_ = 0; - } - - private: - double max_capacity_{1}; - double speed_{1}; - double volume_{max_capacity_}; - - double volume_at_{0}; - double wakeup_at_{0}; - - void update_volume(double now) { - CHECK(now >= volume_at_); - auto passed = now - volume_at_; - volume_ = td::min(volume_ + passed * speed_, max_capacity_); - volume_at_ = now; - } -}; - class FloodControlFast { public: void add_event(double now) { @@ -80,8 +35,53 @@ class FloodControlFast { } private: + class FloodControlBucket { + public: + FloodControlBucket(double duration, double count) + : max_capacity_(count - 1), speed_(count / duration), volume_(max_capacity_) { + } + + void add_event(double now, double size = 1) { + CHECK(now >= wakeup_at_); + update_volume(now); + if (volume_ >= size) { + volume_ -= size; + return; + } + size -= volume_; + volume_ = 0; + wakeup_at_ = volume_at_ + size / speed_; + volume_at_ = wakeup_at_; + } + + double get_wakeup_at() const { + return wakeup_at_; + } + + void clear_events() { + volume_ = max_capacity_; + volume_at_ = 0; + wakeup_at_ = 0; + } + + private: + const double max_capacity_{1}; + const double speed_{1}; + double volume_{1}; + + double volume_at_{0}; + double wakeup_at_{0}; + + void update_volume(double now) { + CHECK(now >= volume_at_); + auto passed = now - volume_at_; + volume_ = td::min(volume_ + passed * speed_, max_capacity_); + volume_at_ = now; + } + }; + double wakeup_at_ = 0; - std::vector buckets_; + vector buckets_; }; } // namespace td diff --git a/tdutils/test/misc.cpp b/tdutils/test/misc.cpp index 9decb679e..c07cb6999 100644 --- a/tdutils/test/misc.cpp +++ b/tdutils/test/misc.cpp @@ -1249,10 +1249,13 @@ TEST(Misc, check_reset_guard) { TEST(FloodControl, Fast) { td::FloodControlFast fc; fc.add_limit(1, 5); + fc.add_limit(5, 10); + + td::int32 count = 0; double now = 0; for (int i = 0; i < 100; i++) { now = fc.get_wakeup_at(); fc.add_event(now); - LOG(INFO) << now; + LOG(INFO) << ++count << ": " << now; } }