From 987e875520c589ece276460e72a22c5b3cdb3f68 Mon Sep 17 00:00:00 2001
From: Arseny Smirnov <arseny30@gmail.com>
Date: Thu, 10 Feb 2022 18:26:11 +0100
Subject: [PATCH] FlatHashMap: support mem_stat->rss in memory benchmark; move
 memory benchmark in other folder, fix build.

---
 CMakeLists.txt                                |  2 +-
 benchmark/CMakeLists.txt                      | 17 +++++
 .../test => benchmark}/hashset_memory.cpp     | 62 +++++++++++++++----
 memprof/memprof.cpp                           | 10 +--
 tdutils/CMakeLists.txt                        |  4 --
 5 files changed, 72 insertions(+), 23 deletions(-)
 rename {tdutils/test => benchmark}/hashset_memory.cpp (70%)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 228dc8f7d..3bd72c996 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -744,7 +744,7 @@ if (MEMPROF)
   endif()
 endif()
 
-add_library(memprof_stat STATIC ${MEMPROF_STAT_SOURCE})
+add_library(memprof_stat EXCLUDE_FROM_ALL STATIC ${MEMPROF_STAT_SOURCE})
 target_include_directories(memprof_stat PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>)
 target_link_libraries(memprof_stat PRIVATE tdutils)
 
diff --git a/benchmark/CMakeLists.txt b/benchmark/CMakeLists.txt
index c4f17cae9..e8608df9b 100644
--- a/benchmark/CMakeLists.txt
+++ b/benchmark/CMakeLists.txt
@@ -72,3 +72,20 @@ if (NOT WIN32 AND NOT CYGWIN)
   add_executable(bench_queue bench_queue.cpp)
   target_link_libraries(bench_queue PRIVATE tdutils)
 endif()
+
+find_package(ABSL QUIET)
+find_package(benchmark QUIET)
+find_package(gflags QUIET)
+find_package(folly QUIET)
+
+if (ABSL_FOUND AND benchmark_FOUND AND gflags_FOUND AND folly_FOUND)
+  add_executable(memory-hashset-memprof EXCLUDE_FROM_ALL ${CMAKE_CURRENT_SOURCE_DIR}/hashset_memory.cpp)
+  target_compile_definitions(memory-hashset-memprof PRIVATE USE_MEMPROF=1)
+  target_link_libraries(memory-hashset-memprof PRIVATE tdutils memprof_stat)
+  target_link_libraries(memory-hashset-memprof PRIVATE SYSTEM Folly::folly absl::flat_hash_map absl::hash)
+
+  add_executable(memory-hashset-os ${CMAKE_CURRENT_SOURCE_DIR}/hashset_memory.cpp)
+  target_compile_definitions(memory-hashset-os PRIVATE USE_MEMPROF=0)
+  target_link_libraries(memory-hashset-os PRIVATE tdutils)
+  target_link_libraries(memory-hashset-os PRIVATE SYSTEM Folly::folly absl::flat_hash_map absl::hash)
+endif()
diff --git a/tdutils/test/hashset_memory.cpp b/benchmark/hashset_memory.cpp
similarity index 70%
rename from tdutils/test/hashset_memory.cpp
rename to benchmark/hashset_memory.cpp
index 55f82f626..a9ee5f9be 100644
--- a/tdutils/test/hashset_memory.cpp
+++ b/benchmark/hashset_memory.cpp
@@ -4,19 +4,43 @@
 // Distributed under the Boost Software License, Version 1.0. (See accompanying
 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
 //
-#include "memprof/memprof.h"
+
+#if USE_MEMPROF
+#include "memprof/memprof_stat.h"
+#endif
 
 #include "td/utils/check.h"
 #include "td/utils/FlatHashMap.h"
 #include "td/utils/format.h"
+#include "td/utils/misc.h"
+#include "td/utils/port/Stat.h"
 #include "td/utils/Slice.h"
-#include "td/utils/UInt.h"
 
 #include <absl/container/flat_hash_map.h>
 #include <folly/container/F14Map.h>
 #include <map>
 #include <unordered_map>
 
