folly::coro::SmallUnboundedQueue

Summary: A smaller version of coro::UnboundedQueue to make AsyncPipe usable in memory-constrained programs.

Reviewed By: yfeldblum

Differential Revision: D29375067

fbshipit-source-id: 5d27e4514312d2748055be764af88c0688d46065

folly/experimental/coro/SmallUnboundedQueue.h

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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/experimental/channels/detail/AtomicQueue.h>
+#include <folly/experimental/coro/Baton.h>
+#include <folly/experimental/coro/Coroutine.h>
+#include <folly/experimental/coro/Mutex.h>
+#include <folly/experimental/coro/Task.h>
+
+#if FOLLY_HAS_COROUTINES
+
+namespace folly {
+namespace coro {
+namespace detail {
+template <bool UseMutex>
+struct SmallUnboundedQueueBase {
+  auto co_scoped_lock() { return ready_awaitable(true); }
+};
+template <>
+struct SmallUnboundedQueueBase<true> {
+  auto co_scoped_lock() { return mutex_.co_scoped_lock(); }
+  folly::coro::Mutex mutex_;
+};
+} // namespace detail
+
+// Alternative to coro::UnboundedQueue with much smaller memory size when empty
+// but lower throughput.
+// Substantially worse in multi-consumer case.
+// Only supports enqueue(T) and dequeue().
+
+template <typename T, bool SingleProducer = false, bool SingleConsumer = false>
+class SmallUnboundedQueue : detail::SmallUnboundedQueueBase<!SingleConsumer> {
+  struct Consumer {
+    void consume() { baton.post(); }
+    void canceled() { std::terminate(); }
+    folly::coro::Baton baton;
+  };
+
+ public:
+  ~SmallUnboundedQueue() { queue_.close(); }
+
+  template <typename U = T>
+  void enqueue(U&& val) {
+    queue_.push(T(std::forward<U>(val)));
+  }
+
+  folly::coro::Task<T> dequeue() {
+    FOLLY_MAYBE_UNUSED auto maybeLock = co_await this->co_scoped_lock();
+    if (buffer_.empty()) {
+      Consumer c;
+      if (queue_.wait(&c)) {
+        bool cancelled = false;
+        CancellationCallback cb(co_await co_current_cancellation_token, [&] {
+          if (queue_.cancelCallback()) {
+            cancelled = true;
+            c.baton.post();
+          }
+        });
+        co_await c.baton;
+        if (cancelled) {
+          co_yield co_cancelled;
+        }
+      }
+      buffer_ = queue_.getMessages();
+      DCHECK(!buffer_.empty());
+    }
+    SCOPE_EXIT { buffer_.pop(); };
+    co_return std::move(buffer_.front());
+  }
+
+ private:
+  folly::channels::detail::AtomicQueue<Consumer, T> queue_;
+  folly::channels::detail::Queue<T> buffer_;
+};
+
+} // namespace coro
+} // namespace folly
+
+#endif // FOLLY_HAS_COROUTINES

