fix use-after-free in addFunctionOnce

Summary: see title

Reviewed By: meyering

Differential Revision: D4441266

fbshipit-source-id: 4a8f589e2995a463e54b3979035a623824acf39e

diff --git a/folly/experimental/FunctionScheduler.cpp b/folly/experimental/FunctionScheduler.cpp
index 6dfb2c2c76f0b9a53c536f749020b8377652336e..83fba55b1627cfdfc48da881cb929f90b0ab1a68 100644 (file)
--- a/folly/experimental/FunctionScheduler.cpp
+++ b/folly/experimental/FunctionScheduler.cpp
@@ -445,6 +445,7 @@ void FunctionScheduler::runOneFunction(std::unique_lock<std::mutex>& lock,
   }
   if (currentFunction_->runOnce) {
     // Don't reschedule if the function only needed to run once.
+    currentFunction_ = nullptr;
     return;
   }
   // Clear currentFunction_

diff --git a/folly/experimental/test/FunctionSchedulerTest.cpp b/folly/experimental/test/FunctionSchedulerTest.cpp
index d881bfd14614212e3089e1aef664100bfc579436..23293d2c5b2dd7b11bee14a9aebfc8db8d64a548 100644 (file)
--- a/folly/experimental/test/FunctionSchedulerTest.cpp
+++ b/folly/experimental/test/FunctionSchedulerTest.cpp
@@ -456,6 +456,13 @@ TEST(FunctionScheduler, AddWithRunOnce) {
   EXPECT_EQ(2, total);
   delay(2);
   EXPECT_EQ(2, total);
+
+  fs.addFunctionOnce([&] { total += 2; }, "add2");
+  delay(1);
+  EXPECT_EQ(4, total);
+  delay(2);
+  EXPECT_EQ(4, total);
+
   fs.shutdown();
 }