Fix TimedMutex deadlock when used both from fiber and main context

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

diff --git a/folly/experimental/FunctionScheduler.h b/folly/experimental/FunctionScheduler.h
index 0f0a848b219867f1af9898b4bebe2e91dca6f763..f3a2d364a6596e92966506c3ee64b5ac0295defb 100644 (file)
--- a/folly/experimental/FunctionScheduler.h
+++ b/folly/experimental/FunctionScheduler.h
@@ -14,9 +14,9 @@
  * limitations under the License.
  */
 
-#ifndef FOLLY_EXPERIMENTAL_FUNCTION_SCHEDULER_H_
-#define FOLLY_EXPERIMENTAL_FUNCTION_SCHEDULER_H_
+#pragma once
 
+#include <folly/Function.h>
 #include <folly/Range.h>
 #include <chrono>
 #include <condition_variable>
@@ -93,7 +93,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 =
@@ -104,18 +104,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,
@@ -125,7 +133,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
@@ -137,8 +145,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);
@@ -155,6 +163,17 @@ class FunctionScheduler {
    */
   void cancelAllFunctions();
 
+  /**
+   * Resets the specified function's timer.
+   * When resetFunctionTimer is called, the specified function's timer will
+   * be reset with the same parameters it was passed initially, including
+   * its startDelay. If the startDelay was 0, the function will be invoked
+   * immediately.
+   *
+   * Returns false if no function exists with the specified name.
+   */
+  bool resetFunctionTimer(StringPiece nameID);
+
   /**
    * Starts the scheduler.
    *
@@ -166,8 +185,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.
@@ -176,24 +196,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;
@@ -207,7 +231,7 @@ class FunctionScheduler {
     }
     void cancel() {
       // Simply reset cb to an empty function.
-      cb = std::function<void()>();
+      cb = {};
     }
     bool isValid() const { return bool(cb); }
   };
@@ -225,6 +249,16 @@ class FunctionScheduler {
                       std::chrono::steady_clock::time_point now);
   void cancelFunction(const std::unique_lock<std::mutex>& lock,
                       FunctionHeap::iterator it);
+  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);
 
   std::thread thread_;
 
@@ -250,5 +284,3 @@ class FunctionScheduler {
 };
 
 }
-
-#endif