From 47d37c8e3d1b861615596283f4483fac33cbcf61 Mon Sep 17 00:00:00 2001
From: Matthew Tolton <mtolton@fb.com>
Date: Fri, 22 Sep 2017 17:32:21 -0700
Subject: [PATCH] Fix deadlock in TimedMutex

Summary:
If a lock is stolen from fiber f1, and then fiber f2 is notified before f1 one
wakes up and discovers the crime, then f1 will clear notifiedFiber_ so that f2
thinks the lock was stolen from it as well and hence recurses back into lock().

Reviewed By: andriigrynenko

Differential Revision: D5896323

fbshipit-source-id: 528ec1ed983175d3e08f3dc07b69bbc783a86cfb
---
 folly/fibers/TimedMutex-inl.h    |  4 +++-
 folly/fibers/test/FibersTest.cpp | 37 ++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/folly/fibers/TimedMutex-inl.h b/folly/fibers/TimedMutex-inl.h
index a42c3bc0..07adf7b4 100644
--- a/folly/fibers/TimedMutex-inl.h
+++ b/folly/fibers/TimedMutex-inl.h
@@ -64,7 +64,9 @@ TimedMutex::LockResult TimedMutex::lockHelper(WaitFunc&& waitFunc) {
       std::lock_guard<folly::SpinLock> lg(lock_);
 
       auto stolen = notifiedFiber_ != &waiter;
-      notifiedFiber_ = nullptr;
+      if (!stolen) {
+        notifiedFiber_ = nullptr;
+      }
       return stolen;
     }();
 
diff --git a/folly/fibers/test/FibersTest.cpp b/folly/fibers/test/FibersTest.cpp
index 73e9b6d6..847149f8 100644
--- a/folly/fibers/test/FibersTest.cpp
+++ b/folly/fibers/test/FibersTest.cpp
@@ -2073,6 +2073,43 @@ TEST(FiberManager, VirtualEventBase) {
   EXPECT_TRUE(done2);
 }
 
+TEST(TimedMutex, ThreadsAndFibersDontDeadlock) {
+  folly::EventBase evb;
+  auto& fm = getFiberManager(evb);
+  TimedMutex mutex;
+  std::thread testThread([&] {
+    for (int i = 0; i < 100; i++) {
+      mutex.lock();
+      mutex.unlock();
+      {
+        Baton b;
+        b.timed_wait(std::chrono::milliseconds(1));
+      }
+    }
+  });
+
+  for (int numFibers = 0; numFibers < 100; numFibers++) {
+    fm.addTask([&] {
+      for (int i = 0; i < 20; i++) {
+        mutex.lock();
+        {
+          Baton b;
+          b.timed_wait(std::chrono::milliseconds(1));
+        }
+        mutex.unlock();
+        {
+          Baton b;
+          b.timed_wait(std::chrono::milliseconds(1));
+        }
+      }
+    });
+  }
+
+  evb.loop();
+  EXPECT_EQ(0, fm.hasTasks());
+  testThread.join();
+}
+
 TEST(TimedMutex, ThreadFiberDeadlockOrder) {
   folly::EventBase evb;
   auto& fm = getFiberManager(evb);
-- 
2.34.1