From 67300dac6ed429facd7c6f5e03b7c245e32f8383 Mon Sep 17 00:00:00 2001 From: Zheng Yan Date: Tue, 19 Dec 2017 19:40:14 -0800 Subject: [PATCH] Add readWithPriority to PriorityMPMCQueue Summary: Add readWithPriority to PriorityMPMCQueue Reviewed By: yfeldblum Differential Revision: D6606604 fbshipit-source-id: 07382e3c6d0f53a8727c059410c0577cf0aac82c --- folly/PriorityMPMCQueue.h | 4 ++++ folly/test/PriorityMPMCQueueTest.cpp | 21 +++++++++++++++++++++ 2 files changed, 25 insertions(+) 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 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)); + } +} -- 2.34.1