From 5b6e2d2c90f96a3077145fd3e209e4bb72a67df2 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Wed, 18 Aug 2021 19:26:03 +0300
Subject: [PATCH] Add ThreadSafeCounter benchmark.

---
 benchmark/bench_misc.cpp             | 77 ++++++++++++++++++++++++++--
 tdutils/td/utils/ThreadSafeCounter.h |  4 ++
 2 files changed, 77 insertions(+), 4 deletions(-)

diff --git a/benchmark/bench_misc.cpp b/benchmark/bench_misc.cpp
index 270de6ea2..21aad77f7 100644
--- a/benchmark/bench_misc.cpp
+++ b/benchmark/bench_misc.cpp
@@ -16,6 +16,7 @@
 #include "td/utils/port/thread.h"
 #include "td/utils/Slice.h"
 #include "td/utils/SliceBuilder.h"
+#include "td/utils/ThreadSafeCounter.h"
 
 #include "td/telegram/telegram_api.h"
 #include "td/telegram/telegram_api.hpp"
@@ -313,7 +314,7 @@ class AtomicReleaseCasIncBench final : public td::Benchmark {
 template <int ThreadN>
 std::atomic<td::uint64> AtomicReleaseCasIncBench<ThreadN>::a_;
 
-template <int ThreadN = 2>
+template <int ThreadN>
 class RwMutexReadBench final : public td::Benchmark {
   td::string get_description() const final {
     return PSTRING() << "RwMutexRead" << ThreadN;
@@ -333,7 +334,8 @@ class RwMutexReadBench final : public td::Benchmark {
     }
   }
 };
-template <int ThreadN = 2>
+
+template <int ThreadN>
 class RwMutexWriteBench final : public td::Benchmark {
   td::string get_description() const final {
     return PSTRING() << "RwMutexWrite" << ThreadN;
@@ -353,19 +355,86 @@ class RwMutexWriteBench final : public td::Benchmark {
     }
   }
 };
+
+class ThreadSafeCounterBench final : public td::Benchmark {
+  static td::ThreadSafeCounter counter_;
+  int thread_count_;
+
+  td::string get_description() const final {
+    return PSTRING() << "ThreadSafeCounter" << thread_count_;
+  }
+  void run(int n) final {
+    counter_.clear();
+    td::vector<td::thread> threads;
+    for (int i = 0; i < thread_count_; i++) {
+      threads.emplace_back([n] {
+        for (int i = 0; i < n; i++) {
+          counter_.add(1);
+        }
+      });
+    }
+    for (auto &thread : threads) {
+      thread.join();
+    }
+    CHECK(counter_.sum() == n * thread_count_);
+  }
+
+ public:
+  explicit ThreadSafeCounterBench(int thread_count) : thread_count_(thread_count) {
+  }
+};
+td::ThreadSafeCounter ThreadSafeCounterBench::counter_;
+
+template <bool StrictOrder>
+class AtomicCounterBench final : public td::Benchmark {
+  static std::atomic<td::int64> counter_;
+  int thread_count_;
+
+  td::string get_description() const final {
+    return PSTRING() << "AtomicCounter" << thread_count_;
+  }
+  void run(int n) final {
+    counter_.store(0);
+    td::vector<td::thread> threads;
+    for (int i = 0; i < thread_count_; i++) {
+      threads.emplace_back([n] {
+        for (int i = 0; i < n; i++) {
+          counter_.fetch_add(1, StrictOrder ? std::memory_order_seq_cst : std::memory_order_relaxed);
+        }
+      });
+    }
+    for (auto &thread : threads) {
+      thread.join();
+    }
+    CHECK(counter_.load() == n * thread_count_);
+  }
+
+ public:
+  explicit AtomicCounterBench(int thread_count) : thread_count_(thread_count) {
+  }
+};
+template <bool StrictOrder>
+std::atomic<td::int64> AtomicCounterBench<StrictOrder>::counter_;
+
 #endif
 
 int main() {
   SET_VERBOSITY_LEVEL(VERBOSITY_NAME(DEBUG));
 #if !TD_THREAD_UNSUPPORTED
+  for (int i = 1; i <= 16; i *= 2) {
+    td::bench(ThreadSafeCounterBench(i));
+    td::bench(AtomicCounterBench<false>(i));
+    td::bench(AtomicCounterBench<true>(i));
+  }
+
   td::bench(AtomicReleaseIncBench<1>());
   td::bench(AtomicReleaseIncBench<2>());
   td::bench(AtomicReleaseCasIncBench<1>());
   td::bench(AtomicReleaseCasIncBench<2>());
   td::bench(RwMutexWriteBench<1>());
   td::bench(RwMutexReadBench<1>());
-  td::bench(RwMutexWriteBench<>());
-  td::bench(RwMutexReadBench<>());
+  td::bench(RwMutexWriteBench<2>());
+  td::bench(RwMutexReadBench<2>());
 #endif
 #if !TD_WINDOWS
   td::bench(UtimeBench());
diff --git a/tdutils/td/utils/ThreadSafeCounter.h b/tdutils/td/utils/ThreadSafeCounter.h
index 45b850286..b1557f29e 100644
--- a/tdutils/td/utils/ThreadSafeCounter.h
+++ b/tdutils/td/utils/ThreadSafeCounter.h
@@ -53,6 +53,10 @@ class ThreadSafeCounter {
     return counter_.sum(0);
   }
 
+  void clear() {
+    counter_.clear();
+  }
+
  private:
   ThreadSafeMultiCounter<1> counter_;
 };