From 2355eee3c3dc8c7c7be50084bd6c3f1ae4d2917e Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Wed, 15 Apr 2015 15:26:35 -0700
Subject: [PATCH] Make sure we can't access LocalData when destroying it

Summary: We can access LocalData while currentFiber is set. We should make sure it's set to null when LocalData::reset is called.

Test Plan: unit test

Reviewed By: alikhtarov@fb.com

Subscribers: folly-diffs@, yfeldblum, chalfant, bwatling

FB internal diff: D1996071

Tasks: 6725667

Signature: t1:1996071:1429135408:d549d577e140ce2867aff4130e73be3884dbd2ed
---
 folly/experimental/fibers/FiberManager-inl.h  | 10 ++++++-
 folly/experimental/fibers/test/FibersTest.cpp | 26 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/folly/experimental/fibers/FiberManager-inl.h b/folly/experimental/fibers/FiberManager-inl.h
index 354622c1..e77ccb45 100644
--- a/folly/experimental/fibers/FiberManager-inl.h
+++ b/folly/experimental/fibers/FiberManager-inl.h
@@ -39,6 +39,11 @@ inline void FiberManager::ensureLoopScheduled() {
 }
 
 inline void FiberManager::runReadyFiber(Fiber* fiber) {
+  SCOPE_EXIT {
+    assert(currentFiber_ == nullptr);
+    assert(activeFiber_ == nullptr);
+  };
+
   assert(fiber->state_ == Fiber::NOT_STARTED ||
          fiber->state_ == Fiber::READY_TO_RUN);
   currentFiber_ = fiber;
@@ -61,6 +66,7 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
   if (fiber->state_ == Fiber::AWAITING) {
     awaitFunc_(*fiber);
     awaitFunc_ = nullptr;
+    currentFiber_ = nullptr;
   } else if (fiber->state_ == Fiber::INVALID) {
     assert(fibersActive_ > 0);
     --fibersActive_;
@@ -77,6 +83,8 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
       }
       fiber->finallyFunc_ = nullptr;
     }
+    // Make sure LocalData is not accessible from its destructor
+    currentFiber_ = nullptr;
     fiber->localData_.reset();
 
     if (fibersPoolSize_ < options_.maxFibersPoolSize) {
@@ -88,10 +96,10 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
       --fibersAllocated_;
     }
   } else if (fiber->state_ == Fiber::YIELDED) {
+    currentFiber_ = nullptr;
     fiber->state_ = Fiber::READY_TO_RUN;
     yieldedFibers_.push_back(*fiber);
   }
-  currentFiber_ = nullptr;
 }
 
 inline bool FiberManager::loopUntilNoReady() {
diff --git a/folly/experimental/fibers/test/FibersTest.cpp b/folly/experimental/fibers/test/FibersTest.cpp
index bc7a7380..41d789ec 100644
--- a/folly/experimental/fibers/test/FibersTest.cpp
+++ b/folly/experimental/fibers/test/FibersTest.cpp
@@ -1281,6 +1281,32 @@ TEST(FiberManager, fiberLocalHeap) {
   testFiberLocal<LargeData>();
 }
 
+TEST(FiberManager, fiberLocalDestructor) {
+  struct CrazyData {
+    size_t data{42};
+
+    ~CrazyData() {
+      if (data == 41) {
+        addTask([]() {
+            EXPECT_EQ(42, local<CrazyData>().data);
+            // Make sure we don't have infinite loop
+            local<CrazyData>().data = 0;
+          });
+      }
+    }
+  };
+
+  FiberManager fm(LocalType<CrazyData>(),
+                  folly::make_unique<SimpleLoopController>());
+
+  fm.addTask([]() {
+      local<CrazyData>().data = 41;
+    });
+
+  fm.loopUntilNoReady();
+  EXPECT_FALSE(fm.hasTasks());
+}
+
 TEST(FiberManager, yieldTest) {
   FiberManager manager(folly::make_unique<SimpleLoopController>());
   auto& loopController =
-- 
2.34.1