PriorityUnboundedBlockingQueue

Summary: [Folly] `PriorityUnboundedBlockingQueue`, paralleling `UnboundedBlockingQueue` and `PriorityLifoSemMPMCQueue`.

Reviewed By: andriigrynenko

Differential Revision: D14536886

fbshipit-source-id: b3ea04c3f67d3e1e9d80626d963be70de1ae79ba

CMakeLists.txt

@@ -529,6 +529,8 @@ if (BUILD_TESTS)
       TEST timed_drivable_executor_test SOURCES TimedDrivableExecutorTest.cpp
 
     DIRECTORY executors/task_queue/test/
+      TEST priority_unbounded_blocking_queue_test
+        SOURCES PriorityUnboundedBlockingQueueTest.cpp
       TEST unbounded_blocking_queue_test SOURCES UnboundedBlockingQueueTest.cpp
 
     DIRECTORY experimental/test/

folly/executors/task_queue/PriorityUnboundedBlockingQueue.h

+/*
+ * Copyright 2019-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/ConstexprMath.h>
+#include <folly/Executor.h>
+#include <folly/concurrency/PriorityUnboundedQueueSet.h>
+#include <folly/executors/task_queue/BlockingQueue.h>
+#include <folly/lang/Exception.h>
+#include <folly/synchronization/LifoSem.h>
+
+namespace folly {
+
+template <class T>
+class PriorityUnboundedBlockingQueue : public BlockingQueue<T> {
+ public:
+  // Note: To use folly::Executor::*_PRI, for numPriorities == 2
+  //       MID_PRI and HI_PRI are treated at the same priority level.
+  explicit PriorityUnboundedBlockingQueue(uint8_t numPriorities)
+      : queue_(numPriorities) {}
+
+  uint8_t getNumPriorities() override {
+    return queue_.priorities();
+  }
+
+  // Add at medium priority by default
+  BlockingQueueAddResult add(T item) override {
+    return addWithPriority(std::move(item), folly::Executor::MID_PRI);
+  }
+
+  BlockingQueueAddResult addWithPriority(T item, int8_t priority) override {
+    queue_.at_priority(translatePriority(priority)).enqueue(std::move(item));
+    return sem_.post();
+  }
+
+  T take() override {
+    sem_.wait();
+    return dequeue();
+  }
+
+  folly::Optional<T> try_take() {
+    if (!sem_.try_wait()) {
+      return none;
+    }
+    return dequeue();
+  }
+
+  folly::Optional<T> try_take_for(std::chrono::milliseconds time) override {
+    if (!sem_.try_wait_for(time)) {
+      return none;
+    }
+    return dequeue();
+  }
+
+  size_t size() override {
+    return queue_.size();
+  }
+
+  size_t sizeGuess() const {
+    return queue_.size();
+  }
+
+ private:
+  size_t translatePriority(int8_t const priority) {
+    int8_t const priorities = queue_.priorities();
+    int8_t const value = (priorities - 1) / 2 + priority;
+    int8_t const lo = 0;
+    int8_t const hi = priorities - 1;
+    return hi - constexpr_clamp(value, lo, hi);
+  }
+
+  T dequeue() {
+    // must follow a successful sem wait
+    if (auto obj = queue_.try_dequeue()) {
+      return std::move(*obj);
+    }
+    terminate_with<std::logic_error>("bug in task queue");
+  }
+
+  LifoSem sem_;
+  PriorityUMPMCQueueSet<T, /* MayBlock = */ true> queue_;
+};
+
+} // namespace folly

folly/executors/task_queue/test/PriorityUnboundedBlockingQueueTest.cpp

+/*
+ * Copyright 2019-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.
+ */
+
+#include <vector>
+
+#include <folly/container/Enumerate.h>
+#include <folly/executors/task_queue/PriorityUnboundedBlockingQueue.h>
+#include <folly/portability/GTest.h>
+#include <folly/synchronization/Baton.h>
+
+using namespace folly;
+
+class PriorityUnboundedBlockingQueueTest : public testing::Test {};
+
+TEST_F(PriorityUnboundedBlockingQueueTest, push_pop) {
+  PriorityUnboundedBlockingQueue<int> q(3);
+  q.add(42);
+  EXPECT_EQ(42, q.take());
+}
+
+TEST_F(PriorityUnboundedBlockingQueueTest, multiple_push_pop) {
+  PriorityUnboundedBlockingQueue<int> q(3);
+  q.add(42);
+  q.add(77);
+  EXPECT_EQ(42, q.take());
+  EXPECT_EQ(77, q.take());
+}
+
+TEST_F(PriorityUnboundedBlockingQueueTest, size) {
+  PriorityUnboundedBlockingQueue<int> q(3);
+  EXPECT_EQ(0, q.size());
+  q.add(42);
+  EXPECT_EQ(1, q.size());
+  q.take();
+  EXPECT_EQ(0, q.size());
+}
+
+TEST_F(PriorityUnboundedBlockingQueueTest, concurrent_push_pop) {
+  PriorityUnboundedBlockingQueue<int> q(3);
+  Baton<> b1, b2;
+  std::thread t([&] {
+    b1.post();
+    EXPECT_EQ(42, q.take());
+    EXPECT_EQ(0, q.size());
+    b2.post();
+  });
+  b1.wait();
+  q.add(42);
+  b2.wait();
+  EXPECT_EQ(0, q.size());
+  t.join();
+}
+
+TEST_F(PriorityUnboundedBlockingQueueTest, priority_order) {
+  PriorityUnboundedBlockingQueue<int> q(3);
+  EXPECT_EQ(0, q.size());
+  q.addWithPriority(27, 0);
+  q.addWithPriority(42, 1);
+  q.addWithPriority(55, 0);
+  q.addWithPriority(12, -1);
+  EXPECT_EQ(4, q.size());
+  EXPECT_EQ(42, q.take());
+  EXPECT_EQ(27, q.take());
+  EXPECT_EQ(55, q.take());
+  EXPECT_EQ(12, q.take());
+  EXPECT_EQ(0, q.size());
+}
+
+// Since PrioritizedBlockingQueueSetTest implements folly::BlockingQueue<T>,
+// addWithPriority method has to accept priority as int_8. This means invalid
+// values for priority (such as negative or very large numbers) might get
+// passed. Verify this behavior.
+TEST_F(PriorityUnboundedBlockingQueueTest, invalid_priorities) {
+  PriorityUnboundedBlockingQueue<int> q(2);
+  q.addWithPriority(1, -1); // expected to be converted to the lowest priority
+  q.addWithPriority(2, 50); // expected to be converted to the highest priority
+
+  EXPECT_EQ(q.take(), 2);
+  EXPECT_EQ(q.take(), 1);
+}