folly/AtomicHashArray-inl.h

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