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 ce58c43c69648ecbaea4bdda8247df87681b3fb5..bb6b62075d13e42508c80511d383ac0f0110505c 100644 (file)
--- 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 73088c27e20705a11e1ba6bf029b97a999135ce9..ec94395851bd52e38f99b69c4a920eae960f4bf6 100644 (file)
--- 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;