AtomicUnorderedMap.h

/*
 * 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.
 */

#pragma once

#include <atomic>
#include <cstdint>
#include <functional>
#include <limits>
#include <stdexcept>
#include <system_error>
#include <type_traits>

#include <folly/Conv.h>
#include <folly/Likely.h>
#include <folly/Random.h>
#include <folly/ScopeGuard.h>
#include <folly/Traits.h>
#include <folly/detail/AtomicUnorderedMapUtils.h>
#include <folly/lang/Bits.h>
#include <folly/portability/SysMman.h>
#include <folly/portability/Unistd.h>

namespace folly {

/// You're probably reading this because you are looking for an
/// AtomicUnorderedMap<K,V> that is fully general, highly concurrent (for
/// reads, writes, and iteration), and makes no performance compromises.
/// We haven't figured that one out yet.  What you will find here is a
/// hash table implementation that sacrifices generality so that it can
/// give you all of the other things.
///
/// LIMITATIONS:
///
/// * Insert only (*) - the only write operation supported directly by
///   AtomicUnorderedInsertMap is findOrConstruct.  There is a (*) because
///   values aren't moved, so you can roll your own concurrency control for
///   in-place updates of values (see MutableData and MutableAtom below),
///   but the hash table itself doesn't help you.
///
/// * No resizing - you must specify the capacity up front, and once
///   the hash map gets full you won't be able to insert.  Insert
///   performance will degrade once the load factor is high.  Insert is
///   O(1/(1-actual_load_factor)).  Note that this is a pretty strong
///   limitation, because you can't remove existing keys.
///
/// * 2^30 maximum default capacity - by default AtomicUnorderedInsertMap
///   uses uint32_t internal indexes (and steals 2 bits), limiting you
///   to about a billion entries.  If you need more you can fill in all
///   of the template params so you change IndexType to uint64_t, or you
///   can use AtomicUnorderedInsertMap64.  64-bit indexes will increase
///   the space over of the map, of course.
///
/// WHAT YOU GET IN EXCHANGE:
///
/// * Arbitrary key and value types - any K and V that can be used in a
///   std::unordered_map can be used here.  In fact, the key and value
///   types don't even have to be copyable or moveable!
///
/// * Keys and values in the map won't be moved - it is safe to keep
///   pointers or references to the keys and values in the map, because
///   they are never moved or destroyed (until the map itself is destroyed).
///
/// * Iterators are never invalidated - writes don't invalidate iterators,
///   so you can scan and insert in parallel.
///
/// * Fast wait-free reads - reads are usually only a single cache miss,
///   even when the hash table is very large.  Wait-freedom means that
///   you won't see latency outliers even in the face of concurrent writes.
///
/// * Lock-free insert - writes proceed in parallel.  If a thread in the
///   middle of a write is unlucky and gets suspended, it doesn't block
///   anybody else.
///
/// COMMENTS ON INSERT-ONLY
///
/// This map provides wait-free linearizable reads and lock-free
/// linearizable inserts.  Inserted values won't be moved, but no
/// concurrency control is provided for safely updating them.  To remind
/// you of that fact they are only provided in const form.  This is the
/// only simple safe thing to do while preserving something like the normal
/// std::map iteration form, which requires that iteration be exposed
/// via std::pair (and prevents encapsulation of access to the value).
///
/// There are a couple of reasonable policies for doing in-place
/// concurrency control on the values.  I am hoping that the policy can
/// be injected via the value type or an extra template param, to keep
/// the core AtomicUnorderedInsertMap insert-only:
///
///   CONST: this is the currently implemented strategy, which is simple,
///   performant, and not that expressive.  You can always put in a value
///   with a mutable field (see MutableAtom below), but that doesn't look
///   as pretty as it should.
///
///   ATOMIC: for integers and integer-size trivially copyable structs
///   (via an adapter like tao/queues/AtomicStruct) the value can be a
///   std::atomic and read and written atomically.
///
///   SEQ-LOCK: attach a counter incremented before and after write.
///   Writers serialize by using CAS to make an even->odd transition,
///   then odd->even after the write.  Readers grab the value with memcpy,
///   checking sequence value before and after.  Readers retry until they
///   see an even sequence number that doesn't change.  This works for
///   larger structs, but still requires memcpy to be equivalent to copy
///   assignment, and it is no longer lock-free.  It scales very well,
///   because the readers are still invisible (no cache line writes).
///
///   LOCK: folly's SharedMutex would be a good choice here.
///
/// MEMORY ALLOCATION
///
/// Underlying memory is allocated as a big anonymous mmap chunk, which
/// might be cheaper than calloc() and is certainly not more expensive
/// for large maps.  If the SkipKeyValueDeletion template param is true
/// then deletion of the map consists of unmapping the backing memory,
/// which is much faster than destructing all of the keys and values.
/// Feel free to override if std::is_trivial_destructor isn't recognizing
/// the triviality of your destructors.
template <
    typename Key,
    typename Value,
    typename Hash = std::hash<Key>,
    typename KeyEqual = std::equal_to<Key>,
    bool SkipKeyValueDeletion =
        (std::is_trivially_destructible<Key>::value &&
         std::is_trivially_destructible<Value>::value),
    template <typename> class Atom = std::atomic,
    typename IndexType = uint32_t,
    typename Allocator = folly::detail::MMapAlloc>

struct AtomicUnorderedInsertMap {
  typedef Key key_type;
  typedef Value mapped_type;
  typedef std::pair<Key, Value> value_type;
  typedef std::size_t size_type;
  typedef std::ptrdiff_t difference_type;
  typedef Hash hasher;
  typedef KeyEqual key_equal;
  typedef const value_type& const_reference;

