New feature support in folly::AtomicHash*

Summary:
AtomicHashMap/AtomicHashArray didn't support a custom equality
function, and didnt support pointer keys. We only ever do
compare-and-swap against the 3 "magic" keys, so as long as we're
careful we can continue to do that, while using the equality
function elsewhere.

Also add a user-configurable growth rate, independent of
maxLoadFactor.

Test Plan: automated tests

Reviewed By: delong.j@fb.com

FB internal diff: D806936

folly/AtomicHashArray-inl.h

@@ -24,8 +24,8 @@
 namespace folly {
 
 // AtomicHashArray private constructor --
-template <class KeyT, class ValueT, class HashFcn>
-AtomicHashArray<KeyT, ValueT, HashFcn>::
+template <class KeyT, class ValueT, class HashFcn, class EqualFcn>
+AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::
 AtomicHashArray(size_t capacity, KeyT emptyKey, KeyT lockedKey,
                 KeyT erasedKey, double maxLoadFactor, size_t cacheSize)
     : capacity_(capacity), maxEntries_(size_t(maxLoadFactor * capacity_ + 0.5)),
@@ -41,9 +41,9 @@ AtomicHashArray(size_t capacity, KeyT emptyKey, KeyT lockedKey,
  *   of key and returns true, or if key does not exist returns false and
  *   ret.index is set to capacity_.
  */
-template <class KeyT, class ValueT, class HashFcn>
-typename AtomicHashArray<KeyT, ValueT, HashFcn>::SimpleRetT
-AtomicHashArray<KeyT, ValueT, HashFcn>::
+template <class KeyT, class ValueT, class HashFcn, class EqualFcn>
+typename AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::SimpleRetT
+AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::
 findInternal(const KeyT key_in) {
   DCHECK_NE(key_in, kEmptyKey_);
   DCHECK_NE(key_in, kLockedKey_);
@@ -52,7 +52,7 @@ findInternal(const KeyT key_in) {
        ;
        idx = probeNext(idx, numProbes)) {
     const KeyT key = acquireLoadKey(cells_[idx]);
-    if (LIKELY(key == key_in)) {
+    if (LIKELY(EqualFcn()(key, key_in))) {
       return SimpleRetT(idx, true);
     }
     if (UNLIKELY(key == kEmptyKey_)) {
@@ -77,10 +77,10 @@ findInternal(const KeyT key_in) {
  *   this will be the previously inserted value, and if the map is full it is
  *   default.
  */
-template <class KeyT, class ValueT, class HashFcn>
+template <class KeyT, class ValueT, class HashFcn, class EqualFcn>
 template <class T>
-typename AtomicHashArray<KeyT, ValueT, HashFcn>::SimpleRetT
-AtomicHashArray<KeyT, ValueT, HashFcn>::
+typename AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::SimpleRetT
+AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::
 insertInternal(KeyT key_in, T&& value) {
   const short NO_NEW_INSERTS = 1;
   const short NO_PENDING_INSERTS = 2;
@@ -140,6 +140,8 @@ insertInternal(KeyT key_in, T&& value) {
             --numPendingEntries_;
             throw;
           }
+          // Direct comparison rather than EqualFcn ok here
+          // (we just inserted it)
           DCHECK(relaxedLoadKey(*cell) == key_in);
           --numPendingEntries_;
           ++numEntries_;  // This is a thread cached atomic increment :)
@@ -159,7 +161,7 @@ insertInternal(KeyT key_in, T&& value) {
     }
 
     const KeyT thisKey = acquireLoadKey(*cell);
-    if (thisKey == key_in) {
+    if (EqualFcn()(thisKey, key_in)) {
       // Found an existing entry for our key, but we don't overwrite the
       // previous value.
       return SimpleRetT(idx, false);
@@ -191,8 +193,8 @@ insertInternal(KeyT key_in, T&& value) {
  *   erased key will never be reused. If there's an associated value, we won't
  *   touch it either.
  */
-template <class KeyT, class ValueT, class HashFcn>
-size_t AtomicHashArray<KeyT, ValueT, HashFcn>::
+template <class KeyT, class ValueT, class HashFcn, class EqualFcn>
+size_t AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::
 erase(KeyT key_in) {
   CHECK_NE(key_in, kEmptyKey_);
   CHECK_NE(key_in, kLockedKey_);
@@ -208,10 +210,10 @@ erase(KeyT key_in) {
       // is similar to how it's handled in find().
       return 0;
     }
-    if (key_in == currentKey) {
+    if (EqualFcn()(currentKey, key_in)) {
       // Found an existing entry for our key, attempt to mark it erased.
       // Some other thread may have erased our key, but this is ok.
-      KeyT expect = key_in;
+      KeyT expect = currentKey;
       if (cellKeyPtr(*cell)->compare_exchange_strong(expect, kErasedKey_)) {
         numErases_.fetch_add(1, std::memory_order_relaxed);
 
@@ -234,13 +236,13 @@ erase(KeyT key_in) {
   }
 }
 
-template <class KeyT, class ValueT, class HashFcn>
-const typename AtomicHashArray<KeyT, ValueT, HashFcn>::Config
-AtomicHashArray<KeyT, ValueT, HashFcn>::defaultConfig;
+template <class KeyT, class ValueT, class HashFcn, class EqualFcn>
+const typename AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::Config
+AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::defaultConfig;
 
-template <class KeyT, class ValueT, class HashFcn>
-typename AtomicHashArray<KeyT, ValueT, HashFcn>::SmartPtr
-AtomicHashArray<KeyT, ValueT, HashFcn>::
+template <class KeyT, class ValueT, class HashFcn, class EqualFcn>
+typename AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::SmartPtr
+AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::
 create(size_t maxSize, const Config& c) {
   CHECK_LE(c.maxLoadFactor, 1.0);
   CHECK_GT(c.maxLoadFactor, 0.0);
@@ -272,8 +274,8 @@ create(size_t maxSize, const Config& c) {
   return map;
 }
 
-template <class KeyT, class ValueT, class HashFcn>
-void AtomicHashArray<KeyT, ValueT, HashFcn>::
+template <class KeyT, class ValueT, class HashFcn, class EqualFcn>
+void AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::
 destroy(AtomicHashArray* p) {
   assert(p);
   FOR_EACH_RANGE(i, 0, p->capacity_) {
@@ -286,8 +288,8 @@ destroy(AtomicHashArray* p) {
 }
 
 // clear -- clears all keys and values in the map and resets all counters
-template <class KeyT, class ValueT, class HashFcn>
-void AtomicHashArray<KeyT, ValueT, HashFcn>::
+template <class KeyT, class ValueT, class HashFcn, class EqualFcn>
+void AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::
 clear() {
   FOR_EACH_RANGE(i, 0, capacity_) {
     if (cells_[i].first != kEmptyKey_) {
@@ -305,9 +307,9 @@ clear() {
 
 // Iterator implementation
 
-template <class KeyT, class ValueT, class HashFcn>
+template <class KeyT, class ValueT, class HashFcn, class EqualFcn>
 template <class ContT, class IterVal>
-struct AtomicHashArray<KeyT, ValueT, HashFcn>::aha_iterator
+struct AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn>::aha_iterator
     : boost::iterator_facade<aha_iterator<ContT,IterVal>,
                              IterVal,
                              boost::forward_traversal_tag>

folly/AtomicHashArray.h

@@ -42,13 +42,16 @@
 
 namespace folly {
 
-template <class KeyT, class ValueT, class HashFcn = std::hash<KeyT>>
+template <class KeyT, class ValueT,
+          class HashFcn = std::hash<KeyT>, class EqualFcn = std::equal_to<KeyT>>
 class AtomicHashMap;
 
-template <class KeyT, class ValueT, class HashFcn = std::hash<KeyT>>
+template <class KeyT, class ValueT,
+          class HashFcn = std::hash<KeyT>, class EqualFcn = std::equal_to<KeyT>>
 class AtomicHashArray : boost::noncopyable {
   static_assert((std::is_convertible<KeyT,int32_t>::value ||
-                 std::is_convertible<KeyT,int64_t>::value),
+                 std::is_convertible<KeyT,int64_t>::value ||
+                 std::is_convertible<KeyT,const void*>::value),
              "You are trying to use AtomicHashArray with disallowed key "
              "types.  You must use atomically compare-and-swappable integer "
              "keys, or a different container class.");
@@ -117,13 +120,15 @@ class AtomicHashArray : boost::noncopyable {
     KeyT   lockedKey;
     KeyT   erasedKey;
     double maxLoadFactor;
+    double growthFactor;
     int    entryCountThreadCacheSize;
     size_t capacity; // if positive, overrides maxLoadFactor
 
-    constexpr Config() : emptyKey(static_cast<KeyT>(-1ul)),
-                         lockedKey(static_cast<KeyT>(-2ul)),
-                         erasedKey(static_cast<KeyT>(-3ul)),
+    constexpr Config() : emptyKey((KeyT)-1),
+                         lockedKey((KeyT)-2),
+                         erasedKey((KeyT)-3),
                          maxLoadFactor(0.8),
+                         growthFactor(-1),
                          entryCountThreadCacheSize(1000),
                          capacity(0) {}
   };
@@ -210,7 +215,7 @@ class AtomicHashArray : boost::noncopyable {
   /* Private data and helper functions... */
 
  private:
-  friend class AtomicHashMap<KeyT,ValueT,HashFcn>;
+  friend class AtomicHashMap<KeyT,ValueT,HashFcn,EqualFcn>;
 
   struct SimpleRetT { size_t idx; bool success;
     SimpleRetT(size_t i, bool s) : idx(i), success(s) {}

folly/AtomicHashMap-inl.h

@@ -22,16 +22,17 @@
 
 namespace folly {
 
-template <class KeyT, class ValueT, class HashFcn>
-const typename AtomicHashMap<KeyT, ValueT, HashFcn>::Config
-AtomicHashMap<KeyT, ValueT, HashFcn>::defaultConfig;
+template <class KeyT, class ValueT, class HashFcn, class EqualFcn>
+const typename AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::Config
+AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::defaultConfig;
 
 // 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>
-AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 AtomicHashMap(size_t size, const Config& config)
-  : kGrowthFrac_(1.0 - config.maxLoadFactor) {
+  : kGrowthFrac_(config.growthFactor < 0 ?
+                 1.0 - config.maxLoadFactor : config.growthFactor) {
   CHECK(config.maxLoadFactor > 0.0 && config.maxLoadFactor < 1.0);
   subMaps_[0].store(SubMap::create(size, config).release(),
     std::memory_order_relaxed);
@@ -43,9 +44,9 @@ AtomicHashMap(size_t size, const Config& config)
 }
 
 // insert --
-template <typename KeyT, typename ValueT, typename HashFcn>
-std::pair<typename AtomicHashMap<KeyT,ValueT,HashFcn>::iterator,bool>
-AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+std::pair<typename AtomicHashMap<KeyT,ValueT,HashFcn,EqualFcn>::iterator,bool>
+AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 insert(key_type k, const mapped_type& v) {
   SimpleRetT ret = insertInternal(k,v);
   SubMap* subMap = subMaps_[ret.i].load(std::memory_order_relaxed);
@@ -53,9 +54,9 @@ insert(key_type k, const mapped_type& v) {
                         ret.success);
 }
 
-template <typename KeyT, typename ValueT, typename HashFcn>
-std::pair<typename AtomicHashMap<KeyT,ValueT,HashFcn>::iterator,bool>
-AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+std::pair<typename AtomicHashMap<KeyT,ValueT,HashFcn,EqualFcn>::iterator,bool>
+AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 insert(key_type k, mapped_type&& v) {
   SimpleRetT ret = insertInternal(k, std::move(v));
   SubMap* subMap = subMaps_[ret.i].load(std::memory_order_relaxed);
@@ -64,15 +65,14 @@ insert(key_type k, mapped_type&& v) {
 }
 
 // insertInternal -- Allocates new sub maps as existing ones fill up.
-template <typename KeyT, typename ValueT, typename HashFcn>
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
 template <class T>
-typename AtomicHashMap<KeyT, ValueT, HashFcn>::SimpleRetT
-AtomicHashMap<KeyT, ValueT, HashFcn>::
+typename AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::SimpleRetT
+AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 insertInternal(key_type key, T&& value) {
  beginInsertInternal:
   int nextMapIdx = // this maintains our state
     numMapsAllocated_.load(std::memory_order_acquire);
-  uint32_t idx = 0;
   typename SubMap::SimpleRetT ret;
   FOR_EACH_RANGE(i, 0, nextMapIdx) {
     // insert in each map successively.  If one succeeds, we're done!
@@ -142,9 +142,9 @@ insertInternal(key_type key, T&& value) {
 }
 
 // find --
-template <typename KeyT, typename ValueT, typename HashFcn>
-typename AtomicHashMap<KeyT, ValueT, HashFcn>::iterator
-AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+typename AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::iterator
+AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 find(KeyT k) {
   SimpleRetT ret = findInternal(k);
   if (ret.i >= numMapsAllocated_.load(std::memory_order_acquire)) {
@@ -154,17 +154,17 @@ find(KeyT k) {
   return iterator(this, ret.i, subMap->makeIter(ret.j));
 }
 
-template <typename KeyT, typename ValueT, typename HashFcn>
-typename AtomicHashMap<KeyT, ValueT, HashFcn>::const_iterator
-AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+typename AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::const_iterator
+AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 find(KeyT k) const {
   return const_cast<AtomicHashMap*>(this)->find(k);
 }
 
 // findInternal --
-template <typename KeyT, typename ValueT, typename HashFcn>
-typename AtomicHashMap<KeyT, ValueT, HashFcn>::SimpleRetT
-AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+typename AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::SimpleRetT
+AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 findInternal(const KeyT k) const {
   SubMap* const primaryMap = subMaps_[0].load(std::memory_order_relaxed);
   typename SubMap::SimpleRetT ret = primaryMap->findInternal(k);
@@ -185,9 +185,9 @@ findInternal(const KeyT k) const {
 }
 
 // findAtInternal -- see encodeIndex() for details.
-template <typename KeyT, typename ValueT, typename HashFcn>
-typename AtomicHashMap<KeyT, ValueT, HashFcn>::SimpleRetT
-AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+typename AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::SimpleRetT
+AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 findAtInternal(uint32_t idx) const {
   uint32_t subMapIdx, subMapOffset;
   if (idx & kSecondaryMapBit_) {
@@ -205,9 +205,9 @@ findAtInternal(uint32_t idx) const {
 }
 
 // erase --
-template <typename KeyT, typename ValueT, typename HashFcn>
-typename AtomicHashMap<KeyT, ValueT, HashFcn>::size_type
-AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+typename AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::size_type
+AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 erase(const KeyT k) {
   int const numMaps = numMapsAllocated_.load(std::memory_order_acquire);
   FOR_EACH_RANGE(i, 0, numMaps) {
@@ -221,8 +221,8 @@ erase(const KeyT k) {
 }
 
 // capacity -- summation of capacities of all submaps
-template <typename KeyT, typename ValueT, typename HashFcn>
-size_t AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+size_t AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 capacity() const {
   size_t totalCap(0);
   int const numMaps = numMapsAllocated_.load(std::memory_order_acquire);
@@ -234,8 +234,8 @@ capacity() const {
 
 // spaceRemaining --
 // number of new insertions until current submaps are all at max load
-template <typename KeyT, typename ValueT, typename HashFcn>
-size_t AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+size_t AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 spaceRemaining() const {
   size_t spaceRem(0);
   int const numMaps = numMapsAllocated_.load(std::memory_order_acquire);
@@ -251,8 +251,8 @@ spaceRemaining() const {
 
 // clear -- Wipes all keys and values from primary map and destroys
 // all secondary maps.  Not thread safe.
-template <typename KeyT, typename ValueT, typename HashFcn>
-void AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+void AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 clear() {
   subMaps_[0].load(std::memory_order_relaxed)->clear();
   int const numMaps = numMapsAllocated_
@@ -267,8 +267,8 @@ clear() {
 }
 
 // size --
-template <typename KeyT, typename ValueT, typename HashFcn>
-size_t AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+size_t AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 size() const {
   size_t totalSize(0);
   int const numMaps = numMapsAllocated_.load(std::memory_order_acquire);
@@ -296,8 +296,8 @@ size() const {
 //         31              1
 //      27-30   which subMap
 //       0-26  subMap offset (index_ret input)
-template <typename KeyT, typename ValueT, typename HashFcn>
-inline uint32_t AtomicHashMap<KeyT, ValueT, HashFcn>::
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
+inline uint32_t AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::
 encodeIndex(uint32_t subMap, uint32_t offset) {
   DCHECK_EQ(offset & kSecondaryMapBit_, 0);  // offset can't be too big
   if (subMap == 0) return offset;
@@ -313,9 +313,9 @@ encodeIndex(uint32_t subMap, uint32_t offset) {
 
 // Iterator implementation
 
-template <typename KeyT, typename ValueT, typename HashFcn>
+template <typename KeyT, typename ValueT, typename HashFcn, typename EqualFcn>
 template<class ContT, class IterVal, class SubIt>
-struct AtomicHashMap<KeyT, ValueT, HashFcn>::ahm_iterator
+struct AtomicHashMap<KeyT, ValueT, HashFcn, EqualFcn>::ahm_iterator
     : boost::iterator_facade<ahm_iterator<ContT,IterVal,SubIt>,
                              IterVal,
                              boost::forward_traversal_tag>

folly/AtomicHashMap.h

@@ -143,10 +143,6 @@ namespace folly {
  * - We don't take the Hash function object as an instance in the
  *   constructor.
  *
- * - We don't take a Compare template parameter (since our keys must
- *   be integers, and the underlying hash array here uses atomic
- *   compare-and-swap instructions, we only allow normal integer
- *   comparisons).
  */
 
 // Thrown when insertion fails due to running out of space for
@@ -157,16 +153,16 @@ struct AtomicHashMapFullError : std::runtime_error {
   {}
 };
 
-template<class KeyT, class ValueT, class HashFcn>
+template<class KeyT, class ValueT, class HashFcn, class EqualFcn>
 class AtomicHashMap : boost::noncopyable {
-  typedef AtomicHashArray<KeyT, ValueT, HashFcn> SubMap;
+  typedef AtomicHashArray<KeyT, ValueT, HashFcn, EqualFcn> SubMap;
 
  public:
   typedef KeyT                key_type;
   typedef ValueT              mapped_type;
   typedef std::pair<const KeyT, ValueT> value_type;
   typedef HashFcn             hasher;
-  typedef std::equal_to<KeyT> key_equal;
+  typedef EqualFcn            key_equal;
   typedef value_type*         pointer;
   typedef value_type&         reference;
   typedef const value_type&   const_reference;
@@ -204,7 +200,7 @@ class AtomicHashMap : boost::noncopyable {
     }
   }
 
-  key_equal key_eq() const { return key_eq(); }
+  key_equal key_eq() const { return key_equal(); }
   hasher hash_function() const { return hasher(); }
 
   // TODO: emplace() support would be nice.