From e715de9e1ce0b2b973dfcca8d7738de205f52500 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Thu, 30 Sep 2021 23:23:41 +0300
Subject: [PATCH] Count number of active clients with a given tqueue_id.

---
 telegram-bot-api/Client.cpp        | 13 ++++++++++---
 telegram-bot-api/ClientManager.cpp | 30 +++++++++++++++++++++++++-----
 telegram-bot-api/ClientManager.h   |  3 +++
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/telegram-bot-api/Client.cpp b/telegram-bot-api/Client.cpp
index d9e3b29..a845eb8 100644
--- a/telegram-bot-api/Client.cpp
+++ b/telegram-bot-api/Client.cpp
@@ -3499,7 +3499,7 @@ void Client::raw_event(const td::Event::Raw &event) {
 }
 
 void Client::loop() {
-  if (logging_out_ || closing_ || was_authorized_) {
+  if (was_authorized_ || logging_out_ || closing_) {
     while (!cmd_queue_.empty()) {
       auto query = std::move(cmd_queue_.front());
       cmd_queue_.pop();
@@ -4081,6 +4081,7 @@ void Client::on_update_authorization_state() {
       if (!was_authorized_) {
         LOG(WARNING) << "Logged in as @" << user_info->username;
         was_authorized_ = true;
+        td::send_event(parent_, td::Event::raw(static_cast<void *>(this)));
         update_shared_unix_time_difference();
         if (!pending_updates_.empty()) {
           LOG(INFO) << "Process " << pending_updates_.size() << " pending updates";
@@ -4096,14 +4097,20 @@ void Client::on_update_authorization_state() {
       if (!logging_out_) {
         LOG(WARNING) << "Logging out";
         logging_out_ = true;
+        if (was_authorized_ && !closing_) {
+          td::send_event(parent_, td::Event::raw(nullptr));
+        }
       }
-      break;
+      return loop();
     case td_api::authorizationStateClosing::ID:
       if (!closing_) {
         LOG(WARNING) << "Closing";
         closing_ = true;
+        if (was_authorized_ && !logging_out_) {
+          td::send_event(parent_, td::Event::raw(nullptr));
+        }
       }
-      break;
+      return loop();
     case td_api::authorizationStateClosed::ID:
       return on_closed();
     default:
diff --git a/telegram-bot-api/ClientManager.cpp b/telegram-bot-api/ClientManager.cpp
index d84eccf..c4730a9 100644
--- a/telegram-bot-api/ClientManager.cpp
+++ b/telegram-bot-api/ClientManager.cpp
@@ -112,13 +112,17 @@ void ClientManager::send(PromisedQueryPtr query) {
       }
       flood_control.add_event(static_cast<td::int32>(now));
     }
+    auto tqueue_id = get_tqueue_id(r_user_id.ok(), query->is_test_dc());
+    if (active_client_count_.find(tqueue_id) != active_client_count_.end()) {
+      // return query->set_retry_after_error(1);
+    }
 
-    auto id = clients_.create(ClientInfo{BotStatActor(stat_.actor_id(&stat_)), token, td::ActorOwn<Client>()});
+    auto id =
+        clients_.create(ClientInfo{BotStatActor(stat_.actor_id(&stat_)), token, tqueue_id, td::ActorOwn<Client>()});
     auto *client_info = clients_.get(id);
-    client_info->client_ =
-        td::create_actor<Client>(PSLICE() << "Client/" << token, actor_shared(this, id), query->token().str(),
-                                 query->is_test_dc(), get_tqueue_id(r_user_id.ok(), query->is_test_dc()), parameters_,
-                                 client_info->stat_.actor_id(&client_info->stat_));
+    client_info->client_ = td::create_actor<Client>(PSLICE() << "Client/" << token, actor_shared(this, id),
+                                                    query->token().str(), query->is_test_dc(), tqueue_id, parameters_,
+                                                    client_info->stat_.actor_id(&client_info->stat_));
 
     auto method = query->method();
     if (method != "deletewebhook" && method != "setwebhook") {
@@ -382,6 +386,21 @@ PromisedQueryPtr ClientManager::get_webhook_restore_query(td::Slice token, td::S
   return PromisedQueryPtr(query.release(), PromiseDeleter(td::PromiseActor<td::unique_ptr<Query>>()));
 }
 
+void ClientManager::raw_event(const td::Event::Raw &event) {
+  auto id = get_link_token();
+  auto *info = clients_.get(id);
+  CHECK(info != nullptr);
+  auto &value = active_client_count_[info->tqueue_id_];
+  if (event.ptr != nullptr) {
+    value++;
+  } else {
+    CHECK(value > 0);
+    if (--value == 0) {
+      active_client_count_.erase(info->tqueue_id_);
+    }
+  }
+}
+
 void ClientManager::hangup_shared() {
   auto id = get_link_token();
   auto *info = clients_.get(id);
@@ -391,6 +410,7 @@ void ClientManager::hangup_shared() {
   clients_.erase(id);
 
   if (close_flag_ && clients_.empty()) {
+    CHECK(active_client_count_.empty());
     close_db();
   }
 }
diff --git a/telegram-bot-api/ClientManager.h b/telegram-bot-api/ClientManager.h
index 8d0c375..e376490 100644
--- a/telegram-bot-api/ClientManager.h
+++ b/telegram-bot-api/ClientManager.h
@@ -52,6 +52,7 @@ class ClientManager final : public td::Actor {
    public:
     BotStatActor stat_;
     td::string token_;
+    td::int64 tqueue_id_;
     td::ActorOwn<Client> client_;
   };
   td::Container<ClientInfo> clients_;
@@ -62,6 +63,7 @@ class ClientManager final : public td::Actor {
 
   std::unordered_map<td::string, td::uint64> token_to_id_;
   std::unordered_map<td::string, td::FloodControlFast> flood_controls_;
+  std::unordered_map<td::int64, td::uint64> active_client_count_;
 
   bool close_flag_ = false;
   td::vector<td::Promise<td::Unit>> close_promises_;
@@ -72,6 +74,7 @@ class ClientManager final : public td::Actor {
                                                     std::shared_ptr<SharedData> shared_data);
 
   void start_up() override;
+  void raw_event(const td::Event::Raw &event) override;
   void hangup_shared() override;
   void close_db();
   void finish_close();