Add coroutine support to fibers::Semaphore

Reviewed By: yfeldblum

Differential Revision: D13515209

fbshipit-source-id: 6d4688242a586b6e5558c62c1c6f3bb7c6595dfb

folly/experimental/coro/test/CoroTest.cpp

@@ -23,6 +23,7 @@
 #include <folly/experimental/coro/BlockingWait.h>
 #include <folly/experimental/coro/Task.h>
 #include <folly/experimental/coro/Utils.h>
+#include <folly/fibers/Semaphore.h>
 #include <folly/io/async/ScopedEventBaseThread.h>
 #include <folly/portability/GTest.h>
 
@@ -379,4 +380,61 @@ TEST(Coro, lambda) {
   executor.run();
   EXPECT_TRUE(coroFuture.isReady());
 }
+
+TEST(Coro, Semaphore) {
+  static constexpr size_t kTasks = 10;
+  static constexpr size_t kIterations = 10000;
+  static constexpr size_t kNumTokens = 10;
+  static constexpr size_t kNumThreads = 16;
+
+  fibers::Semaphore sem(kNumTokens);
+
+  struct Worker {
+    explicit Worker(fibers::Semaphore& s) : sem(s), t([&] { run(); }) {}
+
+    void run() {
+      folly::EventBase evb;
+
+      {
+        std::shared_ptr<folly::EventBase> completionCounter(
+            &evb, [](folly::EventBase* evb_) { evb_->terminateLoopSoon(); });
+
+        for (size_t i = 0; i < kTasks; ++i) {
+          coro::lambda([&, completionCounter]() -> coro::Task<void> {
+            for (size_t j = 0; j < kIterations; ++j) {
+              co_await sem.co_wait();
+              ++counter;
+              sem.signal();
+              --counter;
+
+              EXPECT_LT(counter, kNumTokens);
+              EXPECT_GE(counter, 0);
+            }
+          })
+              .scheduleOn(&evb)
+              .start();
+        }
+      }
+      evb.loopForever();
+    }
+
+    fibers::Semaphore& sem;
+    int counter{0};
+    std::thread t;
+  };
+
+  std::vector<Worker> workers;
+  workers.reserve(kNumThreads);
+  for (size_t i = 0; i < kNumThreads; ++i) {
+    workers.emplace_back(sem);
+  }
+
+  for (auto& worker : workers) {
+    worker.t.join();
+  }
+
+  for (auto& worker : workers) {
+    EXPECT_EQ(0, worker.counter);
+  }
+}
 #endif

folly/fibers/Semaphore.cpp

@@ -57,10 +57,8 @@ void Semaphore::signal() {
       std::memory_order_acquire));
 }
 
-bool Semaphore::waitSlow() {
+bool Semaphore::waitSlow(folly::fibers::Baton& waitBaton) {
   // Slow path, create a baton and acquire a mutex to update the wait list
-  folly::fibers::Baton waitBaton;
-
   {
     auto waitListLock = waitList_.wlock();
     auto& waitList = *waitListLock;
@@ -72,9 +70,7 @@ bool Semaphore::waitSlow() {
     // prepare baton and add to queue
     waitList.push(&waitBaton);
   }
-  // If we managed to create a baton, wait on it
-  // This has to be done here so the mutex has been released
-  waitBaton.wait();
+  // Signal to caller that we managed to push a baton
   return true;
 }
 
@@ -82,9 +78,11 @@ void Semaphore::wait() {
   auto oldVal = tokens_.load(std::memory_order_acquire);
   do {
     while (oldVal == 0) {
+      folly::fibers::Baton waitBaton;
       // If waitSlow fails it is because the token is non-zero by the time
       // the lock is taken, so we can just continue round the loop
-      if (waitSlow()) {
+      if (waitSlow(waitBaton)) {
+        waitBaton.wait();
         return;
       }
       oldVal = tokens_.load(std::memory_order_acquire);
@@ -96,6 +94,30 @@ void Semaphore::wait() {
       std::memory_order_acquire));
 }
 
+#if FOLLY_HAS_COROUTINES
+
+coro::Task<void> Semaphore::co_wait() {
+  auto oldVal = tokens_.load(std::memory_order_acquire);
+  do {
+    while (oldVal == 0) {
+      folly::fibers::Baton waitBaton;
+      // If waitSlow fails it is because the token is non-zero by the time
+      // the lock is taken, so we can just continue round the loop
+      if (waitSlow(waitBaton)) {
+        co_await waitBaton;
+        co_return;
+      }
+      oldVal = tokens_.load(std::memory_order_acquire);
+    }
+  } while (!tokens_.compare_exchange_weak(
+      oldVal,
+      oldVal - 1,
+      std::memory_order_release,
+      std::memory_order_acquire));
+}
+
+#endif
+
 size_t Semaphore::getCapacity() const {
   return capacity_;
 }

folly/fibers/Semaphore.h

@@ -17,6 +17,9 @@
 
 #include <folly/Synchronized.h>
 #include <folly/fibers/Baton.h>
+#if FOLLY_HAS_COROUTINES
+#include <folly/experimental/coro/Task.h>
+#endif
 
 namespace folly {
 namespace fibers {
@@ -45,10 +48,19 @@ class Semaphore {
    */
   void wait();
 
+#if FOLLY_HAS_COROUTINES
+
+  /*
+   * Wait for capacity in the semaphore.
+   */
+  coro::Task<void> co_wait();
+
+#endif
+
   size_t getCapacity() const;
 
  private:
-  bool waitSlow();
+  bool waitSlow(folly::fibers::Baton& waitBaton);
   bool signalSlow();
 
   size_t capacity_;