Add a full drain for folly's ManualExecutor.

Summary: ManualExecutor::run() is stable, which means that in cases where we want to fully drain the executor we have to loop over it. This adds ManualExecutor::drain() which does that internally.

Reviewed By: yfeldblum

Differential Revision: D6126840

fbshipit-source-id: e36cba5c373a57fe01de244977ec852636b58dbd

diff --git a/folly/executors/ManualExecutor.cpp b/folly/executors/ManualExecutor.cpp
index 51ff9c8a8718276ee1d20d1c8812269efa32f703..00ed2f93f59ffe0a7d260c76bb27b372a2c9b8f2 100644 (file)
--- a/folly/executors/ManualExecutor.cpp
+++ b/folly/executors/ManualExecutor.cpp
@@ -69,6 +69,15 @@ size_t ManualExecutor::run() {
   return count;
 }
 
+size_t ManualExecutor::drain() {
+  size_t tasksRun = 0;
+  size_t tasksForSingleRun = 0;
+  while ((tasksForSingleRun = run()) != 0) {
+    tasksRun += tasksForSingleRun;
+  }
+  return tasksRun;
+}
+
 void ManualExecutor::wait() {
   while (true) {
     {

diff --git a/folly/executors/ManualExecutor.h b/folly/executors/ManualExecutor.h
index 028a4ce63af03490c1b0d183cdf6ea725c9e6f2e..bfaaf259472349b58b43641a00d7211f961bad0d 100644 (file)
--- a/folly/executors/ManualExecutor.h
+++ b/folly/executors/ManualExecutor.h
@@ -47,6 +47,14 @@ namespace folly {
     /// moment that this returns.
     size_t run();
 
+    // Do work until there is no more work to do.
+    // Returns the number of functions that were executed (maybe 0).
+    // Unlike run, this method is not stable. It will chase an infinite tail of
+    // work so should be used with care.
+    // There will be no work available to perform at the moment that this
+    // returns.
+    size_t drain();
+
     /// Wait for work to do.
     void wait();
 

diff --git a/folly/executors/test/ExecutorTest.cpp b/folly/executors/test/ExecutorTest.cpp
index 1755a2fc9e210fdcb37a50291b33de6b83c6c38a..2bbc56a1fa270081492f2efa545b7b75ba3959c2 100644 (file)
--- a/folly/executors/test/ExecutorTest.cpp
+++ b/folly/executors/test/ExecutorTest.cpp
@@ -33,6 +33,20 @@ TEST(ManualExecutor, runIsStable) {
   auto f2 = [&]() { x.add(f1); x.add(f1); };
   x.add(f2);
   x.run();
+  EXPECT_EQ(count, 0);
+}
+
+TEST(ManualExecutor, drainIsNotStable) {
+  ManualExecutor x;
+  size_t count = 0;
+  auto f1 = [&]() { count++; };
+  auto f2 = [&]() {
+    x.add(f1);
+    x.add(f1);
+  };
+  x.add(f2);
+  x.drain();
+  EXPECT_EQ(count, 2);
 }
 
 TEST(ManualExecutor, scheduleDur) {