Fix FunctionScheduler::resetFunctionTimer concurrency bug

[folly.git] / folly / experimental / test / FunctionSchedulerTest.cpp

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;