DynamicBoundedQueue.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 <folly/concurrency/CacheLocality.h>
#include <folly/concurrency/UnboundedQueue.h>

#include <glog/logging.h>

#include <atomic>
#include <chrono>

namespace folly {


template <typename T>
struct DefaultWeightFn {
  template <typename Arg>
  uint64_t operator()(Arg&&) const noexcept {
    return 1;
  }
};

template <
    typename T,
    bool SingleProducer,
    bool SingleConsumer,
    bool MayBlock,
    size_t LgSegmentSize = 8,
    size_t LgAlign = 7,
    typename WeightFn = DefaultWeightFn<T>,
    template <typename> class Atom = std::atomic>
class DynamicBoundedQueue {
  using Weight = uint64_t;

  enum WaitingState : uint32_t {
    NOTWAITING = 0,
    WAITING = 1,
  };

  static constexpr bool SPSC = SingleProducer && SingleConsumer;
  static constexpr size_t Align = 1u << LgAlign;

  static_assert(LgAlign < 16, "LgAlign must be < 16");


  // Read mostly by producers
  alignas(Align) Atom<Weight> debit_; // written frequently only by producers
  Atom<Weight> capacity_; // written rarely by capacity resets

  // Read mostly by consumers
  alignas(Align) Atom<Weight> credit_; // written frequently only by consumers
  Atom<Weight> threshold_; // written rarely only by capacity resets

  // Normally written and read rarely by producers and consumers
  // May be read frequently by producers when capacity is full
  alignas(Align) Atom<Weight> transfer_;
  detail::Futex<Atom> waiting_;

  // Underlying unbounded queue
  UnboundedQueue<
      T,
      SingleProducer,
      SingleConsumer,
      MayBlock,
      LgSegmentSize,
      LgAlign,
      Atom>
      q_;

 public:
  explicit DynamicBoundedQueue(Weight capacity)
      : debit_(0),
        capacity_(capacity + threshold(capacity)), // capacity slack
        credit_(0),
        threshold_(threshold(capacity)),
        transfer_(0),
        waiting_(0) {}

  ~DynamicBoundedQueue() {}


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

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

  FOLLY_ALWAYS_INLINE bool try_enqueue(const T& v) {
    return tryEnqueueImpl(v);
  }

