From: Zheng Yan <zyan@fb.com>
Date: Wed, 20 Dec 2017 03:40:14 +0000 (-0800)
Subject: Add readWithPriority to PriorityMPMCQueue
X-Git-Tag: v2017.12.25.00~10
X-Git-Url: http://plrg.eecs.uci.edu/git/?a=commitdiff_plain;h=67300dac6ed429facd7c6f5e03b7c245e32f8383;p=folly.git

Add readWithPriority to PriorityMPMCQueue

Summary: Add readWithPriority to PriorityMPMCQueue

Reviewed By: yfeldblum

Differential Revision: D6606604

fbshipit-source-id: 07382e3c6d0f53a8727c059410c0577cf0aac82c
---

diff --git a/folly/PriorityMPMCQueue.h b/folly/PriorityMPMCQueue.h
index fb4e4645..b12cbf33 100644
--- a/folly/PriorityMPMCQueue.h
+++ b/folly/PriorityMPMCQueue.h
@@ -66,6 +66,10 @@ class PriorityMPMCQueue {
     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_) {
diff --git a/folly/test/PriorityMPMCQueueTest.cpp b/folly/test/PriorityMPMCQueueTest.cpp
index a43cca0b..a129568f 100644
--- a/folly/test/PriorityMPMCQueueTest.cpp
+++ b/folly/test/PriorityMPMCQueueTest.cpp
@@ -73,3 +73,24 @@ TEST(PriorityMPMCQueue, TestPriorities) {
     EXPECT_EQ(6 - i, queue.sizeGuess());
   }
 }
+
+TEST(PriorityMPMCQueue, TestReadWithPriority) {
+  PriorityMPMCQueue<size_t> queue(3, 10);
+  EXPECT_TRUE(queue.isEmpty());
+  EXPECT_EQ(3, queue.getNumPriorities());
+
+  queue.writeWithPriority(2, 2);
+  queue.writeWithPriority(1, 1);
+  queue.writeWithPriority(0, 0);
+
+  EXPECT_FALSE(queue.isEmpty());
+  EXPECT_EQ(3, queue.size());
+  EXPECT_EQ(3, queue.sizeGuess());
+
+  size_t item;
+  for (int i = 0; i < 3; i++) {
+    EXPECT_TRUE(queue.readWithPriority(item, i));
+    EXPECT_EQ(i, item);
+    EXPECT_FALSE(queue.readWithPriority(item, i));
+  }
+}