Fix a race in fibers::Semaphore

Summary: We should never end up doing compare_exchange_weak with oldVal==0, because that may result in increment/decrement from 0.

Reviewed By: yfeldblum

Differential Revision: D13514210

fbshipit-source-id: 0ccffe5d9525389ba02208f3bf37ce14acb9f28e

folly/fibers/Semaphore.cpp

@@ -44,10 +44,11 @@ bool Semaphore::signalSlow() {
 void Semaphore::signal() {
   auto oldVal = tokens_.load(std::memory_order_acquire);
   do {
-    if (oldVal == 0) {
+    while (oldVal == 0) {
       if (signalSlow()) {
-        break;
+        return;
       }
+      oldVal = tokens_.load(std::memory_order_acquire);
     }
   } while (!tokens_.compare_exchange_weak(
       oldVal,
@@ -80,12 +81,13 @@ bool Semaphore::waitSlow() {
 void Semaphore::wait() {
   auto oldVal = tokens_.load(std::memory_order_acquire);
   do {
-    if (oldVal == 0) {
+    while (oldVal == 0) {
       // 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()) {
-        break;
+        return;
       }
+      oldVal = tokens_.load(std::memory_order_acquire);
     }
   } while (!tokens_.compare_exchange_weak(
       oldVal,

folly/fibers/test/FibersTest.cpp

@@ -1555,54 +1555,58 @@ TEST(FiberManager, 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;
 
   Semaphore sem(kNumTokens);
-  int counterA = 0;
-  int counterB = 0;
-
-  auto task = [&sem](int& counter, folly::fibers::Baton& baton) {
-    FiberManager manager(std::make_unique<EventBaseLoopController>());
-    folly::EventBase evb;
-    dynamic_cast<EventBaseLoopController&>(manager.loopController())
-        .attachEventBase(evb);
-
-    {
-      std::shared_ptr<folly::EventBase> completionCounter(
-          &evb, [](folly::EventBase* evb_) { evb_->terminateLoopSoon(); });
-
-      for (size_t i = 0; i < kTasks; ++i) {
-        manager.addTask([&, completionCounter]() {
-          for (size_t j = 0; j < kIterations; ++j) {
-            sem.wait();
-            ++counter;
-            sem.signal();
-            --counter;
-
-            EXPECT_LT(counter, kNumTokens);
-            EXPECT_GE(counter, 0);
-          }
-        });
-      }
 
-      baton.wait();
+  struct Worker {
+    explicit Worker(Semaphore& s) : sem(s), t([&] { run(); }) {}
+
+    void run() {
+      FiberManager manager(std::make_unique<EventBaseLoopController>());
+      folly::EventBase evb;
+      dynamic_cast<EventBaseLoopController&>(manager.loopController())
+          .attachEventBase(evb);
+
+      {
+        std::shared_ptr<folly::EventBase> completionCounter(
+            &evb, [](folly::EventBase* evb_) { evb_->terminateLoopSoon(); });
+
+        for (size_t i = 0; i < kTasks; ++i) {
+          manager.addTask([&, completionCounter]() {
+            for (size_t j = 0; j < kIterations; ++j) {
+              sem.wait();
+              ++counter;
+              sem.signal();
+              --counter;
+
+              EXPECT_LT(counter, kNumTokens);
+              EXPECT_GE(counter, 0);
+            }
+          });
+        }
+      }
+      evb.loopForever();
     }
-    evb.loopForever();
+
+    Semaphore& sem;
+    int counter{0};
+    std::thread t;
   };
 
-  folly::fibers::Baton batonA;
-  folly::fibers::Baton batonB;
-  std::thread threadA([&] { task(counterA, batonA); });
-  std::thread threadB([&] { task(counterB, batonB); });
+  std::vector<Worker> workers;
+  workers.reserve(kNumThreads);
+  for (size_t i = 0; i < kNumThreads; ++i) {
+    workers.emplace_back(sem);
+  }
 
-  batonA.post();
-  batonB.post();
-  threadA.join();
-  threadB.join();
+  for (auto& worker : workers) {
+    worker.t.join();
+  }
 
-  EXPECT_LT(counterA, kNumTokens);
-  EXPECT_LT(counterB, kNumTokens);
-  EXPECT_GE(counterA, 0);
-  EXPECT_GE(counterB, 0);
+  for (auto& worker : workers) {
+    EXPECT_EQ(0, worker.counter);
+  }
 }
 
 template <typename ExecutorT>