  typedef struct ConstIterator {
    ConstIterator(const AtomicUnorderedInsertMap& owner, IndexType slot)
        : owner_(owner), slot_(slot) {}

    ConstIterator(const ConstIterator&) = default;
    ConstIterator& operator=(const ConstIterator&) = default;

    const value_type& operator*() const {
      return owner_.slots_[slot_].keyValue();
    }

    const value_type* operator->() const {
      return &owner_.slots_[slot_].keyValue();
    }

    // pre-increment
    const ConstIterator& operator++() {
      while (slot_ > 0) {
        --slot_;
        if (owner_.slots_[slot_].state() == LINKED) {
          break;
        }
      }
      return *this;
    }

    // post-increment
    ConstIterator operator++(int /* dummy */) {
      auto prev = *this;
      ++*this;
      return prev;
    }

    bool operator==(const ConstIterator& rhs) const {
      return slot_ == rhs.slot_;
    }
    bool operator!=(const ConstIterator& rhs) const { return !(*this == rhs); }

   private:
    const AtomicUnorderedInsertMap& owner_;
    IndexType slot_;
  } const_iterator;

  friend ConstIterator;

  /// Constructs a map that will support the insertion of maxSize key-value
  /// pairs without exceeding the max load factor.  Load factors of greater
  /// than 1 are not supported, and once the actual load factor of the
  /// map approaches 1 the insert performance will suffer.  The capacity
  /// is limited to 2^30 (about a billion) for the default IndexType,
  /// beyond which we will throw invalid_argument.
  explicit AtomicUnorderedInsertMap(
      size_t maxSize,
      float maxLoadFactor = 0.8f,
      const Allocator& alloc = Allocator())
      : allocator_(alloc) {
    size_t capacity = size_t(maxSize / std::min(1.0f, maxLoadFactor) + 128);
    size_t avail = size_t{1} << (8 * sizeof(IndexType) - 2);
    if (capacity > avail && maxSize < avail) {
      // we'll do our best
      capacity = avail;
    }
    if (capacity < maxSize || capacity > avail) {
      throw std::invalid_argument(
          "AtomicUnorderedInsertMap capacity must fit in IndexType with 2 bits "
          "left over");
    }

    numSlots_ = capacity;
    slotMask_ = folly::nextPowTwo(capacity * 4) - 1;
    mmapRequested_ = sizeof(Slot) * capacity;
    slots_ = reinterpret_cast<Slot*>(allocator_.allocate(mmapRequested_));
    zeroFillSlots();
    // mark the zero-th slot as in-use but not valid, since that happens
    // to be our nil value
    slots_[0].stateUpdate(EMPTY, CONSTRUCTING);
  }

  ~AtomicUnorderedInsertMap() {
    if (!SkipKeyValueDeletion) {
      for (size_t i = 1; i < numSlots_; ++i) {
        slots_[i].~Slot();
      }
    }
    allocator_.deallocate(reinterpret_cast<char*>(slots_), mmapRequested_);
  }

