folly/PriorityMPMCQueue.h

   1 /*
   2  * Copyright 2017 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 #pragma once
  17
  18 #include <glog/logging.h>
  19 #include <algorithm>
  20 #include <vector>
  21
  22 #include <folly/MPMCQueue.h>
  23
  24 namespace folly {
  25
  26 /// PriorityMPMCQueue is a thin wrapper on MPMCQueue, providing priorities
  27 /// by managing multiple underlying MPMCQueues. As of now, this does
  28 /// not implement a blocking interface. For the purposes of this
  29 /// class, lower number is higher priority
  30
  31 template <
  32     typename T,
  33     template <typename> class Atom = std::atomic,
  34     bool Dynamic = false>
  35 class PriorityMPMCQueue {
  36  public:
  37   PriorityMPMCQueue(size_t numPriorities, size_t capacity) {
  38     CHECK_GT(numPriorities, 0);
  39     queues_.reserve(numPriorities);
  40     for (size_t i = 0; i < numPriorities; i++) {
  41       queues_.emplace_back(capacity);
  42     }
  43   }
  44
  45   size_t getNumPriorities() {
  46     return queues_.size();
  47   }
  48
  49   // Add at medium priority by default
  50   bool write(T&& item) {
  51     return writeWithPriority(std::move(item), getNumPriorities() / 2);
  52   }
  53
  54   bool writeWithPriority(T&& item, size_t priority) {
  55     size_t queue = std::min(getNumPriorities() - 1, priority);
  56     CHECK_LT(queue, queues_.size());
  57     return queues_.at(queue).write(std::move(item));
  58   }
  59
  60   bool read(T& item) {
  61     for (auto& q : queues_) {
  62       if (q.readIfNotEmpty(item)) {
  63         return true;
  64       }
  65     }
  66     return false;
  67   }
  68
  69   size_t size() const {
  70     size_t total_size = 0;
  71     for (auto& q : queues_) {
  72       // MPMCQueue can have a negative size if there are pending readers.
  73       // Since we don't expose a blocking interface this shouldn't happen,
  74       // But just in case we put a floor at 0
  75       total_size += std::max<ssize_t>(0, q.size());
  76     }
  77     return total_size;
  78   }
  79
  80   size_t sizeGuess() const {
  81     size_t total_size = 0;
  82     for (auto& q : queues_) {
  83       // MPMCQueue can have a negative size if there are pending readers.
  84       // Since we don't expose a blocking interface this shouldn't happen,
  85       // But just in case we put a floor at 0
  86       total_size += std::max<ssize_t>(0, q.sizeGuess());
  87     }
  88     return total_size;
  89   }
  90
  91   /// Returns true if there are no items available for dequeue
  92   bool isEmpty() const {
  93     return size() == 0;
  94   }
  95
  96  private:
  97   std::vector<folly::MPMCQueue<T, Atom, Dynamic>> queues_;
  98 };
  99
 100 } // namespace folly