From: Igor Zinkovsky <igorzi@fb.com>
Date: Fri, 20 Jan 2017 23:45:47 +0000 (-0800)
Subject: fix use-after-free in addFunctionOnce
X-Git-Tag: v2017.03.06.00~88
X-Git-Url: http://demsky.eecs.uci.edu/git/?a=commitdiff_plain;h=61d6143a96e87b2bea87514e0d545799ca3b77de;p=folly.git

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 6dfb2c2c..83fba55b 100644
--- 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 d881bfd1..23293d2c 100644
--- 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();
 }