From f9726a5e19a0df20d10a6d85abeb2c00fb8e61d0 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Tue, 17 Jul 2018 02:46:44 +0300
Subject: [PATCH] Remove global restriction on thread count.

GitOrigin-RevId: 48a3a6dce014fa640a5bb750eacc27304cb550d1
---
 tdactor/td/actor/impl2/Scheduler.h            | 10 ++++++++
 tdutils/td/utils/HazardPointers.h             |  1 +
 .../td/utils/port/detail/ThreadIdGuard.cpp    | 24 +++++++++----------
 tdutils/td/utils/port/thread_local.h          |  4 ----
 4 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/tdactor/td/actor/impl2/Scheduler.h b/tdactor/td/actor/impl2/Scheduler.h
index 9d5783b1..eafc7317 100644
--- a/tdactor/td/actor/impl2/Scheduler.h
+++ b/tdactor/td/actor/impl2/Scheduler.h
@@ -780,6 +780,16 @@ class SchedulerContext
 
 #if !TD_THREAD_UNSUPPORTED
 class Scheduler {
+  static constexpr int32 max_thread_count() {
+    return 256;
+  }
+
+  static int32 get_thread_id() {
+    auto thread_id = ::td::get_thread_id();
+    CHECK(thread_id < max_thread_count());
+    return thread_id;
+  }
+
  public:
   Scheduler(std::shared_ptr<SchedulerGroupInfo> scheduler_group_info, SchedulerId id, size_t cpu_threads_count)
       : scheduler_group_info_(std::move(scheduler_group_info)), cpu_threads_(cpu_threads_count) {
diff --git a/tdutils/td/utils/HazardPointers.h b/tdutils/td/utils/HazardPointers.h
index e13dc802..7e6f27a5 100644
--- a/tdutils/td/utils/HazardPointers.h
+++ b/tdutils/td/utils/HazardPointers.h
@@ -126,6 +126,7 @@ class HazardPointers {
   }
 
   std::atomic<T *> &get_hazard_ptr(size_t thread_id, size_t pos) {
+    CHECK(thread_id < threads_.size());
     return threads_[thread_id].hazard[pos];
   }
 };
diff --git a/tdutils/td/utils/port/detail/ThreadIdGuard.cpp b/tdutils/td/utils/port/detail/ThreadIdGuard.cpp
index d949945e..cacc4610 100644
--- a/tdutils/td/utils/port/detail/ThreadIdGuard.cpp
+++ b/tdutils/td/utils/port/detail/ThreadIdGuard.cpp
@@ -9,8 +9,8 @@
 #include "td/utils/logging.h"
 #include "td/utils/port/thread_local.h"
 
-#include <array>
 #include <mutex>
+#include <set>
 
 namespace td {
 namespace detail {
@@ -18,25 +18,25 @@ class ThreadIdManager {
  public:
   int32 register_thread() {
     std::lock_guard<std::mutex> guard(mutex_);
-    for (size_t i = 0; i < is_id_used_.size(); i++) {
-      if (!is_id_used_[i]) {
-        is_id_used_[i] = true;
-        return static_cast<int32>(i + 1);
-      }
+    if (unused_thread_ids_.empty()) {
+      return ++max_thread_id_;
     }
-    LOG(FATAL) << "Cannot create more than " << max_thread_count() << " threads";
-    return 0;
+    auto it = unused_thread_ids_.begin();
+    auto result = *it;
+    unused_thread_ids_.erase(it);
+    return result;
   }
   void unregister_thread(int32 thread_id) {
-    thread_id--;
     std::lock_guard<std::mutex> guard(mutex_);
-    CHECK(is_id_used_.at(thread_id));
-    is_id_used_[thread_id] = false;
+    CHECK(0 < thread_id && thread_id <= max_thread_id_);
+    bool is_inserted = unused_thread_ids_.insert(thread_id).second;
+    CHECK(is_inserted);
   }
 
  private:
   std::mutex mutex_;
-  std::array<bool, max_thread_count()> is_id_used_{{false}};
+  std::set<int32> unused_thread_ids_;
+  int32 max_thread_id_ = 0;
 };
 static ThreadIdManager thread_id_manager;
 
diff --git a/tdutils/td/utils/port/thread_local.h b/tdutils/td/utils/port/thread_local.h
index 6d8c135e..264650e8 100644
--- a/tdutils/td/utils/port/thread_local.h
+++ b/tdutils/td/utils/port/thread_local.h
@@ -27,10 +27,6 @@ namespace td {
 #endif
 // clang-format on
 
-inline constexpr size_t max_thread_count() {
-  return 256;
-}
-
 // If raw_ptr is not nullptr, allocate T as in std::make_unique<T>(args...) and store pointer into raw_ptr
 template <class T, class P, class... ArgsT>
 bool init_thread_local(P &raw_ptr, ArgsT &&... args);