From 2a7bf4845898201b7d6845502e6b275573cc9337 Mon Sep 17 00:00:00 2001
From: Junlin Zhang <neicullyn@fb.com>
Date: Fri, 16 Dec 2016 15:07:47 -0800
Subject: [PATCH] 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      | 28 ++++++++++++
 folly/experimental/FunctionScheduler.h        |  2 +
 .../test/FunctionSchedulerTest.cpp            | 45 +++++++++++++++++++
 3 files changed, 75 insertions(+)

diff --git a/folly/experimental/FunctionScheduler.cpp b/folly/experimental/FunctionScheduler.cpp
index d2476c6e..0f8aff0e 100644
--- 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 f3a2d364..7c0742b8 100644
--- 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 2de87ed3..4d1b78d7 100644
--- 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();
+}
-- 
2.34.1