Make FiberManager.batonTimedWaitTimeout less flaky

Reviewed By: yfeldblum

Differential Revision: D6957008

fbshipit-source-id: 848932827efaf8529513026b45446da2c1697253

folly/fibers/SimpleLoopController.h

@@ -32,6 +32,10 @@ class SimpleLoopController : public LoopController {
     scheduled_ = false;
   }
 
+  void setTimeFunc(Function<TimePoint()> timeFunc) {
+    timeFunc_ = std::move(timeFunc);
+  }
+
   /**
    * Run FiberManager loop; if no ready task are present,
    * run provided function. Stops after both stop() has been called
@@ -45,7 +49,7 @@ class SimpleLoopController : public LoopController {
     while (LIKELY(waiting || !stopRequested_)) {
       func();
 
-      auto time = Clock::now();
+      auto time = timeFunc_();
 
       for (size_t i = 0; i < scheduledFuncs_.size(); ++i) {
         if (scheduledFuncs_[i].first <= time) {
@@ -104,6 +108,7 @@ class SimpleLoopController : public LoopController {
   std::atomic<int> remoteScheduleCalled_{0};
   int remoteLoopRun_{0};
   std::vector<std::pair<TimePoint, std::function<void()>>> scheduledFuncs_;
+  Function<TimePoint()> timeFunc_{[] { return Clock::now(); }};
 
   /* LoopController interface */
 

folly/fibers/test/FibersTest.cpp

@@ -48,6 +48,9 @@ TEST(FiberManager, batonTimedWaitTimeout) {
   auto& loopController =
       dynamic_cast<SimpleLoopController&>(manager.loopController());
 
+  auto now = SimpleLoopController::Clock::now();
+  loopController.setTimeFunc([&] { return now; });
+
   auto loopFunc = [&]() {
     if (!taskAdded) {
       manager.addTask([&]() {
@@ -72,7 +75,7 @@ TEST(FiberManager, batonTimedWaitTimeout) {
       });
       taskAdded = true;
     } else {
-      std::this_thread::sleep_for(std::chrono::milliseconds(50));
+      now += std::chrono::milliseconds(50);
       iterations++;
     }
   };