/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef FOLLY_ATOMICHASHMAP_H_
#error "This should only be included by AtomicHashMap.h"
#endif
#include <folly/detail/AtomicHashUtils.h>
#include <folly/detail/Iterators.h>
#include <type_traits>
namespace folly {
// AtomicHashMap constructor -- Atomic wrapper that allows growth
// This class has a lot of overhead (184 Bytes) so only use for big maps
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::AtomicHashMap(size_t finalSizeEst, const Config& config)
: kGrowthFrac_(
config.growthFactor < 0 ? 1.0f - config.maxLoadFactor
: config.growthFactor) {
CHECK(config.maxLoadFactor > 0.0f && config.maxLoadFactor < 1.0f);
subMaps_[0].store(
SubMap::create(finalSizeEst, config).release(),
std::memory_order_relaxed);
auto subMapCount = kNumSubMaps_;
FOR_EACH_RANGE (i, 1, subMapCount) {
subMaps_[i].store(nullptr, std::memory_order_relaxed);
}
numMapsAllocated_.store(1, std::memory_order_relaxed);
}
// emplace --
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
template <
typename LookupKeyT,
typename LookupHashFcn,
typename LookupEqualFcn,
typename LookupKeyToKeyFcn,
typename... ArgTs>
std::pair<
typename AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::iterator,
bool>
AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::emplace(LookupKeyT k, ArgTs&&... vCtorArgs) {
SimpleRetT ret = insertInternal<
LookupKeyT,
LookupHashFcn,
LookupEqualFcn,
LookupKeyToKeyFcn>(k, std::forward<ArgTs>(vCtorArgs)...);
SubMap* subMap = subMaps_[ret.i].load(std::memory_order_relaxed);
return std::make_pair(
iterator(this, ret.i, subMap->makeIter(ret.j)), ret.success);
}
// insertInternal -- Allocates new sub maps as existing ones fill up.
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
template <
typename LookupKeyT,
typename LookupHashFcn,
typename LookupEqualFcn,
typename LookupKeyToKeyFcn,
typename... ArgTs>
typename AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::SimpleRetT
AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::insertInternal(LookupKeyT key, ArgTs&&... vCtorArgs) {
beginInsertInternal:
auto nextMapIdx = // this maintains our state
numMapsAllocated_.load(std::memory_order_acquire);
typename SubMap::SimpleRetT ret;
FOR_EACH_RANGE (i, 0, nextMapIdx) {
// insert in each map successively. If one succeeds, we're done!
SubMap* subMap = subMaps_[i].load(std::memory_order_relaxed);
ret = subMap->template insertInternal<
LookupKeyT,
LookupHashFcn,
LookupEqualFcn,
LookupKeyToKeyFcn>(key, std::forward<ArgTs>(vCtorArgs)...);
if (ret.idx == subMap->capacity_) {
continue; // map is full, so try the next one
}
// Either collision or success - insert in either case
return SimpleRetT(i, ret.idx, ret.success);
}
// If we made it this far, all maps are full and we need to try to allocate
// the next one.
SubMap* primarySubMap = subMaps_[0].load(std::memory_order_relaxed);
if (nextMapIdx >= kNumSubMaps_ ||
primarySubMap->capacity_ * kGrowthFrac_ < 1.0) {
// Can't allocate any more sub maps.
throw AtomicHashMapFullError();
}
if (tryLockMap(nextMapIdx)) {
// Alloc a new map and shove it in. We can change whatever
// we want because other threads are waiting on us...
size_t numCellsAllocated =
(size_t)(primarySubMap->capacity_ *
std::pow(1.0 + kGrowthFrac_, nextMapIdx - 1));
size_t newSize = size_t(numCellsAllocated * kGrowthFrac_);
DCHECK(
subMaps_[nextMapIdx].load(std::memory_order_relaxed) ==
(SubMap*)kLockedPtr_);
// create a new map using the settings stored in the first map
Config config;
config.emptyKey = primarySubMap->kEmptyKey_;
config.lockedKey = primarySubMap->kLockedKey_;
config.erasedKey = primarySubMap->kErasedKey_;
config.maxLoadFactor = primarySubMap->maxLoadFactor();
config.entryCountThreadCacheSize =
primarySubMap->getEntryCountThreadCacheSize();
subMaps_[nextMapIdx].store(
SubMap::create(newSize, config).release(), std::memory_order_relaxed);
// Publish the map to other threads.
numMapsAllocated_.fetch_add(1, std::memory_order_release);
DCHECK_EQ(
nextMapIdx + 1, numMapsAllocated_.load(std::memory_order_relaxed));
} else {
// If we lost the race, we'll have to wait for the next map to get
// allocated before doing any insertion here.
detail::atomic_hash_spin_wait([&] {
return nextMapIdx >= numMapsAllocated_.load(std::memory_order_acquire);
});
}
// Relaxed is ok here because either we just created this map, or we
// just did a spin wait with an acquire load on numMapsAllocated_.
SubMap* loadedMap = subMaps_[nextMapIdx].load(std::memory_order_relaxed);
DCHECK(loadedMap && loadedMap != (SubMap*)kLockedPtr_);
ret = loadedMap->insertInternal(key, std::forward<ArgTs>(vCtorArgs)...);
if (ret.idx != loadedMap->capacity_) {
return SimpleRetT(nextMapIdx, ret.idx, ret.success);
}
// We took way too long and the new map is already full...try again from
// the top (this should pretty much never happen).
goto beginInsertInternal;
}
// find --
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
template <class LookupKeyT, class LookupHashFcn, class LookupEqualFcn>
typename AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::iterator
AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::find(LookupKeyT k) {
SimpleRetT ret = findInternal<LookupKeyT, LookupHashFcn, LookupEqualFcn>(k);
if (!ret.success) {
return end();
}
SubMap* subMap = subMaps_[ret.i].load(std::memory_order_relaxed);
return iterator(this, ret.i, subMap->makeIter(ret.j));
}
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
template <class LookupKeyT, class LookupHashFcn, class LookupEqualFcn>
typename AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::const_iterator
AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::find(LookupKeyT k) const {
return const_cast<AtomicHashMap*>(this)
->find<LookupKeyT, LookupHashFcn, LookupEqualFcn>(k);
}
// findInternal --
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
template <class LookupKeyT, class LookupHashFcn, class LookupEqualFcn>
typename AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::SimpleRetT
AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::findInternal(const LookupKeyT k) const {
SubMap* const primaryMap = subMaps_[0].load(std::memory_order_relaxed);
typename SubMap::SimpleRetT ret =
primaryMap
->template findInternal<LookupKeyT, LookupHashFcn, LookupEqualFcn>(k);
if (FOLLY_LIKELY(ret.idx != primaryMap->capacity_)) {
return SimpleRetT(0, ret.idx, ret.success);
}
const unsigned int numMaps =
numMapsAllocated_.load(std::memory_order_acquire);
FOR_EACH_RANGE (i, 1, numMaps) {
// Check each map successively. If one succeeds, we're done!
SubMap* thisMap = subMaps_[i].load(std::memory_order_relaxed);
ret =
thisMap
->template findInternal<LookupKeyT, LookupHashFcn, LookupEqualFcn>(
k);
if (FOLLY_LIKELY(ret.idx != thisMap->capacity_)) {
return SimpleRetT(i, ret.idx, ret.success);
}
}
// Didn't find our key...
return SimpleRetT(numMaps, 0, false);
}
// findAtInternal -- see encodeIndex() for details.
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
typename AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::SimpleRetT
AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::findAtInternal(uint32_t idx) const {
uint32_t subMapIdx, subMapOffset;
if (idx & kSecondaryMapBit_) {
// idx falls in a secondary map
idx &= ~kSecondaryMapBit_; // unset secondary bit
subMapIdx = idx >> kSubMapIndexShift_;
DCHECK_LT(subMapIdx, numMapsAllocated_.load(std::memory_order_relaxed));
subMapOffset = idx & kSubMapIndexMask_;
} else {
// idx falls in primary map
subMapIdx = 0;
subMapOffset = idx;
}
return SimpleRetT(subMapIdx, subMapOffset, true);
}
// erase --
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
typename AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::size_type
AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::erase(const KeyT k) {
int const numMaps = numMapsAllocated_.load(std::memory_order_acquire);
FOR_EACH_RANGE (i, 0, numMaps) {
// Check each map successively. If one succeeds, we're done!
if (subMaps_[i].load(std::memory_order_relaxed)->erase(k)) {
return 1;
}
}
// Didn't find our key...
return 0;
}
// capacity -- summation of capacities of all submaps
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
size_t AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::capacity() const {
size_t totalCap(0);
int const numMaps = numMapsAllocated_.load(std::memory_order_acquire);
FOR_EACH_RANGE (i, 0, numMaps) {
totalCap += subMaps_[i].load(std::memory_order_relaxed)->capacity_;
}
return totalCap;
}
// spaceRemaining --
// number of new insertions until current submaps are all at max load
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
size_t AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::spaceRemaining() const {
size_t spaceRem(0);
int const numMaps = numMapsAllocated_.load(std::memory_order_acquire);
FOR_EACH_RANGE (i, 0, numMaps) {
SubMap* thisMap = subMaps_[i].load(std::memory_order_relaxed);
spaceRem +=
std::max(0, thisMap->maxEntries_ - &thisMap->numEntries_.readFull());
}
return spaceRem;
}
// clear -- Wipes all keys and values from primary map and destroys
// all secondary maps. Not thread safe.
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
void AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::clear() {
subMaps_[0].load(std::memory_order_relaxed)->clear();
int const numMaps = numMapsAllocated_.load(std::memory_order_relaxed);
FOR_EACH_RANGE (i, 1, numMaps) {
SubMap* thisMap = subMaps_[i].load(std::memory_order_relaxed);
DCHECK(thisMap);
SubMap::destroy(thisMap);
subMaps_[i].store(nullptr, std::memory_order_relaxed);
}
numMapsAllocated_.store(1, std::memory_order_relaxed);
}
// size --
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
size_t AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::size() const {
size_t totalSize(0);
int const numMaps = numMapsAllocated_.load(std::memory_order_acquire);
FOR_EACH_RANGE (i, 0, numMaps) {
totalSize += subMaps_[i].load(std::memory_order_relaxed)->size();
}
return totalSize;
}
// encodeIndex -- Encode the submap index and offset into return.
// index_ret must be pre-populated with the submap offset.
//
// We leave index_ret untouched when referring to the primary map
// so it can be as large as possible (31 data bits). Max size of
// secondary maps is limited by what can fit in the low 27 bits.
//
// Returns the following bit-encoded data in index_ret:
// if subMap == 0 (primary map) =>
// bit(s) value
// 31 0
// 0-30 submap offset (index_ret input)
//
// if subMap > 0 (secondary maps) =>
// bit(s) value
// 31 1
// 27-30 which subMap
// 0-26 subMap offset (index_ret input)
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
inline uint32_t AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::encodeIndex(uint32_t subMap, uint32_t offset) {
DCHECK_EQ(offset & kSecondaryMapBit_, 0); // offset can't be too big
if (subMap == 0) {
return offset;
}
// Make sure subMap isn't too big
DCHECK_EQ(subMap >> kNumSubMapBits_, 0);
// Make sure subMap bits of offset are clear
DCHECK_EQ(offset & (~kSubMapIndexMask_ | kSecondaryMapBit_), 0);
// Set high-order bits to encode which submap this index belongs to
return offset | (subMap << kSubMapIndexShift_) | kSecondaryMapBit_;
}
// Iterator implementation
template <
typename KeyT,
typename ValueT,
typename HashFcn,
typename EqualFcn,
typename Allocator,
typename ProbeFcn,
typename KeyConvertFcn>
template <class ContT, class IterVal, class SubIt>
struct AtomicHashMap<
KeyT,
ValueT,
HashFcn,
EqualFcn,
Allocator,
ProbeFcn,
KeyConvertFcn>::ahm_iterator
: detail::IteratorFacade<
ahm_iterator<ContT, IterVal, SubIt>,
IterVal,
std::forward_iterator_tag> {
explicit ahm_iterator() : ahm_(nullptr) {}
// Conversion ctor for interoperability between const_iterator and
// iterator. The enable_if<> magic keeps us well-behaved for
// is_convertible<> (v. the iterator_facade documentation).
template <class OtherContT, class OtherVal, class OtherSubIt>
ahm_iterator(
const ahm_iterator<OtherContT, OtherVal, OtherSubIt>& o,
typename std::enable_if<
std::is_convertible<OtherSubIt, SubIt>::value>::type* = nullptr)
: ahm_(o.ahm_), subMap_(o.subMap_), subIt_(o.subIt_) {}
/*
* Returns the unique index that can be used for access directly
* into the data storage.
*/
uint32_t getIndex() const {
CHECK(!isEnd());
return ahm_->encodeIndex(subMap_, subIt_.getIndex());
}
private:
friend class AtomicHashMap;
explicit ahm_iterator(ContT* ahm, uint32_t subMap, const SubIt& subIt)
: ahm_(ahm), subMap_(subMap), subIt_(subIt) {}
friend class detail::
IteratorFacade<ahm_iterator, IterVal, std::forward_iterator_tag>;
void increment() {
CHECK(!isEnd());
++subIt_;
checkAdvanceToNextSubmap();
}
bool equal(const ahm_iterator& other) const {
if (ahm_ != other.ahm_) {
return false;
}
if (isEnd() || other.isEnd()) {
return isEnd() == other.isEnd();
}
return subMap_ == other.subMap_ && subIt_ == other.subIt_;
}
IterVal& dereference() const { return *subIt_; }
bool isEnd() const { return ahm_ == nullptr; }
void checkAdvanceToNextSubmap() {
if (isEnd()) {
return;
}
SubMap* thisMap = ahm_->subMaps_[subMap_].load(std::memory_order_relaxed);
while (subIt_ == thisMap->end()) {
// This sub iterator is done, advance to next one
if (subMap_ + 1 <
ahm_->numMapsAllocated_.load(std::memory_order_acquire)) {
++subMap_;
thisMap = ahm_->subMaps_[subMap_].load(std::memory_order_relaxed);
subIt_ = thisMap->begin();
} else {
ahm_ = nullptr;
return;
}
}
}
private:
ContT* ahm_;
uint32_t subMap_;
SubIt subIt_;
}; // ahm_iterator
} // namespace folly