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

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.

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.

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();
+}