From: Igor Zinkovsky <igorzi@fb.com>
Date: Tue, 13 Dec 2016 07:21:55 +0000 (-0800)
Subject: allow run-once callbacks
X-Git-Tag: v2016.12.19.00~28
X-Git-Url: http://plrg.eecs.uci.edu/git/?a=commitdiff_plain;h=0dd7a9a6756769a9e2142a40cf351a35c20a413f;p=folly.git

allow run-once callbacks

Summary: Adding new `addFunctionOnce` method that executes provided callback only once.

Reviewed By: yfeldblum

Differential Revision: D4315635

fbshipit-source-id: 4819ec30b4f2e4ab3185a37158404e1c7a96758a
---

diff --git a/folly/experimental/FunctionScheduler.cpp b/folly/experimental/FunctionScheduler.cpp
index 2e0ee6eb..98bd9ed7 100644
--- a/folly/experimental/FunctionScheduler.cpp
+++ b/folly/experimental/FunctionScheduler.cpp
@@ -96,12 +96,13 @@ void FunctionScheduler::addFunction(Function<void()>&& cb,
                                     milliseconds interval,
                                     StringPiece nameID,
                                     milliseconds startDelay) {
-  addFunctionGenericDistribution(
+  addFunctionInternal(
       std::move(cb),
       ConstIntervalFunctor(interval),
       nameID.str(),
       to<std::string>(interval.count(), "ms"),
-      startDelay);
+      startDelay,
+      false /*runOnce*/);
 }
 
 void FunctionScheduler::addFunction(Function<void()>&& cb,
@@ -110,30 +111,45 @@ void FunctionScheduler::addFunction(Function<void()>&& cb,
                                     StringPiece nameID,
                                     milliseconds startDelay) {
   if (latencyDistr.isPoisson) {
-    addFunctionGenericDistribution(
+    addFunctionInternal(
         std::move(cb),
         PoissonDistributionFunctor(latencyDistr.poissonMean),
         nameID.str(),
         to<std::string>(latencyDistr.poissonMean, "ms (Poisson mean)"),
-        startDelay);
+        startDelay,
+        false /*runOnce*/);
   } else {
     addFunction(std::move(cb), interval, nameID, startDelay);
   }
 }
 
+void FunctionScheduler::addFunctionOnce(
+    Function<void()>&& cb,
+    StringPiece nameID,
+    milliseconds startDelay) {
+  addFunctionInternal(
+      std::move(cb),
+      ConstIntervalFunctor(milliseconds::zero()),
+      nameID.str(),
+      "once",
+      startDelay,
+      true /*runOnce*/);
+}
+
 void FunctionScheduler::addFunctionUniformDistribution(
     Function<void()>&& cb,
     milliseconds minInterval,
     milliseconds maxInterval,
     StringPiece nameID,
     milliseconds startDelay) {
-  addFunctionGenericDistribution(
+  addFunctionInternal(
       std::move(cb),
       UniformDistributionFunctor(minInterval, maxInterval),
       nameID.str(),
       to<std::string>(
           "[", minInterval.count(), " , ", maxInterval.count(), "] ms"),
-      startDelay);
+      startDelay,
+      false /*runOnce*/);
 }
 
 void FunctionScheduler::addFunctionGenericDistribution(
@@ -142,6 +158,22 @@ void FunctionScheduler::addFunctionGenericDistribution(
     const std::string& nameID,
     const std::string& intervalDescr,
     milliseconds startDelay) {
+  addFunctionInternal(
+      std::move(cb),
+      std::move(intervalFunc),
+      nameID,
+      intervalDescr,
+      startDelay,
+      false /*runOnce*/);
+}
+
+void FunctionScheduler::addFunctionInternal(
+    Function<void()>&& cb,
+    IntervalDistributionFunc&& intervalFunc,
+    const std::string& nameID,
+    const std::string& intervalDescr,
+    milliseconds startDelay,
+    bool runOnce) {
   if (!cb) {
     throw std::invalid_argument(
         "FunctionScheduler: Scheduled function must be set");
@@ -177,7 +209,8 @@ void FunctionScheduler::addFunctionGenericDistribution(
           std::move(intervalFunc),
           nameID,
           intervalDescr,
-          startDelay));
+          startDelay,
+          runOnce));
 }
 
 bool FunctionScheduler::cancelFunction(StringPiece nameID) {
@@ -382,6 +415,10 @@ void FunctionScheduler::runOneFunction(std::unique_lock<std::mutex>& lock,
     // We shouldn't reschedule it;
     return;
   }
+  if (currentFunction_->runOnce) {
+    // Don't reschedule if the function only needed to run once.
+    return;
+  }
   // Clear currentFunction_
   CHECK_EQ(currentFunction_, &func);
   currentFunction_ = nullptr;
diff --git a/folly/experimental/FunctionScheduler.h b/folly/experimental/FunctionScheduler.h
index 1668c501..f3a2d364 100644
--- a/folly/experimental/FunctionScheduler.h
+++ b/folly/experimental/FunctionScheduler.h
@@ -110,6 +110,14 @@ class FunctionScheduler {
       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
@@ -194,18 +202,22 @@ class FunctionScheduler {
     std::string name;
     std::chrono::milliseconds startDelay;
     std::string intervalDescr;
+    bool runOnce;
 
-    RepeatFunc(Function<void()>&& cback,
-               IntervalDistributionFunc&& intervalFn,
-               const std::string& nameID,
-               const std::string& intervalDistDescription,
-               std::chrono::milliseconds delay)
+    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;
@@ -240,6 +252,14 @@ 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);
+
   std::thread thread_;
 
   // Mutex to protect our member variables.
diff --git a/folly/experimental/test/FunctionSchedulerTest.cpp b/folly/experimental/test/FunctionSchedulerTest.cpp
index eb122d6f..2de87ed3 100644
--- a/folly/experimental/test/FunctionSchedulerTest.cpp
+++ b/folly/experimental/test/FunctionSchedulerTest.cpp
@@ -446,3 +446,15 @@ TEST(FunctionScheduler, GammaIntervalDistribution) {
   delay(2);
   EXPECT_EQ(6, total);
 }
+
+TEST(FunctionScheduler, AddWithRunOnce) {
+  int total = 0;
+  FunctionScheduler fs;
+  fs.addFunctionOnce([&] { total += 2; }, "add2");
+  fs.start();
+  delay(1);
+  EXPECT_EQ(2, total);
+  delay(2);
+  EXPECT_EQ(2, total);
+  fs.shutdown();
+}