add cancelFunctionAndWait and cancelAllFunctionsAndWait to FunctionScheduler

Summary: In FunctionScheduler, cancelFunction and cancelAllFunction will not wait for current running function, which will lead to segmentation fault if the function captures resource to be destructed right after calling cancelFunction or cancelAllFunction. This diff introduces cancelFunctionAndWait and cancelAllFunctionsAndWait, making it possible to wait for completion of current running function.

Reviewed By: yfeldblum

Differential Revision: D4324384

fbshipit-source-id: ac4b272894f753ef3bb175173f10d1ca20c17c41

diff --git a/folly/experimental/FunctionScheduler.cpp b/folly/experimental/FunctionScheduler.cpp
index d2476c6ebbc3a659e1a8c7b5736bfe3ac12473b0..0f8aff0e0780e0629ba749cfff94d55bddf20d89 100644 (file)
--- a/folly/experimental/FunctionScheduler.cpp
+++ b/folly/experimental/FunctionScheduler.cpp
@@ -232,6 +232,25 @@ bool FunctionScheduler::cancelFunction(StringPiece nameID) {
   return false;
 }
 
+bool FunctionScheduler::cancelFunctionAndWait(StringPiece nameID) {
+  std::unique_lock<std::mutex> l(mutex_);
+
+  auto* currentFunction = currentFunction_;
+  if (currentFunction && currentFunction->name == nameID) {
+    runningCondvar_.wait(l, [currentFunction, this]() {
+      return currentFunction != currentFunction_;
+    });
+  }
+
+  for (auto it = functions_.begin(); it != functions_.end(); ++it) {
+    if (it->isValid() && it->name == nameID) {
+      cancelFunction(l, it);
+      return true;
+    }
+  }
+  return false;
+}
+
 void FunctionScheduler::cancelFunction(const std::unique_lock<std::mutex>& l,
                                        FunctionHeap::iterator it) {
   // This function should only be called with mutex_ already locked.
@@ -259,6 +278,14 @@ void FunctionScheduler::cancelAllFunctions() {
   currentFunction_ = nullptr;
 }
 
+void FunctionScheduler::cancelAllFunctionsAndWait() {
+  std::unique_lock<std::mutex> l(mutex_);
+  if (currentFunction_) {
+    runningCondvar_.wait(l, [this]() { return currentFunction_ == nullptr; });
+  }
+  functions_.clear();
+}
+
 bool FunctionScheduler::resetFunctionTimer(StringPiece nameID) {
   std::unique_lock<std::mutex> l(mutex_);
   if (currentFunction_ && currentFunction_->name == nameID) {
@@ -355,6 +382,7 @@ void FunctionScheduler::run() {
     if (sleepTime < milliseconds::zero()) {
       // We need to run this function now
       runOneFunction(lock, now);
+      runningCondvar_.notify_all();
     } else {
       // Re-add the function to the heap, and wait until we actually
       // need to run it.

diff --git a/folly/experimental/FunctionScheduler.h b/folly/experimental/FunctionScheduler.h
index f3a2d364a6596e92966506c3ee64b5ac0295defb..7c0742b87378d799ab89ece5c729c64aeb8144a9 100644 (file)
--- a/folly/experimental/FunctionScheduler.h
+++ b/folly/experimental/FunctionScheduler.h
@@ -157,11 +157,13 @@ class FunctionScheduler {
    * Returns false if no function exists with the specified name.
    */
   bool cancelFunction(StringPiece nameID);
+  bool cancelFunctionAndWait(StringPiece nameID);
 
   /**
    * All functions registered will be canceled.
    */
   void cancelAllFunctions();
+  void cancelAllFunctionsAndWait();
 
   /**
    * Resets the specified function's timer.

diff --git a/folly/experimental/test/FunctionSchedulerTest.cpp b/folly/experimental/test/FunctionSchedulerTest.cpp
index 2de87ed36ab1ebd22ada03c68e6368b59c3fda98..4d1b78d702ee401a2b901c7c338a18c5d0d8bdc6 100644 (file)
--- a/folly/experimental/test/FunctionSchedulerTest.cpp
+++ b/folly/experimental/test/FunctionSchedulerTest.cpp
@@ -458,3 +458,48 @@ TEST(FunctionScheduler, AddWithRunOnce) {
   EXPECT_EQ(2, total);
   fs.shutdown();
 }
+
+TEST(FunctionScheduler, cancelFunctionAndWait) {
+  int total = 0;
+  FunctionScheduler fs;
+  fs.addFunction(
+      [&] {
+        delay(5);
+        total += 2;
+      },
+      testInterval(100),
+      "add2");
+
+  fs.start();
+  delay(1);
+  EXPECT_EQ(0, total); // add2 is still sleeping
+
+  EXPECT_TRUE(fs.cancelFunctionAndWait("add2"));
+  EXPECT_EQ(2, total); // add2 should have completed
+
+  EXPECT_FALSE(fs.cancelFunction("add2")); // add2 has been canceled
+  fs.shutdown();
+}
+
+TEST(FunctionScheduler, cancelAllFunctionsAndWait) {
+  int total = 0;
+  FunctionScheduler fs;
+
+  fs.addFunction(
+      [&] {
+        delay(5);
+        total += 2;
+      },
+      testInterval(100),
+      "add2");
+
+  fs.start();
+  delay(1);
+  EXPECT_EQ(0, total); // add2 is still sleeping
+
+  fs.cancelAllFunctionsAndWait();
+  EXPECT_EQ(2, total);
+
+  EXPECT_FALSE(fs.cancelFunction("add2")); // add2 has been canceled
+  fs.shutdown();
+}