Rewrite FoodControlFast: now it uses a bucket logic

This commit is contained in:
Arseny Smirnov 2022-11-22 15:25:53 +01:00
parent b7e4d567c7
commit b9210f7f6d
2 changed files with 69 additions and 32 deletions

View File

@ -11,52 +11,77 @@
namespace td { namespace td {
class FloodControlFast { class FloodControlBucket {
public: public:
void add_event(int32 now) { FloodControlBucket(double duration, double count)
for (auto &limit : limits_) { : max_capacity_(count), speed_(count / duration), volume_(max_capacity_) {
limit.stat_.add_event(CounterStat::Event(), now);
if (limit.stat_.get_stat(now).count_ > limit.count_) {
wakeup_at_ = max(wakeup_at_, now + limit.duration_ * 2);
}
}
} }
uint32 get_wakeup_at() const { 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_; return wakeup_at_;
} }
void add_limit(uint32 duration, int32 count) { void clear_events() {
limits_.push_back({TimedStat<CounterStat>(duration, 0), duration, count}); 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) {
for (auto &bucket : buckets_) {
bucket.add_event(now);
wakeup_at_ = td::max(wakeup_at_, bucket.get_wakeup_at());
}
}
double get_wakeup_at() const {
return wakeup_at_;
}
void add_limit(double duration, double count) {
buckets_.emplace_back(FloodControlBucket(duration, count));
} }
void clear_events() { void clear_events() {
for (auto &limit : limits_) { for (auto &bucket : buckets_) {
limit.stat_.clear_events(); bucket.clear_events();
} }
wakeup_at_ = 0; wakeup_at_ = 0;
} }
private: private:
class CounterStat { double wakeup_at_ = 0;
public: std::vector<FloodControlBucket> buckets_;
struct Event {};
int32 count_ = 0;
void on_event(Event e) {
count_++;
}
void clear() {
count_ = 0;
}
};
uint32 wakeup_at_ = 0;
struct Limit {
TimedStat<CounterStat> stat_;
uint32 duration_;
int32 count_;
};
std::vector<Limit> limits_;
}; };
} // namespace td } // namespace td

View File

@ -12,6 +12,7 @@
#include "td/utils/CancellationToken.h" #include "td/utils/CancellationToken.h"
#include "td/utils/common.h" #include "td/utils/common.h"
#include "td/utils/ExitGuard.h" #include "td/utils/ExitGuard.h"
#include "td/utils/FloodControlFast.h"
#include "td/utils/Hash.h" #include "td/utils/Hash.h"
#include "td/utils/HashMap.h" #include "td/utils/HashMap.h"
#include "td/utils/HashSet.h" #include "td/utils/HashSet.h"
@ -1244,3 +1245,14 @@ TEST(Misc, serialize) {
TEST(Misc, check_reset_guard) { TEST(Misc, check_reset_guard) {
CheckExitGuard check_exit_guard{false}; CheckExitGuard check_exit_guard{false};
} }
TEST(FloodControl, Fast) {
td::FloodControlFast fc;
fc.add_limit(1, 5);
double now = 0;
for (int i = 0; i < 100; i++) {
now = fc.get_wakeup_at();
fc.add_event(now);
LOG(INFO) << now;
}
}