folly/synchronization/LifoSem: fixed race condition between tryRemoveNode and...

folly/synchronization/LifoSem: fixed race condition between tryRemoveNode and shutdown by checking the lock in shutdown (#989)

Summary:
Pull Request resolved: https://github.com/facebook/folly/pull/989

Original post: https://fb.workplace.com/groups/560979627394613/permalink/1370192126473355/

There was a bug where we weren't checking the isLocked bit when setting the isShutdown bit. Thus, tryRemoveNode would acquire the lock, shutdown would set the shutdown bit, and then tryRemoveNode would release the lock via a direct store which accidentally cleared the shutdown bit.

The new code wait loops in shutdown until the lock is cleared.

Reviewed By: yfeldblum, djwatson

Differential Revision: D13586264

fbshipit-source-id: 52139df8d7880a60039b6dab810898e0546479dc

folly/synchronization/LifoSem.h

@@ -400,6 +400,12 @@ struct LifoSemBase {
     // first set the shutdown bit
     auto h = head_->load(std::memory_order_acquire);
     while (!h.isShutdown()) {
+      if (h.isLocked()) {
+        std::this_thread::yield();
+        h = head_->load(std::memory_order_acquire);
+        continue;
+      }
+
       if (head_->compare_exchange_strong(h, h.withShutdown())) {
         // success
         h = h.withShutdown();

folly/synchronization/test/LifoSemTests.cpp

@@ -310,6 +310,43 @@ TEST_F(LifoSemTest, timeout) {
   }
 }
 
+TEST_F(LifoSemTest, shutdown_try_wait_for) {
+  long seed = folly::randomNumberSeed() % 1000000;
+  LOG(INFO) << "seed=" << seed;
+  DSched sched(DSched::uniform(seed));
+
+  DLifoSem stopped;
+  std::thread worker1 = DSched::thread([&stopped] {
+    while (!stopped.isShutdown()) {
+      // i.e. poll for messages with timeout
+      LOG(INFO) << "thread polled";
+    }
+  });
+  std::thread worker2 = DSched::thread([&stopped] {
+    while (!stopped.isShutdown()) {
+      // Do some work every 1 second
+
+      try {
+        // this is normally 1 second in prod use case.
+        stopped.try_wait_for(std::chrono::milliseconds(1));
+      } catch (folly::ShutdownSemError& e) {
+        LOG(INFO) << "try_wait_for shutdown";
+      }
+    }
+  });
+
+  std::thread shutdown = DSched::thread([&stopped] {
+    LOG(INFO) << "LifoSem shutdown";
+    stopped.shutdown();
+    LOG(INFO) << "LifoSem shutdown done";
+  });
+
+  DSched::join(shutdown);
+  DSched::join(worker1);
+  DSched::join(worker2);
+  LOG(INFO) << "Threads joined";
+}
+
 BENCHMARK(lifo_sem_pingpong, iters) {
   LifoSem a;
   LifoSem b;