UnboundedQueue: Add support for single consumer try_peek operation.

Summary: Add a try_peek() member function valid only if SingleConsumer==true that returns folly::Optional<const T*> which (if has_value() is true) contains a pointer to the first element in the queue.

Reviewed By: djwatson

Differential Revision: D7969290

fbshipit-source-id: 3e15138f6c6523adce4da4cd99b74e0a83cd24d7

folly/concurrency/UnboundedQueue.h

@@ -82,16 +82,21 @@ namespace folly {
 ///         Extracts an element from the front of the queue. Waits
 ///         until an element is available if needed.
 ///     bool try_dequeue(T&);
+///     folly::Optional<T> try_dequeue();
 ///         Tries to extract an element from the front of the queue
-///         if available. Returns true if successful, false otherwise.
+///         if available.
 ///     bool try_dequeue_until(T&, time_point& deadline);
+///     folly::Optional<T> try_dequeue_until(time_point& deadline);
 ///         Tries to extract an element from the front of the queue
-///         if available until the specified deadline.  Returns true
-///         if successful, false otherwise.
+///         if available until the specified deadline.
 ///     bool try_dequeue_for(T&, duration&);
+///     folly::Optional<T> try_dequeue_for(duration&);
 ///         Tries to extract an element from the front of the queue if
-///         available for until the expiration of the specified
-///         duration.  Returns true if successful, false otherwise.
+///         available until the expiration of the specified duration.
+///     const T* try_peek();
+///         Returns pointer to the element at the front of the queue
+///         if available, or nullptr if the queue is empty. Only for
+///         SPSC and MPSC.
 ///
 ///   Secondary functions:
 ///     size_t size();
@@ -324,6 +329,13 @@ class UnboundedQueue {
     return tryDequeueUntil(std::chrono::steady_clock::now() + duration);
   }
 
+  /** try_peek */
+  FOLLY_ALWAYS_INLINE const T* try_peek() noexcept {
+    /* This function is supported only for USPSC and UMPSC queues. */
+    DCHECK(SingleConsumer);
+    return tryPeekUntil(std::chrono::steady_clock::time_point::min());
+  }
+
   /** size */
   size_t size() const noexcept {
     auto p = producerTicket();
@@ -483,6 +495,22 @@ class UnboundedQueue {
     return t < producerTicket();
   }
 
+  /** tryPeekUntil */
+  template <typename Clock, typename Duration>
+  FOLLY_ALWAYS_INLINE const T* tryPeekUntil(
+      const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
+    Segment* s = head();
+    Ticket t = consumerTicket();
+    DCHECK_GE(t, s->minTicket());
+    DCHECK_LT(t, (s->minTicket() + SegmentSize));
+    size_t idx = index(t);
+    Entry& e = s->entry(idx);
+    if (UNLIKELY(!tryDequeueWaitElem(e, t, deadline))) {
+      return nullptr;
+    }
+    return e.peekItem();
+  }
+
   /** findSegment */
   FOLLY_ALWAYS_INLINE
   Segment* findSegment(Segment* s, const Ticket t) const noexcept {
@@ -656,6 +684,11 @@ class UnboundedQueue {
       return getItem();
     }
 
+    FOLLY_ALWAYS_INLINE const T* peekItem() noexcept {
+      flag_.wait();
+      return itemPtr();
+    }
+
     template <typename Clock, typename Duration>
     FOLLY_ALWAYS_INLINE bool tryWaitUntil(
         const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
@@ -674,7 +707,6 @@ class UnboundedQueue {
     FOLLY_ALWAYS_INLINE folly::Optional<T> getItem() noexcept {
       folly::Optional<T> ret = std::move(*(itemPtr()));
       destroyItem();
-
       return ret;
     }
 

folly/concurrency/test/UnboundedQueueTest.cpp

@@ -103,6 +103,30 @@ TEST(UnboundedQueue, timeout) {
   timeout_test<UMPMC, true>();
 }
 
+template <template <typename, bool> class Q, bool MayBlock>
+void peek_test() {
+  Q<int, MayBlock> q;
+  auto res = q.try_peek();
+  ASSERT_FALSE(res);
+  for (int i = 0; i < 1000; ++i) {
+    q.enqueue(i);
+  }
+  for (int i = 0; i < 700; ++i) {
+    int v;
+    q.dequeue(v);
+  }
+  res = q.try_peek();
+  ASSERT_TRUE(res);
+  ASSERT_EQ(*res, 700);
+}
+
+TEST(UnboundedQueue, peek) {
+  peek_test<USPSC, false>();
+  peek_test<UMPSC, false>();
+  peek_test<USPSC, true>();
+  peek_test<UMPSC, true>();
+}
+
 template <typename ProdFunc, typename ConsFunc, typename EndFunc>
 inline uint64_t run_once(
     int nprod,
@@ -184,7 +208,17 @@ void enq_deq_test(const int nprod, const int ncons) {
     uint64_t mysum = 0;
     for (int i = tid; i < ops; i += ncons) {
       int v = -1;
+      int vpeek = -1;
 
+      if (SingleConsumer) {
+        while (true) {
+          auto res = q.try_peek();
+          if (res) {
+            vpeek = *res;
+            break;
+          }
+        }
+      }
       if ((i % 3) == 0) {
         while (!q.try_dequeue(v)) {
           /* keep trying */;
@@ -200,6 +234,9 @@ void enq_deq_test(const int nprod, const int ncons) {
       if (nprod == 1 && ncons == 1) {
         ASSERT_EQ(v, i);
       }
+      if (SingleConsumer) {
+        ASSERT_EQ(v, vpeek);
+      }
       mysum += v;
     }
     sum.fetch_add(mysum);