  /// Searches for the key, returning (iter,false) if it is found.
  /// If it is not found calls the functor Func with a void* argument
  /// that is raw storage suitable for placement construction of a Value
  /// (see raw_value_type), then returns (iter,true).  May call Func and
  /// then return (iter,false) if there are other concurrent writes, in
  /// which case the newly constructed value will be immediately destroyed.
  ///
  /// This function does not block other readers or writers.  If there
  /// are other concurrent writes, many parallel calls to func may happen
  /// and only the first one to complete will win.  The values constructed
  /// by the other calls to func will be destroyed.
  ///
  /// Usage:
  ///
  ///  AtomicUnorderedInsertMap<std::string,std::string> memo;
  ///
  ///  auto value = memo.findOrConstruct(key, [=](void* raw) {
  ///    new (raw) std::string(computation(key));
  ///  })->first;
  template <typename Func>
  std::pair<const_iterator, bool> findOrConstruct(const Key& key, Func&& func) {
    auto const slot = keyToSlotIdx(key);
    auto prev = slots_[slot].headAndState_.load(std::memory_order_acquire);

    auto existing = find(key, slot);
    if (existing != 0) {
      return std::make_pair(ConstIterator(*this, existing), false);
    }

    // The copying of key and the calling of func can throw exceptions. Nothing
    // else in this function can throw an exception. In the event of an
    // exception, deallocate as if the KV was beaten in a concurrent addition.
    const auto idx = allocateNear(slot);
    SCOPE_FAIL {
      slots_[idx].stateUpdate(CONSTRUCTING, EMPTY);
    };
    Key* addr = &slots_[idx].keyValue().first;
    new (addr) Key(key);
    SCOPE_FAIL {
      addr->~Key();
    };
    func(static_cast<void*>(&slots_[idx].keyValue().second));

    while (true) {
      slots_[idx].next_ = prev >> 2;

      // we can merge the head update and the CONSTRUCTING -> LINKED update
      // into a single CAS if slot == idx (which should happen often)
      auto after = idx << 2;
      if (slot == idx) {
        after += LINKED;
      } else {
        after += (prev & 3);
      }

      if (slots_[slot].headAndState_.compare_exchange_strong(prev, after)) {
        // success
        if (idx != slot) {
          slots_[idx].stateUpdate(CONSTRUCTING, LINKED);
        }
        return std::make_pair(ConstIterator(*this, idx), true);
      }
      // compare_exchange_strong updates its first arg on failure, so
      // there is no need to reread prev

      existing = find(key, slot);
      if (existing != 0) {
        // our allocated key and value are no longer needed
        slots_[idx].keyValue().first.~Key();
        slots_[idx].keyValue().second.~Value();
        slots_[idx].stateUpdate(CONSTRUCTING, EMPTY);

        return std::make_pair(ConstIterator(*this, existing), false);
      }
    }
  }

  /// This isn't really emplace, but it is what we need to test.
  /// Eventually we can duplicate all of the std::pair constructor
  /// forms, including a recursive tuple forwarding template
  /// http://functionalcpp.wordpress.com/2013/08/28/tuple-forwarding/).
  template <class K, class V>
  std::pair<const_iterator, bool> emplace(const K& key, V&& value) {
    return findOrConstruct(
        key, [&](void* raw) { new (raw) Value(std::forward<V>(value)); });
  }

  const_iterator find(const Key& key) const {
    return ConstIterator(*this, find(key, keyToSlotIdx(key)));
  }

  const_iterator cbegin() const {
    IndexType slot = numSlots_ - 1;
    while (slot > 0 && slots_[slot].state() != LINKED) {
      --slot;
    }
    return ConstIterator(*this, slot);
  }
  const_iterator begin() const { return cbegin(); }

  const_iterator cend() const { return ConstIterator(*this, 0); }
  const_iterator end() const { return cend(); }

 private:
  enum : IndexType {
    kMaxAllocationTries = 1000, // after this we throw
  };

  enum BucketState : IndexType {
    EMPTY = 0,
    CONSTRUCTING = 1,
    LINKED = 2,
  };

  /// Lock-free insertion is easiest by prepending to collision chains.
  /// A large chaining hash table takes two cache misses instead of
  /// one, however.  Our solution is to colocate the bucket storage and
  /// the head storage, so that even though we are traversing chains we
  /// are likely to stay within the same cache line.  Just make sure to
  /// traverse head before looking at any keys.  This strategy gives us
  /// 32 bit pointers and fast iteration.
  struct Slot {
    /// The bottom two bits are the BucketState, the rest is the index
    /// of the first bucket for the chain whose keys map to this slot.
    /// When things are going well the head usually links to this slot,
    /// but that doesn't always have to happen.
    Atom<IndexType> headAndState_;

    /// The next bucket in the chain
    IndexType next_;

    /// Key and Value
    aligned_storage_for_t<value_type> raw_;

    ~Slot() {
      auto s = state();
      assert(s == EMPTY || s == LINKED);
      if (s == LINKED) {
        keyValue().first.~Key();
        keyValue().second.~Value();
      }
    }

