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,12 +1555,14 @@ 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) {
+  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())
@@ -1583,26 +1585,28 @@ TEST(FiberManager, semaphore) {
             }
           });
         }
-
-      baton.wait();
       }
       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>