fix folly::FunctionScheduler.cancelFunctionAndWait() hanging issue

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

diff --git a/folly/experimental/test/FunctionSchedulerTest.cpp b/folly/experimental/test/FunctionSchedulerTest.cpp
index 23293d2c5b2dd7b11bee14a9aebfc8db8d64a548..5da41d24b5613cd09f3498239de36e49c8f095e4 100644 (file)
--- a/folly/experimental/test/FunctionSchedulerTest.cpp
+++ b/folly/experimental/test/FunctionSchedulerTest.cpp
@@ -13,16 +13,20 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include <folly/experimental/FunctionScheduler.h>
-
 #include <algorithm>
 #include <atomic>
 #include <cassert>
 #include <random>
-#include <folly/Random.h>
 
+#include <folly/Baton.h>
+#include <folly/Random.h>
+#include <folly/experimental/FunctionScheduler.h>
 #include <folly/portability/GTest.h>
 
+#if defined(__linux__)
+#include <dlfcn.h>
+#endif
+
 using namespace folly;
 using std::chrono::milliseconds;
 
@@ -488,6 +492,60 @@ TEST(FunctionScheduler, cancelFunctionAndWait) {
   fs.shutdown();
 }
 
+#if defined(__linux__)
+namespace {
+/**
+ * A helper class that forces our pthread_create() wrapper to fail when
+ * an PThreadCreateFailure object exists.
+ */
+class PThreadCreateFailure {
+ public:
+  PThreadCreateFailure() {
+    ++forceFailure_;
+  }
+  ~PThreadCreateFailure() {
+    --forceFailure_;
+  }
+
+  static bool shouldFail() {
+    return forceFailure_ > 0;
+  }
+
+ private:
+  static std::atomic<int> forceFailure_;
+};
+
+std::atomic<int> PThreadCreateFailure::forceFailure_{0};
+} // unnamed namespce
+
+// Replace the system pthread_create() function with our own stub, so we can
+// trigger failures in the StartThrows() test.
+extern "C" int pthread_create(
+    pthread_t* thread,
+    const pthread_attr_t* attr,
+    void* (*start_routine)(void*),
+    void* arg) {
+  static const auto realFunction = reinterpret_cast<decltype(&pthread_create)>(
+      dlsym(RTLD_NEXT, "pthread_create"));
+  // For sanity, make sure we didn't find ourself,
+  // since that would cause infinite recursion.
+  CHECK_NE(realFunction, pthread_create);
+
+  if (PThreadCreateFailure::shouldFail()) {
+    errno = EINVAL;
+    return -1;
+  }
+  return realFunction(thread, attr, start_routine, arg);
+}
+
+TEST(FunctionScheduler, StartThrows) {
+  FunctionScheduler fs;
+  PThreadCreateFailure fail;
+  EXPECT_ANY_THROW(fs.start());
+  EXPECT_NO_THROW(fs.shutdown());
+}
+#endif
+
 TEST(FunctionScheduler, cancelAllFunctionsAndWait) {
   int total = 0;
   FunctionScheduler fs;
@@ -510,3 +568,77 @@ TEST(FunctionScheduler, cancelAllFunctionsAndWait) {
   EXPECT_FALSE(fs.cancelFunction("add2")); // add2 has been canceled
   fs.shutdown();
 }
+
+TEST(FunctionScheduler, CancelAndWaitOnRunningFunc) {
+  folly::Baton<> baton;
+  std::thread th([&baton]() {
+    FunctionScheduler fs;
+    fs.addFunction([] { delay(10); }, testInterval(2), "func");
+    fs.start();
+    delay(1);
+    EXPECT_TRUE(fs.cancelFunctionAndWait("func"));
+    baton.post();
+  });
+
+  ASSERT_TRUE(baton.timed_wait(testInterval(15)));
+  th.join();
+}
+
+TEST(FunctionScheduler, CancelAllAndWaitWithRunningFunc) {
+  folly::Baton<> baton;
+  std::thread th([&baton]() {
+    FunctionScheduler fs;
+    fs.addFunction([] { delay(10); }, testInterval(2), "func");
+    fs.start();
+    delay(1);
+    fs.cancelAllFunctionsAndWait();
+    baton.post();
+  });
+
+  ASSERT_TRUE(baton.timed_wait(testInterval(15)));
+  th.join();
+}
+
+TEST(FunctionScheduler, CancelAllAndWaitWithOneRunningAndOneWaiting) {
+  folly::Baton<> baton;
+  std::thread th([&baton]() {
+    std::atomic<int> nExecuted(0);
+    FunctionScheduler fs;
+    fs.addFunction(
+        [&nExecuted] {
+          nExecuted++;
+          delay(10);
+        },
+        testInterval(2),
+        "func0");
+    fs.addFunction(
+        [&nExecuted] {
+          nExecuted++;
+          delay(10);
+        },
+        testInterval(2),
+        "func1",
+        testInterval(5));
+    fs.start();
+    delay(1);
+    fs.cancelAllFunctionsAndWait();
+    EXPECT_EQ(nExecuted, 1);
+    baton.post();
+  });
+
+  ASSERT_TRUE(baton.timed_wait(testInterval(15)));
+  th.join();
+}
+
+TEST(FunctionScheduler, ConcurrentCancelFunctionAndWait) {
+  FunctionScheduler fs;
+  fs.addFunction([] { delay(10); }, testInterval(2), "func");
+
+  fs.start();
+  delay(1);
+  std::thread th1([&fs] { EXPECT_TRUE(fs.cancelFunctionAndWait("func")); });
+  delay(1);
+  std::thread th2([&fs] { EXPECT_FALSE(fs.cancelFunctionAndWait("func")); });
+  th1.join();
+  th2.join();
+}