123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114 |
- /*
- * 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 <glog/logging.h>
- #include <algorithm>
- #include <vector>
- #include <folly/MPMCQueue.h>
- namespace folly {
- /// PriorityMPMCQueue is a thin wrapper on MPMCQueue, providing priorities
- /// by managing multiple underlying MPMCQueues. As of now, this does
- /// not implement a blocking interface. For the purposes of this
- /// class, lower number is higher priority
- template <
- typename T,
- template <typename> class Atom = std::atomic,
- bool Dynamic = false>
- class PriorityMPMCQueue {
- public:
- PriorityMPMCQueue(size_t numPriorities, size_t capacity) {
- CHECK_GT(numPriorities, 0);
- queues_.reserve(numPriorities);
- for (size_t i = 0; i < numPriorities; i++) {
- queues_.emplace_back(capacity);
- }
- }
- size_t getNumPriorities() {
- return queues_.size();
- }
- // Add at medium priority by default
- bool write(T&& item) {
- return writeWithPriority(std::move(item), getNumPriorities() / 2);
- }
- bool writeWithPriority(T&& item, size_t priority) {
- size_t queue = std::min(getNumPriorities() - 1, priority);
- CHECK_LT(queue, queues_.size());
- return queues_.at(queue).write(std::move(item));
- }
- bool writeWithPriority(
- T&& item,
- size_t priority,
- std::chrono::milliseconds timeout) {
- size_t queue = std::min(getNumPriorities() - 1, priority);
- CHECK_LT(queue, queues_.size());
- return queues_.at(queue).tryWriteUntil(
- std::chrono::steady_clock::now() + timeout, std::move(item));
- }
- bool read(T& item) {
- for (auto& q : queues_) {
- if (q.readIfNotEmpty(item)) {
- return true;
- }
- }
- return false;
- }
- bool readWithPriority(T& item, size_t priority) {
- return queues_[priority].readIfNotEmpty(item);
- }
- size_t size() const {
- size_t total_size = 0;
- for (auto& q : queues_) {
- // MPMCQueue can have a negative size if there are pending readers.
- // Since we don't expose a blocking interface this shouldn't happen,
- // But just in case we put a floor at 0
- total_size += std::max<ssize_t>(0, q.size());
- }
- return total_size;
- }
- size_t sizeGuess() const {
- size_t total_size = 0;
- for (auto& q : queues_) {
- // MPMCQueue can have a negative size if there are pending readers.
- // Since we don't expose a blocking interface this shouldn't happen,
- // But just in case we put a floor at 0
- total_size += std::max<ssize_t>(0, q.sizeGuess());
- }
- return total_size;
- }
- /// Returns true if there are no items available for dequeue
- bool isEmpty() const {
- return size() == 0;
- }
- private:
- std::vector<folly::MPMCQueue<T, Atom, Dynamic>> queues_;
- };
- } // namespace folly
|