From b9210f7f6dff1f235367e9a2e21bdef0cebe0a20 Mon Sep 17 00:00:00 2001
From: Arseny Smirnov <arseny30@gmail.com>
Date: Tue, 22 Nov 2022 15:25:53 +0100
Subject: [PATCH] Rewrite FoodControlFast: now it uses a bucket logic

---
 tdutils/td/utils/FloodControlFast.h | 89 ++++++++++++++++++-----------
 tdutils/test/misc.cpp               | 12 ++++
 2 files changed, 69 insertions(+), 32 deletions(-)

diff --git a/tdutils/td/utils/FloodControlFast.h b/tdutils/td/utils/FloodControlFast.h
index 2b7cf9930..22f206efd 100644
--- a/tdutils/td/utils/FloodControlFast.h
+++ b/tdutils/td/utils/FloodControlFast.h
@@ -11,52 +11,77 @@
 
 namespace td {
 
-class FloodControlFast {
+class FloodControlBucket {
  public:
-  void add_event(int32 now) {
-    for (auto &limit : limits_) {
-      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);
-      }
-    }
+  FloodControlBucket(double duration, double count)
+      : max_capacity_(count), speed_(count / duration), volume_(max_capacity_) {
   }
 
-  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_;
   }
 
-  void add_limit(uint32 duration, int32 count) {
-    limits_.push_back({TimedStat<CounterStat>(duration, 0), duration, count});
+  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) {
+    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() {
-    for (auto &limit : limits_) {
-      limit.stat_.clear_events();
+    for (auto &bucket : buckets_) {
+      bucket.clear_events();
     }
     wakeup_at_ = 0;
   }
 
  private:
-  class CounterStat {
-   public:
-    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_;
+  double wakeup_at_ = 0;
+  std::vector<FloodControlBucket> buckets_;
 };
 
 }  // namespace td
diff --git a/tdutils/test/misc.cpp b/tdutils/test/misc.cpp
index 6c5ae17d9..9decb679e 100644
--- a/tdutils/test/misc.cpp
+++ b/tdutils/test/misc.cpp
@@ -12,6 +12,7 @@
 #include "td/utils/CancellationToken.h"
 #include "td/utils/common.h"
 #include "td/utils/ExitGuard.h"
+#include "td/utils/FloodControlFast.h"
 #include "td/utils/Hash.h"
 #include "td/utils/HashMap.h"
 #include "td/utils/HashSet.h"
@@ -1244,3 +1245,14 @@ TEST(Misc, serialize) {
 TEST(Misc, check_reset_guard) {
   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;
+  }
+}