Add WaitFreeHashSet.

tdutils/CMakeLists.txt

@@ -336,6 +336,7 @@ set(TDUTILS_TEST_SOURCE
   ${CMAKE_CURRENT_SOURCE_DIR}/test/StealingQueue.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/test/variant.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/test/WaitFreeHashMap.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/test/WaitFreeHashSet.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/test/WaitFreeVector.cpp
   PARENT_SCOPE
 )

tdutils/td/utils/WaitFreeHashSet.h

@@ -0,0 +1,114 @@
+//
+// 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/FlatHashSet.h"
+#include "td/utils/HashTableUtils.h"
+
+#include <functional>
+
+namespace td {
+
+template <class KeyT, class HashT = std::hash<KeyT>, class EqT = std::equal_to<KeyT>>
+class WaitFreeHashSet {
+  using Storage = FlatHashSet<KeyT, 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_set_;
+  struct WaitFreeStorage {
+    Storage sets_[MAX_STORAGE_COUNT];
+  };
+  unique_ptr<WaitFreeStorage> wait_free_storage_;
+
+  Storage &get_wait_free_storage(const KeyT &key) {
+    return wait_free_storage_->sets_[randomize_hash(HashT()(key)) & (MAX_STORAGE_COUNT - 1)];
+  }
+
+  Storage &get_storage(const KeyT &key) {
+    if (wait_free_storage_ == nullptr) {
+      return default_set_;
+    }
+
+    return get_wait_free_storage(key);
+  }
+
+  const Storage &get_storage(const KeyT &key) const {
+    return const_cast<WaitFreeHashSet *>(this)->get_storage(key);
+  }
+
+  void split_storage() {
+    CHECK(wait_free_storage_ == nullptr);
+    wait_free_storage_ = make_unique<WaitFreeStorage>();
+    for (auto &it : default_set_) {
+      get_wait_free_storage(it).insert(it);
+    }
+    default_set_.clear();
+  }
+
+ public:
+  void insert(const KeyT &key) {
+    auto &storage = get_storage(key);
+    storage.insert(key);
+    if (default_set_.size() == MAX_STORAGE_SIZE) {
+      split_storage();
+    }
+  }
+
+  size_t count(const KeyT &key) const {
+    const auto &storage = get_storage(key);
+    return storage.count(key);
+  }
+
+  size_t erase(const KeyT &key) {
+    return get_storage(key).erase(key);
+  }
+
+  void foreach(std::function<void(const KeyT &key)> callback) const {
+    if (wait_free_storage_ == nullptr) {
+      for (auto &it : default_set_) {
+        callback(it);
+      }
+      return;
+    }
+
+    for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) {
+      for (auto &it : wait_free_storage_->sets_[i]) {
+        callback(it);
+      }
+    }
+  }
+
+  size_t size() const {
+    if (wait_free_storage_ == nullptr) {
+      return default_set_.size();
+    }
+
+    size_t result = 0;
+    for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) {
+      result += wait_free_storage_->sets_[i].size();
+    }
+    return result;
+  }
+
+  bool empty() const {
+    if (wait_free_storage_ == nullptr) {
+      return default_set_.empty();
+    }
+
+    for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) {
+      if (!wait_free_storage_->sets_[i].empty()) {
+        return false;
+      }
+    }
+    return true;
+  }
+};
+
+}  // namespace td

tdutils/test/WaitFreeHashSet.cpp

@@ -0,0 +1,71 @@
+//
+// 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/FlatHashSet.h"
+#include "td/utils/Random.h"
+#include "td/utils/tests.h"
+#include "td/utils/WaitFreeHashSet.h"
+
+TEST(WaitFreeHashSet, stress_test) {
+  td::Random::Xorshift128plus rnd(123);
+  td::FlatHashSet<td::uint64> reference;
+  td::WaitFreeHashSet<td::uint64> set;
+
+  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;
+  };
+
+  auto check = [&] {
+    ASSERT_EQ(reference.size(), set.size());
+    ASSERT_EQ(reference.empty(), set.empty());
+
+    if (reference.size() < 100) {
+      td::uint64 result = 0;
+      for (auto &it : reference) {
+        result += it * 101;
+      }
+      set.foreach([&](const td::uint64 &key) { result -= key * 101; });
+      ASSERT_EQ(0u, result);
+    }
+  };
+
+  add_step(2000, [&] {
+    auto key = gen_key();
+    ASSERT_EQ(reference.count(key), set.count(key));
+    reference.insert(key);
+    set.insert(key);
+    ASSERT_EQ(reference.count(key), set.count(key));
+    check();
+  });
+
+  add_step(500, [&] {
+    auto key = gen_key();
+    size_t reference_erased_count = reference.erase(key);
+    size_t set_erased_count = set.erase(key);
+    ASSERT_EQ(reference_erased_count, set_erased_count);
+    check();
+  });
+
+  td::RandomSteps runner(std::move(steps));
+  for (size_t i = 0; i < 1000000; i++) {
+    runner.step(rnd);
+  }
+
+  for (size_t test = 0; test < 1000; test++) {
+    reference = {};
+    set = {};
+
+    for (size_t i = 0; i < 100; i++) {
+      runner.step(rnd);
+    }
+  }
+}