FlatCombiningPriorityQueue.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 <mutex>
#include <queue>

#include <folly/Optional.h>
#include <folly/detail/Futex.h>
#include <folly/experimental/flat_combining/FlatCombining.h>
#include <glog/logging.h>

namespace folly {


template <
    typename T,
    typename PriorityQueue = std::priority_queue<T>,
    typename Mutex = std::mutex,
    template <typename> class Atom = std::atomic>
class FlatCombiningPriorityQueue
    : public folly::FlatCombining<
          FlatCombiningPriorityQueue<T, PriorityQueue, Mutex, Atom>,
          Mutex,
          Atom> {
  using FCPQ = FlatCombiningPriorityQueue<T, PriorityQueue, Mutex, Atom>;
  using FC = folly::FlatCombining<FCPQ, Mutex, Atom>;

 public:
  template <
      typename... PQArgs,
      typename = decltype(PriorityQueue(std::declval<PQArgs>()...))>
  explicit FlatCombiningPriorityQueue(
      // Concurrent priority queue parameter
      const size_t maxSize = 0,
      // Flat combining parameters
      const bool dedicated = true,
      const uint32_t numRecs = 0,
      const uint32_t maxOps = 0,
      // (Sequential) PriorityQueue Parameters
      PQArgs... args)
      : FC(dedicated, numRecs, maxOps),
        maxSize_(maxSize),
        pq_(std::forward<PQArgs>(args)...) {}

  bool empty() const {
    bool res;
    auto fn = [&] { res = pq_.empty(); };
    const_cast<FCPQ*>(this)->requestFC(fn);
    return res;
  }

  size_t size() const {
    size_t res;
    auto fn = [&] { res = pq_.size(); };
    const_cast<FCPQ*>(this)->requestFC(fn);
    return res;
  }

  bool try_push(const T& val) {
    return try_push_impl(
        val, std::chrono::time_point<std::chrono::steady_clock>::min());
  }

  bool try_pop(T& val) {
    return try_pop_impl(
        val, std::chrono::time_point<std::chrono::steady_clock>::min());
  }

  bool try_peek(T& val) {
    return try_peek_impl(
        val, std::chrono::time_point<std::chrono::steady_clock>::min());
  }

  void push(const T& val) {
    try_push_impl(
        val, std::chrono::time_point<std::chrono::steady_clock>::max());
  }

  void pop(T& val) {
    try_pop_impl(
        val, std::chrono::time_point<std::chrono::steady_clock>::max());
  }

  void peek(T& val) {
    try_peek_impl(
        val, std::chrono::time_point<std::chrono::steady_clock>::max());
  }

  folly::Optional<T> try_pop() {
    T val;
    if (try_pop(val)) {
      return std::move(val);
    }
    return folly::none;
  }

  folly::Optional<T> try_peek() {
    T val;
    if (try_peek(val)) {
      return std::move(val);
    }
    return folly::none;
  }

  template <typename Rep, typename Period>
  folly::Optional<T> try_pop_for(
      const std::chrono::duration<Rep, Period>& timeout) {
    T val;
    if (try_pop(val) ||
        try_pop_impl(val, std::chrono::steady_clock::now() + timeout)) {
      return std::move(val);
    }
    return folly::none;
  }

  template <typename Rep, typename Period>
  bool try_push_for(
      const T& val,
      const std::chrono::duration<Rep, Period>& timeout) {
    return (
        try_push(val) ||
        try_push_impl(val, std::chrono::steady_clock::now() + timeout));
  }

  template <typename Rep, typename Period>
  folly::Optional<T> try_peek_for(
      const std::chrono::duration<Rep, Period>& timeout) {
    T val;
    if (try_peek(val) ||
        try_peek_impl(val, std::chrono::steady_clock::now() + timeout)) {
      return std::move(val);
    }
    return folly::none;
  }

  template <typename Clock, typename Duration>
  folly::Optional<T> try_pop_until(
      const std::chrono::time_point<Clock, Duration>& deadline) {
    T val;
    if (try_pop_impl(val, deadline)) {
      return std::move(val);
    }
    return folly::none;
  }

  template <typename Clock, typename Duration>
  bool try_push_until(
      const T& val,
      const std::chrono::time_point<Clock, Duration>& deadline) {
    return try_push_impl(val, deadline);
  }

  template <typename Clock, typename Duration>
  folly::Optional<T> try_peek_until(
      const std::chrono::time_point<Clock, Duration>& deadline) {
    T val;
    if (try_peek_impl(val, deadline)) {
      return std::move(val);
    }
    return folly::none;
  }

 private:
  size_t maxSize_;
  PriorityQueue pq_;
  detail::Futex<Atom> empty_{};
  detail::Futex<Atom> full_{};

  bool isTrue(detail::Futex<Atom>& futex) {
    return futex.load(std::memory_order_relaxed) != 0;
  }

  void setFutex(detail::Futex<Atom>& futex, uint32_t val) {
    futex.store(val, std::memory_order_relaxed);
  }

  bool futexSignal(detail::Futex<Atom>& futex) {
    if (isTrue(futex)) {
      setFutex(futex, 0);
      return true;
    } else {
      return false;
    }
  }

  template <typename Clock, typename Duration>
  bool try_push_impl(
      const T& val,
      const std::chrono::time_point<Clock, Duration>& when);

  template <typename Clock, typename Duration>
  bool try_pop_impl(
      T& val,
      const std::chrono::time_point<Clock, Duration>& when);

  template <typename Clock, typename Duration>
  bool try_peek_impl(
      T& val,
      const std::chrono::time_point<Clock, Duration>& when);
};


template <
    typename T,
    typename PriorityQueue,
    typename Mutex,
    template <typename> class Atom>
template <typename Clock, typename Duration>
inline bool
FlatCombiningPriorityQueue<T, PriorityQueue, Mutex, Atom>::try_push_impl(
    const T& val,
    const std::chrono::time_point<Clock, Duration>& when) {
  while (true) {
    bool res;
    bool wake;

    auto fn = [&] {
      if (maxSize_ > 0 && pq_.size() == maxSize_) {
        setFutex(full_, 1);
        res = false;
        return;
      }
      DCHECK(maxSize_ == 0 || pq_.size() < maxSize_);
      try {
        pq_.push(val);
        wake = futexSignal(empty_);
        res = true;
        return;
      } catch (const std::bad_alloc&) {
        setFutex(full_, 1);
        res = false;
        return;
      }
    };
    this->requestFC(fn);

    if (res) {
      if (wake) {
        detail::futexWake(&empty_);
      }
      return true;
    }
    if (when == std::chrono::time_point<Clock>::min()) {
      return false;
    }
    while (isTrue(full_)) {
      if (when == std::chrono::time_point<Clock>::max()) {
        detail::futexWait(&full_, 1);
      } else {
        if (Clock::now() > when) {
          return false;
        } else {
          detail::futexWaitUntil(&full_, 1, when);
        }
      }
    } // inner while loop
  } // outer while loop
}

template <
    typename T,
    typename PriorityQueue,
    typename Mutex,
    template <typename> class Atom>
template <typename Clock, typename Duration>
inline bool
FlatCombiningPriorityQueue<T, PriorityQueue, Mutex, Atom>::try_pop_impl(
    T& val,
    const std::chrono::time_point<Clock, Duration>& when) {
  while (true) {
    bool res;
    bool wake;

    auto fn = [&] {
      res = !pq_.empty();
      if (res) {
        val = pq_.top();
        pq_.pop();
        wake = futexSignal(full_);
      } else {
        setFutex(empty_, 1);
      }
    };
    this->requestFC(fn);

    if (res) {
      if (wake) {
        detail::futexWake(&full_);
      }
      return true;
    }
    while (isTrue(empty_)) {
      if (when == std::chrono::time_point<Clock>::max()) {
        detail::futexWait(&empty_, 1);
      } else {
        if (Clock::now() > when) {
          return false;
        } else {
          detail::futexWaitUntil(&empty_, 1, when);
        }
      }
    } // inner while loop
  } // outer while loop
}

template <
    typename T,
    typename PriorityQueue,
    typename Mutex,
    template <typename> class Atom>
template <typename Clock, typename Duration>
inline bool
FlatCombiningPriorityQueue<T, PriorityQueue, Mutex, Atom>::try_peek_impl(
    T& val,
    const std::chrono::time_point<Clock, Duration>& when) {
  while (true) {
    bool res;

    auto fn = [&] {
      res = !pq_.empty();
      if (res) {
        val = pq_.top();
      } else {
        setFutex(empty_, 1);
      }
    };
    this->requestFC(fn);

    if (res) {
      return true;
    }
    while (isTrue(empty_)) {
      if (when == std::chrono::time_point<Clock>::max()) {
        detail::futexWait(&empty_, 1);
      } else {
        if (Clock::now() > when) {
          return false;
        } else {
          detail::futexWaitUntil(&empty_, 1, when);
        }
      }
    } // inner while loop
  } // outer while loop
}

} // namespace folly