  FOLLY_ALWAYS_INLINE bool try_enqueue(T&& v) {
    return tryEnqueueImpl(std::move(v));
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE bool try_enqueue_until(
      const T& v,
      const std::chrono::time_point<Clock, Duration>& deadline) {
    return tryEnqueueUntilImpl(v, deadline);
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE bool try_enqueue_until(
      T&& v,
      const std::chrono::time_point<Clock, Duration>& deadline) {
    return tryEnqueueUntilImpl(std::move(v), deadline);
  }

  template <typename Rep, typename Period>
  FOLLY_ALWAYS_INLINE bool try_enqueue_for(
      const T& v,
      const std::chrono::duration<Rep, Period>& duration) {
    return tryEnqueueForImpl(v, duration);
  }

  template <typename Rep, typename Period>
  FOLLY_ALWAYS_INLINE bool try_enqueue_for(
      T&& v,
      const std::chrono::duration<Rep, Period>& duration) {
    return tryEnqueueForImpl(std::move(v), duration);
  }


  FOLLY_ALWAYS_INLINE void dequeue(T& elem) {
    q_.dequeue(elem);
    addCredit(WeightFn()(elem));
  }

  FOLLY_ALWAYS_INLINE bool try_dequeue(T& elem) {
    if (q_.try_dequeue(elem)) {
      addCredit(WeightFn()(elem));
      return true;
    }
    return false;
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE bool try_dequeue_until(
      T& elem,
      const std::chrono::time_point<Clock, Duration>& deadline) {
    if (q_.try_dequeue_until(elem, deadline)) {
      addCredit(WeightFn()(elem));
      return true;
    }
    return false;
  }

  template <typename Rep, typename Period>
  FOLLY_ALWAYS_INLINE bool try_dequeue_for(
      T& elem,
      const std::chrono::duration<Rep, Period>& duration) {
    if (q_.try_dequeue_for(elem, duration)) {
      addCredit(WeightFn()(elem));
      return true;
    }
    return false;
  }


  void reset_capacity(Weight capacity) noexcept {
    Weight thresh = threshold(capacity);
    capacity_.store(capacity + thresh, std::memory_order_release);
    threshold_.store(thresh, std::memory_order_release);
  }

  Weight weight() const noexcept {
    auto d = getDebit();
    auto c = getCredit();
    auto t = getTransfer();
    return d > (c + t) ? d - (c + t) : 0;
  }

  size_t size() const noexcept {
    return q_.size();
  }

  bool empty() const noexcept {
    return q_.empty();
  }

 private:

  // Calculation of threshold to move credits in bulk from consumers
  // to producers
  constexpr Weight threshold(Weight capacity) const noexcept {
    return capacity / 10;
  }

  // Functions called frequently by producers

  template <typename Arg>
  FOLLY_ALWAYS_INLINE void enqueueImpl(Arg&& v) {
    tryEnqueueUntilImpl(
        std::forward<Arg>(v), std::chrono::steady_clock::time_point::max());
  }

  template <typename Arg>
  FOLLY_ALWAYS_INLINE bool tryEnqueueImpl(Arg&& v) {
    return tryEnqueueUntilImpl(
        std::forward<Arg>(v), std::chrono::steady_clock::time_point::min());
  }

  template <typename Clock, typename Duration, typename Arg>
  FOLLY_ALWAYS_INLINE bool tryEnqueueUntilImpl(
      Arg&& v,
      const std::chrono::time_point<Clock, Duration>& deadline) {
    Weight weight = WeightFn()(std::forward<Arg>(v));
    if (LIKELY(tryAddDebit(weight))) {
      q_.enqueue(std::forward<Arg>(v));
      return true;
    }
    return tryEnqueueUntilSlow(std::forward<Arg>(v), deadline);
  }

  template <typename Rep, typename Period, typename Arg>
  FOLLY_ALWAYS_INLINE bool tryEnqueueForImpl(
      Arg&& v,
      const std::chrono::duration<Rep, Period>& duration) {
    if (LIKELY(tryEnqueueImpl(std::forward<Arg>(v)))) {
      return true;
    }
    auto deadline = std::chrono::steady_clock::now() + duration;
    return tryEnqueueUntilSlow(std::forward<Arg>(v), deadline);
  }

  FOLLY_ALWAYS_INLINE bool tryAddDebit(Weight weight) noexcept {
    Weight capacity = getCapacity();
    Weight before = fetchAddDebit(weight);
    if (LIKELY(before + weight <= capacity)) {
      return true;
    } else {
      subDebit(weight);
      return false;
    }
  }

  FOLLY_ALWAYS_INLINE Weight getCapacity() const noexcept {
    return capacity_.load(std::memory_order_acquire);
  }

  FOLLY_ALWAYS_INLINE Weight fetchAddDebit(Weight weight) noexcept {
    Weight before;
    if (SingleProducer) {
      before = getDebit();
      debit_.store(before + weight, std::memory_order_relaxed);
    } else {
      before = debit_.fetch_add(weight, std::memory_order_acq_rel);
    }
    return before;
  }

  FOLLY_ALWAYS_INLINE Weight getDebit() const noexcept {
    return debit_.load(std::memory_order_acquire);
  }

  // Functions called frequently by consumers

  FOLLY_ALWAYS_INLINE void addCredit(Weight weight) noexcept {
    Weight before = fetchAddCredit(weight);
    Weight thresh = getThreshold();
    if (before + weight >= thresh && before < thresh) {
      transferCredit();
    }
  }

  FOLLY_ALWAYS_INLINE Weight fetchAddCredit(Weight weight) noexcept {
    Weight before;
    if (SingleConsumer) {
      before = getCredit();
      credit_.store(before + weight, std::memory_order_relaxed);
    } else {
      before = credit_.fetch_add(weight, std::memory_order_acq_rel);
    }
    return before;
  }

  FOLLY_ALWAYS_INLINE Weight getCredit() const noexcept {
    return credit_.load(std::memory_order_acquire);
  }

  FOLLY_ALWAYS_INLINE Weight getThreshold() const noexcept {
    return threshold_.load(std::memory_order_acquire);
  }

  void subDebit(Weight weight) noexcept {
    Weight before;
    if (SingleProducer) {
      before = getDebit();
      debit_.store(before - weight, std::memory_order_relaxed);
    } else {
      before = debit_.fetch_sub(weight, std::memory_order_acq_rel);
    }
    DCHECK_GE(before, weight);
  }

  template <typename Clock, typename Duration, typename Arg>
  bool tryEnqueueUntilSlow(
      Arg&& v,
      const std::chrono::time_point<Clock, Duration>& deadline) {
    Weight weight = WeightFn()(std::forward<Arg>(v));
    if (canEnqueue(deadline, weight)) {
      q_.enqueue(std::forward<Arg>(v));
      return true;
    } else {
      return false;
    }
  }

  template <typename Clock, typename Duration>
  bool canEnqueue(
      const std::chrono::time_point<Clock, Duration>& deadline,
      Weight weight) noexcept {
    Weight capacity = getCapacity();
    while (true) {
      tryReduceDebit();
      Weight debit = getDebit();
      if ((debit + weight <= capacity) && tryAddDebit(weight)) {
        return true;
      }
      if (deadline < Clock::time_point::max() && Clock::now() >= deadline) {
        return false;
      }
      if (MayBlock) {
        if (canBlock(weight, capacity)) {
          detail::futexWaitUntil(&waiting_, WAITING, deadline);
        }
      } else {
        asm_volatile_pause();
      }
    }
  }

  bool canBlock(Weight weight, Weight capacity) noexcept {
    waiting_.store(WAITING, std::memory_order_relaxed);
    std::atomic_thread_fence(std::memory_order_seq_cst);
    tryReduceDebit();
    Weight debit = getDebit();
    return debit + weight > capacity;
  }

  bool tryReduceDebit() noexcept {
    Weight w = takeTransfer();
    if (w > 0) {
      subDebit(w);
    }
    return w > 0;
  }

  Weight takeTransfer() noexcept {
    Weight w = getTransfer();
    if (w > 0) {
      w = transfer_.exchange(0, std::memory_order_acq_rel);
    }
    return w;
  }

  Weight getTransfer() const noexcept {
    return transfer_.load(std::memory_order_acquire);
  }

  void transferCredit() noexcept {
    Weight credit = takeCredit();
    transfer_.fetch_add(credit, std::memory_order_acq_rel);
    if (MayBlock) {
      std::atomic_thread_fence(std::memory_order_seq_cst);
      waiting_.store(NOTWAITING, std::memory_order_relaxed);
      detail::futexWake(&waiting_);
    }
  }

  Weight takeCredit() noexcept {
    Weight credit;
    if (SingleConsumer) {
      credit = credit_.load(std::memory_order_relaxed);
      credit_.store(0, std::memory_order_relaxed);
    } else {
      credit = credit_.exchange(0, std::memory_order_acq_rel);
    }
    return credit;
  }

}; // DynamicBoundedQueue


template <
    typename T,
    bool MayBlock,
    size_t LgSegmentSize = 8,
    size_t LgAlign = 7,
    typename WeightFn = DefaultWeightFn<T>,
    template <typename> class Atom = std::atomic>
using DSPSCQueue = DynamicBoundedQueue<
    T,
    true,
    true,
    MayBlock,
    LgSegmentSize,
    LgAlign,
    WeightFn,
    Atom>;

template <
    typename T,
    bool MayBlock,
    size_t LgSegmentSize = 8,
    size_t LgAlign = 7,
    typename WeightFn = DefaultWeightFn<T>,
    template <typename> class Atom = std::atomic>
using DMPSCQueue = DynamicBoundedQueue<
    T,
    false,
    true,
    MayBlock,
    LgSegmentSize,
    LgAlign,
    WeightFn,
    Atom>;

template <
    typename T,
    bool MayBlock,
    size_t LgSegmentSize = 8,
    size_t LgAlign = 7,
    typename WeightFn = DefaultWeightFn<T>,
    template <typename> class Atom = std::atomic>
using DSPMCQueue = DynamicBoundedQueue<
    T,
    true,
    false,
    MayBlock,
    LgSegmentSize,
    LgAlign,
    WeightFn,
    Atom>;

template <
    typename T,
    bool MayBlock,
    size_t LgSegmentSize = 8,
    size_t LgAlign = 7,
    typename WeightFn = DefaultWeightFn<T>,
    template <typename> class Atom = std::atomic>
using DMPMCQueue = DynamicBoundedQueue<
    T,
    false,
    false,
    MayBlock,
    LgSegmentSize,
    LgAlign,
    WeightFn,
    Atom>;

} // namespace folly