+int mem_stat_i = -1;
+int mem_stat_cur = 0;
+bool use_memprof() {
+  return mem_stat_i < 0
+#if USE_MEMPROF
+         && is_memprof_on()
+#endif
+      ;
+}
+auto get_memory() {
+#if USE_MEMPROF
+  if (use_memprof()) {
+    return get_used_memory_size();
+  }
+#else
+  CHECK(!use_memprof());
+  return td::mem_stat().ok().resident_size_;
+#endif
+};
+
 template <class T>
 class Generator {
  public:
@@ -64,7 +88,11 @@ class Generator<td::unique_ptr<T>> {
 
 template <class T, class KeyT, class ValueT>
 void measure(td::StringBuilder &sb, td::Slice name, td::Slice key_name, td::Slice value_name) {
-  sb << name << "<" << key_name << "," << value_name << ">:\n";
+  mem_stat_cur++;
+  if (mem_stat_i >= 0 && mem_stat_cur != mem_stat_i) {
+    return;
+  }
+  sb << name << "<" << key_name << "," << value_name << "> " << (use_memprof() ? "memprof" : "os") << "\n";
   size_t ideal_size = sizeof(KeyT) + sizeof(ValueT) + Generator<ValueT>::dyn_size();
 
   sb << "\tempty:" << sizeof(T);
@@ -78,12 +106,11 @@ void measure(td::StringBuilder &sb, td::Slice name, td::Slice key_name, td::Slic
   for (size_t size : {10000000u}) {
     Generator<KeyT> key_generator;
     Generator<ValueT> value_generator;
-    auto start_mem = get_used_memory_size();
+    auto start_mem = get_memory();
     T ht;
     auto ratio = [&]() {
-      auto end_mem = get_used_memory_size();
+      auto end_mem = get_memory();
       auto used_mem = end_mem - start_mem;
-      //return double(used_mem);
       return double(used_mem) / double(ideal_size * ht.size());
     };
     double min_ratio;
@@ -118,19 +145,19 @@ void measure(td::StringBuilder &sb, td::Slice name, td::Slice key_name, td::Slic
   sb << "\n";
 }
 
-
 template <size_t size>
 using Bytes = std::array<uint8_t, size>;
 
-template <template<typename... Args> class T>
+template <template <typename... Args> class T>
 void print_memory_stats(td::Slice name) {
   std::string big_buff(1 << 16, '\0');
   td::StringBuilder sb(big_buff, false);
-#define MEASURE(KeyT, ValueT) \
-  measure<T<KeyT, ValueT>, KeyT, ValueT>(sb, name, #KeyT, #ValueT);
+#define MEASURE(KeyT, ValueT) measure<T<KeyT, ValueT>, KeyT, ValueT>(sb, name, #KeyT, #ValueT);
   MEASURE(uint32_t, uint32_t);
   MEASURE(uint64_t, td::unique_ptr<Bytes<360>>);
-  LOG(ERROR) << "\n" << sb.as_cslice();
+  if (!sb.as_cslice().empty()) {
+    LOG(PLAIN) << "\n" << sb.as_cslice() << "\n";
+  }
 }
 
 #define FOR_EACH_TABLE(F) \
@@ -141,8 +168,17 @@ void print_memory_stats(td::Slice name) {
   F(std::map)
 #define BENCH_MEMORY(T) print_memory_stats<T>(#T);
 
-int main() {
-  CHECK(get_used_memory_size());
+int main(int argc, const char *argv[]) {
+  // Usage:
+  //  % benchmark/memory-hashset-os 0
+  //  max_i = 10
+  //  % for i in {1..10}; do ./benchmark/memory-hashset-os $i; done
+  if (argc > 1) {
+    mem_stat_i = td::to_integer<td::int32>(td::Slice(argv[1]));
+  }
   FOR_EACH_TABLE(BENCH_MEMORY);
+  if (mem_stat_i <= 0) {
+    LOG(PLAIN) << "max_i = " << mem_stat_cur << "\n";
+  }
   return 0;
 }
\ No newline at end of file
diff --git a/memprof/memprof.cpp b/memprof/memprof.cpp
index 519e3f320..d45d5e9b9 100644
--- a/memprof/memprof.cpp
+++ b/memprof/memprof.cpp
@@ -27,17 +27,17 @@ bool is_memprof_on() {
   return true;
 }
 
+#define my_assert(f) \
+  if (!(f)) {        \
+    std::abort();    \
+  }
+
 #if USE_MEMPROF_SAFE
 double get_fast_backtrace_success_rate() {
   return 0;
 }
 #else
 
-#define my_assert(f) \
-  if (!(f)) {        \
-    std::abort();    \
-  }
-
 #if TD_LINUX
 extern void *__libc_stack_end;
 #endif
diff --git a/tdutils/CMakeLists.txt b/tdutils/CMakeLists.txt
index cde74c676..aa366628d 100644
--- a/tdutils/CMakeLists.txt
+++ b/tdutils/CMakeLists.txt
@@ -395,8 +395,4 @@ if (ABSL_FOUND AND benchmark_FOUND AND gflags_FOUND AND folly_FOUND)
   add_executable(benchmark-hashset ${CMAKE_CURRENT_SOURCE_DIR}/test/hashset_benchmark.cpp)
   target_link_libraries(benchmark-hashset PRIVATE tdutils)
   target_link_libraries(benchmark-hashset PRIVATE SYSTEM benchmark::benchmark Folly::folly absl::flat_hash_map absl::hash)
-
-  add_executable(memory-hashset ${CMAKE_CURRENT_SOURCE_DIR}/test/hashset_memory.cpp)
-  target_link_libraries(memory-hashset PRIVATE tdutils memprof_stat)
-  target_link_libraries(memory-hashset PRIVATE SYSTEM Folly::folly absl::flat_hash_map absl::hash)
 endif()