try_wait for semaphore

Summary:
Adding a new "non waitng" version if the `try_wait` that doesn't require a waiter.
This also makes the `BatchSemaphore` aligned with the `Semaphore`.

Reviewed By: Gownta

Differential Revision: D41634112

fbshipit-source-id: c50656b42922b441f6018c275ece3f071c4c910b

folly/fibers/BatchSemaphore.cpp

@@ -31,6 +31,10 @@ bool BatchSemaphore::try_wait(Waiter& waiter, int64_t tokens) {
   return try_wait_common(waiter, tokens);
 }
 
+bool BatchSemaphore::try_wait(int64_t tokens) {
+  return try_wait_common(tokens);
+}
+
 #if FOLLY_HAS_COROUTINES
 
 coro::Task<void> BatchSemaphore::co_wait(int64_t tokens) {

folly/fibers/BatchSemaphore.h

@@ -53,6 +53,12 @@ class BatchSemaphore : public SemaphoreBase {
    */
   bool try_wait(Waiter& waiter, int64_t tokens);
 
+  /**
+   * If the semaphore has capacity, removes a token and returns true. Otherwise
+   * returns false and leaves the semaphore unchanged.
+   */
+  bool try_wait(int64_t tokens);
+
 #if FOLLY_HAS_COROUTINES
 
   /*

folly/fibers/SemaphoreBase.cpp

@@ -121,6 +121,20 @@ bool SemaphoreBase::try_wait_common(Waiter& waiter, int64_t tokens) {
   return true;
 }
 
+bool SemaphoreBase::try_wait_common(int64_t tokens) {
+  auto oldVal = tokens_.load(std::memory_order_acquire);
+  do {
+    if (oldVal < tokens) {
+      return false;
+    }
+  } while (!tokens_.compare_exchange_weak(
+      oldVal,
+      oldVal - tokens,
+      std::memory_order_release,
+      std::memory_order_acquire));
+  return true;
+}
+
 #if FOLLY_HAS_COROUTINES
 
 coro::Task<void> SemaphoreBase::co_wait_common(int64_t tokens) {

folly/fibers/SemaphoreBase.h

@@ -73,6 +73,12 @@ class SemaphoreBase {
    */
   bool try_wait_common(Waiter& waiter, int64_t tokens);
 
+  /**
+   * If the semaphore has capacity, removes a token and returns true. Otherwise
+   * returns false and leaves the semaphore unchanged.
+   */
+  bool try_wait_common(int64_t tokens);
+
 #if FOLLY_HAS_COROUTINES
 
   /*

folly/fibers/test/FibersTest.cpp

@@ -1859,8 +1859,8 @@ TEST(FiberManager, batchSemaphore) {
         for (size_t i = 0; i < kTasks; ++i) {
           manager.addTask([&, completionCounter]() {
             for (size_t j = 0; j < kIterations; ++j) {
-              int tokens = j % 4 + 1;
-              switch (j % 4) {
+              int tokens = j % 5 + 1;
+              switch (j % 5) {
                 case 0:
                   sem.wait(tokens);
                   break;
@@ -1878,6 +1878,12 @@ TEST(FiberManager, batchSemaphore) {
                 case 3:
                   folly::coro::blockingWait(sem.co_wait(tokens));
                   break;
+                case 4: {
+                  if (!sem.try_wait(tokens)) {
+                    sem.wait(tokens);
+                  }
+                  break;
+                }
               }
               counter += tokens;
               sem.signal(tokens);