tdlight/tdutils/td/utils/WaitFreeHashMap.h

191 lines
5.0 KiB
C
Raw Normal View History

2022-05-19 18:20:07 +02:00
//
// 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 = Hash<KeyT>, class EqT = std::equal_to<KeyT>>
2022-05-19 18:20:07 +02:00
class WaitFreeHashMap {
2022-11-18 10:53:26 +01:00
static constexpr size_t MAX_STORAGE_COUNT = 1 << 8;
2022-05-19 18:20:07 +02:00
static_assert((MAX_STORAGE_COUNT & (MAX_STORAGE_COUNT - 1)) == 0, "");
static constexpr uint32 DEFAULT_STORAGE_SIZE = 1 << 12;
2022-05-19 18:20:07 +02:00
2022-11-18 10:53:26 +01:00
FlatHashMap<KeyT, ValueT, HashT, EqT> default_map_;
2022-05-19 18:20:07 +02:00
struct WaitFreeStorage {
2022-11-18 10:53:26 +01:00
WaitFreeHashMap maps_[MAX_STORAGE_COUNT];
2022-05-19 18:20:07 +02:00
};
unique_ptr<WaitFreeStorage> wait_free_storage_;
2022-11-18 10:53:26 +01:00
uint32 hash_mult_ = 1;
uint32 max_storage_size_ = DEFAULT_STORAGE_SIZE;
2022-05-19 18:20:07 +02:00
2022-11-18 10:53:26 +01:00
uint32 get_wait_free_index(const KeyT &key) const {
return randomize_hash(HashT()(key) * hash_mult_) & (MAX_STORAGE_COUNT - 1);
2022-08-03 23:38:13 +02:00
}
2022-11-18 10:53:26 +01:00
WaitFreeHashMap &get_wait_free_storage(const KeyT &key) {
return wait_free_storage_->maps_[get_wait_free_index(key)];
2022-05-19 18:20:07 +02:00
}
2022-11-18 10:53:26 +01:00
const WaitFreeHashMap &get_wait_free_storage(const KeyT &key) const {
return wait_free_storage_->maps_[get_wait_free_index(key)];
2022-08-03 23:38:13 +02:00
}
2022-08-14 14:04:08 +02:00
void split_storage() {
2022-08-03 23:38:13 +02:00
CHECK(wait_free_storage_ == nullptr);
wait_free_storage_ = make_unique<WaitFreeStorage>();
2022-11-18 13:19:01 +01:00
uint32 next_hash_mult = hash_mult_ * 1000000007;
2022-11-18 10:53:26 +01:00
for (uint32 i = 0; i < MAX_STORAGE_COUNT; i++) {
auto &map = wait_free_storage_->maps_[i];
map.hash_mult_ = next_hash_mult;
map.max_storage_size_ = DEFAULT_STORAGE_SIZE + i * next_hash_mult % DEFAULT_STORAGE_SIZE;
}
2022-08-03 23:38:13 +02:00
for (auto &it : default_map_) {
2022-11-18 10:53:26 +01:00
get_wait_free_storage(it.first).set(it.first, std::move(it.second));
2022-08-03 23:38:13 +02:00
}
default_map_.clear();
}
2022-05-19 18:20:07 +02:00
public:
void set(const KeyT &key, ValueT value) {
2022-11-18 10:53:26 +01:00
if (wait_free_storage_ != nullptr) {
return get_wait_free_storage(key).set(key, std::move(value));
}
default_map_[key] = std::move(value);
if (default_map_.size() == max_storage_size_) {
2022-08-14 14:04:08 +02:00
split_storage();
2022-05-19 18:20:07 +02:00
}
}
ValueT get(const KeyT &key) const {
2022-11-18 10:53:26 +01:00
if (wait_free_storage_ != nullptr) {
return get_wait_free_storage(key).get(key);
}
auto it = default_map_.find(key);
if (it == default_map_.end()) {
2022-05-19 18:20:07 +02:00
return {};
}
return it->second;
}
2022-08-04 14:25:46 +02:00
size_t count(const KeyT &key) const {
2022-11-18 10:53:26 +01:00
if (wait_free_storage_ != nullptr) {
return get_wait_free_storage(key).count(key);
}
return default_map_.count(key);
2022-08-04 14:25:46 +02:00
}
// specialization for WaitFreeHashMap<..., unique_ptr<T>>
2022-08-04 21:40:22 +02:00
template <class T = ValueT>
typename T::element_type *get_pointer(const KeyT &key) {
2022-11-18 10:53:26 +01:00
if (wait_free_storage_ != nullptr) {
return get_wait_free_storage(key).get_pointer(key);
}
auto it = default_map_.find(key);
if (it == default_map_.end()) {
return nullptr;
}
return it->second.get();
}
2022-08-04 21:40:22 +02:00
template <class T = ValueT>
const typename T::element_type *get_pointer(const KeyT &key) const {
2022-11-18 10:53:26 +01:00
if (wait_free_storage_ != nullptr) {
return get_wait_free_storage(key).get_pointer(key);
}
auto it = default_map_.find(key);
if (it == default_map_.end()) {
return nullptr;
}
return it->second.get();
}
2022-08-03 23:38:13 +02:00
ValueT &operator[](const KeyT &key) {
if (wait_free_storage_ == nullptr) {
ValueT &result = default_map_[key];
2022-11-18 10:53:26 +01:00
if (default_map_.size() != max_storage_size_) {
2022-08-03 23:38:13 +02:00
return result;
}
2022-08-14 14:04:08 +02:00
split_storage();
2022-08-03 23:38:13 +02:00
}
return get_wait_free_storage(key)[key];
}
2022-05-19 18:20:07 +02:00
size_t erase(const KeyT &key) {
2022-11-18 10:53:26 +01:00
if (wait_free_storage_ != nullptr) {
return get_wait_free_storage(key).erase(key);
}
return default_map_.erase(key);
2022-05-19 18:20:07 +02:00
}
2022-11-18 10:53:26 +01:00
void foreach(const std::function<void(const KeyT &key, ValueT &value)> &callback) {
2022-08-04 14:08:18 +02:00
if (wait_free_storage_ == nullptr) {
for (auto &it : default_map_) {
callback(it.first, it.second);
}
return;
}
2022-11-18 10:53:26 +01:00
for (auto &it : wait_free_storage_->maps_) {
it.foreach(callback);
2022-08-04 14:08:18 +02:00
}
}
2022-11-18 10:53:26 +01:00
void foreach(const std::function<void(const KeyT &key, const ValueT &value)> &callback) const {
2022-08-04 14:25:46 +02:00
if (wait_free_storage_ == nullptr) {
for (auto &it : default_map_) {
callback(it.first, it.second);
}
return;
}
2022-11-18 10:53:26 +01:00
for (auto &it : wait_free_storage_->maps_) {
it.foreach(callback);
2022-08-04 14:25:46 +02:00
}
}
size_t calc_size() const {
if (wait_free_storage_ == nullptr) {
return default_map_.size();
}
size_t result = 0;
for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) {
result += wait_free_storage_->maps_[i].calc_size();
}
return result;
}
bool empty() const {
if (wait_free_storage_ == nullptr) {
return default_map_.empty();
}
for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) {
if (!wait_free_storage_->maps_[i].empty()) {
return false;
}
}
return true;
}
2022-05-19 18:20:07 +02:00
};
} // namespace td