fix folly::FunctionScheduler.cancelFunctionAndWait() hanging issue

[folly.git] / folly / experimental / FunctionScheduler.h

diff --git a/folly/experimental/FunctionScheduler.h b/folly/experimental/FunctionScheduler.h
index e3c0d132e3e83a18b8c16fd94f6ae5547d4aa693..6e5b2198e7946a67d86f17fc996921897d272630 100644 (file)
--- a/folly/experimental/FunctionScheduler.h
+++ b/folly/experimental/FunctionScheduler.h
@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 Facebook, Inc.
+ * Copyright 2017 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <folly/Function.h>
 #include <folly/Range.h>
 #include <chrono>
 #include <condition_variable>
@@ -41,7 +42,9 @@ namespace folly {
  *
  *
  * Note: the class uses only one thread - if you want to use more than one
- *       thread use multiple FunctionScheduler objects
+ *       thread, either use multiple FunctionScheduler objects, or check out
+ *       ThreadedRepeatingFunctionRunner.h for a much simpler contract of
+ *       "run each function periodically in its own thread".
  *
  * start() schedules the functions, while shutdown() terminates further
  * scheduling.
@@ -92,7 +95,7 @@ class FunctionScheduler {
    * Throws an exception on error.  In particular, each function must have a
    * unique name--two functions cannot be added with the same name.
    */
-  void addFunction(const std::function<void()>& cb,
+  void addFunction(Function<void()>&& cb,
                    std::chrono::milliseconds interval,
                    StringPiece nameID = StringPiece(),
                    std::chrono::milliseconds startDelay =
@@ -103,18 +106,26 @@ class FunctionScheduler {
    * LatencyDistribution
    */
   void addFunction(
-      const std::function<void()>& cb,
+      Function<void()>&& cb,
       std::chrono::milliseconds interval,
       const LatencyDistribution& latencyDistr,
       StringPiece nameID = StringPiece(),
       std::chrono::milliseconds startDelay = std::chrono::milliseconds(0));
 
+  /**
+   * Adds a new function to the FunctionScheduler to run only once.
+   */
+  void addFunctionOnce(
+      Function<void()>&& cb,
+      StringPiece nameID = StringPiece(),
+      std::chrono::milliseconds startDelay = std::chrono::milliseconds(0));
+
   /**
     * Add a new function to the FunctionScheduler with the time
     * interval being distributed uniformly within the given interval
     * [minInterval, maxInterval].
     */
-  void addFunctionUniformDistribution(const std::function<void()>& cb,
+  void addFunctionUniformDistribution(Function<void()>&& cb,
                                       std::chrono::milliseconds minInterval,
                                       std::chrono::milliseconds maxInterval,
                                       StringPiece nameID,
@@ -124,7 +135,7 @@ class FunctionScheduler {
    * A type alias for function that is called to determine the time
    * interval for the next scheduled run.
    */
-  using IntervalDistributionFunc = std::function<std::chrono::milliseconds()>;
+  using IntervalDistributionFunc = Function<std::chrono::milliseconds()>;
 
   /**
    * Add a new function to the FunctionScheduler. The scheduling interval
@@ -136,8 +147,8 @@ class FunctionScheduler {
    * @see FunctionScheduler::addFunctionJitterInterval).
    */
   void addFunctionGenericDistribution(
-      const std::function<void()>& cb,
-      const IntervalDistributionFunc& intervalFunc,
+      Function<void()>&& cb,
+      IntervalDistributionFunc&& intervalFunc,
       const std::string& nameID,
       const std::string& intervalDescr,
       std::chrono::milliseconds startDelay);
@@ -148,11 +159,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.
@@ -176,8 +189,9 @@ class FunctionScheduler {
    * Stops the FunctionScheduler.
    *
    * It may be restarted later by calling start() again.
+   * Returns false if the scheduler was not running.
    */
-  void shutdown();
+  bool shutdown();
 
   /**
    * Set the name of the worker thread.
@@ -186,24 +200,28 @@ class FunctionScheduler {
 
  private:
   struct RepeatFunc {
-    std::function<void()> cb;
+    Function<void()> cb;
     IntervalDistributionFunc intervalFunc;
     std::chrono::steady_clock::time_point nextRunTime;
     std::string name;
     std::chrono::milliseconds startDelay;
     std::string intervalDescr;
+    bool runOnce;
 
-    RepeatFunc(const std::function<void()>& cback,
-               const IntervalDistributionFunc& intervalFn,
-               const std::string& nameID,
-               const std::string& intervalDistDescription,
-               std::chrono::milliseconds delay)
-        : cb(cback),
-          intervalFunc(intervalFn),
+    RepeatFunc(
+        Function<void()>&& cback,
+        IntervalDistributionFunc&& intervalFn,
+        const std::string& nameID,
+        const std::string& intervalDistDescription,
+        std::chrono::milliseconds delay,
+        bool once)
+        : cb(std::move(cback)),
+          intervalFunc(std::move(intervalFn)),
           nextRunTime(),
           name(nameID),
           startDelay(delay),
-          intervalDescr(intervalDistDescription) {}
+          intervalDescr(intervalDistDescription),
+          runOnce(once) {}
 
     std::chrono::steady_clock::time_point getNextRunTime() const {
       return nextRunTime;
@@ -217,7 +235,7 @@ class FunctionScheduler {
     }
     void cancel() {
       // Simply reset cb to an empty function.
-      cb = std::function<void()>();
+      cb = {};
     }
     bool isValid() const { return bool(cb); }
   };
@@ -238,6 +256,20 @@ class FunctionScheduler {
   void addFunctionToHeap(const std::unique_lock<std::mutex>& lock,
                          RepeatFunc&& func);
 
+  void addFunctionInternal(
+      Function<void()>&& cb,
+      IntervalDistributionFunc&& intervalFunc,
+      const std::string& nameID,
+      const std::string& intervalDescr,
+      std::chrono::milliseconds startDelay,
+      bool runOnce);
+
+  // Return true if the current function is being canceled
+  bool cancelAllFunctionsWithLock(std::unique_lock<std::mutex>& lock);
+  bool cancelFunctionWithLock(
+      std::unique_lock<std::mutex>& lock,
+      StringPiece nameID);
+
   std::thread thread_;
 
   // Mutex to protect our member variables.
@@ -259,6 +291,7 @@ class FunctionScheduler {
 
   std::string threadName_;
   bool steady_{false};
+  bool cancellingCurrentFunction_{false};
 };
 
 }