fix problem: timeout dosen't take effect for add() of ThreadManager

Summary: timeout dosen't take effect for add() of ThreadManager

Reviewed By: andriigrynenko

Differential Revision: D8674748

fbshipit-source-id: a0336fb78182ad432b2d8e5beb61f785e59f0fdd

folly/PriorityMPMCQueue.h

@@ -57,6 +57,16 @@ class PriorityMPMCQueue {
     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)) {

folly/test/PriorityMPMCQueueTest.cpp

@@ -94,3 +94,36 @@ TEST(PriorityMPMCQueue, TestReadWithPriority) {
     EXPECT_FALSE(queue.readWithPriority(item, i));
   }
 }
+
+TEST(PriorityMPMCQueue, TestWriteWithPriorityAndTimeout) {
+  PriorityMPMCQueue<size_t> queue(5, 1);
+  EXPECT_TRUE(queue.isEmpty());
+  EXPECT_EQ(5, queue.getNumPriorities());
+
+  const auto timeout = std::chrono::milliseconds{30};
+  for (int i = 0; i < 5; i++) {
+    auto time_before = std::chrono::steady_clock::now();
+    EXPECT_TRUE(queue.writeWithPriority(i, i, timeout));
+    auto time_after = std::chrono::steady_clock::now();
+    EXPECT_LE(time_after - time_before, timeout);
+  }
+
+  // check writeWithPriority will wait for at least timeout if the queue is
+  // full.
+  auto time_before = std::chrono::steady_clock::now();
+  EXPECT_FALSE(queue.writeWithPriority(5, 0, timeout));
+  auto time_after = std::chrono::steady_clock::now();
+  EXPECT_GE(time_after - time_before, timeout);
+
+  EXPECT_FALSE(queue.isEmpty());
+  EXPECT_EQ(5, queue.size());
+  EXPECT_EQ(5, queue.sizeGuess());
+
+  size_t item;
+  for (int i = 0; i < 5; i++) {
+    queue.read(item);
+    EXPECT_EQ(i, item);
+    EXPECT_EQ(4 - i, queue.size());
+    EXPECT_EQ(4 - i, queue.sizeGuess());
+  }
+}