folly/AtomicHashArray-inl.h

   1 /*
   2  * Copyright 2015 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #ifndef FOLLY_ATOMICHASHARRAY_H_
  18 #error "This should only be included by AtomicHashArray.h"
  19 #endif
  20
  21 #include <folly/Bits.h>
  22 #include <folly/detail/AtomicHashUtils.h>
  23
  24 namespace folly {
  25
  26 // AtomicHashArray private constructor --
  27 template <class KeyT, class ValueT,
  28           class HashFcn, class EqualFcn, class Allocator, class ProbeFcn>
  29 AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn, Allocator, ProbeFcn>::
  30 AtomicHashArray(size_t capacity, KeyT emptyKey, KeyT lockedKey,
  31                 KeyT erasedKey, double _maxLoadFactor, size_t cacheSize)
  32     : capacity_(capacity),
  33       maxEntries_(size_t(_maxLoadFactor * capacity_ + 0.5)),
  34       kEmptyKey_(emptyKey), kLockedKey_(lockedKey), kErasedKey_(erasedKey),
  35       kAnchorMask_(nextPowTwo(capacity_) - 1), numEntries_(0, cacheSize),
  36       numPendingEntries_(0, cacheSize), isFull_(0), numErases_(0) {
  37 }
  38
  39 /*
  40  * findInternal --
  41  *
  42  *   Sets ret.second to value found and ret.index to index
  43  *   of key and returns true, or if key does not exist returns false and
  44  *   ret.index is set to capacity_.
  45  */
  46 template <class KeyT, class ValueT,
  47           class HashFcn, class EqualFcn, class Allocator, class ProbeFcn>
  48 typename AtomicHashArray<KeyT, ValueT,
  49          HashFcn, EqualFcn, Allocator, ProbeFcn>::SimpleRetT
  50 AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn, Allocator, ProbeFcn>::
  51 findInternal(const KeyT key_in) {
  52   DCHECK_NE(key_in, kEmptyKey_);
  53   DCHECK_NE(key_in, kLockedKey_);
  54   DCHECK_NE(key_in, kErasedKey_);
  55   for (size_t idx = keyToAnchorIdx(key_in), numProbes = 0;
  56        ;
  57        idx = ProbeFcn()(idx, numProbes, capacity_)) {
  58     const KeyT key = acquireLoadKey(cells_[idx]);
  59     if (LIKELY(EqualFcn()(key, key_in))) {
  60       return SimpleRetT(idx, true);
  61     }
  62     if (UNLIKELY(key == kEmptyKey_)) {
  63       // if we hit an empty element, this key does not exist
  64       return SimpleRetT(capacity_, false);
  65     }
  66     // NOTE: the way we count numProbes must be same in find(), insert(),
  67     // and erase(). Otherwise it may break probing.
  68     ++numProbes;
  69     if (UNLIKELY(numProbes >= capacity_)) {
  70       // probed every cell...fail
  71       return SimpleRetT(capacity_, false);
  72     }
  73   }
  74 }
  75
  76 /*
  77  * insertInternal --
  78  *
  79  *   Returns false on failure due to key collision or full.
  80  *   Also sets ret.index to the index of the key.  If the map is full, sets
  81  *   ret.index = capacity_.  Also sets ret.second to cell value, thus if insert
  82  *   successful this will be what we just inserted, if there is a key collision
  83  *   this will be the previously inserted value, and if the map is full it is
  84  *   default.
  85  */
  86 template <class KeyT, class ValueT,
  87           class HashFcn, class EqualFcn, class Allocator, class ProbeFcn>
  88 template <typename... ArgTs>
  89 typename AtomicHashArray<KeyT, ValueT,
  90          HashFcn, EqualFcn, Allocator, ProbeFcn>::SimpleRetT
  91 AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn, Allocator, ProbeFcn>::
  92 insertInternal(KeyT key_in, ArgTs&&... vCtorArgs) {
  93   const short NO_NEW_INSERTS = 1;
  94   const short NO_PENDING_INSERTS = 2;
  95   CHECK_NE(key_in, kEmptyKey_);
  96   CHECK_NE(key_in, kLockedKey_);
  97   CHECK_NE(key_in, kErasedKey_);
  98
  99   size_t idx = keyToAnchorIdx(key_in);
 100   size_t numProbes = 0;
 101   for (;;) {
 102     DCHECK_LT(idx, capacity_);
 103     value_type* cell = &cells_[idx];
 104     if (relaxedLoadKey(*cell) == kEmptyKey_) {
 105       // NOTE: isFull_ is set based on numEntries_.readFast(), so it's
 106       // possible to insert more than maxEntries_ entries. However, it's not
 107       // possible to insert past capacity_.
 108       ++numPendingEntries_;
 109       if (isFull_.load(std::memory_order_acquire)) {
 110         --numPendingEntries_;
 111
 112         // Before deciding whether this insert succeeded, this thread needs to
 113         // wait until no other thread can add a new entry.
 114
 115         // Correctness assumes isFull_ is true at this point. If
 116         // another thread now does ++numPendingEntries_, we expect it
 117         // to pass the isFull_.load() test above. (It shouldn't insert
 118         // a new entry.)
 119         detail::atomic_hash_spin_wait([&] {
 120           return
 121             (isFull_.load(std::memory_order_acquire) != NO_PENDING_INSERTS) &&
 122             (numPendingEntries_.readFull() != 0);
 123         });
 124         isFull_.store(NO_PENDING_INSERTS, std::memory_order_release);
 125
 126         if (relaxedLoadKey(*cell) == kEmptyKey_) {
 127           // Don't insert past max load factor
 128           return SimpleRetT(capacity_, false);
 129         }
 130       } else {
 131         // An unallocated cell. Try once to lock it. If we succeed, insert here.
 132         // If we fail, fall through to comparison below; maybe the insert that
 133         // just beat us was for this very key....
 134         if (tryLockCell(cell)) {
 135           // Write the value - done before unlocking
 136           try {
 137             DCHECK(relaxedLoadKey(*cell) == kLockedKey_);
 138             new (&cell->second) ValueT(std::forward<ArgTs>(vCtorArgs)...);
 139             unlockCell(cell, key_in); // Sets the new key
 140           } catch (...) {
 141             // Transition back to empty key---requires handling
 142             // locked->empty below.
 143             unlockCell(cell, kEmptyKey_);
 144             --numPendingEntries_;
 145             throw;
 146           }
 147           // An erase() can race here and delete right after our insertion
 148           // Direct comparison rather than EqualFcn ok here
 149           // (we just inserted it)
 150           DCHECK(relaxedLoadKey(*cell) == key_in ||
 151                  relaxedLoadKey(*cell) == kErasedKey_);
 152           --numPendingEntries_;
 153           ++numEntries_;  // This is a thread cached atomic increment :)
 154           if (numEntries_.readFast() >= maxEntries_) {
 155             isFull_.store(NO_NEW_INSERTS, std::memory_order_relaxed);
 156           }
 157           return SimpleRetT(idx, true);
 158         }
 159         --numPendingEntries_;
 160       }
 161     }
 162     DCHECK(relaxedLoadKey(*cell) != kEmptyKey_);
 163     if (kLockedKey_ == acquireLoadKey(*cell)) {
 164       detail::atomic_hash_spin_wait([&] {
 165         return kLockedKey_ == acquireLoadKey(*cell);
 166       });
 167     }
 168
 169     const KeyT thisKey = acquireLoadKey(*cell);
 170     if (EqualFcn()(thisKey, key_in)) {
 171       // Found an existing entry for our key, but we don't overwrite the
 172       // previous value.
 173       return SimpleRetT(idx, false);
 174     } else if (thisKey == kEmptyKey_ || thisKey == kLockedKey_) {
 175       // We need to try again (i.e., don't increment numProbes or
 176       // advance idx): this case can happen if the constructor for
 177       // ValueT threw for this very cell (the rethrow block above).
 178       continue;
 179     }
 180
 181
 182     // NOTE: the way we count numProbes must be same in find(),
 183     // insert(), and erase(). Otherwise it may break probing.
 184     ++numProbes;
 185     if (UNLIKELY(numProbes >= capacity_)) {
 186       // probed every cell...fail
 187       return SimpleRetT(capacity_, false);
 188     }
 189
 190     idx = ProbeFcn()(idx, numProbes, capacity_);
 191   }
 192 }
 193
 194
 195 /*
 196  * erase --
 197  *
 198  *   This will attempt to erase the given key key_in if the key is found. It
 199  *   returns 1 iff the key was located and marked as erased, and 0 otherwise.
 200  *
 201  *   Memory is not freed or reclaimed by erase, i.e. the cell containing the
 202  *   erased key will never be reused. If there's an associated value, we won't
 203  *   touch it either.
 204  */
 205 template <class KeyT, class ValueT,
 206           class HashFcn, class EqualFcn, class Allocator, class ProbeFcn>
 207 size_t AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn, Allocator, ProbeFcn>::
 208 erase(KeyT key_in) {
 209   CHECK_NE(key_in, kEmptyKey_);
 210   CHECK_NE(key_in, kLockedKey_);
 211   CHECK_NE(key_in, kErasedKey_);
 212
 213   for (size_t idx = keyToAnchorIdx(key_in), numProbes = 0;
 214        ;
 215        idx = ProbeFcn()(idx, numProbes, capacity_)) {
 216     DCHECK_LT(idx, capacity_);
 217     value_type* cell = &cells_[idx];
 218     KeyT currentKey = acquireLoadKey(*cell);
 219     if (currentKey == kEmptyKey_ || currentKey == kLockedKey_) {
 220       // If we hit an empty (or locked) element, this key does not exist. This
 221       // is similar to how it's handled in find().
 222       return 0;
 223     }
 224     if (EqualFcn()(currentKey, key_in)) {
 225       // Found an existing entry for our key, attempt to mark it erased.
 226       // Some other thread may have erased our key, but this is ok.
 227       KeyT expect = currentKey;
 228       if (cellKeyPtr(*cell)->compare_exchange_strong(expect, kErasedKey_)) {
 229         numErases_.fetch_add(1, std::memory_order_relaxed);
 230
 231         // Even if there's a value in the cell, we won't delete (or even
 232         // default construct) it because some other thread may be accessing it.
 233         // Locking it meanwhile won't work either since another thread may be
 234         // holding a pointer to it.
 235
 236         // We found the key and successfully erased it.
 237         return 1;
 238       }
 239       // If another thread succeeds in erasing our key, we'll stop our search.
 240       return 0;
 241     }
 242
 243     // NOTE: the way we count numProbes must be same in find(), insert(),
 244     // and erase(). Otherwise it may break probing.
 245     ++numProbes;
 246     if (UNLIKELY(numProbes >= capacity_)) {
 247       // probed every cell...fail
 248       return 0;
 249     }
 250   }
 251 }
 252
 253 template <class KeyT, class ValueT,
 254          class HashFcn, class EqualFcn, class Allocator, class ProbeFcn>
 255 typename AtomicHashArray<KeyT, ValueT,
 256          HashFcn, EqualFcn, Allocator, ProbeFcn>::SmartPtr
 257 AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn, Allocator, ProbeFcn>::
 258 create(size_t maxSize, const Config& c) {
 259   CHECK_LE(c.maxLoadFactor, 1.0);
 260   CHECK_GT(c.maxLoadFactor, 0.0);
 261   CHECK_NE(c.emptyKey, c.lockedKey);
 262   size_t capacity = size_t(maxSize / c.maxLoadFactor);
 263   size_t sz = sizeof(AtomicHashArray) + sizeof(value_type) * capacity;
 264
 265   auto const mem = Allocator().allocate(sz);
 266   try {
 267     new (mem) AtomicHashArray(capacity, c.emptyKey, c.lockedKey, c.erasedKey,
 268                               c.maxLoadFactor, c.entryCountThreadCacheSize);
 269   } catch (...) {
 270     Allocator().deallocate(mem, sz);
 271     throw;
 272   }
 273
 274   SmartPtr map(static_cast<AtomicHashArray*>((void *)mem));
 275
 276   /*
 277    * Mark all cells as empty.
 278    *
 279    * Note: we're bending the rules a little here accessing the key
 280    * element in our cells even though the cell object has not been
 281    * constructed, and casting them to atomic objects (see cellKeyPtr).
 282    * (Also, in fact we never actually invoke the value_type
 283    * constructor.)  This is in order to avoid needing to default
 284    * construct a bunch of value_type when we first start up: if you
 285    * have an expensive default constructor for the value type this can
 286    * noticeably speed construction time for an AHA.
 287    */
 288   FOR_EACH_RANGE(i, 0, map->capacity_) {
 289     cellKeyPtr(map->cells_[i])->store(map->kEmptyKey_,
 290       std::memory_order_relaxed);
 291   }
 292   return map;
 293 }
 294
 295 template <class KeyT, class ValueT,
 296           class HashFcn, class EqualFcn, class Allocator, class ProbeFcn>
 297 void AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn, Allocator, ProbeFcn>::
 298 destroy(AtomicHashArray* p) {
 299   assert(p);
 300
 301   size_t sz = sizeof(AtomicHashArray) + sizeof(value_type) * p->capacity_;
 302
 303   FOR_EACH_RANGE(i, 0, p->capacity_) {
 304     if (p->cells_[i].first != p->kEmptyKey_) {
 305       p->cells_[i].~value_type();
 306     }
 307   }
 308   p->~AtomicHashArray();
 309
 310   Allocator().deallocate((char *)p, sz);
 311 }
 312
 313 // clear -- clears all keys and values in the map and resets all counters
 314 template <class KeyT, class ValueT,
 315           class HashFcn, class EqualFcn, class Allocator, class ProbeFcn>
 316 void AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn, Allocator, ProbeFcn>::
 317 clear() {
 318   FOR_EACH_RANGE(i, 0, capacity_) {
 319     if (cells_[i].first != kEmptyKey_) {
 320       cells_[i].~value_type();
 321       *const_cast<KeyT*>(&cells_[i].first) = kEmptyKey_;
 322     }
 323     CHECK(cells_[i].first == kEmptyKey_);
 324   }
 325   numEntries_.set(0);
 326   numPendingEntries_.set(0);
 327   isFull_.store(0, std::memory_order_relaxed);
 328   numErases_.store(0, std::memory_order_relaxed);
 329 }
 330
 331
 332 // Iterator implementation
 333
 334 template <class KeyT, class ValueT,
 335           class HashFcn, class EqualFcn, class Allocator, class ProbeFcn>
 336 template <class ContT, class IterVal>
 337 struct AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn, Allocator, ProbeFcn>::
 338     aha_iterator
 339     : boost::iterator_facade<aha_iterator<ContT,IterVal>,
 340                              IterVal,
 341                              boost::forward_traversal_tag>
 342 {
 343   explicit aha_iterator() : aha_(0) {}
 344
 345   // Conversion ctor for interoperability between const_iterator and
 346   // iterator.  The enable_if<> magic keeps us well-behaved for
 347   // is_convertible<> (v. the iterator_facade documentation).
 348   template<class OtherContT, class OtherVal>
 349   aha_iterator(const aha_iterator<OtherContT,OtherVal>& o,
 350                typename std::enable_if<
 351                std::is_convertible<OtherVal*,IterVal*>::value >::type* = 0)
 352       : aha_(o.aha_)
 353       , offset_(o.offset_)
 354   {}
 355
 356   explicit aha_iterator(ContT* array, size_t offset)
 357       : aha_(array)
 358       , offset_(offset)
 359   {}
 360
 361   // Returns unique index that can be used with findAt().
 362   // WARNING: The following function will fail silently for hashtable
 363   // with capacity > 2^32
 364   uint32_t getIndex() const { return offset_; }
 365
 366   void advancePastEmpty() {
 367     while (offset_ < aha_->capacity_ && !isValid()) {
 368       ++offset_;
 369     }
 370   }
 371
 372  private:
 373   friend class AtomicHashArray;
 374   friend class boost::iterator_core_access;
 375
 376   void increment() {
 377     ++offset_;
 378     advancePastEmpty();
 379   }
 380
 381   bool equal(const aha_iterator& o) const {
 382     return aha_ == o.aha_ && offset_ == o.offset_;
 383   }
 384
 385   IterVal& dereference() const {
 386     return aha_->cells_[offset_];
 387   }
 388
 389   bool isValid() const {
 390     KeyT key = acquireLoadKey(aha_->cells_[offset_]);
 391     return key != aha_->kEmptyKey_  &&
 392       key != aha_->kLockedKey_ &&
 393       key != aha_->kErasedKey_;
 394   }
 395
 396  private:
 397   ContT* aha_;
 398   size_t offset_;
 399 }; // aha_iterator
 400
 401 } // namespace folly