From: Yedidya Feldblum <yfeldblum@fb.com>
Date: Fri, 27 Oct 2017 04:24:48 +0000 (-0700)
Subject: Fix FunctionScheduler::resetFunctionTimer concurrency bug
X-Git-Tag: v2017.10.30.00~9
X-Git-Url: http://plrg.eecs.uci.edu/git/?p=folly.git;a=commitdiff_plain;h=d89c342ab6b905d2a6bfba7317eac41d44077d65

Fix FunctionScheduler::resetFunctionTimer concurrency bug

Summary:
[Folly] Fix `FunctionScheduler::resetFunctionTimer` concurrency bug.

The original code from the blamed diff code has a comment with this correct assessment of the bug:
> TODO: This moves out of RepeatFunc object while folly:Function can potentially be executed. This might be unsafe.

Namely, when the method is invoked with the id of a scheduled function which is concurrently being executed, the function object (including, say, lambda captures) will be moved while possibly being accessed. If the function object is small enough to be placed in-situ within `folly::Function` (48 bytes on x64), then that access to a moved-from object can happen. It might or might not, depending on the particular instance of the race in question. Or, worse, the access might be to a half-moved-from object!

The new test case for `resetFunctionTimer` passes after the fix, but is guaranteed to fail before the fix because we manage to control the concurrency enough to force the bad version of the race to happen. In the test, we just capture a `std::shared_ptr` (we could have capatured, e.g., a `std::unique_ptr` or a long-enough `std::string`) and check that it is not empty - if it is moved from, it will be empty, and the test expectation will fail.

Reviewed By: simpkins

Differential Revision: D6158722

fbshipit-source-id: 33a7ae699bb3b22089fddbebb6d922737668309d
---

diff --git a/folly/experimental/FunctionScheduler.cpp b/folly/experimental/FunctionScheduler.cpp
index ce58c43c..bb6b6207 100644
--- a/folly/experimental/FunctionScheduler.cpp
+++ b/folly/experimental/FunctionScheduler.cpp
@@ -292,14 +292,10 @@ void FunctionScheduler::cancelAllFunctionsAndWait() {
 bool FunctionScheduler::resetFunctionTimer(StringPiece nameID) {
   std::unique_lock<std::mutex> l(mutex_);
   if (currentFunction_ && currentFunction_->name == nameID) {
-    // TODO: This moves out of RepeatFunc object while folly:Function can
-    // potentially be executed. This might be unsafe.
-    auto funcPtrCopy = std::make_unique<RepeatFunc>(std::move(*currentFunction_));
-    // This function is currently being run. Clear currentFunction_
-    // to avoid rescheduling it, and add the function again to honor the
-    // startDelay.
-    currentFunction_ = nullptr;
-    addFunctionToHeap(l, std::move(funcPtrCopy));
+    if (cancellingCurrentFunction_ || currentFunction_->runOnce) {
+      return false;
+    }
+    currentFunction_->resetNextRunTime(steady_clock::now());
     return true;
   }
 
diff --git a/folly/experimental/test/FunctionSchedulerTest.cpp b/folly/experimental/test/FunctionSchedulerTest.cpp
index 73088c27..ec943958 100644
--- a/folly/experimental/test/FunctionSchedulerTest.cpp
+++ b/folly/experimental/test/FunctionSchedulerTest.cpp
@@ -13,11 +13,14 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #include <algorithm>
 #include <atomic>
 #include <cassert>
 #include <random>
 
+#include <boost/thread.hpp>
+
 #include <folly/Baton.h>
 #include <folly/Random.h>
 #include <folly/experimental/FunctionScheduler.h>
@@ -249,6 +252,52 @@ TEST(FunctionScheduler, ResetFunc) {
   EXPECT_EQ(12, total);
 }
 
+TEST(FunctionScheduler, ResetFuncWhileRunning) {
+  struct State {
+    boost::barrier barrier_a{2};
+    boost::barrier barrier_b{2};
+    boost::barrier barrier_c{2};
+    boost::barrier barrier_d{2};
+    bool set = false;
+    size_t count = 0;
+  };
+
+  State state; // held by ref
+  auto mv = std::make_shared<size_t>(); // gets moved
+
+  FunctionScheduler fs;
+  fs.addFunction(
+      [&, mv /* ref + shared_ptr fit in in-situ storage */] {
+        if (!state.set) { // first invocation
+          state.barrier_a.wait();
+          // ensure that resetFunctionTimer is called in this critical section
+          state.barrier_b.wait();
+          ++state.count;
+          EXPECT_TRUE(bool(mv)) << "bug repro: mv was moved-out";
+          state.barrier_c.wait();
+          // main thread checks count here
+          state.barrier_d.wait();
+        } else { // subsequent invocations
+          ++state.count;
+        }
+      },
+      testInterval(3),
+      "nada");
+  fs.start();
+
+  state.barrier_a.wait();
+  state.set = true;
+  fs.resetFunctionTimer("nada");
+  EXPECT_EQ(0, state.count) << "sanity check";
+  state.barrier_b.wait();
+  // fn thread increments count and checks mv here
+  state.barrier_c.wait();
+  EXPECT_EQ(1, state.count) << "sanity check";
+  state.barrier_d.wait();
+  delay(1);
+  EXPECT_EQ(2, state.count) << "sanity check";
+}
+
 TEST(FunctionScheduler, AddInvalid) {
   int total = 0;
   FunctionScheduler fs;