    BucketState state() const {
      return BucketState(headAndState_.load(std::memory_order_acquire) & 3);
    }

    void stateUpdate(BucketState before, BucketState after) {
      assert(state() == before);
      headAndState_ += (after - before);
    }

    value_type& keyValue() {
      assert(state() != EMPTY);
      return *static_cast<value_type*>(static_cast<void*>(&raw_));
    }

    const value_type& keyValue() const {
      assert(state() != EMPTY);
      return *static_cast<const value_type*>(static_cast<const void*>(&raw_));
    }
  };

  // We manually manage the slot memory so we can bypass initialization
  // (by getting a zero-filled mmap chunk) and optionally destruction of
  // the slots

  size_t mmapRequested_;
  size_t numSlots_;

  /// tricky, see keyToSlotIdx
  size_t slotMask_;

  Allocator allocator_;
  Slot* slots_;

  IndexType keyToSlotIdx(const Key& key) const {
    size_t h = hasher()(key);
    h &= slotMask_;
    while (h >= numSlots_) {
      h -= numSlots_;
    }
    return h;
  }

  IndexType find(const Key& key, IndexType slot) const {
    KeyEqual ke = {};
    auto hs = slots_[slot].headAndState_.load(std::memory_order_acquire);
    for (slot = hs >> 2; slot != 0; slot = slots_[slot].next_) {
      if (ke(key, slots_[slot].keyValue().first)) {
        return slot;
      }
    }
    return 0;
  }

  /// Allocates a slot and returns its index.  Tries to put it near
  /// slots_[start].
  IndexType allocateNear(IndexType start) {
    for (IndexType tries = 0; tries < kMaxAllocationTries; ++tries) {
      auto slot = allocationAttempt(start, tries);
      auto prev = slots_[slot].headAndState_.load(std::memory_order_acquire);
      if ((prev & 3) == EMPTY &&
          slots_[slot].headAndState_.compare_exchange_strong(
              prev, prev + CONSTRUCTING - EMPTY)) {
        return slot;
      }
    }
    throw std::bad_alloc();
  }

  /// Returns the slot we should attempt to allocate after tries failed
  /// tries, starting from the specified slot.  This is pulled out so we
  /// can specialize it differently during deterministic testing
  IndexType allocationAttempt(IndexType start, IndexType tries) const {
    if (FOLLY_LIKELY(tries < 8 && start + tries < numSlots_)) {
      return IndexType(start + tries);
    } else {
      IndexType rv;
      if (sizeof(IndexType) <= 4) {
        rv = IndexType(folly::Random::rand32(numSlots_));
      } else {
        rv = IndexType(folly::Random::rand64(numSlots_));
      }
      assert(rv < numSlots_);
      return rv;
    }
  }

  void zeroFillSlots() {
    using folly::detail::GivesZeroFilledMemory;
    if (!GivesZeroFilledMemory<Allocator>::value) {
      memset(static_cast<void*>(slots_), 0, mmapRequested_);
    }
  }
};

/// AtomicUnorderedInsertMap64 is just a type alias that makes it easier
/// to select a 64 bit slot index type.  Use this if you need a capacity
/// bigger than 2^30 (about a billion).  This increases memory overheads,
/// obviously.
template <
    typename Key,
    typename Value,
    typename Hash = std::hash<Key>,
    typename KeyEqual = std::equal_to<Key>,
    bool SkipKeyValueDeletion =
        (std::is_trivially_destructible<Key>::value &&
         std::is_trivially_destructible<Value>::value),
    template <typename> class Atom = std::atomic,
    typename Allocator = folly::detail::MMapAlloc>
using AtomicUnorderedInsertMap64 = AtomicUnorderedInsertMap<
    Key,
    Value,
    Hash,
    KeyEqual,
    SkipKeyValueDeletion,
    Atom,
    uint64_t,
    Allocator>;

/// MutableAtom is a tiny wrapper that gives you the option of atomically
/// updating values inserted into an AtomicUnorderedInsertMap<K,
/// MutableAtom<V>>.  This relies on AtomicUnorderedInsertMap's guarantee
/// that it doesn't move values.
template <typename T, template <typename> class Atom = std::atomic>
struct MutableAtom {
  mutable Atom<T> data;

  explicit MutableAtom(const T& init) : data(init) {}
};

/// MutableData is a tiny wrapper that gives you the option of using an
/// external concurrency control mechanism to updating values inserted
/// into an AtomicUnorderedInsertMap.
template <typename T>
struct MutableData {
  mutable T data;
  explicit MutableData(const T& init) : data(init) {}
};

} // namespace folly