Add future_wait()

Reviewed By: A5he

Differential Revision: D14396454

fbshipit-source-id: a00b1eab527237c885e04413db1a0a413121d91b

folly/fibers/Semaphore.cpp

@@ -116,6 +116,26 @@ coro::Task<void> Semaphore::co_wait() {
       std::memory_order_acquire));
 }
 
+SemiFuture<Unit> Semaphore::future_wait() {
+  auto oldVal = tokens_.load(std::memory_order_acquire);
+  do {
+    while (oldVal == 0) {
+      auto waitBaton = std::make_unique<fibers::Baton>();
+      // 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)) {
+        return futures::wait(std::move(waitBaton));
+      }
+      oldVal = tokens_.load(std::memory_order_acquire);
+    }
+  } while (!tokens_.compare_exchange_weak(
+      oldVal,
+      oldVal - 1,
+      std::memory_order_release,
+      std::memory_order_acquire));
+  return makeSemiFuture();
+}
+
 #endif
 
 size_t Semaphore::getCapacity() const {

folly/fibers/Semaphore.h

@@ -17,6 +17,7 @@
 
 #include <folly/Synchronized.h>
 #include <folly/fibers/Baton.h>
+#include <folly/futures/Future.h>
 #if FOLLY_HAS_COROUTINES
 #include <folly/experimental/coro/Task.h>
 #endif
@@ -57,6 +58,11 @@ class Semaphore {
 
 #endif
 
+  /*
+   * Wait for capacity in the semaphore.
+   */
+  SemiFuture<Unit> future_wait();
+
   size_t getCapacity() const;
 
  private:

folly/fibers/test/FibersTest.cpp

@@ -1594,7 +1594,11 @@ TEST(FiberManager, semaphore) {
         for (size_t i = 0; i < kTasks; ++i) {
           manager.addTask([&, completionCounter]() {
             for (size_t j = 0; j < kIterations; ++j) {
-              sem.wait();
+              if (j % 2) {
+                sem.wait();
+              } else {
+                sem.future_wait().get();
+              }
               ++counter;
               sem.signal();
               --counter;