From 61d6143a96e87b2bea87514e0d545799ca3b77de Mon Sep 17 00:00:00 2001
From: Igor Zinkovsky <igorzi@fb.com>
Date: Fri, 20 Jan 2017 15:45:47 -0800
Subject: [PATCH] fix use-after-free in addFunctionOnce

Summary: see title

Reviewed By: meyering

Differential Revision: D4441266

fbshipit-source-id: 4a8f589e2995a463e54b3979035a623824acf39e
---
 folly/experimental/FunctionScheduler.cpp          | 1 +
 folly/experimental/test/FunctionSchedulerTest.cpp | 7 +++++++
 2 files changed, 8 insertions(+)

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();
 }
 
-- 
2.34.1