UnboundedQueue.h

Go to the documentation of this file.

/*
 * Copyright 2017-present Facebook, Inc.
 *
 * 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 <chrono>
#include <memory>

#include <glog/logging.h>

#include <folly/ConstexprMath.h>
#include <folly/Optional.h>
#include <folly/concurrency/CacheLocality.h>
#include <folly/lang/Align.h>
#include <folly/synchronization/Hazptr.h>
#include <folly/synchronization/SaturatingSemaphore.h>
#include <folly/synchronization/WaitOptions.h>
#include <folly/synchronization/detail/Spin.h>

namespace folly {


template <
    typename T,
    bool SingleProducer,
    bool SingleConsumer,
    bool MayBlock,
    size_t LgSegmentSize = 8,
    size_t LgAlign = constexpr_log2(hardware_destructive_interference_size),
    template <typename> class Atom = std::atomic>
class UnboundedQueue {
  using Ticket = uint64_t;
  class Entry;
  class Segment;

  static constexpr bool SPSC = SingleProducer && SingleConsumer;
  static constexpr size_t Stride = SPSC || (LgSegmentSize <= 1) ? 1 : 27;
  static constexpr size_t SegmentSize = 1u << LgSegmentSize;
  static constexpr size_t Align = 1u << LgAlign;

  static_assert(
      std::is_nothrow_destructible<T>::value,
      "T must be nothrow_destructible");
  static_assert((Stride & 1) == 1, "Stride must be odd");
  static_assert(LgSegmentSize < 32, "LgSegmentSize must be < 32");
  static_assert(LgAlign < 16, "LgAlign must be < 16");

  struct Consumer {
    Atom<Segment*> head;
    Atom<Ticket> ticket;
    explicit Consumer(Segment* s) : head(s), ticket(0) {}
  };
  struct Producer {
    Atom<Segment*> tail;
    Atom<Ticket> ticket;
    explicit Producer(Segment* s) : tail(s), ticket(0) {}
  };

  alignas(Align) Consumer c_;
  alignas(Align) Producer p_;

 public:
  UnboundedQueue()
      : c_(new Segment(0)), p_(c_.head.load(std::memory_order_relaxed)) {}

  ~UnboundedQueue() {
    cleanUpRemainingItems();
    reclaimRemainingSegments();
  }

  FOLLY_ALWAYS_INLINE void enqueue(const T& arg) {
    enqueueImpl(arg);
  }

  FOLLY_ALWAYS_INLINE void enqueue(T&& arg) {
    enqueueImpl(std::move(arg));
  }

  FOLLY_ALWAYS_INLINE void dequeue(T& item) noexcept {
    dequeueImpl(item);
  }

  FOLLY_ALWAYS_INLINE bool try_dequeue(T& item) noexcept {
    auto o = try_dequeue();
    if (LIKELY(o.has_value())) {
      item = std::move(*o);
      return true;
    }
    return false;
  }

  FOLLY_ALWAYS_INLINE folly::Optional<T> try_dequeue() noexcept {
    return tryDequeueUntil(std::chrono::steady_clock::time_point::min());
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE bool try_dequeue_until(
      T& item,
      const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
    folly::Optional<T> o = try_dequeue_until(deadline);

    if (LIKELY(o.has_value())) {
      item = std::move(*o);
      return true;
    }

    return false;
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE folly::Optional<T> try_dequeue_until(
      const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
    return tryDequeueUntil(deadline);
  }

  template <typename Rep, typename Period>
  FOLLY_ALWAYS_INLINE bool try_dequeue_for(
      T& item,
      const std::chrono::duration<Rep, Period>& duration) noexcept {
    folly::Optional<T> o = try_dequeue_for(duration);

    if (LIKELY(o.has_value())) {
      item = std::move(*o);
      return true;
    }

    return false;
  }

  template <typename Rep, typename Period>
  FOLLY_ALWAYS_INLINE folly::Optional<T> try_dequeue_for(
      const std::chrono::duration<Rep, Period>& duration) noexcept {
    folly::Optional<T> o = try_dequeue();
    if (LIKELY(o.has_value())) {
      return o;
    }
    return tryDequeueUntil(std::chrono::steady_clock::now() + duration);
  }

  FOLLY_ALWAYS_INLINE const T* try_peek() noexcept {
    /* This function is supported only for USPSC and UMPSC queues. */
    DCHECK(SingleConsumer);
    return tryPeekUntil(std::chrono::steady_clock::time_point::min());
  }

  size_t size() const noexcept {
    auto p = producerTicket();
    auto c = consumerTicket();
    return p > c ? p - c : 0;
  }

  bool empty() const noexcept {
    auto c = consumerTicket();
    auto p = producerTicket();
    return p <= c;
  }

 private:
  template <typename Arg>
  FOLLY_ALWAYS_INLINE void enqueueImpl(Arg&& arg) {
    if (SPSC) {
      Segment* s = tail();
      enqueueCommon(s, std::forward<Arg>(arg));
    } else {
      // Using hazptr_holder instead of hazptr_local because it is
      // possible that the T ctor happens to use hazard pointers.
      hazptr_holder<Atom> hptr;
      Segment* s = hptr.get_protected(p_.tail);
      enqueueCommon(s, std::forward<Arg>(arg));
    }
  }

  template <typename Arg>
  FOLLY_ALWAYS_INLINE void enqueueCommon(Segment* s, Arg&& arg) {
    Ticket t = fetchIncrementProducerTicket();
    if (!SingleProducer) {
      s = findSegment(s, t);
    }
    DCHECK_GE(t, s->minTicket());
    DCHECK_LT(t, s->minTicket() + SegmentSize);
    size_t idx = index(t);
    Entry& e = s->entry(idx);
    e.putItem(std::forward<Arg>(arg));
    if (responsibleForAlloc(t)) {
      allocNextSegment(s);
    }
    if (responsibleForAdvance(t)) {
      advanceTail(s);
    }
  }

  FOLLY_ALWAYS_INLINE void dequeueImpl(T& item) noexcept {
    if (SPSC) {
      Segment* s = head();
      dequeueCommon(s, item);
    } else {
      // Using hazptr_holder instead of hazptr_local because it is
      // possible to call the T dtor and it may happen to use hazard
      // pointers.
      hazptr_holder<Atom> hptr;
      Segment* s = hptr.get_protected(c_.head);
      dequeueCommon(s, item);
    }
  }

  FOLLY_ALWAYS_INLINE void dequeueCommon(Segment* s, T& item) noexcept {
    Ticket t = fetchIncrementConsumerTicket();
    if (!SingleConsumer) {
      s = findSegment(s, t);
    }
    size_t idx = index(t);
    Entry& e = s->entry(idx);
    e.takeItem(item);
    if (responsibleForAdvance(t)) {
      advanceHead(s);
    }
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE folly::Optional<T> tryDequeueUntil(
      const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
    if (SingleConsumer) {
      Segment* s = head();
      return tryDequeueUntilSC(s, deadline);
    } else {
      // Using hazptr_holder instead of hazptr_local because it is
      //  possible to call ~T() and it may happen to use hazard pointers.
      hazptr_holder<Atom> hptr;
      Segment* s = hptr.get_protected(c_.head);
      return tryDequeueUntilMC(s, deadline);
    }
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE folly::Optional<T> tryDequeueUntilSC(
      Segment* s,
      const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
    Ticket t = consumerTicket();
    DCHECK_GE(t, s->minTicket());
    DCHECK_LT(t, (s->minTicket() + SegmentSize));
    size_t idx = index(t);
    Entry& e = s->entry(idx);
    if (UNLIKELY(!tryDequeueWaitElem(e, t, deadline))) {
      return folly::Optional<T>();
    }
    setConsumerTicket(t + 1);
    auto ret = e.takeItem();
    if (responsibleForAdvance(t)) {
      advanceHead(s);
    }
    return ret;
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE folly::Optional<T> tryDequeueUntilMC(
      Segment* s,
      const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
    while (true) {
      Ticket t = consumerTicket();
      if (UNLIKELY(t >= (s->minTicket() + SegmentSize))) {
        s = getAllocNextSegment(s, t);
        DCHECK(s);
        continue;
      }
      size_t idx = index(t);
      Entry& e = s->entry(idx);
      if (UNLIKELY(!tryDequeueWaitElem(e, t, deadline))) {
        return folly::Optional<T>();
      }
      if (!c_.ticket.compare_exchange_weak(
              t, t + 1, std::memory_order_acq_rel, std::memory_order_acquire)) {
        continue;
      }
      auto ret = e.takeItem();
      if (responsibleForAdvance(t)) {
        advanceHead(s);
      }
      return ret;
    }
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE bool tryDequeueWaitElem(
      Entry& e,
      Ticket t,
      const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
    if (LIKELY(e.tryWaitUntil(deadline))) {
      return true;
    }
    return t < producerTicket();
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE const T* tryPeekUntil(
      const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
    Segment* s = head();
    Ticket t = consumerTicket();
    DCHECK_GE(t, s->minTicket());
    DCHECK_LT(t, (s->minTicket() + SegmentSize));
    size_t idx = index(t);
    Entry& e = s->entry(idx);
    if (UNLIKELY(!tryDequeueWaitElem(e, t, deadline))) {
      return nullptr;
    }
    return e.peekItem();
  }

  FOLLY_ALWAYS_INLINE
  Segment* findSegment(Segment* s, const Ticket t) noexcept {
    while (UNLIKELY(t >= (s->minTicket() + SegmentSize))) {
      s = getAllocNextSegment(s, t);
      DCHECK(s);
    }
    return s;
  }

  Segment* getAllocNextSegment(Segment* s, Ticket t) noexcept {
    Segment* next = s->nextSegment();
    if (!next) {
      DCHECK_GE(t, s->minTicket() + SegmentSize);
      auto diff = t - (s->minTicket() + SegmentSize);
      if (diff > 0) {
        auto dur = std::chrono::microseconds(diff);
        auto deadline = std::chrono::steady_clock::now() + dur;
        WaitOptions opt;
        opt.spin_max(dur);
        detail::spin_pause_until(
            deadline, opt, [s] { return s->nextSegment(); });
        next = s->nextSegment();
        if (next) {
          return next;
        }
      }
      next = allocNextSegment(s);
    }
    DCHECK(next);
    return next;
  }

  Segment* allocNextSegment(Segment* s) {
    auto t = s->minTicket() + SegmentSize;
    Segment* next = new Segment(t);
    next->acquire_ref_safe(); // defined in hazptr_obj_base_linked
    if (!s->casNextSegment(next)) {
      delete next;
      next = s->nextSegment();
    }
    DCHECK(next);
    return next;
  }

  void advanceTail(Segment* s) noexcept {
    if (SPSC) {
      Segment* next = s->nextSegment();
      DCHECK(next);
      setTail(next);
    } else {
      Ticket t = s->minTicket() + SegmentSize;
      advanceTailToTicket(t);
    }
  }

  void advanceTailToTicket(Ticket t) noexcept {
    Segment* s = tail();
    while (s->minTicket() < t) {
      Segment* next = s->nextSegment();
      if (!next) {
        next = allocNextSegment(s);
      }
      DCHECK(next);
      casTail(s, next);
      s = tail();
    }
  }

  void advanceHead(Segment* s) noexcept {
    if (SPSC) {
      while (tail() == s) {
        /* Wait for producer to advance tail. */
        asm_volatile_pause();
      }
      Segment* next = s->nextSegment();
      DCHECK(next);
      setHead(next);
      reclaimSegment(s);
    } else {
      Ticket t = s->minTicket() + SegmentSize;
      advanceHeadToTicket(t);
    }
  }

  void advanceHeadToTicket(Ticket t) noexcept {
    /* Tail must not lag behind head. Otherwise, the algorithm cannot
       be certain about removal of segments. */
    advanceTailToTicket(t);
    Segment* s = head();
    if (SingleConsumer) {
      DCHECK_EQ(s->minTicket() + SegmentSize, t);
      Segment* next = s->nextSegment();
      DCHECK(next);
      setHead(next);
      reclaimSegment(s);
    } else {
      while (s->minTicket() < t) {
        Segment* next = s->nextSegment();
        DCHECK(next);
        if (casHead(s, next)) {
          reclaimSegment(s);
          s = next;
        }
      }
    }
  }

  void reclaimSegment(Segment* s) noexcept {
    if (SPSC) {
      delete s;
    } else {
      s->retire(); // defined in hazptr_obj_base_linked
    }
  }

  void cleanUpRemainingItems() {
    auto end = producerTicket();
    auto s = head();
    for (auto t = consumerTicket(); t < end; ++t) {
      if (t >= s->minTicket() + SegmentSize) {
        s = s->nextSegment();
      }
      DCHECK_LT(t, (s->minTicket() + SegmentSize));
      auto idx = index(t);
      auto& e = s->entry(idx);
      e.destroyItem();
    }
  }

  void reclaimRemainingSegments() {
    auto h = head();
    auto s = h->nextSegment();
    h->setNextSegment(nullptr);
    reclaimSegment(h);
    while (s) {
      auto next = s->nextSegment();
      delete s;
      s = next;
    }
  }

  FOLLY_ALWAYS_INLINE size_t index(Ticket t) const noexcept {
    return (t * Stride) & (SegmentSize - 1);
  }

  FOLLY_ALWAYS_INLINE bool responsibleForAlloc(Ticket t) const noexcept {
    return (t & (SegmentSize - 1)) == 0;
  }

  FOLLY_ALWAYS_INLINE bool responsibleForAdvance(Ticket t) const noexcept {
    return (t & (SegmentSize - 1)) == (SegmentSize - 1);
  }

  FOLLY_ALWAYS_INLINE Segment* head() const noexcept {
    return c_.head.load(std::memory_order_acquire);
  }

  FOLLY_ALWAYS_INLINE Segment* tail() const noexcept {
    return p_.tail.load(std::memory_order_acquire);
  }

  FOLLY_ALWAYS_INLINE Ticket producerTicket() const noexcept {
    return p_.ticket.load(std::memory_order_acquire);
  }

  FOLLY_ALWAYS_INLINE Ticket consumerTicket() const noexcept {
    return c_.ticket.load(std::memory_order_acquire);
  }

  void setHead(Segment* s) noexcept {
    DCHECK(SingleConsumer);
    c_.head.store(s, std::memory_order_relaxed);
  }

  void setTail(Segment* s) noexcept {
    DCHECK(SPSC);
    p_.tail.store(s, std::memory_order_release);
  }

  bool casHead(Segment*& s, Segment* next) noexcept {
    DCHECK(!SingleConsumer);
    return c_.head.compare_exchange_strong(
        s, next, std::memory_order_release, std::memory_order_acquire);
  }

  void casTail(Segment*& s, Segment* next) noexcept {
    DCHECK(!SPSC);
    p_.tail.compare_exchange_strong(
        s, next, std::memory_order_release, std::memory_order_relaxed);
  }

  FOLLY_ALWAYS_INLINE void setProducerTicket(Ticket t) noexcept {
    p_.ticket.store(t, std::memory_order_release);
  }

  FOLLY_ALWAYS_INLINE void setConsumerTicket(Ticket t) noexcept {
    c_.ticket.store(t, std::memory_order_release);
  }

  FOLLY_ALWAYS_INLINE Ticket fetchIncrementConsumerTicket() noexcept {
    if (SingleConsumer) {
      Ticket oldval = consumerTicket();
      setConsumerTicket(oldval + 1);
      return oldval;
    } else { // MC
      return c_.ticket.fetch_add(1, std::memory_order_acq_rel);
    }
  }

  FOLLY_ALWAYS_INLINE Ticket fetchIncrementProducerTicket() noexcept {
    if (SingleProducer) {
      Ticket oldval = producerTicket();
      setProducerTicket(oldval + 1);
      return oldval;
    } else { // MP
      return p_.ticket.fetch_add(1, std::memory_order_acq_rel);
    }
  }

  class Entry {
    folly::SaturatingSemaphore<MayBlock, Atom> flag_;
    typename std::aligned_storage<sizeof(T), alignof(T)>::type item_;

   public:
    template <typename Arg>
    FOLLY_ALWAYS_INLINE void putItem(Arg&& arg) {
      new (&item_) T(std::forward<Arg>(arg));
      flag_.post();
    }

    FOLLY_ALWAYS_INLINE void takeItem(T& item) noexcept {
      flag_.wait();
      getItem(item);
    }

    FOLLY_ALWAYS_INLINE folly::Optional<T> takeItem() noexcept {
      flag_.wait();
      return getItem();
    }

    FOLLY_ALWAYS_INLINE const T* peekItem() noexcept {
      flag_.wait();
      return itemPtr();
    }

    template <typename Clock, typename Duration>
    FOLLY_ALWAYS_INLINE bool tryWaitUntil(
        const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
      // wait-options from benchmarks on contended queues:
      auto const opt =
          flag_.wait_options().spin_max(std::chrono::microseconds(10));
      return flag_.try_wait_until(deadline, opt);
    }

    FOLLY_ALWAYS_INLINE void destroyItem() noexcept {
      itemPtr()->~T();
    }

   private:
    FOLLY_ALWAYS_INLINE void getItem(T& item) noexcept {
      item = std::move(*(itemPtr()));
      destroyItem();
    }

    FOLLY_ALWAYS_INLINE folly::Optional<T> getItem() noexcept {
      folly::Optional<T> ret = std::move(*(itemPtr()));
      destroyItem();
      return ret;
    }

    FOLLY_ALWAYS_INLINE T* itemPtr() noexcept {
      return static_cast<T*>(static_cast<void*>(&item_));
    }
  }; // Entry

  class Segment : public hazptr_obj_base_linked<Segment, Atom> {
    Atom<Segment*> next_{nullptr};
    const Ticket min_;
    alignas(Align) Entry b_[SegmentSize];

   public:
    explicit Segment(const Ticket t) noexcept : min_(t) {}

    Segment* nextSegment() const noexcept {
      return next_.load(std::memory_order_acquire);
    }

    void setNextSegment(Segment* next) {
      next_.store(next, std::memory_order_relaxed);
    }

    bool casNextSegment(Segment* next) noexcept {
      Segment* expected = nullptr;
      return next_.compare_exchange_strong(
          expected, next, std::memory_order_release, std::memory_order_relaxed);
    }

    FOLLY_ALWAYS_INLINE Ticket minTicket() const noexcept {
      DCHECK_EQ((min_ & (SegmentSize - 1)), Ticket(0));
      return min_;
    }

    FOLLY_ALWAYS_INLINE Entry& entry(size_t index) noexcept {
      return b_[index];
    }

    template <typename S>
    void push_links(bool m, S& s) {
      if (m == false) { // next_ is immutable
        auto p = nextSegment();
        if (p) {
          s.push(p);
        }
      }
    }
  }; // Segment

}; // UnboundedQueue

/* Aliases */

template <
    typename T,
    bool MayBlock,
    size_t LgSegmentSize = 8,
    size_t LgAlign = constexpr_log2(hardware_destructive_interference_size),
    template <typename> class Atom = std::atomic>
using USPSCQueue =
    UnboundedQueue<T, true, true, MayBlock, LgSegmentSize, LgAlign, Atom>;

template <
    typename T,
    bool MayBlock,
    size_t LgSegmentSize = 8,
    size_t LgAlign = constexpr_log2(hardware_destructive_interference_size),
    template <typename> class Atom = std::atomic>
using UMPSCQueue =
    UnboundedQueue<T, false, true, MayBlock, LgSegmentSize, LgAlign, Atom>;

template <
    typename T,
    bool MayBlock,
    size_t LgSegmentSize = 8,
    size_t LgAlign = constexpr_log2(hardware_destructive_interference_size),
    template <typename> class Atom = std::atomic>
using USPMCQueue =
    UnboundedQueue<T, true, false, MayBlock, LgSegmentSize, LgAlign, Atom>;

template <
    typename T,
    bool MayBlock,
    size_t LgSegmentSize = 8,
    size_t LgAlign = constexpr_log2(hardware_destructive_interference_size),
    template <typename> class Atom = std::atomic>
using UMPMCQueue =
    UnboundedQueue<T, false, false, MayBlock, LgSegmentSize, LgAlign, Atom>;

} // namespace folly