folly/experimental/coro/test/SmallUnboundedQueueTest.cpp

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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 <folly/Portability.h>
+
+#include <folly/CancellationToken.h>
+#include <folly/experimental/coro/BlockingWait.h>
+#include <folly/experimental/coro/Collect.h>
+#include <folly/experimental/coro/SmallUnboundedQueue.h>
+
+#include <folly/portability/GTest.h>
+
+#include <string>
+#include <thread>
+
+#if FOLLY_HAS_COROUTINES
+
+TEST(SmallUnboundedQueueTest, EnqueueDeque) {
+  folly::coro::SmallUnboundedQueue<std::string, true, true> queue;
+  constexpr auto val = "a string";
+  std::string val1 = val;
+  queue.enqueue(val1);
+  queue.enqueue(std::move(val1));
+  folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
+    for (int i = 0; i < 2; ++i) {
+      auto val2 = co_await queue.dequeue();
+      EXPECT_EQ(val2, val);
+    }
+  }());
+}
+
+TEST(SmallUnboundedQueueTest, DequeueWhileBlocking) {
+  folly::coro::SmallUnboundedQueue<int> queue;
+  folly::ManualExecutor ex;
+
+  auto fut = queue.dequeue().scheduleOn(&ex).start();
+  ex.drain();
+  EXPECT_FALSE(fut.isReady());
+
+  queue.enqueue(0);
+  ex.drain();
+  EXPECT_TRUE(fut.isReady());
+  EXPECT_EQ(std::move(fut).get(), 0);
+}
+
+TEST(SmallUnboundedQueueTest, EnqueueDequeMultiProducer) {
+  folly::coro::SmallUnboundedQueue<int, false, true> queue;
+  std::atomic<int> i = 0;
+
+  std::vector<std::thread> enqueuers;
+  for (int n = 0; n < 5; ++n) {
+    enqueuers.emplace_back([&] {
+      while (true) {
+        int next = i++;
+        if (next >= 100) {
+          break;
+        }
+        queue.enqueue(next);
+      }
+    });
+  }
+
+  folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
+    for (int n = 0; n < 100; ++n) {
+      co_await queue.dequeue();
+    }
+  }());
+
+  for (int n = 0; n < 5; ++n) {
+    enqueuers[n].join();
+  }
+}
+
+TEST(SmallUnboundedQueueTest, EnqueueDequeMultiConsumer) {
+  folly::coro::SmallUnboundedQueue<int, true, false> queue;
+  std::atomic<int> seen = 0;
+
+  std::vector<std::thread> dequeuers;
+  for (int n = 0; n < 5; ++n) {
+    dequeuers.emplace_back([&] {
+      folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
+        while (++seen <= 100) {
+          co_await queue.dequeue();
+        }
+      }());
+    });
+  }
+
+  for (int n = 0; n < 100; ++n) {
+    queue.enqueue(n);
+  }
+  for (int n = 0; n < 5; ++n) {
+    dequeuers[n].join();
+  }
+}
+
+TEST(SmallUnboundedQueueTest, EnqueueDequeMPMC) {
+  folly::coro::SmallUnboundedQueue<int, false, false> queue;
+  std::atomic<int> seen = 0, i = 0;
+
+  std::vector<std::thread> enqueuers;
+  for (int n = 0; n < 5; ++n) {
+    enqueuers.emplace_back([&] {
+      while (true) {
+        int next = i++;
+        if (next >= 100) {
+          break;
+        }
+        queue.enqueue(next);
+      }
+    });
+  }
+
+  std::vector<std::thread> dequeuers;
+  for (int n = 0; n < 5; ++n) {
+    dequeuers.emplace_back([&] {
+      folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
+        while (++seen <= 100) {
+          co_await queue.dequeue();
+        }
+      }());
+    });
+  }
+
+  for (int n = 0; n < 5; ++n) {
+    enqueuers[n].join();
+  }
+  for (int n = 0; n < 5; ++n) {
+    dequeuers[n].join();
+  }
+}
+
+TEST(SmallUnboundedQueueTest, CancelledDequeueThrowsOperationCancelled) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    folly::coro::SmallUnboundedQueue<int> queue;
+    folly::CancellationSource cancelSource;
+
+    co_await folly::coro::collectAll(
+        [&]() -> folly::coro::Task<void> {
+          EXPECT_THROW(
+              (co_await folly::coro::co_withCancellation(
+                  cancelSource.getToken(), queue.dequeue())),
+              folly::OperationCancelled);
+        }(),
+        [&]() -> folly::coro::Task<void> {
+          co_await folly::coro::co_reschedule_on_current_executor;
+          co_await folly::coro::co_reschedule_on_current_executor;
+          cancelSource.requestCancellation();
+        }());
+  }());
+}
+
+TEST(
+    SmallUnboundedQueueTest,
+    CancelledDequeueCompletesNormallyIfAnItemIsAvailable) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    folly::coro::SmallUnboundedQueue<int> queue;
+    folly::CancellationSource cancelSource;
+    cancelSource.requestCancellation();
+
+    queue.enqueue(123);
+
+    int result = co_await folly::coro::co_withCancellation(
+        cancelSource.getToken(), queue.dequeue());
+    EXPECT_EQ(123, result);
+  }());
+}
+#endif