Add WaitFreeHashMap.

tdutils/CMakeLists.txt

@@ -289,6 +289,7 @@ set(TDUTILS_SOURCE
   td/utils/utf8.h
   td/utils/Variant.h
   td/utils/VectorQueue.h
+  td/utils/WaitFreeHashMap.h
 )
 
 if (TDUTILS_MIME_TYPE)
@@ -329,6 +330,7 @@ set(TDUTILS_TEST_SOURCE
   ${CMAKE_CURRENT_SOURCE_DIR}/test/SharedSlice.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/test/StealingQueue.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/test/variant.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/test/WaitFreeHashMap.cpp
   PARENT_SCOPE
 )
 

tdutils/td/utils/WaitFreeHashMap.h

@@ -0,0 +1,66 @@
+//
+// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2022
+//
+// 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)
+//
+#pragma once
+
+#include "td/utils/common.h"
+#include "td/utils/FlatHashMap.h"
+#include "td/utils/HashTableUtils.h"
+
+#include <functional>
+
+namespace td {
+
+template <class KeyT, class ValueT, class HashT = std::hash<KeyT>, class EqT = std::equal_to<KeyT>>
+class WaitFreeHashMap {
+  using Storage = FlatHashMap<KeyT, ValueT, HashT, EqT>;
+  static constexpr size_t MAX_STORAGE_COUNT = 256;
+  static_assert((MAX_STORAGE_COUNT & (MAX_STORAGE_COUNT - 1)) == 0, "");
+  static constexpr size_t MAX_STORAGE_SIZE = MAX_STORAGE_COUNT * MAX_STORAGE_COUNT / 2;
+
+  Storage default_map_;
+  struct WaitFreeStorage {
+    Storage maps_[MAX_STORAGE_COUNT];
+  };
+  unique_ptr<WaitFreeStorage> wait_free_storage_;
+
+  Storage &get_storage(const KeyT &key) {
+    if (wait_free_storage_ == nullptr) {
+      return default_map_;
+    }
+
+    return wait_free_storage_->maps_[randomize_hash(HashT()(key)) & (MAX_STORAGE_COUNT - 1)];
+  }
+
+ public:
+  void set(const KeyT &key, ValueT value) {
+    auto &storage = get_storage(key);
+    storage[key] = std::move(value);
+    if (default_map_.size() == MAX_STORAGE_SIZE) {
+      CHECK(wait_free_storage_ == nullptr);
+      wait_free_storage_ = make_unique<WaitFreeStorage>();
+      for (auto &it : default_map_) {
+        get_storage(it.first).emplace(it.first, std::move(it.second));
+      }
+      default_map_.clear();
+    }
+  }
+
+  ValueT get(const KeyT &key) {
+    auto &storage = get_storage(key);
+    auto it = storage.find(key);
+    if (it == storage.end()) {
+      return {};
+    }
+    return it->second;
+  }
+
+  size_t erase(const KeyT &key) {
+    return get_storage(key).erase(key);
+  }
+};
+
+}  // namespace td

tdutils/test/WaitFreeHashMap.cpp

@@ -0,0 +1,53 @@
+//
+// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2022
+//
+// 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 "td/utils/common.h"
+#include "td/utils/FlatHashMap.h"
+#include "td/utils/Random.h"
+#include "td/utils/tests.h"
+#include "td/utils/WaitFreeHashMap.h"
+
+TEST(WaitFreeHashMap, stress_test) {
+  td::Random::Xorshift128plus rnd(123);
+  td::FlatHashMap<td::uint64, td::uint64> reference;
+  td::WaitFreeHashMap<td::uint64, td::uint64> map;
+
+  td::vector<td::RandomSteps::Step> steps;
+  auto add_step = [&](td::uint32 weight, auto f) {
+    steps.emplace_back(td::RandomSteps::Step{std::move(f), weight});
+  };
+
+  auto gen_key = [&] {
+    return rnd() % 100000 + 1;
+  };
+
+  add_step(2000, [&] {
+    auto key = gen_key();
+    auto value = rnd();
+    reference[key] = value;
+    map.set(key, value);
+    ASSERT_EQ(reference[key], map.get(key));
+  });
+
+  add_step(2000, [&] {
+    auto key = gen_key();
+    auto ref_it = reference.find(key);
+    auto ref_value = ref_it == reference.end() ? 0 : ref_it->second;
+    ASSERT_EQ(ref_value, map.get(key));
+  });
+
+  add_step(500, [&] {
+    auto key = gen_key();
+    size_t reference_erased_count = reference.erase(key);
+    size_t map_erased_count = map.erase(key);
+    ASSERT_EQ(reference_erased_count, map_erased_count);
+  });
+
+  td::RandomSteps runner(std::move(steps));
+  for (size_t i = 0; i < 1000000; i++) {
+    runner.step(rnd);
+  }
+}